3535 required : true
3636 type : string
3737 image_date :
38- required : true
38+ required : false
3939 type : string
40- base_image :
40+ workflow :
4141 required : false
4242 type : string
43- default : ' '
44- is_post_training :
43+ default : ' pre-training '
44+ version_name :
4545 required : false
46- type : boolean
47- default : false
46+ type : string
47+ default : ' '
4848
4949permissions :
5050 contents : read
@@ -66,13 +66,18 @@ jobs:
6666 id : check
6767 shell : bash
6868 run : |
69- if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${{ github.event.inputs.target_device }} " != "all" && "${{ github.event.inputs.target_device }} " != "${{ inputs.device } }" ]]; then
69+ if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${GITHUB_EVENT_INPUTS_TARGET_DEVICE} " != "all" && "${GITHUB_EVENT_INPUTS_TARGET_DEVICE} " != "${INPUTS_DEVICE }" ]]; then
7070 echo "should_run=false" >> $GITHUB_OUTPUT
71- echo "Skipping ${{ inputs.image_name }} build for device: ${{ inputs.device }} in ${{ inputs.build_mode } } mode."
71+ echo "Skipping ${INPUTS_IMAGE_NAME} build for device: ${INPUTS_DEVICE} in ${INPUTS_BUILD_MODE } mode."
7272 else
7373 echo "should_run=true" >> $GITHUB_OUTPUT
74- echo "Building ${{ inputs.image_name }} for device: ${{ inputs.device }} in ${{ inputs.build_mode } } mode."
74+ echo "Building ${INPUTS_IMAGE_NAME} for device: ${INPUTS_DEVICE} in ${INPUTS_BUILD_MODE } mode."
7575 fi
76+ env :
77+ GITHUB_EVENT_INPUTS_TARGET_DEVICE : ${{ github.event.inputs.target_device }}
78+ INPUTS_DEVICE : ${{ inputs.device }}
79+ INPUTS_IMAGE_NAME : ${{ inputs.image_name }}
80+ INPUTS_BUILD_MODE : ${{ inputs.build_mode }}
7681
7782 - name : Checkout MaxText
7883 uses : actions/checkout@v5
@@ -110,40 +115,50 @@ jobs:
110115 push : true
111116 context : .
112117 file : ${{ inputs.dockerfile }}
113- tags : gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
118+ tags : gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}
114119 cache-from : type=gha
115120 outputs : type=image,compression=zstd,force-compression=true
116121 build-args : |
117122 DEVICE=${{ inputs.device }}
118123 MODE=${{ inputs.build_mode }}
124+ WORKFLOW=${{ inputs.workflow }}
119125 JAX_VERSION=NONE
120126 LIBTPU_VERSION=NONE
121127 INCLUDE_TEST_ASSETS=true
122- ${{ inputs.base_image != '' && format('BASEIMAGE={0}', inputs.base_image) || '' }}
123128
124129name : Add tags to Docker image
125130 if : steps.check.outputs.should_run == 'true'
126131 shell : bash
127132 run : |
128- SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${{ inputs.image_name } }"
133+ SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME }"
129134
130- # Add date tag
131- gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${{ inputs.image_date }}" --quiet
135+ if [[ $INPUTS_VERSION_NAME ]]; then
136+ echo "Tagging docker images corresponding to PyPI release..."
137+ gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_VERSION_NAME}" --quiet
138+ else
139+ echo "Tagging docker images corresponding to nightly release..."
132140
133- # Convert date to YYYYMMDD format
134- clean_date=$(echo "$ {{ inputs.image_date }}" | sed 's/[-:]//g' | cut -c1-8)
141+ # Add date tag
142+ gcloud container images add-tag "$SOURCE_IMAGE:$ {{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
135143
136- # Add MaxText tag
137- maxtext_hash=$(git rev-parse --short HEAD)
138- gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
144+ # Convert date to YYYYMMDD format
145+ clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
139146
147+ # Add MaxText tag
148+ maxtext_hash=$(git rev-parse --short HEAD)
149+ gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
140150
141151 # Add post-training dependencies tags
142- if [ "${{ inputs.is_post_training }}" == "true " ]; then
152+ if [ "${{ inputs.workflow }}" == "post-training " ]; then
143153 for dir in tunix vllm tpu-inference; do
144154 if [ -d "./$dir" ]; then
145155 dir_hash=$(git -C "$dir" rev-parse --short HEAD)
146- gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
147- fi
148- done
156+ gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
157+ fi
158+ done
159+ fi
149160 fi
161+ env :
162+ INPUTS_IMAGE_NAME : ${{ inputs.image_name }}
163+ INPUTS_IMAGE_DATE : ${{ inputs.image_date }}
164+ INPUTS_VERSION_NAME : ${{ inputs.version_name }}
0 commit comments