Skip to content

Commit b2c09f9

Browse files
StebossSteboss
committed
write the cache to registry
1 parent c46f0f5 commit b2c09f9

File tree

5 files changed

+43
-129
lines changed

5 files changed

+43
-129
lines changed

‎.github/actions/build-container/__pycache__/generate_bake.cpython-312.pyc‎

5.7 KB
Binary file not shown.

‎.github/actions/build-container/action.yml‎

Lines changed: 10 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -84,26 +84,14 @@ runs:
8484
run: |
8585
echo 'UPLD_IMAGE=ghcr.io/nvidia/jax-toolbox-internal' >> $GITHUB_ENV
8686
echo "BADGE_FILENAME_FULL=${{ inputs.BADGE_FILENAME }}-${{ inputs.ARCHITECTURE }}.json" >> $GITHUB_ENV
87-
echo "BAZEL_EXPORT_DIR=${RUNNER_TEMP}/bazel-cache-export" >> $GITHUB_ENV
8887
echo "BAKE_FILE=${RUNNER_TEMP}/docker-bake.json" >> $GITHUB_ENV
89-
echo "BUILDX_GHA_SCOPE=${{ inputs.CONTAINER_NAME }}-${{ inputs.ARCHITECTURE }}" >> $GITHUB_ENV
90-
91-
- name: Prepare Bazel cache contexts
92-
shell: bash
93-
run: |
94-
repo_context="${{ inputs.bazel-repo-cache-path }}"
95-
if [[ -z "${repo_context}" ]]; then
96-
repo_context="$(mktemp -d)"
97-
fi
98-
disk_context="${{ inputs.bazel-disk-cache-path }}"
99-
if [[ -z "${disk_context}" ]]; then
100-
disk_context="$(mktemp -d)"
88+
echo "BUILDKIT_CACHE_REPO=ghcr.io/nvidia/jax-toolbox-internal-cache" >> $GITHUB_ENV
89+
echo "BUILDKIT_CACHE_KEY_PREFIX=${{ inputs.CONTAINER_NAME }}-${{ inputs.ARCHITECTURE }}" >> $GITHUB_ENV
90+
cache_branch="$(printf '%s' "${GITHUB_REF_NAME}" | tr '/:@ ' '-' | tr -cd '[:alnum:]._-')"
91+
if [[ -z "${cache_branch}" ]]; then
92+
cache_branch="detached-head"
10193
fi
102-
mkdir -p "${repo_context}"
103-
mkdir -p "${disk_context}"
104-
rm -rf "${RUNNER_TEMP}/bazel-cache-export"
105-
echo "BAZEL_REPO_CONTEXT=${repo_context}" >> $GITHUB_ENV
106-
echo "BAZEL_DISK_CONTEXT=${disk_context}" >> $GITHUB_ENV
94+
echo "BUILDKIT_CACHE_BRANCH=${cache_branch}" >> $GITHUB_ENV
10795
10896
- name: Setup SSH
10997
id: setup-ssh
@@ -168,23 +156,21 @@ runs:
168156
BAZEL_CACHE: ${{ inputs.bazel-remote-cache-url }}
169157
BUILD_DATE: ${{ inputs.BUILD_DATE }}
170158
EXTRA_BUILD_ARGS: ${{ inputs.EXTRA_BUILD_ARGS }}
171-
BAZEL_REPO_CONTEXT: ${{ env.BAZEL_REPO_CONTEXT }}
172-
BAZEL_DISK_CONTEXT: ${{ env.BAZEL_DISK_CONTEXT }}
173159
SSH_KNOWN_HOSTS_FILE: ${{ steps.setup-ssh.outputs.known-hosts-file }}
174-
BAZEL_EXPORT_DIR: ${{ env.BAZEL_EXPORT_DIR }}
175-
BUILDX_GHA_SCOPE: ${{ env.BUILDX_GHA_SCOPE }}
160+
BUILDKIT_CACHE_REPO: ${{ env.BUILDKIT_CACHE_REPO }}
161+
BUILDKIT_CACHE_KEY_PREFIX: ${{ env.BUILDKIT_CACHE_KEY_PREFIX }}
162+
BUILDKIT_CACHE_BRANCH: ${{ env.BUILDKIT_CACHE_BRANCH }}
176163
BAKE_FILE: ${{ env.BAKE_FILE }}
177164
run: python3 "${{ github.action_path }}/generate_bake.py"
178165

179-
- name: Build mealkit/final and optionally export Bazel caches
166+
- name: Build mealkit/final
180167
id: final-build
181168
uses: docker/bake-action@v5
182169
with:
183170
files: ${{ env.BAKE_FILE }}
184171
targets: |
185172
mealkit
186173
final
187-
${{ inputs.bazel-save-cache == 'true' && 'cache-export' || '' }}
188174
allow: |
189175
ssh
190176
@@ -202,19 +188,6 @@ runs:
202188
digest="unknown"
203189
fi
204190
echo "digest=${digest}" >> "$GITHUB_OUTPUT"
205-
# this step is needed for saving bazel cache
206-
- name: Copy exported Bazel caches to host
207-
if: ${{ inputs.bazel-save-cache == 'true' && success() }}
208-
shell: bash -euo pipefail {0}
209-
run: |
210-
if [[ -d "${{ env.BAZEL_EXPORT_DIR }}/bazel-repo" && -n "${{ inputs.bazel-repo-cache-path }}" ]]; then
211-
mkdir -p "${{ inputs.bazel-repo-cache-path }}"
212-
cp -a "${{ env.BAZEL_EXPORT_DIR }}/bazel-repo/." "${{ inputs.bazel-repo-cache-path }}/"
213-
fi
214-
if [[ -d "${{ env.BAZEL_EXPORT_DIR }}/bazel-disk" && -n "${{ inputs.bazel-disk-cache-path }}" ]]; then
215-
mkdir -p "${{ inputs.bazel-disk-cache-path }}"
216-
cp -a "${{ env.BAZEL_EXPORT_DIR }}/bazel-disk/." "${{ inputs.bazel-disk-cache-path }}/"
217-
fi
218191
219192
# SITREP GENERATION
220193
- name: Generate sitrep

‎.github/actions/build-container/generate_bake.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,11 @@ def normalize_dockerfile_path(docker_context: str, dockerfile: str) -> str:
6262
def main() -> int:
6363
"""Main function to generate the bake definition for the build container action
6464
Here we are creating the json file for docker/bake.
65-
The bake definition will have 3 targets:
65+
The bake definition will have 2 targets:
6666
- mealkit: this target will build the mealkit image and push it to the
6767
registry with the specified tags and labels
6868
- final: this target will build the final image and push it to the registry
6969
with the specified tags and labels
70-
- cache-export: this target will only export the build cache to be used in future builds
7170
"""
7271
try:
7372
mealkit_tags = non_empty_lines(getenv("MEALKIT_TAGS"))
@@ -82,12 +81,13 @@ def main() -> int:
8281
bazel_cache = getenv("BAZEL_CACHE")
8382
build_date = getenv("BUILD_DATE")
8483
extra_build_args = kv_lines_to_object(getenv("EXTRA_BUILD_ARGS"))
85-
bazel_repo_context = getenv("BAZEL_REPO_CONTEXT")
86-
bazel_disk_context = getenv("BAZEL_DISK_CONTEXT")
8784
ssh_known_hosts_file = getenv("SSH_KNOWN_HOSTS_FILE")
88-
bazel_export_dir = getenv("BAZEL_EXPORT_DIR")
89-
cache_scope = getenv("BUILDX_GHA_SCOPE")
85+
cache_repo = getenv("BUILDKIT_CACHE_REPO")
86+
cache_key_prefix = getenv("BUILDKIT_CACHE_KEY_PREFIX")
87+
cache_branch = getenv("BUILDKIT_CACHE_BRANCH")
9088
bake_file = getenv("BAKE_FILE")
89+
branch_cache_ref = f"{cache_repo}:{cache_key_prefix}-{cache_branch}"
90+
main_cache_ref = f"{cache_repo}:{cache_key_prefix}-main"
9191

9292
build_args: Dict[str, str] = {
9393
"BASE_IMAGE": base_image,
@@ -101,14 +101,13 @@ def main() -> int:
101101
"context": docker_context,
102102
"dockerfile": dockerfile,
103103
"platforms": [f"linux/{architecture}"],
104-
"contexts": {
105-
"bazel_repo": bazel_repo_context,
106-
"bazel_disk": bazel_disk_context,
107-
},
108104
"ssh": ["default"],
109105
"secret": [f"id=SSH_KNOWN_HOSTS,src={ssh_known_hosts_file}"],
110106
"args": build_args,
111-
"cache-from": [f"type=gha,scope={cache_scope}"],
107+
"cache-from": [
108+
f"type=registry,ref={branch_cache_ref},ignore-error=true",
109+
f"type=registry,ref={main_cache_ref},ignore-error=true",
110+
],
112111
}
113112
# create the json file for the bake action
114113
bake_definition = {
@@ -126,11 +125,12 @@ def main() -> int:
126125
"tags": final_tags,
127126
"labels": final_labels,
128127
"output": ["type=image,push=true"],
129-
},
130-
"cache-export": {
131-
**common,
132-
"target": "cache-export",
133-
"output": [f"type=local,dest={bazel_export_dir}"],
128+
"cache-to": [
129+
"type=registry,"
130+
f"ref={branch_cache_ref},"
131+
"mode=max,oci-mediatypes=true,image-manifest=true,"
132+
"compression=zstd,ignore-error=true"
133+
],
134134
},
135135
}
136136
}

‎.github/container/Dockerfile.jax‎

Lines changed: 17 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,15 @@ ARG SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
1818
ARG GIT_USER_NAME="JAX Toolbox"
1919
ARG GIT_USER_EMAIL=jax@nvidia.com
2020

21-
ARG BAZEL_CACHE=/tmp
21+
ARG BAZEL_CACHE=""
2222
ARG BUILD_DATE
2323

2424
###############################################################################
2525
## Build JAX
2626
###############################################################################
2727

2828
FROM ${BASE_IMAGE} AS builder
29+
ARG TARGETARCH
2930
ARG URLREF_JAX
3031
ARG URLREF_TRANSFORMER_ENGINE
3132
ARG URLREF_XLA
@@ -51,28 +52,30 @@ ADD build-jax.sh build-te.sh local_cuda_arch /usr/local/bin/
5152
RUN ARCH="$(dpkg --print-architecture)" && \
5253
wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-${ARCH} && \
5354
chmod +x /usr/local/bin/bazel
54-
# Seed Bazel repository cache from build context supplied by the GitHub workflow.
55-
RUN --mount=type=bind,from=bazel_repo,target=/cache/bazel-repo-seed,ro \
56-
mkdir -p /cache/bazel-repo && \
57-
cp -a /cache/bazel-repo-seed/. /cache/bazel-repo/ 2>/dev/null || true
58-
# Seed Bazel disk cache from build context supplied by the GitHub workflow.
59-
RUN --mount=type=bind,from=bazel_disk,target=/cache/bazel-disk-seed,ro \
60-
mkdir -p /cache/bazel-disk && \
61-
cp -a /cache/bazel-disk-seed/. /cache/bazel-disk/ 2>/dev/null || true
6255
# Populate ${BUILD_PATH_JAXLIB} with editable wheels; --no-install because
6356
# (a) this is the builder stage, and (b) pip-finalize.sh does the install
64-
RUN mkdir -p /builder/extra-targets /cache/bazel-repo /cache/bazel-disk && \
65-
build-jax.sh \
66-
--bazel-cache ${BAZEL_CACHE} \
57+
RUN --mount=type=cache,id=bazel-repo-${TARGETARCH},target=/cache/bazel-repo,sharing=locked \
58+
--mount=type=cache,id=bazel-disk-${TARGETARCH},target=/cache/bazel-disk,sharing=locked \
59+
<<"EOF" bash -euo pipefail
60+
cache_args=()
61+
if [[ -n "${BAZEL_CACHE}" ]]; then
62+
cache_args+=(--bazel-cache "${BAZEL_CACHE}")
63+
fi
64+
65+
mkdir -p /builder/extra-targets
66+
build-jax.sh \
67+
"${cache_args[@]}" \
6768
--build-path-jaxlib ${BUILD_PATH_JAXLIB} \
6869
--extra-targets "${EXTRA_BAZEL_TARGETS}" \
6970
--extra-target-dest /builder/extra-targets \
7071
--no-install \
7172
--src-path-jax ${SRC_PATH_JAX} \
7273
--src-path-xla ${SRC_PATH_XLA} \
7374
--clean \
74-
${EXTRA_BUILD_JAX_ARGS} && \
75-
chmod -R a+rwX /cache/bazel-repo /cache/bazel-disk
75+
--build-param=--bazel_options=--repository_cache=/cache/bazel-repo \
76+
--build-param=--bazel_options=--disk_cache=/cache/bazel-disk \
77+
${EXTRA_BUILD_JAX_ARGS}
78+
EOF
7679

7780
## Transformer engine: check out source and build wheel
7881
RUN --mount=type=ssh \
@@ -147,11 +150,3 @@ EOF
147150

148151
FROM mealkit AS final
149152
RUN pip-finalize.sh
150-
151-
###############################################################################
152-
## Export Bazel caches for GitHub Actions cache/save
153-
###############################################################################
154-
155-
FROM scratch AS cache-export
156-
COPY --from=builder /cache/bazel-repo/ /bazel-repo/
157-
COPY --from=builder /cache/bazel-disk/ /bazel-disk/

‎.github/workflows/_ci.yaml‎

Lines changed: 0 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -67,41 +67,6 @@ jobs:
6767
steps:
6868
- name: Checkout repository
6969
uses: actions/checkout@v4
70-
- name: Detect Dockerfile.jax change in HEAD commit
71-
id: dockerfile-jax-change
72-
shell: bash -euo pipefail {0}
73-
run: |
74-
if git show --name-only --pretty='' HEAD | grep -Fxq '.github/container/Dockerfile.jax'; then
75-
echo "changed=true" >> "$GITHUB_OUTPUT"
76-
else
77-
echo "changed=false" >> "$GITHUB_OUTPUT"
78-
fi
79-
- name: Restore Bazel repository cache
80-
uses: actions/cache/restore@v5
81-
id: cache-restore-repo
82-
with:
83-
path: |
84-
~/.cache/bazel-repo
85-
key: bazel-repo-${{ inputs.ARCHITECTURE }}-${{ github.ref_name }}-${{ hashFiles('.github/container/build-jax.sh', '.github/container/Dockerfile.jax', '.github/actions/build-container/action.yml') }}
86-
restore-keys: |
87-
bazel-repo-${{ inputs.ARCHITECTURE }}-${{ github.ref_name }}-${{ hashFiles('.github/container/build-jax.sh', '.github/container/Dockerfile.jax', '.github/actions/build-container/action.yml') }}
88-
bazel-repo-${{ inputs.ARCHITECTURE }}-${{ github.ref_name }}-
89-
bazel-repo-${{ inputs.ARCHITECTURE }}-main-
90-
bazel-repo-${{ inputs.ARCHITECTURE }}-
91-
- name: Restore Bazel disk cache
92-
uses: actions/cache/restore@v5
93-
id: cache-restore-disk
94-
if: ${{ github.ref_name == 'main' && steps.dockerfile-jax-change.outputs.changed == 'true' }}
95-
with:
96-
path: |
97-
~/.cache/bazel-disk
98-
key: bazel-disk-${{ inputs.ARCHITECTURE }}-main-${{ hashFiles('.github/container/build-jax.sh', '.github/container/Dockerfile.jax', '.github/actions/build-container/action.yml') }}
99-
restore-keys: |
100-
bazel-disk-${{ inputs.ARCHITECTURE }}-main-
101-
- name: Create Bazel cache directories - if missing
102-
run: |
103-
mkdir -p ~/.cache/bazel-repo
104-
mkdir -p ~/.cache/bazel-disk
10570
- name: Build JAX container
10671
id: build-jax
10772
uses: ./.github/actions/build-container
@@ -117,30 +82,11 @@ jobs:
11782
ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
11883
ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
11984
github-token: ${{ secrets.GITHUB_TOKEN }}
120-
bazel-remote-cache-url: ${{ github.ref_name == 'main' && steps.dockerfile-jax-change.outputs.changed == 'true' && '/cache/bazel-disk' || '' }}
121-
bazel-repo-cache-path: /home/runner/.cache/bazel-repo
122-
bazel-disk-cache-path: /home/runner/.cache/bazel-disk
123-
bazel-save-cache: ${{ steps.cache-restore-repo.outputs.cache-hit != 'true' || (github.ref_name == 'main' && steps.dockerfile-jax-change.outputs.changed == 'true' && steps.cache-restore-disk.outputs.cache-hit != 'true') }}
12485
EXTRA_BUILD_ARGS: |
12586
URLREF_JAX=${{ fromJson(inputs.SOURCE_URLREFS).JAX }}
12687
URLREF_XLA=${{ fromJson(inputs.SOURCE_URLREFS).XLA }}
12788
URLREF_FLAX=${{ fromJson(inputs.SOURCE_URLREFS).FLAX }}
12889
URLREF_TRANSFORMER_ENGINE=${{ fromJson(inputs.SOURCE_URLREFS).TRANSFORMER_ENGINE }}
129-
EXTRA_BUILD_JAX_ARGS=--build-param=--bazel_options=--repository_cache=/cache/bazel-repo
130-
- name: Save Bazel repository cache
131-
uses: actions/cache/save@v5
132-
if: ${{ !cancelled() && steps.build-jax.outcome == 'success' && steps.cache-restore-repo.outputs.cache-hit != 'true' }}
133-
with:
134-
path: |
135-
~/.cache/bazel-repo
136-
key: bazel-repo-${{ inputs.ARCHITECTURE }}-${{ github.ref_name }}-${{ hashFiles('.github/container/build-jax.sh', '.github/container/Dockerfile.jax', '.github/actions/build-container/action.yml') }}
137-
- name: Save Bazel disk cache
138-
uses: actions/cache/save@v5
139-
if: ${{ !cancelled() && steps.build-jax.outcome == 'success' && github.ref_name == 'main' && steps.dockerfile-jax-change.outputs.changed == 'true' && steps.cache-restore-disk.outputs.cache-hit != 'true' }}
140-
with:
141-
path: |
142-
~/.cache/bazel-disk
143-
key: bazel-disk-${{ inputs.ARCHITECTURE }}-main-${{ hashFiles('.github/container/build-jax.sh', '.github/container/Dockerfile.jax', '.github/actions/build-container/action.yml') }}
14490
outputs:
14591
DOCKER_TAG_MEALKIT: ${{ steps.build-jax.outputs.DOCKER_TAG_MEALKIT }}
14692
DOCKER_TAG_FINAL: ${{ steps.build-jax.outputs.DOCKER_TAG_FINAL }}

0 commit comments

Comments
 (0)