1+ #! /bin/bash
2+ # Copyright 2024 The JAX Authors.
3+ #
4+ # Licensed under the Apache License, Version 2.0 (the "License");
5+ # you may not use this file except in compliance with the License.
6+ # You may obtain a copy of the License at
7+ #
8+ # http://www.apache.org/licenses/LICENSE-2.0
9+ #
10+ # Unless required by applicable law or agreed to in writing, software
11+ # distributed under the License is distributed on an "AS IS" BASIS,
12+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ # See the License for the specific language governing permissions and
14+ # limitations under the License.
15+ # ==============================================================================
16+ # Set up the build environment for JAX CI jobs. This script depends on the
17+ # "JAXCI_" environment variables set or sourced in the build script.
18+
19+ # Pre-emptively mark the JAX git directory as safe. This is necessary for JAX CI
20+ # jobs running on Linux runners in GitHub Actions. Without this, git complains
21+ # that the directory has dubious ownership and refuses to run any commands.
22+ # Avoid running on Windows runners as git runs into issues with not being able
23+ # to lock the config file. Other git commands seem to work on the Windows
24+ # runners so we can skip this step for Windows.
25+ # TODO(b/375073267): Remove this once we understand why git repositories are
26+ # being marked as unsafe inside the self-hosted runners.
27+ if [[ ! $( uname -s) =~ " MSYS_NT" ; then
28+ git config --global --add safe.directory $JAXCI_JAX_GIT_DIR
29+ fi
30+
31+ function clone_main_xla() {
32+ echo " Cloning XLA at HEAD to $( pwd) "
33+ git clone --depth=1 https://github.com/openxla/xla.git $( pwd)
34+ export JAXCI_XLA_GIT_DIR=$( pwd)
35+ }
36+
37+ # Clone XLA at HEAD if required.
38+ if [[ " $JAXCI_CLONE_MAIN_XLA " == 1 ]]; then
39+ # Clone only if $(pwd)/xla does not exist to avoid failure on re-runs.
40+ if [[ ! -d $( pwd) ; then
41+ clone_main_xla
42+ else
43+ echo " JAXCI_CLONE_MAIN_XLA set but local XLA folder already exists: $( pwd) "
44+ # Set JAXCI_XLA_GIT_DIR if local XLA already exists
45+ export JAXCI_XLA_GIT_DIR=$( pwd)
46+ fi
47+ fi
48+
49+ # If a XLA commit is provided, check out XLA at that commit.
50+ if [[ ! -z " $JAXCI_XLA_COMMIT " ; then
51+ # Clone XLA at HEAD if a path to local XLA is not provided.
52+ if [[ -z " $JAXCI_XLA_GIT_DIR " ; then
53+ clone_main_xla
54+ fi
55+ pushd " $JAXCI_XLA_GIT_DIR "
56+
57+ git fetch --depth=1 origin " $JAXCI_XLA_COMMIT "
58+ echo " JAXCI_XLA_COMMIT is set. Checking out XLA at $JAXCI_XLA_COMMIT "
59+ git checkout " $JAXCI_XLA_COMMIT "
60+
61+ popd
62+ fi
63+
64+ if [[ ! -z ${JAXCI_XLA_GIT_DIR} ]]; then
65+ echo " INFO: Overriding XLA to be read from $JAXCI_XLA_GIT_DIR instead of the"
66+ echo " pinned version in the WORKSPACE."
67+ echo " If you would like to revert this behavior, unset JAXCI_CLONE_MAIN_XLA"
68+ echo " and JAXCI_XLA_COMMIT in your environment. Note that the Bazel RBE test"
69+ echo " commands overrides the XLA repository and thus require a local copy of"
70+ echo " XLA to run."
71+ fi
0 commit comments