← back to Exo
Store custom model cards in State (#2024)
edef8004f8c8751ed62ae74ceb5ae7676164dcbf · 2026-05-07 01:06:39 -0700 · Alex Cheema
## Why
Workers currently update their custom model-card cache by reacting to
`CustomModelCardAdded` / `CustomModelCardDeleted` events directly. That
is another snapshot footgun: a worker restored from State may never see
the historical add/delete event, so the durable State must include the
desired custom-card set.
## How
- Add `State.custom_model_cards`, keyed by `ModelId`.
- Reduce `CustomModelCardAdded` into State.
- Reduce `CustomModelCardDeleted` into State.
- Add focused reducer tests for add and delete.
This PR only makes custom cards durable in State. A follow-up PR will
make workers reconcile their on-disk custom-card cache from this state
instead of relying on those events directly.
## Tests
- `uv run pytest
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py
src/exo/shared/tests/test_state_serialization.py`
- `uv run pytest`
- `uv run ruff check src/exo/shared/types/state.py
src/exo/shared/apply.py
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py`
- `uv run basedpyright`
- `nix fmt`
Files touched
M src/exo/api/main.pyM src/exo/download/coordinator.pyM src/exo/download/impl_shard_downloader.pyM src/exo/shared/apply.pyM src/exo/shared/models/model_cards.pyA src/exo/shared/tests/test_apply/test_apply_custom_model_cards.pyM src/exo/shared/types/state.pyM src/exo/shared/types/text_generation.pyM src/exo/worker/main.pyM src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
Diff
commit edef8004f8c8751ed62ae74ceb5ae7676164dcbf
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Thu May 7 01:06:39 2026 -0700
Store custom model cards in State (#2024)
## Why
Workers currently update their custom model-card cache by reacting to
`CustomModelCardAdded` / `CustomModelCardDeleted` events directly. That
is another snapshot footgun: a worker restored from State may never see
the historical add/delete event, so the durable State must include the
desired custom-card set.
## How
- Add `State.custom_model_cards`, keyed by `ModelId`.
- Reduce `CustomModelCardAdded` into State.
- Reduce `CustomModelCardDeleted` into State.
- Add focused reducer tests for add and delete.
This PR only makes custom cards durable in State. A follow-up PR will
make workers reconcile their on-disk custom-card cache from this state
instead of relying on those events directly.
## Tests
- `uv run pytest
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py
src/exo/shared/tests/test_state_serialization.py`
- `uv run pytest`
- `uv run ruff check src/exo/shared/types/state.py
src/exo/shared/apply.py
src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py`
- `uv run basedpyright`
- `nix fmt`
---
src/exo/api/main.py | 23 ++---
src/exo/download/coordinator.py | 5 +-
src/exo/download/impl_shard_downloader.py | 4 +-
src/exo/shared/apply.py | 28 +++++-
src/exo/shared/models/model_cards.py | 106 +++++++++++----------
.../test_apply/test_apply_custom_model_cards.py | 44 +++++++++
src/exo/shared/types/state.py | 6 +-
src/exo/shared/types/text_generation.py | 4 +-
src/exo/worker/main.py | 23 +++--
.../tests/unittests/test_mlx/test_tokenizers.py | 8 +-
10 files changed, 163 insertions(+), 88 deletions(-)
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index b4c78932..4fb6d2d3 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -133,12 +133,10 @@ from exo.shared.constants import (
)
from exo.shared.election import ElectionMessage
from exo.shared.logging import InterceptLogger
+from exo.shared.models import model_cards
from exo.shared.models.model_cards import (
ModelCard,
ModelId,
- add_to_card_cache,
- get_card,
- get_model_cards,
)
from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
from exo.shared.types.chunks import (
@@ -1635,17 +1633,16 @@ class API:
async def ollama_tags(self) -> OllamaTagsResponse:
"""Returns list of models in Ollama tags format. We return the downloaded ones only."""
- def none_if_empty(value: str) -> str | None:
- return value or None
-
- downloaded_model_ids: set[str] = set()
+ downloaded_model_ids: set[ModelId] = set()
for node_downloads in self.state.downloads.values():
for dl in node_downloads:
if isinstance(dl, DownloadCompleted):
downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
cards = [
- c for c in await get_model_cards() if c.model_id in downloaded_model_ids
+ c
+ for c in await model_cards.card_cache.list_all()
+ if c.model_id in downloaded_model_ids
]
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -1658,8 +1655,8 @@ class API:
size=card.storage_size.in_bytes,
digest="sha256:000000000000",
details=OllamaModelDetails(
- family=none_if_empty(card.family),
- quantization_level=none_if_empty(card.quantization),
+ family=card.family or None,
+ quantization_level=card.quantization or None,
),
)
for card in cards
@@ -1722,7 +1719,7 @@ class API:
async def get_models(self, status: str | None = Query(default=None)) -> ModelList:
"""Returns list of available models, optionally filtered by being downloaded."""
- cards = await get_model_cards()
+ cards = await model_cards.card_cache.list_all()
if status == "downloaded":
downloaded_model_ids: set[str] = set()
@@ -1773,7 +1770,7 @@ class API:
# Immediately update the local cache so the subsequent GET /models
# returns the new model without waiting for the event round-trip.
- add_to_card_cache(card)
+ model_cards.card_cache.cc[card.model_id] = card
return ModelListModel(
id=card.model_id,
@@ -1789,7 +1786,7 @@ class API:
async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
"""Delete a user-added custom model card and sync deletion across the cluster."""
- card = get_card(model_id)
+ card = model_cards.card_cache.get(model_id)
if card is None or not card.is_custom:
raise HTTPException(status_code=404, detail="Custom model card not found")
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py
index 2e9de3a8..de9c4722 100644
--- a/src/exo/download/coordinator.py
+++ b/src/exo/download/coordinator.py
@@ -16,7 +16,8 @@ from exo.download.download_utils import (
)
from exo.download.shard_downloader import ShardDownloader
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
-from exo.shared.models.model_cards import ModelId, get_model_cards
+from exo.shared.models import model_cards
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.commands import (
CancelDownload,
DeleteDownload,
@@ -422,7 +423,7 @@ class DownloadCoordinator:
)
# Scan read-only directories for pre-downloaded models
if EXO_MODELS_READ_ONLY_DIRS:
- for card in await get_model_cards():
+ for card in await model_cards.card_cache.list_all():
mid = card.model_id
if mid in self.active_downloads:
continue
diff --git a/src/exo/download/impl_shard_downloader.py b/src/exo/download/impl_shard_downloader.py
index 14ec83b6..ee85945b 100644
--- a/src/exo/download/impl_shard_downloader.py
+++ b/src/exo/download/impl_shard_downloader.py
@@ -11,11 +11,11 @@ from exo.download.download_utils import (
download_shard,
)
from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models import model_cards
from exo.shared.models.model_cards import (
ModelCard,
ModelId,
ModelTask,
- get_model_cards,
)
from exo.shared.types.memory import Memory
from exo.shared.types.worker.shards import (
@@ -258,7 +258,7 @@ class ResumableShardDownloader(ShardDownloader):
tasks = [
create_task(download_with_semaphore(model_card))
- for model_card in await get_model_cards()
+ for model_card in await model_cards.card_cache.list_all()
]
for task in asyncio.as_completed(tasks):
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index 450c35e7..d5a9d383 100644
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -4,7 +4,8 @@ from datetime import datetime
from loguru import logger
-from exo.shared.types.common import NodeId
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.types.common import ModelId, NodeId
from exo.shared.types.events import (
ChunkGenerated,
CustomModelCardAdded,
@@ -87,10 +88,12 @@ def event_apply(event: Event, state: State) -> State:
| InputChunkReceived()
| TracesCollected()
| TracesMerged()
- | CustomModelCardAdded()
- | CustomModelCardDeleted()
): # Pass-through events that don't modify state
return state
+ case CustomModelCardAdded():
+ return apply_custom_model_card_added(event, state)
+ case CustomModelCardDeleted():
+ return apply_custom_model_card_deleted(event, state)
case InstanceCreated():
return apply_instance_created(event, state)
case InstanceDeleted():
@@ -472,3 +475,22 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
topology.remove_connection(event.conn)
# TODO: Clean up removing the reverse connection
return state.model_copy(update={"topology": topology})
+
+
+def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
+ new_cards: Mapping[ModelId, ModelCard] = {
+ **state.custom_model_cards,
+ event.model_card.model_id: event.model_card,
+ }
+ return state.model_copy(update={"custom_model_cards": new_cards})
+
+
+def apply_custom_model_card_deleted(
+ event: CustomModelCardDeleted, state: State
+) -> State:
+ new_cards: Mapping[ModelId, ModelCard] = {
+ model_id: card
+ for model_id, card in state.custom_model_cards.items()
+ if model_id != event.model_id
+ }
+ return state.model_copy(update={"custom_model_cards": new_cards})
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index 0d9acc1d..0d1648a7 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -39,7 +39,57 @@ _BUILTIN_CARD_DIRS = [
Path(RESOURCES_DIR) / "image_model_cards",
]
-_card_cache: dict[ModelId, "ModelCard"] = {}
+
+class _CardCache:
+ def __init__(self):
+ self.cc: dict[ModelId, "ModelCard"] = {}
+
+ def get(self, model_id: ModelId) -> "ModelCard | None":
+ return self.cc.get(model_id)
+
+ async def save(self, card: "ModelCard"):
+ self.cc[card.model_id] = card
+ try:
+ await card.save_to_custom_dir()
+ except OSError as e:
+ logger.warning(f"failed to save custom model card ({e.strerror})")
+
+ async def pop(self, model_id: ModelId) -> "ModelCard | None":
+ """Delete a user-added custom model card. Returns True if deleted."""
+ card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
+ try:
+ if await card_path.exists():
+ await card_path.unlink()
+ return self.cc.pop(model_id, None)
+ except OSError as e:
+ logger.warning(f"failed to delete custom model card ({e.strerror})")
+
+ async def list_all(self) -> list["ModelCard"]:
+ if len(self.cc) == 0:
+ await self.refresh()
+ if EXO_ENABLE_IMAGE_MODELS:
+ return list(self.cc.values())
+ return [c for c in self.cc.values() if not _is_image_card(c)]
+
+ async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
+ """Load all TOML model cards from a directory into the cache."""
+ async for toml_file in directory.rglob("*.toml"):
+ try:
+ card = await ModelCard.load_from_path(toml_file)
+ if is_custom:
+ card = card.model_copy(update={"is_custom": True})
+ if self.get(card.model_id) is None:
+ self.cc[card.model_id] = card
+ except (ValidationError, TOMLKitError):
+ pass
+
+ async def refresh(self) -> None:
+ for path in _BUILTIN_CARD_DIRS:
+ await self._load_cards_from_dir(path, is_custom=False)
+ await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
+
+
+card_cache = _CardCache()
def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -59,42 +109,10 @@ def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
return None
-async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
- """Load all TOML model cards from a directory into the cache."""
- async for toml_file in directory.rglob("*.toml"):
- try:
- card = await ModelCard.load_from_path(toml_file)
- if is_custom:
- card = card.model_copy(update={"is_custom": True})
- if card.model_id not in _card_cache:
- _card_cache[card.model_id] = card
- except (ValidationError, TOMLKitError):
- pass
-
-
-async def _refresh_card_cache() -> None:
- for path in _BUILTIN_CARD_DIRS:
- await _load_cards_from_dir(path, is_custom=False)
- await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
-
-
def _is_image_card(card: "ModelCard") -> bool:
return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
-def get_card(model_id: ModelId) -> "ModelCard | None":
- """Look up a single model card from the cache by ID."""
- return _card_cache.get(model_id)
-
-
-async def get_model_cards() -> list["ModelCard"]:
- if len(_card_cache) == 0:
- await _refresh_card_cache()
- if EXO_ENABLE_IMAGE_MODELS:
- return list(_card_cache.values())
- return [c for c in _card_cache.values() if not _is_image_card(c)]
-
-
class ModelTask(str, Enum):
TextGeneration = "TextGeneration"
TextToImage = "TextToImage"
@@ -196,14 +214,13 @@ class ModelCard(FrozenModel):
# Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
@staticmethod
async def load(model_id: ModelId) -> "ModelCard":
- if model_id not in _card_cache:
- await _refresh_card_cache()
- if (mc := _card_cache.get(model_id)) is not None:
+ if card_cache.get(model_id) is None:
+ await card_cache.refresh()
+ if (mc := card_cache.get(model_id)) is not None:
return mc
mc = await ModelCard.fetch_from_hf(model_id)
await mc.save_to_custom_dir()
- _card_cache[model_id] = mc
return mc
@staticmethod
@@ -233,21 +250,6 @@ class ModelCard(FrozenModel):
)
-def add_to_card_cache(card: "ModelCard") -> None:
- """Add or update a model card in the in-memory cache."""
- _card_cache[card.model_id] = card
-
-
-async def delete_custom_card(model_id: ModelId) -> bool:
- """Delete a user-added custom model card. Returns True if deleted."""
- card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
- if await card_path.exists():
- await card_path.unlink()
- _card_cache.pop(model_id, None)
- return True
- return False
-
-
class ConfigData(BaseModel):
model_config = {"extra": "ignore"} # Allow unknown fields
diff --git a/src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py b/src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py
new file mode 100644
index 00000000..b5b3066c
--- /dev/null
+++ b/src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py
@@ -0,0 +1,44 @@
+from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelCard, ModelTask
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import (
+ CustomModelCardAdded,
+ CustomModelCardDeleted,
+ IndexedEvent,
+)
+from exo.shared.types.memory import Memory
+from exo.shared.types.state import State
+
+
+def _model_card(model_id: ModelId) -> ModelCard:
+ return ModelCard(
+ model_id=model_id,
+ n_layers=1,
+ storage_size=Memory.from_bytes(1),
+ hidden_size=1,
+ supports_tensor=True,
+ tasks=[ModelTask.TextGeneration],
+ )
+
+
+def test_custom_model_card_added_is_reduced_into_state() -> None:
+ card = _model_card(ModelId("custom/model"))
+
+ state = apply(
+ State(),
+ IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
+ )
+
+ assert state.custom_model_cards == {card.model_id: card}
+
+
+def test_custom_model_card_deleted_removes_card_from_state() -> None:
+ card = _model_card(ModelId("custom/model"))
+ state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
+
+ state = apply(
+ state,
+ IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
+ )
+
+ assert state.custom_model_cards == {}
diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py
index 6c976984..6a60aa08 100644
--- a/src/exo/shared/types/state.py
+++ b/src/exo/shared/types/state.py
@@ -5,8 +5,9 @@ from typing import Any, cast
from pydantic import ConfigDict, Field, field_serializer, field_validator
from pydantic.alias_generators import to_camel
+from exo.shared.models.model_cards import ModelCard
from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import ModelId, NodeId
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
from exo.shared.types.profiling import (
DiskUsage,
@@ -65,6 +66,9 @@ class State(FrozenModel):
instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
prefill_server_ports: Mapping[RunnerId, int] = {}
+ # User-added model cards. Workers can reconcile their on-disk custom card cache
+ custom_model_cards: Mapping[ModelId, ModelCard] = {}
+
@field_serializer("topology", mode="plain")
def _encode_topology(self, value: Topology) -> TopologySnapshot:
return value.to_snapshot()
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
index ccb2512a..29228c36 100644
--- a/src/exo/shared/types/text_generation.py
+++ b/src/exo/shared/types/text_generation.py
@@ -135,9 +135,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
prefill_endpoint: str | None = None
def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
- from exo.shared.models.model_cards import get_card
+ from exo.shared.models import model_cards
- card = get_card(self.model)
+ card = model_cards.card_cache.get(self.model)
if card is None:
return self
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index b35f946a..9d33cc23 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -10,7 +10,7 @@ from exo.api.types import ImageEditsTaskParams
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
from exo.shared.apply import apply
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
-from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
+from exo.shared.models.model_cards import ModelId, card_cache
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.commands import (
DeleteInstance,
@@ -20,8 +20,6 @@ from exo.shared.types.commands import (
)
from exo.shared.types.common import CommandId, NodeId, SystemId
from exo.shared.types.events import (
- CustomModelCardAdded,
- CustomModelCardDeleted,
Event,
IndexedEvent,
InputChunkReceived,
@@ -110,6 +108,8 @@ class Worker:
tg.start_soon(self.plan_step)
tg.start_soon(self._event_applier)
tg.start_soon(self._poll_connection_updates)
+ tg.start_soon(self._reconcile_custom_cards)
+
finally:
# Actual shutdown code - waits for all tasks to complete before executing.
logger.info("Stopping Worker")
@@ -151,7 +151,6 @@ class Worker:
self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
event.chunk
)
-
if (
len(self.input_chunk_buffer[cmd_id])
== self.input_chunk_counts[cmd_id]
@@ -172,12 +171,18 @@ class Worker:
)
] = img
- if isinstance(event, CustomModelCardAdded):
- await event.model_card.save_to_custom_dir()
- add_to_card_cache(event.model_card)
+ async def _reconcile_custom_cards(self) -> None:
+ while True:
+ await anyio.sleep(1)
+ target = dict(self.state.custom_model_cards)
+ for model_id, card in target.items():
+ if card_cache.get(model_id) == card:
+ continue
+ await card_cache.save(card)
- if isinstance(event, CustomModelCardDeleted):
- await delete_custom_card(event.model_id)
+ for card in await card_cache.list_all():
+ if card.model_id not in target:
+ await card_cache.pop(card.model_id)
async def plan_step(self):
while True:
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
index 2f4ca7e6..9572e713 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
@@ -16,7 +16,7 @@ from exo.download.download_utils import (
fetch_file_list_with_cache,
resolve_model_dir,
)
-from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
+from exo.shared.models.model_cards import ModelCard, ModelId, card_cache
from exo.worker.engines.mlx.utils_mlx import (
get_eos_token_ids_for_model,
load_tokenizer_for_model_id,
@@ -76,7 +76,7 @@ def get_test_models() -> list[ModelCard]:
"""Get a representative sample of models to test."""
# Pick one model from each family to test
families: dict[str, ModelCard] = {}
- for card in asyncio.run(get_model_cards()):
+ for card in asyncio.run(card_cache.list_all()):
# Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
parts = card.model_id.short().split("-")
family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
@@ -298,7 +298,7 @@ async def test_tokenizer_special_tokens(model_card: ModelCard) -> None:
async def test_kimi_tokenizer_specifically():
"""Test Kimi tokenizer with its specific patches and quirks."""
kimi_models = [
- card for card in await get_model_cards() if "kimi" in card.model_id.lower()
+ card for card in await card_cache.list_all() if "kimi" in card.model_id.lower()
]
if not kimi_models:
@@ -350,7 +350,7 @@ async def test_glm_tokenizer_specifically():
glm_model_cards = [
card
- for card in await get_model_cards()
+ for card in await card_cache.list_all()
if contains(card, "glm")
and not contains(card, "-5")
and not contains(card, "4.7")
← a0c00f9d fix(placement): gate RDMA on nodeRdmaCtl.enabled at both end
·
back to Exo
·
Use time-weighted power sampling (#2038) 414132ae →