electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Sync custom model cards across nodes (#1768)

Browse Source 
## Motivation

Custom model cards were only saved locally on the node that handled the
API request.

## Changes

- Added AddCustomModelCard and DeleteCustomModelCard commands
- Added CustomModelCardAdded and CustomModelCardDeleted events
- Added custom_model_cards field to cluster State
- Master handles new commands by emitting corresponding events
- Workers persist model cards to disk and update the in-memory cache on
event receipt
- Separated fetch_from_hf (pure fetch) from disk persistence (now
handled by event-sourcing layer)
- Exposed add_to_card_cache() helper for the worker to update the cache

## Why It Works

- Follows the existing event-sourcing pattern: API → Command → Master →
Event → all Workers
- Every node applies the same events, so custom model cards are
consistent across the cluster

## Test Plan

### Manual Testing

Add/delete a custom model via the API on one node, verify it
appears/disappears on all nodes

### Automated Testing

Existing tests cover apply() logic; new event types follow the same
discriminated-union pattern
This commit is contained in:
ciaranbor
2026-03-24 12:51:36 +00:00
committed by GitHub
parent e06e70a835
commit 49951e1b1a
7 changed files with 112 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/api/main.py
+23 -7
View File
@@ -129,9 +129,8 @@ from exo.shared.logging import InterceptLogger
from exo.shared.models.model_cards import (
ModelCard,
ModelId,
delete_custom_card,
get_card,
get_model_cards,
is_custom_card,
)
from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
from exo.shared.types.chunks import (
@@ -143,8 +142,10 @@ from exo.shared.types.chunks import (
ToolCallChunk,
)
from exo.shared.types.commands import (
AddCustomModelCard,
Command,
CreateInstance,
DeleteCustomModelCard,
DeleteDownload,
DeleteInstance,
DownloadCommand,
@@ -1558,7 +1559,7 @@ class API:
storage_size_megabytes=card.storage_size.in_mb,
supports_tensor=card.supports_tensor,
tasks=[task.value for task in card.tasks],
is_custom=is_custom_card(card.model_id),
is_custom=card.is_custom,
family=card.family,
quantization=card.quantization,
base_model=card.base_model,
@@ -1569,7 +1570,7 @@ class API:
)
async def add_custom_model(self, payload: AddCustomModelParams) -> ModelListModel:
"""Fetch a model from HuggingFace and save as a custom model card."""
"""Fetch a model from HuggingFace and save as a custom model card, then sync across the cluster."""
try:
card = await ModelCard.fetch_from_hf(payload.model_id)
except Exception as exc:
@@ -1577,6 +1578,13 @@ class API:
status_code=400, detail=f"Failed to fetch model: {exc}"
) from exc
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=AddCustomModelCard(model_card=card),
)
)
return ModelListModel(
id=card.model_id,
hugging_face_id=card.model_id,
@@ -1590,10 +1598,18 @@ class API:
)
async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
"""Delete a user-added custom model card."""
deleted = await delete_custom_card(model_id)
if not deleted:
"""Delete a user-added custom model card and sync deletion across the cluster."""
card = get_card(model_id)
if card is None or not card.is_custom:
raise HTTPException(status_code=404, detail="Custom model card not found")
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=DeleteCustomModelCard(model_id=model_id),
)
)
return JSONResponse(
{"message": "Model card deleted", "model_id": str(model_id)}
)
src/exo/master/main.py
+12
View File
@@ -13,7 +13,9 @@ from exo.master.placement import (
from exo.shared.apply import apply
from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
from exo.shared.types.commands import (
AddCustomModelCard,
CreateInstance,
DeleteCustomModelCard,
DeleteInstance,
ForwarderCommand,
ForwarderDownloadCommand,
@@ -29,6 +31,8 @@ from exo.shared.types.commands import (
)
from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
from exo.shared.types.events import (
CustomModelCardAdded,
CustomModelCardDeleted,
Event,
GlobalForwarderEvent,
IndexedEvent,
@@ -345,6 +349,14 @@ class Master:
f"Finished command {command.finished_command_id} finished"
)
case AddCustomModelCard():
generated_events.append(
CustomModelCardAdded(model_card=command.model_card)
)
case DeleteCustomModelCard():
generated_events.append(
CustomModelCardDeleted(model_id=command.model_id)
)
case RequestEventLog():
# We should just be able to send everything, since other buffers will ignore old messages
# rate limit to 1000 at a time
src/exo/shared/apply.py
+4
View File
@@ -7,6 +7,8 @@ from loguru import logger
from exo.shared.types.common import NodeId
from exo.shared.types.events import (
ChunkGenerated,
CustomModelCardAdded,
CustomModelCardDeleted,
Event,
IndexedEvent,
InputChunkReceived,
@@ -65,6 +67,8 @@ def event_apply(event: Event, state: State) -> State:
| InputChunkReceived()
| TracesCollected()
| TracesMerged()
| CustomModelCardAdded()
| CustomModelCardDeleted()
): # Pass-through events that don't modify state
return state
case InstanceCreated():
src/exo/shared/models/model_cards.py
+41 -28
View File
@@ -30,30 +30,42 @@ from exo.utils.pydantic_ext import CamelCaseModel
# kinda ugly...
# TODO: load search path from config.toml
_custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
CARD_SEARCH_PATH = [
_BUILTIN_CARD_DIRS = [
Path(RESOURCES_DIR) / "inference_model_cards",
Path(RESOURCES_DIR) / "image_model_cards",
_custom_cards_dir,
]
_card_cache: dict[ModelId, "ModelCard"] = {}
async def _refresh_card_cache():
for path in CARD_SEARCH_PATH:
async for toml_file in path.rglob("*.toml"):
try:
card = await ModelCard.load_from_path(toml_file)
if card.model_id not in _card_cache:
_card_cache[card.model_id] = card
except (ValidationError, TOMLKitError):
pass
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()
@@ -92,6 +104,7 @@ class ModelCard(CamelCaseModel):
capabilities: list[str] = []
uses_cfg: bool = False
trust_remote_code: bool = True
is_custom: bool = False
@field_validator("tasks", mode="before")
@classmethod
@@ -100,7 +113,7 @@ class ModelCard(CamelCaseModel):
async def save(self, path: Path) -> None:
async with await open_file(path, "w") as f:
py = self.model_dump(exclude_none=True)
py = self.model_dump(exclude_none=True, exclude={"is_custom"})
data = tomlkit.dumps(py) # pyright: ignore[reportUnknownMemberType]
await f.write(data)
@@ -122,17 +135,24 @@ class ModelCard(CamelCaseModel):
if (mc := _card_cache.get(model_id)) is not None:
return mc
return await ModelCard.fetch_from_hf(model_id)
mc = await ModelCard.fetch_from_hf(model_id)
await mc.save_to_custom_dir()
_card_cache[model_id] = mc
return mc
@staticmethod
async def fetch_from_hf(model_id: ModelId) -> "ModelCard":
"""Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
"""Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta.
This is a pure fetch — it does NOT save to disk or update the cache.
Persistence is handled by the event-sourcing layer (worker event handler).
"""
# TODO: failure if files do not exist
config_data = await fetch_config_data(model_id)
num_layers = config_data.layer_count
mem_size_bytes = await fetch_safetensors_size(model_id)
mc = ModelCard(
return ModelCard(
model_id=ModelId(model_id),
storage_size=mem_size_bytes,
n_layers=num_layers,
@@ -141,10 +161,13 @@ class ModelCard(CamelCaseModel):
num_key_value_heads=config_data.num_key_value_heads,
tasks=[ModelTask.TextGeneration],
trust_remote_code=False,
is_custom=True,
)
await mc.save_to_custom_dir()
_card_cache[model_id] = mc
return mc
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:
@@ -157,16 +180,6 @@ async def delete_custom_card(model_id: ModelId) -> bool:
return False
def is_custom_card(model_id: ModelId) -> bool:
"""Check if a model card exists in the custom cards directory."""
import os
card_path = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) / (
ModelId(model_id).normalize() + ".toml"
)
return os.path.isfile(str(card_path))
class ConfigData(BaseModel):
model_config = {"extra": "ignore"} # Allow unknown fields
src/exo/shared/types/commands.py
+10
View File
@@ -81,6 +81,14 @@ class CancelDownload(BaseCommand):
model_id: ModelId
class AddCustomModelCard(BaseCommand):
model_card: ModelCard
class DeleteCustomModelCard(BaseCommand):
model_id: ModelId
DownloadCommand = StartDownload | DeleteDownload | CancelDownload
@@ -96,6 +104,8 @@ Command = (
| TaskCancelled
| TaskFinished
| SendInputChunk
| AddCustomModelCard
| DeleteCustomModelCard
)
src/exo/shared/types/events.py
+12 -1
View File
@@ -3,9 +3,10 @@ from typing import final
from pydantic import Field
from exo.shared.models.model_cards import ModelCard
from exo.shared.topology import Connection
from exo.shared.types.chunks import GenerationChunk, InputImageChunk
from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import DownloadProgress
from exo.shared.types.worker.instances import Instance, InstanceId
@@ -106,6 +107,14 @@ class TopologyEdgeDeleted(BaseEvent):
conn: Connection
class CustomModelCardAdded(BaseEvent):
model_card: ModelCard
class CustomModelCardDeleted(BaseEvent):
model_id: ModelId
@final
class TraceEventData(FrozenModel):
name: str
@@ -147,6 +156,8 @@ Event = (
| TopologyEdgeDeleted
| TracesCollected
| TracesMerged
| CustomModelCardAdded
| CustomModelCardDeleted
)
src/exo/worker/main.py
+10 -1
View File
@@ -8,7 +8,7 @@ from loguru import logger
from exo.api.types import ImageEditsTaskParams
from exo.download.download_utils import resolve_model_in_path
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelId
from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
from exo.shared.types.commands import (
ForwarderCommand,
ForwarderDownloadCommand,
@@ -16,6 +16,8 @@ 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,
@@ -130,6 +132,13 @@ class Worker:
event.chunk.data
)
if isinstance(event, CustomModelCardAdded):
await event.model_card.save_to_custom_dir()
add_to_card_cache(event.model_card)
if isinstance(event, CustomModelCardDeleted):
await delete_custom_card(event.model_id)
async def plan_step(self):
while True:
await anyio.sleep(0.1)