diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md index 73a6ee75db065..8e941f5459cfe 100644 --- a/src/lightning_app/CHANGELOG.md +++ b/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 diff --git a/src/lightning_app/utilities/app_helpers.py b/src/lightning_app/utilities/app_helpers.py index bc3d092b280dd..d9efeb6862ba7 100644 --- a/src/lightning_app/utilities/app_helpers.py +++ b/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())