How the TransferQueue controller patch makes clears and lookups cheaper, and what it changes

NVIDIA-NeMo/RL #4769 · NeMo-RL links pin the PR head fc9191b9; TransferQueue links pin the pinned TQ commit 05d4cc33 (v0.1.11).

On each clear, stock FieldMeta.remove_samples builds a set of every remaining sample's shape, just to check whether they are now all equal. That costs one read per live sample, on every clear. The PR's remove_samples counts how many samples were removed since the last scan, and scans only when that count reaches RESCAN_REMOVED_FRACTION (1/8) of what is left. A scan still costs the same; it just runs far less often.

Shapes read per clear: 90 live samples, 2 cleared each time

Clear 5 is the first scan: 10 samples removed since the start, and 80 ÷ 8 = 10. The scan count goes back to 0 and the next scan waits for another eighth. At 300k live samples the author measured 50.8 ms per clear for stock, per nested field; with the patch, a scan still costs that, but it happens about once every 2,300 clears of 16 rows.

What changes, and what stays the same

Where the change lands: RL only. A new module tq_kv_meta_patch.py, called once right before tq.init in _init_tq. No TQ change yet: the same edits are not in an upstream TQ PR, so the copy stays until one lands.

How it gets into the controller: TQ starts its controller from a fixed class, so configure_controller_patch() swaps in a subclass. The subclass carries the new kv_retrieve_meta, and its __init__ installs the new remove_samples inside the controller process only. A version check raises at startup if TQ is not 0.1.11. The lookup change is a single line:

@@ kv_retrieve_meta · stock controller.py L1681-L1688 → tq_kv_meta_patch.py L98-L106
     verified_global_indexes = [idx for idx in global_indexes if idx is not None]

-    # must fetch fields that the requested samples all have
-    col_mask = partition.production_status[verified_global_indexes, :].sum(dim=0).reshape(-1) == len(
-        verified_global_indexes
-    )
-    data_fields = []
-    for field_name, col_idx in partition.field_name_mapping.items():
-        if col_idx < len(col_mask) and col_mask[col_idx]:   # col_mask[i] is a torch read, once per field
-            data_fields.append(field_name)
+    # Fetch the fields every requested sample has. One .tolist(): per-field tensor reads are slow.
+    col_mask = (
+        partition.production_status[verified_global_indexes, :].all(dim=0).tolist()
+    )
+    data_fields = [
+        field_name
+        for field_name, col_idx in partition.field_name_mapping.items()
+        if col_idx < len(col_mask) and col_mask[col_idx]
+    ]

Indentation follows the PR's module-level function; the stock method sits one level deeper inside the class.

production_status only ever holds 0 or 1 (set to 1, reset to 0), so "all ones" and "sum equals row count" pick the same fields, even for zero rows.

The one visible difference is in the clear. Take the PR's own test field: 40 samples, 39 of shape (256,) and one of (128,).

In between, the field stays nested even though all shapes match. Readers get the same tensors either way: TQ storage decides stacked or nested from the stored tensors, not this flag (base.py:661), BatchMeta.get_shapes accepts both forms, and NeMo-RL's _from_wire turns equal-shape rows dense by looking at the rows themselves.