stream stdout and stderr to StateActor (#430)

Summary:
Pull Request resolved: #430

Example to show running an arbitrary command in rust and pipe the stderr and stdout out.

Reviewed By: highker

Differential Revision: D77377307

fbshipit-source-id: 91b4688b2dec78f168f7f2e41b2fdd410a8d037e

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/process.rs

@@ -18,6 +18,7 @@ use std::sync::OnceLock;
 
 use async_trait::async_trait;
 use enum_as_inner::EnumAsInner;
+use hyperactor::ActorRef;
 use hyperactor::ProcId;
 use hyperactor::WorldId;
 use hyperactor::channel;
@@ -28,8 +29,10 @@ 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;
 use ndslice::Shape;
 use tokio::io;
 use tokio::process::Command;
@@ -143,18 +146,51 @@ 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 and remove this flag
+        let use_state_actor = false;
+
+        // Set up stdout and stderr writers
+        let mut stdout_tee: Box<dyn io::AsyncWrite + Send + Unpin + 'static> =
+            Box::new(io::stdout());
+        let mut stderr_tee: Box<dyn io::AsyncWrite + Send + Unpin + 'static> =
+            Box::new(io::stderr());
+
+        // 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);
+                    }
+                }
+            } else {
+                tracing::error!("failed to parse state actor address");
+            }
+        }
 
         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,

hyperactor_state/Cargo.toml

@@ -10,6 +10,8 @@ license = "BSD-3-Clause"
 [dependencies]
 anyhow = "1.0.98"
 async-trait = "0.1.86"
+futures = { version = "0.3.30", features = ["async-await", "compat"] }
+hostname = "0.3"
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 hyperactor_macros = { version = "0.0.0", path = "../hyperactor_macros" }
 serde = { version = "1.0.185", features = ["derive", "rc"] }

hyperactor_state/src/client.rs

@@ -72,8 +72,8 @@ mod tests {
 
     use super::*;
     use crate::create_remote_client;
-    use crate::spawn_actor;
     use crate::test_utils::log_items;
+    use crate::test_utils::spawn_actor;
 
     #[tracing_test::traced_test]
     #[tokio::test]
@@ -83,7 +83,7 @@ mod tests {
         let params = ClientActorParams { sender };
         let client_proc_id =
             hyperactor::reference::ProcId(hyperactor::WorldId("client_server".to_string()), 0);
-        let (client_actor_addr, client_actor_handle) = spawn_actor::<ClientActor>(
+        let (client_actor_addr, client_actor_handle, _client_mailbox) = spawn_actor::<ClientActor>(
             client_actor_addr.clone(),
             client_proc_id,
             "state_client",

hyperactor_state/src/lib.rs

@@ -23,35 +23,10 @@ use hyperactor::mailbox::MailboxServer;
 use hyperactor::proc::Proc;
 
 pub mod client;
+pub mod log_writer;
 pub mod object;
 pub mod state_actor;
 
-/// Creates a state actor server at given address. Returns the server address and a handle to the
-/// state actor.
-#[allow(dead_code)]
-pub async fn spawn_actor<T: Actor + RemoteActor + Binds<T>>(
-    addr: ChannelAddr,
-    proc_id: ProcId,
-    actor_name: &str,
-    params: T::Params,
-) -> Result<(ChannelAddr, ActorHandle<T>)> {
-    // Use the provided ProcId directly
-    let proc = Proc::new(proc_id, BoxedMailboxSender::new(DialMailboxRouter::new()));
-
-    // Set up the channel server
-    let (local_addr, rx) = channel::serve(addr.clone()).await?;
-
-    // Spawn the actor with just a name - the system will generate the full actor ID
-    let actor_handle: ActorHandle<T> = proc.spawn(actor_name, params).await?;
-
-    // Undeliverable messages encountered by the mailbox server
-    // are to be returned to the system actor.
-    let _mailbox_handle = proc.clone().serve(rx, actor_handle.port());
-
-    // Return the address and handle (not a ref)
-    Ok((local_addr, actor_handle))
-}
-
 /// Creates a remote client that can send message to actors in the remote addr.
 /// It is important to keep the client proc alive for the remote_client's lifetime.
 pub async fn create_remote_client(addr: ChannelAddr) -> Result<(Proc, Mailbox)> {
@@ -66,26 +41,57 @@ pub async fn create_remote_client(addr: ChannelAddr) -> Result<(Proc, Mailbox)>
 }
 
 pub mod test_utils {
+    use super::*;
     use crate::object::GenericStateObject;
+    use crate::object::Kind;
     use crate::object::LogSpec;
     use crate::object::LogState;
+    use crate::object::Name;
     use crate::object::StateMetadata;
     use crate::object::StateObject;
 
-    pub fn log_items(seq_low: usize, seq_high: usize) -> Vec<GenericStateObject> {
+    pub fn log_items(seq_low: u64, seq_high: u64) -> Vec<GenericStateObject> {
         let mut log_items = vec![];
         let metadata = StateMetadata {
-            name: "test".to_string(),
-            kind: "log".to_string(),
+            name: Name::StdoutLog(("test_host".to_string(), 12345)),
+            kind: Kind::Log,
         };
         let spec = LogSpec {};
         for seq in seq_low..seq_high {
-            let state = LogState::new(seq, format!("state {}", seq));
+            let state = LogState::from_string(seq, format!("state {}", seq)).unwrap();
             let state_object =
                 StateObject::<LogSpec, LogState>::new(metadata.clone(), spec.clone(), state);
             let generic_state_object = GenericStateObject::try_from(state_object).unwrap();
             log_items.push(generic_state_object);
         }
         log_items
     }
+
+    /// Creates a state actor server at given address. Returns the server address, a handle to the
+    /// state actor, and a client mailbox for sending messages to the actor.
+    pub async fn spawn_actor<T: Actor + RemoteActor + Binds<T>>(
+        addr: ChannelAddr,
+        proc_id: ProcId,
+        actor_name: &str,
+        params: T::Params,
+    ) -> Result<(ChannelAddr, ActorHandle<T>, Mailbox)> {
+        // Use the provided ProcId directly
+        let proc = Proc::new(proc_id, BoxedMailboxSender::new(DialMailboxRouter::new()));
+
+        // Set up the channel server
+        let (local_addr, rx) = channel::serve(addr.clone()).await?;
+
+        // Spawn the actor with just a name - the system will generate the full actor ID
+        let actor_handle: ActorHandle<T> = proc.spawn(actor_name, params).await?;
+
+        // Create a client mailbox for sending messages to the actor
+        let client_mailbox = proc.attach("client")?;
+
+        // Undeliverable messages encountered by the mailbox server
+        // are to be returned to the system actor.
+        let _mailbox_handle = proc.clone().serve(rx, actor_handle.port());
+
+        // Return the address, handle, and client mailbox
+        Ok((local_addr, actor_handle, client_mailbox))
+    }
 }