[JIT] [Bugfix] Make blob codegen cache publication transactional

[Ruye-aa]

Summary

• run blob generators in a unique staging directory with checked subprocesses;
• reject successful generator runs that produce no C++/HIP source;
• publish generated sources only after compilation succeeds and restore the previous
  source cache if publication fails;
• atomically replace the installed module so a failed rebuild keeps the previous
  healthy .so;
• add CPU-only regression tests for interrupted generation and publication.

Problem

build_module() currently writes codegen output directly into the canonical blob/
directory and ignores the os.system() return code. It also removes the installed
module before a replacement has compiled. A failed or interrupted generator can
therefore leave partial sources that are compiled or reused, while also destroying the
last known-good binary.

This was observed on the CK MHA JIT path as a truncated fmha_fwd_api.cpp with no
matching dispatch branch, followed by persistent invalid argument for fmha_fwd
failures until the cache was deleted.

Implementation

The change moves generator execution into aiter/jit/utils/jit_cache.py:

1. all commands for a module write to one unique staging directory;
2. subprocess.run(..., check=True) propagates generator failure;
3. the staging directory must contain at least one .cpp or .cu source;
4. compilation consumes only the staged files;
5. the staged source directory replaces blob/ after successful compilation, with
   rollback if directory publication fails;
6. the compiled module is copied to a temporary sibling and published with
   os.replace().

Tests

python3 -m pytest -q op_tests/test_jit_cache_transaction.py
.....                                                                    [100%]
5 passed

python3 -m py_compile \
  aiter/jit/core.py \
  aiter/jit/utils/jit_cache.py \
  op_tests/test_jit_cache_transaction.py

git diff --check

The regression tests are CPU-only and cover:

• a generator that writes a partial source and exits non-zero;
• a zero-source generator result;
• delayed publication of successful codegen output;
• rollback when source-directory publication fails;
• preservation of an existing .so when copying the replacement fails.

GPU compilation and CK MHA runtime validation are left to the upstream CI because the
local development host has no ROCm GPU.

Closes #4799

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4800 --add-label <label>

[Ruye-aa]

Hi @zufayu and @valarLip, could you please help move #4800 forward?

@zufayu, the required GitHub Actions workflows are still awaiting maintainer approval. Could you approve the pending runs, or point me to a maintainer who can?

@valarLip, when you have time, could you please review this change?

This PR fixes interrupted CK MHA JIT code generation leaving an incomplete persistent cache and causing repeated invalid argument for fmha_fwd failures.

I also applied the patch to the current main in a local scratch worktree. It applies cleanly, and the five CPU regression tests, py_compile, git diff --check, and Black all pass. GPU/ROCm validation is still pending upstream CI.

Thanks!

[Ruye-aa]

Follow-up: I pushed d3c26cf to fix all three Ruff findings from the previous Checks run:

• removed the unused E402 noqa directive;
• combined the two nested context-manager pairs flagged by SIM117.

Local verification now passes for ruff 0.16.0 on the changed file, Black, all five CPU regression tests, py_compile, and git diff --check.

The new push dismissed the previous approval and created a fresh set of workflow runs that are awaiting maintainer approval. @junhaha666, could you please re-approve the updated commit? @zufayu, could you please approve the pending workflows?

[junhaha666]

Hello, it failed the CI test.

[Ruye-aa]

Thanks @junhaha666 for flagging the CI failure. I fixed the reported issue in 449d3e60.

The full Aiter Test was delayed while the workflows were awaiting maintainer approval, so this compatibility gap was detected later than it should have been. The branch also had to be kept in sync with a fast-moving main, and each sync or fix created a new set of workflows requiring approval.

The failure was caused by a valid header-only HSA code generator: it produces asm_*_configs.hpp, while the new transactional validation accepted only .cpp and .cu outputs and incorrectly treated the header as no generated build input. The fix now accepts C++/HIP sources and headers (.cpp, .cu, .h, .hpp, .cuh) while still rejecting a generator that truly produces no build input.

Local validation for the new head:

• the real hsa/codegen.py -m pa header-only path passes transactional staging;
• all 6 CPU regression tests pass;
• Ruff 0.16.0, Black, py_compile, and git diff --check pass;
• the branch is currently up to date with main (behind_by=0).

@zufayu, could you please approve the current pending workflow runs as soon as possible so this fix can be validated by the upstream Python 3.12 wheel/GPU CI? Prompt approval would avoid another long feedback cycle and another branch sync resetting the runs.

@valarLip, could you please review the updated change once the workflows are available and pass?

Thank you.