Fix detection of whether app is running in cloud

src/lightning_app/CHANGELOG.md

@@ -13,7 +13,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
--
+
+- The utility `lightning.app.utilities.cloud.is_running_in_cloud` now returns `True` during loading of the app locally when running with `--cloud` ([#16045](https://github.com/Lightning-AI/lightning/pull/16045))
 
 
 ### Deprecated

src/lightning_app/components/multi_node/base.py

@@ -56,12 +56,12 @@ def run(
         """
         super().__init__()
         if num_nodes > 1 and not is_running_in_cloud():
-            num_nodes = 1
             warnings.warn(
                 f"You set {type(self).__name__}(num_nodes={num_nodes}, ...)` but this app is running locally."
                 " We assume you are debugging and will ignore the `num_nodes` argument."
                 " To run on multiple nodes in the cloud, launch your app with `--cloud`."
             )
+            num_nodes = 1
         self.ws = structures.List(
             *[
                 work_cls(

src/lightning_app/runners/cloud.py

@@ -1,5 +1,6 @@
 import fnmatch
 import json
+import os
 import random
 import string
 import sys
@@ -553,6 +554,10 @@ def _project_has_sufficient_credits(self, project: V1Membership, app: Optional[L
     @classmethod
     def load_app_from_file(cls, filepath: str) -> "LightningApp":
         """Load a LightningApp from a file, mocking the imports."""
+
+        # Pretend we are running in the cloud when loading the app locally
+        os.environ["LAI_RUNNING_IN_CLOUD"] = "1"
+
         try:
             app = load_app_from_file(filepath, raise_exception=True, mock_imports=True)
         except FileNotFoundError as e:
@@ -563,6 +568,8 @@ def load_app_from_file(cls, filepath: str) -> "LightningApp":
             # Create a generic app.
             logger.info("Could not load the app locally. Starting the app directly on the cloud.")
             app = LightningApp(EmptyFlow())
+        finally:
+            del os.environ["LAI_RUNNING_IN_CLOUD"]
         return app
 
     @staticmethod

src/lightning_app/utilities/cloud.py

@@ -39,4 +39,4 @@ def _sigterm_flow_handler(*_, app: "lightning_app.LightningApp"):
 
 def is_running_in_cloud() -> bool:
     """Returns True if the Lightning App is running in the cloud."""
-    return "LIGHTNING_APP_STATE_URL" in os.environ
+    return bool(int(os.environ.get("LAI_RUNNING_IN_CLOUD", "0"))) or "LIGHTNING_APP_STATE_URL" in os.environ

tests/tests_app/components/multi_node/test_base.py

@@ -13,7 +13,7 @@ class Work(LightningWork):
         def run(self):
             pass
 
-    with pytest.warns(UserWarning, match=escape("You set MultiNode(num_nodes=1, ...)` but ")):
+    with pytest.warns(UserWarning, match=escape("You set MultiNode(num_nodes=2, ...)` but ")):
         MultiNode(Work, num_nodes=2, cloud_compute=CloudCompute("gpu"))
 
     with no_warning_call(UserWarning, match=escape("You set MultiNode(num_nodes=1, ...)` but ")):

tests/tests_app/utilities/test_cloud.py

@@ -4,13 +4,18 @@
 from lightning_app.utilities.cloud import is_running_in_cloud
 
 
-@mock.patch.dict(os.environ, clear=True)
-def test_is_running_locally():
-    """We can determine if Lightning is running locally."""
-    assert not is_running_in_cloud()
-
-
-@mock.patch.dict(os.environ, {"LIGHTNING_APP_STATE_URL": "127.0.0.1"})
 def test_is_running_cloud():
     """We can determine if Lightning is running in the cloud."""
-    assert is_running_in_cloud()
+    with mock.patch.dict(os.environ, {}, clear=True):
+        assert not is_running_in_cloud()
+
+    with mock.patch.dict(os.environ, {"LAI_RUNNING_IN_CLOUD": "0"}, clear=True):
+        assert not is_running_in_cloud()
+
+    # in the cloud, LIGHTNING_APP_STATE_URL is defined
+    with mock.patch.dict(os.environ, {"LIGHTNING_APP_STATE_URL": "defined"}, clear=True):
+        assert is_running_in_cloud()
+
+    # LAI_RUNNING_IN_CLOUD is used to fake the value of `is_running_in_cloud` when loading the app for --cloud
+    with mock.patch.dict(os.environ, {"LAI_RUNNING_IN_CLOUD": "1"}):
+        assert is_running_in_cloud()