Commit 4fcffda
authored
[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
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
410 | 410 | | |
411 | 411 | | |
412 | 412 | | |
| 413 | + | |
| 414 | + | |
| 415 | + | |
| 416 | + | |
| 417 | + | |
| 418 | + | |
| 419 | + | |
| 420 | + | |
| 421 | + | |
| 422 | + | |
| 423 | + | |
| 424 | + | |
| 425 | + | |
| 426 | + | |
| 427 | + | |
| 428 | + | |
| 429 | + | |
| 430 | + | |
| 431 | + | |
| 432 | + | |
| 433 | + | |
| 434 | + | |
| 435 | + | |
413 | 436 | | |
414 | 437 | | |
415 | 438 | | |
| |||
Lines changed: 15 additions & 4 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
49 | 49 | | |
50 | 50 | | |
51 | 51 | | |
| 52 | + | |
| 53 | + | |
52 | 54 | | |
53 | 55 | | |
54 | 56 | | |
55 | 57 | | |
56 | 58 | | |
57 | | - | |
58 | | - | |
59 | | - | |
| 59 | + | |
| 60 | + | |
| 61 | + | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
60 | 71 | | |
61 | 72 | | |
62 | | - | |
| 73 | + | |
0 commit comments