App: Fix frontends when using multiprocessing in the cloud (#17324)

Co-authored-by: Ethan Harris <[email protected]>

Author: 2 people

src/lightning/app/CHANGELOG.md

@@ -4,6 +4,14 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+## [2.0.1post] - 2023-04-12
+
+### Fixed
+
+- Fix frontend hosts when running with multi-process in the cloud ([#17324](https://github.com/Lightning-AI/lightning/pull/17324))
+
+
+
 ## [2.0.1] - 2023-03-30
 
 No changes.

src/lightning/app/runners/multiprocess.py

@@ -66,10 +66,17 @@ def dispatch(self, *args: Any, open_ui: bool = True, **kwargs: Any):
 
             _set_frontend_context()
             for frontend in self.app.frontends.values():
-                host = "localhost"
                 port = find_free_network_port()
-                frontend.start_server(host="localhost", port=port)
-                frontend.flow._layout["target"] = f"http://{host}:{port}/{frontend.flow.name}"
+
+                server_host = "0.0.0.0" if in_cloudspace else "localhost"
+                server_target = (
+                    f"https://{port}-{constants.LIGHTNING_CLOUDSPACE_HOST}"
+                    if in_cloudspace
+                    else f"http://localhost:{port}"
+                )
+
+                frontend.start_server(host=server_host, port=port)
+                frontend.flow._layout["target"] = f"{server_target}/{frontend.flow.name}"
 
             _set_flow_context()
 

tests/tests_app/runners/test_multiprocess.py

@@ -5,6 +5,7 @@
 import pytest
 
 from lightning.app import LightningApp, LightningFlow, LightningWork
+from lightning.app.core import constants
 from lightning.app.frontend import StaticWebFrontend, StreamlitFrontend
 from lightning.app.runners import MultiProcessRuntime
 from lightning.app.utilities.component import _get_context
@@ -46,15 +47,34 @@ def run(self):
         self.stop()
 
 
+@pytest.mark.parametrize(
+    "cloudspace_host, port, expected_host, expected_target",
+    [
+        (None, 7000, "localhost", "http://localhost:7000"),
+        ("test.lightning.ai", 7000, "0.0.0.0", "https://7000-test.lightning.ai"),
+    ],
+)
 @mock.patch("lightning.app.runners.multiprocess.find_free_network_port")
-def test_multiprocess_starts_frontend_servers(*_):
+def test_multiprocess_starts_frontend_servers(
+    mock_find_free_network_port, monkeypatch, cloudspace_host, port, expected_host, expected_target
+):
     """Test that the MultiProcessRuntime starts the servers for the frontends in each LightningFlow."""
+
+    monkeypatch.setattr(constants, "LIGHTNING_CLOUDSPACE_HOST", cloudspace_host)
+    mock_find_free_network_port.return_value = port
+
     root = StartFrontendServersTestFlow()
     app = LightningApp(root)
     MultiProcessRuntime(app).dispatch()
 
     app.frontends[root.flow0.name].start_server.assert_called_once()
+    assert app.frontends[root.flow0.name].start_server.call_args.kwargs["host"] == expected_host
+
     app.frontends[root.flow1.name].start_server.assert_called_once()
+    assert app.frontends[root.flow1.name].start_server.call_args.kwargs["host"] == expected_host
+
+    assert app.frontends[root.flow0.name].flow._layout["target"] == f"{expected_target}/{root.flow0.name}"
+    assert app.frontends[root.flow1.name].flow._layout["target"] == f"{expected_target}/{root.flow1.name}"
 
     app.frontends[root.flow0.name].stop_server.assert_called_once()
     app.frontends[root.flow1.name].stop_server.assert_called_once()