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

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,6 +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,
+    VALID_KV_CACHE_BACKENDS,
+)
 
 if TYPE_CHECKING:
     from exo.store.config import ExoConfig
@@ -2563,6 +2568,20 @@ 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."""
+        configured_backend = preferred_env_value(
+            "SKULK_KV_CACHE_BACKEND",
+            "EXO_KV_CACHE_BACKEND",
+            "",
+        )
+        if not configured_backend:
+            return DEFAULT_KV_CACHE_BACKEND
+
+        if configured_backend not in VALID_KV_CACHE_BACKENDS:
+            return DEFAULT_KV_CACHE_BACKEND
+        return configured_backend
+
     async def get_config(self) -> JSONResponse:
         if not self._config_path.exists():
             return JSONResponse(
@@ -2571,10 +2590,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,9 +2607,7 @@ async def get_config(self) -> JSONResponse:
                 "configPath": str(self._config_path),
                 "fileExists": True,
                 "effective": {
-                    "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,
                 },
             }

src/exo/api/tests/test_config_api.py

@@ -0,0 +1,93 @@
+# 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
+
+
+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"
+
+
+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"

src/exo/shared/apply.py

@@ -213,9 +213,39 @@ 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
     }
+    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
     }
@@ -257,9 +287,13 @@ 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,
+            "node_identities": node_identities,
             "node_memory": node_memory,
             "node_disk": node_disk,
             "node_system": node_system,

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")
 
 

src/exo/shared/tests/test_apply/test_apply_node_timed_out.py

@@ -0,0 +1,138 @@
+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.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
+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(),
+        },
+        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)
+
+    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
+
+    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