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 <damccorm@google.com>

* 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 <damccorm@google.com>

Author: akshayjadiyanv

.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
 }

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,22 +192,26 @@ 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
   if not test_pipeline:
     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(

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" ]