chore: bump TQ to v0.19

nemo_rl/data_plane/adapters/transfer_queue.py

@@ -27,8 +27,9 @@
 import socket
 import subprocess
 import time
+import warnings
 from importlib import resources
-from typing import Any
+from typing import Any, cast
 
 import torch
 import transfer_queue as tq
@@ -39,6 +40,7 @@
     DataPlaneConfig,
     KVBatchMeta,
 )
+from nemo_rl.data_plane.schema import PROMOTE_1D_FIELDS
 
 # ──────────────────────────────────────────────────────────────────────────
 # Backend init — lifted from rl-arena/arena/backends.py.
@@ -171,27 +173,28 @@ def patched(*args, **kwargs):
         cls.options = patched  # type: ignore[method-assign]
         return True
 
-    patched_any = False
+    unpatched_classes: list[str] = []
     try:
-        from transfer_queue.storage.simple_backend import SimpleStorageUnit
+        from transfer_queue.storage.simple_storage import SimpleStorageUnit
 
-        patched_any |= _install(SimpleStorageUnit)
+        if not _install(SimpleStorageUnit):
+            unpatched_classes.append("SimpleStorageUnit")
     except ImportError:
-        pass
+        unpatched_classes.append("SimpleStorageUnit")
     try:
         from transfer_queue.controller import TransferQueueController
 
-        patched_any |= _install(TransferQueueController)
+        if not _install(TransferQueueController):
+            unpatched_classes.append("TransferQueueController")
     except ImportError:
-        pass
+        unpatched_classes.append("TransferQueueController")
 
-    if not patched_any:
+    if unpatched_classes:
         # Soft-fail: TQ may have moved its actor classes. The driver will
         # still work; multi-node TQ may need the per-node `uv sync` workaround.
-        import warnings
-
         warnings.warn(
-            "Could not patch TQ actor classes for runtime_env injection. "
+            "Could not patch every TQ actor class for runtime_env injection: "
+            f"unpatched={unpatched_classes}. "
             "Multi-node TQ may fail with ModuleNotFoundError: 'transfer_queue' "
             "on worker nodes. Workaround: run `uv sync` inside each node's "
             "container before the driver runs.",
@@ -322,19 +325,23 @@ def _assert_no_key_loss(src_dict: dict, new_td: TensorDict, fn: str) -> None:
 
 
 def _promote_1d_leaves(td: TensorDict) -> TensorDict:
-    """Unsqueeze 1D tensor leaves to ``(N, 1)`` — mooncake_cpu KV-path workaround.
+    """Promote declared scalar leaves to ``(N, 1)`` for Mooncake.
 
-    Works around TQ's ``KVStorageManager`` 1D schema/data mismatch;
-    :func:`_from_wire` squeezes the trailing 1 back on read. Symmetric
-    with `_from_wire` — callers gate on ``self._promote_1d``.
-    ``NonTensorStack`` / ``NonTensorData`` leaves pass through.
+    The authoritative field list lives in
+    :data:`nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS`. Declared fields must
+    arrive as dense ``(N,)`` tensors. Any other dense 1D tensor is rejected so
+    it cannot silently encounter TQ v0.1.9's schema/data mismatch.
+    ``NonTensorStack`` and ``NonTensorData`` leaves pass through.
 
     Args:
-        td: ``TensorDict`` whose 1D tensor leaves should be promoted.
+        td: TensorDict to validate and encode for the Mooncake wire format.
 
     Returns:
-        ``TensorDict`` with 1D tensor leaves unsqueezed to ``(N, 1)``;
-        all other leaves pass through unchanged.
+        TensorDict with declared scalar leaves promoted to ``(N, 1)``.
+
+    Raises:
+        ValueError: If a declared field is not a dense 1D tensor, or an
+            undeclared field is a dense 1D tensor.
     """
     # td.keys() (top-level) includes NonTensorData / NonTensorStack leaves.
     # keys(include_nested=True, leaves_only=True) enumerates tensor leaves
@@ -343,9 +350,23 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict:
     changed = False
     for k in td.keys():
         v = td.get(k)
-        if isinstance(v, torch.Tensor) and not v.is_nested and v.dim() == 1:
+        field_name = str(k)
+        if field_name in PROMOTE_1D_FIELDS:
+            if not isinstance(v, torch.Tensor) or v.is_nested or v.dim() != 1:
+                shape = tuple(v.shape) if isinstance(v, torch.Tensor) else None
+                raise ValueError(
+                    f"Mooncake scalar field {field_name!r} must be a dense "
+                    f"1D tensor with shape (N,), got {type(v).__name__} "
+                    f"with shape {shape}."
+                )
             new_dict[str(k)] = v.unsqueeze(-1).contiguous()
             changed = True
+        elif isinstance(v, torch.Tensor) and not v.is_nested and v.dim() == 1:
+            raise ValueError(
+                f"Mooncake field {field_name!r} is a dense 1D tensor but is "
+                "not declared in data_plane.schema.PROMOTE_1D_FIELDS. Add "
+                "the field to the schema if it is a per-sample scalar."
+            )
         else:
             new_dict[str(k)] = v
     if not changed:
@@ -356,23 +377,45 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict:
 
 
 def _from_wire(td: TensorDict) -> TensorDict:
-    """Inverse of `_promote_1d_leaves`: squeeze trailing 1 back to (N,)."""
+    """Normalize TQ reads and invert :func:`_promote_1d_leaves` when needed.
+
+    Both TQ v0.1.9 storage managers reconstruct every non-scalar field as a
+    nested tensor, including fields whose rows all have the same shape.
+    Densify those uniform nested tensors first so regular batched inputs retain
+    their dense representation. Truly ragged fields remain nested. Finally,
+    squeeze only singleton dimensions declared in
+    :data:`nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS`.
+    """
     # Same top-level iteration as `_promote_1d_leaves`: NonTensorData /
     # NonTensorStack leaves are only visible via td.keys(), not leaves_only.
     new_dict: dict[str, Any] = {}
     changed = False
     for k in td.keys():
         v = td.get(k)
-        if (
-            isinstance(v, torch.Tensor)
-            and not v.is_nested
-            and v.dim() >= 2
-            and v.shape[-1] == 1
-        ):
-            new_dict[str(k)] = v.squeeze(-1).contiguous()
-            changed = True
+        field_name = str(k)
+        if isinstance(v, torch.Tensor) and v.is_nested:
+            rows = list(v.unbind())
+            if rows and all(row.shape == rows[0].shape for row in rows[1:]):
+                v = torch.stack(rows)
+                changed = True
+        if field_name in PROMOTE_1D_FIELDS:
+            if not isinstance(v, torch.Tensor) or v.is_nested:
+                raise ValueError(
+                    f"Mooncake scalar field {field_name!r} could not be "
+                    "restored as a dense tensor."
+                )
+            if v.dim() == 1:
+                new_dict[field_name] = v
+            elif v.dim() == 2 and v.shape[-1] == 1:
+                new_dict[field_name] = v.squeeze(-1).contiguous()
+                changed = True
+            else:
+                raise ValueError(
+                    f"Mooncake scalar field {field_name!r} must decode as "
+                    f"(N,) or (N, 1), got shape {tuple(v.shape)}."
+                )
         else:
-            new_dict[str(k)] = v
+            new_dict[field_name] = v
     if not changed:
         return td
     new_td = TensorDict(new_dict, batch_size=td.batch_size)
@@ -583,9 +626,9 @@ def put_samples(
             return KVBatchMeta(
                 partition_id=partition_id, task_name=None, sample_ids=[], fields=None
             )
-        if tags is None:
-            tags = [{} for _ in sample_ids]
-
+        user_tags = (
+            [{} for _ in sample_ids] if tags is None else [dict(tag) for tag in tags]
+        )
         wire_fields: TensorDict | None = None
         field_names: list[str] | None = None
         if fields is not None:
@@ -594,25 +637,29 @@ def put_samples(
             # TDs. TQ's encoder forces ``.contiguous()`` per tensor leaf
             # itself, so the call here was redundant for tensors and
             # destructive for non-tensors.
-            wire_fields = fields.detach()  # type: ignore[bad-assignment,missing-argument]
+            detached_fields = cast(
+                TensorDict,
+                fields.detach(),  # type: ignore[missing-argument]
+            )
             if self._promote_1d:
-                wire_fields = _promote_1d_leaves(wire_fields)  # type: ignore[bad-argument-type]
-            field_names = list(wire_fields.keys())
+                detached_fields = _promote_1d_leaves(detached_fields)
+            wire_fields = detached_fields
+            field_names = [str(key) for key in detached_fields.keys()]
 
         # TQ's wire vocabulary is `keys=` — translation point.
         tq.kv_batch_put(
             keys=list(sample_ids),
             partition_id=partition_id,
             fields=wire_fields,
-            tags=tags,
+            tags=user_tags,
         )
 
         return KVBatchMeta(
             partition_id=partition_id,
             task_name=None,
             sample_ids=list(sample_ids),
             fields=field_names,
-            tags=[dict(t) for t in tags] if tags else None,
+            tags=user_tags if user_tags else None,
         )
 
     def get_samples(
@@ -623,15 +670,12 @@ def get_samples(
     ) -> TensorDict:
         if not sample_ids:
             return TensorDict({}, batch_size=(0,))
-        # TQ's wire vocabulary is `keys=` — translation point.
         td = tq.kv_batch_get(
             keys=list(sample_ids),
             partition_id=partition_id,
             select_fields=select_fields,
         )
-        if self._promote_1d:
-            td = _from_wire(td)
-        return td
+        return _from_wire(td)
 
     def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None:
         cleared_via_none = sample_ids is None

nemo_rl/data_plane/schema.py

@@ -63,6 +63,25 @@
 
 ROUTED_EXPERTS_FIELD = "routed_experts"
 
+# Per-sample 1D scalar fields. The TQ adapter promotes these to ``(N, 1)``
+# on write to work around TQ v0.1.9's KVStorageManager schema/data mismatch on
+# the Mooncake backend, and squeezes them back to ``(N,)`` on read. This is the
+# authoritative user-level schema; no per-row shape metadata is carried.
+#
+# Fields listed here must be dense ``(N,)`` tensors when written through the
+# Mooncake adapter. Dense 1D fields not listed here are rejected on that path so
+# a new field cannot silently reintroduce the upstream shape mismatch.
+#
+# Delete this set and the corresponding adapter transforms when upstream TQ
+# fixes 1D field schema extraction.
+PROMOTE_1D_FIELDS: frozenset[str] = frozenset(
+    {
+        INPUT_LENGTHS,
+        "total_reward",
+        SAMPLE_MASK,
+    }
+)
+
 
 def fields_with_optional_routed_experts(
     fields: Sequence[str],

pyproject.toml

@@ -84,11 +84,10 @@ dependencies = [
   # automatically include them. Removes the need for a `[data-plane]`
   # extra and the corresponding plumbing in the per-worker venv builder.
   "tensordict",
-  # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore
-  # refactor: `clear` switched from unanchored `remove_by_regex` to
-  # exact-key `batch_remove`, which fixes a collateral-key-deletion bug
-  # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released.
-  "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39",
+  # TransferQueue v0.1.9 adds full-system checkpoint save/load APIs and
+  # retains the exact-key MooncakeStore clear behavior required by DAPO.
+  # Pin the immutable release commit rather than the mutable version tag.
+  "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@c51614308b68c8d7a87c9b3ef62d59e14c69bde2",
   # Backs data_plane.backend="mooncake_cpu". Default backend is "simple"
   # (in-process), but the mooncake_cpu path needs the `mooncake_master`
   # binary that ships in this wheel at <site-packages>/mooncake/. Bundled
@@ -391,10 +390,9 @@ override-dependencies = [
   "pytest>=9.0.3",
   "langchain>=0.3.28",      # Address CVE-2025-65106
   "langchain-core>=0.3.80", # Address CVE-2025-65106
-  # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs
-  # numpy>=2.1.0 via onnx → ml-dtypes. Override globally so the data-plane
-  # extra composes with mcore/automodel without version-mirroring TQ's
-  # requirements.txt. Forward-compatible across TQ minor bumps.
+  # Forces numpy past tensorrt-llm's `numpy>=2.0.0,<2.4` cap (resolves to
+  # 2.5.x). The original driver — TransferQueue pinning `numpy<2.0.0` — is
+  # gone as of TQ v0.1.9; drop this override once the trtllm cap lifts.
   "numpy>=2.1.0",
   # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame).
   # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it

tests/unit/data_plane/README.md

@@ -30,10 +30,14 @@ Generated audit of every test function under `tests/unit/data_plane/` with a one
 - `test_materialize_default_pad_value_is_zero` — No `pad_value_dict` → pad with 0.
 - `test_response_from_nested_extracts_response_slice` — Worker write-back: jagged (prompt+response) → response only.
 
-## `test_codec_mooncake.py` (4 tests)
-
-- `test_promote_1d_leaves_unsqueezes_1d` — `_promote_1d_leaves` turns 1D `(N,)` leaves into `(N, 1)` for mooncake wire.
-- `test_promote_1d_roundtrip_via_from_wire` — `_promote_1d_leaves` + `_from_wire` restores original `(N,)` shape and values.
+## `test_codec_mooncake.py`
+
+- `test_promote_1d_leaves_unsqueezes_1d` — `_promote_1d_leaves` turns schema-declared scalar fields from `(N,)` into `(N, 1)` for the Mooncake wire.
+- `test_promote_1d_roundtrip_via_from_wire` — `_promote_1d_leaves` + `_from_wire` restores the schema-declared field's original `(N,)` shape and values.
+- `test_from_wire_rejects_invalid_declared_field_shape` — Corrupt or incompatible scalar wire shapes fail at the data-plane boundary.
+- `test_promote_1d_leaves_rejects_undeclared_1d_field` — Unknown dense 1D Mooncake fields fail loudly instead of silently hitting TQ's shape mismatch.
+- `test_put_samples_uses_schema_without_private_shape_tags` — Promotion does not add per-row adapter metadata to user tags.
+- `test_get_samples_uses_static_shape_schema` — Reads restore scalar fields by the shared schema while preserving genuine `(N, 1)` columns.
 - `test_pack_per_token_field_truncates_sp_padding` — pack_per_token_field slices each row to its own length, dropping SP padding.
 - `test_pack_per_token_field_exact_fit_matches_to_nested_by_length` — At exact fit, `pack_per_token_field` matches `to_nested_by_length`.