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

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

Differential Revision: D77848230

Author: highker

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;
@@ -199,7 +204,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.
@@ -296,6 +301,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,46 @@ 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::LogSpec;
+use crate::object::LogState;
+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()?;
+            // TODO: @lky D77377307 do not use raw string to distinguish between stdout and stderr
+            if metadata.name == "stderr" {
+                eprintln!("(remote logging) {}", deserialized_data.state.message);
+            } else {
+                println!("(remote logging) {}", deserialized_data.state.message);
+            }
+            println!();
+        }
+        Ok(())
+    }
+}
 
 /// 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 +66,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 +88,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 +100,36 @@ mod tests {
     use hyperactor::ActorRef;
     use hyperactor::channel;
     use hyperactor::channel::ChannelAddr;
+    use tokio::sync::mpsc::Sender;
 
     use super::*;
     use crate::create_remote_client;
     use crate::spawn_actor;
     use crate::test_utils::log_items;
 
+    /// 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) = spawn_actor::<ClientActor>(

hyperactor_state/src/object.rs

@@ -23,7 +23,7 @@ pub struct StateMetadata {
 pub struct StateObject<S, T> {
     metadata: StateMetadata,
     spec: S,
-    state: T,
+    pub(crate) state: T,
 }
 
 impl<S, T> StateObject<S, T> {
@@ -46,7 +46,7 @@ pub struct LogState {
     /// A monotonically increasing sequence number.
     seq: usize,
     /// The message in the log.
-    message: String,
+    pub(crate) message: String,
 }
 
 impl LogState {

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::spawn_actor;
     use crate::test_utils::log_items;
 
+    /// 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) = spawn_actor::<ClientActor>(