Refactor mechanisms of building TF wheel and storing TF project version.

This change introduces a uniform way of building the TF wheel and controlling the filename version suffixes. A new repository rule `python_wheel_version_suffix_repository` provides information about project and wheel version suffixes. The final value depends on environment variables passed to Bazel command: `_ML_WHEEL_WHEEL_TYPE, _ML_WHEEL_BUILD_DATE, _ML_WHEEL_GIT_HASH, _ML_WHEEL_VERSION_SUFFIX` `tf_version.bzl` defines the TF project version and loads the version suffix information calculated by `python_wheel_version_suffix_repository`. The targets `//tensorflow/core/public:release_version, //tensorflow:tensorflow_bzl //tensorflow/tools/pip_package:setup_py` use the version chunks defined above. The version of the wheel in the build rule output depends on the environment variables. Environment variables combinations for creating wheels with different versions: * snapshot (default build rule behavior): `--repo_env=ML_WHEEL_TYPE=snapshot` * release: `--repo_env=ML_WHEEL_TYPE=release` * release candidate: `--repo_env=ML_WHEEL_TYPE=release --repo_env=ML_WHEEL_VERSION_SUFFIX=-rc1` * nightly build with date as version suffix: `--repo_env=ML_WHEEL_TYPE=nightly --repo_env=ML_WHEEL_BUILD_DATE=<YYYYmmdd>` * build with git data as version suffix: `--repo_env=ML_WHEEL_TYPE=custom --repo_env=ML_WHEEL_BUILD_DATE=$(git show -s --format=%as HEAD) --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)` PiperOrigin-RevId: 733444080
2025-12-06 00:19:58 +01:00 · 2025-03-04 13:25:05 -08:00 · 2025-03-04 13:25:05 -08:00 · 805775fcb5
commit 805775fcb5
parent 9d015d93f4
23 changed files with 319 additions and 169 deletions
--- a/.bazelrc
+++ b/.bazelrc
@ -69,6 +69,16 @@
 # LINT.IfChange
 # Default build options. These are applied first and unconditionally.

+# These are used to generate the ML wheel version string.
+# See the explanation in the file comment of
+# @local_tsl//third_party/py/python_wheel.bzl.
+# The generated version suffix is used in
+# third_party/tensorflow/core/public/release_version.h and
+# third_party/tensorflow/tools/pip_package/setup.oss.py.tpl
+build --repo_env=ML_WHEEL_TYPE="snapshot"
+build --repo_env=ML_WHEEL_BUILD_DATE=""
+build --repo_env=ML_WHEEL_VERSION_SUFFIX=""
+
 # For projects which use TensorFlow as part of a Bazel build process, putting
 # nothing in a bazelrc will default to a monolithic build. The following line
 # opts in to modular op registration support by default.
--- a/1
+++ b/1
@ -2,5 +2,6 @@ exports_files(glob(["requirements*"]) + [
    "configure",
    "configure.py",
    "ACKNOWLEDGEMENTS",
+    "AUTHORS",
    "LICENSE",
 ])
--- a/7
+++ b/7
@ -71,6 +71,13 @@ load("@//tensorflow:workspace0.bzl", "tf_workspace0")

 tf_workspace0()

+load(
+    "@local_tsl//third_party/py:python_wheel.bzl",
+    "python_wheel_version_suffix_repository",
+)
+
+python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
+
 load(
    "@local_tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
    "cuda_json_init_repository",
--- a/ci/official/envs/versions_upload
+++ b/ci/official/envs/versions_upload
@ -26,3 +26,4 @@ TFCI_ARTIFACT_STAGING_GCS_ENABLE=1
 TFCI_ARTIFACT_STAGING_GCS_URI="gs://tensorflow-ci-staging/staging/versions/${_LOUHI_EXECUTION_ID:-$(git rev-parse HEAD)}/"

 TFCI_BAZEL_COMMON_ARGS="$TFCI_BAZEL_COMMON_ARGS --config resultstore"
+TFCI_BUILD_PIP_PACKAGE_ARGS="$TFCI_BUILD_PIP_PACKAGE_ARGS --repo_env=ML_WHEEL_TYPE=release"
--- a/ci/official/installer_wheel.sh
+++ b/ci/official/installer_wheel.sh
@ -21,7 +21,7 @@ if [[ "$TFCI_NIGHTLY_UPDATE_VERSION_ENABLE" == 1 ]]; then
 fi

 # This generates a pure python wheel of the format "*-py3-none-any.whl"
-python3 tensorflow/tools/pip_package/setup.py bdist_wheel --dist-dir "$TFCI_OUTPUT_DIR"
+bazel run --HERMETIC_PYTHON_VERSION=3.13 //tensorflow/tools/pip_package:setup_py_binary -- bdist_wheel --dist-dir "$TFCI_OUTPUT_DIR"

 # Get the name of the pure python wheel that was built. This should
 # resolve to either
--- a/ci/official/upload.sh
+++ b/ci/official/upload.sh
@ -17,13 +17,13 @@
 # job chain to GCS and PyPI.
 source "${BASH_SOURCE%/*}/utilities/setup.sh"

-# Update the version numbers for Nightly only, then fetch the version numbers
-# for choosing the final directory name. This adds "-devYYYYMMDD" to the end of
-# the version string (.devYYYYMMDD for Python; see pypi.org/project/tf-nightly)
+# Calculate the version number for choosing the final directory name. This adds
+# "-devYYYYMMDD" to the end of the version string for nightly builds.
 if [[ "$TFCI_NIGHTLY_UPDATE_VERSION_ENABLE" == 1 ]]; then
-  tfrun python3 tensorflow/tools/ci_build/update_version.py --nightly
+  export TF_VER_FULL="$(tfrun bazel run //tensorflow/tools/ci_build:calculate_full_version -- --wheel-type nightly)"
+else
+  export TF_VER_FULL="$(tfrun bazel run //tensorflow/tools/ci_build:calculate_full_version -- --wheel-type release)"
 fi
-source ci/official/utilities/get_versions.sh

 # Note on gsutil commands:
 # "gsutil cp" always "copies into". It cannot act on the contents of a directory
--- a/ci/official/utilities/get_versions.sh
+++ b/ci/official/utilities/get_versions.sh
@ -1,40 +0,0 @@
-#!/bin/bash
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed 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.
-# ==============================================================================
-# Variables containing version strings extracted from the canonical sources:
-# tensorflow/core/public/version.h and tools/pip_package/setup.py.
-#
-# These variables aren't available by default. Scripts must source this file
-# explicitly, *after* checking if update_version.py needs to be run for a
-# nightly job. update_version.py affects TF_VER_SUFFIX, TF_VER_PYTHON, and
-# TF_VER_FULL.
-
-# Note: in awk, the command '/search/ {commands}' applies the commands to any line that
-# matches the /search/ regular expression. "print $N" prints the Nth "field",
-# where fields are strings separated by whitespace.
-export TF_VER_MAJOR=$(awk '/#define TF_MAJOR_VERSION/ {print $3}' tensorflow/core/public/version.h)
-export TF_VER_MINOR=$(awk '/#define TF_MINOR_VERSION/ {print $3}' tensorflow/core/public/version.h)
-export TF_VER_PATCH=$(awk '/#define TF_PATCH_VERSION/ {print $3}' tensorflow/core/public/version.h)
-
-# Note: in awk, "print $N" prints the Nth "field", where fields are strings separated
-# by whitespace. The flag "-F<x>" changes the behavior such that fields are now strings
-# separated by the character <x>. Therefore, -F\' (escaped single quote) means that field
-# $2 in <Tensor'flow'> is <flow>. This is useful for reading string literals like below.
-export TF_VER_SUFFIX=$(awk -F\" '/#define TF_VERSION_SUFFIX/ {print $2}' tensorflow/core/public/version.h)
-export TF_VER_PYTHON=$(awk -F\' '/_VERSION =/ {print $2}' tensorflow/tools/pip_package/setup.py)
-
-# Derived helper variables.
-export TF_VER_SHORT="${TF_VER_MAJOR}.${TF_VER_MINOR}"
-export TF_VER_FULL="${TF_VER_MAJOR}.${TF_VER_MINOR}.${TF_VER_PATCH}${TF_VER_SUFFIX}"
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@ -1122,11 +1122,22 @@ filegroup(
    ),
 )

+bzl_library(
+    name = "tf_version_bzl",
+    srcs = [
+        "tf_version.bzl",
+        # copybara:uncomment "tf_version.google.bzl",
+    ],
+    # copybara:uncomment parse_tests = False,
+    visibility = ["//visibility:public"],
+)
+
 bzl_library(
    name = "tensorflow_bzl",
    srcs = ["tensorflow.bzl"],
    visibility = ["//visibility:public"],
    deps = [
+        ":tf_version_bzl",
        "//tensorflow/core/platform:build_config_root_bzl",
        "//tensorflow/core/platform:rules_cc_bzl",
        "//third_party/compute_library:build_defs_bzl",
--- a/tensorflow/core/public/BUILD
+++ b/tensorflow/core/public/BUILD
@ -1,4 +1,11 @@
 load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load(
+    "//tensorflow:tf_version.bzl",
+    "MAJOR_VERSION",
+    "MINOR_VERSION",
+    "PATCH_VERSION",
+    "TF_SEMANTIC_VERSION_SUFFIX",
+)
 load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")

 package(
@ -44,5 +51,11 @@ cc_library(
 cc_library(
    name = "release_version",
    hdrs = ["release_version.h"],
+    defines = [
+        "TF_MAJOR_VERSION={}".format(MAJOR_VERSION),
+        "TF_MINOR_VERSION={}".format(MINOR_VERSION),
+        "TF_PATCH_VERSION={}".format(PATCH_VERSION),
+        'TF_VERSION_SUFFIX=\\"{}\\"'.format(TF_SEMANTIC_VERSION_SUFFIX),
+    ],
    visibility = ["//visibility:public"],
 )
--- a/tensorflow/core/public/release_version.h
+++ b/tensorflow/core/public/release_version.h
@ -16,17 +16,12 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PUBLIC_RELEASE_VERSION_H_
 #define TENSORFLOW_CORE_PUBLIC_RELEASE_VERSION_H_

-// TensorFlow uses semantic versioning, see http://semver.org/.
-
-// Also update tensorflow/tensorflow.bzl and
-// tensorflow/tools/pip_package/setup.py
-#define TF_MAJOR_VERSION 2
-#define TF_MINOR_VERSION 20
-#define TF_PATCH_VERSION 0
-
-// TF_VERSION_SUFFIX is non-empty for pre-releases (e.g. "-alpha", "-alpha.1",
-// "-beta", "-rc", "-rc.1")
-#define TF_VERSION_SUFFIX ""
+// A cc_library //third_party/tensorflow/core/public:release_version provides
+// defines with the version data from //third_party/tensorflow/tf_version.bzl.
+// The version suffix can be set by passing the build parameters
+// --repo_env=ML_WHEEL_BUILD_DATE=<date> and
+// --repo_env=ML_WHEEL_VERSION_SUFFIX=<suffix>.
+// To update the project version, update tf_version.bzl.

 #define _TF_STR_HELPER(x) #x
 #define _TF_STR(x) _TF_STR_HELPER(x)
--- a/tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh
+++ b/tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh
@ -20,7 +20,7 @@ PYTHON="${CI_BUILD_PYTHON:-python3}"
 VERSION_SUFFIX=${VERSION_SUFFIX:-}
 export TENSORFLOW_DIR="${SCRIPT_DIR}/../../../.."
 TENSORFLOW_LITE_DIR="${TENSORFLOW_DIR}/tensorflow/lite"
-TENSORFLOW_VERSION=$(grep "_VERSION = " "${TENSORFLOW_DIR}/tensorflow/tools/pip_package/setup.py" | cut -d= -f2 | sed "s/[ '-]//g")
+TENSORFLOW_VERSION=$(grep "TF_VERSION = " "${TENSORFLOW_DIR}/tensorflow/tf_version.bzl" | cut -d= -f2 | sed 's/[ "-]//g')
 export PACKAGE_VERSION="${TENSORFLOW_VERSION}${VERSION_SUFFIX}"
 export PROJECT_NAME=${WHEEL_PROJECT_NAME:-tflite_runtime}
 BUILD_DIR="${SCRIPT_DIR}/gen/tflite_pip/${PYTHON}"
--- a/tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh
+++ b/tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh
@ -20,7 +20,12 @@ PYTHON="${CI_BUILD_PYTHON:-python3}"
 VERSION_SUFFIX=${VERSION_SUFFIX:-}
 export TENSORFLOW_DIR="${SCRIPT_DIR}/../../../.."
 TENSORFLOW_LITE_DIR="${TENSORFLOW_DIR}/tensorflow/lite"
-TENSORFLOW_VERSION=$(grep "_VERSION = " "${TENSORFLOW_DIR}/tensorflow/tools/pip_package/setup.py" | cut -d= -f2 | sed "s/[ '-]//g")
+TENSORFLOW_VERSION=$(grep "TF_VERSION = " "${TENSORFLOW_DIR}/tensorflow/tf_version.bzl" | cut -d= -f2 | sed 's/[ "-]//g')
+IFS='.' read -r -a array <<< "$TENSORFLOW_VERSION"
+TF_MAJOR=${array[0]}
+TF_MINOR=${array[1]}
+TF_PATCH=${array[2]}
+TF_CXX_FLAGS="-DTF_MAJOR_VERSION=${TF_MAJOR} -DTF_MINOR_VERSION=${TF_MINOR} -DTF_PATCH_VERSION=${TF_PATCH} -DTF_VERSION_SUFFIX=''"
 export PACKAGE_VERSION="${TENSORFLOW_VERSION}${VERSION_SUFFIX}"
 export PROJECT_NAME=${WHEEL_PROJECT_NAME:-tflite_runtime}
 BUILD_DIR="${SCRIPT_DIR}/gen/tflite_pip/${PYTHON}"
@ -82,7 +87,7 @@ echo "Building for ${TENSORFLOW_TARGET}"
 case "${TENSORFLOW_TARGET}" in
  armhf)
    eval $(${TENSORFLOW_LITE_DIR}/tools/cmake/download_toolchains.sh "${TENSORFLOW_TARGET}")
-    ARMCC_FLAGS="${ARMCC_FLAGS} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"
+    ARMCC_FLAGS="${ARMCC_FLAGS} ${TF_CXX_FLAGS} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"
    cmake \
      -DCMAKE_C_COMPILER=${ARMCC_PREFIX}gcc \
      -DCMAKE_CXX_COMPILER=${ARMCC_PREFIX}g++ \
@ -96,7 +101,7 @@ case "${TENSORFLOW_TARGET}" in
    ;;
  rpi0)
    eval $(${TENSORFLOW_LITE_DIR}/tools/cmake/download_toolchains.sh "${TENSORFLOW_TARGET}")
-    ARMCC_FLAGS="${ARMCC_FLAGS} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"
+    ARMCC_FLAGS="${ARMCC_FLAGS} ${TF_CXX_FLAGS} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"
    cmake \
      -DCMAKE_C_COMPILER=${ARMCC_PREFIX}gcc \
      -DCMAKE_CXX_COMPILER=${ARMCC_PREFIX}g++ \
@ -110,7 +115,7 @@ case "${TENSORFLOW_TARGET}" in
    ;;
  aarch64)
    eval $(${TENSORFLOW_LITE_DIR}/tools/cmake/download_toolchains.sh "${TENSORFLOW_TARGET}")
-    ARMCC_FLAGS="${ARMCC_FLAGS} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"
+    ARMCC_FLAGS="${ARMCC_FLAGS} ${TF_CXX_FLAGS} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"
    cmake \
      -DCMAKE_C_COMPILER=${ARMCC_PREFIX}gcc \
      -DCMAKE_CXX_COMPILER=${ARMCC_PREFIX}g++ \
@ -123,14 +128,14 @@ case "${TENSORFLOW_TARGET}" in
      "${TENSORFLOW_LITE_DIR}"
    ;;
  native)
-    BUILD_FLAGS=${BUILD_FLAGS:-"-march=native -I${PYTHON_INCLUDE} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"}
+    BUILD_FLAGS=${BUILD_FLAGS:-"-march=native ${TF_CXX_FLAGS} -I${PYTHON_INCLUDE} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"}
    cmake \
      -DCMAKE_C_FLAGS="${BUILD_FLAGS}" \
      -DCMAKE_CXX_FLAGS="${BUILD_FLAGS}" \
      "${TENSORFLOW_LITE_DIR}"
    ;;
  *)
-    BUILD_FLAGS=${BUILD_FLAGS:-"-I${PYTHON_INCLUDE} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"}
+    BUILD_FLAGS=${BUILD_FLAGS:-"${TF_CXX_FLAGS} -I${PYTHON_INCLUDE} -I${PYBIND11_INCLUDE} -I${NUMPY_INCLUDE}"}
    cmake \
      -DCMAKE_C_FLAGS="${BUILD_FLAGS}" \
      -DCMAKE_CXX_FLAGS="${BUILD_FLAGS}" \
--- a/tensorflow/opensource_only.files
+++ b/tensorflow/opensource_only.files
@ -173,7 +173,7 @@ tf_staging/tensorflow/tools/pip_package/MANIFEST.in:
 tf_staging/tensorflow/tools/pip_package/README:
 tf_staging/tensorflow/tools/pip_package/check_load_py_test:.py
 tf_staging/tensorflow/tools/pip_package/pip_smoke_test:.py
-tf_staging/tensorflow/tools/pip_package/setup:.py
+tf_staging/tensorflow/tools/pip_package/setup:.py.tpl
 tf_staging/tensorflow/tools/pip_package/simple_console:.py
 tf_staging/tensorflow/tools/pip_package/simple_console_for_windows:.py
 tf_staging/tensorflow/tools/pip_package/utils/BUILD:
--- a/tensorflow/python/framework/versions_test.py
+++ b/tensorflow/python/framework/versions_test.py
@ -24,8 +24,12 @@ class VersionTest(test.TestCase):
    self.assertEqual(type(versions.__version__), str)
    self.assertEqual(type(versions.VERSION), str)
    # This pattern will need to grow as we include alpha, builds, etc.
-    self.assertRegex(versions.__version__, r'^\d+\.\d+\.(\d+(\-\w+)?|head)$')
-    self.assertRegex(versions.VERSION, r'^\d+\.\d+\.(\d+(\-\w+)?|head)$')
+    self.assertRegex(
+        versions.__version__, r'^\d+\.\d+\.(\d+(\-\w+)?(\+\w+)?|head)$'
+    )
+    self.assertRegex(
+        versions.VERSION, r'^\d+\.\d+\.(\d+(\-\w+)?(\+\w+)?|head)$'
+    )

  def testGraphDefVersion(self):
    version = versions.GRAPH_DEF_VERSION
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@ -25,6 +25,7 @@ load(
    "if_mkldnn_openmp",
    "onednn_v3_define",
 )
+load("//tensorflow:tf_version.bzl", "TF_VERSION")

 #
 # Returns the options to use for a C++ library or binary build.
@ -93,10 +94,7 @@ def register_extension_info(**kwargs):

 # version for the shared libraries, can
 # not contain rc or alpha, only numbers.
-# Also update tensorflow/core/public/version.h
-# and tensorflow/tools/pip_package/setup.py
-WHEEL_VERSION = "2.20.0"
-VERSION = "2.20.0"
+VERSION = TF_VERSION
 VERSION_MAJOR = VERSION.split(".")[0]
 two_gpu_tags = ["requires-gpu-nvidia:2", "manual", "no_pip"]

--- a/tensorflow/tf_version.bzl
+++ b/tensorflow/tf_version.bzl
@ -0,0 +1,14 @@
+"""Provides project and wheel version data for TensorFlow."""
+
+load(
+    "//tensorflow:tf_version.default.bzl",
+    "SEMANTIC_VERSION_SUFFIX",
+    "VERSION_SUFFIX",
+)
+
+# These constants are used by the targets //third_party/tensorflow/core/public:release_version,
+# //third_party/tensorflow:tensorflow_bzl and //third_party/tensorflow/tools/pip_package:setup_py.
+TF_VERSION = "2.20.0"
+MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION = TF_VERSION.split(".")
+TF_WHEEL_VERSION_SUFFIX = VERSION_SUFFIX
+TF_SEMANTIC_VERSION_SUFFIX = SEMANTIC_VERSION_SUFFIX
--- a/tensorflow/tf_version.default.bzl
+++ b/tensorflow/tf_version.default.bzl
@ -0,0 +1,10 @@
+"""Default (OSS) TensorFlow wheel version suffix data."""
+
+load(
+    "@tf_wheel_version_suffix//:wheel_version_suffix.bzl",
+    "SEMANTIC_WHEEL_VERSION_SUFFIX",
+    "WHEEL_VERSION_SUFFIX",
+)
+
+VERSION_SUFFIX = WHEEL_VERSION_SUFFIX
+SEMANTIC_VERSION_SUFFIX = SEMANTIC_WHEEL_VERSION_SUFFIX
--- a/tensorflow/tools/ci_build/BUILD
+++ b/tensorflow/tools/ci_build/BUILD
@ -1,8 +1,17 @@
+load("//tensorflow:strict.default.bzl", "py_strict_binary")
+load("//tensorflow:tf_version.bzl", "TF_VERSION")
+
 package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    licenses = ["notice"],
 )

+py_strict_binary(
+    name = "calculate_full_version",
+    srcs = ["calculate_full_version.py"],
+    args = ["--wheel-version={}".format(TF_VERSION)],
+)
+
 exports_files(
    srcs = glob(["Dockerfile.*"]),
    visibility = ["//visibility:public"],
--- a/tensorflow/tools/ci_build/calculate_full_version.py
+++ b/tensorflow/tools/ci_build/calculate_full_version.py
@ -0,0 +1,48 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed 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.
+# ==============================================================================
+"""Calculates full TensorFlow version."""
+
+import argparse
+import time
+
+
+def parse_args():
+  """Arguments parser."""
+  parser = argparse.ArgumentParser(
+      description="Helper for TensorFlow version calculation",
+      fromfile_prefix_chars="@",
+  )
+  parser.add_argument(
+      "--wheel-type",
+      required=True,
+      choices=["nightly", "release"],
+      help="Type of the wheel",
+  )
+  parser.add_argument("--wheel-version", required=True, help="Wheel version")
+  return parser.parse_args()
+
+
+if __name__ == "__main__":
+  args = parse_args()
+  tensorflow_version_suffix = ""
+  if args.wheel_type == "nightly":
+    tensorflow_version_suffix = "-dev{}".format(time.strftime("%Y%m%d"))
+
+  print(
+      '"{wheel_version}{tensorflow_version_suffix}"'.format(
+          wheel_version=args.wheel_version,
+          tensorflow_version_suffix=tensorflow_version_suffix,
+      )
+  )
--- a/tensorflow/tools/ci_build/update_version.py
+++ b/tensorflow/tools/ci_build/update_version.py
@ -32,30 +32,19 @@ import time

 # File parameters.
 TF_SRC_DIR = "tensorflow"
-RELEASE_VERSION_H = "%s/core/public/release_version.h" % TF_SRC_DIR
-SETUP_PY = "%s/tools/pip_package/setup.py" % TF_SRC_DIR
 README_MD = "./README.md"
-TENSORFLOW_BZL = "%s/tensorflow.bzl" % TF_SRC_DIR
-TF_MAC_ARM64_CI_BUILD = (
-    "%s/tools/ci_build/osx/arm64/tensorflow_as_build_release.Jenkinsfile"
-    % TF_SRC_DIR
-)
-TF_MAC_ARM64_CI_TEST = (
-    "%s/tools/ci_build/osx/arm64/tensorflow_as_test_release.Jenkinsfile"
-    % TF_SRC_DIR
-)
+TF_VERSION_BZL = "%s/tf_version.bzl" % TF_SRC_DIR
+BAZEL_RC = "./.bazelrc"
 RELEVANT_FILES = [
    TF_SRC_DIR,
-    RELEASE_VERSION_H,
-    SETUP_PY,
+    TF_VERSION_BZL,
    README_MD,
-    TF_MAC_ARM64_CI_BUILD,
-    TF_MAC_ARM64_CI_TEST
+    BAZEL_RC,
 ]

 # Version type parameters.
 NIGHTLY_VERSION = 1
-REGULAR_VERSION = 0
+SNAPSHOT_VERSION = 0


 def check_existence(filename):
@ -90,7 +79,7 @@ class Version(object):
      minor: minor string eg. (3)
      patch: patch string eg. (1)
      identifier_string: extension string eg. (-rc0)
-      version_type: version parameter ((REGULAR|NIGHTLY)_VERSION)
+      version_type: version parameter ((SNAPSHOT|NIGHTLY)_VERSION)
    """
    self.major = major
    self.minor = minor
@ -114,7 +103,7 @@ class Version(object):

  @property
  def pep_440_str(self):
-    if self.version_type == REGULAR_VERSION:
+    if self.version_type == SNAPSHOT_VERSION:
      return_string = "%s.%s.%s%s" % (self.major,
                                      self.minor,
                                      self.patch,
@ -151,11 +140,14 @@ class Version(object):
    else:
      identifier_string = ""

-    return Version(major,
-                   minor,
-                   patch,
-                   identifier_string,
-                   version_type)
+    return Version(major, minor, patch, identifier_string, version_type)
+
+
+def _get_regex_match(line, regex, is_last_match=False):
+  match = re.search(regex, line)
+  if match:
+    return (match.group(1), is_last_match)
+  return (None, False)


 def get_current_semver_version():
@ -163,30 +155,58 @@ def get_current_semver_version():

  Returns:
    version: Version object of current SemVer string based on information from
-    core/public/release_version.h
+    .bazelrc and tf_version.bzl files.
  """

  # Get current version information.
-  version_file = open(RELEASE_VERSION_H, "r")
-  for line in version_file:
-    major_match = re.search("^#define TF_MAJOR_VERSION ([0-9]+)", line)
-    minor_match = re.search("^#define TF_MINOR_VERSION ([0-9]+)", line)
-    patch_match = re.search("^#define TF_PATCH_VERSION ([0-9]+)", line)
-    extension_match = re.search("^#define TF_VERSION_SUFFIX \"(.*)\"", line)
-    if major_match:
-      old_major = major_match.group(1)
-    if minor_match:
-      old_minor = minor_match.group(1)
-    if patch_match:
-      old_patch_num = patch_match.group(1)
-    if extension_match:
-      old_extension = extension_match.group(1)
+  bazel_rc_file = open(BAZEL_RC, "r")
+
+  wheel_type = ""
+  wheel_build_date = ""
+  wheel_version_suffix = ""
+
+  for line in bazel_rc_file:
+    wheel_type = (
+        _get_regex_match(line, '^build --repo_env=ML_WHEEL_TYPE="(.+)"')[0]
+        or wheel_type
+    )
+    wheel_build_date = (
+        _get_regex_match(
+            line, '^build --repo_env=ML_WHEEL_BUILD_DATE="([0-9]*)"'
+        )[0]
+        or wheel_build_date
+    )
+    (wheel_version_suffix, is_matched) = _get_regex_match(
+        line,
+        '^build --repo_env=ML_WHEEL_VERSION_SUFFIX="(.*)"',
+        is_last_match=True,
+    )
+    if is_matched:
      break

-  if "dev" in old_extension:
+  tf_version_bzl_file = open(TF_VERSION_BZL, "r")
+  wheel_version = ""
+  for line in tf_version_bzl_file:
+    (wheel_version, is_matched) = _get_regex_match(
+        line, '^TF_VERSION = "([0-9.]+)"', is_last_match=True
+    )
+    if is_matched:
+      break
+  (old_major, old_minor, old_patch_num) = wheel_version.split(".")
+
+  if wheel_type == "nightly":
    version_type = NIGHTLY_VERSION
  else:
-    version_type = REGULAR_VERSION
+    version_type = SNAPSHOT_VERSION
+
+  old_extension = ""
+  if wheel_type == "nightly":
+    old_extension = "-dev{}".format(wheel_build_date)
+  else:
+    if wheel_build_date:
+      old_extension += "-dev{}".format(wheel_build_date)
+    if wheel_version_suffix:
+      old_extension += wheel_version_suffix

  return Version(old_major,
                 old_minor,
@ -195,29 +215,6 @@ def get_current_semver_version():
                 version_type)


-def update_release_version_h(old_version, new_version):
-  """Update tensorflow/core/public/release_version.h."""
-  replace_string_in_line("#define TF_MAJOR_VERSION %s" % old_version.major,
-                         "#define TF_MAJOR_VERSION %s" % new_version.major,
-                         RELEASE_VERSION_H)
-  replace_string_in_line("#define TF_MINOR_VERSION %s" % old_version.minor,
-                         "#define TF_MINOR_VERSION %s" % new_version.minor,
-                         RELEASE_VERSION_H)
-  replace_string_in_line("#define TF_PATCH_VERSION %s" % old_version.patch,
-                         "#define TF_PATCH_VERSION %s" % new_version.patch,
-                         RELEASE_VERSION_H)
-  replace_string_in_line(
-      "#define TF_VERSION_SUFFIX \"%s\"" % old_version.identifier_string,
-      "#define TF_VERSION_SUFFIX \"%s\"" % new_version.identifier_string,
-      RELEASE_VERSION_H)
-
-
-def update_setup_dot_py(old_version, new_version):
-  """Update setup.py."""
-  replace_string_in_line("_VERSION = '%s'" % old_version.string,
-                         "_VERSION = '%s'" % new_version.string, SETUP_PY)
-
-
 def update_readme(old_version, new_version):
  """Update README."""
  pep_440_str = new_version.pep_440_str
@ -226,32 +223,65 @@ def update_readme(old_version, new_version):
                         "%s-" % pep_440_str, README_MD)


-def update_tensorflow_bzl(old_version, new_version):
-  """Update tensorflow.bzl."""
-  old_mmp = "%s.%s.%s" % (old_version.major, old_version.minor,
-                          old_version.patch)
-  new_mmp = "%s.%s.%s" % (new_version.major, new_version.minor,
-                          new_version.patch)
-  replace_string_in_line('VERSION = "%s"' % old_mmp,
-                         'VERSION = "%s"' % new_mmp, TENSORFLOW_BZL)
-  replace_string_in_line(
-      'WHEEL_VERSION = "%s"' % old_version.string,
-      'WHEEL_VERSION = "%s"' % new_version.string,
-      TENSORFLOW_BZL,
+def _get_mmp(version):
+  return "%s.%s.%s" % (version.major, version.minor, version.patch)
+
+
+def _get_wheel_type(version):
+  if version.version_type == NIGHTLY_VERSION:
+    return "nightly"
+  else:
+    return "snapshot"
+
+
+def _get_wheel_build_date(version):
+  date_match = re.search(".*dev([0-9]{8}).*", version.identifier_string)
+  if date_match:
+    return date_match.group(1)
+  return ""
+
+
+def _get_wheel_version_suffix(version):
+  return version.identifier_string.replace(
+      "-dev{}".format(_get_wheel_build_date(version)), ""
  )


-def update_m1_builds(old_version, new_version):
-  """Update M1 builds."""
+def update_tf_version_bzl(old_version, new_version):
+  """Update tf_version.bzl."""
+  old_mmp = _get_mmp(old_version)
+  new_mmp = _get_mmp(new_version)
  replace_string_in_line(
-      "RELEASE_BRANCH = 'r%s.%s'" % (old_version.major, old_version.minor),
-      "RELEASE_BRANCH = 'r%s.%s'" % (new_version.major, new_version.minor),
-      TF_MAC_ARM64_CI_BUILD,
+      'TF_VERSION = "%s"' % old_mmp,
+      'TF_VERSION = "%s"' % new_mmp,
+      TF_VERSION_BZL,
  )
+
+
+def update_bazelrc(old_version, new_version):
+  """Update .bazelrc."""
+  old_wheel_type = _get_wheel_type(old_version)
+  new_wheel_type = _get_wheel_type(new_version)
  replace_string_in_line(
-      "RELEASE_BRANCH = 'r%s.%s'" % (old_version.major, old_version.minor),
-      "RELEASE_BRANCH = 'r%s.%s'" % (new_version.major, new_version.minor),
-      TF_MAC_ARM64_CI_TEST,
+      'build --repo_env=ML_WHEEL_TYPE="%s"' % old_wheel_type,
+      'build --repo_env=ML_WHEEL_TYPE="%s"' % new_wheel_type,
+      BAZEL_RC,
+  )
+
+  old_wheel_build_date = _get_wheel_build_date(old_version)
+  new_wheel_build_date = _get_wheel_build_date(new_version)
+  replace_string_in_line(
+      'build --repo_env=ML_WHEEL_BUILD_DATE="%s"' % old_wheel_build_date,
+      'build --repo_env=ML_WHEEL_BUILD_DATE="%s"' % new_wheel_build_date,
+      BAZEL_RC,
+  )
+
+  old_wheel_suffix = _get_wheel_version_suffix(old_version)
+  new_wheel_suffix = _get_wheel_version_suffix(new_version)
+  replace_string_in_line(
+      'build --repo_env=ML_WHEEL_VERSION_SUFFIX="%s"' % old_wheel_suffix,
+      'build --repo_env=ML_WHEEL_VERSION_SUFFIX="%s"' % new_wheel_suffix,
+      BAZEL_RC,
  )


@ -334,14 +364,11 @@ def main():
                            "-dev" + time.strftime("%Y%m%d"),
                            NIGHTLY_VERSION)
  else:
-    new_version = Version.parse_from_string(args.version, REGULAR_VERSION)
-    # Update Apple Silicon release CI files for release builds only
-    update_m1_builds(old_version, new_version)
+    new_version = Version.parse_from_string(args.version, SNAPSHOT_VERSION)

-  update_release_version_h(old_version, new_version)
-  update_setup_dot_py(old_version, new_version)
+  update_tf_version_bzl(old_version, new_version)
+  update_bazelrc(old_version, new_version)
  update_readme(old_version, new_version)
-  update_tensorflow_bzl(old_version, new_version)

  # Print transition details.
  print("Major: %s -> %s" % (old_version.major, new_version.major))
--- a/tensorflow/tools/pip_package/BUILD
+++ b/tensorflow/tools/pip_package/BUILD
@ -14,6 +14,7 @@ load(
 load("@local_tsl//third_party/py:python_wheel.bzl", "collect_data_files", "transitive_py_deps")
 load("@local_xla//xla/tsl/mkl:build_defs.bzl", "if_enable_mkl", "if_mkl", "if_mkl_ml")
 load("//tensorflow:tensorflow.bzl", "if_wheel_dependency", "if_with_tpu_support", "transitive_hdrs")
+load("//tensorflow:tf_version.bzl", "TF_SEMANTIC_VERSION_SUFFIX", "TF_VERSION")
 load(
    "//tensorflow/core/platform:build_config_root.bzl",
    "if_pywrap",
@ -268,6 +269,27 @@ transitive_py_deps(
    deps = COMMON_PIP_DEPS,
 )

+genrule(
+    name = "setup_py",
+    srcs = ["setup.py.tpl"],
+    outs = ["setup.py"],
+    cmd = """sed -E "s/_VERSION = '0.0.0'/_VERSION = '{wheel_version}{wheel_version_suffix}'/" \
+$(location setup.py.tpl) > $@;""".format(
+        wheel_version = TF_VERSION,
+        wheel_version_suffix = TF_SEMANTIC_VERSION_SUFFIX,
+    ),
+)
+
+py_binary(
+    name = "setup_py_binary",
+    srcs = [":setup_py"],
+    data = [
+        "//:AUTHORS",
+        "//:LICENSE",
+    ],
+    main = "setup.py",
+)
+
 tf_wheel(
    name = "wheel",
    headers = [
@ -290,7 +312,7 @@ tf_wheel(
    source_files = [
        "MANIFEST.in",
        "//tensorflow/tools/pip_package:THIRD_PARTY_NOTICES.txt",
-        "setup.py",
+        ":setup_py",
        ":cc_deps",
        ":py_deps",
        "//tensorflow:tf_python_api_gen_v2",
--- a/tensorflow/tools/pip_package/setup.py.tpl
+++ b/tensorflow/tools/pip_package/setup.py.tpl
@ -43,12 +43,16 @@ from setuptools.command.install import install as InstallCommandBase
 from setuptools.dist import Distribution


+# A genrule //tensorflow/tools/pip_package:setup_py replaces dummy string by
+# by the data provided in //tensorflow/tf_version.bzl.
+# The version suffix can be set by passing the build parameters
+# --repo_env=ML_WHEEL_BUILD_DATE=<date> and
+# --repo_env=ML_WHEEL_VERSION_SUFFIX=<suffix>.
+# To update the project version, update tf_version.bzl.
 # This version string is semver compatible, but incompatible with pip.
 # For pip, we will remove all '-' characters from this string, and use the
 # result for pip.
-# Also update tensorflow/tensorflow.bzl and
-# tensorflow/core/public/release_version.h
-_VERSION = '2.20.0'
+_VERSION = '0.0.0'

 # Update this version when a new libtpu stable version is released.
 LATEST_RELEASE_LIBTPU_VERSION = '0.0.10'
--- a/tensorflow/tools/pip_package/utils/tf_wheel.bzl
+++ b/tensorflow/tools/pip_package/utils/tf_wheel.bzl
@ -34,7 +34,11 @@ load(
    "WHEEL_COLLAB",
    "WHEEL_NAME",
 )
-load("//tensorflow:tensorflow.bzl", "VERSION", "WHEEL_VERSION")
+load(
+    "//tensorflow:tf_version.bzl",
+    "TF_VERSION",
+    "TF_WHEEL_VERSION_SUFFIX",
+)

 def _get_wheel_platform_name(platform_name, platform_tag):
    macos_platform_version = "{}_".format(MACOSX_DEPLOYMENT_TARGET.replace(".", "_")) if MACOSX_DEPLOYMENT_TARGET else ""
@ -49,14 +53,19 @@ def _get_wheel_platform_name(platform_name, platform_tag):
        platform_version = macos_platform_version,
    )

-def _get_full_wheel_name(platform_name, platform_tag):
+def _get_full_wheel_name(
+        platform_name,
+        platform_tag,
+        wheel_version):
    python_version = HERMETIC_PYTHON_VERSION.replace(".", "")
-    wheel_version = WHEEL_VERSION.replace("-dev", ".dev").replace("-", "")
    return "{wheel_name}-{wheel_version}-cp{python_version}-cp{python_version}-{wheel_platform_tag}.whl".format(
        wheel_name = WHEEL_NAME,
        wheel_version = wheel_version,
        python_version = python_version,
-        wheel_platform_tag = _get_wheel_platform_name(platform_name, platform_tag),
+        wheel_platform_tag = _get_wheel_platform_name(
+            platform_name,
+            platform_tag,
+        ),
    )

 def _is_dest_file(basename, dest_files_suffixes):
@ -75,9 +84,11 @@ def _tf_wheel_impl(ctx):
             " `--@local_config_cuda//cuda:override_include_cuda_libs=true`.")
    executable = ctx.executable.wheel_binary

+    full_wheel_version = (TF_VERSION + TF_WHEEL_VERSION_SUFFIX)
    full_wheel_name = _get_full_wheel_name(
        platform_name = ctx.attr.platform_name,
        platform_tag = ctx.attr.platform_tag,
+        wheel_version = full_wheel_version,
    )
    wheel_dir_name = "wheel_house"
    output_file = ctx.actions.declare_file("{wheel_dir}/{wheel_name}".format(
@ -93,7 +104,7 @@ def _tf_wheel_impl(ctx):
    ))
    args.add("--collab", str(WHEEL_COLLAB))
    args.add("--output-name", wheel_dir)
-    args.add("--version", VERSION)
+    args.add("--version", TF_VERSION)

    headers = ctx.files.headers[:]
    for f in headers: