Where #4583's 1,135 lines go, and the pixel copy a hash could replace

NVIDIA-NeMo/RL #4583. Line counts are added lines from the GitHub diff against the parent branch. Permalinks pin 0bf5789.

Two pages on this PR, one question each. This one: what could be smaller.

Where the lines go

The source is not the problem here: 192 lines for a new media path is fair, and most of it reuses the parent PR's code. The weight is in tests. The functional script (nemotron_omni_clevr_vllm_prefix_recovery_single_controller_1n2g.sh) cannot pass until a Gym change lands, so no CI job runs it. Its new completion check (assert_restored_prefixes_completed) repeats the parent's (_assert_prefix_boundary). The 436 unit-test lines are a fair size, but they test the save side. About 10 of the new reject paths have no test.

🔵 review comment — clevr…prefix_recovery…1n2g.sh:187 review comment — vllm_worker_async.py:1184 review comment — tq_token_sink.py:927

The pixel copy is only used to compare — proposed, not in the PR

At the first cut, the call's pixels are written to TQ (:1179-1182). On restore they are read back, compared byte for byte with the freshly processed image, and returned in place of the fresh copy (:1156-1169). By then the two are equal, so returning either one gives the same result. The final row writes the pixels again anyway (:1361-1365), and training reads media only from final rows (rollout_reassembler.py:512). A hash of the pixels does the same check without the copy.

Where the fix lands: RL only. vllm_worker_async.py (write the hash, compare it), captured_media.py (one pure hash helper, next to CapturedMedia), and deletions in tq_token_sink.py and generation_cut_capture.py. No Gym change: the hash rides in the same digest-covered extras dict that already carries media_spans.

@@ captured_media.py · new function, next to CapturedMedia (L137)
+def media_tensor_sha256(tensors: dict[str, torch.Tensor] | None) -> str:
+    # Covers name, dtype and shape too, so it checks as much as torch.equal did.
+    h = hashlib.sha256()
+    for name in sorted(tensors or {}):
+        t = tensors[name].detach().cpu().contiguous()
+        h.update(f"{name}|{t.dtype}|{tuple(t.shape)}|".encode())
+        h.update(t.view(torch.uint8).numpy().tobytes())
+    return h.hexdigest()

@@ vllm_worker_async.py L1179-L1182 · _generation_cut_record_payload
         return (
-            {MEDIA_SPANS_FIELD: [item.to_dict() for item in state.media.items]},
-            state.media.tensors,
+            {
+                MEDIA_SPANS_FIELD: [item.to_dict() for item in state.media.items],
+                "media_tensor_sha256": media_tensor_sha256(state.media.tensors),
+            },
+            None,  # no pixel copy on the prefix row
         )

@@ vllm_worker_async.py L1156-L1169 · _capture_request_media
-        captured_tensors = captured.tensors or {}
-        restored_tensors = restored.tensors or {}
-        if set(captured_tensors) != set(restored_tensors) or any(
-            ...
-        ):
+        if media_tensor_sha256(captured.tensors) != restored_sha256:
             raise MediaCaptureRejected(
                 "Restored generation media tensors differ from reprocessed input",
                 code="retained_media_changed",
             )
-        # Keep the durable TQ-owned tensors for the eventual canonical row.
-        return restored
+        # Equal by hash, so the fresh tensors are the ones the cut saw.
+        return captured
  … plus deletions: the media half of fetch_generation_cut (tq_token_sink.py L914-941),
    FetchedStagedCall.media (L401-403), the tensor rebuild (vllm_worker_async.py L1213-1221),
    and the first-row attachments (generation_cut_capture.py L722-726).

A sketch, not run. restored_sha256 is read from the first row's extras in place of today's pixel read.

🔵 review comment — vllm_worker_async.py:1169 review comment — generation_cut_capture.py:311 review comment — vllm_worker_async.py:1199