Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results. #25126
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_699315679
branch
2 times, most recently
from
c19049a to
e6a2100
Compare
copybara-service
bot
force-pushed
the
test_699315679
branch
from
e6a2100 to
86f8353
Compare
copybara-service
bot
changed the title
py_import targets.
copybara-service
bot
force-pushed
the
test_699315679
branch
3 times, most recently
from
677cd8e to
9228a76
Compare
copybara-service
bot
changed the title
py_import targets.
copybara-service
bot
force-pushed
the
test_699315679
branch
15 times, most recently
from
b67d004 to
2cd817d
Compare
copybara-service
bot
force-pushed
the
test_699315679
branch
7 times, most recently
from
05ade1b to
f96a38a
Compare
copybara-service
bot
force-pushed
the
test_699315679
branch
9 times, most recently
from
04bfda8 to
029b81e
Compare
…tain reproducible wheel content and filename results. This change is a part of the initiative to test the JAX wheels in the presubmit properly. The list of the changes: 1. JAX wheel build rule verifies that `--@local_config_cuda//cuda:include_cuda_libs=false` during the wheel build. There is a way to pass the restriction by providing `--@local_config_cuda//cuda:override_include_cuda_libs=true`. 2. The JAX version number (which is also used in the wheel filenames) is stored in `_version` variable in the file [version.py](https://github.com/jax-ml/jax/blob/main/jax/version.py). The custom repository rule `jax_python_wheel_version_repository` saves this value in `wheel_version.bzl`, so it becomes available in Bazel build phase. 3. The version suffix of the wheel in the build rule output depends on the environment variables. The version suffix chunks that are not reproducible shouldn’t be calculated as a part of the wheel binary: for example, the current date changes every day, thus the wheels built today and tomorrow on the same code version will be technically different. To maintain reproducible wheel content, we need to pass suffix chunks in a form of environment variables. 4. Environment variables combinations for creating wheels with different versions: * `0.5.1.dev0+selfbuilt` (local build, default build rule behavior): `--repo_env=ML_WHEEL_TYPE=snapshot` * `0.5.1` (release): `--repo_env=ML_WHEEL_TYPE=release` * `0.5.1rc1` (release candidate): `--repo_env=ML_WHEEL_TYPE=release --repo_env=ML_WHEEL_VERSION_SUFFIX=rc1` * `0.5.1.dev20250128+3e75e20c7` (nightly build): `--repo_env=ML_WHEEL_TYPE=custom --repo_env=ML_WHEEL_BUILD_DATE=20250128 --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)` PiperOrigin-RevId: 723552265
copybara-service
bot
force-pushed
the
test_699315679
branch
from
029b81e to
d424f5b
Compare
copybara-service bot
pushed a commit
that referenced
this pull request
Feb 13, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in #25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
that referenced
this pull request
Feb 13, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in #25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in #25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in #25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in #25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to google/tsl
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in #25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to google/tsl
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Feb 14, 2025
This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in jax-ml/jax#25126) Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file). You can still build the JAX wheel with `python3 -m build` command. Bazel command example for building nightly JAX wheel: ``` bazel build :jax_wheel \ --config=ci_linux_x86_64 \ --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --repo_env=ML_WHEEL_TYPE=custom \ --repo_env=ML_WHEEL_BUILD_DATE=20250211 \ --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD) ``` Resulting wheel: ``` bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl ``` PiperOrigin-RevId: 724102315
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.
This change is a part of the initiative to test the JAX wheels in the presubmit properly.
The list of the changes:
JAX wheel build rule verifies that
--@local_config_cuda//cuda:include_cuda_libs=falseduring the wheel build. There is a way to pass the restriction by providing--@local_config_cuda//cuda:override_include_cuda_libs=true.The JAX version number (which is also used in the wheel filenames) is stored in
_versionvariable in the file version.py. The custom repository rulejax_python_wheel_version_repositorysaves this value inwheel_version.bzl, so it becomes available in Bazel build phase.The version suffix of the wheel in the build rule output depends on the environment variables.
The version suffix chunks that are not reproducible shouldn’t be calculated as a part of the wheel binary: for example, the current date changes every day, thus the wheels built today and tomorrow on the same code version will be technically different. To maintain reproducible wheel content, we need to pass suffix chunks in a form of environment variables.
Environment variables combinations for creating wheels with different versions:
0.5.1.dev0+selfbuilt(local build, default build rule behavior):--repo_env=ML_WHEEL_TYPE=snapshot0.5.1(release):--repo_env=ML_WHEEL_TYPE=release0.5.1rc1(release candidate):--repo_env=ML_WHEEL_TYPE=release --repo_env=ML_WHEEL_VERSION_SUFFIX=rc10.5.1.dev20250128+3e75e20c7(nightly build):--repo_env=ML_WHEEL_TYPE=custom --repo_env=ML_WHEEL_BUILD_DATE=20250128 --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)