diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 30c5dfe6e..769548997 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -603,7 +603,7 @@ class DiffGenerator: # sends the shutdown command to the server if self.local_scheduler_process and self.owns_scheduler_client: try: - sync_scheduler_client.forward(ShutdownReq()) + sync_scheduler_client.forward(ShutdownReq(), timeout_ms=5000) except Exception: pass @@ -615,12 +615,46 @@ class DiffGenerator: f"Local worker {process.name} did not terminate gracefully, forcing." ) process.terminate() + process.join(timeout=1) + if process.is_alive(): + process.kill() + process.join(timeout=1) self.local_scheduler_process = None if self.owns_scheduler_client: sync_scheduler_client.close() self.owns_scheduler_client = False + def _force_shutdown_local_processes(self) -> None: + local_scheduler_process = getattr(self, "local_scheduler_process", None) + log = globals().get("logger") + if local_scheduler_process: + for process in local_scheduler_process: + if process.is_alive(): + if log is not None: + log.warning( + f"Local worker {process.name} did not terminate gracefully, forcing." + ) + process.terminate() + for process in local_scheduler_process: + process.join(timeout=1) + if process.is_alive(): + if log is not None: + log.warning( + f"Local worker {process.name} did not terminate after terminate(), killing." + ) + process.kill() + process.join(timeout=1) + self.local_scheduler_process = None + + if getattr(self, "owns_scheduler_client", False): + try: + client = globals().get("sync_scheduler_client") + if client is not None: + client.close() + finally: + self.owns_scheduler_client = False + def __enter__(self): return self @@ -630,15 +664,18 @@ class DiffGenerator: def __del__(self): owns_scheduler_client = bool(getattr(self, "owns_scheduler_client", False)) local_scheduler_process = getattr(self, "local_scheduler_process", None) + log = globals().get("logger") if owns_scheduler_client: - logger.warning( - "Generator was garbage collected without being shut down. " - "Attempting to shut down the local server and client." - ) - self.shutdown() + if log is not None: + log.warning( + "Generator was garbage collected without being shut down. " + "Forcing local server and client cleanup." + ) + self._force_shutdown_local_processes() elif local_scheduler_process: - logger.warning( - "Generator was garbage collected without being shut down. " - "Attempting to shut down the local server." - ) - self.shutdown() + if log is not None: + log.warning( + "Generator was garbage collected without being shut down. " + "Forcing local server cleanup." + ) + self._force_shutdown_local_processes() diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index e213f3bb9..2455d5767 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -78,8 +78,12 @@ class SchedulerClient: f"SchedulerClient connected to backend scheduler at {scheduler_endpoint}" ) - def forward(self, batch: Any) -> Any: + def forward(self, batch: Any, timeout_ms: int | None = None) -> Any: """Sends a batch or request to the scheduler and waits for the response.""" + previous_timeout_ms = None + if timeout_ms is not None: + previous_timeout_ms = self.scheduler_socket.getsockopt(zmq.RCVTIMEO) + self.scheduler_socket.setsockopt(zmq.RCVTIMEO, timeout_ms) try: self.scheduler_socket.send_pyobj(batch) output_batch = self.scheduler_socket.recv_pyobj() @@ -88,6 +92,9 @@ class SchedulerClient: except zmq.error.Again: logger.error("Timeout waiting for response from scheduler.") raise TimeoutError("Scheduler did not respond in time.") + finally: + if previous_timeout_ms is not None and self.scheduler_socket is not None: + self.scheduler_socket.setsockopt(zmq.RCVTIMEO, previous_timeout_ms) def ping(self) -> bool: """ diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py new file mode 100644 index 000000000..e43f550dd --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py @@ -0,0 +1,99 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from sglang.multimodal_gen.runtime.entrypoints import diffusion_generator as dg +from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator +from sglang.multimodal_gen.runtime.entrypoints.utils import ShutdownReq + + +class _FakeProcess: + name = "fake-worker" + + def __init__(self): + self.alive = True + self.join_timeouts = [] + self.terminated = False + self.killed = False + + def join(self, timeout=None): + self.join_timeouts.append(timeout) + + def is_alive(self): + return self.alive + + def terminate(self): + self.terminated = True + + def kill(self): + self.killed = True + self.alive = False + + +class TestDiffGeneratorShutdown(unittest.TestCase): + def test_shutdown_uses_bounded_scheduler_timeout_and_forces_worker(self): + generator = object.__new__(DiffGenerator) + process = _FakeProcess() + generator.local_scheduler_process = [process] + generator.owns_scheduler_client = True + + client = SimpleNamespace( + forward=Mock(side_effect=TimeoutError("blocked")), + close=Mock(), + ) + + with patch.object(dg, "sync_scheduler_client", client): + generator.shutdown() + + request = client.forward.call_args.args[0] + self.assertIsInstance(request, ShutdownReq) + self.assertEqual(client.forward.call_args.kwargs, {"timeout_ms": 5000}) + self.assertTrue(process.terminated) + self.assertTrue(process.killed) + self.assertAlmostEqual(process.join_timeouts[0], 10, delta=0.1) + self.assertEqual(process.join_timeouts[1:], [1, 1]) + self.assertIsNone(generator.local_scheduler_process) + self.assertFalse(generator.owns_scheduler_client) + client.close.assert_called_once_with() + + def test_del_forces_cleanup_without_scheduler_request(self): + generator = object.__new__(DiffGenerator) + process = _FakeProcess() + generator.local_scheduler_process = [process] + generator.owns_scheduler_client = True + + client = SimpleNamespace( + forward=Mock(), + close=Mock(), + ) + + with patch.object(dg, "sync_scheduler_client", client): + generator.__del__() + + client.forward.assert_not_called() + client.close.assert_called_once_with() + self.assertTrue(process.terminated) + self.assertTrue(process.killed) + self.assertIsNone(generator.local_scheduler_process) + self.assertFalse(generator.owns_scheduler_client) + + def test_del_tolerates_missing_shutdown_globals(self): + generator = object.__new__(DiffGenerator) + process = _FakeProcess() + generator.local_scheduler_process = [process] + generator.owns_scheduler_client = True + + with ( + patch.object(dg, "logger", None), + patch.object(dg, "sync_scheduler_client", None), + ): + generator.__del__() + + self.assertTrue(process.terminated) + self.assertTrue(process.killed) + self.assertIsNone(generator.local_scheduler_process) + self.assertFalse(generator.owns_scheduler_client) + + +if __name__ == "__main__": + unittest.main()