[PROTOTYPE] E2E airflow tests workflow

[SurbhiJainUSC]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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]

@gemini-cli /investigate

[github-actions]

🤖 Hi @shralex, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

• Job/Matrix: cpu-unit tests / cpu-unit / run (Worker 3)
• Failing Test: N/A (Job failed during container setup)
• Error: FetchError: request to https://34.118.224.1/apis/authorization.k8s.io/v1/selfsubjectaccessreviews failed, reason: connect ETIMEDOUT 34.118.224.1:443

🪵 Error Details & Stack Trace

cpu-unit tests / cpu-unit / run UNKNOWN STEP    2026-05-15T00:19:41.0254735Z ##[error]FetchError: request to https://34.118.224.1/apis/authorization.k8s.io/v1/selfsubjectaccessreviews failed, reason: connect ETIMEDOUT 34.118.224.1:443
cpu-unit tests / cpu-unit / run UNKNOWN STEP    2026-05-15T00:19:41.0310978Z ##[error]Process completed with exit code 1.
cpu-unit tests / cpu-unit / run UNKNOWN STEP    2026-05-15T00:19:41.0334903Z ##[error]Executing the custom container implementation failed. Please contact your self hosted runner administrator.

💡 Root Cause Analysis & Context

Confidence: high

The failure is an infrastructure/environment flake related to the self-hosted GitHub Actions runner. Specifically, the runner timed out while attempting to connect to the Kubernetes API server (34.118.224.1) to perform a selfsubjectaccessreview during the job initialization phase. This happened before any test code was executed.

The code changes in this PR (refactoring maybe_initialize_jax_distributed_system and adding a new E2E workflow) are unrelated to this connectivity issue. Furthermore, the tpu-unit and gpu-unit jobs passed successfully, which confirms that the logical changes in src/maxtext/utils/max_utils.py and the new unit tests in tests/unit/max_utils_test.py are correct and do not introduce regressions for those backends.

🛠️ Recommended Action

Since this is an infrastructure flake, the job should be re-run. No code changes are required to address this specific failure.