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

Summary:
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

@@ -18,6 +18,7 @@ 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;
@@ -84,6 +85,12 @@ impl Alloc for PyAllocWrapper {
         self.inner.transport()
     }
 
+    async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+        LogSource::new_with_local_actor()
+            .await
+            .map_err(AllocatorError::from)
+    }
+
     async fn stop(&mut self) -> Result<(), AllocatorError> {
         self.inner.stop().await
     }

hyperactor_mesh/src/alloc.rs

@@ -224,15 +224,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,
@@ -368,6 +360,12 @@ pub mod test_utils {
             self.alloc.transport()
         }
 
+        async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+            LogSource::new_with_local_actor()
+                .await
+                .map_err(AllocatorError::from)
+        }
+
         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

@@ -50,6 +50,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.
@@ -86,6 +88,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();
@@ -94,6 +99,7 @@ impl Allocator for ProcessAllocator {
             world_id: WorldId(name.to_string()),
             spec: spec.clone(),
             bootstrap_addr,
+            log_source,
             rx,
             index: 0,
             active: HashMap::new(),
@@ -112,6 +118,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>,
@@ -142,6 +149,7 @@ 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();
@@ -155,14 +163,8 @@ impl Child {
             Box<dyn io::AsyncWrite + Send + Unpin + 'static>,
         ) = if use_state_actor {
             // Parse the state actor address
-            let state_actor_addr = match "tcp![::]:3000".parse::<ChannelAddr>() {
-                Ok(addr) => addr,
-                Err(e) => {
-                    tracing::warn!("Failed to parse state actor address: {}", e);
-                    // Create a dummy address that won't be used
-                    ChannelAddr::any(ChannelTransport::Unix)
-                }
-            };
+            // TODO (@lky): do not hardcode actor id here; use state_server_info
+            let state_actor_addr = state_server_info.state_proc_addr;
             (
                 Box::new(StateActorWriter::new(
                     OutputTarget::Stdout,
@@ -386,7 +388,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).
@@ -490,6 +493,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,