[App] Mock missing package imports when launching in the cloud (#15711)

Co-authored-by: manskx <[email protected]>

Author: ethanwharris

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 friendly error message when attempting to run the default cloud compute with a custom base image configured ([#14929](https://github.com/Lightning-AI/lightning/pull/14929))
 
 
+- Improved support for running apps when dependencies aren't installed ([#15711](https://github.com/Lightning-AI/lightning/pull/15711))
+
+
 ### Changed
 
 -

src/lightning_app/runners/cloud.py

@@ -4,7 +4,6 @@
 import string
 import sys
 import time
-import traceback
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, List, Optional, Union
@@ -63,7 +62,7 @@
 from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.dependency_caching import get_hash
-from lightning_app.utilities.load_app import _prettifiy_exception, load_app_from_file
+from lightning_app.utilities.load_app import load_app_from_file
 from lightning_app.utilities.packaging.app_config import _get_config_file, AppConfig
 from lightning_app.utilities.packaging.lightning_utils import _prepare_lightning_wheels_and_requirements
 from lightning_app.utilities.secrets import _names_to_ids
@@ -475,26 +474,17 @@ def _project_has_sufficient_credits(self, project: V1Membership, app: Optional[L
 
     @classmethod
     def load_app_from_file(cls, filepath: str) -> "LightningApp":
-        """This is meant to use only locally for cloud runtime."""
+        """Load a LightningApp from a file, mocking the imports."""
         try:
-            app = load_app_from_file(filepath, raise_exception=True)
-        except ModuleNotFoundError:
-            # this is very generic exception.
-            logger.info("Could not load the app locally. Starting the app directly on the cloud.")
-            # we want to format the exception as if no frame was on top.
-            exp, val, tb = sys.exc_info()
-            listing = traceback.format_exception(exp, val, tb)
-            # remove the entry for the first frame
-            del listing[1]
-            from lightning_app.testing.helpers import EmptyFlow
-
-            # Create a mocking app.
-            app = LightningApp(EmptyFlow())
-
+            app = load_app_from_file(filepath, raise_exception=True, mock_imports=True)
         except FileNotFoundError as e:
             raise e
         except Exception:
-            _prettifiy_exception(filepath)
+            from lightning_app.testing.helpers import EmptyFlow
+
+            # Create a generic app.
+            logger.info("Could not load the app locally. Starting the app directly on the cloud.")
+            app = LightningApp(EmptyFlow())
         return app
 
 

src/lightning_app/testing/testing.py

@@ -361,7 +361,7 @@ def run_app_in_cloud(
         except playwright._impl._api_types.TimeoutError:
             print("'Create Project' dialog not visible, skipping.")
 
-        admin_page.locator(f"text={name}").click()
+        admin_page.locator(f"role=link[name='{name}']").click()
         sleep(5)
         # Scroll to the bottom of the page. Used to capture all logs.
         admin_page.evaluate(

src/lightning_app/utilities/app_helpers.py

@@ -1,5 +1,6 @@
 import abc
 import asyncio
+import builtins
 import enum
 import functools
 import inspect
@@ -10,9 +11,11 @@
 import threading
 import time
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Tuple, Type, TYPE_CHECKING
+from unittest.mock import MagicMock
 
 import websockets
 from deepdiff import Delta
@@ -486,6 +489,29 @@ def _load_state_dict(root_flow: "LightningFlow", state: Dict[str, Any], strict:
                 raise Exception(f"The component {component_name} was re-created during state reloading.")
 
 
+class _MagicMockJsonSerializable(MagicMock):
+    @staticmethod
+    def __json__():
+        return "{}"
+
+
+def _mock_import(*args, original_fn=None):
+    try:
+        return original_fn(*args)
+    except Exception:
+        return _MagicMockJsonSerializable()
+
+
+@contextmanager
+def _mock_missing_imports():
+    original_fn = builtins.__import__
+    builtins.__import__ = functools.partial(_mock_import, original_fn=original_fn)
+    try:
+        yield
+    finally:
+        builtins.__import__ = original_fn
+
+
 def is_static_method(klass_or_instance, attr) -> bool:
     return isinstance(inspect.getattr_static(klass_or_instance, attr), staticmethod)
 

src/lightning_app/utilities/layout.py

@@ -4,6 +4,7 @@
 
 import lightning_app
 from lightning_app.frontend.frontend import Frontend
+from lightning_app.utilities.app_helpers import _MagicMockJsonSerializable
 from lightning_app.utilities.cloud import is_running_in_cloud
 
 
@@ -39,6 +40,9 @@ def _collect_layout(app: "lightning_app.LightningApp", flow: "lightning_app.Ligh
         # When running locally, the target will get overwritten by the dispatcher when launching the frontend servers
         # When running in the cloud, the frontend code will construct the URL based on the flow name
         return flow._layout
+    elif isinstance(layout, _MagicMockJsonSerializable):
+        # Do nothing
+        pass
     elif isinstance(layout, dict):
         layout = _collect_content_layout([layout], flow)
     elif isinstance(layout, (list, tuple)) and all(isinstance(item, dict) for item in layout):
@@ -103,6 +107,9 @@ def _collect_content_layout(layout: List[Dict], flow: "lightning_app.LightningFl
             else:
                 entry["content"] = ""
                 entry["target"] = ""
+        elif isinstance(entry["content"], _MagicMockJsonSerializable):
+            # Do nothing
+            pass
         else:
             m = f"""
             A dictionary returned by `{flow.__class__.__name__}.configure_layout()` contains an unsupported entry.

src/lightning_app/utilities/load_app.py

@@ -4,14 +4,15 @@
 import traceback
 import types
 from contextlib import contextmanager
+from copy import copy
 from typing import Dict, List, TYPE_CHECKING, Union
 
 from lightning_app.utilities.exceptions import MisconfigurationException
 
 if TYPE_CHECKING:
     from lightning_app import LightningApp, LightningFlow, LightningWork
 
-from lightning_app.utilities.app_helpers import Logger
+from lightning_app.utilities.app_helpers import _mock_missing_imports, Logger
 
 logger = Logger(__name__)
 
@@ -30,7 +31,7 @@ def _prettifiy_exception(filepath: str):
     sys.exit(1)
 
 
-def load_app_from_file(filepath: str, raise_exception: bool = False) -> "LightningApp":
+def load_app_from_file(filepath: str, raise_exception: bool = False, mock_imports: bool = False) -> "LightningApp":
     """Load a LightningApp from a file.
 
     Arguments:
@@ -50,7 +51,11 @@ def load_app_from_file(filepath: str, raise_exception: bool = False) -> "Lightni
     module = _create_fake_main_module(filepath)
     try:
         with _patch_sys_argv():
-            exec(code, module.__dict__)
+            if mock_imports:
+                with _mock_missing_imports():
+                    exec(code, module.__dict__)
+            else:
+                exec(code, module.__dict__)
     except Exception as e:
         if raise_exception:
             raise e
@@ -140,7 +145,7 @@ def _patch_sys_argv():
     """
     from lightning_app.cli.lightning_cli import run_app
 
-    original_argv = sys.argv
+    original_argv = copy(sys.argv)
     # 1: Remove the CLI command
     if sys.argv[:3] == ["lightning", "run", "app"]:
         sys.argv = sys.argv[3:]

tests/tests_app/runners/test_cloud.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import sys
 from copy import copy
 from pathlib import Path
 from unittest import mock
@@ -43,7 +44,7 @@
 from lightning_app.runners import backends, cloud, CloudRuntime
 from lightning_app.runners.cloud import _validate_build_spec_and_compute
 from lightning_app.storage import Drive, Mount
-from lightning_app.testing.helpers import EmptyFlow
+from lightning_app.testing.helpers import EmptyFlow, EmptyWork
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.dependency_caching import get_hash
 from lightning_app.utilities.packaging.cloud_compute import CloudCompute
@@ -1230,6 +1231,57 @@ def test_load_app_from_file_module_error():
     assert isinstance(empty_app.root, EmptyFlow)
 
 
+@pytest.mark.parametrize(
+    "lines",
+    [
+        [
+            "import this_package_is_not_real",
+            "from lightning_app import LightningApp",
+            "from lightning_app.testing.helpers import EmptyWork",
+            "app = LightningApp(EmptyWork())",
+        ],
+        [
+            "from this_package_is_not_real import this_module_is_not_real",
+            "from lightning_app import LightningApp",
+            "from lightning_app.testing.helpers import EmptyWork",
+            "app = LightningApp(EmptyWork())",
+        ],
+        [
+            "import this_package_is_not_real",
+            "from this_package_is_not_real import this_module_is_not_real",
+            "from lightning_app import LightningApp",
+            "from lightning_app.testing.helpers import EmptyWork",
+            "app = LightningApp(EmptyWork())",
+        ],
+        [
+            "import this_package_is_not_real",
+            "from lightning_app import LightningApp",
+            "from lightning_app.core.flow import _RootFlow",
+            "from lightning_app.testing.helpers import EmptyWork",
+            "class MyFlow(_RootFlow):",
+            "    def configure_layout(self):",
+            "        return [{'name': 'test', 'content': this_package_is_not_real()}]",
+            "app = LightningApp(MyFlow(EmptyWork()))",
+        ],
+    ],
+)
+@pytest.mark.skipif(sys.platform != "linux", reason="Causing conflicts on non-linux")
+def test_load_app_from_file_mock_imports(tmpdir, lines):
+    path = copy(sys.path)
+    app_file = os.path.join(tmpdir, "app.py")
+
+    with open(app_file, "w") as f:
+        f.write("\n".join(lines))
+
+    app = CloudRuntime.load_app_from_file(app_file)
+    assert isinstance(app, LightningApp)
+    assert isinstance(app.root.work, EmptyWork)
+
+    # Cleanup PATH to prevent conflict with other tests
+    sys.path = path
+    os.remove(app_file)
+
+
 def test_incompatible_cloud_compute_and_build_config():
     """Test that an exception is raised when a build config has a custom image defined, but the cloud compute is
     the default.

tests/tests_app/utilities/test_proxies.py

@@ -67,6 +67,7 @@ def proxy_setattr():
 
 @pytest.mark.parametrize("parallel", [True, False])
 @pytest.mark.parametrize("cache_calls", [False, True])
+@pytest.mark.skipif(sys.platform == "win32", reason="TODO (@ethanwharris): Fix this on Windows")
 def test_work_runner(parallel, cache_calls):
     """This test validates the `WorkRunner` runs the work.run method and properly populates the `delta_queue`,
     `error_queue` and `readiness_queue`."""