Transition OSS Python dependencies and CI builds to hermetic Python.

WORKSPACE

@@ -81,15 +81,6 @@ http_archive(
     ],
 )
 
-http_archive(
-    name = "org_tensorflow_datasets",
-    sha256 = "c6ff4e2306387f0ca45d4f616d9a1c5e79e02ef16d0a8958230a8049ea07fc98",
-    strip_prefix = "datasets-3.2.0",
-    urls = [
-        "https://github.com/tensorflow/datasets/archive/v3.2.0.zip",
-    ],
-)
-
 http_archive(
     name = "pybind11",
     build_file = "//third_party/pybind11:BUILD.bzl",
@@ -140,7 +131,6 @@ py_deps_profile(
     deps_map = {
         "tensorflow": [
             "tf-nightly",
-            "tf_headers",
             "tf_header_lib",
             "libtensorflow_framework",
         ],
@@ -190,6 +180,17 @@ load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
 
 tf_workspace0()
 
+# Toolchains for ML projects hermetic builds.
+# Details: https://github.com/google-ml-infra/rules_ml_toolchain
+http_archive(
+    name = "rules_ml_toolchain",
+    sha256 = "54c1a357f71f611efdb4891ebd4bcbe4aeb6dfa7e473f14fd7ecad5062096616",
+    strip_prefix = "rules_ml_toolchain-d8cb9c2c168cd64000eaa6eda0781a9615a26ffe",
+    urls = [
+        "https://github.com/google-ml-infra/rules_ml_toolchain/archive/d8cb9c2c168cd64000eaa6eda0781a9615a26ffe.tar.gz",
+    ],
+)
+
 # Set up Android.
 load("@org_tensorflow//third_party/android:android_configure.bzl", "android_configure")
 
@@ -206,17 +207,6 @@ load(
 
 python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
 
-# Toolchains for ML projects hermetic builds.
-# Details: https://github.com/google-ml-infra/rules_ml_toolchain
-http_archive(
-    name = "rules_ml_toolchain",
-    sha256 = "54c1a357f71f611efdb4891ebd4bcbe4aeb6dfa7e473f14fd7ecad5062096616",
-    strip_prefix = "rules_ml_toolchain-d8cb9c2c168cd64000eaa6eda0781a9615a26ffe",
-    urls = [
-        "https://github.com/google-ml-infra/rules_ml_toolchain/archive/d8cb9c2c168cd64000eaa6eda0781a9615a26ffe.tar.gz",
-    ],
-)
-
 load(
     "@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
     "cc_toolchain_deps",

docs/tutorials/BUILD

@@ -24,7 +24,7 @@ tf_org_notebook_test(
     deps = [
         "//third_party/py/matplotlib",
         "@pypi//numpy:pkg",
-        #  tensorflow datasets dep,
+        "@pypi//tensorflow_datasets:pkg",
     ],
 )
 
@@ -43,9 +43,9 @@ tf_org_notebook_test(
     tags = ["requires-net:external"],
     deps = [
         "//third_party/py/matplotlib",
-        "@pypi//numpy:pkg",
-        #  tensorflow datasets dep,
         "//third_party/py/tensorflow_text",
+        "@pypi//numpy:pkg",
+        "@pypi//tensorflow_datasets:pkg",
     ],
 )
 
@@ -57,6 +57,6 @@ tf_org_notebook_test(
     deps = [
         "//third_party/py/matplotlib",
         "@pypi//numpy:pkg",
-        #  tensorflow datasets dep,
+        "@pypi//tensorflow_datasets:pkg",
     ],
 )

oss_scripts/build.Dockerfile.aarch64

@@ -22,16 +22,12 @@ RUN --mount=type=cache,target=/var/cache/apt \
     checkinstall \
     libffi-dev
 
-# Install pip dependencies needed for tensorflow-text
+# Install pip dependencies needed for packaging tensorflow-text wheels
 RUN --mount=type=cache,target=/root/.cache \
   ${PYTHON_BIN_PATH} -m pip install -U pip && \
   ${PYTHON_BIN_PATH} -m pip install -U \
-    absl-py \
     auditwheel \
-    etils[epath] \
     patchelf \
-    setuptools \
-    twine \
-    wheel;
+    twine;
 
 WORKDIR "/tmp/tensorflow_text"

oss_scripts/build.Dockerfile.x86_64

@@ -17,16 +17,12 @@ RUN --mount=type=cache,target=/var/cache/apt \
     checkinstall \
     libffi-dev
 
-# Install pip dependencies needed for tensorflow text.
+# Install pip dependencies needed for packaging tensorflow text wheels.
 RUN --mount=type=cache,target=/root/.cache \
   python${PYTHON_VERSION} -m pip install -U pip && \
   python${PYTHON_VERSION} -m pip install -U \
-    absl-py \
     auditwheel \
-    etils[epath] \
     patchelf \
-    setuptools \
-    twine \
-    wheel;
+    twine;
 
 WORKDIR "/tmp/tensorflow_text"

oss_scripts/configure.sh

@@ -14,6 +14,8 @@
 # limitations under the License.
 # ==============================================================================
 
+set -e  # fail and exit on any command erroring
+
 function write_to_bazelrc() {
   echo "$1" >> .bazelrc
 }
@@ -29,28 +31,8 @@ function is_macos() {
   [[ "${osname}" == "darwin" ]]
 }
 
-# Remove .bazelrc if it already exist
-[ -e .bazelrc ] && rm .bazelrc
-
-# Don't install tensorflow locally.
-if (which python3) | grep -q "python3"; then
-  installed_python="python3"
-elif (which python) | grep -q "python"; then
-  installed_python="python"
-fi
-
-# if [[ $($installed_python -m pip show tensorflow) == *tensorflow* ]] ||
-#    [[ $($installed_python -m pip show tensorflow-macos) == *tensorflow-macos* ]] ||
-#    [[ $($installed_python -m pip show tf-nightly) == *tf-nightly* ]]; then
-#   echo 'Using installed tensorflow.'
-# else
-#   echo 'Installing tensorflow.'
-#   if [[ "$IS_NIGHTLY" == "nightly" ]]; then
-#     $installed_python -m pip install --no-deps tf-nightly
-#   else
-#     $installed_python -m pip install --no-deps tensorflow==2.22.0rc0
-#   fi
-# fi
+# Remove .bazelrc if it already exists
+rm -f .bazelrc
 
 # Copy the current bazelversion of TF.
 curl https://raw.githubusercontent.com/tensorflow/tensorflow/v2.22.0-rc0/.bazelversion -o .bazelversion
@@ -63,19 +45,39 @@ sed -i -e 's/build --noincompatible_remove_legacy_whole_archive//' .bazelrc
 
 write_to_bazelrc "build:manylinux2014 --config=release_cpu_linux"
 write_to_bazelrc "test --test_env=TF_USE_LEGACY_KERAS=1"
+write_to_bazelrc "test --test_tag_filters=-benchmark-test"
 
 if [ -z "$HERMETIC_PYTHON_VERSION" ]; then
   if [ -n "$PY_VERSION" ]; then
     HERMETIC_PYTHON_VERSION=$PY_VERSION
   else
-    HERMETIC_PYTHON_VERSION=$($installed_python  -c "import sys; print('.'.join(map(str, sys.version_info[:2])))")
+    if (which python3) | grep -q "python3"; then
+      installed_python="python3"
+    elif (which python) | grep -q "python"; then
+      installed_python="python"
+    fi
+    HERMETIC_PYTHON_VERSION=$($installed_python -c "import sys; print('.'.join(map(str, sys.version_info[:2])))")
   fi
 fi
 export HERMETIC_PYTHON_VERSION
 
-# Only auto-upgrade requirements for nightly.
-if [[ "$IS_NIGHTLY" == "nightly" ]]; then
-  bazel run //oss_scripts/pip_package:requirements.update -- --upgrade
+if [[ -n "$HERMETIC_PYTHON_VERSION" ]]; then
+  write_to_bazelrc "build --repo_env=HERMETIC_PYTHON_VERSION=${HERMETIC_PYTHON_VERSION}"
+fi
+if [[ -n "$IS_NIGHTLY" ]]; then
+  write_to_bazelrc "build --repo_env=IS_NIGHTLY=${IS_NIGHTLY}"
+fi
+if [[ -n "$TF_VERSION" ]]; then
+  write_to_bazelrc "build --repo_env=TF_VERSION=${TF_VERSION}"
+fi
+
+# Auto-upgrade requirements for nightly or when a specific TF_VERSION is requested.
+if [[ "$IS_NIGHTLY" == "nightly" ]] || [[ -n "$TF_VERSION" ]]; then
+  REQUIREMENTS_EXTRA_FLAGS="--upgrade"
+  if [[ "$TF_VERSION" == *"rc"* ]]; then
+    REQUIREMENTS_EXTRA_FLAGS="$REQUIREMENTS_EXTRA_FLAGS --pre"
+  fi
+  bazel run //oss_scripts/pip_package:requirements.update -- $REQUIREMENTS_EXTRA_FLAGS
 fi
 
 TF_ABIFLAG=$(bazel run //oss_scripts/pip_package:tensorflow_build_info -- abi)

oss_scripts/pip_package/BUILD

@@ -23,7 +23,7 @@ py_binary(
     name = "tensorflow_build_info",
     srcs = ["tensorflow_build_info.py"],
     deps = [
-        "@absl_py//absl:app",
+        "@pypi//absl_py:pkg",
         "@release_or_nightly//:tensorflow_pkg",
     ],
 )
@@ -40,6 +40,6 @@ py_binary(
         "//tensorflow_text/core/pybinds:pybinds_binaries",
     ],
     deps = [
-        "@pypi_setuptools//:pkg",
+        "@pypi//setuptools:pkg",
     ],
 )

oss_scripts/prepare_tf_dep.sh

@@ -1,12 +1,6 @@
 #!/bin/bash
 set -e  # fail and exit on any command erroring
 
-if (which python3) | grep -q "python3"; then
-  installed_python="python3"
-elif (which python) | grep -q "python"; then
-  installed_python="python"
-fi
-
 ext=""
 osname="$(uname -s | tr 'A-Z' 'a-z')"
 if [[ "${osname}" == "darwin" ]]; then

oss_scripts/run_build.sh

@@ -13,11 +13,10 @@ fi
 # Run configure.
 source oss_scripts/configure.sh
 
-# Don't overwrite tensorflow version.
-# # Set tensorflow version
-# if [[ $osname != "darwin" ]] || [[ ! $(sysctl -n machdep.cpu.brand_string) =~ "Apple" ]]; then
-#   source oss_scripts/prepare_tf_dep.sh
-# fi
+# Set tensorflow version if nightly or custom TF_VERSION is specified.
+if [[ "$IS_NIGHTLY" == "nightly" ]] || [[ -n "$TF_VERSION" ]]; then
+  source oss_scripts/prepare_tf_dep.sh
+fi
 
 # Build the pip package.
 bazel run ${BUILD_ARGS[@]} --enable_runfiles //oss_scripts/pip_package:build_pip_package -- "$(realpath .)"

oss_scripts/run_tests.sh

@@ -7,16 +7,16 @@ set -e  # fail and exit on any command erroring
 
 # Verify correct version of Bazel
 installed_bazel_version=$(bazel version | grep label | sed -e 's/.*: //')
-tf_bazel_version=$(cat .bazelversion)
+tf_bazel_version=$(head -n 1 .bazelversion)
 if [ "$installed_bazel_version" != "$tf_bazel_version" ]; then
   echo "Incorrect version of Bazel installed."
   echo "Version $tf_bazel_version should be installed, but found version ${installed_bazel_version}."
   echo "Run oss_scripts/install_bazel.sh or manually install the correct version."
   exit 1
 fi
 
-# Set tensorflow version
-if [[ $osname != "Darwin" ]] || [[ ! $(sysctl -n machdep.cpu.brand_string) =~ "Apple" ]]; then
+# Set tensorflow version if nightly or custom TF_VERSION is specified.
+if [[ "$IS_NIGHTLY" == "nightly" ]] || [[ -n "$TF_VERSION" ]]; then
   source oss_scripts/prepare_tf_dep.sh
 fi