fix: merge test/test-full leaf variants per-cell instead of wholesale

scripts/receive_therock/tests/therock_status_document_test.py

@@ -27,6 +27,7 @@
     StatusDocument,
     Variant,
     _merge_variant_leaf,
+    merge_matrix_test_leaf,
 )
 
 
@@ -448,6 +449,63 @@ def test_merge_status_rolls_up_failure() -> None:
     assert merged.status is Status.failure
 
 
+# --- merge_matrix_test_leaf: leaf-header identity ---------------------------
+
+
+def test_matrix_merge_leaf_header_keeps_newer_run_over_stale_loser() -> None:
+    # A stale older-run snapshot loses the push race but still reaches the
+    # merge (the variant path bypasses RunLeaf.should_replace). Cells are held
+    # by Variant.should_replace; the leaf header must not regress with them.
+    existing = _leaf(
+        run_id=200,
+        run_attempt=1,
+        variants=[_variant(matrix={"py": "3.11"}, run_id=200, run_attempt=1)],
+    )
+    new = _leaf(
+        run_id=100,
+        run_attempt=5,
+        variants=[_variant(matrix={"py": "3.11"}, run_id=100, run_attempt=5)],
+    )
+    merged = merge_matrix_test_leaf(existing, new)
+    assert merged.run_id == 200
+    assert merged.run_attempt == 1
+    assert merged.variants is not None
+    assert merged.variants[0].run_id == 200
+
+
+def test_matrix_merge_leaf_header_advances_to_newer_run() -> None:
+    existing = _leaf(
+        run_id=100,
+        run_attempt=2,
+        variants=[_variant(matrix={"py": "3.11"}, run_id=100, run_attempt=2)],
+    )
+    new = _leaf(
+        run_id=200,
+        run_attempt=1,
+        variants=[_variant(matrix={"py": "3.11"}, run_id=200, run_attempt=1)],
+    )
+    merged = merge_matrix_test_leaf(existing, new)
+    # newer run_id wins outright -- its (lower) attempt comes with it.
+    assert merged.run_id == 200
+    assert merged.run_attempt == 1
+
+
+def test_matrix_merge_leaf_header_takes_higher_attempt_within_run() -> None:
+    existing = _leaf(
+        run_id=100,
+        run_attempt=2,
+        variants=[_variant(matrix={"py": "3.11"}, run_id=100, run_attempt=2)],
+    )
+    new = _leaf(
+        run_id=100,
+        run_attempt=1,
+        variants=[_variant(matrix={"py": "3.11"}, run_id=100, run_attempt=1)],
+    )
+    merged = merge_matrix_test_leaf(existing, new)
+    assert merged.run_id == 100
+    assert merged.run_attempt == 2
+
+
 # --- upsert_leaf: build phase -----------------------------------------------
 
 
@@ -592,14 +650,14 @@ def test_upsert_test_guard_is_per_arch() -> None:
 
 def test_upsert_test_replaces_whole_leaf_with_variants_atomically() -> None:
     # A pytorch.test arch leaf is ONE workflow run; its matrix cells are jobs in
-    # that run, carried on the leaf's `variants`. Unlike the build phase, test
-    # leaves are NOT merged cell-by-cell: a higher run_attempt replaces the
-    # entire leaf (and its full variant list) atomically. This is the
-    # rerun-of-failed-jobs case -- GitHub re-stamps every job with the new
-    # attempt, so the producer hands us a complete attempt-2 snapshot and the
-    # old attempt-1 variants are dropped wholesale, not preserved per-cell.
-    # Re-running failed jobs keeps the SAME run_id and only bumps run_attempt;
-    # the variants repeat the leaf's run_id.
+    # that run, carried on the leaf's `variants`. Like the build phase, test
+    # leaves with variants are merged cell-by-cell (see `merge_matrix_test_leaf`)
+    # rather than replaced wholesale. This is the rerun-of-failed-jobs case --
+    # GitHub re-stamps every job with the new attempt, so every cell's higher
+    # run_attempt wins the per-cell guard and the net effect looks atomic: all
+    # attempt-1 values end up superseded. Re-running failed jobs keeps the
+    # SAME run_id and only bumps run_attempt; the variants repeat the leaf's
+    # run_id.
     run_id = 12345900
     doc = StatusDocument()
     doc.upsert_leaf(
@@ -662,9 +720,10 @@ def test_upsert_test_replaces_whole_leaf_with_variants_atomically() -> None:
     assert all(v.status is Status.success for v in leaf.variants)
 
 
-def test_upsert_test_with_variants_rejects_lower_attempt_wholesale() -> None:
-    # The leaf-level guard protects the whole test leaf, variants included: a
-    # stale lower-attempt snapshot cannot clobber a newer one even cell-by-cell.
+def test_upsert_test_with_variants_rejects_lower_attempt_per_cell() -> None:
+    # The variant path merges per-cell (like build) and always reports the
+    # write as accepted, but the per-cell guard still protects the actual
+    # data: a stale lower-attempt snapshot cannot clobber a newer cell.
     doc = StatusDocument()
     doc.upsert_leaf(
         "linux",
@@ -677,7 +736,7 @@ def test_upsert_test_with_variants_rejects_lower_attempt_wholesale() -> None:
             variants=[_variant(matrix={"py": "3.11"}, run_attempt=2)],
         ),
     )
-    assert not doc.upsert_leaf(
+    assert doc.upsert_leaf(
         "linux",
         "gfx942",
         "pytorch",
@@ -689,11 +748,99 @@ def test_upsert_test_with_variants_rejects_lower_attempt_wholesale() -> None:
         ),
     )
     leaf = doc.pipelines.pytorch.test["linux"]["gfx942"]
-    assert leaf.run_attempt == 2
     assert leaf.variants is not None
+    assert leaf.variants[0].run_attempt == 2
     assert leaf.variants[0].status is Status.success
 
 
+def test_upsert_test_with_variants_always_returns_true() -> None:
+    # Mirrors test_upsert_build_with_variants_always_returns_true: the variant
+    # path merges per-cell and bypasses the leaf-level guard.
+    doc = StatusDocument()
+    assert doc.upsert_leaf(
+        "linux",
+        "gfx942",
+        "pytorch",
+        "test",
+        _leaf(variants=[_variant(matrix={"py": "3"})]),
+    )
+
+
+def test_upsert_test_merge_does_not_regress_completed_cell() -> None:
+    # This is the push-race scenario the fix targets: two concurrent
+    # `receive_therock_data.yml` runs fetch fresh job-list snapshots of the
+    # SAME shared entry run at different wall-clock times. Snapshot A sees
+    # py3.11 done but py3.12 still running; snapshot B (fetched slightly
+    # earlier, but whose git push lands second) sees py3.11 still running but
+    # py3.12 done. Regardless of push order, merging cell-by-cell must end up
+    # with BOTH cells advanced -- never one snapshot regressing the other's
+    # progress.
+    doc = StatusDocument()
+    run_id = 555
+    doc.upsert_leaf(
+        "linux",
+        "gfx942",
+        "jax",
+        "test",
+        _leaf(
+            run_id=run_id,
+            run_attempt=1,
+            status=Status.in_progress,
+            completed_at=None,
+            variants=[
+                _variant(
+                    matrix={"py": "3.11"},
+                    run_id=run_id,
+                    status=Status.success,
+                    completed_at="2026-04-08T01:10:00Z",
+                ),
+                _variant(
+                    matrix={"py": "3.12"},
+                    run_id=run_id,
+                    status=Status.in_progress,
+                    completed_at=None,
+                ),
+            ],
+        ),
+    )
+    # A stale snapshot lands next: py3.11 looks in_progress again (fetched
+    # before it finished) but py3.12 has since completed.
+    doc.upsert_leaf(
+        "linux",
+        "gfx942",
+        "jax",
+        "test",
+        _leaf(
+            run_id=run_id,
+            run_attempt=1,
+            status=Status.in_progress,
+            completed_at=None,
+            variants=[
+                _variant(
+                    matrix={"py": "3.11"},
+                    run_id=run_id,
+                    status=Status.in_progress,
+                    completed_at=None,
+                ),
+                _variant(
+                    matrix={"py": "3.12"},
+                    run_id=run_id,
+                    status=Status.success,
+                    completed_at="2026-04-08T01:20:00Z",
+                ),
+            ],
+        ),
+    )
+    leaf = doc.pipelines.jax.test["linux"]["gfx942"]
+    assert leaf.variants is not None
+    by_key = {v.key(): v for v in leaf.variants}
+    assert by_key[(("py", "3.11"),)].status is Status.success
+    assert by_key[(("py", "3.12"),)].status is Status.success
+    # Both cells terminal -> the leaf-level rollup is terminal too.
+    assert leaf.status is Status.success
+    assert leaf.completed_at == "2026-04-08T01:20:00Z"
+
+
 def test_upsert_test_newer_run_id_supersedes_via_upsert() -> None:
     # A fresh re-dispatch mints a NEW workflow run (larger run_id). The guard
     # compares run_id FIRST, so the newer run wins regardless of attempt -- here

scripts/receive_therock/therock_status_document.py

@@ -368,6 +368,65 @@ def _merge_variant_leaf(existing: "RunLeaf | None", new: "RunLeaf") -> "RunLeaf"
     )
 
 
+def merge_matrix_test_leaf(existing: "RunLeaf | None", new: "RunLeaf") -> "RunLeaf":
+    """Merge fan-out test variants by matrix cell instead of replacing the
+    whole leaf (mirrors `_merge_variant_leaf`, but also carries forward
+    `run_id`/`run_attempt`/timestamps so later comparisons -- e.g.
+    `_refresh_same_run_fanout_tests` in therock_update_status_json.py, its
+    other caller -- can still identify the owning run).
+
+    Every completion notification for a shared-entry-run matrix (pytorch/jax
+    py x ref fan-out sharing one `GITHUB_RUN_ID`) re-derives *all* cells from
+    a freshly re-fetched job-list snapshot, and those snapshots race each
+    other under heavy concurrency with no correlation between "wins the push"
+    and "is the freshest/most complete". Wholesale replacement
+    (`arch_map[arch] = new`) lets an earlier, less-complete snapshot silently
+    regress a later, more-complete one whenever it happens to win the race.
+    Merging cell-by-cell through `Variant.should_replace` makes every cell
+    advance monotonically regardless of push-race ordering.
+    """
+    variants: list[Variant] = []
+    positions: dict[tuple[tuple[str, str], ...], int] = {}
+
+    for variant in existing.variants if existing and existing.variants else []:
+        positions[variant.key()] = len(variants)
+        variants.append(variant)
+
+    for variant in new.variants or []:
+        key = variant.key()
+        pos = positions.get(key)
+        if pos is None:
+            positions[key] = len(variants)
+            variants.append(variant)
+        elif variants[pos].should_replace(variant):
+            variants[pos] = variant
+
+    status = Variant.rollup_status(variants, new.status)
+    completed_at: str | None = None
+    if status.is_terminal:
+        ends = [v.completed_at for v in variants if v.completed_at]
+        completed_at = max(ends) if ends else new.completed_at
+
+    starts = [v.started_at for v in variants if v.started_at]
+    started_at = min(starts) if starts else None
+    if not started_at:
+        started_at = (existing.started_at if existing else None) or new.started_at
+
+    if existing is None or existing.should_replace(new):
+        run_id, run_attempt = new.run_id, new.run_attempt
+    else:
+        run_id, run_attempt = existing.run_id, existing.run_attempt
+
+    return RunLeaf(
+        status=status,
+        run_id=run_id,
+        run_attempt=run_attempt,
+        started_at=started_at or None,
+        completed_at=completed_at,
+        variants=variants,
+    )
+
+
 # --- Summary rollup ----------------------------------------------------------
 # `therock_summary.rebuild_summary` derives these from the pipeline detail tree
 # on every update; they are the read-optimized projection consumers render.
@@ -581,6 +640,9 @@ def upsert_leaf(
         )
         arch_map = phase_map.setdefault(platform, {})
         existing = arch_map.get(arch)
+        if leaf.variants:
+            arch_map[arch] = merge_matrix_test_leaf(existing, leaf)
+            return True
         if existing is not None and not existing.should_replace(leaf):
             return False
         arch_map[arch] = leaf

scripts/receive_therock/therock_update_status_json.py

@@ -57,6 +57,7 @@
     Status,
     StatusDocument,
     Variant,
+    merge_matrix_test_leaf,
     rollup_statuses,
 )
 from therock_summary import freeze_requested_architectures, rebuild_summary
@@ -583,27 +584,32 @@ def _refresh_same_run_fanout_tests(
     this function those same-run test leaves would go stale once the build
     itself is done.
 
-    `leaf.status` is this run's own top-level GitHub conclusion. It is not
-    necessarily the worst-of its `variants` (e.g. a matrix cell whose nested
-    test job failed/cancelled does not always flip the run's own conclusion,
-    or a cell can be entirely missing from `variants` if its job never
-    started). When the run only ever reports one architecture, fold it into
-    that architecture's rollup rather than only using it as an
-    empty-variants fallback, so a terminal failure/cancellation at the run
-    level cannot be masked by whatever the individual cells happened to
-    report -- mirroring what `_merge_matrix_build_leaf` does for the build
-    leaf itself.
-
-    One build cell can also nest test jobs for several architectures (see
-    `_variants_from_jobs`), so `leaf.variants` may roll every architecture's
-    outcome together. Re-derive each existing test leaf's variants scoped to
-    its *own* architecture instead of broadcasting `leaf.variants` wholesale,
-    so one architecture's result can never leak into another's. The same
-    reasoning rules out folding in the run-level conclusion once several
-    architectures share the run: that conclusion is a whole-run aggregate,
-    so a failure anywhere -- including in an unrelated architecture's own
-    job -- would otherwise get broadcast onto every architecture, right back
-    into the leakage this function exists to prevent.
+    Each matching test leaf is refreshed under three constraints, all guarding
+    against this coarse build-run snapshot corrupting finer-grained state:
+
+    - Per-architecture scope. One build run can nest test jobs for several
+      architectures, so `leaf.variants` may roll every architecture's outcome
+      together. Each existing test leaf's variants are re-derived scoped to its
+      *own* architecture (`_variants_from_jobs(..., arch=arch)`) rather than
+      broadcasting `leaf.variants` wholesale, so one architecture's result can
+      never leak into another's.
+
+    - Per-cell merge. Those re-derived variants come from the build run's job
+      names, so a blind overwrite could clobber genuinely newer/more-complete
+      per-cell results that already landed from the test leaf's own dedicated
+      completion events. `merge_matrix_test_leaf` merges cell-by-cell, keeping
+      the winning `Variant` per cell, instead of replacing `variants` wholesale.
+
+    - Run-level status only when unambiguous. `leaf.status` is this run's own
+      top-level GitHub conclusion, which is not necessarily the worst-of its
+      variants (a nested test cell that failed/cancelled does not always flip
+      the run's conclusion, and a cell can be missing entirely if its job never
+      started). It is folded into the rollup only when the run reports a single
+      architecture -- mirroring what `_merge_variant_leaf` does for the build
+      leaf itself. Once several architectures share the run that conclusion is a
+      whole-run aggregate, so folding it in would broadcast a failure anywhere
+      -- even in an unrelated architecture -- onto every architecture, the very
+      leakage this scoping exists to prevent.
     """
     cls = workflow_run.classification
     axis_key = _VARIANT_AXIS_KEY.get(cls.pipeline_type)
@@ -645,9 +651,10 @@ def _refresh_same_run_fanout_tests(
                 )
                 if not existing.should_replace(candidate):
                     continue
-                existing.status = candidate.status
-                existing.completed_at = candidate.completed_at
-                existing.variants = candidate.variants
+                merged = merge_matrix_test_leaf(existing, candidate)
+                merged.status = candidate.status
+                merged.completed_at = candidate.completed_at
+                arch_map[arch] = merged
                 wrote = True
     return wrote