Merge pull request #7833 from dmiltr3/dev_new_image

fixing docker build and Vertex AI tests

Author: rwitcher

patches/tfx.patch

@@ -92,48 +92,17 @@ index 19489315a..05b464767 100644
      ]
 
 
-@@ -173,7 +149,6 @@ def make_extra_packages_docker_image():
+@@ -173,7 +173,6 @@ def make_extra_packages_docker_image():
          "mmh>=2.2,<3",
          "python-snappy>=0.7",
          # Required for tfx/examples/penguin/penguin_utils_cloud_tuner.py
 -        "tensorflow-cloud>=0.1,<0.2",
          "tensorflow-io>=0.9.0, <=0.24.0",
      ]
 
-@@ -246,7 +221,6 @@ def make_extra_packages_examples():
+@@ -246,5 +245,4 @@ def make_extra_packages_examples():
          # LINT.ThenChange(
          #     examples/penguin/experimental/penguin_pipeline_sklearn_gcp.py)
          # Required for tfx/examples/penguin/penguin_utils_cloud_tuner.py
 -        "tensorflow-cloud>=0.1,<0.2",
      ]
-
-
-diff --git a/tfx/tools/docker/requirements.txt b/tfx/tools/docker/requirements.txt
-index d7b3e960e..7e7862afc 100644
---- a/tfx/tools/docker/requirements.txt
-+++ b/tfx/tools/docker/requirements.txt
-@@ -300,8 +300,6 @@ tenacity==9.0.0
- tensorboard==2.17.1
- tensorboard-data-server==0.7.2
- tensorflow==2.17.1
--tensorflow-cloud==0.1.16
--tensorflow-data-validation==1.17.0
- tensorflow-datasets==4.9.3
- tensorflow-decision-forests==1.10.1
- tensorflow-estimator==2.15.0
-@@ -312,7 +310,6 @@ tensorflow-metadata>=1.17.1
- # tensorflow-ranking==0.5.5
- tensorflow-serving-api==2.17.1
- tensorflow-text==2.17.0
--tensorflow-transform>=1.17.0
- tensorflow_model_analysis>=0.48.0
- tensorflowjs==4.17.0
- tensorstore==0.1.66
-@@ -320,7 +317,6 @@ termcolor==2.5.0
- terminado==0.18.1
- text-unidecode==1.3
- tflite-support==0.4.4
--tfx-bsl>=1.17.1
- threadpoolctl==3.5.0
- time-machine==2.16.0
- tinycss2==1.3.0

tfx/tools/docker/Dockerfile

@@ -17,61 +17,117 @@ ARG BASE_IMAGE
 # Apache beam version to get Beam binaries.
 ARG BEAM_VERSION
 
-FROM gcr.io/tfx-oss-public/tfx_base:py310-20230921 as wheel-builder
-# docker build command should be run under root directory of github checkout.
-ENV TFX_DIR=/tfx
-ADD . ${TFX_DIR}/src
-WORKDIR ${TFX_DIR}
+FROM ${BASE_IMAGE} AS base-with-gcc13
+RUN /opt/conda/bin/conda install -y --override-channels -c conda-forge \
+    gcc_linux-64=13 \
+    gxx_linux-64=13 \
+    binutils_linux-64=2.40 \
+    ld_impl_linux-64=2.40
+
+ENV CC=/opt/conda/bin/x86_64-conda-linux-gnu-gcc
+ENV CXX=/opt/conda/bin/x86_64-conda-linux-gnu-g++
+ENV LD=/opt/conda/bin/x86_64-conda-linux-gnu-ld
+ENV AR=/opt/conda/bin/x86_64-conda-linux-gnu-ar
+ENV NM=/opt/conda/bin/x86_64-conda-linux-gnu-nm
+ENV OBJCOPY=/opt/conda/bin/x86_64-conda-linux-gnu-objcopy
+ENV OBJDUMP=/opt/conda/bin/x86_64-conda-linux-gnu-objdump
+ENV RANLIB=/opt/conda/bin/x86_64-conda-linux-gnu-ranlib
+ENV STRIP=/opt/conda/bin/x86_64-conda-linux-gnu-strip
+
+ENV BAZEL_COMPILER=gcc
+
+FROM base-with-gcc13 AS wheel-builder
+ARG BASE_IMAGE
+ARG BEAM_VERSION
+# Whether to use C++ wheels (tfdv, tfx_bsl) from the host cache (/tmp/wheels).
+# If true, cached wheels are used if available (this saves time when debugging the docker build issues).
+# If false, wheels are built from source even if they exist in the cache.
+ARG USE_CPP_WHEELS_FROM_TEMP=false
+# Whether to clear the persistent build caches (/tmp/wheels and bazel cache).
+ARG CLEAN_CPP_TEMP_CACHE=false
 
 # Specify what version of dependent libraries will be used. See dependencies.py.
 ARG TFX_DEPENDENCY_SELECTOR
 ENV TFX_DEPENDENCY_SELECTOR=${TFX_DEPENDENCY_SELECTOR}
+
 ENV USE_BAZEL_VERSION=6.5.0
 
-RUN python -m pip install --upgrade pip wheel setuptools==70.0.0 tomli
+RUN apt-get update && apt-get install -y curl git && \
+    (find /opt/conda/bin -name "python3-config" | head -n 1 | xargs -I {} ln -sf {} /usr/bin/python-config) && \
+    (find /opt/conda/bin -name "python3-config" | head -n 1 | xargs -I {} ln -sf {} /opt/conda/bin/python-config)
+RUN mkdir -p /usr/local/lib/bazel/bin && \
+    curl -fLO https://github.com/bazelbuild/bazel/releases/download/${USE_BAZEL_VERSION}/bazel-${USE_BAZEL_VERSION}-linux-x86_64 && \
+    mv bazel-${USE_BAZEL_VERSION}-linux-x86_64 /usr/local/lib/bazel/bin/bazel && \
+    chmod +x /usr/local/lib/bazel/bin/bazel && \
+    ln -sf /usr/local/lib/bazel/bin/bazel /usr/local/bin/bazel
+ENV PATH="/usr/local/lib/bazel/bin:${PATH}"
+
+
+COPY . /tfx/src/
+WORKDIR /tfx/src/
+
+# 1. C++ Wheels (tfdv, tfx_bsl) - Normal production build path
+RUN if [ "$USE_CPP_WHEELS_FROM_TEMP" = "false" ]; then \
+      echo "Rebuild of C++ wheels (tfdv, tfx_bsl)..." && \
+      cp tfx/tools/docker/build_tfdv_wheels.sh /tmp/ && \
+      cp tfx/tools/docker/build_tfx_bsl_wheels.sh /tmp/ && \
+      cp tfx/tools/docker/*.patch /tmp/ && \
+      mkdir -p /tfx/src/dist_wheels && \
+      bash /tmp/build_tfdv_wheels.sh /tfx/src/dist_wheels && \
+      bash /tmp/build_tfx_bsl_wheels.sh /tfx/src/dist_wheels ; \
+    fi
+
+# 2. C++ Wheels (tfdv, tfx_bsl) - Cached Path to avoid any CPP rebuilds
+RUN --mount=type=cache,target=/tmp/wheels --mount=type=cache,target=/root/.cache/bazel \
+    if [ "$USE_CPP_WHEELS_FROM_TEMP" = "true" ]; then \
+      echo "Re-use cached build of C++ wheels (tfdv, tfx_bsl)..." && \
+      cp tfx/tools/docker/build_tfdv_wheels.sh /tmp/ && \
+      cp tfx/tools/docker/build_tfx_bsl_wheels.sh /tmp/ && \
+      cp tfx/tools/docker/*.patch /tmp/ && \
+      if [ ! -f /tmp/wheels/tensorflow_data_validation-*.whl ]; then bash /tmp/build_tfdv_wheels.sh /tmp/wheels; fi && \
+      if [ ! -f /tmp/wheels/tfx_bsl-*.whl ]; then bash /tmp/build_tfx_bsl_wheels.sh /tmp/wheels; fi && \
+      mkdir -p /tfx/src/dist_wheels && \
+      cp /tmp/wheels/tensorflow_data_validation-*.whl /tfx/src/dist_wheels/ && \
+      cp /tmp/wheels/tfx_bsl-*.whl /tfx/src/dist_wheels/ ; \
+    fi
 
 # TODO(b/175089240): clean up conditional checks on whether ml-pipelines-sdk is
 # built after TFX versions <= 0.25 are no longer eligible for cherry-picks.
-RUN cd ${TFX_DIR}/src; \
+RUN echo "Building TFX and SDK wheels from source..." && \
+    mkdir -p /tfx/src/dist_wheels && \
     if [ -e "package_build" ]; then \
-      bash -x package_build/initialize.sh; \
-      cd package_build/ml-pipelines-sdk; \
-      CFLAGS=$(/usr/bin/python-config --cflags) \
-        python setup.py bdist_wheel; \
-      cd ../../package_build/tfx; \
-      CFLAGS=$(/usr/bin/python-config --cflags) \
-        python setup.py bdist_wheel; \
-      cd ../..; \
-      MLSDK_WHEEL=$(find dist -name "ml_pipelines_sdk-*.whl"); \
-      TFX_WHEEL=$(find dist -name "tfx-*.whl"); \
+      bash -x package_build/initialize.sh && \
+      CFLAGS=$(python-config --cflags) python package_build/ml-pipelines-sdk/setup.py bdist_wheel && \
+      CFLAGS=$(python-config --cflags) python package_build/tfx/setup.py bdist_wheel && \
+      cp package_build/ml-pipelines-sdk/dist/*.whl /tfx/src/dist_wheels/ && \
+      cp package_build/tfx/dist/*.whl /tfx/src/dist_wheels/ ; \
     else \
-      CFLAGS=$(/usr/bin/python-config --cflags) \
-        python setup.py bdist_wheel; \
-      MLSDK_WHEEL=; \
-      TFX_WHEEL=$(find dist -name "tfx-*.whl"); \
-    fi; \
-    if [ "${TFX_DEPENDENCY_SELECTOR}" = "NIGHTLY" ]; then \
-      CFLAGS=$(/usr/bin/python-config --cflags) \
-        python -m pip install \
-          --extra-index-url https://pypi-nightly.tensorflow.org/simple \
-          ${MLSDK_WHEEL} ${TFX_WHEEL}[docker-image] -c tfx/tools/docker/requirements.txt -c tfx/tools/docker/build_constraints.txt; \
-    else \
-      CFLAGS=$(/usr/bin/python-config --cflags) \
-        python -m pip install ${MLSDK_WHEEL} ${TFX_WHEEL}[docker-image] -c tfx/tools/docker/requirements.txt -c tfx/tools/docker/build_constraints.txt; \
-    fi;
-
-# We need to name this step for the next COPY --from command.
-FROM apache/beam_python3.10_sdk:${BEAM_VERSION} as beam-worker
+      CFLAGS=$(python-config --cflags) python setup.py bdist_wheel && \
+      cp dist/*.whl /tfx/src/dist_wheels/ ; \
+    fi
 
-# Build stage to extend DLVM image.
-FROM ${BASE_IMAGE} as install
+# --- Stage 2: beam-worker ---
+FROM apache/beam_python3.10_sdk:${BEAM_VERSION} AS beam-worker
 
-# Additional Python packages which will be installed.
+# --- Stage 3: install (Final) ---
+FROM base-with-gcc13 AS install
+ARG BASE_IMAGE
 ARG ADDITIONAL_PACKAGES
-
-# Specify what version of dependent libraries will be used. See dependencies.py.
+ENV TF_CPP_MIN_LOG_LEVEL=3
+# Disables Legacy "Virtual" Devices
+ENV TF_XLA_FLAGS="--tf_xla_enable_xla_devices=false"
 ARG TFX_DEPENDENCY_SELECTOR
 ENV TFX_DEPENDENCY_SELECTOR=${TFX_DEPENDENCY_SELECTOR}
+ENV TF_USE_LEGACY_KERAS=1
+
+# 1. Apply OS security updates and install required system libraries
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    ca-certificates \
+    libsnappy-dev \
+    libcurl4-openssl-dev \
+    && apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 # Copy from image matching installed version of 'apache-beam'.
 COPY --from=beam-worker /opt/apache/beam /opt/apache/beam
@@ -84,60 +140,48 @@ ENTRYPOINT ["/opt/apache/beam/boot"]
 
 LABEL maintainer="tensorflow-extended-dev@googlegroups.com"
 
-RUN sed -i 's/python3/python/g' /usr/bin/pip
-
-RUN python -m pip install --upgrade pip setuptools==70.0.0 wheel
-
-RUN CURRENT_DIR="$(pwd)" && \
-  mkdir -p /usr/local/lib/bazel/bin && \
-  cd "/usr/local/lib/bazel/bin" && \
-  curl -fLO https://releases.bazel.build/6.5.0/release/bazel-6.5.0-linux-x86_64 && \
-  chmod +x bazel-6.5.0-linux-x86_64 && \
-  ln -sf /usr/local/lib/bazel/bin/bazel-6.5.0-linux-x86_64 /usr/local/bin/bazel && \
-  cd "$CURRENT_DIR"
-
-# Build tensorflow-data-validation wheels from source with patch
-RUN  /opt/conda/bin/conda install -y -c conda-forge \
-    gcc_linux-64=13 \
-    gxx_linux-64=13 \
-    binutils_linux-64=2.40 \
-    ld_impl_linux-64=2.40
-
- ENV CC=/opt/conda/bin/x86_64-conda-linux-gnu-gcc
-
- ENV CXX=/opt/conda/bin/x86_64-conda-linux-gnu-g++
-
- ENV LD=/opt/conda/bin/x86_64-conda-linux-gnu-ld
-
-COPY tfx/tools/docker/build_tfdv_wheels.sh /tmp/
-COPY tfx/tools/docker/tfdv.patch /tmp/
-RUN mkdir -p /tmp/wheels && \
-    bash /tmp/build_tfdv_wheels.sh /tmp/wheels
-
-# Build tfx-bsl wheels from source with patch
-COPY tfx/tools/docker/build_tfx_bsl_wheels.sh /tmp/
-COPY tfx/tools/docker/tfx_bsl.patch /tmp/
-RUN bash /tmp/build_tfx_bsl_wheels.sh /tmp/wheels
-
-# Copy constraint file and wheels
-COPY tfx/tools/docker/build_constraints.txt /tmp/
-COPY tfx/tools/docker/wheels/*.whl /tmp/wheels/
-
-# Install tensorflow-data-validation, tfx-bsl, tensorflow-model-analysis, tensorflow-transform, and tensorflow-cloud first
-RUN PIP_CONSTRAINT=/tmp/build_constraints.txt python -m pip install /tmp/wheels/tensorflow_data_validation-*.whl /tmp/wheels/tfx_bsl-*.whl /tmp/wheels/tensorflow_model_analysis-*.whl /tmp/wheels/tensorflow_transform-*.whl /tmp/wheels/tensorflow_cloud-*.whl
-
-COPY --from=wheel-builder /tfx/src/dist/*.whl /tfx/src/dist/
-WORKDIR /tfx/src
-
-RUN MLSDK_WHEEL=$(find dist -name "ml_pipelines_sdk-*.whl"); \
-    TFX_WHEEL=$(find dist -name "tfx-*.whl"); \
-    if [ "${TFX_DEPENDENCY_SELECTOR}" = "NIGHTLY" ]; then \
-      PIP_CONSTRAINT=/tmp/build_constraints.txt python -m pip install \
-        --extra-index-url https://pypi-nightly.tensorflow.org/simple \
-        ${MLSDK_WHEEL} ${TFX_WHEEL}[docker-image] ${ADDITIONAL_PACKAGES} ; \
+COPY --from=wheel-builder /tfx/src /tfx/src
+
+# 2. Upgrade core python build tools and remove unused vulnerable components
+# setuptools==78.1.1 is required for the pkg_resources shim (needed by apache-beam),
+# while providing the security fix for CVE-2025-47273.
+RUN python -m pip install --upgrade pip setuptools==78.1.1 wheel \
+    -c /tfx/src/tfx/tools/docker/requirements.txt \
+    -c /tfx/src/tfx/tools/docker/build_constraints.txt
+
+# 3. Main installation: consolidated to a single RUN for single-pass resolution.
+# We explicitly include setuptools==78.1.1 as a top-level requirement here
+# to ensure the resolver doesn't downgrade it or use a broken version
+# that lacks the pkg_resources shim (needed by apache-beam).
+RUN if [ "${TFX_DEPENDENCY_SELECTOR}" = "NIGHTLY" ]; then \
+	  python -m pip install --no-cache-dir \
+            --extra-index-url https://pypi-nightly.tensorflow.org/simple \
+	    -c /tfx/src/tfx/tools/docker/requirements.txt \
+	    -c /tfx/src/tfx/tools/docker/build_constraints.txt \
+	    /tfx/src/dist_wheels/*.whl \
+	    /tfx/src/tfx/tools/docker/wheels/*.whl \
+	    "$(find /tfx/src/dist_wheels/ \( -name 'tfx_dev-*.whl' -o -name 'tfx-*.whl' \) | head -n 1)[docker-image]" \
+	    tf_keras setuptools==78.1.1 ${ADDITIONAL_PACKAGES} ; \
     else \
-      PIP_CONSTRAINT=/tmp/build_constraints.txt python -m pip install ${MLSDK_WHEEL} ${TFX_WHEEL}[docker-image] \
-      ${ADDITIONAL_PACKAGES} ; \
+	  python -m pip install --no-cache-dir \
+	    -c /tfx/src/tfx/tools/docker/requirements.txt \
+	    -c /tfx/src/tfx/tools/docker/build_constraints.txt \
+	    /tfx/src/dist_wheels/*.whl \
+	    /tfx/src/tfx/tools/docker/wheels/*.whl \
+	    "$(find /tfx/src/dist_wheels/ \( -name 'tfx_dev-*.whl' -o -name 'tfx-*.whl' \) | head -n 1)[docker-image]" \
+	    tf_keras setuptools==78.1.1 ${ADDITIONAL_PACKAGES} ; \
     fi && \
-    echo "Installed python packages:\n" && python -m pip list && \
+    (python -m pip uninstall -y jupyter jupyter-server jupyterlab notebook nbconvert jaraco-context jaraco.context || true)
+
+# 4. Final OS cleanup: remove Go toolchain and other unused tools to fix Go-related CVEs
+# Many High/Critical CVEs are in the Go stdlib/toolchain which we don't need at runtime.
+# We use a loop to avoid build failures if a package name is not found in the repo.
+RUN for pkg in golang-go golang git binutils wget policykit-1 packagekit gnupg2 gcc-12; do apt-get purge -y $pkg || echo "Package $pkg not found, skipping"; done && \
+    rm -rf /usr/local/go && \
+    rm -rf /opt/apache/beam && \
+    find /opt/conda/lib/python3.10/site-packages/apache_beam -type f -name "boot" -delete || true && \
+    apt-get autoremove -y && \
+    apt-get clean
+
+RUN echo "Installed python packages:\n" && python -m pip list && \
     echo "Setuptools version:" && python -c "import setuptools; print(setuptools.__version__)"

tfx/tools/docker/Dockerfile.beam_discovery

@@ -0,0 +1,19 @@
+FROM python:3.10-slim@sha256:c299e10e0070171113f9a1f109dd05e7e634fa94589b056e0e87bb22b2b382a2
+
+WORKDIR /src
+
+COPY tfx/dependencies.py .
+COPY tfx/tools/docker/discover_versions.py .
+COPY tfx/tools/docker/requirements.txt tfx/tools/docker/
+COPY tfx/tools/docker/build_constraints.txt tfx/tools/docker/
+
+ENV PIP_CONSTRAINT="/src/tfx/tools/docker/build_constraints.txt"
+
+RUN pip install --upgrade pip setuptools wheel tomli \
+    -c tfx/tools/docker/requirements.txt \
+    -c tfx/tools/docker/build_constraints.txt
+
+# Ensure dependencies.py can be imported
+ENV PYTHONPATH="/src"
+
+ENTRYPOINT ["python3", "discover_versions.py"]

tfx/tools/docker/build_constraints.txt

@@ -1,2 +1,2 @@
-setuptools==70.0.0
-wheel==0.43.0
+setuptools==78.1.1
+wheel==0.45.1