Add logic for generating ml requirements (#35484)

* Add logic for generating ml requirements

* Add some requirements and loosen torch bounds

* Widen bounds and generate 3.12 containers

* posargs

* Tox config fixes

Author: damccorm

.github/trigger_files/beam_PostCommit_Python_Dependency.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run",
-    "modification": 0
+    "modification": 1
   }

.github/workflows/beam_PostCommit_Python_Dependency.yml

@@ -59,6 +59,7 @@ jobs:
       matrix:
         job_name: [beam_PostCommit_Python_Dependency]
         job_phrase: [Run Python PostCommit Dependency]
+        python_version: ['39','312']
     timeout-minutes: 120
     if: |
       github.event_name == 'workflow_dispatch' ||
@@ -70,9 +71,9 @@ jobs:
       - name: Setup repository
         uses: ./.github/actions/setup-action
         with:
-          comment_phrase: ${{ matrix.job_phrase }}
+          comment_phrase: ${{ matrix.job_phrase }} ${{ matrix.python_version }}
           github_token: ${{ secrets.GITHUB_TOKEN }}
-          github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
+          github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }} ${{ matrix.python_version }})
       - name: Setup environment
         uses: ./.github/actions/setup-environment-action
         with:
@@ -81,7 +82,7 @@ jobs:
       - name: Run postCommitPyDep
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
-          gradle-command: :sdks:python:test-suites:tox:py39:postCommitPyDep
+          gradle-command: :sdks:python:test-suites:tox:py${{ matrix.python_version }}:postCommitPyDep
           arguments: -PuseWheelDistribution
       - name: Archive Python Test Results
         uses: actions/upload-artifact@v4

sdks/python/container/common.gradle

@@ -39,7 +39,20 @@ def generatePythonRequirements = tasks.register("generatePythonRequirements") {
       executable 'sh'
       args '-c', "cd ${rootDir} && ${runScriptsPath} " +
               "${project.ext.pythonVersion} " +
-              "${files(configurations.sdkSourceTarball.files).singleFile} " + "${pipExtraOptions}"
+              "${files(configurations.sdkSourceTarball.files).singleFile} " +
+              "base_image_requirements.txt " +
+              "[gcp,dataframe,test] " +
+              "${pipExtraOptions}"
+    }
+    // Generate versions for ML dependencies
+    exec {
+      executable 'sh'
+      args '-c', "cd ${rootDir} && ${runScriptsPath} " +
+              "${project.ext.pythonVersion} " +
+              "${files(configurations.sdkSourceTarball.files).singleFile} " +
+              "ml_image_requirements.txt " +
+              "[gcp,dataframe,test,tensorflow,torch,transformers] " +
+              "${pipExtraOptions}"
     }
   }
 }

sdks/python/container/py310/base_image_requirements.txt

@@ -92,7 +92,7 @@ hdfs==2.7.3
 httpcore==1.0.9
 httplib2==0.22.0
 httpx==0.28.1
-hypothesis==6.135.17
+hypothesis==6.135.19
 idna==3.10
 importlib_metadata==8.7.0
 iniconfig==2.1.0
@@ -112,7 +112,7 @@ milvus-lite==2.5.1
 mmh3==5.1.0
 mock==5.2.0
 more-itertools==10.7.0
-multidict==6.6.2
+multidict==6.6.3
 mysql-connector-python==9.3.0
 nltk==3.9.1
 numpy==2.2.6

sdks/python/container/py310/ml_image_requirements.txt

@@ -0,0 +1,247 @@
+#    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.0
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.13
+aiosignal==1.3.2
+annotated-types==0.7.0
+anyio==4.9.0
+asn1crypto==1.5.1
+astunparse==1.6.3
+async-timeout==5.0.1
+attrs==25.3.0
+backports.tarfile==1.2.0
+beautifulsoup4==4.13.4
+bs4==0.0.2
+build==1.2.2.post1
+cachetools==5.5.2
+certifi==2025.6.15
+cffi==1.17.1
+charset-normalizer==3.4.2
+click==8.2.1
+cloud-sql-python-connector==1.18.2
+crcmod==1.7
+cryptography==45.0.4
+Cython==3.1.2
+dill==0.3.1.1
+dnspython==2.7.0
+docker==7.1.0
+docopt==0.6.2
+docstring_parser==0.16
+exceptiongroup==1.3.0
+execnet==2.1.1
+fastavro==1.11.1
+fasteners==0.19
+filelock==3.18.0
+flatbuffers==25.2.10
+freezegun==1.5.2
+frozenlist==1.7.0
+fsspec==2025.5.1
+future==1.0.0
+gast==0.6.0
+google-api-core==2.25.1
+google-api-python-client==2.174.0
+google-apitools==0.5.31
+google-auth==2.40.3
+google-auth-httplib2==0.2.0
+google-cloud-aiplatform==1.100.0
+google-cloud-bigquery==3.34.0
+google-cloud-bigquery-storage==2.32.0
+google-cloud-bigtable==2.31.0
+google-cloud-core==2.4.3
+google-cloud-datastore==2.21.0
+google-cloud-dlp==3.31.0
+google-cloud-language==2.17.2
+google-cloud-profiler==4.1.0
+google-cloud-pubsub==2.30.0
+google-cloud-pubsublite==1.12.0
+google-cloud-recommendations-ai==0.10.18
+google-cloud-resource-manager==1.14.2
+google-cloud-spanner==3.55.0
+google-cloud-storage==2.19.0
+google-cloud-videointelligence==2.16.2
+google-cloud-vision==3.10.2
+google-crc32c==1.7.1
+google-genai==1.23.0
+google-pasta==0.2.0
+google-resumable-media==2.7.2
+googleapis-common-protos==1.70.0
+greenlet==3.2.3
+grpc-google-iam-v1==0.14.2
+grpc-interceptor==0.15.4
+grpcio==1.65.5
+grpcio-status==1.63.0rc1
+guppy3==3.1.5
+h11==0.16.0
+h5py==3.14.0
+hdfs==2.7.3
+hf-xet==1.1.5
+httpcore==1.0.9
+httplib2==0.22.0
+httpx==0.28.1
+huggingface-hub==0.33.1
+hypothesis==6.135.19
+idna==3.10
+importlib_metadata==8.7.0
+iniconfig==2.1.0
+jaraco.classes==3.4.0
+jaraco.context==6.0.1
+jaraco.functools==4.2.1
+jeepney==0.9.0
+Jinja2==3.1.6
+joblib==1.5.1
+jsonpickle==3.4.2
+jsonschema==4.24.0
+jsonschema-specifications==2025.4.1
+keras==3.10.0
+keyring==25.6.0
+keyrings.google-artifactregistry-auth==1.1.2
+libclang==18.1.1
+Markdown==3.8.2
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+mdurl==0.1.2
+milvus-lite==2.5.1
+ml-dtypes==0.3.2
+mmh3==5.1.0
+mock==5.2.0
+more-itertools==10.7.0
+mpmath==1.3.0
+multidict==6.6.3
+mysql-connector-python==9.3.0
+namex==0.1.0
+networkx==3.4.2
+nltk==3.9.1
+numpy==1.26.4
+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-nvtx-cu12==12.6.77
+oauth2client==4.1.3
+objsize==0.7.1
+opentelemetry-api==1.34.1
+opentelemetry-sdk==1.34.1
+opentelemetry-semantic-conventions==0.55b1
+opt_einsum==3.4.0
+optree==0.16.0
+oracledb==3.2.0
+orjson==3.10.18
+overrides==7.7.0
+packaging==25.0
+pandas==2.2.3
+parameterized==0.9.0
+pg8000==1.31.2
+pip==25.1.1
+pluggy==1.6.0
+propcache==0.3.2
+proto-plus==1.26.1
+protobuf==4.25.8
+psycopg2-binary==2.9.10
+pyarrow==18.1.0
+pyarrow-hotfix==0.7
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
+pycparser==2.22
+pydantic==2.11.7
+pydantic_core==2.33.2
+pydot==1.4.2
+Pygments==2.19.2
+PyHamcrest==2.1.0
+PyJWT==2.9.0
+pymilvus==2.5.11
+pymongo==4.13.2
+PyMySQL==1.1.1
+pyparsing==3.2.3
+pyproject_hooks==1.2.0
+pytest==7.4.4
+pytest-timeout==2.4.0
+pytest-xdist==3.7.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.1.1
+python-tds==1.16.1
+pytz==2025.2
+PyYAML==6.0.2
+redis==5.3.0
+referencing==0.36.2
+regex==2024.11.6
+requests==2.32.4
+requests-mock==1.12.1
+rich==14.0.0
+rpds-py==0.25.1
+rsa==4.9.1
+safetensors==0.5.3
+scikit-learn==1.7.0
+scipy==1.15.3
+scramp==1.4.5
+SecretStorage==3.3.3
+setuptools==80.9.0
+shapely==2.1.1
+six==1.17.0
+sniffio==1.3.1
+sortedcontainers==2.4.0
+soupsieve==2.7
+SQLAlchemy==2.0.41
+sqlalchemy_pytds==1.0.2
+sqlparse==0.5.3
+sympy==1.14.0
+tenacity==8.5.0
+tensorboard==2.16.2
+tensorboard-data-server==0.7.2
+tensorflow==2.16.2
+tensorflow-cpu-aws==2.16.2;platform_machine=="aarch64"
+tensorflow-io-gcs-filesystem==0.37.1
+termcolor==3.1.0
+testcontainers==4.10.0
+threadpoolctl==3.6.0
+tokenizers==0.21.2
+tomli==2.2.1
+torch==2.7.1
+tqdm==4.67.1
+transformers==4.48.3
+triton==3.3.1
+typing-inspection==0.4.1
+typing_extensions==4.14.0
+tzdata==2025.2
+ujson==5.10.0
+uritemplate==4.2.0
+urllib3==2.5.0
+virtualenv-clone==0.5.7
+websockets==15.0.1
+Werkzeug==3.1.3
+wheel==0.45.1
+wrapt==1.17.2
+yarl==1.20.1
+zipp==3.23.0
+zstandard==0.23.0

sdks/python/container/py311/base_image_requirements.txt

@@ -90,7 +90,7 @@ hdfs==2.7.3
 httpcore==1.0.9
 httplib2==0.22.0
 httpx==0.28.1
-hypothesis==6.135.17
+hypothesis==6.135.19
 idna==3.10
 importlib_metadata==8.7.0
 iniconfig==2.1.0
@@ -110,7 +110,7 @@ milvus-lite==2.5.1
 mmh3==5.1.0
 mock==5.2.0
 more-itertools==10.7.0
-multidict==6.6.2
+multidict==6.6.3
 mysql-connector-python==9.3.0
 nltk==3.9.1
 numpy==2.2.6