Skip to content

Commit 2253bf8

Browse files
committed
refactor: simplify EngineCoreClient.make_client in AsyncLLM
Signed-off-by: googs1025 <googs1025@gmail.com>
1 parent 4e88723 commit 2253bf8

File tree

2 files changed

+22
-17
lines changed

2 files changed

+22
-17
lines changed

vllm/v1/engine/async_llm.py

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,7 @@
2727
from vllm.usage.usage_lib import UsageContext
2828
from vllm.utils import Device, cdiv
2929
from vllm.v1.engine import EngineCoreRequest
30-
from vllm.v1.engine.core_client import (AsyncMPClient, DPAsyncMPClient,
31-
RayDPClient)
30+
from vllm.v1.engine.core_client import EngineCoreClient
3231
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
3332
from vllm.v1.engine.output_processor import (OutputProcessor,
3433
RequestOutputCollector)
@@ -120,15 +119,8 @@ def __init__(
120119
log_stats=self.log_stats)
121120

122121
# EngineCore (starts the engine in background process).
123-
core_client_class: type[AsyncMPClient]
124-
if vllm_config.parallel_config.data_parallel_size == 1:
125-
core_client_class = AsyncMPClient
126-
elif vllm_config.parallel_config.data_parallel_backend == "ray":
127-
core_client_class = RayDPClient
128-
else:
129-
core_client_class = DPAsyncMPClient
130-
131-
self.engine_core = core_client_class(
122+
123+
self.engine_core = EngineCoreClient.make_async_mp_client(
132124
vllm_config=vllm_config,
133125
executor_class=executor_class,
134126
log_stats=self.log_stats,

vllm/v1/engine/core_client.py

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -67,18 +67,31 @@ def make_client(
6767
"is not currently supported.")
6868

6969
if multiprocess_mode and asyncio_mode:
70-
if vllm_config.parallel_config.data_parallel_size > 1:
71-
if vllm_config.parallel_config.data_parallel_backend == "ray":
72-
return RayDPClient(vllm_config, executor_class, log_stats)
73-
return DPAsyncMPClient(vllm_config, executor_class, log_stats)
74-
75-
return AsyncMPClient(vllm_config, executor_class, log_stats)
70+
return EngineCoreClient.make_async_mp_client(
71+
vllm_config, executor_class, log_stats)
7672

7773
if multiprocess_mode and not asyncio_mode:
7874
return SyncMPClient(vllm_config, executor_class, log_stats)
7975

8076
return InprocClient(vllm_config, executor_class, log_stats)
8177

78+
@staticmethod
79+
def make_async_mp_client(
80+
vllm_config: VllmConfig,
81+
executor_class: type[Executor],
82+
log_stats: bool,
83+
client_addresses: Optional[dict[str, str]] = None,
84+
client_index: int = 0,
85+
) -> "MPClient":
86+
if vllm_config.parallel_config.data_parallel_size > 1:
87+
if vllm_config.parallel_config.data_parallel_backend == "ray":
88+
return RayDPClient(vllm_config, executor_class, log_stats,
89+
client_addresses, client_index)
90+
return DPAsyncMPClient(vllm_config, executor_class, log_stats,
91+
client_addresses, client_index)
92+
return AsyncMPClient(vllm_config, executor_class, log_stats,
93+
client_addresses, client_index)
94+
8295
@abstractmethod
8396
def shutdown(self):
8497
...

0 commit comments

Comments
 (0)