diff --git a/README.md b/README.md index 7ef26a8b..3ebc4691 100644 --- a/README.md +++ b/README.md @@ -383,6 +383,25 @@ 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. +Each existing per-unit checkpoint file streams the same logical pickle through +zstd level 1, using the existing `TQ_NUM_THREADS` budget for parallel compression. +The versioned header records the compressed length and SHA256; the frame also +has a zstd checksum. Readers validate the complete frame before replacing live +storage and accept legacy raw pickle shards. New compressed files require the +new reader; route reducers alone remain compatible with ordinary Torch pickle. +The existing concurrent unit fanout, async ownership, fsync and directory +publication stay in charge: a failed or incomplete unit cannot publish a cut. +This does not establish a production compression ratio or commit-time bound. + +The existing `GET_METRICS` request accepts optional `route_stats: true` and +returns `route_storage`: packed/unpacked/unknown row counts, known logical and +stored view bytes, and deduplicated backing-store bytes for routed experts. +It reads tensor metadata only, does not unpack values, and excludes other +fields, allocator overhead and RSS. Unknown rows remain explicit. Ordinary +requests and background collection leave this scan off. The existing +`TQMetricsExporter._query_storage_unit(..., route_stats=True)` sends the bounded +query without starting a monitor; its default remains unchanged. + For detached SimpleStorage saves, committed `metadata.json` includes an optional `storage_checkpoint` receipt: aggregate async/fallback mode and the existing per-unit ID, IP, preparation flag, accounted retained bytes (or `null` when diff --git a/requirements.txt b/requirements.txt index 088da4b3..6c84af5b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ msgspec psutil omegaconf prometheus_client>=0.20.0 +zstandard>=0.25.0 diff --git a/tests/e2e/test_checkpoint_ownership_e2e.py b/tests/e2e/test_checkpoint_ownership_e2e.py index c5900750..191cc4c4 100644 --- a/tests/e2e/test_checkpoint_ownership_e2e.py +++ b/tests/e2e/test_checkpoint_ownership_e2e.py @@ -25,6 +25,7 @@ import transfer_queue as tq from transfer_queue import interface +from transfer_queue.storage.simple_storage import _read_storage_checkpoint @pytest.fixture(scope="module") @@ -90,8 +91,7 @@ def physical_state(path): states = [] for entry in sorted(entries, key=lambda entry: entry["position"]): file = path / "simple_storage" / f"su_{entry['position']}_{entry['storage_unit_id']}.pkl" - with file.open("rb") as f: - states.append(pickle.load(f)) + states.append(_read_storage_checkpoint(file)) return states diff --git a/tests/e2e/test_live_route_checkpoint_e2e.py b/tests/e2e/test_live_route_checkpoint_e2e.py index 3dab8ac9..e56cd540 100644 --- a/tests/e2e/test_live_route_checkpoint_e2e.py +++ b/tests/e2e/test_live_route_checkpoint_e2e.py @@ -12,6 +12,7 @@ from tensordict import NonTensorStack, TensorDict import transfer_queue as tq +from transfer_queue.metrics import TQMetricsExporter @pytest.fixture @@ -77,6 +78,21 @@ def test_public_kv_metadata_and_async_checkpoint_roundtrip(system, tmp_path): 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 + reader = TQMetricsExporter(role="diagnostic") + try: + reports = [ + reader._query_storage_unit(info, name, route_stats=True)["route_storage"] + for name, info in tq.get_client().storage_manager.get_zmq_server_info().items() + ] + finally: + for socket in reader._zmq_sockets.values(): + socket.close(linger=0) + if reader._zmq_ctx is not None: + reader._zmq_ctx.term() + assert sum(report["packed_rows"] for report in reports) == 2 + assert sum(report["stored_dtype_counts"].get("torch.uint8", 0) for report in reports) == 2 + assert sum(report["logical_view_bytes_known"] for report in reports) == routes.nbytes + assert sum(report["stored_view_bytes_known"] for report in reports) == routes.numel() 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" diff --git a/tests/test_checkpoint_compression.py b/tests/test_checkpoint_compression.py new file mode 100644 index 00000000..c3eeb734 --- /dev/null +++ b/tests/test_checkpoint_compression.py @@ -0,0 +1,229 @@ +# Copyright 2026 Zyphra Technologies Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Lossless compressed shards keep the existing storage cut and failure contract.""" + +import hashlib +import io +import json +import pickle +import struct +import time +from threading import Event + +import numpy as np +import pytest +import torch +import zstandard + +from transfer_queue.storage.simple_storage import ( + SimpleStorageUnit, + StorageUnitData, + _CheckpointHashStream, + _CheckpointPickler, + _read_storage_checkpoint, + _write_storage_checkpoint, +) +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 = "compressed-test", 8 + actor.storage_data = StorageUnitData(8) + return actor + + +def request(path, **body): + return ZMQMessage.create(ZMQRequestType.SAVE_STORAGE_CHECKPOINT, "test", {"path": str(path), **body}) + + +def state(): + backing = torch.arange(24, dtype=torch.int64) + return { + "storage_unit_id": "compressed-test", + "storage_unit_size": 8, + "field_data": { + "input_ids": {0: backing, 1: backing[2:8]}, + "position_ids": {0: backing}, + "response_mask": {0: torch.tensor([False, True, True])}, + "routed_experts": {0: torch.arange(160, dtype=torch.int16).reshape(4, 40, 1) % 16}, + "log_probs": {0: torch.tensor([-0.0, float("inf"), float("nan")], dtype=torch.float32)}, + "metadata": {0: {"route": "train-route", "text": "\u03bb\n", "raw": bytes(range(256))}}, + "numpy": {0: np.arange(12, dtype=np.int32).reshape(3, 4)}, + }, + "active_keys": {0, 1}, + } + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_exact_tensor_bytes_dtypes_shapes_objects_and_aliases(tmp_path, legacy): + original = state() + path = tmp_path / "checkpoint.pkl" + if legacy: + path.write_bytes(pickle.dumps(original, protocol=pickle.HIGHEST_PROTOCOL)) + else: + _write_storage_checkpoint(path, original) + actual = _read_storage_checkpoint(path) + baseline = pickle.loads(pickle.dumps(original, protocol=pickle.HIGHEST_PROTOCOL)) + fields = actual["field_data"] + for name, values in original["field_data"].items(): + for index, value in values.items(): + restored = fields[name][index] + if isinstance(value, torch.Tensor): + assert restored.dtype == value.dtype and restored.shape == value.shape + assert torch.equal(restored.contiguous().view(torch.uint8), value.contiguous().view(torch.uint8)) + elif isinstance(value, np.ndarray): + assert restored.dtype == value.dtype and np.array_equal(restored, value) + else: + assert restored == value + assert fields["input_ids"][0] is fields["position_ids"][0] + fields["input_ids"][1][0] = 91 + baseline["field_data"]["input_ids"][1][0] = 91 + assert torch.equal(fields["input_ids"][0], baseline["field_data"]["input_ids"][0]) + assert actual["active_keys"] == original["active_keys"] + + +@pytest.mark.parametrize( + "change", ["truncate", "flip", "append", "header", "checksumless", "trailing_pickle", "hash", "truncate_relength"] +) +def test_corruption_refuses_before_replacing_live_storage(tmp_path, change): + path = tmp_path / "broken.pkl" + _write_storage_checkpoint(path, state()) + payload = bytearray(path.read_bytes()) + if change == "truncate": + del payload[-5:] + elif change == "truncate_relength": + del payload[-2:] + payload[8:16] = struct.pack(" zmq.Socke self._zmq_sockets[su_id] = sock return sock - def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, Any] | None: + def _query_storage_unit( + self, su_info: ZMQServerInfo, su_id: str, *, route_stats: bool = False + ) -> dict[str, Any] | None: """Send a synchronous GET_METRICS request to a single storage unit.""" + if type(route_stats) is not bool: + raise ValueError("route_stats must be a boolean") try: sock = self._get_or_create_socket(su_id, su_info) request_msg = ZMQMessage.create( request_type=ZMQRequestType.GET_METRICS, sender_id="metrics_collector", - body={}, + body={"route_stats": True} if route_stats else {}, ) sock.send_multipart(request_msg.serialize()) response_frames = sock.recv_multipart(copy=False) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index d1ff5f72..3be50559 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib import os import pickle import sys @@ -28,6 +29,7 @@ import ray import torch import zmq +import zstandard from tensordict import TensorDictBase from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads @@ -52,6 +54,8 @@ TQ_STORAGE_POLLER_TIMEOUT = int(os.environ.get("TQ_STORAGE_POLLER_TIMEOUT", 5)) # in seconds TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) +_CHECKPOINT_MAGIC = b"TQZSTD1\0" +_CHECKPOINT_HEADER_SIZE = len(_CHECKPOINT_MAGIC) + 8 + 32 class _CheckpointPickler(pickle.Pickler): @@ -81,6 +85,63 @@ def reducer_override(self, value): return torch._utils._rebuild_device_tensor_from_numpy, (array.astype(np.uint8), torch.int16, "cpu", False) +class _CheckpointHashStream: + """Hash compressed bytes during the existing write/read, without a second I/O pass.""" + + def __init__(self, stream): + self.stream, self.digest, self.size = stream, hashlib.sha256(), 0 + + def write(self, data): + count = self.stream.write(data) + self.digest.update(memoryview(data)[:count]) + self.size += count + return count + + def read(self, size=-1): + data = self.stream.read(size) + self.digest.update(data) + self.size += len(data) + return data + + def flush(self): + self.stream.flush() + + +def _write_storage_checkpoint(path, state): + with open(path, "wb") as raw: + raw.write(_CHECKPOINT_MAGIC + bytes(_CHECKPOINT_HEADER_SIZE - len(_CHECKPOINT_MAGIC))) + hashed = _CheckpointHashStream(raw) + compressor = zstandard.ZstdCompressor(level=1, threads=max(0, TQ_NUM_THREADS - 1), write_checksum=True) + with compressor.stream_writer(hashed, closefd=False) as stream: + _CheckpointPickler(stream, state).dump(state) + size = raw.tell() + raw.seek(len(_CHECKPOINT_MAGIC)) + raw.write(size.to_bytes(8, "little") + hashed.digest.digest()) + raw.flush() + os.fsync(raw.fileno()) + + +def _read_storage_checkpoint(path): + with open(path, "rb") as raw: + if raw.read(len(_CHECKPOINT_MAGIC)) != _CHECKPOINT_MAGIC: + raw.seek(0) + return pickle.load(raw) + size, expected_digest = int.from_bytes(raw.read(8), "little"), raw.read(32) + if size != os.fstat(raw.fileno()).st_size: + raise ValueError("Compressed checkpoint size mismatch") + if not zstandard.get_frame_parameters(raw.read(18)).has_checksum: + raise ValueError("Compressed checkpoint lacks a checksum") + raw.seek(_CHECKPOINT_HEADER_SIZE) + hashed = _CheckpointHashStream(raw) + with zstandard.ZstdDecompressor().stream_reader(hashed, closefd=False, read_across_frames=True) as stream: + state = pickle.load(stream) + if stream.read(1): + raise ValueError("Compressed checkpoint has trailing payload") + if hashed.size != size - _CHECKPOINT_HEADER_SIZE or hashed.digest.digest() != expected_digest: + raise ValueError("Compressed checkpoint digest mismatch") + return state + + class StorageUnitData: """Storage unit for managing 2D data structure (samples × fields). @@ -107,6 +168,45 @@ def active_key_count(self) -> int: """Number of active keys currently stored.""" return len(self._active_keys) + def route_storage_stats(self) -> dict: + stats: dict[str, Any] = dict( + row_count=0, + packed_rows=0, + unpacked_rows=0, + unknown_rows=0, + logical_view_bytes_known=0, + stored_view_bytes_known=0, + unique_backing_bytes_known=0, + ) + stats["logical_dtype_counts"], stats["stored_dtype_counts"] = {}, {} + stores = set() + for value in self.field_data.get("routed_experts", {}).values(): + stats["row_count"] += 1 + packed = isinstance(value, PackedRoutes) + tensor = value.data if packed else value + if ( + type(tensor) is not torch.Tensor + or tensor.device.type != "cpu" + or tensor.layout != torch.strided + or tensor.is_nested + ): + stats["unknown_rows"] += 1 + continue + storage = tensor.untyped_storage() + identity = (storage.data_ptr(), storage.nbytes()) + if identity not in stores: + stores.add(identity) + stats["unique_backing_bytes_known"] += storage.nbytes() + stats["packed_rows" if packed else "unpacked_rows"] += 1 + for key, dtype in ( + ("logical_dtype_counts", "torch.int16" if packed else str(tensor.dtype)), + ("stored_dtype_counts", str(tensor.dtype)), + ): + stats[key][dtype] = stats[key].get(dtype, 0) + 1 + stats["stored_view_bytes_known"] += tensor.nbytes + stats["logical_view_bytes_known"] += tensor.numel() * (2 if packed else tensor.element_size()) + return stats + def get_data(self, fields: list[str], global_indexes: list, *, packed_routes: bool = False) -> dict[str, list]: """Get data by global index keys. @@ -405,7 +505,7 @@ def _worker_routine(self) -> None: with monitor.measure(op_type="CLEAR_DATA"): response_msg = self._handle_clear(request_msg) elif operation == ZMQRequestType.GET_METRICS: # type: ignore[arg-type] - response_msg = self._handle_get_metrics() + response_msg = self._handle_get_metrics(route_stats=request_msg.body.get("route_stats", False)) elif operation == ZMQRequestType.SAVE_STORAGE_CHECKPOINT: # type: ignore[arg-type] response_msg = self._handle_save_checkpoint(request_msg) elif operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] @@ -595,13 +695,15 @@ def _handle_clear(self, data_parts: ZMQMessage) -> ZMQMessage: ) return response_msg - def _handle_get_metrics(self) -> ZMQMessage: + def _handle_get_metrics(self, *, route_stats: bool = False) -> ZMQMessage: """Handle GET_METRICS request by returning storage unit statistics. Returns: ZMQMessage containing storage unit ID, capacity, active keys, process RSS memory, and per-operation request stats. """ + if type(route_stats) is not bool: + raise ValueError("route_stats must be a boolean") try: process_rss = psutil.Process().memory_info().rss except Exception: @@ -613,6 +715,8 @@ def _handle_get_metrics(self) -> ZMQMessage: "active_keys": self.storage_data.active_key_count, "process_rss_bytes": process_rss, } + if route_stats: + metrics["route_storage"] = self.storage_data.route_storage_stats() # Include per-operation stats if Prometheus metrics are enabled if self._metrics is not None: @@ -701,10 +805,7 @@ def failure_details(error): raise ValueError("unknown checkpoint phase") def write(): - with open(path, "wb") as f: - _CheckpointPickler(f, state).dump(state) - f.flush() - os.fsync(f.fileno()) + _write_storage_checkpoint(path, state) prepared = ( phase == "prepare" @@ -781,8 +882,7 @@ def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: """ path = data_parts.body["path"] try: - with open(path, "rb") as f: - data = pickle.load(f) + data = _read_storage_checkpoint(path) if data["storage_unit_size"] != self.storage_unit_size: logger.warning( diff --git a/transfer_queue/version/version b/transfer_queue/version/version index 6b956112..5e54c775 100644 --- a/transfer_queue/version/version +++ b/transfer_queue/version/version @@ -1 +1 @@ -0.1.10+zyphra.2 +0.1.10+zyphra.3