check if we are in ipython env

Summary:
In ipython notebook, a cell can end fast. It may lose unflushed logs
unable to print to stdout/stderr. The patch register the flush function upon a
cell exiting.

Differential Revision: D79982702

Author: highker

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: