Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Allow PR refs for jax or xla #116

Closed
wants to merge 6 commits into from
Closed

Allow PR refs for jax or xla #116

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e04678a
Adds ability to build xla or jax from a PR
terrykong Jul 5, 2023
7ee0d03
Updates ci.yaml to properly reflect updated dependencies where pax/t5x
terrykong Jul 5, 2023
64cc09c
Allow PRs building jax
terrykong Jul 11, 2023
365a276
Remove dependency change in ci.yaml since we're moving back to using jax
terrykong Jul 12, 2023
538fcc0
Allow parsing of PR numbers instead of the full PR ref
terrykong Jul 19, 2023
739ea48
Merge branch 'main' into allow-pr-refs
yhtang Jul 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 23 additions & 3 deletions .github/container/Dockerfile.jax
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# syntax=docker/dockerfile:1-labs
ARG BASE_IMAGE=ghcr.io/nvidia/jax-toolbox:base
ARG REPO_JAX="https://github.com/google/jax.git"
ARG REPO_XLA="https://github.com/openxla/xla.git"
Expand All @@ -21,10 +22,29 @@ ARG SRC_PATH_JAX
ARG SRC_PATH_XLA
ARG BAZEL_CACHE

RUN git clone "${REPO_JAX}" "${SRC_PATH_JAX}" && cd "${SRC_PATH_JAX}" && git checkout ${REF_JAX}
RUN <<"EOF" bash -e
git clone "${REPO_JAX}" "${SRC_PATH_JAX}"
cd "${SRC_PATH_JAX}"
orig_ref=${REF_JAX}
if [[ "${orig_ref}" =~ ^pull/ ]]; then
PR_ID=$(cut -d/ -f2 <<<"${orig_ref}")
REF_JAX=PR-${PR_ID}
git fetch origin ${orig_ref}:${REF_JAX}
fi
git checkout ${REF_JAX}
EOF
RUN --mount=type=ssh \
--mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
git clone "${REPO_XLA}" "${SRC_PATH_XLA}" && cd "${SRC_PATH_XLA}" && git checkout ${REF_XLA}
--mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts <<"EOF" bash -e
git clone "${REPO_XLA}" "${SRC_PATH_XLA}"
cd "${SRC_PATH_XLA}"
orig_ref=${REF_XLA}
if [[ "${orig_ref}" =~ ^pull/ ]]; then
PR_ID=$(cut -d/ -f2 <<<"${orig_ref}")
REF_XLA=PR-${PR_ID}
git fetch origin ${orig_ref}:${REF_XLA}
fi
git checkout ${REF_XLA}
EOF

ADD build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
RUN build-jax.sh \
Expand Down
10 changes: 7 additions & 3 deletions .github/workflows/ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,12 @@ on:
required: false
default: 'nvidia/cuda:12.1.1-devel-ubuntu22.04'
SRC_JAX:
description: 'JAX source: <repo>#<branch|tag|commit>'
description: 'JAX source: <repo>#<branch|tag|commit|PR no.>'
type: string
required: true
default: 'https://github.com/google/jax.git#main'
SRC_XLA:
description: 'XLA source: <repo>#<branch|tag|commit>'
description: 'XLA source: <repo>#<branch|tag|commit|PR no.>'
type: string
required: true
default: 'https://github.com/openxla/xla.git#main'
Expand Down Expand Up @@ -88,7 +88,11 @@ jobs:
DEFAULT="$3"
SRC="${INPUT:-${DEFAULT}}"
echo "REPO_${PACKAGE}=$(echo "${SRC}" | cut -f1 -d#)" >> $GITHUB_OUTPUT
echo "REF_${PACKAGE}=$(echo "${SRC}" | cut -f2 -d#)" >> $GITHUB_OUTPUT
GIT_REF="$(echo "${SRC}" | cut -f2 -d#)"
if [[ $GIT_REF =~ ^[1-9][0-9]*$ ]]; then
GIT_REF=pull/$GIT_REF/head
fi
echo "REF_${PACKAGE}=${GIT_REF}" >> $GITHUB_OUTPUT
}

# default values are for `pull_request`` event types
Expand Down
Loading