Skip to content

Commit 49a4470

Browse files
SurbhiJainUSCSurbhiJainUSC
committed
Add job to build and upload docker image to pypi_release workflow
1 parent 05c5083 commit 49a4470

2 files changed

Lines changed: 105 additions & 28 deletions

File tree

‎.github/workflows/build_and_push_docker_image.yml‎

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -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,36 @@ 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
166+
167+
# Drop the "${{ github.run_id }}" tag from the image
168+
gcloud container images untag "$SOURCE_IMAGE:${{ github.run_id }}" --quiet
156169
env:
157170
INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
158171
INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
172+
INPUTS_VERSION_NAME: ${{ inputs.version_name }}

‎.github/workflows/pypi_release.yml‎

Lines changed: 80 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ name: Publish MaxText to PyPI
1818

1919
# Triggers when a new "release" is published in the GitHub UI
2020
on:
21+
pull_request: # TODO: Remove this trigger after testing, it's added to test the workflow on pull requests.
2122
release:
2223
types: [published]
2324

@@ -41,23 +42,85 @@ jobs:
4142
name: Build and Test MaxText Package
4243
needs: [release_approval]
4344
uses: ./.github/workflows/build_and_test_maxtext.yml
45+
secrets: inherit
4446

45-
publish_maxtext_package_to_pypi:
46-
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]
47+
#publish_maxtext_package_to_pypi:
48+
# name: Publish MaxText package to PyPI
49+
# needs: [build_and_test_maxtext_package]
50+
# runs-on: ubuntu-latest
51+
# environment: release
52+
# steps:
53+
# - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
54+
# - name: Download MaxText wheel
55+
# uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0
56+
# with:
57+
# name: maxtext-wheel
58+
# path: dist/
59+
# - name: Publish MaxText wheel to PyPI
60+
# # Official action for PyPI Trusted Publishing
61+
# uses: pypa/gh-action-pypi-publish@release/v1
62+
# with:
63+
# packages-dir: dist/
64+
65+
get_latest_maxtext_pypi_version:
66+
name: Get latest MaxText PyPI version
67+
needs: [build_and_test_maxtext_package] # TODO: change to publish_maxtext_package_to_pypi after testing
5068
runs-on: ubuntu-latest
51-
environment: release
69+
outputs:
70+
latest_pypi_version: ${{ steps.get_version.outputs.version }}
5271
steps:
53-
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
54-
- name: Download MaxText wheel
55-
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0
56-
with:
57-
name: maxtext-wheel
58-
path: dist/
59-
- name: Publish MaxText wheel to PyPI
60-
# Official action for PyPI Trusted Publishing
61-
uses: pypa/gh-action-pypi-publish@release/v1
62-
with:
63-
packages-dir: dist/
72+
- name: Install jq
73+
run: sudo apt-get update && sudo apt-get install -y jq
74+
- name: Fetch latest version of maxtext from PyPI
75+
id: get_version
76+
run: |
77+
# Fetch JSON from PyPI for 'maxtext'
78+
echo "Fetching latest version from https://pypi.org/pypi/maxtext/json"
79+
pypi_json=$(curl -s https://pypi.org/pypi/maxtext/json)
80+
81+
# Extract the version from the "info" section using jq
82+
latest_version=$(echo "$pypi_json" | jq -r ".info.version")
83+
84+
if [ -z "$latest_version" ] || [ "$latest_version" == "null" ]; then
85+
echo "Error: Could not parse latest version from PyPI JSON."
86+
exit 1
87+
fi
88+
89+
echo "Successfully fetched latest MaxText version on PyPI: $latest_version"
90+
# Set the output variable for other jobs to consume
91+
echo "version=$latest_version" >> "$GITHUB_OUTPUT"
92+
93+
# This job builds and pushes MaxText stable Docker images for both TPU and GPU devices.
94+
# It runs only after the release is approved, ensuring that the Docker images are built and pushed only for approved releases.
95+
# Creates docker image for MaxText commit corresponding to the release.
96+
upload_maxtext_docker_images:
97+
name: Build and upload MaxText stable Docker images
98+
needs: [get_latest_maxtext_pypi_version]
99+
strategy:
100+
fail-fast: false
101+
matrix:
102+
include:
103+
- device: tpu
104+
build_mode: stable
105+
image_name: maxtext_jax_stable
106+
workflow: pre-training
107+
dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
108+
- device: gpu
109+
build_mode: stable
110+
image_name: maxtext_gpu_jax_stable
111+
workflow: pre-training
112+
dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
113+
- device: tpu
114+
build_mode: stable
115+
image_name: maxtext_post_training_stable
116+
workflow: post-training
117+
dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
118+
uses: ./.github/workflows/build_and_push_docker_image.yml
119+
with:
120+
image_name: ${{ matrix.image_name }}
121+
device: ${{ matrix.device }}
122+
build_mode: ${{ matrix.build_mode }}
123+
workflow: ${{ matrix.workflow }}
124+
dockerfile: ${{ matrix.dockerfile }}
125+
maxtext_sha: ${{ github.sha }}
126+
version_name: ${{ github.event.release.tag_name }}

0 commit comments

Comments
 (0)