move runner ack after status update

src/exo/worker/plan.py

@@ -292,14 +292,9 @@ def _pending_tasks(
             if task.instance_id != runner.bound_instance.instance.instance_id:
                 continue
 
-            # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough
-            # however, realistically the task status should be set to completed by the LAST runner, so this is a true race
-            # the actual solution is somewhat deeper than this bypass - TODO!
             if task.task_id in runner.completed:
                 continue
 
-            # TODO: Check ordering aligns with MLX distributeds expectations.
-
             if isinstance(runner.status, RunnerReady) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard

src/exo/worker/runner/runner.py

@@ -140,12 +140,12 @@ def main(
     with task_receiver as tasks:
         for task in tasks:
             if task.task_id in seen:
-                logger.warning("repeat task - potential error")
+                logger.warning("repeat task - currently a logic bug, please report")
+                continue
             seen.add(task.task_id)
             event_sender.send(
                 TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
             )
-            event_sender.send(TaskAcknowledged(task_id=task.task_id))
             match task:
                 case ConnectToGroup() if isinstance(
                     current_status, (RunnerIdle, RunnerFailed)
@@ -157,6 +157,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
                     group = initialize_mlx(bound_instance)
 
                     logger.info("runner connected")
@@ -173,6 +174,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
                     def on_model_load_timeout() -> None:
                         event_sender.send(
@@ -215,6 +217,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
                     logger.info(f"warming up inference for instance: {instance}")
                     if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
@@ -254,6 +257,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
                     assert model and not isinstance(model, DistributedImageModel)
                     assert tokenizer
 
@@ -385,6 +389,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
                     try:
                         image_index = 0
@@ -447,6 +452,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
                     try:
                         image_index = 0
@@ -502,6 +508,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
                     current_status = RunnerShutdown()
                 case _:
                     raise ValueError(

src/exo/worker/runner/runner_supervisor.py

@@ -49,6 +49,7 @@ class RunnerSupervisor:
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
+    _tg: TaskGroup | None = field(default=None, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
@@ -93,30 +94,27 @@ class RunnerSupervisor:
     async def run(self):
         self.runner_process.start()
         async with create_task_group() as tg:
+            self._tg = tg
             tg.start_soon(self._forward_events)
 
-
-    def shutdown(self):
-        logger.info("Runner supervisor shutting down")
         self._ev_recv.close()
         self._task_sender.close()
         self._event_sender.close()
-        self.runner_process.join(1)
+        await to_thread.run_sync(self.runner_process.join, 30)
         if not self.runner_process.is_alive():
-            logger.info("Runner process succesfully terminated")
             return
 
         # This is overkill but it's not technically bad, just unnecessary.
         logger.warning("Runner process didn't shutdown succesfully, terminating")
         self.runner_process.terminate()
-        self.runner_process.join(5)
+        await to_thread.run_sync(self.runner_process.join, 5)
         if not self.runner_process.is_alive():
             return
 
         logger.critical("Runner process didn't respond to SIGTERM, killing")
         self.runner_process.kill()
 
-        self.runner_process.join(5)
+        await to_thread.run_sync(self.runner_process.join, 5)
         if not self.runner_process.is_alive():
             return
 
@@ -124,6 +122,10 @@ class RunnerSupervisor:
             "Runner process didn't respond to SIGKILL. System resources may have leaked"
         )
 
+    def shutdown(self):
+        assert self._tg
+        self._tg.cancel_scope.cancel()
+
     async def start_task(self, task: Task):
         if task.task_id in self.pending:
             logger.warning(

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -201,29 +201,29 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
             TaskStatusUpdated(
                 task_id=INITIALIZATION_TASK_ID, task_status=TaskStatus.Running
             ),
-            TaskAcknowledged(task_id=INITIALIZATION_TASK_ID),
             RunnerStatusUpdated(
                 runner_id=RUNNER_1_ID, runner_status=RunnerConnecting()
             ),
+            TaskAcknowledged(task_id=INITIALIZATION_TASK_ID),
             TaskStatusUpdated(
                 task_id=INITIALIZATION_TASK_ID, task_status=TaskStatus.Complete
             ),
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerConnected()),
             TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Running),
-            TaskAcknowledged(task_id=LOAD_TASK_ID),
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoading()),
+            TaskAcknowledged(task_id=LOAD_TASK_ID),
             TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Complete),
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoaded()),
             TaskStatusUpdated(task_id=WARMUP_TASK_ID, task_status=TaskStatus.Running),
-            TaskAcknowledged(task_id=WARMUP_TASK_ID),
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerWarmingUp()),
+            TaskAcknowledged(task_id=WARMUP_TASK_ID),
             TaskStatusUpdated(task_id=WARMUP_TASK_ID, task_status=TaskStatus.Complete),
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerReady()),
             TaskStatusUpdated(
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Running
             ),
-            TaskAcknowledged(task_id=CHAT_COMPLETION_TASK_ID),
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerRunning()),
+            TaskAcknowledged(task_id=CHAT_COMPLETION_TASK_ID),
             expected_chunk,
             TaskStatusUpdated(
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Complete
@@ -231,10 +231,10 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
             # CHAT COMPLETION TASK SHOULD COMPLETE BEFORE RUNNER READY
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerReady()),
             TaskStatusUpdated(task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Running),
-            TaskAcknowledged(task_id=SHUTDOWN_TASK_ID),
             RunnerStatusUpdated(
                 runner_id=RUNNER_1_ID, runner_status=RunnerShuttingDown()
             ),
+            TaskAcknowledged(task_id=SHUTDOWN_TASK_ID),
             TaskStatusUpdated(
                 task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete
             ),

tests/run_exo_on.sh

@@ -19,23 +19,36 @@ git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
 }
 
 echo "Deploying $commit to $# hosts..."
+hosts=("$@")
 cleanup() {
+  for host in "${hosts[@]}"; do
+    ssh -T -o BatchMode=yes "$host@$host" "pkill -SIGINT -of exo-env" &
+  done
+  wait
   jobs -pr | xargs -r kill 2>/dev/null || true
 }
-trap 'jobs -pr | xargs -r kill 2>/dev/null || true' EXIT INT TERM
+trap 'cleanup' EXIT INT TERM
 
 colours=($'\e[31m' $'\e[32m' $'\e[33m' $'\e[34m')
 reset=$'\e[0m'
 i=0
 for host; do
   colour=${colours[i++ % 4]}
-  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo" \
-    |& awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
+  {
+    ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+      "/nix/var/nix/profiles/default/bin/nix shell nixpkgs#git -c bash -s -- '$commit'" \
+      2>&1 | awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
+  } <<'EOF'
+        set -euo pipefail
+        cd exo
+        git fetch -q origin
+        git checkout -q "$1"
+        EXO_LIBP2P_NAMESPACE="$1" /nix/var/nix/profiles/default/bin/nix run .#exo
+EOF
 done
 
 for host; do
   echo "Waiting for $host..."
-  until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done
+  until curl -sf "http://$host:52415/models"; do sleep 1; done
 done
 wait