python simple bootstrap (#1415)

Summary:
Pull Request resolved: #1415

Adds the python implementation for simple bootstrapping.

This lets us create host meshes out of any list of workers running the worker loop.
ghstack-source-id: 314295600
exported-using-ghexport

Reviewed By: colin2328

Differential Revision: D83721433

fbshipit-source-id: 52b4cc0e556e78a1bfe861ec1d3aab0838bebc94

Author: zdevito

hyperactor/src/channel.rs

@@ -577,7 +577,6 @@ impl FromStr for ChannelAddr {
     type Err = anyhow::Error;
 
     fn from_str(addr: &str) -> Result<Self, Self::Err> {
-        // "!" is the legacy delimiter; ":" is preferred
         match addr.split_once('!').or_else(|| addr.split_once(':')) {
             Some(("local", rest)) => rest
                 .parse::<u64>()
@@ -596,6 +595,102 @@ impl FromStr for ChannelAddr {
     }
 }
 
+impl ChannelAddr {
+    /// Parse ZMQ-style URL format: scheme://address
+    /// Supports:
+    /// - tcp://hostname:port or tcp://*:port (wildcard binding)
+    /// - inproc://endpoint-name (equivalent to local)
+    /// - ipc://path (equivalent to unix)
+    /// - metatls://hostname:port or metatls://*:port
+    pub fn from_zmq_url(address: &str) -> Result<Self, anyhow::Error> {
+        // Try ZMQ-style URL format first (scheme://...)
+        let (scheme, address) = address.split_once("://").ok_or_else(|| {
+            anyhow::anyhow!("address must be in url form scheme://endppoint {}", address)
+        })?;
+
+        match scheme {
+            "tcp" => {
+                let (host, port) = Self::split_host_port(address)?;
+
+                if host == "*" {
+                    // Wildcard binding - use IPv6 unspecified address
+                    Ok(Self::Tcp(SocketAddr::new("::".parse().unwrap(), port)))
+                } else {
+                    // Resolve hostname to IP address for proper SocketAddr creation
+                    let socket_addr = Self::resolve_hostname_to_socket_addr(host, port)?;
+                    Ok(Self::Tcp(socket_addr))
+                }
+            }
+            "inproc" => {
+                // inproc://port -> local:port
+                // Port must be a valid u64 number
+                let port = address.parse::<u64>().map_err(|_| {
+                    anyhow::anyhow!("inproc endpoint must be a valid port number: {}", address)
+                })?;
+                Ok(Self::Local(port))
+            }
+            "ipc" => {
+                // ipc://path -> unix:path
+                Ok(Self::Unix(net::unix::SocketAddr::from_str(address)?))
+            }
+            "metatls" => {
+                let (host, port) = Self::split_host_port(address)?;
+
+                if host == "*" {
+                    // Wildcard binding - use IPv6 unspecified address directly without hostname resolution
+                    Ok(Self::MetaTls(MetaTlsAddr::Host {
+                        hostname: std::net::Ipv6Addr::UNSPECIFIED.to_string(),
+                        port,
+                    }))
+                } else {
+                    Ok(Self::MetaTls(MetaTlsAddr::Host {
+                        hostname: host.to_string(),
+                        port,
+                    }))
+                }
+            }
+            scheme => Err(anyhow::anyhow!("unsupported ZMQ scheme: {}", scheme)),
+        }
+    }
+
+    /// Split host:port string, supporting IPv6 addresses
+    fn split_host_port(address: &str) -> Result<(&str, u16), anyhow::Error> {
+        if let Some((host, port_str)) = address.rsplit_once(':') {
+            let port: u16 = port_str
+                .parse()
+                .map_err(|_| anyhow::anyhow!("invalid port: {}", port_str))?;
+            Ok((host, port))
+        } else {
+            Err(anyhow::anyhow!("invalid address format: {}", address))
+        }
+    }
+
+    /// Resolve hostname to SocketAddr, handling both IP addresses and hostnames
+    fn resolve_hostname_to_socket_addr(host: &str, port: u16) -> Result<SocketAddr, anyhow::Error> {
+        // Handle IPv6 addresses in brackets by stripping the brackets
+        let host_clean = if host.starts_with('[') && host.ends_with(']') {
+            &host[1..host.len() - 1]
+        } else {
+            host
+        };
+
+        // First try to parse as an IP address directly
+        if let Ok(ip_addr) = host_clean.parse::<IpAddr>() {
+            return Ok(SocketAddr::new(ip_addr, port));
+        }
+
+        // If not an IP, try hostname resolution
+        use std::net::ToSocketAddrs;
+        let mut addrs = (host_clean, port)
+            .to_socket_addrs()
+            .map_err(|e| anyhow::anyhow!("failed to resolve hostname '{}': {}", host_clean, e))?;
+
+        addrs
+            .next()
+            .ok_or_else(|| anyhow::anyhow!("no addresses found for hostname '{}'", host_clean))
+    }
+}
+
 /// Universal channel transmitter.
 #[derive(Debug)]
 pub struct ChannelTx<M: RemoteMessage> {
@@ -832,6 +927,78 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_zmq_style_channel_addr() {
+        // Test TCP addresses
+        assert_eq!(
+            ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080").unwrap(),
+            ChannelAddr::Tcp("127.0.0.1:8080".parse().unwrap())
+        );
+
+        // Test TCP wildcard binding
+        assert_eq!(
+            ChannelAddr::from_zmq_url("tcp://*:5555").unwrap(),
+            ChannelAddr::Tcp("[::]:5555".parse().unwrap())
+        );
+
+        // Test inproc (maps to local with numeric endpoint)
+        assert_eq!(
+            ChannelAddr::from_zmq_url("inproc://12345").unwrap(),
+            ChannelAddr::Local(12345)
+        );
+
+        // Test ipc (maps to unix)
+        assert_eq!(
+            ChannelAddr::from_zmq_url("ipc:///tmp/my-socket").unwrap(),
+            ChannelAddr::Unix(unix::SocketAddr::from_pathname("/tmp/my-socket").unwrap())
+        );
+
+        // Test metatls with hostname
+        assert_eq!(
+            ChannelAddr::from_zmq_url("metatls://example.com:443").unwrap(),
+            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: "example.com".to_string(),
+                port: 443
+            })
+        );
+
+        // Test metatls with IP address (should be normalized)
+        assert_eq!(
+            ChannelAddr::from_zmq_url("metatls://192.168.1.1:443").unwrap(),
+            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: "192.168.1.1".to_string(),
+                port: 443
+            })
+        );
+
+        // Test metatls with wildcard (should use IPv6 unspecified address)
+        assert_eq!(
+            ChannelAddr::from_zmq_url("metatls://*:8443").unwrap(),
+            ChannelAddr::MetaTls(MetaTlsAddr::Host {
+                hostname: "::".to_string(),
+                port: 8443
+            })
+        );
+
+        // Test TCP hostname resolution (should resolve hostname to IP)
+        // Note: This test may fail in environments without proper DNS resolution
+        // We test that it at least doesn't fail to parse
+        let tcp_hostname_result = ChannelAddr::from_zmq_url("tcp://localhost:8080");
+        assert!(tcp_hostname_result.is_ok());
+
+        // Test IPv6 address
+        assert_eq!(
+            ChannelAddr::from_zmq_url("tcp://[::1]:1234").unwrap(),
+            ChannelAddr::Tcp("[::1]:1234".parse().unwrap())
+        );
+
+        // Test error cases
+        assert!(ChannelAddr::from_zmq_url("invalid://scheme").is_err());
+        assert!(ChannelAddr::from_zmq_url("tcp://invalid-port").is_err());
+        assert!(ChannelAddr::from_zmq_url("metatls://no-port").is_err());
+        assert!(ChannelAddr::from_zmq_url("inproc://not-a-number").is_err());
+    }
+
     #[tokio::test]
     async fn test_multiple_connections() {
         for addr in ChannelTransport::all().map(ChannelAddr::any) {

monarch_hyperactor/src/bootstrap.rs

@@ -6,7 +6,15 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use futures::future::try_join_all;
+use hyperactor::channel::ChannelAddr;
+use hyperactor_mesh::Bootstrap;
+use hyperactor_mesh::bootstrap::BootstrapCommand;
 use hyperactor_mesh::bootstrap_or_die;
+use hyperactor_mesh::v1::HostMeshRef;
+use hyperactor_mesh::v1::Name;
+use hyperactor_mesh::v1::host_mesh::HostMesh;
+use monarch_types::MapPyErr;
 use pyo3::Bound;
 use pyo3::PyAny;
 use pyo3::PyResult;
@@ -17,6 +25,9 @@ use pyo3::types::PyModule;
 use pyo3::types::PyModuleMethods;
 use pyo3::wrap_pyfunction;
 
+use crate::pytokio::PyPythonTask;
+use crate::v1::host_mesh::PyHostMesh;
+
 #[pyfunction]
 #[pyo3(signature = ())]
 pub fn bootstrap_main(py: Python) -> PyResult<Bound<PyAny>> {
@@ -36,6 +47,93 @@ pub fn bootstrap_main(py: Python) -> PyResult<Bound<PyAny>> {
     })
 }
 
+#[pyfunction]
+pub fn run_worker_loop_forever(py: Python<'_>, address: &str) -> PyResult<PyPythonTask> {
+    let addr = ChannelAddr::from_zmq_url(address)?;
+
+    // Check if we're running in a PAR/XAR build by looking for FB_XAR_INVOKED_NAME environment variable
+    let invoked_name = std::env::var("FB_XAR_INVOKED_NAME");
+
+    let mut env: std::collections::HashMap<String, String> = std::env::vars().collect();
+
+    let command = Some(if let Ok(invoked_name) = invoked_name {
+        // For PAR/XAR builds: use argv[0] from Python's sys.argv as the current executable
+        let current_exe = std::path::PathBuf::from(&invoked_name);
+
+        // For PAR/XAR builds: set PAR_MAIN_OVERRIDE and no additional args
+        env.insert(
+            "PAR_MAIN_OVERRIDE".to_string(),
+            "monarch._src.actor.bootstrap_main".to_string(),
+        );
+        BootstrapCommand {
+            program: current_exe,
+            arg0: Some(invoked_name),
+            args: vec![],
+            env,
+        }
+    } else {
+        // For regular Python builds: use current_exe() and -m arguments
+        let current_exe = std::env::current_exe().map_err(|e| {
+            pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
+                "Failed to get current executable: {}",
+                e
+            ))
+        })?;
+        let current_exe_str = current_exe.to_string_lossy().to_string();
+        BootstrapCommand {
+            program: current_exe,
+            arg0: Some(current_exe_str),
+            args: vec![
+                "-m".to_string(),
+                "monarch._src.actor.bootstrap_main".to_string(),
+            ],
+            env,
+        }
+    });
+
+    let boot = Bootstrap::Host {
+        addr,
+        config: None,
+        command,
+    };
+
+    PyPythonTask::new(async {
+        let err = boot.bootstrap().await;
+        Err(err).map_pyerr()?;
+        Ok(())
+    })
+}
+
+#[pyfunction]
+pub fn attach_to_workers<'py>(
+    workers: Vec<Bound<'py, PyPythonTask>>,
+    name: Option<&str>,
+) -> PyResult<PyPythonTask> {
+    let tasks = workers
+        .into_iter()
+        .map(|x| x.borrow_mut().take_task())
+        .collect::<PyResult<Vec<_>>>()?;
+
+    let name = Name::new(name.unwrap_or("hosts"));
+    PyPythonTask::new(async move {
+        let results = try_join_all(tasks).await?;
+
+        let addresses: Result<Vec<ChannelAddr>, anyhow::Error> = Python::with_gil(|py| {
+            results
+                .into_iter()
+                .map(|result| {
+                    let url_str: String = result.bind(py).extract()?;
+                    ChannelAddr::from_zmq_url(&url_str)
+                })
+                .collect()
+        });
+        let addresses = addresses?;
+
+        let host_mesh = HostMesh::take(name, HostMeshRef::from_hosts(addresses));
+        Ok(PyHostMesh::new_owned(host_mesh))
+    })
+}
+
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     let f = wrap_pyfunction!(bootstrap_main, hyperactor_mod)?;
     f.setattr(
@@ -44,5 +142,19 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     )?;
     hyperactor_mod.add_function(f)?;
 
+    let f = wrap_pyfunction!(run_worker_loop_forever, hyperactor_mod)?;
+    f.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_hyperactor.bootstrap",
+    )?;
+    hyperactor_mod.add_function(f)?;
+
+    let f = wrap_pyfunction!(attach_to_workers, hyperactor_mod)?;
+    f.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_hyperactor.bootstrap",
+    )?;
+    hyperactor_mod.add_function(f)?;
+
     Ok(())
 }

monarch_hyperactor/src/v1/host_mesh.rs

@@ -87,13 +87,13 @@ impl PyBootstrapCommand {
     name = "HostMesh",
     module = "monarch._rust_bindings.monarch_hyperactor.v1.host_mesh"
 )]
-enum PyHostMesh {
+pub(crate) enum PyHostMesh {
     Owned(PyHostMeshImpl),
     Ref(PyHostMeshRefImpl),
 }
 
 impl PyHostMesh {
-    fn new_owned(inner: HostMesh) -> Self {
+    pub(crate) fn new_owned(inner: HostMesh) -> Self {
         Self::Owned(PyHostMeshImpl(SharedCell::from(inner)))
     }
 

monarch_types/src/lib.rs

@@ -45,7 +45,7 @@ pub trait MapPyErr<T> {
 }
 impl<T, E> MapPyErr<T> for Result<T, E>
 where
-    E: Error,
+    E: ToString,
 {
     fn map_pyerr(self) -> Result<T, PyErr> {
         self.map_err(|err| PyErr::new::<PyValueError, _>(err.to_string()))

python/monarch/_rust_bindings/monarch_hyperactor/bootstrap.pyi

@@ -6,4 +6,17 @@
 
 # pyre-strict
 
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+PrivateKey = Union[bytes, Path, None]
+CA = Union[bytes, Path, Literal["trust_all_connections"]]
+
+from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
+from monarch._rust_bindings.monarch_hyperactor.v1.host_mesh import HostMesh
+
 def bootstrap_main() -> None: ...
+def run_worker_loop_forever(address: str) -> PythonTask[None]: ...
+def attach_to_workers(
+    workers: List[PythonTask[str]], name: Optional[str] = None
+) -> PythonTask[HostMesh]: ...