Skip to content

Commit 55e57ec

Browse files
SurbhiJainUSCSurbhiJainUSC
committed
Add job to build and upload docker image to pypi_release workflow
1 parent 6b6dbc2 commit 55e57ec

2 files changed

Lines changed: 88 additions & 15 deletions

File tree

‎.github/workflows/build_and_push_docker_image.yml‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ on:
3535
required: true
3636
type: string
3737
image_date:
38-
required: true
38+
required: false
3939
type: string
4040
base_image:
4141
required: false
@@ -45,6 +45,10 @@ on:
4545
required: false
4646
type: string
4747
default: 'pre-training'
48+
version_name:
49+
required: false
50+
type: string
51+
default: ''
4852

4953
permissions:
5054
contents: read
@@ -115,7 +119,7 @@ jobs:
115119
push: true
116120
context: .
117121
file: ${{ inputs.dockerfile }}
118-
tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
122+
tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}
119123
cache-from: type=gha
120124
outputs: type=image,compression=zstd,force-compression=true
121125
build-args: |
@@ -133,26 +137,33 @@ jobs:
133137
run: |
134138
SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
135139
136-
# Add date tag
137-
gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
140+
if [[ $INPUTS_VERSION_NAME ]]; then
141+
echo "Tagging docker images corresponding to PyPI release..."
142+
gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_VERSION_NAME}" --quiet
143+
else
144+
echo "Tagging docker images corresponding to nightly release..."
138145
139-
# Convert date to YYYYMMDD format
140-
clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
146+
# Add date tag
147+
gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
141148
142-
# Add MaxText tag
143-
maxtext_hash=$(git rev-parse --short HEAD)
144-
gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
149+
# Convert date to YYYYMMDD format
150+
clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
145151
152+
# Add MaxText tag
153+
maxtext_hash=$(git rev-parse --short HEAD)
154+
gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
146155
147156
# Add post-training dependencies tags
148157
if [ "${{ inputs.workflow }}" == "post-training" ]; then
149158
for dir in tunix vllm tpu-inference; do
150159
if [ -d "./$dir" ]; then
151160
dir_hash=$(git -C "$dir" rev-parse --short HEAD)
152-
gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
153-
fi
154-
done
161+
gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
162+
fi
163+
done
164+
fi
155165
fi
156166
env:
157167
INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
158168
INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
169+
INPUTS_VERSION_NAME: ${{ inputs.version_name }}

‎.github/workflows/pypi_release.yml‎

Lines changed: 65 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,11 @@ jobs:
4141
name: Build and Test MaxText Package
4242
needs: [release_approval]
4343
uses: ./.github/workflows/build_and_test_maxtext.yml
44+
secrets: inherit
4445

4546
publish_maxtext_package_to_pypi:
4647
name: Publish MaxText package to PyPI
47-
# Temporarily only require release_approval for a one-time upload.
48-
# Immediately revert this to `needs: [build_and_test_maxtext_package]`.
49-
needs: [release_approval]
48+
needs: [build_and_test_maxtext_package]
5049
runs-on: ubuntu-latest
5150
environment: release
5251
steps:
@@ -61,3 +60,66 @@ jobs:
6160
uses: pypa/gh-action-pypi-publish@release/v1
6261
with:
6362
packages-dir: dist/
63+
64+
get_latest_maxtext_pypi_version:
65+
name: Get latest MaxText PyPI version
66+
needs: [publish_maxtext_package_to_pypi]
67+
runs-on: ubuntu-latest
68+
outputs:
69+
latest_pypi_version: ${{ steps.get_version.outputs.version }}
70+
steps:
71+
- name: Install jq
72+
run: sudo apt-get update && sudo apt-get install -y jq
73+
- name: Fetch latest version of maxtext from PyPI
74+
id: get_version
75+
run: |
76+
# Fetch JSON from PyPI for 'maxtext'
77+
echo "Fetching latest version from https://pypi.org/pypi/maxtext/json"
78+
pypi_json=$(curl -s https://pypi.org/pypi/maxtext/json)
79+
80+
# Extract the version from the "info" section using jq
81+
latest_version=$(echo "$pypi_json" | jq -r ".info.version")
82+
83+
if [ -z "$latest_version" ] || [ "$latest_version" == "null" ]; then
84+
echo "Error: Could not parse latest version from PyPI JSON."
85+
exit 1
86+
fi
87+
88+
echo "Successfully fetched latest MaxText version on PyPI: $latest_version"
89+
# Set the output variable for other jobs to consume
90+
echo "version=$latest_version" >> "$GITHUB_OUTPUT"
91+
92+
# This job builds and pushes MaxText stable Docker images for both TPU and GPU devices.
93+
# It runs only after a new release is published to PyPI.
94+
# Creates docker image for MaxText commit corresponding to the release.
95+
upload_maxtext_docker_images:
96+
name: ${{ matrix.image_name }}
97+
needs: [get_latest_maxtext_pypi_version]
98+
strategy:
99+
fail-fast: false
100+
matrix:
101+
include:
102+
- device: tpu
103+
build_mode: stable
104+
image_name: maxtext_jax_stable
105+
workflow: pre-training
106+
dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
107+
- device: gpu
108+
build_mode: stable
109+
image_name: maxtext_gpu_jax_stable
110+
workflow: pre-training
111+
dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
112+
- device: tpu
113+
build_mode: stable
114+
image_name: maxtext_post_training_stable
115+
workflow: post-training
116+
dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
117+
uses: ./.github/workflows/build_and_push_docker_image.yml
118+
with:
119+
image_name: ${{ matrix.image_name }}
120+
device: ${{ matrix.device }}
121+
build_mode: ${{ matrix.build_mode }}
122+
workflow: ${{ matrix.workflow }}
123+
dockerfile: ${{ matrix.dockerfile }}
124+
maxtext_sha: ${{ github.sha }}
125+
version_name: ${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}

0 commit comments

Comments
 (0)