Feat: multiple endpoints using a list of LitServer

src/litserve/server.py

@@ -27,7 +27,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from queue import Empty
-from typing import Callable, Dict, Optional, Sequence, Tuple, Union
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import uvicorn
 from fastapi import Depends, FastAPI, HTTPException, Request, Response
@@ -472,3 +472,68 @@ def setup_auth(self):
         if LIT_SERVER_API_KEY:
             return api_key_auth
         return no_auth
+
+
+def run_all(
+    servers: List[LitServer],
+    port: Union[str, int] = 8000,
+    num_api_servers: Optional[int] = None,
+    log_level: str = "info",
+    generate_client_file: bool = True,
+    api_server_worker_type: Optional[str] = None,
+    **kwargs,
+):
+    """Run multiple LitServers on the same port."""
+    if not servers:
+        raise ValueError("No servers provided to run_all")
+
+    if any(not isinstance(server, LitServer) for server in servers):
+        raise ValueError("All elements in the servers list must be instances of LitServer")
+
+    if generate_client_file:
+        for server in servers:
+            server.generate_client_file()
+
+    port_msg = f"port must be a value from 1024 to 65535 but got {port}"
+    try:
+        port = int(port)
+    except ValueError:
+        raise ValueError(port_msg)
+    if not (1024 <= port <= 65535):
+        raise ValueError(port_msg)
+
+    if num_api_servers is None:
+        num_api_servers = sum(len(server.workers) for server in servers)
+    if num_api_servers < 1:
+        raise ValueError("num_api_servers must be greater than 0")
+
+    if sys.platform == "win32":
+        print("Windows does not support forking. Using threads api_server_worker_type will be set to 'thread'")
+        api_server_worker_type = "thread"
+    elif api_server_worker_type is None:
+        api_server_worker_type = "process"
+
+    app = servers[0].app
+    config = uvicorn.Config(app=app, host="0.0.0.0", port=port, log_level=log_level, **kwargs)
+    sockets = [config.bind_socket()]
+
+    managers, all_workers = [], []
+    try:
+        all_servers = []
+        for server in servers:
+            manager, litserve_workers = server.launch_inference_worker(num_api_servers)
+            managers.append(manager)
+            all_workers.extend(litserve_workers)
+
+            _servers = server._start_server(port, num_api_servers, log_level, sockets, api_server_worker_type, **kwargs)
+            all_servers.extend(_servers)
+        print(f"Swagger UI is available at http://0.0.0.0:{port}/docs")
+        for s in all_servers:
+            s.join()
+    finally:
+        print("Shutting down LitServe")
+        for w in all_workers:
+            w.terminate()
+            w.join()
+        for manager in managers:
+            manager.shutdown()

tests/test_lit_server_multi_endpoints.py

@@ -0,0 +1,28 @@
+
+
+from litserve.server import LitServer
+
+
+async def test_lit_server_with_multi_endpoints(simple_litapi):
+    server1 = LitServer(simple_litapi, api_path="/predict-1", timeout=10)
+    server2 = LitServer(simple_litapi, api_path="/predict-2", timeout=10)
+    servers = [server1, server2]
+    # TODO: update test to use run_all
+
+    # Run the servers in a separate thread
+    port = 8000
+    # server_thread = threading.Thread(
+    #     target=run_all, args=(servers,), kwargs={"port": port, "num_api_servers": 2, "log_level": "debug"}
+    # )
+    # server_thread.start()
+
+    # async with AsyncClient(base_url=f"http://localhost:{port}") as client:
+    #     # Test server1 endpoint
+    #     response1 = await client.post("/predict-1", json={"input": 1})
+    #     assert response1.status_code == 200
+    #     assert response1.json() == {"output": 1}
+
+    #     # Test server2 endpoint
+    #     response2 = await client.post("/predict-2", json={"input": 2})
+    #     assert response2.status_code == 200
+    #     assert response2.json() == {"output": 2}