[XPU] Optimize XPU worker shutdown logic to prevent resource leak (#46433)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Jun 30, 2026 at 15:27 UTC
8cc242335de805cac390580f0dcd9e69b6ed86c0
7 files changed
+41
-5
.buildkite/intel_jobs/basic_correctness.yaml
+1
@@ -23,4 +23,5 @@ steps:
23
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
24
'cd tests &&
25
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
26
+ pytest -v -s basic_correctness/test_cpu_offload.py &&
27
pytest -v -s basic_correctness/test_mem.py::test_end_to_end'
.buildkite/intel_jobs/lora_intel.yaml
+1
-1
@@ -128,10 +128,10 @@ steps:
128
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
129
'cd tests &&
130
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
131
- (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
131
pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
132
pytest -v -s lora/test_transformers_model.py &&
133
pytest -v -s lora/test_chatglm3_tp.py &&
134
+ pytest -v -s lora/test_llama_tp.py::test_llama_lora &&
135
pytest -s -v lora/test_minicpmv_tp.py'
136
137
- label: LoRA Multimodal
tests/utils.py
+7
@@ -604,6 +604,13 @@ class RemoteVLLMServer:
604
mem_info = nvmlDeviceGetMemoryInfo(handle)
605
total_used += mem_info.used
606
return total_used
607
+ elif current_platform.is_xpu():
608
+ total_used = 0
609
+ device_count = current_platform.device_count()
610
+ for i in range(device_count):
611
+ free, total = torch.xpu.mem_get_info(i)
612
+ total_used += total - free
613
+ return total_used
614
except Exception as e:
615
print(f"[RemoteOpenAIServer] Could not query GPU memory: {e}")
616
return None
vllm/platforms/xpu.py
+10
@@ -243,6 +243,16 @@ class XPUPlatform(Platform):
243
if "VLLM_WORKER_MULTIPROC_METHOD" not in os.environ:
244
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
245
246
+ # XPU requires graceful shutdown to allow oneCCL/Level Zero resources
247
+ # to be properly released. Without this, subsequent server startups on
248
+ # the same devices may hang during CCL initialization.
249
+ if vllm_config.shutdown_timeout == 0:
250
+ vllm_config.shutdown_timeout = 5
251
+ logger.info(
252
+ "XPU platform: set server shutdown_timeout=%d.",
253
+ vllm_config.shutdown_timeout,
254
+ )
255
+
256
@classmethod
257
def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
258
super().update_block_size_for_backend(vllm_config)
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -6385,7 +6385,7 @@ class GPUModelRunner(
6385
_ROPE_DICT.clear()
6386
6387
reset_workspace_manager()
6388
- if current_platform.is_rocm():
6388
+ if current_platform.is_rocm() or current_platform.is_xpu():
6389
gc.collect()
6390
torch.accelerator.empty_cache()
6391
torch.accelerator.synchronize()
vllm/v1/worker/gpu_worker.py
+4
-3
@@ -1294,10 +1294,11 @@ class Worker(WorkerBase):
1294
# Release kept-alive cumem pools while the pluggable allocator wrappers
1295
# and callbacks are still alive, so MemPool teardown is not deferred to
1296
# interpreter finalization (pytorch/pytorch#145168).
1297
- from vllm.device_allocator.cumem import CuMemAllocator
1297
+ if current_platform.is_cuda_alike():
1298
+ from vllm.device_allocator.cumem import CuMemAllocator
1299
1299
- if CuMemAllocator.instance is not None:
1300
- CuMemAllocator.instance.release_pools()
1300
+ if CuMemAllocator.instance is not None:
1301
+ CuMemAllocator.instance.release_pools()
1302
1303
def elastic_ep_execute(self, execute_method: str, *args, **kwargs):
1304
return self.elastic_ep_executor.execute(execute_method, *args, **kwargs)
vllm/v1/worker/xpu_worker.py
+17
@@ -162,3 +162,20 @@ class XPUWorker(Worker):
162
logger.debug("Starting torch profiler with trace name: %s", trace_name)
163
164
super().profile(is_start=is_start, profile_prefix=profile_prefix)
165
+
166
+ def shutdown(self) -> None:
167
+ logger.info(
168
+ "XPUWorker shutdown: cleaning up (rank=%d, local_rank=%d)",
169
+ self.rank,
170
+ self.local_rank,
171
+ )
172
+ super().shutdown()
173
+ from vllm.device_allocator.xpumem import XpuMemAllocator
174
+
175
+ if XpuMemAllocator.instance is not None:
176
+ XpuMemAllocator.instance.release_pools()
177
+ logger.info(
178
+ "XPUWorker shutdown: done (rank=%d, local_rank=%d)",
179
+ self.rank,
180
+ self.local_rank,
181
+ )