Add dump and inference-engine-checksum comparison helpers for FT tests

[fzyzcjy]

Add comparison helpers used by fault-tolerance tests to compare dumped tensors and
inference-engine checksums offline (generic comparators, dump comparison, and an
inference-engine checksum comparison built on the event-analyzer checksum rule).

• miles/utils/test_utils/comparisons/{comparators,dumps,inference_engine_checksums}.py and tests.

[gemini-code-assist]

Code Review

This pull request introduces test comparison utilities, including a runner for the tensor comparator, a helper to compare leaf dump directories, and a utility to compare inference engine weight checksums, along with their respective unit tests. The review feedback highlights two key issues: first, a bug in run_comparator where the 'rank' key is incorrectly omitted from skip_keys when grouping_skip_keys is provided; second, a potential assertion failure in compare_inference_engine_checksums during fault-tolerance tests with retries due to duplicate rollout events, which should be resolved by only keeping the latest event per rollout.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

There is a bug in how skip_keys is constructed. If grouping_skip_keys is provided, "rank" is completely omitted from the list. According to the comments, "rank" must always be skipped to avoid baseline_load_failed errors when comparing FT and non-FT runs. We should ensure "rank" is always included in skip_keys regardless of whether grouping_skip_keys is provided.

Suggested change

	skip_keys: list[str] = list(grouping_skip_keys) if grouping_skip_keys is not None else ["rank"]
	skip_keys: list[str] = ["rank"] + (list(grouping_skip_keys) if grouping_skip_keys is not None else [])

[Shi-Dong]

This looks like a bug.

[gemini-code-assist]

In fault-tolerance (FT) tests, rollouts can be retried or re-run after a failure, which appends duplicate InferenceEngineWeightChecksumEvents for the same rollout_id to the event log. The current implementation asserts that there are no duplicate events for any rollout_id (lines 51-53), which will cause the comparison helper to fail during any test run that involves a retry. Additionally, running the internal consistency check on all events (including aborted/failed attempts) can lead to false positives. We should filter the events to only keep the latest (successful) attempt for each rollout_id before performing the consistency check and the baseline-vs-target comparison.

    baseline_by_rollout = _latest_events_by_rollout(baseline)
    target_by_rollout = _latest_events_by_rollout(target)

    # Each side's engines must already agree internally (same invariant as the production rule), so
    # one representative engine per rollout then proves baseline == target regardless of engine count.
    assert not inference_engine_weight_checksum_consistency.check(
        list(baseline_by_rollout.values())
    ), "Baseline engines disagree with each other"
    assert not inference_engine_weight_checksum_consistency.check(
        list(target_by_rollout.values())
    ), "Target engines disagree with each other"

    assert baseline_by_rollout.keys() == target_by_rollout.keys(), (
        f"Engine checksum rollout_id sets differ: "
        f"baseline={sorted(baseline_by_rollout)} "
        f"vs target={sorted(target_by_rollout)}"
    )

    mismatches: list[ChecksumMismatchIssue] = []
    for rollout_id in sorted(baseline_by_rollout):
        mismatches += list(
            compare_flat_dicts(
                a=baseline_by_rollout[rollout_id].engine_checksums[0],
                b=target_by_rollout[rollout_id].engine_checksums[0],
                label_a=f"baseline/rollout_{rollout_id}",
                label_b=f"target/rollout_{rollout_id}",
            )
        )
    assert not mismatches, "Engine weight checksum baseline-vs-target mismatch:\n" + "\n".join(
        f"  - {m.label_a} vs {m.label_b} key {m.key}: {m.value_a} != {m.value_b}" for m in mismatches
    )
    print(f"Engine weight checksum comparison passed: {len(baseline_by_rollout)} rollout(s) compared")


def _latest_events_by_rollout(
    events: list[InferenceEngineWeightChecksumEvent],
) -> dict[int, InferenceEngineWeightChecksumEvent]:
    by_rollout: dict[int, InferenceEngineWeightChecksumEvent] = {}
    for event in events:
        if event.rollout_id is None:
            continue
        assert event.engine_checksums, f"No engine checksums for rollout {event.rollout_id}"
        by_rollout[event.rollout_id] = event
    return by_rollout

[Shi-Dong]

LGTM