[diffusion] chore: fix unclean shutdown and resource leaks (#18477)

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():

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]],

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

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.")

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."""