[App] Improve pdb for multiprocessing (#15950)

Co-authored-by: thomas <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
(cherry picked from commit 482b279)

Author: 3 people

examples/app_dag/requirements.txt

@@ -1,2 +1,2 @@
-sklearn
+scikit-learn
 pandas

pyproject.toml

@@ -94,6 +94,7 @@ module = [
     "lightning_app.frontend.streamlit_base",
     "lightning_app.frontend.utils",
     "lightning_app.frontend.web",
+    "lightning_app.perf.pdb",
     "lightning_app.runners.backends.__init__",
     "lightning_app.runners.backends.backend",
     "lightning_app.runners.backends.cloud",

src/lightning_app/CHANGELOG.md

@@ -46,6 +46,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed a bug where using `L.app.structures` would cause multiple apps to be opened and fail with an error in the cloud ([#15911](https://github.com/Lightning-AI/lightning/pull/15911))
 - Fixed PythonServer generating noise on M1 ([#15949](https://github.com/Lightning-AI/lightning/pull/15949))
 
+- Fixed multiprocessing breakpoint ([#15950](https://github.com/Lightning-AI/lightning/pull/15950))
 
 ## [1.8.3] - 2022-11-22
 

src/lightning_app/perf/pdb.py

@@ -1,19 +1,36 @@
+import multiprocessing
+import os
 import pdb
 import sys
-from typing import Any
 
+_stdin = [None]
+_stdin_lock = multiprocessing.Lock()
+try:
+    _stdin_fd = sys.stdin.fileno()
+except Exception:
+    _stdin_fd = None
 
+
+# Taken from https://github.com/facebookresearch/metaseq/blob/main/metaseq/pdb.py
 class MPPdb(pdb.Pdb):
-    """debugger for forked programs."""
+    """A Pdb wrapper that works in a multiprocessing environment."""
+
+    def __init__(self) -> None:
+        pdb.Pdb.__init__(self, nosigint=True)
 
-    def interaction(self, *args: Any, **kwargs: Any) -> None:
-        _stdin = sys.stdin
-        try:
-            sys.stdin = open("/dev/stdin")
-            pdb.Pdb.interaction(self, *args, **kwargs)
-        finally:
-            sys.stdin = _stdin
+    def _cmdloop(self) -> None:
+        stdin_back = sys.stdin
+        with _stdin_lock:
+            try:
+                if _stdin_fd is not None:
+                    if not _stdin[0]:
+                        _stdin[0] = os.fdopen(_stdin_fd)
+                    sys.stdin = _stdin[0]
+                self.cmdloop()
+            finally:
+                sys.stdin = stdin_back
 
 
-def set_trace(*args: Any, **kwargs: Any) -> None:
-    MPPdb().set_trace(*args, **kwargs)
+def set_trace() -> None:
+    pdb = MPPdb()
+    pdb.set_trace(sys._getframe().f_back)