diff --git a/docs/data_dump.md b/docs/data_dump.md index 4365ff63..4857e35b 100644 --- a/docs/data_dump.md +++ b/docs/data_dump.md @@ -55,10 +55,12 @@ 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`, 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: ```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/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 572336a8..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) @@ -1389,7 +1395,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 new file mode 100644 index 00000000..ac18bee9 --- /dev/null +++ b/tests/e2e/test_kv_link_e2e.py @@ -0,0 +1,395 @@ +# 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 PUBLIC_PARTITION payload without copies.""" + +import shutil + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from tensordict import TensorDict + +import transfer_queue as tq +from transfer_queue.data_dump import _read_row_index + +PUBLIC = tq.PUBLIC_PARTITION +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_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, length: int = 3) -> torch.Tensor: + return torch.arange(p, p + length) + + +def _put_group(p: int) -> list[str]: + """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 + + +def _links(key: str) -> dict: + return tq.kv_describe(key, PARTITION)[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_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.""" + for i in range(4): + 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 ((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] + + +def test_linked_keys_read_the_public_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))] + 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): + 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="must not be empty"): + tq.kv_link([], PARTITION, "p0", "prompt") + 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)}) + + 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_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()}) + + 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_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"] + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + 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_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) + + 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") + 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("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:]) + + +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"] + + 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 + assert tq.kv_clear_unlinked_public() == [] + + +def test_link_refuses_a_key_marked_by_clear_unlinked_public(tq_system): + 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()) + + with pytest.raises(RuntimeError, match="must have produced"): + 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_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)) + + 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_refuses_a_linked_field(tq_system): + keys = _put_group(0) + client = tq.get_client() + meta = client.kv_retrieve_meta(keys, PARTITION) + + 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): + 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): + 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_linked_public_key_is_read_only(tq_system): + keys = _put_group(0) + + 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 + + # 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 + + +def test_dump_and_load_relink(tq_system, dump_test_root): + keys = _put_group(0) + 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) + + with pytest.raises(RuntimeError, match="does not exist|not found"): + 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) + 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 + + +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/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_controller_data_partitions.py b/tests/test_controller_data_partitions.py index fd458ed1..67e44568 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1492,3 +1492,77 @@ 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] + + +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 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 _public(self): + """Public row 0 holds ``x``.""" + from transfer_queue.controller import PUBLIC_PARTITION, DataPartitionStatus + + public = DataPartitionStatus(partition_id=PUBLIC_PARTITION) + assert public.update_production_status([0], ["x"], field_schema=self.SCHEMA) + return public + + def test_source_index_follows_link_put_and_clear(self): + from transfer_queue.controller import DataPartitionStatus + + public = self._public() + partition = DataPartitionStatus(partition_id="train") + partition.link_fields([1, 2], public, 0, ["x"]) + assert partition.linked_by == {0: {1, 2}} + 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 == {} + assert partition.linked_sources([0]) == set() + + 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)) + + partition.link_fields([1], self._public(), 0, ["x"]) + assert partition.linked_by == {0: {1}} diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py index b2de9cfe..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": [], "tag": tag}}, + "rows": {"key": {"global_index": 0, "fields": [], "unfinalized": [], "links": {}, "tag": tag}}, "field_schema": {}, } ) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 6fc49d2e..907342b9 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -256,6 +256,25 @@ 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_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/__init__.py b/transfer_queue/__init__.py index 05e63cdf..722e8ba9 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, @@ -23,7 +24,9 @@ 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, async_kv_put, async_kv_update, @@ -37,7 +40,9 @@ kv_batch_put, kv_batch_update, kv_clear, + kv_clear_unlinked_public, kv_describe, + kv_link, kv_list, kv_put, kv_update, @@ -76,6 +81,8 @@ "kv_list", "kv_describe", "kv_clear", + "kv_link", + "kv_clear_unlinked_public", "kv_update", "async_kv_put", "async_kv_batch_put", @@ -85,8 +92,11 @@ "async_kv_list", "async_kv_describe", "async_kv_clear", + "async_kv_link", + "async_kv_clear_unlinked_public", "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 73f458f4..af7d977f 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 ( @@ -412,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: @@ -475,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" @@ -1175,7 +1184,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( @@ -1186,6 +1195,48 @@ 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, + public_keys: list[str], + fields: list[str], + socket: zmq.asyncio.Socket | None = None, + ) -> None: + """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") + await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_LINK, + response_type=ZMQRequestType.KV_LINK_RESPONSE, + body={"keys": keys, "partition_id": partition_id, "public_keys": public_keys, "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. @@ -1501,6 +1552,8 @@ 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._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) @@ -1642,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: @@ -1978,6 +2032,14 @@ 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_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.""" + 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 c44ac236..d2af1e92 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 @@ -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. @@ -75,18 +79,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 +122,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 +170,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 +197,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]: """ @@ -367,6 +394,17 @@ 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 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": {}, "ever_linked": set(), **state}) + # Dynamic configuration - these are computed from the current state @property def total_samples_num(self) -> int: @@ -565,6 +603,9 @@ def update_production_status( # Save these global_indexes self.global_indexes.update(global_indices) + # 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 except Exception as e: @@ -583,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], @@ -941,6 +995,8 @@ 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) + 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) @@ -970,6 +1026,63 @@ 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], + 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 = public.field_metadata[name] + shape = meta.per_sample_shapes.get(source) if meta.is_nested else meta.shape + schema[name] = { + "dtype": meta.dtype, + "shape": shape, + "is_nested": False, + "is_non_tensor": meta.is_non_tensor, + } + # Before recording links: update_production_status drops links on the cells it marks. + 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}") + for idx in global_indexes: + row_links = self.links.setdefault(idx, {}) + for name in field_names: + row_links[name] = source + self.linked_by.setdefault(source, set()).update(global_indexes) + + 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 sources + for idx in self.linked_by.get(source, ()) + if field_names is None or any(self.links[idx].get(name) == source for name in field_names) + } + + 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 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] def kv_retrieve_indexes(self, keys: list[str]) -> list[int | None]: """Translate the user-specified keys to global_indexes""" @@ -1172,6 +1285,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}: " @@ -1507,6 +1629,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 +1644,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 +1670,9 @@ 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 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) if not partition: @@ -1553,6 +1686,37 @@ 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 or rewrite public payloads that keys of other partitions still link to.""" + 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 " + f"({len(linkers)} keys, e.g. {sorted(linkers)[:20]}); 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). @@ -1573,6 +1737,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) @@ -1657,6 +1823,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) @@ -1696,6 +1864,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: @@ -1749,7 +1921,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_clear those fields first to store private values" + ) metadata = self.kv_retrieve_meta(keys, partition_id, create=True) partition = self._get_partition(partition_id) assert partition is not None @@ -1777,6 +1965,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() + 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) @@ -1830,10 +2020,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 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. @@ -1847,8 +2038,10 @@ 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, {}) fields = sorted( field_name for field_name, field_meta in partition.field_metadata.items() @@ -1862,16 +2055,109 @@ 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: 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. 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} + 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``. + + ``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: A public key does not exist. + ValueError: ``partition_id`` is ``PUBLIC_PARTITION``, ``keys`` or ``fields`` is empty, + ``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) + 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. + 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: + # 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_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) + 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): + 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. + + 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(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) + 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() @@ -2099,6 +2385,8 @@ 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.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, } @@ -2304,6 +2592,17 @@ 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["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: + 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( @@ -2361,6 +2660,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.""" @@ -2414,9 +2715,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"]) diff --git a/transfer_queue/data_dump.py b/transfer_queue/data_dump.py index ad3791fb..e3b6fdbf 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 @@ -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" diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 0b4ab04b..7f6bc4b6 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 @@ -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 @@ -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 @@ -694,6 +696,59 @@ 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 | 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 + 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. ``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. + partition_id: Partition holding ``keys``; not ``PUBLIC_PARTITION``. + 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 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)) + + +def kv_clear_unlinked_public() -> 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. + + 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, @@ -725,7 +780,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: @@ -1165,8 +1221,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"}}``. ``fields`` lists every stored field of the - key and ``unfinalized`` the subset that is not ready yet. + ``{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. @@ -1181,7 +1239,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() } @@ -1223,6 +1282,23 @@ 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 | 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) + 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]: + """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: @@ -1480,7 +1556,10 @@ 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. Fields linked + 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: @@ -1508,6 +1587,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"]} @@ -1548,10 +1631,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.") @@ -1566,6 +1651,29 @@ 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") + # 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 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[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()} diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py index ab395e07..14c86a79 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. @@ -475,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 @@ -523,6 +536,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 +566,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 +579,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 +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=[ + {**{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 @@ -826,6 +846,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 +873,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..e017d44c 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,52 @@ 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, 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) + 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, 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={list(fields)}", partial( self._get_from_single_storage_unit, group.global_indexes, - metadata.field_names, + 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], ), 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 +876,19 @@ 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[tuple[tuple[str, ...], int], str], Any] = {} + for (_, fields, group), (_, 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, row in enumerate(group.batch_positions): + values[(fields, row), field] = su_data[field][i] + ordered_data = { + 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()} diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 9c765cfb..d2d525be 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -129,6 +129,10 @@ 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" + MARK_UNLINKED_PUBLIC_CLEARING = "MARK_UNLINKED_PUBLIC_CLEARING" + MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE = "MARK_UNLINKED_PUBLIC_CLEARING_RESPONSE" # METRICS GET_METRICS = "GET_METRICS"