(beam) 01/01: Add vllm extra

[da...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

damccorm pushed a commit to branch users/damccorm/vllmExtra
in repository https://gitbox.apache.org/repos/asf/beam.git

commit ae8245b4be8c4102686b4cd7292fbf3e90de7ff2
Author: Danny Mccormick <da...@google.com>
AuthorDate: Tue Oct 7 10:09:39 2025 -0400

    Add vllm extra
---
 .../apache_beam/ml/inference/test_resources/vllm.dockerfile    | 10 ++--------
 sdks/python/setup.py                                           |  3 ++-
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile b/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile
index 5727437809c..389ae1fe7d6 100644
--- a/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile
+++ b/sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile
@@ -54,15 +54,9 @@ ENV PYTHONPATH=/opt/apache/beam:$PYTHONPATH
 
 # 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.
+# and copied into the build context. This will include vLLM dependencies
 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
+RUN python3 -m pip install --no-cache-dir "/tmp/beam.tar.gz[gcp,vllm]"
 
 # 8) Use the Beam boot script as entrypoint
 ENTRYPOINT ["/opt/apache/beam/boot"]
\ No newline at end of file
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 719d188ed26..14e95f7ccc2 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -610,7 +610,8 @@ if __name__ == '__main__':
           ],
           'xgboost': ['xgboost>=1.6.0,<2.1.3', 'datatable==1.0.0'],
           'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0'],
-          'milvus': milvus_dependency
+          'milvus': milvus_dependency,
+          'vllm': ['openai>=1.52.2', 'vllm>=0.6.3', 'triton>=3.1.0']
       },
       zip_safe=False,
       # PyPI package information.