From e54f4f58be3bdf13d592c46ca799a99992155931 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Mon, 19 Dec 2022 01:39:21 +0100 Subject: [PATCH] Fix detection of whether app is running in cloud (#16045) --- src/lightning_app/CHANGELOG.md | 3 ++- src/lightning_app/runners/cloud.py | 7 +++++++ src/lightning_app/utilities/cloud.py | 2 +- tests/tests_app/utilities/test_cloud.py | 21 +++++++++++++-------- 4 files changed, 23 insertions(+), 10 deletions(-) diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md index 1ffd4f1efb55f..dcf3a73744d26 100644 --- a/src/lightning_app/CHANGELOG.md +++ b/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 diff --git a/src/lightning_app/runners/cloud.py b/src/lightning_app/runners/cloud.py index 0011c29ce0e6b..265a47919b870 100644 --- a/src/lightning_app/runners/cloud.py +++ b/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 diff --git a/src/lightning_app/utilities/cloud.py b/src/lightning_app/utilities/cloud.py index 20ab6d14827c9..6db634649fdbb 100644 --- a/src/lightning_app/utilities/cloud.py +++ b/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 diff --git a/tests/tests_app/utilities/test_cloud.py b/tests/tests_app/utilities/test_cloud.py index 573ec46106b84..6e93ad1e68d57 100644 --- a/tests/tests_app/utilities/test_cloud.py +++ b/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()