@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 61aa4b2901

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[P/D] Add a shutdown method to the Connector API (#22699)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>

Chauncey committed Sep 8, 2025 at 14:07 UTC 61aa4b2901fa2a1376bcd781da253d83526e3090

10 files changed +52 -12

vllm/distributed/kv_transfer/__init__.py

+4 -3

@@ -2,11 +2,12 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from vllm.distributed.kv_transfer.kv_transfer_state import (

5 - KVConnectorBaseType, ensure_kv_transfer_initialized, get_kv_transfer_group,

6 - has_kv_transfer_group, is_v1_kv_transfer_group)

5 + KVConnectorBaseType, ensure_kv_transfer_initialized,

6 + ensure_kv_transfer_shutdown, get_kv_transfer_group, has_kv_transfer_group,

7 + is_v1_kv_transfer_group)

9 __all__ = [

10 "get_kv_transfer_group", "has_kv_transfer_group",

11 "is_v1_kv_transfer_group", "ensure_kv_transfer_initialized",

11 - "KVConnectorBaseType"

12 + "ensure_kv_transfer_shutdown", "KVConnectorBaseType"

13 ]

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -226,6 +226,14 @@ class KVConnectorBase_V1(ABC):

226 """

227 return None, None

228

229 + def shutdown(self):

230 + """

231 + Shutdown the connector. This is called when the worker process

232 + is shutting down to ensure that all the async operations are

233 + completed and the connector is cleaned up properly.

234 + """

235 + return None

236 +

237 # ==============================

238 # Scheduler-side methods

239 # ==============================

vllm/distributed/kv_transfer/kv_transfer_state.py

@@ -64,3 +64,10 @@ def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:

64 config=vllm_config, role=KVConnectorRole.WORKER)

65 else:

66 raise ValueError("V0 is no longer supported")

67 +

68 +

69 +def ensure_kv_transfer_shutdown() -> None:

70 + global _KV_CONNECTOR_AGENT

71 + if _KV_CONNECTOR_AGENT is not None:

72 + _KV_CONNECTOR_AGENT.shutdown()

73 + _KV_CONNECTOR_AGENT = None

vllm/executor/executor_base.py

+1 -1

@@ -231,7 +231,7 @@ class ExecutorBase(ABC):

231

232 def shutdown(self) -> None:

233 """Shutdown the executor."""

234 - return

234 + self.collective_rpc("shutdown")

235

236 def __del__(self):

237 self.shutdown()

vllm/v1/core/sched/scheduler.py

@@ -1188,6 +1188,8 @@ class Scheduler(SchedulerInterface):

1188 def shutdown(self) -> None:

1189 if self.kv_event_publisher:

1190 self.kv_event_publisher.shutdown()

1191 + if self.connector is not None:

1192 + self.connector.shutdown()

1193

1194 ########################################################################

1195 # KV Connector Related Methods

vllm/v1/executor/multiproc_executor.py

+9 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import multiprocessing

4 -import os

4 import pickle

5 import queue

6 import signal

@@ -507,6 +506,7 @@ class WorkerProc:

506 return cast(list[WorkerProcHandle], ready_proc_handles)

507

508 def shutdown(self):

509 + self.worker.shutdown()

510 self.rpc_broadcast_mq = None

511 self.worker_response_mq = None

512 destroy_model_parallel()

@@ -536,7 +536,7 @@ class WorkerProc:

536 # tuple[Connection, Connection]

537 reader, ready_writer = kwargs.pop("ready_pipe")

538 death_pipe = kwargs.pop("death_pipe", None)

539 -

539 + shutdown_event = threading.Event()

540 # Start death monitoring thread if death_pipe is provided

541 if death_pipe is not None:

542

@@ -548,7 +548,7 @@ class WorkerProc:

548 # Parent process has exited, terminate this worker

549 logger.info("Parent process exited, terminating worker")

550 # Send signal to self to trigger clean shutdown

551 - os.kill(os.getpid(), signal.SIGTERM)

551 + shutdown_event.set()

552 except Exception as e:

553 logger.warning("Death monitoring error: %s", e)

554

@@ -576,7 +576,7 @@ class WorkerProc:

576 ready_writer.close()

577 ready_writer = None

578

579 - worker.worker_busy_loop()

579 + worker.worker_busy_loop(cancel=shutdown_event)

580

581 except Exception:

582 # NOTE: if an Exception arises in busy_loop, we send

@@ -586,6 +586,8 @@ class WorkerProc:

586

587 if ready_writer is not None:

588 logger.exception("WorkerProc failed to start.")

589 + elif shutdown_event.is_set():

590 + logger.info("WorkerProc shutting down.")

591 else:

592 logger.exception("WorkerProc failed.")

593

@@ -637,11 +639,11 @@ class WorkerProc:

639 output = self.async_output_queue.get()

640 self.enqueue_output(output)

641

640 - def worker_busy_loop(self):

642 + def worker_busy_loop(self, cancel: Optional[threading.Event] = None):

643 """Main busy loop for Multiprocessing Workers"""

644 while True:

643 - method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue()

644 -

645 + method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue(

646 + cancel=cancel)

647 try:

648 if isinstance(method, str):

649 func = getattr(self.worker, method)

vllm/v1/worker/gpu_worker.py

@@ -601,6 +601,9 @@ class Worker(WorkerBase):

601 self.model_runner.save_tensorized_model(

602 tensorizer_config=tensorizer_config, )

603

604 + def shutdown(self) -> None:

605 + self.model_runner.ensure_kv_transfer_shutdown()

606 +

607

608 def init_worker_distributed_environment(

609 vllm_config: VllmConfig,

vllm/v1/worker/kv_connector_model_runner_mixin.py

+7 -1

@@ -9,7 +9,8 @@ from typing import Generator # noqa: UP035

9 from typing import TYPE_CHECKING, Optional

11 from vllm.config import VllmConfig

12 -from vllm.distributed.kv_transfer import (get_kv_transfer_group,

12 +from vllm.distributed.kv_transfer import (ensure_kv_transfer_shutdown,

13 + get_kv_transfer_group,

14 has_kv_transfer_group)

15 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase

16 from vllm.forward_context import get_forward_context, set_forward_context

@@ -42,6 +43,11 @@ class KVConnectorModelRunnerMixin:

43 # Do this here to save a collective_rpc.

44 kv_connector.start_load_kv(get_forward_context())

46 + @staticmethod

47 + def ensure_kv_transfer_shutdown() -> None:

48 + if has_kv_transfer_group():

49 + ensure_kv_transfer_shutdown()

50 +

51 @staticmethod

52 def maybe_wait_for_kv_save() -> None:

53 if has_kv_transfer_group():

vllm/v1/worker/tpu_worker.py

@@ -330,6 +330,9 @@ class TPUWorker:

330

331 ensure_kv_transfer_initialized(vllm_config)

332

333 + def shutdown(self) -> None:

334 + self.model_runner.ensure_kv_transfer_shutdown()

335 +

336

337 if USE_TPU_COMMONS:

338 from tpu_commons.worker import TPUWorker as TPUCommonsWorker

vllm/worker/worker_base.py

@@ -129,6 +129,10 @@ class WorkerBase:

129 """Get vocabulary size from model configuration."""

130 return self.model_config.get_vocab_size()

131

132 + def shutdown(self) -> None:

133 + """Clean up resources held by the worker."""

134 + return

135 +

136

137 class DelegateWorkerBase(WorkerBase):

138 """

@@ -519,6 +523,10 @@ class WorkerWrapperBase:

523 from vllm.utils import init_cached_hf_modules

524 init_cached_hf_modules()

525

526 + def shutdown(self) -> None:

527 + if self.worker is not None:

528 + self.worker.shutdown()

529 +

530 def adjust_rank(self, rank_mapping: Dict[int, int]) -> None:

531 """

532 Adjust the rpc_rank based on the given mapping.