Merge branch 'main' into fern-migration

Author: zhenghax

.github/workflows/ci.yaml

@@ -4,8 +4,6 @@ on:
   schedule:
     - cron: '0 7 * * *'
       timezone: "Europe/London"
-    - cron: '0 6 * * 6' # 06:00 Saturday Europe/London; workflow gates to biweekly scale-training runs
-      timezone: "Europe/London"
   push:
     branches:
       - 'pull-request/[0-9]+'
@@ -53,15 +51,10 @@ on:
           - axlearn - run only the tests for axlearn
         options: [full, jax, te, maxtext, axlearn]
         default: full
-      SCALE_TRAINING:
-        type: boolean
-        description: "Whether to run scale-training specific tests."
-        default: false
-        required: false
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.event.schedule || github.run_id }}
-  cancel-in-progress: true
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
 
 permissions:
   contents: write       # to fetch code and push branch
@@ -72,7 +65,6 @@ permissions:
 
 env:
   DEFAULT_MANIFEST_ARTIFACT_NAME: bumped-manifest
-  SCALE_TRAINING_BIWEEKLY_ANCHOR_BUILD_DATE: 2026-04-18
 
 jobs:
 
@@ -86,8 +78,6 @@ jobs:
       MANIFEST_BRANCH: ${{ steps.manifest-branch.outputs.MANIFEST_BRANCH }}
       MERGE_BUMPED_MANIFEST: ${{ steps.manifest-branch.outputs.MERGE_BUMBED_MANIFEST }}
       CUDA_IMAGE: ${{ steps.cuda-image.outputs.CUDA_IMAGE }}
-      SHOULD_RUN_WORKFLOW: ${{ steps.run-type.outputs.SHOULD_RUN_WORKFLOW }}
-      IS_SCALE_TRAINING: ${{ steps.run-type.outputs.IS_SCALE_TRAINING }}
       SUFFIX: ${{ steps.run-type.outputs.SUFFIX }}
     steps:
       - name: Cancel workflow run if the trigger is a draft PR
@@ -158,41 +148,8 @@ jobs:
           fi
           echo "CUDA_IMAGE=${CUDA_IMAGE}" >> $GITHUB_OUTPUT
 
-      - name: Determine if scale-training run
-        id: run-type
-        shell: bash -x -e {0}
-        run: |
-          if [[ "${{ inputs.SCALE_TRAINING }}" == "true" ]]; then
-            echo "Manual scale-training run requested"
-            echo "SHOULD_RUN_WORKFLOW=true" >> $GITHUB_OUTPUT
-            echo "IS_SCALE_TRAINING=true" >> $GITHUB_OUTPUT
-            echo "SUFFIX=-scale-training" >> $GITHUB_OUTPUT
-          elif [[ "${{ github.event.schedule }}" == "0 6 * * 6" ]]; then
-            anchor_epoch=$(date -u -d "${{ env.SCALE_TRAINING_BIWEEKLY_ANCHOR_BUILD_DATE }}" +%s)
-            build_epoch=$(date -u -d "${{ steps.date.outputs.BUILD_DATE }}" +%s)
-            days_since_anchor=$(( (build_epoch - anchor_epoch) / 86400 ))
-
-            if (( days_since_anchor >= 0 && days_since_anchor % 14 == 0 )); then
-              echo "Scheduled biweekly scale-training run for build date ${{ steps.date.outputs.BUILD_DATE }}"
-              echo "SHOULD_RUN_WORKFLOW=true" >> $GITHUB_OUTPUT
-              echo "IS_SCALE_TRAINING=true" >> $GITHUB_OUTPUT
-              echo "SUFFIX=-scale-training" >> $GITHUB_OUTPUT
-            else
-              echo "Skipping off-week scale-training schedule for build date ${{ steps.date.outputs.BUILD_DATE }}"
-              echo "SHOULD_RUN_WORKFLOW=false" >> $GITHUB_OUTPUT
-              echo "IS_SCALE_TRAINING=false" >> $GITHUB_OUTPUT
-              echo "SUFFIX=" >> $GITHUB_OUTPUT
-            fi
-          else
-            echo "SHOULD_RUN_WORKFLOW=true" >> $GITHUB_OUTPUT
-            echo "IS_SCALE_TRAINING=false" >> $GITHUB_OUTPUT
-            echo "SUFFIX=" >> $GITHUB_OUTPUT
-          fi
-
-
   bump-manifest:
     needs: metadata
-    if: ${{ needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' }}
     runs-on: ubuntu-22.04
     outputs:
       SOURCE_URLREFS: ${{ steps.source-urlrefs.outputs.SOURCE_URLREFS }}
@@ -249,8 +206,6 @@ jobs:
       - name: Create URL ref build args
         id: source-urlrefs
         shell: bash -x -e {0}
-        env:
-          GH_TOKEN: ${{ github.token}}
         run: |
           # converts manifest yaml to a json object of {SOFTWARE_NAME: URL#REF, ...}
           urlrefs=$(
@@ -266,27 +221,12 @@ jobs:
             urlrefs=$(echo "$urlrefs" | jq -c ". + {\"$PACKAGE\": \"$URLREF\"}")
           done
 
-          # for scale training we need specific XLA and JAX refs
-          if [[ "${{ needs.metadata.outputs.IS_SCALE_TRAINING }}" == "true" ]]; then
-            BRANCH_INFO=$(python3 .github/workflows/scripts/extract_commits.py)
-            echo ${BRANCH_INFO} | jq
-            XLA_BRANCH=$(jq -r '.name' <<< "$BRANCH_INFO")
-            XLA_SHA=$(jq -r '.commit.sha' <<< "$BRANCH_INFO")
-            JAX_REPO=$(jq -r '.jax.repo' <<< "$BRANCH_INFO")
-            JAX_COMMIT=$(jq -r '.jax.commit' <<< "$BRANCH_INFO")
-            urlrefs=$(echo "$urlrefs" | jq -c ". + {\"XLA\": \"https://github.com/openxla/xla.git#${XLA_SHA}\"}")
-            urlrefs=$(echo "$urlrefs" | jq -c ". + {\"JAX\": \"${JAX_REPO}#${JAX_COMMIT}\"}")
-
-            echo "Using XLA branch ${XLA_BRANCH} at sha ${XLA_SHA} for scale-training run"
-            echo "Using JAX repo ${JAX_REPO} at commit ${JAX_COMMIT} from STAGING.md for scale-training run"
-          fi
           echo "SOURCE_URLREFS=${urlrefs}" >> $GITHUB_OUTPUT
           echo "Final URLREF going to be used"
           echo ${urlrefs} | jq -c '.'
 
   amd64:
     needs: [metadata, bump-manifest]
-    if: ${{ needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' }}
     uses: ./.github/workflows/_ci.yaml
     with:
       ARCHITECTURE: amd64
@@ -299,7 +239,6 @@ jobs:
 
   arm64:
     needs: [metadata, bump-manifest]
-    if: ${{ needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' }}
     uses: ./.github/workflows/_ci.yaml
     with:
       ARCHITECTURE: arm64
@@ -313,7 +252,7 @@ jobs:
   # Only merge if everything succeeds
   merge-new-manifest:
     runs-on: ubuntu-22.04
-    if: ${{ !cancelled() && needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' && needs.metadata.outputs.MERGE_BUMPED_MANIFEST == 'true' && needs.metadata.outputs.MANIFEST_BRANCH != github.sha }}
+    if: ${{ !cancelled() && needs.metadata.outputs.MERGE_BUMPED_MANIFEST == 'true' && needs.metadata.outputs.MANIFEST_BRANCH != github.sha }}
     needs:
       - metadata
       - amd64
@@ -424,7 +363,7 @@ jobs:
 
   make-publish-configs:
     runs-on: ubuntu-22.04
-    if:  ${{ !cancelled() && needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' }}
+    if:  ${{ !cancelled() }}
     env:
       MEALKIT_IMAGE_REPO: ${{ needs.metadata.outputs.PUBLISH == 'true' && 'jax-mealkit' || 'mock-jax-mealkit' }}
       FINAL_IMAGE_REPO: ${{ needs.metadata.outputs.PUBLISH == 'true' && 'jax' || 'mock-jax' }}
@@ -513,7 +452,7 @@ jobs:
     needs:
       - metadata
       - make-publish-configs
-    if:  ${{ !cancelled() && needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' && needs.make-publish-configs.outputs.PUBLISH_CONFIGS.config != '{"config":[]}' }}
+    if:  ${{ !cancelled() && needs.make-publish-configs.outputs.PUBLISH_CONFIGS.config != '{"config":[]}' }}
     strategy:
       fail-fast: false
       matrix: ${{ fromJson(needs.make-publish-configs.outputs.PUBLISH_CONFIGS) }}
@@ -529,7 +468,7 @@ jobs:
 
   finalize:
     needs: [metadata, amd64, arm64, publish-containers]
-    if: ${{ !cancelled() && needs.metadata.outputs.SHOULD_RUN_WORKFLOW == 'true' }}
+    if: ${{ !cancelled() }}
     uses: ./.github/workflows/_finalize.yaml
     with:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}