(3/n) Specialize state actor bootstrap for each alloc (#459)

Summary:
Pull Request resolved: #459

Different alloc should have different setup of state actor.
The remote allocator should bootstrap the state actor inside the initializer.
This will be done in the follow-up diffs.

Differential Revision: D77914042

Author: James Sun

hyperactor_extension/src/alloc.rs

@@ -0,0 +1,174 @@
+/*
+ * 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 std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use async_trait::async_trait;
+use hyperactor::WorldId;
+use hyperactor::channel::ChannelTransport;
+use hyperactor_mesh::alloc::Alloc;
+use hyperactor_mesh::alloc::AllocConstraints;
+use hyperactor_mesh::alloc::AllocSpec;
+use hyperactor_mesh::alloc::AllocatorError;
+use hyperactor_mesh::alloc::ProcState;
+use hyperactor_mesh::log_source::LogSource;
+use hyperactor_mesh::shape::Shape;
+use ndslice::Slice;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+
+/// A python class that wraps a Rust Alloc trait object. It represents what
+/// is shown on the python side. Internals are not exposed.
+/// It ensures that the Alloc is only used once (i.e. moved) in rust.
+#[pyclass(
+    name = "Alloc",
+    module = "monarch._rust_bindings.hyperactor_extension.alloc"
+)]
+pub struct PyAlloc {
+    pub inner: Arc<Mutex<Option<PyAllocWrapper>>>,
+}
+
+impl PyAlloc {
+    /// Create a new PyAlloc with provided boxed trait.
+    pub fn new(inner: Box<dyn Alloc + Sync + Send>) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(Some(PyAllocWrapper { inner }))),
+        }
+    }
+
+    /// Take the internal Alloc object.
+    pub fn take(&self) -> Option<PyAllocWrapper> {
+        self.inner.lock().unwrap().take()
+    }
+}
+
+#[pymethods]
+impl PyAlloc {
+    fn __repr__(&self) -> PyResult<String> {
+        let data = self.inner.lock().unwrap();
+        match &*data {
+            None => Ok("Alloc(None)".to_string()),
+            Some(wrapper) => Ok(format!("Alloc({})", wrapper.shape())),
+        }
+    }
+}
+
+/// Internal wrapper to translate from a dyn Alloc to an impl Alloc. Used
+/// to support polymorphism in the Python bindings.
+pub struct PyAllocWrapper {
+    inner: Box<dyn Alloc + Sync + Send>,
+}
+
+#[async_trait]
+impl Alloc for PyAllocWrapper {
+    async fn next(&mut self) -> Option<ProcState> {
+        self.inner.next().await
+    }
+
+    fn shape(&self) -> &Shape {
+        self.inner.shape()
+    }
+
+    fn world_id(&self) -> &WorldId {
+        self.inner.world_id()
+    }
+
+    fn transport(&self) -> ChannelTransport {
+        self.inner.transport()
+    }
+
+    async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+        self.inner.log_source().await
+    }
+
+    async fn stop(&mut self) -> Result<(), AllocatorError> {
+        self.inner.stop().await
+    }
+}
+
+#[pyclass(
+    name = "AllocConstraints",
+    module = "monarch._rust_bindings.hyperactor_extension.alloc"
+)]
+pub struct PyAllocConstraints {
+    inner: AllocConstraints,
+}
+
+#[pymethods]
+impl PyAllocConstraints {
+    #[new]
+    #[pyo3(signature = (match_labels=None))]
+    fn new(match_labels: Option<HashMap<String, String>>) -> PyResult<Self> {
+        let mut constraints = AllocConstraints::default();
+        if let Some(match_lables) = match_labels {
+            constraints.match_labels = match_lables;
+        }
+
+        Ok(Self { inner: constraints })
+    }
+}
+
+#[pyclass(
+    name = "AllocSpec",
+    module = "monarch._rust_bindings.hyperactor_extension.alloc"
+)]
+pub struct PyAllocSpec {
+    pub inner: AllocSpec,
+}
+
+#[pymethods]
+impl PyAllocSpec {
+    #[new]
+    #[pyo3(signature = (constraints, **kwargs))]
+    fn new(constraints: &PyAllocConstraints, kwargs: Option<&Bound<'_, PyAny>>) -> PyResult<Self> {
+        let Some(kwargs) = kwargs else {
+            return Err(PyValueError::new_err(
+                "Shape must have at least one dimension",
+            ));
+        };
+        let shape_dict = kwargs.downcast::<PyDict>()?;
+
+        let mut keys = Vec::new();
+        let mut values = Vec::new();
+        for (key, value) in shape_dict {
+            keys.push(key.clone());
+            values.push(value.clone());
+        }
+
+        let shape = Shape::new(
+            keys.into_iter()
+                .map(|key| key.extract::<String>())
+                .collect::<PyResult<Vec<String>>>()?,
+            Slice::new_row_major(
+                values
+                    .into_iter()
+                    .map(|key| key.extract::<usize>())
+                    .collect::<PyResult<Vec<usize>>>()?,
+            ),
+        )
+        .map_err(|e| PyValueError::new_err(format!("Invalid shape: {:?}", e)))?;
+
+        Ok(Self {
+            inner: AllocSpec {
+                shape,
+                constraints: constraints.inner.clone(),
+            },
+        })
+    }
+}
+
+pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add_class::<PyAlloc>()?;
+    module.add_class::<PyAllocConstraints>()?;
+    module.add_class::<PyAllocSpec>()?;
+
+    Ok(())
+}

hyperactor_mesh/src/alloc.rs

@@ -223,15 +223,7 @@ pub trait Alloc {
     /// It allows remote processes to stream stdout and stderr back to the client.
     /// A client can connect to the log source to obtain the streamed logs.
     /// A log source is allocation specific. Each allocator can decide how to stream the logs back.
-    async fn log_source(&self) -> Result<LogSource, AllocatorError> {
-        // TODO: this should be implemented based on different allocators.
-        // Having this temporarily here so that the client can connect to the log source.
-        // But the client will not get anything.
-        // The following diffs will gradually implement this for different allocators.
-        LogSource::new_with_local_actor()
-            .await
-            .map_err(AllocatorError::from)
-    }
+    async fn log_source(&self) -> Result<LogSource, AllocatorError>;
 
     /// Stop this alloc, shutting down all of its procs. A clean
     /// shutdown should result in Stop events from all allocs,
@@ -367,6 +359,10 @@ pub mod test_utils {
             self.alloc.transport()
         }
 
+        async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+            self.alloc.log_source().await
+        }
+
         async fn stop(&mut self) -> Result<(), AllocatorError> {
             self.alloc.stop().await
         }

hyperactor_mesh/src/alloc/local.rs

@@ -33,6 +33,7 @@ use crate::alloc::AllocSpec;
 use crate::alloc::Allocator;
 use crate::alloc::AllocatorError;
 use crate::alloc::ProcState;
+use crate::log_source::LogSource;
 use crate::proc_mesh::mesh_agent::MeshAgent;
 use crate::shortuuid::ShortUuid;
 
@@ -252,6 +253,14 @@ impl Alloc for LocalAlloc {
         ChannelTransport::Local
     }
 
+    async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+        // Local alloc does not need to stream logs back.
+        // The client can subscribe to it but local actors will not stream logs into it.
+        LogSource::new_with_local_actor()
+            .await
+            .map_err(AllocatorError::from)
+    }
+
     async fn stop(&mut self) -> Result<(), AllocatorError> {
         for rank in 0..self.size() {
             self.todo_tx

hyperactor_mesh/src/alloc/process.rs

@@ -29,7 +29,6 @@ use hyperactor::channel::ChannelTx;
 use hyperactor::channel::Rx;
 use hyperactor::channel::Tx;
 use hyperactor::channel::TxStatus;
-use hyperactor::id;
 use hyperactor::sync::flag;
 use hyperactor::sync::monitor;
 use hyperactor_state::state_actor::StateActor;
@@ -53,6 +52,8 @@ use crate::bootstrap;
 use crate::bootstrap::Allocator2Process;
 use crate::bootstrap::Process2Allocator;
 use crate::bootstrap::Process2AllocatorMessage;
+use crate::log_source::LogSource;
+use crate::log_source::StateServerInfo;
 use crate::shortuuid::ShortUuid;
 
 /// The maximum number of log lines to tail keep for managed processes.
@@ -89,6 +90,9 @@ impl Allocator for ProcessAllocator {
         let (bootstrap_addr, rx) = channel::serve(ChannelAddr::any(ChannelTransport::Unix))
             .await
             .map_err(anyhow::Error::from)?;
+        let log_source = LogSource::new_with_local_actor()
+            .await
+            .map_err(AllocatorError::from)?;
 
         let name = ShortUuid::generate();
         let n = spec.shape.slice().len();
@@ -97,6 +101,7 @@ impl Allocator for ProcessAllocator {
             world_id: WorldId(name.to_string()),
             spec: spec.clone(),
             bootstrap_addr,
+            log_source,
             rx,
             index: 0,
             active: HashMap::new(),
@@ -115,6 +120,7 @@ pub struct ProcessAlloc {
     world_id: WorldId, // to provide storage
     spec: AllocSpec,
     bootstrap_addr: ChannelAddr,
+    log_source: LogSource,
     rx: channel::ChannelRx<Process2Allocator>,
     index: usize,
     active: HashMap<usize, Child>,
@@ -145,13 +151,14 @@ struct Child {
 impl Child {
     fn monitored(
         mut process: tokio::process::Child,
+        state_server_info: StateServerInfo,
     ) -> (Self, impl Future<Output = ProcStopReason>) {
         let (group, handle) = monitor::group();
         let (exit_flag, exit_guard) = flag::guarded();
         let stop_reason = Arc::new(OnceLock::new());
 
         // TODO(lky): enable state actor branch and remove this flag
-        let use_state_actor = false;
+        let use_state_actor = true;
 
         // Set up stdout and stderr writers
         let mut stdout_tee: Box<dyn io::AsyncWrite + Send + Unpin + 'static> =
@@ -161,24 +168,20 @@ impl Child {
 
         // If state actor is enabled, try to set up LogWriter instances
         if use_state_actor {
-            let state_actor_ref = ActorRef::<StateActor>::attest(id!(state_server[0].state[0]));
-            // Parse the state actor address
-            if let Ok(state_actor_addr) = "tcp![::]:3000".parse::<ChannelAddr>() {
-                // Use the helper function to create both writers at once
-                match hyperactor_state::log_writer::create_log_writers(
-                    state_actor_addr,
-                    state_actor_ref,
-                ) {
-                    Ok((stdout_writer, stderr_writer)) => {
-                        stdout_tee = stdout_writer;
-                        stderr_tee = stderr_writer;
-                    }
-                    Err(e) => {
-                        tracing::error!("failed to create log writers: {}", e);
-                    }
+            let state_actor_ref = ActorRef::<StateActor>::attest(state_server_info.state_actor_id);
+            let state_actor_addr = state_server_info.state_proc_addr;
+            // Use the helper function to create both writers at once
+            match hyperactor_state::log_writer::create_log_writers(
+                state_actor_addr,
+                state_actor_ref,
+            ) {
+                Ok((stdout_writer, stderr_writer)) => {
+                    stdout_tee = stdout_writer;
+                    stderr_tee = stderr_writer;
+                }
+                Err(e) => {
+                    tracing::error!("failed to create log writers: {}", e);
                 }
-            } else {
-                tracing::error!("failed to parse state actor address");
             }
         }
 
@@ -394,7 +397,8 @@ impl ProcessAlloc {
                         None
                     }
                     Ok(rank) => {
-                        let (handle, monitor) = Child::monitored(process);
+                        let (handle, monitor) =
+                            Child::monitored(process, self.log_source.server_info());
                         self.children.spawn(async move { (index, monitor.await) });
                         self.active.insert(index, handle);
                         // Adjust for shape slice offset for non-zero shapes (sub-shapes).
@@ -498,6 +502,10 @@ impl Alloc for ProcessAlloc {
         ChannelTransport::Unix
     }
 
+    async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+        Ok(self.log_source.clone())
+    }
+
     async fn stop(&mut self) -> Result<(), AllocatorError> {
         // We rely on the teardown here, and that the process should
         // exit on its own. We should have a hard timeout here as well,