(monarch_hyperactor) Create python binding for a RemoteAllocator that takes a list of remote channel addresses (#170)

Summary:
Pull Request resolved: #170

To support multi-node actor meshes in OSS without having to write a custom allocator for each scheduler (e.g. `SlurmAllocator`, `KubernetesAllocator`) we take advantage of the infrastructure we already have in TorchX and TorchElastic.

This Diff creates Python bindings for `RemoteAllocatorBase` that takes a list of server addresses (in channel_addr format - e.g. `metatls!devgpu032.nha1.facebook.com:26600` or `tcp!devgpu032.nha1.facebook.com:26601`) of remote-process-allocator server and connects to it.

The internals reuse existing `RemoteProcessAlloc` with a custom `PyRemoteProcessAllocInitializer` that simply returns a `Vec<RemoteProcessAllocHost>` given the user provided list of server addresses.

Recommended to start the review at `monarch‎/python‎/tests‎/test_allocator.py‎` to get a sense of what the API/Usage looks like.

The next diff will provide a function that gets the list of server addresses given a job-id (more specifically a monarch server handle of the form `{scheduler}://{namespace}/{job_id}` e.g. `slurm://default/monarch-kiuk-123`) and returns an Allocator that can be used to create a `ProcMesh` as usual.

NOTE: WIP fixing type-checking failures so ignore those...

Differential Revision: D75928565

Author: kiukchung

.github/workflows/test.yml

@@ -54,6 +54,9 @@ jobs:
         # Install test dependencies
         pip install -r python/tests/requirements.txt
 
+        # Install remote process_allocator binary (some tests use it)
+        cargo install --path monarch_hyperactor
+
         # Build and install monarch
         # NB: monarch is currently can't be built in isolated builds (e.g not PEP519 compatible)
         # because 'torch-sys' needs to be compiled against 'torch' in the main python environment

monarch_extension/src/lib.rs

@@ -97,6 +97,10 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         module,
         "monarch_hyperactor.alloc",
     )?)?;
+    monarch_hyperactor::channel::register_python_bindings(&get_or_add_new_module(
+        module,
+        "monarch_hyperactor.channel",
+    )?)?;
     monarch_hyperactor::actor_mesh::register_python_bindings(&get_or_add_new_module(
         module,
         "monarch_hyperactor.actor_mesh",

monarch_hyperactor/src/alloc.rs

@@ -7,13 +7,22 @@
  */
 
 use std::collections::HashMap;
+use std::str::FromStr;
 use std::sync::Arc;
+use std::time::Duration;
 
+use async_trait::async_trait;
+use hyperactor::WorldId;
+use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::ChannelTransport;
 use hyperactor_extension::alloc::PyAlloc;
 use hyperactor_extension::alloc::PyAllocSpec;
 use hyperactor_mesh::alloc::Allocator;
 use hyperactor_mesh::alloc::LocalAllocator;
 use hyperactor_mesh::alloc::ProcessAllocator;
+use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAlloc;
+use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocHost;
+use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocInitializer;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use tokio::process::Command;
@@ -132,9 +141,159 @@ impl PyProcessAllocator {
     }
 }
 
+struct PyRemoteProcessAllocInitializer {
+    addrs: Vec<String>,
+}
+
+#[async_trait]
+impl RemoteProcessAllocInitializer for PyRemoteProcessAllocInitializer {
+    async fn initialize_alloc(&mut self) -> Result<Vec<RemoteProcessAllocHost>, anyhow::Error> {
+        self.addrs
+            .iter()
+            .map(|channel_addr| {
+                let addr = ChannelAddr::from_str(channel_addr)?;
+                let remote_host = match addr {
+                    ChannelAddr::Tcp(socket_addr) => RemoteProcessAllocHost {
+                        id: socket_addr.ip().to_string(),
+                        hostname: socket_addr.ip().to_string(),
+                    },
+                    ChannelAddr::MetaTls(hostname, _) => RemoteProcessAllocHost {
+                        id: hostname.clone(),
+                        hostname: hostname.clone(),
+                    },
+                    ChannelAddr::Unix(_) => RemoteProcessAllocHost {
+                        id: addr.to_string(),
+                        hostname: addr.to_string(),
+                    },
+                    _ => {
+                        anyhow::bail!("Unsupported transport for channel address: `{addr:?}`")
+                    }
+                };
+                Ok(remote_host)
+            })
+            .collect()
+    }
+}
+
+#[pyclass(
+    name = "RemoteAllocatorBase",
+    module = "monarch._rust_bindings.monarch_hyperactor.alloc",
+    subclass
+)]
+pub struct PyRemoteAllocator {
+    world_id: String,
+    addrs: Vec<String>,
+    heartbeat_interval_millis: u64,
+}
+
+const DEFAULT_REMOTE_ALLOCATOR_PORT: u16 = 26600;
+const DEFAULT_REMOTE_ALLOCATOR_HEARTBEAT_INTERVAL_MILLIS: u64 = 5000;
+
+#[pymethods]
+impl PyRemoteAllocator {
+    #[classattr]
+    const DEFAULT_PORT: u16 = DEFAULT_REMOTE_ALLOCATOR_PORT;
+
+    #[classattr]
+    const DEFAULT_HEARTBEAT_INTERVAL_MILLIS: u64 =
+        DEFAULT_REMOTE_ALLOCATOR_HEARTBEAT_INTERVAL_MILLIS;
+
+    #[new]
+    #[pyo3(signature = (
+        world_id,
+        addrs,
+        heartbeat_interval_millis = DEFAULT_REMOTE_ALLOCATOR_HEARTBEAT_INTERVAL_MILLIS,
+    ))]
+    fn new(world_id: String, addrs: Vec<String>, heartbeat_interval_millis: u64) -> PyResult<Self> {
+        Ok(Self {
+            world_id,
+            addrs,
+            heartbeat_interval_millis,
+        })
+    }
+
+    fn allocate_nonblocking<'py>(
+        &self,
+        py: Python<'py>,
+        spec: &PyAllocSpec,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let addrs = self.addrs.clone();
+        let world_id = self.world_id.clone();
+        let spec_inner = spec.inner.clone();
+        let heartbeat_interval_millis = self.heartbeat_interval_millis;
+
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            // all addrs expected to have the same transport; use the first one
+            let first_addr = addrs.first().expect("addrs should not be empty");
+            let first_addr = ChannelAddr::from_str(first_addr)?;
+            let transport = first_addr.transport();
+            let port = match first_addr {
+                ChannelAddr::Tcp(socket_addr) => socket_addr.port(),
+                ChannelAddr::MetaTls(_, port) => port,
+                ChannelAddr::Unix(_) => 0,
+                ChannelAddr::Local(_) => 0,
+                ChannelAddr::Sim(_) => {
+                    return Err(PyRuntimeError::new_err(format!(
+                        "Unsupported channel_addr: {first_addr:?}"
+                    )));
+                }
+            };
+
+            let alloc = RemoteProcessAlloc::new(
+                spec_inner,
+                WorldId(world_id),
+                transport,
+                port,
+                Duration::from_millis(heartbeat_interval_millis),
+                PyRemoteProcessAllocInitializer { addrs },
+            )
+            .await?;
+
+            Ok(PyAlloc::new(Box::new(alloc)))
+        })
+    }
+    fn allocate_blocking<'py>(&self, py: Python<'py>, spec: &PyAllocSpec) -> PyResult<PyAlloc> {
+        let addrs = self.addrs.clone();
+        let world_id = self.world_id.clone();
+        let spec_inner = spec.inner.clone();
+        let heartbeat_interval_millis = self.heartbeat_interval_millis;
+
+        signal_safe_block_on(py, async move {
+            // all addrs expected to have the same transport; use the first one
+            let first_addr = addrs.first().expect("addrs should not be empty");
+            let first_addr = ChannelAddr::from_str(first_addr)?;
+            let transport = first_addr.transport();
+            let port = match first_addr {
+                ChannelAddr::Tcp(socket_addr) => socket_addr.port(),
+                ChannelAddr::MetaTls(_, port) => port,
+                ChannelAddr::Unix(_) => 0,
+                ChannelAddr::Local(_) => 0,
+                ChannelAddr::Sim(_) => {
+                    return Err(PyRuntimeError::new_err(format!(
+                        "Unsupported channel_addr: {first_addr:?}"
+                    )));
+                }
+            };
+
+            let alloc = RemoteProcessAlloc::new(
+                spec_inner,
+                WorldId(world_id),
+                transport,
+                port,
+                Duration::from_millis(heartbeat_interval_millis),
+                PyRemoteProcessAllocInitializer { addrs },
+            )
+            .await?;
+
+            Ok(PyAlloc::new(Box::new(alloc)))
+        })?
+    }
+}
+
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     hyperactor_mod.add_class::<PyProcessAllocator>()?;
     hyperactor_mod.add_class::<PyLocalAllocator>()?;
+    hyperactor_mod.add_class::<PyRemoteAllocator>()?;
 
     Ok(())
 }

monarch_hyperactor/src/bin/process_allocator/common.rs

@@ -13,22 +13,24 @@ use clap::command;
 use hyperactor::channel::ChannelAddr;
 use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocator;
 use tokio::process::Command;
+
 #[derive(Parser, Debug)]
 #[command(about = "Runs hyperactor's process allocator")]
 pub struct Args {
     #[arg(
         long,
-        default_value = "[::]",
-        help = "The address bind to. The process allocator runs on `bind_addr:port`"
+        default_value_t = 26600,
+        help = "The port to bind to on [::] (all network interfaces on this host). Same as specifying `--addr=[::]:{port}`"
     )]
-    pub addr: String,
+    pub port: u16,
 
     #[arg(
         long,
-        default_value_t = 26600,
-        help = "Port to bind to. The process allocator runs on `bind_addr:port`"
+        help = "The address to bind to in the form: \
+                `{transport}!{address}:{port}` (e.g. `tcp!127.0.0.1:26600`). \
+                If specified, `--port` argument is ignored"
     )]
-    pub port: u16,
+    pub addr: Option<String>,
 
     #[arg(
         long,
@@ -72,8 +74,8 @@ mod tests {
 
         let parsed_args = Args::parse_from(args);
 
-        assert_eq!(parsed_args.addr, "[::]");
         assert_eq!(parsed_args.port, 26600);
+        assert_eq!(parsed_args.addr, None);
         assert_eq!(parsed_args.program, "monarch_bootstrap");
         Ok(())
     }
@@ -82,15 +84,13 @@ mod tests {
     async fn test_args() -> Result<(), anyhow::Error> {
         let args = vec![
             "process_allocator",
-            "--addr=127.0.0.1",
-            "--port=29500",
+            "--addr=tcp!127.0.0.1:29501",
             "--program=/bin/echo",
         ];
 
         let parsed_args = Args::parse_from(args);
 
-        assert_eq!(parsed_args.addr, "127.0.0.1");
-        assert_eq!(parsed_args.port, 29500);
+        assert_eq!(parsed_args.addr, Some("tcp!127.0.0.1:29501".to_string()));
         assert_eq!(parsed_args.program, "/bin/echo");
         Ok(())
     }

monarch_hyperactor/src/bin/process_allocator/main.rs

@@ -8,6 +8,8 @@
 
 mod common;
 
+use std::str::FromStr;
+
 use clap::Parser;
 use common::Args;
 use common::main_impl;
@@ -18,9 +20,11 @@ async fn main() {
     let args = Args::parse();
     hyperactor::initialize();
 
-    let bind = format!("{}:{}", args.addr, args.port);
-    let socket_addr: std::net::SocketAddr = bind.parse().unwrap();
-    let serve_address = ChannelAddr::Tcp(socket_addr);
+    let bind = args
+        .addr
+        .unwrap_or_else(|| format!("tcp![::]:{}", args.port));
+
+    let serve_address = ChannelAddr::from_str(&bind).unwrap();
 
     let _ = main_impl(serve_address, args.program).await.unwrap();
 }

monarch_hyperactor/src/channel.rs

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::ChannelTransport;
+use pyo3::prelude::*;
+
+/// Python binding for [`hyperactor::channel::ChannelTransport`]
+#[pyclass(
+    name = "ChannelTransport",
+    module = "monarch._rust_bindings.monarch_hyperactor.channel",
+    eq
+)]
+#[derive(PartialEq, Clone, Copy)]
+pub enum PyChannelTransport {
+    Tcp,
+    MetaTls,
+    Local,
+    Unix,
+    // Sim(/*proxy address:*/ ChannelAddr), TODO kiuk@ add support
+}
+
+#[pyclass(
+    name = "ChannelAddr",
+    module = "monarch._rust_bindings.monarch_hyperactor.channel"
+)]
+pub struct PyChannelAddr;
+
+#[pymethods]
+impl PyChannelAddr {
+    /// Returns an "any" address for the given transport type.
+    /// Primarily used to bind servers
+    #[staticmethod]
+    fn any(transport: PyChannelTransport) -> PyResult<String> {
+        Ok(ChannelAddr::any(transport.into()).to_string())
+    }
+}
+
+impl From<PyChannelTransport> for ChannelTransport {
+    fn from(val: PyChannelTransport) -> Self {
+        match val {
+            PyChannelTransport::Tcp => ChannelTransport::Tcp,
+            PyChannelTransport::MetaTls => ChannelTransport::MetaTls,
+            PyChannelTransport::Local => ChannelTransport::Local,
+            PyChannelTransport::Unix => ChannelTransport::Unix,
+        }
+    }
+}
+
+pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
+    hyperactor_mod.add_class::<PyChannelTransport>()?;
+    hyperactor_mod.add_class::<PyChannelAddr>()?;
+    Ok(())
+}

monarch_hyperactor/src/lib.rs

@@ -12,6 +12,7 @@ pub mod actor;
 pub mod actor_mesh;
 pub mod alloc;
 pub mod bootstrap;
+pub mod channel;
 pub mod mailbox;
 pub mod ndslice;
 pub mod proc;