Pack live R3 route storage and frames with exact int16 restoration

README.md

@@ -367,14 +367,21 @@ Please kindly cite our paper if you find this repo is useful:
 
 ### Bounded asynchronous checkpoint storage
 
-SimpleStorage checkpoint writers store nonoverlapping contiguous CPU `routed_experts`
-int16 tensors as uint8 when every value is in 0..255. Ordinary pickle/Torch
-readers reconstruct the original int16 dtype, shape and values, including older
-readers; existing checkpoints remain readable. Live storage and snapshot byte
-accounting stay unchanged. Disjoint route rows from decoded dense/jagged batches
-are eligible. Overlapping views, cross-field backing, out-of-range values and
-uncertain ownership use ordinary serialization. Packing adds a range scan and
-a temporary uint8 copy in the writer; it is not a bound on process RSS.
+SimpleStorage packs independent CPU `routed_experts` cells into live uint8
+storage and transport frames while retaining int16 controller schemas and
+restoring int16 consumer values. IDs in 0..255 fit; a -1 sentinel also fits
+when all other IDs are at most 254. Wider ranges, unsupported values,
+overlapping or cross-field aliases retain their previous representation.
+No token, mask, log probability or route entries are dropped. All clients
+and storage processes must use the same package version; new readers accept
+old unpacked frames and checkpoints. Checkpoint packing restores the logical
+dtype through ordinary Torch reducers, including for the shifted sentinel.
+Disjoint route rows from decoded dense/jagged batches remain eligible.
+
+Snapshot ownership accounting traverses TensorDict metadata through its public
+dictionary projection. Unknown payload ownership still selects synchronous
+storage. Each storage cut logs its accounted bytes, budget and async decision
+before starting a write; a larger budget alone cannot admit unknown ownership.
 
 For detached SimpleStorage saves, committed `metadata.json` includes an optional
 `storage_checkpoint` receipt: aggregate async/fallback mode and the existing

tests/e2e/test_live_route_checkpoint_e2e.py

@@ -0,0 +1,93 @@
+# Copyright 2026 Zyphra Technologies Inc.
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""Public KV route schemas and sampled fields survive a real save/restore."""
+
+import hashlib
+import tempfile
+
+import pytest
+import ray
+import torch
+from omegaconf import OmegaConf
+from tensordict import NonTensorStack, TensorDict
+
+import transfer_queue as tq
+
+
+@pytest.fixture
+def system():
+    assert not ray.is_initialized()
+    directory = tempfile.TemporaryDirectory(prefix="tq-rc-")
+    ray.init(
+        address="local",
+        num_cpus=3,
+        num_gpus=0,
+        include_dashboard=False,
+        object_store_memory=128 * 1024**2,
+        _node_ip_address="127.0.0.1",
+        _temp_dir=directory.name,
+    )
+    config = OmegaConf.create(
+        {
+            "controller": {"polling_mode": True},
+            "metrics": {"enabled": False},
+            "backend": {
+                "storage_backend": "SimpleStorage",
+                "SimpleStorage": {
+                    "total_storage_size": 32,
+                    "num_data_storage_units": 2,
+                },
+            },
+        }
+    )
+    try:
+        tq.init(config)
+        yield
+    finally:
+        tq.close()
+        ray.shutdown()
+        directory.cleanup()
+
+
+def tensor_hash(value):
+    digest = hashlib.sha256()
+    for row in value.unbind(0):
+        digest.update(str((row.dtype, tuple(row.shape))).encode())
+        digest.update(row.contiguous().view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def test_public_kv_metadata_and_async_checkpoint_roundtrip(system, tmp_path):
+    keys, partition = ["row-a", "row-b"], "live-route"
+    routes = torch.arange(640, dtype=torch.int16).reshape(2, 8, 40, 1) % 16
+    routes[0, 0, 0, 0] = -1
+    fields = TensorDict(
+        {
+            "routed_experts": routes,
+            "input_ids": torch.arange(16, dtype=torch.int64).reshape(2, 8),
+            "position_ids": torch.arange(8, dtype=torch.int64).expand(2, 8).clone(),
+            "response_mask": torch.tensor([[0, 0, 1, 1, 1, 1, 0, 0]] * 2),
+            "rollout_log_probs": torch.linspace(-3, -0.1, 16).reshape(2, 8),
+            "metrics": NonTensorStack({"reward": 1.0}, {"reward": 0.0}),
+        },
+        batch_size=[2],
+    )
+    expected = {name: tensor_hash(value) for name, value in fields.items() if isinstance(value, torch.Tensor)}
+    tags = [{"status": "completed", "session_id": index, "global_steps": 9} for index in range(2)]
+    tq.kv_batch_put(keys=keys, partition_id=partition, fields=fields, tags=tags)
+    meta = tq.get_client().kv_retrieve_meta(keys, partition, create=False)
+    assert meta.field_schema["routed_experts"]["dtype"] == torch.int16
+    before = tq.kv_batch_get(keys=keys, partition_id=partition)
+    assert {name: tensor_hash(before[name]) for name in expected} == expected
+    path = tmp_path / "checkpoint"
+    future = tq.save_checkpoint(path, async_storage=True, max_snapshot_bytes=16 * 1024**2, timeout_s=30)
+    assert future is not None
+    future.result(timeout=30)
+    assert (path / "metadata.json").is_file()
+    tq.kv_clear(keys=keys, partition_id=partition)
+    tq.load_checkpoint(path)
+    after = tq.kv_batch_get(keys=keys, partition_id=partition)
+    restored_meta = tq.get_client().kv_retrieve_meta(keys, partition, create=False)
+    assert restored_meta.field_schema["routed_experts"] == meta.field_schema["routed_experts"]
+    assert {name: tensor_hash(after[name]) for name in expected} == expected
+    assert tq.kv_list(partition_id=partition)[partition] == dict(zip(keys, tags, strict=True))

tests/test_checkpoint_route_packing.py

@@ -16,7 +16,11 @@
 import pytest
 import torch
 
-from transfer_queue.storage.simple_storage import SimpleStorageUnit, StorageUnitData, _CheckpointPickler
+from transfer_queue.storage.simple_storage import (
+    SimpleStorageUnit,
+    StorageUnitData,
+    _CheckpointPickler,
+)
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
 
 
@@ -144,8 +148,8 @@ def test_real_save_load_handler_and_unmodified_reader(tmp_path, legacy):
         assert reply.body["success"] and reply.body["failure"] is None
     loaded = unit()
     assert loaded._handle_load_checkpoint(request(path)).body["success"]
-    assert torch.equal(loaded.storage_data.field_data["routed_experts"][0], route)
-    assert loaded.storage_data.field_data["routed_experts"][0].dtype == torch.int16
+    restored = loaded.storage_data.get_data(["routed_experts"], [0])["routed_experts"][0]
+    assert torch.equal(restored, route) and restored.dtype == torch.int16
     # Isolated Python imports no candidate SDK; the installed Torch/pickle reader
     # must reconstruct the exact original tensor without a new decoding branch.
     code = (
@@ -154,7 +158,9 @@ def test_real_save_load_handler_and_unmodified_reader(tmp_path, legacy):
         "assert v.dtype==torch.int16 and tuple(v.shape)==(4,40,1); "
         "assert torch.equal(v,torch.arange(160,dtype=torch.int16).reshape(4,40,1)%16)"
     )
-    subprocess.run([sys.executable, "-I", "-c", code, str(path)], check=True, timeout=30)
+    packed_path = tmp_path / "ordinary-reader.pkl"
+    packed_path.write_bytes(packed(state_for(route)))
+    subprocess.run([sys.executable, "-I", "-c", code, str(packed_path)], check=True, timeout=30)
 
 
 def test_async_capture_survives_clear_and_reuse(tmp_path, monkeypatch):
@@ -265,7 +271,7 @@ def test_public_transport_put_save_route_eligibility(tmp_path, layout):
     started = time.perf_counter()
     assert actor._handle_put(received).request_type == ZMQRequestType.PUT_DATA_RESPONSE
     costs["put_s"] = time.perf_counter() - started
-    values = actor.storage_data.field_data["routed_experts"]
+    values = dict(enumerate(actor.storage_data.get_data(["routed_experts"], list(range(4)))["routed_experts"]))
     state = state_for(values[0])
     state["field_data"]["routed_experts"], state["active_keys"] = values, set(values)
     started = time.perf_counter()
@@ -286,9 +292,9 @@ def test_public_transport_put_save_route_eligibility(tmp_path, layout):
     assert actor._handle_save_checkpoint(request(path)).body["failure"] is None
     costs["packed_handler_serialize_write_fsync_s"] = time.perf_counter() - started
     payload = path.read_bytes()
-    assert b"_rebuild_device_tensor_from_numpy" in payload
+    assert b"_rebuild_device_tensor_from_numpy" in direct
     assert len(payload) < len(legacy) * 0.6
-    restored = pickle.loads(payload)["field_data"]["routed_experts"]
+    restored = pickle.loads(path.read_bytes())["field_data"]["routed_experts"]
     ordinary = pickle.loads(legacy)["field_data"]["routed_experts"]
     for index, expected in enumerate(rows):
         assert restored[index].dtype == expected.dtype and restored[index].shape == expected.shape
@@ -305,7 +311,7 @@ def test_public_transport_put_save_route_eligibility(tmp_path, layout):
     loaded = unit()
     assert loaded._handle_load_checkpoint(request(path)).body["success"]
     for index, expected in enumerate(rows):
-        assert torch.equal(loaded.storage_data.field_data["routed_experts"][index], expected)
+        assert torch.equal(loaded.storage_data.get_data(["routed_experts"], [index])["routed_experts"][0], expected)
     legacy_path = tmp_path / "legacy.pkl"
     legacy_path.write_bytes(legacy)
     code = (
@@ -315,7 +321,9 @@ def test_public_transport_put_save_route_eligibility(tmp_path, layout):
         "assert a.keys()==b.keys(); "
         "assert all(a[k].dtype==b[k].dtype and a[k].shape==b[k].shape and torch.equal(a[k],b[k]) for k in a)"
     )
-    subprocess.run([sys.executable, "-I", "-c", code, str(path), str(legacy_path)], check=True, timeout=30)
+    packed_path = tmp_path / "ordinary-reader.pkl"
+    packed_path.write_bytes(direct)
+    subprocess.run([sys.executable, "-I", "-c", code, str(packed_path), str(legacy_path)], check=True, timeout=30)
     for mode, data in (("legacy", legacy), ("packed", direct)):
         root = os.environ.get("CHECKPOINT_BENCH_WRITE_ROOT", str(tmp_path))
         with tempfile.TemporaryDirectory(prefix="route-wire-", dir=root) as directory:
@@ -332,7 +340,7 @@ def test_public_transport_put_save_route_eligibility(tmp_path, layout):
                 "layout": layout,
                 "legacy_bytes": len(legacy),
                 "saved_bytes": len(payload),
-                "packed": b"_rebuild_device_tensor_from_numpy" in payload,
+                "packed": b"_rebuild_device_tensor_from_numpy" in direct,
                 "costs": costs,
                 "rows": [
                     {

tests/test_live_route_codec.py

@@ -0,0 +1,143 @@
+# Copyright 2026 Zyphra Technologies Inc.
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""Public wire/storage roundtrips keep logical fields while retaining half-width routes."""
+
+import pickle
+import subprocess
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+from tensordict import NonTensorData, NonTensorStack, TensorDict
+
+from transfer_queue.storage.simple_storage import SimpleStorageUnit, StorageUnitData
+from transfer_queue.utils.route_codec import PackedRoutes, pack_route_fields
+from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
+
+
+def unit():
+    actor = object.__new__(SimpleStorageUnit.__ray_metadata__.modified_class)
+    actor.storage_unit_id, actor.storage_unit_size = "live-route-test", 8
+    actor.storage_data = StorageUnitData(8)
+    return actor
+
+
+def put(fields):
+    return ZMQMessage.create(ZMQRequestType.PUT_DATA, "test", {"global_indexes": [0, 1], "data": fields})
+
+
+def get():
+    return ZMQMessage.create(
+        ZMQRequestType.GET_DATA, "test", {"global_indexes": [0, 1], "fields": ["routed_experts", "input_ids"]}
+    )
+
+
+@pytest.mark.parametrize("layout", ["list", "dense", "jagged"])
+def test_real_wire_put_physical_bytes_get_and_restore(tmp_path, layout):
+    rows = [torch.arange(320, dtype=torch.int16).reshape(8, 40, 1) % 16 for _ in range(2)]
+    routes = (
+        rows
+        if layout == "list"
+        else (torch.stack(rows) if layout == "dense" else torch.nested.as_nested_tensor(rows, layout=torch.jagged))
+    )
+    tokens = torch.arange(16).reshape(2, 8)
+    original = {"routed_experts": routes, "input_ids": tokens}
+    old_frames = put(original).serialize()
+    frames = put(pack_route_fields(original)).serialize()
+    assert sum(memoryview(x).nbytes for x in frames[1:]) < sum(memoryview(x).nbytes for x in old_frames[1:])
+    request = ZMQMessage.deserialize(frames, retain_packed_routes=True)
+    actor = unit()
+    assert actor._handle_put(request).request_type == ZMQRequestType.PUT_DATA_RESPONSE
+    physical = actor.storage_data.field_data["routed_experts"]
+    assert all(isinstance(value, PackedRoutes) and value.data.dtype == torch.uint8 for value in physical.values())
+    assert sum(value.data.nbytes for value in physical.values()) == sum(value.nbytes for value in rows) // 2
+    received = ZMQMessage.deserialize(actor._handle_get(get()).serialize()).body["data"]
+    for index, expected in enumerate(rows):
+        assert received["routed_experts"][index].dtype == torch.int16
+        assert torch.equal(received["routed_experts"][index], expected)
+        assert torch.equal(received["input_ids"][index], tokens[index])
+    path = tmp_path / "storage.pkl"
+    save = ZMQMessage.create(ZMQRequestType.SAVE_STORAGE_CHECKPOINT, "test", {"path": str(path)})
+    assert actor._handle_save_checkpoint(save).body["failure"] is None
+    restored = unit()
+    assert restored._handle_load_checkpoint(save).body["success"]
+    again = ZMQMessage.deserialize(restored._handle_get(get()).serialize()).body["data"]
+    assert all(torch.equal(a, b) for a, b in zip(again["routed_experts"], received["routed_experts"], strict=True))
+    assert all(isinstance(value, PackedRoutes) for value in restored.storage_data.field_data["routed_experts"].values())
+
+
+@pytest.mark.parametrize("values", [[0, 255], [-1, 0], [-1, 254], [-1, 255], [0, 256], [-32768, 32767], []])
+def test_range_guards_and_mixed_old_wire(values):
+    route = torch.tensor(values, dtype=torch.int16)
+    fields = pack_route_fields({"routed_experts": [route]})
+    physical = fields["routed_experts"][0]
+    eligible = all(0 <= value <= 255 for value in values) or (min(values, default=0) == -1 and max(values) <= 254)
+    assert isinstance(physical, PackedRoutes) == eligible
+    message = ZMQMessage.create(ZMQRequestType.PUT_DATA, "test", {"data": fields})
+    restored = ZMQMessage.deserialize(message.serialize()).body["data"]["routed_experts"][0]
+    assert restored.dtype == route.dtype and torch.equal(restored, route)
+    legacy = ZMQMessage.create(ZMQRequestType.PUT_DATA, "test", {"data": {"routed_experts": [route]}})
+    old = ZMQMessage.deserialize(legacy.serialize(), retain_packed_routes=True).body["data"]["routed_experts"][0]
+    assert old.dtype == torch.int16 and torch.equal(old, route)
+
+
+def test_cross_field_and_overlapping_aliases_keep_original_values():
+    route = torch.arange(16, dtype=torch.int16)
+    fields = {"routed_experts": [route], "alias": [route]}
+    packed = pack_route_fields(fields)
+    assert packed["routed_experts"][0] is packed["alias"][0] is route
+    fields = {"routed_experts": [route[:10], route[5:]], "input_ids": [torch.arange(2)]}
+    packed = pack_route_fields(fields)
+    assert all(type(value) is torch.Tensor for value in packed["routed_experts"])
+    packed["routed_experts"][0][5] = 98
+    assert packed["routed_experts"][1][0] == 98
+
+
+def test_known_tensordict_metadata_allows_accounted_async_cut():
+    rows = NonTensorStack({"field": "one"}, {"field": "two"})
+    actor = unit()
+    request = put({"routed_experts": [torch.arange(16, dtype=torch.int16)] * 2, "metadata": rows})
+    decoded = ZMQMessage.deserialize(request.serialize(), retain_packed_routes=True)
+    assert actor._handle_put(decoded).request_type == ZMQRequestType.PUT_DATA_RESPONSE
+    assert all(isinstance(value, TensorDict) for value in actor.storage_data.field_data["metadata"].values())
+    snapshot, size = actor.storage_data.checkpoint_snapshot({0, 1})
+    assert size is not None and size > 32
+    ordinary = pickle.loads(pickle.dumps(snapshot))
+    assert ordinary["active_keys"] == {0, 1}
+
+
+@pytest.mark.parametrize("dtype", [torch.uint8, torch.int32, torch.float32])
+def test_other_route_dtypes_are_unchanged(dtype):
+    value = torch.arange(16, dtype=dtype)
+    assert pack_route_fields({"routed_experts": [value]})["routed_experts"][0] is value
+
+
+@pytest.mark.parametrize("offset", [-1, 0])
+def test_packed_checkpoint_uses_ordinary_torch_reader(tmp_path, offset):
+    value = PackedRoutes(torch.tensor([0, 1, 254], dtype=torch.uint8), offset=offset)
+    path = tmp_path / "route.pkl"
+    path.write_bytes(pickle.dumps(value))
+    code = (
+        "import pickle,sys,torch;v=pickle.load(open(sys.argv[1],'rb'));"
+        "assert v.dtype==torch.int16;"
+        "assert torch.equal(v,torch.tensor([0,1,254],dtype=torch.int16)+int(sys.argv[2]))"
+    )
+    subprocess.run([sys.executable, "-I", "-c", code, str(path), str(offset)], check=True, timeout=30)
+
+
+@pytest.mark.parametrize("dtype,offset", [(torch.int16, 0), (torch.uint8, 1)])
+def test_invalid_physical_encoding_is_refused(dtype, offset):
+    with pytest.raises(ValueError):
+        PackedRoutes(torch.tensor([0, 1], dtype=dtype), offset=offset)
+
+
+def test_many_tensordict_projections_do_not_hide_unknown_ownership():
+    metadata = [TensorDict({"value": index}, batch_size=[]) for index in range(128)]
+    metadata.append(TensorDict({"opaque": NonTensorData(SimpleNamespace(value=1))}, batch_size=[]))
+    data = StorageUnitData(256)
+    data.put_data({"metadata": metadata}, list(range(len(metadata))))
+    _, retained = data.checkpoint_snapshot(set(range(len(metadata))))
+    assert retained is None
+    fields = {"routed_experts": [torch.arange(16, dtype=torch.int16)], "metadata": metadata}
+    assert type(pack_route_fields(fields)["routed_experts"][0]) is torch.Tensor

transfer_queue/storage/managers/simple_storage_manager.py

@@ -32,6 +32,7 @@
 from transfer_queue.metadata import BatchMeta, extract_field_schema
 from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory
 from transfer_queue.utils.logging_utils import get_logger
+from transfer_queue.utils.route_codec import pack_route_fields
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
     ZMQRequestType,
@@ -389,7 +390,11 @@ async def _put_to_single_storage_unit(
             request_type=ZMQRequestType.PUT_DATA,  # type: ignore[arg-type]
             sender_id=self.storage_manager_id,
             receiver_id=target_storage_unit,
-            body={"global_indexes": global_indexes, "data": storage_data, "data_parser": data_parser},
+            body={
+                "global_indexes": global_indexes,
+                "data": pack_route_fields(storage_data) if data_parser is None else storage_data,
+                "data_parser": data_parser,
+            },
         )
 
         try: