diff --git a/src/litserve/utils.py b/src/litserve/utils.py index 2e05ee280..464ab37b8 100644 --- a/src/litserve/utils.py +++ b/src/litserve/utils.py @@ -68,10 +68,13 @@ def wrap_litserve_start(server: "LitServer"): server.lit_spec.response_queue_id = 0 manager, processes = server.launch_inference_worker(num_uvicorn_servers=1) server._prepare_app_run(server.app) - yield server - for p in processes: - p.terminate() - manager.shutdown() + try: + yield server + finally: + for p in processes: + p.terminate() + p.join() + manager.shutdown() async def call_after_stream(streamer: AsyncIterator, callback, *args, **kwargs):