diff --git a/rust/exo_rs/src/task.rs b/rust/exo_rs/src/task.rs index 08b5fd31a..0900a084b 100644 --- a/rust/exo_rs/src/task.rs +++ b/rust/exo_rs/src/task.rs @@ -67,7 +67,12 @@ impl TaskRequester { ) -> PyResult> { let session = self.session.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - let receiver = declare_task_stream(&session, &command_id)?; + let receiver = session + .declare_subscriber(task_chunks_key(command_id.as_str())) + .wait() + .map_err(|e| { + PyConnectionError::new_err(format!("failed to declare task stream: {e}")) + })?; request_task_admission(&session, instance_id, command_id, command).await?; Ok(TaskStream { receiver }) @@ -92,10 +97,6 @@ impl TaskRequester { } } -fn task_key(instance_id: &str, command_id: &str) -> String { - format!("task/instances/{instance_id}/tasks/{command_id}") -} - fn task_chunks_key(command_id: &str) -> String { format!("task/commands/{command_id}/chunks") } @@ -104,16 +105,6 @@ fn task_assignment_key(instance_id: &str, task_id: &str) -> String { format!("task_assignments/{instance_id}/{task_id}") } -fn declare_task_stream( - session: &ZSession, - command_id: &str, -) -> PyResult>> { - session - .declare_subscriber(task_chunks_key(command_id)) - .wait() - .map_err(|e| PyConnectionError::new_err(format!("failed to declare task stream: {e}"))) -} - async fn request_task_admission( session: &ZSession, instance_id: String, @@ -121,7 +112,7 @@ async fn request_task_admission( command: String, ) -> PyResult<()> { let replies = session - .get(task_key(&instance_id, &command_id)) + .get(format!("task/instances/{instance_id}/tasks/{command_id}")) .payload(command) .congestion_control(CongestionControl::Block) .consolidation(ConsolidationMode::None) diff --git a/src/exo/api/main.py b/src/exo/api/main.py index 9903a364b..0bfb92290 100644 --- a/src/exo/api/main.py +++ b/src/exo/api/main.py @@ -23,6 +23,13 @@ from hypercorn.utils import LifespanTimeoutError, ShutdownError from loguru import logger from pydantic import TypeAdapter, ValidationError +from exo.master.placement import ( + add_instance_to_placements, + cancel_unnecessary_downloads, + delete_instance, + get_transition_events, + place_instance, +) from exo.api.adapters.chat_completions import ( chat_request_to_text_generation, collect_chat_response, @@ -128,7 +135,6 @@ from exo.api.types.openai_responses import ( ResponsesResponse, ) from exo.master.image_store import ImageStore -from exo.master.placement import place_instance as get_instance_placements from exo.shared.apply import apply from exo.shared.constants import ( DASHBOARD_DIR, @@ -423,13 +429,24 @@ class API: ) from e async def place_instance(self, payload: PlaceInstanceParams): + state = self.state.with_aggregator(self.aggregator) command = PlaceInstance( model_card=await ModelCard.load(payload.model_id), sharding=payload.sharding, instance_meta=payload.instance_meta, min_nodes=payload.min_nodes, ) - await self._send(command) + new_instance = place_instance( + command, + state.topology, + state.node_memory, + state.node_network, + state.node_backends, + download_status=state.downloads, + node_rdma_ctl=state.node_rdma_ctl, + ) + dialing = new_instance.primary_output_node() + return CreateInstanceResponse( message="Command received.", diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 5bec14ef7..60b5d7a12 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -362,9 +362,10 @@ class Master: selected_instance.shard_assignments.shards ) case DeleteInstance(): - placement = delete_instance(command, self.state.instances) + state = self.state.with_aggregator(self.aggregator) + placement = delete_instance(command, state.instances) transition_events = get_transition_events( - self.state.instances, placement, self.state.tasks + state.instances, placement, state.tasks ) for cmd in cancel_unnecessary_downloads( placement, self.state.downloads @@ -392,13 +393,14 @@ class Master: ) generated_events.extend(transition_events) case CreateInstance(): + state = self.state.with_aggregator(self.aggregator) placement = add_instance_to_placements( command, - self.state.topology, - self.state.instances, + state.topology, + state.instances, ) transition_events = get_transition_events( - self.state.instances, placement, self.state.tasks + state.instances, placement, state.tasks ) generated_events.extend(transition_events) case TaskCancelled(): diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index 0e9c2a33e..ed5b76bf7 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -106,14 +106,13 @@ def _cycle_download_score( def place_instance( command: PlaceInstance, topology: Topology, - current_instances: Mapping[InstanceId, Instance], node_memory: Mapping[NodeId, MemoryUsage], node_network: Mapping[NodeId, NodeNetworkInfo], node_backends: Mapping[NodeId, list[Backend]], required_nodes: set[NodeId] | None = None, download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None, -) -> dict[InstanceId, Instance]: +) -> Instance: cycles = topology.get_cycles() candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles)) @@ -258,7 +257,6 @@ def place_instance( cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids) instance_id = InstanceId() - target_instances = dict(deepcopy(current_instances)) match command.instance_meta: case InstanceMeta.MlxJaccl: @@ -274,7 +272,7 @@ def place_instance( cycle_digraph=cycle_digraph, node_network=node_network, ) - target_instances[instance_id] = MlxJacclInstance( + return MlxJacclInstance( instance_id=instance_id, shard_assignments=shard_assignments, jaccl_devices=mlx_jaccl_devices, @@ -288,14 +286,13 @@ def place_instance( ephemeral_port=ephemeral_port, node_network=node_network, ) - target_instances[instance_id] = MlxRingInstance( + return MlxRingInstance( instance_id=instance_id, shard_assignments=shard_assignments, hosts_by_node=hosts_by_node, ephemeral_port=ephemeral_port, ) - return target_instances def delete_instance( diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py index 61c7be312..ccc141364 100644 --- a/src/exo/shared/types/worker/instances.py +++ b/src/exo/shared/types/worker/instances.py @@ -32,6 +32,9 @@ class BaseInstance(TaggedModel): if nid == node_id: yield rid + def primary_output_node(self) -> NodeId: + return self.shard_assignments.shards[self.shard_assignments.primary_output_node].node_id + class MlxRingInstance(BaseInstance): hosts_by_node: dict[NodeId, list[Host]]