[CI/UT][Graph] Add ut for torchair graph mode (#1103)

### What this PR does / why we need it?
Add ut for torchair graph mode on DeepSeekV3

### How was this patch tested?
CI passed with new added test.

---------

Signed-off-by: MengqingCao <cmq0113@163.com>
Signed-off-by: Mengqing Cao <cmq0113@163.com>
This commit is contained in:
Mengqing Cao
2025-06-14 16:59:00 +08:00
committed by GitHub
parent 94a52cf577
commit a3b5af8307
4 changed files with 100 additions and 12 deletions

View File

@@ -17,6 +17,7 @@
# Adapted from vllm-project/vllm/blob/main/tests/conftest.py
#
import contextlib
import gc
from typing import List, Optional, Tuple, TypeVar, Union
@@ -53,11 +54,17 @@ PromptAudioInput = _PromptMultiModalInput[Tuple[np.ndarray, int]]
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
def cleanup_dist_env_and_memory():
def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
destroy_model_parallel()
destroy_distributed_environment()
with contextlib.suppress(AssertionError):
torch.distributed.destroy_process_group()
if shutdown_ray:
import ray # Lazy import Ray
ray.shutdown()
gc.collect()
torch.npu.empty_cache()
torch.npu.reset_peak_memory_stats()
class VllmRunner: