Skip to content

Commit 4fcffda

Browse files
[BugFix][WeightTransfer] Lazy-load weight transfer engines to avoid eager ray import (vllm-project#10816)
### What this PR does / why we need it? PR vllm-project#10592 (`[Feature] WeightTransfer: Add NPUIPCWeightTransferEngine backend for Ascend NPU`) passed all gates but broke the smoke test with: ``` ModuleNotFoundError: No module named 'ray' ``` **Root cause** `vllm_ascend/patch/platform/patch_weight_transfer_engine.py` runs during *global plugin patching* — extremely early in startup, triggered even by `AsyncEngineArgs.add_cli_args(parser)` → `load_general_plugins()`. This happens long before any weight-transfer backend is selected. The patch eagerly imported the engine classes at module top level: ```python from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import ( NPUIPCWeightTransferEngine, ) ``` `npu_ipc_engine` transitively imports `vllm.distributed.weight_transfer.ipc_engine`, which does `import ray` at module top level. `ray` is an optional dependency, so when it is not installed the entire `vllm_ascend` plugin load aborts and **every** `vllm serve` invocation crashes — even workloads that never use weight transfer. Import chain that crashed: ``` vllm_ascend/__init__.py:_ensure_global_patch() → utils.py:adapt_patch(is_global_patch=True) → patch/platform/__init__.py → patch_weight_transfer_engine.py → distributed/weight_transfer/npu_ipc_engine.py → vllm/distributed/weight_transfer/ipc_engine.py → import ray # ModuleNotFoundError ``` **Fix** Replace the eager top-level imports with lazy loader functions. `WeightTransferEngineFactory._registry[backend]` is a zero-arg callable that returns the engine class and is only invoked from `create_engine()` when the backend is actually requested. Deferring the import into those callables matches the factory's existing lazy-loading contract, so the `ray`-importing module is only loaded when the `ipc` backend is selected (at which point ray is genuinely needed): ```python def _load_npu_ipc_engine() -> "type[WeightTransferEngine]": from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import ( NPUIPCWeightTransferEngine, ) return NPUIPCWeightTransferEngine WeightTransferEngineFactory._registry["ipc"] = _load_npu_ipc_engine ``` Only the `ipc` registration is made lazy, since the IPC engine is the one that transitively imports `ray`. The `nccl` → HCCL registration keeps its eager import, as `hccl_engine` has no optional-dependency import problem. ### Does this PR introduce _any_ user-facing change? No. The `nccl`/`ipc` backend resolution behaves identically; only the timing of the underlying module import changes (deferred until the backend is actually used). ### How was this patch tested? - `python -m py_compile` and `ruff check` / `ruff format --check` pass on the changed file. - The fix removes the only eager import of the IPC engine on the global-patch path, so plugin loading no longer requires `ray`. The other registration site (`distributed/weight_transfer/__init__.py:register_engine`) already uses lazy string-based registration and was unaffected. - vLLM version: v0.22.1 - vLLM main: vllm-project/vllm@967c5c3 --------- Signed-off-by: learning-sketch <learning-sketch@users.noreply.github.com> Co-authored-by: learning-sketch <learning-sketch@users.noreply.github.com>
1 parent f68273d commit 4fcffda

2 files changed

Lines changed: 38 additions & 4 deletions

File tree

vllm_ascend/patch/__init__.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,29 @@
410410
# Future Plan:
411411
# Remove this patch when upstream vLLM relaxes the Literal type to str or
412412
# provides an extension point for out-of-tree weight transfer backends.
413+
# 2. `vllm.distributed.weight_transfer.factory.WeightTransferEngineFactory._registry["ipc"]`
414+
# Why:
415+
# The "ipc" backend must resolve to NPUIPCWeightTransferEngine on Ascend NPU.
416+
# However, this patch runs during global plugin patching - extremely early in
417+
# startup, before any weight transfer backend is selected. Importing the IPC
418+
# engine eagerly pulls in vllm.distributed.weight_transfer.ipc_engine, which
419+
# does `import ray` at module top level. Since ray is an optional dependency,
420+
# its absence aborts the whole vllm_ascend plugin load and crashes every
421+
# `vllm serve` invocation - even workloads that never use weight transfer.
422+
# How:
423+
# Register a lazy loader function (instead of an eager import) that imports
424+
# and returns NPUIPCWeightTransferEngine only when create_engine() is invoked
425+
# for the "ipc" backend. This matches the factory's zero-arg-callable
426+
# lazy-loading contract, so the ray-importing module is loaded only when ipc
427+
# is actually requested. (HCCL keeps its eager import - it never imports ray.)
428+
# Related PR (if no, explain why):
429+
# No. The eager `import ray` lives in upstream vLLM's ipc_engine module; the
430+
# lazy loader is a local workaround until upstream defers that import.
431+
# Future Plan:
432+
# Remove this workaround once upstream vLLM stops importing ray at module top
433+
# level in vllm.distributed.weight_transfer.ipc_engine (e.g. defers it into
434+
# the code path that actually needs ray), so importing the IPC engine no
435+
# longer requires the optional ray dependency.
413436
#
414437
# ** 15. File: platform/patch_kv_cache_coordinator.py**
415438
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

vllm_ascend/patch/platform/patch_weight_transfer_engine.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,25 @@
4949
# Remove this patch when upstream vllm relaxes the Literal type to str
5050
# or provides an extension point for out-of-tree backends.
5151

52+
from typing import TYPE_CHECKING
53+
5254
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
5355

5456
from vllm_ascend.distributed.weight_transfer.hccl_engine import (
5557
HCCLWeightTransferEngine,
5658
)
57-
from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import (
58-
NPUIPCWeightTransferEngine,
59-
)
59+
60+
if TYPE_CHECKING:
61+
from vllm.distributed.weight_transfer.base import WeightTransferEngine
62+
63+
64+
def _load_npu_ipc_engine() -> "type[WeightTransferEngine]":
65+
from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import (
66+
NPUIPCWeightTransferEngine,
67+
)
68+
69+
return NPUIPCWeightTransferEngine
70+
6071

6172
WeightTransferEngineFactory._registry["nccl"] = lambda: HCCLWeightTransferEngine
62-
WeightTransferEngineFactory._registry["ipc"] = lambda: NPUIPCWeightTransferEngine
73+
WeightTransferEngineFactory._registry["ipc"] = _load_npu_ipc_engine

0 commit comments

Comments
 (0)