mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 11:56:07 +00:00
d424f5b5b3
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
126 lines
3.3 KiB
Python
126 lines
3.3 KiB
Python
# The XLA commit is determined by third_party/xla/workspace.bzl.
|
|
load("//third_party/xla:workspace.bzl", jax_xla_workspace = "repo")
|
|
jax_xla_workspace()
|
|
|
|
# Initialize hermetic Python
|
|
load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
|
|
python_init_rules()
|
|
|
|
load("@xla//third_party/py:python_init_repositories.bzl", "python_init_repositories")
|
|
python_init_repositories(
|
|
requirements = {
|
|
"3.10": "//build:requirements_lock_3_10.txt",
|
|
"3.11": "//build:requirements_lock_3_11.txt",
|
|
"3.12": "//build:requirements_lock_3_12.txt",
|
|
"3.13": "//build:requirements_lock_3_13.txt",
|
|
"3.13-ft": "//build:requirements_lock_3_13_ft.txt",
|
|
},
|
|
local_wheel_inclusion_list = [
|
|
"jaxlib*",
|
|
"jax_cuda*",
|
|
"jax-cuda*",
|
|
],
|
|
local_wheel_workspaces = ["//jaxlib:jax.bzl"],
|
|
local_wheel_dist_folder = "../dist",
|
|
default_python_version = "system",
|
|
)
|
|
|
|
load("@xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
|
|
python_init_toolchains()
|
|
|
|
load("@xla//third_party/py:python_init_pip.bzl", "python_init_pip")
|
|
python_init_pip()
|
|
|
|
load("@pypi//:requirements.bzl", "install_deps")
|
|
install_deps()
|
|
|
|
# Optional, to facilitate testing against newest versions of Python
|
|
load("@xla//third_party/py:python_repo.bzl", "custom_python_interpreter")
|
|
custom_python_interpreter(
|
|
name = "python_dev",
|
|
urls = ["https://www.python.org/ftp/python/{version}/Python-{version_variant}.tgz"],
|
|
strip_prefix = "Python-{version_variant}",
|
|
version = "3.13.0",
|
|
version_variant = "3.13.0rc2",
|
|
)
|
|
|
|
load("@xla//:workspace4.bzl", "xla_workspace4")
|
|
xla_workspace4()
|
|
|
|
load("@xla//:workspace3.bzl", "xla_workspace3")
|
|
xla_workspace3()
|
|
|
|
load("@xla//:workspace2.bzl", "xla_workspace2")
|
|
xla_workspace2()
|
|
|
|
load("@xla//:workspace1.bzl", "xla_workspace1")
|
|
xla_workspace1()
|
|
|
|
load("@xla//:workspace0.bzl", "xla_workspace0")
|
|
xla_workspace0()
|
|
|
|
load("//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
|
|
flatbuffers()
|
|
|
|
load("//jaxlib:jax_python_wheel.bzl", "jax_python_wheel_repository")
|
|
jax_python_wheel_repository(
|
|
name = "jax_wheel",
|
|
version_key = "_version",
|
|
version_source = "//jax:version.py",
|
|
)
|
|
|
|
load(
|
|
"@tsl//third_party/py:python_wheel.bzl",
|
|
"python_wheel_version_suffix_repository",
|
|
)
|
|
python_wheel_version_suffix_repository(
|
|
name = "jax_wheel_version_suffix",
|
|
)
|
|
|
|
load(
|
|
"@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
|
|
"cuda_json_init_repository",
|
|
)
|
|
|
|
cuda_json_init_repository()
|
|
|
|
load(
|
|
"@cuda_redist_json//:distributions.bzl",
|
|
"CUDA_REDISTRIBUTIONS",
|
|
"CUDNN_REDISTRIBUTIONS",
|
|
)
|
|
load(
|
|
"@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
|
|
"cuda_redist_init_repositories",
|
|
"cudnn_redist_init_repository",
|
|
)
|
|
|
|
cuda_redist_init_repositories(
|
|
cuda_redistributions = CUDA_REDISTRIBUTIONS,
|
|
)
|
|
|
|
cudnn_redist_init_repository(
|
|
cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
|
|
)
|
|
|
|
load(
|
|
"@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
|
|
"cuda_configure",
|
|
)
|
|
|
|
cuda_configure(name = "local_config_cuda")
|
|
|
|
load(
|
|
"@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
|
|
"nccl_redist_init_repository",
|
|
)
|
|
|
|
nccl_redist_init_repository()
|
|
|
|
load(
|
|
"@tsl//third_party/nccl/hermetic:nccl_configure.bzl",
|
|
"nccl_configure",
|
|
)
|
|
|
|
nccl_configure(name = "local_config_nccl")
|