[App] Enable debugger with LightningApp (#15590)

examples/app_boring/app_dynamic.py

@@ -64,4 +64,4 @@ class BoringApp(L.LightningFlow):
         return {"name": "Boring Tab", "content": self.dict["dst_w"].url + "/file" if "dst_w" in self.dict else ""}
 
 
-app = L.LightningApp(BoringApp(), debug=True)
+app = L.LightningApp(BoringApp(), log_level="debug")

examples/app_commands_and_api/app.py

@@ -50,4 +50,4 @@ class FlowCommands(LightningFlow):
         ]
 
 
-app = LightningApp(FlowCommands(), debug=True)
+app = LightningApp(FlowCommands(), log_level="debug")

examples/app_mount/app.py

@@ -32,4 +32,4 @@ class Flow(L.LightningFlow):
         self.work_1.run()
 
 
-app = L.LightningApp(Flow(), debug=True)
+app = L.LightningApp(Flow(), log_level="debug")

examples/app_template_streamlit_ui/app.py

@@ -45,4 +45,4 @@ class HelloWorld(LightningFlow):
         return [{"name": "StreamLitUI", "content": self.streamlit_ui}]
 
 
-app = LightningApp(HelloWorld(), debug=True)
+app = LightningApp(HelloWorld(), log_level="debug")

examples/app_v0/app.py

@@ -46,4 +46,4 @@ class V0App(L.LightningFlow):
         return [tab1, tab2, tab3]
 
 
-app = L.LightningApp(V0App(), debug=True)
+app = L.LightningApp(V0App(), log_level="debug")

examples/app_works_on_default_machine/app_v2.py

@@ -50,4 +50,4 @@ class Flow(LightningFlow):
         return [{"name": w.name, "content": w} for i, w in enumerate(self.works())]
 
 
-app = LightningApp(Flow(), debug=True)
+app = LightningApp(Flow(), log_level="debug")

src/lightning/__init__.py

@@ -45,6 +45,9 @@ import lightning.app  # isort: skip # noqa: E402
 
 lightning.app._PROJECT_ROOT = os.path.dirname(lightning.app._PROJECT_ROOT)
 
+# Enable breakpoint within forked processes.
+__builtins__["breakpoint"] = pdb.set_trace
+
 __all__ = [
     "LightningApp",
     "LightningFlow",

src/lightning_app/CHANGELOG.md

@@ -18,6 +18,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added a `start_with_flow` flag to the `LightningWork` which can be disabled to prevent the work from starting at the same time as the flow ([#15591](https://github.com/Lightning-AI/lightning/pull/15591))
 
+- Added support for running Lightning App with VSCode IDE debugger ([#15590](https://github.com/Lightning-AI/lightning/pull/15590))
+
+
 ### Changed
 
 - Changed the `flow.flows` to be recursive wont to align the behavior with the `flow.works` ([#15466](https://github.com/Lightning-AI/lightning/pull/15466))

src/lightning_app/core/app.py

@@ -27,7 +27,12 @@ from lightning_app.frontend import Frontend
 from lightning_app.storage import Drive, Path
 from lightning_app.storage.path import _storage_root_dir
 from lightning_app.utilities import frontend
-from lightning_app.utilities.app_helpers import _delta_to_app_state_delta, _LightningAppRef, Logger
+from lightning_app.utilities.app_helpers import (
+    _delta_to_app_state_delta,
+    _LightningAppRef,
+    _should_dispatch_app,
+    Logger,
+)
 from lightning_app.utilities.commands.base import _process_requests
 from lightning_app.utilities.component import _convert_paths_after_init, _validate_root_flow
 from lightning_app.utilities.enum import AppStage, CacheCallsKeys
@@ -52,7 +57,7 @@ class LightningApp:
         self,
         root: Union["LightningFlow", "LightningWork"],
         flow_cloud_compute: Optional["lightning_app.CloudCompute"] = None,
-        debug: bool = False,
+        log_level: str = "info",
         info: frontend.AppInfo = None,
         root_path: str = "",
     ):
@@ -70,7 +75,7 @@ class LightningApp:
             root: The root ``LightningFlow`` or ``LightningWork`` component, that defines all the app's nested
                  components, running infinitely. It must define a `run()` method that the app can call.
             flow_cloud_compute: The default Cloud Compute used for flow, Rest API and frontend's.
-            debug: Whether to activate the Lightning Logger debug mode.
+            log_level: The log level for the app, one of [`info`, `debug`].
                 This can be helpful when reporting bugs on Lightning repo.
             info: Provide additional info about the app which will be used to update html title,
                 description and image meta tags and specify any additional tags as list of html strings.
@@ -150,8 +155,11 @@ class LightningApp:
         # is only available after all Flows and Works have been instantiated.
         _convert_paths_after_init(self.root)
 
+        if log_level not in ("debug", "info"):
+            raise Exception(f"Log Level should be in ['debug', 'info']. Found {log_level}")
+
         # Lazily enable debugging.
-        if debug or DEBUG_ENABLED:
+        if log_level == "debug" or DEBUG_ENABLED:
             if not DEBUG_ENABLED:
                 os.environ["LIGHTNING_DEBUG"] = "2"
             _console.setLevel(logging.DEBUG)
@@ -162,6 +170,12 @@ class LightningApp:
         # this should happen once for all apps before the ui server starts running.
         frontend.update_index_file(FRONTEND_DIR, info=info, root_path=root_path)
 
+        if _should_dispatch_app():
+            os.environ["LIGHTNING_DISPATCHED"] = "1"
+            from lightning_app.runners import MultiProcessRuntime
+
+            MultiProcessRuntime(self).dispatch()
+
     def get_component_by_name(self, component_name: str):
         """Returns the instance corresponding to the given component name."""
         from lightning_app.structures import Dict as LightningDict

src/lightning_app/runners/runtime.py

@@ -1,4 +1,5 @@
 import multiprocessing
+import os
 import sys
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -54,6 +55,9 @@ def dispatch(
     from lightning_app.runners.runtime_type import RuntimeType
     from lightning_app.utilities.component import _set_flow_context
 
+    # Used to indicate Lightning has been dispatched
+    os.environ["LIGHTNING_DISPATCHED"] = "1"
+
     _set_flow_context()
 
     runtime_type = RuntimeType(runtime_type)

src/lightning_app/testing/testing.py

@@ -167,7 +167,7 @@ class _SingleWorkFlow(LightningFlow):
 def run_work_isolated(work, *args, start_server: bool = False, **kwargs):
     """This function is used to run a work a single time with multiprocessing runtime."""
     MultiProcessRuntime(
-        LightningApp(_SingleWorkFlow(work, args, kwargs), debug=True),
+        LightningApp(_SingleWorkFlow(work, args, kwargs), log_level="debug"),
         start_server=start_server,
     ).dispatch()
     # pop the stopped status.

src/lightning_app/utilities/app_helpers.py

@@ -488,3 +488,16 @@ def _load_state_dict(root_flow: "LightningFlow", state: Dict[str, Any], strict:
 
 def is_static_method(klass_or_instance, attr) -> bool:
     return isinstance(inspect.getattr_static(klass_or_instance, attr), staticmethod)
+
+
+def _debugger_is_active() -> bool:
+    """Return if the debugger is currently active."""
+    return hasattr(sys, "gettrace") and sys.gettrace() is not None
+
+
+def _should_dispatch_app() -> bool:
+    return (
+        _debugger_is_active()
+        and not bool(int(os.getenv("LIGHTNING_DISPATCHED", "0")))
+        and "LIGHTNING_APP_STATE_URL" not in os.environ
+    )

tests/tests_app/conftest.py

@@ -20,6 +20,8 @@ GITHUB_APP_URLS = {
     "template_react_ui": "https://github.com/Lightning-AI/lightning-template-react.git",
 }
 
+os.environ["LIGHTNING_DISPATCHED"] = "1"
+
 
 def pytest_sessionstart(*_):
     """Pytest hook that get called after the Session object has been created and before performing collection and

tests/tests_app/core/test_lightning_api.py

@@ -75,7 +75,7 @@ class _A(LightningFlow):
 @pytest.mark.parametrize("runtime_cls", [MultiProcessRuntime])
 def test_app_state_api(runtime_cls):
     """This test validates the AppState can properly broadcast changes from work within its own process."""
-    app = LightningApp(_A(), debug=True)
+    app = LightningApp(_A(), log_level="debug")
     runtime_cls(app, start_server=True).dispatch()
     assert app.root.work_a.var_a == -1
     _set_work_context()
@@ -110,7 +110,7 @@ class A2(LightningFlow):
 @pytest.mark.parametrize("runtime_cls", [SingleProcessRuntime])
 def test_app_state_api_with_flows(runtime_cls, tmpdir):
     """This test validates the AppState can properly broadcast changes from flows."""
-    app = LightningApp(A2(), debug=True)
+    app = LightningApp(A2(), log_level="debug")
     runtime_cls(app, start_server=True).dispatch()
     assert app.root.var_a == -1
 
@@ -185,7 +185,7 @@ class AppStageTestingApp(LightningApp):
 def test_app_stage_from_frontend(runtime_cls):
     """This test validates that delta from the `api_delta_queue` manipulating the ['app_state']['stage'] would
     start and stop the app."""
-    app = AppStageTestingApp(FlowA(), debug=True)
+    app = AppStageTestingApp(FlowA(), log_level="debug")
     app.stage = AppStage.BLOCKING
     runtime_cls(app, start_server=True).dispatch()
 
@@ -197,7 +197,7 @@ def test_update_publish_state_and_maybe_refresh_ui():
     - receives a notification to refresh the UI and makes a GET Request (streamlit).
     """
 
-    app = AppStageTestingApp(FlowA(), debug=True)
+    app = AppStageTestingApp(FlowA(), log_level="debug")
     publish_state_queue = _MockQueue("publish_state_queue")
     api_response_queue = _MockQueue("api_response_queue")
 
@@ -224,7 +224,7 @@ async def test_start_server(x_lightning_type, monkeypatch):
         def get(self, timeout: int = 0):
             return self._queue[0]
 
-    app = AppStageTestingApp(FlowA(), debug=True)
+    app = AppStageTestingApp(FlowA(), log_level="debug")
     app._update_layout()
     app.stage = AppStage.BLOCKING
     publish_state_queue = InfiniteQueue("publish_state_queue")

tests/tests_app/core/test_lightning_app.py

@@ -107,7 +107,7 @@ class SimpleFlow(LightningFlow):
 @pytest.mark.parametrize("runtime_cls", [SingleProcessRuntime])
 def test_simple_app(component_cls, runtime_cls, tmpdir):
     comp = component_cls()
-    app = LightningApp(comp, debug=True)
+    app = LightningApp(comp, log_level="debug")
     assert app.root == comp
     expected = {
         "app_state": ANY,
@@ -249,7 +249,7 @@ def test_get_component_by_name_raises():
 
 @pytest.mark.parametrize("runtime_cls", [SingleProcessRuntime, MultiProcessRuntime])
 def test_nested_component(runtime_cls):
-    app = LightningApp(A(), debug=True)
+    app = LightningApp(A(), log_level="debug")
     runtime_cls(app, start_server=False).dispatch()
     assert app.root.w_a.c == 1
     assert app.root.b.w_b.c == 1
@@ -361,7 +361,7 @@ class SimpleApp2(LightningApp):
 @pytest.mark.parametrize("runtime_cls", [SingleProcessRuntime, MultiProcessRuntime])
 def test_app_restarting_move_to_blocking(runtime_cls, tmpdir):
     """Validates sending restarting move the app to blocking again."""
-    app = SimpleApp2(CounterFlow(), debug=True)
+    app = SimpleApp2(CounterFlow(), log_level="debug")
     runtime_cls(app, start_server=False).dispatch()
 
 
@@ -395,7 +395,7 @@ class AppWithFrontend(LightningApp):
 @mock.patch("lightning_app.frontend.stream_lit.StreamlitFrontend.stop_server")
 def test_app_starts_with_complete_state_copy(_, __):
     """Test that the LightningApp captures the initial state in a separate copy when _run() gets called."""
-    app = AppWithFrontend(FlowWithFrontend(), debug=True)
+    app = AppWithFrontend(FlowWithFrontend(), log_level="debug")
     MultiProcessRuntime(app, start_server=False).dispatch()
     assert app.run_once_call_count == 3
 
@@ -992,7 +992,7 @@ def test_debug_mode_logging():
 
     from lightning_app.core.app import _console
 
-    app = LightningApp(A4(), debug=True)
+    app = LightningApp(A4(), log_level="debug")
     assert _console.level == logging.DEBUG
     assert os.getenv("LIGHTNING_DEBUG") == "2"
 

tests/tests_app/runners/test_cloud.py

@@ -761,10 +761,10 @@ class TestAppCreationClient:
                             ),
                             drives=[],
                             user_requested_compute_config=V1UserRequestedComputeConfig(
-                                name="default", count=1, disk_size=0, shm_size=0
+                                name="default", count=1, disk_size=0, shm_size=0, preemptible=mock.ANY
                             ),
                             network_config=[V1NetworkConfig(name=mock.ANY, host=None, port=8080)],
-                            cluster_id="test",
+                            cluster_id=mock.ANY,
                         ),
                     )
                 ],

tests/tests_app/storage/test_path.py

@@ -377,7 +377,7 @@ class SourceToDestFlow(LightningFlow):
 
 def test_multiprocess_path_in_work_and_flow(tmpdir):
     root = SourceToDestFlow(tmpdir)
-    app = LightningApp(root, debug=True)
+    app = LightningApp(root, log_level="debug")
     MultiProcessRuntime(app, start_server=False).dispatch()
 
 
@@ -551,7 +551,7 @@ class OverwriteFolderFlow(LightningFlow):
 def test_path_get_overwrite(tmpdir):
     """Test that .get(overwrite=True) overwrites the entire directory and replaces all files."""
     root = OverwriteFolderFlow(tmpdir)
-    app = LightningApp(root, debug=True)
+    app = LightningApp(root, log_level="debug")
     MultiProcessRuntime(app, start_server=False).dispatch()
 
 

tests/tests_app/storage/test_payload.py

@@ -146,7 +146,7 @@ class Flow(LightningFlow):
 def test_payload_works(tmpdir):
     """This tests validates the payload api can be used to transfer return values from a work to another."""
     with mock.patch("lightning_app.storage.path._storage_root_dir", lambda: pathlib.Path(tmpdir)):
-        app = LightningApp(Flow(), debug=True)
+        app = LightningApp(Flow(), log_level="debug")
         MultiProcessRuntime(app, start_server=False).dispatch()
 
     os.remove("value_all")

tests/tests_app/structures/test_structures.py

@@ -494,6 +494,6 @@ class FlowPayload(LightningFlow):
 
 
 def test_structures_with_payload():
-    app = LightningApp(FlowPayload(), debug=True)
+    app = LightningApp(FlowPayload(), log_level="debug")
     MultiProcessRuntime(app, start_server=False).dispatch()
     os.remove("payload")

tests/tests_app/utilities/test_proxies.py

@@ -266,7 +266,7 @@ class WorkRunnerPatch(WorkRunner):
 
 @mock.patch("lightning_app.runners.backends.mp_process.WorkRunner", WorkRunnerPatch)
 def test_proxy_timeout():
-    app = LightningApp(FlowTimeout(), debug=True)
+    app = LightningApp(FlowTimeout(), log_level="debug")
     MultiProcessRuntime(app, start_server=False).dispatch()
 
     call_hash = app.root.work._calls[CacheCallsKeys.LATEST_CALL_HASH]

tests/tests_app_examples/collect_failures/app.py

@@ -43,4 +43,4 @@ class RootFlow(LightningFlow):
 
 
 if __name__ == "__main__":
-    app = LightningApp(RootFlow(), debug=True)
+    app = LightningApp(RootFlow(), log_level="debug")

tests/tests_app_examples/conftest.py

@@ -11,6 +11,8 @@ from lightning_app.utilities.packaging import cloud_compute
 from lightning_app.utilities.packaging.app_config import _APP_CONFIG_FILENAME
 from lightning_app.utilities.state import AppState
 
+os.environ["LIGHTNING_DISPATCHED"] = "1"
+
 
 def pytest_sessionfinish(session, exitstatus):
     """Pytest hook that get called after whole test run finished, right before returning the exit status to the

tests/tests_app_examples/custom_work_dependencies/app.py

@@ -24,7 +24,8 @@ class WorkWithCustomBaseImage(LightningWork):
     def __init__(self, cloud_compute: CloudCompute = CloudCompute(), **kwargs):
         # this image has been created from ghcr.io/gridai/base-images:v1.8-cpu
         # by just adding an empty file at /content/.e2e_test
-        image_tag = os.getenv("LIGHTNING_E2E_TEST_IMAGE_VERSION", "v1.12")
+        image_tag = "v1.15"
+        # image_tag = os.getenv("LIGHTNING_E2E_TEST_IMAGE_VERSION", "v1.15")
         custom_image = f"ghcr.io/gridai/image-for-testing-custom-images-in-e2e:{image_tag}"
         build_config = BuildConfig(image=custom_image)
         super().__init__(parallel=True, **kwargs, cloud_compute=cloud_compute, cloud_build_config=build_config)
@@ -50,4 +51,4 @@ class CustomWorkBuildConfigChecker(LightningFlow):
             self._exit()
 
 
-app = LightningApp(CustomWorkBuildConfigChecker(), debug=True)
+app = LightningApp(CustomWorkBuildConfigChecker())

tests/tests_app_examples/idle_timeout/app.py

@@ -68,4 +68,4 @@ class RootFlow(LightningFlow):
             self._exit()
 
 
-app = LightningApp(RootFlow(), debug=True)
+app = LightningApp(RootFlow(), log_level="debug")