Refactor Github Action per b/485167538 (#3181)

Co-authored-by: Ben Knutson <benknutson@google.com>

Author: google-admin

.github/workflows/build_and_push_docker_image.yml

@@ -66,13 +66,18 @@ jobs:
         id: check
         shell: bash
         run: |
-          if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${{ github.event.inputs.target_device }}" != "all" && "${{ github.event.inputs.target_device }}" != "${{ inputs.device }}" ]]; then
+          if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${GITHUB_EVENT_INPUTS_TARGET_DEVICE}" != "all" && "${GITHUB_EVENT_INPUTS_TARGET_DEVICE}" != "${INPUTS_DEVICE}" ]]; then
             echo "should_run=false" >> $GITHUB_OUTPUT
-            echo "Skipping ${{ inputs.image_name }} build for device: ${{ inputs.device }} in ${{ inputs.build_mode }} mode."
+            echo "Skipping ${INPUTS_IMAGE_NAME} build for device: ${INPUTS_DEVICE} in ${INPUTS_BUILD_MODE} mode."
           else
             echo "should_run=true" >> $GITHUB_OUTPUT
-            echo "Building ${{ inputs.image_name }} for device: ${{ inputs.device }} in ${{ inputs.build_mode }} mode."
+            echo "Building ${INPUTS_IMAGE_NAME} for device: ${INPUTS_DEVICE} in ${INPUTS_BUILD_MODE} mode."
           fi
+        env:
+          GITHUB_EVENT_INPUTS_TARGET_DEVICE: ${{ github.event.inputs.target_device }}
+          INPUTS_DEVICE: ${{ inputs.device }}
+          INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
+          INPUTS_BUILD_MODE: ${{ inputs.build_mode }}
 
       - name: Checkout MaxText
         uses: actions/checkout@v5
@@ -125,13 +130,13 @@ jobs:
         if: steps.check.outputs.should_run == 'true'
         shell: bash
         run: |
-          SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}"
+          SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
 
           # Add date tag
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${{ inputs.image_date }}" --quiet
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
 
           # Convert date to YYYYMMDD format
-          clean_date=$(echo "${{ inputs.image_date }}" | sed 's/[-:]//g' | cut -c1-8)
+          clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
 
           # Add MaxText tag
           maxtext_hash=$(git rev-parse --short HEAD)
@@ -147,3 +152,6 @@ jobs:
               fi
             done
           fi
+        env:
+          INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
+          INPUTS_IMAGE_DATE: ${{ inputs.image_date }}

.github/workflows/build_and_test_maxtext.yml

@@ -57,8 +57,8 @@ jobs:
             exit 0
           fi
 
-          git fetch origin ${{ github.base_ref }}
-          CHANGED_FILES=$(git diff --name-only origin/${{ github.base_ref }}...HEAD)
+          git fetch origin ${GITHUB_BASE_REF}
+          CHANGED_FILES=$(git diff --name-only origin/${GITHUB_BASE_REF}...HEAD)
 
           echo "Changed files:"
           echo "$CHANGED_FILES"
@@ -273,20 +273,20 @@ jobs:
       - name: Check test results
         run: |
           # If doc-only, all tests should be skipped
-          if [ "${{ needs.doc_only_check.outputs.run_tests }}" == "false" ]; then
+          if [ "${NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS}" == "false" ]; then
             echo "Documentation-only changes detected, tests were skipped"
             exit 0
           fi
 
           # Otherwise, check that build and all tests passed or were skipped
-          echo "Build result: ${{ needs.build_and_upload_maxtext_package.result }}"
-          echo "CPU tests: ${{ needs.maxtext_cpu_unit_tests.result }}"
-          echo "TPU tests: ${{ needs.maxtext_tpu_unit_tests.result }}"
-          echo "TPU integration: ${{ needs.maxtext_tpu_integration_tests.result }}"
-          echo "TPU pathways: ${{ needs.maxtext_tpu_pathways_unit_tests.result }}"
-          echo "TPU pathways integration: ${{ needs.maxtext_tpu_pathways_integration_tests.result }}"
-          echo "GPU tests: ${{ needs.maxtext_gpu_unit_tests.result }}"
-          echo "GPU integration: ${{ needs.maxtext_gpu_integration_tests.result }}"
+          echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}"
+          echo "CPU tests: ${NEEDS_MAXTEXT_CPU_UNIT_TESTS_RESULT}"
+          echo "TPU tests: ${NEEDS_MAXTEXT_TPU_UNIT_TESTS_RESULT}"
+          echo "TPU integration: ${NEEDS_MAXTEXT_TPU_INTEGRATION_TESTS_RESULT}"
+          echo "TPU pathways: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
+          echo "TPU pathways integration: ${NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT}"
+          echo "GPU tests: ${NEEDS_MAXTEXT_GPU_UNIT_TESTS_RESULT}"
+          echo "GPU integration: ${NEEDS_MAXTEXT_GPU_INTEGRATION_TESTS_RESULT}"
 
           # Fail only if any job failed or was cancelled (skipped is OK)
           if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
@@ -295,6 +295,16 @@ jobs:
           fi
 
           echo "All required tests passed successfully"
+        env:
+          NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS: ${{ needs.doc_only_check.outputs.run_tests }}
+          NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
+          NEEDS_MAXTEXT_CPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_cpu_unit_tests.result }}
+          NEEDS_MAXTEXT_TPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_unit_tests.result }}
+          NEEDS_MAXTEXT_TPU_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_integration_tests.result }}
+          NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_unit_tests.result }}
+          NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_integration_tests.result }}
+          NEEDS_MAXTEXT_GPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_gpu_unit_tests.result }}
+          NEEDS_MAXTEXT_GPU_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_gpu_integration_tests.result }}
 
   all_notebooks_passed:
     name: All Notebooks Passed
@@ -304,14 +314,14 @@ jobs:
     steps:
       - name: Check notebooks results
         run: |
-          if [ "${{ needs.doc_only_check.outputs.run_notebooks }}" == "false" ]; then
+          if [ "${NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_NOTEBOOKS}" == "false" ]; then
             echo "Non-notebook changes detected, runs were skipped"
             exit 0
           fi
 
           # Otherwise, check that build and notebooks run passed or were skipped
-          echo "Build result: ${{ needs.build_and_upload_maxtext_package.result }}"
-          echo "Jupyter Notebooks result: ${{ needs.maxtext_jupyter_notebooks.result }}"
+          echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}"
+          echo "Jupyter Notebooks result: ${NEEDS_MAXTEXT_JUPYTER_NOTEBOOKS_RESULT}"
 
           # Fail only if any job failed or was cancelled (skipped is OK)
           if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
@@ -320,6 +330,10 @@ jobs:
           fi
 
           echo "All required notebooks passed successfully"
+        env:
+          NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_NOTEBOOKS: ${{ needs.doc_only_check.outputs.run_notebooks }}
+          NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
+          NEEDS_MAXTEXT_JUPYTER_NOTEBOOKS_RESULT: ${{ needs.maxtext_jupyter_notebooks.result }}
 
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build

.github/workflows/run_tests_against_package.yml

@@ -101,18 +101,18 @@ jobs:
       - name: Run Tests
         shell: bash
         run: |
-          if [ "${{ inputs.is_scheduled_run }}" == "true" ]; then
-            FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
+          if [ "${INPUTS_IS_SCHEDULED_RUN}" == "true" ]; then
+            FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER}"
           else
-            FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only"
+            FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER} and not scheduled_only"
           fi
           # TODO: Use package data for testing and remove the env vars
           export MAXTEXT_REPO_ROOT=$(pwd)
           export MAXTEXT_ASSETS_ROOT=$(pwd)/src/maxtext/assets
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
           export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           # omit this libtpu init args for gpu tests
-          if [ "${{ inputs.device_type }}" != "cuda12" ]; then
+          if [ "${INPUTS_DEVICE_TYPE}" != "cuda12" ]; then
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
           fi
           if [ "${{ inputs.total_workers }}" -gt 1 ]; then
@@ -122,7 +122,7 @@ jobs:
             SPLIT_ARGS=""
           fi
           # TODO: Fix the skipped tests and remove the deselect flags
-          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} \
+          .venv/bin/python3 -m pytest ${INPUTS_PYTEST_ADDOPTS} \
             -v \
             -m "${FINAL_PYTEST_MARKER}" \
             --durations=0 \
@@ -134,6 +134,10 @@ jobs:
             $SPLIT_ARGS
         env:
           PYTHONPATH: "${{ github.workspace }}/src"
+          INPUTS_IS_SCHEDULED_RUN: ${{ inputs.is_scheduled_run }}
+          INPUTS_PYTEST_MARKER: ${{ inputs.pytest_marker }}
+          INPUTS_DEVICE_TYPE: ${{ inputs.device_type }}
+          INPUTS_PYTEST_ADDOPTS: ${{ inputs.pytest_addopts }}
       - name: Upload results to Codecov
         uses: codecov/codecov-action@v5
         continue-on-error: true