From 6c0aff88ed518203d8e82f775f8d797200c1d93c Mon Sep 17 00:00:00 2001 From: bhimrazy Date: Tue, 18 Feb 2025 11:34:31 +0545 Subject: [PATCH] fix: ensure proper cleanup of processes in wrap_litserve_start --- src/litserve/utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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):