vLLM Gemma Benchmarks (GPU) (#35809)

* add vLLM benchmarks

* update base image version

* apply yapf

* update vLLM workflow

* fix Dockerfile & add sdkLocationOverride flag to gradle

* minor adjustments to vLLM gemma benchmark

* apply rat,yapf,pylint

* rm newline

Author: A1K28

.github/actions/build-push-docker-action/action.yml

@@ -0,0 +1,45 @@
+# This is a composite action to build and push a Docker image.
+name: 'Docker Build and Push'
+description: 'Builds and pushes a Docker image to a container registry.'
+
+inputs:
+  dockerfile_path:
+    description: 'Path to the Dockerfile'
+    required: true
+  image_name:
+    description: 'Base name for the Docker image (e.g., gcr.io/my-project/my-app)'
+    required: true
+  image_tag:
+    description: 'Tag for the Docker image (e.g., latest, or a git sha)'
+    required: true
+  build_context:
+    description: 'The build context for the Docker build command'
+    required: false
+    default: '.'
+
+outputs:
+  image_url:
+    description: "The full URL of the pushed image, including the tag"
+    value: ${{ steps.build-push.outputs.image_url }} # the value is set from a step's output
+
+runs:
+  using: "composite"
+  steps:
+    - name: Configure Docker to use Google Cloud credentials
+      shell: bash
+      run: gcloud auth configure-docker --quiet
+
+    - name: Build and Push Docker Image
+      id: build-push # give the step an ID to reference its output
+      shell: bash
+      run: |
+        # Construct the full image URL from the inputs
+        FULL_IMAGE_URL="${{ inputs.image_name }}:${{ inputs.image_tag }}"
+        echo "Building image: $FULL_IMAGE_URL"
+        
+        # Build the image
+        docker build -t $FULL_IMAGE_URL -f ${{ inputs.dockerfile_path }} ${{ inputs.build_context }}
+        # Push the image
+        docker push $FULL_IMAGE_URL
+        # Set the output value for this action
+        echo "image_url=$FULL_IMAGE_URL" >> $GITHUB_OUTPUT

.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml

@@ -55,7 +55,7 @@ jobs:
       (github.event_name == 'schedule' && github.repository == 'apache/beam') ||
       github.event.comment.body == 'Run Inference Benchmarks'
     runs-on: [self-hosted, ubuntu-20.04, main]
-    timeout-minutes: 900
+    timeout-minutes: 1000
     name: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
     strategy:
       matrix:
@@ -72,7 +72,12 @@ jobs:
       - name: Setup Python environment
         uses: ./.github/actions/setup-environment-action
         with:
+          java-version: default
           python-version: '3.10'
+      - name: Package Python SDK using Gradle
+        run: ./gradlew :sdks:python:sdist -PpythonVersion=3.10
+      - name: Configure Docker for Artifact Registry
+        run: gcloud auth configure-docker us-docker.pkg.dev
       - name: Prepare test arguments
         uses: ./.github/actions/test-arguments-action
         with:
@@ -86,9 +91,28 @@ jobs:
             ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Imagenet_Classification_Resnet_152_Tesla_T4_GPU.txt
             ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Streaming_DistilBert_Base_Uncased.txt
             ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Batch_DistilBert_Base_Uncased.txt
+            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_VLLM_Gemma_Batch.txt
       # The env variables are created and populated in the test-arguments-action as "<github.job>_test_arguments_<argument_file_paths_index>"
       - name: get current time
         run: echo "NOW_UTC=$(date '+%m%d%H%M%S' --utc)" >> $GITHUB_ENV
+      - name: Build VLLM Development Image
+        id: build_vllm_image
+        uses: ./.github/actions/build-push-docker-action
+        with:
+          dockerfile_path: 'sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile'
+          image_name: 'us-docker.pkg.dev/apache-beam-testing/beam-temp/beam-vllm-gpu-base'
+          image_tag: ${{ github.sha }}
+      - name: Run VLLM Gemma Batch Test
+        uses: ./.github/actions/gradle-command-self-hosted-action
+        timeout-minutes: 180
+        with:
+          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
+          arguments: |
+            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.vllm_gemma_benchmarks \
+            -Prunner=DataflowRunner \
+            -PsdkLocationOverride=false \
+            -PpythonVersion=3.10 \
+            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/vllm_tests_requirements.txt '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_8 }} --mode=batch --job_name=benchmark-tests-vllm-with-gemma-2b-it-batch-${{env.NOW_UTC}} --sdk_container_image=${{ steps.build_vllm_image.outputs.image_url }}'
       - name: run Pytorch Sentiment Streaming using Hugging Face distilbert-base-uncased model
         uses: ./.github/actions/gradle-command-self-hosted-action
         timeout-minutes: 180

.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_VLLM_Gemma_Batch.txt

@@ -0,0 +1,36 @@
+# 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.
+
+--runner=DataflowRunner
+--region=us-central1
+--temp_location=gs://temp-storage-for-perf-tests/loadtests
+--staging_location=gs://temp-storage-for-perf-tests/loadtests
+--input=gs://apache-beam-ml/testing/inputs/sentences_50k.txt
+--machine_type=n1-standard-8
+--worker_zone=us-central1-b
+--disk_size_gb=50
+--input_options={}
+--num_workers=8
+--max_num_workers=25
+--autoscaling_algorithm=THROUGHPUT_BASED
+--publish_to_big_query=true
+--sdk_location=container
+--output_table=apache-beam-testing.beam_run_inference.result_gemma_vllm_batch
+--metrics_dataset=beam_run_inference
+--metrics_table=gemma_vllm_batch
+--influx_measurement=gemma_vllm_batch
+--model_gcs_path=gs://apache-beam-ml/models/gemma-2b-it
+--dataflow_service_options=worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver
+--experiments=use_runner_v2

.test-infra/tools/refresh_looker_metrics.py

@@ -42,6 +42,7 @@
     ("80", ["253", "254", "255", "256", "257"]),  # PyTorch Resnet 152 Tesla T4
     ("82", ["263", "264", "265", "266", "267"]),  # PyTorch Sentiment Streaming DistilBERT base uncased
     ("85", ["268", "269", "270", "271", "272"]),  # PyTorch Sentiment Batch DistilBERT base uncased
+    ("86", ["284", "285", "286", "287", "288"]),  # VLLM Batch Gemma
 ]
 
 

sdks/python/apache_beam/examples/inference/vllm_gemma_batch.py

@@ -0,0 +1,130 @@
+#
+# 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.
+#
+
+from __future__ import annotations
+
+import logging
+import os
+import tempfile
+
+import apache_beam as beam
+from apache_beam.io.filesystems import FileSystems
+from apache_beam.ml.inference.base import RunInference
+from apache_beam.ml.inference.vllm_inference import VLLMCompletionsModelHandler
+from apache_beam.ml.inference.vllm_inference import _VLLMModelServer
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import SetupOptions
+
+
+class GemmaVLLMOptions(PipelineOptions):
+  """Custom pipeline options for the Gemma vLLM batch inference job."""
+  @classmethod
+  def _add_argparse_args(cls, parser):
+    parser.add_argument(
+        "--input",
+        dest="input_file",
+        required=True,
+        help="Input file gs://path containing prompts.",
+    )
+    parser.add_argument(
+        "--output_table",
+        required=True,
+        help="BigQuery table to write to in the form project:dataset.table.",
+    )
+    parser.add_argument(
+        "--model_gcs_path",
+        required=True,
+        help="GCS path to the directory containing model files.",
+    )
+
+
+class FormatOutput(beam.DoFn):
+  def process(self, element):
+    prompt = element.example
+    comp = element.inference
+
+    if hasattr(comp, 'choices'):
+      completion = comp.choices[0].text
+    # fallback to a single .text field
+    elif hasattr(comp, 'text'):
+      completion = comp.text
+    # final fallback
+    else:
+      completion = str(comp)
+
+    yield {'prompt': prompt, 'completion': completion}
+
+
+class GcsVLLMCompletionsModelHandler(VLLMCompletionsModelHandler):
+  def __init__(self, model_name, vllm_server_kwargs=None):
+    super().__init__(model_name, vllm_server_kwargs)
+    self._local_model_dir = None
+
+  def _download_gcs_directory(self, gcs_path: str, local_path: str):
+    logging.info("Downloading model from %s to %s…", gcs_path, local_path)
+    matches = FileSystems.match([os.path.join(gcs_path, "**")])[0].metadata_list
+    for md in matches:
+      rel = os.path.relpath(md.path, gcs_path)
+      dst = os.path.join(local_path, rel)
+      os.makedirs(os.path.dirname(dst), exist_ok=True)
+      with FileSystems.open(md.path) as src, open(dst, "wb") as dstf:
+        dstf.write(src.read())
+    logging.info("Download complete.")
+
+  def load_model(self) -> _VLLMModelServer:
+    uri = self._model_name
+    if uri.startswith("gs://"):
+      self._local_model_dir = tempfile.mkdtemp(prefix="vllm_model_")
+      self._download_gcs_directory(uri, self._local_model_dir)
+      logging.info("Loading vLLM from local dir %s", self._local_model_dir)
+      return _VLLMModelServer(self._local_model_dir, self._vllm_server_kwargs)
+    else:
+      logging.info("Loading vLLM from HF hub: %s", uri)
+      return super().load_model()
+
+
+def run(argv=None, save_main_session=True, test_pipeline=None):
+  # Build pipeline options
+  opts = PipelineOptions(argv)
+
+  gem = opts.view_as(GemmaVLLMOptions)
+  opts.view_as(SetupOptions).save_main_session = save_main_session
+
+  logging.info("Pipeline starting with model path: %s", gem.model_gcs_path)
+  handler = GcsVLLMCompletionsModelHandler(
+      model_name=gem.model_gcs_path,
+      vllm_server_kwargs={"served-model-name": gem.model_gcs_path})
+
+  with (test_pipeline or beam.Pipeline(options=opts)) as p:
+    _ = (
+        p
+        | "Read" >> beam.io.ReadFromText(gem.input_file)
+        | "InferBatch" >> RunInference(handler, inference_batch_size=32)
+        | "FormatForBQ" >> beam.ParDo(FormatOutput())
+        | "WriteToBQ" >> beam.io.WriteToBigQuery(
+            gem.output_table,
+            schema="prompt:STRING,completion:STRING",
+            write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
+            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
+            method=beam.io.WriteToBigQuery.Method.FILE_LOADS,
+        ))
+  return p.result
+
+
+if __name__ == "__main__":
+  logging.getLogger().setLevel(logging.INFO)
+  run()

sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile

@@ -15,33 +15,54 @@
 # limitations under the License.
 
 # Used for any vLLM integration test
+# Dockerfile — Beam dev harness + install dev SDK from LOCAL source package
 
 FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
 
-RUN apt update
-RUN apt install software-properties-common -y
-RUN add-apt-repository ppa:deadsnakes/ppa
-RUN apt update
+# 1) Non-interactive + timezone
+ENV DEBIAN_FRONTEND=noninteractive \
+    TZ=Etc/UTC
 
-ARG DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+      curl \
+      tzdata \
+      software-properties-common \
+      python3.10-full \
+      python3.10-distutils \
+      build-essential \
+      python3.10-dev \
+      cython3 && \
+    ln -fs /usr/share/zoneinfo/$TZ /etc/localtime && \
+    dpkg-reconfigure --frontend noninteractive tzdata && \
+    rm -rf /var/lib/apt/lists/*
 
-RUN apt install python3.12 -y
-RUN apt install python3.12-venv -y
-RUN apt install python3.12-dev -y
-RUN rm /usr/bin/python3
-RUN ln -s python3.12 /usr/bin/python3
-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
+# 2) Symlink python3 to 3.10
+RUN ln -sf /usr/bin/python3.10 /usr/bin/python3 && \
+    ln -sf /usr/bin/python3.10 /usr/bin/python
 
-RUN pip install --no-cache-dir -vvv apache-beam[gcp]==2.58.1
-RUN pip install openai vllm
+# 3) Install pip, setuptools & wheel
+RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
+    python3 -m pip install --upgrade pip setuptools wheel
 
-RUN apt install libcairo2-dev pkg-config python3-dev -y
-RUN pip install pycairo
+# 4) Copy the Beam SDK harness (for Dataflow workers)
+COPY --from=gcr.io/apache-beam-testing/beam-sdk/beam_python3.10_sdk:2.68.0.dev \
+     /opt/apache/beam /opt/apache/beam
 
-# 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
+# 5) Make sure the harness is discovered first
+ENV PYTHONPATH=/opt/apache/beam:$PYTHONPATH
 
-# Set the entrypoint to Apache Beam SDK worker launcher.
-ENTRYPOINT [ "/opt/apache/beam/boot" ]
+# 6) Install the Beam dev SDK from the local source package.
+# This .tar.gz file will be created by GitHub Actions workflow
+# and copied into the build context.
+COPY ./sdks/python/build/apache-beam.tar.gz /tmp/beam.tar.gz
+RUN python3 -m pip install --no-cache-dir "/tmp/beam.tar.gz[gcp]"
+
+# 7) Install vLLM, and other dependencies
+RUN python3 -m pip install --no-cache-dir \
+      openai>=1.52.2 \
+      vllm>=0.6.3 \
+      triton>=3.1.0
+
+# 8) Use the Beam boot script as entrypoint
+ENTRYPOINT ["/opt/apache/beam/boot"]