feat(memory-pool): cross-machine zero-copy tensor transport for end-edge-cloud collaboration

apis/rust/node/src/daemon_connection/interactive.rs

@@ -81,7 +81,11 @@ impl InteractiveEvents {
             }
             DaemonRequest::RegisterPinnedMemory { .. }
             | DaemonRequest::ReadPinnedMemory { .. }
-            | DaemonRequest::FreePinnedMemory { .. } => DaemonReply::Result(Ok(())),
+            | DaemonRequest::FreePinnedMemory { .. }
+            | DaemonRequest::WritePinnedMemory { .. } => DaemonReply::Result(Ok(())),
+            DaemonRequest::RegisterCrossMachinePool { .. } => {
+                eyre::bail!("cross-machine pool registration is not supported in interactive mode")
+            }
             DaemonRequest::NodeConfig { .. } => {
                 eyre::bail!("unexpected NodeConfig in interactive mode")
             }

apis/rust/node/src/daemon_connection/node_integration_testing.rs

@@ -134,7 +134,13 @@ impl IntegrationTestingEvents {
             }
             DaemonRequest::RegisterPinnedMemory { .. }
             | DaemonRequest::ReadPinnedMemory { .. }
-            | DaemonRequest::FreePinnedMemory { .. } => DaemonReply::Result(Ok(())),
+            | DaemonRequest::FreePinnedMemory { .. }
+            | DaemonRequest::WritePinnedMemory { .. } => DaemonReply::Result(Ok(())),
+            DaemonRequest::RegisterCrossMachinePool { .. } => {
+                eyre::bail!(
+                    "cross-machine pool registration is not supported in integration-testing mode"
+                )
+            }
             DaemonRequest::NodeConfig { .. } => {
                 eyre::bail!("unexpected NodeConfig in interactive mode")
             }

apis/rust/node/src/node/control_channel.rs

@@ -199,4 +199,65 @@ impl ControlChannel {
             other => bail!("unexpected FreePinnedMemory reply: {other:?}"),
         }
     }
+
+    pub fn write_pinned_memory(
+        &mut self,
+        shared_memory_id: String,
+        tensor_data: Vec<u8>,
+        size: usize,
+    ) -> eyre::Result<()> {
+        let request = DaemonRequest::WritePinnedMemory {
+            shared_memory_id,
+            tensor_data,
+            size,
+        };
+        let reply = self
+            .channel
+            .request(&Timestamped {
+                inner: request,
+                timestamp: self.clock.new_timestamp(),
+            })
+            .wrap_err("failed to send WritePinnedMemory request to dora-daemon")?;
+        match reply {
+            DaemonReply::Result(Ok(())) => Ok(()),
+            DaemonReply::Result(Err(e)) => bail!("{e}"),
+            other => bail!("unexpected WritePinnedMemory reply: {other:?}"),
+        }
+    }
+
+    /// Register a pool on a remote machine via the daemon (the daemon
+    /// resolves the machine through the coordinator and mirrors the
+    /// pool there with a synchronous confirmation).
+    #[allow(clippy::too_many_arguments)]
+    pub fn register_cross_machine_pool(
+        &mut self,
+        shared_memory_id: String,
+        shmem_name: String,
+        size: usize,
+        dtype: String,
+        shape: Vec<i64>,
+        device: String,
+        machine_id: String,
+    ) -> eyre::Result<(Result<(), String>, bool)> {
+        let request = DaemonRequest::RegisterCrossMachinePool {
+            shared_memory_id,
+            shmem_name,
+            size,
+            dtype,
+            shape,
+            device,
+            machine_id,
+        };
+        let reply = self
+            .channel
+            .request(&Timestamped {
+                inner: request,
+                timestamp: self.clock.new_timestamp(),
+            })
+            .wrap_err("failed to send RegisterCrossMachinePool request to dora-daemon")?;
+        match reply {
+            DaemonReply::CrossMachinePoolRegistered { result, direct } => Ok((result, direct)),
+            other => bail!("unexpected RegisterCrossMachinePool reply: {other:?}"),
+        }
+    }
 }

apis/rust/node/src/node/mod.rs

@@ -2346,6 +2346,47 @@ impl DoraNode {
     pub fn free_pinned_memory(&mut self, shared_memory_id: String) -> Result<(), eyre::Error> {
         self.control_channel.free_pinned_memory(shared_memory_id)
     }
+
+    /// Write tensor bytes to a pinned memory pool via the daemon. The
+    /// daemon forwards the payload to remote daemons so the mirror pool
+    /// is updated in place.
+    pub fn write_pinned_memory(
+        &mut self,
+        shared_memory_id: String,
+        tensor_data: Vec<u8>,
+        size: usize,
+    ) -> Result<(), eyre::Error> {
+        self.control_channel
+            .write_pinned_memory(shared_memory_id, tensor_data, size)
+    }
+
+    /// Register a memory pool on a remote machine via the daemon. The
+    /// daemon resolves the machine through the coordinator and mirrors
+    /// the pool there with a synchronous confirmation, returning
+    /// `Ok(Ok(()))` on success or `Ok(Err(msg))` when the mirror failed
+    /// (unresolved machine, remote pool creation failure, or ack
+    /// timeout).
+    #[allow(clippy::too_many_arguments)]
+    pub fn register_cross_machine_pool(
+        &mut self,
+        shared_memory_id: String,
+        shmem_name: String,
+        size: usize,
+        dtype: String,
+        shape: Vec<i64>,
+        device: String,
+        machine_id: String,
+    ) -> Result<(Result<(), String>, bool), eyre::Error> {
+        self.control_channel.register_cross_machine_pool(
+            shared_memory_id,
+            shmem_name,
+            size,
+            dtype,
+            shape,
+            device,
+            machine_id,
+        )
+    }
 }
 
 /// Builder for initializing a node with custom connection parameters.

binaries/cli/src/command/record.rs

@@ -468,6 +468,11 @@ fn run_record_proxy(args: Record) -> eyre::Result<()> {
                         node_id, output_id, ..
                     } => (node_id.to_string(), output_id.to_string()),
                     InterDaemonEvent::OutputClosed { .. } => continue,
+                    InterDaemonEvent::MemoryPoolWrite { .. } => continue,
+                    InterDaemonEvent::RegisterPool { .. }
+                    | InterDaemonEvent::RegisterPoolAck { .. }
+                    | InterDaemonEvent::MemoryPoolWriteAck { .. }
+                    | InterDaemonEvent::FreePool { .. } => continue,
                 };
 
                 let now_nanos = SystemTime::now()

binaries/cli/src/command/topic/echo.rs

@@ -243,6 +243,11 @@ fn inspect(
             } => {
                 eprintln!("Output {node_id}/{output_id} closed");
             }
+            InterDaemonEvent::MemoryPoolWrite { .. } => {}
+            InterDaemonEvent::RegisterPool { .. }
+            | InterDaemonEvent::RegisterPoolAck { .. }
+            | InterDaemonEvent::MemoryPoolWriteAck { .. }
+            | InterDaemonEvent::FreePool { .. } => {}
         }
     }
 

binaries/cli/src/command/topic/hz.rs

@@ -337,6 +337,11 @@ fn run_hz(
                     }
                 }
                 InterDaemonEvent::OutputClosed { .. } => {}
+                InterDaemonEvent::MemoryPoolWrite { .. } => {}
+                InterDaemonEvent::RegisterPool { .. }
+                | InterDaemonEvent::RegisterPoolAck { .. }
+                | InterDaemonEvent::MemoryPoolWriteAck { .. }
+                | InterDaemonEvent::FreePool { .. } => {}
             }
         }
     });

binaries/cli/src/command/topic/info.rs

@@ -195,6 +195,11 @@ fn info(
                             stats_clone.record(data_size, data_type, Instant::now());
                         }
                         InterDaemonEvent::OutputClosed { .. } => break,
+                        InterDaemonEvent::MemoryPoolWrite { .. } => {}
+                        InterDaemonEvent::RegisterPool { .. }
+                        | InterDaemonEvent::RegisterPoolAck { .. }
+                        | InterDaemonEvent::MemoryPoolWriteAck { .. }
+                        | InterDaemonEvent::FreePool { .. } => {}
                     }
                 }
                 Ok(Err(_)) => continue,

binaries/coordinator/src/lib.rs

@@ -250,12 +250,20 @@ async fn start_with_events(
     #[cfg(feature = "metrics")]
     let otel_metrics = otel_metrics::new_shared();
 
+    // DaemonId -> WS peer address, shared with the WS server so the
+    // ResolveMachine reply can carry the target daemon's direct-TCP data
+    // listener address.
+    let daemon_peer_addrs: Arc<
+        std::sync::RwLock<std::collections::HashMap<String, std::net::SocketAddr>>,
+    > = Arc::new(std::sync::RwLock::new(std::collections::HashMap::new()));
     let (port, ws_shutdown, ws_future) = ws_server::serve(
         bind,
         ws_event_tx.clone(),
         clock.clone(),
         auth_token,
         artifact_store,
+        store.clone(),
+        daemon_peer_addrs.clone(),
     )
     .await
     .wrap_err("failed to start WS server")?;
@@ -274,6 +282,7 @@ async fn start_with_events(
             clock,
             store,
             span_store,
+            daemon_peer_addrs,
             #[cfg(feature = "metrics")]
             otel_metrics,
         )
@@ -299,6 +308,9 @@ async fn start_inner(
     clock: Arc<HLC>,
     store: Arc<dyn CoordinatorStore>,
     span_store: SpanStore,
+    daemon_peer_addrs: Arc<
+        std::sync::RwLock<std::collections::HashMap<String, std::net::SocketAddr>>,
+    >,
     #[cfg(feature = "metrics")] otel_metrics: otel_metrics::SharedMetrics,
 ) -> eyre::Result<()> {
     let daemon_heartbeat_interval =
@@ -427,6 +439,12 @@ async fn start_inner(
                     match version_check_result.map_err(|e| eyre!(e)).and(send_result) {
                         Ok(()) => {
                             let _ = daemon_id_tx.send(daemon_id.clone());
+                            if let Some(peer_addr) = connection.peer_addr {
+                                daemon_peer_addrs
+                                    .write()
+                                    .unwrap_or_else(|e| e.into_inner())
+                                    .insert(daemon_id.to_string(), peer_addr);
+                            }
                             daemon_connections.add(daemon_id.clone(), connection);
                             if let Err(e) =
                                 store.register_daemon(dora_coordinator_store::DaemonInfo {

binaries/coordinator/src/state.rs

@@ -98,6 +98,10 @@ pub(crate) struct DaemonConnection {
     pub(crate) sender: mpsc::Sender<String>,
     /// Shared with the ws_daemon handler task to resolve correlation-based replies.
     pub(crate) pending_replies: Arc<Mutex<HashMap<Uuid, oneshot::Sender<String>>>>,
+    /// The daemon's WS peer address as seen by the coordinator (set at
+    /// registration). Lets other daemons reach this daemon's direct-TCP
+    /// memory-pool data listener.
+    pub(crate) peer_addr: Option<std::net::SocketAddr>,
     pub(crate) last_heartbeat: Instant,
     pub(crate) labels: BTreeMap<String, String>,
     /// Latest fault tolerance stats from this daemon (updated on each heartbeat).
@@ -131,6 +135,7 @@ impl DaemonConnection {
         Self {
             sender,
             pending_replies,
+            peer_addr: None,
             last_heartbeat: Instant::now(),
             labels,
             ft_stats: None,

binaries/coordinator/src/ws_daemon.rs

@@ -3,10 +3,12 @@ use crate::{
     state::DaemonConnection,
 };
 use axum::extract::ws::{Message, WebSocket};
+use dora_coordinator_store::CoordinatorStore;
 use dora_core::uhlc::HLC;
 use dora_message::{
     common::DaemonId,
-    daemon_to_coordinator::{CoordinatorRequest, DaemonEvent},
+    coordinator_to_daemon::ResolveMachineReply,
+    daemon_to_coordinator::{CoordinatorRequest, DaemonEvent, Timestamped},
     ws_protocol::WsResponse,
 };
 use futures::{SinkExt, StreamExt};
@@ -22,6 +24,11 @@ pub(crate) async fn handle_daemon_ws(
     socket: WebSocket,
     event_tx: mpsc::Sender<Event>,
     clock: Arc<HLC>,
+    store: Arc<dyn CoordinatorStore>,
+    peer_addr: std::net::SocketAddr,
+    daemon_peer_addrs: Arc<
+        std::sync::RwLock<std::collections::HashMap<String, std::net::SocketAddr>>,
+    >,
 ) {
     let (mut ws_tx, mut ws_rx) = socket.split();
 
@@ -74,9 +81,14 @@ pub(crate) async fn handle_daemon_ws(
                         &clock,
                         &cmd_tx,
                         &pending_replies,
+                        &store,
                         &mut tracked_daemon_id,
                         &mut tracked_connection_id,
-                    ).await {
+                        peer_addr,
+                        daemon_peer_addrs.clone(),
+                    )
+                    .await
+                    {
                         break;
                     }
                 } else {
@@ -112,20 +124,29 @@ pub(crate) async fn handle_daemon_ws(
 /// instead of going through the `serde_json::Value` used for routing.
 #[derive(serde::Deserialize)]
 struct DaemonWsRequestRaw {
+    /// Request id from the daemon envelope — echoed back in replies so
+    /// the daemon can route the reply to its pending caller.
+    id: Uuid,
     params: dora_message::daemon_to_coordinator::Timestamped<
         dora_message::daemon_to_coordinator::CoordinatorRequest,
     >,
 }
 
 /// Handle a daemon request (event or register). Returns false if the event channel closed.
+#[allow(clippy::too_many_arguments)]
 async fn handle_daemon_request(
     raw_text: &str,
     event_tx: &mpsc::Sender<Event>,
     clock: &HLC,
     cmd_tx: &mpsc::Sender<String>,
     pending_replies: &Arc<Mutex<HashMap<Uuid, oneshot::Sender<String>>>>,
+    store: &Arc<dyn CoordinatorStore>,
     tracked_daemon_id: &mut Option<DaemonId>,
     tracked_connection_id: &mut Option<Uuid>,
+    peer_addr: std::net::SocketAddr,
+    daemon_peer_addrs: Arc<
+        std::sync::RwLock<std::collections::HashMap<String, std::net::SocketAddr>>,
+    >,
 ) -> bool {
     let parsed: DaemonWsRequestRaw = match serde_json::from_str(raw_text) {
         Ok(m) => m,
@@ -135,6 +156,7 @@ async fn handle_daemon_request(
         }
     };
     let message = parsed.params;
+    let request_id = parsed.id;
 
     if let Err(err) = clock.update_with_timestamp(&message.timestamp) {
         tracing::warn!("failed to update coordinator clock: {err}");
@@ -165,6 +187,7 @@ async fn handle_daemon_request(
             let mut connection =
                 DaemonConnection::new(cmd_tx.clone(), pending_replies.clone(), labels.clone());
             connection.supports_hub_sources = supports_hub_sources;
+            connection.peer_addr = Some(peer_addr);
             // Capture the connection_id before moving `connection` into the event.
             let connection_id = connection.connection_id;
             let (daemon_id_tx, daemon_id_rx) = oneshot::channel();
@@ -209,6 +232,49 @@ async fn handle_daemon_request(
                 true
             }
         }
+        CoordinatorRequest::ResolveMachine { machine_id } => {
+            // Resolve the machine id against the registered-daemon store;
+            // unknown machines (or store errors) resolve to `found: false`.
+            // Also report the target daemon's WS peer address so the
+            // requesting daemon can reach its direct-TCP data listener.
+            let (found, address) = match store.get_daemon_by_machine(&machine_id) {
+                Ok(Some(d)) => (
+                    true,
+                    daemon_peer_addrs
+                        .read()
+                        .unwrap_or_else(|e| e.into_inner())
+                        .get(&d.to_string())
+                        .copied(),
+                ),
+                Ok(None) => (false, None),
+                Err(e) => {
+                    tracing::warn!("failed to resolve machine `{machine_id}`: {e}");
+                    (false, None)
+                }
+            };
+            // Reply over the same WS envelope the Register flow uses
+            // (`{"id", "method": "daemon_event", "params": <Timestamped<...>>}`),
+            // mirroring `DaemonConnection::send`.
+            let reply = Timestamped {
+                inner: ResolveMachineReply::ResolveMachineResult { found, address },
+                timestamp: clock.new_timestamp(),
+            };
+            let params = match serde_json::to_string(&reply) {
+                Ok(params) => params,
+                Err(err) => {
+                    tracing::warn!("failed to serialize ResolveMachine reply: {err}");
+                    return true;
+                }
+            };
+            // Echo the request id so the daemon can route this reply to
+            // its pending caller (COORDINATOR_PENDING).
+            let json =
+                format!(r#"{{"id":"{request_id}","method":"daemon_event","params":{params}}}"#);
+            if cmd_tx.send(json).await.is_err() {
+                return false;
+            }
+            true
+        }
     }
 }