[App] Fix support for streamlit > 1.14 (#16139)

src/lightning_app/CHANGELOG.md

@@ -66,6 +66,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed the endpoint info tab not showing up in `AutoScaler` UI ([#16128](https://github.com/Lightning-AI/lightning/pull/16128))
 
 
+- Fixed an issue where an exception would be raised in the logs when using a recent version of streamlit ([#16139](https://github.com/Lightning-AI/lightning/pull/16139))
+
+
 ## [1.8.5] - 2022-12-15
 
 ### Added

src/lightning_app/utilities/app_helpers.py

@@ -184,11 +184,22 @@ def render_non_authorized(self):
 
 
 def target_fn():
-    from streamlit.server.server import Server
+    try:
+        # streamlit >= 1.14.0
+        from streamlit import runtime
+
+        get_instance = runtime.get_instance
+        exists = runtime.exists()
+    except ImportError:
+        # Older versions
+        from streamlit.server.server import Server
+
+        get_instance = Server.get_current
+        exists = bool(Server._singleton)
 
     async def update_fn():
-        server = Server.get_current()
-        sessions = list(server._session_info_by_id.values())
+        runtime_instance = get_instance()
+        sessions = list(runtime_instance._session_info_by_id.values())
         url = (
             "localhost:8080"
             if "LIGHTNING_APP_STATE_URL" in os.environ
@@ -198,15 +209,20 @@ async def update_fn():
         last_updated = time.time()
         async with websockets.connect(ws_url) as websocket:
             while True:
-                _ = await websocket.recv()
-                while (time.time() - last_updated) < 1:
-                    time.sleep(0.1)
-                for session in sessions:
-                    session = session.session
-                    session.request_rerun(session._client_state)
-                last_updated = time.time()
-
-    if Server._singleton:
+                try:
+                    _ = await websocket.recv()
+
+                    while (time.time() - last_updated) < 1:
+                        time.sleep(0.1)
+                    for session in sessions:
+                        session = session.session
+                        session.request_rerun(session._client_state)
+                    last_updated = time.time()
+                except websockets.exceptions.ConnectionClosedOK:
+                    # The websocket is not enabled
+                    break
+
+    if exists:
         asyncio.run(update_fn())