diff --git a/README.md b/README.md index 4df55109..a5745395 100644 --- a/README.md +++ b/README.md @@ -116,15 +116,26 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API - **(async_)kv_put**: Insert/Update a multi-column sample by key, with an optional metadata tag. - **(async_)kv_batch_put**: Put multiple key-value pairs efficiently in batches. - **(async_)kv_batch_get**: Retrieve samples (by keys), supporting column selection (by fields). -- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into already-produced fields with `merge_fn(old, new)`. SimpleStorage only. +- **(async_)kv_update / (async_)kv_batch_update**: Merge new values into fields with `merge_fn(old, new)`, adding missing keys and fields. SimpleStorage only. - **(async_)kv_list**: List keys and tags (metadata) in a partition. +- **(async_)kv_describe**: Show each key's stored fields, including those `kv_update(finalize=False)` keeps hidden. - **(async_)kv_clear**: Remove key-value pairs from storage, or only selected `fields` while keeping the keys. -`kv_update` preserves each field's tensor/non-tensor type and tensor dtype. A -timeout has an unknown outcome because a storage unit may commit after the -caller stops waiting, so non-idempotent merge operations must not be retried -blindly. Updates are atomic within one storage-unit request; a batch spanning -multiple units is not a distributed transaction. +`kv_update` merges only the rows the controller records as holding the field +(including hidden ones); every other row takes the new value as given, without +calling `merge_fn`, and overwrites anything a storage unit kept after a failed +clear. Every stored value must keep the field's tensor/non-tensor type and +tensor dtype in the partition. A timeout has an unknown outcome because a +storage unit may commit after the caller stops waiting, so non-idempotent merge +operations must not be retried blindly. Updates are atomic within one +storage-unit request; a batch spanning multiple units is not a distributed +transaction. + +To build a field over several updates without exposing intermediate values, pass +`finalize=False` to every `kv_update` except the last one, which uses the default +`finalize=True`. Until then the field is stored but not ready, so samplers, the +dataloader and `kv_batch_get` do not see it. Updates to one key must come from +one writer at a time; use `kv_global_lock` if several processes write the same key. **Key Features** diff --git a/docs/data_dump.md b/docs/data_dump.md index f3829b5b..4365ff63 100644 --- a/docs/data_dump.md +++ b/docs/data_dump.md @@ -55,7 +55,10 @@ These count application reads, not filesystem read-ahead or physical disk traffi ## Format and compatibility -Dumps use `format_version: 3`, the only version this build reads: +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: ```text dump_info.json diff --git a/tests/e2e/test_data_dump_e2e.py b/tests/e2e/test_data_dump_e2e.py index a4fd5b9c..c60b50bf 100644 --- a/tests/e2e/test_data_dump_e2e.py +++ b/tests/e2e/test_data_dump_e2e.py @@ -243,6 +243,54 @@ def test_heterogeneous_field_sets_are_grouped(self, tq_system, dump_dir, control retrieved = tq.kv_batch_get(keys=["h1"], partition_id=partition_id, select_fields=["routed_experts"]) _assert_rows_equal(retrieved["routed_experts"], [torch.tensor([7, 8])]) + def test_restore_keeps_unfinalized_fields_hidden(self, tq_system, dump_dir, controller): + partition_id = "d_unfinalized" + _put_rows(partition_id, ["u0", "u1"]) + + def concat(old, new): + return torch.cat([old, new]) + + tq.kv_update( + key="u1", + partition_id=partition_id, + fields={"response": torch.tensor([1, 2])}, + merge_fn=concat, + finalize=False, + ) + + tq.dump_data_by_key(dump_dir, ["u0", "u1"], partition_id) + assert _read_row_index(dump_dir)["rows"]["u1"]["unfinalized"] == ["response"] + ray.get(controller.clear_partition.remote(partition_id)) + tq.load_data_by_key(dump_dir) + + snapshot = ray.get(controller.get_partition_snapshot.remote(partition_id)) + row = snapshot.keys_mapping["u1"] + assert snapshot.production_status[row, snapshot.field_name_mapping["input_ids"]] == 1 + assert snapshot.production_status[row, snapshot.field_name_mapping["response"]] == 0 + + tq.kv_update(key="u1", partition_id=partition_id, fields={"response": torch.tensor([3])}, merge_fn=concat) + 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" + _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. + row_index = _read_row_index(dump_dir) + 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_path.write_text(json.dumps(info), encoding="utf-8") + + ray.get(controller.clear_partition.remote(partition_id)) + tq.load_data_by_key(dump_dir) + retrieved = tq.kv_batch_get(keys="w0", partition_id=partition_id, select_fields="input_ids") + _assert_rows_equal(retrieved["input_ids"], [_row_input_ids(0)]) + def test_other_partitions_are_untouched(self, tq_system, dump_dir, controller): """Restoring merges by key, unlike a checkpoint load which replaces everything.""" # Define test data @@ -681,10 +729,10 @@ async def publish_all(): started, ready = [], asyncio.Event() async def load(shards, target_storage_unit): - updates = [{"global_indexes": [i], "field_schema": {}} for i in range(num_updates)] + updates = [{"global_indexes": [i], "field_schema": {}, "ready": True} for i in range(num_updates)] return {"bytes_read": 0, "updates": updates} - async def notify(partition_id, global_indexes, field_schema): + async def notify(partition_id, global_indexes, field_schema, **_kwargs): started.append(global_indexes) if len(started) == num_updates: ready.set() diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index fa5d488b..8edc63c4 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -1244,25 +1244,62 @@ def test_batch_update_multiple_fields_across_storage_units(self, controller, tq_ assert "score" not in partition.field_metadata assert set(tq_api.kv_list(partition_id=partition_id)[partition_id]) == set(keys) - def test_rejects_missing_key_unproduced_field_and_incompatible_result(self, tq_api, backend_name): + def test_unfinalized_updates_stay_hidden_until_finalize(self, controller, tq_api, backend_name): if backend_name != "SimpleStorage": pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") - with pytest.raises(ValueError, match="not found"): - tq_api.kv_update( - key="no_such", - partition_id="test_partition", - fields={"tokens": torch.tensor([1])}, - merge_fn=lambda _old, new: new, - ) + partition_id = "test_partition" + key = "streamed" + + def concat(old, new): + return torch.cat([old, new]) + + def tokens_ready(): + partition = get_controller_partition(controller, partition_id) + row, col = partition.keys_mapping[key], partition.field_name_mapping["tokens"] + return bool(partition.production_status[row, col]) + + tq_api.kv_put(key=key, partition_id=partition_id, fields={"prompt": torch.tensor([7])}) + for chunk in ([1], [2]): + fields = {"tokens": torch.tensor(chunk)} + meta = tq_api.kv_update(key=key, partition_id=partition_id, fields=fields, merge_fn=concat, finalize=False) + assert not tokens_ready() + assert set(tq_api.kv_batch_get(keys=key, partition_id=partition_id).keys()) == {"prompt"} + with pytest.raises(ValueError, match="not finalized"): + tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") + 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"]} + with pytest.raises(RuntimeError, match="not found"): + tq_api.kv_describe(keys=[key, "absent"], partition_id=partition_id) + + tq_api.kv_update(key=key, partition_id=partition_id, fields={"tokens": torch.tensor([3])}, merge_fn=concat) + assert tokens_ready() + assert tq_api.kv_describe(keys=key, partition_id=partition_id)[key]["unfinalized"] == [] + data = tq_api.kv_batch_get(keys=key, partition_id=partition_id, select_fields="tokens") + assert_tensor_equal(data["tokens"][0], torch.tensor([1, 2, 3])) + tq_api.kv_clear(keys=key, partition_id=partition_id) + + def test_creates_missing_key_and_field_and_rejects_incompatible_result(self, tq_api, backend_name): + if backend_name != "SimpleStorage": + pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") + + def fail(_old, _new): + raise AssertionError("merge_fn must not run for a missing value") + + meta = tq_api.kv_update( + key="new_key", partition_id="test_partition", fields={"tokens": torch.tensor([1, 2])}, merge_fn=fail + ) + assert meta.fields == ["tokens"] + tq_api.kv_update( + key="new_key", partition_id="test_partition", fields={"fresh": torch.tensor([9])}, merge_fn=fail + ) + created = tq_api.kv_batch_get(keys="new_key", partition_id="test_partition") + assert_tensor_equal(created["tokens"][0], torch.tensor([1, 2])) + assert_tensor_equal(created["fresh"][0], torch.tensor([9])) + tq_api.kv_put(key="sample", partition_id="test_partition", fields={"tokens": torch.tensor([1, 2])}) - with pytest.raises(ValueError, match="not produced"): - tq_api.kv_update( - key="sample", - partition_id="test_partition", - fields={"fresh": torch.tensor([1])}, - merge_fn=lambda _old, new: new, - ) with pytest.raises(RuntimeError, match="keep dtype"): tq_api.kv_update( key="sample", @@ -1273,6 +1310,62 @@ def test_rejects_missing_key_unproduced_field_and_incompatible_result(self, tq_a unchanged = tq_api.kv_batch_get(keys="sample", partition_id="test_partition", select_fields="tokens") assert_tensor_equal(unchanged["tokens"][0], torch.tensor([1, 2])) + def test_new_row_must_match_column_type_held_by_other_keys(self, tq_api, backend_name): + if backend_name != "SimpleStorage": + pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") + + def fail(_old, _new): + raise AssertionError("merge_fn must not run for a missing value") + + partition_id = "test_partition" + # The column is known only through a hidden write on another key. + tq_api.kv_update( + key="k1", partition_id=partition_id, fields={"reward": torch.tensor([0.5])}, merge_fn=fail, finalize=False + ) + with pytest.raises(RuntimeError, match="keep dtype"): + tq_api.kv_update(key="k2", partition_id=partition_id, fields={"reward": torch.tensor([1])}, merge_fn=fail) + + # Nothing was stored for k2, so a compatible value is still a first write. + tq_api.kv_update(key="k2", partition_id=partition_id, fields={"reward": torch.tensor([0.25])}, merge_fn=fail) + data = tq_api.kv_batch_get(keys="k2", partition_id=partition_id, select_fields="reward") + assert_tensor_equal(data["reward"][0], torch.tensor([0.25])) + tq_api.kv_clear(keys=["k1", "k2"], partition_id=partition_id) + + def test_storage_values_the_controller_does_not_hold_are_overwritten(self, tq_api, backend_name, monkeypatch): + if backend_name != "SimpleStorage": + pytest.skip("merge-backed kv_update is implemented only for SimpleStorage") + + from transfer_queue import interface + + def concat(old, new): + return torch.cat([old, new]) + + async def fail_clear(*_args, **_kwargs): + raise RuntimeError("simulated storage unit timeout") + + async def reject_notify(*_args, **_kwargs): + return False + + partition_id = "test_partition" + manager = interface._maybe_create_tq_client().storage_manager + fields = {"response": torch.tensor([100, 101]), "prompt": torch.tensor([7])} + tq_api.kv_put(key="k", partition_id=partition_id, fields=fields) + with monkeypatch.context() as patched: + patched.setattr(manager, "_clear_single_storage_unit", fail_clear) + tq_api.kv_clear(keys="k", partition_id=partition_id, fields="response") + tq_api.kv_update(key="k", partition_id=partition_id, fields={"response": torch.tensor([1])}, merge_fn=concat) + data = tq_api.kv_batch_get(keys="k", partition_id=partition_id, select_fields="response") + assert_tensor_equal(data["response"][0], torch.tensor([1])) + + # A first chunk the controller never recorded is rewritten on retry, not merged twice. + with monkeypatch.context() as patched, pytest.raises(RuntimeError, match="did not accept"): + patched.setattr(manager, "notify_data_update", reject_notify) + tq_api.kv_update(key="r", partition_id=partition_id, fields={"tokens": torch.tensor([5])}, merge_fn=concat) + tq_api.kv_update(key="r", partition_id=partition_id, fields={"tokens": torch.tensor([5])}, merge_fn=concat) + data = tq_api.kv_batch_get(keys="r", partition_id=partition_id, select_fields="tokens") + assert_tensor_equal(data["tokens"][0], torch.tensor([5])) + tq_api.kv_clear(keys=["k", "r"], partition_id=partition_id) + def test_kv_backend_rejects(self, tq_api, backend_name): if backend_name == "SimpleStorage": pytest.skip("rejection is the KV-backend contract") diff --git a/tests/test_controller.py b/tests/test_controller.py index 32c1603a..edbfdff2 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -1191,6 +1191,29 @@ def test_controller_kv_retrieve_meta_with_production_status(self, ray_setup): # Clean up ray.get(tq_controller.clear_partition.remote(partition_id)) + def test_controller_kv_prepare_update_reports_held_rows_from_field_metadata(self, ray_setup): + """Hidden cells count as held; new keys and fields the partition lacks do not.""" + tq_controller = TransferQueueController.remote() + partition_id = "kv_prepare_update_test" + existing = ray.get( + tq_controller.kv_retrieve_meta.remote(keys=["s0", "s1"], partition_id=partition_id, create=True) + ) + s0, s1 = existing.global_indexes + response = {"response": {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False}} + prompt = {"prompt": {"dtype": torch.int64, "shape": (1,), "is_nested": False, "is_non_tensor": False}} + ray.get(tq_controller.update_production_status.remote(partition_id, [s0], response, ready=False)) + ray.get(tq_controller.update_production_status.remote(partition_id, [s1], prompt)) + + metadata, merge_plan = ray.get( + tq_controller.kv_prepare_update.remote(["s0", "s1", "s2"], partition_id, ["response", "fresh"]) + ) + + assert metadata.global_indexes[:2] == [s0, s1] + assert merge_plan["response"] == {"dtype": torch.int64, "is_non_tensor": False, "held": [True, False, False]} + assert merge_plan["fresh"] == {"dtype": None, "is_non_tensor": None, "held": [False, False, False]} + + ray.get(tq_controller.clear_partition.remote(partition_id)) + def test_controller_kv_retrieve_meta_with_custom_meta(self, ray_setup): """Test kv_retrieve_meta preserves custom_meta through retrieve.""" tq_controller = TransferQueueController.remote() diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index ef8538cc..3d92e175 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -1350,6 +1350,28 @@ def test_clear_fields_drops_fieldmeta_when_no_sample_holds_it(self): assert "pv" not in partition.field_metadata assert partition.global_indexes == {0, 1} + def test_unfinalized_write_records_metadata_but_stays_hidden(self): + partition = self._make_partition() + t4 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False} + partition.update_production_status([0, 1], ["pv"], {"pv": t4}) + + assert partition.update_production_status([1], ["pv"], {"pv": t4}, ready=False) + assert partition.field_metadata["pv"].global_indexes == {0, 1} + assert partition.scan_data_status(["pv"], task_name="trainer") == [0] + + partition.update_production_status([1], ["pv"], {"pv": t4}) + assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1] + + def test_dtype_conflict_leaves_readiness_unchanged(self): + partition = self._make_partition() + f32 = {"dtype": "torch.float32", "shape": (4,), "is_nested": False, "is_non_tensor": False} + i64 = {**f32, "dtype": "torch.int64"} + partition.update_production_status([0, 1], ["pv"], {"pv": f32}) + + assert not partition.update_production_status([1], ["pv"], {"pv": i64}, ready=False) + assert partition.scan_data_status(["pv"], task_name="trainer") == [0, 1] + assert partition.field_metadata["pv"].dtype == "torch.float32" + def test_update_production_status_updates_field_metadata(self): """Test that update_production_status correctly updates field_metadata via _update_field_metadata.""" partition = self._make_partition() diff --git a/tests/test_data_dump.py b/tests/test_data_dump.py index b6eecb55..b2de9cfe 100644 --- a/tests/test_data_dump.py +++ b/tests/test_data_dump.py @@ -149,7 +149,14 @@ def read(self, size=-1): for source, target in [(1, 101), (6, 106)]: offset, length = reply.body["row_offsets"][source] records.append( - {"source_index": source, "target_index": target, "fields": ["x"], "offset": offset, "length": length} + { + "source_index": source, + "target_index": target, + "fields": ["x"], + "unfinalized": [], + "offset": offset, + "length": length, + } ) loaded = unit._handle_load_rows( ZMQMessage.create( @@ -182,7 +189,14 @@ def read(self, size=-1): def test_unit_rejects_invalid_records(unit, tmp_path, problem): path = tmp_path / "row.pkl" path.write_bytes(pickle.dumps({"global_index": 1, "fields": {"x": "value"}})) - record = {"source_index": 1, "target_index": 2, "offset": 0, "length": path.stat().st_size, "fields": ["x"]} + record = { + "source_index": 1, + "target_index": 2, + "offset": 0, + "length": path.stat().st_size, + "fields": ["x"], + "unfinalized": [], + } if problem == "wrong_index": record["source_index"] = 3 elif problem == "truncated": diff --git a/tests/test_kv_update.py b/tests/test_kv_update.py index 567722ff..85328f1d 100644 --- a/tests/test_kv_update.py +++ b/tests/test_kv_update.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import errno from unittest.mock import MagicMock, patch @@ -21,22 +22,22 @@ import torch from tensordict import TensorDict -from transfer_queue.interface import _single_update_batch +from transfer_queue.interface import _single_update_batch, async_kv_batch_update from transfer_queue.metadata import BatchMeta from transfer_queue.storage.managers.base import KVStorageManager from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager from transfer_queue.storage.managers.ray_storage_manager import RayStorageManager from transfer_queue.storage.managers.simple_storage_manager import _build_update_field_schema from transfer_queue.storage.managers.yuanrong_manager import YuanrongStorageManager -from transfer_queue.storage.simple_storage import HybridStorageUnitData, StorageUnitData +from transfer_queue.storage.simple_storage import HybridStorageUnitData, StorageKeyNotFoundError, StorageUnitData def _concat(old, new): return torch.cat([old, new]) -def _tensor_schema(*fields, dtype=torch.int64): - return {field: {"dtype": dtype, "shape": None, "is_nested": True, "is_non_tensor": False} for field in fields} +def _column(*held, dtype=torch.int64, is_non_tensor=False): + return {"dtype": dtype, "is_non_tensor": is_non_tensor, "held": list(held)} def test_single_update_batch_wraps_tensor_and_python_values(): @@ -53,15 +54,20 @@ def test_single_update_batch_rejects_invalid_fields(fields): _single_update_batch(fields) +@patch("transfer_queue.interface._maybe_create_tq_client", side_effect=AssertionError("must fail before any RPC")) +def test_batch_update_rejects_duplicate_keys_before_registering(_client): + fields = TensorDict({"tokens": torch.tensor([[1], [2], [3]])}, batch_size=3) + + with pytest.raises(ValueError, match=r"duplicated: \['k'\]"): + asyncio.run(async_kv_batch_update(["k", "other", "k"], "p", fields, _concat)) + + def test_apply_update_concatenates_and_keeps_field_name(): data = StorageUnitData() data.put_data({"sequence_ids": [torch.tensor([10, 11, 12])]}, [0]) described = data.apply_update( - [0], - {"sequence_ids": [torch.tensor([20, 21])]}, - _concat, - _tensor_schema("sequence_ids"), + [0], {"sequence_ids": [torch.tensor([20, 21])]}, _concat, {"sequence_ids": _column(True)} ) assert described["sequence_ids"] == {"dtype": torch.int64, "shapes": [(5,)]} @@ -78,17 +84,31 @@ def fail(_old, _new): raise RuntimeError("merge failed") with pytest.raises(RuntimeError, match="merge failed"): - data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, _tensor_schema("tokens")) + data.apply_update([7], {"tokens": [torch.tensor([9])]}, fail, {"tokens": _column(True)}) assert torch.equal(data.field_data["tokens"][7], original) -def test_apply_update_rejects_unproduced_field(): +def test_apply_update_stores_missing_values_without_merging(): data = StorageUnitData() data.put_data({"tokens": [torch.tensor([1])]}, [3]) + merged = [] + + def concat(old, new): + merged.append(old) + return torch.cat([old, new]) + + described = data.apply_update( + [3, 4], + {"tokens": [torch.tensor([2]), torch.tensor([5])], "fresh": [{"a": 1}, {"a": 2}]}, + concat, + {"tokens": _column(True, False), "fresh": _column(False, False, dtype=None, is_non_tensor=None)}, + ) - with pytest.raises(ValueError, match="unproduced field"): - data.apply_update([3], {"fresh": [torch.tensor([2])]}, lambda _old, new: new, _tensor_schema("fresh")) - assert "fresh" not in data.field_data + assert len(merged) == 1 + assert torch.equal(data.field_data["tokens"][3], torch.tensor([1, 2])) + assert torch.equal(data.field_data["tokens"][4], torch.tensor([5])) + assert data.field_data["fresh"] == {3: {"a": 1}, 4: {"a": 2}} + assert described["tokens"]["shapes"] == [(2,), (1,)] @pytest.mark.parametrize( @@ -104,20 +124,66 @@ def test_apply_update_rejects_incompatible_tensor_results_before_writing(merge_f data.put_data({"tokens": [original.clone()]}, [0]) with pytest.raises(TypeError, match=error): - data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, _tensor_schema("tokens")) + data.apply_update([0], {"tokens": [torch.tensor([3])]}, merge_fn, {"tokens": _column(True)}) assert torch.equal(data.field_data["tokens"][0], original) def test_apply_update_rejects_tensor_result_for_non_tensor_field(): data = StorageUnitData() data.put_data({"meta": [{"step": 1}]}, [0]) - schema = {"meta": {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": True}} with pytest.raises(TypeError, match="must remain non-tensor"): - data.apply_update([0], {"meta": [{"step": 2}]}, lambda _old, _new: torch.tensor([2]), schema) + data.apply_update( + [0], + {"meta": [{"step": 2}]}, + lambda _old, _new: torch.tensor([2]), + {"meta": _column(True, dtype=None, is_non_tensor=True)}, + ) assert data.field_data["meta"][0] == {"step": 1} +@pytest.mark.parametrize( + "column, value, error", + [ + (_column(True, False, dtype=torch.float32), torch.tensor([1]), "keep dtype"), + (_column(True, False, dtype=torch.float32), {"not": "a tensor"}, "keep dtype"), + (_column(False, False, dtype=None, is_non_tensor=True), torch.tensor([1.0]), "must remain non-tensor"), + ], +) +def test_apply_update_checks_new_rows_against_column_type_before_writing(column, value, error): + data = StorageUnitData() + data.put_data({"score": [torch.tensor([0.5])]}, [0]) + + with pytest.raises(TypeError, match=error): + data.apply_update([0, 1], {"score": [torch.tensor([0.25]), value]}, _concat, {"score": column}) + assert set(data.field_data["score"]) == {0} + assert torch.equal(data.field_data["score"][0], torch.tensor([0.5])) + + +def test_apply_update_overwrites_a_local_value_the_controller_does_not_hold(): + data = StorageUnitData() + data.put_data({"response": [torch.tensor([100, 101])]}, [0]) + + def must_not_merge(_old, _new): + raise AssertionError("a value the controller does not hold must not be merged") + + data.apply_update([0], {"response": [torch.tensor([1])]}, must_not_merge, {"response": _column(False)}) + + assert torch.equal(data.field_data["response"][0], torch.tensor([1])) + + +def test_apply_update_rejects_held_row_missing_locally_before_writing(): + data = StorageUnitData() + data.put_data({"tokens": [torch.tensor([1])]}, [0]) + + with pytest.raises(StorageKeyNotFoundError): + data.apply_update( + [0, 1], {"tokens": [torch.tensor([2]), torch.tensor([3])]}, _concat, {"tokens": _column(True, True)} + ) + assert set(data.field_data["tokens"]) == {0} + assert torch.equal(data.field_data["tokens"][0], torch.tensor([1])) + + def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path): data = HybridStorageUnitData( storage_size=4, threshold_bytes=64, ssd_path=str(tmp_path), run_id="run", unit_id="unit" @@ -126,7 +192,7 @@ def test_apply_update_decodes_ssd_offloaded_old_value(tmp_path): data.put_data({"tokens": [prompt]}, [0]) assert data.ssd_active_values == 1 - data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, _tensor_schema("tokens")) + data.apply_update([0], {"tokens": [torch.tensor([99])]}, _concat, {"tokens": _column(True)}) assert torch.equal(data.get_data(["tokens"], [0])["tokens"][0], torch.cat([prompt, torch.tensor([99])])) assert data.ssd_active_values == 1 @@ -152,7 +218,7 @@ def fail_second_field(encoded_values): monkeypatch.setattr(data._ssd_store, "write_values", fail_second_field) new_data = {"tokens": [torch.tensor([99])], "mask": [torch.tensor([1])]} with pytest.raises(OSError, match="No space left"): - data.apply_update([0], new_data, _concat, _tensor_schema("tokens", "mask")) + data.apply_update([0], new_data, _concat, {"tokens": _column(True), "mask": _column(True)}) stored = data.get_data(["tokens", "mask"], [0]) assert torch.equal(stored["tokens"][0], prompt) @@ -216,4 +282,4 @@ async def test_kv_backends_reject_update(mock_create, manager_cls): values = TensorDict({"x": torch.ones(1, 1, dtype=torch.int64)}, batch_size=1) with pytest.raises(NotImplementedError, match=f"not supported by {manager_cls.__name__}"): - await manager.update_data(meta, values, lambda _old, new: new) + await manager.update_data(meta, values, lambda _old, new: new, {"x": _column(True)}) diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py index 70be8985..ff514992 100644 --- a/tests/test_storage_request_retry.py +++ b/tests/test_storage_request_retry.py @@ -107,6 +107,9 @@ async def test_put_retries_only_without_a_parser(data_parser, expected_attempts) manager.notify_data_update.assert_not_awaited() +_HELD_INPUT_IDS = {"input_ids": {"dtype": torch.int64, "is_non_tensor": False, "held": [True]}} + + @pytest.mark.asyncio async def test_update_reports_controller_rejection_after_storage_commit(): manager = _manager() @@ -117,7 +120,7 @@ async def test_update_reports_controller_rejection_after_storage_commit(): values, metadata = _single_sample_batch() with pytest.raises(RuntimeError, match="Storage update committed.*do not retry"): - await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new])) + await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]), _HELD_INPUT_IDS) manager._update_to_single_storage_unit.assert_awaited_once() manager.notify_data_update.assert_awaited_once() @@ -131,7 +134,7 @@ async def test_update_timeout_is_not_retried_or_published(): values, metadata = _single_sample_batch() with patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), pytest.raises(StorageUnitTimeout): - await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new])) + await manager.update_data(metadata, values, lambda old, new: torch.cat([old, new]), _HELD_INPUT_IDS) manager._update_to_single_storage_unit.assert_awaited_once() manager.notify_data_update.assert_not_awaited() diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 6c514b2b..05e63cdf 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -23,6 +23,7 @@ async_kv_batch_put, async_kv_batch_update, async_kv_clear, + async_kv_describe, async_kv_list, async_kv_put, async_kv_update, @@ -36,6 +37,7 @@ kv_batch_put, kv_batch_update, kv_clear, + kv_describe, kv_list, kv_put, kv_update, @@ -72,6 +74,7 @@ "kv_batch_get", "kv_batch_get_by_meta", "kv_list", + "kv_describe", "kv_clear", "kv_update", "async_kv_put", @@ -80,6 +83,7 @@ "async_kv_batch_get", "async_kv_batch_get_by_meta", "async_kv_list", + "async_kv_describe", "async_kv_clear", "async_kv_update", "KVBatchMeta", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 41ff29d4..7a088b99 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -497,13 +497,17 @@ async def async_update( metadata: BatchMeta, values: TensorDict, merge_fn: Callable[[Any, Any], Any], + merge_plan: dict[str, dict[str, Any]], + finalize: bool = True, ) -> BatchMeta: - """Merge new values into produced fields on SimpleStorage units. + """Merge new values into fields on SimpleStorage units, storing them where absent. Args: - metadata: Samples to update. The key must already exist. + metadata: Samples to update, already registered with the controller. values: New values aligned with ``metadata``. - merge_fn: Called per sample per field as ``merge_fn(old, new)``. + merge_fn: Called per sample per field as ``merge_fn(old, new)`` on held rows. + merge_plan: Per-field column type and held rows, from ``async_kv_prepare_update``. + finalize: If False, keep the updated fields hidden from consumers. Returns: The same metadata with field_schema and production status updated. @@ -519,8 +523,10 @@ async def async_update( if not metadata or metadata.size == 0: raise ValueError("metadata cannot be none or empty") - field_schema = await self.storage_manager.update_data(metadata, values, merge_fn) - return metadata.apply_field_schema(field_schema) + field_schema = await self.storage_manager.update_data(metadata, values, merge_fn, merge_plan, finalize=finalize) + if not finalize: + metadata.production_status[:] = 0 + return metadata.apply_field_schema(field_schema, set_all_ready=finalize) async def async_get_data(self, metadata: BatchMeta) -> TensorDict: """Asynchronously fetch data from storage units and organize into TensorDict. @@ -1010,6 +1016,23 @@ async def async_kv_retrieve_meta( except Exception as e: raise RuntimeError(f"[{self.client_id}] Failed in async_kv_retrieve_meta: {e}") from e + @with_controller_socket + async def async_kv_prepare_update( + self, keys: list[str], partition_id: str, fields: list[str], socket=None + ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]: + """Register missing keys and get, per field, its column type and which rows hold it. + + Returns: + The keys' BatchMeta and the ``merge_plan`` that ``async_update`` takes. + """ + response_msg = await self._request_controller( + socket=socket, + request_type=ZMQRequestType.KV_PREPARE_UPDATE, + response_type=ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE, + body={"keys": keys, "partition_id": partition_id, "fields": fields}, + ) + return response_msg.body["metadata"], response_msg.body["merge_plan"] + async def async_kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) -> None: """Clear fields of existing keys while keeping the keys, tags and consumption status. @@ -1133,8 +1156,8 @@ async def async_kv_describe( key raises instead of being dropped. Returns: - ``{"rows": {key: {"global_index", "fields", "tag"}}, "field_schema": {field: - {"is_nested", "is_non_tensor"}}}``. + ``{"rows": {key: {"global_index", "fields", "unfinalized", "tag"}}, "field_schema": + {field: {"is_nested", "is_non_tensor"}}}``. """ response = await self._request_controller( socket=socket, @@ -1455,6 +1478,7 @@ def wrapper(*args, **kwargs): self._set_custom_meta = _make_sync(self.async_set_custom_meta) self._reset_consumption = _make_sync(self.async_reset_consumption) self._kv_retrieve_meta = _make_sync(self.async_kv_retrieve_meta) + self._kv_prepare_update = _make_sync(self.async_kv_prepare_update) 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) @@ -1638,12 +1662,14 @@ def update( metadata: BatchMeta, values: TensorDict, merge_fn: Callable[[Any, Any], Any], + merge_plan: dict[str, dict[str, Any]], + finalize: bool = True, ) -> BatchMeta: - """Synchronously merge new values into produced fields on SimpleStorage units. + """Synchronously merge or add values on SimpleStorage units. See ``async_update``. """ - return self._update(metadata, values, merge_fn) + return self._update(metadata, values, merge_fn, merge_plan, finalize=finalize) def get_data(self, metadata: BatchMeta) -> TensorDict: """Synchronously fetch data from storage units and organize into TensorDict. @@ -1859,6 +1885,15 @@ def kv_retrieve_meta( """ return self._kv_retrieve_meta(keys=keys, partition_id=partition_id, create=create) + def kv_prepare_update( + self, keys: list[str], partition_id: str, fields: list[str] + ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]: + """Synchronously register missing keys and get the ``merge_plan`` for ``update``. + + See ``async_kv_prepare_update``. + """ + return self._kv_prepare_update(keys, partition_id, fields) + def kv_retrieve_keys( self, global_indexes: list[int] | int, diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index c2ba668f..d19152cd 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -500,6 +500,7 @@ def update_production_status( field_names: list[str], field_schema: dict[str, dict[str, Any]], custom_backend_meta: dict[int, dict[str, Any]] | None = None, + ready: bool = True, ) -> bool: """ Update production status for specific samples and fields. @@ -515,11 +516,20 @@ def update_production_status( field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}} custom_backend_meta: Optional per-sample per-field custom metadata provided by storage backend + ready: Whether consumers may see these cells. False records the stored + value (field metadata) but hides it until a finalizing kv_update. Returns: True if update was successful, False on error """ try: + # Same dtype check as FieldMeta.update, run first so a conflicting write changes nothing. + for name, incoming in field_schema.items(): + existing = self.field_metadata.get(name) + dtype = incoming.get("dtype") + if existing is not None and None not in (existing.dtype, dtype) and existing.dtype != dtype: + raise ValueError(f"Field {name!r} dtype mismatch: {existing.dtype} != {dtype}") + # Derive field_names from field_schema to guarantee consistency field_names = list(field_schema.keys()) @@ -543,7 +553,7 @@ def update_production_status( # Update production status if self.production_status is not None and global_indices and field_names: field_indices = [self.field_name_mapping.get(f) for f in field_names] - self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = 1 + self.production_status[torch.tensor(global_indices)[:, None], torch.tensor(field_indices)] = int(ready) # Update field metadata self._update_field_metadata(global_indices, field_schema, custom_backend_meta) @@ -1130,6 +1140,7 @@ def update_production_status( global_indexes: list[int], field_schema: dict[str, dict[str, Any]], custom_backend_meta: dict[int, dict[str, Any]] | None = None, + ready: bool = True, ) -> bool: """ Update production status for specific samples and fields in a partition. @@ -1140,6 +1151,7 @@ def update_production_status( global_indexes: List of sample indices to update field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}} custom_backend_meta: Optional custom backend metadata + ready: False keeps the written cells hidden from consumers (unfinalized kv_update) Returns: True if update was successful, False otherwise @@ -1150,11 +1162,13 @@ def update_production_status( logger.error(f"Partition {partition_id} not found") return False - success = partition.update_production_status(global_indexes, field_names, field_schema, custom_backend_meta) + success = partition.update_production_status( + global_indexes, field_names, field_schema, custom_backend_meta, ready=ready + ) if success: logger.debug( f"[{self.controller_id}]: Updated production status for partition {partition_id}: " - f"samples={global_indexes}, fields={field_names}" + f"samples={global_indexes}, fields={field_names}, ready={ready}" ) return success @@ -1717,6 +1731,32 @@ def kv_retrieve_meta( return self.generate_batch_meta(partition_id, verified_global_indexes, data_fields, mode="force_fetch") + def kv_prepare_update( + self, keys: list[str], partition_id: str, fields: list[str] + ) -> tuple[BatchMeta, dict[str, dict[str, Any]]]: + """Register missing keys and tell kv_update, per field, which rows to merge into. + + Storage units may keep values the controller no longer tracks (e.g. after a failed + clear), so ``held`` comes from ``FieldMeta.global_indexes``, which covers hidden cells. + + 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. + """ + metadata = self.kv_retrieve_meta(keys, partition_id, create=True) + partition = self._get_partition(partition_id) + assert partition is not None + merge_plan = {} + for name in fields: + meta = partition.field_metadata.get(name) + held = meta.global_indexes if meta is not None else set() + merge_plan[name] = { + "dtype": meta.dtype if meta is not None else None, + "is_non_tensor": meta.is_non_tensor if meta is not None else None, + "held": [idx in held for idx in metadata.global_indexes], + } + return metadata, merge_plan + def kv_clear_fields(self, keys: list[str], partition_id: str, fields: list[str]) -> BatchMeta: """Clear fields of existing keys and return the cleared cells' metadata for storage deletion. @@ -1783,9 +1823,10 @@ 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", "tag"}}, "field_schema": {field: - {"is_nested", "is_non_tensor"}}}``. ``fields`` is empty for a row with no - produced field yet; ``field_schema`` covers the fields of the described rows. + ``{"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. Raises: KeyError: The partition or any key does not exist. @@ -1799,18 +1840,23 @@ 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}") - rows: dict[str, dict[str, Any]] = { - key: { + rows: dict[str, dict[str, Any]] = {} + for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True): + fields = sorted( + field_name + for field_name, field_meta in partition.field_metadata.items() + if global_index in field_meta.global_indexes + ) + rows[key] = { "global_index": global_index, - "fields": sorted( - field_name - for field_name, field_meta in partition.field_metadata.items() - if global_index in field_meta.global_indexes - ), + "fields": fields, + "unfinalized": [ + name + for name in fields + if not partition.production_status[global_index, partition.field_name_mapping[name]] + ], "tag": partition.custom_meta.get(global_index, {}), } - for key, global_index in zip(keys, cast(list[int], global_indexes), strict=True) - } # 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)} @@ -2043,6 +2089,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_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.SAVE_CONTROLLER_CHECKPOINT: self._handle_save_controller_checkpoint_request, @@ -2097,6 +2144,7 @@ def _handle_notify_data_update_request(self, request_msg: ZMQMessage) -> ZMQMess global_indexes=global_indexes, field_schema=params.get("field_schema", {}), custom_backend_meta=params.get("custom_backend_meta", {}), + ready=params.get("ready", True), ) if success: if self._metrics is not None: @@ -2247,6 +2295,13 @@ 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_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"]) + return self._make_response( + request_msg, ZMQRequestType.KV_PREPARE_UPDATE_RESPONSE, {"metadata": metadata, "merge_plan": merge_plan} + ) + 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 80fa53ca..ad3791fb 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, tag} + row_index.pt # key -> {global_index, fields, unfinalized, tag} shards/ shard_info.json # file, unit, row count, source index -> [offset, length] shard__.pkl # independent {global_index, fields} records @@ -42,7 +42,10 @@ logger = get_logger(__name__) -DUMP_FORMAT_VERSION = 3 +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) SHARD_SUBDIR = "shards" _DUMP_INFO_FILE = "dump_info.json" @@ -171,10 +174,10 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: raise FileNotFoundError(f"{_DUMP_INFO_FILE} not found in {dump_dir}") with open(info_path, encoding="utf-8") as f: dump_info = json.load(f) - if dump_info["format_version"] != DUMP_FORMAT_VERSION: + if dump_info["format_version"] not in _READABLE_FORMAT_VERSIONS: raise ValueError( f"Unsupported dump format version {dump_info['format_version']} in {dump_dir}; " - f"this build reads version {DUMP_FORMAT_VERSION}" + f"this build reads versions {list(_READABLE_FORMAT_VERSIONS)}" ) row_index = _read_row_index(dump_dir) @@ -215,6 +218,7 @@ def read_dump(dump_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: "key": key, "source_index": source_index, "fields": rows[key]["fields"], + "unfinalized": rows[key].get("unfinalized", []), "offset": offset, "length": length, } diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index b3e0ede4..1ba2f36b 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -539,6 +539,7 @@ def kv_batch_get(keys: list[str] | str, partition_id: str, select_fields: list[s Raises: ValueError: If keys or partition are not found ValueError: If empty fields exist in any key (sample) + ValueError: If a field in ``select_fields`` is not ready in every found key Example: >>> import transfer_queue as tq @@ -596,6 +597,21 @@ def kv_list(partition_id: str | None = None) -> dict[str, dict[str, Any]]: return tq_client._run_coroutine(async_kv_list(partition_id=partition_id)) +def kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, Any]]: + """Describe which fields each key holds, including unfinalized ones. + + See ``async_kv_describe``. + + Example: + >>> import transfer_queue as tq + >>> tq.init() + >>> rows = tq.kv_describe(keys=["sample_1", "sample_2"], partition_id="train") + >>> hidden = {key: row["unfinalized"] for key, row in rows.items() if row["unfinalized"]} + """ + tq_client = _maybe_create_tq_client() + return tq_client._run_coroutine(async_kv_describe(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. @@ -629,22 +645,33 @@ def kv_update( partition_id: str, fields: dict[str, Any], merge_fn: Callable[[Any, Any], Any], + finalize: bool = True, ) -> KVBatchMeta: - """Merge new values into produced fields of one existing key. + """Merge new values into the fields of one key, adding the key or fields if missing. + + A missing key is created and a field the controller does not record for the key + yet is stored as given, without calling ``merge_fn``. Unlike ``kv_put``, it takes + no ``tags`` or ``data_parser``. Args: - key: Existing user-specified key. The key is not created if it is missing. + key: User-specified key. Created if missing. partition_id: Partition that holds the key. - fields: New values keyed by an already-produced field name. - merge_fn: ``merge_fn(old, new) -> stored``. SimpleStorage runs it once - per field on the unit that owns the sample. + fields: New values keyed by field name. + merge_fn: ``merge_fn(old, new) -> stored`` for fields the key already holds, + hidden ones included. SimpleStorage runs it once per field on the unit + that owns the sample. + finalize: If False, the updated fields stay hidden from consumers until an update + with True (the default), so a field built over several updates never shows an + intermediate value. Hiding a field that is already ready is not guaranteed: + readers that already hold its metadata can still see the new value. Returns: Metadata for the updated sample. Raises: - ValueError: If the key or a field is missing. 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. NotImplementedError: If the storage backend is not SimpleStorage. Note: @@ -658,6 +685,7 @@ def kv_update( partition_id=partition_id, fields=fields, merge_fn=merge_fn, + finalize=finalize, ) ) @@ -667,15 +695,17 @@ def kv_batch_update( partition_id: str, fields: TensorDict, merge_fn: Callable[[Any, Any], Any], + finalize: bool = True, ) -> KVBatchMeta: - """Merge batched values into produced fields of existing keys. + """Merge batched values into the fields of keys, adding keys or fields if missing. - A timeout has an unknown outcome because a unit may commit after the caller - stops waiting. Do not blindly retry a merge. + See ``kv_update``. ``keys`` must be unique; a repeated key raises ``ValueError`` + before anything is registered or written. A timeout has an unknown outcome because + a unit may commit after the caller stops waiting. Do not blindly retry a merge. """ tq_client = _maybe_create_tq_client() return tq_client._run_coroutine( - async_kv_batch_update(keys=keys, partition_id=partition_id, fields=fields, merge_fn=merge_fn) + async_kv_batch_update(keys=keys, partition_id=partition_id, fields=fields, merge_fn=merge_fn, finalize=finalize) ) @@ -955,6 +985,7 @@ async def async_kv_batch_get( Raises: ValueError: If keys or partition are not found ValueError: If empty fields exist in any key (sample) + ValueError: If a field in ``select_fields`` is not ready in every found key Example: >>> import transfer_queue as tq @@ -980,6 +1011,13 @@ async def async_kv_batch_get( fields_to_fetch = [select_fields] else: fields_to_fetch = select_fields + # The meta keeps only fields ready in every key; select_fields would drop the rest silently. + not_ready = [f for f in fields_to_fetch if f not in batch_meta.field_names] + if not_ready: + raise ValueError( + f"Fields {not_ready} are not ready in all the requested keys: missing, or written by " + "kv_update(finalize=False) and not finalized yet" + ) batch_meta = batch_meta.select_fields(fields_to_fetch) if not batch_meta.is_ready: @@ -1031,6 +1069,37 @@ async def async_kv_list(partition_id: str | None = None) -> dict[str, dict[str, return partition_info +async def async_kv_describe(keys: list[str] | str, partition_id: str) -> dict[str, dict[str, Any]]: + """Asynchronously describe which fields each key holds, including unfinalized ones. + + ``kv_batch_get`` returns only fields that are ready in every key, so this is the way + to find fields written by ``kv_update(finalize=False)`` and not yet finalized. + + Args: + keys: Single key or list of keys to describe. + 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. + + Raises: + RuntimeError: If the partition or any key does not exist. + + Example: + >>> import transfer_queue as tq + >>> tq.init() + >>> rows = await tq.async_kv_describe(keys="sample_1", partition_id="train") + >>> print(rows["sample_1"]["unfinalized"]) + """ + keys = [keys] if isinstance(keys, str) else keys + 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() + } + + 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. @@ -1091,13 +1160,15 @@ async def async_kv_update( partition_id: str, fields: dict[str, Any], merge_fn: Callable[[Any, Any], Any], + finalize: bool = True, ) -> KVBatchMeta: - """Asynchronously update fields of an existing key. See ``kv_update``.""" + """Asynchronously merge or add fields of one key. See ``kv_update``.""" return await async_kv_batch_update( keys=[key], partition_id=partition_id, fields=_single_update_batch(fields), merge_fn=merge_fn, + finalize=finalize, ) @@ -1106,10 +1177,16 @@ async def async_kv_batch_update( partition_id: str, fields: TensorDict, merge_fn: Callable[[Any, Any], Any], + finalize: bool = True, ) -> KVBatchMeta: - """Asynchronously merge batched values into produced fields of existing keys.""" + """Asynchronously merge or add batched fields of keys. See ``kv_update``.""" if not isinstance(keys, list) or not keys or not all(isinstance(key, str) for key in keys): raise TypeError("keys must be a non-empty list of strings") + # A repeated missing key would get one row per occurrence, orphaning all but the last; + # a repeated existing key would merge twice from the same old value, losing one. + if len(set(keys)) != len(keys): + duplicates = sorted({key for key in keys if keys.count(key) > 1}) + raise ValueError(f"keys must be unique in one kv_batch_update; duplicated: {duplicates}") if not isinstance(fields, TensorDict) or not fields.keys(): raise TypeError("fields must be a non-empty TensorDict") if fields.batch_size != torch.Size([len(keys)]): @@ -1117,15 +1194,22 @@ async def async_kv_batch_update( if not callable(merge_fn): raise TypeError("merge_fn must be callable") tq_client = _maybe_create_tq_client() - batch_meta = await tq_client.async_kv_retrieve_meta(keys=keys, partition_id=partition_id, create=False) + # Checked up front: missing keys are registered before storage is asked to merge. + if not isinstance(tq_client.storage_manager, AsyncSimpleStorageManager): + raise NotImplementedError( + f"kv_update is not supported by {type(tq_client.storage_manager).__name__}; " + "it requires the SimpleStorage backend" + ) + batch_meta, merge_plan = await tq_client.async_kv_prepare_update( + keys=keys, partition_id=partition_id, fields=list(fields.keys()) + ) if batch_meta.size != len(keys): - raise ValueError("Some keys or the partition were not found") - missing_fields = sorted(set(fields.keys()) - set(batch_meta.field_names)) - if missing_fields: - raise ValueError(f"Fields are not produced for every key: {missing_fields}") + raise RuntimeError(f"Retrieved BatchMeta size {batch_meta.size} does not match {len(keys)} keys") - batch_meta = await tq_client.async_update(metadata=batch_meta, values=fields, merge_fn=merge_fn) + batch_meta = await tq_client.async_update( + metadata=batch_meta, values=fields, merge_fn=merge_fn, merge_plan=merge_plan, finalize=finalize + ) return KVBatchMeta( keys=keys, diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 1015172e..f881a0d5 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -226,15 +226,17 @@ async def notify_data_update( global_indexes: list[int], field_schema: dict[str, dict[str, Any]], custom_backend_meta: dict[int, dict[str, Any]] | None = None, + ready: bool = True, ) -> bool: """ - Notify controller that new data is ready. + Notify controller that new data is stored. Args: partition_id: Current data partition id. global_indexes: Data update related global_indexes. field_schema: Columnar field schema {field_name: {dtype, shape, is_nested, ...}}. custom_backend_meta: Per-field custom_meta for each sample, in {global_index: {field: custom_meta}} format. + ready: Whether consumers may see the data now; False for a kv_update with finalize=False. """ if not self.controller_info: @@ -264,6 +266,7 @@ async def notify_data_update( "global_indexes": global_indexes, "field_schema": normalized_field_schema, "custom_backend_meta": custom_backend_meta, + "ready": ready, }, ).serialize() @@ -362,13 +365,18 @@ async def update_data( metadata: BatchMeta, values: TensorDict, merge_fn: Callable[[Any, Any], Any], + merge_plan: dict[str, dict[str, Any]], + finalize: bool = True, ) -> dict[str, dict[str, Any]]: - """Merge new values on the unit that holds each sample. + """Merge new values on the unit that holds each sample, or store them if absent. Args: metadata: Samples to update. values: New values as a TensorDict aligned with ``metadata``. - merge_fn: Called per sample per field as ``merge_fn(old, new)``. + merge_fn: Called per sample per field as ``merge_fn(old, new)`` on held rows. + merge_plan: ``{field: {"dtype", "is_non_tensor", "held"}}`` from the controller; + ``held`` is aligned with ``metadata`` and decides which rows merge. + finalize: If False, keep the updated fields hidden from consumers. Returns: field_schema of the stored values, keyed by field name. diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index b5bf88b2..1ce03bf5 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -544,11 +544,13 @@ async def update_data( metadata: BatchMeta, values: TensorDict, merge_fn: Callable[[Any, Any], Any], + merge_plan: dict[str, dict[str, Any]], + finalize: bool = True, ) -> dict[str, dict[str, Any]]: - """Merge values on each owner unit and notify the controller. + """Merge or add values on each owner unit and notify the controller. Each unit reports the values it stored; this assembles one batch-level - field_schema so production status stays ready. + field_schema. ``finalize`` decides whether the fields become ready. """ logger.debug(f"[{self.storage_manager_id}]: receive update_data request, updating {metadata.size} samples.") @@ -564,6 +566,10 @@ async def update_data( tasks = [] for su_id, group in routing.items(): storage_data = {f: self._select_by_positions(values[f], group.batch_positions) for f in field_names} + unit_columns = { + f: {**merge_plan[f], "held": [merge_plan[f]["held"][p] for p in group.batch_positions]} + for f in field_names + } tasks.append( self._request_with_retry( "update", @@ -574,7 +580,7 @@ async def update_data( group.global_indexes, storage_data, merge_fn, - {name: metadata.field_schema[name] for name in field_names}, + unit_columns, target_storage_unit=su_id, ), # A merge reads stored state, so a replay after a lost answer @@ -603,6 +609,7 @@ async def update_data( metadata.partition_ids[0], metadata.global_indexes, field_schema, + ready=finalize, ) if not published: raise RuntimeError( @@ -674,7 +681,7 @@ async def _update_to_single_storage_unit( global_indexes: list[int], storage_data: dict[str, Any], merge_fn: Callable[[Any, Any], Any], - field_schema: dict[str, dict[str, Any]], + merge_plan: dict[str, dict[str, Any]], target_storage_unit: str, socket: zmq.Socket = None, ) -> dict[str, dict[str, Any]]: @@ -683,7 +690,7 @@ async def _update_to_single_storage_unit( "global_indexes": global_indexes, "data": storage_data, "merge_fn": merge_fn, - "field_schema": field_schema, + "merge_plan": merge_plan, } request_msg = ZMQMessage.create( @@ -1139,7 +1146,9 @@ async def load_rows_by_index(self, partition_id: str, shards: list[dict[str, Any # dominate a wide cross-topology load. The notify pool bounds the concurrency. await asyncio.gather( *( - self.notify_data_update(partition_id, update["global_indexes"], update["field_schema"]) + self.notify_data_update( + partition_id, update["global_indexes"], update["field_schema"], ready=update["ready"] + ) for result in loaded for update in result["updates"] ) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index ca8ad853..09a245f1 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -112,6 +112,16 @@ class _SSDValueRef: TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0)) +def _check_value_type(field: str, value: Any, is_tensor: bool, dtype: Any) -> None: + """Reject a value that would change its field's tensor/non-tensor kind or tensor dtype.""" + if not is_tensor: + if isinstance(value, torch.Tensor): + raise TypeError(f"Value for non-tensor field {field!r} must remain non-tensor") + elif not isinstance(value, torch.Tensor) or (dtype is not None and value.dtype != dtype): + actual = value.dtype if isinstance(value, torch.Tensor) else type(value).__name__ + raise TypeError(f"Value for tensor field {field!r} must keep dtype {dtype}, got {actual}") + + def _describe_stored_values(values: list) -> dict[str, Any]: """Describe per-sample stored values for the manager. @@ -235,36 +245,31 @@ def apply_update( global_indexes: list[int], new_data: dict[str, Any], merge_fn: Callable[[Any, Any], Any], - field_schema: dict[str, dict[str, Any]], + merge_plan: dict[str, dict[str, Any]], ) -> dict[str, dict[str, Any]]: - """Merge produced values, validate their types, then write once. + """Merge into held values, store the other rows as is, validate, then write once. - All merge calls and type checks finish before one ``put_data`` call, so - a merge error or incompatible result leaves storage unchanged. + ``merge_plan[field]["held"]`` comes from the controller, not from this unit: a local + value the controller no longer tracks (e.g. after a failed clear) is overwritten, + never merged. Every final value must match the field's known column type. All + merges and checks finish before one ``put_data`` call, so a merge error, an + incompatible value or a held row missing here leaves storage unchanged. Returns: Per-sample description of the stored values, keyed by field name. """ computed: dict[str, list] = {} for field, new_values in new_data.items(): - stored = self.field_data.get(field, {}) - missing = [idx for idx in global_indexes if idx not in stored] - if missing: - raise ValueError(f"Cannot update unproduced field {field!r} for indexes {missing[:20]}") - old_values = self.get_data([field], global_indexes)[field] - computed[field] = [merge_fn(old, new_values[i]) for i, old in enumerate(old_values)] - - expected = field_schema[field] - if expected.get("is_non_tensor"): - if any(isinstance(value, torch.Tensor) for value in computed[field]): - raise TypeError(f"Merge result for non-tensor field {field!r} must remain non-tensor") - else: - for value in computed[field]: - if not isinstance(value, torch.Tensor) or value.dtype != expected["dtype"]: - actual = value.dtype if isinstance(value, torch.Tensor) else type(value).__name__ - raise TypeError( - f"Merge result for tensor field {field!r} must keep dtype {expected['dtype']}, got {actual}" - ) + column = merge_plan[field] + present = [i for i, held in enumerate(column["held"]) if held] + old_values = self.get_data([field], [global_indexes[i] for i in present])[field] if present else [] + values = list(new_values) + for i, old in zip(present, old_values, strict=True): + values[i] = merge_fn(old, new_values[i]) + if column["is_non_tensor"] is not None: + for value in values: + _check_value_type(field, value, not column["is_non_tensor"], column["dtype"]) + computed[field] = values self.put_data(computed, global_indexes) return {field: _describe_stored_values(values) for field, values in computed.items()} @@ -1243,12 +1248,12 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: ) def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage: - """Merge existing fields, write once, and describe what was stored.""" + """Merge or add fields, write once, and describe what was stored.""" try: global_indexes = data_parts.body["global_indexes"] merge_fn = data_parts.body.get("merge_fn") new_data = data_parts.body.get("data") - field_schema = data_parts.body.get("field_schema") + merge_plan = data_parts.body.get("merge_plan") with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_update" @@ -1257,9 +1262,9 @@ def _handle_update(self, data_parts: ZMQMessage) -> ZMQMessage: raise TypeError(f"merge_fn must be callable, got {type(merge_fn).__name__}") if not isinstance(new_data, dict) or not new_data: raise TypeError("update data must be a non-empty dict") - if not isinstance(field_schema, dict) or set(field_schema) != set(new_data): - raise ValueError("field_schema must describe every update field") - stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, field_schema) + if not isinstance(merge_plan, dict) or set(merge_plan) != set(new_data): + raise ValueError("merge_plan must describe every update field") + stored_shapes = self.storage_data.apply_update(global_indexes, new_data, merge_fn, merge_plan) return ZMQMessage.create( request_type=ZMQRequestType.UPDATE_DATA_RESPONSE, # type: ignore[arg-type] @@ -1496,7 +1501,7 @@ def _handle_dump_rows(self, request: ZMQMessage) -> ZMQMessage: def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage: """Seek to assigned records and merge them into local storage at current indexes.""" - updates = [] + updates: list[dict[str, Any]] = [] bytes_read = 0 try: with limit_pytorch_auto_parallel_threads(TQ_NUM_THREADS): @@ -1511,19 +1516,23 @@ def _handle_load_rows(self, request: ZMQMessage) -> ZMQMessage: f, row["offset"], row["length"], row["source_index"], row["fields"] ) validate_dump_values(fields, shard["field_schema"], row["source_index"]) - groups[tuple(row["fields"])].append((row, fields)) + groups[(tuple(row["fields"]), tuple(row["unfinalized"]))].append((row, fields)) bytes_read += row["length"] - for signature, rows in groups.items(): + for (signature, unfinalized), rows in groups.items(): indexes = [row["target_index"] for row, _ in rows] values = {name: [fields[name] for _, fields in rows] for name in signature} - schema = select_dump_schema( - shard["field_schema"], - [row["source_index"] for row, _ in rows], - indexes, - signature, - ) self.storage_data.put_data(values, indexes) - updates.append({"global_indexes": indexes, "field_schema": schema}) + source_indexes = [row["source_index"] for row, _ in rows] + finalized = tuple(name for name in signature if name not in unfinalized) + # Publish each field with the readiness it had when dumped. + for names, ready in ((finalized, True), (unfinalized, False)): + if names: + schema = select_dump_schema( + shard["field_schema"], source_indexes, indexes, names + ) + updates.append( + {"global_indexes": indexes, "field_schema": schema, "ready": ready} + ) logger.info( "[%s]: loaded %s rows (%s bytes)", self.storage_unit_id, diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index af30be40..9c765cfb 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -125,6 +125,8 @@ class ZMQRequestType(ExplicitEnum): KV_LIST_RESPONSE = "KV_LIST_RESPONSE" KV_DESCRIBE = "KV_DESCRIBE" KV_DESCRIBE_RESPONSE = "KV_DESCRIBE_RESPONSE" + KV_PREPARE_UPDATE = "KV_PREPARE_UPDATE" + KV_PREPARE_UPDATE_RESPONSE = "KV_PREPARE_UPDATE_RESPONSE" KV_CLEAR_FIELDS = "KV_CLEAR_FIELDS" KV_CLEAR_FIELDS_RESPONSE = "KV_CLEAR_FIELDS_RESPONSE" diff --git a/tutorial/02_kv_interface.py b/tutorial/02_kv_interface.py index 999a8193..ef9e62b4 100644 --- a/tutorial/02_kv_interface.py +++ b/tutorial/02_kv_interface.py @@ -215,7 +215,7 @@ def main(): Key Methods: 1. (async_)kv_put - Insert/Update a multi-column sample by key, with optional metadata tag 2. (async_)kv_batch_put - Put multiple key-value pairs efficiently in batch - 3. (async_)kv_update - Merge new values into produced fields with merge_fn(old, new) + 3. (async_)kv_update - Merge into fields with merge_fn(old, new), adding missing keys/fields 4. (async_)kv_batch_get - Retrieve samples (by keys), supporting column selection (by fields) 5. (async_)kv_list - List keys and tags (metadata) in a partition 6. (async_)kv_clear - Remove key-value pairs, or only selected fields, from storage