diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index 972ebe6512..a4ed498169 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -3,7 +3,6 @@ """DeepEP test utilities.""" import dataclasses -import os import traceback from collections.abc import Callable from typing import Concatenate @@ -14,7 +13,7 @@ from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUs from typing_extensions import ParamSpec from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2 -from vllm.utils.network_utils import get_open_port +from vllm.utils.network_utils import get_file_store_init_method if has_deep_ep(): from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( @@ -105,7 +104,7 @@ def parallel_launch( world_size, world_size, 0, - f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}", + get_file_store_init_method(), worker, ) + args,