From 6f1db8519f4f0b9d26c452a72cb5b021cdba9531 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sat, 17 Oct 2026 00:27:52 -0700 Subject: [PATCH 1/7] perf(data_plane): cheaper TransferQueue kv_retrieve_meta on the controller kv_retrieve_meta reads a torch tensor once per field (col_mask[col_idx]). On the single-threaded TQ controller that is about a quarter of a put+get cycle. Build the field mask with one .all(dim=0).tolist() instead. The controller is a Ray actor, so the method is overridden in a controller subclass started by tq.init. The override is a plain copy of the stock method; a sha256 of the stock source makes it fail loudly when the TQ pin moves. Measured on one controller (real ZMQ, 20k live rows): +15% saturation throughput (5.6k -> 6.5k req/s) over geometric growth alone. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Zhiyu Li --- .../data_plane/adapters/tq_kv_meta_patch.py | 115 ++++++++++++++++++ nemo_rl/data_plane/adapters/transfer_queue.py | 2 + .../unit/data_plane/test_tq_kv_meta_patch.py | 51 ++++++++ 3 files changed, 168 insertions(+) create mode 100644 nemo_rl/data_plane/adapters/tq_kv_meta_patch.py create mode 100644 tests/unit/data_plane/test_tq_kv_meta_patch.py diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py new file mode 100644 index 00000000000..7818c6d65bc --- /dev/null +++ b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py @@ -0,0 +1,115 @@ +# 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`` for the TransferQueue controller. + +The stock method 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 +controller is a Ray actor, so it is overridden in a subclass started by ``tq.init``. +""" + +import hashlib +import inspect +import logging +from typing import Any + +from transfer_queue.metadata import BatchMeta + +logger = logging.getLogger(__name__) + +# sha256 of the stock method's source at the pinned TQ revision (see pyproject.toml). +_STOCK_SHA256 = "25635dcd71647ad7a03e284976e0b7aebc04ba7b6829cad59df1532d46e15301" + + +def kv_retrieve_meta(self, keys: list[str], partition_id: str, create: bool = False) -> BatchMeta: + 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 assert_stock_source(cls: type) -> None: + if hashlib.sha256(inspect.getsource(cls.kv_retrieve_meta).encode()).hexdigest() != _STOCK_SHA256: + raise RuntimeError( + "transfer_queue's kv_retrieve_meta changed, so the tq_kv_meta_patch override would replace newer " + "code. The TQ pin in pyproject.toml has moved: re-copy the method." + ) + + +def make_patched_controller(actor_class: Any) -> Any: + import ray + + base = actor_class.__ray_metadata__.modified_class + assert_stock_source(base) + + class PatchedController(base): + _nrl_kv_meta_patched = True + kv_retrieve_meta = kv_retrieve_meta + + return ray.remote(**actor_class._default_options)(PatchedController) + + +def configure_controller_patch() -> None: + """Make ``tq.init`` start the patched controller; call before it.""" + import transfer_queue.interface as interface + + 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..e13e176d373 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_kv_meta_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/tests/unit/data_plane/test_tq_kv_meta_patch.py b/tests/unit/data_plane/test_tq_kv_meta_patch.py new file mode 100644 index 00000000000..c94d3e8c739 --- /dev/null +++ b/tests/unit/data_plane/test_tq_kv_meta_patch.py @@ -0,0 +1,51 @@ +# 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 nemo_rl.data_plane.adapters import tq_kv_meta_patch as patch # noqa: E402 +from transfer_queue.controller import TransferQueueController # noqa: E402 + +CONTROLLER = TransferQueueController.__ray_metadata__.modified_class + + +def test_stock_source_matches_the_pinned_hash(): + patch.assert_stock_source(CONTROLLER) # fails when the TQ pin moves: re-copy the method + + +def test_changed_source_is_rejected(): + class Changed: + def kv_retrieve_meta(self, keys, partition_id, create=False): + return None + + with pytest.raises(RuntimeError, match="changed"): + patch.assert_stock_source(Changed) + + +def test_override_matches_stock_kv_retrieve_meta(): + ctl = CONTROLLER(polling_mode=True) + schema = lambda fields: { # noqa: E731 + f: {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False} for f in fields + } + 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"] From f164e77bf73b5b50de53a6958c35b44ca569a16b Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 21 Oct 2026 01:45:33 -0700 Subject: [PATCH 2/7] refactor(data_plane): hoist tq_kv_meta_patch imports, type-check the module Move the in-function ray and transfer_queue.interface imports to module top and add the module to pyrefly project-includes. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/adapters/tq_kv_meta_patch.py | 6 ++---- pyrefly.toml | 1 + 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py index 7818c6d65bc..4c3b378b26f 100644 --- a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py +++ b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py @@ -24,6 +24,8 @@ import logging from typing import Any +import ray +import transfer_queue.interface as interface from transfer_queue.metadata import BatchMeta logger = logging.getLogger(__name__) @@ -94,8 +96,6 @@ def assert_stock_source(cls: type) -> None: def make_patched_controller(actor_class: Any) -> Any: - import ray - base = actor_class.__ray_metadata__.modified_class assert_stock_source(base) @@ -108,8 +108,6 @@ class PatchedController(base): def configure_controller_patch() -> None: """Make ``tq.init`` start the patched controller; call before it.""" - import transfer_queue.interface as interface - 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/pyrefly.toml b/pyrefly.toml index ca658115f82..e6da50cdb37 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_kv_meta_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", From 3215a868c865ba7635caa528faddbf4b85883ea6 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 21 Oct 2026 17:44:48 -0700 Subject: [PATCH 3/7] perf(data_plane): throttle the nested-shape rescan in the TQ controller FieldMeta.remove_samples rescans every remaining shape of a nested field on each clear, so CLEAR_META grows with the number of live samples. Override it in the controller subclass to rescan only after an eighth of the live samples are gone. A source hash of the stock method fails loudly when the TQ pin moves, as for kv_retrieve_meta. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Zhiyu Li --- .../data_plane/adapters/tq_kv_meta_patch.py | 63 ++++++++++++++++--- .../unit/data_plane/test_tq_kv_meta_patch.py | 49 ++++++++++++++- 2 files changed, 99 insertions(+), 13 deletions(-) diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py index 4c3b378b26f..92df5b5b997 100644 --- a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py +++ b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py @@ -11,12 +11,17 @@ # 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`` for the TransferQueue controller. +"""Cheaper ``kv_retrieve_meta`` and ``FieldMeta.remove_samples`` for the TransferQueue controller. -The stock method reads a torch tensor once per field, about a quarter of a put+get +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 -controller is a Ray actor, so it is overridden in a subclass started by ``tq.init``. +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 an eighth of them is gone. + +The controller is a Ray actor, so both are applied in a subclass started by ``tq.init``. """ import hashlib @@ -26,12 +31,14 @@ import ray import transfer_queue.interface as interface +from transfer_queue.controller import FieldMeta from transfer_queue.metadata import BatchMeta logger = logging.getLogger(__name__) -# sha256 of the stock method's source at the pinned TQ revision (see pyproject.toml). -_STOCK_SHA256 = "25635dcd71647ad7a03e284976e0b7aebc04ba7b6829cad59df1532d46e15301" +# sha256 of the stock methods' source at the pinned TQ revision (see pyproject.toml). +_KV_RETRIEVE_META_SHA256 = "25635dcd71647ad7a03e284976e0b7aebc04ba7b6829cad59df1532d46e15301" +_REMOVE_SAMPLES_SHA256 = "847e2e96135d7be868cfcc0865e6a69649e586e1dbe60178f2a3f5ef20afb09f" def kv_retrieve_meta(self, keys: list[str], partition_id: str, create: bool = False) -> BatchMeta: @@ -87,22 +94,58 @@ def kv_retrieve_meta(self, keys: list[str], partition_id: str, create: bool = Fa return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch") -def assert_stock_source(cls: type) -> None: - if hashlib.sha256(inspect.getsource(cls.kv_retrieve_meta).encode()).hexdigest() != _STOCK_SHA256: +def remove_samples(self, indexes: list[int]) -> None: + 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 an eighth of them is gone. + if self.is_nested and self.removed_since_scan * 8 >= 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: + FieldMeta.remove_samples = remove_samples + FieldMeta.removed_since_scan = 0 + + +def assert_stock_source(method: Any, sha256: str) -> None: + if hashlib.sha256(inspect.getsource(method).encode()).hexdigest() != sha256: raise RuntimeError( - "transfer_queue's kv_retrieve_meta changed, so the tq_kv_meta_patch override would replace newer " + f"transfer_queue's {method.__qualname__} changed, so the tq_kv_meta_patch override would replace newer " "code. The TQ pin in pyproject.toml has moved: re-copy the method." ) def make_patched_controller(actor_class: Any) -> Any: base = actor_class.__ray_metadata__.modified_class - assert_stock_source(base) + assert_stock_source(base.kv_retrieve_meta, _KV_RETRIEVE_META_SHA256) + assert_stock_source(FieldMeta.remove_samples, _REMOVE_SAMPLES_SHA256) 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) diff --git a/tests/unit/data_plane/test_tq_kv_meta_patch.py b/tests/unit/data_plane/test_tq_kv_meta_patch.py index c94d3e8c739..16906be837a 100644 --- a/tests/unit/data_plane/test_tq_kv_meta_patch.py +++ b/tests/unit/data_plane/test_tq_kv_meta_patch.py @@ -16,13 +16,15 @@ pytest.importorskip("transfer_queue") from nemo_rl.data_plane.adapters import tq_kv_meta_patch as patch # noqa: E402 -from transfer_queue.controller import TransferQueueController # noqa: E402 +from transfer_queue.controller import FieldMeta, TransferQueueController # noqa: E402 CONTROLLER = TransferQueueController.__ray_metadata__.modified_class def test_stock_source_matches_the_pinned_hash(): - patch.assert_stock_source(CONTROLLER) # fails when the TQ pin moves: re-copy the method + # fails when the TQ pin moves: re-copy the methods + patch.assert_stock_source(CONTROLLER.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256) + patch.assert_stock_source(FieldMeta.remove_samples, patch._REMOVE_SAMPLES_SHA256) def test_changed_source_is_rejected(): @@ -31,7 +33,48 @@ def kv_retrieve_meta(self, keys, partition_id, create=False): return None with pytest.raises(RuntimeError, match="changed"): - patch.assert_stock_source(Changed) + patch.assert_stock_source(Changed.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256) + + +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 test_override_matches_stock_kv_retrieve_meta(): From 525b883c82c24520db749e6f4172492955cc1155 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 21 Oct 2026 18:11:09 -0700 Subject: [PATCH 4/7] style(data_plane): ruff-format the TQ controller patch, add docstrings and a hook test Format both files with the repo's ruff and sort the test imports, which the pre-commit lint-check job requires. Add docstrings to the module's public functions, and a test that the patched controller class carries the kv_retrieve_meta override and installs FieldMeta.remove_samples when started. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhiyu Li --- .../data_plane/adapters/tq_kv_meta_patch.py | 60 +++++++++++++---- .../unit/data_plane/test_tq_kv_meta_patch.py | 66 +++++++++++++++---- 2 files changed, 102 insertions(+), 24 deletions(-) diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py index 92df5b5b997..60571f4d9f9 100644 --- a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py +++ b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py @@ -37,17 +37,28 @@ logger = logging.getLogger(__name__) # sha256 of the stock methods' source at the pinned TQ revision (see pyproject.toml). -_KV_RETRIEVE_META_SHA256 = "25635dcd71647ad7a03e284976e0b7aebc04ba7b6829cad59df1532d46e15301" -_REMOVE_SAMPLES_SHA256 = "847e2e96135d7be868cfcc0865e6a69649e586e1dbe60178f2a3f5ef20afb09f" - - -def kv_retrieve_meta(self, keys: list[str], partition_id: str, create: bool = False) -> BatchMeta: - logger.debug(f"[{self.controller_id}] Retrieve keys {keys} in partition {partition_id}") +_KV_RETRIEVE_META_SHA256 = ( + "25635dcd71647ad7a03e284976e0b7aebc04ba7b6829cad59df1532d46e15301" +) +_REMOVE_SAMPLES_SHA256 = ( + "847e2e96135d7be868cfcc0865e6a69649e586e1dbe60178f2a3f5ef20afb09f" +) + + +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.") + logger.warning( + f"[{self.controller_id}]: Partition {partition_id} not found. Returning empty BatchMeta." + ) return BatchMeta.empty() self.create_partition(partition_id) @@ -64,7 +75,9 @@ def kv_retrieve_meta(self, keys: list[str], partition_id: str, create: bool = Fa f"They will be excluded from the retrieved BatchMeta." ) else: - batch_global_indexes = partition.activate_pre_allocated_indexes(len(none_indexes)) + 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( @@ -77,24 +90,31 @@ def kv_retrieve_meta(self, keys: list[str], partition_id: str, create: bool = Fa 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.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() + 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") + 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 an eighth is removed.""" before = len(self.global_indexes) for idx in indexes: self.per_sample_shapes.pop(idx, None) @@ -112,7 +132,8 @@ def remove_samples(self, indexes: list[int]) -> None: if self.is_nested and self.removed_since_scan * 8 >= 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() + 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 @@ -121,11 +142,21 @@ def remove_samples(self, indexes: list[int]) -> None: 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_source(method: Any, sha256: str) -> None: + """Check that ``method``'s source still matches the copy this module overrides. + + Args: + method: The stock TransferQueue method. + sha256: The sha256 of its source at the pinned TQ revision. + + Raises: + RuntimeError: If the source changed, which means the TQ pin moved. + """ if hashlib.sha256(inspect.getsource(method).encode()).hexdigest() != sha256: raise RuntimeError( f"transfer_queue's {method.__qualname__} changed, so the tq_kv_meta_patch override would replace newer " @@ -134,6 +165,7 @@ def assert_stock_source(method: Any, sha256: str) -> None: def make_patched_controller(actor_class: Any) -> Any: + """Return a Ray actor class that runs the controller with both overrides.""" base = actor_class.__ray_metadata__.modified_class assert_stock_source(base.kv_retrieve_meta, _KV_RETRIEVE_META_SHA256) assert_stock_source(FieldMeta.remove_samples, _REMOVE_SAMPLES_SHA256) @@ -152,5 +184,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: 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): + if not getattr( + actor_class.__ray_metadata__.modified_class, "_nrl_kv_meta_patched", False + ): interface.TransferQueueController = make_patched_controller(actor_class) diff --git a/tests/unit/data_plane/test_tq_kv_meta_patch.py b/tests/unit/data_plane/test_tq_kv_meta_patch.py index 16906be837a..11f774d4ca2 100644 --- a/tests/unit/data_plane/test_tq_kv_meta_patch.py +++ b/tests/unit/data_plane/test_tq_kv_meta_patch.py @@ -15,15 +15,18 @@ import torch pytest.importorskip("transfer_queue") -from nemo_rl.data_plane.adapters import tq_kv_meta_patch as patch # noqa: E402 from transfer_queue.controller import FieldMeta, TransferQueueController # noqa: E402 +from nemo_rl.data_plane.adapters import tq_kv_meta_patch as patch # noqa: E402 + CONTROLLER = TransferQueueController.__ray_metadata__.modified_class def test_stock_source_matches_the_pinned_hash(): # fails when the TQ pin moves: re-copy the methods - patch.assert_stock_source(CONTROLLER.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256) + patch.assert_stock_source( + CONTROLLER.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256 + ) patch.assert_stock_source(FieldMeta.remove_samples, patch._REMOVE_SAMPLES_SHA256) @@ -33,7 +36,9 @@ def kv_retrieve_meta(self, keys, partition_id, create=False): return None with pytest.raises(RuntimeError, match="changed"): - patch.assert_stock_source(Changed.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256) + patch.assert_stock_source( + Changed.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256 + ) def _nested_field(n): @@ -47,7 +52,10 @@ def _nested_field(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] + 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): @@ -55,7 +63,9 @@ def test_remove_samples_defers_the_rescan_until_an_eighth_is_gone(monkeypatch): 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 + 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 @@ -74,21 +84,55 @@ def test_remove_samples_override_serves_the_same_shapes_as_stock(monkeypatch): 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) + assert (patched.is_nested, patched.shape, patched.dtype) == ( + stock.is_nested, + stock.shape, + stock.dtype, + ) def test_override_matches_stock_kv_retrieve_meta(): ctl = CONTROLLER(polling_mode=True) schema = lambda fields: { # noqa: E731 - f: {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False} for f in fields + f: { + "dtype": torch.int64, + "shape": (2,), + "is_nested": False, + "is_non_tensor": False, + } + for f in fields } - 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" + 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"] + assert sorted(patch.kv_retrieve_meta(ctl, ["a", "b"], "p").field_schema) == [ + "f0", + "f1", + ] + + +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 From 54e1e039e8ee6a90ce678bf5bba2106cb0f58bc2 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 21 Oct 2026 18:17:40 -0700 Subject: [PATCH 5/7] refactor(data_plane): name the nested-shape rescan fraction Replace the literal 8 in remove_samples with RESCAN_REMOVED_FRACTION = 1 / 8. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhiyu Li --- nemo_rl/data_plane/adapters/tq_kv_meta_patch.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py index 60571f4d9f9..0341bda8f2d 100644 --- a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py +++ b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py @@ -19,7 +19,7 @@ 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 an eighth of them is gone. +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``. """ @@ -44,6 +44,9 @@ "847e2e96135d7be868cfcc0865e6a69649e586e1dbe60178f2a3f5ef20afb09f" ) +# 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 @@ -114,7 +117,7 @@ def kv_retrieve_meta( def remove_samples(self, indexes: list[int]) -> None: - """Stock ``FieldMeta.remove_samples`` that rescans nested shapes only after an eighth is removed.""" + """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) @@ -128,8 +131,10 @@ def remove_samples(self, indexes: list[int]) -> None: self.dtype = None self.per_sample_shapes.clear() else: - # Scanning every remaining shape is O(live), so only rescan after an eighth of them is gone. - if self.is_nested and self.removed_since_scan * 8 >= len(self.global_indexes): + # 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 From fc9191b97b2034ac0e1757e53c9cfc7eefda26b8 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 21 Oct 2026 18:31:20 -0700 Subject: [PATCH 6/7] refactor(data_plane): guard the TQ overrides on the TQ version, not source hashes The two sha256 source hashes tripped the secrets detector as high-entropy hex strings. Check transfer_queue.__version__ against STOCK_TQ_VERSION instead. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhiyu Li --- .../data_plane/adapters/tq_kv_meta_patch.py | 32 +++++++------------ .../unit/data_plane/test_tq_kv_meta_patch.py | 20 ++++-------- 2 files changed, 17 insertions(+), 35 deletions(-) diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py index 0341bda8f2d..752adbdbd64 100644 --- a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py +++ b/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py @@ -24,25 +24,19 @@ The controller is a Ray actor, so both are applied in a subclass started by ``tq.init``. """ -import hashlib -import inspect import logging 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 logger = logging.getLogger(__name__) -# sha256 of the stock methods' source at the pinned TQ revision (see pyproject.toml). -_KV_RETRIEVE_META_SHA256 = ( - "25635dcd71647ad7a03e284976e0b7aebc04ba7b6829cad59df1532d46e15301" -) -_REMOVE_SAMPLES_SHA256 = ( - "847e2e96135d7be868cfcc0865e6a69649e586e1dbe60178f2a3f5ef20afb09f" -) +# 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 @@ -152,28 +146,24 @@ def patch_field_meta() -> None: FieldMeta.removed_since_scan = 0 -def assert_stock_source(method: Any, sha256: str) -> None: - """Check that ``method``'s source still matches the copy this module overrides. - - Args: - method: The stock TransferQueue method. - sha256: The sha256 of its source at the pinned TQ revision. +def assert_stock_tq_version() -> None: + """Check that the installed TransferQueue is the version the overrides were copied from. Raises: - RuntimeError: If the source changed, which means the TQ pin moved. + RuntimeError: If the version differs, which means the TQ pin moved. """ - if hashlib.sha256(inspect.getsource(method).encode()).hexdigest() != sha256: + if transfer_queue.__version__ != STOCK_TQ_VERSION: raise RuntimeError( - f"transfer_queue's {method.__qualname__} changed, so the tq_kv_meta_patch override would replace newer " - "code. The TQ pin in pyproject.toml has moved: re-copy the method." + f"transfer_queue is {transfer_queue.__version__}, but tq_kv_meta_patch copies methods from " + f"{STOCK_TQ_VERSION}. Re-copy kv_retrieve_meta and FieldMeta.remove_samples, 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 - assert_stock_source(base.kv_retrieve_meta, _KV_RETRIEVE_META_SHA256) - assert_stock_source(FieldMeta.remove_samples, _REMOVE_SAMPLES_SHA256) class PatchedController(base): _nrl_kv_meta_patched = True diff --git a/tests/unit/data_plane/test_tq_kv_meta_patch.py b/tests/unit/data_plane/test_tq_kv_meta_patch.py index 11f774d4ca2..e18e242cea8 100644 --- a/tests/unit/data_plane/test_tq_kv_meta_patch.py +++ b/tests/unit/data_plane/test_tq_kv_meta_patch.py @@ -22,23 +22,15 @@ CONTROLLER = TransferQueueController.__ray_metadata__.modified_class -def test_stock_source_matches_the_pinned_hash(): +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_source( - CONTROLLER.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256 - ) - patch.assert_stock_source(FieldMeta.remove_samples, patch._REMOVE_SAMPLES_SHA256) - + patch.assert_stock_tq_version() -def test_changed_source_is_rejected(): - class Changed: - def kv_retrieve_meta(self, keys, partition_id, create=False): - return None - with pytest.raises(RuntimeError, match="changed"): - patch.assert_stock_source( - Changed.kv_retrieve_meta, patch._KV_RETRIEVE_META_SHA256 - ) +def test_other_tq_version_is_rejected(monkeypatch): + monkeypatch.setattr(patch.transfer_queue, "__version__", "0.0.0") + with pytest.raises(RuntimeError, match="Re-copy"): + patch.assert_stock_tq_version() def _nested_field(n): From 612cdfc284f583df2d0999f28e4ac9a2fb9b9889 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Sun, 25 Oct 2026 22:20:22 -0700 Subject: [PATCH 7/7] refactor(data_plane): rename to tq_controller_patch, log via TQ, test create and re-config Review on #4769: - "tq_controller_patch.py fits better": rename the module and its test. - C3 "Use TQ's get_logger in the copy": TQ_LOGGING_LEVEL controls the copied log lines again. - C2 "reword the version error": say when to delete the module and when only to bump STOCK_TQ_VERSION, and name the upstream change in the module docstring. - C1 "Add tests for the copied create=True branch and for calling configure_controller_patch() twice". Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhiyu Li --- ...v_meta_patch.py => tq_controller_patch.py} | 14 ++++-- nemo_rl/data_plane/adapters/transfer_queue.py | 2 +- pyrefly.toml | 2 +- ...a_patch.py => test_tq_controller_patch.py} | 49 ++++++++++++++++--- 4 files changed, 54 insertions(+), 13 deletions(-) rename nemo_rl/data_plane/adapters/{tq_kv_meta_patch.py => tq_controller_patch.py} (91%) rename tests/unit/data_plane/{test_tq_kv_meta_patch.py => test_tq_controller_patch.py} (72%) diff --git a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py b/nemo_rl/data_plane/adapters/tq_controller_patch.py similarity index 91% rename from nemo_rl/data_plane/adapters/tq_kv_meta_patch.py rename to nemo_rl/data_plane/adapters/tq_controller_patch.py index 752adbdbd64..14103c8187c 100644 --- a/nemo_rl/data_plane/adapters/tq_kv_meta_patch.py +++ b/nemo_rl/data_plane/adapters/tq_controller_patch.py @@ -22,9 +22,11 @@ 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. """ -import logging from typing import Any import ray @@ -32,8 +34,10 @@ 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 -logger = logging.getLogger(__name__) +# 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" @@ -154,8 +158,10 @@ def assert_stock_tq_version() -> None: """ if transfer_queue.__version__ != STOCK_TQ_VERSION: raise RuntimeError( - f"transfer_queue is {transfer_queue.__version__}, but tq_kv_meta_patch copies methods from " - f"{STOCK_TQ_VERSION}. Re-copy kv_retrieve_meta and FieldMeta.remove_samples, then update " + 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." ) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index e13e176d373..af979be0cb1 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -53,7 +53,7 @@ from tensordict import TensorDict from nemo_rl.data_plane.adapters import mooncake_ssd_offload -from nemo_rl.data_plane.adapters.tq_kv_meta_patch import configure_controller_patch +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 ( diff --git a/pyrefly.toml b/pyrefly.toml index e6da50cdb37..cea283db457 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -159,7 +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_kv_meta_patch.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_kv_meta_patch.py b/tests/unit/data_plane/test_tq_controller_patch.py similarity index 72% rename from tests/unit/data_plane/test_tq_kv_meta_patch.py rename to tests/unit/data_plane/test_tq_controller_patch.py index e18e242cea8..64540a154cb 100644 --- a/tests/unit/data_plane/test_tq_kv_meta_patch.py +++ b/tests/unit/data_plane/test_tq_controller_patch.py @@ -17,7 +17,7 @@ pytest.importorskip("transfer_queue") from transfer_queue.controller import FieldMeta, TransferQueueController # noqa: E402 -from nemo_rl.data_plane.adapters import tq_kv_meta_patch as patch # noqa: E402 +from nemo_rl.data_plane.adapters import tq_controller_patch as patch # noqa: E402 CONTROLLER = TransferQueueController.__ray_metadata__.modified_class @@ -29,7 +29,7 @@ def test_installed_tq_is_the_version_the_overrides_were_copied_from(): def test_other_tq_version_is_rejected(monkeypatch): monkeypatch.setattr(patch.transfer_queue, "__version__", "0.0.0") - with pytest.raises(RuntimeError, match="Re-copy"): + with pytest.raises(RuntimeError, match="STOCK_TQ_VERSION"): patch.assert_stock_tq_version() @@ -83,9 +83,8 @@ def test_remove_samples_override_serves_the_same_shapes_as_stock(monkeypatch): ) -def test_override_matches_stock_kv_retrieve_meta(): - ctl = CONTROLLER(polling_mode=True) - schema = lambda fields: { # noqa: E731 +def _schema(fields): + return { f: { "dtype": torch.int64, "shape": (2,), @@ -94,14 +93,18 @@ def test_override_matches_stock_kv_retrieve_meta(): } 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} + "p", indexes, _schema(["f0", "f1"]), {g: {} for g in indexes} ) ctl.update_production_status( - "p", indexes[:1], schema(["f2"]), {indexes[0]: {}} + "p", indexes[:1], _schema(["f2"]), {indexes[0]: {}} ) # only row "a" for keys in (["a"], ["a", "b"], ["b", "c"], ["a", "b", "c"], ["missing"], []): @@ -115,6 +118,38 @@ def test_override_matches_stock_kv_retrieve_meta(): ] +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)