create StateActor to stream logs for hyperactor (#211)

Summary:
Pull Request resolved: #211

This is an initial commit for the state actor. This diff creates a circular buffer to store latest log messages. If the circular buffer is full, messages will be discarded with a FIFO rule.

Next:
1. register the StateActor with actors in the proc_mesh so logs can be streamed in.
2. stream the logs out to loggregator so we can see the logs in `hyper log`.

Differential Revision: D76343950

Author: kaiyuan-li

hyperactor_state/Cargo.toml

@@ -0,0 +1,19 @@
+# @generated by autocargo from //monarch/hyperactor_state:hyperactor_state
+
+[package]
+name = "hyperactor_state"
+version = "0.0.0"
+authors = ["Meta"]
+edition = "2021"
+license = "BSD-3-Clause"
+
+[dependencies]
+anyhow = "1.0.95"
+async-trait = "0.1.86"
+hyperactor = { version = "0.0.0", path = "../hyperactor" }
+hyperactor_macros = { version = "0.0.0", path = "../hyperactor_macros" }
+serde = { version = "1.0.185", features = ["derive", "rc"] }
+serde_json = { version = "1.0.140", features = ["float_roundtrip", "unbounded_depth"] }
+tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
+tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
+unicode-ident = "1.0.12"

hyperactor_state/src/client.rs

@@ -0,0 +1,123 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use anyhow::Result;
+use async_trait::async_trait;
+use hyperactor::Actor;
+use hyperactor::Handler;
+use hyperactor::Instance;
+use hyperactor::Named;
+use hyperactor_macros::HandleClient;
+use hyperactor_macros::RefClient;
+use serde::Deserialize;
+use serde::Serialize;
+use tokio::sync::mpsc::Sender;
+
+use crate::object::GenericStateObject;
+
+/// A client to interact with the state actor.
+#[derive(Debug)]
+#[hyperactor::export(ClientMessage)]
+pub struct ClientActor {
+    sender: Sender<GenericStateObject>,
+}
+
+/// Endpoints for the client actor.
+#[derive(Handler, HandleClient, RefClient, Debug, Serialize, Deserialize, Named)]
+pub enum ClientMessage {
+    /// Push a batch of logs to the logs buffer
+    Logs {
+        /// Object to set.
+        logs: Vec<GenericStateObject>,
+    },
+}
+
+pub struct ClientActorParams {
+    pub sender: Sender<GenericStateObject>,
+}
+
+#[async_trait]
+impl Actor for ClientActor {
+    type Params = ClientActorParams;
+
+    async fn new(ClientActorParams { sender }: ClientActorParams) -> Result<Self, anyhow::Error> {
+        Ok(Self { sender })
+    }
+}
+
+#[async_trait]
+#[hyperactor::forward(ClientMessage)]
+impl ClientMessageHandler for ClientActor {
+    async fn logs(
+        &mut self,
+        _this: &Instance<Self>,
+        logs: Vec<GenericStateObject>,
+    ) -> Result<(), anyhow::Error> {
+        for log in logs {
+            self.sender.send(log).await?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use hyperactor::channel;
+    use hyperactor::channel::ChannelAddr;
+    use hyperactor::id;
+
+    use super::*;
+    use crate::create_remote_client;
+    use crate::spawn_actor;
+    use crate::test_utils::log_items;
+
+    #[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::<GenericStateObject>(10);
+        let params = ClientActorParams { sender };
+        let (client_actor_addr, client_actor_ref) = spawn_actor::<ClientActor>(
+            client_actor_addr.clone(),
+            id![state_client[0].state_client],
+            params,
+        )
+        .await
+        .unwrap();
+
+        let (_client_proc, remote_client) = create_remote_client(client_actor_addr).await.unwrap();
+
+        let log_items_0_10 = log_items(0, 10);
+        client_actor_ref
+            .logs(&remote_client, log_items_0_10.clone())
+            .await
+            .unwrap();
+
+        // Collect received messages with timeout
+        let mut fetched_logs = vec![];
+        for _ in 0..10 {
+            // Timeout prevents hanging if a message is missing
+            let log = tokio::time::timeout(Duration::from_secs(1), receiver.recv())
+                .await
+                .expect("timed out waiting for message")
+                .expect("channel closed unexpectedly");
+
+            fetched_logs.push(log);
+        }
+
+        // Verify we received all expected logs
+        assert_eq!(fetched_logs.len(), 10);
+        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;
+        assert!(extra.is_err(), "expected no more messages");
+    }
+}

hyperactor_state/src/lib.rs

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use anyhow::Result;
+use hyperactor::Actor;
+use hyperactor::ActorHandle;
+use hyperactor::ActorId;
+use hyperactor::ActorRef;
+use hyperactor::Mailbox;
+use hyperactor::actor::Binds;
+use hyperactor::actor::RemoteActor;
+use hyperactor::channel;
+use hyperactor::channel::ChannelAddr;
+use hyperactor::id;
+use hyperactor::mailbox::BoxedMailboxSender;
+use hyperactor::mailbox::DialMailboxRouter;
+use hyperactor::mailbox::MailboxClient;
+use hyperactor::mailbox::MailboxServer;
+use hyperactor::mailbox::MessageEnvelope;
+use hyperactor::mailbox::Undeliverable;
+use hyperactor::proc::Proc;
+
+mod client;
+mod object;
+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(crate) async fn spawn_actor<T: Actor + RemoteActor + Binds<T>>(
+    addr: ChannelAddr,
+    actor_id: ActorId,
+    params: T::Params,
+) -> Result<(ChannelAddr, ActorRef<T>)> {
+    let proc_id = actor_id.proc_id();
+    let proc = Proc::new(
+        proc_id.clone(),
+        BoxedMailboxSender::new(DialMailboxRouter::new()),
+    );
+    let (local_addr, rx) = channel::serve(addr.clone()).await?;
+    let actor_handle: ActorHandle<T> = proc.spawn(actor_id.name(), params).await?;
+    actor_handle.bind::<T>();
+
+    // Undeliverable messages encountered by the mailbox server
+    // are to be returned to the system actor.
+    let return_handle = actor_handle.port::<Undeliverable<MessageEnvelope>>();
+    let _mailbox_handle = proc.clone().serve(rx, return_handle);
+
+    Ok((local_addr, actor_handle.bind()))
+}
+
+/// 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(crate) async fn create_remote_client(addr: ChannelAddr) -> Result<(Proc, Mailbox)> {
+    let remote_sender = MailboxClient::new(channel::dial(addr).unwrap());
+    let client_proc_id = id!(client).random_user_proc();
+    let client_proc = Proc::new(
+        client_proc_id.clone(),
+        BoxedMailboxSender::new(remote_sender),
+    );
+    let remote_client = client_proc.attach("client").unwrap();
+    Ok((client_proc, remote_client))
+}
+
+#[cfg(test)]
+pub(crate) mod test_utils {
+    use crate::object::GenericStateObject;
+    use crate::object::LogSpec;
+    use crate::object::LogState;
+    use crate::object::StateMetadata;
+    use crate::object::StateObject;
+
+    pub(crate) fn log_items(seq_low: usize, seq_high: usize) -> Vec<GenericStateObject> {
+        let mut log_items = vec![];
+        let metadata = StateMetadata {
+            name: "test".to_string(),
+            kind: "log".to_string(),
+        };
+        let spec = LogSpec {};
+        for seq in seq_low..seq_high {
+            let state = LogState::new(seq, format!("state {}", seq));
+            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
+    }
+}

hyperactor_state/src/object.rs

@@ -0,0 +1,90 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use hyperactor::Named;
+use serde::Deserialize;
+use serde::Serialize;
+use serde_json::Value;
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct StateMetadata {
+    /// Name of the state object owner.
+    pub name: String,
+    /// Kind of the object.
+    pub kind: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct StateObject<S, T> {
+    metadata: StateMetadata,
+    spec: S,
+    state: T,
+}
+
+impl<S, T> StateObject<S, T> {
+    #[allow(dead_code)]
+    pub fn new(metadata: StateMetadata, spec: S, state: T) -> Self {
+        Self {
+            metadata,
+            spec,
+            state,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct LogSpec;
+
+#[allow(dead_code)]
+#[derive(Debug, Serialize, Deserialize)]
+pub struct LogState {
+    /// A monotonically increasing sequence number.
+    seq: usize,
+    /// The message in the log.
+    message: String,
+}
+
+impl LogState {
+    #[allow(dead_code)]
+    pub fn new(seq: usize, message: String) -> Self {
+        Self { seq, message }
+    }
+}
+
+/// A generic state object which is the partially serialized version of a
+/// [`StateObject`]. Since [`StateObject`] takes generic types, those type information
+/// can be retrieved from the metadata to deserialize [`GenericStateObject`] into
+/// a [`StateObject<S, T>`].
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Named)]
+pub struct GenericStateObject {
+    metadata: StateMetadata,
+    data: String,
+}
+
+impl<S, T> TryFrom<StateObject<S, T>> for GenericStateObject
+where
+    S: Spec,
+    T: State,
+{
+    type Error = serde_json::Error;
+
+    fn try_from(obj: StateObject<S, T>) -> Result<Self, Self::Error> {
+        Ok(Self {
+            metadata: obj.metadata.clone(),
+            data: serde_json::to_string(&obj)?,
+        })
+    }
+}
+
+/// Spec is the define the desired state of an object, defined by the user.
+pub trait Spec: Serialize + for<'de> Deserialize<'de> {}
+/// State is the current state of an object.
+pub trait State: Serialize + for<'de> Deserialize<'de> {}
+
+impl Spec for LogSpec {}
+impl State for LogState {}