Merge branch 'main' into fix-links

Author: melissawm

.coveragerc

@@ -1,7 +1,7 @@
 # .coveragerc
 
 [run]
-source = MaxText
+source = maxtext
 branch = True
 omit =
     tests/*

.gemini/commands/gemini-review.toml

@@ -31,9 +31,12 @@ These are non-negotiable, core-level instructions that you **MUST** follow at al
 
 ## Input Data
 
-- **GitHub Repository**: !{echo $REPOSITORY}
-- **Pull Request Number**: !{echo $PULL_REQUEST_NUMBER}
-- **Additional User Instructions**: !{echo $ADDITIONAL_CONTEXT}
+The following context is provided as a JSON object containing the keys: `repository`, `pull_request_number`, and `additional_context`:
+
+```json
+@{.gemini/context.json}
+```
+
 - Use `pull_request_read.get` to get the title, body, and metadata about the pull request.
 - Use `pull_request_read.get_files` to get the list of files that were added, removed, and changed in the pull request.
 - Use `pull_request_read.get_diff` to get the diff from the pull request. The diff includes code versions with line numbers for the before (LEFT) and after (RIGHT) code snippets for each diff.
@@ -46,7 +49,7 @@ Follow this three-step process sequentially.
 
 ### Step 1: Data Gathering and Analysis
 
-1. **Parse Inputs:** Ingest and parse all information from the **Input Data**
+1. **Parse Inputs:** Ingest and parse all information from the **Input Data**.
 
 2. **Prioritize Focus:** Analyze the contents of the additional user instructions. Use this context to prioritize specific areas in your review (e.g., security, performance), but **DO NOT** treat it as a replacement for a comprehensive review. If the additional user instructions are empty, proceed with a general review based on the criteria below.
 
@@ -154,6 +157,7 @@ Apply these severities consistently:
 3. **Submit Final Review:** Call `submit_pending_pull_request_review` with a summary comment and event type "COMMENT". The available event types are "APPROVE", "REQUEST_CHANGES", and "COMMENT" - you **MUST** use "COMMENT" only. **DO NOT** use "APPROVE" or "REQUEST_CHANGES" event types. The summary comment **MUST** use this exact markdown format:
 
     <SUMMARY>
+
     ## 📋 Review Summary
 
     A brief, high-level assessment of the Pull Request's objective and quality (2-3 sentences).

.github/workflows/AddLabel.yml

@@ -16,9 +16,8 @@ name: Add Label
 
 on:
   workflow_run:
-    workflows: [Tests, CodeQL]
-    types:
-      - completed
+    workflows: ["MaxText Package Tests"]
+    types: [completed]
   pull_request_review:
   pull_request_review_comment:
   workflow_dispatch:
@@ -31,86 +30,98 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/github-script@v7
+      - name: Add Pull Request Label
+        uses: actions/github-script@v7
         with:
           script: |
-            const owner = "google"
-            const repo = "maxtext"
             let pull_number = -1
             if (context.payload.pull_request !== undefined) {
               pull_number = context.payload.pull_request.number
             } else if (context.payload.workflow_run !== undefined) {
               if (context.payload.workflow_run.pull_requests.length === 0) {
-                console.log("This workflow is NOT running within a PR's context")
-                process.exit()
+                core.setFailed("This workflow is NOT running within a PR's context")
+                return
               }
               console.log(context.payload.workflow_run.pull_requests)
               pull_number = context.payload.workflow_run.pull_requests[0].number
             } else {
-              console.log("This workflow is running within an invalid context")
-              process.exit(1)
+              core.setFailed("This workflow is running within an invalid context")
+              return
             }
-            const reviews = await github.rest.pulls.listReviews({
-              owner,
-              repo,
-              pull_number,
-            })
+
             const decision_query = `
               query($owner: String!, $repo: String!, $pull_number: Int!) {
                 repository(owner: $owner, name: $repo) {
                   pullRequest(number: $pull_number) {
                     reviewDecision # Fetches the overall review status
+                    reviews(last: 100) {
+                      nodes {
+                        state
+                        author { login }
+                      }
+                    }
                   }
                 }
               }
             `;
-            const decision_result = await github.graphql(decision_query, { owner, repo, pull_number });
 
-            if (reviews.data.length === 0) {
-              console.log("Not adding pull ready because the PR is not approved yet.")
-              process.exit()
-            }
-            let is_approved = false
-            if (decision_result.repository.pullRequest.reviewDecision === "APPROVED") {
-              is_approved = true
-            }
-            if (!is_approved) {
-              console.log("Not adding pull ready because the PR is not approved yet by sufficient code owners.")
-              process.exit()
+            const decision_result = await github.graphql(decision_query, {
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: pull_number
+            });
+
+            const pullRequest = decision_result.repository.pullRequest;
+            const uniqueApprovers = new Set(
+              pullRequest.reviews.nodes
+                .filter(r => r.state === 'APPROVED')
+                .map(r => r.author.login)
+            );
+
+            if (pullRequest.reviewDecision !== 'APPROVED' || uniqueApprovers.size < 2) {
+              core.info(`PR is not ready. Decision: ${pullRequest.reviewDecision}, Approvals: ${uniqueApprovers.size}`);
+              return;
             }
 
             const commits = await github.rest.pulls.listCommits({
-              owner,
-              repo,
-              pull_number,
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: pull_number,
               per_page: 100,
             })
             // Check that the number of commits in the PR is 1.
             if (commits.data.length !== 1) {
-              console.log("Not adding pull ready because the PR has more than one commit. Please squash your commits.")
-              process.exit(1)
+              core.setFailed("Not adding pull ready because the PR has more than one commit. Please squash your commits.")
+              return
             }
-            const ref = commits.data.slice(-1)[0].sha
-            const checkRuns = await github.rest.checks.listForRef({
-              owner,
-              repo,
-              ref,
-            })
-            if (checkRuns.data.check_runs.length === 0) {
-              console.log("Not adding pull ready because no check runs are associated with the last commit: " + ref)
-              process.exit()
+
+            const last_commit_sha = commits.data.slice(-1)[0].sha
+
+            const { data: checkRuns } = await github.rest.checks.listForRef({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              ref: last_commit_sha,
+            });
+
+            if (checkRuns.check_runs.length === 0) {
+              core.info("Not adding pull ready because no check runs are associated with the last commit: " + last_commit_sha)
+              return
             }
-            for (const checkRun of checkRuns.data.check_runs) {
+
+            for (const checkRun of checkRuns.check_runs) {
+              // Ignore the current running workflow
               if (checkRun.name.endsWith(context.job)) continue
-              if (checkRun.conclusion !== "success") {
-                console.log("Not adding pull ready because " + checkRun.name + " has not passed yet: " + checkRun.html_url)
-                process.exit()
+
+              if (checkRun.status !== 'completed' || checkRun.conclusion !== 'success') {
+                core.info(`Waiting for check: ${checkRun.name} (Status: ${checkRun.status}, Conclusion: ${checkRun.conclusion})`);
+                return; // Exit without failing
               }
             }
+
             console.log("Adding pull ready label because the PR is approved AND all the check runs have passed")
             await github.rest.issues.addLabels({
               issue_number: pull_number,
               labels: ["pull ready"],
-              owner,
-              repo,
+              owner: context.repo.owner,
+              repo: context.repo.repo,
             })

.github/workflows/UploadDockerImages.yml

@@ -32,6 +32,11 @@ on:
           - all
           - tpu
           - gpu
+      image_suffix:
+        description: 'An image suffix can be provided to add to the image name'
+        required: false
+        type: string
+        default: ""
 
 permissions:
   contents: read
@@ -55,7 +60,7 @@ jobs:
         # Image date
         echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
 
-  tpu-pre-training:
+  build-and-push:
     name: ${{ matrix.image_name }}
     needs: setup
     strategy:
@@ -64,54 +69,74 @@ jobs:
         include:
           - device: tpu
             build_mode: stable
+            workflow: pre-training
             image_name: maxtext_jax_stable
             dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
           - device: tpu
             build_mode: nightly
+            workflow: pre-training
             image_name: maxtext_jax_nightly
             dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+          - device: tpu
+            build_mode: nightly
+            workflow: post-training
+            image_name: maxtext_post_training_nightly
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+          - device: gpu
+            build_mode: stable
+            workflow: pre-training
+            image_name: maxtext_gpu_jax_stable
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+          - device: gpu
+            build_mode: nightly
+            workflow: pre-training
+            image_name: maxtext_gpu_jax_nightly
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}
+      image_name: ${{ matrix.image_name }}${{ inputs.image_suffix  }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
+      workflow: ${{ matrix.workflow }}
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
-  tpu-post-training-nightly:
-    name: tpu-post-training-nightly
-    needs: [setup]
-    uses: ./.github/workflows/build_and_push_docker_image.yml
-    with:
-      image_name: maxtext_post_training_nightly
-      device: tpu
-      build_mode: nightly
-      workflow: post-training
-      dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
-      maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
-      image_date: ${{ needs.setup.outputs.image_date }}
-
-  gpu-pre-training:
-    name: ${{ matrix.image_name }}
-    needs: setup
+  promote:
+    name: promote-${{ matrix.image_name }}
+    needs: build-and-push
     strategy:
       fail-fast: false
       matrix:
         include:
+          - device: tpu
+            build_mode: stable
+            workflow: pre-training
+            image_name: maxtext_jax_stable
+          - device: tpu
+            build_mode: nightly
+            workflow: pre-training
+            image_name: maxtext_jax_nightly
+          - device: tpu
+            build_mode: nightly
+            workflow: post-training
+            image_name: maxtext_post_training_nightly
           - device: gpu
             build_mode: stable
+            workflow: pre-training
             image_name: maxtext_gpu_jax_stable
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
           - device: gpu
             build_mode: nightly
+            workflow: pre-training
             image_name: maxtext_gpu_jax_nightly
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
-    uses: ./.github/workflows/build_and_push_docker_image.yml
+
+    uses: ./.github/workflows/promote_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}
+      image_name: ${{ matrix.image_name }}${{ inputs.image_suffix  }}
+      image_tag: ${{ github.run_id }}
       device: ${{ matrix.device }}
-      build_mode: ${{ matrix.build_mode }}
-      dockerfile: ${{ matrix.dockerfile }}
-      maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
-      image_date: ${{ needs.setup.outputs.image_date }}
+      workflow: ${{ matrix.workflow }}
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}

.github/workflows/build_and_push_docker_image.yml

@@ -45,6 +45,9 @@ on:
         required: false
         type: string
         default: ''
+    secrets:
+      HF_TOKEN:
+        required: true
 
 permissions:
   contents: read
@@ -80,6 +83,14 @@ jobs:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
           INPUTS_BUILD_MODE: ${{ inputs.build_mode }}
 
+      - name: Matrix Debugger
+        run: |
+          echo "device: ${{ inputs.device }}"
+          echo "workflow: ${{ inputs.workflow }}"
+          echo "build_mode: ${{ inputs.build_mode }}"
+          echo "image_name: ${{ inputs.image_name }}"
+          echo "dockerfile: ${{ inputs.dockerfile }}"
+
       - name: Checkout MaxText
         uses: actions/checkout@v5
         if: steps.check.outputs.should_run == 'true'
@@ -126,27 +137,25 @@ jobs:
         shell: bash
         run: |
           SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
+          TEMP_IMG="${SOURCE_IMAGE}:${{ github.run_id }}"
 
           if [[ $INPUTS_VERSION_NAME ]]; then
             echo "Tagging docker images corresponding to PyPI release..."
-            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_VERSION_NAME}" --quiet
+            gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:${INPUTS_VERSION_NAME}" --quiet
           else
             echo "Tagging docker images corresponding to nightly release..."
 
             # Add date tag
-            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
+            gcloud container images add-tag "${TEMP_IMG}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
 
             # Convert date to YYYYMMDD format
             clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
 
             # Add MaxText tag
-            maxtext_hash=$(git rev-parse --short HEAD)
-            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
-
-            # Add latest tag (TODO: add this tag only after tests pass)
-            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:latest" --quiet
+            gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:maxtext_${MAXTEXT_SHA}_${clean_date}" --quiet
           fi
         env:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
           INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
           INPUTS_VERSION_NAME: ${{ inputs.version_name }}
+          MAXTEXT_SHA: ${{ inputs.maxtext_sha }}