create StateActor to store and serve 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/src/lib.rs

@@ -84,6 +84,7 @@ mod parse;
 pub mod proc;
 pub mod reference;
 pub mod simnet;
+mod state;
 pub mod supervision;
 pub mod sync;
 /// Test utilities

hyperactor/src/state/mod.rs

@@ -0,0 +1,10 @@
+/*
+ * 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.
+ */
+
+mod object;
+mod state_actor;

hyperactor/src/state/object.rs

@@ -0,0 +1,83 @@
+/*
+ * 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 serde::Deserialize;
+use serde::Serialize;
+
+use crate as hyperactor;
+use crate::Named;
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct StateMetadata {
+    /// Name of the actor.
+    pub name: String,
+    /// Kind of the object.
+    pub kind: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct StateObject<S, T> {
+    metadata: StateMetadata,
+    spec: S,
+    status: T,
+}
+
+impl<S, T> StateObject<S, T> {
+    #[allow(dead_code)]
+    pub fn new(metadata: StateMetadata, spec: S, status: T) -> Self {
+        Self {
+            metadata,
+            spec,
+            status,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct LogSpec;
+
+#[allow(dead_code)]
+#[derive(Debug, Serialize, Deserialize)]
+pub struct LogStatus {
+    /// A monotonically increasing sequence number.
+    seq: usize,
+    /// The message in the log.
+    message: String,
+}
+
+impl LogStatus {
+    pub fn new(seq: usize, message: String) -> Self {
+        Self { seq, message }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Named)]
+pub struct GenericStateObject {
+    metadata: StateMetadata,
+    data: String,
+}
+
+impl<S, T> From<StateObject<S, T>> for GenericStateObject
+where
+    S: Spec,
+    T: Status,
+{
+    fn from(value: StateObject<S, T>) -> Self {
+        Self {
+            metadata: value.metadata.clone(),
+            data: serde_json::to_string(&value).unwrap(),
+        }
+    }
+}
+
+pub trait Spec: Serialize + for<'de> Deserialize<'de> {}
+
+pub trait Status: Serialize + for<'de> Deserialize<'de> {}
+
+impl Spec for LogSpec {}
+impl Status for LogStatus {}

hyperactor/src/state/state_actor.rs

@@ -0,0 +1,162 @@
+/*
+ * 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 std::collections::VecDeque;
+
+use async_trait::async_trait;
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate as hyperactor;
+use crate::Actor;
+use crate::HandleClient;
+use crate::Handler;
+use crate::Instance;
+use crate::Named;
+use crate::OncePortRef;
+use crate::RefClient;
+use crate::state::object::GenericStateObject;
+
+/// Maximum number of logs to store in the buffer. Once the buffer is full, the oldest logs will be dropped.
+const LOG_CAPACITY: usize = 1 << 10;
+
+/// A state actor which serves as a centralized store for state.
+#[derive(Debug)]
+#[hyperactor::export(StateMessage)]
+pub struct StateActor {
+    /// Logs buffer.
+    logs: VecDeque<GenericStateObject>,
+}
+
+/// Endpoints for the state actor.
+#[derive(Handler, HandleClient, RefClient, Debug, Serialize, Deserialize, Named)]
+pub enum StateMessage {
+    /// Push a batch of logs to the logs buffer
+    SetLogs {
+        /// Object to set.
+        logs: Vec<GenericStateObject>,
+    },
+    /// Fetch logs from the buffer. This will drain the fetched logs from buffer.
+    GetLogs {
+        #[reply]
+        ret: OncePortRef<Vec<GenericStateObject>>,
+    },
+}
+
+#[async_trait]
+impl Actor for StateActor {
+    type Params = ();
+
+    async fn new(_params: ()) -> Result<Self, anyhow::Error> {
+        Ok(Self {
+            logs: VecDeque::new(),
+        })
+    }
+}
+
+#[async_trait]
+#[hyperactor::forward(StateMessage)]
+impl StateMessageHandler for StateActor {
+    async fn set_logs(
+        &mut self,
+        _this: &Instance<Self>,
+        logs: Vec<GenericStateObject>,
+    ) -> Result<(), anyhow::Error> {
+        self.logs.extend(logs);
+        while self.logs.len() > LOG_CAPACITY {
+            self.logs.pop_front();
+        }
+        Ok(())
+    }
+
+    async fn get_logs(
+        &mut self,
+        _this: &Instance<Self>,
+    ) -> Result<Vec<GenericStateObject>, anyhow::Error> {
+        let logs = self.logs.drain(..).collect::<Vec<_>>();
+        Ok(logs)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use crate::ActorHandle;
+    use crate::proc::Proc;
+    use crate::state::object::LogSpec;
+    use crate::state::object::LogStatus;
+    use crate::state::object::StateMetadata;
+    use crate::state::object::StateObject;
+
+    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 status = LogStatus::new(seq, format!("status {}", seq));
+            let state_object =
+                StateObject::<LogSpec, LogStatus>::new(metadata.clone(), spec.clone(), status);
+            let generic_state_object = GenericStateObject::from(state_object);
+            log_items.push(generic_state_object);
+        }
+        log_items
+    }
+
+    #[tokio::test]
+    async fn test_state_actor_set_get_logs() {
+        let proc = Proc::local();
+        let state_actor: ActorHandle<StateActor> = proc.spawn("state", ()).await.unwrap();
+        let client = proc.attach("client").unwrap();
+        let log_items_0_10 = log_items(0, 10);
+        // Set once and get once.
+        state_actor
+            .set_logs(&client, log_items_0_10.clone())
+            .await
+            .unwrap();
+        let fetched_log_items = state_actor.get_logs(&client).await.unwrap();
+        assert_eq!(fetched_log_items, log_items_0_10);
+
+        // Another get should return nothing.
+        let fetched_log_items = state_actor.get_logs(&client).await.unwrap();
+        assert_eq!(fetched_log_items, vec![]);
+
+        // Set twice and get should get all logs.
+        let log_items_10_20 = log_items(10, 20);
+        let log_items_20_30 = log_items(20, 30);
+        let log_items_10_30 = log_items(10, 30);
+        state_actor
+            .set_logs(&client, log_items_10_20.clone())
+            .await
+            .unwrap();
+        state_actor
+            .set_logs(&client, log_items_20_30.clone())
+            .await
+            .unwrap();
+        let fetched_log_items = state_actor.get_logs(&client).await.unwrap();
+        assert_eq!(fetched_log_items, log_items_10_30);
+
+        // over capacity logs should be dropped.
+        let log_items_0_max_capacity = log_items(0, LOG_CAPACITY);
+        let log_items_max_capacity_plus_10 = log_items(LOG_CAPACITY, LOG_CAPACITY + 1);
+        let log_items_10_max_capacity_plus_10 = log_items(1, LOG_CAPACITY + 1);
+        state_actor
+            .set_logs(&client, log_items_0_max_capacity.clone())
+            .await
+            .unwrap();
+        state_actor
+            .set_logs(&client, log_items_max_capacity_plus_10.clone())
+            .await
+            .unwrap();
+        let fetched_log_items = state_actor.get_logs(&client).await.unwrap();
+        assert_eq!(fetched_log_items, log_items_10_max_capacity_plus_10);
+    }
+}