[hyperactor] mesh: implement local proc bypasss

Pull Request resolved: #1442

This change allows procs local to a host to bypass the host multiplexer. This should improve multicast delivery, as local comm actors do not need to traverse the same (frontend) bottleneck.
ghstack-source-id: 319797871
@exported-using-ghexport

Differential Revision: [D83996264](https://our.internmc.facebook.com/intern/diff/D83996264/)

Author: mariusae

hyperactor/src/mailbox.rs

@@ -164,7 +164,8 @@ pub type Data = Vec<u8>;
     Deserialize,
     Named,
     Clone,
-    PartialEq
+    PartialEq,
+    Eq
 )]
 pub enum DeliveryError {
     /// The destination address is not reachable.

hyperactor/src/mailbox/undeliverable.rs

@@ -32,6 +32,13 @@ use crate::supervision::ActorSupervisionEvent;
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Named)]
 pub struct Undeliverable<M: Message>(pub M);
 
+impl<M: Message> Undeliverable<M> {
+    /// Return the inner M-typed message.
+    pub fn into_inner(self) -> M {
+        self.0
+    }
+}
+
 // Port handle and receiver for undeliverable messages.
 pub(crate) fn new_undeliverable_port() -> (
     PortHandle<Undeliverable<MessageEnvelope>>,

hyperactor_mesh/Cargo.toml

@@ -75,6 +75,7 @@ serde_bytes = "0.11"
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 serde_multipart = { version = "0.0.0", path = "../serde_multipart" }
 strum = { version = "0.27.1", features = ["derive"] }
+tempfile = "3.22"
 thiserror = "2.0.12"
 tokio = { version = "1.47.1", features = ["full", "test-util", "tracing"] }
 tokio-stream = { version = "0.1.17", features = ["fs", "io-util", "net", "signal", "sync", "time"] }
@@ -88,7 +89,6 @@ criterion = { version = "0.5.1", features = ["async_tokio", "csv_output"] }
 itertools = "0.14.0"
 maplit = "1.0"
 proptest = "1.5"
-tempfile = "3.22"
 timed_test = { version = "0.0.0", path = "../timed_test" }
 tracing-test = { version = "0.2.3", features = ["no-env-filter"] }
 

hyperactor_mesh/src/bootstrap.rs

@@ -15,6 +15,7 @@ use std::future;
 use std::io;
 use std::io::Write;
 use std::os::unix::process::ExitStatusExt;
+use std::path::Path;
 use std::path::PathBuf;
 use std::process::Stdio;
 use std::sync::Arc;
@@ -35,6 +36,7 @@ use hyperactor::ProcId;
 use hyperactor::attrs::Attrs;
 use hyperactor::channel;
 use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::ChannelError;
 use hyperactor::channel::ChannelTransport;
 use hyperactor::channel::Rx;
 use hyperactor::channel::Tx;
@@ -50,10 +52,13 @@ use hyperactor::host::HostError;
 use hyperactor::host::ProcHandle;
 use hyperactor::host::ProcManager;
 use hyperactor::host::TerminateSummary;
+use hyperactor::mailbox::IntoBoxedMailboxSender;
+use hyperactor::mailbox::MailboxClient;
 use hyperactor::mailbox::MailboxServer;
 use hyperactor::proc::Proc;
 use serde::Deserialize;
 use serde::Serialize;
+use tempfile::TempDir;
 use tokio::process::Child;
 use tokio::process::ChildStderr;
 use tokio::process::ChildStdout;
@@ -68,6 +73,8 @@ use crate::v1;
 use crate::v1::host_mesh::mesh_agent::HostAgentMode;
 use crate::v1::host_mesh::mesh_agent::HostMeshAgent;
 
+mod mailbox;
+
 declare_attrs! {
     /// If enabled (default), bootstrap child processes install
     /// `PR_SET_PDEATHSIG(SIGKILL)` so the kernel reaps them if the
@@ -228,6 +235,10 @@ pub enum Bootstrap {
         backend_addr: ChannelAddr,
         /// The callback address used to indicate successful spawning.
         callback_addr: ChannelAddr,
+        /// Directory for storing proc socket files. Procs place their sockets
+        /// in this directory, so that they can be looked up by other procs
+        /// for direct transfer.
+        socket_dir_path: PathBuf,
         /// Optional config snapshot (`hyperactor::config::Attrs`)
         /// captured by the parent. If present, the child installs it
         /// as the `Runtime` layer so the parent's effective config
@@ -340,6 +351,7 @@ impl Bootstrap {
                 proc_id,
                 backend_addr,
                 callback_addr,
+                socket_dir_path,
                 config,
             } => {
                 if let Some(attrs) = config {
@@ -359,15 +371,39 @@ impl Bootstrap {
                     eprintln!("(bootstrap) PDEATHSIG disabled via config");
                 }
 
-                let result =
-                    host::spawn_proc(proc_id, backend_addr, callback_addr, |proc| async move {
-                        ProcMeshAgent::boot_v1(proc).await
-                    })
-                    .await;
-                match result {
-                    Ok(_proc) => halt().await,
-                    Err(e) => e.into(),
-                }
+                let (local_addr, name) = ok!(proc_id
+                    .as_direct()
+                    .ok_or_else(|| anyhow::anyhow!("invalid proc id type: {}", proc_id)));
+                // TODO provide a direct way to construct these
+                let serve_addr = format!("unix:{}", socket_dir_path.join(name).display());
+                let serve_addr = serve_addr.parse().unwrap();
+
+                // The following is a modified host::spawn_proc to support direct
+                // dialing between local procs: 1) we bind each proc to a deterministic
+                // address in socket_dir_path; 2) we use LocalProcDialer to dial these
+                // addresses for local procs.
+                let proc_sender = mailbox::LocalProcDialer::new(
+                    local_addr.clone(),
+                    socket_dir_path,
+                    ok!(MailboxClient::dial(backend_addr)),
+                );
+
+                let proc = Proc::new(proc_id.clone(), proc_sender.into_boxed());
+
+                let agent_handle = ok!(ProcMeshAgent::boot_v1(proc.clone())
+                    .await
+                    .map_err(|e| HostError::AgentSpawnFailure(proc_id, e)));
+
+                // Finally serve the proc on the same transport as the backend address,
+                // and call back.
+                let (proc_addr, proc_rx) = ok!(channel::serve(serve_addr, "proc_backend"));
+                proc.clone().serve(proc_rx);
+                ok!(ok!(channel::dial(callback_addr))
+                    .send((proc_addr, agent_handle.bind::<ProcMeshAgent>()))
+                    .await
+                    .map_err(ChannelError::from));
+
+                halt().await
             }
             Bootstrap::Host {
                 addr,
@@ -385,7 +421,8 @@ impl Bootstrap {
                     Some(command) => command,
                     None => ok!(BootstrapCommand::current()),
                 };
-                let manager = BootstrapProcManager::new(command);
+                let manager = BootstrapProcManager::new(command).unwrap();
+
                 let (host, _handle) = ok!(Host::serve(manager, addr).await);
                 let addr = host.addr().clone();
                 let host_mesh_agent = ok!(host
@@ -1437,6 +1474,11 @@ pub struct BootstrapProcManager {
     /// FileMonitor that aggregates logs from all children.
     /// None if file monitor creation failed.
     file_appender: Option<Arc<crate::logging::FileAppender>>,
+
+    /// Directory for storing proc socket files. Procs place their sockets
+    /// in this directory, so that they can be looked up by other procs
+    /// for direct transfer.
+    socket_dir: TempDir,
 }
 
 impl Drop for BootstrapProcManager {
@@ -1486,7 +1528,7 @@ impl BootstrapProcManager {
     /// This is the general entry point when you want to manage procs
     /// backed by a specific binary path (e.g. a bootstrap
     /// trampoline).
-    pub(crate) fn new(command: BootstrapCommand) -> Self {
+    pub(crate) fn new(command: BootstrapCommand) -> Result<Self, io::Error> {
         let log_monitor = match crate::logging::FileAppender::new() {
             Some(fm) => {
                 tracing::info!("log monitor created successfully");
@@ -1498,19 +1540,25 @@ impl BootstrapProcManager {
             }
         };
 
-        Self {
+        Ok(Self {
             command,
             children: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
             pid_table: Arc::new(std::sync::Mutex::new(HashMap::new())),
             file_appender: log_monitor,
-        }
+            socket_dir: runtime_dir()?,
+        })
     }
 
     /// The bootstrap command used to launch processes.
     pub fn command(&self) -> &BootstrapCommand {
         &self.command
     }
 
+    /// The socket directory, where per-proc Unix sockets are placed.
+    pub fn socket_dir(&self) -> &Path {
+        self.socket_dir.path()
+    }
+
     /// Return the current [`ProcStatus`] for the given [`ProcId`], if
     /// the proc is known to this manager.
     ///
@@ -1688,6 +1736,7 @@ impl ProcManager for BootstrapProcManager {
             proc_id: proc_id.clone(),
             backend_addr,
             callback_addr,
+            socket_dir_path: self.socket_dir.path().to_owned(),
             config: Some(config.client_config_override),
         };
         let mut cmd = self.command.new();
@@ -2114,6 +2163,18 @@ impl Write for Debug {
     }
 }
 
+/// Create a new runtime [`TempDir`]. The directory is created in
+/// `$XDG_RUNTIME_DIR`, otherwise falling back to the system tempdir.
+fn runtime_dir() -> io::Result<TempDir> {
+    match std::env::var_os("XDG_RUNTIME_DIR") {
+        Some(runtime_dir) => {
+            let path = PathBuf::from(runtime_dir);
+            tempfile::tempdir_in(path)
+        }
+        None => tempfile::tempdir(),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::path::PathBuf;
@@ -2157,6 +2218,7 @@ mod tests {
                 proc_id: id!(foo[0]),
                 backend_addr: ChannelAddr::any(ChannelTransport::Tcp(TcpMode::Hostname)),
                 callback_addr: ChannelAddr::any(ChannelTransport::Unix),
+                socket_dir_path: PathBuf::from("notexist"),
                 config: None,
             },
         ];
@@ -2214,13 +2276,16 @@ mod tests {
         attrs[MESH_TAIL_LOG_LINES] = 123;
         attrs[MESH_BOOTSTRAP_ENABLE_PDEATHSIG] = false;
 
+        let socket_dir = runtime_dir().unwrap();
+
         // Proc case
         {
             let original = Bootstrap::Proc {
                 proc_id: id!(foo[42]),
                 backend_addr: ChannelAddr::any(ChannelTransport::Unix),
                 callback_addr: ChannelAddr::any(ChannelTransport::Unix),
                 config: Some(attrs.clone()),
+                socket_dir_path: socket_dir.path().to_owned(),
             };
             let env_str = original.to_env_safe_string().expect("encode bootstrap");
             let decoded = Bootstrap::from_env_safe_string(&env_str).expect("decode bootstrap");
@@ -2260,14 +2325,13 @@ mod tests {
         use std::process::Stdio;
 
         use tokio::process::Command;
-        use tokio::time::Duration;
 
         // Manager; program path is irrelevant for this test.
         let command = BootstrapCommand {
             program: PathBuf::from("/bin/true"),
             ..Default::default()
         };
-        let manager = BootstrapProcManager::new(command);
+        let manager = BootstrapProcManager::new(command).unwrap();
 
         // Spawn a long-running child process (sleep 30) with
         // kill_on_drop(true).
@@ -2647,7 +2711,7 @@ mod tests {
             program: PathBuf::from("/bin/true"),
             ..Default::default()
         };
-        let manager = BootstrapProcManager::new(command);
+        let manager = BootstrapProcManager::new(command).unwrap();
 
         // Spawn a fast-exiting child.
         let mut cmd = Command::new("true");
@@ -2681,7 +2745,7 @@ mod tests {
             program: PathBuf::from("/bin/sleep"),
             ..Default::default()
         };
-        let manager = BootstrapProcManager::new(command);
+        let manager = BootstrapProcManager::new(command).unwrap();
 
         // Spawn a process that will live long enough to kill.
         let mut cmd = Command::new("/bin/sleep");
@@ -2798,7 +2862,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: PathBuf::from("/bin/true"),
             ..Default::default()
-        });
+        })
+        .unwrap();
         let unknown = ProcId::Direct(ChannelAddr::any(ChannelTransport::Unix), "nope".into());
         assert!(manager.status(&unknown).await.is_none());
     }
@@ -2808,7 +2873,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: PathBuf::from("/bin/sleep"),
             ..Default::default()
-        });
+        })
+        .unwrap();
 
         // Long-ish child so it's alive while we "steal" it.
         let mut cmd = Command::new("/bin/sleep");
@@ -2847,7 +2913,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: PathBuf::from("/bin/sleep"),
             ..Default::default()
-        });
+        })
+        .unwrap();
 
         let mut cmd = Command::new("/bin/sleep");
         cmd.arg("5").stdout(Stdio::null()).stderr(Stdio::null());
@@ -3200,8 +3267,6 @@ mod tests {
         instance: &hyperactor::Instance<()>,
         _tag: &str,
     ) -> (ProcId, ChannelAddr) {
-        let proc_id = id!(bootstrap_child[0]);
-
         // Serve a Unix channel as the "backend_addr" and hook it into
         // this test proc.
         let (backend_addr, rx) =
@@ -3212,6 +3277,9 @@ mod tests {
         // router.
         instance.proc().clone().serve(rx);
 
+        // We return an arbitrary (but unbound!) unix direct proc id here;
+        // it is okay, as we're not testing connectivity.
+        let proc_id = ProcId::Direct(ChannelTransport::Unix.any(), "test".to_string());
         (proc_id, backend_addr)
     }
 
@@ -3223,7 +3291,7 @@ mod tests {
             .unwrap();
         let (instance, _handle) = root.instance("client").unwrap();
 
-        let mgr = BootstrapProcManager::new(BootstrapCommand::test());
+        let mgr = BootstrapProcManager::new(BootstrapCommand::test()).unwrap();
         let (proc_id, backend_addr) = make_proc_id_and_backend_addr(&instance, "t_term").await;
         let handle = mgr
             .spawn(
@@ -3286,7 +3354,7 @@ mod tests {
             .unwrap();
         let (instance, _handle) = root.instance("client").unwrap();
 
-        let mgr = BootstrapProcManager::new(BootstrapCommand::test());
+        let mgr = BootstrapProcManager::new(BootstrapCommand::test()).unwrap();
 
         // Proc identity + host backend channel the child will dial.
         let (proc_id, backend_addr) = make_proc_id_and_backend_addr(&instance, "t_kill").await;
@@ -3509,7 +3577,8 @@ mod tests {
         let manager = BootstrapProcManager::new(BootstrapCommand {
             program: std::path::PathBuf::from("/bin/true"), // unused in this test
             ..Default::default()
-        });
+        })
+        .unwrap();
 
         // Give the monitors time to start
         RealClock.sleep(Duration::from_millis(1000)).await;