diff --git a/src/exo/api/main.py b/src/exo/api/main.py index 08d95880..32900c78 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -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)} ) diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 94701c6c..1511897e 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -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 diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 38132a32..fb9f6d9a 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -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(): diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index 048c469f..fef1efaa 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -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 diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index fb55632e..a6c988a5 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -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 ) diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index c9c2fc3a..d9799313 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -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 ) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 5389e78d..4bfed825 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -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)