From 0150f92b18e967f7d0bb86b4c01023289b217c9f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 14:58:19 +0800 Subject: [PATCH 01/38] [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. Signed-off-by: OutstanderWang --- 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 28ab8e1896c5107f4c20b406539a452ce3b8990f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:04:21 +0800 Subject: [PATCH 02/38] [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. Signed-off-by: OutstanderWang --- 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 8d7380e01f3d95d5c89428d6100f5211d31437cd Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:23:21 +0800 Subject: [PATCH 03/38] [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. Signed-off-by: OutstanderWang --- 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 c79e95fc11b4f3d6db927ba37c0e0d48b04b20a8 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:35:12 +0800 Subject: [PATCH 04/38] [perf] index kv_link links by source so clear checks scale with the clear Signed-off-by: OutstanderWang --- 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 260e1a04463b67dbb0dd4c37e92879115f375815 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:36:36 +0800 Subject: [PATCH 05/38] [fix] linked keys follow their source's readiness and shape Signed-off-by: OutstanderWang --- 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 3d1334f2b8e9217b06479dc4044ba1ceb523edce Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:37:34 +0800 Subject: [PATCH 06/38] [fix] refuse kv_update on fields linked by kv_link with a clear error Signed-off-by: OutstanderWang --- 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 a8b905310cc21204143645f1da9efecb6cbee565 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:38:41 +0800 Subject: [PATCH 07/38] [fix] refuse kv_link that relinks a source key or links a key to itself Signed-off-by: OutstanderWang --- 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 35854fa9027eb07dc9ad5cf64f12773c5b09bb2e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:40:27 +0800 Subject: [PATCH 08/38] [fix] a put through BatchMeta drops its local links on the written fields Signed-off-by: OutstanderWang --- 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 f9baa65072cb334b5260f83d1059cfb3cc1cdae6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:41:34 +0800 Subject: [PATCH 09/38] [feat] kv_describe reports the fields each key links by kv_link Signed-off-by: OutstanderWang --- 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 e5b24541602095ddeb0e64cb397ac09940e0f94a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:42:17 +0800 Subject: [PATCH 10/38] [doc] document and test that data_parser refuses fields linked by kv_link Signed-off-by: OutstanderWang --- 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 1990f79e7359fae529613f283c933639f28bf4ae Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:42:26 +0800 Subject: [PATCH 11/38] [doc] describe kv_link readiness, kv_update and source-sample behavior Signed-off-by: OutstanderWang --- 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 7bb47f8af5196686e910dda577dc0be43830fcc2 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:45:35 +0800 Subject: [PATCH 12/38] [fix] load controller checkpoints saved before kv_link Signed-off-by: OutstanderWang --- 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 e2c4ec5a845b79ec9a813c0eb0b5a0b88653b796 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:47:12 +0800 Subject: [PATCH 13/38] [fix] refuse dumping a linked source field that is not finalized Signed-off-by: OutstanderWang --- 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 25dbbbdb7009f8fd4679bb9e29f32a5ba29c9014 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 16 Oct 2026 15:47:47 +0800 Subject: [PATCH 14/38] [fix] BatchMeta.union takes links from the side whose fields win Signed-off-by: OutstanderWang --- 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 2d99b29eda75f6fd76703e8305de24aa17a8354e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 10:49:14 +0800 Subject: [PATCH 15/38] [feat] balance each partition's rows across SimpleStorage units Signed-off-by: OutstanderWang --- 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 498d7d137fe9c802e88a8c59d1c2207e8cfcfc1f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 10:54:03 +0800 Subject: [PATCH 16/38] [feat] kv_link reads shared payloads from a __public__ partition Signed-off-by: OutstanderWang --- 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 ccfe6bb9d51f2483fc2b42f9f72066bb8a7da4f1 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 11:41:43 +0800 Subject: [PATCH 17/38] feat: add kv_clear_unlinked_public to free public keys no key links to Signed-off-by: OutstanderWang --- 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 88d7ee215c446497d836fbb557c48e649cdfa314 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 12:07:28 +0800 Subject: [PATCH 18/38] [fix] drop links to public rows released while a racing put made them ready again Signed-off-by: OutstanderWang --- 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 0befabfc9a5994d722dccd86251c99c02f70c1e1 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 12:19:52 +0800 Subject: [PATCH 19/38] [fix] kv_link validates field types before creating keys Signed-off-by: OutstanderWang --- 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 bb0cc155a5069aeaad22d9113c4261d50b4419bc Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 19 Oct 2026 12:33:19 +0800 Subject: [PATCH 20/38] [test] controller checkpoints restore kv_link links and their clear guard Signed-off-by: OutstanderWang --- 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 c2addb839fb1213aac9db5f66ae9e086ae114620 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 12:52:43 +0800 Subject: [PATCH 21/38] [fix] repeated keys keep their own parser_args when reading Signed-off-by: OutstanderWang --- 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 674f388664b6190828954675e43135383ab0523e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 12:55:46 +0800 Subject: [PATCH 22/38] [fix] bump the dump format to 5 so older readers reject dumps with links Signed-off-by: OutstanderWang --- 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 50cf460381996c7116375d9c39953c0a745bd982 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:02:44 +0800 Subject: [PATCH 23/38] [fix] kv_link refuses an empty key list Signed-off-by: OutstanderWang --- 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 b8d560dddd1b1a2403f51ce5b36042c5f3cb642e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:02:47 +0800 Subject: [PATCH 24/38] [fix] kv_link refuses fields a key already stores itself Signed-off-by: OutstanderWang --- 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 734ba88cec25530d24299270eb92013d7f4062d0 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:03:40 +0800 Subject: [PATCH 25/38] [fix] check the SimpleStorage backend in TransferQueueClient.kv_link Signed-off-by: OutstanderWang --- 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 915c0652ba20a209e54db1b2fcdec3aafb4f0c0f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:04:41 +0800 Subject: [PATCH 26/38] [fix] load relinks before writing payloads so an incompatible link writes nothing Signed-off-by: OutstanderWang --- 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 e485b7b4f50994335fa769ac1c74d53f382cc392 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:07:08 +0800 Subject: [PATCH 27/38] [feat] a linked public key is read-only until no key links to it Signed-off-by: OutstanderWang --- 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 7d030ffa217dd914f6666ddc943b1e0326f85956 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:09:39 +0800 Subject: [PATCH 28/38] [feat] puts refuse fields linked by kv_link instead of replacing the link Signed-off-by: OutstanderWang --- 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 91d73a284e6dc03405f44c49b93b4a1e7b7e56df Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:12:07 +0800 Subject: [PATCH 29/38] [feat] kv_describe lists linked fields as readable, and link refusals name the linking keys Signed-off-by: OutstanderWang --- 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 8fa7890001aa08b275d0635fec008f91de1892ba Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:13:14 +0800 Subject: [PATCH 30/38] [feat] kv_clear_unlinked_public frees only public keys that were linked before Signed-off-by: OutstanderWang --- 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 1372eb9b984cc282b6e325d0a7930e4147c0ee83 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 13:15:00 +0800 Subject: [PATCH 31/38] [feat] kv_link takes one public key per key to link a batch in one atomic call Signed-off-by: OutstanderWang --- 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 9e371d046b25e715828a149564a1a2282e1cccdd Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 20 Oct 2026 22:28:11 +0800 Subject: [PATCH 32/38] [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. Signed-off-by: OutstanderWang --- 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) From 3304345b5d9e74df687bc0650d20abe3f356a5ae Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 22 Oct 2026 17:31:17 +0800 Subject: [PATCH 33/38] [feat] kv_exists reports ready fields without logging; kv_clear_unlinked_public frees stale never-linked keys kv_exists returns the ready fields of each existing key and omits missing keys, so a caller can probe for a key without the controller's not-found warnings. kv_clear_unlinked_public(stale_keys=...) also frees the listed public keys that no key ever linked, for producers whose every consumer went away. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_link_e2e.py | 34 ++++++++++++++++++++ transfer_queue/__init__.py | 4 +++ transfer_queue/client.py | 30 +++++++++++++---- transfer_queue/controller.py | 44 +++++++++++++++++++++++-- transfer_queue/interface.py | 53 ++++++++++++++++++++++++++----- transfer_queue/utils/zmq_utils.py | 2 ++ 6 files changed, 150 insertions(+), 17 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index ac18bee9..c156ac9c 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -15,6 +15,7 @@ """End-to-end tests for kv_link: n keys reading one PUBLIC_PARTITION payload without copies.""" +import asyncio import shutil import pytest @@ -144,6 +145,26 @@ def test_link_a_batch_of_public_keys_in_one_call(tq_system): assert [_links(key) for key in keys] == [{"prompt": "p0"}, {"prompt": "p0"}, {"prompt": "p1"}] +def test_exists_reports_ready_fields_and_omits_missing_keys(tq_system): + assert tq.kv_exists(["p0"], "never_created") == {} + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"prompt": _prompt(0)}) + tq.kv_put(key="p0", partition_id=PUBLIC, fields={"meta": _prompt(1)}) + tq.kv_link("r0", PARTITION, "p0", "prompt") + + assert tq.kv_exists(["p0", "missing", "p0"], PUBLIC) == {"p0": ["meta", "prompt"]} + assert tq.kv_exists("r0", PARTITION) == {"r0": ["prompt"]} + assert asyncio.run(tq.async_kv_exists(["missing", "p0"], PUBLIC)) == {"p0": ["meta", "prompt"]} + + +def test_exists_reports_no_ready_field_for_a_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], [PUBLIC])) + + assert tq.kv_exists("p0", PUBLIC) == {"p0": []} + + 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)}) @@ -252,6 +273,19 @@ def test_clear_unlinked_public_frees_only_keys_no_longer_linked(tq_system): assert tq.kv_clear_unlinked_public() == [] +def test_clear_unlinked_public_frees_stale_keys_no_key_links(tq_system): + keys0 = _put_group(0) + # Never linked, as when every waiter of a producer was cancelled before kv_link. + tq.kv_put(key="p1", partition_id=PUBLIC, fields={"prompt": _prompt(1)}) + tq.kv_put(key="p2", partition_id=PUBLIC, fields={"prompt": _prompt(2)}) + + assert tq.kv_clear_unlinked_public(stale_keys=["p0", "p1", "missing"]) == ["p1"] + + assert tq.kv_list(PUBLIC)[PUBLIC].keys() == {"p0", "p2"} + assert _rows(tq.kv_batch_get(keys0, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N + assert asyncio.run(tq.async_kv_clear_unlinked_public(stale_keys=["p2"])) == ["p2"] + + def test_link_refuses_a_key_marked_by_clear_unlinked_public(tq_system): keys = _put_group(0) tq.kv_clear(keys, PARTITION) diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 722e8ba9..25da4320 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -26,6 +26,7 @@ async_kv_clear, async_kv_clear_unlinked_public, async_kv_describe, + async_kv_exists, async_kv_link, async_kv_list, async_kv_put, @@ -42,6 +43,7 @@ kv_clear, kv_clear_unlinked_public, kv_describe, + kv_exists, kv_link, kv_list, kv_put, @@ -80,6 +82,7 @@ "kv_batch_get_by_meta", "kv_list", "kv_describe", + "kv_exists", "kv_clear", "kv_link", "kv_clear_unlinked_public", @@ -91,6 +94,7 @@ "async_kv_batch_get_by_meta", "async_kv_list", "async_kv_describe", + "async_kv_exists", "async_kv_clear", "async_kv_link", "async_kv_clear_unlinked_public", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index af7d977f..d180424a 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1195,6 +1195,22 @@ async def async_kv_describe( ) return {"rows": response.body["rows"], "field_schema": response.body["field_schema"]} + @with_controller_socket + async def async_kv_exists( + self, + keys: list[str], + partition_id: str, + socket: zmq.asyncio.Socket | None = None, + ) -> dict[str, list[str]]: + """Asynchronously map each existing key to its ready fields; see ``kv_exists``.""" + response = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_EXISTS, + response_type=ZMQRequestType.KV_EXISTS_RESPONSE, + body={"keys": keys, "partition_id": partition_id}, + ) + return response.body["ready"] + @with_controller_socket async def async_kv_link( self, @@ -1215,25 +1231,27 @@ async def async_kv_link( body={"keys": keys, "partition_id": partition_id, "public_keys": public_keys, "fields": fields}, ) - async def async_clear_unlinked_public(self) -> list[str]: + async def async_clear_unlinked_public(self, stale_keys: list[str] | None = None) -> 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() + metadata, keys = await self._mark_unlinked_public_clearing(stale_keys) 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]]: + async def _mark_unlinked_public_clearing( + self, stale_keys: list[str] | None, 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={}, + body={"stale_keys": stale_keys}, ) return response.body["metadata"], response.body["keys"] @@ -2036,9 +2054,9 @@ def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fi """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]: + def clear_unlinked_public(self, stale_keys: list[str] | None = None) -> list[str]: """Synchronously clear every ``PUBLIC_PARTITION`` key no key links to; return the cleared keys.""" - return self._clear_unlinked_public() + return self._clear_unlinked_public(stale_keys) # ==================== Selective Data Dump API ==================== diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index d2af1e92..e04630b1 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -1910,6 +1910,33 @@ def kv_retrieve_meta( return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch") + def kv_exists(self, keys: list[str], partition_id: str) -> dict[str, list[str]]: + """Map each existing key to its ready fields, linked ones included, without logging the missing. + + Unlike ``kv_retrieve_meta``, a missing key or partition is an expected answer here, so a + caller can probe keys that are usually absent without flooding the controller log. + + Args: + keys: Keys to look up. + partition_id: Partition holding ``keys``. + + Returns: + ``{key: sorted ready field names}`` for every key that exists; a missing key, or every + key of a missing partition, is left out. A key with no ready field maps to ``[]``. + """ + partition = self._get_partition(partition_id) + if partition is None: + return {} + ready: dict[str, list[str]] = {} + for key, global_index in zip(keys, partition.kv_retrieve_indexes(keys), strict=True): + if global_index is None: + continue + row = partition.production_status[global_index] + ready[key] = sorted( + name for name, col in partition.field_name_mapping.items() if col < len(row) and row[col] == 1 + ) + return ready + def kv_prepare_update( self, keys: list[str], partition_id: str, fields: list[str] ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]: @@ -2138,12 +2165,16 @@ def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fi partition.link_fields(indexes, public, source, fields) public.ever_linked.add(source) - def mark_unlinked_public_clearing(self) -> tuple[BatchMeta, list[str]]: + def mark_unlinked_public_clearing(self, stale_keys: list[str] | None = None) -> tuple[BatchMeta, list[str]]: """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. + Args: + stale_keys: Public keys the caller gave up waiting on; each one no key links to now + is marked too, even if none ever linked it. Missing keys are ignored. + Returns: The marked rows' BatchMeta, for storage deletion and clear_meta, and their keys. """ @@ -2151,7 +2182,8 @@ 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(public.revert_keys_mapping[idx] for idx in public.ever_linked - linked) + stale = {public.keys_mapping[key] for key in stale_keys or () if key in public.keys_mapping} + keys = sorted(public.revert_keys_mapping[idx] for idx in (public.ever_linked | stale) - linked) if not keys: return BatchMeta.empty(), [] metadata = self.kv_retrieve_meta(keys, PUBLIC_PARTITION) @@ -2382,6 +2414,7 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | ZMQRequestType.KV_RETRIEVE_KEYS: self._handle_kv_retrieve_keys_request, ZMQRequestType.KV_LIST: self._handle_kv_list_request, ZMQRequestType.KV_DESCRIBE: self._handle_kv_describe_request, + ZMQRequestType.KV_EXISTS: self._handle_kv_exists_request, 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, @@ -2585,6 +2618,11 @@ def _handle_kv_describe_request(self, request_msg: ZMQMessage) -> ZMQMessage: description = self.kv_describe(params["keys"], params["partition_id"]) return self._make_response(request_msg, ZMQRequestType.KV_DESCRIBE_RESPONSE, description) + def _handle_kv_exists_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + ready = self.kv_exists(keys=params["keys"], partition_id=params["partition_id"]) + return self._make_response(request_msg, ZMQRequestType.KV_EXISTS_RESPONSE, {"ready": ready}) + def _handle_kv_prepare_update_request(self, request_msg: ZMQMessage) -> ZMQMessage: params = request_msg.body metadata, merge_plan = self.kv_prepare_update(params["keys"], params["partition_id"], params["fields"]) @@ -2598,7 +2636,7 @@ def _handle_kv_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: 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() + metadata, keys = self.mark_unlinked_public_clearing(request_msg.body.get("stale_keys")) return self._make_response( request_msg, ZMQRequestType.MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE, {"metadata": metadata, "keys": keys} ) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 7f6bc4b6..875cd54e 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -668,6 +668,31 @@ def kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, return tq_client._run_coroutine(async_kv_describe(keys=keys, partition_id=partition_id)) +def kv_exists(keys: list[str] | str, partition_id: str) -> dict[str, list[str]]: + """Report which keys exist and the fields each has ready, without logging the missing ones. + + ``kv_describe`` raises and ``kv_batch_get`` warns for a missing key or partition; this + treats absence as an ordinary answer, for probing keys that are often not written yet, + e.g. whether another worker has already put a shared prompt. + + Args: + keys: Single key or list of keys to look up. + partition_id: Partition containing the keys; it need not exist. + + Returns: + ``{key: sorted ready field names}`` for every key that exists, linked fields included. + Missing keys are left out; a key with no ready field (e.g. being cleared) maps to ``[]``. + + Example: + >>> import transfer_queue as tq + >>> tq.init() + >>> ready = tq.kv_exists(keys=["prompt_0", "prompt_1"], partition_id=tq.PUBLIC_PARTITION) + >>> missing = [key for key in ("prompt_0", "prompt_1") if key not in ready] + """ + tq_client = _maybe_create_tq_client() + return tq_client._run_coroutine(async_kv_exists(keys=keys, partition_id=partition_id)) + + def kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None: """Clear key-value pairs from TransferQueue. @@ -732,21 +757,26 @@ def kv_link(keys: list[str] | str, partition_id: str, public_key: str | list[str tq_client._run_coroutine(async_kv_link(keys, partition_id, public_key, fields)) -def kv_clear_unlinked_public() -> list[str]: +def kv_clear_unlinked_public(stale_keys: list[str] | None = None) -> list[str]: """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 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. + between its ``kv_put`` and ``kv_link`` is safe; pass it in ``stale_keys`` once its + producer is given up on. 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. + + Args: + stale_keys: Public keys to clear as well if no key links to them now, even if none + ever did, e.g. the keys a previous call left behind. Missing keys are ignored. Example: >>> tq.kv_clear(trajectory_keys, partition_id="train") - >>> freed = tq.kv_clear_unlinked_public() + >>> freed = tq.kv_clear_unlinked_public(stale_keys=keys_left_by_the_previous_step) """ tq_client = _maybe_create_tq_client() - return tq_client._run_coroutine(async_kv_clear_unlinked_public()) + return tq_client._run_coroutine(async_kv_clear_unlinked_public(stale_keys)) def kv_update( @@ -1244,6 +1274,13 @@ async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[st } +async def async_kv_exists(keys: list[str] | str, partition_id: str) -> dict[str, list[str]]: + """Asynchronously report which keys exist and their ready fields. See ``kv_exists``.""" + keys = [keys] if isinstance(keys, str) else keys + tq_client = _maybe_create_tq_client() + return await tq_client.async_kv_exists(list(dict.fromkeys(keys)), partition_id) + + async def async_kv_clear(keys: list[str] | str, partition_id: str, fields: list[str] | str | None = None) -> None: """Asynchronously clear key-value pairs from TransferQueue. @@ -1293,10 +1330,10 @@ async def async_kv_link( 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]: +async def async_kv_clear_unlinked_public(stale_keys: list[str] | None = None) -> 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() + return await tq_client.async_clear_unlinked_public(stale_keys) def _single_update_batch(fields: dict[str, Any]) -> TensorDict: diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index d2d525be..33a9de56 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -119,6 +119,8 @@ class ZMQRequestType(ExplicitEnum): # KV_INTERFACE KV_RETRIEVE_META = "KV_RETRIEVE_META" KV_RETRIEVE_META_RESPONSE = "KV_RETRIEVE_META_RESPONSE" + KV_EXISTS = "KV_EXISTS" + KV_EXISTS_RESPONSE = "KV_EXISTS_RESPONSE" KV_RETRIEVE_KEYS = "KV_RETRIEVE_KEYS" KV_RETRIEVE_KEYS_RESPONSE = "KV_RETRIEVE_KEYS_RESPONSE" KV_LIST = "KV_LIST" From 7f44d64b85adde4c94db7753ec5530b8324aa78d Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 23 Oct 2026 11:48:11 +0800 Subject: [PATCH 34/38] [feat] BatchMeta rows can link, write and clear PUBLIC_PARTITION rows like keys - TransferQueueClient.link / async_link(metadata, public_metadata, fields) links rows by global index and returns a BatchMeta whose _links point at the public rows. - A BatchMeta put into PUBLIC_PARTITION is refused while other rows link to the target rows, so a linked public row is read-only on both write paths. - kv_clear_unlinked_public sweeps public rows by index, so rows without keys are freed too. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_link_e2e.py | 77 +++++++++++- transfer_queue/client.py | 81 ++++++++++++- transfer_queue/controller.py | 189 ++++++++++++++++++++++-------- transfer_queue/utils/zmq_utils.py | 4 + 4 files changed, 297 insertions(+), 54 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index c156ac9c..37f7babc 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -383,13 +383,11 @@ def test_linked_public_key_is_read_only(tq_system): write() assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N - # 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 = client.kv_retrieve_meta(keys, PARTITION) - assert [tuple(shape) for shape in meta.get_shapes("prompt")] == [(5,)] * N + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + client.put(TensorDict({"prompt": _prompt(5, 5).unsqueeze(0)}, batch_size=1), metadata=meta) + assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N def test_dump_and_load_relink(tq_system, dump_test_root): @@ -427,3 +425,72 @@ def test_load_with_an_incompatible_linked_field_writes_nothing(tq_system, dump_t shutil.rmtree(linked_dir) assert tq.kv_list(PARTITION)[PARTITION].keys() == {"other"} + + +def _meta_group(p: int): + """Put prompt ``p`` into the public partition and N response rows, all through a BatchMeta with no keys.""" + client = tq.get_client() + public_meta = client.put(TensorDict({"prompt": _prompt(p).unsqueeze(0)}, batch_size=1), partition_id=PUBLIC) + rows_meta = client.put(TensorDict({"response": torch.arange(N).unsqueeze(1)}, batch_size=N), partition_id=PARTITION) + return public_meta, rows_meta + + +def test_batch_meta_rows_read_a_public_row_put_through_a_batch_meta(tq_system): + client = tq.get_client() + public_meta, rows_meta = _meta_group(0) + + linked = client.link(rows_meta, public_meta.select_samples([0] * N), "prompt") + + data = client.get_data(linked) + assert _rows(data["prompt"]) == [_prompt(0).tolist()] * N + assert _rows(data["response"]) == [[i] for i in range(N)] + assert linked._links == [{"prompt": public_meta.global_indexes[0]}] * N + + +def test_batch_meta_put_refuses_a_linked_public_row_until_no_row_links_it(tq_system): + client = tq.get_client() + public_meta, rows_meta = _meta_group(0) + linked = client.link(rows_meta, public_meta.select_samples([0] * N), "prompt") + rewrite = TensorDict({"prompt": _prompt(7).unsqueeze(0)}, batch_size=1) + + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + client.put(rewrite, metadata=public_meta) + assert _rows(client.get_data(linked)["prompt"]) == [_prompt(0).tolist()] * N + + client.clear_samples(linked) + client.put(rewrite, metadata=public_meta) + assert _rows(client.get_data(public_meta)["prompt"]) == [_prompt(7).tolist()] + + +def test_unlinked_public_rows_without_keys_are_freed_once_their_rows_are_cleared(tq_system): + client = tq.get_client() + public_meta, rows_meta = _meta_group(0) + linked = client.link(rows_meta, public_meta.select_samples([0] * N), "prompt") + + with pytest.raises(RuntimeError, match="linked by keys of other partitions"): + client.clear_samples(public_meta) + + client.clear_samples(linked) + assert tq.kv_clear_unlinked_public() == [] + _, more_rows = _meta_group(1) + with pytest.raises(RuntimeError, match="not found"): + client.link(more_rows, public_meta.select_samples([0] * N), "prompt") + + +def test_batch_meta_link_refuses_what_kv_link_refuses(tq_system): + client = tq.get_client() + public_meta, rows_meta = _meta_group(0) + public_rows = public_meta.select_samples([0] * N) + + with pytest.raises(RuntimeError, match="must have produced"): + client.link(rows_meta, public_rows, "response") + with pytest.raises(RuntimeError, match="cannot link"): + client.link(public_meta, public_meta, "prompt") + with pytest.raises(RuntimeError, match=f"1 public rows for {N} rows"): + client.link(rows_meta, public_meta, "prompt") + with pytest.raises(ValueError, match=f"only partition {PUBLIC!r}"): + client.link(rows_meta, rows_meta, "prompt") + + client.put(TensorDict({"prompt": _prompt(5).unsqueeze(0).repeat(N, 1)}, batch_size=N), metadata=rows_meta) + with pytest.raises(RuntimeError, match="already store"): + client.link(rows_meta, public_rows, "prompt") diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d180424a..29618027 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -24,6 +24,7 @@ import zmq.asyncio from tensordict import TensorDict +from transfer_queue.controller import PUBLIC_PARTITION from transfer_queue.metadata import BatchMeta from transfer_queue.storage import StorageManagerFactory from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager @@ -415,7 +416,8 @@ async def async_put( Raises: 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 + RuntimeError: If storage operation fails, or if ``metadata`` names ``PUBLIC_PARTITION`` rows that + other rows link to: those stay read-only until no row links to them Example: >>> batch_size = 4 @@ -464,6 +466,7 @@ async def async_put( f"You may receive 2D tensors in key-value based backend." ) + inserted = metadata is None if metadata is None: if partition_id is None: raise ValueError("partition_id must be provided if metadata is not given") @@ -484,6 +487,13 @@ async def async_put( 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" ) + public_rows = [ + idx + for idx, pid in zip(metadata.global_indexes, metadata.partition_ids, strict=True) + if pid == PUBLIC_PARTITION + ] + if public_rows and not inserted: + await self._check_public_write(public_rows) with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.client_id}] async_put" @@ -1231,6 +1241,67 @@ async def async_kv_link( body={"keys": keys, "partition_id": partition_id, "public_keys": public_keys, "fields": fields}, ) + @with_controller_socket + async def async_link( + self, + metadata: BatchMeta, + public_metadata: BatchMeta, + fields: list[str] | str, + socket: zmq.asyncio.Socket | None = None, + ) -> BatchMeta: + """Asynchronously make ``fields`` of each row of ``metadata`` read the aligned row of ``public_metadata``. + + The BatchMeta counterpart of ``kv_link``: neither side needs a key, so rows put through a + BatchMeta, e.g. ``put(data, partition_id=PUBLIC_PARTITION)``, can be shared too. A linked + public row is read-only, and ``clear_samples`` refuses it, until no row links to it; + ``kv_clear_unlinked_public`` then frees it. + + Args: + metadata: Rows of one partition, not ``PUBLIC_PARTITION``, that will read ``fields``. + public_metadata: Rows of ``PUBLIC_PARTITION`` that produced ``fields``, aligned with ``metadata``. + fields: Single field name or list of field names to link. + + Returns: + ``metadata``'s rows over its fields and ``fields``, carrying the new links; reads and + puts through it see the links. + + Raises: + ValueError: ``metadata`` spans several partitions or ``public_metadata`` is not all in + ``PUBLIC_PARTITION``. + RuntimeError: The controller refused the link; see ``TransferQueueController.link``. + NotImplementedError: If the storage backend is not SimpleStorage. + """ + if not isinstance(self.storage_manager, AsyncSimpleStorageManager): + raise NotImplementedError(f"{type(self.storage_manager).__name__} does not support link") + fields = [fields] if isinstance(fields, str) else list(fields) + partition_ids = set(metadata.partition_ids) + if len(partition_ids) != 1: + raise ValueError(f"metadata must cover exactly one partition, got {sorted(partition_ids)}") + if set(public_metadata.partition_ids) != {PUBLIC_PARTITION}: + raise ValueError(f"public_metadata must cover only partition {PUBLIC_PARTITION!r}") + response = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.LINK, + response_type=ZMQRequestType.LINK_RESPONSE, + body={ + "global_indexes": metadata.global_indexes, + "partition_id": partition_ids.pop(), + "public_indexes": public_metadata.global_indexes, + "fields": fields, + "data_fields": list(dict.fromkeys([*metadata.field_names, *fields])), + }, + ) + return response.body["metadata"] + + @with_controller_socket + async def _check_public_write(self, global_indexes: list[int], socket=None) -> None: + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.CHECK_PUBLIC_WRITE, + response_type=ZMQRequestType.CHECK_PUBLIC_WRITE_RESPONSE, + body={"global_indexes": global_indexes}, + ) + async def async_clear_unlinked_public(self, stale_keys: list[str] | None = None) -> list[str]: """Asynchronously clear every ``PUBLIC_PARTITION`` key no key links to; see ``kv_clear_unlinked_public``. @@ -1572,6 +1643,7 @@ def wrapper(*args, **kwargs): 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._link = _make_sync(self.async_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) @@ -1715,7 +1787,8 @@ def put( Raises: 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 + RuntimeError: If storage operation fails, or if ``metadata`` names ``PUBLIC_PARTITION`` rows that + other rows link to: those stay read-only until no row links to them Example: >>> batch_size = 4 @@ -2054,6 +2127,10 @@ def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fi """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 link(self, metadata: BatchMeta, public_metadata: BatchMeta, fields: list[str] | str) -> BatchMeta: + """Synchronously make ``fields`` of ``metadata``'s rows read ``public_metadata``'s; see ``async_link``.""" + return self._link(metadata, public_metadata, fields) + def clear_unlinked_public(self, stale_keys: list[str] | None = None) -> list[str]: """Synchronously clear every ``PUBLIC_PARTITION`` key no key links to; return the cleared keys.""" return self._clear_unlinked_public(stale_keys) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index e04630b1..262c3b21 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -64,6 +64,13 @@ # shared by its n responses. Its rows are balanced over storage units like any partition's. PUBLIC_PARTITION = "__public__" + +def _fields_ready_in_all(partition: "DataPartitionStatus", global_indexes: list[int]) -> list[str]: + """Return the fields every one of ``global_indexes`` has produced, all fields for no rows.""" + col_mask = partition.production_status[global_indexes, :].sum(dim=0).reshape(-1) == len(global_indexes) + return [name for name, col in partition.field_name_mapping.items() if col < len(col_mask) and col_mask[col]] + + # 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. @@ -1900,14 +1907,7 @@ def kv_retrieve_meta( verified_global_indexes = [idx for idx in global_indexes if idx is not None] # must fetch fields that the requested samples all have - col_mask = partition.production_status[verified_global_indexes, :].sum(dim=0).reshape(-1) == len( - verified_global_indexes - ) - data_fields = [] - for field_name, col_idx in partition.field_name_mapping.items(): - if col_idx < len(col_mask) and col_mask[col_idx]: - data_fields.append(field_name) - + data_fields = _fields_ready_in_all(partition, verified_global_indexes) return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch") def kv_exists(self, keys: list[str], partition_id: str) -> dict[str, list[str]]: @@ -2117,49 +2117,123 @@ def kv_link(self, keys: list[str], partition_id: str, public_keys: list[str], fi 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. - 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}") + partition = self._get_partition(partition_id) + rows = [None] * len(keys) if partition is None else partition.kv_retrieve_indexes(keys) # Before creating keys, so an incompatible link leaves none behind. + unproduced, stored = self._link_conflicts(partition, rows, public, sources, fields) + if not fields or unproduced: + raise ValueError( + f"public keys {sorted({public_keys[i] for i in unproduced})[:20]} must have produced fields {fields}" + ) + if stored: + raise ValueError( + f"keys {[keys[i] for i in stored][:20]} already store some of fields {fields}; " + "kv_clear those fields before linking them" + ) + metadata = self.kv_retrieve_meta(keys, partition_id, create=True) + self._link_rows(partition_id, metadata.global_indexes, public, sources, fields) + + def link( + self, + global_indexes: list[int], + partition_id: str, + public_indexes: list[int], + fields: list[str], + data_fields: list[str], + ) -> BatchMeta: + """Make ``fields`` of these rows read the payload of the aligned rows of ``PUBLIC_PARTITION``. + + The BatchMeta counterpart of ``kv_link``: rows are named by global index, so neither side + needs a key. Every check runs before anything changes, so a refused call links nothing. + + Returns: + The linked rows' BatchMeta over ``data_fields``, carrying the new links. + + Raises: + KeyError: A row is not in its partition. + ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, the rows are empty, repeated or not + aligned, a public row has not produced all ``fields``, or a row stores one of them itself. + """ + if partition_id == PUBLIC_PARTITION: + raise ValueError(f"rows of partition {PUBLIC_PARTITION!r} hold shared payloads and cannot link") + if not global_indexes: + raise ValueError("rows must not be empty") + if len(public_indexes) != len(global_indexes): + raise ValueError(f"{len(public_indexes)} public rows for {len(global_indexes)} rows") + if len(set(global_indexes)) != len(global_indexes): + raise ValueError(f"rows must be unique, got {len(global_indexes) - len(set(global_indexes))} repeats") + public = self._get_partition(PUBLIC_PARTITION) + absent = sorted({idx for idx in public_indexes if public is None or idx not in public.global_indexes}) + if absent: + raise KeyError(f"rows {absent[:20]} not found in partition {PUBLIC_PARTITION!r}") 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" - ) + missing = sorted({idx for idx in global_indexes if partition is None or idx not in partition.global_indexes}) + if missing: + raise KeyError(f"rows {missing[:20]} not found in partition {partition_id!r}") + assert public is not None and partition is not None + unproduced, stored = self._link_conflicts(partition, global_indexes, public, public_indexes, fields) + if not fields or unproduced: + raise ValueError( + f"public rows {sorted({public_indexes[i] for i in unproduced})[:20]} must have produced fields {fields}" + ) + if stored: + raise ValueError( + f"rows {[global_indexes[i] for i in stored][:20]} already store some of fields {fields}; " + "clear those fields before linking them" + ) + self._link_rows(partition_id, global_indexes, public, public_indexes, fields) + return self.generate_batch_meta(partition_id, global_indexes, data_fields, mode="force_fetch") + + def _link_conflicts( + self, + partition: "DataPartitionStatus | None", + rows: list[int | None], + public: "DataPartitionStatus", + sources: list[int], + fields: list[str], + ) -> tuple[list[int], list[int]]: + """Return the positions whose source lacks a field, and those whose row stores one itself. + + Also refuses fields whose dtype the linking partition already holds differently. + """ + # Production status, not field metadata: mark_clearing zeroes it before storage deletes the payload. + unproduced = [ + pos + for pos, source in enumerate(sources) + 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 partition is None or unproduced or not fields: + return unproduced, [] + # Linking only redirects reads, so a value the row stores itself would stay in storage unread. + stored = [ + pos + for pos, idx in enumerate(rows) + 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 not stored: metas = {name: public.field_metadata[name] for name in fields} 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) + return unproduced, stored + + def _link_rows( + self, partition_id: str, rows: list[int], public: "DataPartitionStatus", sources: list[int], fields: list[str] + ) -> None: partition = self._get_partition(partition_id) assert partition is not None rows_by_source: dict[int, list[int]] = defaultdict(list) - for idx, source in zip(metadata.global_indexes, sources, strict=True): + for idx, source in zip(rows, sources, strict=True): rows_by_source[source].append(idx) for source, indexes in rows_by_source.items(): partition.link_fields(indexes, public, source, fields) @@ -2176,19 +2250,22 @@ def mark_unlinked_public_clearing(self, stale_keys: list[str] | None = None) -> is marked too, even if none ever linked it. Missing keys are ignored. Returns: - The marked rows' BatchMeta, for storage deletion and clear_meta, and their keys. + The marked rows' BatchMeta, for storage deletion and clear_meta, and the keys of those + rows; a row put through a BatchMeta has no key, so it is in the BatchMeta only. """ 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())) stale = {public.keys_mapping[key] for key in stale_keys or () if key in public.keys_mapping} - keys = sorted(public.revert_keys_mapping[idx] for idx in (public.ever_linked | stale) - linked) - if not keys: + indexes = sorted((public.ever_linked | stale) - linked) + if not indexes: return BatchMeta.empty(), [] - metadata = self.kv_retrieve_meta(keys, PUBLIC_PARTITION) + metadata = self.generate_batch_meta( + PUBLIC_PARTITION, indexes, _fields_ready_in_all(public, indexes), mode="force_fetch" + ) self.mark_clearing(metadata.global_indexes, metadata.partition_ids) - return metadata, keys + return metadata, sorted(public.revert_keys_mapping[idx] for idx in indexes if idx in public.revert_keys_mapping) def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" @@ -2420,6 +2497,8 @@ def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | 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.LINK: self._handle_link_request, + ZMQRequestType.CHECK_PUBLIC_WRITE: self._handle_check_public_write_request, ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request, ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: self._handle_load_controller_checkpoint_request, } @@ -2635,6 +2714,22 @@ def _handle_kv_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: 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_link_request(self, request_msg: ZMQMessage) -> ZMQMessage: + params = request_msg.body + metadata = self.link( + params["global_indexes"], + params["partition_id"], + params["public_indexes"], + params["fields"], + params["data_fields"], + ) + return self._make_response(request_msg, ZMQRequestType.LINK_RESPONSE, {"metadata": metadata}) + + def _handle_check_public_write_request(self, request_msg: ZMQMessage) -> ZMQMessage: + # A BatchMeta put names its rows directly, so it must be refused here, before storage is written. + self._check_public_not_linked(request_msg.body["global_indexes"]) + return self._make_response(request_msg, ZMQRequestType.CHECK_PUBLIC_WRITE_RESPONSE, {}) + def _handle_mark_unlinked_public_clearing_request(self, request_msg: ZMQMessage) -> ZMQMessage: metadata, keys = self.mark_unlinked_public_clearing(request_msg.body.get("stale_keys")) return self._make_response( diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 33a9de56..c5023b78 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -135,6 +135,10 @@ class ZMQRequestType(ExplicitEnum): KV_LINK_RESPONSE = "KV_LINK_RESPONSE" MARK_UNLINKED_PUBLIC_CLEARING = "MARK_UNLINKED_PUBLIC_CLEARING" MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE = "MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE" + LINK = "LINK" + LINK_RESPONSE = "LINK_RESPONSE" + CHECK_PUBLIC_WRITE = "CHECK_PUBLIC_WRITE" + CHECK_PUBLIC_WRITE_RESPONSE = "CHECK_PUBLIC_WRITE_RESPONSE" # METRICS GET_METRICS = "GET_METRICS" From a0c339749ab8dd76a76138da0098bdd5ac9e00c0 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 23 Oct 2026 12:45:05 +0800 Subject: [PATCH 35/38] [feat] kv_batch_get data_parser applies to linked fields on the public row's storage unit A linked field used to refuse a data_parser. It is now read from its public row once per distinct plain parser_args entry (scalars and lists of them) among the keys linking to it, and the public row's storage unit parses it, so only the parsed part is sent back. Any other parser_args entry, such as a tensor, is parsed once per key. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_link_e2e.py | 32 +++++++++--- tests/test_async_simple_storage_manager.py | 52 +++++++++++++++++++ transfer_queue/interface.py | 7 +-- .../managers/simple_storage_manager.py | 49 +++++++++++++---- 4 files changed, 120 insertions(+), 20 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 37f7babc..5a6e8f05 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -343,19 +343,39 @@ def test_put_through_meta_refuses_a_linked_field(tq_system): assert _rows(tq.kv_batch_get(keys, PARTITION)["prompt"]) == [_prompt(0).tolist()] * N -def test_data_parser_refuses_linked_fields(tq_system): +def _per_value(parse): + """A data_parser applying ``parse(value, arg)`` to each value of every field.""" + return lambda data, args: { + field: [parse(value, arg) for value, arg in zip(values, args, strict=True)] for field, values in data.items() + } + + +def test_data_parser_slices_a_linked_field_per_key(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 + data = tq.kv_batch_get( + keys, PARTITION, data_parser=_per_value(lambda value, length: value[:length]), parser_args=[1, 3, 1] + ) - 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["prompt"]) == [_prompt(0)[:1].tolist(), _prompt(0).tolist(), _prompt(0)[:1].tolist()] assert _rows(data["response"]) == [[i] for i in range(N)] +def test_data_parser_selects_elements_of_a_linked_field_per_key(tq_system): + keys = _put_group(0) + _put_group(10) + + data = tq.kv_batch_get( + keys, + PARTITION, + select_fields="prompt", + data_parser=_per_value(lambda value, positions: value[positions]), + parser_args=[[0, 2], [1], [2, 2, 0], [0], [0], [1, 2]], + ) + + assert _rows(data["prompt"]) == [[0, 2], [1], [2, 2, 0], [10], [10], [11, 12]] + + def test_update_refuses_a_linked_field(tq_system): keys = _put_group(0) diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py index a5389db1..bc3c4136 100644 --- a/tests/test_async_simple_storage_manager.py +++ b/tests/test_async_simple_storage_manager.py @@ -25,6 +25,7 @@ from transfer_queue.metadata import BatchMeta from transfer_queue.storage import AsyncSimpleStorageManager +from transfer_queue.storage.managers.simple_storage_manager import _parser_arg_identity from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo @@ -664,3 +665,54 @@ def test_all_none_to_nontensorstack(self): assert isinstance(result, NonTensorStack) assert result.tolist() == [None, None] + + +@pytest.mark.asyncio +async def test_get_data_parses_a_linked_field_on_its_source_unit_once_per_distinct_arg(): + manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "test_parser_links" + manager.storage_unit_infos = { + su: ZMQServerInfo(role=Role.STORAGE, id=su, ip="127.0.0.1", ports={"put_get_socket": port}) + for su, port in (("storage_0", 19014), ("storage_1", 19015)) + } + non_tensor = {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True} + # The linking rows are odd, so on storage_1; their public source 4 is on storage_0. + batch_meta = BatchMeta( + global_indexes=[1, 3, 5, 7], + partition_ids=["p0"] * 4, + field_schema={"images": non_tensor, "response": non_tensor}, + production_status=np.ones(4, dtype=np.int8), + _links=[{"images": 4}] * 4, + ) + tensor_arg = torch.tensor([1]) + parser_args = [[0], [1], [0], tensor_arg] + requests = [] + + async def fake_get(global_indexes, fields, target_storage_unit=None, data_parser=None, parser_args=None, **_): + requests.append((target_storage_unit, list(global_indexes), list(fields), parser_args)) + rows = list(zip(global_indexes, parser_args, strict=True)) + return fields, {f: [f"{f}_{gi}_{arg}" for gi, arg in rows] for f in fields} + + manager._get_from_single_storage_unit = fake_get + + def parser(data, args): + return data + + data = await manager.get_data(batch_meta, data_parser=parser, parser_args=parser_args) + + assert sorted(requests, key=lambda r: r[0]) == [ + ("storage_0", [4, 4, 4], ["images"], [[0], [1], tensor_arg]), + ("storage_1", [1, 3, 5, 7], ["response"], parser_args), + ] + assert list(data["images"]) == ["images_4_[0]", "images_4_[1]", "images_4_[0]", f"images_4_{tensor_arg}"] + assert list(data["response"]) == [f"response_{gi}_{arg}" for gi, arg in zip([1, 3, 5, 7], parser_args, strict=True)] + + +def test_parser_arg_identity_shares_only_equal_plain_args(): + assert _parser_arg_identity([0, 2], 0) == _parser_arg_identity([0, 2], 1) + assert _parser_arg_identity(None, 0) == _parser_arg_identity(None, 1) + assert _parser_arg_identity([0, 2], 0) != _parser_arg_identity((0, 2), 1) + assert _parser_arg_identity(1, 0) != _parser_arg_identity(True, 1) + assert _parser_arg_identity(1, 0) != _parser_arg_identity(1.0, 1) + assert _parser_arg_identity(torch.tensor([0]), 0) != _parser_arg_identity(torch.tensor([0]), 1) + assert _parser_arg_identity([torch.tensor([0])], 0) != _parser_arg_identity([torch.tensor([0])], 1) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 875cd54e..21e5a116 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -561,8 +561,10 @@ 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``, and so - does a selected field linked by ``kv_link``. + A field linked by ``kv_link`` is parsed on its public row's unit, which holds it once + per distinct plain ``parser_args`` entry (scalars and lists of them) among the keys + linking to it, so only the parsed part is sent back. Only supported by SimpleStorage; + KV backends raise ``NotImplementedError``. 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. @@ -577,7 +579,6 @@ 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 diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index e017d44c..337e7ab0 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -63,6 +63,28 @@ TQ_SIMPLE_STORAGE_DUMP_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_DUMP_TIMEOUT", 3600)) +_SHAREABLE_PARSER_ARG_TYPES = (type(None), bool, int, float, str) + + +def _parser_arg_identity(parser_arg: Any, pos: int) -> tuple: + """Return a key equal for two plain parser_args (scalars and lists of them) of equal type and value. + + Any other argument, such as a tensor, gets one unique to batch position ``pos``: a linked + row is then parsed once per such key instead of being shared across keys. + """ + + def identity(value: Any) -> tuple | None: + if isinstance(value, _SHAREABLE_PARSER_ARG_TYPES): + return (type(value), value) + if isinstance(value, list | tuple): + items = [identity(item) for item in value] + return None if None in items else (type(value), tuple(items)) + return None + + key = identity(parser_arg) + return ("position", pos) if key is None else ("value", key) + + class StorageUnitTimeout(RuntimeError): """A storage unit did not answer within the send/recv timeout. @@ -801,7 +823,8 @@ async def get_data( Retrieve data from remote StorageUnit based on metadata. 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. + from its source row, once however many samples link to it, or with ``parser_args`` once + per distinct plain parser_arg of those samples. Args: metadata: BatchMeta that contains metadata for data retrieval. @@ -819,30 +842,32 @@ async def get_data( if metadata.size == 0: return TensorDict({}, batch_size=0) - # 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, 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. + # keep their parser_args, and a linked field once per source row and parser_arg however + # many link to it, so the source's unit parses it and replies with the parsed part only. rows_by_fields: dict[tuple[str, ...], list[int]] = defaultdict(list) + args_by_fields: dict[tuple[str, ...], list[Any]] = defaultdict(list) 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 + linked_slots: dict[tuple, tuple[tuple[str, ...], int]] = {} # (source, field, arg) -> slot for pos, (global_index, links) in enumerate(zip(metadata.global_indexes, metadata._links, strict=True)): + parser_arg = None if parser_args is None else parser_args[pos] 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) + args_by_fields[own].append(parser_arg) 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,])) + linked = (links[field], field, _parser_arg_identity(parser_arg, pos)) + if linked not in linked_slots: + linked_slots[linked] = ((field,), len(rows_by_fields[field,])) rows_by_fields[field,].append(links[field]) - slots[pos, field] = linked_slots[links[field], field] + args_by_fields[field,].append(parser_arg) + slots[pos, field] = linked_slots[linked] requests = [ (su_id, fields, group) for fields, rows in rows_by_fields.items() @@ -860,7 +885,9 @@ async def get_data( 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], + parser_args=None + if parser_args is None + else [args_by_fields[fields][p] for p in group.batch_positions], ), max_attempts=max_attempts, ) From 1c48464910a759192536ee4364babc91dcf3bc2d Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 23 Oct 2026 14:39:07 +0800 Subject: [PATCH 36/38] [chore] pass the docstring sanity check Signed-off-by: OutstanderWang --- transfer_queue/controller.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 262c3b21..ff837cc8 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -98,6 +98,7 @@ def __init__(self): @property def reusable_indexes(self) -> list[int]: + """Every reusable index, unit by unit and FIFO within each unit.""" 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: From a22f20b50eaf32f9c77367f7cdc1bb67354e8a3a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 23 Oct 2026 18:23:48 +0800 Subject: [PATCH 37/38] [fix] type _link_conflicts rows as a Sequence so mypy accepts list[int] Signed-off-by: OutstanderWang --- transfer_queue/controller.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index ff837cc8..d500c027 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -18,6 +18,7 @@ import pickle import time from collections import defaultdict, deque +from collections.abc import Sequence from dataclasses import dataclass, field from itertools import groupby from operator import itemgetter @@ -2188,7 +2189,7 @@ def link( def _link_conflicts( self, partition: "DataPartitionStatus | None", - rows: list[int | None], + rows: Sequence[int | None], public: "DataPartitionStatus", sources: list[int], fields: list[str], From 590d873680ea50d106473dea5f0140726e3a5e1c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 25 Oct 2026 16:04:38 +0800 Subject: [PATCH 38/38] [fix] linked fields from one public row are read together; link keeps caller metadata - get_data reads the fields a row links to one public row in one unit request, so a data_parser sees all of them together, as it does when reading the public row itself. - TransferQueueClient.link returns the controller's links and schema with the caller's extra_info and custom_meta, which the controller does not hold. - put / async_put, async_link and BatchMeta.select_fields document that the written fields are data.keys(), that select_fields filters only local metadata, and that a public-row overwrite must not overlap with linking that row. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_link_e2e.py | 31 ++++++++++++++ tests/test_async_simple_storage_manager.py | 35 ++++++++++++++++ transfer_queue/client.py | 41 +++++++++++++++++-- transfer_queue/metadata.py | 4 ++ .../managers/simple_storage_manager.py | 23 +++++++---- 5 files changed, 122 insertions(+), 12 deletions(-) diff --git a/tests/e2e/test_kv_link_e2e.py b/tests/e2e/test_kv_link_e2e.py index 5a6e8f05..d5ecc61f 100644 --- a/tests/e2e/test_kv_link_e2e.py +++ b/tests/e2e/test_kv_link_e2e.py @@ -514,3 +514,34 @@ def test_batch_meta_link_refuses_what_kv_link_refuses(tq_system): client.put(TensorDict({"prompt": _prompt(5).unsqueeze(0).repeat(N, 1)}, batch_size=N), metadata=rows_meta) with pytest.raises(RuntimeError, match="already store"): client.link(rows_meta, public_rows, "prompt") + + +def test_data_parser_sees_every_field_a_row_links_to_one_public_row(tq_system): + client = tq.get_client() + public_meta = client.put( + TensorDict({"prompt": torch.tensor([[1, 2, 3]]), "length": torch.tensor([[2]])}, batch_size=1), + partition_id=PUBLIC, + ) + rows_meta = client.put(TensorDict({"response": torch.tensor([[0]])}, batch_size=1), partition_id=PARTITION) + linked = client.link(rows_meta, public_meta, ["prompt", "length"]).select_fields(["prompt", "length"]) + + def trim_prompt_to_length(data, args): + prompts = zip(data["prompt"], data["length"], strict=True) + return {"prompt": [value[: int(length.item())] for value, length in prompts], "length": data["length"]} + + data = client.get_data(linked, data_parser=trim_prompt_to_length) + + assert _rows(data["prompt"]) == [[1, 2]] + assert _rows(data["length"]) == [[2]] + + +def test_batch_meta_link_keeps_the_callers_extra_info_and_custom_meta(tq_system): + client = tq.get_client() + public_meta, rows_meta = _meta_group(0) + rows_meta.extra_info["epoch"] = 3 + rows_meta.update_custom_meta([{"group": f"g{i}"} for i in range(N)]) + + linked = client.link(rows_meta, public_meta.select_samples([0] * N), "prompt") + + assert linked.extra_info == {"epoch": 3} + assert linked.custom_meta == [{"group": f"g{i}"} for i in range(N)] diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py index bc3c4136..42e7609a 100644 --- a/tests/test_async_simple_storage_manager.py +++ b/tests/test_async_simple_storage_manager.py @@ -716,3 +716,38 @@ def test_parser_arg_identity_shares_only_equal_plain_args(): assert _parser_arg_identity(1, 0) != _parser_arg_identity(1.0, 1) assert _parser_arg_identity(torch.tensor([0]), 0) != _parser_arg_identity(torch.tensor([0]), 1) assert _parser_arg_identity([torch.tensor([0])], 0) != _parser_arg_identity([torch.tensor([0])], 1) + + +@pytest.mark.asyncio +async def test_get_data_reads_the_fields_a_row_links_to_one_source_together(): + manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "test_links_by_source" + manager.storage_unit_infos = { + su: ZMQServerInfo(role=Role.STORAGE, id=su, ip="127.0.0.1", ports={"put_get_socket": port}) + for su, port in (("storage_0", 19016), ("storage_1", 19017)) + } + non_tensor = {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True} + # Rows 1 and 3 link prompt and length to source 4; row 5 links them to sources 4 and 6. + batch_meta = BatchMeta( + global_indexes=[1, 3, 5], + partition_ids=["p0"] * 3, + field_schema={"prompt": non_tensor, "length": non_tensor}, + production_status=np.ones(3, dtype=np.int8), + _links=[{"prompt": 4, "length": 4}, {"prompt": 4, "length": 4}, {"prompt": 4, "length": 6}], + ) + requests = [] + + async def fake_get(global_indexes, fields, target_storage_unit=None, data_parser=None, parser_args=None, **_): + requests.append((list(global_indexes), sorted(fields))) + return fields, {f: [f"{f}_{gi}" for gi in global_indexes] for f in fields} + + manager._get_from_single_storage_unit = fake_get + + def parser(data, args): + return data + + data = await manager.get_data(batch_meta, data_parser=parser) + + assert sorted(requests) == [([4], ["length", "prompt"]), ([4], ["prompt"]), ([6], ["length"])] + assert list(data["prompt"]) == ["prompt_4"] * 3 + assert list(data["length"]) == ["length_4", "length_4", "length_6"] diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 29618027..4a5dd775 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -409,6 +409,18 @@ async def async_put( Do not change the inner order of values within each column. Only supported by SimpleStorage. + Warning: + The fields written are ``data.keys()``; ``metadata.field_names`` does not restrict + them, and ``BatchMeta.select_fields`` filters neither ``data`` nor the controller's + links. A linked field is refused only by the links ``metadata`` still carries, so a + put through ``linked.select_fields([...])`` of a batch read through ``linked`` stores + private copies of its linked fields and drops those links. To write back a batch + holding linked fields, select the fields to write from ``data``, e.g. + ``put(batch.select("response"), metadata=linked)``. + + Do not overwrite a ``PUBLIC_PARTITION`` row while any worker or client may be linking + to it: the check before writing does not reserve the row; see ``async_link``. + Returns: BatchMeta: The metadata used for the put operation (currently returns the input metadata or auto-retrieved metadata; will be updated in a future version to reflect the post-put state) @@ -1261,9 +1273,16 @@ async def async_link( public_metadata: Rows of ``PUBLIC_PARTITION`` that produced ``fields``, aligned with ``metadata``. fields: Single field name or list of field names to link. + Warning: + Do not overwrite a ``PUBLIC_PARTITION`` row while any worker or client may be linking + to it. A put checks the row's links before writing storage but does not reserve it, + so a link accepted in between reads the old payload and then the new one. Write each + public row once before linking it; to replace one, wait until no row links to it, + overwrite it, wait for the put to return, and only then link to it again. + Returns: - ``metadata``'s rows over its fields and ``fields``, carrying the new links; reads and - puts through it see the links. + ``metadata``'s rows over its fields and ``fields``, carrying the new links and keeping + ``metadata``'s ``extra_info`` and ``custom_meta``; reads and puts through it see the links. Raises: ValueError: ``metadata`` spans several partitions or ``public_metadata`` is not all in @@ -1291,7 +1310,11 @@ async def async_link( "data_fields": list(dict.fromkeys([*metadata.field_names, *fields])), }, ) - return response.body["metadata"] + linked: BatchMeta = response.body["metadata"] + # The controller knows only synced custom_meta, and extra_info never leaves the caller. + linked.extra_info = dict(metadata.extra_info) + linked.custom_meta = [dict(meta) for meta in metadata.custom_meta] + return linked @with_controller_socket async def _check_public_write(self, global_indexes: list[int], socket=None) -> None: @@ -1780,6 +1803,18 @@ def put( Do not change the inner order of values within each column. Only supported by SimpleStorage. + Warning: + The fields written are ``data.keys()``; ``metadata.field_names`` does not restrict + them, and ``BatchMeta.select_fields`` filters neither ``data`` nor the controller's + links. A linked field is refused only by the links ``metadata`` still carries, so a + put through ``linked.select_fields([...])`` of a batch read through ``linked`` stores + private copies of its linked fields and drops those links. To write back a batch + holding linked fields, select the fields to write from ``data``, e.g. + ``put(batch.select("response"), metadata=linked)``. + + Do not overwrite a ``PUBLIC_PARTITION`` row while any worker or client may be linking + to it: the check before writing does not reserve the row; see ``async_link``. + Returns: BatchMeta: The metadata used for the put operation (currently returns the input metadata or auto-retrieved metadata; will be updated in a future version to reflect the post-put state) diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py index 14c86a79..46158c90 100644 --- a/transfer_queue/metadata.py +++ b/transfer_queue/metadata.py @@ -548,6 +548,10 @@ def select_fields(self, field_names: list[str]) -> "BatchMeta": Returns: BatchMeta: A new BatchMeta instance containing only the specified fields. + + Note: + This filters only this local metadata, links included; it neither unlinks fields on the + controller nor selects the data a put writes. Select those fields from the data too. """ new_field_schema = {} for field_name in field_names: diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 337e7ab0..d14e67c3 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -847,12 +847,13 @@ async def get_data( # 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 and parser_arg however - # many link to it, so the source's unit parses it and replies with the parsed part only. + # keep their parser_args, and the fields it links to one source row together, once per + # source, field set and parser_arg however many link to it. The source's unit then + # parses those fields in one call and replies with the parsed part only. rows_by_fields: dict[tuple[str, ...], list[int]] = defaultdict(list) args_by_fields: dict[tuple[str, ...], list[Any]] = defaultdict(list) slots: dict[tuple[int, str], tuple[tuple[str, ...], int]] = {} # (sample, field) -> (fields, row) - linked_slots: dict[tuple, tuple[tuple[str, ...], int]] = {} # (source, field, arg) -> slot + linked_slots: dict[tuple, tuple[tuple[str, ...], int]] = {} # (source, fields, arg) -> slot for pos, (global_index, links) in enumerate(zip(metadata.global_indexes, metadata._links, strict=True)): parser_arg = None if parser_args is None else parser_args[pos] own = tuple(f for f in metadata.field_names if f not in links) @@ -860,14 +861,18 @@ async def get_data( slots.update({(pos, f): (own, len(rows_by_fields[own])) for f in own}) rows_by_fields[own].append(global_index) args_by_fields[own].append(parser_arg) + fields_by_source: dict[int, list[str]] = defaultdict(list) for field in metadata.field_names: if field in links: - linked = (links[field], field, _parser_arg_identity(parser_arg, pos)) - if linked not in linked_slots: - linked_slots[linked] = ((field,), len(rows_by_fields[field,])) - rows_by_fields[field,].append(links[field]) - args_by_fields[field,].append(parser_arg) - slots[pos, field] = linked_slots[linked] + fields_by_source[links[field]].append(field) + for source, source_fields in fields_by_source.items(): + fields = tuple(source_fields) + linked = (source, fields, _parser_arg_identity(parser_arg, pos)) + if linked not in linked_slots: + linked_slots[linked] = (fields, len(rows_by_fields[fields])) + rows_by_fields[fields].append(source) + args_by_fields[fields].append(parser_arg) + slots.update({(pos, f): linked_slots[linked] for f in fields}) requests = [ (su_id, fields, group) for fields, rows in rows_by_fields.items()