From d21c676e09055b6e2b9daa75b214c168688ac82b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 14:58:19 +0800 Subject: [PATCH 01/32] [feat] add kv_link so keys can read another key's fields without copies GRPO-style rollouts put one prompt and link its n responses to it, so the prompt payload is stored once. Links are controller metadata only; the SimpleStorage read path resolves them and reads each source row once. A put on a linked field replaces the link, clearing a linked key drops only the link, and a source cannot be cleared while other keys still link to it. dump_data_by_key records links and load_data_by_key relinks. --- tests/e2e/test_kv_link_e2e.py | 140 ++++++++++++++++++ tests/test_async_simple_storage_manager.py | 40 +++++ tests/test_data_dump.py | 2 +- transfer_queue/__init__.py | 4 + transfer_queue/client.py | 22 +++ transfer_queue/controller.py | 118 ++++++++++++++- transfer_queue/data_dump.py | 2 +- transfer_queue/interface.py | 58 +++++++- transfer_queue/metadata.py | 15 ++ .../managers/simple_storage_manager.py | 49 ++++-- transfer_queue/utils/zmq_utils.py | 2 + 11 files changed, 430 insertions(+), 22 deletions(-) create mode 100644 tests/e2e/test_kv_link_e2e.py diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py new file mode 100644 index 00000000..bb503c90 --- /dev/null +++ b/tests/e2e/test_kv_link_e2e.py @@ -0,0 +1,140 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# 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. + +"""End-to-end tests for kv_link: n keys reading one prompt payload without copies.""" + +import shutil + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from tensordict import TensorDict + +import transfer_queue as tq + +PARTITION = "link" +N = 3 + + +@pytest.fixture(scope="module") +def tq_system(): + if not ray.is_initialized(): + ray.init(namespace="TestKVLinkE2E") + tq.init( + OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": 200, "num_data_storage_units": 2}, + }, + } + ) + ) + yield + tq.close() + ray.shutdown() + + +@pytest.fixture(autouse=True) +def fresh_partition(tq_system): + yield + tq.get_client().clear_partition(PARTITION) + + +def _prompt(p: int) -> torch.Tensor: + return torch.tensor([p, p + 1, p + 2]) + + +def _put_group(p: int) -> list[str]: + """Put prompt ``p`` once and link its N responses to it; return the response keys.""" + tq.kv_put(key=f"p{p}", partition_id=PARTITION, fields={"prompt": _prompt(p)}) + keys = [f"p{p}_r{i}" for i in range(N)] + tq.kv_link(keys, PARTITION, f"p{p}", "prompt") + return keys + + +def _links(key: str) -> dict: + return tq.get_client().kv_describe([key], PARTITION)["rows"][key]["links"] + + +def _rows(field: torch.Tensor) -> list[list]: + """A batched field comes back as a nested tensor, which has no ``tolist``.""" + return [row.tolist() for row in field.unbind()] + + +def test_linked_keys_read_the_source_payload(tq_system): + keys = _put_group(0) + _put_group(10) + tq.kv_batch_put( + keys=keys, + partition_id=PARTITION, + fields=TensorDict({"response": torch.arange(len(keys)).unsqueeze(1)}, batch_size=len(keys)), + ) + + data = tq.kv_batch_get(keys, PARTITION) + + assert _rows(data["prompt"]) == [_prompt(0).tolist()] * N + [_prompt(10).tolist()] * N + assert _rows(data["response"]) == [[i] for i in range(len(keys))] + + +def test_link_requires_a_produced_source_field(tq_system): + tq.kv_put(key="p0", partition_id=PARTITION, fields={"prompt": _prompt(0)}) + + with pytest.raises(RuntimeError, match="not found"): + tq.kv_link("r", PARTITION, "missing", "prompt") + with pytest.raises(RuntimeError, match="must have produced"): + tq.kv_link("r", PARTITION, "p0", "image") + + +def test_source_cannot_be_cleared_while_linked(tq_system): + keys = _put_group(0) + + with pytest.raises(RuntimeError, match="linked by other keys"): + tq.kv_clear("p0", PARTITION) + with pytest.raises(RuntimeError, match="linked by other keys"): + tq.kv_clear("p0", PARTITION, fields="prompt") + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + + tq.kv_clear(keys[1:], PARTITION, fields="prompt") + tq.kv_clear([keys[0], "p0"], PARTITION) + + assert tq.kv_list(PARTITION)[PARTITION].keys() == set(keys[1:]) + assert all(not _links(key) for key in keys[1:]) + + +def test_put_replaces_the_link(tq_system): + keys = _put_group(0) + + tq.kv_put(key=keys[0], partition_id=PARTITION, fields={"prompt": _prompt(7)}) + + assert _links(keys[0]) == {} + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(7).tolist()] + [_prompt(0).tolist()] * 2 + + +def test_dump_and_load_relink(tq_system, dump_test_root): + keys = _put_group(0) + dump_dir = dump_test_root / "link_dump" + + with pytest.raises(ValueError, match="must be dumped"): + tq.dump_data_by_key(dump_dir, keys, PARTITION) + + tq.dump_data_by_key(dump_dir, ["p0", *keys], PARTITION) + tq.get_client().clear_partition(PARTITION) + tq.load_data_by_key(dump_dir) + shutil.rmtree(dump_dir) + + assert all(_links(key) == {"prompt": "p0"} for key in keys) + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py index b4dbfbf5..a5389db1 100644 --- a/tests/test_async_simple_storage_manager.py +++ b/tests/test_async_simple_storage_manager.py @@ -317,6 +317,46 @@ def parser(data, args): assert called_with == {"storage_0": ([0, 2], parser, ["b", "d"]), "storage_1": ([3, 1], parser, ["a", "c"])} +@pytest.mark.asyncio +async def test_get_data_reads_linked_field_once_from_source(): + """Rows 0 and 1 link ``prompt`` to row 2: it is read once from row 2, never from 0 or 1.""" + storage_unit_infos = { + f"storage_{i}": ZMQServerInfo( + role=Role.STORAGE, id=f"storage_{i}", ip="127.0.0.1", ports={"put_get_socket": 19012 + i} + ) + for i in range(2) + } + with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"): + manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "test_get_links" + manager.storage_unit_infos = storage_unit_infos + manager.controller_info = None + manager.controller_handshake_socket = None + manager.zmq_context = None + + schema = {"dtype": torch.float32, "shape": (1,), "is_nested": False, "is_non_tensor": False} + batch_meta = BatchMeta( + global_indexes=[0, 1], + partition_ids=["p0"] * 2, + field_schema={"prompt": schema, "response": schema}, + production_status=np.ones(2, dtype=np.int8), + _links=[{"prompt": 2}, {"prompt": 2}], + ) + calls = [] + + async def fake_get(global_indexes, fields, target_storage_unit=None, **kwargs): + calls.append((list(global_indexes), list(fields))) + return fields, {f: [torch.tensor([gi * 10.0 + (f == "response")]) for gi in global_indexes] for f in fields} + + manager._get_from_single_storage_unit = fake_get + + result = await manager.get_data(batch_meta) + + assert sorted(calls) == [([0], ["response"]), ([1], ["response"]), ([2], ["prompt"])] + assert [row.tolist() for row in result["prompt"].unbind()] == [[20.0], [20.0]] + assert [row.tolist() for row in result["response"].unbind()] == [[1.0], [11.0]] + + @pytest.mark.asyncio async def test_clear_data_routes_from_hash(caplog): """clear_data should route using global_idx % num_su (hash routing).""" diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py index b2de9cfe..134de21c 100644 --- a/tests/test_data_dump.py +++ b/tests/test_data_dump.py @@ -67,7 +67,7 @@ def test_row_index_compacts_tensors_inside_tags(monkeypatch, tmp_path): tag = {"nested": [SimpleNamespace(value=batch[0])]} client = SimpleNamespace( kv_describe=lambda *_: { - "rows": {"key": {"global_index": 0, "fields": [], "tag": tag}}, + "rows": {"key": {"global_index": 0, "fields": [], "links": {}, "tag": tag}}, "field_schema": {}, } ) diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 05e63cdf..364c1279 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -24,6 +24,7 @@ async_kv_batch_update, async_kv_clear, async_kv_describe, + async_kv_link, async_kv_list, async_kv_put, async_kv_update, @@ -38,6 +39,7 @@ kv_batch_update, kv_clear, kv_describe, + kv_link, kv_list, kv_put, kv_update, @@ -76,6 +78,7 @@ "kv_list", "kv_describe", "kv_clear", + "kv_link", "kv_update", "async_kv_put", "async_kv_batch_put", @@ -85,6 +88,7 @@ "async_kv_list", "async_kv_describe", "async_kv_clear", + "async_kv_link", "async_kv_update", "KVBatchMeta", ] diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 73f458f4..09b3fd8a 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1186,6 +1186,23 @@ async def async_kv_describe( ) return {"rows": response.body["rows"], "field_schema": response.body["field_schema"]} + @with_controller_socket + async def async_kv_link( + self, + keys: list[str], + partition_id: str, + source_key: str, + fields: list[str], + socket: zmq.asyncio.Socket | None = None, + ) -> None: + """Asynchronously make ``fields`` of ``keys`` read the payload of ``source_key``; see ``kv_link``.""" + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_LINK, + response_type=ZMQRequestType.KV_LINK_RESPONSE, + body={"keys": keys, "partition_id": partition_id, "source_key": source_key, "fields": fields}, + ) + def close(self) -> None: """Close the client and cleanup resources including storage manager. @@ -1501,6 +1518,7 @@ def wrapper(*args, **kwargs): self._kv_retrieve_keys = _make_sync(self.async_kv_retrieve_keys) self._kv_list = _make_sync(self.async_kv_list) self._kv_describe = _make_sync(self.async_kv_describe) + self._kv_link = _make_sync(self.async_kv_link) self._validate_dump_schema = _make_sync(self.async_validate_dump_schema) self._dump_rows_by_index = _make_sync(self.async_dump_rows_by_index) self._load_rows_by_key = _make_sync(self.async_load_rows_by_key) @@ -1978,6 +1996,10 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: """ return self._kv_describe(keys, partition_id) + def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: list[str]) -> None: + """Synchronously make ``fields`` of ``keys`` read the payload of ``source_key``; see ``kv_link``.""" + return self._kv_link(keys, partition_id, source_key, fields) + # ==================== Selective Data Dump API ==================== def validate_dump_schema(self, partition_id: str, field_schema: dict) -> None: diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index c44ac236..ac2ff5d6 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -367,6 +367,9 @@ class DataPartitionStatus: keys_mapping: dict[str, int] = field(default_factory=dict) # key -> global_idx revert_keys_mapping: dict[int, str] = field(default_factory=dict) # global_idx -> key + # Linked fields hold no payload of their own and read the source row's instead + links: dict[int, dict[str, int]] = field(default_factory=dict) # global_idx -> {field: source global_idx} + # Dynamic configuration - these are computed from the current state @property def total_samples_num(self) -> int: @@ -565,6 +568,9 @@ def update_production_status( # Save these global_indexes self.global_indexes.update(global_indices) + # A written value replaces the link it lands on + self._drop_links(global_indices, field_names) + return True except Exception as e: @@ -944,6 +950,7 @@ def clear_data(self, indexes_to_release: list[int], clear_consumption: bool = Tr for idx in indexes_to_release: self.field_custom_backend_meta.pop(idx, None) self.custom_meta.pop(idx, None) + self.links.pop(idx, None) if idx in self.revert_keys_mapping: self.keys_mapping.pop(self.revert_keys_mapping[idx], None) @@ -970,6 +977,60 @@ def clear_fields(self, global_indexes: list[int], field_names: list[str]): backend_meta = self.field_custom_backend_meta.get(idx, {}) for field_name in field_names: backend_meta.pop(field_name, None) + self._drop_links(global_indexes, field_names) + + # ==================== Field Links ==================== + def link_fields(self, global_indexes: list[int], source: int, field_names: list[str]) -> None: + """Mark ``field_names`` produced on these rows, read from the payload of row ``source``.""" + schema = {} + for name in field_names: + meta = self.field_metadata[name] + schema[name] = { + "dtype": meta.dtype, + "shape": meta.shape, + "is_nested": meta.is_nested, + "is_non_tensor": meta.is_non_tensor, + } + if meta.is_nested: + schema[name]["per_sample_shapes"] = {idx: meta.per_sample_shapes[source] for idx in global_indexes} + # Before recording links: update_production_status drops links on the cells it marks. + if not self.update_production_status(global_indexes, field_names, schema): + raise RuntimeError(f"failed to link fields {field_names} in partition {self.partition_id}") + # Point at the row holding the payload, so a chain of links is one hop. + source_links = self.links.get(source, {}) + for idx in global_indexes: + row_links = self.links.setdefault(idx, {}) + for name in field_names: + row_links[name] = source_links.get(name, source) + + def check_not_linked(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: + """Refuse to clear payloads that rows outside this clear still read through links.""" + clearing = set(global_indexes) + sources = { + source + for idx, row_links in self.links.items() + if idx not in clearing + for name, source in row_links.items() + if source in clearing and (field_names is None or name in field_names) + } + if sources: + keys = sorted(str(self.revert_keys_mapping.get(idx, idx)) for idx in sources) + raise ValueError( + f"keys {keys} in partition {self.partition_id!r} are linked by other keys; " + "clear the linked keys first or in the same call" + ) + + def _drop_links(self, global_indexes: list[int], field_names: list[str]) -> None: + if not self.links: + return + for idx in global_indexes: + row_links = self.links.get(idx) + if row_links is None: + continue + for name in field_names: + row_links.pop(name, None) + if not row_links: + del self.links[idx] def kv_retrieve_indexes(self, keys: list[str]) -> list[int | None]: """Translate the user-specified keys to global_indexes""" @@ -1507,6 +1568,13 @@ def generate_batch_meta( # Convert controller dict[int, dict] → BatchMeta list[dict] (aligned with batch_global_indexes) custom_meta_list = [custom_meta_dict.get(global_index, {}) for global_index in batch_global_indexes] custom_backend_meta_list = [custom_backend_meta.get(global_index, {}) for global_index in batch_global_indexes] + links_list = None + if partition.links: + wanted = set(data_fields) + links_list = [ + {name: source for name, source in partition.links.get(global_index, {}).items() if name in wanted} + for global_index in batch_global_indexes + ] batch_meta = BatchMeta( global_indexes=batch_global_indexes, @@ -1515,6 +1583,7 @@ def generate_batch_meta( production_status=production_status, custom_meta=custom_meta_list, _custom_backend_meta=custom_backend_meta_list, + _links=links_list, ) return batch_meta @@ -1540,6 +1609,12 @@ def mark_clearing(self, global_indexes: list[int], partition_ids: list[str]) -> combined = list(zip(partition_ids, global_indexes, strict=True)) combined.sort(key=itemgetter(0)) + # Validate every partition first so a refused clear marks nothing. + for partition_id, group in groupby(combined, key=itemgetter(0)): + partition = self._get_partition(partition_id) + if partition: + partition.check_not_linked([idx for _, idx in group]) + for partition_id, group in groupby(combined, key=itemgetter(0)): partition = self._get_partition(partition_id) if not partition: @@ -1777,6 +1852,7 @@ def kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) fields = [f for f in fields if f in partition.field_metadata] if not global_indexes or not fields: return BatchMeta.empty() + partition.check_not_linked(global_indexes, fields) metadata = self.generate_batch_meta(partition_id, global_indexes, fields, mode="force_fetch") partition.clear_fields(global_indexes, fields) @@ -1830,10 +1906,11 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: partition_id: Partition that owns ``keys``. Returns: - ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema": - {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` is empty for a row with no - produced field yet; ``unfinalized`` lists the stored fields not yet ready (written - by ``kv_update(finalize=False)``); ``field_schema`` covers the described rows' fields. + ``{"rows": {key: {"global_index", "fields", "unfinalized", "links", "tag"}}, "field_schema": + {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` lists the row's stored fields and is + empty for a row with no produced field yet; ``unfinalized`` lists the stored fields not yet + ready (written by ``kv_update(finalize=False)``); ``links`` maps each linked field to its + source key; ``field_schema`` covers the described rows' stored fields. Raises: KeyError: The partition or any key does not exist. @@ -1849,10 +1926,11 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: rows: dict[str, dict[str, Any]] = {} for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True): + links = partition.links.get(global_index, {}) fields = sorted( field_name for field_name, field_meta in partition.field_metadata.items() - if global_index in field_meta.global_indexes + if global_index in field_meta.global_indexes and field_name not in links ) rows[key] = { "global_index": global_index, @@ -1862,6 +1940,7 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: for name in fields if not partition.production_status[global_index, partition.field_name_mapping[name]] ], + "links": {name: partition.revert_keys_mapping[source] for name, source in links.items()}, "tag": partition.custom_meta.get(global_index, {}), } # Only the declared types: row dtypes and shapes live with the storage units. @@ -1872,6 +1951,29 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: } return {"rows": rows, "field_schema": field_schema} + def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: list[str]) -> None: + """Make ``fields`` of ``keys`` read the payload of ``source_key`` without copying it. + + Missing ``keys`` are created. The source must already have produced ``fields``. + + Raises: + KeyError: The partition or ``source_key`` does not exist. + ValueError: ``fields`` is empty or ``source_key`` has not produced all of them. + """ + partition = self._get_partition(partition_id) + if partition is None or source_key not in partition.keys_mapping: + raise KeyError(f"source key {source_key!r} not found in partition {partition_id!r}") + source = partition.keys_mapping[source_key] + missing = [ + name + for name in fields + if name not in partition.field_metadata or source not in partition.field_metadata[name].global_indexes + ] + if not fields or missing: + raise ValueError(f"source key {source_key!r} must have produced fields {missing or fields}") + metadata = self.kv_retrieve_meta(keys, partition_id, create=True) + partition.link_fields(metadata.global_indexes, source, fields) + def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" self.zmq_context = zmq.Context() @@ -2099,6 +2201,7 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | ZMQRequestType.KV_PREPARE_UPDATE: self._handle_kv_prepare_update_request, ZMQRequestType.VALIDATE_DUMP_SCHEMA: self._handle_validate_dump_schema_request, ZMQRequestType.KV_CLEAR_FIELDS: self._handle_kv_clear_fields_request, + ZMQRequestType.KV_LINK: self._handle_kv_link_request, ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request, ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: self._handle_load_controller_checkpoint_request, } @@ -2304,6 +2407,11 @@ def _handle_kv_prepare_update_request(self, request_msg: ZMQMessage) -> ZMQMessa request_msg, ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE, {"metadata": metadata, "merge_plan": merge_plan} ) + def _handle_kv_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + self.kv_link(params["keys"], params["partition_id"], params["source_key"], params["fields"]) + return self._make_response(request_msg, ZMQRequestType.KV_LINK_RESPONSE, {}) + def _handle_kv_clear_fields_request(self, request_msg: ZMQMessage) -> ZMQMessage: params = request_msg.body metadata = self.kv_clear_fields( diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py index ad3791fb..e39a8e30 100644 --- a/transfer_queue/data_dump.py +++ b/transfer_queue/data_dump.py @@ -23,7 +23,7 @@ / dump_info.json # version, partition, counts - row_index.pt # key -> {global_index, fields, unfinalized, tag} + row_index.pt # key -> {global_index, fields, unfinalized, links, tag} shards/ shard_info.json # file, unit, row count, source index -> [offset, length] shard__.pkl # independent {global_index, fields} records diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 0b4ab04b..8ae30a0a 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -33,7 +33,7 @@ import shutil import subprocess import time -from collections import Counter +from collections import Counter, defaultdict from importlib import resources from pathlib import Path from typing import Any, Callable @@ -694,6 +694,33 @@ def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id, fields=fields)) +def kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: list[str] | str) -> None: + """Let ``keys`` read ``fields`` from ``source_key`` without storing another copy. + + For example, the n responses of a GRPO prompt can link the prompt's fields instead + of putting the prompt n times. Missing ``keys`` are created, and linked fields read + as produced. Putting a linked field stores a private value in place of the link; + clearing a linked key or field drops only the link. ``source_key`` cannot be cleared + while other keys link to it, unless they are cleared in the same call. + + Args: + keys: Single key or list of keys that will read the linked fields. + partition_id: Partition holding both ``keys`` and ``source_key``. + source_key: Existing key that has produced ``fields``. + fields: Single field name or list of field names to link. + + Raises: + RuntimeError: If ``source_key`` does not exist or has not produced all of ``fields``. + NotImplementedError: If the storage backend is not SimpleStorage. + + Example: + >>> tq.kv_put(key="prompt_0", partition_id="train", fields={"input_ids": prompt_ids}) + >>> tq.kv_link([f"prompt_0_{i}" for i in range(n)], "train", "prompt_0", "input_ids") + """ + tq_client = _maybe_create_tq_client() + tq_client._run_coroutine(async_kv_link(keys, partition_id, source_key, fields)) + + def kv_update( key: str, partition_id: str, @@ -1223,6 +1250,20 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[ await tq_client.async_clear_samples(batch_meta) +async def async_kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: list[str] | str) -> None: + """Asynchronously let ``keys`` read ``fields`` from ``source_key``. See ``kv_link``.""" + tq_client = _maybe_create_tq_client() + # Only SimpleStorage resolves links on read; other backends would read the linking key's own cells. + if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): + raise NotImplementedError(f"{type(tq_client.storage_manager).__name__} does not support kv_link") + await tq_client.async_kv_link( + [keys] if isinstance(keys, str) else keys, + partition_id, + source_key, + [fields] if isinstance(fields, str) else fields, + ) + + def _single_update_batch(fields: dict[str, Any]) -> TensorDict: """Wrap one sample's field mapping in a one-row TensorDict.""" if not isinstance(fields, dict) or not fields: @@ -1480,7 +1521,8 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - Args: dump_dir: Directory to write the dump into. - keys: Keys to dump. Duplicates are dropped, first occurrence wins. + keys: Keys to dump. Duplicates are dropped, first occurrence wins. Must include + the source key of every field linked by ``kv_link``. partition_id: Partition that owns ``keys``. Returns: @@ -1488,6 +1530,7 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - Raises: FileExistsError: ``dump_dir`` already exists. + ValueError: A key links to a source key that is not dumped. RuntimeError: TransferQueue is not initialized, the partition or a key does not exist, or a storage unit holds no data for a row it was asked to dump. NotImplementedError: The storage backend does not support selective dumps. @@ -1508,6 +1551,9 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - description = client.kv_describe(unique_keys, partition_id) if unique_keys else {"rows": {}, "field_schema": {}} row_index = {"partition_id": partition_id, **description} rows = row_index["rows"] + missing_sources = sorted({source for row in rows.values() for source in row["links"].values()} - rows.keys()) + if missing_sources: + raise ValueError(f"linked source keys {missing_sources} must be dumped with the keys linking to them") # A row whose fields are all still unproduced has nothing for a storage unit to # dump, but it keeps its key and tag so the restore can recreate the row. fields_by_index = {row["global_index"]: row["fields"] for row in rows.values() if row["fields"]} @@ -1567,5 +1613,13 @@ def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: rows = row_index["rows"] client.validate_dump_schema(partition_id, row_index["field_schema"]) bytes_read = client.load_rows_by_key(partition_id, rows, shards) + # Relink once the sources hold their payloads, one call per source and field set. + link_groups: dict[tuple[str, tuple[str, ...]], list[str]] = defaultdict(list) + for key, row in rows.items(): + links = row.get("links", {}) + for source in set(links.values()): + link_groups[source, tuple(sorted(name for name, s in links.items() if s == source))].append(key) + for (source, names), keys in link_groups.items(): + client.kv_link(keys, partition_id, source, list(names)) logger.info(f"Restored {len(rows)} keys into partition {partition_id} from {dump_dir}, reading {bytes_read} bytes") return {key: row["tag"] for key, row in rows.items()} diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py index ab395e07..3d88146c 100644 --- a/transfer_queue/metadata.py +++ b/transfer_queue/metadata.py @@ -202,6 +202,7 @@ class BatchMeta: extra_info: Additional batch-level information. custom_meta: Per-sample user-defined metadata, list aligned with global_indexes. _custom_backend_meta: Per-sample per-field storage backend metadata, list aligned with global_indexes. + _links: Per-sample {field: source global_index} for linked fields, list aligned with global_indexes. """ __slots__ = ( @@ -215,6 +216,7 @@ class BatchMeta: "_size", "_field_names", "_is_ready", + "_links", ) def __init__( @@ -226,6 +228,7 @@ def __init__( extra_info: dict[str, Any] | None = None, custom_meta: list[dict[str, Any]] | None = None, _custom_backend_meta: list[dict[str, Any]] | None = None, + _links: list[dict[str, int]] | None = None, ) -> None: if field_schema is None: field_schema = {} @@ -288,6 +291,12 @@ def __init__( raise ValueError( f"_custom_backend_meta length {len(self._custom_backend_meta)} != batch_size {batch_size}" ) + if not _links: + self._links: list[dict[str, int]] = [{} for _ in range(batch_size)] + else: + self._links = [dict(d) for d in _links] + if len(self._links) != batch_size: + raise ValueError(f"_links length {len(self._links)} != batch_size {batch_size}") def __getstate__(self): """Serialize for pickle/Ray. @@ -523,6 +532,7 @@ def select_samples(self, sample_indices: list[int]) -> "BatchMeta": extra_info=self.extra_info, custom_meta=new_custom_meta, _custom_backend_meta=new_custom_backend_meta, + _links=[self._links[i] for i in sample_indices], ) def select_fields(self, field_names: list[str]) -> "BatchMeta": @@ -552,6 +562,7 @@ def select_fields(self, field_names: list[str]) -> "BatchMeta": extra_info=copy.deepcopy(self.extra_info), custom_meta=copy.deepcopy(self.custom_meta), _custom_backend_meta=selected_custom_backend_meta, + _links=[{f: s for f, s in m.items() if f in field_names} for m in self._links], ) def copy(self) -> "BatchMeta": @@ -564,6 +575,7 @@ def copy(self) -> "BatchMeta": extra_info=copy.deepcopy(self.extra_info), custom_meta=copy.deepcopy(self.custom_meta), _custom_backend_meta=copy.deepcopy(self._custom_backend_meta), + _links=self._links, ) def __len__(self) -> int: @@ -701,6 +713,7 @@ def union(self, other: "BatchMeta") -> "BatchMeta": extra_info=merged_extra_info, custom_meta=merged_custom_meta, _custom_backend_meta=merged_custom_backend_meta, + _links=[{**mine, **theirs} for mine, theirs in zip(self._links, other._links, strict=True)], ) @classmethod @@ -826,6 +839,7 @@ def concat(cls, data: list["BatchMeta"], validate: bool = True) -> "BatchMeta": extra_info=merged_extra_info, custom_meta=all_custom_meta, _custom_backend_meta=all_custom_backend_meta, + _links=[links for chunk in data for links in chunk._links], ) def reorder(self, indices: list[int]): @@ -852,6 +866,7 @@ def reorder(self, indices: list[int]): self.custom_meta = [self.custom_meta[i] for i in indices] self._custom_backend_meta = [self._custom_backend_meta[i] for i in indices] + self._links = [self._links[i] for i in indices] @classmethod def empty(cls, extra_info: dict[str, Any] | None = None) -> "BatchMeta": diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index e505a732..732f21fb 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -800,7 +800,8 @@ async def get_data( """ Retrieve data from remote StorageUnit based on metadata. - Routes to each SU using global_idx % num_su (hash routing). + Routes to each SU using global_idx % num_su (hash routing). A linked field is read + from its source row, once however many samples link to it. Args: metadata: BatchMeta that contains metadata for data retrieval. @@ -818,26 +819,44 @@ async def get_data( if metadata.size == 0: return TensorDict({}, batch_size=0) - routing = self._group_by_hash(metadata.global_indexes) + # A parser sees every selected field of each sample on one unit; a linked field lives on another row. + if data_parser is not None and any(metadata._links): + raise NotImplementedError("data_parser does not support samples with fields linked by kv_link") # Parser-backed gets are not replayed, matching put: the public API does not constrain parser side effects. max_attempts = 1 if data_parser is not None else None + # A unit request reads the same fields for every row in it, and a row stores only its + # unlinked fields, so rows are grouped by the field set they physically hold. Without + # links this is one group in batch order, so batch_positions index parser_args. + fields_by_row: dict[int, set[str]] = defaultdict(set) + for global_index, links in zip(metadata.global_indexes, metadata._links, strict=True): + for field in metadata.field_names: + fields_by_row[links.get(field, global_index)].add(field) + rows_by_fields: dict[tuple[str, ...], list[int]] = defaultdict(list) + for global_index, fields in fields_by_row.items(): + rows_by_fields[tuple(sorted(fields))].append(global_index) + requests = [ + (su_id, list(fields), group) + for fields, rows in rows_by_fields.items() + for su_id, group in self._group_by_hash(rows).items() + ] + tasks = [ self._request_with_retry( "get", su_id, - f"samples={len(group.global_indexes)} fields={list(metadata.field_names)}", + f"samples={len(group.global_indexes)} fields={fields}", partial( self._get_from_single_storage_unit, group.global_indexes, - metadata.field_names, + fields, target_storage_unit=su_id, data_parser=data_parser, parser_args=None if parser_args is None else [parser_args[p] for p in group.batch_positions], ), max_attempts=max_attempts, ) - for su_id, group in routing.items() + for su_id, fields, group in requests ] try: results = await asyncio.gather(*tasks) @@ -849,19 +868,23 @@ async def get_data( f"[{self.storage_manager_id}]: get_data failed. " f"partition_id={metadata.partition_ids[0]}, " f"num_samples={metadata.size}, " - f"num_storage_units={len(routing)}, " + f"num_requests={len(requests)}, " f"error={type(e).__name__}: {e}" ) raise - # Scatter results directly to batch positions — no intermediate per-sample dict - n = len(metadata.global_indexes) - ordered_data: dict[str, list] = {field: [None] * n for field in metadata.field_names} - - for (su_id, group), (fields, su_data) in zip(routing.items(), results, strict=True): + values: dict[tuple[int, str], Any] = {} + for (_, _, group), (fields, su_data) in zip(requests, results, strict=True): for field in fields: - for i, pos in enumerate(group.batch_positions): - ordered_data[field][pos] = su_data[field][i] + for i, global_index in enumerate(group.global_indexes): + values[global_index, field] = su_data[field][i] + ordered_data = { + field: [ + values[links.get(field, global_index), field] + for global_index, links in zip(metadata.global_indexes, metadata._links, strict=True) + ] + for field in metadata.field_names + } tensor_data = {field: self._pack_field_values(v) for field, v in ordered_data.items()} diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 9c765cfb..4ff5047b 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -129,6 +129,8 @@ class ZMQRequestType(ExplicitEnum): KV_PREPARE_UPDATE_RESPONSE = "KV_PREPARE_UPDATE_RESPONSE" KV_CLEAR_FIELDS = "KV_CLEAR_FIELDS" KV_CLEAR_FIELDS_RESPONSE = "KV_CLEAR_FIELDS_RESPONSE" + KV_LINK = "KV_LINK" + KV_LINK_RESPONSE = "KV_LINK_RESPONSE" # METRICS GET_METRICS = "GET_METRICS" From dd244ae214fdab7a87c86e94a9e25ecd40a3043e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:04:21 +0800 Subject: [PATCH 02/32] [fix] refuse kv_link to a source that is being cleared kv_link checked field metadata, which stays until the clear finishes, so a link made between mark_clearing and the storage delete pointed at a payload about to vanish. Check production status, which mark_clearing zeroes first. --- tests/e2e/test_kv_link_e2e.py | 10 ++++++++++ transfer_queue/controller.py | 4 +++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index bb503c90..34019f44 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -99,6 +99,16 @@ def test_link_requires_a_produced_source_field(tq_system): tq.kv_link("r", PARTITION, "p0", "image") +def test_link_refuses_a_source_being_cleared(tq_system): + tq.kv_put(key="p0", partition_id=PARTITION, fields={"prompt": _prompt(0)}) + source = tq.get_client().kv_describe(["p0"], PARTITION)["rows"]["p0"]["global_index"] + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + ray.get(controller.mark_clearing.remote([source], [PARTITION])) + + with pytest.raises(RuntimeError, match="must have produced"): + tq.kv_link("r", PARTITION, "p0", "prompt") + + def test_source_cannot_be_cleared_while_linked(tq_system): keys = _put_group(0) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index ac2ff5d6..c1bfcab0 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1964,10 +1964,12 @@ def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: l if partition is None or source_key not in partition.keys_mapping: raise KeyError(f"source key {source_key!r} not found in partition {partition_id!r}") source = partition.keys_mapping[source_key] + # Production status, not field metadata: mark_clearing zeroes it before storage deletes the payload. missing = [ name for name in fields - if name not in partition.field_metadata or source not in partition.field_metadata[name].global_indexes + if name not in partition.field_name_mapping + or partition.production_status[source, partition.field_name_mapping[name]] != 1 ] if not fields or missing: raise ValueError(f"source key {source_key!r} must have produced fields {missing or fields}") From 6fa7ac146f62afda09a1bf3bc1c8545a6b9dedc6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:23:21 +0800 Subject: [PATCH 03/32] [fix] reject duplicate keys in kv_link A repeated missing key was registered twice, leaving a keyless row that held a link and could never be cleared. Reject duplicates before any RPC, as the batch writes do. --- tests/e2e/test_kv_link_e2e.py | 3 +++ transfer_queue/interface.py | 5 ++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 34019f44..5c7d9afe 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -97,6 +97,9 @@ def test_link_requires_a_produced_source_field(tq_system): tq.kv_link("r", PARTITION, "missing", "prompt") with pytest.raises(RuntimeError, match="must have produced"): tq.kv_link("r", PARTITION, "p0", "image") + with pytest.raises(ValueError, match="duplicates"): + tq.kv_link(["r", "r"], PARTITION, "p0", "prompt") + assert tq.kv_list(PARTITION)[PARTITION].keys() == {"p0"} def test_link_refuses_a_source_being_cleared(tq_system): diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 8ae30a0a..1fc3c02e 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -711,6 +711,7 @@ def kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: l Raises: RuntimeError: If ``source_key`` does not exist or has not produced all of ``fields``. + ValueError: If ``keys`` contains duplicates. NotImplementedError: If the storage backend is not SimpleStorage. Example: @@ -1256,8 +1257,10 @@ async def async_kv_link(keys: list[str] | str, partition_id: str, source_key: st # Only SimpleStorage resolves links on read; other backends would read the linking key's own cells. if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): raise NotImplementedError(f"{type(tq_client.storage_manager).__name__} does not support kv_link") + keys = [keys] if isinstance(keys, str) else keys + _check_unique_keys(keys) await tq_client.async_kv_link( - [keys] if isinstance(keys, str) else keys, + keys, partition_id, source_key, [fields] if isinstance(fields, str) else fields, From 53b5e729e225c7e6c661679ad9e792bd2ee87d9b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:35:12 +0800 Subject: [PATCH 04/32] [perf] index kv_link links by source so clear checks scale with the clear --- tests/test_controller_data_partitions.py | 29 ++++++++++++++++++++++++ transfer_queue/controller.py | 23 +++++++++++-------- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index fd458ed1..981ad3c7 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1492,3 +1492,32 @@ def test_fieldmeta_to_batch_schema_nested_missing_sample(self): schema = field_meta.to_batch_schema([0, 1]) assert schema["per_sample_shapes"] == [(3,), None] + + +class TestFieldLinks: + """Rows linking to a source row through DataPartitionStatus.link_fields.""" + + SCHEMA = {"x": {"dtype": "torch.int64", "shape": (3,), "is_nested": False, "is_non_tensor": False}} + + def _linked_partition(self): + """Source row 0 holds ``x``; rows 1 and 2 link to it.""" + from transfer_queue.controller import DataPartitionStatus + + partition = DataPartitionStatus(partition_id="links") + assert partition.update_production_status([0], ["x"], field_schema=self.SCHEMA) + partition.link_fields([1, 2], 0, ["x"]) + return partition + + def test_source_index_follows_put_and_clear(self): + import pytest + + partition = self._linked_partition() + assert partition.linked_by == {0: {1, 2}} + with pytest.raises(ValueError, match="linked by other keys"): + partition.check_not_linked([0, 1]) + + assert partition.update_production_status([1], ["x"], field_schema=self.SCHEMA) + partition.clear_data([2]) + + assert partition.links == {} and partition.linked_by == {} + partition.check_not_linked([0]) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index c1bfcab0..1c5579b0 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -369,6 +369,7 @@ class DataPartitionStatus: # Linked fields hold no payload of their own and read the source row's instead links: dict[int, dict[str, int]] = field(default_factory=dict) # global_idx -> {field: source global_idx} + linked_by: dict[int, set[int]] = field(default_factory=dict) # source global_idx -> global_idxs linking to it # Dynamic configuration - these are computed from the current state @property @@ -947,10 +948,10 @@ def clear_data(self, indexes_to_release: list[int], clear_consumption: bool = Tr # only clear empty fields for field_name in empty_fields: self.field_metadata.pop(field_name) + self._drop_links(indexes_to_release) for idx in indexes_to_release: self.field_custom_backend_meta.pop(idx, None) self.custom_meta.pop(idx, None) - self.links.pop(idx, None) if idx in self.revert_keys_mapping: self.keys_mapping.pop(self.revert_keys_mapping[idx], None) @@ -1002,16 +1003,17 @@ def link_fields(self, global_indexes: list[int], source: int, field_names: list[ row_links = self.links.setdefault(idx, {}) for name in field_names: row_links[name] = source_links.get(name, source) + self.linked_by.setdefault(row_links[name], set()).add(idx) def check_not_linked(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: """Refuse to clear payloads that rows outside this clear still read through links.""" clearing = set(global_indexes) sources = { source - for idx, row_links in self.links.items() + for source in clearing + for idx in self.linked_by.get(source, ()) if idx not in clearing - for name, source in row_links.items() - if source in clearing and (field_names is None or name in field_names) + and any(s == source and (field_names is None or n in field_names) for n, s in self.links[idx].items()) } if sources: keys = sorted(str(self.revert_keys_mapping.get(idx, idx)) for idx in sources) @@ -1020,15 +1022,18 @@ def check_not_linked(self, global_indexes: list[int], field_names: list[str] | N "clear the linked keys first or in the same call" ) - def _drop_links(self, global_indexes: list[int], field_names: list[str]) -> None: - if not self.links: - return + def _drop_links(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: + """Drop the links of these rows on ``field_names`` (all fields if None).""" for idx in global_indexes: row_links = self.links.get(idx) if row_links is None: continue - for name in field_names: - row_links.pop(name, None) + for name in list(row_links) if field_names is None else field_names: + source = row_links.pop(name, None) + if source is not None and source not in row_links.values(): + self.linked_by[source].discard(idx) + if not self.linked_by[source]: + del self.linked_by[source] if not row_links: del self.links[idx] From 7f6f570ec1c2b1af32df336fec9053c3513ed72c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:36:36 +0800 Subject: [PATCH 05/32] [fix] linked keys follow their source's readiness and shape --- tests/e2e/test_kv_link_e2e.py | 11 +++++++++++ tests/test_controller_data_partitions.py | 11 +++++++++++ transfer_queue/controller.py | 11 +++++++++++ 3 files changed, 33 insertions(+) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 5c7d9afe..2a1a4919 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -137,6 +137,17 @@ def test_put_replaces_the_link(tq_system): assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(7).tolist()] + [_prompt(0).tolist()] * 2 +def test_unfinalized_source_update_is_hidden_from_linked_keys(tq_system): + keys = _put_group(0) + + tq.kv_update("p0", PARTITION, {"prompt": _prompt(5)}, merge_fn=lambda old, new: new, finalize=False) + with pytest.raises(ValueError, match="not ready"): + tq.kv_batch_get(keys, PARTITION, select_fields="prompt") + + tq.kv_update("p0", PARTITION, {"prompt": _prompt(5)}, merge_fn=lambda old, new: new) + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(5).tolist()] * N + + def test_dump_and_load_relink(tq_system, dump_test_root): keys = _put_group(0) dump_dir = dump_test_root / "link_dump" diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index 981ad3c7..dc5abc01 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1521,3 +1521,14 @@ def test_source_index_follows_put_and_clear(self): assert partition.links == {} and partition.linked_by == {} partition.check_not_linked([0]) + + def test_linked_rows_follow_source_readiness_and_shape(self): + partition = self._linked_partition() + wider = {"x": {**self.SCHEMA["x"], "shape": (5,)}} + + assert partition.update_production_status([0], ["x"], field_schema=wider, ready=False) + assert partition.production_status[[0, 1, 2], 0].tolist() == [0, 0, 0] + assert partition.field_metadata["x"].per_sample_shapes == {0: (5,), 1: (5,), 2: (5,)} + + assert partition.update_production_status([0], ["x"], field_schema=wider) + assert partition.production_status[[0, 1, 2], 0].tolist() == [1, 1, 1] diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 1c5579b0..39caff14 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -572,6 +572,17 @@ def update_production_status( # A written value replaces the link it lands on self._drop_links(global_indices, field_names) + # Rows linking to a written cell read its new payload, so they take on its readiness + # (an unfinalized kv_update stays hidden behind links too) and its shape. + for source in global_indices: + for idx in self.linked_by.get(source, ()): + for name, linked_source in self.links[idx].items(): + if linked_source == source and name in field_names: + self.production_status[idx, self.field_name_mapping[name]] = int(ready) + meta = self.field_metadata[name] + if meta.is_nested: + meta.per_sample_shapes[idx] = meta.per_sample_shapes[source] + return True except Exception as e: From 3fb2f1458f1a95bad0737f66921f80f320fd2bd6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:37:34 +0800 Subject: [PATCH 06/32] [fix] refuse kv_update on fields linked by kv_link with a clear error --- tests/e2e/test_kv_link_e2e.py | 15 +++++++++++++++ transfer_queue/controller.py | 16 ++++++++++++++++ transfer_queue/interface.py | 3 ++- 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 2a1a4919..b369e139 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -137,6 +137,21 @@ def test_put_replaces_the_link(tq_system): assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(7).tolist()] + [_prompt(0).tolist()] * 2 +def test_update_refuses_a_linked_field(tq_system): + keys = _put_group(0) + + with pytest.raises(RuntimeError, match="cannot merge into fields linked"): + tq.kv_batch_update( + ["new", keys[0]], + PARTITION, + TensorDict({"prompt": torch.stack([_prompt(5)] * 2)}, batch_size=2), + merge_fn=lambda old, new: new, + ) + + assert "new" not in tq.kv_list(PARTITION)[PARTITION] + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + + def test_unfinalized_source_update_is_hidden_from_linked_keys(tq_system): keys = _put_group(0) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 39caff14..05da7f2a 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1840,7 +1840,23 @@ def kv_prepare_update( Returns: The keys' BatchMeta and ``{field: {"dtype", "is_non_tensor", "held"}}``, where ``held`` is aligned with the BatchMeta rows. The types are None for a new field. + + Raises: + ValueError: A key links one of ``fields`` by kv_link; nothing is registered. """ + partition = self._get_partition(partition_id) + if partition is not None: + # A unit merges into the old value on the sample's own row; a linked value lives on the source's. + linked = [ + key + for key, idx in zip(keys, partition.kv_retrieve_indexes(keys), strict=True) + if idx is not None and any(name in partition.links.get(idx, {}) for name in fields) + ] + if linked: + raise ValueError( + f"kv_update cannot merge into fields linked by kv_link on keys {linked}; " + "kv_put a value to replace the link first" + ) metadata = self.kv_retrieve_meta(keys, partition_id, create=True) partition = self._get_partition(partition_id) assert partition is not None diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 1fc3c02e..2670bf9e 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -753,7 +753,8 @@ def kv_update( Raises: TypeError: If ``fields`` is not a non-empty dict or ``merge_fn`` is not callable. RuntimeError: If ``merge_fn`` fails or a stored value would change the field's - tensor/non-tensor type or dtype; the failing storage unit writes nothing. + tensor/non-tensor type or dtype; the failing storage unit writes nothing. Also + if a field is linked by ``kv_link``, before anything is registered or written. NotImplementedError: If the storage backend is not SimpleStorage. Note: From 3230d9ffaddfe454cdaaad5cdfb06749d10a8771 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:38:41 +0800 Subject: [PATCH 07/32] [fix] refuse kv_link that relinks a source key or links a key to itself --- tests/e2e/test_kv_link_e2e.py | 13 +++++++++++++ transfer_queue/controller.py | 8 +++++++- transfer_queue/interface.py | 3 ++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index b369e139..cc9065ab 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -99,9 +99,22 @@ def test_link_requires_a_produced_source_field(tq_system): tq.kv_link("r", PARTITION, "p0", "image") with pytest.raises(ValueError, match="duplicates"): tq.kv_link(["r", "r"], PARTITION, "p0", "prompt") + with pytest.raises(RuntimeError, match="cannot link to itself"): + tq.kv_link(["r", "p0"], PARTITION, "p0", "prompt") assert tq.kv_list(PARTITION)[PARTITION].keys() == {"p0"} +def test_link_refuses_to_relink_a_source(tq_system): + keys = _put_group(0) + tq.kv_put(key="q0", partition_id=PARTITION, fields={"prompt": _prompt(9)}) + + with pytest.raises(RuntimeError, match="linked by other keys"): + tq.kv_link("p0", PARTITION, "q0", "prompt") + + assert _links("p0") == {} + assert _rows(tq.kv_batch_get(["p0", *keys], PARTITION)["prompt"]) == [_prompt(0).tolist()] * (N + 1) + + def test_link_refuses_a_source_being_cleared(tq_system): tq.kv_put(key="p0", partition_id=PARTITION, fields={"prompt": _prompt(0)}) source = tq.get_client().kv_describe(["p0"], PARTITION)["rows"]["p0"]["global_index"] diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 05da7f2a..5bb4987a 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1990,7 +1990,8 @@ def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: l Raises: KeyError: The partition or ``source_key`` does not exist. - ValueError: ``fields`` is empty or ``source_key`` has not produced all of them. + ValueError: ``fields`` is empty, ``source_key`` has not produced all of them, or + ``keys`` include ``source_key`` or a key other keys link these fields to. """ partition = self._get_partition(partition_id) if partition is None or source_key not in partition.keys_mapping: @@ -2005,6 +2006,11 @@ def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: l ] if not fields or missing: raise ValueError(f"source key {source_key!r} must have produced fields {missing or fields}") + existing = [idx for idx in partition.kv_retrieve_indexes(keys) if idx is not None] + if source in existing: + raise ValueError(f"source key {source_key!r} cannot link to itself") + # Keys linking to a relinked key would keep reading its old payload, not the new source. + partition.check_not_linked(existing, fields) metadata = self.kv_retrieve_meta(keys, partition_id, create=True) partition.link_fields(metadata.global_indexes, source, fields) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 2670bf9e..ca367c8e 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -710,7 +710,8 @@ def kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: l fields: Single field name or list of field names to link. Raises: - RuntimeError: If ``source_key`` does not exist or has not produced all of ``fields``. + RuntimeError: If ``source_key`` does not exist or has not produced all of ``fields``, + or ``keys`` include ``source_key`` or a key that other keys link these fields to. ValueError: If ``keys`` contains duplicates. NotImplementedError: If the storage backend is not SimpleStorage. From 28fe2411477e5581447e808c52d6e8d945cd6c42 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:40:27 +0800 Subject: [PATCH 08/32] [fix] a put through BatchMeta drops its local links on the written fields --- tests/e2e/test_kv_link_e2e.py | 11 +++++++++++ tests/test_metadata.py | 11 +++++++++++ transfer_queue/metadata.py | 4 ++++ 3 files changed, 26 insertions(+) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index cc9065ab..245895ca 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -150,6 +150,17 @@ def test_put_replaces_the_link(tq_system): assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(7).tolist()] + [_prompt(0).tolist()] * 2 +def test_put_through_meta_reads_back_the_written_value(tq_system): + keys = _put_group(0) + client = tq.get_client() + meta = client.kv_retrieve_meta(keys, PARTITION) + + client.put(TensorDict({"prompt": torch.stack([_prompt(7)] * N)}, batch_size=N), metadata=meta) + + assert _rows(client.get_data(meta)["prompt"]) == [_prompt(7).tolist()] * N + assert _rows(tq.kv_batch_get("p0", PARTITION)["prompt"]) == [_prompt(0).tolist()] + + def test_update_refuses_a_linked_field(tq_system): keys = _put_group(0) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 6fc49d2e..3433271c 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -256,6 +256,17 @@ def test_setstate_readonly_production_status(self): restored.add_fields(td) # Should not raise assert restored.is_ready + def test_add_fields_replaces_links_on_written_fields(self): + """A put through this meta stores its own values, so later reads must not follow links.""" + from tensordict import TensorDict + + batch = self._make_batch(batch_size=2) + batch._links = [{"field_a": 7, "field_b": 7}, {"field_a": 7}] + + batch.add_fields(TensorDict({"field_a": torch.randn(2, 2)}, batch_size=2)) + + assert batch._links == [{"field_b": 7}, {}] + def test_shallow_copy_isolation_global_indexes(self): """Modifying the original global_indexes list does not affect BatchMeta.""" original_indexes = [0, 1, 2] diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py index 3d88146c..690ea50e 100644 --- a/transfer_queue/metadata.py +++ b/transfer_queue/metadata.py @@ -484,6 +484,10 @@ def apply_field_schema(self, field_schema: dict[str, dict[str, Any]], set_all_re set_all_ready (bool): If True, set all production_status to READY_FOR_CONSUME. Default is True. """ self.field_schema.update(field_schema) + # The written values replace links on these fields, as they do in the controller. + for links in self._links: + for name in field_schema: + links.pop(name, None) if set_all_ready: self.production_status[:] = 1 From bf8cce637b0e3759a7be288c30387f353a3c81ef Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:41:34 +0800 Subject: [PATCH 09/32] [feat] kv_describe reports the fields each key links by kv_link --- tests/e2e/test_kv_interface_e2e.py | 2 +- tests/e2e/test_kv_link_e2e.py | 2 +- transfer_queue/client.py | 2 +- transfer_queue/interface.py | 8 +++++--- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 572336a8..de3f611a 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -1389,7 +1389,7 @@ def tokens_ready(): with pytest.raises(ValueError, match="not finalized"): tq_api.kv_batch_get_by_meta(meta) row = tq_api.kv_describe(keys=key, partition_id=partition_id)[key] - assert row == {"fields": ["prompt", "tokens"], "unfinalized": ["tokens"]} + assert row == {"fields": ["prompt", "tokens"], "unfinalized": ["tokens"], "links": {}} with pytest.raises(RuntimeError, match="not found"): tq_api.kv_describe(keys=[key, "absent"], partition_id=partition_id) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 245895ca..f0e8cb03 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -68,7 +68,7 @@ def _put_group(p: int) -> list[str]: def _links(key: str) -> dict: - return tq.get_client().kv_describe([key], PARTITION)["rows"][key]["links"] + return tq.kv_describe(key, PARTITION)[key]["links"] def _rows(field: torch.Tensor) -> list[list]: diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 09b3fd8a..d3bcbbd0 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1175,7 +1175,7 @@ async def async_kv_describe( key raises instead of being dropped. Returns: - ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema": + ``{"rows": {key: {"global_index", "fields", "unfinalized", "links", "tag"}}, "field_schema": {field: {"is_nested", "is_non_tensor"}}}``. """ response = await self._request_controller( diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index ca367c8e..f3b12b24 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -1195,8 +1195,9 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st partition_id: Partition containing the keys. Returns: - ``{key: {"fields", "unfinalized"}}``. ``fields`` lists every stored field of the - key and ``unfinalized`` the subset that is not ready yet. + ``{key: {"fields", "unfinalized", "links"}}``. ``fields`` lists every stored field of + the key and ``unfinalized`` the subset that is not ready yet. ``links`` maps each field + linked by ``kv_link`` to its source key; such a field is ready when the source's is. Raises: RuntimeError: If the partition or any key does not exist. @@ -1211,7 +1212,8 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st tq_client = _maybe_create_tq_client() description = await tq_client.async_kv_describe(list(dict.fromkeys(keys)), partition_id) return { - key: {"fields": row["fields"], "unfinalized": row["unfinalized"]} for key, row in description["rows"].items() + key: {"fields": row["fields"], "unfinalized": row["unfinalized"], "links": row["links"]} + for key, row in description["rows"].items() } From 7b3172a38f76d37b3716caac1e0147075c52a6e9 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:42:17 +0800 Subject: [PATCH 10/32] [doc] document and test that data_parser refuses fields linked by kv_link --- tests/e2e/test_kv_link_e2e.py | 13 +++++++++++++ transfer_queue/interface.py | 4 +++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index f0e8cb03..6f30e009 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -90,6 +90,19 @@ def test_linked_keys_read_the_source_payload(tq_system): assert _rows(data["response"]) == [[i] for i in range(len(keys))] +def test_data_parser_refuses_linked_fields(tq_system): + keys = _put_group(0) + tq.kv_batch_put(keys, PARTITION, fields=TensorDict({"response": torch.arange(N).unsqueeze(1)}, batch_size=N)) + + def identity(data, args): + return data + + with pytest.raises(NotImplementedError, match="kv_link"): + tq.kv_batch_get(keys, PARTITION, data_parser=identity) + data = tq.kv_batch_get(keys, PARTITION, select_fields="response", data_parser=identity) + assert _rows(data["response"]) == [[i] for i in range(N)] + + def test_link_requires_a_produced_source_field(tq_system): tq.kv_put(key="p0", partition_id=PARTITION, fields={"prompt": _prompt(0)}) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index f3b12b24..f9fd2cad 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -561,7 +561,8 @@ def kv_batch_get( - It is sent with cloudpickle, so a module-level function must be importable on the storage units. SSD-offloaded values are read in full before parsing. - Only supported by SimpleStorage; KV backends raise ``NotImplementedError``. + Only supported by SimpleStorage; KV backends raise ``NotImplementedError``, and so + does a selected field linked by ``kv_link``. parser_args: Optional per-key arguments for ``data_parser``, aligned with ``keys``. Requires ``data_parser``. They are serialized like request data, so tuples arrive as lists. @@ -576,6 +577,7 @@ def kv_batch_get( ValueError: If ``parser_args`` is given without ``data_parser`` or does not match the number of keys RuntimeError: If ``data_parser`` fails or breaks its contract on a storage unit + NotImplementedError: If ``data_parser`` is given and a selected field is linked by ``kv_link`` Example: >>> import transfer_queue as tq From c1275d6c1cec5b38577ba56ce2192b39c95317a4 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:42:26 +0800 Subject: [PATCH 11/32] [doc] describe kv_link readiness, kv_update and source-sample behavior --- transfer_queue/interface.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index f9fd2cad..c0c62b31 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -701,10 +701,14 @@ def kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: l For example, the n responses of a GRPO prompt can link the prompt's fields instead of putting the prompt n times. Missing ``keys`` are created, and linked fields read - as produced. Putting a linked field stores a private value in place of the link; + the source's current value and are ready whenever the source's are. Putting a linked + field stores a private value in place of the link, and ``kv_update`` refuses it; clearing a linked key or field drops only the link. ``source_key`` cannot be cleared while other keys link to it, unless they are cleared in the same call. + ``source_key`` stays an ordinary sample: ``get_meta`` and dataloaders that ask only + for linked fields also return it, so give consumers a field the source lacks. + Args: keys: Single key or list of keys that will read the linked fields. partition_id: Partition holding both ``keys`` and ``source_key``. From 579e99c54bdab16d160d560632cb65eb42fb409a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:45:35 +0800 Subject: [PATCH 12/32] [fix] load controller checkpoints saved before kv_link --- tests/test_controller_data_partitions.py | 13 +++++++++++++ transfer_queue/controller.py | 4 ++++ 2 files changed, 17 insertions(+) diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index dc5abc01..79ba9c6f 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1522,6 +1522,19 @@ def test_source_index_follows_put_and_clear(self): assert partition.links == {} and partition.linked_by == {} partition.check_not_linked([0]) + def test_checkpoint_saved_before_links_still_loads(self): + import pickle + + from transfer_queue.controller import DataPartitionStatus + + partition = DataPartitionStatus(partition_id="old") + del partition.links, partition.linked_by + partition = pickle.loads(pickle.dumps(partition)) + + assert partition.update_production_status([0], ["x"], field_schema=self.SCHEMA) + partition.link_fields([1], 0, ["x"]) + assert partition.linked_by == {0: {1}} + def test_linked_rows_follow_source_readiness_and_shape(self): partition = self._linked_partition() wider = {"x": {**self.SCHEMA["x"], "shape": (5,)}} diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 5bb4987a..91d62651 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -371,6 +371,10 @@ class DataPartitionStatus: links: dict[int, dict[str, int]] = field(default_factory=dict) # global_idx -> {field: source global_idx} linked_by: dict[int, set[int]] = field(default_factory=dict) # source global_idx -> global_idxs linking to it + def __setstate__(self, state: dict[str, Any]) -> None: + # Controller checkpoints pickle partitions; ones saved before kv_link have no links. + self.__dict__.update({"links": {}, "linked_by": {}, **state}) + # Dynamic configuration - these are computed from the current state @property def total_samples_num(self) -> int: From bd4d9b80d9d88dde2b72f51112093933f31a7f46 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:47:12 +0800 Subject: [PATCH 13/32] [fix] refuse dumping a linked source field that is not finalized --- tests/e2e/test_kv_link_e2e.py | 5 +++++ transfer_queue/interface.py | 9 ++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 6f30e009..25b257b0 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -207,6 +207,11 @@ def test_dump_and_load_relink(tq_system, dump_test_root): with pytest.raises(ValueError, match="must be dumped"): tq.dump_data_by_key(dump_dir, keys, PARTITION) + tq.kv_update("p0", PARTITION, {"prompt": _prompt(0)}, merge_fn=lambda old, new: new, finalize=False) + with pytest.raises(ValueError, match="must be finalized"): + tq.dump_data_by_key(dump_dir, ["p0", *keys], PARTITION) + tq.kv_update("p0", PARTITION, {"prompt": _prompt(0)}, merge_fn=lambda old, new: new) + tq.dump_data_by_key(dump_dir, ["p0", *keys], PARTITION) tq.get_client().clear_partition(PARTITION) tq.load_data_by_key(dump_dir) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index c0c62b31..9d8ce138 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -1543,7 +1543,8 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - Raises: FileExistsError: ``dump_dir`` already exists. - ValueError: A key links to a source key that is not dumped. + ValueError: A key links to a source key that is not dumped, or to a source field + that is not finalized. RuntimeError: TransferQueue is not initialized, the partition or a key does not exist, or a storage unit holds no data for a row it was asked to dump. NotImplementedError: The storage backend does not support selective dumps. @@ -1567,6 +1568,12 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - missing_sources = sorted({source for row in rows.values() for source in row["links"].values()} - rows.keys()) if missing_sources: raise ValueError(f"linked source keys {missing_sources} must be dumped with the keys linking to them") + # Load restores unfinalized fields hidden, and kv_link refuses a hidden source field. + hidden_sources = sorted( + {s for row in rows.values() for name, s in row["links"].items() if name in rows[s]["unfinalized"]} + ) + if hidden_sources: + raise ValueError(f"fields that other keys link to on {hidden_sources} must be finalized before dumping") # A row whose fields are all still unproduced has nothing for a storage unit to # dump, but it keeps its key and tag so the restore can recreate the row. fields_by_index = {row["global_index"]: row["fields"] for row in rows.values() if row["fields"]} From 78852038807f992bfa035556cc563e38baeded07 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:47:47 +0800 Subject: [PATCH 14/32] [fix] BatchMeta.union takes links from the side whose fields win --- tests/test_metadata.py | 8 ++++++++ transfer_queue/metadata.py | 5 ++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 3433271c..907342b9 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -267,6 +267,14 @@ def test_add_fields_replaces_links_on_written_fields(self): assert batch._links == [{"field_b": 7}, {}] + def test_union_takes_links_with_the_overriding_fields(self): + linked = self._make_batch(batch_size=1) + linked._links = [{"field_a": 7, "field_b": 7}] + private = self._make_batch(batch_size=1, field_names=["field_a"]) + + assert linked.union(private)._links == [{"field_b": 7}] + assert private.union(linked)._links == [{"field_a": 7, "field_b": 7}] + def test_shallow_copy_isolation_global_indexes(self): """Modifying the original global_indexes list does not affect BatchMeta.""" original_indexes = [0, 1, 2] diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py index 690ea50e..14c86a79 100644 --- a/transfer_queue/metadata.py +++ b/transfer_queue/metadata.py @@ -717,7 +717,10 @@ def union(self, other: "BatchMeta") -> "BatchMeta": extra_info=merged_extra_info, custom_meta=merged_custom_meta, _custom_backend_meta=merged_custom_backend_meta, - _links=[{**mine, **theirs} for mine, theirs in zip(self._links, other._links, strict=True)], + _links=[ + {**{f: s for f, s in mine.items() if f not in other.field_schema}, **theirs} + for mine, theirs in zip(self._links, other._links, strict=True) + ], ) @classmethod From 96ea8da7d36514998d128be2e49d2c7c0db9e204 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 10:49:14 +0800 Subject: [PATCH 15/32] [feat] balance each partition's rows across SimpleStorage units --- tests/e2e/test_kv_link_e2e.py | 14 ++++ tests/test_controller_data_partitions.py | 32 ++++++++ transfer_queue/controller.py | 95 +++++++++++++++--------- 3 files changed, 106 insertions(+), 35 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 25b257b0..006fd6a3 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -76,6 +76,20 @@ def _rows(field: torch.Tensor) -> list[list]: return [row.tolist() for row in field.unbind()] +def test_each_partition_spreads_evenly_over_storage_units(tq_system): + """Two storage units hold global_idx % 2; alternating puts used to send each partition to one.""" + other = f"{PARTITION}_other" + for i in range(4): + tq.kv_put(key=f"a{i}", partition_id=PARTITION, fields={"prompt": _prompt(i)}) + tq.kv_put(key=f"b{i}", partition_id=other, fields={"prompt": _prompt(i)}) + + client = tq.get_client() + for partition_id, prefix in ((PARTITION, "a"), (other, "b")): + rows = client.kv_describe([f"{prefix}{i}" for i in range(4)], partition_id)["rows"] + assert sorted(row["global_index"] % 2 for row in rows.values()) == [0, 0, 1, 1] + client.clear_partition(other) + + def test_linked_keys_read_the_source_payload(tq_system): keys = _put_group(0) + _put_group(10) tq.kv_batch_put( diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index 79ba9c6f..f8cb6779 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1494,6 +1494,38 @@ def test_fieldmeta_to_batch_schema_nested_missing_sample(self): assert schema["per_sample_shapes"] == [(3,), None] +def test_index_manager_balances_each_partition_across_storage_units(): + """Alternating one-row allocations put every row of a partition on one unit before.""" + from transfer_queue.controller import PartitionIndexManager + + manager = PartitionIndexManager() + manager.set_num_storage_units(2) + for _ in range(4): + manager.allocate_indexes("__public__", 1) + manager.allocate_indexes("train", 1) + for partition_id in ("__public__", "train"): + assert sorted(idx % 2 for idx in manager.partition_to_indexes[partition_id]) == [0, 0, 1, 1] + + manager.release_indexes("train", [idx for idx in manager.partition_to_indexes["train"] if idx % 2 == 0]) + manager.allocate_indexes("val", 3) + assert sorted(idx % 2 for idx in manager.partition_to_indexes["val"]) == [0, 0, 1] + assert manager.unit_loads["train"] == [0, 2] + + reusable = manager.reusable_indexes + manager.set_num_storage_units(2, reusable) + assert sorted(manager.reusable_indexes) == sorted(reusable) + assert manager.unit_loads == {"__public__": [2, 2], "train": [0, 2], "val": [2, 1]} + + +def test_index_manager_with_one_unit_reuses_fifo(): + from transfer_queue.controller import PartitionIndexManager + + manager = PartitionIndexManager() + assert manager.allocate_indexes("p", 3) == [0, 1, 2] + manager.release_indexes("p", [2, 0]) + assert manager.allocate_indexes("q", 3) == [2, 0, 3] + + class TestFieldLinks: """Rows linking to a source row through DataPartitionStatus.link_fields.""" diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 91d62651..40b8ccc4 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -17,7 +17,7 @@ import os import pickle import time -from collections import defaultdict +from collections import defaultdict, deque from dataclasses import dataclass, field from itertools import groupby from operator import itemgetter @@ -75,18 +75,38 @@ def __init__(self): # Records the set of global_indexes used by each partition self.partition_to_indexes = defaultdict(set) - # Reusable global_index pool - stored using list - self.reusable_indexes = [] - # Global index counter for allocating new indexes self.global_index_counter = 0 # Track all active indexes self.allocated_indexes = set() + # SimpleStorage stores global_idx on unit global_idx % num_storage_units, so picking the + # index picks the unit. Reusable indexes are kept FIFO per unit. + self.set_num_storage_units(1, []) + + @property + def reusable_indexes(self) -> list[int]: + return [idx for pool in self.reusable_by_unit for idx in pool] + + def set_num_storage_units(self, num_storage_units: int, reusable_indexes: list[int] | None = None) -> None: + """Regroup the reusable indexes and per-partition unit loads for ``num_storage_units`` units.""" + reusable = self.reusable_indexes if reusable_indexes is None else reusable_indexes + self.num_storage_units = num_storage_units + self.reusable_by_unit: list[deque[int]] = [deque() for _ in range(num_storage_units)] + for idx in reusable: + self.reusable_by_unit[idx % num_storage_units].append(idx) + # partition_id -> rows of the partition on each unit + self.unit_loads: dict[str, list[int]] = {} + for partition_id, indexes in self.partition_to_indexes.items(): + loads = self.unit_loads[partition_id] = [0] * num_storage_units + for idx in indexes: + loads[idx % num_storage_units] += 1 + def allocate_indexes(self, partition_id, count=1) -> list[int]: """ - Allocate global_indexes for the specified partition. + Allocate global_indexes for the specified partition, each on the unit where the + partition has the fewest rows, so every partition stays balanced across units. Prioritizes obtaining from reusable pool, allocates new indexes when insufficient. Args: @@ -98,39 +118,40 @@ def allocate_indexes(self, partition_id, count=1) -> list[int]: """ if count <= 0: raise ValueError(f"Number of indexes needed must be larger than 0, but got {count}") + units = self.num_storage_units + loads = self.unit_loads.setdefault(partition_id, [0] * units) indexes = [] - - # Get indexes from reusable pool - if self.reusable_indexes: - # Calculate number of indexes needed from reusable pool - num_reuse = min(count, len(self.reusable_indexes)) - - # Use slice operation to get multiple elements at once (FIFO principle) - indexes.extend(self.reusable_indexes[:num_reuse]) - del self.reusable_indexes[:num_reuse] - - # If reusable pool doesn't have enough indexes, allocate new ones - if len(indexes) < count: - # Ensure newly allocated indexes don't conflict with existing ones - needed = count - len(indexes) - # Batch allocate consecutive index ranges - start_index = self.global_index_counter - end_index = start_index + needed - - # Directly generate consecutive index list - new_indexes = list(range(start_index, end_index)) - - # Batch update status - self.allocated_indexes.update(new_indexes) - self.global_index_counter = end_index - - indexes.extend(new_indexes) - + for _ in range(count): + # Among the least-loaded units, prefer reusing an index, then the counter's next one. + unit = min( + range(units), + key=lambda u: (loads[u], not self.reusable_by_unit[u], u != self.global_index_counter % units), + ) + if self.reusable_by_unit[unit]: + idx = self.reusable_by_unit[unit].popleft() + else: + # Indexes skipped to reach this unit wait in their own unit's pool. + while self.global_index_counter % units != unit: + self.reusable_by_unit[self.global_index_counter % units].append(self.global_index_counter) + self.global_index_counter += 1 + idx = self.global_index_counter + self.global_index_counter += 1 + loads[unit] += 1 + indexes.append(idx) + + self.allocated_indexes.update(indexes) # Record partition-index relationship self.partition_to_indexes[partition_id].update(indexes) return indexes + def _reuse(self, partition_id, indexes) -> None: + loads = self.unit_loads.get(partition_id) + for idx in indexes: + self.reusable_by_unit[idx % self.num_storage_units].append(idx) + if loads is not None: + loads[idx % self.num_storage_units] -= 1 + def release_partition(self, partition_id) -> list[int]: """ Release all global_indexes of the specified partition, adding them to reusable pool. @@ -145,7 +166,8 @@ def release_partition(self, partition_id) -> list[int]: indexes = self.partition_to_indexes.pop(partition_id) # Add released indexes to reusable pool - self.reusable_indexes.extend(indexes) + self._reuse(partition_id, indexes) + self.unit_loads.pop(partition_id, None) # Remove these indexes from allocated_indexes for idx in indexes: @@ -171,12 +193,13 @@ def release_indexes(self, partition_id: str, indexes_to_release: list[int]): raise ValueError("Some indexes to release do not belong to the specified partition.") partition_indexes.difference_update(indexes_to_release) - self.reusable_indexes.extend(indexes_to_release) + self._reuse(partition_id, indexes_to_release) self.allocated_indexes.difference_update(indexes_to_release) # If partition has no more indexes, remove it from the mapping if not partition_indexes: self.partition_to_indexes.pop(partition_id, None) + self.unit_loads.pop(partition_id, None) def get_indexes_for_partition(self, partition_id) -> list[int]: """ @@ -2513,6 +2536,8 @@ def get_zmq_server_info(self) -> ZMQServerInfo: def store_config(self, conf: DictConfig) -> None: """Store the global config of TransferQueue.""" self.tq_config = conf + if conf.backend.storage_backend == "SimpleStorage": + self.index_manager.set_num_storage_units(int(conf.backend.SimpleStorage.num_data_storage_units)) def get_config(self) -> DictConfig: """Retrieve the global config of TransferQueue.""" @@ -2566,9 +2591,9 @@ def load_checkpoint(self, path: str) -> None: im = state["index_manager"] self.index_manager.partition_to_indexes = defaultdict(set, im["partition_to_indexes"]) - self.index_manager.reusable_indexes = im["reusable_indexes"] self.index_manager.global_index_counter = im["global_index_counter"] self.index_manager.allocated_indexes = im["allocated_indexes"] + self.index_manager.set_num_storage_units(self.index_manager.num_storage_units, im["reusable_indexes"]) self.sampler.load_checkpoint(state["sampler"]) From ee979379b9ba1b325424cbb112ca51efcb464db6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 10:54:03 +0800 Subject: [PATCH 16/32] [feat] kv_link reads shared payloads from a __public__ partition --- tests/e2e/test_kv_link_e2e.py | 133 ++++++++++----------- tests/test_controller_data_partitions.py | 43 +++---- transfer_queue/__init__.py | 2 + transfer_queue/client.py | 12 +- transfer_queue/controller.py | 146 ++++++++++++----------- transfer_queue/interface.py | 95 ++++++++------- 6 files changed, 214 insertions(+), 217 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 006fd6a3..20391b8f 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""End-to-end tests for kv_link: n keys reading one prompt payload without copies.""" +"""End-to-end tests for kv_link: n keys reading one PUBLIC_PARTITION payload without copies.""" import shutil @@ -25,6 +25,7 @@ import transfer_queue as tq +PUBLIC = tq.PUBLIC_PARTITION PARTITION = "link" N = 3 @@ -50,18 +51,20 @@ def tq_system(): @pytest.fixture(autouse=True) -def fresh_partition(tq_system): +def fresh_partitions(tq_system): yield + # Linking keys first: a public key cannot be cleared while keys link to it. tq.get_client().clear_partition(PARTITION) + tq.get_client().clear_partition(PUBLIC) -def _prompt(p: int) -> torch.Tensor: - return torch.tensor([p, p + 1, p + 2]) +def _prompt(p: int, length: int = 3) -> torch.Tensor: + return torch.arange(p, p + length) def _put_group(p: int) -> list[str]: - """Put prompt ``p`` once and link its N responses to it; return the response keys.""" - tq.kv_put(key=f"p{p}", partition_id=PARTITION, fields={"prompt": _prompt(p)}) + """Put prompt ``p`` once into the public partition and link N response keys to it.""" + tq.kv_put(key=f"p{p}", partition_id=PUBLIC, fields={"prompt": _prompt(p)}) keys = [f"p{p}_r{i}" for i in range(N)] tq.kv_link(keys, PARTITION, f"p{p}", "prompt") return keys @@ -78,19 +81,17 @@ def _rows(field: torch.Tensor) -> list[list]: def test_each_partition_spreads_evenly_over_storage_units(tq_system): """Two storage units hold global_idx % 2; alternating puts used to send each partition to one.""" - other = f"{PARTITION}_other" for i in range(4): - tq.kv_put(key=f"a{i}", partition_id=PARTITION, fields={"prompt": _prompt(i)}) - tq.kv_put(key=f"b{i}", partition_id=other, fields={"prompt": _prompt(i)}) + tq.kv_put(key=f"p{i}", partition_id=PUBLIC, fields={"prompt": _prompt(i)}) + tq.kv_put(key=f"r{i}", partition_id=PARTITION, fields={"response": _prompt(i)}) client = tq.get_client() - for partition_id, prefix in ((PARTITION, "a"), (other, "b")): + for partition_id, prefix in ((PUBLIC, "p"), (PARTITION, "r")): rows = client.kv_describe([f"{prefix}{i}" for i in range(4)], partition_id)["rows"] assert sorted(row["global_index"] % 2 for row in rows.values()) == [0, 0, 1, 1] - client.clear_partition(other) -def test_linked_keys_read_the_source_payload(tq_system): +def test_linked_keys_read_the_public_payload(tq_system): keys = _put_group(0) + _put_group(10) tq.kv_batch_put( keys=keys, @@ -102,68 +103,49 @@ def test_linked_keys_read_the_source_payload(tq_system): assert _rows(data["prompt"]) == [_prompt(0).tolist()] * N + [_prompt(10).tolist()] * N assert _rows(data["response"]) == [[i] for i in range(len(keys))] + assert _links(keys[0]) == {"prompt": "p0"} -def test_data_parser_refuses_linked_fields(tq_system): - keys = _put_group(0) - tq.kv_batch_put(keys, PARTITION, fields=TensorDict({"response": torch.arange(N).unsqueeze(1)}, batch_size=N)) - - def identity(data, args): - return data - - with pytest.raises(NotImplementedError, match="kv_link"): - tq.kv_batch_get(keys, PARTITION, data_parser=identity) - data = tq.kv_batch_get(keys, PARTITION, select_fields="response", data_parser=identity) - assert _rows(data["response"]) == [[i] for i in range(N)] - - -def test_link_requires_a_produced_source_field(tq_system): - tq.kv_put(key="p0", partition_id=PARTITION, fields={"prompt": _prompt(0)}) +def test_link_requires_a_produced_public_field(tq_system): + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) with pytest.raises(RuntimeError, match="not found"): tq.kv_link("r", PARTITION, "missing", "prompt") with pytest.raises(RuntimeError, match="must have produced"): tq.kv_link("r", PARTITION, "p0", "image") + with pytest.raises(RuntimeError, match="cannot link"): + tq.kv_link("r", PUBLIC, "p0", "prompt") with pytest.raises(ValueError, match="duplicates"): tq.kv_link(["r", "r"], PARTITION, "p0", "prompt") - with pytest.raises(RuntimeError, match="cannot link to itself"): - tq.kv_link(["r", "p0"], PARTITION, "p0", "prompt") - assert tq.kv_list(PARTITION)[PARTITION].keys() == {"p0"} - - -def test_link_refuses_to_relink_a_source(tq_system): - keys = _put_group(0) - tq.kv_put(key="q0", partition_id=PARTITION, fields={"prompt": _prompt(9)}) - - with pytest.raises(RuntimeError, match="linked by other keys"): - tq.kv_link("p0", PARTITION, "q0", "prompt") + assert tq.kv_list().get(PARTITION, {}) == {} - assert _links("p0") == {} - assert _rows(tq.kv_batch_get(["p0", *keys], PARTITION)["prompt"]) == [_prompt(0).tolist()] * (N + 1) - -def test_link_refuses_a_source_being_cleared(tq_system): - tq.kv_put(key="p0", partition_id=PARTITION, fields={"prompt": _prompt(0)}) - source = tq.get_client().kv_describe(["p0"], PARTITION)["rows"]["p0"]["global_index"] +def test_link_refuses_a_public_key_being_cleared(tq_system): + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + source = tq.get_client().kv_describe(["p0"], PUBLIC)["rows"]["p0"]["global_index"] controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") - ray.get(controller.mark_clearing.remote([source], [PARTITION])) + ray.get(controller.mark_clearing.remote([source], [PUBLIC])) with pytest.raises(RuntimeError, match="must have produced"): tq.kv_link("r", PARTITION, "p0", "prompt") -def test_source_cannot_be_cleared_while_linked(tq_system): +def test_public_key_cannot_be_cleared_while_linked(tq_system): keys = _put_group(0) - with pytest.raises(RuntimeError, match="linked by other keys"): - tq.kv_clear("p0", PARTITION) - with pytest.raises(RuntimeError, match="linked by other keys"): - tq.kv_clear("p0", PARTITION, fields="prompt") + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + tq.kv_clear("p0", PUBLIC) + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + tq.kv_clear("p0", PUBLIC, fields="prompt") + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + tq.get_client().clear_partition(PUBLIC) assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + tq.kv_clear(keys[0], PARTITION) tq.kv_clear(keys[1:], PARTITION, fields="prompt") - tq.kv_clear([keys[0], "p0"], PARTITION) + tq.kv_clear("p0", PUBLIC) + assert tq.kv_list().get(PUBLIC, {}) == {} assert tq.kv_list(PARTITION)[PARTITION].keys() == set(keys[1:]) assert all(not _links(key) for key in keys[1:]) @@ -185,7 +167,20 @@ def test_put_through_meta_reads_back_the_written_value(tq_system): client.put(TensorDict({"prompt": torch.stack([_prompt(7)] * N)}, batch_size=N), metadata=meta) assert _rows(client.get_data(meta)["prompt"]) == [_prompt(7).tolist()] * N - assert _rows(tq.kv_batch_get("p0", PARTITION)["prompt"]) == [_prompt(0).tolist()] + assert _rows(tq.kv_batch_get("p0", PUBLIC)["prompt"]) == [_prompt(0).tolist()] + + +def test_data_parser_refuses_linked_fields(tq_system): + keys = _put_group(0) + tq.kv_batch_put(keys, PARTITION, fields=TensorDict({"response": torch.arange(N).unsqueeze(1)}, batch_size=N)) + + def identity(data, args): + return data + + with pytest.raises(NotImplementedError, match="kv_link"): + tq.kv_batch_get(keys, PARTITION, data_parser=identity) + data = tq.kv_batch_get(keys, PARTITION, select_fields="response", data_parser=identity) + assert _rows(data["response"]) == [[i] for i in range(N)] def test_update_refuses_a_linked_field(tq_system): @@ -203,33 +198,35 @@ def test_update_refuses_a_linked_field(tq_system): assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N -def test_unfinalized_source_update_is_hidden_from_linked_keys(tq_system): +def test_linked_keys_follow_public_readiness_and_shape(tq_system): keys = _put_group(0) - tq.kv_update("p0", PARTITION, {"prompt": _prompt(5)}, merge_fn=lambda old, new: new, finalize=False) + tq.kv_update("p0", PUBLIC, {"prompt": _prompt(5, 5)}, merge_fn=lambda old, new: new, finalize=False) with pytest.raises(ValueError, match="not ready"): tq.kv_batch_get(keys, PARTITION, select_fields="prompt") - tq.kv_update("p0", PARTITION, {"prompt": _prompt(5)}, merge_fn=lambda old, new: new) - assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(5).tolist()] * N + tq.kv_update("p0", PUBLIC, {"prompt": _prompt(5, 5)}, merge_fn=lambda old, new: new) + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(5, 5).tolist()] * N + meta = tq.get_client().kv_retrieve_meta(keys, PARTITION) + assert [tuple(shape) for shape in meta.get_shapes("prompt")] == [(5,)] * N def test_dump_and_load_relink(tq_system, dump_test_root): keys = _put_group(0) - dump_dir = dump_test_root / "link_dump" - - with pytest.raises(ValueError, match="must be dumped"): - tq.dump_data_by_key(dump_dir, keys, PARTITION) + public_dir, linked_dir = dump_test_root / "public_dump", dump_test_root / "link_dump" + tq.dump_data_by_key(public_dir, ["p0"], PUBLIC) + tq.dump_data_by_key(linked_dir, keys, PARTITION) + tq.get_client().clear_partition(PARTITION) + tq.get_client().clear_partition(PUBLIC) - tq.kv_update("p0", PARTITION, {"prompt": _prompt(0)}, merge_fn=lambda old, new: new, finalize=False) - with pytest.raises(ValueError, match="must be finalized"): - tq.dump_data_by_key(dump_dir, ["p0", *keys], PARTITION) - tq.kv_update("p0", PARTITION, {"prompt": _prompt(0)}, merge_fn=lambda old, new: new) + with pytest.raises(RuntimeError, match="does not exist|not found"): + tq.load_data_by_key(linked_dir) + assert tq.kv_list().get(PARTITION, {}) == {} - tq.dump_data_by_key(dump_dir, ["p0", *keys], PARTITION) - tq.get_client().clear_partition(PARTITION) - tq.load_data_by_key(dump_dir) - shutil.rmtree(dump_dir) + tq.load_data_by_key(public_dir) + tq.load_data_by_key(linked_dir) + shutil.rmtree(public_dir) + shutil.rmtree(linked_dir) assert all(_links(key) == {"prompt": "p0"} for key in keys) assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index f8cb6779..67e44568 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1527,32 +1527,33 @@ def test_index_manager_with_one_unit_reuses_fifo(): class TestFieldLinks: - """Rows linking to a source row through DataPartitionStatus.link_fields.""" + """Rows of a partition linking to a public row through DataPartitionStatus.link_fields.""" SCHEMA = {"x": {"dtype": "torch.int64", "shape": (3,), "is_nested": False, "is_non_tensor": False}} - def _linked_partition(self): - """Source row 0 holds ``x``; rows 1 and 2 link to it.""" - from transfer_queue.controller import DataPartitionStatus + def _public(self): + """Public row 0 holds ``x``.""" + from transfer_queue.controller import PUBLIC_PARTITION, DataPartitionStatus - partition = DataPartitionStatus(partition_id="links") - assert partition.update_production_status([0], ["x"], field_schema=self.SCHEMA) - partition.link_fields([1, 2], 0, ["x"]) - return partition + public = DataPartitionStatus(partition_id=PUBLIC_PARTITION) + assert public.update_production_status([0], ["x"], field_schema=self.SCHEMA) + return public - def test_source_index_follows_put_and_clear(self): - import pytest + def test_source_index_follows_link_put_and_clear(self): + from transfer_queue.controller import DataPartitionStatus - partition = self._linked_partition() + public = self._public() + partition = DataPartitionStatus(partition_id="train") + partition.link_fields([1, 2], public, 0, ["x"]) assert partition.linked_by == {0: {1, 2}} - with pytest.raises(ValueError, match="linked by other keys"): - partition.check_not_linked([0, 1]) + assert partition.field_metadata["x"].shape == (3,) and partition.production_status[[1, 2], 0].tolist() == [1, 1] + assert partition.linked_sources([0], ["x"]) == {0} and partition.linked_sources([0], ["y"]) == set() assert partition.update_production_status([1], ["x"], field_schema=self.SCHEMA) partition.clear_data([2]) assert partition.links == {} and partition.linked_by == {} - partition.check_not_linked([0]) + assert partition.linked_sources([0]) == set() def test_checkpoint_saved_before_links_still_loads(self): import pickle @@ -1563,17 +1564,5 @@ def test_checkpoint_saved_before_links_still_loads(self): del partition.links, partition.linked_by partition = pickle.loads(pickle.dumps(partition)) - assert partition.update_production_status([0], ["x"], field_schema=self.SCHEMA) - partition.link_fields([1], 0, ["x"]) + partition.link_fields([1], self._public(), 0, ["x"]) assert partition.linked_by == {0: {1}} - - def test_linked_rows_follow_source_readiness_and_shape(self): - partition = self._linked_partition() - wider = {"x": {**self.SCHEMA["x"], "shape": (5,)}} - - assert partition.update_production_status([0], ["x"], field_schema=wider, ready=False) - assert partition.production_status[[0, 1, 2], 0].tolist() == [0, 0, 0] - assert partition.field_metadata["x"].per_sample_shapes == {0: (5,), 1: (5,), 2: (5,)} - - assert partition.update_production_status([0], ["x"], field_schema=wider) - assert partition.production_status[[0, 1, 2], 0].tolist() == [1, 1, 1] diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 364c1279..bbbd8284 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -16,6 +16,7 @@ import os from .client import TransferQueueClient +from .controller import PUBLIC_PARTITION from .dataloader import StreamingDataLoader, StreamingDataset from .interface import ( async_kv_batch_get, @@ -91,6 +92,7 @@ "async_kv_link", "async_kv_update", "KVBatchMeta", + "PUBLIC_PARTITION", ] + [ # KV key locks: process-local and cluster-wide diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d3bcbbd0..b49f64ff 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1191,16 +1191,16 @@ async def async_kv_link( self, keys: list[str], partition_id: str, - source_key: str, + public_key: str, fields: list[str], socket: zmq.asyncio.Socket | None = None, ) -> None: - """Asynchronously make ``fields`` of ``keys`` read the payload of ``source_key``; see ``kv_link``.""" + """Asynchronously make ``fields`` of ``keys`` read the payload of ``public_key``; see ``kv_link``.""" await self._request_controller( socket=socket, request_type=ZMQRequestType.KV_LINK, response_type=ZMQRequestType.KV_LINK_RESPONSE, - body={"keys": keys, "partition_id": partition_id, "source_key": source_key, "fields": fields}, + body={"keys": keys, "partition_id": partition_id, "public_key": public_key, "fields": fields}, ) def close(self) -> None: @@ -1996,9 +1996,9 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: """ return self._kv_describe(keys, partition_id) - def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: list[str]) -> None: - """Synchronously make ``fields`` of ``keys`` read the payload of ``source_key``; see ``kv_link``.""" - return self._kv_link(keys, partition_id, source_key, fields) + def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: list[str]) -> None: + """Synchronously make ``fields`` of ``keys`` read the payload of ``public_key``; see ``kv_link``.""" + return self._kv_link(keys, partition_id, public_key, fields) # ==================== Selective Data Dump API ==================== diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 40b8ccc4..3704e052 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -60,6 +60,10 @@ # storage units, so the budget covers a fleet rather than the two ROUTERs bound here. TQ_CONTROLLER_ZMQ_MAX_SOCKETS = int(os.environ.get("TQ_CONTROLLER_ZMQ_MAX_SOCKETS", 4096)) +# Holds payloads that keys of other partitions read through kv_link, e.g. a prompt video +# shared by its n responses. Its rows are balanced over storage units like any partition's. +PUBLIC_PARTITION = "__public__" + # Sample pre-allocation for StreamingDataLoader compatibility. # By pre-allocating sample indices (typically global_batch_size), consumers can accurately # determine consumption status even before producers have generated the samples. @@ -390,9 +394,9 @@ class DataPartitionStatus: keys_mapping: dict[str, int] = field(default_factory=dict) # key -> global_idx revert_keys_mapping: dict[int, str] = field(default_factory=dict) # global_idx -> key - # Linked fields hold no payload of their own and read the source row's instead - links: dict[int, dict[str, int]] = field(default_factory=dict) # global_idx -> {field: source global_idx} - linked_by: dict[int, set[int]] = field(default_factory=dict) # source global_idx -> global_idxs linking to it + # Linked fields hold no payload of their own and read a PUBLIC_PARTITION row's instead + links: dict[int, dict[str, int]] = field(default_factory=dict) # global_idx -> {field: public global_idx} + linked_by: dict[int, set[int]] = field(default_factory=dict) # public global_idx -> global_idxs linking to it def __setstate__(self, state: dict[str, Any]) -> None: # Controller checkpoints pickle partitions; ones saved before kv_link have no links. @@ -599,17 +603,6 @@ def update_production_status( # A written value replaces the link it lands on self._drop_links(global_indices, field_names) - # Rows linking to a written cell read its new payload, so they take on its readiness - # (an unfinalized kv_update stays hidden behind links too) and its shape. - for source in global_indices: - for idx in self.linked_by.get(source, ()): - for name, linked_source in self.links[idx].items(): - if linked_source == source and name in field_names: - self.production_status[idx, self.field_name_mapping[name]] = int(ready) - meta = self.field_metadata[name] - if meta.is_nested: - meta.per_sample_shapes[idx] = meta.per_sample_shapes[source] - return True except Exception as e: @@ -1019,46 +1012,45 @@ def clear_fields(self, global_indexes: list[int], field_names: list[str]): self._drop_links(global_indexes, field_names) # ==================== Field Links ==================== - def link_fields(self, global_indexes: list[int], source: int, field_names: list[str]) -> None: - """Mark ``field_names`` produced on these rows, read from the payload of row ``source``.""" + def link_fields( + self, + global_indexes: list[int], + public: "DataPartitionStatus", + source: int, + field_names: list[str], + ready: bool = True, + ) -> None: + """Mark ``field_names`` of these rows as read from row ``source`` of the public partition. + + The rows take the source row's dtype and shape, and ``ready`` as their readiness. + """ schema = {} for name in field_names: - meta = self.field_metadata[name] + meta = public.field_metadata[name] + shape = meta.per_sample_shapes.get(source) if meta.is_nested else meta.shape schema[name] = { "dtype": meta.dtype, - "shape": meta.shape, - "is_nested": meta.is_nested, + "shape": shape, + "is_nested": False, "is_non_tensor": meta.is_non_tensor, } - if meta.is_nested: - schema[name]["per_sample_shapes"] = {idx: meta.per_sample_shapes[source] for idx in global_indexes} # Before recording links: update_production_status drops links on the cells it marks. - if not self.update_production_status(global_indexes, field_names, schema): + if not self.update_production_status(global_indexes, field_names, schema, ready=ready): raise RuntimeError(f"failed to link fields {field_names} in partition {self.partition_id}") - # Point at the row holding the payload, so a chain of links is one hop. - source_links = self.links.get(source, {}) for idx in global_indexes: row_links = self.links.setdefault(idx, {}) for name in field_names: - row_links[name] = source_links.get(name, source) - self.linked_by.setdefault(row_links[name], set()).add(idx) + row_links[name] = source + self.linked_by.setdefault(source, set()).update(global_indexes) - def check_not_linked(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: - """Refuse to clear payloads that rows outside this clear still read through links.""" - clearing = set(global_indexes) - sources = { + def linked_sources(self, sources: list[int], field_names: list[str] | None = None) -> set[int]: + """The public rows among ``sources`` that rows here link ``field_names`` (all if None) to.""" + return { source - for source in clearing + for source in sources for idx in self.linked_by.get(source, ()) - if idx not in clearing - and any(s == source and (field_names is None or n in field_names) for n, s in self.links[idx].items()) + if field_names is None or any(self.links[idx].get(name) == source for name in field_names) } - if sources: - keys = sorted(str(self.revert_keys_mapping.get(idx, idx)) for idx in sources) - raise ValueError( - f"keys {keys} in partition {self.partition_id!r} are linked by other keys; " - "clear the linked keys first or in the same call" - ) def _drop_links(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: """Drop the links of these rows on ``field_names`` (all fields if None).""" @@ -1276,6 +1268,15 @@ def update_production_status( success = partition.update_production_status( global_indexes, field_names, field_schema, custom_backend_meta, ready=ready ) + if success and partition_id == PUBLIC_PARTITION: + # Linking rows read the rewritten payload, so they take on its readiness (an + # unfinalized kv_update stays hidden behind links too) and its shape. + for other in self.partitions.values(): + for source in set(global_indexes) & other.linked_by.keys(): + for idx in list(other.linked_by[source]): + names = [name for name in field_names if other.links[idx].get(name) == source] + if names: + other.link_fields([idx], partition, source, names, ready=ready) if success: logger.debug( f"[{self.controller_id}]: Updated production status for partition {partition_id}: " @@ -1652,11 +1653,8 @@ def mark_clearing(self, global_indexes: list[int], partition_ids: list[str]) -> combined = list(zip(partition_ids, global_indexes, strict=True)) combined.sort(key=itemgetter(0)) - # Validate every partition first so a refused clear marks nothing. - for partition_id, group in groupby(combined, key=itemgetter(0)): - partition = self._get_partition(partition_id) - if partition: - partition.check_not_linked([idx for _, idx in group]) + # Validate first so a refused clear marks nothing. + self._check_public_not_linked([idx for partition_id, idx in combined if partition_id == PUBLIC_PARTITION]) for partition_id, group in groupby(combined, key=itemgetter(0)): partition = self._get_partition(partition_id) @@ -1671,6 +1669,17 @@ def mark_clearing(self, global_indexes: list[int], partition_ids: list[str]) -> if existing and partition.production_status is not None: partition.production_status[existing, :] = 0 + def _check_public_not_linked(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: + """Refuse to clear public payloads that keys of other partitions still link to.""" + linked = set().union(*(p.linked_sources(global_indexes, field_names) for p in self.partitions.values())) + if linked: + public = self.partitions[PUBLIC_PARTITION] + keys = sorted(str(public.revert_keys_mapping.get(idx, idx)) for idx in linked) + raise ValueError( + f"keys {keys} in partition {PUBLIC_PARTITION!r} are linked by keys of other partitions; " + "clear those keys first" + ) + def clear_partition(self, partition_id: str, clear_consumption: bool = True): """ Clear data for a specific partition (delete the whole partition). @@ -1911,7 +1920,8 @@ def kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) fields = [f for f in fields if f in partition.field_metadata] if not global_indexes or not fields: return BatchMeta.empty() - partition.check_not_linked(global_indexes, fields) + if partition_id == PUBLIC_PARTITION: + self._check_public_not_linked(global_indexes, fields) metadata = self.generate_batch_meta(partition_id, global_indexes, fields, mode="force_fetch") partition.clear_fields(global_indexes, fields) @@ -1969,7 +1979,7 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` lists the row's stored fields and is empty for a row with no produced field yet; ``unfinalized`` lists the stored fields not yet ready (written by ``kv_update(finalize=False)``); ``links`` maps each linked field to its - source key; ``field_schema`` covers the described rows' stored fields. + key in ``PUBLIC_PARTITION``; ``field_schema`` covers the described rows' stored fields. Raises: KeyError: The partition or any key does not exist. @@ -1983,6 +1993,7 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: if missing_keys: raise KeyError(f"keys not found in partition {partition_id!r}: {missing_keys}") + public_keys = self.partitions[PUBLIC_PARTITION].revert_keys_mapping if partition.links else {} rows: dict[str, dict[str, Any]] = {} for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True): links = partition.links.get(global_index, {}) @@ -1999,7 +2010,7 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: for name in fields if not partition.production_status[global_index, partition.field_name_mapping[name]] ], - "links": {name: partition.revert_keys_mapping[source] for name, source in links.items()}, + "links": {name: public_keys[source] for name, source in links.items()}, "tag": partition.custom_meta.get(global_index, {}), } # Only the declared types: row dtypes and shapes live with the storage units. @@ -2010,36 +2021,35 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: } return {"rows": rows, "field_schema": field_schema} - def kv_link(self, keys: list[str], partition_id: str, source_key: str, fields: list[str]) -> None: - """Make ``fields`` of ``keys`` read the payload of ``source_key`` without copying it. + def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: list[str]) -> None: + """Make ``fields`` of ``keys`` read the payload of ``public_key`` in ``PUBLIC_PARTITION``. - Missing ``keys`` are created. The source must already have produced ``fields``. + Missing ``keys`` are created. ``public_key`` must already have produced ``fields``. Raises: - KeyError: The partition or ``source_key`` does not exist. - ValueError: ``fields`` is empty, ``source_key`` has not produced all of them, or - ``keys`` include ``source_key`` or a key other keys link these fields to. - """ - partition = self._get_partition(partition_id) - if partition is None or source_key not in partition.keys_mapping: - raise KeyError(f"source key {source_key!r} not found in partition {partition_id!r}") - source = partition.keys_mapping[source_key] + KeyError: ``public_key`` does not exist. + ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, or ``fields`` is empty or not + all produced by ``public_key``. + """ + if partition_id == PUBLIC_PARTITION: + raise ValueError(f"keys of partition {PUBLIC_PARTITION!r} hold shared payloads and cannot link") + public = self._get_partition(PUBLIC_PARTITION) + if public is None or public_key not in public.keys_mapping: + raise KeyError(f"key {public_key!r} not found in partition {PUBLIC_PARTITION!r}") + source = public.keys_mapping[public_key] # Production status, not field metadata: mark_clearing zeroes it before storage deletes the payload. missing = [ name for name in fields - if name not in partition.field_name_mapping - or partition.production_status[source, partition.field_name_mapping[name]] != 1 + if name not in public.field_name_mapping + or public.production_status[source, public.field_name_mapping[name]] != 1 ] if not fields or missing: - raise ValueError(f"source key {source_key!r} must have produced fields {missing or fields}") - existing = [idx for idx in partition.kv_retrieve_indexes(keys) if idx is not None] - if source in existing: - raise ValueError(f"source key {source_key!r} cannot link to itself") - # Keys linking to a relinked key would keep reading its old payload, not the new source. - partition.check_not_linked(existing, fields) + raise ValueError(f"public key {public_key!r} must have produced fields {missing or fields}") metadata = self.kv_retrieve_meta(keys, partition_id, create=True) - partition.link_fields(metadata.global_indexes, source, fields) + partition = self._get_partition(partition_id) + assert partition is not None + partition.link_fields(metadata.global_indexes, public, source, fields) def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" @@ -2476,7 +2486,7 @@ def _handle_kv_prepare_update_request(self, request_msg: ZMQMessage) -> ZMQMessa def _handle_kv_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: params = request_msg.body - self.kv_link(params["keys"], params["partition_id"], params["source_key"], params["fields"]) + self.kv_link(params["keys"], params["partition_id"], params["public_key"], params["fields"]) return self._make_response(request_msg, ZMQRequestType.KV_LINK_RESPONSE, {}) def _handle_kv_clear_fields_request(self, request_msg: ZMQMessage) -> ZMQMessage: diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 9d8ce138..8520713b 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -47,7 +47,7 @@ from transfer_queue import data_dump from transfer_queue.client import TransferQueueClient -from transfer_queue.controller import TransferQueueController +from transfer_queue.controller import PUBLIC_PARTITION, TransferQueueController from transfer_queue.kv_lock import _close_lock_managers, _lock_settings, _start_lock_managers from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 @@ -696,37 +696,35 @@ def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id, fields=fields)) -def kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: list[str] | str) -> None: - """Let ``keys`` read ``fields`` from ``source_key`` without storing another copy. +def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: list[str] | str) -> None: + """Let ``keys`` read ``fields`` from ``public_key`` in ``PUBLIC_PARTITION`` without another copy. - For example, the n responses of a GRPO prompt can link the prompt's fields instead - of putting the prompt n times. Missing ``keys`` are created, and linked fields read - the source's current value and are ready whenever the source's are. Putting a linked - field stores a private value in place of the link, and ``kv_update`` refuses it; - clearing a linked key or field drops only the link. ``source_key`` cannot be cleared - while other keys link to it, unless they are cleared in the same call. - - ``source_key`` stays an ordinary sample: ``get_meta`` and dataloaders that ask only - for linked fields also return it, so give consumers a field the source lacks. + Put a payload shared by many keys, e.g. the video of a GRPO prompt read by its n + responses, once into ``tq.PUBLIC_PARTITION`` and link the keys to it. Its rows are + spread evenly over the storage units, apart from the rows of every other partition. + Missing ``keys`` are created, and linked fields read the public key's current value + and are ready whenever its fields are. Putting a linked field stores a private value + in place of the link, and ``kv_update`` refuses it; clearing a linked key or field + drops only the link. ``public_key`` cannot be cleared while keys link to it. Args: keys: Single key or list of keys that will read the linked fields. - partition_id: Partition holding both ``keys`` and ``source_key``. - source_key: Existing key that has produced ``fields``. + partition_id: Partition holding ``keys``; not ``PUBLIC_PARTITION``. + public_key: Key in ``PUBLIC_PARTITION`` that has produced ``fields``. fields: Single field name or list of field names to link. Raises: - RuntimeError: If ``source_key`` does not exist or has not produced all of ``fields``, - or ``keys`` include ``source_key`` or a key that other keys link these fields to. + RuntimeError: If ``public_key`` does not exist or has not produced all of ``fields``, + or ``partition_id`` is ``PUBLIC_PARTITION``. ValueError: If ``keys`` contains duplicates. NotImplementedError: If the storage backend is not SimpleStorage. Example: - >>> tq.kv_put(key="prompt_0", partition_id="train", fields={"input_ids": prompt_ids}) - >>> tq.kv_link([f"prompt_0_{i}" for i in range(n)], "train", "prompt_0", "input_ids") + >>> tq.kv_put(key="prompt_0", partition_id=tq.PUBLIC_PARTITION, fields={"video": video}) + >>> tq.kv_link([f"prompt_0_{i}" for i in range(n)], "train", "prompt_0", "video") """ tq_client = _maybe_create_tq_client() - tq_client._run_coroutine(async_kv_link(keys, partition_id, source_key, fields)) + tq_client._run_coroutine(async_kv_link(keys, partition_id, public_key, fields)) def kv_update( @@ -1203,7 +1201,7 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st Returns: ``{key: {"fields", "unfinalized", "links"}}``. ``fields`` lists every stored field of the key and ``unfinalized`` the subset that is not ready yet. ``links`` maps each field - linked by ``kv_link`` to its source key; such a field is ready when the source's is. + linked by ``kv_link`` to its key in ``PUBLIC_PARTITION``, whose readiness it shares. Raises: RuntimeError: If the partition or any key does not exist. @@ -1261,20 +1259,15 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[ await tq_client.async_clear_samples(batch_meta) -async def async_kv_link(keys: list[str] | str, partition_id: str, source_key: str, fields: list[str] | str) -> None: - """Asynchronously let ``keys`` read ``fields`` from ``source_key``. See ``kv_link``.""" +async def async_kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: list[str] | str) -> None: + """Asynchronously let ``keys`` read ``fields`` from ``public_key``. See ``kv_link``.""" tq_client = _maybe_create_tq_client() # Only SimpleStorage resolves links on read; other backends would read the linking key's own cells. if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): raise NotImplementedError(f"{type(tq_client.storage_manager).__name__} does not support kv_link") keys = [keys] if isinstance(keys, str) else keys _check_unique_keys(keys) - await tq_client.async_kv_link( - keys, - partition_id, - source_key, - [fields] if isinstance(fields, str) else fields, - ) + await tq_client.async_kv_link(keys, partition_id, public_key, [fields] if isinstance(fields, str) else fields) def _single_update_batch(fields: dict[str, Any]) -> TensorDict: @@ -1534,8 +1527,9 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - Args: dump_dir: Directory to write the dump into. - keys: Keys to dump. Duplicates are dropped, first occurrence wins. Must include - the source key of every field linked by ``kv_link``. + keys: Keys to dump. Duplicates are dropped, first occurrence wins. Fields linked + by ``kv_link`` are dumped as links; dump their ``PUBLIC_PARTITION`` keys + separately and load that dump first. partition_id: Partition that owns ``keys``. Returns: @@ -1543,8 +1537,6 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - Raises: FileExistsError: ``dump_dir`` already exists. - ValueError: A key links to a source key that is not dumped, or to a source field - that is not finalized. RuntimeError: TransferQueue is not initialized, the partition or a key does not exist, or a storage unit holds no data for a row it was asked to dump. NotImplementedError: The storage backend does not support selective dumps. @@ -1565,15 +1557,6 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - description = client.kv_describe(unique_keys, partition_id) if unique_keys else {"rows": {}, "field_schema": {}} row_index = {"partition_id": partition_id, **description} rows = row_index["rows"] - missing_sources = sorted({source for row in rows.values() for source in row["links"].values()} - rows.keys()) - if missing_sources: - raise ValueError(f"linked source keys {missing_sources} must be dumped with the keys linking to them") - # Load restores unfinalized fields hidden, and kv_link refuses a hidden source field. - hidden_sources = sorted( - {s for row in rows.values() for name, s in row["links"].items() if name in rows[s]["unfinalized"]} - ) - if hidden_sources: - raise ValueError(f"fields that other keys link to on {hidden_sources} must be finalized before dumping") # A row whose fields are all still unproduced has nothing for a storage unit to # dump, but it keeps its key and tag so the restore can recreate the row. fields_by_index = {row["global_index"]: row["fields"] for row in rows.values() if row["fields"]} @@ -1614,10 +1597,12 @@ def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: ``{key: tag}`` for every restored key. Raises: - RuntimeError: TransferQueue is not initialized or a storage unit fails. + RuntimeError: TransferQueue is not initialized, a storage unit fails, or a key + linked to ``PUBLIC_PARTITION`` names a public key that does not exist. NotImplementedError: The storage backend is not SimpleStorage. FileNotFoundError: The dump is incomplete. - ValueError: The manifest or a row disagrees with the row index. + ValueError: The manifest or a row disagrees with the row index, or a linked public + field is not ready. """ if _TQ_CONTROLLER is None: raise RuntimeError("TransferQueue is not initialized. Call tq.init() first.") @@ -1632,14 +1617,28 @@ def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: partition_id = row_index["partition_id"] rows = row_index["rows"] client.validate_dump_schema(partition_id, row_index["field_schema"]) + # Checked before loading anything: relinking needs the public fields, restored by their own load. + public_keys = sorted({key for row in rows.values() for key in row.get("links", {}).values()}) + if public_keys: + public_rows = client.kv_describe(public_keys, PUBLIC_PARTITION)["rows"] + unready = sorted( + { + key + for row in rows.values() + for name, key in row.get("links", {}).items() + if name not in public_rows[key]["fields"] or name in public_rows[key]["unfinalized"] + } + ) + if unready: + raise ValueError(f"linked fields of {PUBLIC_PARTITION!r} keys {unready} are not ready; load them first") bytes_read = client.load_rows_by_key(partition_id, rows, shards) - # Relink once the sources hold their payloads, one call per source and field set. + # One kv_link per public key and field set. link_groups: dict[tuple[str, tuple[str, ...]], list[str]] = defaultdict(list) for key, row in rows.items(): links = row.get("links", {}) - for source in set(links.values()): - link_groups[source, tuple(sorted(name for name, s in links.items() if s == source))].append(key) - for (source, names), keys in link_groups.items(): - client.kv_link(keys, partition_id, source, list(names)) + for public_key in set(links.values()): + link_groups[public_key, tuple(sorted(name for name, k in links.items() if k == public_key))].append(key) + for (public_key, names), keys in link_groups.items(): + client.kv_link(keys, partition_id, public_key, list(names)) logger.info(f"Restored {len(rows)} keys into partition {partition_id} from {dump_dir}, reading {bytes_read} bytes") return {key: row["tag"] for key, row in rows.items()} From 31e606a8fcd0af4e5632a8fe2edc03023817c7b1 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 11:41:43 +0800 Subject: [PATCH 17/32] feat: add kv_clear_unlinked_public to free public keys no key links to --- tests/e2e/test_kv_link_e2e.py | 24 ++++++++++++++++++++++++ transfer_queue/__init__.py | 4 ++++ transfer_queue/client.py | 27 +++++++++++++++++++++++++++ transfer_queue/controller.py | 27 +++++++++++++++++++++++++++ transfer_queue/interface.py | 22 ++++++++++++++++++++++ transfer_queue/utils/zmq_utils.py | 2 ++ 6 files changed, 106 insertions(+) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 20391b8f..eafd8531 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -150,6 +150,30 @@ def test_public_key_cannot_be_cleared_while_linked(tq_system): assert all(not _links(key) for key in keys[1:]) +def test_clear_unlinked_public_frees_only_unlinked_keys(tq_system): + keys0, keys1 = _put_group(0), _put_group(1) + tq.kv_put(key="p2", partition_id=PUBLIC, fields={"prompt": _prompt(2)}) + tq.kv_clear(keys1[:2], PARTITION) + tq.kv_clear(keys1[2:], PARTITION, fields="prompt") + + assert tq.kv_clear_unlinked_public() == ["p1", "p2"] + + assert tq.kv_list(PUBLIC)[PUBLIC].keys() == {"p0"} + with pytest.raises(ValueError, match="not found"): + tq.kv_batch_get("p1", PUBLIC) + assert _rows(tq.kv_batch_get(keys0, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + assert tq.kv_clear_unlinked_public() == [] + + +def test_link_refuses_a_key_marked_by_clear_unlinked_public(tq_system): + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + ray.get(controller.mark_unlinked_public_clearing.remote()) + + with pytest.raises(RuntimeError, match="must have produced"): + tq.kv_link("r", PARTITION, "p0", "prompt") + + def test_put_replaces_the_link(tq_system): keys = _put_group(0) diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index bbbd8284..722e8ba9 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -24,6 +24,7 @@ async_kv_batch_put, async_kv_batch_update, async_kv_clear, + async_kv_clear_unlinked_public, async_kv_describe, async_kv_link, async_kv_list, @@ -39,6 +40,7 @@ kv_batch_put, kv_batch_update, kv_clear, + kv_clear_unlinked_public, kv_describe, kv_link, kv_list, @@ -80,6 +82,7 @@ "kv_describe", "kv_clear", "kv_link", + "kv_clear_unlinked_public", "kv_update", "async_kv_put", "async_kv_batch_put", @@ -90,6 +93,7 @@ "async_kv_describe", "async_kv_clear", "async_kv_link", + "async_kv_clear_unlinked_public", "async_kv_update", "KVBatchMeta", "PUBLIC_PARTITION", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index b49f64ff..4461457b 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1203,6 +1203,28 @@ async def async_kv_link( body={"keys": keys, "partition_id": partition_id, "public_key": public_key, "fields": fields}, ) + async def async_clear_unlinked_public(self) -> list[str]: + """Asynchronously clear every ``PUBLIC_PARTITION`` key no key links to; see ``kv_clear_unlinked_public``. + + Returns: + The cleared keys. + """ + metadata, keys = await self._mark_unlinked_public_clearing() + if metadata.size > 0: + await self.storage_manager.clear_data(metadata) + await self._clear_meta_in_controller(metadata) + return keys + + @with_controller_socket + async def _mark_unlinked_public_clearing(self, socket=None) -> tuple[BatchMeta, list[str]]: + response = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.MARK_UNLINKED_PUBLIC_CLEARING, + response_type=ZMQRequestType.MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE, + body={}, + ) + return response.body["metadata"], response.body["keys"] + def close(self) -> None: """Close the client and cleanup resources including storage manager. @@ -1519,6 +1541,7 @@ def wrapper(*args, **kwargs): self._kv_list = _make_sync(self.async_kv_list) self._kv_describe = _make_sync(self.async_kv_describe) self._kv_link = _make_sync(self.async_kv_link) + self._clear_unlinked_public = _make_sync(self.async_clear_unlinked_public) self._validate_dump_schema = _make_sync(self.async_validate_dump_schema) self._dump_rows_by_index = _make_sync(self.async_dump_rows_by_index) self._load_rows_by_key = _make_sync(self.async_load_rows_by_key) @@ -2000,6 +2023,10 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l """Synchronously make ``fields`` of ``keys`` read the payload of ``public_key``; see ``kv_link``.""" return self._kv_link(keys, partition_id, public_key, fields) + def clear_unlinked_public(self) -> list[str]: + """Synchronously clear every ``PUBLIC_PARTITION`` key no key links to; return the cleared keys.""" + return self._clear_unlinked_public() + # ==================== Selective Data Dump API ==================== def validate_dump_schema(self, partition_id: str, field_schema: dict) -> None: diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 3704e052..6521ed8a 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2051,6 +2051,26 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l assert partition is not None partition.link_fields(metadata.global_indexes, public, source, fields) + def mark_unlinked_public_clearing(self) -> tuple[BatchMeta, list[str]]: + """Mark every ``PUBLIC_PARTITION`` key that no key links to as clearing. + + Selecting and marking in one controller call means a concurrent kv_link cannot link + a row this clear deletes: kv_link refuses fields whose production status was zeroed. + + Returns: + The marked rows' BatchMeta, for storage deletion and clear_meta, and their keys. + """ + public = self._get_partition(PUBLIC_PARTITION) + if public is None: + return BatchMeta.empty(), [] + linked = set().union(*(p.linked_by.keys() for p in self.partitions.values())) + keys = sorted(key for key, idx in public.keys_mapping.items() if idx not in linked) + if not keys: + return BatchMeta.empty(), [] + metadata = self.kv_retrieve_meta(keys, PUBLIC_PARTITION) + self.mark_clearing(metadata.global_indexes, metadata.partition_ids) + return metadata, keys + def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" self.zmq_context = zmq.Context() @@ -2279,6 +2299,7 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | ZMQRequestType.VALIDATE_DUMP_SCHEMA: self._handle_validate_dump_schema_request, ZMQRequestType.KV_CLEAR_FIELDS: self._handle_kv_clear_fields_request, ZMQRequestType.KV_LINK: self._handle_kv_link_request, + ZMQRequestType.MARK_UNLINKED_PUBLIC_CLEARING: self._handle_mark_unlinked_public_clearing_request, ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request, ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: self._handle_load_controller_checkpoint_request, } @@ -2489,6 +2510,12 @@ def _handle_kv_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: self.kv_link(params["keys"], params["partition_id"], params["public_key"], params["fields"]) return self._make_response(request_msg, ZMQRequestType.KV_LINK_RESPONSE, {}) + def _handle_mark_unlinked_public_clearing_request(self, request_msg: ZMQMessage) -> ZMQMessage: + metadata, keys = self.mark_unlinked_public_clearing() + return self._make_response( + request_msg, ZMQRequestType.MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE, {"metadata": metadata, "keys": keys} + ) + def _handle_kv_clear_fields_request(self, request_msg: ZMQMessage) -> ZMQMessage: params = request_msg.body metadata = self.kv_clear_fields( diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 8520713b..63b4b0d4 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -727,6 +727,22 @@ def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: l tq_client._run_coroutine(async_kv_link(keys, partition_id, public_key, fields)) +def kv_clear_unlinked_public() -> list[str]: + """Clear every key in ``PUBLIC_PARTITION`` that no key links to, and return those keys. + + For example, after clearing a step's trajectories, this frees the prompt videos no + trajectory reads any more. A public key that is put but not linked yet is cleared too, + so call it when no writer is between its ``kv_put`` and ``kv_link``. A ``kv_link`` + that races with it is refused for the keys it clears. + + Example: + >>> tq.kv_clear(trajectory_keys, partition_id="train") + >>> freed = tq.kv_clear_unlinked_public() + """ + tq_client = _maybe_create_tq_client() + return tq_client._run_coroutine(async_kv_clear_unlinked_public()) + + def kv_update( key: str, partition_id: str, @@ -1270,6 +1286,12 @@ async def async_kv_link(keys: list[str] | str, partition_id: str, public_key: st await tq_client.async_kv_link(keys, partition_id, public_key, [fields] if isinstance(fields, str) else fields) +async def async_kv_clear_unlinked_public() -> list[str]: + """Asynchronously clear every ``PUBLIC_PARTITION`` key no key links to. See ``kv_clear_unlinked_public``.""" + tq_client = _maybe_create_tq_client() + return await tq_client.async_clear_unlinked_public() + + def _single_update_batch(fields: dict[str, Any]) -> TensorDict: """Wrap one sample's field mapping in a one-row TensorDict.""" if not isinstance(fields, dict) or not fields: diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 4ff5047b..d2d525be 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -131,6 +131,8 @@ class ZMQRequestType(ExplicitEnum): KV_CLEAR_FIELDS_RESPONSE = "KV_CLEAR_FIELDS_RESPONSE" KV_LINK = "KV_LINK" KV_LINK_RESPONSE = "KV_LINK_RESPONSE" + MARK_UNLINKED_PUBLIC_CLEARING = "MARK_UNLINKED_PUBLIC_CLEARING" + MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE = "MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE" # METRICS GET_METRICS = "GET_METRICS" From b3849e9d326da3ba1de41d0c91aa2a46847b438a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 12:07:28 +0800 Subject: [PATCH 18/32] [fix] drop links to public rows released while a racing put made them ready again --- tests/e2e/test_kv_link_e2e.py | 15 +++++++++++++++ transfer_queue/controller.py | 15 +++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index eafd8531..86a44f25 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -130,6 +130,21 @@ def test_link_refuses_a_public_key_being_cleared(tq_system): tq.kv_link("r", PARTITION, "p0", "prompt") +def test_link_made_while_its_public_key_is_cleared_is_dropped(tq_system): + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + source = tq.get_client().kv_describe(["p0"], PUBLIC)["rows"]["p0"]["global_index"] + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + ray.get(controller.mark_clearing.remote([source], [PUBLIC])) + # A concurrent put makes p0 ready again, so a link slips in before the clear finishes. + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + tq.kv_link("r", PARTITION, "p0", "prompt") + + ray.get(controller.clear_meta.remote([source], [PUBLIC])) + + assert _links("r") == {} + assert "prompt" not in tq.kv_describe("r", PARTITION)["r"]["fields"] + + def test_public_key_cannot_be_cleared_while_linked(tq_system): keys = _put_group(0) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 6521ed8a..b9fd0851 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1680,6 +1680,17 @@ def _check_public_not_linked(self, global_indexes: list[int], field_names: list[ "clear those keys first" ) + def _unlink_released_public(self, global_indexes: list[int]) -> None: + """Mark fields linked to released public rows as never written. + + mark_clearing refuses linked rows, but a put during the clear makes a row ready + again and lets kv_link through; such a link would point at a reusable index. + """ + for other in self.partitions.values(): + for source in set(global_indexes) & other.linked_by.keys(): + for idx in list(other.linked_by[source]): + other.clear_fields([idx], [name for name, src in other.links[idx].items() if src == source]) + def clear_partition(self, partition_id: str, clear_consumption: bool = True): """ Clear data for a specific partition (delete the whole partition). @@ -1700,6 +1711,8 @@ def clear_partition(self, partition_id: str, clear_consumption: bool = True): global_indexes_range = list(self.index_manager.get_indexes_for_partition(partition_id)) partition.clear_data(global_indexes_range, clear_consumption) + if partition_id == PUBLIC_PARTITION: + self._unlink_released_public(global_indexes_range) self.index_manager.release_partition(partition_id) self.partitions.pop(partition_id) self.sampler.clear_cache(partition_id) @@ -1784,6 +1797,8 @@ def clear_meta( # Clear data from partition partition.clear_data(global_indexes_to_clear, clear_consumption) + if partition_id == PUBLIC_PARTITION: + self._unlink_released_public(global_indexes_to_clear) # Release the specific indexes from index manager self.index_manager.release_indexes(partition_id, global_indexes_to_clear) From c2048cb50aa9799d1b6709e1c2b27da04ce357da Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 12:19:52 +0800 Subject: [PATCH 19/32] [fix] kv_link validates field types before creating keys --- tests/e2e/test_kv_link_e2e.py | 9 +++++++++ transfer_queue/controller.py | 7 +++++++ 2 files changed, 16 insertions(+) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 86a44f25..f268471d 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -120,6 +120,15 @@ def test_link_requires_a_produced_public_field(tq_system): assert tq.kv_list().get(PARTITION, {}) == {} +def test_link_with_conflicting_dtype_creates_no_keys(tq_system): + tq.kv_put(key="a", partition_id=PARTITION, fields={"prompt": _prompt(0)}) + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0).float()}) + + with pytest.raises(RuntimeError, match="dtype mismatch"): + tq.kv_link(["r0", "r1"], PARTITION, "p0", "prompt") + assert tq.kv_list(PARTITION)[PARTITION].keys() == {"a"} + + def test_link_refuses_a_public_key_being_cleared(tq_system): tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) source = tq.get_client().kv_describe(["p0"], PUBLIC)["rows"]["p0"]["global_index"] diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index b9fd0851..378f3599 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2061,6 +2061,13 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l ] if not fields or missing: raise ValueError(f"public key {public_key!r} must have produced fields {missing or fields}") + # Before creating keys, so an incompatible link leaves none behind. + partition = self._get_partition(partition_id) + if partition is not None: + metas = {name: public.field_metadata[name] for name in fields} + partition.validate_field_schema( + {name: {"dtype": meta.dtype, "is_non_tensor": meta.is_non_tensor} for name, meta in metas.items()} + ) metadata = self.kv_retrieve_meta(keys, partition_id, create=True) partition = self._get_partition(partition_id) assert partition is not None From 5837879910f1a2e2c24421415f1a4d955f280e0b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 12:33:19 +0800 Subject: [PATCH 20/32] [test] controller checkpoints restore kv_link links and their clear guard --- tests/e2e/test_kv_link_e2e.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index f268471d..48b250cb 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -198,6 +198,22 @@ def test_link_refuses_a_key_marked_by_clear_unlinked_public(tq_system): tq.kv_link("r", PARTITION, "p0", "prompt") +def test_checkpoint_restores_links_and_their_clear_guard(tq_system, tmp_path): + keys = _put_group(0) + tq.save_checkpoint(tmp_path / "ck") + tq.kv_clear(keys, PARTITION) + tq.kv_clear("p0", PUBLIC) + + tq.load_checkpoint(tmp_path / "ck") + + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + assert all(_links(key) == {"prompt": "p0"} for key in keys) + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + tq.kv_clear("p0", PUBLIC) + tq.kv_clear(keys, PARTITION) + assert tq.kv_clear_unlinked_public() == ["p0"] + + def test_put_replaces_the_link(tq_system): keys = _put_group(0) From 4c441ca35ad42f4267cd0236fa69ca0712dfc142 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 12:52:43 +0800 Subject: [PATCH 21/32] [fix] repeated keys keep their own parser_args when reading --- tests/e2e/test_kv_interface_e2e.py | 6 +++ .../managers/simple_storage_manager.py | 46 ++++++++++--------- 2 files changed, 31 insertions(+), 21 deletions(-) diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index de3f611a..03a03298 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -812,6 +812,12 @@ def parse(data, args): assert_nested_tensor_equal(retrieved["input_ids"], [input_ids[0, :2], input_ids[1, :3], input_ids[2, :1]]) assert_tensor_equal(retrieved["scores"], scores * 10) + # A repeated key reads its own argument, and the keys after it keep theirs. + retrieved = tq_api.kv_batch_get( + keys=["parse_0", "parse_0", "parse_1"], partition_id=partition_id, data_parser=parse, parser_args=[1, 2, 3] + ) + assert_nested_tensor_equal(retrieved["input_ids"], [input_ids[0, :1], input_ids[0, :2], input_ids[1, :3]]) + retrieved = tq_api.kv_batch_get(keys=keys, partition_id=partition_id) assert_tensor_equal(retrieved["input_ids"], input_ids) assert_tensor_equal(retrieved["scores"], scores) diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 732f21fb..e017d44c 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -825,18 +825,26 @@ async def get_data( # Parser-backed gets are not replayed, matching put: the public API does not constrain parser side effects. max_attempts = 1 if data_parser is not None else None - # A unit request reads the same fields for every row in it, and a row stores only its - # unlinked fields, so rows are grouped by the field set they physically hold. Without - # links this is one group in batch order, so batch_positions index parser_args. - fields_by_row: dict[int, set[str]] = defaultdict(set) - for global_index, links in zip(metadata.global_indexes, metadata._links, strict=True): - for field in metadata.field_names: - fields_by_row[links.get(field, global_index)].add(field) + # A unit request reads the same fields for every row in it, so rows are grouped by the + # fields read from them: a sample's own fields once per occurrence, so repeated keys + # keep their parser_args, and a linked field once per source row however many link to it. + # Without links this is one group in batch order, so batch_positions index parser_args. rows_by_fields: dict[tuple[str, ...], list[int]] = defaultdict(list) - for global_index, fields in fields_by_row.items(): - rows_by_fields[tuple(sorted(fields))].append(global_index) + slots: dict[tuple[int, str], tuple[tuple[str, ...], int]] = {} # (sample, field) -> (fields, row) + linked_slots: dict[tuple[int, str], tuple[tuple[str, ...], int]] = {} # (source, field) -> slot + for pos, (global_index, links) in enumerate(zip(metadata.global_indexes, metadata._links, strict=True)): + own = tuple(f for f in metadata.field_names if f not in links) + if own: + slots.update({(pos, f): (own, len(rows_by_fields[own])) for f in own}) + rows_by_fields[own].append(global_index) + for field in metadata.field_names: + if field in links: + if (links[field], field) not in linked_slots: + linked_slots[links[field], field] = ((field,), len(rows_by_fields[field,])) + rows_by_fields[field,].append(links[field]) + slots[pos, field] = linked_slots[links[field], field] requests = [ - (su_id, list(fields), group) + (su_id, fields, group) for fields, rows in rows_by_fields.items() for su_id, group in self._group_by_hash(rows).items() ] @@ -845,11 +853,11 @@ async def get_data( self._request_with_retry( "get", su_id, - f"samples={len(group.global_indexes)} fields={fields}", + f"samples={len(group.global_indexes)} fields={list(fields)}", partial( self._get_from_single_storage_unit, group.global_indexes, - fields, + list(fields), target_storage_unit=su_id, data_parser=data_parser, parser_args=None if parser_args is None else [parser_args[p] for p in group.batch_positions], @@ -873,17 +881,13 @@ async def get_data( ) raise - values: dict[tuple[int, str], Any] = {} - for (_, _, group), (fields, su_data) in zip(requests, results, strict=True): + values: dict[tuple[tuple[tuple[str, ...], int], str], Any] = {} + for (_, fields, group), (_, su_data) in zip(requests, results, strict=True): for field in fields: - for i, global_index in enumerate(group.global_indexes): - values[global_index, field] = su_data[field][i] + for i, row in enumerate(group.batch_positions): + values[(fields, row), field] = su_data[field][i] ordered_data = { - field: [ - values[links.get(field, global_index), field] - for global_index, links in zip(metadata.global_indexes, metadata._links, strict=True) - ] - for field in metadata.field_names + field: [values[slots[pos, field], field] for pos in range(metadata.size)] for field in metadata.field_names } tensor_data = {field: self._pack_field_values(v) for field, v in ordered_data.items()} From 8fdf3ff14545fb7de0fb4ccb2c8e1eb86a1f8b3e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 12:55:46 +0800 Subject: [PATCH 22/32] [fix] bump the dump format to 5 so older readers reject dumps with links --- docs/data_dump.md | 9 +++++---- tests/e2e/test_data_dump_e2e.py | 13 ++++++++----- transfer_queue/data_dump.py | 9 +++++---- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/docs/data_dump.md b/docs/data_dump.md index 4365ff63..5ecfa33b 100644 --- a/docs/data_dump.md +++ b/docs/data_dump.md @@ -55,10 +55,11 @@ These count application reads, not filesystem read-ahead or physical disk traffi ## Format and compatibility -Dumps use `format_version: 4`. Version 4 adds each row's `unfinalized` fields, which -restore keeps hidden. This build also reads version 3 dumps, whose fields restore as -ready. Builds that only read version 3 reject version 4 dumps instead of exposing -hidden fields. The layout: +Dumps use `format_version: 5`. Version 5 adds each row's `links`, the fields it reads +from `PUBLIC_PARTITION` keys through `kv_link`; version 4 added `unfinalized`, the +fields restore keeps hidden. This build also reads versions 3 and 4: version 4 rows +have no links, and version 3 fields restore as ready. Older builds reject newer dumps +instead of dropping linked fields or exposing hidden ones. The layout: ```text dump_info.json diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py index c60b50bf..95816533 100644 --- a/tests/e2e/test_data_dump_e2e.py +++ b/tests/e2e/test_data_dump_e2e.py @@ -272,18 +272,21 @@ def concat(old, new): retrieved = tq.kv_batch_get(keys="u1", partition_id=partition_id, select_fields="response") _assert_rows_equal(retrieved["response"], [torch.tensor([1, 2, 3])]) - def test_restores_a_version_3_dump_as_ready(self, tq_system, dump_dir, controller): - partition_id = "d_version3" + @pytest.mark.parametrize("version", [3, 4]) + def test_restores_an_older_dump_as_ready(self, tq_system, dump_dir, controller, version): + partition_id = f"d_version{version}" _put_rows(partition_id, ["w0"]) tq.dump_data_by_key(dump_dir, ["w0"], partition_id) - # Rewrite it as version 3 wrote it: no ``unfinalized`` in the row index. + # Rewrite it as that version wrote it: no ``links``, and before 4 no ``unfinalized``. row_index = _read_row_index(dump_dir) - del row_index["rows"]["w0"]["unfinalized"] + del row_index["rows"]["w0"]["links"] + if version == 3: + del row_index["rows"]["w0"]["unfinalized"] torch.save(row_index, dump_dir / "row_index.pt") info_path = dump_dir / "dump_info.json" info = json.loads(info_path.read_text(encoding="utf-8")) - info["format_version"] = 3 + info["format_version"] = version info_path.write_text(json.dumps(info), encoding="utf-8") ray.get(controller.clear_partition.remote(partition_id)) diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py index e39a8e30..e3b6fdbf 100644 --- a/transfer_queue/data_dump.py +++ b/transfer_queue/data_dump.py @@ -42,10 +42,11 @@ logger = get_logger(__name__) -DUMP_FORMAT_VERSION = 4 -# Version 3 predates ``unfinalized``; every field it holds was ready. A version-3 reader -# rejects version 4 instead of restoring hidden fields as ready. -_READABLE_FORMAT_VERSIONS = (3, DUMP_FORMAT_VERSION) +DUMP_FORMAT_VERSION = 5 +# Version 3 predates ``unfinalized``, so every field it holds was ready; version 4 predates +# ``links``. Older readers reject newer dumps instead of restoring hidden fields as ready or +# dropping linked fields, which ``fields`` leaves out. +_READABLE_FORMAT_VERSIONS = (3, 4, DUMP_FORMAT_VERSION) SHARD_SUBDIR = "shards" _DUMP_INFO_FILE = "dump_info.json" From 5162c4fd8425d465d521c26e1e141f106a9669f6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:02:44 +0800 Subject: [PATCH 23/32] [fix] kv_link refuses an empty key list --- tests/e2e/test_kv_link_e2e.py | 2 ++ transfer_queue/controller.py | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 48b250cb..b5a9ab50 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -117,6 +117,8 @@ def test_link_requires_a_produced_public_field(tq_system): tq.kv_link("r", PUBLIC, "p0", "prompt") with pytest.raises(ValueError, match="duplicates"): tq.kv_link(["r", "r"], PARTITION, "p0", "prompt") + with pytest.raises(RuntimeError, match="must not be empty"): + tq.kv_link([], PARTITION, "p0", "prompt") assert tq.kv_list().get(PARTITION, {}) == {} diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 378f3599..82bbc00e 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2043,11 +2043,13 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l Raises: KeyError: ``public_key`` does not exist. - ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, or ``fields`` is empty or not - all produced by ``public_key``. + ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, ``keys`` or ``fields`` is empty, + or ``public_key`` has not produced all ``fields``. """ if partition_id == PUBLIC_PARTITION: raise ValueError(f"keys of partition {PUBLIC_PARTITION!r} hold shared payloads and cannot link") + if not keys: + raise ValueError("keys must not be empty") public = self._get_partition(PUBLIC_PARTITION) if public is None or public_key not in public.keys_mapping: raise KeyError(f"key {public_key!r} not found in partition {PUBLIC_PARTITION!r}") From 78134b56b644788bb6c24fe7398134b91e5ca82d Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:02:47 +0800 Subject: [PATCH 24/32] [fix] kv_link refuses fields a key already stores itself --- tests/e2e/test_kv_link_e2e.py | 13 +++++++++++++ transfer_queue/controller.py | 19 ++++++++++++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index b5a9ab50..f6322568 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -122,6 +122,19 @@ def test_link_requires_a_produced_public_field(tq_system): assert tq.kv_list().get(PARTITION, {}) == {} +def test_link_refuses_a_field_the_key_stores_itself(tq_system): + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + tq.kv_put(key="r0", partition_id=PARTITION, fields={"prompt": _prompt(7)}) + + with pytest.raises(RuntimeError, match="already store"): + tq.kv_link(["r0", "r1"], PARTITION, "p0", "prompt") + assert tq.kv_list(PARTITION)[PARTITION].keys() == {"r0"} + + tq.kv_clear("r0", PARTITION, fields="prompt") + tq.kv_link(["r0", "r1"], PARTITION, "p0", "prompt") + assert _links("r0") == {"prompt": "p0"} + + def test_link_with_conflicting_dtype_creates_no_keys(tq_system): tq.kv_put(key="a", partition_id=PARTITION, fields={"prompt": _prompt(0)}) tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0).float()}) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 82bbc00e..fe9ca88b 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2044,7 +2044,7 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l Raises: KeyError: ``public_key`` does not exist. ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, ``keys`` or ``fields`` is empty, - or ``public_key`` has not produced all ``fields``. + ``public_key`` has not produced all ``fields``, or a key stores one of them itself. """ if partition_id == PUBLIC_PARTITION: raise ValueError(f"keys of partition {PUBLIC_PARTITION!r} hold shared payloads and cannot link") @@ -2066,6 +2066,23 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l # Before creating keys, so an incompatible link leaves none behind. partition = self._get_partition(partition_id) if partition is not None: + # Linking only redirects reads, so a value the key stores itself would stay in storage unread. + stored = [ + key + for key, idx in zip(keys, partition.kv_retrieve_indexes(keys), strict=True) + if idx is not None + and any( + name not in partition.links.get(idx, {}) + and name in partition.field_metadata + and idx in partition.field_metadata[name].global_indexes + for name in fields + ) + ] + if stored: + raise ValueError( + f"keys {stored[:20]} already store some of fields {fields}; " + "kv_clear those fields before linking them" + ) metas = {name: public.field_metadata[name] for name in fields} partition.validate_field_schema( {name: {"dtype": meta.dtype, "is_non_tensor": meta.is_non_tensor} for name, meta in metas.items()} From ae87d25737b974d247fb04e33fb947169b1df834 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:03:40 +0800 Subject: [PATCH 25/32] [fix] check the SimpleStorage backend in TransferQueueClient.kv_link --- tests/e2e/test_kv_link_e2e.py | 13 +++++++++++++ transfer_queue/client.py | 4 ++++ transfer_queue/interface.py | 3 --- 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index f6322568..ad65bc89 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -135,6 +135,19 @@ def test_link_refuses_a_field_the_key_stores_itself(tq_system): assert _links("r0") == {"prompt": "p0"} +def test_link_requires_simple_storage_on_the_client_too(tq_system, monkeypatch): + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + client = tq.get_client() + monkeypatch.setattr(client, "storage_manager", object()) + + with pytest.raises(NotImplementedError, match="does not support kv_link"): + client.kv_link(["r"], PARTITION, "p0", ["prompt"]) + with pytest.raises(NotImplementedError, match="does not support kv_link"): + tq.kv_link("r", PARTITION, "p0", "prompt") + monkeypatch.undo() + assert tq.kv_list().get(PARTITION, {}) == {} + + def test_link_with_conflicting_dtype_creates_no_keys(tq_system): tq.kv_put(key="a", partition_id=PARTITION, fields={"prompt": _prompt(0)}) tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0).float()}) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 4461457b..61fc2902 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -26,6 +26,7 @@ from transfer_queue.metadata import BatchMeta from transfer_queue.storage import StorageManagerFactory +from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( @@ -1196,6 +1197,9 @@ async def async_kv_link( socket: zmq.asyncio.Socket | None = None, ) -> None: """Asynchronously make ``fields`` of ``keys`` read the payload of ``public_key``; see ``kv_link``.""" + # Only SimpleStorage resolves links on read; other backends would read the linking key's own cells. + if not isinstance(self.storage_manager, AsyncSimpleStorageManager): + raise NotImplementedError(f"{type(self.storage_manager).__name__} does not support kv_link") await self._request_controller( socket=socket, request_type=ZMQRequestType.KV_LINK, diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 63b4b0d4..86044d88 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -1278,9 +1278,6 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[ async def async_kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: list[str] | str) -> None: """Asynchronously let ``keys`` read ``fields`` from ``public_key``. See ``kv_link``.""" tq_client = _maybe_create_tq_client() - # Only SimpleStorage resolves links on read; other backends would read the linking key's own cells. - if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): - raise NotImplementedError(f"{type(tq_client.storage_manager).__name__} does not support kv_link") keys = [keys] if isinstance(keys, str) else keys _check_unique_keys(keys) await tq_client.async_kv_link(keys, partition_id, public_key, [fields] if isinstance(fields, str) else fields) From b0d2acbb673e41e00022344fdbb95a322a23bf5f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:04:41 +0800 Subject: [PATCH 26/32] [fix] load relinks before writing payloads so an incompatible link writes nothing --- tests/e2e/test_kv_link_e2e.py | 15 +++++++++++++++ transfer_queue/interface.py | 5 +++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index ad65bc89..04e1cbbb 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -322,3 +322,18 @@ def test_dump_and_load_relink(tq_system, dump_test_root): assert all(_links(key) == {"prompt": "p0"} for key in keys) assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + + +def test_load_with_an_incompatible_linked_field_writes_nothing(tq_system, dump_test_root): + keys = _put_group(0) + tq.kv_batch_put(keys, PARTITION, fields=TensorDict({"response": torch.arange(N).unsqueeze(1)}, batch_size=N)) + linked_dir = dump_test_root / "incompatible_link_dump" + tq.dump_data_by_key(linked_dir, keys, PARTITION) + tq.get_client().clear_partition(PARTITION) + tq.kv_put(key="other", partition_id=PARTITION, fields={"prompt": _prompt(0).float()}) + + with pytest.raises(RuntimeError, match="dtype mismatch"): + tq.load_data_by_key(linked_dir) + shutil.rmtree(linked_dir) + + assert tq.kv_list(PARTITION)[PARTITION].keys() == {"other"} diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 86044d88..4bada4db 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -1650,8 +1650,8 @@ def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: ) if unready: raise ValueError(f"linked fields of {PUBLIC_PARTITION!r} keys {unready} are not ready; load them first") - bytes_read = client.load_rows_by_key(partition_id, rows, shards) - # One kv_link per public key and field set. + # Links first: kv_link checks field types before creating any key, so an incompatible + # destination fails before any payload is written. One kv_link per public key and field set. link_groups: dict[tuple[str, tuple[str, ...]], list[str]] = defaultdict(list) for key, row in rows.items(): links = row.get("links", {}) @@ -1659,5 +1659,6 @@ def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: link_groups[public_key, tuple(sorted(name for name, k in links.items() if k == public_key))].append(key) for (public_key, names), keys in link_groups.items(): client.kv_link(keys, partition_id, public_key, list(names)) + bytes_read = client.load_rows_by_key(partition_id, rows, shards) logger.info(f"Restored {len(rows)} keys into partition {partition_id} from {dump_dir}, reading {bytes_read} bytes") return {key: row["tag"] for key, row in rows.items()} From bba58db492a2c6fed9612ee0ed0efbdd90c15768 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:07:08 +0800 Subject: [PATCH 27/32] [feat] a linked public key is read-only until no key links to it --- docs/data_dump.md | 3 ++- tests/e2e/test_kv_link_e2e.py | 20 ++++++++++++++------ transfer_queue/controller.py | 6 +++++- transfer_queue/interface.py | 13 +++++++------ 4 files changed, 28 insertions(+), 14 deletions(-) diff --git a/docs/data_dump.md b/docs/data_dump.md index 5ecfa33b..4857e35b 100644 --- a/docs/data_dump.md +++ b/docs/data_dump.md @@ -56,7 +56,8 @@ These count application reads, not filesystem read-ahead or physical disk traffi ## Format and compatibility Dumps use `format_version: 5`. Version 5 adds each row's `links`, the fields it reads -from `PUBLIC_PARTITION` keys through `kv_link`; version 4 added `unfinalized`, the +from `PUBLIC_PARTITION` keys through `kv_link`, recorded by public key name: load +relinks to whatever those keys hold then, so load their dump first. Version 4 added `unfinalized`, the fields restore keeps hidden. This build also reads versions 3 and 4: version 4 rows have no links, and version 3 fields restore as ready. Older builds reject newer dumps instead of dropping linked fields or exposing hidden ones. The layout: diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 04e1cbbb..77ac5727 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -290,16 +290,24 @@ def test_update_refuses_a_linked_field(tq_system): assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N -def test_linked_keys_follow_public_readiness_and_shape(tq_system): +def test_linked_public_key_is_read_only(tq_system): keys = _put_group(0) - tq.kv_update("p0", PUBLIC, {"prompt": _prompt(5, 5)}, merge_fn=lambda old, new: new, finalize=False) - with pytest.raises(ValueError, match="not ready"): - tq.kv_batch_get(keys, PARTITION, select_fields="prompt") + for write in ( + lambda: tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(5, 5)}), + lambda: tq.kv_put(key="p0", partition_id=PUBLIC, tag={"step": 1}), + lambda: tq.kv_update("p0", PUBLIC, {"prompt": _prompt(5, 5)}, merge_fn=lambda old, new: new), + ): + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + write() + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N - tq.kv_update("p0", PUBLIC, {"prompt": _prompt(5, 5)}, merge_fn=lambda old, new: new) + # A BatchMeta put skips the KV write path; linked keys still follow the new payload's shape. + client = tq.get_client() + meta = client.kv_retrieve_meta(["p0"], PUBLIC) + client.put(TensorDict({"prompt": _prompt(5, 5).unsqueeze(0)}, batch_size=1), metadata=meta) assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(5, 5).tolist()] * N - meta = tq.get_client().kv_retrieve_meta(keys, PARTITION) + meta = client.kv_retrieve_meta(keys, PARTITION) assert [tuple(shape) for shape in meta.get_shapes("prompt")] == [(5,)] * N diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index fe9ca88b..d23ade77 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1670,7 +1670,7 @@ def mark_clearing(self, global_indexes: list[int], partition_ids: list[str]) -> partition.production_status[existing, :] = 0 def _check_public_not_linked(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: - """Refuse to clear public payloads that keys of other partitions still link to.""" + """Refuse to clear or rewrite public payloads that keys of other partitions still link to.""" linked = set().union(*(p.linked_sources(global_indexes, field_names) for p in self.partitions.values())) if linked: public = self.partitions[PUBLIC_PARTITION] @@ -1838,6 +1838,10 @@ def kv_retrieve_meta( assert partition is not None global_indexes = partition.kv_retrieve_indexes(keys) + if create and partition_id == PUBLIC_PARTITION: + # create=True is the KV write path (kv_put, kv_update, restore); a linked public row + # stays read-only so linking keys never see their shared payload change. + self._check_public_not_linked([idx for idx in global_indexes if idx is not None]) none_indexes = [idx for idx, value in enumerate(global_indexes) if value is None] if len(none_indexes) > 0: diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 4bada4db..9ccea451 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -702,10 +702,10 @@ def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: l Put a payload shared by many keys, e.g. the video of a GRPO prompt read by its n responses, once into ``tq.PUBLIC_PARTITION`` and link the keys to it. Its rows are spread evenly over the storage units, apart from the rows of every other partition. - Missing ``keys`` are created, and linked fields read the public key's current value - and are ready whenever its fields are. Putting a linked field stores a private value - in place of the link, and ``kv_update`` refuses it; clearing a linked key or field - drops only the link. ``public_key`` cannot be cleared while keys link to it. + Missing ``keys`` are created, and linked fields read the public key's payload and are + ready whenever its fields are. Putting a linked field stores a private value in place + of the link, and ``kv_update`` refuses it; clearing a linked key or field drops only + the link. ``public_key`` cannot be written or cleared while keys link to it. Args: keys: Single key or list of keys that will read the linked fields. @@ -1547,8 +1547,9 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - Args: dump_dir: Directory to write the dump into. keys: Keys to dump. Duplicates are dropped, first occurrence wins. Fields linked - by ``kv_link`` are dumped as links; dump their ``PUBLIC_PARTITION`` keys - separately and load that dump first. + by ``kv_link`` are dumped as links by public key name; dump those + ``PUBLIC_PARTITION`` keys too and load that dump first, since load relinks to + whatever the names hold then. partition_id: Partition that owns ``keys``. Returns: From 11d6a38a7b073c899c6c817a08a269fba3ad836d Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:09:39 +0800 Subject: [PATCH 28/32] [feat] puts refuse fields linked by kv_link instead of replacing the link --- tests/e2e/test_kv_link_e2e.py | 25 ++++++++++++++++++------- transfer_queue/client.py | 13 +++++++++++-- transfer_queue/controller.py | 4 ++-- transfer_queue/interface.py | 7 ++++--- 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 77ac5727..906f8080 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -242,24 +242,35 @@ def test_checkpoint_restores_links_and_their_clear_guard(tq_system, tmp_path): assert tq.kv_clear_unlinked_public() == ["p0"] -def test_put_replaces_the_link(tq_system): +def test_put_refuses_a_linked_field(tq_system): keys = _put_group(0) + tq.kv_batch_put(keys, PARTITION, fields=TensorDict({"response": torch.arange(N).unsqueeze(1)}, batch_size=N)) - tq.kv_put(key=keys[0], partition_id=PARTITION, fields={"prompt": _prompt(7)}) + with pytest.raises(ValueError, match="linked by kv_link"): + tq.kv_put(key=keys[0], partition_id=PARTITION, fields={"prompt": _prompt(7)}) + # A read-modify-write of the whole batch would turn every link into a private copy. + data = tq.kv_batch_get(keys, PARTITION) + data["response"] += 100 + with pytest.raises(ValueError, match="linked by kv_link"): + tq.kv_batch_put(keys, PARTITION, fields=data) + assert all(_links(key) == {"prompt": "p0"} for key in keys) + tq.kv_batch_put(keys, PARTITION, fields=data.select("response")) + tq.kv_clear(keys[0], PARTITION, fields="prompt") + tq.kv_put(key=keys[0], partition_id=PARTITION, fields={"prompt": _prompt(7)}) assert _links(keys[0]) == {} assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(7).tolist()] + [_prompt(0).tolist()] * 2 + assert _rows(tq.kv_batch_get(keys, PARTITION)["response"]) == [[i + 100] for i in range(N)] -def test_put_through_meta_reads_back_the_written_value(tq_system): +def test_put_through_meta_refuses_a_linked_field(tq_system): keys = _put_group(0) client = tq.get_client() meta = client.kv_retrieve_meta(keys, PARTITION) - client.put(TensorDict({"prompt": torch.stack([_prompt(7)] * N)}, batch_size=N), metadata=meta) - - assert _rows(client.get_data(meta)["prompt"]) == [_prompt(7).tolist()] * N - assert _rows(tq.kv_batch_get("p0", PUBLIC)["prompt"]) == [_prompt(0).tolist()] + with pytest.raises(ValueError, match="linked by kv_link"): + client.put(TensorDict({"prompt": torch.stack([_prompt(7)] * N)}, batch_size=N), metadata=meta) + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N def test_data_parser_refuses_linked_fields(tq_system): diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 61fc2902..6806bb71 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -413,7 +413,8 @@ async def async_put( metadata; will be updated in a future version to reflect the post-put state) Raises: - ValueError: If metadata is None or empty, or if partition_id is None when metadata is not provided + ValueError: If metadata is None or empty, if partition_id is None when metadata is not provided, + or if ``data`` writes a field that ``metadata`` marks as linked by ``kv_link`` RuntimeError: If storage operation fails Example: @@ -476,6 +477,13 @@ async def async_put( if not metadata or metadata.size == 0: raise ValueError("metadata cannot be none or empty") + # A read-modify-write would otherwise turn every link into a private copy without notice. + linked = sorted({name for links in metadata._links for name in links if name in data.keys()}) + if linked: + raise ValueError( + f"fields {linked} are linked by kv_link on some samples; leave them out of the put, " + "or kv_clear those fields first to store private values" + ) with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.client_id}] async_put" @@ -1687,7 +1695,8 @@ def put( metadata; will be updated in a future version to reflect the post-put state) Raises: - ValueError: If metadata is None or empty, or if partition_id is None when metadata is not provided + ValueError: If metadata is None or empty, if partition_id is None when metadata is not provided, + or if ``data`` writes a field that ``metadata`` marks as linked by ``kv_link`` RuntimeError: If storage operation fails Example: diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index d23ade77..43483696 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -600,7 +600,7 @@ def update_production_status( # Save these global_indexes self.global_indexes.update(global_indices) - # A written value replaces the link it lands on + # A written value replaces the link it lands on; KV puts refuse linked fields up front. self._drop_links(global_indices, field_names) return True @@ -1910,7 +1910,7 @@ def kv_prepare_update( if linked: raise ValueError( f"kv_update cannot merge into fields linked by kv_link on keys {linked}; " - "kv_put a value to replace the link first" + "kv_clear those fields first to store private values" ) metadata = self.kv_retrieve_meta(keys, partition_id, create=True) partition = self._get_partition(partition_id) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 9ccea451..316754d5 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -703,9 +703,10 @@ def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: l responses, once into ``tq.PUBLIC_PARTITION`` and link the keys to it. Its rows are spread evenly over the storage units, apart from the rows of every other partition. Missing ``keys`` are created, and linked fields read the public key's payload and are - ready whenever its fields are. Putting a linked field stores a private value in place - of the link, and ``kv_update`` refuses it; clearing a linked key or field drops only - the link. ``public_key`` cannot be written or cleared while keys link to it. + ready whenever its fields are. ``kv_put`` and ``kv_update`` refuse a linked field, so + a read-modify-write must leave it out; clearing a linked key or field drops only the + link, after which the field takes a private value. ``public_key`` cannot be written + or cleared while keys link to it. Args: keys: Single key or list of keys that will read the linked fields. From 514ee429001a94953d55bb7ad73e353507a41bfe Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:12:07 +0800 Subject: [PATCH 29/32] [feat] kv_describe lists linked fields as readable, and link refusals name the linking keys --- tests/e2e/test_kv_link_e2e.py | 10 ++++++++-- tests/test_data_dump.py | 2 +- transfer_queue/controller.py | 27 ++++++++++++++++++--------- transfer_queue/interface.py | 11 ++++++++--- 4 files changed, 35 insertions(+), 15 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 906f8080..4a0af121 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -24,6 +24,7 @@ from tensordict import TensorDict import transfer_queue as tq +from transfer_queue.data_dump import _read_row_index PUBLIC = tq.PUBLIC_PARTITION PARTITION = "link" @@ -103,7 +104,11 @@ def test_linked_keys_read_the_public_payload(tq_system): assert _rows(data["prompt"]) == [_prompt(0).tolist()] * N + [_prompt(10).tolist()] * N assert _rows(data["response"]) == [[i] for i in range(len(keys))] - assert _links(keys[0]) == {"prompt": "p0"} + assert tq.kv_describe(keys[0], PARTITION)[keys[0]] == { + "fields": ["prompt", "response"], + "unfinalized": [], + "links": {"prompt": "p0"}, + } def test_link_requires_a_produced_public_field(tq_system): @@ -185,7 +190,7 @@ def test_link_made_while_its_public_key_is_cleared_is_dropped(tq_system): def test_public_key_cannot_be_cleared_while_linked(tq_system): keys = _put_group(0) - with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + with pytest.raises(RuntimeError, match=r"linked by keys of other partitions \(3 keys, e.g. \['link:p0_r0'"): tq.kv_clear("p0", PUBLIC) with pytest.raises(RuntimeError, match="linked by keys of other partitions"): tq.kv_clear("p0", PUBLIC, fields="prompt") @@ -334,6 +339,7 @@ def test_dump_and_load_relink(tq_system, dump_test_root): tq.load_data_by_key(linked_dir) assert tq.kv_list().get(PARTITION, {}) == {} + assert _read_row_index(linked_dir)["rows"][keys[0]]["fields"] == [] tq.load_data_by_key(public_dir) tq.load_data_by_key(linked_dir) shutil.rmtree(public_dir) diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py index 134de21c..5472f75a 100644 --- a/tests/test_data_dump.py +++ b/tests/test_data_dump.py @@ -67,7 +67,7 @@ def test_row_index_compacts_tensors_inside_tags(monkeypatch, tmp_path): tag = {"nested": [SimpleNamespace(value=batch[0])]} client = SimpleNamespace( kv_describe=lambda *_: { - "rows": {"key": {"global_index": 0, "fields": [], "links": {}, "tag": tag}}, + "rows": {"key": {"global_index": 0, "fields": [], "unfinalized": [], "links": {}, "tag": tag}}, "field_schema": {}, } ) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 43483696..09a89b65 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1671,13 +1671,22 @@ def mark_clearing(self, global_indexes: list[int], partition_ids: list[str]) -> def _check_public_not_linked(self, global_indexes: list[int], field_names: list[str] | None = None) -> None: """Refuse to clear or rewrite public payloads that keys of other partitions still link to.""" - linked = set().union(*(p.linked_sources(global_indexes, field_names) for p in self.partitions.values())) + linked: set[int] = set() + linkers: list[str] = [] + for p in self.partitions.values(): + for source in p.linked_sources(global_indexes, field_names): + linked.add(source) + linkers.extend( + f"{p.partition_id}:{p.revert_keys_mapping.get(idx, idx)}" + for idx in p.linked_by[source] + if field_names is None or any(p.links[idx].get(name) == source for name in field_names) + ) if linked: public = self.partitions[PUBLIC_PARTITION] keys = sorted(str(public.revert_keys_mapping.get(idx, idx)) for idx in linked) raise ValueError( - f"keys {keys} in partition {PUBLIC_PARTITION!r} are linked by keys of other partitions; " - "clear those keys first" + f"keys {keys} in partition {PUBLIC_PARTITION!r} are linked by keys of other partitions " + f"({len(linkers)} keys, e.g. {sorted(linkers)[:20]}); clear those keys first" ) def _unlink_released_public(self, global_indexes: list[int]) -> None: @@ -1995,10 +2004,10 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: Returns: ``{"rows": {key: {"global_index", "fields", "unfinalized", "links", "tag"}}, "field_schema": - {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` lists the row's stored fields and is - empty for a row with no produced field yet; ``unfinalized`` lists the stored fields not yet - ready (written by ``kv_update(finalize=False)``); ``links`` maps each linked field to its - key in ``PUBLIC_PARTITION``; ``field_schema`` covers the described rows' stored fields. + {field: {"is_nested", "is_non_tensor"}}}``. ``fields`` lists the row's readable fields, linked + ones included, and is empty for a row with no produced field yet; ``unfinalized`` lists those + not yet ready (written by ``kv_update(finalize=False)``); ``links`` maps each linked field to + its key in ``PUBLIC_PARTITION``; ``field_schema`` covers the described rows' stored fields. Raises: KeyError: The partition or any key does not exist. @@ -2019,7 +2028,7 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: fields = sorted( field_name for field_name, field_meta in partition.field_metadata.items() - if global_index in field_meta.global_indexes and field_name not in links + if global_index in field_meta.global_indexes ) rows[key] = { "global_index": global_index, @@ -2036,7 +2045,7 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: field_schema = { name: {"is_nested": bool(meta.is_nested), "is_non_tensor": bool(meta.is_non_tensor)} for name, meta in partition.field_metadata.items() - if any(name in row["fields"] for row in rows.values()) + if any(name in row["fields"] and name not in row["links"] for row in rows.values()) } return {"rows": rows, "field_schema": field_schema} diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 316754d5..f9ecb6f2 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -1216,9 +1216,10 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st partition_id: Partition containing the keys. Returns: - ``{key: {"fields", "unfinalized", "links"}}``. ``fields`` lists every stored field of - the key and ``unfinalized`` the subset that is not ready yet. ``links`` maps each field - linked by ``kv_link`` to its key in ``PUBLIC_PARTITION``, whose readiness it shares. + ``{key: {"fields", "unfinalized", "links"}}``. ``fields`` lists every field the key + can read, linked ones included, and ``unfinalized`` the subset that is not ready yet. + ``links`` maps each field linked by ``kv_link`` to its key in ``PUBLIC_PARTITION``, + whose readiness it shares. Raises: RuntimeError: If the partition or any key does not exist. @@ -1578,6 +1579,10 @@ def dump_data_by_key(dump_dir: str | Path, keys: list[str], partition_id: str) - description = client.kv_describe(unique_keys, partition_id) if unique_keys else {"rows": {}, "field_schema": {}} row_index = {"partition_id": partition_id, **description} rows = row_index["rows"] + # Linked payloads stay with their public keys; the row index records them as links. + for row in rows.values(): + row["fields"] = [name for name in row["fields"] if name not in row["links"]] + row["unfinalized"] = [name for name in row["unfinalized"] if name not in row["links"]] # A row whose fields are all still unproduced has nothing for a storage unit to # dump, but it keeps its key and tag so the restore can recreate the row. fields_by_index = {row["global_index"]: row["fields"] for row in rows.values() if row["fields"]} From 2e812ec15c22d1121252539652dc8ff77e838ef5 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:13:14 +0800 Subject: [PATCH 30/32] [feat] kv_clear_unlinked_public frees only public keys that were linked before --- tests/e2e/test_kv_link_e2e.py | 10 ++++++---- transfer_queue/controller.py | 11 ++++++++--- transfer_queue/interface.py | 9 +++++---- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 4a0af121..5908e4a4 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -207,15 +207,16 @@ def test_public_key_cannot_be_cleared_while_linked(tq_system): assert all(not _links(key) for key in keys[1:]) -def test_clear_unlinked_public_frees_only_unlinked_keys(tq_system): +def test_clear_unlinked_public_frees_only_keys_no_longer_linked(tq_system): keys0, keys1 = _put_group(0), _put_group(1) + # Put but not linked yet, as by a producer between kv_put and kv_link. tq.kv_put(key="p2", partition_id=PUBLIC, fields={"prompt": _prompt(2)}) tq.kv_clear(keys1[:2], PARTITION) tq.kv_clear(keys1[2:], PARTITION, fields="prompt") - assert tq.kv_clear_unlinked_public() == ["p1", "p2"] + assert tq.kv_clear_unlinked_public() == ["p1"] - assert tq.kv_list(PUBLIC)[PUBLIC].keys() == {"p0"} + assert tq.kv_list(PUBLIC)[PUBLIC].keys() == {"p0", "p2"} with pytest.raises(ValueError, match="not found"): tq.kv_batch_get("p1", PUBLIC) assert _rows(tq.kv_batch_get(keys0, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N @@ -223,7 +224,8 @@ def test_clear_unlinked_public_frees_only_unlinked_keys(tq_system): def test_link_refuses_a_key_marked_by_clear_unlinked_public(tq_system): - tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + keys = _put_group(0) + tq.kv_clear(keys, PARTITION) controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") ray.get(controller.mark_unlinked_public_clearing.remote()) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 09a89b65..239fceeb 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -397,10 +397,13 @@ class DataPartitionStatus: # Linked fields hold no payload of their own and read a PUBLIC_PARTITION row's instead links: dict[int, dict[str, int]] = field(default_factory=dict) # global_idx -> {field: public global_idx} linked_by: dict[int, set[int]] = field(default_factory=dict) # public global_idx -> global_idxs linking to it + # Public rows some key has linked to; kv_clear_unlinked_public frees only these, so a + # payload put but not linked yet is safe from it. + ever_linked: set[int] = field(default_factory=set) def __setstate__(self, state: dict[str, Any]) -> None: # Controller checkpoints pickle partitions; ones saved before kv_link have no links. - self.__dict__.update({"links": {}, "linked_by": {}, **state}) + self.__dict__.update({"links": {}, "linked_by": {}, "ever_linked": set(), **state}) # Dynamic configuration - these are computed from the current state @property @@ -980,6 +983,7 @@ def clear_data(self, indexes_to_release: list[int], clear_consumption: bool = Tr for field_name in empty_fields: self.field_metadata.pop(field_name) self._drop_links(indexes_to_release) + self.ever_linked.difference_update(indexes_to_release) for idx in indexes_to_release: self.field_custom_backend_meta.pop(idx, None) self.custom_meta.pop(idx, None) @@ -2104,9 +2108,10 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l partition = self._get_partition(partition_id) assert partition is not None partition.link_fields(metadata.global_indexes, public, source, fields) + public.ever_linked.add(source) def mark_unlinked_public_clearing(self) -> tuple[BatchMeta, list[str]]: - """Mark every ``PUBLIC_PARTITION`` key that no key links to as clearing. + """Mark every ``PUBLIC_PARTITION`` key that keys linked to before but none links to now as clearing. Selecting and marking in one controller call means a concurrent kv_link cannot link a row this clear deletes: kv_link refuses fields whose production status was zeroed. @@ -2118,7 +2123,7 @@ def mark_unlinked_public_clearing(self) -> tuple[BatchMeta, list[str]]: if public is None: return BatchMeta.empty(), [] linked = set().union(*(p.linked_by.keys() for p in self.partitions.values())) - keys = sorted(key for key, idx in public.keys_mapping.items() if idx not in linked) + keys = sorted(public.revert_keys_mapping[idx] for idx in public.ever_linked - linked) if not keys: return BatchMeta.empty(), [] metadata = self.kv_retrieve_meta(keys, PUBLIC_PARTITION) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index f9ecb6f2..ea61aa3b 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -729,12 +729,13 @@ def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: l def kv_clear_unlinked_public() -> list[str]: - """Clear every key in ``PUBLIC_PARTITION`` that no key links to, and return those keys. + """Clear every key in ``PUBLIC_PARTITION`` that keys linked to but none links to now; return them. For example, after clearing a step's trajectories, this frees the prompt videos no - trajectory reads any more. A public key that is put but not linked yet is cleared too, - so call it when no writer is between its ``kv_put`` and ``kv_link``. A ``kv_link`` - that races with it is refused for the keys it clears. + trajectory reads any more. A public key never linked yet is left alone, so a producer + between its ``kv_put`` and ``kv_link`` is safe; clear such a key with ``kv_clear``. A + ``kv_link`` that races with it is refused for the keys it clears. ``PUBLIC_PARTITION`` + is shared by every job on this TransferQueue, so prefix public keys with a job id. Example: >>> tq.kv_clear(trajectory_keys, partition_id="train") From 1166b53329ad7e39b615cf49562c492f25bab4a4 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:15:00 +0800 Subject: [PATCH 31/32] [feat] kv_link takes one public key per key to link a batch in one atomic call --- tests/e2e/test_kv_link_e2e.py | 19 ++++++++++++- transfer_queue/client.py | 12 ++++---- transfer_queue/controller.py | 53 ++++++++++++++++++++++------------- transfer_queue/interface.py | 29 +++++++++++-------- 4 files changed, 76 insertions(+), 37 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 5908e4a4..ab9e93bb 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -127,6 +127,23 @@ def test_link_requires_a_produced_public_field(tq_system): assert tq.kv_list().get(PARTITION, {}) == {} +def test_link_a_batch_of_public_keys_in_one_call(tq_system): + for p in range(2): + tq.kv_put(key=f"p{p}", partition_id=PUBLIC, fields={"prompt": _prompt(p)}) + keys = ["p0_r0", "p0_r1", "p1_r0"] + + with pytest.raises(RuntimeError, match="2 public keys for 3 keys"): + tq.kv_link(keys, PARTITION, ["p0", "p1"], "prompt") + # Atomic: one missing public key links and creates nothing. + with pytest.raises(RuntimeError, match=r"keys \['missing'\] not found"): + tq.kv_link(keys, PARTITION, ["p0", "p0", "missing"], "prompt") + assert tq.kv_list().get(PARTITION, {}) == {} + + tq.kv_link(keys, PARTITION, ["p0", "p0", "p1"], "prompt") + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * 2 + [_prompt(1).tolist()] + assert [_links(key) for key in keys] == [{"prompt": "p0"}, {"prompt": "p0"}, {"prompt": "p1"}] + + def test_link_refuses_a_field_the_key_stores_itself(tq_system): tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) tq.kv_put(key="r0", partition_id=PARTITION, fields={"prompt": _prompt(7)}) @@ -146,7 +163,7 @@ def test_link_requires_simple_storage_on_the_client_too(tq_system, monkeypatch): monkeypatch.setattr(client, "storage_manager", object()) with pytest.raises(NotImplementedError, match="does not support kv_link"): - client.kv_link(["r"], PARTITION, "p0", ["prompt"]) + client.kv_link(["r"], PARTITION, ["p0"], ["prompt"]) with pytest.raises(NotImplementedError, match="does not support kv_link"): tq.kv_link("r", PARTITION, "p0", "prompt") monkeypatch.undo() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 6806bb71..af7d977f 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1200,11 +1200,11 @@ async def async_kv_link( self, keys: list[str], partition_id: str, - public_key: str, + public_keys: list[str], fields: list[str], socket: zmq.asyncio.Socket | None = None, ) -> None: - """Asynchronously make ``fields`` of ``keys`` read the payload of ``public_key``; see ``kv_link``.""" + """Asynchronously make ``fields`` of each key read its public key's payload; see ``kv_link``.""" # Only SimpleStorage resolves links on read; other backends would read the linking key's own cells. if not isinstance(self.storage_manager, AsyncSimpleStorageManager): raise NotImplementedError(f"{type(self.storage_manager).__name__} does not support kv_link") @@ -1212,7 +1212,7 @@ async def async_kv_link( socket=socket, request_type=ZMQRequestType.KV_LINK, response_type=ZMQRequestType.KV_LINK_RESPONSE, - body={"keys": keys, "partition_id": partition_id, "public_key": public_key, "fields": fields}, + body={"keys": keys, "partition_id": partition_id, "public_keys": public_keys, "fields": fields}, ) async def async_clear_unlinked_public(self) -> list[str]: @@ -2032,9 +2032,9 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: """ return self._kv_describe(keys, partition_id) - def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: list[str]) -> None: - """Synchronously make ``fields`` of ``keys`` read the payload of ``public_key``; see ``kv_link``.""" - return self._kv_link(keys, partition_id, public_key, fields) + def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fields: list[str]) -> None: + """Synchronously make ``fields`` of each key read its public key's payload; see ``kv_link``.""" + return self._kv_link(keys, partition_id, public_keys, fields) def clear_unlinked_public(self) -> list[str]: """Synchronously clear every ``PUBLIC_PARTITION`` key no key links to; return the cleared keys.""" diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 239fceeb..d91cb2d4 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2053,33 +2053,44 @@ def kv_describe(self, keys: list[str], partition_id: str) -> dict[str, Any]: } return {"rows": rows, "field_schema": field_schema} - def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: list[str]) -> None: - """Make ``fields`` of ``keys`` read the payload of ``public_key`` in ``PUBLIC_PARTITION``. + def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fields: list[str]) -> None: + """Make ``fields`` of each key read the payload of its key in ``PUBLIC_PARTITION``. - Missing ``keys`` are created. ``public_key`` must already have produced ``fields``. + ``public_keys`` is aligned with ``keys``. Missing ``keys`` are created. Every check + runs before anything changes, so a refused call creates no key and links nothing. Raises: - KeyError: ``public_key`` does not exist. + KeyError: A public key does not exist. ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, ``keys`` or ``fields`` is empty, - ``public_key`` has not produced all ``fields``, or a key stores one of them itself. + ``public_keys`` is not aligned with ``keys``, a public key has not produced all + ``fields``, or a key stores one of them itself. """ if partition_id == PUBLIC_PARTITION: raise ValueError(f"keys of partition {PUBLIC_PARTITION!r} hold shared payloads and cannot link") if not keys: raise ValueError("keys must not be empty") + if len(public_keys) != len(keys): + raise ValueError(f"{len(public_keys)} public keys for {len(keys)} keys") public = self._get_partition(PUBLIC_PARTITION) - if public is None or public_key not in public.keys_mapping: - raise KeyError(f"key {public_key!r} not found in partition {PUBLIC_PARTITION!r}") - source = public.keys_mapping[public_key] + absent = sorted({k for k in public_keys if public is None or k not in public.keys_mapping}) + if absent: + raise KeyError(f"keys {absent[:20]} not found in partition {PUBLIC_PARTITION!r}") + assert public is not None + sources = [public.keys_mapping[k] for k in public_keys] # Production status, not field metadata: mark_clearing zeroes it before storage deletes the payload. - missing = [ - name - for name in fields - if name not in public.field_name_mapping - or public.production_status[source, public.field_name_mapping[name]] != 1 - ] - if not fields or missing: - raise ValueError(f"public key {public_key!r} must have produced fields {missing or fields}") + unproduced = sorted( + { + public_key + for public_key, source in zip(public_keys, sources, strict=True) + if any( + name not in public.field_name_mapping + or public.production_status[source, public.field_name_mapping[name]] != 1 + for name in fields + ) + } + ) + if not fields or unproduced: + raise ValueError(f"public keys {unproduced[:20]} must have produced fields {fields}") # Before creating keys, so an incompatible link leaves none behind. partition = self._get_partition(partition_id) if partition is not None: @@ -2107,8 +2118,12 @@ def kv_link(self, keys: list[str], partition_id: str, public_key: str, fields: l metadata = self.kv_retrieve_meta(keys, partition_id, create=True) partition = self._get_partition(partition_id) assert partition is not None - partition.link_fields(metadata.global_indexes, public, source, fields) - public.ever_linked.add(source) + rows_by_source: dict[int, list[int]] = defaultdict(list) + for idx, source in zip(metadata.global_indexes, sources, strict=True): + rows_by_source[source].append(idx) + for source, indexes in rows_by_source.items(): + partition.link_fields(indexes, public, source, fields) + public.ever_linked.add(source) def mark_unlinked_public_clearing(self) -> tuple[BatchMeta, list[str]]: """Mark every ``PUBLIC_PARTITION`` key that keys linked to before but none links to now as clearing. @@ -2566,7 +2581,7 @@ def _handle_kv_prepare_update_request(self, request_msg: ZMQMessage) -> ZMQMessa def _handle_kv_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: params = request_msg.body - self.kv_link(params["keys"], params["partition_id"], params["public_key"], params["fields"]) + self.kv_link(params["keys"], params["partition_id"], params["public_keys"], params["fields"]) return self._make_response(request_msg, ZMQRequestType.KV_LINK_RESPONSE, {}) def _handle_mark_unlinked_public_clearing_request(self, request_msg: ZMQMessage) -> ZMQMessage: diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index ea61aa3b..7f6bc4b6 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -696,7 +696,7 @@ def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | tq_client._run_coroutine(async_kv_clear(keys=keys, partition_id=partition_id, fields=fields)) -def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: list[str] | str) -> None: +def kv_link(keys: list[str] | str, partition_id: str, public_key: str | list[str], fields: list[str] | str) -> None: """Let ``keys`` read ``fields`` from ``public_key`` in ``PUBLIC_PARTITION`` without another copy. Put a payload shared by many keys, e.g. the video of a GRPO prompt read by its n @@ -711,18 +711,22 @@ def kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: l Args: keys: Single key or list of keys that will read the linked fields. partition_id: Partition holding ``keys``; not ``PUBLIC_PARTITION``. - public_key: Key in ``PUBLIC_PARTITION`` that has produced ``fields``. + public_key: Key in ``PUBLIC_PARTITION`` that has produced ``fields``, or a list of + such keys aligned with ``keys`` to link a whole batch in one atomic call. fields: Single field name or list of field names to link. Raises: - RuntimeError: If ``public_key`` does not exist or has not produced all of ``fields``, - or ``partition_id`` is ``PUBLIC_PARTITION``. + RuntimeError: If a public key does not exist or has not produced all of ``fields``, + a list ``public_key`` is not aligned with ``keys``, or ``partition_id`` is + ``PUBLIC_PARTITION``. Nothing is created or linked then. ValueError: If ``keys`` contains duplicates. NotImplementedError: If the storage backend is not SimpleStorage. Example: >>> tq.kv_put(key="prompt_0", partition_id=tq.PUBLIC_PARTITION, fields={"video": video}) >>> tq.kv_link([f"prompt_0_{i}" for i in range(n)], "train", "prompt_0", "video") + >>> # A batch of prompts at once: each response key with its own prompt key. + >>> tq.kv_link(response_keys, "train", [key.rsplit("_", 1)[0] for key in response_keys], "video") """ tq_client = _maybe_create_tq_client() tq_client._run_coroutine(async_kv_link(keys, partition_id, public_key, fields)) @@ -1278,12 +1282,15 @@ async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[ await tq_client.async_clear_samples(batch_meta) -async def async_kv_link(keys: list[str] | str, partition_id: str, public_key: str, fields: list[str] | str) -> None: +async def async_kv_link( + keys: list[str] | str, partition_id: str, public_key: str | list[str], fields: list[str] | str +) -> None: """Asynchronously let ``keys`` read ``fields`` from ``public_key``. See ``kv_link``.""" tq_client = _maybe_create_tq_client() keys = [keys] if isinstance(keys, str) else keys _check_unique_keys(keys) - await tq_client.async_kv_link(keys, partition_id, public_key, [fields] if isinstance(fields, str) else fields) + public_keys = [public_key] * len(keys) if isinstance(public_key, str) else public_key + await tq_client.async_kv_link(keys, partition_id, public_keys, [fields] if isinstance(fields, str) else fields) async def async_kv_clear_unlinked_public() -> list[str]: @@ -1659,14 +1666,14 @@ def load_data_by_key(dump_dir: str | Path) -> dict[str, dict]: if unready: raise ValueError(f"linked fields of {PUBLIC_PARTITION!r} keys {unready} are not ready; load them first") # Links first: kv_link checks field types before creating any key, so an incompatible - # destination fails before any payload is written. One kv_link per public key and field set. - link_groups: dict[tuple[str, tuple[str, ...]], list[str]] = defaultdict(list) + # destination fails before any payload is written. One kv_link per linked field set. + link_groups: dict[tuple[str, ...], list[tuple[str, str]]] = defaultdict(list) for key, row in rows.items(): links = row.get("links", {}) for public_key in set(links.values()): - link_groups[public_key, tuple(sorted(name for name, k in links.items() if k == public_key))].append(key) - for (public_key, names), keys in link_groups.items(): - client.kv_link(keys, partition_id, public_key, list(names)) + link_groups[tuple(sorted(name for name, k in links.items() if k == public_key))].append((key, public_key)) + for names, pairs in link_groups.items(): + client.kv_link([key for key, _ in pairs], partition_id, [public_key for _, public_key in pairs], list(names)) bytes_read = client.load_rows_by_key(partition_id, rows, shards) logger.info(f"Restored {len(rows)} keys into partition {partition_id} from {dump_dir}, reading {bytes_read} bytes") return {key: row["tag"] for key, row in rows.items()} From c8e999a09f2e9827eacded933aa7907802a8b3de Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 22:28:11 +0800 Subject: [PATCH 32/32] [fix] kv_link merges into a field whose schema is non-tensor A put of a non-tensor value turns the partition's field non-tensor, and a later put of a tensor still lands. kv_link checked the strict restore schema instead, so linking a tensor prompt into a field that had held a text-only prompt failed. Reject only links that the merge itself would fail on: two tensor dtypes. --- tests/e2e/test_kv_link_e2e.py | 12 ++++++++++++ transfer_queue/controller.py | 15 ++++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index ab9e93bb..ac18bee9 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -179,6 +179,18 @@ def test_link_with_conflicting_dtype_creates_no_keys(tq_system): assert tq.kv_list(PARTITION)[PARTITION].keys() == {"a"} +def test_link_a_tensor_into_a_field_put_as_non_tensor(tq_system): + # A put of a non-tensor value turns the field non-tensor for the partition, and later tensor puts still land. + tq.kv_put(key="text", partition_id=PARTITION, fields={"prompt": "no image"}) + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + + tq.kv_link(["r0", "r1"], PARTITION, "p0", "prompt") + + data = tq.kv_batch_get(["text", "r0", "r1"], PARTITION)["prompt"] + assert data[0] == "no image" + assert [row.tolist() for row in data[1:]] == [_prompt(0).tolist()] * 2 + + def test_link_refuses_a_public_key_being_cleared(tq_system): tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) source = tq.get_client().kv_describe(["p0"], PUBLIC)["rows"]["p0"]["global_index"] diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index d91cb2d4..d2af1e92 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -624,6 +624,19 @@ def validate_field_schema(self, field_schema: dict[str, dict[str, Any]]) -> None if dtype is not None and existing.dtype is not None and dtype != existing.dtype: raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}") + def validate_link_schema(self, field_schema: dict[str, dict[str, Any]]) -> None: + """Reject only what linking would fail on: two tensor schemas of different dtypes. + + A link merges like a put does, so a non-tensor schema on either side makes the field non-tensor. + """ + for name, incoming in field_schema.items(): + existing = self.field_metadata.get(name) + if existing is None or existing.is_non_tensor or incoming.get("is_non_tensor", False): + continue + dtype = incoming.get("dtype") + if dtype is not None and existing.dtype is not None and dtype != existing.dtype: + raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}") + def _update_field_metadata( self, global_indexes: list[int], @@ -2112,7 +2125,7 @@ def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fi "kv_clear those fields before linking them" ) metas = {name: public.field_metadata[name] for name in fields} - partition.validate_field_schema( + partition.validate_link_schema( {name: {"dtype": meta.dtype, "is_non_tensor": meta.is_non_tensor} for name, meta in metas.items()} ) metadata = self.kv_retrieve_meta(keys, partition_id, create=True)