Trigger unit tests for docker images upload

[xibinliu]

Description

Trigger the CI with the stable and nightly images

Tests

Image Build and Test workflows run by this PR: here

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[shralex]

Should we split this file or add an option to distinguish stable image building and testing from nightly building and testing ? Assuming that nightly might be frequently broken, my concern is that bundling the two would prevent us from publishing stable images to pypi

[xibinliu]

The build and tests jobs are independent for stable and nightly images, so they won't impact each other. If we want to keep them in one file depends on how the pypi publishing process use the output from this workflow. @SurbhiJainUSC could you comment here?

[SurbhiJainUSC]

This workflow uploads the docker image to PyPI. This is independent of pypi release process. Also, we will be able to release stable docker image to GCR even if nightly docker image fails.

[shralex]

this should probably just be true throughout this file. The idea for this flag is to run more tests when this isn't affecting a PR

[xibinliu]

The trigger could be scheduled or manual for this workflow. But it probably okay to run full tests even with manual trigger. I now changed is_scheduled_run to true throughout the file.

[shralex]

why do we need to make worker_group and total_workers parameters ? do these change across invocations from uploadDockerImages and build_and_test_maxtext ?

[xibinliu]

Good point. I removed the "worker_group" and "total_workers" as parameters from the coordinator, and set them based on the other input (cpu or not).

[shralex]

Instead of using complex selection operators here consider creating a setup job that maps the flavor to specific outputs that would then be used as parameters. This centralizes the logic and provides implicit validation. Something like this:

configure:
outputs:
device_type: ${{ steps.map.outputs.device_type }}
device_name: ${{ steps.map.outputs.device_name }}
# ... other outputs
steps:
- id: map
run: |
case "${{ inputs.flavor }}" in
"tpu-unit")
echo "device_type=tpu" >> $GITHUB_OUTPUT
echo "device_name=v6e-4" >> $GITHUB_OUTPUT
# also for pytest markers etc
;;
"gpu-unit")
echo "device_type=cuda12" >> $GITHUB_OUTPUT
echo "device_name=a100-40gb-4" >> $GITHUB_OUTPUT
;;
*)
echo "::error::Unsupported flavor: ${{ inputs.flavor }}"
exit 1
;;
esac

execute-test-package:
needs: configure
uses: ./.github/workflows/run_tests_against_package.yml
with:
device_type: ${{ needs.configure.outputs.device_type }}
# ...

[xibinliu]

The "configure" job has to be run a VM, isn't this heavier than just using selection operators?