better cancellation

src/exo/worker/runner/runner_supervisor.py

@@ -49,7 +49,6 @@ 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)
@@ -94,27 +93,30 @@ 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()
-        await to_thread.run_sync(self.runner_process.join, 30)
+        self.runner_process.join(1)
         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()
-        await to_thread.run_sync(self.runner_process.join, 5)
+        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()
 
-        await to_thread.run_sync(self.runner_process.join, 5)
+        self.runner_process.join(5)
         if not self.runner_process.is_alive():
             return
 
@@ -122,10 +124,6 @@ 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(

tests/run_exo_on.sh

@@ -19,36 +19,23 @@ 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 'cleanup' EXIT INT TERM
+trap 'jobs -pr | xargs -r kill 2>/dev/null || true' 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 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
+  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() }' &
 done
 
 for host; do
   echo "Waiting for $host..."
-  until curl -sf "http://$host:52415/models"; do sleep 1; done
+  until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done
 done
 wait