Fix pipeline prefill short-prompt wedge + dashboard kv-backend display

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,
                 },

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,

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

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 = "<no Python frame available>"
+            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,36 +278,57 @@ 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:
                 _pending_prefill_sends.append(
                     (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)