From efbf9850eb248a6f309a08825ce3320bec5dfd8b Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Sun, 15 Feb 2026 16:55:38 -0800 Subject: [PATCH] style: apply nix fmt to new edge case test file Co-Authored-By: Claude Opus 4.6 --- .../test_continuous_batching_edge_cases.py | 70 ++++++++++++++----- 1 file changed, 53 insertions(+), 17 deletions(-) diff --git a/src/exo/worker/tests/unittests/test_runner/test_continuous_batching_edge_cases.py b/src/exo/worker/tests/unittests/test_runner/test_continuous_batching_edge_cases.py index a6081ec5..672a19e7 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_continuous_batching_edge_cases.py +++ b/src/exo/worker/tests/unittests/test_runner/test_continuous_batching_edge_cases.py @@ -70,14 +70,21 @@ class ScriptedBatchEngine: """ def __init__(self, *_args: Any, **_kwargs: Any): - self._active: dict[int, tuple[CommandId, TaskId, list[tuple[str, FinishReason | None]]]] = {} + self._active: dict[ + int, tuple[CommandId, TaskId, list[tuple[str, FinishReason | None]]] + ] = {} self._pending: list[tuple[CommandId, TaskId, TextGenerationTaskParams]] = [] self._uid = 0 self.rank = 0 # map command_id -> scripted tokens, set externally before tasks arrive self.scripts: dict[str, list[tuple[str, FinishReason | None]]] = {} - def queue_request(self, command_id: CommandId, task_id: TaskId, task_params: TextGenerationTaskParams) -> None: + def queue_request( + self, + command_id: CommandId, + task_id: TaskId, + task_params: TextGenerationTaskParams, + ) -> None: self._pending.append((command_id, task_id, task_params)) def sync_and_insert_pending(self) -> list[int]: @@ -138,11 +145,18 @@ class FakeBatchEngineWithTokens: def __init__(self, *_args: Any, **_kwargs: Any): self._active_requests: dict[int, tuple[CommandId, TaskId, int, int]] = {} - self._pending_inserts: list[tuple[CommandId, TaskId, TextGenerationTaskParams]] = [] + self._pending_inserts: list[ + tuple[CommandId, TaskId, TextGenerationTaskParams] + ] = [] self._uid_counter = 0 self.rank = 0 - def queue_request(self, command_id: CommandId, task_id: TaskId, task_params: TextGenerationTaskParams) -> None: + def queue_request( + self, + command_id: CommandId, + task_id: TaskId, + task_params: TextGenerationTaskParams, + ) -> None: self._pending_inserts.append((command_id, task_id, task_params)) def sync_and_insert_pending(self) -> list[int]: @@ -163,7 +177,9 @@ class FakeBatchEngineWithTokens: def step(self) -> list[BatchedGenerationResponse]: results: list[BatchedGenerationResponse] = [] done: list[int] = [] - for uid, (cmd_id, task_id, tokens_gen, max_tokens) in list(self._active_requests.items()): + for uid, (cmd_id, task_id, tokens_gen, max_tokens) in list( + self._active_requests.items() + ): tokens_gen += 1 finish = "stop" if tokens_gen >= max_tokens else None results.append( @@ -171,7 +187,10 @@ class FakeBatchEngineWithTokens: command_id=cmd_id, task_id=task_id, response=GenerationResponse( - token=tokens_gen, text=f"token{tokens_gen}", finish_reason=finish, usage=None + token=tokens_gen, + text=f"token{tokens_gen}", + finish_reason=finish, + usage=None, ), ) ) @@ -263,7 +282,9 @@ WARMUP_TASK = StartWarmup(task_id=TaskId("warmup"), instance_id=INSTANCE_1_ID) SETUP_TASKS: list[Task] = [INIT_TASK, LOAD_TASK, WARMUP_TASK] -def make_chat_task(task_id: str, command_id: str, max_tokens: int = 3) -> TextGeneration: +def make_chat_task( + task_id: str, command_id: str, max_tokens: int = 3 +) -> TextGeneration: return TextGeneration( task_id=TaskId(task_id), command_id=CommandId(command_id), @@ -340,7 +361,8 @@ def _run_with_tasks( def chunks_for(events: list[Event], command_id: str) -> list[ChunkGenerated]: return [ - e for e in events + e + for e in events if isinstance(e, ChunkGenerated) and e.command_id == CommandId(command_id) ] @@ -530,7 +552,11 @@ def test_all_requests_finish_on_same_step(): ] events = _run_with_tasks([*tasks], engine_instance=engine) - for cmd_id, expected_text in [("cmd_a", "alpha"), ("cmd_b", "beta"), ("cmd_c", "gamma")]: + for cmd_id, expected_text in [ + ("cmd_a", "alpha"), + ("cmd_b", "beta"), + ("cmd_c", "gamma"), + ]: c = chunks_for(events, cmd_id) assert len(c) == 1, f"Expected 1 chunk for {cmd_id}, got {len(c)}" assert isinstance(c[0].chunk, TokenChunk) @@ -544,13 +570,13 @@ def test_all_requests_finish_on_same_step(): # Runner should be back to RunnerReady after all completions last_status = [ - e for e in events + e + for e in events if isinstance(e, RunnerStatusUpdated) and not isinstance(e.runner_status, (RunnerRunning,)) ] ready_after_gen = [ - e for e in last_status - if isinstance(e.runner_status, RunnerReady) + e for e in last_status if isinstance(e.runner_status, RunnerReady) ] assert len(ready_after_gen) >= 2, ( "Expected RunnerReady after warmup and after generation completes" @@ -567,7 +593,12 @@ def test_staggered_completions_in_batch(): engine.scripts["c1"] = [("a", "stop")] # finishes step 1 engine.scripts["c2"] = [("a", None), ("b", "stop")] # finishes step 2 engine.scripts["c3"] = [("a", None), ("b", None), ("c", "stop")] # finishes step 3 - engine.scripts["c4"] = [("a", None), ("b", None), ("c", None), ("d", "stop")] # finishes step 4 + engine.scripts["c4"] = [ + ("a", None), + ("b", None), + ("c", None), + ("d", "stop"), + ] # finishes step 4 tasks = [ *SETUP_TASKS, @@ -596,7 +627,9 @@ def test_staggered_completions_in_batch(): @pytest.fixture def patch_batch_engine(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(FakeGroup())) - monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MockTokenizer))) + monkeypatch.setattr( + mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MockTokenizer)) + ) monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1)) monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", make_nothin(None)) monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngineWithTokens) @@ -637,11 +670,14 @@ def test_eight_requests_staggered(patch_batch_engine: None): # Verify runner transitions back to ready after all requests complete # Find the last RunnerReady before shutdown ready_events = [ - (idx, e) for idx, e in enumerate(events) - if isinstance(e, RunnerStatusUpdated) and isinstance(e.runner_status, RunnerReady) + (idx, e) + for idx, e in enumerate(events) + if isinstance(e, RunnerStatusUpdated) + and isinstance(e.runner_status, RunnerReady) ] shutdown_idx = next( - idx for idx, e in enumerate(events) + idx + for idx, e in enumerate(events) if isinstance(e, TaskStatusUpdated) and e.task_id == TaskId("shutdown") and e.task_status == TaskStatus.Running