From 32593224d86caac95516982f4a455eccae9d7d4f Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Tue, 7 Apr 2026 22:51:00 -0500 Subject: [PATCH 01/32] Fix pipeline_parallel_prefill short-prompt wedge + dashboard kv-backend display Two small bugs surfaced during the rotorquant rollout, neither caused by rotorquant itself but both blocking real usage: 1. Pipeline-parallel prefill wedges when ``n_real == 1`` (the prompt fits inside a single per-rank chunk). PR #101's ``cb2dc68e`` removed the prior ``num_tokens >= prefill_step_size`` gate to fix a Gemma warmup hang in stream_generate; the unintended consequence was that *every* short prompt now routed through pipeline_parallel_prefill and hit a different wedge there. Reproduced on Gemma 4 26B during 23-token warmup. Fix: route a prompt through pipeline_parallel_prefill only when it produces at least two effective per-rank chunks. Long prompts still use the pipeline path PR #101 was preserving; short prompts go through stream_generate, which is known to work for that shape. 2. ``GET /config`` reads only ``EXO_KV_CACHE_BACKEND`` when ``skulk.yaml`` exists on disk, ignoring the newer ``SKULK_KV_CACHE_BACKEND`` env var. The ``fileExists=False`` branch already handles both, so the dashboard silently displayed "default" even though the runtime was honoring the SKULK_ override correctly. One-line fix to match the other branch. Tests: - Updated test_prefill_uses_pipeline_parallel_path_for_long_pipeline_prompts to use a 5000-token prompt (4 chunks at group_size=3, qualifies for the pipeline path) - Added test_prefill_uses_stream_generate_for_short_pipeline_prompts as a regression test for the Gemma 4 23-token warmup wedge Co-Authored-By: Claude Opus 4.6 (1M context) --- src/exo/api/main.py | 3 +- .../worker/engines/mlx/generator/generate.py | 29 ++++++-- .../test_mlx/test_prefill_path_selection.py | 68 ++++++++++++++++++- 3 files changed, 92 insertions(+), 8 deletions(-) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index 6fbe43bef..953b32dc1 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -2592,7 +2592,8 @@ async def get_config(self) -> JSONResponse: "fileExists": True, "effective": { "kv_cache_backend": os.environ.get( - "EXO_KV_CACHE_BACKEND", "default" + "SKULK_KV_CACHE_BACKEND", + os.environ.get("EXO_KV_CACHE_BACKEND", "default"), ), "has_hf_token": has_hf_token or "HF_TOKEN" in os.environ, }, diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index bbac32de2..07ba27f3c 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -466,16 +466,37 @@ def combined_progress_callback(processed: int, total: int) -> None: effective_prefill_step_size = ( prefill_step_size // min(4, group_size) if is_pipeline else prefill_step_size ) + + # Pipeline-parallel prefill currently wedges when the prompt fits inside a + # single per-rank chunk (i.e. ``n_real == 1`` in the prefill loop). This + # was observed during Gemma 4 warmup with a 23-token prompt and reproduced + # in the issue tracker for any short prompt under pipeline parallelism. + # PR #101's pipeline-prefill widening (cb2dc68e) removed the prior + # ``num_tokens >= prefill_step_size`` gate to fix a different Gemma warmup + # hang in stream_generate; the result was that *every* short prompt now + # routed through pipeline_parallel_prefill and hit the new wedge. The + # narrower gate below preserves PR #101's fix for non-trivial prompts + # (where pipeline_parallel_prefill is the right path) while routing + # short / warmup prompts back through stream_generate, which is known to + # work for that shape. See PR review thread on PR #103 for the full + # diagnosis. + pipeline_chunks = ( + (num_tokens + effective_prefill_step_size - 1) // effective_prefill_step_size + if effective_prefill_step_size > 0 + else 0 + ) + use_pipeline_prefill = is_pipeline and pipeline_chunks >= 2 logger.info( "Prefill path selected: " - f"{'pipeline_parallel_prefill' if is_pipeline else 'stream_generate'} " - f"(rank={rank}, prompt_tokens={num_tokens}, " + f"{'pipeline_parallel_prefill' if use_pipeline_prefill else 'stream_generate'} " + f"(rank={rank}, prompt_tokens={num_tokens}, is_pipeline={is_pipeline}, " f"prefill_step_size_input={prefill_step_size}, " - f"prefill_step_size_effective={effective_prefill_step_size})" + f"prefill_step_size_effective={effective_prefill_step_size}, " + f"pipeline_chunks={pipeline_chunks})" ) try: - if is_pipeline: + if use_pipeline_prefill: set_pipeline_queue_sends(model, queue_sends=True) assert group is not None, "Pipeline prefill requires a distributed group" with _hang_debug_watch( diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py b/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py index 54e77f768..0b8c00067 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py @@ -34,9 +34,15 @@ def __init__(self) -> None: self.layers: list[object] = [] -def test_prefill_uses_pipeline_parallel_path_for_short_pipeline_prompts( +def test_prefill_uses_pipeline_parallel_path_for_long_pipeline_prompts( monkeypatch: pytest.MonkeyPatch, ) -> None: + """Long prompts on pipeline-parallel models should use pipeline_parallel_prefill. + + The threshold is two effective per-rank chunks. With group_size=3 and + prefill_step_size=4096, the effective per-rank chunk size is 1365, so a + 5000-token prompt produces 4 chunks and qualifies. + """ calls: list[str] = [] fake_cache = _FakeCache() @@ -70,7 +76,7 @@ def _pipeline_parallel_prefill(**kwargs: object) -> None: model=_FakeModel(), tokenizer=object(), sampler=object(), - prompt_tokens=list(range(128)), + prompt_tokens=list(range(5000)), cache=[fake_cache], group=_FakeGroup(), on_prefill_progress=None, @@ -78,7 +84,63 @@ def _pipeline_parallel_prefill(**kwargs: object) -> None: ) assert calls == ["pipeline_parallel_prefill"] - assert prefill_tokens == 128 + assert prefill_tokens == 5000 + assert snapshots == [] + assert prefill_tps >= 0.0 + assert fake_cache.trim_calls == [2] + + +def test_prefill_uses_stream_generate_for_short_pipeline_prompts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Short prompts that fit in a single per-rank chunk must avoid pipeline prefill. + + Regression for the Gemma 4 warmup wedge: pipeline_parallel_prefill hangs + when ``n_real == 1`` (the prompt fits inside a single per-rank chunk), + so any such prompt — including the 23-token warmup prompt — must route + through stream_generate even on pipeline-parallel models. PR #101's + pipeline-prefill widening removed this guard and reintroduced the hang + on Gemma 4; this test locks the guard back in. + """ + calls: list[str] = [] + fake_cache = _FakeCache() + + monkeypatch.setattr(generate_module, "mx_barrier", lambda _group: None) + monkeypatch.setattr( + generate_module, + "_has_pipeline_communication_layer", + lambda _model: True, + ) + + def _pipeline_parallel_prefill(**kwargs: object) -> None: + calls.append("pipeline_parallel_prefill") + + monkeypatch.setattr( + generate_module, + "pipeline_parallel_prefill", + _pipeline_parallel_prefill, + ) + + def _stream_generate(*args: object, **kwargs: object) -> Iterator[object]: + calls.append("stream_generate") + yield object() + + monkeypatch.setattr(generate_module, "stream_generate", _stream_generate) + + prefill_tps, prefill_tokens, snapshots = generate_module.prefill( + model=_FakeModel(), + tokenizer=object(), + sampler=object(), + prompt_tokens=list(range(23)), # exactly the gemma 4 warmup prompt size + cache=[fake_cache], + group=_FakeGroup(), + on_prefill_progress=None, + distributed_prompt_progress_callback=None, + ) + + assert calls == ["stream_generate"] + assert "pipeline_parallel_prefill" not in calls + assert prefill_tokens == 23 assert snapshots == [] assert prefill_tps >= 0.0 assert fake_cache.trim_calls == [2] From edf7145a0de61b61d1803e9f84b08d42d1696b4d Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Tue, 7 Apr 2026 23:25:15 -0500 Subject: [PATCH 02/32] Temporarily bypass LLM warmup for debugging --- src/exo/worker/runner/llm_inference/runner.py | 16 ++++++++++- .../test_runner/test_event_ordering.py | 27 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index ca5be4eaa..fa94cf737 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -69,6 +69,13 @@ from .tool_parsers import make_mlx_parser +def _should_skip_llm_warmup() -> bool: + # Temporary experiment for Gemma 4 pipeline bring-up: let runners report + # ready without synthetic warmup so we can probe the first real prompt path. + # Set SKULK_FORCE_LLM_WARMUP=1 to restore the normal behavior. + return os.environ.get("SKULK_FORCE_LLM_WARMUP") != "1" + + class ExitCode(str, Enum): AllTasksComplete = "AllTasksComplete" Shutdown = "Shutdown" @@ -243,7 +250,14 @@ def on_layer_loaded(layers_loaded: int, total: int) -> None: self.update_status(RunnerWarmingUp()) self.acknowledge_task(task) - self.generator.warmup() + if _should_skip_llm_warmup(): + logger.warning( + "Skipping LLM warmup and marking runner ready " + "(temporary debug bypass; set SKULK_FORCE_LLM_WARMUP=1 " + "to restore synthetic warmup)" + ) + else: + self.generator.warmup() logger.info( f"runner initialized in {time.time() - self.setup_start_time} seconds" 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 cd4c18d70..b9cff269e 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 @@ -344,3 +344,30 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch): RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()), ], ) + + +def test_warmup_task_can_be_bypassed_while_runner_still_becomes_ready( + patch_out_mlx: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("SKULK_FORCE_LLM_WARMUP", raising=False) + + def fail_if_called(*_args: object, **_kwargs: object) -> None: + raise AssertionError("warmup() should be bypassed in this debug branch") + + monkeypatch.setattr(mlx_batch_generator.BatchGenerator, "warmup", fail_if_called) + monkeypatch.setattr( + mlx_batch_generator.SequentialGenerator, "warmup", fail_if_called + ) + + events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, SHUTDOWN_TASK]) + + assert any( + isinstance(event, RunnerStatusUpdated) + and isinstance(event.runner_status, RunnerWarmingUp) + for event in events + ) + assert any( + isinstance(event, RunnerStatusUpdated) + and isinstance(event.runner_status, RunnerReady) + for event in events + ) From 70e4192adcb4b4ae080517f948d85ba6414637e7 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Tue, 7 Apr 2026 23:27:25 -0500 Subject: [PATCH 03/32] Add pipeline wrapper hang instrumentation --- src/exo/worker/engines/mlx/auto_parallel.py | 151 +++++++++++++++++--- 1 file changed, 130 insertions(+), 21 deletions(-) diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index 22154a7e7..6ab5f4bfc 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -1,9 +1,14 @@ +import contextlib import os +import sys import threading +import time +import traceback from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Generator from functools import partial from inspect import signature +from types import FrameType from typing import TYPE_CHECKING, Any, Literal, Protocol, cast import mlx.core as mx @@ -70,10 +75,83 @@ _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = [] +def _mlx_hang_debug_enabled() -> bool: + value = os.environ.get("SKULK_MLX_HANG_DEBUG") or os.environ.get( + "EXO_MLX_HANG_DEBUG" + ) + if value is None: + return False + return value.lower() not in {"", "0", "false", "no", "off"} + + +def _mlx_hang_debug_interval_seconds() -> float: + raw = os.environ.get("SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS") or os.environ.get( + "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS" + ) + if raw is None: + return 30.0 + with contextlib.suppress(ValueError): + return max(float(raw), 1.0) + return 30.0 + + +@contextlib.contextmanager +def _hang_debug_watch(label: str) -> Generator[None]: + """Emit periodic stack-rich logs while a pipeline wrapper stage is stuck.""" + if not _mlx_hang_debug_enabled(): + yield + return + + interval_seconds = _mlx_hang_debug_interval_seconds() + started_at = time.monotonic() + monitored_thread_id = threading.get_ident() + finished = threading.Event() + + logger.info( + f"[hang-debug] Entering {label} (watchdog interval={interval_seconds:.0f}s)" + ) + + def watchdog() -> None: + while not finished.wait(timeout=interval_seconds): + elapsed = time.monotonic() - started_at + current_frames = cast( + Callable[[], dict[int, FrameType]] | None, + getattr(sys, "_current_frames", None), + ) + frame = ( + None + if current_frames is None + else current_frames().get(monitored_thread_id) + ) + if frame is None: + stack_text = "" + else: + stack_text = "".join(traceback.format_stack(frame)) + logger.warning( + f"[hang-debug] Still in {label} after {elapsed:.1f}s\n{stack_text}" + ) + + watchdog_thread = threading.Thread( + target=watchdog, + name=f"hang-debug:{label}", + daemon=True, + ) + watchdog_thread.start() + + try: + yield + finally: + finished.set() + elapsed = time.monotonic() - started_at + logger.info(f"[hang-debug] Leaving {label} after {elapsed:.1f}s") + + def flush_prefill_sends() -> None: for output, dst, group in _pending_prefill_sends: - sent = mx.distributed.send(output, dst, group=group) - mx.async_eval(sent) + with _hang_debug_watch(f"flush_prefill_sends send dst={dst}"): + sent = mx.distributed.send(output, dst, group=group) + with _hang_debug_watch(f"flush_prefill_sends async_eval dst={dst}"): + mx.async_eval(sent) _pending_prefill_sends.clear() @@ -163,10 +241,20 @@ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: if self.r != 0: # We want to avoid GPU timeout errors by evalling the distributed operation # so that it stays on CPU, which does not have a timeout. - mx.eval(x) - x = mx.distributed.recv_like(x, (self.r - 1), group=self.group) - mx.eval(x) - return self.original_layer(x, *args, **kwargs) + with _hang_debug_watch( + f"pipeline_first eval_input rank={self.r} src={self.r - 1}" + ): + mx.eval(x) + with _hang_debug_watch( + f"pipeline_first recv_like rank={self.r} src={self.r - 1}" + ): + x = mx.distributed.recv_like(x, (self.r - 1), group=self.group) + with _hang_debug_watch( + f"pipeline_first eval_recv rank={self.r} src={self.r - 1}" + ): + mx.eval(x) + with _hang_debug_watch(f"pipeline_first original_layer rank={self.r}"): + return self.original_layer(x, *args, **kwargs) class PipelineLastLayer(CustomMlxLayer): @@ -190,11 +278,17 @@ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: x, *args, **kwargs ).arguments.get("cache", None) - output: mx.array = self.original_layer(x, *args, **kwargs) + with _hang_debug_watch( + f"pipeline_last original_layer rank={self.r} world={self.s} prefill={self.is_prefill}" + ): + output: mx.array = self.original_layer(x, *args, **kwargs) # Eval layer output to materialize it before send — this splits the graph # so the send is isolated and the receiving rank's recv can complete. - mx.eval(output) + with _hang_debug_watch( + f"pipeline_last eval_output rank={self.r} world={self.s} prefill={self.is_prefill}" + ): + mx.eval(output) if self.r != self.s - 1: if self.queue_sends: @@ -202,24 +296,39 @@ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: (output, (self.r + 1) % self.s, self.group) ) else: - output = mx.distributed.send( - output, (self.r + 1) % self.s, group=self.group - ) + with _hang_debug_watch( + f"pipeline_last send rank={self.r} dst={(self.r + 1) % self.s} prefill={self.is_prefill}" + ): + output = mx.distributed.send( + output, (self.r + 1) % self.s, group=self.group + ) if cache is not None: # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA) # doesn't have .keys directly; access via first sub-cache. _cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore if hasattr(_cache, "keys"): # pyright: ignore[reportAny] _cache.keys = mx.depends(_cache.keys, output) # type: ignore - mx.eval(output) + with _hang_debug_watch( + f"pipeline_last eval_send rank={self.r} dst={(self.r + 1) % self.s} prefill={self.is_prefill}" + ): + mx.eval(output) if cache is not None and hasattr(_cache, "keys"): # type: ignore - mx.eval(_cache.keys) # type: ignore + with _hang_debug_watch( + f"pipeline_last eval_cache_dep rank={self.r} dst={(self.r + 1) % self.s} prefill={self.is_prefill}" + ): + mx.eval(_cache.keys) # type: ignore if not self.is_prefill: - output = mx.distributed.all_gather(output, group=self.group)[ - -output.shape[0] : - ] - mx.eval(output) + with _hang_debug_watch( + f"pipeline_last all_gather rank={self.r} world={self.s}" + ): + output = mx.distributed.all_gather(output, group=self.group)[ + -output.shape[0] : + ] + with _hang_debug_watch( + f"pipeline_last eval_all_gather rank={self.r} world={self.s}" + ): + mx.eval(output) return output @@ -436,8 +545,8 @@ def _patch_one(target: object) -> None: if getattr(cls, "_exo_pipeline_patched", False): return - original_call = cls.__call__ # type: ignore[attr-defined] - call_signature = signature(original_call) # type: ignore[arg-type] + original_call = cls.__call__ + call_signature = signature(original_call) def patched_call( self: object, @@ -463,7 +572,7 @@ def patched_call( return logits - cls.__call__ = patched_call # type: ignore[method-assign] + cls.__call__ = patched_call cls._exo_pipeline_patched = True # type: ignore[attr-defined] _patch_one(model) From 430462e9dc039bba29dc9d266960d95140bcbec3 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Tue, 7 Apr 2026 23:41:51 -0500 Subject: [PATCH 04/32] Handle closed worker info event stream --- src/exo/worker/main.py | 21 +++++++---- .../unittests/test_worker_info_forwarding.py | 35 +++++++++++++++++++ 2 files changed, 49 insertions(+), 7 deletions(-) create mode 100644 src/exo/worker/tests/unittests/test_worker_info_forwarding.py diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 99911b76a..6358e25b6 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -6,7 +6,7 @@ from pathlib import Path import anyio -from anyio import fail_after +from anyio import BrokenResourceError, ClosedResourceError, fail_after from loguru import logger from PIL import Image @@ -173,13 +173,20 @@ async def run(self): async def _forward_info(self, recv: Receiver[GatheredInfo]): with recv as info_stream: async for info in info_stream: - await self.event_sender.send( - NodeGatheredInfo( - node_id=self.node_id, - when=str(datetime.now(tz=timezone.utc)), - info=info, + try: + await self.event_sender.send( + NodeGatheredInfo( + node_id=self.node_id, + when=str(datetime.now(tz=timezone.utc)), + info=info, + ) ) - ) + except (ClosedResourceError, BrokenResourceError): + logger.info( + "Worker info forwarding stopped because the event stream " + "was already closed" + ) + return async def _event_applier(self): with self.event_receiver as events: diff --git a/src/exo/worker/tests/unittests/test_worker_info_forwarding.py b/src/exo/worker/tests/unittests/test_worker_info_forwarding.py new file mode 100644 index 000000000..4c6c5ae90 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_worker_info_forwarding.py @@ -0,0 +1,35 @@ +import pytest + +from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand +from exo.shared.types.common import NodeId +from exo.shared.types.events import Event, IndexedEvent +from exo.utils.channels import channel +from exo.utils.info_gatherer.info_gatherer import GatheredInfo, MiscData +from exo.worker.main import Worker + + +@pytest.mark.asyncio +async def test_forward_info_ignores_closed_event_sender() -> None: + indexed_event_sender, indexed_event_receiver = channel[IndexedEvent]() + event_sender, _ = channel[Event]() + command_sender, _ = channel[ForwarderCommand]() + download_sender, _ = channel[ForwarderDownloadCommand]() + info_sender, info_receiver = channel[GatheredInfo]() + + worker = Worker( + node_id=NodeId("node-a"), + event_receiver=indexed_event_receiver, + event_sender=event_sender, + command_sender=command_sender, + download_command_sender=download_sender, + ) + + event_sender.close() + await info_sender.send(MiscData(friendly_name="kite3")) + info_sender.close() + + await worker._forward_info(info_receiver) # pyright: ignore[reportPrivateUsage] + + indexed_event_sender.close() + command_sender.close() + download_sender.close() From cc0c0332ac4cce1ffc9b8e78206101d64ed459cd Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Tue, 7 Apr 2026 23:58:07 -0500 Subject: [PATCH 05/32] Allow warmup planning to tolerate ready peers --- src/exo/worker/plan.py | 7 +- .../tests/unittests/test_plan/test_warmup.py | 86 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 1c2900eb8..f7dfdc2bb 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -251,14 +251,17 @@ def _ready_to_warmup( accepting_ranks_ready = device_rank > 0 and all( isinstance( all_runners.get(global_runner_id, None), - (RunnerLoaded, RunnerWarmingUp), + (RunnerLoaded, RunnerWarmingUp, RunnerReady), ) for global_runner_id in shard_assignments.runner_to_shard ) # Rank = 0 connecting_rank_ready = device_rank == 0 and all( - isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp) + isinstance( + all_runners.get(global_runner_id, None), + (RunnerWarmingUp, RunnerReady), + ) for global_runner_id in shard_assignments.runner_to_shard if global_runner_id != runner_id ) diff --git a/src/exo/worker/tests/unittests/test_plan/test_warmup.py b/src/exo/worker/tests/unittests/test_plan/test_warmup.py index 528484467..a40e4abc8 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py +++ b/src/exo/worker/tests/unittests/test_plan/test_warmup.py @@ -5,6 +5,7 @@ RunnerIdle, RunnerLoaded, RunnerLoading, + RunnerReady, RunnerWarmingUp, ) from exo.worker.tests.constants import ( @@ -67,6 +68,50 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming(): assert result.instance_id == INSTANCE_1_ID +def test_plan_starts_warmup_for_accepting_rank_when_peer_is_already_ready(): + """ + Non-zero ranks should still start warmup when a peer races through the + temporary debug warmup bypass and reaches Ready before they observe + WarmingUp. + """ + shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=3) + shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=3) + shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=2, world_size=3) + instance = get_mlx_ring_instance( + instance_id=INSTANCE_1_ID, + model_id=MODEL_A_ID, + node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID, NODE_C: RUNNER_3_ID}, + runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1, RUNNER_3_ID: shard2}, + ) + + bound_instance = BoundInstance( + instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B + ) + local_runner = FakeRunnerSupervisor( + bound_instance=bound_instance, status=RunnerLoaded() + ) + + runners = {RUNNER_2_ID: local_runner} + instances = {INSTANCE_1_ID: instance} + all_runners = { + RUNNER_1_ID: RunnerReady(), + RUNNER_2_ID: RunnerLoaded(), + RUNNER_3_ID: RunnerLoaded(), + } + + result = plan_mod.plan( + node_id=NODE_B, + runners=runners, # type: ignore + global_download_status={NODE_A: []}, + instances=instances, + all_runners=all_runners, + tasks={}, + ) + + assert isinstance(result, StartWarmup) + assert result.instance_id == INSTANCE_1_ID + + def test_plan_starts_warmup_for_rank_zero_after_others_warming(): """ For device_rank == 0, StartWarmup should only be emitted once all the @@ -108,6 +153,47 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming(): assert result.instance_id == INSTANCE_1_ID +def test_plan_starts_warmup_for_rank_zero_after_other_rank_is_already_ready(): + """ + Rank zero should still start warmup when another rank races from WarmingUp + to Ready before the planner observes the intermediate state. + """ + shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) + shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) + instance = get_mlx_ring_instance( + instance_id=INSTANCE_1_ID, + model_id=MODEL_A_ID, + node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, + runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, + ) + + bound_instance = BoundInstance( + instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A + ) + local_runner = FakeRunnerSupervisor( + bound_instance=bound_instance, status=RunnerLoaded() + ) + + runners = {RUNNER_1_ID: local_runner} + instances = {INSTANCE_1_ID: instance} + all_runners = { + RUNNER_1_ID: RunnerLoaded(), + RUNNER_2_ID: RunnerReady(), + } + + result = plan_mod.plan( + node_id=NODE_A, + runners=runners, # type: ignore + global_download_status={NODE_A: []}, + instances=instances, + all_runners=all_runners, + tasks={}, + ) + + assert isinstance(result, StartWarmup) + assert result.instance_id == INSTANCE_1_ID + + def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warming(): """ Non-zero rank should not start warmup while any shard is not Loaded/WarmingUp. From b91a99b62d206944ad658b710228ba8036ab8a74 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 00:43:24 -0500 Subject: [PATCH 06/32] Make Gemma warmup more representative --- .../worker/engines/mlx/generator/generate.py | 17 +++++- src/exo/worker/runner/llm_inference/runner.py | 12 ++-- .../unittests/test_mlx/test_warmup_request.py | 61 +++++++++++++++++++ .../test_runner/test_event_ordering.py | 2 +- 4 files changed, 82 insertions(+), 10 deletions(-) create mode 100644 src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 07ba27f3c..3bc2e457f 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -576,13 +576,24 @@ def warmup_inference( ) -> int: logger.info(f"warming up inference for instance: {model_id}") - content = "Prompt to warm up the inference engine. Repeat this." + instructions = ( + "You are a helpful assistant. Answer the user in one short sentence." + ) + content = ( + "Summarize this status update in one sentence: " + "Skulk is starting a distributed Gemma 4 model runner, verifying " + "pipeline communication between nodes, and preparing to serve chat requests." + ) warmup_task_params = TextGenerationTaskParams( model=model_id, + instructions=instructions, input=[InputMessage(role="user", content=content)], - max_output_tokens=50, - temperature=0.0, + max_output_tokens=32, + temperature=1.0, + top_p=0.95, + top_k=64, + enable_thinking=False, ) with _hang_debug_watch(f"warmup apply_chat_template model={model_id}"): diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index fa94cf737..bd65d7746 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -70,10 +70,10 @@ def _should_skip_llm_warmup() -> bool: - # Temporary experiment for Gemma 4 pipeline bring-up: let runners report - # ready without synthetic warmup so we can probe the first real prompt path. - # Set SKULK_FORCE_LLM_WARMUP=1 to restore the normal behavior. - return os.environ.get("SKULK_FORCE_LLM_WARMUP") != "1" + # Temporary escape hatch for debug sessions where we want runners to reach + # Ready without issuing the synthetic warmup request. Normal behavior keeps + # warmup enabled unless SKULK_SKIP_LLM_WARMUP=1 is set explicitly. + return os.environ.get("SKULK_SKIP_LLM_WARMUP") == "1" class ExitCode(str, Enum): @@ -253,8 +253,8 @@ def on_layer_loaded(layers_loaded: int, total: int) -> None: if _should_skip_llm_warmup(): logger.warning( "Skipping LLM warmup and marking runner ready " - "(temporary debug bypass; set SKULK_FORCE_LLM_WARMUP=1 " - "to restore synthetic warmup)" + "(temporary debug bypass; unset SKULK_SKIP_LLM_WARMUP " + "or set it to 0 to restore synthetic warmup)" ) else: self.generator.warmup() diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py new file mode 100644 index 000000000..7124be74b --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -0,0 +1,61 @@ +from typing import cast + +import pytest + +import exo.worker.engines.mlx.generator.generate as generate_mod +from exo.shared.types.common import ModelId +from exo.shared.types.text_generation import TextGenerationTaskParams + + +class _FakeGroup: + def size(self) -> int: + return 3 + + +def test_warmup_inference_uses_realistic_gemma_style_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def fake_apply_chat_template(*, tokenizer: object, task_params: object, model_card: object): + del tokenizer, model_card + captured["task_params"] = task_params + return "warmup prompt" + + def fake_mx_barrier(_group: object) -> None: + return None + + def fake_mlx_generate(**_kwargs: object): + if False: + yield None + return + + def fake_all_gather(array: object, *, group: object): + del group + return array + + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) + monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) + monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) + + check_every = generate_mod.warmup_inference( + model=object(), # type: ignore[arg-type] + tokenizer=object(), # type: ignore[arg-type] + group=cast(object, _FakeGroup()), # type: ignore[arg-type] + model_id=ModelId("mlx-community/gemma-4-26b-a4b-it-4bit"), + model_card=None, + ) + + task_params = cast(TextGenerationTaskParams, captured["task_params"]) + assert check_every == 0 + assert task_params.instructions == ( + "You are a helpful assistant. Answer the user in one short sentence." + ) + assert task_params.enable_thinking is False + assert task_params.temperature == 1.0 + assert task_params.top_p == 0.95 + assert task_params.top_k == 64 + assert task_params.max_output_tokens == 32 + first_message = task_params.input[0] + assert "Summarize this status update in one sentence" in first_message.content 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 b9cff269e..d966b8861 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 @@ -349,7 +349,7 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch): def test_warmup_task_can_be_bypassed_while_runner_still_becomes_ready( patch_out_mlx: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch ): - monkeypatch.delenv("SKULK_FORCE_LLM_WARMUP", raising=False) + monkeypatch.setenv("SKULK_SKIP_LLM_WARMUP", "1") def fail_if_called(*_args: object, **_kwargs: object) -> None: raise AssertionError("warmup() should be bypassed in this debug branch") From 4c6000f09e8ee0ac57d65c4d0e24f4889c86e0e5 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 01:16:32 -0500 Subject: [PATCH 07/32] Test Gemma 4 prompts without empty thought suffix --- src/exo/worker/engines/mlx/gemma4_prompt.py | 12 +++---- .../unittests/test_mlx/test_gemma4_prompt.py | 36 +++++++++++++++++++ 2 files changed, 41 insertions(+), 7 deletions(-) create mode 100644 src/exo/worker/tests/unittests/test_mlx/test_gemma4_prompt.py diff --git a/src/exo/worker/engines/mlx/gemma4_prompt.py b/src/exo/worker/engines/mlx/gemma4_prompt.py index 82640a3dc..4f8ee7a52 100644 --- a/src/exo/worker/engines/mlx/gemma4_prompt.py +++ b/src/exo/worker/engines/mlx/gemma4_prompt.py @@ -8,8 +8,6 @@ from typing import Any -_GEMMA4_EMPTY_THOUGHT_CHANNEL = "<|channel>thought\n" - def strip_gemma4_thinking(text: str) -> str: """Remove Gemma 4 thinking blocks from assistant history.""" @@ -91,11 +89,11 @@ def render_gemma4_prompt( prompt_parts.append("\n") if add_generation_prompt: + # Gemma's published chat template appends an empty thought channel even + # when thinking is disabled. We intentionally omit that suffix here + # because it appears to trigger a pipeline warmup wedge in our MLX + # distributed path, and real traffic succeeds without this synthetic + # prefix. prompt_parts.append("<|turn>model\n") - if not enable_thinking: - # Match the shipped Gemma 4 chat template exactly. The model's own - # tokenizer appends an empty thought channel before assistant - # generation, and our output parser now understands that format. - prompt_parts.append(_GEMMA4_EMPTY_THOUGHT_CHANNEL) return "".join(prompt_parts) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_gemma4_prompt.py b/src/exo/worker/tests/unittests/test_mlx/test_gemma4_prompt.py new file mode 100644 index 000000000..6bb0331c5 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_gemma4_prompt.py @@ -0,0 +1,36 @@ +from exo.worker.engines.mlx.gemma4_prompt import render_gemma4_prompt + + +def test_render_gemma4_prompt_does_not_append_empty_thought_channel_when_thinking_disabled(): + prompt = render_gemma4_prompt( + [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ], + add_generation_prompt=True, + enable_thinking=False, + ) + + assert prompt == ( + "" + "<|turn>system\n" + "You are helpful." + "\n" + "<|turn>user\n" + "Hello" + "\n" + "<|turn>model\n" + ) + assert "<|channel>thought" not in prompt + assert "" not in prompt + + +def test_render_gemma4_prompt_keeps_think_marker_when_thinking_enabled(): + prompt = render_gemma4_prompt( + [{"role": "user", "content": "Hello"}], + add_generation_prompt=True, + enable_thinking=True, + ) + + assert prompt.startswith("<|turn>system\n<|think|>\n") + assert prompt.endswith("<|turn>model\n") From 2e4fdf8580118fbb2489fdda5e5141aea2d51ec5 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 01:51:02 -0500 Subject: [PATCH 08/32] Trace warmup and first live request shapes --- .../worker/engines/mlx/generator/generate.py | 10 +++++ src/exo/worker/engines/mlx/utils_mlx.py | 43 +++++++++++++++++++ .../runner/llm_inference/batch_generator.py | 29 +++++++++++++ .../unittests/test_mlx/test_warmup_request.py | 20 +++++++++ .../test_runner/test_event_ordering.py | 33 ++++++++++++++ 5 files changed, 135 insertions(+) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 3bc2e457f..c771230c9 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -62,6 +62,7 @@ from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, fix_unmatched_think_end_tokens, + log_request_shape, mx_barrier, system_prompt_token_count, ) @@ -606,6 +607,15 @@ def warmup_inference( "Warmup prompt prepared " f"(model={model_id}, prompt_chars={len(warmup_prompt)}, group_size={group.size() if group is not None else 1})" ) + log_request_shape( + "warmup", + warmup_task_params, + warmup_prompt, + extra={ + "group_size": group.size() if group is not None else 1, + "model_id": str(model_id), + }, + ) tokens_generated = 0 diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 2bb663edf..4ec293858 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -78,6 +78,49 @@ Group = mx.distributed.Group + +def _request_shape_debug_enabled() -> bool: + """Return whether request-shape tracing is enabled for prompt debugging.""" + value = os.environ.get("SKULK_TRACE_REQUEST_SHAPES") or os.environ.get( + "EXO_TRACE_REQUEST_SHAPES" + ) + if value is None: + return False + return value.strip().lower() not in {"0", "false", "no", "off"} + + +def log_request_shape( + label: str, + task_params: TextGenerationTaskParams, + prompt: str, + *, + extra: dict[str, Any] | None = None, +) -> None: + """Emit an exact request-shape trace for prompt/rendering comparisons. + + This is intentionally opt-in because it logs full rendered prompts and the + complete internal request payload. We use it to compare synthetic warmup + requests against real traffic when a model wedges only during startup. + """ + if not _request_shape_debug_enabled(): + return + + payload: dict[str, Any] = { + "label": label, + "task_params": task_params.model_dump(mode="json"), + "prompt_chars": len(prompt), + "prompt_lines": prompt.count("\n") + 1, + } + if extra: + payload.update(extra) + + logger.info( + "[request-shape] " + f"{json.dumps(payload, ensure_ascii=True, sort_keys=True)}" + ) + logger.info(f"[request-shape] prompt label={label}\n{prompt}") + + def _gemma4_output_length_for_pixel_values( pixel_values: mx.array, patch_size: int, diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py index 9379c672f..d0dc44a11 100644 --- a/src/exo/worker/runner/llm_inference/batch_generator.py +++ b/src/exo/worker/runner/llm_inference/batch_generator.py @@ -27,6 +27,7 @@ ) from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, + log_request_shape, mx_all_gather_tasks, mx_any, ) @@ -143,6 +144,7 @@ class SequentialGenerator(InferenceGenerator): ] | None ) = field(default=None, init=False) + _logged_first_request_shape: bool = field(default=False, init=False) def warmup(self): self.check_for_cancel_every = warmup_inference( @@ -270,6 +272,19 @@ def _build_generator(self, task: TextGeneration) -> Generator[GenerationResponse prompt = apply_chat_template( self.tokenizer, task.task_params, model_card=self.model_card ) + if not self._logged_first_request_shape: + log_request_shape( + "first-live-request", + task.task_params, + prompt, + extra={ + "command_id": str(task.command_id), + "device_rank": self.device_rank, + "generator": "sequential", + "task_id": str(task.task_id), + }, + ) + self._logged_first_request_shape = True def on_prefill_progress(processed: int, total: int) -> None: if self.device_rank == 0: @@ -350,6 +365,7 @@ class BatchGenerator(InferenceGenerator): Generator[GenerationResponse | ToolCallResponse | None], ], ] = field(default_factory=dict, init=False) + _logged_first_request_shape: bool = field(default=False, init=False) def __post_init__(self) -> None: self._mlx_gen = ExoBatchGenerator( @@ -510,6 +526,19 @@ def _start_task(self, task: TextGeneration) -> int: prompt = apply_chat_template( self.tokenizer, task.task_params, model_card=self.model_card ) + if not self._logged_first_request_shape: + log_request_shape( + "first-live-request", + task.task_params, + prompt, + extra={ + "command_id": str(task.command_id), + "device_rank": self.device_rank, + "generator": "batch", + "task_id": str(task.task_id), + }, + ) + self._logged_first_request_shape = True def on_prefill_progress(processed: int, total: int) -> None: if self.device_rank == 0: diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 7124be74b..81c94261e 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -34,9 +34,22 @@ def fake_all_gather(array: object, *, group: object): del group return array + def fake_log_request_shape( + label: str, + task_params: object, + prompt: str, + *, + extra: object = None, + ) -> None: + captured["label"] = label + captured["logged_task_params"] = task_params + captured["logged_prompt"] = prompt + captured["logged_extra"] = extra + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod, "log_request_shape", fake_log_request_shape) monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) check_every = generate_mod.warmup_inference( @@ -59,3 +72,10 @@ def fake_all_gather(array: object, *, group: object): assert task_params.max_output_tokens == 32 first_message = task_params.input[0] assert "Summarize this status update in one sentence" in first_message.content + assert captured["label"] == "warmup" + assert captured["logged_task_params"] == task_params + assert captured["logged_prompt"] == "warmup prompt" + assert captured["logged_extra"] == { + "group_size": 3, + "model_id": "mlx-community/gemma-4-26b-a4b-it-4bit", + } 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 d966b8861..02fec3d72 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 @@ -371,3 +371,36 @@ def fail_if_called(*_args: object, **_kwargs: object) -> None: and isinstance(event.runner_status, RunnerReady) for event in events ) + + +def test_first_live_request_shape_is_logged_once( + patch_out_mlx: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch +): + logged: list[tuple[str, TextGenerationTaskParams, str, object]] = [] + + def fake_log_request_shape( + label: str, + task_params: TextGenerationTaskParams, + prompt: str, + *, + extra: object = None, + ) -> None: + logged.append((label, task_params, prompt, extra)) + + monkeypatch.setattr(mlx_batch_generator, "log_request_shape", fake_log_request_shape) + + _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK], send_after_ready=[SHUTDOWN_TASK]) + + assert logged == [ + ( + "first-live-request", + CHAT_PARAMS, + "test prompt", + { + "command_id": str(COMMAND_1_ID), + "device_rank": 0, + "generator": "batch", + "task_id": str(CHAT_COMPLETION_TASK_ID), + }, + ) + ] From 66392dfffbb2b26c6759b56f55cebd1f582049c3 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 02:24:23 -0500 Subject: [PATCH 09/32] Use hello warmup prompt for MLX isolation --- .../worker/engines/mlx/generator/generate.py | 19 ++++--------------- .../unittests/test_mlx/test_warmup_request.py | 16 +++++++--------- 2 files changed, 11 insertions(+), 24 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index c771230c9..f14009617 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -577,23 +577,12 @@ def warmup_inference( ) -> int: logger.info(f"warming up inference for instance: {model_id}") - instructions = ( - "You are a helpful assistant. Answer the user in one short sentence." - ) - content = ( - "Summarize this status update in one sentence: " - "Skulk is starting a distributed Gemma 4 model runner, verifying " - "pipeline communication between nodes, and preparing to serve chat requests." - ) - warmup_task_params = TextGenerationTaskParams( model=model_id, - instructions=instructions, - input=[InputMessage(role="user", content=content)], - max_output_tokens=32, - temperature=1.0, - top_p=0.95, - top_k=64, + # Keep warmup as close as possible to the smallest known-good live + # request while still exercising the full prompt-render/generation path. + input=[InputMessage(role="user", content="hello")], + max_output_tokens=8, enable_thinking=False, ) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 81c94261e..77407b93a 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_realistic_gemma_style_request( +def test_warmup_inference_uses_minimal_hello_request( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -62,16 +62,14 @@ def fake_log_request_shape( task_params = cast(TextGenerationTaskParams, captured["task_params"]) assert check_every == 0 - assert task_params.instructions == ( - "You are a helpful assistant. Answer the user in one short sentence." - ) + assert task_params.instructions is None assert task_params.enable_thinking is False - assert task_params.temperature == 1.0 - assert task_params.top_p == 0.95 - assert task_params.top_k == 64 - assert task_params.max_output_tokens == 32 + assert task_params.temperature is None + assert task_params.top_p is None + assert task_params.top_k is None + assert task_params.max_output_tokens == 8 first_message = task_params.input[0] - assert "Summarize this status update in one sentence" in first_message.content + assert first_message.content == "hello" assert captured["label"] == "warmup" assert captured["logged_task_params"] == task_params assert captured["logged_prompt"] == "warmup prompt" From e252ccc50e8fd68bf20ee4a88455cab83d481b7f Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 08:59:58 -0500 Subject: [PATCH 10/32] Restore sampler settings for warmup bisect --- src/exo/worker/engines/mlx/generator/generate.py | 7 +++++-- .../tests/unittests/test_mlx/test_warmup_request.py | 8 ++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index f14009617..7a1f3e2d5 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -579,11 +579,14 @@ def warmup_inference( warmup_task_params = TextGenerationTaskParams( model=model_id, - # Keep warmup as close as possible to the smallest known-good live - # request while still exercising the full prompt-render/generation path. + # Reintroduce warmup dimensions one at a time so we can isolate which + # request characteristic destabilizes the MLX/GPU path. input=[InputMessage(role="user", content="hello")], max_output_tokens=8, enable_thinking=False, + temperature=1.0, + top_p=0.95, + top_k=64, ) with _hang_debug_watch(f"warmup apply_chat_template model={model_id}"): diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 77407b93a..b4ec47f9e 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_minimal_hello_request( +def test_warmup_inference_uses_hello_request_with_explicit_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -64,9 +64,9 @@ def fake_log_request_shape( assert check_every == 0 assert task_params.instructions is None assert task_params.enable_thinking is False - assert task_params.temperature is None - assert task_params.top_p is None - assert task_params.top_k is None + assert task_params.temperature == 1.0 + assert task_params.top_p == 0.95 + assert task_params.top_k == 64 assert task_params.max_output_tokens == 8 first_message = task_params.input[0] assert first_message.content == "hello" From 1b4d3ee415944706ab4262a775d592e00ed9df6f Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 09:10:23 -0500 Subject: [PATCH 11/32] Restore warmup instructions for bisect --- src/exo/worker/engines/mlx/generator/generate.py | 1 + .../worker/tests/unittests/test_mlx/test_warmup_request.py | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 7a1f3e2d5..f5b34d147 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -581,6 +581,7 @@ def warmup_inference( model=model_id, # Reintroduce warmup dimensions one at a time so we can isolate which # request characteristic destabilizes the MLX/GPU path. + instructions="You are a helpful assistant. Answer the user in one short sentence.", input=[InputMessage(role="user", content="hello")], max_output_tokens=8, enable_thinking=False, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index b4ec47f9e..cb5012a8b 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_hello_request_with_explicit_sampler_settings( +def test_warmup_inference_uses_hello_request_with_instructions_and_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -62,7 +62,10 @@ def fake_log_request_shape( task_params = cast(TextGenerationTaskParams, captured["task_params"]) assert check_every == 0 - assert task_params.instructions is None + assert ( + task_params.instructions + == "You are a helpful assistant. Answer the user in one short sentence." + ) assert task_params.enable_thinking is False assert task_params.temperature == 1.0 assert task_params.top_p == 0.95 From f55a421b203babd64fa0210ddc70b7ac80ea1323 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 09:34:01 -0500 Subject: [PATCH 12/32] Revert "Restore warmup instructions for bisect" This reverts commit 1b4d3ee415944706ab4262a775d592e00ed9df6f. --- src/exo/worker/engines/mlx/generator/generate.py | 1 - .../worker/tests/unittests/test_mlx/test_warmup_request.py | 7 ++----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index f5b34d147..7a1f3e2d5 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -581,7 +581,6 @@ def warmup_inference( model=model_id, # Reintroduce warmup dimensions one at a time so we can isolate which # request characteristic destabilizes the MLX/GPU path. - instructions="You are a helpful assistant. Answer the user in one short sentence.", input=[InputMessage(role="user", content="hello")], max_output_tokens=8, enable_thinking=False, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index cb5012a8b..b4ec47f9e 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_hello_request_with_instructions_and_sampler_settings( +def test_warmup_inference_uses_hello_request_with_explicit_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -62,10 +62,7 @@ def fake_log_request_shape( task_params = cast(TextGenerationTaskParams, captured["task_params"]) assert check_every == 0 - assert ( - task_params.instructions - == "You are a helpful assistant. Answer the user in one short sentence." - ) + assert task_params.instructions is None assert task_params.enable_thinking is False assert task_params.temperature == 1.0 assert task_params.top_p == 0.95 From b206f87a81cdd04314dd5852a3e672e4df24a45e Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 09:55:00 -0500 Subject: [PATCH 13/32] Disable pipeline prefill mode for short warmup path --- .../worker/engines/mlx/generator/generate.py | 20 +++++++++++-------- .../test_mlx/test_prefill_path_selection.py | 14 +++++++++++++ 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 7a1f3e2d5..e4a07a9dd 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -453,14 +453,6 @@ def combined_progress_callback(processed: int, total: int) -> None: distributed_prompt_progress_callback() progress_callback(processed, total) - set_pipeline_prefill(model, is_prefill=True) - - with _hang_debug_watch(f"prefill barrier rank={rank} group_size={group_size}"): - mx_barrier(group) - logger.info( - f"Starting prefill (rank={rank}, group_size={group_size}, prompt_tokens={num_tokens})" - ) - is_pipeline = _has_pipeline_communication_layer(model) prefill_step_size = 4096 @@ -495,6 +487,18 @@ def combined_progress_callback(processed: int, total: int) -> None: f"prefill_step_size_effective={effective_prefill_step_size}, " f"pipeline_chunks={pipeline_chunks})" ) + # Only enable pipeline-prefill mode for the true pipeline prefill path. + # Short prompts routed through stream_generate must behave like normal + # generation; leaving pipeline wrappers in prefill mode there can strand + # non-zero ranks in the pipeline send/recv path without the coordinated + # queue/flush behavior used by pipeline_parallel_prefill. + set_pipeline_prefill(model, is_prefill=use_pipeline_prefill) + + with _hang_debug_watch(f"prefill barrier rank={rank} group_size={group_size}"): + mx_barrier(group) + logger.info( + f"Starting prefill (rank={rank}, group_size={group_size}, prompt_tokens={num_tokens})" + ) try: if use_pipeline_prefill: diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py b/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py index 0b8c00067..e0ddb8b53 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py @@ -45,8 +45,14 @@ def test_prefill_uses_pipeline_parallel_path_for_long_pipeline_prompts( """ calls: list[str] = [] fake_cache = _FakeCache() + prefill_mode_calls: list[bool] = [] monkeypatch.setattr(generate_module, "mx_barrier", lambda _group: None) + monkeypatch.setattr( + generate_module, + "set_pipeline_prefill", + lambda _model, *, is_prefill: prefill_mode_calls.append(is_prefill), + ) monkeypatch.setattr( generate_module, "_has_pipeline_communication_layer", @@ -88,6 +94,7 @@ def _pipeline_parallel_prefill(**kwargs: object) -> None: assert snapshots == [] assert prefill_tps >= 0.0 assert fake_cache.trim_calls == [2] + assert prefill_mode_calls == [True, False] def test_prefill_uses_stream_generate_for_short_pipeline_prompts( @@ -104,8 +111,14 @@ def test_prefill_uses_stream_generate_for_short_pipeline_prompts( """ calls: list[str] = [] fake_cache = _FakeCache() + prefill_mode_calls: list[bool] = [] monkeypatch.setattr(generate_module, "mx_barrier", lambda _group: None) + monkeypatch.setattr( + generate_module, + "set_pipeline_prefill", + lambda _model, *, is_prefill: prefill_mode_calls.append(is_prefill), + ) monkeypatch.setattr( generate_module, "_has_pipeline_communication_layer", @@ -144,6 +157,7 @@ def _stream_generate(*args: object, **kwargs: object) -> Iterator[object]: assert snapshots == [] assert prefill_tps >= 0.0 assert fake_cache.trim_calls == [2] + assert prefill_mode_calls == [False, False] def test_prefill_uses_stream_generate_when_model_is_not_pipeline( From a7be93fbead7567df265645b37c8d5c8ffae8b70 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 10:10:46 -0500 Subject: [PATCH 14/32] Retry warmup instructions after prefill fix --- src/exo/worker/engines/mlx/generator/generate.py | 1 + .../worker/tests/unittests/test_mlx/test_warmup_request.py | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index e4a07a9dd..ba95f62cd 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -585,6 +585,7 @@ def warmup_inference( model=model_id, # Reintroduce warmup dimensions one at a time so we can isolate which # request characteristic destabilizes the MLX/GPU path. + instructions="You are a helpful assistant. Answer the user in one short sentence.", input=[InputMessage(role="user", content="hello")], max_output_tokens=8, enable_thinking=False, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index b4ec47f9e..cb5012a8b 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_hello_request_with_explicit_sampler_settings( +def test_warmup_inference_uses_hello_request_with_instructions_and_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -62,7 +62,10 @@ def fake_log_request_shape( task_params = cast(TextGenerationTaskParams, captured["task_params"]) assert check_every == 0 - assert task_params.instructions is None + assert ( + task_params.instructions + == "You are a helpful assistant. Answer the user in one short sentence." + ) assert task_params.enable_thinking is False assert task_params.temperature == 1.0 assert task_params.top_p == 0.95 From dab0b00f1fe779005507a3fb670787a18072ccd6 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 10:27:01 -0500 Subject: [PATCH 15/32] Retry longer warmup content after prefill fix --- src/exo/worker/engines/mlx/generator/generate.py | 10 +++++++++- .../tests/unittests/test_mlx/test_warmup_request.py | 7 +++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index ba95f62cd..3f6cc3b2b 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -586,7 +586,15 @@ def warmup_inference( # Reintroduce warmup dimensions one at a time so we can isolate which # request characteristic destabilizes the MLX/GPU path. instructions="You are a helpful assistant. Answer the user in one short sentence.", - input=[InputMessage(role="user", content="hello")], + input=[ + InputMessage( + role="user", + content=( + "Summarize this status update in one sentence: " + "the pipeline prefill change improved startup stability." + ), + ) + ], max_output_tokens=8, enable_thinking=False, temperature=1.0, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index cb5012a8b..b083652f9 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_hello_request_with_instructions_and_sampler_settings( +def test_warmup_inference_uses_longer_user_content_with_instructions_and_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -72,7 +72,10 @@ def fake_log_request_shape( assert task_params.top_k == 64 assert task_params.max_output_tokens == 8 first_message = task_params.input[0] - assert first_message.content == "hello" + assert ( + first_message.content + == "Summarize this status update in one sentence: the pipeline prefill change improved startup stability." + ) assert captured["label"] == "warmup" assert captured["logged_task_params"] == task_params assert captured["logged_prompt"] == "warmup prompt" From 12e3cdd4d0f503efe9da54c03849a5e6ad982cf3 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 10:41:09 -0500 Subject: [PATCH 16/32] Revert "Retry longer warmup content after prefill fix" This reverts commit dab0b00f1fe779005507a3fb670787a18072ccd6. --- src/exo/worker/engines/mlx/generator/generate.py | 10 +--------- .../tests/unittests/test_mlx/test_warmup_request.py | 7 ++----- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 3f6cc3b2b..ba95f62cd 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -586,15 +586,7 @@ def warmup_inference( # Reintroduce warmup dimensions one at a time so we can isolate which # request characteristic destabilizes the MLX/GPU path. instructions="You are a helpful assistant. Answer the user in one short sentence.", - input=[ - InputMessage( - role="user", - content=( - "Summarize this status update in one sentence: " - "the pipeline prefill change improved startup stability." - ), - ) - ], + input=[InputMessage(role="user", content="hello")], max_output_tokens=8, enable_thinking=False, temperature=1.0, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index b083652f9..cb5012a8b 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_longer_user_content_with_instructions_and_sampler_settings( +def test_warmup_inference_uses_hello_request_with_instructions_and_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -72,10 +72,7 @@ def fake_log_request_shape( assert task_params.top_k == 64 assert task_params.max_output_tokens == 8 first_message = task_params.input[0] - assert ( - first_message.content - == "Summarize this status update in one sentence: the pipeline prefill change improved startup stability." - ) + assert first_message.content == "hello" assert captured["label"] == "warmup" assert captured["logged_task_params"] == task_params assert captured["logged_prompt"] == "warmup prompt" From 16d5a1cb0a73e5b829c7b8c7968490e3c40cf6f7 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 10:42:05 -0500 Subject: [PATCH 17/32] Try neutral padded warmup content --- src/exo/worker/engines/mlx/generator/generate.py | 10 +++++++++- .../tests/unittests/test_mlx/test_warmup_request.py | 7 +++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index ba95f62cd..f96537a13 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -586,7 +586,15 @@ def warmup_inference( # Reintroduce warmup dimensions one at a time so we can isolate which # request characteristic destabilizes the MLX/GPU path. instructions="You are a helpful assistant. Answer the user in one short sentence.", - input=[InputMessage(role="user", content="hello")], + input=[ + InputMessage( + role="user", + content=( + "hello hello hello hello hello hello hello hello " + "hello hello hello hello hello hello hello hello" + ), + ) + ], max_output_tokens=8, enable_thinking=False, temperature=1.0, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index cb5012a8b..dc9211a9e 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_hello_request_with_instructions_and_sampler_settings( +def test_warmup_inference_uses_neutral_padded_user_content_with_instructions_and_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -72,7 +72,10 @@ def fake_log_request_shape( assert task_params.top_k == 64 assert task_params.max_output_tokens == 8 first_message = task_params.input[0] - assert first_message.content == "hello" + assert ( + first_message.content + == "hello hello hello hello hello hello hello hello hello hello hello hello hello hello hello hello" + ) assert captured["label"] == "warmup" assert captured["logged_task_params"] == task_params assert captured["logged_prompt"] == "warmup prompt" From 1fd68247bd16bda605321c74be353f06abae3704 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 12:47:12 -0500 Subject: [PATCH 18/32] Clean up stale runners on node timeout --- src/exo/shared/apply.py | 28 ++++ .../test_apply/test_apply_node_timed_out.py | 128 ++++++++++++++++++ 2 files changed, 156 insertions(+) create mode 100644 src/exo/shared/tests/test_apply/test_apply_node_timed_out.py diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 8c9dafd8b..b525477a3 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -213,6 +213,31 @@ def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> Sta def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: topology = copy.deepcopy(state.topology) topology.remove_node(event.node_id) + affected_instance_ids = { + instance_id + for instance_id, instance in state.instances.items() + if event.node_id in instance.shard_assignments.node_to_runner + } + affected_runner_ids = { + runner_id + for instance_id in affected_instance_ids + for runner_id in state.instances[instance_id].shard_assignments.runner_to_shard + } + instances = { + instance_id: instance + for instance_id, instance in state.instances.items() + if instance_id not in affected_instance_ids + } + runners = { + runner_id: runner_status + for runner_id, runner_status in state.runners.items() + if runner_id not in affected_runner_ids + } + tasks = { + task_id: task + for task_id, task in state.tasks.items() + if task.instance_id not in affected_instance_ids + } last_seen = { key: value for key, value in state.last_seen.items() if key != event.node_id } @@ -257,6 +282,9 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: ) return state.model_copy( update={ + "instances": instances, + "runners": runners, + "tasks": tasks, "downloads": downloads, "topology": topology, "last_seen": last_seen, diff --git a/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py b/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py new file mode 100644 index 000000000..a7d11c500 --- /dev/null +++ b/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py @@ -0,0 +1,128 @@ +from datetime import datetime + +from exo.shared.apply import apply_node_timed_out +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.types.common import NodeId +from exo.shared.types.events import NodeTimedOut +from exo.shared.types.memory import Memory +from exo.shared.types.state import State +from exo.shared.types.tasks import StartWarmup, TaskId, TaskStatus +from exo.shared.types.worker.instances import InstanceId, MlxRingInstance +from exo.shared.types.worker.runners import ( + RunnerId, + RunnerIdle, + RunnerReady, + RunnerWarmingUp, + ShardAssignments, +) +from exo.shared.types.worker.shards import PipelineShardMetadata + + +def _make_pipeline_shard(model_id: ModelId, device_rank: int, world_size: int) -> PipelineShardMetadata: + return PipelineShardMetadata( + model_card=ModelCard( + model_id=model_id, + storage_size=Memory.from_mb(100000), + n_layers=32, + hidden_size=2048, + supports_tensor=False, + tasks=[ModelTask.TextGeneration], + ), + device_rank=device_rank, + world_size=world_size, + start_layer=0, + end_layer=32, + n_layers=32, + ) + + +def test_apply_node_timed_out_removes_affected_instances_runners_and_tasks() -> None: + node_a = NodeId("node-a") + node_b = NodeId("node-b") + node_c = NodeId("node-c") + + affected_instance_id = InstanceId("affected-instance") + unaffected_instance_id = InstanceId("unaffected-instance") + + affected_runner_a = RunnerId("affected-runner-a") + affected_runner_b = RunnerId("affected-runner-b") + unaffected_runner = RunnerId("unaffected-runner") + + model_id = ModelId("mlx-community/gemma-4-26b-a4b-it-4bit") + + affected_instance = MlxRingInstance( + instance_id=affected_instance_id, + shard_assignments=ShardAssignments( + model_id=model_id, + node_to_runner={ + node_a: affected_runner_a, + node_b: affected_runner_b, + }, + runner_to_shard={ + affected_runner_a: _make_pipeline_shard(model_id, device_rank=0, world_size=2), + affected_runner_b: _make_pipeline_shard(model_id, device_rank=1, world_size=2), + }, + ), + hosts_by_node={}, + ephemeral_port=50000, + ) + + unaffected_instance = MlxRingInstance( + instance_id=unaffected_instance_id, + shard_assignments=ShardAssignments( + model_id=model_id, + node_to_runner={node_c: unaffected_runner}, + runner_to_shard={ + unaffected_runner: _make_pipeline_shard(model_id, device_rank=0, world_size=1), + }, + ), + hosts_by_node={}, + ephemeral_port=50001, + ) + + affected_task_id = TaskId("affected-task") + unaffected_task_id = TaskId("unaffected-task") + state = State( + instances={ + affected_instance_id: affected_instance, + unaffected_instance_id: unaffected_instance, + }, + runners={ + affected_runner_a: RunnerWarmingUp(), + affected_runner_b: RunnerReady(), + unaffected_runner: RunnerIdle(), + }, + tasks={ + affected_task_id: StartWarmup( + task_id=affected_task_id, + instance_id=affected_instance_id, + task_status=TaskStatus.Pending, + ), + unaffected_task_id: StartWarmup( + task_id=unaffected_task_id, + instance_id=unaffected_instance_id, + task_status=TaskStatus.Pending, + ), + }, + last_seen={ + node_a: datetime.now(), + node_b: datetime.now(), + node_c: datetime.now(), + }, + ) + + new_state = apply_node_timed_out(NodeTimedOut(node_id=node_a), state) + + assert affected_instance_id not in new_state.instances + assert unaffected_instance_id in new_state.instances + + assert affected_runner_a not in new_state.runners + assert affected_runner_b not in new_state.runners + assert unaffected_runner in new_state.runners + + assert affected_task_id not in new_state.tasks + assert unaffected_task_id in new_state.tasks + + assert node_a not in new_state.last_seen + assert node_b in new_state.last_seen + assert node_c in new_state.last_seen From bc0c7ae85328ccc8571b1ae6ac36c4b72e36cff4 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 13:15:43 -0500 Subject: [PATCH 19/32] Raise warmup output budget for bisect --- src/exo/worker/engines/mlx/generator/generate.py | 2 +- src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index f96537a13..130e24e5f 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -595,7 +595,7 @@ def warmup_inference( ), ) ], - max_output_tokens=8, + max_output_tokens=1024, enable_thinking=False, temperature=1.0, top_p=0.95, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index dc9211a9e..5a90eb369 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -70,7 +70,7 @@ def fake_log_request_shape( assert task_params.temperature == 1.0 assert task_params.top_p == 0.95 assert task_params.top_k == 64 - assert task_params.max_output_tokens == 8 + assert task_params.max_output_tokens == 1024 first_message = task_params.input[0] assert ( first_message.content From 3d54e878368753699c38b5dec40936bb4b2f3868 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 13:31:16 -0500 Subject: [PATCH 20/32] Add runner lifecycle debug logging --- src/exo/worker/main.py | 72 +++++++++++++++++-- src/exo/worker/runner/llm_inference/runner.py | 21 ++++-- 2 files changed, 84 insertions(+), 9 deletions(-) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 6358e25b6..5c8e10e8a 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -42,7 +42,9 @@ CreateRunner, DownloadModel, ImageEdits, + LoadModel, Shutdown, + StartWarmup, Task, TaskStatus, TextGeneration, @@ -62,6 +64,51 @@ from exo.worker.runner.runner_supervisor import RunnerSupervisor +def _summarize_worker_task(task: Task) -> str: + """Return a compact task summary for worker lifecycle logs.""" + if isinstance(task, CreateRunner): + shard = task.bound_instance.bound_shard + return ( + "CreateRunner(" + f"instance_id={task.instance_id!r}, " + f"runner_id={task.bound_instance.bound_runner_id!r}, " + f"node_id={task.bound_instance.bound_node_id!r}, " + f"device_rank={shard.device_rank}, " + f"world_size={shard.world_size}, " + f"layers={shard.start_layer}:{shard.end_layer})" + ) + if isinstance(task, DownloadModel): + shard = task.shard_metadata + return ( + "DownloadModel(" + f"instance_id={task.instance_id!r}, " + f"model={shard.model_card.model_id!r}, " + f"device_rank={shard.device_rank}, " + f"world_size={shard.world_size}, " + f"layers={shard.start_layer}:{shard.end_layer})" + ) + if isinstance(task, Shutdown): + return ( + "Shutdown(" + f"instance_id={task.instance_id!r}, " + f"runner_id={task.runner_id!r})" + ) + if isinstance(task, LoadModel): + return f"LoadModel(instance_id={task.instance_id!r})" + if isinstance(task, StartWarmup): + return f"StartWarmup(instance_id={task.instance_id!r})" + if isinstance(task, CancelTask): + return ( + "CancelTask(" + f"instance_id={task.instance_id!r}, " + f"runner_id={task.runner_id!r}, " + f"cancelled_task_id={task.cancelled_task_id!r})" + ) + if isinstance(task, (TextGeneration, ImageEdits)): + return repr(task) + return task.__class__.__name__ + + def _log_image_transport(message: str) -> None: """Emit image transport logs only at INFO when explicitly requested. @@ -245,7 +292,7 @@ async def plan_step(self): if not self._download_backoff.should_proceed(model_id): continue - logger.info(f"Worker plan: {task.__class__.__name__}") + logger.info(f"Worker plan: {_summarize_worker_task(task)}") assert task.task_status await self.event_sender.send(TaskCreated(task_id=task.task_id, task=task)) @@ -451,12 +498,29 @@ async def shutdown(self): async def _start_runner_task(self, task: Task): if (instance := self.state.instances.get(task.instance_id)) is not None: - await self.runners[ - instance.shard_assignments.node_to_runner[self.node_id] - ].start_task(task) + runner_id = instance.shard_assignments.node_to_runner[self.node_id] + shard = instance.shard(runner_id) + logger.info( + "Dispatching worker task " + f"({_summarize_worker_task(task)}, " + f"target_runner_id={runner_id}, " + f"device_rank={shard.device_rank if shard is not None else 'unknown'}, " + f"world_size={shard.world_size if shard is not None else 'unknown'})" + ) + await self.runners[runner_id].start_task(task) def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor: """Creates and stores a new AssignedRunner with initial downloading status.""" + shard = task.bound_instance.bound_shard + logger.info( + "Creating runner supervisor " + f"(instance_id={task.instance_id}, " + f"runner_id={task.bound_instance.bound_runner_id}, " + f"node_id={task.bound_instance.bound_node_id}, " + f"device_rank={shard.device_rank}, " + f"world_size={shard.world_size}, " + f"layers={shard.start_layer}:{shard.end_layer})" + ) runner = RunnerSupervisor.create( bound_instance=task.bound_instance, event_sender=self.event_sender.clone(), diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index bd65d7746..f17812150 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -146,6 +146,17 @@ def _summarize_text_generation_task(task: TextGeneration) -> str: f"enable_thinking={params.enable_thinking!r})" ) + def _lifecycle_context(self) -> str: + """Return stable runner identity fields for lifecycle logs.""" + return ( + f"instance_id={self.instance.instance_id}, " + f"runner_id={self.runner_id}, " + f"node_id={self.bound_instance.bound_node_id}, " + f"device_rank={self.shard_metadata.device_rank}, " + f"world_size={self.shard_metadata.world_size}, " + f"layers={self.shard_metadata.start_layer}:{self.shard_metadata.end_layer}" + ) + def update_status(self, status: RunnerStatus): self.current_status = status self.event_sender.send( @@ -181,7 +192,7 @@ def handle_first_task(self, task: Task): self.current_status, (RunnerIdle, RunnerFailed) ): assert isinstance(self.generator, Builder) - logger.info("runner connecting") + logger.info(f"runner connecting ({self._lifecycle_context()})") self.update_status(RunnerConnecting()) self.acknowledge_task(task) @@ -205,7 +216,7 @@ def handle_first_task(self, task: Task): total_layers = ( self.shard_metadata.end_layer - self.shard_metadata.start_layer ) - logger.info("runner loading") + logger.info(f"runner loading ({self._lifecycle_context()})") self.update_status( RunnerLoading(layers_loaded=0, total_layers=total_layers) @@ -241,11 +252,11 @@ def on_layer_loaded(layers_loaded: int, total: int) -> None: self.send_task_status(task.task_id, TaskStatus.Complete) self.update_status(RunnerLoaded()) - logger.info("runner loaded") + logger.info(f"runner loaded ({self._lifecycle_context()})") case StartWarmup() if isinstance(self.current_status, RunnerLoaded): assert isinstance(self.generator, InferenceGenerator) - logger.info("runner warming up") + logger.info(f"runner warming up ({self._lifecycle_context()})") self.update_status(RunnerWarmingUp()) self.acknowledge_task(task) @@ -265,7 +276,7 @@ def on_layer_loaded(layers_loaded: int, total: int) -> None: self.send_task_status(task.task_id, TaskStatus.Complete) self.update_status(RunnerReady()) - logger.info("runner ready") + logger.info(f"runner ready ({self._lifecycle_context()})") case TextGeneration() if isinstance(self.current_status, RunnerReady): return_code = self.handle_generation_tasks(starting_task=task) From a22effd367d54347c20360a0e0d8d4b11640ac9d Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 14:08:38 -0500 Subject: [PATCH 21/32] Make pipeline warmup prompt length configurable --- .../worker/engines/mlx/generator/generate.py | 32 ++++++++++--- .../unittests/test_mlx/test_warmup_request.py | 47 +++++++++++++++++-- 2 files changed, 68 insertions(+), 11 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 130e24e5f..e36c2bbff 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -102,6 +102,28 @@ def _mlx_hang_debug_interval_seconds() -> float: return 30.0 +def _warmup_repeat_count() -> int: + """Return the neutral warmup token repeat count used for debugging.""" + raw = os.environ.get("SKULK_DEBUG_WARMUP_REPEAT_COUNT") or os.environ.get( + "EXO_DEBUG_WARMUP_REPEAT_COUNT" + ) + if raw is None: + return 1 + with contextlib.suppress(ValueError): + return max(int(raw), 1) + return 1 + + +def _warmup_user_content() -> str: + """Return the synthetic warmup user content. + + A short prompt is the safest default for pipeline models. When we need to + debug prompt-length hangs, an environment override lets us scale the same + neutral content without another code push. + """ + return " ".join(["hello"] * _warmup_repeat_count()) + + @contextlib.contextmanager def _hang_debug_watch(label: str) -> Generator[None]: """Emit periodic stack-rich logs while the current thread is stuck in one phase.""" @@ -583,16 +605,14 @@ def warmup_inference( warmup_task_params = TextGenerationTaskParams( model=model_id, - # Reintroduce warmup dimensions one at a time so we can isolate which - # request characteristic destabilizes the MLX/GPU path. + # Keep the default pipeline warmup prompt tiny so clusters boot + # reliably; use SKULK_DEBUG_WARMUP_REPEAT_COUNT to stretch the same + # neutral content when bisecting prompt-length hangs. instructions="You are a helpful assistant. Answer the user in one short sentence.", input=[ InputMessage( role="user", - content=( - "hello hello hello hello hello hello hello hello " - "hello hello hello hello hello hello hello hello" - ), + content=_warmup_user_content(), ) ], max_output_tokens=1024, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 5a90eb369..c76b9cd64 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_neutral_padded_user_content_with_instructions_and_sampler_settings( +def test_warmup_inference_uses_safe_default_user_content_with_instructions_and_sampler_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -72,10 +72,7 @@ def fake_log_request_shape( assert task_params.top_k == 64 assert task_params.max_output_tokens == 1024 first_message = task_params.input[0] - assert ( - first_message.content - == "hello hello hello hello hello hello hello hello hello hello hello hello hello hello hello hello" - ) + assert first_message.content == "hello" assert captured["label"] == "warmup" assert captured["logged_task_params"] == task_params assert captured["logged_prompt"] == "warmup prompt" @@ -83,3 +80,43 @@ def fake_log_request_shape( "group_size": 3, "model_id": "mlx-community/gemma-4-26b-a4b-it-4bit", } + + +def test_warmup_inference_honors_repeat_count_override( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def fake_apply_chat_template(*, tokenizer: object, task_params: object, model_card: object): + del tokenizer, model_card + captured["task_params"] = task_params + return "warmup prompt" + + def fake_mx_barrier(_group: object) -> None: + return None + + def fake_mlx_generate(**_kwargs: object): + if False: + yield None + return + + def fake_all_gather(array: object, *, group: object): + del group + return array + + monkeypatch.setenv("SKULK_DEBUG_WARMUP_REPEAT_COUNT", "4") + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) + monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) + monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) + + generate_mod.warmup_inference( + model=object(), # type: ignore[arg-type] + tokenizer=object(), # type: ignore[arg-type] + group=cast(object, _FakeGroup()), # type: ignore[arg-type] + model_id=ModelId("mlx-community/gemma-4-26b-a4b-it-4bit"), + model_card=None, + ) + + task_params = cast(TextGenerationTaskParams, captured["task_params"]) + assert task_params.input[0].content == "hello hello hello hello" From 4520e6134ba8cad10fad05f2ba15b3a3fe915461 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 15:32:51 -0500 Subject: [PATCH 22/32] Make warmup instructions opt-in for debugging --- .../worker/engines/mlx/generator/generate.py | 20 ++++++-- .../unittests/test_mlx/test_warmup_request.py | 50 +++++++++++++++++-- 2 files changed, 62 insertions(+), 8 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index e36c2bbff..a71a959ae 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -124,6 +124,19 @@ def _warmup_user_content() -> str: return " ".join(["hello"] * _warmup_repeat_count()) +def _warmup_instructions() -> str | None: + """Return optional warmup instructions for prompt-shape debugging.""" + raw = os.environ.get("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS") or os.environ.get( + "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS" + ) + if raw is None: + return None + include = raw.strip().lower() not in {"0", "false", "no", "off"} + if not include: + return None + return "You are a helpful assistant. Answer the user in one short sentence." + + @contextlib.contextmanager def _hang_debug_watch(label: str) -> Generator[None]: """Emit periodic stack-rich logs while the current thread is stuck in one phase.""" @@ -606,9 +619,10 @@ def warmup_inference( warmup_task_params = TextGenerationTaskParams( model=model_id, # Keep the default pipeline warmup prompt tiny so clusters boot - # reliably; use SKULK_DEBUG_WARMUP_REPEAT_COUNT to stretch the same - # neutral content when bisecting prompt-length hangs. - instructions="You are a helpful assistant. Answer the user in one short sentence.", + # reliably; use SKULK_DEBUG_WARMUP_REPEAT_COUNT and + # SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS when bisecting prompt-shape + # hangs without another code push. + instructions=_warmup_instructions(), input=[ InputMessage( role="user", diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index c76b9cd64..9758b2527 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,7 +12,7 @@ def size(self) -> int: return 3 -def test_warmup_inference_uses_safe_default_user_content_with_instructions_and_sampler_settings( +def test_warmup_inference_uses_safe_default_user_content_without_instructions( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -62,10 +62,7 @@ def fake_log_request_shape( task_params = cast(TextGenerationTaskParams, captured["task_params"]) assert check_every == 0 - assert ( - task_params.instructions - == "You are a helpful assistant. Answer the user in one short sentence." - ) + assert task_params.instructions is None assert task_params.enable_thinking is False assert task_params.temperature == 1.0 assert task_params.top_p == 0.95 @@ -120,3 +117,46 @@ def fake_all_gather(array: object, *, group: object): task_params = cast(TextGenerationTaskParams, captured["task_params"]) assert task_params.input[0].content == "hello hello hello hello" + + +def test_warmup_inference_honors_instruction_override( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def fake_apply_chat_template(*, tokenizer: object, task_params: object, model_card: object): + del tokenizer, model_card + captured["task_params"] = task_params + return "warmup prompt" + + def fake_mx_barrier(_group: object) -> None: + return None + + def fake_mlx_generate(**_kwargs: object): + if False: + yield None + return + + def fake_all_gather(array: object, *, group: object): + del group + return array + + monkeypatch.setenv("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "1") + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) + monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) + monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) + + generate_mod.warmup_inference( + model=object(), # type: ignore[arg-type] + tokenizer=object(), # type: ignore[arg-type] + group=cast(object, _FakeGroup()), # type: ignore[arg-type] + model_id=ModelId("mlx-community/gemma-4-26b-a4b-it-4bit"), + model_card=None, + ) + + task_params = cast(TextGenerationTaskParams, captured["task_params"]) + assert ( + task_params.instructions + == "You are a helpful assistant. Answer the user in one short sentence." + ) From f1df1091d6a68049d1e02cd03f69ca772613b2a6 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 17:15:26 -0500 Subject: [PATCH 23/32] Force minimal warmup for pipeline models --- .../worker/engines/mlx/generator/generate.py | 45 ++++++++++++---- .../unittests/test_mlx/test_warmup_request.py | 53 +++++++++++++++++-- 2 files changed, 84 insertions(+), 14 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index a71a959ae..f57ad410d 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -114,18 +114,27 @@ def _warmup_repeat_count() -> int: return 1 -def _warmup_user_content() -> str: +def _is_distributed_pipeline_warmup(group: mx.distributed.Group | None) -> bool: + """Return whether warmup is running for a multi-node pipeline model.""" + return group is not None and group.size() > 1 + + +def _warmup_user_content(group: mx.distributed.Group | None) -> str: """Return the synthetic warmup user content. - A short prompt is the safest default for pipeline models. When we need to - debug prompt-length hangs, an environment override lets us scale the same - neutral content without another code push. + Distributed pipeline warmup intentionally stays minimal because richer + synthetic prompts have been observed to wedge stream_generate prefill. + Single-node debugging may still scale the neutral content via environment. """ + if _is_distributed_pipeline_warmup(group): + return "hello" return " ".join(["hello"] * _warmup_repeat_count()) -def _warmup_instructions() -> str | None: +def _warmup_instructions(group: mx.distributed.Group | None) -> str | None: """Return optional warmup instructions for prompt-shape debugging.""" + if _is_distributed_pipeline_warmup(group): + return None raw = os.environ.get("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS") or os.environ.get( "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS" ) @@ -614,19 +623,33 @@ def warmup_inference( model_id: ModelId, model_card: ModelCard | None = None, ) -> int: + """Run a conservative synthetic warmup request to validate the MLX path. + + Distributed pipeline warmup intentionally uses a minimal prompt because + richer synthetic prompts have been observed to wedge short-prompt + ``stream_generate`` prefill. Debug shaping overrides remain available for + single-node investigation, but pipeline warmup always stays on the minimal + sanity-check path. + """ logger.info(f"warming up inference for instance: {model_id}") + if _is_distributed_pipeline_warmup(group) and ( + _warmup_repeat_count() != 1 or _warmup_instructions(None) is not None + ): + logger.info( + "Distributed pipeline warmup is forcing the minimal sanity-check " + "prompt and ignoring warmup shaping overrides" + ) + warmup_task_params = TextGenerationTaskParams( model=model_id, - # Keep the default pipeline warmup prompt tiny so clusters boot - # reliably; use SKULK_DEBUG_WARMUP_REPEAT_COUNT and - # SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS when bisecting prompt-shape - # hangs without another code push. - instructions=_warmup_instructions(), + # Distributed pipeline warmup always uses the minimal sanity-check + # shape; single-node debugging can still opt into prompt shaping. + instructions=_warmup_instructions(group), input=[ InputMessage( role="user", - content=_warmup_user_content(), + content=_warmup_user_content(group), ) ], max_output_tokens=1024, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 9758b2527..1df4f429b 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -12,6 +12,11 @@ def size(self) -> int: return 3 +class _SingleNodeGroup: + def size(self) -> int: + return 1 + + def test_warmup_inference_uses_safe_default_user_content_without_instructions( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -79,7 +84,7 @@ def fake_log_request_shape( } -def test_warmup_inference_honors_repeat_count_override( +def test_warmup_inference_ignores_repeat_count_override_for_pipeline_groups( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -116,10 +121,10 @@ def fake_all_gather(array: object, *, group: object): ) task_params = cast(TextGenerationTaskParams, captured["task_params"]) - assert task_params.input[0].content == "hello hello hello hello" + assert task_params.input[0].content == "hello" -def test_warmup_inference_honors_instruction_override( +def test_warmup_inference_ignores_instruction_override_for_pipeline_groups( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -156,6 +161,48 @@ def fake_all_gather(array: object, *, group: object): ) task_params = cast(TextGenerationTaskParams, captured["task_params"]) + assert task_params.instructions is None + + +def test_warmup_inference_honors_repeat_and_instruction_overrides_for_single_node( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def fake_apply_chat_template(*, tokenizer: object, task_params: object, model_card: object): + del tokenizer, model_card + captured["task_params"] = task_params + return "warmup prompt" + + def fake_mx_barrier(_group: object) -> None: + return None + + def fake_mlx_generate(**_kwargs: object): + if False: + yield None + return + + def fake_all_gather(array: object, *, group: object): + del group + return array + + monkeypatch.setenv("SKULK_DEBUG_WARMUP_REPEAT_COUNT", "4") + monkeypatch.setenv("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "1") + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) + monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) + monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) + + generate_mod.warmup_inference( + model=object(), # type: ignore[arg-type] + tokenizer=object(), # type: ignore[arg-type] + group=cast(object, _SingleNodeGroup()), # type: ignore[arg-type] + model_id=ModelId("mlx-community/gemma-4-26b-a4b-it-4bit"), + model_card=None, + ) + + task_params = cast(TextGenerationTaskParams, captured["task_params"]) + assert task_params.input[0].content == "hello hello hello hello" assert ( task_params.instructions == "You are a helpful assistant. Answer the user in one short sentence." From 097685eec15a0a7cf8c4e14d390be95096398cb1 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 17:17:41 -0500 Subject: [PATCH 24/32] Document pipeline warmup policy and debug env vars --- website/docs/architecture.md | 37 ++++++++++++++++++++++++++ website/docs/model-behaviors/gemma4.md | 26 ++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/website/docs/architecture.md b/website/docs/architecture.md index 763201f07..89f3c5a1e 100644 --- a/website/docs/architecture.md +++ b/website/docs/architecture.md @@ -213,6 +213,43 @@ When enabled, the runner logs: `EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS` remain accepted as compatibility fallbacks while older scripts are updated. +## Pipeline Warmup Policy + +Distributed pipeline models now use an intentionally minimal warmup request. + +The goal of warmup is only to validate that MLX prefill and first-token +generation are functional without risking richer synthetic prompt shapes that +have been observed to wedge `stream_generate` prefill on multi-node pipeline +models. + +For distributed pipeline warmup, Skulk always uses: + +- no synthetic instructions +- a single user message with content `hello` +- `enable_thinking=False` +- the normal sampler defaults used by the warmup path + +The following debug-only environment variables are available: + +- `SKULK_DEBUG_WARMUP_REPEAT_COUNT` +- `SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS` + +These are only honored for single-node debugging. For distributed pipeline +warmup, Skulk intentionally ignores them and stays on the minimal sanity-check +prompt. + +Legacy compatibility aliases are also accepted: + +- `EXO_DEBUG_WARMUP_REPEAT_COUNT` +- `EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS` + +There is also an emergency bypass for debugging: + +- `SKULK_SKIP_LLM_WARMUP=1` + +That bypass marks the runner ready without issuing the synthetic warmup request +and should only be used for diagnosis, not normal operation. + ## When to Read More If you are: diff --git a/website/docs/model-behaviors/gemma4.md b/website/docs/model-behaviors/gemma4.md index 7f4b60235..f5e6742de 100644 --- a/website/docs/model-behaviors/gemma4.md +++ b/website/docs/model-behaviors/gemma4.md @@ -100,6 +100,32 @@ This causes the runner to log: The older `EXO_MLX_HANG_DEBUG` names still work as compatibility fallbacks, but new scripts and docs should prefer the `SKULK_` prefix. +For distributed Gemma 4 pipeline warmup specifically, Skulk now forces a +minimal synthetic prompt by design: + +- no synthetic instructions +- a single user message with content `hello` + +This is intentional. During debugging, richer synthetic warmup prompts were +observed to wedge short-prompt `stream_generate` prefill on multi-node pipeline +setups. + +Two debug-only warmup shaping env vars exist: + +- `SKULK_DEBUG_WARMUP_REPEAT_COUNT` +- `SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS` + +However, distributed pipeline warmup intentionally ignores them and stays on the +minimal sanity-check prompt. They remain useful for single-node investigation. + +If you need to bypass synthetic warmup entirely during diagnosis, you can use: + +```bash +export SKULK_SKIP_LLM_WARMUP=1 +``` + +That bypass should be treated as a temporary debugging escape hatch only. + ## Why This Matters Gemma 4 is the proof point for the model capability system. From 6ae70786e1e58260b1e19ce07d0a54e30526bdf2 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 17:51:36 -0500 Subject: [PATCH 25/32] Address high-severity PR review findings --- src/exo/shared/apply.py | 6 + .../test_apply/test_apply_node_timed_out.py | 10 ++ .../worker/engines/mlx/generator/generate.py | 38 +++---- src/exo/worker/engines/mlx/utils_mlx.py | 2 +- src/exo/worker/main.py | 34 +++++- .../test_mlx/test_prefill_path_selection.py | 105 ++++++++++++++++++ .../test_mlx/test_request_shape_debug.py | 15 +++ .../unittests/test_worker_task_summary.py | 47 ++++++++ website/docs/architecture.md | 7 +- website/docs/model-behaviors/gemma4.md | 5 +- 10 files changed, 238 insertions(+), 31 deletions(-) create mode 100644 src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py create mode 100644 src/exo/worker/tests/unittests/test_worker_task_summary.py diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index b525477a3..45448475e 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -241,6 +241,11 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: last_seen = { key: value for key, value in state.last_seen.items() if key != event.node_id } + node_identities = { + key: value + for key, value in state.node_identities.items() + if key != event.node_id + } downloads = { key: value for key, value in state.downloads.items() if key != event.node_id } @@ -288,6 +293,7 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: "downloads": downloads, "topology": topology, "last_seen": last_seen, + "node_identities": node_identities, "node_memory": node_memory, "node_disk": node_disk, "node_system": node_system, diff --git a/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py b/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py index a7d11c500..759824d1c 100644 --- a/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py +++ b/src/exo/shared/tests/test_apply/test_apply_node_timed_out.py @@ -5,6 +5,7 @@ from exo.shared.types.common import NodeId from exo.shared.types.events import NodeTimedOut from exo.shared.types.memory import Memory +from exo.shared.types.profiling import NodeIdentity from exo.shared.types.state import State from exo.shared.types.tasks import StartWarmup, TaskId, TaskStatus from exo.shared.types.worker.instances import InstanceId, MlxRingInstance @@ -109,6 +110,11 @@ def test_apply_node_timed_out_removes_affected_instances_runners_and_tasks() -> node_b: datetime.now(), node_c: datetime.now(), }, + node_identities={ + node_a: NodeIdentity(friendly_name="kite1"), + node_b: NodeIdentity(friendly_name="kite2"), + node_c: NodeIdentity(friendly_name="kite3"), + }, ) new_state = apply_node_timed_out(NodeTimedOut(node_id=node_a), state) @@ -126,3 +132,7 @@ def test_apply_node_timed_out_removes_affected_instances_runners_and_tasks() -> assert node_a not in new_state.last_seen assert node_b in new_state.last_seen assert node_c in new_state.last_seen + + assert node_a not in new_state.node_identities + assert node_b in new_state.node_identities + assert node_c in new_state.node_identities diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index f57ad410d..6f6dc3d33 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -504,21 +504,16 @@ def combined_progress_callback(processed: int, total: int) -> None: prefill_step_size // min(4, group_size) if is_pipeline else prefill_step_size ) - # Pipeline-parallel prefill currently wedges when the prompt fits inside a - # single per-rank chunk (i.e. ``n_real == 1`` in the prefill loop). This - # was observed during Gemma 4 warmup with a 23-token prompt and reproduced - # in the issue tracker for any short prompt under pipeline parallelism. - # PR #101's pipeline-prefill widening (cb2dc68e) removed the prior - # ``num_tokens >= prefill_step_size`` gate to fix a different Gemma warmup - # hang in stream_generate; the result was that *every* short prompt now - # routed through pipeline_parallel_prefill and hit the new wedge. The - # narrower gate below preserves PR #101's fix for non-trivial prompts - # (where pipeline_parallel_prefill is the right path) while routing - # short / warmup prompts back through stream_generate, which is known to - # work for that shape. See PR review thread on PR #103 for the full - # diagnosis. + # Pipeline-parallel prefill can wedge when the prompt fits inside a single + # effective per-rank chunk (i.e. ``n_real == 1`` in the prefill loop). + # Mirror pipeline_parallel_prefill's chunking math here by reserving the + # final token for the post-loop pass before computing the real chunk count. + # That keeps multi-chunk prompts on pipeline_parallel_prefill while routing + # single-chunk short / warmup prompts through stream_generate. + real_prefill_tokens = max(num_tokens - 1, 0) pipeline_chunks = ( - (num_tokens + effective_prefill_step_size - 1) // effective_prefill_step_size + (real_prefill_tokens + effective_prefill_step_size - 1) + // effective_prefill_step_size if effective_prefill_step_size > 0 else 0 ) @@ -585,12 +580,10 @@ def combined_progress_callback(processed: int, total: int) -> None: ) break # Stop after first iteration - cache is now filled except PrefillCancelled: + raise + finally: set_pipeline_queue_sends(model, queue_sends=False) set_pipeline_prefill(model, is_prefill=False) - raise - - set_pipeline_queue_sends(model, queue_sends=False) - set_pipeline_prefill(model, is_prefill=False) # stream_generate added 1 extra generated token to the cache, so we should trim it. # Because of needing to roll back arrays cache, we will generate on 2 tokens so trim 1 more. @@ -625,11 +618,10 @@ def warmup_inference( ) -> int: """Run a conservative synthetic warmup request to validate the MLX path. - Distributed pipeline warmup intentionally uses a minimal prompt because - richer synthetic prompts have been observed to wedge short-prompt - ``stream_generate`` prefill. Debug shaping overrides remain available for - single-node investigation, but pipeline warmup always stays on the minimal - sanity-check path. + Distributed pipeline warmup intentionally stays on the minimal sanity-check + prompt shape. Richer synthetic prompt templates have been observed to hang + warmup prefill in that distributed path, so prompt-shaping overrides remain + reserved for single-node investigation. """ logger.info(f"warming up inference for instance: {model_id}") diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 4ec293858..cf1710c8d 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -86,7 +86,7 @@ def _request_shape_debug_enabled() -> bool: ) if value is None: return False - return value.strip().lower() not in {"0", "false", "no", "off"} + return value.strip().lower() not in {"", "0", "false", "no", "off"} def log_request_shape( diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 5c8e10e8a..45adb7254 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -104,8 +104,38 @@ def _summarize_worker_task(task: Task) -> str: f"runner_id={task.runner_id!r}, " f"cancelled_task_id={task.cancelled_task_id!r})" ) - if isinstance(task, (TextGeneration, ImageEdits)): - return repr(task) + if isinstance(task, TextGeneration): + params = task.task_params + return ( + "TextGeneration(" + f"task_id={task.task_id!r}, " + f"command_id={task.command_id!r}, " + f"instance_id={task.instance_id!r}, " + f"model={params.model!r}, " + f"input_messages={len(params.input)}, " + f"chat_template_messages={len(params.chat_template_messages or [])}, " + f"images={len(params.images)}, " + f"cached_image_indices={sorted(params.image_hashes.keys())}, " + f"total_input_chunks={params.total_input_chunks}, " + f"image_count={params.image_count}, " + f"stream={params.stream}, " + f"reasoning_effort={params.reasoning_effort!r}, " + f"enable_thinking={params.enable_thinking!r})" + ) + if isinstance(task, ImageEdits): + params = task.task_params + return ( + "ImageEdits(" + f"task_id={task.task_id!r}, " + f"command_id={task.command_id!r}, " + f"instance_id={task.instance_id!r}, " + f"model={params.model!r}, " + f"total_input_chunks={params.total_input_chunks}, " + f"has_inline_image_data={bool(params.image_data)}, " + f"n={params.n!r}, " + f"size={params.size!r}, " + f"stream={params.stream!r})" + ) return task.__class__.__name__ diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py b/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py index e0ddb8b53..7049aa9fe 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_prefill_path_selection.py @@ -160,6 +160,111 @@ def _stream_generate(*args: object, **kwargs: object) -> Iterator[object]: assert prefill_mode_calls == [False, False] +def test_prefill_uses_stream_generate_at_single_chunk_boundary( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Prompts with only one real pipeline chunk must stay off pipeline prefill.""" + calls: list[str] = [] + fake_cache = _FakeCache() + prefill_mode_calls: list[bool] = [] + + monkeypatch.setattr(generate_module, "mx_barrier", lambda _group: None) + monkeypatch.setattr( + generate_module, + "set_pipeline_prefill", + lambda _model, *, is_prefill: prefill_mode_calls.append(is_prefill), + ) + monkeypatch.setattr( + generate_module, + "_has_pipeline_communication_layer", + lambda _model: True, + ) + monkeypatch.setattr( + generate_module, + "pipeline_parallel_prefill", + lambda **kwargs: calls.append("pipeline_parallel_prefill"), + ) + + def _stream_generate(*args: object, **kwargs: object) -> Iterator[object]: + calls.append("stream_generate") + yield object() + + monkeypatch.setattr(generate_module, "stream_generate", _stream_generate) + + prefill_tps, prefill_tokens, snapshots = generate_module.prefill( + model=_FakeModel(), + tokenizer=object(), + sampler=object(), + prompt_tokens=list(range(1366)), + cache=[fake_cache], + group=_FakeGroup(), + on_prefill_progress=None, + distributed_prompt_progress_callback=None, + ) + + assert calls == ["stream_generate"] + assert prefill_tokens == 1366 + assert snapshots == [] + assert prefill_tps >= 0.0 + assert fake_cache.trim_calls == [2] + assert prefill_mode_calls == [False, False] + + +def test_prefill_resets_pipeline_flags_after_unexpected_exception( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Unexpected prefill errors must restore pipeline-mode flags before bubbling.""" + fake_cache = _FakeCache() + prefill_mode_calls: list[bool] = [] + queue_send_calls: list[bool] = [] + + monkeypatch.setattr(generate_module, "mx_barrier", lambda _group: None) + monkeypatch.setattr( + generate_module, + "set_pipeline_prefill", + lambda _model, *, is_prefill: prefill_mode_calls.append(is_prefill), + ) + monkeypatch.setattr( + generate_module, + "set_pipeline_queue_sends", + lambda _model, *, queue_sends: queue_send_calls.append(queue_sends), + ) + monkeypatch.setattr( + generate_module, + "_has_pipeline_communication_layer", + lambda _model: True, + ) + monkeypatch.setattr( + generate_module, + "stream_generate", + lambda *args, **kwargs: (_ for _ in () if False), + ) + + def _pipeline_parallel_prefill(**kwargs: object) -> None: + raise RuntimeError("boom") + + monkeypatch.setattr( + generate_module, + "pipeline_parallel_prefill", + _pipeline_parallel_prefill, + ) + + with pytest.raises(RuntimeError, match="boom"): + generate_module.prefill( + model=_FakeModel(), + tokenizer=object(), + sampler=object(), + prompt_tokens=list(range(5000)), + cache=[fake_cache], + group=_FakeGroup(), + on_prefill_progress=None, + distributed_prompt_progress_callback=None, + ) + + assert queue_send_calls == [True, False] + assert prefill_mode_calls == [True, False] + + def test_prefill_uses_stream_generate_when_model_is_not_pipeline( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py b/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py new file mode 100644 index 000000000..04f99dbb7 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py @@ -0,0 +1,15 @@ +import pytest + +from exo.worker.engines.mlx import utils_mlx + + +@pytest.mark.parametrize("env_name", ["SKULK_TRACE_REQUEST_SHAPES", "EXO_TRACE_REQUEST_SHAPES"]) +def test_request_shape_debug_blank_env_is_disabled( + monkeypatch: pytest.MonkeyPatch, + env_name: str, +) -> None: + monkeypatch.delenv("SKULK_TRACE_REQUEST_SHAPES", raising=False) + monkeypatch.delenv("EXO_TRACE_REQUEST_SHAPES", raising=False) + monkeypatch.setenv(env_name, "") + + assert utils_mlx._request_shape_debug_enabled() is False diff --git a/src/exo/worker/tests/unittests/test_worker_task_summary.py b/src/exo/worker/tests/unittests/test_worker_task_summary.py new file mode 100644 index 000000000..08e91ce26 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_worker_task_summary.py @@ -0,0 +1,47 @@ +from exo.api.types import ImageEditsTaskParams +from exo.shared.types.common import CommandId +from exo.shared.types.tasks import TaskId, TextGeneration +from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams +from exo.worker.main import _summarize_worker_task + + +def test_summarize_worker_task_redacts_text_generation_payloads() -> None: + task = TextGeneration( + task_id=TaskId("task-1"), + command_id=CommandId("command-1"), + instance_id="instance-1", + task_params=TextGenerationTaskParams( + model="model-a", + input=[InputMessage(role="user", content="secret prompt")], + stream=True, + ), + ) + + summary = _summarize_worker_task(task) + + assert "secret prompt" not in summary + assert "input_messages=1" in summary + assert "model='model-a'" in summary + + +def test_summarize_worker_task_redacts_image_edit_payloads() -> None: + from exo.shared.types.tasks import ImageEdits + + task = ImageEdits( + task_id=TaskId("task-2"), + command_id=CommandId("command-2"), + instance_id="instance-2", + task_params=ImageEditsTaskParams( + image_data="A" * 128, + total_input_chunks=3, + prompt="remove background", + model="image-model", + ), + ) + + summary = _summarize_worker_task(task) + + assert "remove background" not in summary + assert "AAAA" not in summary + assert "has_inline_image_data=True" in summary + assert "total_input_chunks=3" in summary diff --git a/website/docs/architecture.md b/website/docs/architecture.md index 89f3c5a1e..4d7302f86 100644 --- a/website/docs/architecture.md +++ b/website/docs/architecture.md @@ -218,9 +218,10 @@ fallbacks while older scripts are updated. Distributed pipeline models now use an intentionally minimal warmup request. The goal of warmup is only to validate that MLX prefill and first-token -generation are functional without risking richer synthetic prompt shapes that -have been observed to wedge `stream_generate` prefill on multi-node pipeline -models. +generation are functional while matching the current `generate.prefill()` +routing behavior: short single-chunk prompts are sent through +`stream_generate` because `pipeline_parallel_prefill` has been observed to +wedge for that case on multi-node pipeline models. For distributed pipeline warmup, Skulk always uses: diff --git a/website/docs/model-behaviors/gemma4.md b/website/docs/model-behaviors/gemma4.md index f5e6742de..3ddf52bd7 100644 --- a/website/docs/model-behaviors/gemma4.md +++ b/website/docs/model-behaviors/gemma4.md @@ -107,8 +107,9 @@ minimal synthetic prompt by design: - a single user message with content `hello` This is intentional. During debugging, richer synthetic warmup prompts were -observed to wedge short-prompt `stream_generate` prefill on multi-node pipeline -setups. +observed to trigger the distributed warmup hang path on multi-node pipeline +setups. Short single-chunk prompts are therefore routed through +`stream_generate`, while richer synthetic warmup shapes are avoided entirely. Two debug-only warmup shaping env vars exist: From 431559c0cbc6ba6ba0b89e6e4d3d8c429a16b473 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 18:21:21 -0500 Subject: [PATCH 26/32] Tighten warmup skip and add review regression tests --- src/exo/api/tests/test_config_api.py | 56 +++++++++++++++++++ src/exo/worker/engines/mlx/auto_parallel.py | 2 +- src/exo/worker/runner/llm_inference/runner.py | 22 +++++++- .../test_mlx/test_auto_parallel_debug.py | 11 ++++ .../test_runner/test_event_ordering.py | 34 +++++++++++ website/docs/architecture.md | 3 +- website/docs/model-behaviors/gemma4.md | 2 + 7 files changed, 125 insertions(+), 5 deletions(-) create mode 100644 src/exo/api/tests/test_config_api.py create mode 100644 src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py diff --git a/src/exo/api/tests/test_config_api.py b/src/exo/api/tests/test_config_api.py new file mode 100644 index 000000000..d5efe6103 --- /dev/null +++ b/src/exo/api/tests/test_config_api.py @@ -0,0 +1,56 @@ +# pyright: reportUnusedFunction=false, reportPrivateUsage=false +"""Tests for the GET /config API endpoint.""" + +from pathlib import Path +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from exo.api.main import API +from exo.shared.types.common import NodeId +from exo.utils.channels import channel + + +def _build_api(node_id: str = "test-node") -> API: + """Create a minimal API instance for config endpoint testing.""" + command_sender, _ = channel() + download_sender, _ = channel() + _, event_receiver = channel() + _, election_receiver = channel() + return API( + NodeId(node_id), + port=52415, + event_receiver=event_receiver, + command_sender=command_sender, + download_command_sender=download_sender, + election_receiver=election_receiver, + enable_event_log=False, + mount_dashboard=False, + ) + + +def test_get_config_reports_effective_kv_backend_when_file_exists( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config_path = tmp_path / "exo.yaml" + config_path.write_text( + "inference:\n kv_cache_backend: default\nhf_token: secret-token\n", + encoding="utf-8", + ) + monkeypatch.setenv("SKULK_KV_CACHE_BACKEND", "optiq") + + api = _build_api() + api._config_path = config_path # pyright: ignore[reportPrivateUsage] + client = TestClient(api.app) + + response = client.get("/config") + + assert response.status_code == 200 + data: dict[str, Any] = response.json() + assert data["fileExists"] is True + assert data["configPath"] == str(config_path) + assert data["config"]["inference"]["kv_cache_backend"] == "default" + assert data["config"].get("hf_token") is None + assert data["effective"]["kv_cache_backend"] == "optiq" + assert data["effective"]["has_hf_token"] is True diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index 6ab5f4bfc..d84448dc3 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -81,7 +81,7 @@ def _mlx_hang_debug_enabled() -> bool: ) if value is None: return False - return value.lower() not in {"", "0", "false", "no", "off"} + return value.strip().lower() not in {"", "0", "false", "no", "off"} def _mlx_hang_debug_interval_seconds() -> float: diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index f17812150..49075cae5 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -69,11 +69,26 @@ from .tool_parsers import make_mlx_parser -def _should_skip_llm_warmup() -> bool: +def _should_skip_llm_warmup(group_size: int) -> bool: + """Return whether the synthetic warmup request should be bypassed. + + Distributed warmup uses barriers and collectives, so skipping it on only + one rank can strand peers inside warmup coordination. The debug bypass is + therefore limited to single-node groups. + """ # Temporary escape hatch for debug sessions where we want runners to reach # Ready without issuing the synthetic warmup request. Normal behavior keeps # warmup enabled unless SKULK_SKIP_LLM_WARMUP=1 is set explicitly. - return os.environ.get("SKULK_SKIP_LLM_WARMUP") == "1" + if os.environ.get("SKULK_SKIP_LLM_WARMUP") != "1": + return False + if group_size > 1: + logger.warning( + "Ignoring SKULK_SKIP_LLM_WARMUP=1 for distributed warmup " + f"(group_size={group_size}); synthetic warmup must stay consistent " + "across all ranks" + ) + return False + return True class ExitCode(str, Enum): @@ -261,7 +276,8 @@ def on_layer_loaded(layers_loaded: int, total: int) -> None: self.update_status(RunnerWarmingUp()) self.acknowledge_task(task) - if _should_skip_llm_warmup(): + group_size = self.generator.group.size() if self.generator.group else 1 + if _should_skip_llm_warmup(group_size): logger.warning( "Skipping LLM warmup and marking runner ready " "(temporary debug bypass; unset SKULK_SKIP_LLM_WARMUP " diff --git a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py new file mode 100644 index 000000000..49eb7a707 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py @@ -0,0 +1,11 @@ +import pytest + +from exo.worker.engines.mlx import auto_parallel + + +def test_mlx_hang_debug_blank_env_value_is_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SKULK_MLX_HANG_DEBUG", "") + + assert auto_parallel._mlx_hang_debug_enabled() is False # pyright: ignore[reportPrivateUsage] 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 02fec3d72..f95d7c819 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 @@ -224,6 +224,14 @@ def size(self) -> int: return 1 +class MockDistributedGroup: + def rank(self) -> int: + return 0 + + def size(self) -> int: + return 2 + + def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None): bound_instance = get_bound_mlx_ring_instance( instance_id=INSTANCE_1_ID, @@ -373,6 +381,32 @@ def fail_if_called(*_args: object, **_kwargs: object) -> None: ) +def test_distributed_warmup_skip_env_is_ignored( + patch_out_mlx: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("SKULK_SKIP_LLM_WARMUP", "1") + monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockDistributedGroup())) + + warmup_calls = 0 + + def fake_warmup(self: object) -> None: + del self + nonlocal warmup_calls + warmup_calls += 1 + + monkeypatch.setattr(mlx_batch_generator.BatchGenerator, "warmup", fake_warmup) + monkeypatch.setattr(mlx_batch_generator.SequentialGenerator, "warmup", fake_warmup) + + events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, SHUTDOWN_TASK]) + + assert warmup_calls == 1 + assert any( + isinstance(event, RunnerStatusUpdated) + and isinstance(event.runner_status, RunnerReady) + for event in events + ) + + def test_first_live_request_shape_is_logged_once( patch_out_mlx: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch ): diff --git a/website/docs/architecture.md b/website/docs/architecture.md index 4d7302f86..0a3cee28f 100644 --- a/website/docs/architecture.md +++ b/website/docs/architecture.md @@ -249,7 +249,8 @@ There is also an emergency bypass for debugging: - `SKULK_SKIP_LLM_WARMUP=1` That bypass marks the runner ready without issuing the synthetic warmup request -and should only be used for diagnosis, not normal operation. +and should only be used for diagnosis, not normal operation. Distributed +groups ignore it so warmup coordination cannot diverge across ranks. ## When to Read More diff --git a/website/docs/model-behaviors/gemma4.md b/website/docs/model-behaviors/gemma4.md index 3ddf52bd7..6b1894d59 100644 --- a/website/docs/model-behaviors/gemma4.md +++ b/website/docs/model-behaviors/gemma4.md @@ -126,6 +126,8 @@ export SKULK_SKIP_LLM_WARMUP=1 ``` That bypass should be treated as a temporary debugging escape hatch only. +Distributed pipeline groups ignore it so one rank cannot skip warmup while +peers are still participating in warmup coordination. ## Why This Matters From 2302c410b9332e8f87f85de0706a2c5f1f1768f6 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 18:49:44 -0500 Subject: [PATCH 27/32] Bound warmup decode to the first token --- .../worker/engines/mlx/generator/generate.py | 4 ++ .../unittests/test_mlx/test_warmup_request.py | 39 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 6f6dc3d33..b29cf7bf3 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -690,6 +690,10 @@ def warmup_inference( group=group, ): tokens_generated += 1 + # Warmup is only meant to validate prefill plus the first decode + # step. Stopping after the first token keeps runner readiness + # bounded even if the model would otherwise continue sampling. + break check_for_cancel_every = min( math.ceil(tokens_generated / max(time.monotonic() - t, 0.001)), 100 diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 1df4f429b..0df333d36 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -207,3 +207,42 @@ def fake_all_gather(array: object, *, group: object): task_params.instructions == "You are a helpful assistant. Answer the user in one short sentence." ) + + +def test_warmup_inference_stops_after_first_generated_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + generated_tokens = 0 + + def fake_apply_chat_template(*, tokenizer: object, task_params: object, model_card: object): + del tokenizer, task_params, model_card + return "warmup prompt" + + def fake_mx_barrier(_group: object) -> None: + return None + + def fake_mlx_generate(**_kwargs: object): + nonlocal generated_tokens + for token in ("first", "second", "third"): + generated_tokens += 1 + yield token + + def fake_all_gather(array: object, *, group: object): + del group + return array + + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) + monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) + monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) + + check_every = generate_mod.warmup_inference( + model=object(), # type: ignore[arg-type] + tokenizer=object(), # type: ignore[arg-type] + group=cast(object, _SingleNodeGroup()), # type: ignore[arg-type] + model_id=ModelId("mlx-community/gemma-4-26b-a4b-it-4bit"), + model_card=None, + ) + + assert generated_tokens == 1 + assert check_every == 100 From e7a1d591a4361ac37f572f7b1d3250bed9a04dd7 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 19:56:41 -0500 Subject: [PATCH 28/32] Tighten warmup planner and debug helpers --- src/exo/worker/engines/mlx/gemma4_prompt.py | 10 ++++++---- .../worker/engines/mlx/generator/generate.py | 12 +++++------ src/exo/worker/main.py | 2 +- src/exo/worker/plan.py | 4 ++-- .../tests/unittests/test_plan/test_warmup.py | 20 +++++++++---------- 5 files changed, 24 insertions(+), 24 deletions(-) diff --git a/src/exo/worker/engines/mlx/gemma4_prompt.py b/src/exo/worker/engines/mlx/gemma4_prompt.py index 4f8ee7a52..1cf44ce41 100644 --- a/src/exo/worker/engines/mlx/gemma4_prompt.py +++ b/src/exo/worker/engines/mlx/gemma4_prompt.py @@ -1,9 +1,11 @@ """Gemma 4 prompt rendering helpers. -This module mirrors the Gemma 4 chat structure used by the reference -Hugging Face template and Ollama's dedicated Gemma 4 renderer. We use it -to avoid relying on generic tokenizer chat templating for Gemma 4 because -we need exact control over the multimodal prompt the model sees. +This module follows the Gemma 4 chat structure used by the reference Hugging +Face template and Ollama's dedicated Gemma 4 renderer, while preserving one +intentional divergence: when thinking is disabled we omit the empty synthetic +thought-channel suffix because that shape wedges our distributed MLX warmup +path. We keep a dedicated renderer so multimodal prompts stay under explicit +repo control instead of relying on generic tokenizer chat templating. """ from typing import Any diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index b29cf7bf3..f550648c9 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -114,8 +114,8 @@ def _warmup_repeat_count() -> int: return 1 -def _is_distributed_pipeline_warmup(group: mx.distributed.Group | None) -> bool: - """Return whether warmup is running for a multi-node pipeline model.""" +def _is_distributed_warmup(group: mx.distributed.Group | None) -> bool: + """Return whether warmup is running with a multi-node distributed group.""" return group is not None and group.size() > 1 @@ -126,21 +126,21 @@ def _warmup_user_content(group: mx.distributed.Group | None) -> str: synthetic prompts have been observed to wedge stream_generate prefill. Single-node debugging may still scale the neutral content via environment. """ - if _is_distributed_pipeline_warmup(group): + if _is_distributed_warmup(group): return "hello" return " ".join(["hello"] * _warmup_repeat_count()) def _warmup_instructions(group: mx.distributed.Group | None) -> str | None: """Return optional warmup instructions for prompt-shape debugging.""" - if _is_distributed_pipeline_warmup(group): + if _is_distributed_warmup(group): return None raw = os.environ.get("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS") or os.environ.get( "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS" ) if raw is None: return None - include = raw.strip().lower() not in {"0", "false", "no", "off"} + include = raw.strip().lower() not in {"", "0", "false", "no", "off"} if not include: return None return "You are a helpful assistant. Answer the user in one short sentence." @@ -625,7 +625,7 @@ def warmup_inference( """ logger.info(f"warming up inference for instance: {model_id}") - if _is_distributed_pipeline_warmup(group) and ( + if _is_distributed_warmup(group) and ( _warmup_repeat_count() != 1 or _warmup_instructions(None) is not None ): logger.info( diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 45adb7254..b9d733a86 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -259,7 +259,7 @@ async def _forward_info(self, recv: Receiver[GatheredInfo]): ) ) except (ClosedResourceError, BrokenResourceError): - logger.info( + logger.debug( "Worker info forwarding stopped because the event stream " "was already closed" ) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index f7dfdc2bb..28f34a1c3 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -251,7 +251,7 @@ def _ready_to_warmup( accepting_ranks_ready = device_rank > 0 and all( isinstance( all_runners.get(global_runner_id, None), - (RunnerLoaded, RunnerWarmingUp, RunnerReady), + (RunnerLoaded, RunnerWarmingUp), ) for global_runner_id in shard_assignments.runner_to_shard ) @@ -260,7 +260,7 @@ def _ready_to_warmup( connecting_rank_ready = device_rank == 0 and all( isinstance( all_runners.get(global_runner_id, None), - (RunnerWarmingUp, RunnerReady), + RunnerWarmingUp, ) for global_runner_id in shard_assignments.runner_to_shard if global_runner_id != runner_id diff --git a/src/exo/worker/tests/unittests/test_plan/test_warmup.py b/src/exo/worker/tests/unittests/test_plan/test_warmup.py index a40e4abc8..b513a20ab 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py +++ b/src/exo/worker/tests/unittests/test_plan/test_warmup.py @@ -68,11 +68,11 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming(): assert result.instance_id == INSTANCE_1_ID -def test_plan_starts_warmup_for_accepting_rank_when_peer_is_already_ready(): +def test_plan_does_not_start_warmup_for_accepting_rank_when_peer_is_ready(): """ - Non-zero ranks should still start warmup when a peer races through the - temporary debug warmup bypass and reaches Ready before they observe - WarmingUp. + Non-zero ranks should not start distributed warmup once any peer has + already reached Ready because that implies the coordinated warmup phase has + moved past the point where this runner can safely join. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=3) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=3) @@ -108,8 +108,7 @@ def test_plan_starts_warmup_for_accepting_rank_when_peer_is_already_ready(): tasks={}, ) - assert isinstance(result, StartWarmup) - assert result.instance_id == INSTANCE_1_ID + assert result is None def test_plan_starts_warmup_for_rank_zero_after_others_warming(): @@ -153,10 +152,10 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming(): assert result.instance_id == INSTANCE_1_ID -def test_plan_starts_warmup_for_rank_zero_after_other_rank_is_already_ready(): +def test_plan_does_not_start_warmup_for_rank_zero_after_other_rank_is_ready(): """ - Rank zero should still start warmup when another rank races from WarmingUp - to Ready before the planner observes the intermediate state. + Rank zero should not start distributed warmup once another rank has already + reached Ready because the synchronized warmup phase is no longer joinable. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) @@ -190,8 +189,7 @@ def test_plan_starts_warmup_for_rank_zero_after_other_rank_is_already_ready(): tasks={}, ) - assert isinstance(result, StartWarmup) - assert result.instance_id == INSTANCE_1_ID + assert result is None def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warming(): From 9163d8d6595e5f05b953e63a3b9d742afe449490 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 20:08:45 -0500 Subject: [PATCH 29/32] Tighten env precedence for warmup and debug config --- src/exo/api/main.py | 16 +++++----- src/exo/api/tests/test_config_api.py | 19 ++++++++++++ src/exo/worker/engines/mlx/auto_parallel.py | 16 +++++++--- src/exo/worker/engines/mlx/constants.py | 10 +++--- .../worker/engines/mlx/generator/generate.py | 28 +++++++++++------ src/exo/worker/engines/mlx/utils_mlx.py | 11 +++++-- .../test_mlx/test_auto_parallel_debug.py | 10 ++++++ .../test_mlx/test_request_shape_debug.py | 9 ++++++ .../unittests/test_mlx/test_warmup_request.py | 31 +++++++++++++++++++ 9 files changed, 120 insertions(+), 30 deletions(-) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index 953b32dc1..f8c1da9f2 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -2563,6 +2563,12 @@ async def complete_onboarding(self) -> JSONResponse: # Config & Store endpoints # ------------------------------------------------------------------ + def _effective_kv_cache_backend(self) -> str: + """Return the effective KV backend after SKULK/EXO env precedence is applied.""" + if "SKULK_KV_CACHE_BACKEND" in os.environ: + return os.environ["SKULK_KV_CACHE_BACKEND"] or "default" + return os.environ.get("EXO_KV_CACHE_BACKEND") or "default" + async def get_config(self) -> JSONResponse: if not self._config_path.exists(): return JSONResponse( @@ -2571,10 +2577,7 @@ async def get_config(self) -> JSONResponse: "configPath": str(self._config_path), "fileExists": False, "effective": { - "kv_cache_backend": os.environ.get( - "SKULK_KV_CACHE_BACKEND", - os.environ.get("EXO_KV_CACHE_BACKEND", "default"), - ), + "kv_cache_backend": self._effective_kv_cache_backend(), }, } ) @@ -2591,10 +2594,7 @@ async def get_config(self) -> JSONResponse: "configPath": str(self._config_path), "fileExists": True, "effective": { - "kv_cache_backend": os.environ.get( - "SKULK_KV_CACHE_BACKEND", - os.environ.get("EXO_KV_CACHE_BACKEND", "default"), - ), + "kv_cache_backend": self._effective_kv_cache_backend(), "has_hf_token": has_hf_token or "HF_TOKEN" in os.environ, }, } diff --git a/src/exo/api/tests/test_config_api.py b/src/exo/api/tests/test_config_api.py index d5efe6103..3cf952ff6 100644 --- a/src/exo/api/tests/test_config_api.py +++ b/src/exo/api/tests/test_config_api.py @@ -54,3 +54,22 @@ def test_get_config_reports_effective_kv_backend_when_file_exists( assert data["config"].get("hf_token") is None assert data["effective"]["kv_cache_backend"] == "optiq" assert data["effective"]["has_hf_token"] is True + + +def test_get_config_treats_blank_skulk_kv_backend_as_default( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config_path = tmp_path / "exo.yaml" + config_path.write_text("inference:\n kv_cache_backend: default\n", encoding="utf-8") + monkeypatch.setenv("SKULK_KV_CACHE_BACKEND", "") + monkeypatch.setenv("EXO_KV_CACHE_BACKEND", "optiq") + + api = _build_api() + api._config_path = config_path # pyright: ignore[reportPrivateUsage] + client = TestClient(api.app) + + response = client.get("/config") + + assert response.status_code == 200 + data: dict[str, Any] = response.json() + assert data["effective"]["kv_cache_backend"] == "default" diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index d84448dc3..dd9079d63 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -75,18 +75,24 @@ _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = [] +def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: + """Return the SKULK env value when present, else the legacy EXO value.""" + if skulk_key in os.environ: + return os.environ[skulk_key] + return os.environ.get(exo_key) + + def _mlx_hang_debug_enabled() -> bool: - value = os.environ.get("SKULK_MLX_HANG_DEBUG") or os.environ.get( - "EXO_MLX_HANG_DEBUG" - ) + value = _preferred_env_value("SKULK_MLX_HANG_DEBUG", "EXO_MLX_HANG_DEBUG") if value is None: return False return value.strip().lower() not in {"", "0", "false", "no", "off"} def _mlx_hang_debug_interval_seconds() -> float: - raw = os.environ.get("SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS") or os.environ.get( - "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS" + raw = _preferred_env_value( + "SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS", + "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS", ) if raw is None: return 30.0 diff --git a/src/exo/worker/engines/mlx/constants.py b/src/exo/worker/engines/mlx/constants.py index 912897f94..3e7388d28 100644 --- a/src/exo/worker/engines/mlx/constants.py +++ b/src/exo/worker/engines/mlx/constants.py @@ -23,12 +23,14 @@ "optiq", ] DEFAULT_KV_CACHE_BACKEND: KVCacheBackend = "default" +_kv_cache_backend_value = ( + os.environ["SKULK_KV_CACHE_BACKEND"] + if "SKULK_KV_CACHE_BACKEND" in os.environ + else os.environ.get("EXO_KV_CACHE_BACKEND", DEFAULT_KV_CACHE_BACKEND) +) KV_CACHE_BACKEND: KVCacheBackend = cast( KVCacheBackend, - os.environ.get( - "SKULK_KV_CACHE_BACKEND", - os.environ.get("EXO_KV_CACHE_BACKEND", DEFAULT_KV_CACHE_BACKEND), - ), + _kv_cache_backend_value if _kv_cache_backend_value else DEFAULT_KV_CACHE_BACKEND, ) TURBOQUANT_K_BITS: int | None = ( int(os.environ.get("SKULK_TQ_K_BITS", os.environ.get("EXO_TQ_K_BITS", ""))) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index f550648c9..554611e48 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -80,20 +80,26 @@ _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 +def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: + """Return the SKULK env value when present, else the legacy EXO value.""" + if skulk_key in os.environ: + return os.environ[skulk_key] + return os.environ.get(exo_key) + + def _mlx_hang_debug_enabled() -> bool: """Return whether verbose warmup/prefill hang diagnostics are enabled.""" - value = os.environ.get("SKULK_MLX_HANG_DEBUG") or os.environ.get( - "EXO_MLX_HANG_DEBUG" - ) + value = _preferred_env_value("SKULK_MLX_HANG_DEBUG", "EXO_MLX_HANG_DEBUG") if value is None: return False - return value.strip().lower() not in {"0", "false", "no", "off"} + return value.strip().lower() not in {"", "0", "false", "no", "off"} def _mlx_hang_debug_interval_seconds() -> float: """Return the periodic interval for hang-debug watchdog logs.""" - raw = os.environ.get("SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS") or os.environ.get( - "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS" + raw = _preferred_env_value( + "SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS", + "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS", ) if raw is None: return 30.0 @@ -104,8 +110,9 @@ def _mlx_hang_debug_interval_seconds() -> float: def _warmup_repeat_count() -> int: """Return the neutral warmup token repeat count used for debugging.""" - raw = os.environ.get("SKULK_DEBUG_WARMUP_REPEAT_COUNT") or os.environ.get( - "EXO_DEBUG_WARMUP_REPEAT_COUNT" + raw = _preferred_env_value( + "SKULK_DEBUG_WARMUP_REPEAT_COUNT", + "EXO_DEBUG_WARMUP_REPEAT_COUNT", ) if raw is None: return 1 @@ -135,8 +142,9 @@ def _warmup_instructions(group: mx.distributed.Group | None) -> str | None: """Return optional warmup instructions for prompt-shape debugging.""" if _is_distributed_warmup(group): return None - raw = os.environ.get("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS") or os.environ.get( - "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS" + raw = _preferred_env_value( + "SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", + "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", ) if raw is None: return None diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index cf1710c8d..520a292fe 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -79,11 +79,16 @@ Group = mx.distributed.Group +def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: + """Return the SKULK env value when present, else the legacy EXO value.""" + if skulk_key in os.environ: + return os.environ[skulk_key] + return os.environ.get(exo_key) + + def _request_shape_debug_enabled() -> bool: """Return whether request-shape tracing is enabled for prompt debugging.""" - value = os.environ.get("SKULK_TRACE_REQUEST_SHAPES") or os.environ.get( - "EXO_TRACE_REQUEST_SHAPES" - ) + value = _preferred_env_value("SKULK_TRACE_REQUEST_SHAPES", "EXO_TRACE_REQUEST_SHAPES") if value is None: return False return value.strip().lower() not in {"", "0", "false", "no", "off"} diff --git a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py index 49eb7a707..c0c9b2e70 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel_debug.py @@ -6,6 +6,16 @@ def test_mlx_hang_debug_blank_env_value_is_disabled( monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.delenv("EXO_MLX_HANG_DEBUG", raising=False) monkeypatch.setenv("SKULK_MLX_HANG_DEBUG", "") assert auto_parallel._mlx_hang_debug_enabled() is False # pyright: ignore[reportPrivateUsage] + + +def test_mlx_hang_debug_blank_skulk_value_does_not_fallback_to_legacy_true_value( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SKULK_MLX_HANG_DEBUG", "") + monkeypatch.setenv("EXO_MLX_HANG_DEBUG", "1") + + assert auto_parallel._mlx_hang_debug_enabled() is False # pyright: ignore[reportPrivateUsage] diff --git a/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py b/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py index 04f99dbb7..41c8b582b 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_request_shape_debug.py @@ -13,3 +13,12 @@ def test_request_shape_debug_blank_env_is_disabled( monkeypatch.setenv(env_name, "") assert utils_mlx._request_shape_debug_enabled() is False + + +def test_request_shape_debug_blank_skulk_env_does_not_fallback_to_legacy_true_value( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SKULK_TRACE_REQUEST_SHAPES", "") + monkeypatch.setenv("EXO_TRACE_REQUEST_SHAPES", "1") + + assert utils_mlx._request_shape_debug_enabled() is False diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index 0df333d36..cea453fdc 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -17,6 +17,17 @@ def size(self) -> int: return 1 +def _clear_warmup_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Reset warmup debug env vars so tests are independent from the shell.""" + for env_name in ( + "SKULK_DEBUG_WARMUP_REPEAT_COUNT", + "EXO_DEBUG_WARMUP_REPEAT_COUNT", + "SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", + "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", + ): + monkeypatch.delenv(env_name, raising=False) + + def test_warmup_inference_uses_safe_default_user_content_without_instructions( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -51,6 +62,7 @@ def fake_log_request_shape( captured["logged_prompt"] = prompt captured["logged_extra"] = extra + _clear_warmup_env(monkeypatch) monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) @@ -106,6 +118,7 @@ def fake_all_gather(array: object, *, group: object): del group return array + _clear_warmup_env(monkeypatch) monkeypatch.setenv("SKULK_DEBUG_WARMUP_REPEAT_COUNT", "4") monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) @@ -146,6 +159,7 @@ def fake_all_gather(array: object, *, group: object): del group return array + _clear_warmup_env(monkeypatch) monkeypatch.setenv("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "1") monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) @@ -186,6 +200,7 @@ def fake_all_gather(array: object, *, group: object): del group return array + _clear_warmup_env(monkeypatch) monkeypatch.setenv("SKULK_DEBUG_WARMUP_REPEAT_COUNT", "4") monkeypatch.setenv("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "1") monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) @@ -231,6 +246,7 @@ def fake_all_gather(array: object, *, group: object): del group return array + _clear_warmup_env(monkeypatch) monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) @@ -246,3 +262,18 @@ def fake_all_gather(array: object, *, group: object): assert generated_tokens == 1 assert check_every == 100 + + +def test_warmup_helpers_prefer_blank_skulk_values_over_legacy_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_warmup_env(monkeypatch) + monkeypatch.setenv("SKULK_DEBUG_WARMUP_REPEAT_COUNT", "") + monkeypatch.setenv("EXO_DEBUG_WARMUP_REPEAT_COUNT", "4") + monkeypatch.setenv("SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "") + monkeypatch.setenv("EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "1") + + assert generate_mod._warmup_repeat_count() == 1 # pyright: ignore[reportPrivateUsage] + assert ( + generate_mod._warmup_instructions(cast(object, _SingleNodeGroup())) is None + ) # pyright: ignore[reportPrivateUsage] From 33f8c0f358601c0e94ae72b8a3532e4c18040493 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 20:44:43 -0500 Subject: [PATCH 30/32] Normalize invalid KV backend values in config --- src/exo/api/main.py | 19 +++++++++++++++++-- src/exo/api/tests/test_config_api.py | 18 ++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index f8c1da9f2..95ba2b496 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -200,6 +200,7 @@ from exo.utils.disk_event_log import DiskEventLog from exo.utils.power_sampler import PowerSampler from exo.utils.task_group import TaskGroup +from exo.worker.engines.mlx.constants import DEFAULT_KV_CACHE_BACKEND, KVCacheBackend if TYPE_CHECKING: from exo.store.config import ExoConfig @@ -2566,8 +2567,22 @@ async def complete_onboarding(self) -> JSONResponse: def _effective_kv_cache_backend(self) -> str: """Return the effective KV backend after SKULK/EXO env precedence is applied.""" if "SKULK_KV_CACHE_BACKEND" in os.environ: - return os.environ["SKULK_KV_CACHE_BACKEND"] or "default" - return os.environ.get("EXO_KV_CACHE_BACKEND") or "default" + configured_backend = os.environ["SKULK_KV_CACHE_BACKEND"] + else: + configured_backend = os.environ.get("EXO_KV_CACHE_BACKEND", "") + if not configured_backend: + return DEFAULT_KV_CACHE_BACKEND + + valid_backends: tuple[KVCacheBackend, ...] = ( + "default", + "mlx_quantized", + "turboquant", + "turboquant_adaptive", + "optiq", + ) + if configured_backend not in valid_backends: + return DEFAULT_KV_CACHE_BACKEND + return configured_backend async def get_config(self) -> JSONResponse: if not self._config_path.exists(): diff --git a/src/exo/api/tests/test_config_api.py b/src/exo/api/tests/test_config_api.py index 3cf952ff6..161ac5208 100644 --- a/src/exo/api/tests/test_config_api.py +++ b/src/exo/api/tests/test_config_api.py @@ -73,3 +73,21 @@ def test_get_config_treats_blank_skulk_kv_backend_as_default( assert response.status_code == 200 data: dict[str, Any] = response.json() assert data["effective"]["kv_cache_backend"] == "default" + + +def test_get_config_treats_invalid_skulk_kv_backend_as_default( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config_path = tmp_path / "exo.yaml" + config_path.write_text("inference:\n kv_cache_backend: default\n", encoding="utf-8") + monkeypatch.setenv("SKULK_KV_CACHE_BACKEND", "typo-backend") + + api = _build_api() + api._config_path = config_path # pyright: ignore[reportPrivateUsage] + client = TestClient(api.app) + + response = client.get("/config") + + assert response.status_code == 200 + data: dict[str, Any] = response.json() + assert data["effective"]["kv_cache_backend"] == "default" From 3312c29a3ab3d6ee4b5dbeb468c28c9b155fc4c7 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 20:54:11 -0500 Subject: [PATCH 31/32] Keep warmup cancel checks from becoming too frequent --- .../worker/engines/mlx/generator/generate.py | 16 ++++++-- .../unittests/test_mlx/test_warmup_request.py | 37 +++++++++++++++++++ 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 554611e48..982980b8b 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -78,6 +78,7 @@ generation_stream = mx.new_stream(mx.default_device()) _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 +_MIN_CANCEL_CHECK_INTERVAL = 10 def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: @@ -703,9 +704,18 @@ def warmup_inference( # bounded even if the model would otherwise continue sampling. break - check_for_cancel_every = min( - math.ceil(tokens_generated / max(time.monotonic() - t, 0.001)), 100 - ) + # The single-token warmup path intentionally samples only the first decode + # step, so cold-start compile/prefill latency can dominate the elapsed + # measurement here. Keep cancellation checks reasonably spaced even when + # that first-token latency is slow, while still capping the interval for + # fast models. + if tokens_generated == 0: + check_for_cancel_every = 0 + else: + check_for_cancel_every = max( + _MIN_CANCEL_CHECK_INTERVAL, + min(math.ceil(tokens_generated / max(time.monotonic() - t, 0.001)), 100), + ) with _hang_debug_watch( f"warmup final barrier model={model_id} tokens_generated={tokens_generated}" diff --git a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py index cea453fdc..dcf1f7e1d 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_warmup_request.py @@ -277,3 +277,40 @@ def test_warmup_helpers_prefer_blank_skulk_values_over_legacy_env( assert ( generate_mod._warmup_instructions(cast(object, _SingleNodeGroup())) is None ) # pyright: ignore[reportPrivateUsage] + + +def test_warmup_inference_enforces_minimum_cancel_check_interval_on_slow_start( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monotonic_values = iter([100.0, 112.0]) + + def fake_apply_chat_template(*, tokenizer: object, task_params: object, model_card: object): + del tokenizer, task_params, model_card + return "warmup prompt" + + def fake_mx_barrier(_group: object) -> None: + return None + + def fake_mlx_generate(**_kwargs: object): + yield "first" + + def fake_all_gather(array: object, *, group: object): + del group + return array + + _clear_warmup_env(monkeypatch) + monkeypatch.setattr(generate_mod, "apply_chat_template", fake_apply_chat_template) + monkeypatch.setattr(generate_mod, "mx_barrier", fake_mx_barrier) + monkeypatch.setattr(generate_mod, "mlx_generate", fake_mlx_generate) + monkeypatch.setattr(generate_mod.mx.distributed, "all_gather", fake_all_gather) + monkeypatch.setattr(generate_mod.time, "monotonic", lambda: next(monotonic_values)) + + check_every = generate_mod.warmup_inference( + model=object(), # type: ignore[arg-type] + tokenizer=object(), # type: ignore[arg-type] + group=cast(object, _SingleNodeGroup()), # type: ignore[arg-type] + model_id=ModelId("mlx-community/gemma-4-26b-a4b-it-4bit"), + model_card=None, + ) + + assert check_every == 10 From 2cb8e00ee9f79a36d5158f75bafabdebfa587966 Mon Sep 17 00:00:00 2001 From: Thomas Tupper Date: Wed, 8 Apr 2026 21:49:30 -0500 Subject: [PATCH 32/32] Deduplicate env precedence and KV backend validation --- src/exo/api/main.py | 24 +++++++++---------- src/exo/shared/constants.py | 17 +++++++++++++ src/exo/worker/engines/mlx/auto_parallel.py | 12 +++------- src/exo/worker/engines/mlx/cache.py | 10 ++------ src/exo/worker/engines/mlx/constants.py | 17 +++++++++---- .../worker/engines/mlx/generator/generate.py | 16 ++++--------- src/exo/worker/engines/mlx/utils_mlx.py | 13 ++++------ 7 files changed, 56 insertions(+), 53 deletions(-) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index 95ba2b496..7e68b1f7f 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -141,6 +141,7 @@ EXO_IMAGE_TRANSPORT_DEBUG, EXO_MAX_CHUNK_SIZE, EXO_TRACING_CACHE_DIR, + preferred_env_value, ) from exo.shared.election import ElectionMessage from exo.shared.logging import InterceptLogger @@ -200,7 +201,10 @@ from exo.utils.disk_event_log import DiskEventLog from exo.utils.power_sampler import PowerSampler from exo.utils.task_group import TaskGroup -from exo.worker.engines.mlx.constants import DEFAULT_KV_CACHE_BACKEND, KVCacheBackend +from exo.worker.engines.mlx.constants import ( + DEFAULT_KV_CACHE_BACKEND, + VALID_KV_CACHE_BACKENDS, +) if TYPE_CHECKING: from exo.store.config import ExoConfig @@ -2566,21 +2570,15 @@ async def complete_onboarding(self) -> JSONResponse: def _effective_kv_cache_backend(self) -> str: """Return the effective KV backend after SKULK/EXO env precedence is applied.""" - if "SKULK_KV_CACHE_BACKEND" in os.environ: - configured_backend = os.environ["SKULK_KV_CACHE_BACKEND"] - else: - configured_backend = os.environ.get("EXO_KV_CACHE_BACKEND", "") + configured_backend = preferred_env_value( + "SKULK_KV_CACHE_BACKEND", + "EXO_KV_CACHE_BACKEND", + "", + ) if not configured_backend: return DEFAULT_KV_CACHE_BACKEND - valid_backends: tuple[KVCacheBackend, ...] = ( - "default", - "mlx_quantized", - "turboquant", - "turboquant_adaptive", - "optiq", - ) - if configured_backend not in valid_backends: + if configured_backend not in VALID_KV_CACHE_BACKENDS: return DEFAULT_KV_CACHE_BACKEND return configured_backend diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py index f1794ae9d..10d44ca28 100644 --- a/src/exo/shared/constants.py +++ b/src/exo/shared/constants.py @@ -14,6 +14,23 @@ def _env(skulk_key: str, exo_key: str, default: str | None = None) -> str | None return os.environ.get(skulk_key, os.environ.get(exo_key, default)) +def preferred_env_value( + skulk_key: str, + exo_key: str, + default: str | None = None, +) -> str | None: + """Return the SKULK value when its key exists, else the legacy EXO value. + + This variant preserves explicitly blank `SKULK_*` values instead of + treating them as absent, which is important for debug/config toggles where + a blank value should disable the feature instead of falling back to a + legacy `EXO_*` env var. + """ + if skulk_key in os.environ: + return os.environ[skulk_key] + return os.environ.get(exo_key, default) + + _SKULK_HOME_ENV = _env("SKULK_HOME", "EXO_HOME") diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index dd9079d63..6094bb2d8 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -62,6 +62,7 @@ from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel +from exo.shared.constants import preferred_env_value from exo.shared.types.worker.shards import PipelineShardMetadata from exo.worker.runner.bootstrap import logger @@ -75,22 +76,15 @@ _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = [] -def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: - """Return the SKULK env value when present, else the legacy EXO value.""" - if skulk_key in os.environ: - return os.environ[skulk_key] - return os.environ.get(exo_key) - - def _mlx_hang_debug_enabled() -> bool: - value = _preferred_env_value("SKULK_MLX_HANG_DEBUG", "EXO_MLX_HANG_DEBUG") + value = preferred_env_value("SKULK_MLX_HANG_DEBUG", "EXO_MLX_HANG_DEBUG") if value is None: return False return value.strip().lower() not in {"", "0", "false", "no", "off"} def _mlx_hang_debug_interval_seconds() -> float: - raw = _preferred_env_value( + raw = preferred_env_value( "SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS", "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS", ) diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 2911f81df..2e4c4859d 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -26,6 +26,7 @@ TURBOQUANT_FP16_LAYERS, TURBOQUANT_K_BITS, TURBOQUANT_V_BITS, + VALID_KV_CACHE_BACKENDS, KVCacheBackend, ) from exo.worker.engines.mlx.turboquant import ( @@ -588,14 +589,7 @@ def _is_power_of_two(n: int) -> bool: def get_kv_cache_backend() -> KVCacheBackend: backend = KV_CACHE_BACKEND - valid_backends: tuple[KVCacheBackend, ...] = ( - "default", - "mlx_quantized", - "turboquant", - "turboquant_adaptive", - "optiq", - ) - if backend not in valid_backends: + if backend not in VALID_KV_CACHE_BACKENDS: logger.warning( f"Unknown EXO_KV_CACHE_BACKEND={backend!r}; falling back to {DEFAULT_KV_CACHE_BACKEND!r}" ) diff --git a/src/exo/worker/engines/mlx/constants.py b/src/exo/worker/engines/mlx/constants.py index 3e7388d28..87d039542 100644 --- a/src/exo/worker/engines/mlx/constants.py +++ b/src/exo/worker/engines/mlx/constants.py @@ -1,6 +1,8 @@ import os from typing import Literal, cast +from exo.shared.constants import preferred_env_value + # TODO: Do we want so many constants? # I think we want a lot of these as parameters? @@ -23,10 +25,17 @@ "optiq", ] DEFAULT_KV_CACHE_BACKEND: KVCacheBackend = "default" -_kv_cache_backend_value = ( - os.environ["SKULK_KV_CACHE_BACKEND"] - if "SKULK_KV_CACHE_BACKEND" in os.environ - else os.environ.get("EXO_KV_CACHE_BACKEND", DEFAULT_KV_CACHE_BACKEND) +VALID_KV_CACHE_BACKENDS: tuple[KVCacheBackend, ...] = ( + "default", + "mlx_quantized", + "turboquant", + "turboquant_adaptive", + "optiq", +) +_kv_cache_backend_value = preferred_env_value( + "SKULK_KV_CACHE_BACKEND", + "EXO_KV_CACHE_BACKEND", + DEFAULT_KV_CACHE_BACKEND, ) KV_CACHE_BACKEND: KVCacheBackend = cast( KVCacheBackend, diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 982980b8b..5e40a7124 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -28,6 +28,7 @@ TopLogprobItem, Usage, ) +from exo.shared.constants import preferred_env_value from exo.shared.models.model_cards import ModelCard from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory @@ -81,16 +82,9 @@ _MIN_CANCEL_CHECK_INTERVAL = 10 -def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: - """Return the SKULK env value when present, else the legacy EXO value.""" - if skulk_key in os.environ: - return os.environ[skulk_key] - return os.environ.get(exo_key) - - def _mlx_hang_debug_enabled() -> bool: """Return whether verbose warmup/prefill hang diagnostics are enabled.""" - value = _preferred_env_value("SKULK_MLX_HANG_DEBUG", "EXO_MLX_HANG_DEBUG") + value = preferred_env_value("SKULK_MLX_HANG_DEBUG", "EXO_MLX_HANG_DEBUG") if value is None: return False return value.strip().lower() not in {"", "0", "false", "no", "off"} @@ -98,7 +92,7 @@ def _mlx_hang_debug_enabled() -> bool: def _mlx_hang_debug_interval_seconds() -> float: """Return the periodic interval for hang-debug watchdog logs.""" - raw = _preferred_env_value( + raw = preferred_env_value( "SKULK_MLX_HANG_DEBUG_INTERVAL_SECONDS", "EXO_MLX_HANG_DEBUG_INTERVAL_SECONDS", ) @@ -111,7 +105,7 @@ def _mlx_hang_debug_interval_seconds() -> float: def _warmup_repeat_count() -> int: """Return the neutral warmup token repeat count used for debugging.""" - raw = _preferred_env_value( + raw = preferred_env_value( "SKULK_DEBUG_WARMUP_REPEAT_COUNT", "EXO_DEBUG_WARMUP_REPEAT_COUNT", ) @@ -143,7 +137,7 @@ def _warmup_instructions(group: mx.distributed.Group | None) -> str | None: """Return optional warmup instructions for prompt-shape debugging.""" if _is_distributed_warmup(group): return None - raw = _preferred_env_value( + raw = preferred_env_value( "SKULK_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", "EXO_DEBUG_WARMUP_INCLUDE_INSTRUCTIONS", ) diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 520a292fe..f1a411e85 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -27,6 +27,7 @@ from mlx_lm.models.deepseek_v3 import DeepseekV3Model from mlx_lm.tokenizer_utils import TokenizerWrapper +from exo.shared.constants import preferred_env_value from exo.shared.models.capabilities import resolve_model_capability_profile from exo.shared.models.model_cards import ( ModelCard, @@ -79,16 +80,12 @@ Group = mx.distributed.Group -def _preferred_env_value(skulk_key: str, exo_key: str) -> str | None: - """Return the SKULK env value when present, else the legacy EXO value.""" - if skulk_key in os.environ: - return os.environ[skulk_key] - return os.environ.get(exo_key) - - def _request_shape_debug_enabled() -> bool: """Return whether request-shape tracing is enabled for prompt debugging.""" - value = _preferred_env_value("SKULK_TRACE_REQUEST_SHAPES", "EXO_TRACE_REQUEST_SHAPES") + value = preferred_env_value( + "SKULK_TRACE_REQUEST_SHAPES", + "EXO_TRACE_REQUEST_SHAPES", + ) if value is None: return False return value.strip().lower() not in {"", "0", "false", "no", "off"}