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 new file mode 100644 index 00000000..a6081ec5 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_runner/test_continuous_batching_edge_cases.py @@ -0,0 +1,673 @@ +""" +Edge-case tests for continuous batching in the runner. + +Tests cover: +1. Concurrent requests with overlapping tool calls +2. Requests that finish mid-generation with 'length' reason +3. Multiple requests finishing on the same step() call +4. Batch of 5+ simultaneous completions +""" + +# ruff: noqa: E402 +# pyright: reportAny=false +# pyright: reportUnknownArgumentType=false +# pyright: reportUnknownMemberType=false +# pyright: reportAttributeAccessIssue=false +# pyright: reportInvalidTypeVarUse=false +# pyright: reportPrivateUsage=false + +import json +from typing import Any +from unittest.mock import MagicMock + +import pytest + +import exo.worker.runner.runner as mlx_runner +from exo.shared.types.api import FinishReason +from exo.shared.types.chunks import TokenChunk, ToolCallChunk +from exo.shared.types.common import CommandId, NodeId +from exo.shared.types.events import ( + ChunkGenerated, + Event, + RunnerStatusUpdated, + TaskStatusUpdated, +) +from exo.shared.types.tasks import ( + ConnectToGroup, + LoadModel, + Shutdown, + StartWarmup, + Task, + TaskId, + TaskStatus, + TextGeneration, +) +from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams +from exo.shared.types.worker.runner_response import GenerationResponse +from exo.shared.types.worker.runners import RunnerReady, RunnerRunning +from exo.utils.channels import mp_channel +from exo.worker.engines.mlx.generator.batch_engine import ( + BatchedGenerationResponse, +) +from exo.worker.tests.constants import ( + INSTANCE_1_ID, + MODEL_A_ID, + NODE_A, + RUNNER_1_ID, +) +from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance + +# --------------------------------------------------------------------------- +# Fake batch engines +# --------------------------------------------------------------------------- + + +class ScriptedBatchEngine: + """Batch engine driven by scripted per-request token sequences. + + Each request produces a predefined list of (text, finish_reason) pairs. + One step() call pops one token per active request. + """ + + def __init__(self, *_args: Any, **_kwargs: Any): + 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: + self._pending.append((command_id, task_id, task_params)) + + def sync_and_insert_pending(self) -> list[int]: + uids: list[int] = [] + for cmd_id, task_id, _params in self._pending: + uid = self._uid + self._uid += 1 + script = list(self.scripts.get(str(cmd_id), [("tok", "stop")])) + self._active[uid] = (cmd_id, task_id, script) + uids.append(uid) + self._pending.clear() + return uids + + @property + def has_pending_inserts(self) -> bool: + return bool(self._pending) + + def step(self) -> list[BatchedGenerationResponse]: + results: list[BatchedGenerationResponse] = [] + done: list[int] = [] + for uid, (cmd_id, task_id, script) in self._active.items(): + if not script: + continue + text, finish_reason = script.pop(0) + results.append( + BatchedGenerationResponse( + command_id=cmd_id, + task_id=task_id, + response=GenerationResponse( + token=0, text=text, finish_reason=finish_reason, usage=None + ), + ) + ) + if finish_reason is not None: + done.append(uid) + for uid in done: + del self._active[uid] + return results + + @property + def has_active_requests(self) -> bool: + return bool(self._active) + + @property + def active_count(self) -> int: + return len(self._active) + + def sync_completions(self) -> None: + pass + + @property + def is_distributed(self) -> bool: + return False + + +class FakeBatchEngineWithTokens: + """Generates N tokens per request (reused from the main test file).""" + + 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._uid_counter = 0 + self.rank = 0 + + 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]: + uids: list[int] = [] + for command_id, task_id, task_params in self._pending_inserts: + uid = self._uid_counter + self._uid_counter += 1 + max_tokens = task_params.max_output_tokens or 3 + self._active_requests[uid] = (command_id, task_id, 0, max_tokens) + uids.append(uid) + self._pending_inserts.clear() + return uids + + @property + def has_pending_inserts(self) -> bool: + return bool(self._pending_inserts) + + 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()): + tokens_gen += 1 + finish = "stop" if tokens_gen >= max_tokens else None + results.append( + BatchedGenerationResponse( + command_id=cmd_id, + task_id=task_id, + response=GenerationResponse( + token=tokens_gen, text=f"token{tokens_gen}", finish_reason=finish, usage=None + ), + ) + ) + if finish: + done.append(uid) + else: + self._active_requests[uid] = (cmd_id, task_id, tokens_gen, max_tokens) + for uid in done: + del self._active_requests[uid] + return results + + @property + def has_active_requests(self) -> bool: + return bool(self._active_requests) + + @property + def active_count(self) -> int: + return len(self._active_requests) + + def sync_completions(self) -> None: + pass + + @property + def is_distributed(self) -> bool: + return False + + +# --------------------------------------------------------------------------- +# Mock tokenizers +# --------------------------------------------------------------------------- + + +class MockTokenizer: + tool_parser = None + tool_call_start = None + tool_call_end = None + has_tool_calling = False + + +class MockToolTokenizer: + """Tokenizer with tool calling enabled for testing.""" + + has_tool_calling = True + tool_call_start = "" + tool_call_end = "" + + @staticmethod + def _tool_parser(text: str) -> dict[str, Any]: + return json.loads(text) + + +class FakeGroup: + def rank(self) -> int: + return 0 + + def size(self) -> int: + return 1 + + +# --------------------------------------------------------------------------- +# Event collector & runner helper +# --------------------------------------------------------------------------- + + +class EventCollector: + def __init__(self) -> None: + self.events: list[Event] = [] + + def send(self, event: Event) -> None: + self.events.append(event) + + def close(self) -> None: + pass + + def join(self) -> None: + pass + + +def make_nothin[T, U, V](res: T): + def nothin(*_1: U, **_2: V) -> T: + return res + + return nothin + + +INIT_TASK = ConnectToGroup(task_id=TaskId("init"), instance_id=INSTANCE_1_ID) +LOAD_TASK = LoadModel(task_id=TaskId("load"), instance_id=INSTANCE_1_ID) +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: + return TextGeneration( + task_id=TaskId(task_id), + command_id=CommandId(command_id), + task_params=TextGenerationTaskParams( + model=MODEL_A_ID, + input=[InputMessage(role="user", content="hello")], + stream=True, + max_output_tokens=max_tokens, + ), + instance_id=INSTANCE_1_ID, + ) + + +def _run_with_tasks( + tasks: list[Task], + engine_cls: type = FakeBatchEngineWithTokens, + tokenizer_cls: type = MockTokenizer, + engine_instance: Any | None = None, +) -> list[Event]: + """Run tasks through the runner with configurable engine and tokenizer.""" + bound = get_bound_mlx_ring_instance( + instance_id=INSTANCE_1_ID, + model_id=MODEL_A_ID, + runner_id=RUNNER_1_ID, + node_id=NodeId(NODE_A), + ) + task_sender, task_receiver = mp_channel[Task]() + collector = EventCollector() + shutdown = Shutdown( + task_id=TaskId("shutdown"), + instance_id=INSTANCE_1_ID, + runner_id=RUNNER_1_ID, + ) + + import exo.worker.runner.runner as r + + orig_init_mlx = r.initialize_mlx + orig_load = r.load_mlx_items + orig_warmup = r.warmup_inference + orig_check = r._check_for_debug_prompts + orig_engine = r.BatchGenerationEngine + + r.initialize_mlx = make_nothin(FakeGroup()) + r.load_mlx_items = make_nothin((MagicMock(), tokenizer_cls)) + r.warmup_inference = make_nothin(1) + r._check_for_debug_prompts = make_nothin(None) + if engine_instance is not None: + r.BatchGenerationEngine = lambda *_a, **_kw: engine_instance # pyright: ignore[reportUnknownLambdaType] + else: + r.BatchGenerationEngine = engine_cls + + try: + with task_sender: + for t in tasks: + task_sender.send(t) + task_sender.send(shutdown) + task_receiver.close = lambda: None + task_receiver.join = lambda: None + r.main(bound, collector, task_receiver) # pyright: ignore[reportArgumentType] + finally: + r.initialize_mlx = orig_init_mlx + r.load_mlx_items = orig_load + r.warmup_inference = orig_warmup + r._check_for_debug_prompts = orig_check + r.BatchGenerationEngine = orig_engine + + return collector.events + + +# --------------------------------------------------------------------------- +# Helpers for querying events +# --------------------------------------------------------------------------- + + +def chunks_for(events: list[Event], command_id: str) -> list[ChunkGenerated]: + return [ + e for e in events + if isinstance(e, ChunkGenerated) and e.command_id == CommandId(command_id) + ] + + +def completed_task_ids(events: list[Event]) -> set[TaskId]: + return { + e.task_id + for e in events + if isinstance(e, TaskStatusUpdated) and e.task_status == TaskStatus.Complete + } + + +# =========================================================================== +# Test 1: Concurrent requests with overlapping tool calls +# =========================================================================== + + +def test_concurrent_tool_calls_and_normal_text(): + """Two concurrent requests: one emits normal text, the other a tool call. + + Verifies that: + - The normal request produces TokenChunks with its text + - The tool-call request produces a ToolCallChunk + - Both tasks complete + """ + engine = ScriptedBatchEngine() + # cmd_normal: 2 normal tokens then stop + engine.scripts["cmd_normal"] = [ + ("hello", None), + (" world", "stop"), + ] + # cmd_tool: tool_start, body, tool_end (suppressed), then finish + engine.scripts["cmd_tool"] = [ + ("", None), # swallowed by tracker + ('{"name":"get_weather","arguments":{"city":"SF"}}', None), # accumulated + ("", None), # triggers ToolCallChunk emission + ("done", "stop"), # normal trailing token + ] + + chat_normal = make_chat_task("t_normal", "cmd_normal", max_tokens=100) + chat_tool = make_chat_task("t_tool", "cmd_tool", max_tokens=100) + + events = _run_with_tasks( + [*SETUP_TASKS, chat_normal, chat_tool], + tokenizer_cls=MockToolTokenizer, + engine_instance=engine, + ) + + # Normal request: all chunks should be TokenChunk + normal_chunks = chunks_for(events, "cmd_normal") + assert len(normal_chunks) == 2 + assert all(isinstance(c.chunk, TokenChunk) for c in normal_chunks) + assert normal_chunks[-1].chunk.finish_reason == "stop" + + # Tool-call request + tool_chunks = chunks_for(events, "cmd_tool") + # → swallowed, body → accumulated, → ToolCallChunk, "done" → TokenChunk + tool_call_events = [c for c in tool_chunks if isinstance(c.chunk, ToolCallChunk)] + token_events = [c for c in tool_chunks if isinstance(c.chunk, TokenChunk)] + + assert len(tool_call_events) == 1, ( + f"Expected 1 ToolCallChunk, got {len(tool_call_events)}" + ) + tc_chunk = tool_call_events[0].chunk + assert isinstance(tc_chunk, ToolCallChunk) + assert tc_chunk.tool_calls[0].name == "get_weather" + assert json.loads(tc_chunk.tool_calls[0].arguments) == {"city": "SF"} + + assert len(token_events) == 1, "Expected 1 trailing TokenChunk after tool call" + assert token_events[0].chunk.finish_reason == "stop" + + # Both tasks should complete + done = completed_task_ids(events) + assert TaskId("t_normal") in done + assert TaskId("t_tool") in done + + +def test_tool_call_interrupted_by_finish_reason(): + """Tool call in progress when finish_reason fires — partial text emitted.""" + engine = ScriptedBatchEngine() + engine.scripts["cmd1"] = [ + ("", None), + ('{"name":"f"', "stop"), # finish while inside tool call + ] + + chat = make_chat_task("t1", "cmd1", max_tokens=100) + events = _run_with_tasks( + [*SETUP_TASKS, chat], + tokenizer_cls=MockToolTokenizer, + engine_instance=engine, + ) + + chunks = chunks_for(events, "cmd1") + assert len(chunks) == 1 + chunk = chunks[0].chunk + assert isinstance(chunk, TokenChunk) + # The interrupted tool call should be emitted as raw text + assert "" in chunk.text + assert '{"name":"f"' in chunk.text + assert chunk.finish_reason == "stop" + + assert TaskId("t1") in completed_task_ids(events) + + +# =========================================================================== +# Test 2: Request finishing with 'length' reason (timeout mid-generation) +# =========================================================================== + + +def test_request_finishes_with_length_reason(): + """Request that hits max_tokens limit and finishes with 'length'.""" + engine = ScriptedBatchEngine() + engine.scripts["cmd1"] = [ + ("tok1", None), + ("tok2", None), + ("tok3", "length"), # hit the token limit + ] + + chat = make_chat_task("t1", "cmd1", max_tokens=100) + events = _run_with_tasks( + [*SETUP_TASKS, chat], + engine_instance=engine, + ) + + chunks = chunks_for(events, "cmd1") + assert len(chunks) == 3 + + # Last chunk should have finish_reason="length" + assert isinstance(chunks[-1].chunk, TokenChunk) + assert chunks[-1].chunk.finish_reason == "length" + + # Earlier chunks should have no finish_reason + for c in chunks[:-1]: + assert isinstance(c.chunk, TokenChunk) + assert c.chunk.finish_reason is None + + assert TaskId("t1") in completed_task_ids(events) + + +def test_mixed_finish_reasons_across_requests(): + """Two requests finishing with different reasons: 'stop' and 'length'.""" + engine = ScriptedBatchEngine() + engine.scripts["cmd_stop"] = [("a", None), ("b", "stop")] + engine.scripts["cmd_len"] = [("x", None), ("y", "length")] + + chat1 = make_chat_task("t_stop", "cmd_stop", max_tokens=100) + chat2 = make_chat_task("t_len", "cmd_len", max_tokens=100) + + events = _run_with_tasks( + [*SETUP_TASKS, chat1, chat2], + engine_instance=engine, + ) + + stop_chunks = chunks_for(events, "cmd_stop") + len_chunks = chunks_for(events, "cmd_len") + + assert stop_chunks[-1].chunk.finish_reason == "stop" + assert len_chunks[-1].chunk.finish_reason == "length" + + done = completed_task_ids(events) + assert TaskId("t_stop") in done + assert TaskId("t_len") in done + + +# =========================================================================== +# Test 3: Multiple finish reasons in rapid succession (same step) +# =========================================================================== + + +def test_all_requests_finish_on_same_step(): + """Three requests that all finish on the same step() call. + + This tests that the runner and _process_generation_results correctly + handle multiple completions in a single step. + """ + engine = ScriptedBatchEngine() + # All three produce exactly 1 token and finish + engine.scripts["cmd_a"] = [("alpha", "stop")] + engine.scripts["cmd_b"] = [("beta", "stop")] + engine.scripts["cmd_c"] = [("gamma", "stop")] + + tasks = [ + *SETUP_TASKS, + make_chat_task("ta", "cmd_a", max_tokens=100), + make_chat_task("tb", "cmd_b", max_tokens=100), + make_chat_task("tc", "cmd_c", max_tokens=100), + ] + events = _run_with_tasks([*tasks], engine_instance=engine) + + 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) + assert c[0].chunk.text == expected_text + assert c[0].chunk.finish_reason == "stop" + + done = completed_task_ids(events) + assert TaskId("ta") in done + assert TaskId("tb") in done + assert TaskId("tc") in done + + # Runner should be back to RunnerReady after all completions + last_status = [ + 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) + ] + assert len(ready_after_gen) >= 2, ( + "Expected RunnerReady after warmup and after generation completes" + ) + + +def test_staggered_completions_in_batch(): + """Four requests with different token counts — they complete at different steps. + + Verifies each request gets the right number of chunks and the runner + tracks active_requests correctly as requests drain. + """ + engine = ScriptedBatchEngine() + 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 + + tasks = [ + *SETUP_TASKS, + make_chat_task("t1", "c1", max_tokens=100), + make_chat_task("t2", "c2", max_tokens=100), + make_chat_task("t3", "c3", max_tokens=100), + make_chat_task("t4", "c4", max_tokens=100), + ] + events = _run_with_tasks([*tasks], engine_instance=engine) + + assert len(chunks_for(events, "c1")) == 1 + assert len(chunks_for(events, "c2")) == 2 + assert len(chunks_for(events, "c3")) == 3 + assert len(chunks_for(events, "c4")) == 4 + + done = completed_task_ids(events) + for tid in ["t1", "t2", "t3", "t4"]: + assert TaskId(tid) in done, f"Task {tid} should be complete" + + +# =========================================================================== +# Test 4: Batch of 5+ simultaneous completions +# =========================================================================== + + +@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, "warmup_inference", make_nothin(1)) + monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", make_nothin(None)) + monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngineWithTokens) + + +def test_five_simultaneous_completions(patch_batch_engine: None): + """Five requests submitted together, all generating tokens and completing.""" + chats = [make_chat_task(f"t{i}", f"cmd{i}", max_tokens=2) for i in range(5)] + events = _run_with_tasks([*SETUP_TASKS, *chats]) + + for i in range(5): + c = chunks_for(events, f"cmd{i}") + assert len(c) == 2, f"Expected 2 chunks for cmd{i}, got {len(c)}" + assert c[-1].chunk.finish_reason == "stop" + + done = completed_task_ids(events) + for i in range(5): + assert TaskId(f"t{i}") in done + + +def test_eight_requests_staggered(patch_batch_engine: None): + """Eight requests with varying token counts, verifying all complete correctly.""" + chats = [make_chat_task(f"t{i}", f"cmd{i}", max_tokens=i + 1) for i in range(8)] + events = _run_with_tasks([*SETUP_TASKS, *chats]) + + for i in range(8): + c = chunks_for(events, f"cmd{i}") + expected = i + 1 + assert len(c) == expected, ( + f"Expected {expected} chunks for cmd{i}, got {len(c)}" + ) + assert c[-1].chunk.finish_reason == "stop" + + done = completed_task_ids(events) + for i in range(8): + assert TaskId(f"t{i}") in done + + # 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) + ] + shutdown_idx = next( + idx for idx, e in enumerate(events) + if isinstance(e, TaskStatusUpdated) + and e.task_id == TaskId("shutdown") + and e.task_status == TaskStatus.Running + ) + # There should be a RunnerReady event between generation and shutdown + ready_before_shutdown = [idx for idx, _ in ready_events if idx < shutdown_idx] + assert len(ready_before_shutdown) >= 1, ( + "Expected RunnerReady between generation completion and shutdown" + ) + + +def test_ten_simultaneous_single_token(): + """Ten requests that each produce exactly one token — all finish on step 1.""" + engine = ScriptedBatchEngine() + for i in range(10): + engine.scripts[f"cmd{i}"] = [(f"word{i}", "stop")] + + chats = [make_chat_task(f"t{i}", f"cmd{i}", max_tokens=100) for i in range(10)] + events = _run_with_tasks([*SETUP_TASKS, *chats], engine_instance=engine) + + for i in range(10): + c = chunks_for(events, f"cmd{i}") + assert len(c) == 1 + assert isinstance(c[0].chunk, TokenChunk) + assert c[0].chunk.text == f"word{i}" + assert c[0].chunk.finish_reason == "stop" + + done = completed_task_ids(events) + assert len(done & {TaskId(f"t{i}") for i in range(10)}) == 10