Migrate TensorFlow on newest Hermetic Python set of rules

This makes Tensorflow consistent with JAX and XLA

PiperOrigin-RevId: 647832813
This commit is contained in:
Vadym Matsishevskyi 2024-06-28 16:54:53 -07:00 committed by TensorFlower Gardener
parent e851bbfeba
commit ec90bea08d
12 changed files with 90 additions and 435 deletions

View File

@ -12,6 +12,3 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# The requirements updater has its own WORKSPACE file
ci/official/requirements_updater

2
BUILD
View File

@ -1,4 +1,4 @@
exports_files([
exports_files(glob(["requirements*"]) + [
"configure",
"configure.py",
"ACKNOWLEDGEMENTS",

View File

@ -7,79 +7,12 @@ workspace(name = "org_tensorflow")
# We must initialize hermetic python first.
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
http_archive(
name = "bazel_skylib",
sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
urls = [
"https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
"https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
],
)
http_archive(
name = "rules_java",
sha256 = "c73336802d0b4882e40770666ad055212df4ea62cfa6edf9cb0f9d29828a0934",
url = "https://github.com/bazelbuild/rules_java/releases/download/5.3.5/rules_java-5.3.5.tar.gz",
)
http_archive(
name = "rules_python",
sha256 = "9d04041ac92a0985e344235f5d946f71ac543f1b1565f2cdbc9a2aaee8adf55b",
strip_prefix = "rules_python-0.26.0",
url = "https://github.com/bazelbuild/rules_python/releases/download/0.26.0/rules_python-0.26.0.tar.gz",
)
# buildifier: disable=same-origin-load
load("@rules_python//python:repositories.bzl", "py_repositories")
py_repositories()
load("@rules_python//python:repositories.bzl", "python_register_toolchains") # buildifier: disable=same-origin-load
load(
"//tensorflow/tools/toolchains/python:python_repo.bzl",
"python_repository",
)
python_repository(name = "python_version_repo")
load("@python_version_repo//:py_version.bzl", "TF_PYTHON_VERSION")
python_register_toolchains(
name = "python",
ignore_root_user_error = True,
python_version = TF_PYTHON_VERSION,
)
load("@python//:defs.bzl", "interpreter")
load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
NUMPY_ANNOTATIONS = {
"numpy": package_annotation(
additive_build_content = """\
filegroup(
name = "includes",
srcs = glob(["site-packages/numpy/core/include/**/*.h"]),
)
cc_library(
name = "numpy_headers",
hdrs = [":includes"],
strip_include_prefix="site-packages/numpy/core/include/",
)
""",
),
}
pip_parse(
name = "pypi",
annotations = NUMPY_ANNOTATIONS,
python_interpreter_target = interpreter,
requirements = "//:requirements_lock_" + TF_PYTHON_VERSION.replace(".", "_") + ".txt",
)
load("@pypi//:requirements.bzl", "install_deps")
install_deps()
# Initialize the TensorFlow repository and all dependencies.
#
# The cascade of load() statements and tf_workspace?() calls works around the
@ -90,6 +23,36 @@ load("@//tensorflow:workspace3.bzl", "tf_workspace3")
tf_workspace3()
# Initialize hermetic Python
load("@local_xla//third_party/py:python_init_rules.bzl", "python_init_rules")
python_init_rules()
load("@local_xla//third_party/py:python_init_repositories.bzl", "python_init_repositories")
python_init_repositories(
default_python_version = "system",
requirements = {
"3.9": "//:requirements_lock_3_9.txt",
"3.10": "//:requirements_lock_3_10.txt",
"3.11": "//:requirements_lock_3_11.txt",
"3.12": "//:requirements_lock_3_12.txt",
},
)
load("@local_xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
python_init_toolchains()
load("@local_xla//third_party/py:python_init_pip.bzl", "python_init_pip")
python_init_pip()
load("@pypi//:requirements.bzl", "install_deps")
install_deps()
# End hermetic Python initialization
load("@//tensorflow:workspace2.bzl", "tf_workspace2")
tf_workspace2()

View File

@ -1 +0,0 @@
6.5.0

View File

@ -13,80 +13,17 @@
# limitations under the License.
# ==============================================================================
load("@python//3.10:defs.bzl", compile_pip_requirements_3_10 = "compile_pip_requirements")
load("@python//3.11:defs.bzl", compile_pip_requirements_3_11 = "compile_pip_requirements")
load("@python//3.12:defs.bzl", compile_pip_requirements_3_12 = "compile_pip_requirements")
load("@python//3.9:defs.bzl", compile_pip_requirements_3_9 = "compile_pip_requirements")
load("@updater_config_repository//:updater_config_repository.bzl", "REQUIREMENTS_FILE_NAME")
load("@python//:defs.bzl", "compile_pip_requirements")
load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
compile_pip_requirements_3_9(
name = "requirements_3_9",
extra_args = ["--allow-unsafe"],
requirements_in = REQUIREMENTS_FILE_NAME,
requirements_txt = "requirements_lock_3_9.txt",
)
compile_pip_requirements_3_10(
name = "requirements_3_10",
extra_args = ["--allow-unsafe"],
requirements_in = REQUIREMENTS_FILE_NAME,
requirements_txt = "requirements_lock_3_10.txt",
)
compile_pip_requirements_3_11(
name = "requirements_3_11",
extra_args = ["--allow-unsafe"],
requirements_in = REQUIREMENTS_FILE_NAME,
requirements_txt = "requirements_lock_3_11.txt",
)
compile_pip_requirements_3_12(
name = "requirements_3_12",
extra_args = ["--allow-unsafe"],
requirements_in = REQUIREMENTS_FILE_NAME,
requirements_txt = "requirements_lock_3_12.txt",
)
compile_pip_requirements_3_9(
name = "requirements_3_9_release",
compile_pip_requirements(
name = "requirements",
extra_args = [
"--allow-unsafe",
"-P keras-nightly",
"-P tb-nightly",
"--build-isolation",
"--rebuild",
],
generate_hashes = True,
requirements_in = "requirements.in",
requirements_txt = "requirements_lock_3_9.txt",
)
compile_pip_requirements_3_10(
name = "requirements_3_10_release",
extra_args = [
"--allow-unsafe",
"-P keras-nightly",
"-P tb-nightly",
],
requirements_in = "requirements.in",
requirements_txt = "requirements_lock_3_10.txt",
)
compile_pip_requirements_3_11(
name = "requirements_3_11_release",
extra_args = [
"--allow-unsafe",
"-P keras-nightly",
"-P tb-nightly",
],
requirements_in = "requirements.in",
requirements_txt = "requirements_lock_3_11.txt",
)
compile_pip_requirements_3_12(
name = "requirements_3_12_release",
extra_args = [
"--allow-unsafe",
"-P keras-nightly",
"-P tb-nightly",
],
requirements_in = "requirements.in",
requirements_txt = "requirements_lock_3_12.txt",
requirements_txt = REQUIREMENTS,
)

View File

@ -1,127 +1,71 @@
# Hermetic Python
## Managing hermetic Python
Hermetic Python allows not to rely on system-installed Python, and
system-installed Python packages. \
Instead, an independent Python toolchain is registered, ensuring the right
dependencies are always used. \
See https://github.com/bazelbuild/rules_python/ for more details.
To make sure that TensorFlow's build is reproducible, behaves uniformly across
supported platforms (Linux, Windows, MacOS) and is properly isolated from
specifics of a local system, we rely on hermetic Python (see
[rules_python](https://github.com/bazelbuild/rules_python)) for all build
and test commands executed via Bazel. This means that your system Python
installation will be ignored during the build and Python interpreter itself
as well as all the Python dependencies will be managed by bazel directly.
### Specifying the Python version
### Specifying Python version
Note: Only a number of minor Python versions are supported at any given time.
By default, the lowest supported version is used.
To set a different version, use the `TF_PYTHON_VERSION` environment variable,
e.g.
The hermetic Python version is controlled by `HERMETIC_PYTHON_VERSION`
environment variable, which could be setin one of the following ways:
```
export TF_PYTHON_VERSION=3.11
# Either add an entry to your `.bazelrc` file
build --repo_env=HERMETIC_PYTHON_VERSION=3.12
# OR pass it directly to your specific build command
bazel build <target> --repo_env=HERMETIC_PYTHON_VERSION=3.12
# OR set the environment variable globally in your shell:
export HERMETIC_PYTHON_VERSION=3.12
```
To specify the version via a Bazel command argument, use the following:
You may run builds and tests against different versions of Python sequentially
on the same machine by simply switching the value of `HERMETIC_PYTHON_VERSION`
between the runs. All the python-agnostic parts of the build cache from the
previous build will be preserved and reused for the subsequent builds.
### Specifying Python dependencies
During bazel build all TensorFlow's Python dependencies are pinned to their
specific versions. This is necessary to ensure reproducibility of the build.
The pinned versions of the full transitive closure of TensorFlow's dependencies
together with their corresponding hashes are specified in
`requirements_lock_<python version>.txt` files (e.g.
`requirements_lock_3_12.txt` for `Python 3.12`).
To update the lock files, make sure
`ci/official/requirements_updater/requirements.in` contains the desired direct
dependencies list and then execute the following command (which will call
[pip-compile](https://pypi.org/project/pip-tools/) under the hood):
```
--repo_env=TF_PYTHON_VERSION=3.11
bazel run //ci/official/requirements_updater:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12
```
## Requirements updater
where `3.12` is the `Python` version you wish to update.
Requirements updater is a standalone tool, intended to simplify process of
updating requirements for multiple minor versions of Python.
It takes in a file with a set of dependencies, and produces a more detailed
requirements file for each version, with hashes specified for each
dependency required, as well as their sub-dependencies.
### How to update/add requirements
By default, the name of the base requirements file is `requirements.in`, but it
can be set using the `REQUIREMENTS_FILE_NAME` variable. \
For example:
Note, since it is still `pip` and `pip-compile` tools used under the hood, so
most of the command line arguments and features supported by those tools will be
acknowledged by the Bazel requirements updater command as well. For example, if
you wish the updater to consider pre-release versions simply pass `--pre`
argument to the bazel command:
```
export REQUIREMENTS_FILE_NAME=my_requirements.in
bazel run //ci/official/requirements_updater:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12 -- --pre
```
To specify the file via a Bazel command argument, use the following:
If you need to upgrade all of the packages in requirements lock file, just pass
the `--upgrade` parameter:
```
--repo_env=REQUIREMENTS_FILE_NAME=my_requirements.in
bazel run //ci/official/requirements_updater:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12 -- --upgrade
```
### How to run the updater
```
bash updater.sh
```
## How to add a new Python version
Note: Updating the
[rules-python](https://github.com/bazelbuild/rules_python/releases) version may
be required before going through the steps below. This is due to the new Python
versions becoming available through `rules-python`. \
See
[here](https://github.com/tensorflow/tensorflow/commit/f91457f258fdd78f693044a57efa63a38335d1de),
and
[here](https://github.com/tensorflow/tensorflow/commit/052445e04ce20fd747657e0198a1bcec2b6dff5b),
for an example.
See
[this commit](https://github.com/tensorflow/tensorflow/commit/5f7f05a80aac9b01325a78ec3fcff0dbedb1cc23)
as a rough example of the steps below.
All the files referenced below are located in the same directory as this README,
unless indicated otherwise.
1) Add the new version to the `VERSIONS` variable inside
`tensorflow/tools/toolchains/python/python_repo.bzl`. \
While this isn't necessary for running the updater, it is required for
actually using the new version with Tensorflow.
2) In the `WORKSPACE` file, add the new version to the `python_versions`
parameter of the `python_register_multi_toolchains` function.
3) In the `BUILD.bazel` file, add a load statement for the new version, e.g.
```
load("@python//3.11:defs.bzl",
compile_pip_requirements_3_11 = "compile_pip_requirements")
```
Add a new entry for the loaded `compile_pip_requirements`, e.g.
```
compile_pip_requirements_3_11(
name = "requirements_3_11",
extra_args = ["--allow-unsafe"],
requirements_in = "requirements.in",
requirements_txt = "requirements_lock_3_11.txt",
)
```
```
compile_pip_requirements_3_11(
name = "requirements_3_11_release",
extra_args = [
"--allow-unsafe",
"-P keras-nightly",
"-P tb-nightly",
],
requirements_in = "requirements.in",
requirements_txt = "requirements_lock_3_11.txt",
)
```
4) Add the version to `SUPPORTED_VERSIONS` in `updater.sh`, and
`release_updater.sh`
5) Run the `updater.sh` shell script. \
If the base requirements file hasn't yet been updated to account for the new
Python version, which will require different versions for at least some
dependencies, it will need to be updated now, for the script to run
successfully.
6) A new `requirements_lock_3_11.txt` file should appear under the root of the
`tensorflow` directory.
For the full set of supported parameters please check
[pip-compile](https://pip-tools.readthedocs.io/en/latest/cli/pip-compile/)
documentation

View File

@ -1,55 +0,0 @@
# buildifier: disable=load-on-top
workspace(name = "requirements_updater")
# buildifier: disable=load-on-top
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
http_archive(
name = "bazel_skylib",
sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
urls = [
"https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
"https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
],
)
http_archive(
name = "rules_python",
sha256 = "9d04041ac92a0985e344235f5d946f71ac543f1b1565f2cdbc9a2aaee8adf55b",
strip_prefix = "rules_python-0.26.0",
url = "https://github.com/bazelbuild/rules_python/releases/download/0.26.0/rules_python-0.26.0.tar.gz",
)
# buildifier: disable=same-origin-load
load("@rules_python//python:repositories.bzl", "py_repositories")
py_repositories()
load("@rules_python//python:repositories.bzl", "python_register_multi_toolchains") # buildifier: disable=same-origin-load
load("@rules_python//python/pip_install:repositories.bzl", "pip_install_dependencies")
default_python_version = "3.10"
load(
"//:updater_config_repository.bzl",
"updater_config_repository", # @unused
)
updater_config_repository(name = "updater_config_repository")
python_register_multi_toolchains(
name = "python",
default_version = default_python_version,
ignore_root_user_error = True,
python_versions = [
"3.9",
"3.10",
"3.11",
"3.12",
],
register_coverage_tool = True,
)
pip_install_dependencies()

View File

@ -1,36 +0,0 @@
#!/usr/bin/env 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.
# ==============================================================================
# script to run pip-compile for keras, tensorboard, estimator deps.
# if there is a change in requirements.in then all lock files will be updated
# accordingly.
mv BUILD.bazel BUILD
SUPPORTED_VERSIONS=("3_9" "3_10" "3_11" "3_12")
for VERSION in "${SUPPORTED_VERSIONS[@]}"
do
cp ../../../requirements_lock_"$VERSION".txt "requirements_lock_"$VERSION".txt"
bazel run \
--experimental_convenience_symlinks=ignore \
--enable_bzlmod=false \
//:requirements_"$VERSION"_release.update
sed -i '/^#/d' requirements_lock_"$VERSION".txt
mv "requirements_lock_"$VERSION".txt" ../../../requirements_lock_"$VERSION".txt
done
mv BUILD BUILD.bazel

View File

@ -1,32 +0,0 @@
# TODO(b/304751256): Adjust the numpy pin to a single version, when ready
numpy ~= 1.23.5 ; python_version <= "3.11"
numpy ~= 1.26.0 ; python_version >= "3.12"
wheel ~= 0.41.2
h5py >= 3.10.0
lit ~= 17.0.2
opt_einsum == 3.3.0
astunparse == 1.6.3
dill == 0.3.7
astor == 0.7.1
typing_extensions == 4.8.0
gast == 0.4.0
termcolor == 2.3.0
wrapt == 1.16.0
tblib == 2.0.0
ml_dtypes >= 0.4.0, < 0.5.0
# Install tensorboard, and keras
# Note that here we want the latest version that matches TF major.minor version
# Note that we must use nightly here as these are used in nightly jobs
# For release jobs, we will pin these on the release branch
keras-nightly ~= 3.0.0.dev
tb-nightly ~= 2.18.0.a
# Test dependencies
grpcio >= 1.24.3, < 2.0
portpicker == 1.6.0
scipy == 1.11.3
requests >= 2.31.0
packaging==23.2
setuptools==68.2.2
jax==0.4.7

View File

@ -1,39 +0,0 @@
#!/usr/bin/env 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.
# ==============================================================================
# script to run pip-compile for each requirement.
# if there is a change in requirements.in then all lock files will be updated
# accordingly
# All commands run relative to this directory
cd "$(dirname "${BASH_SOURCE[0]}")"
mv BUILD.bazel BUILD
SUPPORTED_VERSIONS=("3_9" "3_10" "3_11" "3_12")
for VERSION in "${SUPPORTED_VERSIONS[@]}"
do
touch "requirements_lock_$VERSION.txt"
bazel run \
--experimental_convenience_symlinks=ignore \
--enable_bzlmod=false \
//:requirements_"$VERSION".update
sed -i '/^#/d' requirements_lock_"$VERSION".txt
mv requirements_lock_"$VERSION".txt ../../../requirements_lock_"$VERSION".txt
done
mv BUILD BUILD.bazel

View File

@ -1,22 +0,0 @@
'''
Repository rule to set some environment variables.
Can be set via build parameter "--repo_env=<VARIABLE_NAME>=<value>"
e.g "--repo_env=REQUIREMENTS_FILE_NAME=requirements.in"
List of variables:
REQUIREMENTS_FILE_NAME
'''
def _updater_config_repository_impl(repository_ctx):
repository_ctx.file("BUILD", "")
requirements_file_name = repository_ctx.os.environ.get("REQUIREMENTS_FILE_NAME", "requirements.in")
repository_ctx.file(
"updater_config_repository.bzl",
"REQUIREMENTS_FILE_NAME = \"%s\"" %
requirements_file_name,
)
updater_config_repository = repository_rule(
implementation = _updater_config_repository_impl,
environ = ["REQUIREMENTS_FILE_NAME"],
)

View File

@ -20,7 +20,6 @@
LATEST_BAZEL_VERSION=6.5.0
# LINT.ThenChange(
# //tensorflow/opensource_only/.bazelversion,
# //tensorflow/opensource_only/ci/official/requirements_updater/.bazelversion
# //tensorflow/tools/ci_build/install/install_bazel.sh,
# //tensorflow/tools/ci_build/install/install_bazel_from_source.sh,
# //tensorflow/tools/toolchains/cross_compile/cc/cc_toolchain_config.bzl)