From 7c563eb990286695eb303be601a3235d7427a2f7 Mon Sep 17 00:00:00 2001 From: Robert Washbourne Date: Fri, 23 Oct 2026 08:41:39 +0000 Subject: [PATCH 1/2] Pack live route payloads while preserving logical int16 schemas Signed-off-by: Robert Washbourne --- README.md | 23 +-- tests/e2e/test_live_route_checkpoint_e2e.py | 93 +++++++++++++ tests/test_checkpoint_route_packing.py | 28 ++-- tests/test_live_route_codec.py | 131 ++++++++++++++++++ .../managers/simple_storage_manager.py | 7 +- transfer_queue/storage/simple_storage.py | 67 ++++----- transfer_queue/utils/route_codec.py | 107 ++++++++++++++ transfer_queue/utils/serial_utils.py | 17 ++- transfer_queue/utils/zmq_utils.py | 4 +- transfer_queue/version/version | 2 +- 10 files changed, 412 insertions(+), 67 deletions(-) create mode 100644 tests/e2e/test_live_route_checkpoint_e2e.py create mode 100644 tests/test_live_route_codec.py create mode 100644 transfer_queue/utils/route_codec.py diff --git a/README.md b/README.md index 63e26f6a..7ef26a8b 100644 --- a/README.md +++ b/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 diff --git a/tests/e2e/test_live_route_checkpoint_e2e.py b/tests/e2e/test_live_route_checkpoint_e2e.py new file mode 100644 index 00000000..3dab8ac9 --- /dev/null +++ b/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)) diff --git a/tests/test_checkpoint_route_packing.py b/tests/test_checkpoint_route_packing.py index 28950d75..b2208756 100644 --- a/tests/test_checkpoint_route_packing.py +++ b/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": [ { diff --git a/tests/test_live_route_codec.py b/tests/test_live_route_codec.py new file mode 100644 index 00000000..3f221d4b --- /dev/null +++ b/tests/test_live_route_codec.py @@ -0,0 +1,131 @@ +# 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 + +import pytest +import torch +from tensordict import 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) diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 5cad6126..cd646ca3 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/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: diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 24e7d537..2d4595f9 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -28,11 +28,13 @@ import ray import torch import zmq +from tensordict import TensorDictBase from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.perf_utils import IntervalPerfMonitor +from transfer_queue.utils.route_codec import PackedRoutes, pack_route_fields, packable_route_ids from transfer_queue.utils.zmq_utils import ( ZMQMessage, ZMQRequestType, @@ -57,42 +59,7 @@ class _CheckpointPickler(pickle.Pickler): def __init__(self, stream, state): super().__init__(stream, protocol=pickle.HIGHEST_PROTOCOL) - self.routes = {id(value) for value in state["field_data"].get("routed_experts", {}).values()} - if not self.routes: - return - pending, seen, owners = [state], set(), {} - while pending: - value = pending.pop() - if id(value) in seen: - continue - seen.add(id(value)) - if type(value) is torch.Tensor: - if value.device.type != "cpu" or value.layout != torch.strided or value.is_nested: - self.routes.clear() - break - storage = value.untyped_storage() - identity = (storage.data_ptr(), storage.nbytes()) - owners.setdefault(identity, []).append(value) - elif type(value) is dict: - pending.extend(value.keys()) - pending.extend(value.values()) - elif type(value) in (list, tuple, set, frozenset): - pending.extend(value) - elif type(value) not in (str, bytes, bytearray, int, float, bool, type(None)): - # Opaque objects and NumPy views can conceal shared tensor storage. - self.routes.clear() - break - for aliases in owners.values(): - if len(aliases) <= 1: - continue - # Storage puts unbind dense/jagged batches into disjoint row views. - # Overlap or another field's view must retain ordinary serialization. - route_rows = all( - id(value) in self.routes and value.dtype == torch.int16 and value.is_contiguous() for value in aliases - ) - spans = sorted((value.storage_offset(), value.storage_offset() + value.numel()) for value in aliases) - if not route_rows or any(left[1] > right[0] for left, right in zip(spans, spans[1:], strict=False)): - self.routes.difference_update(id(value) for value in aliases) + self.routes = packable_route_ids(state["field_data"]) def reducer_override(self, value): if id(value) not in self.routes: @@ -140,7 +107,7 @@ def active_key_count(self) -> int: """Number of active keys currently stored.""" return len(self._active_keys) - def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: + def get_data(self, fields: list[str], global_indexes: list, *, packed_routes: bool = False) -> dict[str, list]: """Get data by global index keys. Args: @@ -158,6 +125,10 @@ def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: ) try: result[field] = [self.field_data[field][k] for k in global_indexes] + if not packed_routes: + result[field] = [ + value.restore() if isinstance(value, PackedRoutes) else value for value in result[field] + ] except KeyError as e: raise KeyError(f"StorageUnitData get_data: key {e} not found in field '{field}'") from e return result @@ -220,7 +191,11 @@ def checkpoint_snapshot(self, indexes: set[int]) -> tuple[dict, int | None]: continue seen.add(id(value)) size += sys.getsizeof(value) - if isinstance(value, torch.Tensor): + if isinstance(value, PackedRoutes): + pending.append(value.data) + elif isinstance(value, TensorDictBase): + pending.append(value.to_dict()) + elif isinstance(value, torch.Tensor): if type(value) is not torch.Tensor or value.device.type != "cpu" or value.is_nested: return {"field_data": fields, "active_keys": active}, None try: @@ -411,7 +386,7 @@ def _worker_routine(self) -> None: identity = messages[0] serialized_msg = messages[1:] - request_msg = ZMQMessage.deserialize(serialized_msg) + request_msg = ZMQMessage.deserialize(serialized_msg, retain_packed_routes=True) operation = request_msg.request_type try: @@ -523,7 +498,7 @@ def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: f"data_parser changed the number of elements for key '{k}': " f"expected {orig_len}, got {new_len}" ) - self.storage_data.put_data(field_data, global_indexes) + self.storage_data.put_data(pack_route_fields(field_data), global_indexes) # After put operation finish, send a message to the client response_msg = ZMQMessage.create( @@ -560,7 +535,7 @@ def _handle_get(self, data_parts: ZMQMessage) -> ZMQMessage: with limit_pytorch_auto_parallel_threads( target_num_threads=TQ_NUM_THREADS, info=f"[{self.storage_unit_id}] _handle_get" ): - result_data = self.storage_data.get_data(fields, global_indexes) + result_data = self.storage_data.get_data(fields, global_indexes, packed_routes=True) response_msg = ZMQMessage.create( request_type=ZMQRequestType.GET_DATA_RESPONSE, # type: ignore[arg-type] @@ -734,6 +709,14 @@ def write(): and retained_bytes is not None and retained_bytes <= data_parts.body["max_snapshot_bytes"] ) + logger.info( + "[%s]: checkpoint storage cut phase=%s async=%s retained_bytes=%s max_snapshot_bytes=%s", + self.storage_unit_id, + phase, + prepared, + retained_bytes, + data_parts.body.get("max_snapshot_bytes"), + ) future: Future[None] = Future() self._checkpoint_future, self._checkpoint_path = future, path if prepared: @@ -824,7 +807,7 @@ def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: if retained: candidate[field_name] = retained active_keys = set().union(*(values.keys() for values in candidate.values())) - self.storage_data.field_data = candidate + self.storage_data.field_data = pack_route_fields(candidate) self.storage_data._active_keys = active_keys logger.info( diff --git a/transfer_queue/utils/route_codec.py b/transfer_queue/utils/route_codec.py new file mode 100644 index 00000000..135a320f --- /dev/null +++ b/transfer_queue/utils/route_codec.py @@ -0,0 +1,107 @@ +# Copyright 2026 Zyphra Technologies Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Physical route bytes retain the public int16 field contract.""" + +import torch +from tensordict import TensorDictBase + + +class PackedRoutes: + __slots__ = ("_data", "_offset") + + def __init__(self, data, offset=0): + if type(data) is not torch.Tensor or data.dtype != torch.uint8 or data.device.type != "cpu": + raise ValueError("Packed routes require a CPU uint8 tensor") + if type(offset) is not int or offset not in (0, -1): + raise ValueError("Packed route offset must be zero or the -1 sentinel shift") + self._data, self._offset = data, offset + + @property + def data(self): + return self._data + + @property + def offset(self): + return self._offset + + def restore(self): + value = self.data.to(torch.int16) + return value.add_(self.offset) if self.offset else value + + def __reduce__(self): + # Ordinary Torch readers restore the original logical dtype. + if self.offset: + return torch.add, (PackedRoutes(self.data), self.offset) + return torch._utils._rebuild_device_tensor_from_numpy, (self.data.numpy(), torch.int16, "cpu", False) + + +def packable_route_ids(fields, *, reject_identity_aliases=False): + rows = fields.get("routed_experts", {}) + routes = {id(value) for value in (rows.values() if isinstance(rows, dict) else rows) if type(value) is torch.Tensor} + if not routes: + return set() + pending, seen, owners, counts = [fields], set(), {}, {} + while pending: + value = pending.pop() + if type(value) is torch.Tensor: + counts[id(value)] = counts.get(id(value), 0) + 1 + if id(value) in seen: + continue + seen.add(id(value)) + if isinstance(value, PackedRoutes): + pending.append(value.data) + elif type(value) is torch.Tensor: + if value.device.type != "cpu" or value.layout != torch.strided or value.is_nested: + return set() + storage = value.untyped_storage() + identity = (storage.data_ptr(), storage.nbytes()) + owners.setdefault(identity, []).append(value) + elif isinstance(value, TensorDictBase): + pending.append(value.to_dict()) + elif type(value) is dict: + pending.extend(value.keys()) + pending.extend(value.values()) + elif type(value) in (list, tuple, set, frozenset): + pending.extend(value) + elif type(value) not in (str, bytes, bytearray, int, float, bool, type(None)): + return set() + for aliases in owners.values(): + if len(aliases) <= 1: + continue + route_rows = all( + id(value) in routes and value.dtype == torch.int16 and value.is_contiguous() for value in aliases + ) + spans = sorted((value.storage_offset(), value.storage_offset() + value.numel()) for value in aliases) + if not route_rows or any(left[1] > right[0] for left, right in zip(spans, spans[1:], strict=False)): + routes.difference_update(id(value) for value in aliases) + if reject_identity_aliases: + routes.difference_update(identity for identity, count in counts.items() if count > 1) + return routes + + +def pack_route_fields(fields): + """Pack independent route cells; unsupported values retain their old path.""" + routes = fields.get("routed_experts") + if routes is None: + return fields + rows = dict(routes) if isinstance(routes, dict) else dict(enumerate(routes)) + normalized = {**fields, "routed_experts": rows} + eligible = packable_route_ids(normalized, reject_identity_aliases=True) + for index, value in rows.items(): + if ( + id(value) not in eligible + or type(value) is not torch.Tensor + or value.dtype != torch.int16 + or not value.is_contiguous() + or value.is_neg() + or value.__dict__ + or any(name is not None for name in getattr(value, "names", ())) + ): + continue + array = value.numpy() + low, high = (int(array.min()), int(array.max())) if array.size else (0, 0) + if low >= 0 and high <= 255: + rows[index] = PackedRoutes(value.to(torch.uint8)) + elif low == -1 and high <= 254: + rows[index] = PackedRoutes((value + 1).to(torch.uint8), offset=-1) + return {**fields, "routed_experts": rows if isinstance(routes, dict) else list(rows.values())} diff --git a/transfer_queue/utils/serial_utils.py b/transfer_queue/utils/serial_utils.py index cfd9a022..8c193ee3 100644 --- a/transfer_queue/utils/serial_utils.py +++ b/transfer_queue/utils/serial_utils.py @@ -36,12 +36,14 @@ from tensordict import TensorDictBase from transfer_queue.utils.logging_utils import get_logger +from transfer_queue.utils.route_codec import PackedRoutes CUSTOM_TYPE_PICKLE = 1 CUSTOM_TYPE_CLOUDPICKLE = 2 CUSTOM_TYPE_TENSOR = 3 # For tensor with buffer reference CUSTOM_TYPE_NESTED_TENSOR = 4 # For nested tensor (strided or jagged) CUSTOM_TYPE_NUMPY = 5 # For numpy ndarray with buffer reference +CUSTOM_TYPE_PACKED_ROUTES = 6 # 0xC1 is permanently reserved (invalid) in msgpack spec — safe to use as pickle fallback sentinel. _PICKLE_FALLBACK_SENTINEL = b"\xc1\xfe\xed" @@ -106,6 +108,9 @@ def enc_hook(self, obj: Any) -> Any: - numpy.ndarray: Convert to tensor for unified handling """ + if isinstance(obj, PackedRoutes): + meta = (self._encode_regular_tensor_meta(obj.data), obj.offset) + return msgpack.Ext(CUSTOM_TYPE_PACKED_ROUTES, pickle.dumps(meta, protocol=pickle.HIGHEST_PROTOCOL)) if isinstance(obj, torch.Tensor): return self._encode_tensor(obj) @@ -248,7 +253,8 @@ class MsgpackDecoder: threads and async coroutines. """ - def __init__(self): + def __init__(self, *, retain_packed_routes=False): + self.retain_packed_routes = retain_packed_routes self.decoder = msgpack.Decoder(ext_hook=self.ext_hook) @property @@ -378,12 +384,17 @@ def ext_hook(self, code: int, data: memoryview) -> Any: if code == CUSTOM_TYPE_NUMPY: meta = pickle.loads(data) return self._decode_numpy(meta) + if code == CUSTOM_TYPE_PACKED_ROUTES: + meta, offset = pickle.loads(data) + value = PackedRoutes(self._decode_tensor(meta), offset=offset) + return value if self.retain_packed_routes else value.restore() raise NotImplementedError(f"Extension type code {code} is not supported") _encoder = MsgpackEncoder() _decoder = MsgpackDecoder() +_storage_decoder = MsgpackDecoder(retain_packed_routes=True) # Values msgpack cannot represent. None of these derive from ValueError: OverflowError is @@ -430,7 +441,7 @@ def encode(obj: Any) -> list[bytestr]: return [_PICKLE_FALLBACK_SENTINEL, pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)] -def decode(frames: list) -> Any: +def decode(frames: list, *, retain_packed_routes=False) -> Any: """Decode frames produced by encode. Transparently handles both the msgpack zero-copy path and the pickle @@ -438,7 +449,7 @@ def decode(frames: list) -> Any: """ if len(frames) >= 2 and _is_pickle_fallback(frames[0]): return pickle.loads(frames[1]) - return _decoder.decode(frames) + return (_storage_decoder if retain_packed_routes else _decoder).decode(frames) # Packed buffer layout: diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 2ca63d7f..d976c6da 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -215,7 +215,7 @@ def serialize(self) -> list: return encode(msg_dict) @classmethod - def deserialize(cls, frames: list) -> "ZMQMessage": + def deserialize(cls, frames: list, *, retain_packed_routes=False) -> "ZMQMessage": """Deserialize: choose decoding path based on the first frame marker (zero-copy or pickle fallback).""" if not frames: raise ValueError("Empty frames received") @@ -228,7 +228,7 @@ def deserialize(cls, frames: list) -> "ZMQMessage": raise ZMQMessageDecodeError(f"leading frame is empty; {describe_frames(frames)}") try: - result = decode(frames) + result = decode(frames, retain_packed_routes=retain_packed_routes) except Exception as e: raise ZMQMessageDecodeError(f"{type(e).__name__}: {e}; {describe_frames(frames)}") from e diff --git a/transfer_queue/version/version b/transfer_queue/version/version index 1be458f8..6b956112 100644 --- a/transfer_queue/version/version +++ b/transfer_queue/version/version @@ -1 +1 @@ -0.1.10+zyphra.1 +0.1.10+zyphra.2 From 417c699b0975a54ec415ea94ac1c3237b6c4443c Mon Sep 17 00:00:00 2001 From: Robert Washbourne Date: Fri, 23 Oct 2026 08:49:53 +0000 Subject: [PATCH 2/2] Retain TensorDict projections throughout ownership walks Signed-off-by: Robert Washbourne --- tests/test_live_route_codec.py | 14 +++++++++++++- transfer_queue/storage/simple_storage.py | 6 ++++-- transfer_queue/utils/route_codec.py | 6 ++++-- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/test_live_route_codec.py b/tests/test_live_route_codec.py index 3f221d4b..3fca148e 100644 --- a/tests/test_live_route_codec.py +++ b/tests/test_live_route_codec.py @@ -5,10 +5,11 @@ import pickle import subprocess import sys +from types import SimpleNamespace import pytest import torch -from tensordict import NonTensorStack, TensorDict +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 @@ -129,3 +130,14 @@ def test_packed_checkpoint_uses_ordinary_torch_reader(tmp_path, offset): 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 diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 2d4595f9..d1ff5f72 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -184,7 +184,7 @@ def checkpoint_snapshot(self, indexes: set[int]) -> tuple[dict, int | None]: } active = self._active_keys & indexes stores: set[tuple[int, int] | int] = set() - pending, seen, size = [fields, active], set(), 0 + pending, seen, size, projections = [fields, active], set(), 0, [] while pending: value = pending.pop() if id(value) in seen: @@ -194,7 +194,9 @@ def checkpoint_snapshot(self, indexes: set[int]) -> tuple[dict, int | None]: if isinstance(value, PackedRoutes): pending.append(value.data) elif isinstance(value, TensorDictBase): - pending.append(value.to_dict()) + # Projected containers must outlive the identity-based walk. + projections.append(value.to_dict()) + pending.append(projections[-1]) elif isinstance(value, torch.Tensor): if type(value) is not torch.Tensor or value.device.type != "cpu" or value.is_nested: return {"field_data": fields, "active_keys": active}, None diff --git a/transfer_queue/utils/route_codec.py b/transfer_queue/utils/route_codec.py index 135a320f..aa233a39 100644 --- a/transfer_queue/utils/route_codec.py +++ b/transfer_queue/utils/route_codec.py @@ -40,7 +40,7 @@ def packable_route_ids(fields, *, reject_identity_aliases=False): routes = {id(value) for value in (rows.values() if isinstance(rows, dict) else rows) if type(value) is torch.Tensor} if not routes: return set() - pending, seen, owners, counts = [fields], set(), {}, {} + pending, seen, owners, counts, projections = [fields], set(), {}, {}, [] while pending: value = pending.pop() if type(value) is torch.Tensor: @@ -57,7 +57,9 @@ def packable_route_ids(fields, *, reject_identity_aliases=False): identity = (storage.data_ptr(), storage.nbytes()) owners.setdefault(identity, []).append(value) elif isinstance(value, TensorDictBase): - pending.append(value.to_dict()) + # Keep projected containers alive while their identities are in seen. + projections.append(value.to_dict()) + pending.append(projections[-1]) elif type(value) is dict: pending.extend(value.keys()) pending.extend(value.values())