[Raiden Weight Sync 1/7] Add cross-repo drift guard and FFI resolution tests

[YixuanWang-99]

Overview

Part of the stacked Raiden weight-sync enablement PR chain replacing #5089.

Stack:

• [M7] AI-Hypercomputer/maxtext (this PR): Cross-repo drift guard
• [M1] AI-Hypercomputer/maxtext: Inhomogeneous layer cycle support in raiden_unscan
• [M2] AI-Hypercomputer/maxtext: MoE 128-lane layout and padding (pairs with Tunix T5)
• [M3] AI-Hypercomputer/maxtext: Target-free streaming weight conversion
• [M5] AI-Hypercomputer/maxtext: Shared expert gate weight_dtype fix
• [M4] AI-Hypercomputer/maxtext: TrainingEngine Raiden FFI integration (pairs with Tunix T2a)
• [M6] AI-Hypercomputer/maxtext: Rollout cleanup

Details

• Test-only drift guard to prevent drift against Tunix interfaces.
• Includes table-driven test verifying FFI transport resolution across Pathways proxy, CPU, wheel presence, and host staging combinations.

Verification

• pytest tests/post_training/unit/cross_repo_drift_test.py (5/5 passed).

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.

[gemini-code-assist]

Code Review

This pull request introduces a new test file, cross_repo_drift_test.py, to verify that vendored MoE weight interleaving functions and constants match across the maxtext and tunix repositories, and to test the Raiden FFI resolution matrix. Feedback on the tests suggests comparing both parameter names and default values in the signature drift guard to prevent silent behavior drift, and avoiding clear=True when patching os.environ to prevent potential test flakiness from missing ambient environment variables.

[gemini-code-assist]

Comparing only the parameter names (keys()) in the signature drift guard is insufficient. If a default value of a parameter changes in one of the repositories, the parameter names will still match, but the behavior will silently drift. To prevent this, compare both the parameter names and their default values.

Suggested change

	self.assertEqual(
	list(maxtext_sig.parameters.keys()),
	list(tunix_sig.parameters.keys()),
	)
	self.assertEqual(
	[(p.name, p.default) for p in maxtext_sig.parameters.values()],
	[(p.name, p.default) for p in tunix_sig.parameters.values()],
	)

[gemini-code-assist]

Using clear=True in mock.patch.dict completely empties os.environ during the test execution. This can cause unexpected failures or flakiness if JAX, pytest, or other system/test runner components rely on ambient environment variables (such as PATH, PYTHONPATH, or JAX configuration variables). Instead, patch the environment without clearing it, and explicitly remove RAIDEN_USE_FFI if env_val is None.

      with mock.patch.dict("os.environ", env_dict), mock.patch.object(
          raiden_synchronizer, "_get_raiden_ffi", return_value=fake_wheel
      ):
        if env_val is None:
          os.environ.pop("RAIDEN_USE_FFI", None)

[SurbhiJainUSC]

Why do we need this test?

[khatwanimohit]

+1

[khatwanimohit]

can you move the import on top

[khatwanimohit]

+1

[YixuanWang-99]

Closing as per maintainer feedback: cross-repo drift guard test is unnecessary.