Use vLLM extra to generate GPU requirements files (#36420)

* Add vllm extra

* Push requirements files

* preinstall torch

* Fix up requirement generation

* Revert "Fix up requirement generation"

This reverts commit 1a5d907.

* install rust

* Highmem

* Exclude 3.13 for now

* CPU torch wheels

* remove 3.9

* Add requirements files

* Clarify comment

Author: damccorm

.github/workflows/update_python_dependencies.yml

@@ -36,7 +36,7 @@ env:
 
 jobs:
   set-properties:
-    runs-on: [self-hosted, ubuntu-20.04]
+    runs-on: [self-hosted, ubuntu-20.04, highmem]
     outputs:
       properties: ${{ steps.test-properties.outputs.properties }}
     steps:
@@ -64,7 +64,6 @@ jobs:
           java-version: default
           go-version: default
           disable-cache: true
-
       - name: Update Python Dependencies
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:

sdks/python/container/common.gradle

@@ -56,6 +56,27 @@ def generatePythonRequirements = tasks.register("generatePythonRequirements") {
               "[gcp,dataframe,test,ml_cpu] " +
               "${pipExtraOptions}"
     }
+    // TODO(https://github.com/apache/beam/issues/36637)
+    // Skip generating Python 3.13 requirements for now since not all 3.13
+    // wheels are available/buildable.
+    // Also skip 3.9 because there are some dependency version conflicts. This
+    // is fine since 3.9 will be EoL by the next release, and we can remove
+    // this condition once we remove support entirely.
+    if ("${project.ext.pythonVersion}" != "3.13" && "${project.ext.pythonVersion}" != "3.9") {
+      // GPU requirements not used for any containers directly due to
+      // licensing, but can be picked up by customers or other consumers for
+      // use.
+      exec {
+        executable 'sh'
+        args '-c', "cd ${rootDir} && ${runScriptsPath} " +
+                "${project.ext.pythonVersion} " +
+                "${files(configurations.sdkSourceTarball.files).singleFile} " +
+                "gpu_image_requirements.txt " +
+                "container/ml " +
+                "[gcp,dataframe,test,tensorflow,torch,transformers,vllm] " +
+                "${pipExtraOptions}"
+      }
+    }
   }
 }
 

sdks/python/container/ml/py310/gpu_image_requirements.txt

@@ -0,0 +1,313 @@
+#    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.
+
+# Autogenerated requirements file for Apache Beam py310 container image.
+# Run ./gradlew :sdks:python:container:generatePythonRequirementsAll to update.
+# Do not edit manually, adjust ../base_image_requirements_manual.txt or
+# Apache Beam's setup.py instead, and regenerate the list.
+# You will need Python interpreters for all versions supported by Beam, see:
+# https://s.apache.org/beam-python-dev-wiki
+# Reach out to a committer if you need help.
+
+absl-py==2.3.1
+aiofiles==25.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.13.1
+aiosignal==1.4.0
+annotated-doc==0.0.3
+annotated-types==0.7.0
+anyio==4.11.0
+asn1crypto==1.5.1
+astor==0.8.1
+astunparse==1.6.3
+async-timeout==5.0.1
+attrs==25.4.0
+backports.tarfile==1.2.0
+beartype==0.21.0
+beautifulsoup4==4.14.2
+blake3==1.0.8
+bs4==0.0.2
+build==1.3.0
+cachetools==6.2.1
+cbor2==5.7.1
+certifi==2025.10.5
+cffi==2.0.0
+charset-normalizer==3.4.4
+click==8.2.1
+cloud-sql-python-connector==1.18.5
+cloudpickle==3.1.1
+compressed-tensors==0.10.2
+crcmod==1.7
+cryptography==46.0.3
+cupy-cuda12x==13.6.0
+Cython==3.1.6
+depyf==0.19.0
+dill==0.3.1.1
+diskcache==5.6.3
+distro==1.9.0
+dnspython==2.8.0
+docker==7.1.0
+docopt==0.6.2
+docstring_parser==0.17.0
+einops==0.8.1
+email-validator==2.3.0
+exceptiongroup==1.3.0
+execnet==2.1.1
+fastapi==0.120.1
+fastapi-cli==0.0.14
+fastapi-cloud-cli==0.3.1
+fastavro==1.12.1
+fasteners==0.20
+fastrlock==0.8.3
+filelock==3.20.0
+flatbuffers==25.9.23
+freezegun==1.5.5
+frozenlist==1.8.0
+fsspec==2025.9.0
+future==1.0.0
+gast==0.6.0
+gguf==0.17.1
+google-api-core==2.28.0
+google-api-python-client==2.185.0
+google-apitools==0.5.31
+google-auth==2.41.1
+google-auth-httplib2==0.2.0
+google-cloud-aiplatform==1.122.0
+google-cloud-bigquery==3.38.0
+google-cloud-bigquery-storage==2.33.1
+google-cloud-bigtable==2.34.0
+google-cloud-core==2.4.3
+google-cloud-datastore==2.21.0
+google-cloud-dlp==3.33.0
+google-cloud-language==2.18.0
+google-cloud-profiler==4.1.0
+google-cloud-pubsub==2.31.1
+google-cloud-pubsublite==1.12.0
+google-cloud-recommendations-ai==0.10.18
+google-cloud-resource-manager==1.15.0
+google-cloud-secret-manager==2.25.0
+google-cloud-spanner==3.59.0
+google-cloud-storage==2.19.0
+google-cloud-videointelligence==2.17.0
+google-cloud-vision==3.11.0
+google-crc32c==1.7.1
+google-genai==1.46.0
+google-pasta==0.2.0
+google-resumable-media==2.7.2
+googleapis-common-protos==1.71.0
+greenlet==3.2.4
+grpc-google-iam-v1==0.14.3
+grpc-interceptor==0.15.4
+grpcio==1.65.5
+grpcio-status==1.65.5
+guppy3==3.1.5
+h11==0.16.0
+h5py==3.15.1
+hdfs==2.7.3
+hf-xet==1.2.0
+httpcore==1.0.9
+httplib2==0.22.0
+httptools==0.7.1
+httpx==0.28.1
+huggingface-hub==0.36.0
+hypothesis==6.142.4
+idna==3.11
+importlib_metadata==8.7.0
+iniconfig==2.3.0
+interegular==0.3.3
+jaraco.classes==3.4.0
+jaraco.context==6.0.1
+jaraco.functools==4.3.0
+jeepney==0.9.0
+Jinja2==3.1.6
+jiter==0.11.1
+joblib==1.5.2
+jsonpickle==3.4.2
+jsonschema==4.25.1
+jsonschema-specifications==2025.9.1
+keras==3.12.0
+keyring==25.6.0
+keyrings.google-artifactregistry-auth==1.1.2
+lark==1.2.2
+libclang==18.1.1
+llguidance==0.7.30
+llvmlite==0.44.0
+lm-format-enforcer==0.10.12
+Markdown==3.9
+markdown-it-py==4.0.0
+MarkupSafe==3.0.3
+mdurl==0.1.2
+milvus-lite==2.5.1
+mistral_common==1.8.5
+ml_dtypes==0.5.3
+mmh3==5.2.0
+mock==5.2.0
+more-itertools==10.8.0
+mpmath==1.3.0
+msgpack==1.1.2
+msgspec==0.19.0
+multidict==6.7.0
+namex==0.1.0
+networkx==3.4.2
+ninja==1.13.0
+nltk==3.9.2
+numba==0.61.2
+numpy==2.2.6
+nvidia-cublas-cu12==12.6.4.1
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cudnn-cu12==9.5.1.17
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cufile-cu12==1.11.1.6
+nvidia-curand-cu12==10.3.7.77
+nvidia-cusolver-cu12==11.7.1.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cusparselt-cu12==0.6.3
+nvidia-nccl-cu12==2.26.2
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-nvshmem-cu12==3.3.20
+nvidia-nvtx-cu12==12.6.77
+oauth2client==4.1.3
+objsize==0.7.1
+openai==1.107.1
+openai-harmony==0.0.4
+opencv-python-headless==4.12.0.88
+opentelemetry-api==1.38.0
+opentelemetry-sdk==1.38.0
+opentelemetry-semantic-conventions==0.59b0
+opt_einsum==3.4.0
+optree==0.17.0
+oracledb==3.4.0
+orjson==3.11.4
+outlines_core==0.2.10
+overrides==7.7.0
+packaging==25.0
+pandas==2.2.3
+parameterized==0.9.0
+partial-json-parser==0.2.1.1.post6
+pg8000==1.31.5
+pillow==12.0.0
+pip==25.3
+pluggy==1.6.0
+prometheus-fastapi-instrumentator==7.1.0
+prometheus_client==0.23.1
+propcache==0.4.1
+proto-plus==1.26.1
+protobuf==5.29.5
+psutil==7.1.2
+psycopg2-binary==2.9.11
+py-cpuinfo==9.0.0
+pyarrow==18.1.0
+pyarrow-hotfix==0.7
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
+pybase64==1.4.2
+pycountry==24.6.1
+pycparser==2.23
+pydantic==2.12.3
+pydantic-extra-types==2.10.6
+pydantic_core==2.41.4
+pydot==1.4.2
+Pygments==2.19.2
+PyHamcrest==2.1.0
+PyJWT==2.10.1
+pymilvus==2.5.16
+pymongo==4.15.3
+PyMySQL==1.1.2
+pyparsing==3.2.5
+pyproject_hooks==1.2.0
+pytest==8.4.2
+pytest-timeout==2.4.0
+pytest-xdist==3.8.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.2.1
+python-json-logger==4.0.0
+python-multipart==0.0.20
+python-tds==1.17.1
+pytz==2025.2
+PyYAML==6.0.3
+pyzmq==27.1.0
+ray==2.50.1
+redis==5.3.1
+referencing==0.37.0
+regex==2025.10.23
+requests==2.32.5
+requests-mock==1.12.1
+rich==14.2.0
+rich-toolkit==0.15.1
+rignore==0.7.1
+rpds-py==0.28.0
+rsa==4.9.1
+safetensors==0.6.2
+scikit-learn==1.7.2
+scipy==1.15.3
+scramp==1.4.6
+SecretStorage==3.4.0
+sentencepiece==0.2.1
+sentry-sdk==2.42.1
+setproctitle==1.3.7
+setuptools==80.9.0
+shapely==2.1.2
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+sortedcontainers==2.4.0
+soundfile==0.13.1
+soupsieve==2.8
+soxr==1.0.0
+SQLAlchemy==2.0.44
+sqlalchemy_pytds==1.0.2
+sqlparse==0.5.3
+starlette==0.49.0
+sympy==1.14.0
+tenacity==8.5.0
+tensorboard==2.20.0
+tensorboard-data-server==0.7.2
+tensorflow==2.20.0
+tensorflow-cpu-aws==2.20.0;platform_machine=="aarch64"
+termcolor==3.2.0
+testcontainers==4.13.2
+threadpoolctl==3.6.0
+tiktoken==0.12.0
+tokenizers==0.21.4
+tomli==2.3.0
+torch==2.7.1
+torchaudio==2.7.1
+torchvision==0.22.1
+tqdm==4.67.1
+transformers==4.55.4
+triton==3.3.1
+typer==0.20.0
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+tzdata==2025.2
+ujson==5.11.0
+uritemplate==4.2.0
+urllib3==2.5.0
+uvicorn==0.38.0
+uvloop==0.22.1
+virtualenv-clone==0.5.7
+vllm==0.10.1.1
+watchfiles==1.1.1
+websockets==15.0.1
+Werkzeug==3.1.3
+wheel==0.45.1
+wrapt==2.0.0
+xformers==0.0.31
+xgrammar==0.1.21
+yarl==1.22.0
+zipp==3.23.0
+zstandard==0.25.0