This is an automated email from the ASF dual-hosted git repository. damccorm pushed a commit to branch users/damccorm/mlContainers in repository https://gitbox.apache.org/repos/asf/beam.git
commit e5fd2eb4a1af1be9b0301f9cc2770b4235b1dfcd Author: Danny Mccormick <[email protected]> AuthorDate: Mon Jun 30 11:10:08 2025 -0400 Add logic for generating ml requirements --- sdks/python/container/common.gradle | 15 ++++++++++++++- sdks/python/container/run_generate_requirements.sh | 16 +++++++++++++--- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/sdks/python/container/common.gradle b/sdks/python/container/common.gradle index 0175778a630..6c3de38d56c 100644 --- a/sdks/python/container/common.gradle +++ b/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}" } } } diff --git a/sdks/python/container/run_generate_requirements.sh b/sdks/python/container/run_generate_requirements.sh index 6c160bc6ac9..b54b0a2020c 100755 --- a/sdks/python/container/run_generate_requirements.sh +++ b/sdks/python/container/run_generate_requirements.sh @@ -38,10 +38,12 @@ fi PY_VERSION=$1 SDK_TARBALL=$2 +REQUIREMENTS_FILE_NAME=$3 +EXTRAS=$4 # Use the PIP_EXTRA_OPTIONS environment variable to pass additional flags to the pip install command. # For example, you can include the --pre flag in $PIP_EXTRA_OPTIONS to download pre-release versions of packages. # Note that you can modify the behavior of the pip install command in this script by passing in your own $PIP_EXTRA_OPTIONS. -PIP_EXTRA_OPTIONS=$3 +PIP_EXTRA_OPTIONS=$5 if ! python"$PY_VERSION" --version > /dev/null 2>&1 ; then echo "Please install a python${PY_VERSION} interpreter. See s.apache.org/beam-python-dev-wiki for Python installation tips." @@ -53,6 +55,14 @@ if ! python"$PY_VERSION" -m venv --help > /dev/null 2>&1 ; then exit 1 fi +if [ -z "$REQUIREMENTS_FILE_NAME" ]; then + REQUIREMENTS_FILE_NAME="base_image_requirements.txt" +fi + +if [ -z "$EXTRAS" ]; then + EXTRAS="[gcp,dataframe,test]" +fi + set -ex ENV_PATH="$PWD/build/python${PY_VERSION/./}_requirements_gen" @@ -65,7 +75,7 @@ pip install --upgrade pip setuptools wheel # Install dataframe deps to add have Dataframe support in released images. # Install test deps since some integration tests need dependencies, # such as pytest, installed in the runner environment. -pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir "$SDK_TARBALL"[gcp,dataframe,test] +pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir "$SDK_TARBALL""$EXTRAS" pip install ${PIP_EXTRA_OPTIONS:+"$PIP_EXTRA_OPTIONS"} --no-cache-dir -r "$PWD"/sdks/python/container/base_image_requirements_manual.txt pip uninstall -y apache-beam @@ -75,7 +85,7 @@ echo "Installed dependencies:" pip freeze PY_IMAGE="py${PY_VERSION//.}" -REQUIREMENTS_FILE=$PWD/sdks/python/container/$PY_IMAGE/base_image_requirements.txt +REQUIREMENTS_FILE=$PWD/sdks/python/container/$PY_IMAGE/$REQUIREMENTS_FILE_NAME cat <<EOT > "$REQUIREMENTS_FILE" # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with
