(1/n) hook state client to client for logging (#445)

Summary:

Launch logging client through proc mesh. State actor bootstrap needs to be
manually setup. The following diffs will make things automated.

Reviewed By: shayne-fletcher

Differential Revision: D77848230

Author: James Sun

hyperactor_mesh/src/proc_mesh.rs

@@ -25,6 +25,7 @@ use hyperactor::actor::remote::Remote;
 use hyperactor::cap;
 use hyperactor::channel;
 use hyperactor::channel::ChannelAddr;
+use hyperactor::id;
 use hyperactor::mailbox;
 use hyperactor::mailbox::BoxableMailboxSender;
 use hyperactor::mailbox::BoxedMailboxSender;
@@ -38,6 +39,10 @@ use hyperactor::proc::Proc;
 use hyperactor::reference::ProcId;
 use hyperactor::reference::Reference;
 use hyperactor::supervision::ActorSupervisionEvent;
+use hyperactor_state::client::ClientActor;
+use hyperactor_state::client::ClientActorParams;
+use hyperactor_state::state_actor::StateActor;
+use hyperactor_state::state_actor::StateMessageClient;
 use ndslice::Range;
 use ndslice::Shape;
 use ndslice::ShapeError;
@@ -200,7 +205,7 @@ impl ProcMesh {
         client_proc
             .clone()
             .serve(client_rx, mailbox::monitored_return_handle());
-        router.bind(client_proc_id.clone().into(), client_proc_addr);
+        router.bind(client_proc_id.clone().into(), client_proc_addr.clone());
 
         // Bind this router to the global router, to enable cross-mesh routing.
         // TODO: unbind this when we incorporate mesh destruction too.
@@ -297,6 +302,42 @@ impl ProcMesh {
                 .map_err(anyhow::Error::from)?;
         }
 
+        // Get a reference to the state actor for streaming logs.
+        // TODO: bind logging options to python API so that users can choose if to stream (optionally aggregated) logs back or not.
+        // TODO: spin up state actor locally and remotely with names and addresses passed in here.
+        let state_actor_id = id!(state_server[0].state[0]);
+        let state_actor_ref = ActorRef::<StateActor>::attest(state_actor_id.clone());
+        let state_actor_addr = "tcp![::]:3000".parse::<ChannelAddr>().unwrap();
+        router.bind(state_actor_id.into(), state_actor_addr.clone());
+
+        let log_handler = Box::new(hyperactor_state::client::StdlogHandler {});
+        let params = ClientActorParams { log_handler };
+
+        let client_logging_actor: ActorRef<ClientActor> = client_proc
+            .spawn::<ClientActor>("logging_client", params)
+            .await
+            .unwrap()
+            .bind();
+
+        match state_actor_ref
+            .subscribe_logs(
+                &client,
+                client_proc_addr.clone(),
+                client_logging_actor.clone(),
+            )
+            .await
+        {
+            Ok(_) => {}
+            Err(err) => {
+                // TODO: talking to the state actor is not on the critical path.
+                // However, if the state actor is not reachable, we will receive spams of mailbox warnings.
+                tracing::warn!(
+                    "failed to subscribe to state actor logs; remote logging will not be available on the client: {}",
+                    err
+                );
+            }
+        }
+
         let shape = alloc.shape().clone();
         let world_id = alloc.world_id().clone();
 

hyperactor_state/src/client.rs

@@ -16,17 +16,82 @@ use hyperactor_macros::HandleClient;
 use hyperactor_macros::RefClient;
 use serde::Deserialize;
 use serde::Serialize;
-use tokio::sync::mpsc::Sender;
 
 use crate::object::GenericStateObject;
+use crate::object::Kind;
+use crate::object::LogSpec;
+use crate::object::LogState;
+use crate::object::Name;
+use crate::object::StateObject;
+
+pub trait LogHandler: Sync + Send + std::fmt::Debug + 'static {
+    // we cannot call it handle here as it conflicts with hyperactor macro
+    fn handle_log(&self, logs: Vec<GenericStateObject>) -> Result<()>;
+}
+
+/// A log handler that flushes GenericStateObject to stdout.
+#[derive(Debug)]
+pub struct StdlogHandler;
+
+impl LogHandler for StdlogHandler {
+    fn handle_log(&self, logs: Vec<GenericStateObject>) -> Result<()> {
+        for log in logs {
+            let metadata = log.metadata();
+            let deserialized_data: StateObject<LogSpec, LogState> = log.data().deserialized()?;
+
+            // Deserialize the message and process line by line with UTF-8
+            let message_lines = deserialize_message_lines(&deserialized_data.state.message)?;
+
+            // TODO: @lky D77377307 do not use raw string to distinguish between stdout and stderr
+            if metadata.kind != Kind::Log {
+                continue;
+            }
+            match &metadata.name {
+                Name::StdoutLog((hostname, pid)) => {
+                    for line in message_lines {
+                        // TODO: @lky hostname and pid should only be printed for non-aggregated logs. =
+                        // For aggregated logs, we should leave as is for better aggregation.
+                        println!("[{} {}] {}", hostname, pid, line);
+                    }
+                }
+                Name::StderrLog((hostname, pid)) => {
+                    for line in message_lines {
+                        eprintln!("[{} {}] {}", hostname, pid, line);
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+/// Deserialize a Serialized message and split it into UTF-8 lines
+fn deserialize_message_lines(
+    serialized_message: &hyperactor::data::Serialized,
+) -> Result<Vec<String>> {
+    // Try to deserialize as String first
+    if let Ok(message_str) = serialized_message.deserialized::<String>() {
+        return Ok(message_str.lines().map(|s| s.to_string()).collect());
+    }
+
+    // If that fails, try to deserialize as Vec<u8> and convert to UTF-8
+    if let Ok(message_bytes) = serialized_message.deserialized::<Vec<u8>>() {
+        let message_str = String::from_utf8(message_bytes)?;
+        return Ok(message_str.lines().map(|s| s.to_string()).collect());
+    }
+
+    // If both fail, return an error
+    anyhow::bail!("Failed to deserialize message as either String or Vec<u8>")
+}
 
 /// A client to interact with the state actor.
 #[derive(Debug)]
 #[hyperactor::export(
     handlers = [ClientMessage],
 )]
 pub struct ClientActor {
-    sender: Sender<Vec<GenericStateObject>>,
+    // TODO: extend hyperactor macro to support a generic to avoid using Box here.
+    log_handler: Box<dyn LogHandler>,
 }
 
 /// Endpoints for the client actor.
@@ -37,15 +102,17 @@ pub enum ClientMessage {
 }
 
 pub struct ClientActorParams {
-    pub sender: Sender<Vec<GenericStateObject>>,
+    pub log_handler: Box<dyn LogHandler>,
 }
 
 #[async_trait]
 impl Actor for ClientActor {
     type Params = ClientActorParams;
 
-    async fn new(ClientActorParams { sender }: ClientActorParams) -> Result<Self, anyhow::Error> {
-        Ok(Self { sender })
+    async fn new(
+        ClientActorParams { log_handler }: ClientActorParams,
+    ) -> Result<Self, anyhow::Error> {
+        Ok(Self { log_handler })
     }
 }
 
@@ -57,7 +124,7 @@ impl ClientMessageHandler for ClientActor {
         _cx: &Context<Self>,
         logs: Vec<GenericStateObject>,
     ) -> Result<(), anyhow::Error> {
-        self.sender.send(logs).await?;
+        self.log_handler.handle_log(logs)?;
         Ok(())
     }
 }
@@ -69,18 +136,38 @@ mod tests {
     use hyperactor::ActorRef;
     use hyperactor::channel;
     use hyperactor::channel::ChannelAddr;
+    use hyperactor::clock::Clock;
+    use hyperactor::data::Serialized;
+    use tokio::sync::mpsc::Sender;
 
     use super::*;
     use crate::create_remote_client;
     use crate::test_utils::log_items;
     use crate::test_utils::spawn_actor;
 
+    /// A log handler that flushes GenericStateObject to a mpsc channel.
+    #[derive(Debug)]
+    struct MpscLogHandler {
+        sender: Sender<Vec<GenericStateObject>>,
+    }
+
+    impl LogHandler for MpscLogHandler {
+        fn handle_log(&self, logs: Vec<GenericStateObject>) -> Result<()> {
+            let sender = self.sender.clone();
+            tokio::spawn(async move {
+                sender.send(logs).await.unwrap();
+            });
+            Ok(())
+        }
+    }
+
     #[tracing_test::traced_test]
     #[tokio::test]
     async fn test_client_basics() {
         let client_actor_addr = ChannelAddr::any(channel::ChannelTransport::Unix);
         let (sender, mut receiver) = tokio::sync::mpsc::channel::<Vec<GenericStateObject>>(10);
-        let params = ClientActorParams { sender };
+        let log_handler = Box::new(MpscLogHandler { sender });
+        let params = ClientActorParams { log_handler };
         let client_proc_id =
             hyperactor::reference::ProcId(hyperactor::WorldId("client_server".to_string()), 0);
         let (client_actor_addr, client_actor_handle, _client_mailbox) = spawn_actor::<ClientActor>(
@@ -102,7 +189,9 @@ mod tests {
             .unwrap();
 
         // Collect received messages with timeout
-        let fetched_logs = tokio::time::timeout(Duration::from_secs(1), receiver.recv())
+        let clock = hyperactor::clock::ClockKind::default();
+        let fetched_logs = clock
+            .timeout(Duration::from_secs(1), receiver.recv())
             .await
             .expect("timed out waiting for message")
             .expect("channel closed unexpectedly");
@@ -112,7 +201,53 @@ mod tests {
         assert_eq!(fetched_logs, log_items_0_10);
 
         // Now test that no extra message is waiting
-        let extra = tokio::time::timeout(Duration::from_millis(100), receiver.recv()).await;
+        let extra = clock
+            .timeout(Duration::from_millis(100), receiver.recv())
+            .await;
         assert!(extra.is_err(), "expected no more messages");
     }
+
+    #[test]
+    fn test_deserialize_message_lines_string() {
+        // Test deserializing a String message with multiple lines
+        let message = "Line 1\nLine 2\nLine 3".to_string();
+        let serialized = Serialized::serialize_anon(&message).unwrap();
+
+        let result = deserialize_message_lines(&serialized).unwrap();
+
+        assert_eq!(result, vec!["Line 1", "Line 2", "Line 3"]);
+
+        // Test deserializing a Vec<u8> message with UTF-8 content
+        let message_bytes = "Hello\nWorld\nUTF-8 \u{1F980}".as_bytes().to_vec();
+        let serialized = Serialized::serialize_anon(&message_bytes).unwrap();
+
+        let result = deserialize_message_lines(&serialized).unwrap();
+
+        assert_eq!(result, vec!["Hello", "World", "UTF-8 \u{1F980}"]);
+
+        // Test deserializing a single line message
+        let message = "Single line message".to_string();
+        let serialized = Serialized::serialize_anon(&message).unwrap();
+
+        let result = deserialize_message_lines(&serialized).unwrap();
+
+        assert_eq!(result, vec!["Single line message"]);
+
+        // Test deserializing an empty lines
+        let message = "\n\n".to_string();
+        let serialized = Serialized::serialize_anon(&message).unwrap();
+
+        let result = deserialize_message_lines(&serialized).unwrap();
+
+        assert_eq!(result, vec!["", ""]);
+
+        // Test error handling for invalid UTF-8 bytes
+        let invalid_utf8_bytes = vec![0xFF, 0xFE, 0xFD]; // Invalid UTF-8 sequence
+        let serialized = Serialized::serialize_anon(&invalid_utf8_bytes).unwrap();
+
+        let result = deserialize_message_lines(&serialized);
+
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("invalid utf-8"));
+    }
 }

hyperactor_state/src/state_actor.rs

@@ -92,13 +92,31 @@ mod tests {
     use std::time::Duration;
 
     use hyperactor::channel;
+    use tokio::sync::mpsc::Sender;
 
     use super::*;
     use crate::client::ClientActorParams;
+    use crate::client::LogHandler;
     use crate::create_remote_client;
     use crate::test_utils::log_items;
     use crate::test_utils::spawn_actor;
 
+    /// A log handler that flushes GenericStateObject to a mpsc channel.
+    #[derive(Debug)]
+    struct MpscLogHandler {
+        sender: Sender<Vec<GenericStateObject>>,
+    }
+
+    impl LogHandler for MpscLogHandler {
+        fn handle_log(&self, logs: Vec<GenericStateObject>) -> Result<()> {
+            let sender = self.sender.clone();
+            tokio::spawn(async move {
+                sender.send(logs).await.unwrap();
+            });
+            Ok(())
+        }
+    }
+
     #[tokio::test]
     async fn test_subscribe_logs() {
         let state_actor_addr = ChannelAddr::any(channel::ChannelTransport::Unix);
@@ -112,7 +130,8 @@ mod tests {
 
         let client_actor_addr = ChannelAddr::any(channel::ChannelTransport::Unix);
         let (sender, mut receiver) = tokio::sync::mpsc::channel::<Vec<GenericStateObject>>(20);
-        let params = ClientActorParams { sender };
+        let log_handler = Box::new(MpscLogHandler { sender });
+        let params = ClientActorParams { log_handler };
         let client_proc_id =
             hyperactor::reference::ProcId(hyperactor::WorldId("client_server".to_string()), 0);
         let (client_actor_addr, client_actor_handle, _client_mailbox) = spawn_actor::<ClientActor>(