App: Fix AppState, streamlit example (#17452)

(cherry picked from commit e826724)

Author: ethanwharris

.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"

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

src/lightning/app/utilities/state.py

@@ -26,7 +26,7 @@
 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 @@ def __init__(
                 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 @@ def __init__(
         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

tests/tests_app/utilities/test_state.py

@@ -4,6 +4,7 @@
 
 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, {})