rust example to capture stdout and stderr

Differential Revision: D77377307

Author: kaiyuan-li

hyperactor_mesh/Cargo.toml

@@ -35,6 +35,7 @@ erased-serde = "0.3.27"
 futures = { version = "0.3.30", features = ["async-await", "compat"] }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 hyperactor_mesh_macros = { version = "0.0.0", path = "../hyperactor_mesh_macros" }
+hyperactor_state = { version = "0.0.0", path = "../hyperactor_state" }
 hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
 libc = "0.2.139"
 mockall = "0.13.1"

hyperactor_mesh/src/alloc.rs

@@ -13,6 +13,7 @@ pub mod local;
 pub(crate) mod logtailer;
 pub mod process;
 pub mod remoteprocess;
+pub mod state_actor_writer;
 
 use std::collections::HashMap;
 use std::fmt;

hyperactor_mesh/src/alloc/process.rs

@@ -43,6 +43,8 @@ use super::AllocatorError;
 use super::ProcState;
 use super::ProcStopReason;
 use super::logtailer::LogTailer;
+use super::state_actor_writer::OutputTarget;
+use super::state_actor_writer::StateActorWriter;
 use crate::assign::Ranks;
 use crate::bootstrap;
 use crate::bootstrap::Allocator2Process;
@@ -143,18 +145,49 @@ impl Child {
     ) -> (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
+        let use_state_actor = true;
+
+        let (stdout_tee, stderr_tee): (
+            Box<dyn io::AsyncWrite + Send + Unpin + 'static>,
+            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)
+                }
+            };
+            (
+                Box::new(StateActorWriter::new(
+                    OutputTarget::Stdout,
+                    state_actor_addr.clone(),
+                )),
+                Box::new(StateActorWriter::new(
+                    OutputTarget::Stderr,
+                    state_actor_addr,
+                )),
+            )
+        } else {
+            (Box::new(io::stdout()), Box::new(io::stderr()))
+        };
 
         let stdout = LogTailer::tee(
             MAX_TAIL_LOG_LINES,
             process.stdout.take().unwrap(),
-            io::stdout(),
+            stdout_tee,
         );
+
         let stderr = LogTailer::tee(
             MAX_TAIL_LOG_LINES,
             process.stderr.take().unwrap(),
-            io::stderr(),
+            stderr_tee,
         );
-        let stop_reason = Arc::new(OnceLock::new());
 
         let child = Self {
             channel: ChannelState::NotConnected,