Skip to content

Commit 2282e04

Browse files
SurbhiJainUSCSurbhiJainUSC
committed
Separate TPU and GPU nightly docker image build workflows
1 parent c7a5a39 commit 2282e04

2 files changed

Lines changed: 95 additions & 60 deletions

File tree

‎.github/workflows/gpu_nightly_images_pipeline.yml‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# Copyright 2023–2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# This workflow builds and pushes GPU MaxText images.
16+
# It runs automatically daily at 12am UTC, or manually via Workflow Dispatch.
17+
18+
name: Build GPU Nightly Docker Images
19+
20+
on:
21+
schedule:
22+
# Run the job daily at 12AM UTC
23+
- cron: '0 0 * * *'
24+
workflow_dispatch:
25+
inputs:
26+
image_suffix:
27+
description: 'An image suffix can be provided to add to the image name'
28+
required: false
29+
type: string
30+
default: ""
31+
32+
permissions:
33+
contents: read
34+
35+
jobs:
36+
build_maxtext_package:
37+
name: Build MaxText Package
38+
uses: ./.github/workflows/build_package.yml
39+
with:
40+
device_type: tpu
41+
device_name: v4-8
42+
cloud_runner: linux-x86-n2-16-buildkit
43+
44+
build_and_push_docker_images:
45+
name: Build ${{ matrix.name }} Docker Image
46+
needs: build_maxtext_package
47+
strategy:
48+
fail-fast: false
49+
matrix:
50+
include:
51+
- name: "GPU Pre-Training Stable"
52+
build_mode: stable
53+
workflow: pre-training
54+
image_name: maxtext_gpu_jax_stable
55+
- name: "GPU Pre-Training Nightly"
56+
build_mode: nightly
57+
workflow: pre-training
58+
image_name: maxtext_gpu_jax_nightly
59+
uses: ./.github/workflows/build_and_push_docker_image.yml
60+
with:
61+
image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
62+
device: gpu
63+
build_mode: ${{ matrix.build_mode }}
64+
workflow: ${{ matrix.workflow }}
65+
dockerfile: maxtext_gpu_dependencies.Dockerfile
66+
maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
67+
include_test_assets: true
68+
secrets:
69+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
70+
71+
notify_failure:
72+
name: Notify failed build
73+
needs: [build_and_push_docker_images]
74+
if: ${{ failure() && inputs.image_suffix == '' }}
75+
runs-on: ubuntu-latest
76+
permissions:
77+
issues: write
78+
steps:
79+
- name: Create issue on failure
80+
uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b
81+
with:
82+
github-token: ${{ secrets.GITHUB_TOKEN }}
83+
title-template: "MaxText Docker Image Build Failure"
84+
label-name: "docker-image-build-failure"

.github/workflows/nightly_images_pipeline.yml renamed to .github/workflows/tpu_nightly_images_pipeline.yml

Lines changed: 11 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -12,26 +12,17 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
# This workflow builds and pushes MaxText images for both TPU and GPU devices.
15+
# This workflow builds and pushes TPU MaxText images.
1616
# It runs automatically daily at 12am UTC, or manually via Workflow Dispatch.
1717

18-
name: Build Nightly Docker Images
18+
name: Build TPU Nightly Docker Images
1919

2020
on:
2121
schedule:
2222
# Run the job daily at 12AM UTC
2323
- cron: '0 0 * * *'
2424
workflow_dispatch:
2525
inputs:
26-
target_device:
27-
description: 'Specify target device (all, tpu, or gpu)'
28-
required: true
29-
type: choice
30-
default: 'all'
31-
options:
32-
- all
33-
- tpu
34-
- gpu
3526
image_suffix:
3627
description: 'An image suffix can be provided to add to the image name'
3728
required: false
@@ -50,104 +41,64 @@ jobs:
5041
device_name: v4-8
5142
cloud_runner: linux-x86-n2-16-buildkit
5243

53-
build_and_push_tpu_docker_images:
44+
build_and_push_docker_images:
5445
name: Build ${{ matrix.name }} Docker Image
5546
needs: build_maxtext_package
5647
strategy:
5748
fail-fast: false
5849
matrix:
5950
include:
6051
- name: "TPU Pre-Training Stable"
61-
device: tpu
6252
build_mode: stable
6353
workflow: pre-training
6454
image_name: maxtext_jax_stable
65-
dockerfile: maxtext_tpu_dependencies.Dockerfile
6655
- name: "TPU Pre-Training Nightly"
67-
device: tpu
6856
build_mode: nightly
6957
workflow: pre-training
7058
image_name: maxtext_jax_nightly
71-
dockerfile: maxtext_tpu_dependencies.Dockerfile
7259
- name: "TPU Post-Training Nightly"
73-
device: tpu
7460
build_mode: nightly
7561
workflow: post-training
7662
image_name: maxtext_post_training_nightly
77-
dockerfile: maxtext_tpu_dependencies.Dockerfile
7863
uses: ./.github/workflows/build_and_push_docker_image.yml
7964
with:
8065
image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
81-
device: ${{ matrix.device }}
66+
device: tpu
8267
build_mode: ${{ matrix.build_mode }}
8368
workflow: ${{ matrix.workflow }}
84-
dockerfile: ${{ matrix.dockerfile }}
69+
dockerfile: maxtext_tpu_dependencies.Dockerfile
8570
maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
8671
include_test_assets: true
8772
run_tests: false
8873
secrets:
8974
HF_TOKEN: ${{ secrets.HF_TOKEN }}
9075

91-
build_and_push_gpu_docker_images:
92-
name: Build ${{ matrix.name }} Docker Image
93-
needs: build_maxtext_package
94-
strategy:
95-
fail-fast: false
96-
matrix:
97-
include:
98-
- name: "GPU Pre-Training Stable"
99-
device: gpu
100-
build_mode: stable
101-
workflow: pre-training
102-
image_name: maxtext_gpu_jax_stable
103-
dockerfile: maxtext_gpu_dependencies.Dockerfile
104-
- name: "GPU Pre-Training Nightly"
105-
device: gpu
106-
build_mode: nightly
107-
workflow: pre-training
108-
image_name: maxtext_gpu_jax_nightly
109-
dockerfile: maxtext_gpu_dependencies.Dockerfile
110-
uses: ./.github/workflows/build_and_push_docker_image.yml
111-
with:
112-
image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
113-
device: ${{ matrix.device }}
114-
build_mode: ${{ matrix.build_mode }}
115-
workflow: ${{ matrix.workflow }}
116-
dockerfile: ${{ matrix.dockerfile }}
117-
maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
118-
include_test_assets: true
119-
secrets:
120-
HF_TOKEN: ${{ secrets.HF_TOKEN }}
121-
122-
run_tpu_ci_tests:
123-
name: Run TPU ${{ matrix.name }} CI Tests
124-
needs: build_and_push_tpu_docker_images
76+
run_ci_tests:
77+
name: Run ${{ matrix.name }} CI Tests
78+
needs: build_and_push_docker_images
12579
strategy:
12680
fail-fast: false
12781
matrix:
12882
include:
12983
- name: "Pre-Training Stable"
13084
image_name: maxtext_jax_stable
131-
device: tpu
13285
workflow: pre-training
13386
- name: "Pre-Training Nightly"
13487
image_name: maxtext_jax_nightly
135-
device: tpu
13688
workflow: pre-training
13789
- name: "Post-Training Nightly"
13890
image_name: maxtext_post_training_nightly
139-
device: tpu
14091
workflow: post-training
14192
uses: ./.github/workflows/run_ci_tests.yml
14293
with:
14394
image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
14495
image_tag: ${{ github.run_id }}
145-
device: ${{ matrix.device }}
96+
device: tpu
14697
workflow: ${{ matrix.workflow }}
14798

14899
run_e2e_tests:
149100
name: Run E2E tests
150-
needs: build_and_push_tpu_docker_images
101+
needs: build_and_push_docker_images
151102
if: github.event_name == 'schedule'
152103
uses: ./.github/workflows/run_e2e_tests.yml
153104
with:
@@ -158,7 +109,7 @@ jobs:
158109

159110
notify_failure:
160111
name: Notify failed build
161-
needs: [build_and_push_tpu_docker_images, build_and_push_gpu_docker_images, run_tpu_ci_tests, run_e2e_tests]
112+
needs: [build_and_push_docker_images, run_ci_tests, run_e2e_tests]
162113
if: ${{ failure() && inputs.image_suffix == '' }}
163114
runs-on: ubuntu-latest
164115
permissions:

0 commit comments

Comments
 (0)