Acknowledge task after runner status is updated

add resources directory to the nix exo package, and fixes the env for the dashboard dir

.github/workflows/pipeline.yml

@@ -142,4 +142,6 @@ jobs:
           # Run pytest outside sandbox (needs GPU access for MLX)
           export HOME="$RUNNER_TEMP"
           export EXO_TESTS=1
-          EXO_RESOURCES_DIR="$PWD/resources" $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib
+          export EXO_DASHBOARD_DIR="$PWD/dashboard/" 
+          export EXO_RESOURCES_DIR="$PWD/resources" 
+          $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib

python/parts.nix

@@ -69,7 +69,8 @@
           # Create wrapper scripts
           for script in exo exo-master exo-worker; do
             makeWrapper ${exoVenv}/bin/$script $out/bin/$script \
-              --set DASHBOARD_DIR ${self'.packages.dashboard} \
+              --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
+              --set EXO_RESOURCES_DIR ${inputs.self + "/resources"} \
               ${lib.optionalString pkgs.stdenv.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
           done
         '';

src/exo/shared/constants.py

@@ -39,7 +39,7 @@ RESOURCES_DIR = (
 )
 _DASHBOARD_DIR_ENV = os.environ.get("EXO_DASHBOARD_DIR", None)
 DASHBOARD_DIR = (
-    find_dashboard() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
+    find_dashboard() if _DASHBOARD_DIR_ENV is None else Path.home() / _DASHBOARD_DIR_ENV
 )
 
 # Log files (data/logs or cache)

src/exo/worker/runner/runner.py

@@ -125,7 +125,6 @@ def main(
             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)
@@ -137,6 +136,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")
@@ -153,6 +153,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(
@@ -195,6 +196,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:
@@ -234,6 +236,8 @@ 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
 
@@ -365,6 +369,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
                     try:
                         # Generate images using the image generation backend
@@ -430,6 +435,7 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
+                    event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
                     try:
                         image_index = 0
@@ -488,6 +494,8 @@ 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/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
             ),