diff --git a/nemo_rl/data_plane/adapters/tq_controller_patch.py b/nemo_rl/data_plane/adapters/tq_controller_patch.py new file mode 100644 index 00000000000..14103c8187c --- /dev/null +++ b/nemo_rl/data_plane/adapters/tq_controller_patch.py @@ -0,0 +1,191 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Cheaper ``kv_retrieve_meta`` and ``FieldMeta.remove_samples`` for the TransferQueue controller. + +The stock ``kv_retrieve_meta`` reads a torch tensor once per field, about a quarter of a put+get +cycle on the single-threaded controller. ``kv_retrieve_meta`` below is the stock +method with one ``.all(dim=0).tolist()`` instead; results are identical. + +The stock ``FieldMeta.remove_samples`` rescans every remaining shape of a nested field on each +clear, so ``CLEAR_META`` grows with the number of live samples. ``remove_samples`` below rescans +only after ``RESCAN_REMOVED_FRACTION`` of them is gone. + +The controller is a Ray actor, so both are applied in a subclass started by ``tq.init``. + +Delete this module once the pinned TransferQueue includes both changes, proposed upstream on +https://github.com/ZhiyuLi-Nvidia/TransferQueue/tree/perf/throttle-nested-shape-rescan. +""" + +from typing import Any + +import ray +import transfer_queue +import transfer_queue.interface as interface +from transfer_queue.controller import FieldMeta +from transfer_queue.metadata import BatchMeta +from transfer_queue.utils.logging_utils import get_logger + +# TQ's logger, so TQ_LOGGING_LEVEL still controls the copied log lines. +logger = get_logger(__name__) + +# The TransferQueue version the overrides below were copied from (see the pin in pyproject.toml). +STOCK_TQ_VERSION = "0.1.11" + +# A nested field rescans its shapes once this fraction of its live samples was removed. +RESCAN_REMOVED_FRACTION = 1 / 8 + + +def kv_retrieve_meta( + self, keys: list[str], partition_id: str, create: bool = False +) -> BatchMeta: + """Stock ``TransferQueueController.kv_retrieve_meta`` with one field-mask read.""" + logger.debug( + f"[{self.controller_id}] Retrieve keys {keys} in partition {partition_id}" + ) + + partition = self._get_partition(partition_id) + if partition is None: + if not create: + logger.warning( + f"[{self.controller_id}]: Partition {partition_id} not found. Returning empty BatchMeta." + ) + return BatchMeta.empty() + + self.create_partition(partition_id) + partition = self._get_partition(partition_id) + + assert partition is not None + global_indexes = partition.kv_retrieve_indexes(keys) + + none_indexes = [idx for idx, value in enumerate(global_indexes) if value is None] + if len(none_indexes) > 0: + if not create: + logger.warning( + f"Keys {[keys[i] for i in none_indexes]} were not found in partition {partition_id}. " + f"They will be excluded from the retrieved BatchMeta." + ) + else: + batch_global_indexes = partition.activate_pre_allocated_indexes( + len(none_indexes) + ) + + if len(batch_global_indexes) < len(none_indexes): + new_global_indexes = self.index_manager.allocate_indexes( + partition_id, count=(len(none_indexes) - len(batch_global_indexes)) + ) + batch_global_indexes.extend(new_global_indexes) + + partition.global_indexes.update(batch_global_indexes) + + for i in range(len(none_indexes)): + global_indexes[none_indexes[i]] = batch_global_indexes[i] + partition.keys_mapping[keys[none_indexes[i]]] = batch_global_indexes[i] + partition.revert_keys_mapping[batch_global_indexes[i]] = keys[ + none_indexes[i] + ] + + partition.ensure_samples_capacity(max(batch_global_indexes) + 1) + + verified_global_indexes = [idx for idx in global_indexes if idx is not None] + + # 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] + ] + + return self.generate_batch_meta( + partition_id, verified_global_indexes, data_fields, mode="force_fetch" + ) + + +def remove_samples(self, indexes: list[int]) -> None: + """Stock ``FieldMeta.remove_samples`` that rescans nested shapes only after ``RESCAN_REMOVED_FRACTION``.""" + before = len(self.global_indexes) + for idx in indexes: + self.per_sample_shapes.pop(idx, None) + self.global_indexes.discard(idx) + self.removed_since_scan += before - len(self.global_indexes) + + if len(self.global_indexes) == 0: + self.is_nested = None + self.is_non_tensor = None + self.shape = None + self.dtype = None + self.per_sample_shapes.clear() + else: + # Scanning every remaining shape is O(live), so only rescan after a fraction of them is gone. + if self.is_nested and self.removed_since_scan >= RESCAN_REMOVED_FRACTION * len( + self.global_indexes + ): + self.removed_since_scan = 0 + remaining_shapes = set( + tuple(shape) if isinstance(shape, list) else shape + for shape in self.per_sample_shapes.values() + ) + if len(remaining_shapes) == 1: + self.is_nested = False + self.shape = next(iter(remaining_shapes)) + self.per_sample_shapes.clear() + + +def patch_field_meta() -> None: + """Install ``remove_samples`` on ``FieldMeta`` in this process; the controller calls it on start.""" + FieldMeta.remove_samples = remove_samples + FieldMeta.removed_since_scan = 0 + + +def assert_stock_tq_version() -> None: + """Check that the installed TransferQueue is the version the overrides were copied from. + + Raises: + RuntimeError: If the version differs, which means the TQ pin moved. + """ + if transfer_queue.__version__ != STOCK_TQ_VERSION: + raise RuntimeError( + f"transfer_queue is {transfer_queue.__version__}, but tq_controller_patch copies methods from " + f"{STOCK_TQ_VERSION}. If this TQ includes the upstream change named in the module docstring, " + "delete this module and its call in _init_tq. Otherwise diff kv_retrieve_meta and " + "FieldMeta.remove_samples against this TQ, re-copy them only if they changed, then update " + "STOCK_TQ_VERSION." + ) + + +def make_patched_controller(actor_class: Any) -> Any: + """Return a Ray actor class that runs the controller with both overrides.""" + assert_stock_tq_version() + base = actor_class.__ray_metadata__.modified_class + + class PatchedController(base): + _nrl_kv_meta_patched = True + kv_retrieve_meta = kv_retrieve_meta + + def __init__(self, *args: Any, **kwargs: Any) -> None: + patch_field_meta() + super().__init__(*args, **kwargs) + + return ray.remote(**actor_class._default_options)(PatchedController) + + +def configure_controller_patch() -> None: + """Make ``tq.init`` start the patched controller; call before it.""" + actor_class = interface.TransferQueueController + if not getattr( + actor_class.__ray_metadata__.modified_class, "_nrl_kv_meta_patched", False + ): + interface.TransferQueueController = make_patched_controller(actor_class) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index c851c79cd05..af979be0cb1 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -53,6 +53,7 @@ from tensordict import TensorDict from nemo_rl.data_plane.adapters import mooncake_ssd_offload +from nemo_rl.data_plane.adapters.tq_controller_patch import configure_controller_patch from nemo_rl.data_plane.adapters.transfer_queue_env import rail_link_layers from nemo_rl.data_plane.codec import timed_codec from nemo_rl.data_plane.interfaces import ( @@ -1019,6 +1020,7 @@ def _init_tq( conf = OmegaConf.merge(base, overlay) + configure_controller_patch() # pyrefly: ignore # bad-argument-type tq.init(conf=conf) diff --git a/pyrefly.toml b/pyrefly.toml index ca658115f82..cea283db457 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -159,6 +159,7 @@ project-includes = [ "nemo_rl/data_plane/adapters/local.py", "nemo_rl/data_plane/adapters/mooncake_ssd_offload.py", "nemo_rl/data_plane/adapters/noop.py", + "nemo_rl/data_plane/adapters/tq_controller_patch.py", "nemo_rl/data_plane/adapters/tq_mooncake_checkpoint.py", "nemo_rl/data_plane/adapters/transfer_queue.py", "nemo_rl/data_plane/adapters/transfer_queue_env.py", diff --git a/tests/unit/data_plane/test_tq_controller_patch.py b/tests/unit/data_plane/test_tq_controller_patch.py new file mode 100644 index 00000000000..64540a154cb --- /dev/null +++ b/tests/unit/data_plane/test_tq_controller_patch.py @@ -0,0 +1,165 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest +import torch + +pytest.importorskip("transfer_queue") +from transfer_queue.controller import FieldMeta, TransferQueueController # noqa: E402 + +from nemo_rl.data_plane.adapters import tq_controller_patch as patch # noqa: E402 + +CONTROLLER = TransferQueueController.__ray_metadata__.modified_class + + +def test_installed_tq_is_the_version_the_overrides_were_copied_from(): + # fails when the TQ pin moves: re-copy the methods + patch.assert_stock_tq_version() + + +def test_other_tq_version_is_rejected(monkeypatch): + monkeypatch.setattr(patch.transfer_queue, "__version__", "0.0.0") + with pytest.raises(RuntimeError, match="STOCK_TQ_VERSION"): + patch.assert_stock_tq_version() + + +def _nested_field(n): + return FieldMeta( + global_indexes=set(range(n)), + dtype="torch.float32", + is_nested=True, + is_non_tensor=False, + per_sample_shapes={i: (256,) if i < n - 1 else (128,) for i in range(n)}, + ) + + +def _served_shapes(field_meta, indexes): + return [ + field_meta.per_sample_shapes[i] if field_meta.is_nested else field_meta.shape + for i in indexes + ] + + +def test_remove_samples_defers_the_rescan_until_an_eighth_is_gone(monkeypatch): + monkeypatch.setattr(FieldMeta, "remove_samples", patch.remove_samples) + monkeypatch.setattr(FieldMeta, "removed_since_scan", 0, raising=False) + field_meta = _nested_field(40) + + field_meta.remove_samples( + [39] + ) # the odd shape is gone, but 1 of 39 is under an eighth + assert field_meta.is_nested is True + + field_meta.remove_samples([0, 1, 2, 3]) # 5 of 35 reaches an eighth + assert field_meta.is_nested is False + assert field_meta.shape == (256,) + assert field_meta.per_sample_shapes == {} + + +def test_remove_samples_override_serves_the_same_shapes_as_stock(monkeypatch): + stock, patched = _nested_field(40), _nested_field(40) + monkeypatch.setattr(FieldMeta, "removed_since_scan", 0, raising=False) + remaining = list(range(40)) + while remaining: + cleared, remaining = remaining[:4], remaining[4:] + FieldMeta.remove_samples(stock, cleared) + patch.remove_samples(patched, cleared) + assert _served_shapes(patched, remaining) == _served_shapes(stock, remaining) + assert patched.global_indexes == stock.global_indexes + assert (patched.is_nested, patched.shape, patched.dtype) == ( + stock.is_nested, + stock.shape, + stock.dtype, + ) + + +def _schema(fields): + return { + f: { + "dtype": torch.int64, + "shape": (2,), + "is_nested": False, + "is_non_tensor": False, + } + for f in fields + } + + +def test_override_matches_stock_kv_retrieve_meta(): + ctl = CONTROLLER(polling_mode=True) + indexes = list( + ctl.kv_retrieve_meta(["a", "b", "c"], "p", create=True).global_indexes + ) + ctl.update_production_status( + "p", indexes, _schema(["f0", "f1"]), {g: {} for g in indexes} + ) + ctl.update_production_status( + "p", indexes[:1], _schema(["f2"]), {indexes[0]: {}} + ) # only row "a" + + for keys in (["a"], ["a", "b"], ["b", "c"], ["a", "b", "c"], ["missing"], []): + want = CONTROLLER.kv_retrieve_meta(ctl, keys, "p") + got = patch.kv_retrieve_meta(ctl, keys, "p") + assert list(got.global_indexes) == list(want.global_indexes) + assert sorted(got.field_schema) == sorted(want.field_schema) + assert sorted(patch.kv_retrieve_meta(ctl, ["a", "b"], "p").field_schema) == [ + "f0", + "f1", + ] + + +def test_override_matches_stock_kv_retrieve_meta_with_create(): + stock, patched = CONTROLLER(polling_mode=True), CONTROLLER(polling_mode=True) + for ctl, retrieve in ( + (stock, CONTROLLER.kv_retrieve_meta), + (patched, patch.kv_retrieve_meta), + ): + indexes = list(retrieve(ctl, ["a", "b"], "p", create=True).global_indexes) + ctl.update_production_status( + "p", indexes, _schema(["f0"]), {g: {} for g in indexes} + ) + for keys in (["a", "x"], ["y", "b", "z"], ["new"]): + want = CONTROLLER.kv_retrieve_meta(stock, keys, "p", create=True) + got = patch.kv_retrieve_meta(patched, keys, "p", create=True) + assert list(got.global_indexes) == list(want.global_indexes) + assert sorted(got.field_schema) == sorted(want.field_schema) + assert ( + patched._get_partition("p").keys_mapping + == stock._get_partition("p").keys_mapping + ) + + +def test_configure_controller_patch_twice_wraps_once(monkeypatch): + monkeypatch.setattr( + patch.interface, "TransferQueueController", TransferQueueController + ) + patch.configure_controller_patch() + wrapped = patch.interface.TransferQueueController + patch.configure_controller_patch() + assert patch.interface.TransferQueueController is wrapped + assert wrapped.__ray_metadata__.modified_class.__mro__[1] is CONTROLLER + + +def test_patched_controller_installs_both_overrides_on_start(monkeypatch): + monkeypatch.setattr(FieldMeta, "remove_samples", FieldMeta.remove_samples) + monkeypatch.setattr(FieldMeta, "removed_since_scan", 0, raising=False) + patched = patch.make_patched_controller( + TransferQueueController + ).__ray_metadata__.modified_class + # Ray wraps actor methods, so compare where the method comes from. + assert patched.kv_retrieve_meta.__module__ == patch.__name__ + # Building the class patches nothing; starting the controller does. + assert FieldMeta.remove_samples is not patch.remove_samples + + patched(polling_mode=True) + assert FieldMeta.remove_samples is patch.remove_samples