ci: install published FlagOS wheels across accelerator jobs

[lvyufeng]

AI Agent Information

• Agent/Tool: OpenAI Codex desktop
• Model: GPT-6
• Human Reviewer: @lvyufeng
• Session Summary: Investigated recurring one-hour accelerator CI timeouts. The request was to use published FlagOS wheels for all seven platforms and prepare matching toolchains where host driver versions remain unknown.

Summary

An observed GCU integration job spent about 40 minutes retrying a FlagGems Git clone and reached its one-hour timeout shortly after tests began. This PR installs published FlagGems 5.4.0 and vendor-specific FlagTree 0.7.0 RC wheels across all seven accelerator scripts; a later GCU run reduced environment preparation from 54m47s to 3m20s and passed integration tests. DCU uses its matching FlagCX RC wheel, while CUDA, MetaX, GCU, and MUSA wheel paths remain gated by toolkit compatibility. CUDA skips the only published native FlagGems C++ wheel because it fails to load with the CI PyTorch 2.10 ABI; it also disables FlagTree's optional TLE distributed loader on CUDA 13.0, since the library that loader discovers requires libcudart.so.12. MUSA stays on its previously running image until the runner driver and image delivery can support MUSA 5.2.

Change Type

• Bug Fix
• Performance Optimization
• CI/Infrastructure

Platforms Affected

• CUDA
• MetaX
• Ascend
• PPU
• Platform-agnostic (all platforms, including DCU, GCU, and MUSA)

Problem Analysis

What was broken/missing?

The one-hour integration budget was consumed largely by environment setup rather than tests. The GCU job in run 35949082553 took about 55 minutes to prepare, then ran tests for about two minutes before cancellation.

Why did it happen?

The first FlagGems Git clone ran from 06:08:03 to 06:26:59 and failed with an HTTP/2 error. Its retry ran from 06:27:11 to 06:48:04, consuming about 40 minutes in total. Other platform scripts also cloned FlagGems. Published Python wheels remove those clone steps; the unavailable compatible CUDA native wheel is handled by disabling that optional C++ route.

Investigation process:

1. Inspected the timed-out job log and separated environment setup from test execution.
2. Examined all seven setup scripts, shared pins, platform configs, and dedicated workflows.
3. Checked the published wheel indexes, versions, Python tags, and toolkit variants; inspected the CUDA native wheel payload and its PyTorch 2.10 loader failure.
4. Compared wheel toolkit requirements with current images and runner driver reports, and measured GCU setup time on a successful PR run.

Solution Design

Implementation approach:

Install exact binary versions from the appropriate FlagOS index using --only-binary=:all: and --no-deps, retaining isolated CPU PyTorch environments. Check the installed FlagGems version and importability. Use the matching DCU FlagCX wheel; keep CUDA, MetaX, GCU, and MUSA FlagCX installation conditional on compatible image/toolkit upgrades.

Key design decisions:

• Keep current CUDA 13.0, MetaX 3.8.0, GCU 1.9.7, and previously working MUSA 5.1 CI images while host driver compatibility remains unconfirmed; the matching FlagCX paths remain prepared.
• Disable the optional CUDA FlagGems C++ route on the bootstrapped PyTorch 2.10 path. The only published native wheel loads libtriton_jit.so with an undefined PyTorch C++ symbol, while the Python FlagGems wheel remains usable.
• Disable FlagTree's optional TLE distributed loader in the CUDA 13.0 environment with USE_TLE_DIST=0. The CUDA job has no compatible FlagCX wheel, yet FlagTree discovers a libflagcx.so that requires libcudart.so.12; importing FlagGems would otherwise fail before tests run. Preserve the gate across CI steps and smoke-test the TLE import during setup.
• Ascend's published CANN 9 FlagCX wheel needs libtorch_npu.so, which the intentionally torch_npu-free environment does not have. No compatible PPU FlagCX wheel was found. Both platforms still receive FlagTree and FlagGems wheels.
• Preserve the original PR #417 for the common helper refactor. This branch stacks the wheel migration on its commit.

Code changes by file:

• .github/version-pins.env: Pin FlagGems 5.4.0, seven FlagTree RC builds, and available FlagCX variants.
• .github/scripts/lib/set_env_common.sh: Replace VCS installation with checked, binary-only wheel installation.
• .github/scripts/set_env_ascend.sh: Install published Ascend FlagTree and FlagGems wheels and matching runtime dependencies.
• .github/scripts/set_env_cuda.sh: Install the published Python FlagGems wheel, disable the incompatible native C++ route on bootstrap and the CUDA 12-linked TLE distributed loader on CUDA 13.0, and prepare CUDA 13.3 FlagCX behind a toolkit check.
• .github/scripts/set_env_dcu.sh: Install published Hygon FlagTree, FlagGems, and DTK FlagCX wheels, reusing valid preinstalled versions; persist the autoload gate across CI steps.
• .github/scripts/set_env_gcu.sh: Install Enflame wheels and prepare FlagCX for a future TOPS 1.9.10 image, including its autoload gate.
• .github/scripts/set_env_metax.sh: Install MetaX wheels and prepare the isolated venv and FlagCX path for a future MACA 3.8.1.3 image, including its autoload gate.
• .github/scripts/set_env_musa.sh: Install MThreads FlagTree and FlagGems wheels; keep MUSA 5.2 FlagCX as an explicit opt-in for a matching future image/driver.
• .github/scripts/set_env_ppu.sh: Install published PPU FlagTree and FlagGems wheels through the runner's direct route.
• tests/integration/ops/test_dtype_route_fallback.py: Ignore only the known Ascend FlagTree startup warning in the subprocess-output assertion.
• tests/integration/test_compute_device_index.py: Use exactly representable addmm inputs so this device-guard test does not fail on normal vendor matmul rounding.
• .github/configs/ascend.yml: Update the expected FlagTree release.
• .github/configs/cuda.yml: Update the expected FlagTree release.
• .github/configs/musa.yml: Pin the previously running MUSA 5.1 image by digest until a 5.2 runner upgrade is ready.
• .github/configs/ppu.yml: Update the expected FlagTree release.
• .github/workflows/agnostic-checks.yml: Included from the prerequisite ci: share the set_env_* helper functions in one library #417 helper refactor; unchanged by the wheel migration commit.
• .github/workflows/integration-test-cuda.yml: Correct the setup comment after removing the native source build.
• .github/workflows/integration-test-musa.yml: Use the same pinned MUSA 5.1 image as the platform config.
• pyproject.toml: Document the published, vendor-specific FlagCX distribution.
• tests/unit/test_ci_version_pins.py: Enforce shared wheel pins rather than source revisions.
• tests/unit/test_set_env_common.py: Keep the CUDA interpreter override contract check valid as calls change.

Changes by commit:

1. cafd32f - ci: share the set_env_* helper functions in one library: prerequisite from ci: share the set_env_* helper functions in one library #417.
2. 8f25bed - ci: install published FlagOS wheels across accelerator jobs: binary wheel and toolchain changes in this PR.
3. 5176d84 - fix: handle FlagCX autoload and Ascend wheel warning in CI: repair two failures found in the first PR CI run.
4. 76bb7e2 - fix: persist FlagCX autoload gate across CI steps: prevent the isolated wheel build backend from auto-importing FlagCX on DCU and MUSA.
5. a540ce0 - fix: keep CI on compatible wheel and image paths: avoid the CUDA C++ ABI mismatch, make MetaX's device-guard assertion numerically stable, and defer the MUSA 5.2 image until its host driver and pull time are addressed.
6. bc0278d - fix: gate FlagTree TLE distributed loader on CUDA 13.0: prevent FlagTree from loading a CUDA 12-linked FlagCX library in the CUDA 13.0 CI environment.

Verification

Pre-submission Checklist

• Linting passed (ruff check ., ruff format --check .)
• Type checking passed (not applicable to shell/YAML changes)
• All tests pass (full suite cannot collect without torch on this macOS host)
• Manual testing completed (log reproduction, published wheel and payload checks, shell/YAML/static verification; hardware CI remains to run)
• No debug/temporary code
• Documentation updated (dependency comments and setup documentation in changed files)
• Commit messages follow conventions
• All text in English

Linting Results

$ ruff check .
All checks passed!
$ ruff format --check .
303 files already formatted
$ bash -n .github/scripts/lib/set_env_common.sh .github/scripts/set_env_{ascend,cuda,dcu,gcu,metax,musa,ppu}.sh .github/version-pins.env
(exit 0)
$ git diff --check
(exit 0)
Ruby YAML parse of modified YAML files: YAML parse passed
$ bash -n .github/scripts/set_env_{cuda,dcu,gcu,metax,musa}.sh
(exit 0, after follow-up fix)

Test Results

$ python -m pytest tests/unit/test_ci_version_pins.py tests/unit/test_dcu_env_script.py tests/unit/test_set_env_common.py -v -p no:cacheprovider
============================== 14 passed in 0.03s ==============================

$ python -m pytest tests/unit/test_setup_build.py tests/unit/test_ci_version_pins.py tests/unit/test_set_env_common.py -q -p no:cacheprovider
27 passed in 0.21s

$ python -m pytest tests/unit/ -v -p no:cacheprovider
1 skipped, 11 warnings, 21 errors during collection: local Python 3.14.6 environment has no torch.

$ python -m pytest tests/integration/ops/ -v -m 'anyplatform or cuda' -p no:cacheprovider
Collection/configuration failed: ModuleNotFoundError: No module named 'torch'.

Manual Verification

Before: GCU run 35949082553 spent 06:08:03-06:26:59 on a failing FlagGems clone and 06:27:11-06:48:04 retrying it; the one-hour job was canceled after about two minutes of tests.
After: all seven setup scripts request published FlagGems 5.4.0 wheels and binary-only FlagTree. GCU run 35997727477 prepared its integration environment in 3m20s and passed its tests. The incompatible CUDA native wheel is omitted; CUDA's Python FlagGems path remains enabled. The latest CUDA TLE loader fix still needs hardware CI validation.
PR CI run 35997727477: DCU, Ascend, PPU, and both GCU jobs passed. CUDA failed loading the only published FlagGems C++ wheel with an undefined c10 symbol. MetaX failed one unseeded addmm comparison by 0.00209 against a 0.001 tolerance; an earlier run passed, and commit a540ce0 uses exactly representable values. MUSA timed out while pulling its 10.5 GB compressed 5.2 image, before container initialization. Commit a540ce0 returns to the previously running MUSA 5.1 image and leaves FlagCX 5.2 as an opt-in. The new CI run must validate those changes.
PR CI run 36047006888: Ascend, DCU, GCU, MetaX, MUSA, and PPU passed. CUDA passed device availability but failed importing FlagGems: FlagTree's optional TLE distributed loader found a FlagCX library linked to `libcudart.so.12`, absent from the CUDA 13.0 environment. The follow-up disables only that optional loader on the current CUDA toolkit path; a new CI run must validate it.

Performance Impact

Environment setup timing

Metric	Before	After
GCU FlagGems source clone/retry, observed run 35949082553	About 40 minutes	Removed from script
GCU integration environment preparation	54m47s	3m20s in run 35997727477

The published Enflame FlagGems wheel is about 9.4 MB. The measured GCU improvement is from separate CI runs on shared hardware, so network and runner variation can affect it.

Code Quality Verification

Style Consistency

• Matched existing code style and naming conventions
• Used the shared pip_retry and setup library
• Kept explanatory comments around toolkit gates and ABI constraints

Edge Cases Considered

1. Stale image packages: verify the exact FlagGems version and an importable module; DCU also validates preinstalled FlagTree and FlagCX.
2. Vendor ABI mismatch: use --no-deps, preserve CPU PyTorch isolation, and gate incompatible FlagCX wheels by toolkit/image.
3. Failed wheel resolution: --only-binary=:all: fails rather than falling back to a source build.

Potential Risks

1. The CUDA TLE distributed-loader fix still needs on-device CI validation.
2. MUSA 5.2 requires a confirmed host driver upgrade and a pre-cached or faster-delivered image. CUDA 13.3 and MetaX 3.8.1.3 image switches also require driver confirmation.
3. The optional CUDA FlagGems C++ route is disabled on the bootstrap path until a PyTorch 2.10-compatible wheel is published.

Rollback Plan

Revert the wheel-migration and follow-up commits starting at 8f25bed if the remaining CI paths cannot be made compatible; leave prerequisite #417 independent.

Related Work

• Depends on #417 for the shared setup helpers.
• Timeout evidence: GCU run 35949082553.

Explicitly Not Included

• Switching CUDA, MetaX, or MUSA to newer images before compatible host drivers are confirmed, or switching GCU before its matching toolchain is deployed.
• CUDA FlagGems C++ dispatch on the bootstrap path until a compatible native wheel exists.
• A PPU FlagCX wheel or CANN 9 FlagCX in the torch_npu-free Ascend environment; neither is compatible with the current runners.

Human Review Notes

Areas needing special attention:

1. MUSA runner driver upgrade and image delivery before enabling the prepared 5.2 FlagCX path.
2. CUDA C++ wheel ABI against PyTorch 2.10 and MetaX /flagos upgrade path.
3. Follow-up image/driver alignment before enabling CUDA 13.3 and MetaX 3.8.1.3 FlagCX wheels.

Questions for reviewer:

1. When will the MUSA runner's driver and image cache be ready for 5.2 and its FlagCX wheel?
2. Which CUDA and MetaX host driver versions are deployed, so their prepared image paths can be enabled later?

Additional Context

The local environment lacks torch, so full unit and integration collection could not complete here. The measured GCU setup improvement and chip job results are linked above; the latest commit needs another CI run.