[XPU] fix_test_worker_memory_snapshot_for_xpu (#50097)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Qiming Zhang committed
Sep 15, 2026 at 20:05 UTC
4ace5cd9d80a8d56bf935c4a058a1369f1b94da3
1 file changed
+21
-6
tests/v1/worker/test_worker_memory_snapshot.py
+21
-6
@@ -12,6 +12,7 @@ import torch
12
13
from vllm.config import set_current_vllm_config
14
from vllm.engine.arg_utils import EngineArgs
15
+from vllm.platforms import current_platform
16
from vllm.utils.mem_utils import MemorySnapshot
17
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
18
@@ -66,8 +67,22 @@ def worker_process(
67
model="facebook/opt-125m", tensor_parallel_size=2, load_format="dummy"
68
).create_engine_config()
69
70
+ # Select worker class and patch target based on platform
71
+ if current_platform.is_xpu():
72
+ from vllm.v1.worker.xpu_worker import XPUWorker
73
+
74
+ worker_cls = XPUWorker
75
+ init_patch_target = (
76
+ "vllm.v1.worker.xpu_worker.init_worker_distributed_environment"
77
+ )
78
+ else:
79
+ worker_cls = Worker
80
+ init_patch_target = (
81
+ "vllm.v1.worker.gpu_worker.init_worker_distributed_environment"
82
+ )
83
+
84
# Create worker
70
- worker = Worker(
85
+ worker = worker_cls(
86
vllm_config=vllm_config,
87
local_rank=rank,
88
rank=rank,
@@ -81,7 +96,7 @@ def worker_process(
96
97
# Apply minimal patches to track operation order
98
init_patch = patch(
84
- "vllm.v1.worker.gpu_worker.init_worker_distributed_environment",
99
+ init_patch_target,
100
side_effect=make_operation_tracker(
101
"init_distributed", original_init_worker
102
),
@@ -93,7 +108,7 @@ def worker_process(
108
)
109
all_reduce_patch = patch(
110
"torch.distributed.all_reduce",
96
- side_effect=make_operation_tracker("nccl_all_reduce", original_all_reduce),
111
+ side_effect=make_operation_tracker("dist_all_reduce", original_all_reduce),
112
)
113
114
with (
@@ -182,13 +197,13 @@ def test_init_distributed_is_called_before_memory_snapshot():
197
198
# Raises ValueError if the operation is not found
199
init_distributed = rank_ops.index("init_distributed")
185
- nccl_all_reduce = rank_ops.index("nccl_all_reduce")
200
+ dist_all_reduce = rank_ops.index("dist_all_reduce")
201
memory_snapshot = rank_ops.index("memory_snapshot")
202
203
# Verify order: init_distributed should happen before memory_snapshot
189
- assert init_distributed < nccl_all_reduce < memory_snapshot, (
204
+ assert init_distributed < dist_all_reduce < memory_snapshot, (
205
f"Rank {rank}: init_distributed (index {init_distributed}) "
191
- f"must happen before nccl_all_reduce (index {nccl_all_reduce}) "
206
+ f"must happen before dist_all_reduce (index {dist_all_reduce}) "
207
f"and memory_snapshot (index {memory_snapshot})"
208
)
209