[WIP] Feat/kv share

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

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))

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)