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

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,10 @@ impl Alloc for PyAllocWrapper {
         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
     }

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;
@@ -51,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.
@@ -87,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();
@@ -95,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(),
@@ -113,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>,
@@ -143,13 +149,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> =
@@ -159,24 +166,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");
             }
         }
 
@@ -389,7 +392,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).
@@ -493,6 +497,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,