Skip to content

Commit 6fee0db

Browse files
googs1025leoli1208
authored andcommitted
[Misc] refactor: simplify EngineCoreClient.make_async_mp_client in AsyncLLM (vllm-project#18817)
Signed-off-by: googs1025 <googs1025@gmail.com>
1 parent bde2426 commit 6fee0db

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
@@ -28,8 +28,7 @@
2828
from vllm.usage.usage_lib import UsageContext
2929
from vllm.utils import Device, cdiv
3030
from vllm.v1.engine import EngineCoreRequest
31-
from vllm.v1.engine.core_client import (AsyncMPClient, DPAsyncMPClient,
32-
RayDPClient)
31+
from vllm.v1.engine.core_client import EngineCoreClient
3332
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
3433
from vllm.v1.engine.output_processor import (OutputProcessor,
3534
RequestOutputCollector)
@@ -121,15 +120,8 @@ def __init__(
121120
log_stats=self.log_stats)
122121

123122
# EngineCore (starts the engine in background process).
124-
core_client_class: type[AsyncMPClient]
125-
if vllm_config.parallel_config.data_parallel_size == 1:
126-
core_client_class = AsyncMPClient
127-
elif vllm_config.parallel_config.data_parallel_backend == "ray":
128-
core_client_class = RayDPClient
129-
else:
130-
core_client_class = DPAsyncMPClient
131-
132-
self.engine_core = core_client_class(
123+
124+
self.engine_core = EngineCoreClient.make_async_mp_client(
133125
vllm_config=vllm_config,
134126
executor_class=executor_class,
135127
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
@@ -68,18 +68,31 @@ def make_client(
6868
"is not currently supported.")
6969

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

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

8177
return InprocClient(vllm_config, executor_class, log_stats)
8278

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

0 commit comments

Comments
 (0)