sync flush of proc mesh (meta-pytorch#823)

Summary:

Provide sync flush so it is guaranteed all the flushed logs on the
remote procs will be streamed back and flushed on client's stdout/stderr.

Differential Revision: D80051803

Author: James Sun

hyperactor_mesh/src/logging.rs

@@ -22,12 +22,15 @@ use chrono::DateTime;
 use chrono::Local;
 use hyperactor::Actor;
 use hyperactor::ActorRef;
+use hyperactor::Bind;
 use hyperactor::Context;
 use hyperactor::HandleClient;
 use hyperactor::Handler;
 use hyperactor::Instance;
 use hyperactor::Named;
+use hyperactor::OncePortRef;
 use hyperactor::RefClient;
+use hyperactor::Unbind;
 use hyperactor::channel;
 use hyperactor::channel::ChannelAddr;
 use hyperactor::channel::ChannelRx;
@@ -39,9 +42,6 @@ use hyperactor::channel::TxStatus;
 use hyperactor::clock::Clock;
 use hyperactor::clock::RealClock;
 use hyperactor::data::Serialized;
-use hyperactor::message::Bind;
-use hyperactor::message::Bindings;
-use hyperactor::message::Unbind;
 use hyperactor_telemetry::env;
 use hyperactor_telemetry::log_file_path;
 use serde::Deserialize;
@@ -235,6 +235,24 @@ impl fmt::Display for Aggregator {
     }
 }
 
+/// Messages that can be sent to the LogClientActor remotely.
+#[derive(
+    Debug,
+    Clone,
+    Serialize,
+    Deserialize,
+    Named,
+    Handler,
+    HandleClient,
+    RefClient,
+    Bind,
+    Unbind
+)]
+pub enum LogFlushMessage {
+    /// Flush the log
+    ForceSyncFlush {},
+}
+
 /// Messages that can be sent to the LogClientActor remotely.
 #[derive(
     Debug,
@@ -260,7 +278,10 @@ pub enum LogMessage {
     },
 
     /// Flush the log
-    Flush {},
+    Flush {
+        /// If true, force a flush sync barrier across all procs
+        synced: bool,
+    },
 }
 
 /// Messages that can be sent to the LogClient locally.
@@ -279,6 +300,14 @@ pub enum LogClientMessage {
         /// The time window in seconds to aggregate logs. If None, aggregation is disabled.
         aggregate_window_sec: Option<u64>,
     },
+
+    /// Synchronously flush all the logs from all the procs. This is for client to call.
+    StartSyncFlush {
+        /// Expect these many procs to ack the flush message.
+        expected_procs: usize,
+        /// Return once we have received the acks from all the procs
+        reply: OncePortRef<()>,
+    },
 }
 
 /// Trait for sending logs
@@ -351,7 +380,9 @@ impl LogSender for LocalLogSender {
     async fn flush(&mut self) -> anyhow::Result<()> {
         // send will make sure message is delivered
         if TxStatus::Active == *self.status.borrow() {
-            match self.tx.send(LogMessage::Flush {}).await {
+            // this is just to make sure the log line is sent to the other side of the channel.
+            // it is up to the forwarder to decide when to flush the log.
+            match self.tx.send(LogMessage::Flush { synced: false }).await {
                 Ok(()) => Ok(()),
                 Err(e) => {
                     tracing::error!("log sender {} error sending flush message: {}", self.pid, e);
@@ -570,7 +601,9 @@ impl<T: LogSender + Unpin + 'static, S: io::AsyncWrite + Send + Unpin + 'static>
     Named,
     Handler,
     HandleClient,
-    RefClient
+    RefClient,
+    Bind,
+    Unbind
 )]
 pub enum LogForwardMessage {
     /// Receive the log from the parent process and forward ti to the client.
@@ -580,18 +613,6 @@ pub enum LogForwardMessage {
     SetMode { stream_to_client: bool },
 }
 
-impl Bind for LogForwardMessage {
-    fn bind(&mut self, _bindings: &mut Bindings) -> anyhow::Result<()> {
-        Ok(())
-    }
-}
-
-impl Unbind for LogForwardMessage {
-    fn unbind(&self, _bindings: &mut Bindings) -> anyhow::Result<()> {
-        Ok(())
-    }
-}
-
 /// A log forwarder that receives the log from its parent process and forward it back to the client
 #[derive(Debug)]
 #[hyperactor::export(
@@ -659,17 +680,28 @@ impl Actor for LogForwardActor {
 #[hyperactor::forward(LogForwardMessage)]
 impl LogForwardMessageHandler for LogForwardActor {
     async fn forward(&mut self, ctx: &Context<Self>) -> Result<(), anyhow::Error> {
-        if let Ok(LogMessage::Log {
-            hostname,
-            pid,
-            output_target,
-            payload,
-        }) = self.rx.recv().await
-        {
-            if self.stream_to_client {
-                self.logging_client_ref
-                    .log(ctx, hostname, pid, output_target, payload)
-                    .await?;
+        match self.rx.recv().await {
+            Ok(LogMessage::Flush { synced }) => {
+                if synced {
+                    self.logging_client_ref.flush(ctx, true).await?;
+                } else {
+                    // no need to do anything. The previous messages have already been sent to the client.
+                }
+            }
+            Ok(LogMessage::Log {
+                hostname,
+                pid,
+                output_target,
+                payload,
+            }) => {
+                if self.stream_to_client {
+                    self.logging_client_ref
+                        .log(ctx, hostname, pid, output_target, payload)
+                        .await?;
+                }
+            }
+            Err(e) => {
+                return Err(e.into());
             }
         }
 
@@ -708,6 +740,54 @@ fn deserialize_message_lines(
     anyhow::bail!("Failed to deserialize message as either String or Vec<u8>")
 }
 
+/// An actor that send flush message to the log forwarder actor.
+/// The reason we need an extra actor instead of reusing the log forwarder actor
+/// is because the log forwarder can be blocked on the rx.recv() that listens on the new log lines.
+/// Thus, we need to create anew channel as a tx to send the flush message to the log forwarder
+/// So we do not get into a deadlock.
+#[derive(Debug)]
+#[hyperactor::export(
+    spawn = true,
+    handlers = [LogFlushMessage {cast = true}],
+)]
+pub struct LogFlushActor {
+    tx: ChannelTx<LogMessage>,
+}
+
+#[async_trait]
+impl Actor for LogFlushActor {
+    type Params = ();
+
+    async fn new(_: ()) -> Result<Self, anyhow::Error> {
+        let log_channel: ChannelAddr = match std::env::var(BOOTSTRAP_LOG_CHANNEL) {
+            Ok(channel) => channel.parse()?,
+            Err(err) => {
+                tracing::debug!(
+                    "log forwarder actor failed to read env var {}: {}",
+                    BOOTSTRAP_LOG_CHANNEL,
+                    err
+                );
+                // TODO: this should error out; it can only happen with local proc; we need to fix it.
+                ChannelAddr::any(ChannelTransport::Unix)
+            }
+        };
+        let tx = channel::dial::<LogMessage>(log_channel)?;
+
+        Ok(Self { tx })
+    }
+}
+
+#[async_trait]
+#[hyperactor::forward(LogFlushMessage)]
+impl LogFlushMessageHandler for LogFlushActor {
+    async fn force_sync_flush(&mut self, _cx: &Context<Self>) -> Result<(), anyhow::Error> {
+        self.tx
+            .send(LogMessage::Flush { synced: true })
+            .await
+            .map_err(anyhow::Error::from)
+    }
+}
+
 /// A client to receive logs from remote processes
 #[derive(Debug)]
 #[hyperactor::export(
@@ -719,6 +799,8 @@ pub struct LogClientActor {
     aggregators: HashMap<OutputTarget, Aggregator>,
     last_flush_time: SystemTime,
     next_flush_deadline: Option<SystemTime>,
+    ongoing_flush_port: Option<OncePortRef<()>>,
+    unflushed_procs: usize,
 }
 
 impl LogClientActor {
@@ -748,6 +830,12 @@ impl LogClientActor {
             OutputTarget::Stderr => eprintln!("{}", message),
         }
     }
+
+    fn flush_internal(&mut self) {
+        self.print_aggregators();
+        self.last_flush_time = RealClock.system_time_now();
+        self.next_flush_deadline = None;
+    }
 }
 
 #[async_trait]
@@ -766,6 +854,8 @@ impl Actor for LogClientActor {
             aggregators,
             last_flush_time: RealClock.system_time_now(),
             next_flush_deadline: None,
+            ongoing_flush_port: None,
+            unflushed_procs: 0,
         })
     }
 }
@@ -817,20 +907,23 @@ impl LogMessageHandler for LogClientActor {
                 let new_deadline = self.last_flush_time + Duration::from_secs(window);
                 let now = RealClock.system_time_now();
                 if new_deadline <= now {
-                    self.flush(cx).await?;
+                    self.flush_internal();
                 } else {
                     let delay = new_deadline.duration_since(now)?;
                     match self.next_flush_deadline {
                         None => {
                             self.next_flush_deadline = Some(new_deadline);
-                            cx.self_message_with_delay(LogMessage::Flush {}, delay)?;
+                            cx.self_message_with_delay(LogMessage::Flush { synced: false }, delay)?;
                         }
                         Some(deadline) => {
                             // Some early log lines have alrady triggered the flush.
                             if new_deadline < deadline {
                                 // This can happen if the user has adjusted the aggregation window.
                                 self.next_flush_deadline = Some(new_deadline);
-                                cx.self_message_with_delay(LogMessage::Flush {}, delay)?;
+                                cx.self_message_with_delay(
+                                    LogMessage::Flush { synced: false },
+                                    delay,
+                                )?;
                             }
                         }
                     }
@@ -841,10 +934,21 @@ impl LogMessageHandler for LogClientActor {
         Ok(())
     }
 
-    async fn flush(&mut self, _cx: &Context<Self>) -> Result<(), anyhow::Error> {
-        self.print_aggregators();
-        self.last_flush_time = RealClock.system_time_now();
-        self.next_flush_deadline = None;
+    async fn flush(&mut self, cx: &Context<Self>, synced: bool) -> Result<(), anyhow::Error> {
+        if synced {
+            if self.unflushed_procs == 0 || self.ongoing_flush_port.is_none() {
+                anyhow::bail!("found no ongoing flush request");
+            }
+            self.unflushed_procs -= 1;
+            if self.unflushed_procs == 0 {
+                self.flush_internal();
+                let reply = self.ongoing_flush_port.take().unwrap();
+                self.ongoing_flush_port = None;
+                reply.send(cx, ()).map_err(anyhow::Error::from)?;
+            }
+        } else {
+            self.flush_internal();
+        }
 
         Ok(())
     }
@@ -865,6 +969,21 @@ impl LogClientMessageHandler for LogClientActor {
         self.aggregate_window_sec = aggregate_window_sec;
         Ok(())
     }
+
+    async fn start_sync_flush(
+        &mut self,
+        _cx: &Context<Self>,
+        expected_procs_flushed: usize,
+        reply: OncePortRef<()>,
+    ) -> Result<(), anyhow::Error> {
+        if self.unflushed_procs > 0 || self.ongoing_flush_port.is_some() {
+            anyhow::bail!("forcing a flush while the ongoing flush has not finished yet");
+        }
+
+        self.ongoing_flush_port = Some(reply.clone());
+        self.unflushed_procs = expected_procs_flushed;
+        Ok(())
+    }
 }
 
 #[cfg(test)]

monarch_extension/src/logging.rs

@@ -13,6 +13,8 @@ use hyperactor_mesh::RootActorMesh;
 use hyperactor_mesh::actor_mesh::ActorMesh;
 use hyperactor_mesh::logging::LogClientActor;
 use hyperactor_mesh::logging::LogClientMessage;
+use hyperactor_mesh::logging::LogFlushActor;
+use hyperactor_mesh::logging::LogFlushMessage;
 use hyperactor_mesh::logging::LogForwardActor;
 use hyperactor_mesh::logging::LogForwardMessage;
 use hyperactor_mesh::selection::Selection;
@@ -33,6 +35,9 @@ use pyo3::types::PyModule;
 pub struct LoggingMeshClient {
     // handles remote process log forwarding; no python runtime
     forwarder_mesh: SharedCell<RootActorMesh<'static, LogForwardActor>>,
+    // because forwarder mesh keeps listening to the new coming logs,
+    // the flush mesh is a way to unblock it from busy waiting a log and do sync flush.
+    flush_mesh: SharedCell<RootActorMesh<'static, LogFlushActor>>,
     // handles python logger; has python runtime
     logger_mesh: SharedCell<RootActorMesh<'static, LoggerRuntimeActor>>,
     client_actor: ActorHandle<LogClientActor>,
@@ -47,9 +52,11 @@ impl LoggingMeshClient {
             let client_actor = proc_mesh.client_proc().spawn("log_client", ()).await?;
             let client_actor_ref = client_actor.bind();
             let forwarder_mesh = proc_mesh.spawn("log_forwarder", &client_actor_ref).await?;
+            let flush_mesh = proc_mesh.spawn("log_flusher", &()).await?;
             let logger_mesh = proc_mesh.spawn("logger", &()).await?;
             Ok(Self {
                 forwarder_mesh,
+                flush_mesh,
                 logger_mesh,
                 client_actor,
             })
@@ -97,6 +104,36 @@ impl LoggingMeshClient {
 
         Ok(())
     }
+
+    // A sync flush mechanism for the client make sure all the stdout/stderr are streamed back and flushed.
+    fn flush(&self, proc_mesh: &PyProcMesh) -> PyResult<PyPythonTask> {
+        let inner_mesh = proc_mesh.try_inner()?;
+        let (tx, rx) = inner_mesh.client().open_once_port::<()>();
+        // First initialize a sync flush.
+        self.client_actor
+            .send(LogClientMessage::StartSyncFlush {
+                expected_procs: inner_mesh.shape().slice().len(),
+                reply: tx.bind(),
+            })
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+
+        // Then ask all the flushers to ask the log forwarders to sync flush
+        let flush_inner_mesh = self.flush_mesh.borrow().map_err(anyhow::Error::msg)?;
+        flush_inner_mesh
+            .cast(
+                inner_mesh.client(),
+                Selection::True,
+                LogFlushMessage::ForceSyncFlush {},
+            )
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+
+        // Finally the forwarder will send sync point back to the client, flush, and return.
+        PyPythonTask::new(async move {
+            rx.recv()
+                .await
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+        })
+    }
 }
 
 impl Drop for LoggingMeshClient {

python/monarch/_rust_bindings/monarch_extension/logging.pyi

@@ -21,3 +21,4 @@ class LoggingMeshClient:
     def set_mode(
         self, stream_to_client: bool, aggregate_window_sec: int | None, level: int
     ) -> None: ...
+    def flush(self, proc_mesh: ProcMesh) -> PythonTask[None]: ...