Skip to content

Commit

Permalink
Run default _AsyncLLMEngine._run_workers_async in threadpool (#1628)
Browse files Browse the repository at this point in the history
  • Loading branch information
dominik-schwabe authored Nov 11, 2023
1 parent 0d57822 commit 1b290ac
Showing 1 changed file with 6 additions and 7 deletions.
13 changes: 6 additions & 7 deletions vllm/engine/async_llm_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,18 +206,17 @@ async def _run_workers_async(
**kwargs,
) -> Any:
"""Runs the given method on all workers."""
all_outputs = []
coros = []
for worker in self.workers:
if self.parallel_config.worker_use_ray:
executor = partial(worker.execute_method.remote, method)
coros.append(
worker.execute_method.remote(method, *args, **kwargs))
else:
executor = getattr(worker, method)
coros.append(asyncio.get_event_loop().run_in_executor(
None, partial(executor, *args, **kwargs)))

output = executor(*args, **kwargs)
all_outputs.append(output)

if self.parallel_config.worker_use_ray:
all_outputs = await asyncio.gather(*all_outputs)
all_outputs = await asyncio.gather(*coros)

if get_all_outputs:
return all_outputs
Expand Down

0 comments on commit 1b290ac

Please sign in to comment.