diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 3c397f486..810e030f0 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, + ShutdownReq, UnmergeLoraWeightsReq, format_lora_message, ) @@ -483,8 +484,13 @@ class DiffGenerator: If in local mode, it also shuts down the scheduler server. """ # sends the shutdown command to the server + if self.local_scheduler_process and self.owns_scheduler_client: + try: + sync_scheduler_client.forward(ShutdownReq()) + except Exception: + pass + if self.local_scheduler_process: - logger.info("Waiting for local worker processes to terminate...") for process in self.local_scheduler_process: process.join(timeout=10) if process.is_alive(): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index a52643c55..337a04fbf 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -47,6 +47,11 @@ class ListLorasReq: pass +@dataclasses.dataclass +class ShutdownReq: + pass + + def format_lora_message( lora_nickname: Union[str, List[str]], target: Union[str, List[str]], diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index 272c4ae16..318ed6042 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -182,6 +182,8 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): else: launch_http_server_only(server_args) + return processes + def launch_http_server_only(server_args): # set for endpoints to access global_server_args diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 6ce092d85..869068e9b 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1,6 +1,7 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 +import gc import multiprocessing as mp import os import time @@ -404,4 +405,12 @@ def run_scheduler_process( print(OOM_MSG) raise finally: + # Clean up resources to speed up shutdown + if "scheduler" in locals(): + del scheduler + gc.collect() + if torch.cuda.is_initialized(): + torch.cuda.empty_cache() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() logger.info(f"Worker {rank}: Shutdown complete.") diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 1944ad789..041f6b7fe 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -15,6 +15,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, + ShutdownReq, UnmergeLoraWeightsReq, _parse_size, save_image_to_path, @@ -87,6 +88,7 @@ class Scheduler: Req: self._handle_generation, List[Req]: self._handle_generation, ListLorasReq: self._handle_list_loras, + ShutdownReq: self._handle_shutdown, } # FIFO, new reqs are appended @@ -123,6 +125,10 @@ class Scheduler: def _handle_list_loras(self, _reqs: List[Any]) -> OutputBatch: return self.worker.list_loras() + def _handle_shutdown(self, _reqs: List[Any]) -> OutputBatch: + self._running = False + return OutputBatch() + def _handle_generation(self, reqs: List[Req]): warmup_reqs = [req for req in reqs if req.is_warmup] if warmup_reqs: @@ -381,10 +387,9 @@ class Scheduler: logger.error(f"ZMQ error sending reply: {e}") continue - logger.info("Scheduler event loop terminated.") if self.receiver is not None: self.receiver.close() - self.context.term() + self.context.destroy(linger=0) def _broadcast_task(self, payload: dict[str, Any]) -> None: """Broadcast a task to all slave worker processes."""