flush log upon ipython notebook cell exit (#816)

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 registers the flush function upon a cell exiting.

Differential Revision: D79982702

Author: James Sun

python/monarch/_src/actor/proc_mesh.py

@@ -67,6 +67,16 @@
 from monarch._src.actor.future import DeprecatedNotAFuture, Future
 from monarch._src.actor.shape import MeshTrait
 
+HAS_IPYTHON = False
+try:
+    # Check if we are in ipython environment
+    # pyre-ignore[21]
+    from IPython import get_ipython
+
+    HAS_IPYTHON = True
+except ImportError:
+    pass
+
 HAS_TENSOR_ENGINE = False
 try:
     # Torch is needed for tensor engine
@@ -163,7 +173,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 +183,21 @@ async def _init_manager_actors_coro(
             aggregate_window_sec=3,
             level=logging.INFO,
         )
+        if HAS_IPYTHON and get_ipython() is not None:
+            # 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

@@ -724,6 +724,104 @@ 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
+
+                with unittest.mock.patch(
+                    "monarch._src.actor.proc_mesh.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())
+
+                    # 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: