From e82672414edce344ffa22ee2bfccf0c415c92f2f Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Mon, 24 Apr 2023 12:27:35 +0100 Subject: [PATCH] App: Fix AppState, streamlit example (#17452) --- .azure/app-cloud-e2e.yml | 7 +++---- requirements/app/base.txt | 2 +- src/lightning/app/utilities/state.py | 27 ++++++++++++++++++++----- tests/tests_app/utilities/test_state.py | 27 ++++++++++++++++++++----- 4 files changed, 48 insertions(+), 15 deletions(-) diff --git a/.azure/app-cloud-e2e.yml b/.azure/app-cloud-e2e.yml index 100c8d64fe..260c7fb4b2 100644 --- a/.azure/app-cloud-e2e.yml +++ b/.azure/app-cloud-e2e.yml @@ -70,10 +70,9 @@ jobs: 'App: boring_app': name: "boring_app" dir: "public" - # TODO: RESOLVE ME ASAP - # 'App: template_streamlit_ui': - # name: "template_streamlit_ui" - # dir: "public" + 'App: template_streamlit_ui': + name: "template_streamlit_ui" + dir: "public" 'App: template_react_ui': name: "template_react_ui" dir: "public" diff --git a/requirements/app/base.txt b/requirements/app/base.txt index 090c0915cf..f0d24ce89c 100644 --- a/requirements/app/base.txt +++ b/requirements/app/base.txt @@ -1,4 +1,4 @@ -lightning-cloud >=0.5.33 +lightning-cloud >=0.5.34 packaging typing-extensions >=4.0.0, <=4.4.0 deepdiff >=5.7.0, <6.2.4 diff --git a/src/lightning/app/utilities/state.py b/src/lightning/app/utilities/state.py index 17b740476f..4fbe58600b 100644 --- a/src/lightning/app/utilities/state.py +++ b/src/lightning/app/utilities/state.py @@ -26,7 +26,7 @@ from requests.exceptions import ConnectionError from lightning.app.core.constants import APP_SERVER_HOST, APP_SERVER_PORT from lightning.app.storage.drive import _maybe_create_drive from lightning.app.utilities.app_helpers import AppStatePlugin, BaseStatePlugin, Logger -from lightning.app.utilities.network import _configure_session +from lightning.app.utilities.network import _configure_session, LightningClient logger = Logger(__name__) @@ -50,6 +50,7 @@ def headers_for(context: Dict[str, str]) -> Dict[str, str]: class AppState: _APP_PRIVATE_KEYS: Tuple[str, ...] = ( + "_use_localhost", "_host", "_session_id", "_state", @@ -93,10 +94,9 @@ class AppState: on this AppState, this affiliation will be used to reduce the scope of the given state. plugin: A plugin to handle authorization. """ - use_localhost = "LIGHTNING_APP_STATE_URL" not in os.environ - self._host = host or APP_SERVER_HOST - self._port = port or (APP_SERVER_PORT if use_localhost else None) - self._url = f"{self._host}:{self._port}" if use_localhost else self._host + self._use_localhost = "LIGHTNING_APP_STATE_URL" not in os.environ + self._host = host or ("http://127.0.0.1" if self._use_localhost else None) + self._port = port or (APP_SERVER_PORT if self._use_localhost else None) self._last_state = last_state self._state = state self._session_id = "1234" @@ -105,6 +105,23 @@ class AppState: self._attach_plugin(plugin) self._session = self._configure_session() + @property + def _url(self) -> str: + if self._host is None: + app_ip = "" + + if "LIGHTNING_CLOUD_PROJECT_ID" in os.environ and "LIGHTNING_CLOUD_APP_ID" in os.environ: + client = LightningClient() + app_instance = client.lightningapp_instance_service_get_lightningapp_instance( + os.environ.get("LIGHTNING_CLOUD_PROJECT_ID"), + os.environ.get("LIGHTNING_CLOUD_APP_ID"), + ) + app_ip = app_instance.status.ip_address + + # TODO: Don't hard code port 8080 here + self._host = f"http://{app_ip}:8080" if app_ip else APP_SERVER_HOST + return f"{self._host}:{self._port}" if self._use_localhost else self._host + def _attach_plugin(self, plugin: Optional[BaseStatePlugin]) -> None: if plugin is not None: plugin = plugin diff --git a/tests/tests_app/utilities/test_state.py b/tests/tests_app/utilities/test_state.py index 3abe93d437..d4bd688327 100644 --- a/tests/tests_app/utilities/test_state.py +++ b/tests/tests_app/utilities/test_state.py @@ -4,6 +4,7 @@ from unittest import mock import pytest import requests +from lightning_cloud.openapi import Externalv1LightningappInstance, V1LightningappInstanceStatus import lightning.app from lightning.app import LightningApp, LightningFlow, LightningWork @@ -266,13 +267,29 @@ def test_get_send_request(monkeypatch): state.w.counter = 1 -@mock.patch("lightning.app.utilities.state.APP_SERVER_HOST", "https://lightning-cloud.com") -@mock.patch.dict(os.environ, {"LIGHTNING_APP_STATE_URL": "https://lightning-cloud.com"}) -def test_app_state_with_env_var(**__): +@mock.patch.dict( + os.environ, + { + "LIGHTNING_APP_STATE_URL": "https://lightning-cloud.com", + "LIGHTNING_CLOUD_PROJECT_ID": "test-project-id", + "LIGHTNING_CLOUD_APP_ID": "test-app-id", + }, +) +@mock.patch("lightning.app.utilities.state.LightningClient") +def test_app_state_with_env_var(mock_client): + mock_client().lightningapp_instance_service_get_lightningapp_instance.return_value = Externalv1LightningappInstance( + status=V1LightningappInstanceStatus(ip_address="test-ip"), + ) state = AppState() - assert state._host == "https://lightning-cloud.com" + url = state._url + + mock_client().lightningapp_instance_service_get_lightningapp_instance.assert_called_once_with( + "test-project-id", + "test-app-id", + ) + + assert url == "http://test-ip:8080" assert not state._port - assert state._url == "https://lightning-cloud.com" @mock.patch.dict(os.environ, {})