check if we are in ipython env

Summary:
In ipython notebook, a cell can end fast. Yet the process can still run in the background.
However, the background process will not flush logs to the existing cell anymore.
The patch register the flush function upon a cell exiting.

Differential Revision: D79982702

Author: James Sun

hyperactor_mesh/src/logging.rs

@@ -27,6 +27,7 @@ use hyperactor::HandleClient;
 use hyperactor::Handler;
 use hyperactor::Instance;
 use hyperactor::Named;
+use hyperactor::OncePortRef;
 use hyperactor::RefClient;
 use hyperactor::channel;
 use hyperactor::channel::ChannelAddr;
@@ -279,6 +280,11 @@ pub enum LogClientMessage {
         /// The time window in seconds to aggregate logs. If None, aggregation is disabled.
         aggregate_window_sec: Option<u64>,
     },
+
+    Flush {
+        /// Synchronously flush all the logs
+        reply: OncePortRef<()>,
+    },
 }
 
 /// Trait for sending logs
@@ -748,6 +754,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]
@@ -817,7 +829,7 @@ 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 {
@@ -842,9 +854,7 @@ impl LogMessageHandler for LogClientActor {
     }
 
     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;
+        self.flush_internal();
 
         Ok(())
     }
@@ -865,6 +875,15 @@ impl LogClientMessageHandler for LogClientActor {
         self.aggregate_window_sec = aggregate_window_sec;
         Ok(())
     }
+
+    async fn flush(
+        &mut self,
+        cx: &Context<Self>,
+        reply: OncePortRef<()>,
+    ) -> Result<(), anyhow::Error> {
+        self.flush_internal();
+        reply.send(cx, ()).map_err(anyhow::Error::from)
+    }
 }
 
 #[cfg(test)]

monarch_extension/src/logging.rs

@@ -9,12 +9,14 @@
 #![allow(unsafe_op_in_unsafe_fn)]
 
 use hyperactor::ActorHandle;
+use hyperactor::PortRef;
 use hyperactor_mesh::RootActorMesh;
 use hyperactor_mesh::actor_mesh::ActorMesh;
 use hyperactor_mesh::logging::LogClientActor;
 use hyperactor_mesh::logging::LogClientMessage;
 use hyperactor_mesh::logging::LogForwardActor;
 use hyperactor_mesh::logging::LogForwardMessage;
+use hyperactor_mesh::logging::LogMessage;
 use hyperactor_mesh::selection::Selection;
 use hyperactor_mesh::shared_cell::SharedCell;
 use monarch_hyperactor::logging::LoggerRuntimeActor;
@@ -90,6 +92,18 @@ impl LoggingMeshClient {
 
         Ok(())
     }
+
+    fn flush(&self, proc_mesh: &PyProcMesh) -> PyResult<PyPythonTask> {
+        let (tx, rx) = proc_mesh.try_inner()?.client().open_once_port::<()>();
+        self.client_actor
+            .send(LogClientMessage::Flush { reply: tx.bind() })
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        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]: ...

python/monarch/_src/actor/proc_mesh.py

@@ -67,6 +67,13 @@
 from monarch._src.actor.future import DeprecatedNotAFuture, Future
 from monarch._src.actor.shape import MeshTrait
 
+try:
+    # Check if we are in ipython environment
+    # pyre-ignore[21]
+    from IPython import get_ipython
+except ImportError:
+    get_ipython = None
+
 HAS_TENSOR_ENGINE = False
 try:
     # Torch is needed for tensor engine
@@ -163,7 +170,7 @@ async def _init_manager_actors_coro(
         proc_mesh_: "Shared[HyProcMesh]",
         setup: Callable[[], None] | None = None,
     ) -> "HyProcMesh":
-        proc_mesh = await proc_mesh_
+        proc_mesh: HyProcMesh = await proc_mesh_
         # WARNING: it is unsafe to await self._proc_mesh here
         # because self._proc_mesh is the result of this function itself!
 
@@ -173,6 +180,21 @@ async def _init_manager_actors_coro(
             aggregate_window_sec=3,
             level=logging.INFO,
         )
+        if get_ipython:
+            # For ipython environment, a cell can end fast with threads running in background.
+            # Flush all the ongoing logs proactively to avoid missing logs.
+            assert self._logging_mesh_client is not None
+            logging_client: LoggingMeshClient = self._logging_mesh_client
+            ipython = get_ipython()
+
+            # pyre-ignore[21]
+            from IPython.core.interactiveshell import ExecutionResult
+
+            # pyre-ignore[11]
+            def flush_logs(_: ExecutionResult) -> None:
+                return Future(coro=logging_client.flush(proc_mesh).spawn().task()).get()
+
+            ipython.events.register("post_run_cell", flush_logs)
 
         _rdma_manager = (
             # type: ignore[16]

python/tests/test_python_actors.py

@@ -715,6 +715,109 @@ async def test_logging_option_defaults() -> None:
             pass
 
 
+@pytest.mark.timeout(60)
+async def test_flush_logs_ipython() -> None:
+    """Test that logs are flushed when get_ipython is available and post_run_cell event is triggered."""
+    # Save original file descriptors
+    original_stdout_fd = os.dup(1)  # stdout
+
+    try:
+        # Create temporary files to capture output
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False) as stdout_file:
+            stdout_path = stdout_file.name
+
+            # Redirect file descriptors to our temp files
+            os.dup2(stdout_file.fileno(), 1)
+
+            # Also redirect Python's sys.stdout
+            original_sys_stdout = sys.stdout
+            sys.stdout = stdout_file
+
+            try:
+                # Mock IPython environment
+                class MockExecutionResult:
+                    pass
+
+                class MockEvents:
+                    def __init__(self):
+                        self.callbacks = {}
+
+                    def register(self, event_name, callback):
+                        if event_name not in self.callbacks:
+                            self.callbacks[event_name] = []
+                        self.callbacks[event_name].append(callback)
+
+                    def trigger(self, event_name, *args, **kwargs):
+                        if event_name in self.callbacks:
+                            for callback in self.callbacks[event_name]:
+                                callback(*args, **kwargs)
+
+                class MockIPython:
+                    def __init__(self):
+                        self.events = MockEvents()
+
+                mock_ipython = MockIPython()
+
+                # Patch get_ipython to return our mock using unittest.mock
+                import unittest.mock
+
+                import monarch._src.actor.proc_mesh as proc_mesh_module
+
+                with unittest.mock.patch.object(
+                    proc_mesh_module, "get_ipython", lambda: mock_ipython
+                ):
+                    pm = await proc_mesh(gpus=2)
+                    am = await pm.spawn("printer", Printer)
+
+                    # Set aggregation window to ensure logs are buffered
+                    await pm.logging_option(
+                        stream_to_client=True, aggregate_window_sec=600
+                    )
+                    await asyncio.sleep(1)
+
+                    # Generate some logs that will be aggregated
+                    for _ in range(5):
+                        await am.print.call("ipython test log")
+
+                    # Trigger the post_run_cell event which should flush logs
+                    mock_ipython.events.trigger("post_run_cell", MockExecutionResult())
+
+                    # Wait a bit to ensure flush completes
+                    await asyncio.sleep(1)
+
+                    # Flush all outputs
+                    stdout_file.flush()
+                    os.fsync(stdout_file.fileno())
+
+            finally:
+                # Restore Python's sys.stdout
+                sys.stdout = original_sys_stdout
+
+        # Restore original file descriptors
+        os.dup2(original_stdout_fd, 1)
+
+        # Read the captured output
+        with open(stdout_path, "r") as f:
+            stdout_content = f.read()
+
+        # Clean up temp files
+        os.unlink(stdout_path)
+
+        # Verify that logs were flushed when the post_run_cell event was triggered
+        # We should see the aggregated logs in the output
+        assert re.search(
+            r"\[10 similar log lines\].*ipython test log", stdout_content
+        ), stdout_content
+
+    finally:
+        # Ensure file descriptors are restored even if something goes wrong
+        try:
+            os.dup2(original_stdout_fd, 1)
+            os.close(original_stdout_fd)
+        except OSError:
+            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: