custom discovery take 2

Cargo.lock

@@ -319,6 +319,26 @@ version = "0.6.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e"
 
+[[package]]
+name = "bytemuck"
+version = "1.25.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec"
+dependencies = [
+ "bytemuck_derive",
+]
+
+[[package]]
+name = "bytemuck_derive"
+version = "1.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.117",
+]
+
 [[package]]
 name = "byteorder"
 version = "1.5.0"
@@ -798,7 +818,6 @@ dependencies = [
  "rand 0.10.1",
  "serde_json",
  "tokio",
- "tracing",
  "zenoh",
 ]
 
@@ -1821,6 +1840,7 @@ name = "networking"
 version = "0.0.1"
 dependencies = [
  "async-stream",
+ "bytemuck",
  "futures-lite",
  "log",
  "netwatcher",
@@ -3671,7 +3691,6 @@ dependencies = [
  "signal-hook-registry",
  "socket2 0.6.1",
  "tokio-macros",
- "tracing",
  "windows-sys 0.61.2",
 ]
 

Cargo.toml

@@ -21,10 +21,21 @@ opt-level = 3
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
 
-# Macro dependecies
+# pyo3
+pyo3 = "0.27.2"
+pyo3-async-runtimes = "0.27.0"
+pyo3-log = "0.13.2"
+pyo3-stub-gen = "0.22.2"
+
+# util
 extend = "1.2"
 tokio = "1.46"
 futures-lite = "2.6.1"
+async-stream = "0.3.6"
+pin-project = "1.1.10"
+serde_json = "1.0.149"
+rand = "0.10.1"
+parking_lot = "0.12.5"
 
 # Tracing/logging
 log = "0.4"
@@ -32,19 +43,10 @@ env_logger = "0.11.10"
 
 # networking
 zenoh = "=1.9.0"
-async-stream = "0.3.6"
-netwatcher = "0.6.0"
-parking_lot = "0.12.5"
-pin-project = "1.1.10"
-pyo3 = "0.27.2"
-pyo3-async-runtimes = "0.27.0"
-pyo3-log = "0.13.2"
-pyo3-stub-gen = "0.22.2"
-rand = "0.10.1"
-serde_json = "1.0.149"
-tracing = "0.1.44"
 zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
 zenoh-plugin-trait = "=1.9.0"
+netwatcher = "0.6.0"
+bytemuck = "1.25.0"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

rust/exo_net/Cargo.toml

@@ -36,7 +36,7 @@ pyo3-async-runtimes = { workspace = true, features = [
 pyo3-log.workspace = true
 
 # async runtime
-tokio = { workspace = true, features = ["full", "tracing"] }
+tokio = { workspace = true, features = ["full"] }
 futures-lite.workspace = true
 pin-project.workspace = true
 
@@ -49,4 +49,3 @@ zenoh.workspace = true
 rand.workspace = true
 serde_json.workspace = true
 parking_lot.workspace = true
-tracing.workspace = true

rust/networking/Cargo.toml

@@ -14,6 +14,7 @@ zenoh-plugin-storage-manager.workspace = true
 zenoh-plugin-trait.workspace = true
 rand.workspace = true
 log.workspace = true
+bytemuck = { workspace = true, features = ["derive"] }
 
 [lints]
 workspace = true

rust/networking/examples/put_string.rs

@@ -7,7 +7,7 @@ async fn main() -> Result<()> {
     zenoh::init_log_from_env_or("info");
     info!("Opening session...");
     let cfg = networking::cfg(rand::random(), 0)?;
-    let session = networking::open(cfg).await?;
+    let session = networking::open(cfg, 52414).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/examples/serve_storage.rs

@@ -6,8 +6,8 @@ use zenoh::Result;
 async fn main() -> Result<()> {
     zenoh::init_log_from_env_or("info");
     info!("Opening session...");
-    let cfg = networking::cfg(rand::random(), 0)?;
-    let session = networking::open(cfg).await?;
+    let cfg = networking::cfg(rand::random(), 52414)?;
+    let session = networking::open(cfg, 52414).await?;
     let _tok = session
         .z
         .liveliness()

rust/networking/src/discovery.rs

@@ -0,0 +1,311 @@
+use std::{
+    io::{self, ErrorKind},
+    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
+    sync::Arc,
+    time::Duration,
+};
+
+use bytemuck::{Pod, Zeroable};
+use log::{debug, trace, warn};
+use netwatcher::WatchHandle;
+use parking_lot::Mutex;
+use tokio::{
+    net::UdpSocket,
+    time::{Interval, interval},
+};
+use zenoh::config::ZenohId;
+
+const GROUP: Ipv6Addr = Ipv6Addr::new(0xff12, 0, 0, 0, 0, 0, 0xe0a1, 0xde89);
+
+pub struct Discovery {
+    sock: Arc<UdpSocket>,
+    ifaces: Arc<Mutex<Vec<SocketAddr>>>,
+    last_nonce: Mutex<[u8; 8]>,
+    /// the port of the service we are doing discovery for - transmitted to peers
+    listen_port: u16,
+    zid: ZenohId,
+    tick: Interval,
+    _sync: Mutex<WatchHandle>,
+}
+
+impl Discovery {
+    pub async fn new(zid: ZenohId, listen_port: u16) -> io::Result<Self> {
+        let discovery_port = 52413;
+        let sock = Arc::new(UdpSocket::bind(format!("[::]:{discovery_port}")).await?);
+        //sock.set_multicast_loop_v6(false)?;
+        let ifaces: Arc<Mutex<Vec<SocketAddr>>> = Default::default();
+        let _sync = Mutex::new(
+            netwatcher::watch_interfaces_with_callback({
+                let sock = sock.clone();
+                let ifaces = ifaces.clone();
+                move |update| {
+                    for (iface_idx, iface) in update.interfaces.iter() {
+                        if iface
+                            .ipv6_ips()
+                            .all(|addr| addr.is_loopback() || addr.is_unspecified())
+                        {
+                            continue;
+                        }
+                        if let Err(e) = sock.join_multicast_v6(&GROUP, *iface_idx).inspect(|_| {
+                            ifaces.lock().push(SocketAddr::V6(SocketAddrV6::new(
+                                GROUP, 52413, 0, *iface_idx,
+                            )))
+                        }) {
+                            if let Some(iface) = update.interfaces.get(&iface_idx) {
+                                warn!(
+                                    "failed to join multicast v6 for interface {}: {e}",
+                                    iface.name
+                                )
+                            }
+                        }
+                    }
+                    for iface_idx in update.diff.removed {
+                        ifaces.lock().retain(|addr| {
+                            if let SocketAddr::V6(v6) = addr {
+                                v6.scope_id() != iface_idx
+                            } else {
+                                true
+                            }
+                        });
+                        if let Err(e) = sock.leave_multicast_v6(&GROUP, iface_idx) {
+                            if let Some(iface) = update.interfaces.get(&iface_idx) {
+                                warn!(
+                                    "failed to leave multicast v6 for interface {}: {e}",
+                                    iface.name
+                                )
+                            }
+                        }
+                    }
+                }
+            })
+            // todo: better error handling here
+            .expect("failed to bind discovery watcher"),
+        );
+        Ok(Self {
+            sock,
+            ifaces,
+            last_nonce: Default::default(),
+            listen_port,
+            zid,
+            tick: interval(Duration::from_secs(1)),
+            _sync,
+        })
+    }
+
+    pub async fn next(&mut self) -> io::Result<Discovered> {
+        let mut buf = [0u8; Hello::buf_size() + WhatsUp::buf_size() + 1];
+        loop {
+            tokio::select! {
+                _ = self.tick.tick() => {
+                    self.announce().await?;
+                }
+                res = self.sock.recv_from(&mut buf) => {
+                    let Ok((bytes_read, addr)) = res else { continue; };
+                    if let Some(discovered) = self.respond(bytes_read, addr, &mut buf).await? {
+                        return Ok(discovered)
+                    }
+                }
+            }
+        }
+    }
+
+    async fn respond(
+        &self,
+        bytes_read: usize,
+        addr: SocketAddr,
+        buf: &mut [u8],
+    ) -> io::Result<Option<Discovered>> {
+        trace!(
+            "raw recv: {bytes_read} bytes from {addr}: {:02x?}",
+            &buf[..bytes_read]
+        );
+        if bytes_read < size_of::<Header>() {
+            trace!("dropped: early EOF");
+            return Ok(None);
+        };
+        let header: &Header = bytemuck::from_bytes(&buf[0..size_of::<Header>()]);
+        if header.magic != *b"EXO" {
+            trace!("dropped: wrong magic");
+            return Ok(None);
+        };
+        let Ok(kind) = header.kind.try_into() else {
+            trace!("dropped: unknown message kind {}", header.kind);
+            return Ok(None);
+        };
+        match kind {
+            Kind::Hello => {
+                let total = Hello::buf_size();
+                if bytes_read != total {
+                    trace!("dropped: hello wrong size");
+                    return Ok(None);
+                }
+                let hello: &Hello = bytemuck::from_bytes(&buf[size_of::<Header>()..total]);
+                if hello.nonce == *self.last_nonce.lock() {
+                    trace!("dropped: local hello nonce");
+                    return Ok(None);
+                }
+
+                // reply
+                let mut reply_buf = [0u8; WhatsUp::buf_size()];
+                let reply = WhatsUp {
+                    nonce: hello.nonce,
+                    zid: self.zid.to_le_bytes(),
+                    port_le: self.listen_port.to_le_bytes(),
+                };
+                reply.write_into(&mut reply_buf);
+
+                for i in 0..4 {
+                    if self
+                        .sock
+                        .send_to(&reply_buf, addr)
+                        .await
+                        .inspect_err(|e| debug!("send to {addr} failed: {e}"))
+                        .is_ok_and(|sent| sent == WhatsUp::buf_size())
+                    {
+                        trace!(
+                            "sent {} bytes to {addr} after {} attempt(s)",
+                            WhatsUp::buf_size(),
+                            i + 1
+                        );
+                        break;
+                    };
+                    tokio::time::sleep(Duration::from_millis(300)).await;
+                }
+                Ok(None)
+            }
+            Kind::WhatsUp => {
+                let total = WhatsUp::buf_size();
+                if bytes_read != total {
+                    trace!("dropped: whatsup wrong size");
+                    return Ok(None);
+                }
+                let whats_up: &WhatsUp = bytemuck::from_bytes(&buf[size_of::<Header>()..total]);
+                if whats_up.nonce == [0u8; 8] || whats_up.nonce != *self.last_nonce.lock() {
+                    trace!("dropped: stale nonce");
+                    return Ok(None);
+                }
+                let SocketAddr::V6(v6) = addr else {
+                    trace!("dropped: v4 addr used");
+                    return Ok(None);
+                };
+                let Ok(zid) = ZenohId::try_from(&whats_up.zid[..]) else {
+                    trace!("dropped: zenoh conversion failed");
+                    return Ok(None);
+                };
+                if zid == self.zid {
+                    trace!("dropped: self zenoh id");
+                    return Ok(None);
+                }
+                // discovered
+                let addr = {
+                    let mut x = v6.clone();
+                    x.set_port(u16::from_le_bytes(whats_up.port_le));
+                    x
+                };
+                Ok(Some(Discovered { addr, zid }))
+            }
+        }
+    }
+
+    async fn announce(&self) -> io::Result<()> {
+        let nonce = rand::random();
+        *self.last_nonce.lock() = nonce;
+        let hello = Hello { nonce };
+
+        let mut buf = [0u8; Hello::buf_size()];
+        hello.write_into(&mut buf);
+
+        let addrs = self.ifaces.lock().clone();
+        debug!("announcing {hello:?} to {addrs:?}");
+        // rev so .remove() doesn't break things
+        for (i, addr) in addrs.into_iter().enumerate().rev() {
+            match self.sock.send_to(&buf, addr).await {
+                Ok(bytes) => trace!("sent {bytes} to {addr}"),
+                Err(e) if e.kind() == ErrorKind::HostUnreachable => {
+                    debug!("disabling discovery address {addr}: {e}");
+                    _ = self.ifaces.lock().swap_remove(i);
+                }
+                Err(e) => debug!("failed to reach {addr}: {e}"),
+            }
+        }
+        Ok(())
+    }
+}
+
+pub trait Message: Pod {
+    const KIND: Kind;
+    fn header() -> Header {
+        Header {
+            magic: *b"EXO",
+            kind: Self::KIND as u8,
+        }
+    }
+    fn write_into(&self, buf: &mut [u8]) {
+        let total = size_of::<Header>() + size_of::<Self>();
+        assert!(total <= buf.len());
+        buf[0..size_of::<Header>()].copy_from_slice(bytemuck::bytes_of(&Self::header()));
+        buf[size_of::<Header>()..total].copy_from_slice(bytemuck::bytes_of(self));
+    }
+}
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy)]
+// packet & version
+pub enum Kind {
+    Hello = 0,
+    WhatsUp = 1,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Discovered {
+    pub zid: ZenohId,
+    pub addr: SocketAddrV6,
+}
+
+pub struct UnknownKind;
+impl TryFrom<u8> for Kind {
+    type Error = UnknownKind;
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Kind::Hello),
+            1 => Ok(Kind::WhatsUp),
+            _ => Err(UnknownKind),
+        }
+    }
+}
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable)]
+pub struct Header {
+    magic: [u8; 3],
+    kind: u8,
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable)]
+pub struct Hello {
+    pub nonce: [u8; 8],
+}
+impl Hello {
+    const fn buf_size() -> usize {
+        size_of::<Header>() + size_of::<Self>()
+    }
+}
+impl Message for Hello {
+    const KIND: Kind = Kind::Hello;
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Pod, Zeroable)]
+pub struct WhatsUp {
+    pub nonce: [u8; 8],
+    pub zid: [u8; 16],
+    pub port_le: [u8; 2],
+}
+impl WhatsUp {
+    const fn buf_size() -> usize {
+        size_of::<Header>() + size_of::<Self>()
+    }
+}
+impl Message for WhatsUp {
+    const KIND: Kind = Kind::WhatsUp;
+}

rust/networking/src/lib.rs

@@ -1,24 +1,26 @@
-use std::{env, panic, sync::Arc};
+use std::env;
 
-use netwatcher::WatchHandle;
-use parking_lot::Mutex;
-use tokio::{sync::mpsc, task::JoinHandle};
-use zenoh::{Result, Session as ZSession, config::WhatAmI, internal::runtime::Runtime};
+use tokio::task::JoinHandle;
+use zenoh::{Result, Session as ZSession, config::Locator};
 use zenoh_plugin_storage_manager::StoragesPlugin;
 use zenoh_plugin_trait::PluginsManager;
 
 pub use zenoh::{Config, config::ZenohId};
 
+use crate::discovery::Discovery;
+
+pub mod discovery;
 pub mod swarm;
 
 pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
+    assert!(listen_port != 0, "must used defined listen port port");
     let namespace = env::var("EXO_ZENOH_NAMESPACE").unwrap_or_else(|_| "exo".to_string());
     let mut cfg = zenoh::Config::default();
     // todo: cleanup
     cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
-    cfg.insert_json5("mode", "\"peer\"")?;
+    cfg.insert_json5("mode", "\"router\"")?;
     cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
-    cfg.insert_json5("scouting/multicast/enabled", "true")?;
+    cfg.insert_json5("scouting/multicast/enabled", "false")?;
     cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
     cfg.insert_json5("scouting/gossip/multihop", "true")?;
     cfg.insert_json5("namespace", &format!("{namespace:?}"))?;
@@ -39,7 +41,8 @@ pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
     Ok(cfg)
 }
 
-pub async fn open(cfg: zenoh::Config) -> Result<Session> {
+pub async fn open(cfg: zenoh::Config, listen_port: u16) -> Result<Session> {
+    assert!(listen_port != 0, "must used defined listen port");
     let mut plugins = PluginsManager::static_plugins_only();
     plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
     let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
@@ -48,57 +51,42 @@ pub async fn open(cfg: zenoh::Config) -> Result<Session> {
         .await?;
     let z = zenoh::session::init(runtime.clone().into()).await?;
     runtime.start().await?;
-    let _watch_all_handle = watch_all(runtime).await?;
-    Ok(Session {
-        z,
-        _watch_all_handle,
-    })
-}
-async fn watch_all(runtime: Runtime) -> Result<WatchAllHandle> {
-    log::info!("spawning scout");
-    let mut cfg = Config::default();
-    cfg.insert_json5("scouting/multicast/ttl", "3")?;
-    cfg.insert_json5("scouting/multicast/interface", "\"auto\"")?;
-    let mut scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
-    let (send, mut recv) = mpsc::unbounded_channel();
-    let _sync = Arc::new(Mutex::new(netwatcher::watch_interfaces_with_callback(
-        move |u| _ = send.send(u),
-    )?));
-    let _async = tokio::task::spawn(async move {
+    let mut discovery = Discovery::new(z.zid(), listen_port).await?;
+    let _jh = tokio::task::spawn(async move {
         loop {
-            tokio::select! {
-                u = recv.recv() => {
-                    if u.is_none() {
-                        return Ok(());
-                    }
-                    log::info!("reloading scout");
-                    scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
-                }
-                hello = scout.recv_async() => {
-                    if let Ok(hello) = hello {
-                        // todo: auth
-                        runtime
-                            .connect_peer(&hello.zid().into(), hello.locators())
-                            .await;
-                    }
-                }
+            let Ok(discovered) = discovery.next().await.inspect_err(|e| {
+                log::warn!("discovery error {e}");
+            }) else {
+                continue;
+            };
+
+            if discovered.zid > runtime.zid() {
+                log::debug!("not connecting to peer with greater zid");
+                continue;
             }
+
+            let Ok(locator) =
+                Locator::new("tcp", discovered.addr.to_string(), "").inspect_err(|e| {
+                    log::warn!("failed to pass locator from addr: {e}");
+                })
+            else {
+                continue;
+            };
+
+            runtime
+                .connect_peer(&discovered.zid.into(), &[locator])
+                .await;
         }
     });
-    Ok(WatchAllHandle { _sync, _async })
+    Ok(Session { z, _jh })
 }
 
 pub struct Session {
     pub z: ZSession,
-    _watch_all_handle: WatchAllHandle,
+    _jh: JoinHandle<()>,
 }
-impl Drop for WatchAllHandle {
+impl Drop for Session {
     fn drop(&mut self) {
-        self._async.abort();
+        self._jh.abort();
     }
 }
-
-pub struct WatchAllHandle {
-    _sync: Arc<Mutex<WatchHandle>>,
-    _async: JoinHandle<Result<()>>,
-}

rust/networking/src/swarm.rs

@@ -179,8 +179,8 @@ pub async fn create_swarm(
     if !bootstrap_peers.is_empty() || listen_port != 0 {
         todo!();
     }
-    let cfg = crate::cfg(identity, listen_port)?;
-    let session = crate::open(cfg).await?;
+    let cfg = crate::cfg(identity, 52414)?;
+    let session = crate::open(cfg, 52414).await?;
     Ok(Swarm {
         session,
         from_client,

src/exo/api/main.py

@@ -199,7 +199,7 @@ from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
-from exo.utils.channels import Receiver, Sender
+from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.power_sampler import PowerSampler
 from exo.utils.task_group import TaskGroup
@@ -241,11 +241,16 @@ class Transport:
     command_sender: NetSender = field(init=False)
     paused: bool = field(init=False, default=False)
     paused_ev: anyio.Event = field(init=False, default_factory=anyio.Event)
+    tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     def __post_init__(self):
         # TODO: retire root keyspace
         self.command_sender = self.session.net_sender("orchestrator")
 
+    async def run(self):
+        async with self.tg:
+            await anyio.sleep_forever()
+
     async def send_command(self, command: Command) -> bool:
         while self.paused:
             await self.paused_ev.wait()
@@ -270,7 +275,14 @@ class Transport:
     async def stream(
         self,
         command_id: CommandId,
-    ) -> AsyncGenerator[Chunk,]:
+    ) -> AsyncGenerator[Chunk]:
+        send, recv = channel[Chunk]()
+        self.tg.start_soon(self._run_stream, command_id, send)
+        async with recv:
+            async for item in recv:
+                yield item
+
+    async def _run_stream(self, command_id: CommandId, send: Sender[Chunk]):
         try:
             with anyio.CancelScope() as cs:
                 self.cancel_scopes[command_id] = cs
@@ -285,12 +297,11 @@ class Transport:
                             "stream terminated early without finish reason EOF"
                         )
                         break
-                    yield (
-                        chunk := cast(
-                            Chunk,
-                            TypeAdapter(Chunk).validate_json(
+                    await send.send(
+                        chunk := (
+                            TypeAdapter[Chunk](Chunk).validate_json(
                                 data, strict=True, extra="forbid"
-                            ),
+                            )
                         )
                     )
                     if (
@@ -298,7 +309,7 @@ class Transport:
                         and chunk.finish_reason is not None
                     ):
                         break
-        except anyio.get_cancelled_exc_class():
+        except (anyio.get_cancelled_exc_class(), anyio.BrokenResourceError, anyio.ClosedResourceError):
             with anyio.CancelScope(shield=True):
                 await self.command_sender.send(
                     TaskCancelled(cancelled_command_id=command_id)
@@ -315,7 +326,7 @@ class Transport:
                 )
 
     def cancel(self, command_id: CommandId) -> bool:
-        if (cs := self.cancel_scopes.get(command_id, None)) is not None:
+        if (cs := self.cancel_scopes.pop(command_id, None)) is not None:
             cs.cancel()
             return True
         return False
@@ -1839,6 +1850,7 @@ class API:
         try:
             async with self._tg as tg:
                 logger.info("Starting API")
+                tg.start_soon(self.transport.run)
                 tg.start_soon(self._apply_state)
                 tg.start_soon(self._pause_on_new_election)
                 tg.start_soon(self._cleanup_expired_images)

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")
+            command_receiver=session.net_receiver("orchestrator"),
         )
 
         er_send, er_recv = channel[ElectionResult]()

src/exo/master/main.py

@@ -123,7 +123,7 @@ class Master:
         node_id: NodeId,
         session_id: SessionId,
         *,
-        command_receiver: NetReceiver, # todo: not this type
+        command_receiver: NetReceiver,  # todo: not this type
         event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
@@ -196,9 +196,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 +205,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 +243,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 +252,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()
@@ -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()

src/exo/worker/runner/runner.py

@@ -86,7 +86,7 @@ class Runner:
         self,
         bound_instance: BoundInstance,
         builder: Builder,
-        event_sender: MpSender[Event | ChunkGenerated],
+        event_sender: MpSender[Event],
         task_receiver: MpReceiver[Task],
     ):
         self.event_sender = event_sender