startit

rust/exo_rs/src/task.rs

@@ -67,7 +67,12 @@ impl TaskRequester {
     ) -> PyResult<Bound<'py, PyAny>> {
         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<Subscriber<FifoChannelHandler<Sample>>> {
-    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)

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.",

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():

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(

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]]