Skip to content

Commit b5170bb

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request #3993 from AI-Hypercomputer:run_e2e_tests
PiperOrigin-RevId: 933956662
2 parents 45b0f2b + 8538100 commit b5170bb

5 files changed

Lines changed: 230 additions & 66 deletions

File tree

‎.github/workflows/UploadDockerImages.yml‎

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -42,41 +42,47 @@ permissions:
4242
contents: read
4343

4444
jobs:
45-
build_and_upload_maxtext_package:
45+
build_maxtext_package:
46+
name: Build MaxText Package
4647
uses: ./.github/workflows/build_package.yml
4748
with:
4849
device_type: tpu
4950
device_name: v4-8
5051
cloud_runner: linux-x86-n2-16-buildkit
5152

52-
build-and-push:
53-
name: ${{ matrix.image_name }}
54-
needs: build_and_upload_maxtext_package
53+
build_and_push_docker_image:
54+
name: Build ${{ matrix.name }} Docker Image
55+
needs: build_maxtext_package
5556
strategy:
5657
fail-fast: false
5758
matrix:
5859
include:
59-
- device: tpu
60+
- name: "TPU Pre-Training Stable"
61+
device: tpu
6062
build_mode: stable
6163
workflow: pre-training
6264
image_name: maxtext_jax_stable
6365
dockerfile: maxtext_tpu_dependencies.Dockerfile
64-
- device: tpu
66+
- name: "TPU Pre-Training Nightly"
67+
device: tpu
6568
build_mode: nightly
6669
workflow: pre-training
6770
image_name: maxtext_jax_nightly
6871
dockerfile: maxtext_tpu_dependencies.Dockerfile
69-
- device: tpu
72+
- name: "TPU Post-Training Nightly"
73+
device: tpu
7074
build_mode: nightly
7175
workflow: post-training
7276
image_name: maxtext_post_training_nightly
7377
dockerfile: maxtext_tpu_dependencies.Dockerfile
74-
- device: gpu
78+
- name: "GPU Pre-Training Stable"
79+
device: gpu
7580
build_mode: stable
7681
workflow: pre-training
7782
image_name: maxtext_gpu_jax_stable
7883
dockerfile: maxtext_gpu_dependencies.Dockerfile
79-
- device: gpu
84+
- name: "GPU Pre-Training Nightly"
85+
device: gpu
8086
build_mode: nightly
8187
workflow: pre-training
8288
image_name: maxtext_gpu_jax_nightly
@@ -93,9 +99,19 @@ jobs:
9399
secrets:
94100
HF_TOKEN: ${{ secrets.HF_TOKEN }}
95101

102+
run_e2e_tests:
103+
name: Run E2E tests
104+
needs: build_and_push_docker_image
105+
uses: ./.github/workflows/run_e2e_tests.yml
106+
with:
107+
mode: ${{ inputs.image_suffix != '' && format('{0}_{1}', 'nightly', inputs.image_suffix) || 'nightly' }}
108+
run_id: ${{ github.run_id }}
109+
maxtext_sha: ${{ github.sha }}
110+
secrets: inherit
111+
96112
notify_failure:
97113
name: Notify failed build
98-
needs: [build-and-push]
114+
needs: [build_and_push_docker_image, run_e2e_tests]
99115
if: ${{ failure() && inputs.image_suffix == '' }}
100116
runs-on: ubuntu-latest
101117
permissions:

‎.github/workflows/build_and_push_docker_image.yml‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ permissions:
5555

5656
jobs:
5757
pre_build_check:
58+
name: Pre-build Check
5859
runs-on: ubuntu-latest
5960
outputs:
6061
should_run: ${{ steps.check.outputs.should_run }}
@@ -85,6 +86,7 @@ jobs:
8586
fi
8687
8788
build_and_push:
89+
name: Build and Push Docker Image
8890
needs: pre_build_check
8991
runs-on: linux-x86-n2-16-buildkit
9092
container: google/cloud-sdk:524.0.0
@@ -185,10 +187,11 @@ jobs:
185187
INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
186188
INPUTS_VERSION_NAME: ${{ inputs.version_name }}
187189

188-
promote_image:
190+
run_ci_tests:
191+
name: Run Unit and Integration Tests
189192
needs: [pre_build_check, build_and_push]
190193
if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true
191-
uses: ./.github/workflows/promote_docker_image.yml
194+
uses: ./.github/workflows/run_ci_tests.yml
192195
with:
193196
image_name: ${{ inputs.image_name }}
194197
image_tag: ${{ github.run_id }}

‎.github/workflows/promote_docker_image.yml‎

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

15-
# This workflow will test and promote a MaxText Docker image.
15+
# This workflow acts as the final gate in the CI/CD pipeline for MaxText Docker images.
16+
# Triggered by: Airflow DAG completion (repository_dispatch)
17+
# 1. Validates the result of the triggered Airflow DAG (E2E tests).
18+
# 2. If the DAG succeeded, it promotes the Docker images by tagging the artifact with the latest tag.
1619

1720
name: Promote MaxText Docker Image
1821

1922
on:
20-
workflow_call:
21-
inputs:
22-
image_name:
23-
required: true
24-
type: string
25-
image_tag:
26-
required: true
27-
type: string
28-
device:
29-
required: true
30-
type: string
31-
workflow:
32-
required: true
33-
type: string
23+
repository_dispatch:
24+
types: [airflow-dag-complete]
25+
3426
permissions:
3527
contents: read
3628

3729
jobs:
38-
check_image:
39-
runs-on: linux-x86-n2-16-buildkit
40-
container: google/cloud-sdk:524.0.0
41-
outputs:
42-
image_exists: ${{ steps.check.outputs.exists }}
30+
handle_result:
31+
name: Handle Airflow DAG Result - ${{ github.event.client_payload.dag_id }}
32+
runs-on: ubuntu-latest
4333
steps:
44-
- name: Check if image exists
45-
id: check
46-
shell: bash
34+
- name: Report DAG result
4735
run: |
48-
if gcloud container images describe "gcr.io/${{ vars.PROJECT_NAME }}/${{ inputs.image_name }}:${{ inputs.image_tag }}" >/dev/null 2>&1; then
49-
echo "exists=true" >> $GITHUB_OUTPUT
50-
else
51-
echo "exists=false" >> $GITHUB_OUTPUT
52-
fi
36+
STATE="${{ github.event.client_payload.state }}"
37+
DAG_ID="${{ github.event.client_payload.dag_id }}"
38+
DAG_RUN_ID="${{ github.event.client_payload.dag_run_id }}"
39+
SHA="${{ github.event.client_payload.sha }}"
5340
54-
test:
55-
needs: check_image
56-
if: needs.check_image.outputs.image_exists == 'true'
57-
strategy:
58-
fail-fast: false
59-
matrix:
60-
flavor: >-
61-
${{ fromJSON('{
62-
"gpu-pre-training": ["gpu-unit", "gpu-integration"],
63-
"tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit"],
64-
"tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit"]
65-
}')[format('{0}-{1}', inputs.device, inputs.workflow)] }}
66-
uses: ./.github/workflows/run_tests_coordinator.yml
67-
with:
68-
flavor: ${{ matrix.flavor }}
69-
base_image: ${{ inputs.image_name }}:${{ inputs.image_tag }}
70-
is_scheduled_run: true
71-
maxtext_installed: true
72-
additional_pytest_args: ${{ contains(inputs.image_name, 'maxtext_jax_nightly') && '--ignore=tests/integration/hlo_diff_test.py' || '' }}
41+
echo "================================"
42+
echo "DAG ID: ${DAG_ID}"
43+
echo "DAG Run ID: ${DAG_RUN_ID}"
44+
echo "Commit SHA: ${SHA}"
45+
echo "State: ${STATE}"
46+
echo "================================"
7347
74-
tagging:
75-
needs: [test]
76-
if: needs.test.result == 'success'
48+
case "$STATE" in
49+
success) echo "E2E tests PASSED." ;;
50+
failed|upstream_failed) echo "E2E tests FAILED."; exit 1 ;;
51+
*) echo "DAG ended with unexpected state: ${STATE}"; exit 1 ;;
52+
esac
53+
54+
tag_docker_image:
55+
name: Promote ${{ matrix.image_name }} Docker Image
56+
needs: handle_result
57+
if: needs.handle_result.result == 'success'
7758
runs-on: linux-x86-n2-16-buildkit
7859
container: google/cloud-sdk:524.0.0
60+
strategy:
61+
fail-fast: false
62+
matrix:
63+
image_name:
64+
- maxtext_jax_nightly
65+
- maxtext_gpu_jax_nightly
66+
- maxtext_post_training_nightly
7967
steps:
8068
- name: Configure Docker
8169
run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
82-
8370
- name: Add tags to Docker image
8471
shell: bash
8572
run: |
86-
# Add Latest Tag
87-
SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${{ inputs.image_name }}"
88-
gcloud container images add-tag "${SOURCE_IMAGE}:${{ inputs.image_tag }}" "${SOURCE_IMAGE}:latest" --quiet
73+
SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${{ matrix.image_name }}"
74+
GITHUB_RUN_ID="${{ github.event.client_payload.github_run_id }}"
75+
76+
# Add the traceability tag to confirm it passed validation suite
77+
SHORT_DAG_RUN_ID=$(echo "${{ github.event.client_payload.dag_run_id }}" | cut -c1-7)
78+
gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
79+
"${SOURCE_IMAGE}:verified-${SHORT_DAG_RUN_ID}" --quiet
80+
81+
# Add "latest" tag
82+
gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" "${SOURCE_IMAGE}:latest" --quiet

‎.github/workflows/run_ci_tests.yml‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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 will run unit and integration tests on MaxText Docker images.
16+
17+
name: MaxText Unit and Integration Tests
18+
19+
on:
20+
workflow_call:
21+
inputs:
22+
image_name:
23+
required: true
24+
type: string
25+
image_tag:
26+
required: true
27+
type: string
28+
device:
29+
required: true
30+
type: string
31+
workflow:
32+
required: true
33+
type: string
34+
permissions:
35+
contents: read
36+
37+
jobs:
38+
check_image:
39+
runs-on: linux-x86-n2-16-buildkit
40+
container: google/cloud-sdk:524.0.0
41+
outputs:
42+
image_exists: ${{ steps.check.outputs.exists }}
43+
steps:
44+
- name: Check if image exists
45+
id: check
46+
shell: bash
47+
run: |
48+
if gcloud container images describe "gcr.io/${{ vars.PROJECT_NAME }}/${{ inputs.image_name }}:${{ inputs.image_tag }}" >/dev/null 2>&1; then
49+
echo "exists=true" >> $GITHUB_OUTPUT
50+
else
51+
echo "exists=false" >> $GITHUB_OUTPUT
52+
fi
53+
54+
test:
55+
needs: check_image
56+
if: needs.check_image.outputs.image_exists == 'true'
57+
strategy:
58+
fail-fast: false
59+
matrix:
60+
flavor: >-
61+
${{ fromJSON('{
62+
"gpu-pre-training": ["gpu-unit", "gpu-integration"],
63+
"tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit"],
64+
"tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit"]
65+
}')[format('{0}-{1}', inputs.device, inputs.workflow)] }}
66+
uses: ./.github/workflows/run_tests_coordinator.yml
67+
with:
68+
flavor: ${{ matrix.flavor }}
69+
base_image: ${{ inputs.image_name }}:${{ inputs.image_tag }}
70+
is_scheduled_run: true
71+
maxtext_installed: true
72+
additional_pytest_args: ${{ contains(inputs.image_name, 'maxtext_jax_nightly') && '--ignore=tests/integration/hlo_diff_test.py' || '' }}

‎.github/workflows/run_e2e_tests.yml‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
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 will trigger the MaxText E2E Airflow DAG to run the E2E tests.
16+
17+
name: MaxText E2E Airflow Tests
18+
19+
on:
20+
workflow_call:
21+
inputs:
22+
mode:
23+
description: 'Build mode: stable or nightly'
24+
required: false
25+
type: string
26+
default: 'stable'
27+
run_id:
28+
description: 'GitHub workflow run_id to be passed to the Airflow DAG run.'
29+
required: true
30+
type: string
31+
maxtext_sha:
32+
description: 'GitHub SHA of MaxText to be passed to the Airflow DAG run.'
33+
required: true
34+
type: string
35+
36+
permissions:
37+
contents: read
38+
39+
jobs:
40+
e2e_airflow_tests:
41+
name: E2E Airflow Tests
42+
runs-on: linux-x86-n2-16-buildkit
43+
container: google/cloud-sdk:524.0.0
44+
steps:
45+
- name: Get Airflow URI
46+
id: info
47+
run: |
48+
AIRFLOW_URI=$(gcloud composer environments describe ml-automation-solutions \
49+
--location us-central1 \
50+
--project cloud-ml-auto-solutions \
51+
--format "value(config.airflowUri)")
52+
echo "airflow_uri=${AIRFLOW_URI}" >> "$GITHUB_OUTPUT"
53+
54+
- name: Trigger DAG
55+
run: |
56+
IAP_TOKEN=$(gcloud auth print-access-token)
57+
58+
RESPONSE=$(curl -s -w "\n%{http_code}" -X POST \
59+
"${{ steps.info.outputs.airflow_uri }}/api/v1/dags/maxtext_e2e_tests/dagRuns" \
60+
-H "Authorization: Bearer ${IAP_TOKEN}" \
61+
-H "Content-Type: application/json" \
62+
-d "{
63+
\"conf\": {
64+
\"build_mode\": \"${{ inputs.mode }}\",
65+
\"maxtext_sha\": \"${{ inputs.maxtext_sha }}\",
66+
\"github_run_id\": \"${{ inputs.run_id }}\",
67+
\"github_repo\": \"${{ github.repository }}\",
68+
\"github_callback_token\": \"${{ secrets.AIRFLOW_CALLBACK_TOKEN }}\"
69+
}
70+
}")
71+
72+
HTTP_STATUS=$(echo "$RESPONSE" | tail -1)
73+
BODY=$(echo "$RESPONSE" | sed '$d')
74+
echo "Response body: ${BODY}"
75+
76+
if [ "${HTTP_STATUS}" -lt 200 ] || [ "${HTTP_STATUS}" -ge 300 ]; then
77+
echo "Error: Airflow API returned HTTP ${HTTP_STATUS}"
78+
exit 1
79+
fi

0 commit comments

Comments
 (0)