flush logs in the end no matter what (meta-pytorch#761)

Summary:
Pull Request resolved: meta-pytorch#761

make sure we are able to flush whatever remaining when program ends.
However, we could still miss logs as log through channel is not always
guaranteed to be streamed back when a program ends.

Differential Revision: D79620358

Author: James Sun

hyperactor_mesh/src/logging.rs

@@ -12,6 +12,7 @@ use std::path::Path;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::RwLock;
 use std::task::Context as TaskContext;
 use std::task::Poll;
 use std::time::Duration;
@@ -668,6 +669,30 @@ pub struct LogClientActor {
     /// The watch sender for the aggregation window in seconds
     aggregate_window_tx: watch::Sender<u64>,
     should_aggregate: bool,
+    // Store aggregators directly in the actor for access in Drop
+    aggregators: Arc<RwLock<HashMap<OutputTarget, Aggregator>>>,
+}
+
+impl LogClientActor {
+    fn print_aggregators(aggregators: &RwLock<HashMap<OutputTarget, Aggregator>>) {
+        let mut aggregators_guard = aggregators.write().unwrap();
+        for (output_target, aggregator) in aggregators_guard.iter_mut() {
+            if aggregator.is_empty() {
+                continue;
+            }
+            match output_target {
+                OutputTarget::Stdout => {
+                    println!("{}", aggregator);
+                }
+                OutputTarget::Stderr => {
+                    eprintln!("{}", aggregator);
+                }
+            }
+
+            // Reset the aggregator
+            aggregator.reset();
+        }
+    }
 }
 
 #[async_trait]
@@ -683,26 +708,42 @@ impl Actor for LogClientActor {
         let (aggregate_window_tx, aggregate_window_rx) =
             watch::channel(DEFAULT_AGGREGATE_WINDOW_SEC);
 
+        // Initialize aggregators
+        let mut aggregators = HashMap::new();
+        aggregators.insert(OutputTarget::Stderr, Aggregator::new());
+        aggregators.insert(OutputTarget::Stdout, Aggregator::new());
+        let aggregators = Arc::new(RwLock::new(aggregators));
+
+        // Clone aggregators for the aggregator task
+        let aggregators_for_task = Arc::clone(&aggregators);
+
         // Start the loggregator
-        let aggregator_handle =
-            { tokio::spawn(async move { start_aggregator(log_rx, aggregate_window_rx).await }) };
+        let aggregator_handle = tokio::spawn(async move {
+            start_aggregator(log_rx, aggregate_window_rx, aggregators_for_task).await
+        });
 
         Ok(Self {
             log_tx,
             aggregator_handle,
             aggregate_window_tx,
-            should_aggregate: false,
+            should_aggregate: true,
+            aggregators,
         })
     }
 }
 
+impl Drop for LogClientActor {
+    fn drop(&mut self) {
+        // Flush the remaining logs before shutting down
+        Self::print_aggregators(&self.aggregators);
+    }
+}
+
 async fn start_aggregator(
     mut log_rx: mpsc::Receiver<(OutputTarget, String)>,
     mut interval_sec_rx: watch::Receiver<u64>,
+    aggregators: Arc<RwLock<HashMap<OutputTarget, Aggregator>>>,
 ) -> anyhow::Result<()> {
-    let mut aggregators = HashMap::new();
-    aggregators.insert(OutputTarget::Stderr, Aggregator::new());
-    aggregators.insert(OutputTarget::Stdout, Aggregator::new());
     let mut interval =
         tokio::time::interval(tokio::time::Duration::from_secs(*interval_sec_rx.borrow()));
 
@@ -711,7 +752,8 @@ async fn start_aggregator(
         tokio::select! {
             // Process incoming log messages
             Some((output_target, log_line)) = log_rx.recv() => {
-                if let Some(aggregator) = aggregators.get_mut(&output_target) {
+                let mut aggregators_guard = aggregators.write().unwrap();
+                if let Some(aggregator) = aggregators_guard.get_mut(&output_target) {
                     if let Err(e) = aggregator.add_line(&log_line) {
                         tracing::error!("error adding log line: {}", e);
                     }
@@ -726,24 +768,14 @@ async fn start_aggregator(
 
             // Every interval tick, print and reset the aggregator
             _ = interval.tick() => {
-                for (output_target, aggregator) in aggregators.iter_mut() {
-                    if aggregator.is_empty() {
-                        continue;
-                    }
-                    if output_target == &OutputTarget::Stdout {
-                        println!("{}", aggregator);
-                    } else {
-                        eprintln!("{}", aggregator);
-                    }
-
-                    // Reset the aggregator
-                    aggregator.reset();
-                }
+                LogClientActor::print_aggregators(&aggregators);
             }
 
             // Exit if the channel is closed
             else => {
                 tracing::error!("log channel closed, exiting aggregator");
+                // Print final aggregated logs before shutting down
+                LogClientActor::print_aggregators(&aggregators);
                 break;
             }
         }

monarch_extension/src/logging.rs

@@ -92,6 +92,12 @@ impl LoggingMeshClient {
     }
 }
 
+impl Drop for LoggingMeshClient {
+    fn drop(&mut self) {
+        let _ = self.client_actor.drain_and_stop().unwrap();
+    }
+}
+
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<LoggingMeshClient>()?;
     Ok(())

python/tests/actor_log_flush_test.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# 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.
+
+# pyre-strict
+
+import asyncio
+import logging
+import sys
+
+from monarch.actor import Actor, endpoint, proc_mesh
+
+
+class Printer(Actor):
+    def __init__(self) -> None:
+        self.logger: logging.Logger = logging.getLogger()
+
+    @endpoint
+    async def print(self, content: str) -> None:
+        print(f"{content}", flush=True)
+        sys.stdout.flush()
+        sys.stderr.flush()
+
+
+async def main() -> None:
+    pm = await proc_mesh(gpus=2)
+    # never flush
+    await pm.logging_option(aggregate_window_sec=1000)
+    am = await pm.spawn("printer", Printer)
+
+    # These should be streamed to client
+    for _ in range(5):
+        await am.print.call("has print streaming")
+
+    # Sleep a tiny so we allow the logs to stream back to the client
+    await asyncio.sleep(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/tests/test_python_actors.py

@@ -6,10 +6,12 @@
 
 # pyre-unsafe
 import asyncio
+import importlib
 import logging
 import operator
 import os
 import re
+import subprocess
 import sys
 import tempfile
 import threading
@@ -749,6 +751,37 @@ async def test_logging_option_defaults() -> None:
             pass
 
 
+# oss_skip: importlib not pulling resource correctly in git CI, needs to be revisited
+@pytest.mark.oss_skip
+async def test_flush_logs_fast_exit() -> None:
+    # We use a subprocess to run the test so we can handle the flushed logs at the end.
+    # Otherwise, it is hard to restore the original stdout/stderr.
+    with (
+        importlib.resources.as_file(
+            importlib.resources.files(__package__).joinpath("actor_log_flush_test")
+        ) as test_log_main,
+    ):
+        if not test_log_main.exists():
+            raise ImportError("cannot find file binary test_log_main")
+
+        # Run the binary in a separate process and capture stdout and stderr
+        process = subprocess.Popen(
+            [test_log_main],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+
+        # Wait for the process to complete and get output
+        stdout, _ = process.communicate(timeout=60)
+
+        # Check if the process ended without error
+        if process.returncode != 0:
+            raise RuntimeError(f"Process ended with error code {process.returncode}. ")
+
+        # Assertions on the captured output
+        assert "has print streaming" in stdout, stdout
+
+
 class SendAlot(Actor):
     @endpoint
     async def send(self, port: Port[int]):