Fix detection of whether app is running in cloud (#16045)

Author: awaelchli

src/lightning_app/CHANGELOG.md

@@ -21,7 +21,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/runners/cloud.py

@@ -1,5 +1,6 @@
 import fnmatch
 import json
+import os
 import random
 import re
 import string
@@ -589,6 +590,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:
@@ -599,6 +604,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/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()