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

Summary:
Pull Request resolved: #170

See: P1833193535 for how the user-facing UX would look like.

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

NOTE: hyperactor's ChannelAddr can be represented in string format such as: `tcp!127.0.0.1:26600` or `metatls!devgpu001.pci.facebook.com:26600` which includes all the necessary information to create a `Channel`. Unfortunately, the current `RemoteAllocator` related interfaces take `transport` (`ChannelTransport`), `port`, and a list of `hostnames`, and applies the same transport and port to all. This isn't ideal (especially for flexibility in deployment and testing). So the python bindings take a list of channel address strings rather than a list of hostnames.

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.

## Next Steps:
1. [1/2] Add hostnames to `monarch.tools.mesh_spec.MeshSpec` struct + ability to query hostnames given a running job.

2. [2/2] Make it possible to run 2x remote process allocators (each on its own port) on MAST

Reviewed By: technicianted

Differential Revision: D75928565

fbshipit-source-id: 2dac97fbdabe11c1578c5930501ac49b93906727

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

@@ -142,6 +142,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,17 +7,30 @@
  */
 
 use std::collections::HashMap;
+use std::str::FromStr;
 use std::sync::Arc;
+use std::time::Duration;
 
+use anyhow::anyhow;
+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::AllocSpec;
 use hyperactor_mesh::alloc::Allocator;
+use hyperactor_mesh::alloc::AllocatorError;
 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;
 
+use crate::channel::PyChannelAddr;
 use crate::runtime::signal_safe_block_on;
 
 #[pyclass(
@@ -48,7 +61,7 @@ impl PyLocalAllocator {
                 .allocate(spec)
                 .await
                 .map(|inner| PyAlloc::new(Box::new(inner)))
-                .map_err(|e| PyRuntimeError::new_err(format!("{:?}", e)))
+                .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
         })
     }
 
@@ -132,9 +145,174 @@ impl PyProcessAllocator {
     }
 }
 
+/// A `[hyperactor_mesh::alloc::RemoteProcessAllocInitializer]` wrapper to enable subclassing from Python.
+///
+/// Basically follows https://pyo3.rs/v0.25.0/trait-bounds.html.
+/// The Python subclass should implement `def initialize_alloc(self) -> list[str]`.
+pub struct PyRemoteProcessAllocInitializer {
+    // instance of a Python subclass of `monarch._rust_bindings.monarch_hyperactor.alloc.RemoteProcessAllocInitializer`.
+    py_inner: Py<PyAny>,
+}
+
+impl Clone for PyRemoteProcessAllocInitializer {
+    fn clone(&self) -> Self {
+        Self {
+            py_inner: Python::with_gil(|py| Py::clone_ref(&self.py_inner, py)),
+        }
+    }
+}
+impl PyRemoteProcessAllocInitializer {
+    /// calls the initializer's `initialize_alloc()` as implemented in python
+    ///
+    /// NOTE: changes to python method calls must be made in sync with
+    ///   the method signature of `RemoteAllocInitializer` in
+    ///   `monarch/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi`
+    async fn py_initialize_alloc(&self) -> PyResult<Vec<String>> {
+        // call the function as implemented in python
+        let future = Python::with_gil(|py| -> PyResult<_> {
+            let coroutine = self.py_inner.bind(py).call_method0("initialize_alloc")?;
+            pyo3_async_runtimes::tokio::into_future(coroutine)
+        })?;
+
+        let addrs = future.await?;
+        Python::with_gil(|py| -> PyResult<Vec<String>> { addrs.extract(py) })
+    }
+
+    async fn get_transport_and_port(&self) -> PyResult<(ChannelTransport, u16)> {
+        // NOTE: the upstream RemoteAllocator APIs take (transport, port, hostnames)
+        //   (e.g. assumes the same transport and port for all servers).
+        //   Until that is fixed we have to assume the same here.
+        //   Get the transport and port from the first address
+        // TODO T227130269
+        let addrs = self.py_initialize_alloc().await?;
+        let addr = addrs
+            .first()
+            .ok_or_else(|| anyhow!("initializer must return non-empty list of addresses"))?;
+        let channel_addr = PyChannelAddr::parse(addr)?;
+        let port = channel_addr.get_port()?;
+        let transport = channel_addr.get_transport()?;
+        Ok((transport.into(), port))
+    }
+}
+
+#[async_trait]
+impl RemoteProcessAllocInitializer for PyRemoteProcessAllocInitializer {
+    async fn initialize_alloc(&mut self) -> Result<Vec<RemoteProcessAllocHost>, anyhow::Error> {
+        // call the function as implemented in python
+        let addrs = self.py_initialize_alloc().await?;
+        addrs
+            .iter()
+            .map(|channel_addr| {
+                let addr = ChannelAddr::from_str(channel_addr)?;
+                let (id, hostname) = match addr {
+                    ChannelAddr::Tcp(socket) => (socket.ip().to_string(), socket.ip().to_string()),
+                    ChannelAddr::MetaTls(hostname, _) => (hostname.clone(), hostname.clone()),
+                    ChannelAddr::Unix(_) => (addr.to_string(), addr.to_string()),
+                    _ => anyhow::bail!("unsupported transport for channel address: `{addr}`"),
+                };
+                Ok(RemoteProcessAllocHost { id, hostname })
+            })
+            .collect()
+    }
+}
+
+#[pyclass(
+    name = "RemoteAllocatorBase",
+    module = "monarch._rust_bindings.monarch_hyperactor.alloc",
+    subclass
+)]
+#[derive(Clone)]
+pub struct PyRemoteAllocator {
+    // IMPORTANT: other than the `initializer` this struct should not hold any non-trivially
+    //   clonable data (e.g. such that the Clone derive-attribute would not work).
+    //   This allows us to avoid having yet-another-wrapper for PyRemoteAllocator since
+    //   PyRemoteProcessAllocInitializer is already a wrapper and its wrapped Py<PyAny> is
+    //   shared by reference.
+    world_id: String,
+    initializer: PyRemoteProcessAllocInitializer,
+    heartbeat_interval: Duration,
+}
+
+#[async_trait]
+impl Allocator for PyRemoteAllocator {
+    type Alloc = RemoteProcessAlloc;
+
+    async fn allocate(&mut self, spec: AllocSpec) -> Result<Self::Alloc, AllocatorError> {
+        let initializer = self.initializer.clone();
+        let (transport, port) = initializer
+            .get_transport_and_port()
+            .await
+            .map_err(|e| AllocatorError::Other(e.into()))?;
+
+        let alloc = RemoteProcessAlloc::new(
+            spec,
+            WorldId(self.world_id.clone()),
+            transport,
+            port,
+            self.heartbeat_interval,
+            initializer,
+        )
+        .await?;
+        Ok(alloc)
+    }
+}
+
+#[pymethods]
+impl PyRemoteAllocator {
+    #[new]
+    #[pyo3(signature = (
+        world_id,
+        initializer,
+        heartbeat_interval = Duration::from_secs(5),
+    ))]
+    fn new(
+        world_id: String,
+        initializer: Py<PyAny>,
+        heartbeat_interval: Duration,
+    ) -> PyResult<Self> {
+        Ok(Self {
+            world_id,
+            initializer: PyRemoteProcessAllocInitializer {
+                py_inner: initializer,
+            },
+            heartbeat_interval,
+        })
+    }
+
+    fn allocate_nonblocking<'py>(
+        &self,
+        py: Python<'py>,
+        spec: &PyAllocSpec,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let spec = spec.inner.clone();
+        let mut cloned = self.clone();
+
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            cloned
+                .allocate(spec)
+                .await
+                .map(|alloc| PyAlloc::new(Box::new(alloc)))
+                .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
+        })
+    }
+    fn allocate_blocking<'py>(&self, py: Python<'py>, spec: &PyAllocSpec) -> PyResult<PyAlloc> {
+        let spec = spec.inner.clone();
+        let mut cloned = self.clone();
+
+        signal_safe_block_on(py, async move {
+            cloned
+                .allocate(spec)
+                .await
+                .map(|alloc| PyAlloc::new(Box::new(alloc)))
+                .map_err(|e| PyRuntimeError::new_err(format!("{:?}", e)))
+        })?
+    }
+}
+
 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();
 }