api streams

rust/exo_net/exo_net.pyi

@@ -16,7 +16,7 @@ __all__ = [
 
 @typing.final
 class NetReceiver:
-    def recv(self) -> collections.abc.Awaitable[bytes | None]: ...
+    def recv(self) -> collections.abc.Awaitable[bytes]: ...
 
 @typing.final
 class NetSender:
@@ -25,23 +25,27 @@ class NetSender:
 @typing.final
 class NetworkingHandle:
     @staticmethod
-    def new(identity: bytes, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> tuple[NetworkingHandle, PySession]: ...
+    def new(
+        identity: bytes,
+        bootstrap_peers: typing.Sequence[builtins.str],
+        listen_port: builtins.int,
+    ) -> tuple[NetworkingHandle, PySession]: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
-        
+
         Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
         """
     async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Unsubscribes from a `GossipSub` topic.
-        
+
         Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
         """
     async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
         r"""
         Publishes a message with multiple topics to the `GossipSub` network.
-        
+
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
     async def recv(self) -> PyFromSwarm: ...
@@ -53,16 +57,18 @@ class PyFromSwarm:
         @property
         def connected(self) -> builtins.bool: ...
         def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
-    
+
     @typing.final
     class Message(PyFromSwarm):
-        __match_args__ = ("topic", "data",)
+        __match_args__ = (
+            "topic",
+            "data",
+        )
         @property
         def topic(self) -> builtins.str: ...
         @property
         def data(self) -> bytes: ...
         def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
-    
 
 @typing.final
 class PySession:
@@ -73,4 +79,3 @@ class PySession:
 @typing.final
 class StateProxy:
     def snapshot(self) -> collections.abc.Awaitable[str]: ...
-

rust/exo_net/src/lib.rs

@@ -9,12 +9,12 @@ mod allow_threading;
 mod networking;
 mod point_to_point;
 mod session;
-mod state;
+// mod state;
 
 use crate::networking::networking_submodule;
 use crate::point_to_point::{NetReceiver, NetSender};
 use crate::session::PySession;
-use crate::state::StateProxy;
+//use crate::state::StateProxy;
 use pyo3::prelude::PyModule;
 use pyo3::types::PyModuleMethods;
 use pyo3::{Bound, PyResult, pymodule};
@@ -164,7 +164,7 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     //       too many importing issues...
     // m.add_class::<PyKeypair>()?;
     // networking_submodule(m)?;
-    m.add_class::<StateProxy>()?;
+    // m.add_class::<StateProxy>()?;
     m.add_class::<PySession>()?;
     m.add_class::<NetReceiver>()?;
     m.add_class::<NetSender>()?;

rust/exo_net/src/point_to_point.rs

@@ -1,9 +1,8 @@
 use std::sync::Arc;
 
 use pyo3::exceptions::PyConnectionError;
-use pyo3::types::PyBytes;
-use pyo3::types::PyNone;
-use pyo3::{BoundObject, prelude::*};
+use pyo3::prelude::*;
+use pyo3::{exceptions::PyStopAsyncIteration, types::PyBytes};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 use zenoh::Result;
 use zenoh::{
@@ -23,7 +22,7 @@ pub struct NetReceiver {
 #[pymethods]
 impl NetReceiver {
     #[gen_stub(override_return_type(
-        type_repr="collections.abc.Awaitable[bytes | None]",
+        type_repr="collections.abc.Awaitable[bytes]",
         imports=("collections.abc")
     ))]
     pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
@@ -37,9 +36,9 @@ impl NetReceiver {
                 match subscriber.recv_async().await {
                     Err(_) => {
                         // stream closed;
-                        Ok(Python::attach(|py| PyNone::get(py).unbind()).into_any())
+                        Err(PyStopAsyncIteration::new_err(()))
                     }
-                    Ok(sample) => Ok(sample.payload().to_bytes().to_vec().pybytes().into_any()),
+                    Ok(sample) => Ok(sample.payload().to_bytes().to_vec().pybytes()),
                 }
             }
         })
@@ -72,18 +71,23 @@ impl NetSender {
             let bytes = data.as_bytes().to_vec();
             async move {
                 if is_first {
+                    log::warn!("sender waiting for listener");
                     wait_for_listener(&*publisher)
                         .await
                         .map_err(|e| PyConnectionError::new_err(e.to_string()))?;
+                    log::warn!("listener found");
                 }
+                log::warn!("checking for matcher");
                 if !publisher
                     .matching_status()
                     .await
                     .map_err(|e| PyConnectionError::new_err(e.to_string()))?
                     .matching()
                 {
+                    log::warn!("no matcher found");
                     return Ok(false);
                 }
+                log::warn!("publishing");
                 publisher
                     .put(&bytes)
                     .await

rust/exo_net/src/session.rs

@@ -5,11 +5,10 @@ use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 
 use zenoh::Session;
 use zenoh::Wait;
-use zenoh::qos::CongestionControl;
 
 use crate::{
     point_to_point::{NetReceiver, NetSender},
-    state::StateProxy,
+    //state::StateProxy,
 };
 
 #[gen_stub_pyclass]
@@ -55,7 +54,6 @@ impl PySession {
             publisher: Arc::new(
                 self.session
                     .declare_publisher(key)
-                    .congestion_control(CongestionControl::Block)
                     .wait()
                     // C5: key format error, could be declaration error
                     .map_err(|e| PyValueError::new_err(e.to_string()))?,
@@ -64,9 +62,9 @@ impl PySession {
         })
     }
 
-    pub fn state_proxy(&self) -> StateProxy {
-        StateProxy {
-            session: self.session.clone(),
-        }
-    }
+    //pub fn state_proxy(&self) -> StateProxy {
+        //StateProxy {
+            //session: self.session.clone(),
+        //}
+    //}
 }

rust/exo_pyo3_bindings/Cargo.toml

@@ -0,0 +1,62 @@
+[package]
+name = "exo_pyo3_bindings"
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
+
+[lib]
+doctest = false
+path = "src/lib.rs"
+name = "exo_pyo3_bindings"
+
+# "cdylib" needed to produce shared library for Python to import
+# "rlib" needed for stub-gen to run
+crate-type = ["cdylib", "rlib"]
+
+[[bin]]
+path = "src/bin/stub_gen.rs"
+name = "stub_gen"
+doc = false
+
+[lints]
+workspace = true
+
+[dependencies]
+networking.workspace = true
+extend.workspace = true
+
+# interop
+pyo3 = { version = "0.27.2", features = [
+  # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
+  # "nightly", # enables better-supported GIL integration
+  "experimental-async", # async support in #[pyfunction] & #[pymethods]
+  #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
+  #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
+  # "multiple-pymethods", # allows multiple #[pymethods] sections per class
+
+  # integrations with other libraries
+  # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
+  # "ordered-float", "rust_decimal", "smallvec",
+  # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
+] }
+pyo3-stub-gen = { version = "0.17.2" }
+pyo3-async-runtimes = { version = "0.27.0", features = [
+  "attributes",
+  "tokio-runtime",
+  "testing",
+] }
+pyo3-log = "0.13.2"
+
+# async runtime
+tokio = { workspace = true, features = ["full", "tracing"] }
+futures-lite = { workspace = true }
+pin-project = "1.1.10"
+
+# Tracing
+log.workspace = true
+env_logger.workspace = true
+
+# Networking
+zenoh.workspace = true
+zerompk.workspace = true
+rand = "0.10.1"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -0,0 +1,72 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+import builtins
+import typing
+
+@typing.final
+class Keypair:
+    r"""
+    Identity keypair of a node.
+    """
+    @staticmethod
+    def generate() -> Keypair:
+        r"""
+        Generate a new Ed25519 keypair.
+        """
+    @staticmethod
+    def from_bytes(bytes: bytes) -> Keypair:
+        r"""
+        Construct an Ed25519 keypair from secret key bytes
+        """
+    def to_bytes(self) -> bytes:
+        r"""
+        Get the secret key bytes underlying the keypair
+        """
+    def to_node_id(self) -> builtins.str:
+        r"""
+        Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
+        """
+
+@typing.final
+class NetworkingHandle:
+    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
+    async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
+        r"""
+        Subscribe to a `GossipSub` topic.
+        
+        Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
+        """
+    async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
+        r"""
+        Unsubscribes from a `GossipSub` topic.
+        
+        Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
+        """
+    async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
+        r"""
+        Publishes a message with multiple topics to the `GossipSub` network.
+        
+        If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
+        """
+    async def recv(self) -> PyFromSwarm: ...
+
+class PyFromSwarm:
+    @typing.final
+    class Connection(PyFromSwarm):
+        __match_args__ = ("connected",)
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(PyFromSwarm):
+        __match_args__ = ("topic", "data",)
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
+    
+    ...
+

rust/networking/src/lib.rs

@@ -94,6 +94,7 @@ pub struct Session {
 }
 impl Drop for WatchAllHandle {
     fn drop(&mut self) {
+        log::error!("aborting iface watcher");
         self._async.abort();
     }
 }

rust/networking/src/swarm.rs

@@ -50,6 +50,7 @@ impl Swarm {
             mut from_client,
         } = self;
         let stream = async_stream::stream! {
+            // very important!
             let mut session = session;
             let (mut to_topics, mut from_topics) = mpsc::channel(1024);
             let mut topics = Topics::new();

src/exo/api/main.py

@@ -234,7 +234,7 @@ def _require_disaggregation_enabled() -> None:
 
 @dataclass
 class Transport:
-    session: PySession
+    z: PySession
     cancel_scopes: dict[CommandId, anyio.CancelScope] = field(
         init=False, default_factory=dict
     )
@@ -244,7 +244,7 @@ class Transport:
 
     def __post_init__(self):
         # TODO: retire root keyspace
-        self.command_sender = self.session.net_sender("orchestrator")
+        self.command_sender = self.z.net_sender("orchestrator")
 
     async def send_command(self, command: Command) -> bool:
         while self.paused:
@@ -275,16 +275,11 @@ class Transport:
             with anyio.CancelScope() as cs:
                 self.cancel_scopes[command_id] = cs
                 # recv from any node
-                receiver = self.session.net_receiver(
+                receiver = self.z.net_receiver(
                     f"runners/*/active_tasks/{command_id}/chunks"
                 )
                 while True:
                     data = await receiver.recv()
-                    if data is None:
-                        logger.warning(
-                            "stream terminated early without finish reason EOF"
-                        )
-                        break
                     yield (
                         chunk := cast(
                             Chunk,

src/exo/api/tests/test_cancel_command.py

@@ -1,11 +1,11 @@
 # pyright: reportUnusedFunction=false, reportAny=false
 from typing import Any
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from exo.api.main import API, Transport
+from exo.api.main import API
 from exo.shared.types.common import CommandId
 
 
@@ -15,9 +15,9 @@ def _make_api() -> Any:
     app = FastAPI()
     api = object.__new__(API)
     api.app = app
-    api.transport = object.__new__(Transport)
-    api.transport.cancel = AsyncMock()
-    api.transport.send_command = AsyncMock()
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._send = AsyncMock()  # pyright: ignore[reportPrivateUsage]
     api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
     app.post("/v1/cancel/{command_id}")(api.cancel_command)
     return api
@@ -43,14 +43,16 @@ def test_cancel_active_text_generation() -> None:
     client = TestClient(api.app)
 
     cid = CommandId("text-cmd-123")
+    sender = MagicMock()
+    api._text_generation_queues[cid] = sender
 
     response = client.post(f"/v1/cancel/{cid}")
     assert response.status_code == 200
     data: dict[str, Any] = response.json()
     assert data["message"] == "Command cancelled."
     assert data["command_id"] == str(cid)
-    api.transport.cancel.assert_called_once()
-    api.transport.send_command.assert_called_once()
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
     task_cancelled = api._send.call_args[0][0]
     assert task_cancelled.cancelled_command_id == cid
 
@@ -61,13 +63,15 @@ def test_cancel_active_image_generation() -> None:
     client = TestClient(api.app)
 
     cid = CommandId("img-cmd-456")
+    sender = MagicMock()
+    api._image_generation_queues[cid] = sender
 
     response = client.post(f"/v1/cancel/{cid}")
     assert response.status_code == 200
     data: dict[str, Any] = response.json()
     assert data["message"] == "Command cancelled."
     assert data["command_id"] == str(cid)
-    api.transport.cancel.assert_called_once()
-    api.transport.send_command.assert_called_once()
-    task_cancelled = api.transport.send_command.call_args[0][0]
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
     assert task_cancelled.cancelled_command_id == cid

src/exo/api/tests/test_instance_deleted_stream_cleanup.py

@@ -1,10 +1,9 @@
 # pyright: reportUnusedFunction=false, reportAny=false
 """Tests that InstanceDeleted events close active generation streams."""
 
-from typing import Any
 from unittest.mock import MagicMock
 
-from exo.api.main import API, Transport
+from exo.api.main import API
 from exo.api.types import ImageGenerationTaskParams
 from exo.shared.types.common import CommandId, ModelId
 from exo.shared.types.state import State
@@ -17,11 +16,12 @@ from exo.shared.types.text_generation import (
 from exo.shared.types.worker.instances import InstanceId
 
 
-def _make_api_with_state(state: State) -> Any:
+def _make_api_with_state(state: State) -> API:
     """Create a minimal API instance with pre-set state."""
     api = object.__new__(API)
     api.state = state
-    api.transport = object.__new__(Transport)
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
     return api
 
 
@@ -47,10 +47,13 @@ def test_close_streams_for_deleted_instance() -> None:
     state = State(tasks={task.task_id: task})
     api = _make_api_with_state(state)
 
-    api._close_streams_for_instance(instance_id)
+    sender = MagicMock()
+    api._text_generation_queues[command_id] = sender  # pyright: ignore[reportPrivateUsage]
 
-    api.transport.cancel.assert_called_once()
-    assert api.transport.cancel.call_args[0][0] == command_id
+    api._close_streams_for_instance(instance_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_called_once()
+    assert command_id not in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
 
 
 def test_close_streams_ignores_unrelated_instances() -> None:
@@ -69,6 +72,7 @@ def test_close_streams_ignores_unrelated_instances() -> None:
     api._close_streams_for_instance(target_id)  # pyright: ignore[reportPrivateUsage]
 
     sender.close.assert_not_called()
+    assert other_cmd in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
 
 
 def test_close_streams_for_deleted_instance_image_generation() -> None:

src/exo/main.py

@@ -114,7 +114,7 @@ class Node:
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
-            command_receiver=session.net_receiver("orchestrator")
+            session=session,
         )
 
         er_send, er_recv = channel[ElectionResult]()
@@ -214,7 +214,7 @@ class Node:
                         download_command_sender=self.router.sender(
                             topics.DOWNLOAD_COMMANDS
                         ),
-                        command_receiver=self.session.net_receiver("orchestrator"),
+                        session=self.session,
                     )
                     self._tg.start_soon(self.master.run)
                 elif (

src/exo/master/main.py

@@ -1,7 +1,8 @@
 from datetime import datetime, timedelta, timezone
+from typing import cast
 
 import anyio
-from exo_net import NetReceiver
+from exo_net import PySession
 from loguru import logger
 from pydantic import TypeAdapter
 
@@ -123,18 +124,18 @@ class Master:
         node_id: NodeId,
         session_id: SessionId,
         *,
-        command_receiver: NetReceiver, # todo: not this type
         event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
         download_command_sender: Sender[ForwarderDownloadCommand],
+        session: PySession,
     ):
         self.node_id = node_id
         self.session_id = session_id
         self.state = State()
         self._tg: TaskGroup = TaskGroup()
         self.command_task_mapping: dict[CommandId, TaskId] = {}
-        self.command_receiver = command_receiver
+        self.session = session
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
         self.download_command_sender = download_command_sender
@@ -163,12 +164,12 @@ class Master:
         self._tg.cancel_tasks()
 
     async def _command_processor(self) -> None:
+        receiver = self.session.net_receiver("orchestrator")
         while True:
-            data = await self.command_receiver.recv()
-            if not data:
-                break
+            command = cast(
+                Command, TypeAdapter(Command).validate_json(await receiver.recv())
+            )
             try:
-                command = TypeAdapter[Command](Command).validate_json(data)
                 logger.info(f"Executing command: {command}")
 
                 generated_events: list[Event] = []
@@ -196,9 +197,7 @@ class Master:
                                     if task.instance_id == instance.instance_id
                                     and task.task_status in in_flight
                                 )
-                                instance_task_counts[instance.instance_id] = (
-                                    task_count
-                                )
+                                instance_task_counts[instance.instance_id] = task_count
 
                         if not instance_task_counts:
                             raise ValueError(
@@ -207,9 +206,7 @@ class Master:
 
                         available_instance_ids = sorted(
                             instance_task_counts.keys(),
-                            key=lambda instance_id: instance_task_counts[
-                                instance_id
-                            ],
+                            key=lambda instance_id: instance_task_counts[instance_id],
                         )
 
                         decode_instance_id = available_instance_ids[0]
@@ -247,9 +244,7 @@ class Master:
                                     if task.instance_id == instance.instance_id
                                     and task.task_status in in_flight
                                 )
-                                instance_task_counts[instance.instance_id] = (
-                                    task_count
-                                )
+                                instance_task_counts[instance.instance_id] = task_count
 
                         if not instance_task_counts:
                             raise ValueError(
@@ -258,9 +253,7 @@ class Master:
 
                         available_instance_ids = sorted(
                             instance_task_counts.keys(),
-                            key=lambda instance_id: instance_task_counts[
-                                instance_id
-                            ],
+                            key=lambda instance_id: instance_task_counts[instance_id],
                         )
 
                         task_id = TaskId()
@@ -285,11 +278,10 @@ class Master:
                                 selected_instance_id
                             )
                             if selected_instance:
-                                ranks = set(
+                                self._expected_ranks[task_id] = set(
                                     shard.device_rank
                                     for shard in selected_instance.shard_assignments.runner_to_shard.values()
                                 )
-                                self._expected_ranks[task_id] = ranks
                     case ImageEdits():
                         for instance in self.state.instances.values():
                             if (
@@ -303,9 +295,7 @@ class Master:
                                     if task.instance_id == instance.instance_id
                                     and task.task_status in in_flight
                                 )
-                                instance_task_counts[instance.instance_id] = (
-                                    task_count
-                                )
+                                instance_task_counts[instance.instance_id] = task_count
 
                         if not instance_task_counts:
                             raise ValueError(
@@ -314,9 +304,7 @@ class Master:
 
                         available_instance_ids = sorted(
                             instance_task_counts.keys(),
-                            key=lambda instance_id: instance_task_counts[
-                                instance_id
-                            ],
+                            key=lambda instance_id: instance_task_counts[instance_id],
                         )
 
                         task_id = TaskId()
@@ -341,11 +329,11 @@ class Master:
                                 selected_instance_id
                             )
                             if selected_instance:
-                                ranks = set(
+                                self._expected_ranks[task_id] = set(
                                     shard.device_rank
                                     for shard in selected_instance.shard_assignments.runner_to_shard.values()
                                 )
-                                self._expected_ranks[task_id] = ranks
+
                     case DeleteInstance():
                         placement = delete_instance(command, self.state.instances)
                         transition_events = get_transition_events(
@@ -443,7 +431,9 @@ class Master:
                             InstanceLinkDeleted(link_id=command.link_id)
                         )
                     case RequestEventLog():
-                        end = len(self._event_log)
+                        # We should just be able to send everything, since other buffers will ignore old messages
+                        # rate limit to 1000 at a time
+                        end = min(command.since_idx + 1000, len(self._event_log))
                         for i, event in enumerate(
                             self._event_log.read_range(command.since_idx, end),
                             start=command.since_idx,

src/exo/master/tests/test_master.py

@@ -6,6 +6,7 @@ import pytest
 from loguru import logger
 
 from exo.master.main import Master
+from exo.routing.router import get_node_id_keypair
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.commands import (
     CommandId,
@@ -46,7 +47,8 @@ from exo.utils.channels import channel
 
 @pytest.mark.asyncio
 async def test_master():
-    node_id = NodeId("master test")
+    keypair = get_node_id_keypair()
+    node_id = NodeId(keypair.to_node_id())
     session_id = SessionId(master_node_id=node_id, election_clock=0)
 
     ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()

src/exo/routing/event_router.py

@@ -4,7 +4,6 @@ from random import random
 import anyio
 from anyio import BrokenResourceError, ClosedResourceError
 from anyio.abc import CancelScope
-from exo_net import NetSender
 from loguru import logger
 
 from exo.shared.types.commands import RequestEventLog
@@ -20,6 +19,8 @@ from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.task_group import TaskGroup
 
+from exo_net import NetSender
+
 
 @dataclass
 class EventRouter:

src/exo/shared/types/chunks.py

@@ -32,21 +32,15 @@ class TokenChunk(BaseChunk):
 
 class ErrorChunk(BaseChunk):
     error_message: str
-
-    @property
-    def finish_reason(self) -> Literal["error"]:
-        return "error"
+    finish_reason: Literal["error"] = "error"
 
 
 class ToolCallChunk(BaseChunk):
     tool_calls: list[ToolCallItem]
     usage: Usage | None
+    finish_reason: Literal["tool_calls"] = "tool_calls"
     stats: GenerationStats | None = None
 
-    @property
-    def finish_reason(self) -> Literal["tool_calls"]:
-        return "tool_calls"
-
 
 class ImageChunk(BaseChunk):
     data: str
@@ -90,10 +84,6 @@ class PrefillProgressChunk(BaseChunk):
     processed_tokens: int
     total_tokens: int
 
-    @property
-    def finish_reason(self) -> FinishReason | None:
-        return None
-
 
 StatusChunk = PrefillProgressChunk
 GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk

src/exo/shared/types/events.py

@@ -161,7 +161,6 @@ Event = (
     | NodeDownloadProgress
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
-    | ChunkGenerated
     | InputChunkReceived
     | TracesCollected
     | TracesMerged

src/exo/worker/engines/mlx/builder.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.common import ModelId
-from exo.shared.types.events import Event
+from exo.shared.types.events import ChunkGenerated, Event
 from exo.shared.types.tasks import TaskId
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runner_response import ModelLoadingResponse
@@ -32,7 +32,7 @@ from .vision import VisionProcessor
 @dataclass
 class MlxBuilder(Builder):
     model_id: ModelId
-    event_sender: MpSender[Event]
+    event_sender: MpSender[Event | ChunkGenerated]
     cancel_receiver: MpReceiver[TaskId]
     inference_model: Model | None = None
     tokenizer: TokenizerWrapper | None = None

src/exo/worker/main.py

@@ -15,6 +15,7 @@ from exo.shared.models.model_cards import ModelId, card_cache
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
+    ForwarderCommand,
     ForwarderDownloadCommand,
     StartDownload,
 )

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -96,7 +96,7 @@ class SequentialGenerator(Engine):
     model_id: ModelId
     device_rank: int
     cancel_receiver: MpReceiver[TaskId]
-    event_sender: MpSender[Event]
+    event_sender: MpSender[Event | ChunkGenerated]
     vision_processor: VisionProcessor | None = None
     check_for_cancel_every: int = 50
 
@@ -327,7 +327,7 @@ class BatchGenerator(Engine):
     model_id: ModelId
     device_rank: int
     cancel_receiver: MpReceiver[TaskId]
-    event_sender: MpSender[Event]
+    event_sender: MpSender[Event | ChunkGenerated]
     check_for_cancel_every: int = 50
     vision_processor: VisionProcessor | None = None
 

src/exo/worker/runner/supervisor.py

@@ -53,7 +53,7 @@ class RunnerSupervisor:
     shard_metadata: ShardMetadata
     bound_instance: BoundInstance
     runner_process: mp.Process
-    _ev_recv: MpReceiver[Event]
+    _ev_recv: MpReceiver[Event | ChunkGenerated]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
     _cancel_sender: MpSender[TaskId]
@@ -77,7 +77,7 @@ class RunnerSupervisor:
         event_sender: Sender[Event],
         session: PySession,
     ) -> Self:
-        ev_send, ev_recv = mp_channel[Event]()
+        ev_send, ev_recv = mp_channel[Event | ChunkGenerated]()
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
@@ -215,19 +215,12 @@ class RunnerSupervisor:
             with self._ev_recv as events:
                 async for event in events:
                     if isinstance(event, ChunkGenerated):
-                        if (pub := pubs.get(event.command_id, None)) is None:
-                            pub = pubs[event.command_id] = self.session.net_sender(
+                        if event.command_id not in pubs:
+                            pubs[event.command_id] = self.session.net_sender(
                                 f"runners/{self.bound_instance.bound_runner_id}/active_tasks/{event.command_id}/chunks"
                             )
-                        sent = await pub.send(
-                            event.chunk.model_dump_json().encode("utf-8")
-                        )
-                        if not sent:
-                            logger.warning(
-                                "api node closed communication, dropping chunk"
-                            )
-                        if event.chunk.finish_reason is not None:
-                            pubs.pop(event.command_id, None)
+                        pub = pubs[event.command_id]
+                        await pub.send(event.chunk.model_dump_json().encode("utf-8"))
                         continue
 
                     if isinstance(event, RunnerStatusUpdated):