← back to Exo
Acknowledge task after runner status is updated (#1381)
221640a65b4554985c99300775ea5695ccfd2238 · 2026-02-05 12:00:37 +0000 · rltakashige
## Motivation
Duplicate tasks are still observed.
## Changes
Moved task acknowledgement to after the runner has changed its status.
## Why It Works
Tasks now remain pending until the runner has updated its status.
## Test Plan
### Manual Testing
Seems to work fine from manual testing. Hard to test a race condition
though.
### Automated Testing
Updated the event ordering test.
Files touched
M src/exo/worker/runner/runner.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
Diff
commit 221640a65b4554985c99300775ea5695ccfd2238
Author: rltakashige <rl.takashige@gmail.com>
Date: Thu Feb 5 12:00:37 2026 +0000
Acknowledge task after runner status is updated (#1381)
## Motivation
Duplicate tasks are still observed.
## Changes
Moved task acknowledgement to after the runner has changed its status.
## Why It Works
Tasks now remain pending until the runner has updated its status.
## Test Plan
### Manual Testing
Seems to work fine from manual testing. Hard to test a race condition
though.
### Automated Testing
Updated the event ordering test.
---
src/exo/worker/runner/runner.py | 10 +++++++++-
.../worker/tests/unittests/test_runner/test_event_ordering.py | 10 +++++-----
2 files changed, 14 insertions(+), 6 deletions(-)
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 3232e0e6..109ea219 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -145,7 +145,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)
@@ -157,6 +156,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 +173,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 +216,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 +256,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
@@ -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,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(
diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
index 16a43f2b..edf5ef3a 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/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
),
← 6177550c Ciaran/parallel cfg (#1361)
·
back to Exo
·
feat: add uncertainty visualization with token-level logprob 01b86a9e →