feat: semaphores: backwards compatibility by falling back to `run_wor…

…ker`
ZettaAI · Nov 8, 2023 · cf5b2f1 · cf5b2f1
1 parent 5470773
commit cf5b2f1
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/zetta_utils/cloud_management/resource_allocation/k8s/common.py b/zetta_utils/cloud_management/resource_allocation/k8s/common.py
@@ -65,15 +65,23 @@ def get_mazepa_worker_command(
     num_procs: int = 1,
     semaphores_spec: dict[SemaphoreType, int] | None = None,
 ):
+    if num_procs == 1 and semaphores_spec is None:
+        command = "mazepa.run_worker"
+    else:
+        command = "mazepa.run_worker_manager"
+    if semaphores_spec is None:
+        semaphores_line = ""
+    else:
+        semaphores_line = f"semaphores_spec: {json.dumps(semaphores_spec)}\n"
     result = (
         """
     zetta -vv -l try run -s '{
-        "@type": "mazepa.run_worker_manager"
     """
+        + f'"@type": "{command}"\n'
         + f"task_queue: {json.dumps(task_queue_spec)}\n"
         + f"outcome_queue: {json.dumps(outcome_queue_spec)}\n"
-        + f"semaphores_spec: {json.dumps(semaphores_spec)}\n"
         + f"num_procs: {num_procs}\n"
+        + semaphores_line
         + """
         sleep_sec: 5
     }'