refactor(data-plane): make kv_batch_get(select_fields) required

Silent over-fetch was possible when callers omitted select_fields: the
noop adapter fetched every registered field via set intersection; the
TQ adapter forwarded None to the backend. Bulk schemas are wide and
fetching everything is the most expensive shape the wire can take.

select_fields is now a required list[str] on DataPlaneClient.kv_batch_get
and all concrete implementations. Callers must name what they read;
fetch-all is still possible by passing list(meta.fields) explicitly.

Also: worker_mixin internal call sites use list(meta.fields) directly
(fail-loud TypeError if meta.fields is None, rather than silently
producing an empty TensorDict).

Per yuki-97 PR review (#6).

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

Author: ZhiyuLi-Nvidia

nemo_rl/data_plane/adapters/noop.py

@@ -199,16 +199,12 @@ def kv_batch_get(
         self,
         keys: list[str],
         partition_id: str,
-        select_fields: list[str] | None = None,
+        select_fields: list[str],
     ) -> TensorDict:
         rec = self._partitions[partition_id]
         if not keys:
             return TensorDict({}, batch_size=(0,))
 
-        if select_fields is None:
-            available = set.intersection(*(set(rec.rows[k].keys()) for k in keys))
-            select_fields = sorted(available)
-
         out: dict[str, list[torch.Tensor]] = {f: [] for f in select_fields}
         for key in keys:
             row = rec.rows[key]

nemo_rl/data_plane/adapters/transfer_queue.py

@@ -587,14 +587,14 @@ def kv_batch_get(
         self,
         keys: list[str],
         partition_id: str,
-        select_fields: list[str] | None = None,
+        select_fields: list[str],
     ) -> TensorDict:
         if not keys:
             return TensorDict({}, batch_size=(0,))
         td = self._tq.kv_batch_get(
             keys=list(keys),
             partition_id=partition_id,
-            select_fields=list(select_fields) if select_fields else None,
+            select_fields=select_fields,
         )
         if self._promote_1d:
             td = _from_wire(td)

nemo_rl/data_plane/interfaces.py

@@ -312,17 +312,22 @@ def kv_batch_get(
         self,
         keys: list[str],
         partition_id: str,
-        select_fields: list[str] | None = None,
+        select_fields: list[str],
     ) -> TensorDict:
         """Direct fetch by uids.
 
         Used by per-DP-rank slice fetches. Does NOT advance any per-task
         consumption cursor — that only happens via :meth:`claim_meta`.
 
+        ``select_fields`` is required (no implicit "fetch every field"
+        fallback): bulk schemas are wide and silent over-fetch is the
+        most expensive shape the wire can take. Callers must name what
+        they read.
+
         Args:
             keys: Uids to fetch.
             partition_id: Partition the keys live in.
-            select_fields: Subset of fields; ``None`` fetches every registered field.
+            select_fields: Subset of fields to fetch.
 
         Returns:
             ``TensorDict`` keyed by field name, batched along ``keys``.

nemo_rl/data_plane/observability.py

@@ -308,7 +308,7 @@ def kv_batch_put(self, keys, partition_id, fields=None, tags=None):
         self._record_put(partition_id, keys_list, n_bytes)
         return out
 
-    def kv_batch_get(self, keys, partition_id, select_fields=None):
+    def kv_batch_get(self, keys, partition_id, select_fields):
         return self._run(
             "get",
             partition_id,

nemo_rl/data_plane/worker_mixin.py

@@ -221,7 +221,7 @@ def _fetch(
                 td = self._require_dp_client().kv_batch_get(
                     keys=meta.keys,
                     partition_id=meta.partition_id,
-                    select_fields=list(meta.fields) if meta.fields else None,
+                    select_fields=list(meta.fields),
                 )
                 data = materialize(
                     td,
@@ -246,7 +246,7 @@ def _fetch(
         td = self._require_dp_client().kv_batch_get(
             keys=meta.keys,
             partition_id=meta.partition_id,
-            select_fields=list(meta.fields) if meta.fields else None,
+            select_fields=list(meta.fields),
         )
         data = materialize(
             td,