mirror of
https://github.com/ROCm/jax.git
synced 2025-04-25 09:56:08 +00:00
721d1a3211
List of changes: 1. Allow us to build a RC wheel when building release artifacts. This is done by modifying the build CLI to use the new JAX build rule and passing in the build options that control the wheel tag. A new build argument `use_new_wheel_build_rule` is introduced to the build CLI to avoid breaking anyone that uses the CLI and the old build rule. Note that this option will go way in the future when the build CLI migrates fully to the new build rule. 2. Change the upload script to upload both rc and release tagged wheels (changes internal) PiperOrigin-RevId: 733464219
82 lines
4.0 KiB
Bash
82 lines
4.0 KiB
Bash
# Copyright 2024 The JAX Authors.
|
|
#
|
|
# 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.
|
|
# ==============================================================================
|
|
# This file contains all the default values for the "JAXCI_" environment
|
|
# variables used in the CI scripts. These variables are used to control the
|
|
# behavior of the CI scripts such as the Python version used, path to JAX/XLA
|
|
# repo, if to clone XLA repo, etc.
|
|
|
|
# The path to the JAX git repository.
|
|
export JAXCI_JAX_GIT_DIR=$(pwd)
|
|
|
|
# Controls the version of Hermetic Python to use. Use system default if not
|
|
# set.
|
|
export JAXCI_HERMETIC_PYTHON_VERSION=${JAXCI_HERMETIC_PYTHON_VERSION:-$(python3 -V | awk '{print $2}' | awk -F. '{print $1"."$2}')}
|
|
|
|
# Set JAXCI_XLA_GIT_DIR to the root of the XLA git repository to use a local
|
|
# copy of XLA instead of the pinned version in the WORKSPACE. When
|
|
# JAXCI_CLONE_MAIN_XLA=1, this gets set automatically.
|
|
export JAXCI_XLA_GIT_DIR=${JAXCI_XLA_GIT_DIR:-}
|
|
|
|
# If set to 1, the builds will clone the XLA repository at HEAD and set its
|
|
# path in JAXCI_XLA_GIT_DIR.
|
|
export JAXCI_CLONE_MAIN_XLA=${JAXCI_CLONE_MAIN_XLA:-0}
|
|
|
|
# Allows overriding the XLA commit that is used.
|
|
export JAXCI_XLA_COMMIT=${JAXCI_XLA_COMMIT:-}
|
|
|
|
# Controls the location where the artifacts are written to.
|
|
export JAXCI_OUTPUT_DIR="$(pwd)/dist"
|
|
|
|
# When enabled, artifacts will be built with RBE. Requires gcloud authentication
|
|
# and only certain platforms support RBE. Therefore, this flag is enabled only
|
|
# for CI builds where RBE is supported.
|
|
export JAXCI_BUILD_ARTIFACT_WITH_RBE=${JAXCI_BUILD_ARTIFACT_WITH_RBE:-0}
|
|
|
|
# On platforms where RBE is not supported, we use Bazel remote cache to speed up
|
|
# builds. When this flag is enabled, Bazel will also try to push new cache
|
|
# entries to the bucket. Since writes to the bucket require authentication, this
|
|
# flag is enabled only for CI builds.
|
|
export JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE=${JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE:-0}
|
|
|
|
# Type of artifacts to build. Valid values are "default", "release", "nightly".
|
|
# This affects the wheel naming/tag.
|
|
export JAXCI_ARTIFACT_TYPE=${JAXCI_ARTIFACT_TYPE:-"default"}
|
|
|
|
# When building release artifacts, we build a release candidate wheel ("rc"
|
|
# tagged wheel) in addition to the release wheel. This environment variable
|
|
# sets the version of the release candidate ("RC") artifact to build.
|
|
export JAXCI_WHEEL_RC_VERSION=${JAXCI_WHEEL_RC_VERSION:-}
|
|
|
|
# #############################################################################
|
|
# Test script specific environment variables.
|
|
# #############################################################################
|
|
# Sets the value of `JAX_ENABLE_X64` in the test scripts. CI builds override
|
|
# this value in the Github action workflow files.
|
|
export JAXCI_ENABLE_X64=${JAXCI_ENABLE_X64:-0}
|
|
|
|
# Pytest specific environment variables below. Used in run_pytest_*.sh scripts.
|
|
# Sets the number of TPU cores for the TPU machine type. These values are
|
|
# defined in the TPU GitHub Actions workflow.
|
|
export JAXCI_TPU_CORES=${JAXCI_TPU_CORES:-}
|
|
|
|
# JAXCI_PYTHON points to the Python interpreter to use for installing JAX wheels
|
|
# on the system. By default, it is set to match the version of the hermetic
|
|
# Python used by Bazel for building the wheels.
|
|
export JAXCI_PYTHON=${JAXCI_PYTHON:-python${JAXCI_HERMETIC_PYTHON_VERSION}}
|
|
|
|
# Installs the JAX package in editable mode at the current commit. Enabled by
|
|
# default. Nightly/Release builds disable this flag in the Github action
|
|
# workflow files.
|
|
export JAXCI_INSTALL_JAX_CURRENT_COMMIT=${JAXCI_INSTALL_JAX_CURRENT_COMMIT:-"1"} |