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 2, 14:45 UTC
[EC Connector] Add shutdown API to EC Connector. (#42423)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
omerpaz95 committed
May 28, 2026 at 15:28 UTC
811d80519552e2964c2e0e808b66d9a24931bf1e
5 files changed
+25
-1
vllm/distributed/ec_transfer/__init__.py
+2
@@ -3,6 +3,7 @@
3
4
from vllm.distributed.ec_transfer.ec_transfer_state import (
5
ensure_ec_transfer_initialized,
6
+ ensure_ec_transfer_shutdown,
7
get_ec_transfer,
8
has_ec_transfer,
9
)
@@ -10,5 +11,6 @@ from vllm.distributed.ec_transfer.ec_transfer_state import (
11
__all__ = [
12
"get_ec_transfer",
13
"ensure_ec_transfer_initialized",
14
+ "ensure_ec_transfer_shutdown",
15
"has_ec_transfer",
16
]
vllm/distributed/ec_transfer/ec_connector/base.py
+8
@@ -79,6 +79,14 @@ class ECConnectorBase(ABC):
79
def is_consumer(self) -> bool:
80
return self._is_consumer
81
82
+ def shutdown(self) -> None:
83
+ """
84
+ Shutdown the connector. This is called when the process
85
+ is shutting down to ensure that all the async operations are
86
+ completed and the connector is cleaned up properly.
87
+ """
88
+ return None
89
+
90
# ==============================
91
# Worker-side methods
92
# ==============================
vllm/distributed/ec_transfer/ec_transfer_state.py
+7
@@ -40,3 +40,10 @@ def ensure_ec_transfer_initialized(vllm_config: "VllmConfig") -> None:
40
_EC_CONNECTOR_AGENT = ECConnectorFactory.create_connector(
41
config=vllm_config, role=ECConnectorRole.WORKER
42
)
43
+
44
+
45
+def ensure_ec_transfer_shutdown() -> None:
46
+ global _EC_CONNECTOR_AGENT
47
+ if _EC_CONNECTOR_AGENT is not None:
48
+ _EC_CONNECTOR_AGENT.shutdown()
49
+ _EC_CONNECTOR_AGENT = None
vllm/v1/core/sched/scheduler.py
+2
@@ -2041,6 +2041,8 @@ class Scheduler(SchedulerInterface):
2041
self.kv_event_publisher.shutdown()
2042
if self.connector is not None:
2043
self.connector.shutdown()
2044
+ if self.ec_connector is not None:
2045
+ self.ec_connector.shutdown()
2046
2047
########################################################################
2048
# KV Connector Related Methods
vllm/v1/worker/gpu_worker.py
+6
-1
@@ -23,7 +23,10 @@ from vllm.distributed import (
23
init_distributed_environment,
24
set_custom_all_reduce,
25
)
26
-from vllm.distributed.ec_transfer import ensure_ec_transfer_initialized
26
+from vllm.distributed.ec_transfer import (
27
+ ensure_ec_transfer_initialized,
28
+ ensure_ec_transfer_shutdown,
29
+)
30
from vllm.distributed.eplb.eplb_utils import override_envs_for_eplb
31
from vllm.distributed.kv_transfer import (
32
ensure_kv_transfer_initialized,
@@ -1101,6 +1104,8 @@ class Worker(WorkerBase):
1104
# has_kv_transfer_group can be None during interpreter shutdown.
1105
if ensure_kv_transfer_shutdown is not None:
1106
ensure_kv_transfer_shutdown()
1107
+ if ensure_ec_transfer_shutdown is not None:
1108
+ ensure_ec_transfer_shutdown()
1109
if self.profiler is not None:
1110
self.profiler.shutdown()
1111