[XPU][CI] fix hf runner (#52730)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Qiming Zhang committed
Aug 18, 2026 at 20:30 UTC
e575b5f1a967b7ebb44eb65c4a3203e269cfc54d
1 file changed
+10
-7
tests/conftest.py
+10
-7
@@ -893,12 +893,15 @@ class HfRunner:
893
return self.model.predict(prompts, *args, convert_to_tensor=True, **kwargs)
894
895
def __enter__(self):
896
- if current_platform.is_rocm():
897
- # Record starting memory usage stats on ROCm so that we can wait for
898
- # memory to roughly settle back below these levels on shutdown. This is
899
- # helpful in cases where the HfRunner is initialized after significant GPU
900
- # memory is already occupied, e.g. in
896
+ if current_platform.is_rocm() or current_platform.is_xpu():
897
+ # Record starting memory usage stats on ROCm/XPU so that we can
898
+ # wait for memory to roughly settle back below these levels on
899
+ # shutdown. This is helpful in cases where the HfRunner is
900
+ # initialized after significant GPU memory is already occupied,
901
+ # e.g. in
902
# tests/basic_correctness/test_basic_correctness.py::test_models_distributed
903
+ # where vllm worker processes are still alive and holding GPU
904
+ # memory when hf_runner.__exit__ is called.
905
from tests.utils import (
906
get_physical_device_indices,
907
record_gpu_memory_usage_stats,
@@ -918,8 +921,8 @@ class HfRunner:
921
922
del self.model
923
cleanup_dist_env_and_memory()
921
- # ROCm frees VRAM lazily; wait so a runner started right after this HF
922
- # model exits does not OOM on its startup memory guard.
924
+ # ROCm/XPU free VRAM lazily; wait so a runner started right after this
925
+ # HF model exits does not OOM on its startup memory guard.
926
wait_for_memory_to_settle(
927
threshold_ratio=getattr(self, "threshold_ratios", None)
928
)