From 525aaab8085e41935ed90663da1db7cab3baa9a8 Mon Sep 17 00:00:00 2001 From: Evan Date: Wed, 25 Feb 2026 18:44:45 +0000 Subject: [PATCH] fix --- src/exo/master/api.py | 12 ++++----- src/exo/shared/types/worker/instances.py | 32 +++++++++++++++++++++++- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 85d964b7..af7cfa73 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -524,15 +524,15 @@ class API: if ( model_card.model_id, - sharding, - instance_meta, + instance.sharding(), + instance.instance_meta(), len(placement_node_ids), ) not in seen: previews.append( PlacementPreview( model_id=model_card.model_id, - sharding=sharding, - instance_meta=instance_meta, + sharding=instance.sharding(), + instance_meta=instance.instance_meta(), instance=instance, memory_delta_by_node=memory_delta_by_node or None, error=None, @@ -541,8 +541,8 @@ class API: seen.add( ( model_card.model_id, - sharding, - instance_meta, + instance.sharding(), + instance.instance_meta(), len(placement_node_ids), ) ) diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py index 76bd6fd4..25e2b753 100644 --- a/src/exo/shared/types/worker/instances.py +++ b/src/exo/shared/types/worker/instances.py @@ -4,7 +4,13 @@ from pydantic import model_validator from exo.shared.models.model_cards import ModelTask from exo.shared.types.common import Host, Id, NodeId -from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata +from exo.shared.types.worker.runners import RunnerId, ShardAssignments +from exo.shared.types.worker.shards import ( + PipelineShardMetadata, + Sharding, + ShardMetadata, + TensorShardMetadata, +) from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel @@ -24,16 +30,40 @@ class BaseInstance(TaggedModel): def shard(self, runner_id: RunnerId) -> ShardMetadata | None: return self.shard_assignments.runner_to_shard.get(runner_id, None) + @staticmethod + def instance_meta() -> InstanceMeta: ... + + def sharding(self) -> Sharding: + if all( + isinstance(sm, PipelineShardMetadata) + for sm in self.shard_assignments.runner_to_shard.values() + ): + return Sharding.Pipeline + if all( + isinstance(sm, TensorShardMetadata) + for sm in self.shard_assignments.runner_to_shard.values() + ): + return Sharding.Tensor + raise ValueError("shard metadata malformed") + class MlxRingInstance(BaseInstance): hosts_by_node: dict[NodeId, list[Host]] ephemeral_port: int + @staticmethod + def instance_meta() -> InstanceMeta: + return InstanceMeta.MlxRing + class MlxJacclInstance(BaseInstance): jaccl_devices: list[list[str | None]] jaccl_coordinators: dict[NodeId, str] + @staticmethod + def instance_meta() -> InstanceMeta: + return InstanceMeta.MlxJaccl + # TODO: Single node instance Instance = MlxRingInstance | MlxJacclInstance