This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new aec5b5da813 Add embedded NVIDIA Dynamo support to vLLM ModelHandler
(#38701)
aec5b5da813 is described below
commit aec5b5da8135d009701975a3cc07a4c7a175ae0d
Author: akshayjadiyanv <[email protected]>
AuthorDate: Tue Jun 30 07:46:04 2026 -0700
Add embedded NVIDIA Dynamo support to vLLM ModelHandler (#38701)
* Add embedded NVIDIA Dynamo support to vLLM ModelHandler
VLLMCompletionsModelHandler and VLLMChatModelHandler gain two
keyword-only parameters, use_dynamo (default False) and
dynamo_frontend_kwargs. When use_dynamo=True, the handler launches a
dynamo.frontend process as the OpenAI-compatible local endpoint plus a
separate dynamo.vllm worker, instead of vllm.entrypoints.openai.api_server.
The existing native-vLLM path is unchanged when the flag is absent.
The example pipeline vllm_text_completion.py gains --use_dynamo and
--max_tokens flags. validate_inference_args is now a no-op on both
handlers so OpenAI-style request kwargs (e.g. max_tokens) can be passed
through RunInference. A new unit-test module covers process-launch
behaviour for both paths.
This supersedes #36966 (now closed) and rebases the embedded-Dynamo
approach onto current master, preserving the recent batching-kwargs
additions to the ModelHandler base.
Co-authored-by: Danny McCormick <[email protected]>
* Harden _VLLMModelServer process lifecycle per code review
Apply five robustness fixes flagged on PR #38701:
- Track the temporary etcd data dir as self._etcd_data_dir and
shutil.rmtree(..., ignore_errors=True) it in _stop_processes so worker
restarts don't leak /tmp directories.
- Wrap process.terminate() / process.wait() / process.kill() in a single
try/except OSError to absorb the ProcessLookupError race when a process
exits between poll() and the signal call.
- Switch the ETCD_ENDPOINTS removal from `del os.environ[...]` to
`os.environ.pop(..., None)` to be idempotent.
- Wrap __del__ in try/except Exception so cleanup never raises during
interpreter shutdown.
- Add the embedded etcd process to the check_connectivity() poll loop so
an etcd death fails fast instead of waiting out the 10-minute timeout.
* Enable Dataflow IT for embedded Dynamo on T4
Bump vllm.dockerfile.old to apache-beam[gcp]==2.71.0 (and the
COPY-from beam_python3.12_sdk image to 2.71.0), install
ai-dynamo[vllm], and add the etcd binary required by embedded
Dynamo's runtime discovery.
Uncomment the Dynamo IT block in common.gradle. Drop the unused
machine_type override so it inherits n1-standard-4 from argMap,
and switch nvidia-l4 -> nvidia-tesla-t4 to match the existing
native vLLM ITs and the local Dataflow validation (per @damccorm
review).
Validated end-to-end on Dataflow with Qwen/Qwen3-0.6B; the
nvext.timing field present on every PredictionResult confirms the
Dynamo frontend served the requests.
* Trigger Python PostCommit for Dynamo IT
Bump the beam_PostCommit_Python trigger file so the postcommit
suite (inferencePostCommitITPy312 -> vllmTests) runs the embedded
Dynamo IT against the rebuilt apache-beam-testing vLLM image.
* fix: run Dynamo vLLM IT separately in py312 PostCommit
Split vllmDynamoTests from vllmTests so py312 validates Dynamo without
blocking on the pre-existing native opt-125m hang in apache-beam-testing.
* Restore full vLLM postcommit suite
* fix: fold Dynamo IT back into vllmTests for py312 PostCommit
Revert the Option A split now that the native opt-125m vLLM hang is
fixed (3.12 PostCommit passed in ~2.5h). vllmTests again runs
completion -> chat -> Dynamo as a single suite; the separate
vllmDynamoTests task is removed. Bump PostCommit trigger to re-run.
---------
Co-authored-by: Danny McCormick <[email protected]>
---
.github/trigger_files/beam_PostCommit_Python.json | 4 +-
.../examples/inference/vllm_text_completion.py | 24 +-
.../inference/test_resources/vllm.dockerfile.old | 14 +-
.../apache_beam/ml/inference/vllm_inference.py | 331 ++++++++++++++++++---
.../ml/inference/vllm_inference_test.py | 157 ++++++++++
sdks/python/test-suites/dataflow/common.gradle | 40 ++-
6 files changed, 510 insertions(+), 60 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python.json
b/.github/trigger_files/beam_PostCommit_Python.json
index 2bb052d5f71..00c0cdc3f9c 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.",
- "pr": "37345",
- "modification": 53
+ "pr": "38701",
+ "modification": 55
}
diff --git a/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
b/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
index a7468f521eb..21f0d41d3cb 100644
--- a/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
+++ b/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
@@ -138,6 +138,20 @@ def parse_known_args(argv):
'Passed to the vLLM OpenAI server as --gpu-memory-utilization '
'(fraction of total GPU memory for KV cache). Lower this if the '
'engine fails to start with CUDA out of memory.'))
+ parser.add_argument(
+ '--use_dynamo',
+ dest='use_dynamo',
+ action='store_true',
+ help=(
+ 'Use embedded NVIDIA Dynamo as the vLLM engine. Requires '
+ 'ai-dynamo[vllm] and the etcd binary in the runtime environment. '
+ 'See VLLMCompletionsModelHandler for limitations of embedded mode.'))
+ parser.add_argument(
+ '--max_tokens',
+ dest='max_tokens',
+ type=int,
+ default=16,
+ help='Maximum number of tokens to generate for each example.')
return parser.parse_known_args(argv)
@@ -178,14 +192,17 @@ def run(
build_vllm_server_kwargs(known_args))
model_handler = VLLMCompletionsModelHandler(
- model_name=known_args.model, vllm_server_kwargs=effective_vllm_kwargs)
+ model_name=known_args.model,
+ vllm_server_kwargs=effective_vllm_kwargs,
+ use_dynamo=known_args.use_dynamo)
input_examples = COMPLETION_EXAMPLES
if known_args.chat:
model_handler = VLLMChatModelHandler(
model_name=known_args.model,
chat_template_path=known_args.chat_template,
- vllm_server_kwargs=dict(effective_vllm_kwargs))
+ vllm_server_kwargs=dict(effective_vllm_kwargs),
+ use_dynamo=known_args.use_dynamo)
input_examples = CHAT_EXAMPLES
pipeline = test_pipeline
@@ -193,7 +210,8 @@ def run(
pipeline = beam.Pipeline(options=pipeline_options)
examples = pipeline | "Create examples" >> beam.Create(input_examples)
- predictions = examples | "RunInference" >> RunInference(model_handler)
+ predictions = examples | "RunInference" >> RunInference(
+ model_handler, inference_args={'max_tokens': known_args.max_tokens})
process_output = predictions | "Process Predictions" >> beam.ParDo(
PostProcessor())
_ = process_output | "WriteOutput" >> beam.io.WriteToText(
diff --git
a/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old
b/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old
index b9c99e49e02..d0080debc09 100644
--- a/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old
+++ b/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old
@@ -34,14 +34,22 @@ RUN python3 --version
RUN apt-get install -y curl
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12 && pip install
--upgrade pip
-RUN pip install --no-cache-dir -vvv apache-beam[gcp]==2.58.1
-RUN pip install openai vllm
+RUN pip install --no-cache-dir -vvv apache-beam[gcp]==2.71.0
+RUN pip install --no-cache-dir openai vllm ai-dynamo[vllm]
RUN apt install libcairo2-dev pkg-config python3-dev -y
RUN pip install pycairo
+# etcd binary required by embedded NVIDIA Dynamo for runtime discovery.
+ENV ETCD_VERSION=v3.5.13
+RUN curl -L
https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-linux-amd64.tar.gz
-o /tmp/etcd.tar.gz && \
+ tar xzf /tmp/etcd.tar.gz -C /tmp && \
+ mv /tmp/etcd-${ETCD_VERSION}-linux-amd64/etcd /usr/local/bin/etcd && \
+ chmod +x /usr/local/bin/etcd && \
+ rm -rf /tmp/etcd*
+
# Copy the Apache Beam worker dependencies from the Beam Python 3.12 SDK image.
-COPY --from=apache/beam_python3.12_sdk:2.58.1 /opt/apache/beam /opt/apache/beam
+COPY --from=apache/beam_python3.12_sdk:2.71.0 /opt/apache/beam /opt/apache/beam
# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]
\ No newline at end of file
diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py
b/sdks/python/apache_beam/ml/inference/vllm_inference.py
index 38283f1efd4..e5d8918d5ec 100644
--- a/sdks/python/apache_beam/ml/inference/vllm_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py
@@ -20,10 +20,12 @@
import asyncio
import logging
import os
+import shutil
import subprocess
import sys
import threading
import time
+import urllib.request
import uuid
from collections.abc import Callable
from collections.abc import Iterable
@@ -109,36 +111,216 @@ def getAsyncVLLMClient(port) -> AsyncOpenAI:
)
+# Embedded Dynamo runtime defaults proven on the smoke test: etcd discovery,
+# TCP request plane, ZMQ event plane, KV events disabled. KV-aware routing,
+# disaggregated prefill/decode, and the Planner are not active in this mode.
+_DYNAMO_FRONTEND_DEFAULT_KWARGS: dict[str, Optional[str]] = {
+ 'discovery-backend': 'etcd',
+ 'request-plane': 'tcp',
+ 'event-plane': 'zmq',
+ 'router-mode': 'round-robin',
+ 'no-router-kv-events': None,
+}
+
+_DYNAMO_ENGINE_DEFAULT_KWARGS: dict[str, Optional[str]] = {
+ 'discovery-backend': 'etcd',
+ 'request-plane': 'tcp',
+ 'event-plane': 'zmq',
+ 'kv-events-config': '{"enable_kv_cache_events": false}',
+}
+
+
+def _append_kwargs(cmd: list[str], kwargs: dict[str, Optional[str]]) -> None:
+ for k, v in kwargs.items():
+ cmd.append(f'--{k}')
+ # Only add values for commands with value part.
+ if v is not None:
+ cmd.append(v)
+
+
+def _uses_etcd_discovery(kwargs: dict[str, Optional[str]]) -> bool:
+ return kwargs.get('discovery-backend') == 'etcd'
+
+
class _VLLMModelServer():
- def __init__(self, model_name: str, vllm_server_kwargs: dict[str, str]):
+ def __init__(
+ self,
+ model_name: str,
+ vllm_server_kwargs: dict[str, Optional[str]],
+ dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None,
+ use_dynamo: bool = False):
self._model_name = model_name
self._vllm_server_kwargs = vllm_server_kwargs
+ self._dynamo_frontend_kwargs = dynamo_frontend_kwargs or {}
self._server_started = False
self._server_process = None
+ self._dynamo_process = None
+ self._etcd_process = None
+ self._etcd_data_dir: Optional[str] = None
+ self._managed_etcd_endpoint = None
self._server_port: int = -1
self._server_process_lock = threading.RLock()
+ self._use_dynamo = use_dynamo
self.start_server()
+ @staticmethod
+ def _stop_process(process: Optional[subprocess.Popen]) -> None:
+ if process is None or process.poll() is not None:
+ return
+ # A process may exit between poll() and terminate() / kill(), in which
+ # case the OS raises ProcessLookupError (or another OSError). Treat that
+ # as already-stopped so we don't bail out of the broader cleanup.
+ try:
+ process.terminate()
+ try:
+ process.wait(timeout=10)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ process.wait()
+ except OSError:
+ pass
+
+ def _stop_processes(self) -> None:
+ self._stop_process(self._dynamo_process)
+ self._stop_process(self._server_process)
+ self._stop_process(self._etcd_process)
+ if (self._managed_etcd_endpoint is not None and
+ os.environ.get('ETCD_ENDPOINTS') == self._managed_etcd_endpoint):
+ os.environ.pop('ETCD_ENDPOINTS', None)
+ if self._etcd_data_dir is not None:
+ shutil.rmtree(self._etcd_data_dir, ignore_errors=True)
+ self._etcd_data_dir = None
+ self._dynamo_process = None
+ self._server_process = None
+ self._etcd_process = None
+ self._managed_etcd_endpoint = None
+ self._server_started = False
+ self._server_port = -1
+
+ def _process_status(self) -> str:
+ process_status = []
+ if self._server_process is not None:
+ process_status.append(
+ 'frontend/server exit code: %s' % self._server_process.poll())
+ if self._dynamo_process is not None:
+ process_status.append(
+ 'dynamo worker exit code: %s' % self._dynamo_process.poll())
+ if self._etcd_process is not None:
+ process_status.append('etcd exit code: %s' % self._etcd_process.poll())
+ return ', '.join(process_status) or 'no process status available'
+
+ def __del__(self):
+ # __del__ may run during interpreter shutdown when module globals can
+ # already be torn down; swallow any cleanup failures so we don't print
+ # a noisy traceback.
+ try:
+ self._stop_processes()
+ except Exception: # pylint: disable=broad-except
+ pass
+
+ def _uses_embedded_etcd(self) -> bool:
+ return (
+ self._use_dynamo and
+ _uses_etcd_discovery(self._dynamo_frontend_kwargs) and
+ _uses_etcd_discovery(self._vllm_server_kwargs) and
+ 'ETCD_ENDPOINTS' not in os.environ)
+
+ def _wait_for_etcd(self, endpoint: str, timeout_secs=30) -> None:
+ deadline = time.time() + timeout_secs
+ health_url = endpoint.rstrip('/') + '/health'
+ while time.time() < deadline and self._etcd_process.poll() is None:
+ try:
+ with urllib.request.urlopen(health_url, timeout=2) as response:
+ if response.status < 500:
+ return
+ except Exception: # pylint: disable=broad-except
+ time.sleep(1)
+
+ process_status = self._process_status()
+ self._stop_processes()
+ raise RuntimeError(
+ "Failed to start embedded etcd for Dynamo. Process status: "
+ f"{process_status}. Install etcd in the worker container or set "
+ "ETCD_ENDPOINTS to an external etcd service.")
+
+ def _ensure_etcd(self) -> None:
+ if not self._uses_embedded_etcd():
+ return
+ if shutil.which('etcd') is None:
+ raise RuntimeError(
+ "Embedded Dynamo mode requires etcd when ETCD_ENDPOINTS is not "
+ "set. Install etcd in the worker container or set ETCD_ENDPOINTS "
+ "to an external etcd service.")
+
+ etcd_name = f'beam-dynamo-etcd-{uuid.uuid4().hex}'
+ self._etcd_data_dir = f'/tmp/{etcd_name}'
+ peer_port, = subprocess_server.pick_port(None)
+ etcd_cmd = [
+ 'etcd',
+ '--name',
+ etcd_name,
+ '--listen-client-urls',
+ 'http://127.0.0.1:{{PORT}}',
+ '--advertise-client-urls',
+ 'http://127.0.0.1:{{PORT}}',
+ '--listen-peer-urls',
+ f'http://127.0.0.1:{peer_port}',
+ '--initial-advertise-peer-urls',
+ f'http://127.0.0.1:{peer_port}',
+ '--initial-cluster',
+ f'{etcd_name}=http://127.0.0.1:{peer_port}',
+ '--data-dir',
+ self._etcd_data_dir,
+ '--log-level',
+ 'warn',
+ ]
+ self._etcd_process, etcd_port = start_process(etcd_cmd)
+ endpoint = f'http://127.0.0.1:{etcd_port}'
+ os.environ['ETCD_ENDPOINTS'] = endpoint
+ self._managed_etcd_endpoint = endpoint
+ self._wait_for_etcd(endpoint)
+
def start_server(self, retries=3):
with self._server_process_lock:
if not self._server_started:
- server_cmd = [
- sys.executable,
- '-m',
- 'vllm.entrypoints.openai.api_server',
- '--model',
- self._model_name,
- '--port',
- '{{PORT}}',
- ]
- for k, v in self._vllm_server_kwargs.items():
- server_cmd.append(f'--{k}')
- # Only add values for commands with value part.
- if v is not None:
- server_cmd.append(v)
+ self._stop_processes()
+ self._ensure_etcd()
+ if self._use_dynamo:
+ # Dynamo embedded mode uses the frontend as its OpenAI-compatible
+ # local endpoint and a separate vLLM worker process.
+ server_cmd = [
+ sys.executable,
+ '-m',
+ 'dynamo.frontend',
+ '--http-port',
+ '{{PORT}}',
+ ]
+ _append_kwargs(server_cmd, self._dynamo_frontend_kwargs)
+ else:
+ server_cmd = [
+ sys.executable,
+ '-m',
+ 'vllm.entrypoints.openai.api_server',
+ '--model',
+ self._model_name,
+ '--port',
+ '{{PORT}}',
+ ]
+ _append_kwargs(server_cmd, self._vllm_server_kwargs)
self._server_process, self._server_port = start_process(server_cmd)
+ if self._use_dynamo:
+ server_cmd = [
+ sys.executable,
+ '-m',
+ 'dynamo.vllm',
+ '--model',
+ self._model_name,
+ ]
+ _append_kwargs(server_cmd, self._vllm_server_kwargs)
+ self._dynamo_process, _ = start_process(server_cmd)
+
self.check_connectivity(retries)
def get_server_port(self) -> int:
@@ -146,9 +328,14 @@ class _VLLMModelServer():
self.start_server()
return self._server_port
- def check_connectivity(self, retries=3):
+ def check_connectivity(self, retries=3, timeout_secs=600):
+ start_time = time.time()
with getVLLMClient(self._server_port) as client:
- while self._server_process.poll() is None:
+ while (time.time() - start_time < timeout_secs and
+ self._server_process.poll() is None and
+ (self._dynamo_process is None or
+ self._dynamo_process.poll() is None) and
+ (self._etcd_process is None or self._etcd_process.poll() is
None)):
try:
models = client.models.list().data
logging.info('models: %s' % models)
@@ -160,12 +347,13 @@ class _VLLMModelServer():
# Sleep while bringing up the process
time.sleep(5)
+ process_status = self._process_status()
+ self._stop_processes()
if retries == 0:
- self._server_started = False
raise Exception(
- "Failed to start vLLM server, polling process exited with code " +
- "%s. Next time a request is tried, the server will be restarted" %
- self._server_process.poll())
+ "Failed to start vLLM server. Process status: "
+ f"{process_status}. Next time a request is tried, the server "
+ "will be restarted")
else:
self.start_server(retries - 1)
@@ -176,8 +364,10 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
def __init__(
self,
model_name: str,
- vllm_server_kwargs: Optional[dict[str, str]] = None,
+ vllm_server_kwargs: Optional[dict[str, Optional[str]]] = None,
*,
+ use_dynamo: bool = False,
+ dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
max_batch_duration_secs: Optional[int] = None,
@@ -197,15 +387,28 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
https://docs.vllm.ai/en/latest/models/supported_models.html for
supported models.
vllm_server_kwargs: Any additional kwargs to be passed into your vllm
- server when it is being created. Will be invoked using
- `python -m vllm.entrypoints.openai.api_serverv <beam provided args>
- <vllm_server_kwargs>`. For example, you could pass
- `{'echo': 'true'}` to prepend new messages with the previous message.
- On ~16GB GPUs, pass lower ``max-num-seqs`` and
- ``gpu-memory-utilization`` values (see
- ``apache_beam.examples.inference.vllm_text_completion``). For a list of
- possible kwargs, see
+ server when it is being created. When ``use_dynamo`` is disabled,
+ this is invoked using ``python -m vllm.entrypoints.openai.api_server
+ <beam provided args> <vllm_server_kwargs>``. When ``use_dynamo`` is
+ enabled, these kwargs are passed to the ``dynamo.vllm`` worker
+ process. For example, you could pass ``{'echo': 'true'}`` to prepend
+ new messages with the previous message. On ~16GB GPUs, pass lower
+ ``max-num-seqs`` and ``gpu-memory-utilization`` values (see
+ ``apache_beam.examples.inference.vllm_text_completion``). For a list
+ of possible kwargs, see
https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-completions-api
+ use_dynamo: Whether to use NVIDIA Dynamo as the underlying vLLM engine.
+ Requires installing Dynamo in your runtime environment
+ (``pip install ai-dynamo[vllm]``). This is an opt-in single-worker
+ embedded mode; KV-aware routing, disaggregated prefill/decode, KVBM
+ offload across nodes, the Planner, and Grove are not active in
+ embedded mode. Dynamo also requires an etcd-style discovery service:
+ when ``ETCD_ENDPOINTS`` is unset, Beam starts a local etcd, which
+ requires the ``etcd`` binary in the worker environment.
+ dynamo_frontend_kwargs: Additional kwargs to be passed to the
+ ``dynamo.frontend`` process when ``use_dynamo`` is enabled. By
+ default, embedded Dynamo uses etcd discovery, TCP request plane, ZMQ
+ event plane, round-robin routing, and disables router KV events.
min_batch_size: optional. the minimum batch size to use when batching
inputs.
max_batch_size: optional. the maximum batch size to use when batching
@@ -229,10 +432,20 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
batch_length_fn=batch_length_fn,
batch_bucket_boundaries=batch_bucket_boundaries)
self._model_name = model_name
- self._vllm_server_kwargs: dict[str, str] = vllm_server_kwargs or {}
+ self._vllm_server_kwargs: dict[str, Optional[str]] = ({
+ **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {})
+ } if use_dynamo else vllm_server_kwargs or {})
+ self._dynamo_frontend_kwargs: dict[str, Optional[str]] = {
+ **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {})
+ }
+ self._use_dynamo = use_dynamo
def load_model(self) -> _VLLMModelServer:
- return _VLLMModelServer(self._model_name, self._vllm_server_kwargs)
+ return _VLLMModelServer(
+ self._model_name,
+ self._vllm_server_kwargs,
+ self._dynamo_frontend_kwargs,
+ self._use_dynamo)
async def _async_run_inference(
self,
@@ -274,6 +487,12 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
"""
return asyncio.run(self._async_run_inference(batch, model, inference_args))
+ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]):
+ # Override the base validator so OpenAI-compatible request kwargs such as
+ # ``max_tokens`` can be passed through ``RunInference`` to the vLLM /
+ # Dynamo server.
+ pass
+
def share_model_across_processes(self) -> bool:
return True
@@ -285,8 +504,10 @@ class
VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
self,
model_name: str,
chat_template_path: Optional[str] = None,
- vllm_server_kwargs: Optional[dict[str, str]] = None,
+ vllm_server_kwargs: Optional[dict[str, Optional[str]]] = None,
*,
+ use_dynamo: bool = False,
+ dynamo_frontend_kwargs: Optional[dict[str, Optional[str]]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
max_batch_duration_secs: Optional[int] = None,
@@ -311,12 +532,26 @@ class
VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
For info on chat templates, see:
https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#chat-template
vllm_server_kwargs: Any additional kwargs to be passed into your vllm
- server when it is being created. Will be invoked using
- `python -m vllm.entrypoints.openai.api_serverv <beam provided args>
- <vllm_server_kwargs>`. For example, you could pass
- `{'echo': 'true'}` to prepend new messages with the previous message.
- For a list of possible kwargs, see
+ server when it is being created. When ``use_dynamo`` is disabled,
+ this is invoked using ``python -m vllm.entrypoints.openai.api_server
+ <beam provided args> <vllm_server_kwargs>``. When ``use_dynamo`` is
+ enabled, these kwargs are passed to the ``dynamo.vllm`` worker
+ process. For example, you could pass ``{'echo': 'true'}`` to prepend
+ new messages with the previous message. For a list of possible
+ kwargs, see
https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-chat-api
+ use_dynamo: Whether to use NVIDIA Dynamo as the underlying vLLM engine.
+ Requires installing Dynamo in your runtime environment
+ (``pip install ai-dynamo[vllm]``). This is an opt-in single-worker
+ embedded mode; KV-aware routing, disaggregated prefill/decode, KVBM
+ offload across nodes, the Planner, and Grove are not active in
+ embedded mode. Dynamo also requires an etcd-style discovery service:
+ when ``ETCD_ENDPOINTS`` is unset, Beam starts a local etcd, which
+ requires the ``etcd`` binary in the worker environment.
+ dynamo_frontend_kwargs: Additional kwargs to be passed to the
+ ``dynamo.frontend`` process when ``use_dynamo`` is enabled. By
+ default, embedded Dynamo uses etcd discovery, TCP request plane, ZMQ
+ event plane, round-robin routing, and disables router KV events.
min_batch_size: optional. the minimum batch size to use when batching
inputs.
max_batch_size: optional. the maximum batch size to use when batching
@@ -340,9 +575,15 @@ class
VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
batch_length_fn=batch_length_fn,
batch_bucket_boundaries=batch_bucket_boundaries)
self._model_name = model_name
- self._vllm_server_kwargs: dict[str, str] = vllm_server_kwargs or {}
+ self._vllm_server_kwargs: dict[str, Optional[str]] = ({
+ **_DYNAMO_ENGINE_DEFAULT_KWARGS, **(vllm_server_kwargs or {})
+ } if use_dynamo else vllm_server_kwargs or {})
+ self._dynamo_frontend_kwargs: dict[str, Optional[str]] = {
+ **_DYNAMO_FRONTEND_DEFAULT_KWARGS, **(dynamo_frontend_kwargs or {})
+ }
self._chat_template_path = chat_template_path
self._chat_file = f'template-{uuid.uuid4().hex}.jinja'
+ self._use_dynamo = use_dynamo
def load_model(self) -> _VLLMModelServer:
chat_template_contents = ''
@@ -355,7 +596,11 @@ class
VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
f.write(chat_template_contents)
self._vllm_server_kwargs['chat_template'] = local_chat_template_path
- return _VLLMModelServer(self._model_name, self._vllm_server_kwargs)
+ return _VLLMModelServer(
+ self._model_name,
+ self._vllm_server_kwargs,
+ self._dynamo_frontend_kwargs,
+ self._use_dynamo)
async def _async_run_inference(
self,
@@ -400,5 +645,11 @@ class
VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
"""
return asyncio.run(self._async_run_inference(batch, model, inference_args))
+ def validate_inference_args(self, inference_args: Optional[dict[str, Any]]):
+ # Override the base validator so OpenAI-compatible request kwargs such as
+ # ``max_tokens`` can be passed through ``RunInference`` to the vLLM /
+ # Dynamo server.
+ pass
+
def share_model_across_processes(self) -> bool:
return True
diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference_test.py
b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py
new file mode 100644
index 00000000000..4ff5186178e
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/vllm_inference_test.py
@@ -0,0 +1,157 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import os
+import sys
+import types
+import unittest
+from unittest import mock
+
+# Protect against environments where the OpenAI python library is not
+# available. The command-construction tests below do not actually need a
+# real OpenAI client; stubbing the module is enough for vllm_inference to
+# import cleanly.
+# pylint: disable=wrong-import-order, wrong-import-position
+try:
+ import openai # pylint: disable=unused-import
+except ImportError:
+ openai = types.ModuleType('openai')
+
+ class _FakeOpenAI:
+ pass
+
+ openai.AsyncOpenAI = _FakeOpenAI
+ openai.OpenAI = _FakeOpenAI
+ sys.modules['openai'] = openai
+
+from apache_beam.ml.inference import vllm_inference
+
+
+class _FakeProcess:
+ def __init__(self):
+ self.returncode = None
+
+ def poll(self):
+ return self.returncode
+
+ def terminate(self):
+ self.returncode = 0
+
+ def wait(self, timeout=None):
+ return self.returncode
+
+ def kill(self):
+ self.returncode = -9
+
+
+class _FakeModels:
+ def list(self):
+ return types.SimpleNamespace(data=[object()])
+
+
+class _FakeClient:
+ def __init__(self):
+ self.models = _FakeModels()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ return False
+
+
+def _record_start_process(commands):
+ def start_process(cmd):
+ commands.append(list(cmd))
+ return _FakeProcess(), 10000 + len(commands)
+
+ return start_process
+
+
+class VLLMInferenceTest(unittest.TestCase):
+ def test_native_vllm_starts_single_server_process(self):
+ commands = []
+ with mock.patch.object(vllm_inference,
+ 'start_process',
+ _record_start_process(commands)):
+ with mock.patch.object(vllm_inference, 'getVLLMClient'):
+ vllm_inference.getVLLMClient.return_value = _FakeClient()
+ vllm_inference.VLLMCompletionsModelHandler(
+ model_name='test-model',
+ vllm_server_kwargs={
+ 'gpu-memory-utilization': '0.9'
+ }).load_model()
+ self.assertEqual(1, len(commands))
+ self.assertIn('vllm.entrypoints.openai.api_server', commands[0])
+ self.assertIn('--model', commands[0])
+ self.assertIn('test-model', commands[0])
+ self.assertIn('--gpu-memory-utilization', commands[0])
+ self.assertIn('0.9', commands[0])
+ self.assertNotIn('dynamo.frontend', commands[0])
+ self.assertNotIn('dynamo.vllm', commands[0])
+
+ def test_dynamo_starts_frontend_and_engine_with_separate_kwargs(self):
+ commands = []
+ with mock.patch.dict(os.environ,
+ {'ETCD_ENDPOINTS': 'http://127.0.0.1:2379'}):
+ with mock.patch.object(vllm_inference,
+ 'start_process',
+ _record_start_process(commands)):
+ with mock.patch.object(vllm_inference, 'getVLLMClient'):
+ vllm_inference.getVLLMClient.return_value = _FakeClient()
+ vllm_inference.VLLMCompletionsModelHandler(
+ model_name='test-model',
+ vllm_server_kwargs={
+ 'tensor-parallel-size': '1'
+ },
+ use_dynamo=True,
+ dynamo_frontend_kwargs={
+ 'router-mode': 'round-robin'
+ }).load_model()
+ self.assertEqual(2, len(commands))
+ frontend_cmd = commands[0]
+ engine_cmd = commands[1]
+ self.assertIn('dynamo.frontend', frontend_cmd)
+ self.assertIn('--http-port', frontend_cmd)
+ self.assertIn('--discovery-backend', frontend_cmd)
+ self.assertIn('--request-plane', frontend_cmd)
+ self.assertIn('--event-plane', frontend_cmd)
+ self.assertIn('--router-mode', frontend_cmd)
+ self.assertIn('--no-router-kv-events', frontend_cmd)
+ self.assertNotIn('--model', frontend_cmd)
+ self.assertNotIn('--tensor-parallel-size', frontend_cmd)
+ self.assertNotIn('--kv-events-config', frontend_cmd)
+ self.assertIn('dynamo.vllm', engine_cmd)
+ self.assertIn('--model', engine_cmd)
+ self.assertIn('test-model', engine_cmd)
+ self.assertIn('--discovery-backend', engine_cmd)
+ self.assertIn('--request-plane', engine_cmd)
+ self.assertIn('--event-plane', engine_cmd)
+ self.assertIn('--kv-events-config', engine_cmd)
+ self.assertIn('--tensor-parallel-size', engine_cmd)
+ self.assertNotIn('--http-port', engine_cmd)
+ self.assertNotIn('--router-mode', engine_cmd)
+ self.assertNotIn('--no-router-kv-events', engine_cmd)
+
+ def test_validate_inference_args_accepts_openai_request_kwargs(self):
+ vllm_inference.VLLMCompletionsModelHandler(
+ 'test-model').validate_inference_args({'max_tokens': 8})
+ vllm_inference.VLLMChatModelHandler('test-model').validate_inference_args(
+ {'max_tokens': 8})
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/test-suites/dataflow/common.gradle
b/sdks/python/test-suites/dataflow/common.gradle
index 480e2a62a2e..c450eb3612f 100644
--- a/sdks/python/test-suites/dataflow/common.gradle
+++ b/sdks/python/test-suites/dataflow/common.gradle
@@ -450,25 +450,28 @@ def tensorRTTests = tasks.create("tensorRTtests") {
}
}
+def vllmBaseArgMap = [
+ "runner": "DataflowRunner",
+ "machine_type":"n1-standard-4",
+ // TODO(https://github.com/apache/beam/issues/22651): Build docker image for
VLLM tests during Run time.
+ // This would also enable to use wheel "--sdk_location" as other tasks, and
eliminate distTarBall dependency
+ // declaration for this project.
+ // Right now, this is built from
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old
+ "sdk_container_image":
"us.gcr.io/apache-beam-testing/python-postcommit-it/vllm:latest",
+ "sdk_location": files(configurations.distTarBall.files).singleFile,
+ "project": "apache-beam-testing",
+ "region": "us-central1",
+ "disk_size_gb": 75
+]
+
def vllmTests = tasks.create("vllmTests") {
dependsOn 'installGcpTest'
dependsOn ':sdks:python:sdist'
doLast {
def testOpts = basicPytestOpts
- def argMap = [
- "runner": "DataflowRunner",
- "machine_type":"n1-standard-4",
- // TODO(https://github.com/apache/beam/issues/22651): Build docker image
for VLLM tests during Run time.
- // This would also enable to use wheel "--sdk_location" as other tasks,
and eliminate distTarBall dependency
- // declaration for this project.
- // Right now, this is built from
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile.old
- "sdk_container_image":
"us.gcr.io/apache-beam-testing/python-postcommit-it/vllm:latest",
- "sdk_location": files(configurations.distTarBall.files).singleFile,
- "project": "apache-beam-testing",
- "region": "us-central1",
+ def argMap = vllmBaseArgMap + [
"model": "facebook/opt-125m",
"output": "gs://apache-beam-ml/outputs/vllm_predictions.txt",
- "disk_size_gb": 75
]
def cmdArgs = mapToArgString(argMap)
// Exec one version with and one version without the chat option
@@ -480,6 +483,19 @@ def vllmTests = tasks.create("vllmTests") {
executable 'sh'
args '-c', ". ${envdir}/bin/activate && pip install openai && python -m
apache_beam.examples.inference.vllm_text_completion $cmdArgs --chat true
--chat_template
'gs://apache-beam-ml/additional_files/sample_chat_template.jinja'
--experiment='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver:5xx'"
}
+ // Embedded NVIDIA Dynamo path. Reuses the same sdk_container_image
+ // (vllm.dockerfile.old now installs etcd and ai-dynamo[vllm]) and the
+ // same nvidia-tesla-t4 accelerator as the native vLLM ITs above.
+ // Validated end-to-end on Dataflow with Qwen/Qwen3-0.6B on T4.
+ def dynamoArgMap = vllmBaseArgMap + [
+ "model": "Qwen/Qwen3-0.6B",
+ "output": "gs://apache-beam-ml/outputs/vllm_dynamo_predictions.txt",
+ ]
+ def dynamoCmdArgs = mapToArgString(dynamoArgMap)
+ exec {
+ executable 'sh'
+ args '-c', ". ${envdir}/bin/activate && pip install openai && python -m
apache_beam.examples.inference.vllm_text_completion $dynamoCmdArgs --use_dynamo
--max_tokens 8
--experiment='worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver:5xx'"
+ }
}
}
