You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2024/10/07 20:01:28 UTC
(beam) 01/01: vLLM model handler efficiency improvements
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch users/damccorm/vllmEfficiency
in repository https://gitbox.apache.org/repos/asf/beam.git
commit 40c78161d8c86cb0401b9f3e06fc7175b2bcfca1
Author: Danny McCormick <da...@google.com>
AuthorDate: Mon Oct 7 16:01:10 2024 -0400
vLLM model handler efficiency improvements
---
.github/trigger_files/beam_PostCommit_Python.json | 2 +-
.../apache_beam/ml/inference/vllm_inference.py | 109 ++++++++++++++-------
2 files changed, 75 insertions(+), 36 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json
index 1eb60f6e495..9e1d1e1b80d 100644
--- a/.github/trigger_files/beam_PostCommit_Python.json
+++ b/.github/trigger_files/beam_PostCommit_Python.json
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run.",
- "modification": 3
+ "modification": 4
}
diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py b/sdks/python/apache_beam/ml/inference/vllm_inference.py
index 28890083d93..630642f57f7 100644
--- a/sdks/python/apache_beam/ml/inference/vllm_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py
@@ -17,6 +17,7 @@
# pytype: skip-file
+import asyncio
import logging
import os
import subprocess
@@ -35,7 +36,7 @@ from apache_beam.io.filesystems import FileSystems
from apache_beam.ml.inference.base import ModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.utils import subprocess_server
-from openai import OpenAI
+from openai import AsyncOpenAI, OpenAI
try:
import vllm # pylint: disable=unused-import
@@ -94,6 +95,15 @@ def getVLLMClient(port) -> OpenAI:
)
+def getAsyncVLLMClient(port) -> AsyncOpenAI:
+ openai_api_key = "EMPTY"
+ openai_api_base = f"http://localhost:{port}/v1"
+ return AsyncOpenAI(
+ api_key=openai_api_key,
+ base_url=openai_api_base,
+ )
+
+
class _VLLMModelServer():
def __init__(self, model_name: str, vllm_server_kwargs: Dict[str, str]):
self._model_name = model_name
@@ -142,6 +152,7 @@ class _VLLMModelServer():
time.sleep(5)
if retries == 0:
+ raise Exception('test')
self._server_started = False
raise Exception(
"Failed to start vLLM server, polling process exited with code " +
@@ -150,6 +161,8 @@ class _VLLMModelServer():
else:
self.start_server(retries - 1)
+ raise Exception('test2')
+
class VLLMCompletionsModelHandler(ModelHandler[str,
PredictionResult,
@@ -183,6 +196,34 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
def load_model(self) -> _VLLMModelServer:
return _VLLMModelServer(self._model_name, self._vllm_server_kwargs)
+
+ async def _async_run_inference(
+ self,
+ batch: Sequence[Sequence[OpenAIChatMessage]],
+ model: _VLLMModelServer,
+ inference_args: Optional[Dict[str, Any]] = None
+ ) -> Iterable[PredictionResult]:
+ client = getAsyncVLLMClient(model.get_server_port())
+ inference_args = inference_args or {}
+ async_predictions = []
+ for prompt in batch:
+ try:
+ completion = client.completions.create(
+ model=self._model_name, prompt=prompt, **inference_args)
+ async_predictions.append(completion)
+ except Exception as e:
+ model.check_connectivity()
+ raise e
+
+ predictions = []
+ for p in async_predictions:
+ try:
+ predictions.append(await p)
+ except Exception as e:
+ model.check_connectivity()
+ raise e
+
+ return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
def run_inference(
self,
@@ -200,22 +241,7 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
Returns:
An Iterable of type PredictionResult.
"""
- client = getVLLMClient(model.get_server_port())
- inference_args = inference_args or {}
- predictions = []
- # TODO(https://github.com/apache/beam/issues/32528): We should add support
- # for taking in batches and doing a bunch of async calls. That will end up
- # being more efficient when we can do in bundle batching.
- for prompt in batch:
- try:
- completion = client.completions.create(
- model=self._model_name, prompt=prompt, **inference_args)
- predictions.append(completion)
- except Exception as e:
- model.check_connectivity()
- raise e
-
- return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
+ return asyncio.run(self._async_run_inference(batch, model, inference_args))
def share_model_across_processes(self) -> bool:
return True
@@ -271,29 +297,16 @@ class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
self._vllm_server_kwargs['chat_template'] = local_chat_template_path
return _VLLMModelServer(self._model_name, self._vllm_server_kwargs)
-
- def run_inference(
+
+ async def _async_run_inference(
self,
batch: Sequence[Sequence[OpenAIChatMessage]],
model: _VLLMModelServer,
inference_args: Optional[Dict[str, Any]] = None
) -> Iterable[PredictionResult]:
- """Runs inferences on a batch of text strings.
-
- Args:
- batch: A sequence of examples as OpenAI messages.
- model: A _VLLMModelServer for connecting to the spun up server.
- inference_args: Any additional arguments for an inference.
-
- Returns:
- An Iterable of type PredictionResult.
- """
- client = getVLLMClient(model.get_server_port())
+ client = getAsyncVLLMClient(model.get_server_port())
inference_args = inference_args or {}
- predictions = []
- # TODO(https://github.com/apache/beam/issues/32528): We should add support
- # for taking in batches and doing a bunch of async calls. That will end up
- # being more efficient when we can do in bundle batching.
+ async_predictions = []
for messages in batch:
formatted = []
for message in messages:
@@ -301,12 +314,38 @@ class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
try:
completion = client.chat.completions.create(
model=self._model_name, messages=formatted, **inference_args)
- predictions.append(completion)
+ async_predictions.append(completion)
+ except Exception as e:
+ model.check_connectivity()
+ raise e
+
+ predictions = []
+ for p in async_predictions:
+ try:
+ predictions.append(await p)
except Exception as e:
model.check_connectivity()
raise e
return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
+ def run_inference(
+ self,
+ batch: Sequence[Sequence[OpenAIChatMessage]],
+ model: _VLLMModelServer,
+ inference_args: Optional[Dict[str, Any]] = None
+ ) -> Iterable[PredictionResult]:
+ """Runs inferences on a batch of text strings.
+
+ Args:
+ batch: A sequence of examples as OpenAI messages.
+ model: A _VLLMModelServer for connecting to the spun up server.
+ inference_args: Any additional arguments for an inference.
+
+ Returns:
+ An Iterable of type PredictionResult.
+ """
+ return asyncio.run(self._async_run_inference(batch, model, inference_args))
+
def share_model_across_processes(self) -> bool:
return True