Install and build task-sdk in prod images

Since the scheduler needs it (at least for now), we need to ensure that we include it in our prod-image builds in CI
apache · Oct 30, 2024 · a94684a · a94684a
1 parent 48f3f66
commit a94684a
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 38 deletions.
diff --git a/.github/workflows/prod-image-build.yml b/.github/workflows/prod-image-build.yml
@@ -181,6 +181,11 @@ jobs:
         run: >
           breeze release-management prepare-airflow-package --package-format wheel
         if: inputs.do-build  == 'true' && inputs.upload-package-artifact == 'true'
+      - name: "Prepare task-sdk package"
+        shell: bash
+        run: >
+          breeze release-management prepare-task-sdk-package --package-format wheel
+        if: inputs.do-build  == 'true' && inputs.upload-package-artifact == 'true'
       - name: "Upload prepared packages as artifacts"
         uses: actions/upload-artifact@v4
         with:

diff --git a/Dockerfile b/Dockerfile
@@ -718,6 +718,7 @@ COPY <<"EOF" /install_from_docker_context_files.sh
 
 
 function install_airflow_and_providers_from_docker_context_files(){
+    local flags=()
     if [[ ${INSTALL_MYSQL_CLIENT} != "true" ]]; then
         AIRFLOW_EXTRAS=${AIRFLOW_EXTRAS/mysql,}
     fi
@@ -756,10 +757,10 @@ function install_airflow_and_providers_from_docker_context_files(){
         install_airflow_package=("apache-airflow[${AIRFLOW_EXTRAS}]==${AIRFLOW_VERSION}")
     fi
 
-    # Find Provider packages in docker-context files
-    readarray -t installing_providers_packages< <(python /scripts/docker/get_package_specs.py /docker-context-files/apache?airflow?providers*.{whl,tar.gz} 2>/dev/null || true)
+    # Find Provider/TaskSDK packages in docker-context files
+    readarray -t airflow_packages< <(python /scripts/docker/get_package_specs.py /docker-context-files/apache?airflow?{providers,task?sdk}*.{whl,tar.gz} 2>/dev/null || true)
     echo
-    echo "${COLOR_BLUE}Found provider packages in docker-context-files folder: ${installing_providers_packages[*]}${COLOR_RESET}"
+    echo "${COLOR_BLUE}Found provider packages in docker-context-files folder: ${airflow_packages[*]}${COLOR_RESET}"
     echo
 
     if [[ ${USE_CONSTRAINTS_FOR_CONTEXT_PACKAGES=} == "true" ]]; then
@@ -772,11 +773,7 @@ function install_airflow_and_providers_from_docker_context_files(){
             echo "${COLOR_BLUE}Installing docker-context-files packages with constraints found in ${local_constraints_file}${COLOR_RESET}"
             echo
             # force reinstall all airflow + provider packages with constraints found in
-            set -x
-            ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade \
-                ${ADDITIONAL_PIP_INSTALL_FLAGS} --constraint "${local_constraints_file}" \
-                "${install_airflow_package[@]}" "${installing_providers_packages[@]}"
-            set +x
+            flags=(--upgrade --constraint "${local_constraints_file}")
             echo
             echo "${COLOR_BLUE}Copying ${local_constraints_file} to ${HOME}/constraints.txt${COLOR_RESET}"
             echo
@@ -785,23 +782,21 @@ function install_airflow_and_providers_from_docker_context_files(){
             echo
             echo "${COLOR_BLUE}Installing docker-context-files packages with constraints from GitHub${COLOR_RESET}"
             echo
-            set -x
-            ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} \
-                ${ADDITIONAL_PIP_INSTALL_FLAGS} \
-                --constraint "${HOME}/constraints.txt" \
-                "${install_airflow_package[@]}" "${installing_providers_packages[@]}"
-            set +x
+            flags=(--constraint "${HOME}/constraints.txt")
         fi
     else
         echo
         echo "${COLOR_BLUE}Installing docker-context-files packages without constraints${COLOR_RESET}"
         echo
-        set -x
-        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} \
-            ${ADDITIONAL_PIP_INSTALL_FLAGS} \
-            "${install_airflow_package[@]}" "${installing_providers_packages[@]}"
-        set +x
+        flags=()
     fi
+
+    set -x
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} \
+        ${ADDITIONAL_PIP_INSTALL_FLAGS} \
+        "${flags[@]}" \
+        "${install_airflow_package[@]}" "${airflow_packages[@]}"
+    set +x
     common::install_packaging_tools
     pip check
 }

diff --git a/scripts/docker/install_from_docker_context_files.sh b/scripts/docker/install_from_docker_context_files.sh
@@ -27,6 +27,7 @@
 # TODO: rewrite it all in Python (and all other scripts in scripts/docker)
 
 function install_airflow_and_providers_from_docker_context_files(){
+    local flags=()
     if [[ ${INSTALL_MYSQL_CLIENT} != "true" ]]; then
         AIRFLOW_EXTRAS=${AIRFLOW_EXTRAS/mysql,}
     fi
@@ -65,10 +66,10 @@ function install_airflow_and_providers_from_docker_context_files(){
         install_airflow_package=("apache-airflow[${AIRFLOW_EXTRAS}]==${AIRFLOW_VERSION}")
     fi
 
-    # Find Provider packages in docker-context files
-    readarray -t installing_providers_packages< <(python /scripts/docker/get_package_specs.py /docker-context-files/apache?airflow?providers*.{whl,tar.gz} 2>/dev/null || true)
+    # Find Provider/TaskSDK packages in docker-context files
+    readarray -t airflow_packages< <(python /scripts/docker/get_package_specs.py /docker-context-files/apache?airflow?{providers,task?sdk}*.{whl,tar.gz} 2>/dev/null || true)
     echo
-    echo "${COLOR_BLUE}Found provider packages in docker-context-files folder: ${installing_providers_packages[*]}${COLOR_RESET}"
+    echo "${COLOR_BLUE}Found provider packages in docker-context-files folder: ${airflow_packages[*]}${COLOR_RESET}"
     echo
 
     if [[ ${USE_CONSTRAINTS_FOR_CONTEXT_PACKAGES=} == "true" ]]; then
@@ -81,11 +82,7 @@ function install_airflow_and_providers_from_docker_context_files(){
             echo "${COLOR_BLUE}Installing docker-context-files packages with constraints found in ${local_constraints_file}${COLOR_RESET}"
             echo
             # force reinstall all airflow + provider packages with constraints found in
-            set -x
-            ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade \
-                ${ADDITIONAL_PIP_INSTALL_FLAGS} --constraint "${local_constraints_file}" \
-                "${install_airflow_package[@]}" "${installing_providers_packages[@]}"
-            set +x
+            flags=(--upgrade --constraint "${local_constraints_file}")
             echo
             echo "${COLOR_BLUE}Copying ${local_constraints_file} to ${HOME}/constraints.txt${COLOR_RESET}"
             echo
@@ -94,23 +91,21 @@ function install_airflow_and_providers_from_docker_context_files(){
             echo
             echo "${COLOR_BLUE}Installing docker-context-files packages with constraints from GitHub${COLOR_RESET}"
             echo
-            set -x
-            ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} \
-                ${ADDITIONAL_PIP_INSTALL_FLAGS} \
-                --constraint "${HOME}/constraints.txt" \
-                "${install_airflow_package[@]}" "${installing_providers_packages[@]}"
-            set +x
+            flags=(--constraint "${HOME}/constraints.txt")
         fi
     else
         echo
         echo "${COLOR_BLUE}Installing docker-context-files packages without constraints${COLOR_RESET}"
         echo
-        set -x
-        ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} \
-            ${ADDITIONAL_PIP_INSTALL_FLAGS} \
-            "${install_airflow_package[@]}" "${installing_providers_packages[@]}"
-        set +x
+        flags=()
     fi
+
+    set -x
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} \
+        ${ADDITIONAL_PIP_INSTALL_FLAGS} \
+        "${flags[@]}" \
+        "${install_airflow_package[@]}" "${airflow_packages[@]}"
+    set +x
     common::install_packaging_tools
     pip check
 }