← back to Exo
switch to ModelCard.load outside of download log (#1339)
46a14153dd8d9c614e0d1762dbb60380ac02d1d2 · 2026-01-30 11:20:20 +0000 · Evan Quiney
some attempts to load model cards (i.e. build_base_shard) always went
through networking rather than using downloaded model cards. we should
always default to ModelCard.load in these scenarios
Files touched
M src/exo/download/impl_shard_downloader.pyM src/exo/master/api.py
Diff
commit 46a14153dd8d9c614e0d1762dbb60380ac02d1d2
Author: Evan Quiney <evanev7@gmail.com>
Date: Fri Jan 30 11:20:20 2026 +0000
switch to ModelCard.load outside of download log (#1339)
some attempts to load model cards (i.e. build_base_shard) always went
through networking rather than using downloaded model cards. we should
always default to ModelCard.load in these scenarios
---
src/exo/download/impl_shard_downloader.py | 2 +-
src/exo/master/api.py | 42 +++++++++++--------------------
2 files changed, 16 insertions(+), 28 deletions(-)
diff --git a/src/exo/download/impl_shard_downloader.py b/src/exo/download/impl_shard_downloader.py
index 68c6fdcf..4c9b02ed 100644
--- a/src/exo/download/impl_shard_downloader.py
+++ b/src/exo/download/impl_shard_downloader.py
@@ -21,7 +21,7 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
async def build_base_shard(model_id: ModelId) -> ShardMetadata:
- model_card = await ModelCard.from_hf(model_id)
+ model_card = await ModelCard.load(model_id)
return PipelineShardMetadata(
model_card=model_card,
device_rank=0,
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index de2b3a3e..c2323618 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -146,18 +146,6 @@ def chunk_to_response(
)
-async def resolve_model_card(model_id: ModelId) -> ModelCard:
- if model_id in MODEL_CARDS:
- model_card = MODEL_CARDS[model_id]
- return model_card
-
- for card in MODEL_CARDS.values():
- if card.model_id == ModelId(model_id):
- return card
-
- return await ModelCard.from_hf(model_id)
-
-
class API:
def __init__(
self,
@@ -279,7 +267,7 @@ class API:
async def place_instance(self, payload: PlaceInstanceParams):
command = PlaceInstance(
- model_card=await resolve_model_card(payload.model_id),
+ model_card=await ModelCard.load(payload.model_id),
sharding=payload.sharding,
instance_meta=payload.instance_meta,
min_nodes=payload.min_nodes,
@@ -296,7 +284,7 @@ class API:
self, payload: CreateInstanceParams
) -> CreateInstanceResponse:
instance = payload.instance
- model_card = await resolve_model_card(instance.shard_assignments.model_id)
+ model_card = await ModelCard.load(instance.shard_assignments.model_id)
required_memory = model_card.storage_size
available_memory = self._calculate_total_available_memory()
@@ -324,7 +312,7 @@ class API:
instance_meta: InstanceMeta = InstanceMeta.MlxRing,
min_nodes: int = 1,
) -> Instance:
- model_card = await resolve_model_card(model_id)
+ model_card = await ModelCard.load(model_id)
try:
placements = get_instance_placements(
@@ -565,7 +553,7 @@ class API:
text_parts: list[str] = []
tool_calls: list[ToolCall] = []
- model: str | None = None
+ model: ModelId | None = None
finish_reason: FinishReason | None = None
usage: Usage | None = None
@@ -624,7 +612,7 @@ class API:
) -> BenchChatCompletionResponse:
text_parts: list[str] = []
tool_calls: list[ToolCall] = []
- model: str | None = None
+ model: ModelId | None = None
finish_reason: FinishReason | None = None
stats: GenerationStats | None = None
@@ -677,7 +665,7 @@ class API:
)
return resp
- async def _trigger_notify_user_to_download_model(self, model_id: str) -> None:
+ async def _trigger_notify_user_to_download_model(self, model_id: ModelId) -> None:
logger.warning(
"TODO: we should send a notification to the user to download the model"
)
@@ -686,7 +674,7 @@ class API:
self, payload: ChatCompletionTaskParams
) -> ChatCompletionResponse | StreamingResponse:
"""Handle chat completions, supporting both streaming and non-streaming responses."""
- model_card = await resolve_model_card(ModelId(payload.model))
+ model_card = await ModelCard.load(ModelId(payload.model))
payload.model = model_card.model_id
if not any(
@@ -713,7 +701,7 @@ class API:
async def bench_chat_completions(
self, payload: BenchChatCompletionTaskParams
) -> BenchChatCompletionResponse:
- model_card = await resolve_model_card(ModelId(payload.model))
+ model_card = await ModelCard.load(ModelId(payload.model))
payload.model = model_card.model_id
if not any(
@@ -733,12 +721,12 @@ class API:
response = await self._collect_chat_completion_with_stats(command.command_id)
return response
- async def _validate_image_model(self, model: str) -> ModelId:
+ async def _validate_image_model(self, model: ModelId) -> ModelId:
"""Validate model exists and return resolved model ID.
Raises HTTPException 404 if no instance is found for the model.
"""
- model_card = await resolve_model_card(ModelId(model))
+ model_card = await ModelCard.load(model)
resolved_model = model_card.model_id
if not any(
instance.shard_assignments.model_id == resolved_model
@@ -784,7 +772,7 @@ class API:
When stream=True and partial_images > 0, returns a StreamingResponse
with SSE-formatted events for partial and final images.
"""
- payload.model = await self._validate_image_model(payload.model)
+ payload.model = await self._validate_image_model(ModelId(payload.model))
command = ImageGeneration(
request_params=payload,
@@ -1029,7 +1017,7 @@ class API:
async def bench_image_generations(
self, request: Request, payload: BenchImageGenerationTaskParams
) -> BenchImageGenerationResponse:
- payload.model = await self._validate_image_model(payload.model)
+ payload.model = await self._validate_image_model(ModelId(payload.model))
payload.stream = False
payload.partial_images = 0
@@ -1050,7 +1038,7 @@ class API:
self,
image: UploadFile,
prompt: str,
- model: str,
+ model: ModelId,
n: int,
size: str,
response_format: Literal["url", "b64_json"],
@@ -1145,7 +1133,7 @@ class API:
command = await self._send_image_edits_command(
image=image,
prompt=prompt,
- model=model,
+ model=ModelId(model),
n=n,
size=size,
response_format=response_format,
@@ -1201,7 +1189,7 @@ class API:
command = await self._send_image_edits_command(
image=image,
prompt=prompt,
- model=model,
+ model=ModelId(model),
n=n,
size=size,
response_format=response_format,
← 9ba61f37 improve log message in shard downloader
·
back to Exo
·
make node-ids unique per-session (#1338) 2a4e0d46 →