diff --git a/scripts/performance_test/README_PERFTEST.md b/scripts/performance_test/README_PERFTEST.md index 16b4b6bf..2f925bb1 100644 --- a/scripts/performance_test/README_PERFTEST.md +++ b/scripts/performance_test/README_PERFTEST.md @@ -43,6 +43,8 @@ python perftest.py \ | `--worker_node_ip` | Worker node IP address (required for Yuanrong) | None | No | | `--output_csv` | Path to output CSV file | None | No | | `--use_complex_case` | Use complex test case with nested tensors and NonTensorStack fields | False | No | +| `--ssd_offload` | Enable SimpleStorage SSD offload; requires `--ssd_path` | False | No | +| `--ssd_path` | Existing local SSD directory used by `--ssd_offload` | None | No | ## Backend Configuration @@ -58,6 +60,10 @@ backend: num_data_storage_units: 16 ``` +`--ssd_offload` injects the SSD configuration and writes `--ssd_path` into the +in-memory configuration for that run. The path must already exist and be a +directory on every storage node. + ### Yuanrong Configuration ```yaml @@ -171,6 +177,49 @@ HEAD_NODE_IP=192.168.0.1 WORKER_NODE_IP=192.168.0.2 DEVICE=npu ./run_perf_test.s After running the tests, `draw_figure.py` reads all CSV files from `results/` and generates a grouped bar chart comparing total throughput (Gbps) across backends and data sizes. +## Running the SSD Offload Performance Test + +`run_ssd_offload_perf_test.sh` is the dedicated entry point for storage-tier +features such as SimpleStorage SSD offload. The script keeps backend and size +matrices as arrays for future extension; currently it runs only: + +- **Backend**: SimpleStorage +- **Workloads**: + - Sample1MiB (batch=512, fields=8, seq=262144, total=4 GiB) + - Sample2MiB (batch=512, fields=8, seq=524288, total=8 GiB) + - Sample4MiB (batch=512, fields=8, seq=1048576, total=16 GiB) + +The workload names describe the size of one float32 field sample. All three +sizes meet the SSD offload threshold, while the batch size and field count stay +fixed for a direct comparison. +For each backend and workload, the script runs SSD offload first and then runs +the same workload without offload as the host-memory baseline. + +```bash +HEAD_NODE_IP=192.168.0.1 \ +WORKER_NODE_IP=192.168.0.2 \ +SSD_OFFLOAD_PATH=/path/to/local/ssd \ +./run_ssd_offload_perf_test.sh +``` + +Configuration variables: + +| Variable | Description | Default | +|----------|-------------|---------| +| `HEAD_NODE_IP` | Head node IP address | `127.0.0.1` | +| `WORKER_NODE_IP` | Worker node IP address | `127.0.0.1` | +| `DEVICE` | Device type (`cpu`, `npu`, `gpu`) | `cpu` | +| `NUM_TEST_ITERATIONS` | Number of iterations | `4` | +| `USE_COMPLEX_CASE` | Use complex test data | `false` | +| `SSD_OFFLOAD_PATH` | Existing local SSD directory | Required | + +Results are written to: + +- SSD offload: `results/simplestorage_ssd_{sample-size}.csv` +- Host memory: `results/simplestorage_{sample-size}.csv` + +Here, `{sample-size}` is `sample1mib`, `sample2mib`, or `sample4mib`. + ## Examples ### SimpleStorage backend (simple case) @@ -185,6 +234,24 @@ python perftest.py --backend_config=perftest_config.yaml --backend=SimpleStorage --head_node_ip=192.168.0.1 --use_complex_case ``` +### SimpleStorage SSD offload + +Run the host-memory and SSD cases separately with identical workload options: + +```bash +# Host-memory baseline +python perftest.py --backend_config=perftest_config.yaml --backend=SimpleStorage \ + --head_node_ip=192.168.0.1 --global_batch_size=64 --field_num=8 \ + --seq_len=262144 --num_test_iterations=3 \ + --output_csv=results/simple_storage_memory.csv + +# SSD offload +python perftest.py --backend_config=perftest_config.yaml --backend=SimpleStorage \ + --head_node_ip=192.168.0.1 --global_batch_size=64 --field_num=8 \ + --seq_len=262144 --num_test_iterations=3 --ssd_offload --ssd_path=/path/to/local/ssd \ + --output_csv=results/simple_storage_ssd.csv +``` + ### Yuanrong backend (inter-node) ```bash python perftest.py --backend_config=perftest_config.yaml --backend=Yuanrong \ @@ -225,6 +292,7 @@ Throughput is shown in both Gb/s (gigabits per second) and GB/s (gigabytes per s |--------|-------------| | `backend` | Backend name | | `device` | Device type | +| `ssd_offload` | Whether SimpleStorage SSD offload was enabled | | `total_data_size_gb` | Data size in GB | | `put_time` | PUT duration (seconds) | | `get_time` | GET duration (seconds) | diff --git a/scripts/performance_test/perftest.py b/scripts/performance_test/perftest.py index 08459600..32d84ddd 100644 --- a/scripts/performance_test/perftest.py +++ b/scripts/performance_test/perftest.py @@ -19,6 +19,7 @@ import logging import os import time +from pathlib import Path from typing import Any import ray @@ -293,6 +294,8 @@ def __init__( worker_node_ip: str | None = None, output_csv: str | None = None, use_complex_case: bool = False, + ssd_offload: bool = False, + ssd_path: str | None = None, ): """Initialize the throughput tester. @@ -308,6 +311,8 @@ def __init__( worker_node_ip: Worker node IP address (required for Yuanrong) output_csv: Path to output CSV file (optional) use_complex_case: Whether to use complex test case (nested + nontensor fields) + ssd_offload: Enable SimpleStorage SSD offload + ssd_path: Existing local SSD directory used when SSD offload is enabled """ self.backend_config_path = backend_config_path self.backend_override = backend @@ -320,12 +325,18 @@ def __init__( self.worker_node_ip = worker_node_ip self.output_csv = output_csv self.use_complex_case = use_complex_case + self.ssd_offload_requested = ssd_offload + self.ssd_path_override = ssd_path # Prepare full config for tq.init() self.full_config = self._prepare_config() # Get backend from config self.backend = self.full_config["backend"]["storage_backend"] + simple_storage_config = self.full_config["backend"].get("SimpleStorage", {}) + ssd_config = simple_storage_config.get("ssd_offload", {}) + self.ssd_offload_enabled = self.backend == "SimpleStorage" and bool(ssd_config.get("enabled", False)) + self.ssd_path = str(ssd_config["path"]) if self.ssd_offload_enabled else None # GDR is configured via backend.MooncakeStore.use_gdr (no separate CLI flag). self.use_gdr = bool(self.full_config["backend"].get("MooncakeStore", {}).get("use_gdr", False)) @@ -376,6 +387,26 @@ def _prepare_config(self) -> dict[str, Any]: if config.backend.storage_backend == "SimpleStorage": config.backend.SimpleStorage.total_storage_size = total_storage_size + if self.ssd_offload_requested: + if config.backend.storage_backend != "SimpleStorage": + raise ValueError("--ssd_offload is only supported by the SimpleStorage backend") + if not self.ssd_path_override: + raise ValueError("--ssd_offload requires --ssd_path") + config.backend.SimpleStorage.ssd_offload = { + "enabled": True, + "path": self.ssd_path_override, + } + + if config.backend.storage_backend == "SimpleStorage": + ssd_config = config.backend.SimpleStorage.get("ssd_offload", None) + if ssd_config is not None and ssd_config.get("enabled", False): + ssd_path = ssd_config.get("path", None) + if not ssd_path or not Path(str(ssd_path)).is_dir(): + raise ValueError( + f"SSD offload directory does not exist or is not a directory: {ssd_path!r}. " + "Update --ssd_path to an existing local SSD directory." + ) + return OmegaConf.to_container(config, resolve=True) def _initialize_clients(self) -> None: @@ -490,6 +521,9 @@ def run_throughput_test(self, skip_dataset_create=False) -> dict[str, Any]: logger.info(f"Backend: {self.backend}") logger.info(f"Device: {self.device}") logger.info(f"GDR: {self.use_gdr}") + logger.info(f"SSD Offload: {self.ssd_offload_enabled}") + if self.ssd_offload_enabled: + logger.info(f"SSD Path: {self.ssd_path}") logger.info(f"Total Data Size: {self.total_data_size_gb:.6f} GB") logger.info(f"PUT Time: {put_time:.8f}s") logger.info(f"GET Time: {get_time:.8f}s") @@ -499,10 +533,11 @@ def run_throughput_test(self, skip_dataset_create=False) -> dict[str, Any]: logger.info("=" * 60) # Return results (only Gb/s for CSV, not GB/s) - return { + result = { "backend": self.backend, "device": self.device, "use_gdr": self.use_gdr, + "ssd_offload": self.ssd_offload_enabled, "total_data_size_gb": self.total_data_size_gb, "put_time": put_time, "get_time": get_time, @@ -510,6 +545,7 @@ def run_throughput_test(self, skip_dataset_create=False) -> dict[str, Any]: "get_gbit_per_sec": get_gbit_per_sec, "total_gbit_per_sec": total_gbit_per_sec, } + return result def close(self) -> None: """Close the transfer_queue clients.""" @@ -607,6 +643,18 @@ def main() -> None: default=False, help="Use complex test case with nested tensors and nontensor fields (default: False, simple case)", ) + parser.add_argument( + "--ssd_offload", + action="store_true", + default=False, + help=("Enable SimpleStorage SSD offload. Requires --ssd_path."), + ) + parser.add_argument( + "--ssd_path", + type=str, + default=None, + help="Existing local SSD directory used by --ssd_offload.", + ) args = parser.parse_args() @@ -623,6 +671,8 @@ def main() -> None: worker_node_ip=args.worker_node_ip, output_csv=args.output_csv, use_complex_case=args.use_complex_case, + ssd_offload=args.ssd_offload, + ssd_path=args.ssd_path, ) # Run test multiple times for consistent results using a for loop diff --git a/scripts/performance_test/run_ssd_offload_perf_test.sh b/scripts/performance_test/run_ssd_offload_perf_test.sh new file mode 100755 index 00000000..29607ab3 --- /dev/null +++ b/scripts/performance_test/run_ssd_offload_perf_test.sh @@ -0,0 +1,83 @@ +#!/bin/bash + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +RESULTS_DIR="${SCRIPT_DIR}/results" +PERFTEST_PY="${SCRIPT_DIR}/perftest.py" +CONFIG_YAML="${SCRIPT_DIR}/perftest_config.yaml" + +mkdir -p "${RESULTS_DIR}" + +# ========== User Configuration ========== +HEAD_NODE_IP="${HEAD_NODE_IP:-127.0.0.1}" +WORKER_NODE_IP="${WORKER_NODE_IP:-127.0.0.1}" +DEVICE="${DEVICE:-cpu}" +NUM_TEST_ITERATIONS="${NUM_TEST_ITERATIONS:-4}" +USE_COMPLEX_CASE="${USE_COMPLEX_CASE:-false}" +SSD_OFFLOAD_PATH="${SSD_OFFLOAD_PATH:-}" +# ======================================== + +if [[ -z "${SSD_OFFLOAD_PATH}" ]]; then + echo "SSD_OFFLOAD_PATH must point to an existing local SSD directory." >&2 + exit 1 +fi + +if [[ ! -d "${SSD_OFFLOAD_PATH}" ]]; then + echo "SSD_OFFLOAD_PATH does not exist or is not a directory: ${SSD_OFFLOAD_PATH}" >&2 + exit 1 +fi + +# Extension points for additional SSD-capable backends and workload sizes. +# Currently, SSD offload is a SimpleStorage feature. Workload names describe +# the per-field sample size that controls memory-versus-SSD routing. +BACKENDS=("SimpleStorage") +declare -a SETTINGS=( + # batch_size, field_num, seq_len, name + "512,8,262144,Sample1MiB" + "512,8,524288,Sample2MiB" + "512,8,1048576,Sample4MiB" +) + +COMPLEX_ARGS=() +if [[ "${USE_COMPLEX_CASE}" == "true" ]]; then + COMPLEX_ARGS=(--use_complex_case) +fi + +for backend in "${BACKENDS[@]}"; do + echo "==========================================" + echo "Testing SSD offload comparison: ${backend}" + echo "SSD path: ${SSD_OFFLOAD_PATH}" + echo "==========================================" + + for setting in "${SETTINGS[@]}"; do + IFS=',' read -r batch_size field_num seq_len name <<< "${setting}" + ssd_output_csv="${RESULTS_DIR}/${backend,,}_ssd_${name,,}.csv" + memory_output_csv="${RESULTS_DIR}/${backend,,}_${name,,}.csv" + + echo " SSD offload: ${name} (batch=${batch_size}, fields=${field_num}, seq=${seq_len})" + python "${PERFTEST_PY}" --backend_config="${CONFIG_YAML}" --backend="${backend}" \ + --device="${DEVICE}" \ + --global_batch_size="${batch_size}" --field_num="${field_num}" --seq_len="${seq_len}" \ + --num_test_iterations="${NUM_TEST_ITERATIONS}" \ + --head_node_ip="${HEAD_NODE_IP}" --worker_node_ip="${WORKER_NODE_IP}" \ + --output_csv="${ssd_output_csv}" --ssd_offload --ssd_path="${SSD_OFFLOAD_PATH}" \ + "${COMPLEX_ARGS[@]}" + + sleep 10 + + echo " Host memory: ${name} (batch=${batch_size}, fields=${field_num}, seq=${seq_len})" + python "${PERFTEST_PY}" --backend_config="${CONFIG_YAML}" --backend="${backend}" \ + --device="${DEVICE}" \ + --global_batch_size="${batch_size}" --field_num="${field_num}" --seq_len="${seq_len}" \ + --num_test_iterations="${NUM_TEST_ITERATIONS}" \ + --head_node_ip="${HEAD_NODE_IP}" --worker_node_ip="${WORKER_NODE_IP}" \ + --output_csv="${memory_output_csv}" \ + "${COMPLEX_ARGS[@]}" + + sleep 10 + done +done + +echo "" +echo "All SSD offload comparison tests completed!" diff --git a/tests/e2e/test_ssd_offload_e2e.py b/tests/e2e/test_ssd_offload_e2e.py new file mode 100644 index 00000000..d6b99fe1 --- /dev/null +++ b/tests/e2e/test_ssd_offload_e2e.py @@ -0,0 +1,127 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end coverage for SimpleStorage local SSD offload.""" + +import pytest +import ray +import torch +from omegaconf import OmegaConf + +import transfer_queue as tq + + +@pytest.fixture(scope="module") +def ssd_root(tmp_path_factory): + """Return an isolated SSD root shared by this E2E module.""" + return tmp_path_factory.mktemp("tq-ssd-offload") + + +@pytest.fixture(scope="module") +def tq_system(ssd_root): + """Initialize a public TransferQueue system with SSD offload enabled.""" + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + + config = OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 20, + "num_data_storage_units": 2, + "ssd_offload": { + "enabled": True, + "path": str(ssd_root), + }, + }, + }, + } + ) + tq.init(config) + yield + tq.close() + assert not list(ssd_root.rglob("*.bin")) + assert not any(path.is_dir() for path in ssd_root.iterdir()) + if ray.is_initialized(): + ray.shutdown() + + +@pytest.fixture +def controller(tq_system): + """Return the controller used to clean test partitions.""" + return ray.get_actor("TransferQueueController", namespace="transfer_queue") + + +@pytest.fixture(autouse=True) +def cleanup_partitions(controller): + """Remove every partition created by an E2E test.""" + yield + for partition_id in ray.get(controller.list_partitions.remote()): + ray.get(controller.clear_partition.remote(partition_id)) + + +def test_public_api_routes_per_sample_and_migrates_on_overwrite(tq_system, ssd_root): + """Public KV operations preserve values while samples migrate between tiers.""" + partition_id = "ssd-routing" + small = torch.arange(16, dtype=torch.float32) + large = torch.arange(262144, dtype=torch.float32) + + tq.kv_put(key="small", partition_id=partition_id, fields={"value": small}) + tq.kv_put(key="large", partition_id=partition_id, fields={"value": large}) + + torch.testing.assert_close( + tq.kv_batch_get(keys=["small"], partition_id=partition_id)["value"][0], + small, + ) + torch.testing.assert_close( + tq.kv_batch_get(keys=["large"], partition_id=partition_id)["value"][0], + large, + ) + assert len(list(ssd_root.rglob("*.bin"))) == 1 + + tq.kv_put(key="small", partition_id=partition_id, fields={"value": large}) + assert len(list(ssd_root.rglob("*.bin"))) == 2 + + tq.kv_put(key="large", partition_id=partition_id, fields={"value": small}) + assert len(list(ssd_root.rglob("*.bin"))) == 1 + torch.testing.assert_close( + tq.kv_batch_get(keys=["large"], partition_id=partition_id)["value"][0], + small, + ) + + tq.kv_clear(keys=["small"], partition_id=partition_id) + assert not list(ssd_root.rglob("*.bin")) + + +def test_checkpoint_round_trip_recreates_ssd_data(tq_system, ssd_root, tmp_path): + """Checkpoint restore recreates logical data without depending on old SSD files.""" + partition_id = "ssd-checkpoint" + key = "large" + value = torch.arange(262144, dtype=torch.float32) + checkpoint_dir = tmp_path / "checkpoint" + + tq.kv_put(key=key, partition_id=partition_id, fields={"value": value}) + assert list(ssd_root.rglob("*.bin")) + tq.save_checkpoint(checkpoint_dir) + + tq.kv_clear(keys=[key], partition_id=partition_id) + assert not list(ssd_root.rglob("*.bin")) + tq.load_checkpoint(checkpoint_dir) + + restored = tq.kv_batch_get(keys=[key], partition_id=partition_id)["value"][0] + torch.testing.assert_close(restored, value) + assert list(ssd_root.rglob("*.bin")) diff --git a/tests/e2e/test_ssd_offload_multinode_e2e.py b/tests/e2e/test_ssd_offload_multinode_e2e.py new file mode 100644 index 00000000..eae0cb90 --- /dev/null +++ b/tests/e2e/test_ssd_offload_multinode_e2e.py @@ -0,0 +1,104 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Opt-in multi-node E2E coverage for node-local SimpleStorage SSD offload.""" + +import os +from pathlib import Path + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy +from tensordict import TensorDict + +import transfer_queue as tq + +_SSD_ROOT = os.environ.get("TQ_SSD_E2E_ROOT") + + +@ray.remote +def _local_ssd_file_count(root: str) -> int: + """Count SSD sample files on the current Ray node.""" + return len(list(Path(root).rglob("*.bin"))) + + +def _counts_by_node(root: str) -> list[int]: + active_nodes = [node for node in ray.nodes() if node["Alive"]] + return ray.get( + [ + _local_ssd_file_count.options( + scheduling_strategy=NodeAffinitySchedulingStrategy( + node_id=node["NodeID"], + soft=False, + ) + ).remote(root) + for node in active_nodes + ] + ) + + +@pytest.mark.skipif( + _SSD_ROOT is None, + reason="set TQ_SSD_E2E_ROOT to a node-local SSD path", +) +def test_multinode_ssd_offload_uses_and_cleans_each_node(): + """Exercise SSD-backed public KV operations across two Ray nodes.""" + ray.init(ignore_reinit_error=True) + alive_nodes = [node for node in ray.nodes() if node["Alive"]] + assert len(alive_nodes) >= 2 + + config = OmegaConf.create( + { + "controller": {"polling_mode": True}, + "backend": { + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": 20, + "num_data_storage_units": 2, + "ssd_offload": { + "enabled": True, + "path": _SSD_ROOT, + }, + }, + }, + } + ) + keys = [f"sample-{index}" for index in range(4)] + values = torch.arange(4 * 262144, dtype=torch.float32).view(4, 262144) + partition_id = "ssd-multinode" + + try: + tq.init(config) + tq.kv_batch_put( + keys=keys, + partition_id=partition_id, + fields=TensorDict({"value": values}, batch_size=4), + ) + restored = tq.kv_batch_get(keys=keys, partition_id=partition_id)["value"] + for actual, expected in zip(restored.unbind(), values.unbind(), strict=True): + torch.testing.assert_close(actual, expected) + + counts = _counts_by_node(_SSD_ROOT) + assert sum(counts) == len(keys) + assert sum(count > 0 for count in counts) >= 2 + + tq.kv_clear(keys=keys, partition_id=partition_id) + assert sum(_counts_by_node(_SSD_ROOT)) == 0 + finally: + tq.close() + + assert sum(_counts_by_node(_SSD_ROOT)) == 0 diff --git a/tests/test_serial_utils_on_cpu.py b/tests/test_serial_utils_on_cpu.py index 29ad9c73..95ebf4f1 100644 --- a/tests/test_serial_utils_on_cpu.py +++ b/tests/test_serial_utils_on_cpu.py @@ -97,6 +97,36 @@ def test_zmq_msg_serialization(): ) +def test_zmq_deserialize_returns_storage_buffer_info(): + from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType + + dense = torch.arange(24, dtype=torch.int64).view(3, 8) + variable = [torch.ones(2, dtype=torch.float32), torch.ones(5, dtype=torch.float32)] + msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_DATA, + sender_id="test", + body={"global_indexes": [1, 2, 3], "data": {"dense": dense, "variable": variable}}, + ) + + decoded, storage_info = ZMQMessage.deserialize_with_storage_info(msg.serialize()) + dense_buffer = storage_info.get_buffer(decoded.body["data"]["dense"]) + assert dense_buffer is not None + assert dense_buffer.encoding == "tensor" + assert dense_buffer.dtype == "int64" + assert dense_buffer.shape == (3, 8) + assert memoryview(dense_buffer.buffer).nbytes == dense.numel() * dense.element_size() + + for decoded_sample, original_sample in zip( + decoded.body["data"]["variable"], + variable, + strict=True, + ): + sample_buffer = storage_info.get_buffer(decoded_sample) + assert sample_buffer is not None + assert sample_buffer.shape == tuple(original_sample.shape) + assert memoryview(sample_buffer.buffer).nbytes == (original_sample.numel() * original_sample.element_size()) + + @pytest.mark.parametrize( "make_view", [ diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index e36063ee..f40dd120 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -14,14 +14,20 @@ # limitations under the License. import time +from concurrent.futures import ThreadPoolExecutor +import numpy as np import pytest import ray import tensordict import torch import zmq -from transfer_queue.storage.simple_storage import SimpleStorageUnit +from transfer_queue.storage.simple_storage import ( + HybridStorageUnitData, + SimpleStorageUnit, + SSDFieldStore, +) from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket @@ -457,6 +463,189 @@ def test_storage_unit_data_capacity_uses_active_keys(): assert storage._active_keys == {0, 1, 3} +def test_hybrid_storage_routes_each_sample_by_payload_size_and_clears(tmp_path): + """Each large sample gets an SSD file while small samples remain in memory.""" + storage = HybridStorageUnitData( + storage_size=2, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + try: + storage.put_data( + { + "mixed": [b"x" * 100, b"y" * 10], + "small": [1, 2], + "complex": [{"payload": "x" * 100}, {"payload": "y" * 100}], + }, + [1, 2], + ) + + assert storage._locations == { + "mixed": {1: "ssd", 2: "mem"}, + "small": {1: "mem", 2: "mem"}, + "complex": {1: "ssd", 2: "ssd"}, + } + assert storage.get_data(["mixed", "small", "complex"], [1, 2]) == { + "small": [1, 2], + "complex": [{"payload": "x" * 100}, {"payload": "y" * 100}], + "mixed": [b"x" * 100, b"y" * 10], + } + sample_files = list(tmp_path.rglob("*.bin")) + assert len(sample_files) == 3 + assert len({path.name for path in sample_files}) == 3 + + storage.clear([1]) + assert storage.active_key_count == 1 + with pytest.raises(KeyError): + storage.get_data(["mixed"], [1]) + assert len(list(tmp_path.rglob("*.bin"))) == 1 + + with pytest.raises(ValueError, match="Storage capacity exceeded"): + storage.put_data({"mixed": [b"a" * 100, b"b" * 100]}, [3, 4]) + finally: + storage.close() + + assert not (tmp_path / "test-run").exists() + + +def test_hybrid_storage_rolls_back_failed_ssd_put(tmp_path, monkeypatch): + """A failed SSD append must not expose memory fields or routing state.""" + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + + def fail_write(_fd, _data): + raise OSError("injected SSD failure") + + monkeypatch.setattr(storage._ssd_store, "_write_many", fail_write) + try: + with pytest.raises(OSError, match="injected SSD failure"): + storage.put_data({"small": [1], "large": [b"x" * 100]}, [1]) + + assert storage.active_key_count == 0 + assert storage._locations == {} + assert storage._mem_store.field_data == {} + finally: + storage.close() + + +def test_failed_ssd_overwrite_preserves_old_value(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + old_value = b"o" * 100 + try: + storage.put_data({"value": [old_value]}, [1]) + old_path = next(tmp_path.rglob("*.bin")) + + def fail_write(_fd, _data): + raise OSError("injected overwrite failure") + + monkeypatch.setattr(storage._ssd_store, "_write_many", fail_write) + with pytest.raises(OSError, match="injected overwrite failure"): + storage.put_data({"value": [b"n" * 100]}, [1]) + + assert storage.get_data(["value"], [1])["value"] == [old_value] + assert old_path.exists() + assert list(tmp_path.rglob("*.bin")) == [old_path] + finally: + storage.close() + + +def test_hybrid_storage_default_threshold_is_one_mib(tmp_path): + storage = HybridStorageUnitData( + storage_size=10, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + try: + storage.put_data( + {"value": [b"a" * (1024 * 1024 - 1), b"b" * (1024 * 1024)]}, + [1, 2], + ) + assert storage._locations["value"] == {1: "mem", 2: "ssd"} + finally: + storage.close() + + +def test_hybrid_storage_round_trips_supported_codecs(tmp_path): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + values = torch.arange(64, dtype=torch.float32).view(2, 32) + variable = [ + torch.arange(20, dtype=torch.float32), + torch.arange(40, dtype=torch.float32), + ] + arrays = np.arange(64, dtype=np.float32).reshape(2, 32) + objects = [{"payload": "a" * 100}, {"payload": "b" * 100}] + try: + fields = {"tensor": values, "variable": variable, "array": arrays, "object": objects} + storage.put_data(fields, [1, 2]) + assert storage._locations == {field: {1: "ssd", 2: "ssd"} for field in fields} + assert {entry.codec for entry in storage._ssd_store._offset_index["tensor"].values()} == {"tensor"} + assert {entry.codec for entry in storage._ssd_store._offset_index["array"].values()} == {"numpy"} + assert {entry.codec for entry in storage._ssd_store._offset_index["object"].values()} == {"pickle"} + + result = storage.get_data(list(fields), [1, 2]) + torch.testing.assert_close(result["tensor"][0], values[0]) + torch.testing.assert_close(result["tensor"][1], values[1]) + torch.testing.assert_close(result["variable"][0], variable[0]) + torch.testing.assert_close(result["variable"][1], variable[1]) + np.testing.assert_array_equal(result["array"][0], arrays[0]) + np.testing.assert_array_equal(result["array"][1], arrays[1]) + assert result["object"] == objects + finally: + storage.close() + + +def test_ssd_storage_cleans_orphans_without_removing_active_units(tmp_path): + """Startup cleanup removes unlocked owned directories and preserves live units.""" + active = SSDFieldStore(str(tmp_path), "active-run", "active-unit") + orphan = SSDFieldStore(str(tmp_path), "orphan-run", "orphan-unit") + orphan_path = orphan._base_path + orphan._owner_lock.close() + orphan._owner_lock = None + try: + scanner = SSDFieldStore(str(tmp_path), "scanner-run", "scanner-unit") + try: + assert active._base_path.exists() + assert not orphan_path.exists() + finally: + scanner.close() + finally: + active.close() + orphan.close() + + +def test_ssd_storage_units_can_initialize_concurrently(tmp_path): + def create_store(rank): + return SSDFieldStore(str(tmp_path), "shared-run", f"unit-{rank}") + + with ThreadPoolExecutor(max_workers=8) as executor: + stores = list(executor.map(create_store, range(8))) + try: + assert len([path for path in (tmp_path / "shared-run").iterdir() if path.is_dir()]) == 8 + finally: + for store in stores: + store.close() + + def test_storage_unit_data_parser(storage_setup): """Test data_parser functionality in SimpleStorageUnit. diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index bd83a599..448e1250 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -35,6 +35,17 @@ backend: # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null + # SSD offload reduces the long-lived host-memory footprint by keeping large + # fields on local SSD. Small fields remain in memory. + # Each offloaded (field, global_index) sample owns one temporary file so + # CLEAR can reclaim its disk space immediately. + ssd_offload: + # Master switch. Set to true to enable SSD offload. + enabled: false + # Required when enabled. Its parent must exist on a local NVMe SSD; + # the final directory is created if needed. Example: /path/to/ssd/offload/ + path: null + # MooncakeStore: high-performance KV-based hierarchical storage # that supports RDMA transport between GPU and DRAM. MooncakeStore: diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 98a12954..2091a7b2 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,6 +45,10 @@ _TQ_STORAGE: Any = None _TQ_CONTROLLER: Any = None +# Storage worker and proxy joins may take up to 10 seconds; leave time for Ray +# dispatch and SSD cleanup without allowing close() to block indefinitely. +_SIMPLE_STORAGE_SHUTDOWN_TIMEOUT_S = 15 + def _maybe_create_tq_client(conf: DictConfig | None = None) -> TransferQueueClient: global _TQ_CLIENT @@ -239,8 +243,17 @@ def close(): for key, value in _TQ_STORAGE.items(): if key == "SimpleStorage": # only the process that do first-time init can clean the distributed storage - for storage in value.values(): - ray.kill(storage) + storage_handles = list(value.values()) + try: + ray.get( + [storage.shutdown.remote() for storage in storage_handles], + timeout=_SIMPLE_STORAGE_SHUTDOWN_TIMEOUT_S, + ) + except Exception as e: + logger.warning(f"Failed to gracefully shut down SimpleStorage units: {e}") + finally: + for storage in storage_handles: + ray.kill(storage) elif key == "MooncakeStore": check = subprocess.run(["pgrep", "-f", "mooncake_master"], stdout=subprocess.PIPE, text=True) if check.returncode == 0: diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 381341ec..6e7a3e42 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -15,6 +15,7 @@ import math from typing import Any +from uuid import uuid4 from omegaconf import DictConfig @@ -44,12 +45,17 @@ def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: math.ceil(total_storage_size / num_data_storage_units) if total_storage_size is not None else None ) + ssd_config = conf.backend.SimpleStorage.get("ssd_offload", None) + ssd_run_id = uuid4().hex if ssd_config is not None and ssd_config.get("enabled", False) else None + for storage_unit_rank in range(num_data_storage_units): storage_node = SimpleStorageUnit.options( # type: ignore[attr-defined] scheduling_strategy=scheduling_strategies[storage_unit_rank], name=f"TransferQueueStorageUnit#{storage_unit_rank}", ).remote( storage_unit_size=storage_unit_size, + ssd_config=ssd_config, + ssd_run_id=ssd_run_id, ) simple_storage_handles[f"TransferQueueStorageUnit#{storage_unit_rank}"] = storage_node logger.info( diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index f80f905e..6586f29b 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -13,22 +13,48 @@ # See the License for the specific language governing permissions and # limitations under the License. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import fcntl +import hashlib +import json import os import pickle +import re +import shutil import time import weakref +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path from threading import Event, Thread from typing import TYPE_CHECKING, Any from uuid import uuid4 +import numpy as np import psutil import ray +import torch import zmq 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.serial_utils import DecodedBuffer, DecodeStorageInfo from transfer_queue.utils.zmq_utils import ( ZMQMessage, ZMQRequestType, @@ -46,6 +72,35 @@ 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)) +TQ_SSD_READ_THREADS = int(os.environ.get("TQ_SSD_READ_THREADS", 32)) +TQ_SSD_WRITE_THREADS = int(os.environ.get("TQ_SSD_WRITE_THREADS", 8)) +DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES = 1024 * 1024 + +_SSD_OWNER_FORMAT = "transfer_queue_ssd_offload_v1" +_SSD_OWNER_FILE = "owner.json" +_SSD_OWNER_LOCK_FILE = "owner.lock" + + +@dataclass(frozen=True) +class SSDEncodedSample: + """One sample represented in a form that can be written directly to SSD.""" + + payload: memoryview + codec: str + dtype: str | None = None + shape: tuple[int, ...] | None = None + + +@dataclass(frozen=True) +class SSDIndexEntry: + """Location and reconstruction metadata for one SSD-backed sample.""" + + path: Path + length: int + codec: str + dtype: str | None = None + shape: tuple[int, ...] | None = None + # Marks a GET_ERROR reply as "the key is gone" so the caller can tell it apart from a real fault. KEY_NOT_FOUND_MARKER = "TQKeyNotFound" @@ -148,6 +203,758 @@ def clear(self, keys: list[int]) -> None: self._active_keys -= set(keys) +def _field_to_filename(field: str) -> str: + """Map an arbitrary field name to a filesystem-safe, collision-resistant stem.""" + safe = re.sub(r"[^\w.-]", "_", field).strip("._")[:64] or "field" + digest = hashlib.sha256(field.encode()).hexdigest()[:16] + return f"{safe}_{digest}" + + +def _validate_path_component(value: str, name: str) -> None: + """Reject values that could escape or ambiguously address an owned directory.""" + if not value or value in {".", ".."} or Path(value).name != value: + raise ValueError(f"{name} must be a non-empty filesystem path component") + + +def _cleanup_orphaned_ssd_units(ssd_root: Path) -> None: + """Remove SSD unit directories whose owning process no longer holds its lock.""" + if not ssd_root.is_dir(): + return + + try: + run_dirs = list(ssd_root.iterdir()) + except OSError as e: + logger.warning(f"Failed to scan SSD offload root {ssd_root}: {e}") + return + + for run_dir in run_dirs: + if not run_dir.is_dir() or run_dir.is_symlink(): + continue + try: + unit_dirs = list(run_dir.iterdir()) + except OSError: + continue + for unit_dir in unit_dirs: + if not unit_dir.is_dir() or unit_dir.is_symlink(): + continue + + owner_path = unit_dir / _SSD_OWNER_FILE + lock_path = unit_dir / _SSD_OWNER_LOCK_FILE + try: + owner = json.loads(owner_path.read_text()) + if owner.get("format") != _SSD_OWNER_FORMAT: + continue + lock_file = open(lock_path, "a+b") + except (OSError, ValueError, TypeError): + continue + + try: + try: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + continue + shutil.rmtree(unit_dir) + except OSError as e: + logger.warning(f"Failed to clean orphaned SSD offload directory {unit_dir}: {e}") + finally: + lock_file.close() + + try: + run_dir.rmdir() + except OSError: + pass + + +class SSDFieldStore: + """One-file-per-sample SSD storage with atomically published indexes.""" + + def __init__(self, ssd_path: str, run_id: str, unit_id: str) -> None: + _validate_path_component(run_id, "run_id") + _validate_path_component(unit_id, "unit_id") + self._ssd_root = Path(ssd_path).resolve() + self._run_id = run_id + self._unit_id = unit_id + self._offset_index: dict[str, dict[int, SSDIndexEntry]] = {} + self._active_keys: set[int] = set() + self._closed = False + self._owner_lock: Any = None + + if not self._ssd_root.parent.is_dir(): + raise ValueError(f"SSD offload parent directory does not exist: {self._ssd_root.parent}") + self._ssd_root.mkdir(exist_ok=True) + cleanup_lock_path = self._ssd_root / ".cleanup.lock" + with open(cleanup_lock_path, "a+b") as cleanup_lock: + fcntl.flock(cleanup_lock.fileno(), fcntl.LOCK_EX) + _cleanup_orphaned_ssd_units(self._ssd_root) + self._base_path = self._create_owned_directory() + self._read_pool = ThreadPoolExecutor( + max_workers=max(TQ_SSD_READ_THREADS, 1), + thread_name_prefix="tq-ssd-read", + ) + self._write_pool = ThreadPoolExecutor( + max_workers=max(TQ_SSD_WRITE_THREADS, 1), + thread_name_prefix="tq-ssd-write", + ) + + def _create_owned_directory(self) -> Path: + run_dir = self._ssd_root / self._run_id + run_dir.mkdir(parents=True, exist_ok=True) + base_path = run_dir / self._unit_id + temp_path = run_dir / f".tmp-{self._unit_id}-{uuid4().hex}" + temp_path.mkdir() + + lock_file = None + try: + lock_file = open(temp_path / _SSD_OWNER_LOCK_FILE, "a+b") + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + owner = { + "format": _SSD_OWNER_FORMAT, + "run_id": self._run_id, + "unit_id": self._unit_id, + "pid": os.getpid(), + } + (temp_path / _SSD_OWNER_FILE).write_text(json.dumps(owner)) + temp_path.rename(base_path) + self._owner_lock = lock_file + return base_path + except Exception: + if lock_file is not None: + lock_file.close() + shutil.rmtree(temp_path, ignore_errors=True) + raise + + def _sample_directory(self, field: str, global_index: int) -> Path: + digest = hashlib.sha256(f"{field}\0{global_index}".encode()).hexdigest() + directory = self._base_path / _field_to_filename(field) / digest[:2] + directory.mkdir(parents=True, exist_ok=True) + return directory + + @staticmethod + def _write_many(fd: int, payloads: list[memoryview]) -> None: + """Write multiple payloads without concatenating them in host memory.""" + views = [payload.cast("B") for payload in payloads if payload.nbytes] + configured_iov_max = os.sysconf("SC_IOV_MAX") if "SC_IOV_MAX" in os.sysconf_names else 1024 + iov_max = max(int(configured_iov_max), 1) + while views: + batch = views[:iov_max] + written = os.writev(fd, batch) + if written <= 0: + raise OSError("SSDFieldStore writev returned no progress") + + consumed = 0 + while consumed < len(batch) and written >= batch[consumed].nbytes: + written -= batch[consumed].nbytes + consumed += 1 + views = views[consumed:] + if written: + views[0] = views[0][written:] + + @property + def active_key_count(self) -> int: + """Return the number of global indexes with SSD-backed samples.""" + return len(self._active_keys) + + def _write_sample( + self, + field: str, + global_index: int, + sample: SSDEncodedSample, + ) -> SSDIndexEntry: + directory = self._sample_directory(field, global_index) + token = uuid4().hex + temp_path = directory / f".tmp-{token}" + final_path = directory / f"{token}.bin" + try: + fd = os.open(temp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + self._write_many(fd, [sample.payload]) + finally: + os.close(fd) + temp_path.rename(final_path) + except Exception: + for path in (temp_path, final_path): + try: + path.unlink(missing_ok=True) + except OSError: + pass + raise + return SSDIndexEntry( + path=final_path, + length=sample.payload.nbytes, + codec=sample.codec, + dtype=sample.dtype, + shape=sample.shape, + ) + + def prepare_encoded( + self, + encoded_fields: dict[str, dict[int, SSDEncodedSample]], + ) -> dict[str, dict[int, SSDIndexEntry]]: + """Write new sample files without publishing them to readers.""" + prepared: dict[str, dict[int, SSDIndexEntry]] = {} + futures = { + self._write_pool.submit( + self._write_sample, + field, + global_index, + sample, + ): (field, global_index) + for field, samples in encoded_fields.items() + for global_index, sample in samples.items() + } + first_error: Exception | None = None + for future, (field, global_index) in futures.items(): + try: + entry = future.result() + prepared.setdefault(field, {})[global_index] = entry + except Exception as e: + if first_error is None: + first_error = e + if first_error is not None: + self.discard_prepared(prepared) + raise first_error + return prepared + + @staticmethod + def discard_prepared(prepared: dict[str, dict[int, SSDIndexEntry]]) -> None: + """Delete prepared files after a PUT fails before publication.""" + for entries in prepared.values(): + for entry in entries.values(): + try: + entry.path.unlink(missing_ok=True) + except OSError: + pass + + def commit_prepared(self, prepared: dict[str, dict[int, SSDIndexEntry]]) -> None: + """Publish prepared files, then remove superseded sample files.""" + old_entries: list[SSDIndexEntry] = [] + for field, field_entries in prepared.items(): + index = self._offset_index.setdefault(field, {}) + for global_index, new_entry in field_entries.items(): + old_entry = index.get(global_index) + index[global_index] = new_entry + if old_entry is not None: + old_entries.append(old_entry) + self._active_keys.add(global_index) + + for entry in old_entries: + self._unlink_entry(entry) + + @staticmethod + def _unlink_entry(entry: SSDIndexEntry) -> None: + try: + entry.path.unlink(missing_ok=True) + except OSError as e: + logger.warning(f"Failed to delete superseded SSD sample {entry.path}: {e}") + + def remove(self, field: str, global_index: int) -> None: + """Remove one SSD-backed field sample if it exists.""" + index = self._offset_index.get(field) + if index is None: + return + entry = index.pop(global_index, None) + if entry is not None: + self._unlink_entry(entry) + if not index: + self._offset_index.pop(field, None) + if not any(global_index in entries for entries in self._offset_index.values()): + self._active_keys.discard(global_index) + + def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: + """Read per-sample values from their independently owned files.""" + result: dict[str, list] = {} + for field in fields: + if field not in self._offset_index: + raise ValueError( + f"SSDFieldStore get_data: field '{field}' not found. Available: {list(self._offset_index.keys())}" + ) + idx_map = self._offset_index[field] + entries = [] + for gidx in global_indexes: + if gidx not in idx_map: + raise KeyError(f"SSDFieldStore get_data: key {gidx} not found in field '{field}'") + entries.append((field, gidx, idx_map[gidx])) + result[field] = list(self._read_pool.map(self._read_entry, entries)) + return result + + @classmethod + def _read_entry(cls, item: tuple[str, int, SSDIndexEntry]) -> Any: + field, global_index, entry = item + raw = entry.path.read_bytes() + if len(raw) != entry.length: + raise OSError( + f"SSDFieldStore short read for field '{field}', key {global_index}: " + f"expected {entry.length} bytes, got {len(raw)}" + ) + return cls._decode_sample(raw, entry) + + @staticmethod + def _decode_sample(raw: bytes, entry: SSDIndexEntry) -> Any: + if entry.codec == "tensor": + if entry.dtype is None or entry.shape is None: + raise ValueError("Tensor SSD entry is missing dtype or shape") + dtype = getattr(torch, entry.dtype) + if not raw: + return torch.empty(entry.shape, dtype=dtype) + return torch.frombuffer(raw, dtype=dtype).view(entry.shape) + if entry.codec == "numpy": + if entry.dtype is None or entry.shape is None: + raise ValueError("NumPy SSD entry is missing dtype or shape") + if not raw: + return np.empty(entry.shape, dtype=np.dtype(entry.dtype)) + return np.frombuffer(raw, dtype=np.dtype(entry.dtype)).reshape(entry.shape) + if entry.codec == "bytes": + return raw + if entry.codec == "pickle": + return pickle.loads(raw) + raise ValueError(f"Unsupported SSD codec: {entry.codec}") + + def clear(self, keys: list) -> None: + """Remove logical entries and delete their sample files.""" + for field, idx_map in list(self._offset_index.items()): + for key in keys: + extent = idx_map.pop(key, None) + if extent is not None: + self._unlink_entry(extent) + if not idx_map: + self._offset_index.pop(field, None) + self._active_keys -= set(keys) + + def get_state(self) -> dict[str, dict]: + """Load all SSD data into memory; used by checkpoint serialisation.""" + state: dict[str, dict] = {} + for field in self._offset_index: + indexes = list(self._offset_index[field]) + values = self.get_data([field], indexes)[field] + state[field] = dict(zip(indexes, values, strict=True)) + return state + + def close(self) -> None: + """Stop I/O workers and delete the storage directory.""" + if self._closed: + return + self._closed = True + self._write_pool.shutdown(wait=True) + self._read_pool.shutdown(wait=True) + cleanup_lock_path = self._ssd_root / ".cleanup.lock" + with open(cleanup_lock_path, "a+b") as cleanup_lock: + fcntl.flock(cleanup_lock.fileno(), fcntl.LOCK_EX) + try: + shutil.rmtree(self._base_path, ignore_errors=True) + finally: + if self._owner_lock is not None: + self._owner_lock.close() + self._owner_lock = None + try: + self._base_path.parent.rmdir() + except OSError: + pass + + +class HybridStorageUnitData: + """Route fields to memory or SSD while preserving StorageUnitData semantics.""" + + def __init__( + self, + storage_size: int | None, + ssd_path: str, + run_id: str, + unit_id: str, + threshold_bytes: int = DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES, + ) -> None: + self._mem_store = StorageUnitData(storage_size=None) + self._ssd_store = SSDFieldStore(ssd_path, run_id, unit_id) + self._threshold = threshold_bytes + self._ssd_path = ssd_path + self._run_id = run_id + self._unit_id = unit_id + self._storage_size = storage_size + self._locations: dict[str, dict[int, str]] = {} + self._active_keys: set[int] = set() + + @property + def active_key_count(self) -> int: + """Return the number of active global indexes across both tiers.""" + return len(self._active_keys) + + @staticmethod + def _split_batched_buffer( + decoded_buffer: DecodedBuffer, + sample_count: int, + ) -> list[SSDEncodedSample] | None: + if ( + decoded_buffer.buffer is None + or decoded_buffer.dtype is None + or decoded_buffer.shape is None + or not decoded_buffer.shape + or decoded_buffer.shape[0] != sample_count + ): + return None + + payload = memoryview(decoded_buffer.buffer).cast("B") + if payload.nbytes % sample_count: + return None + sample_bytes = payload.nbytes // sample_count + codec = "tensor" if decoded_buffer.encoding == "tensor" else "numpy" + return [ + SSDEncodedSample( + payload=payload[position * sample_bytes : (position + 1) * sample_bytes], + codec=codec, + dtype=decoded_buffer.dtype, + shape=decoded_buffer.shape[1:], + ) + for position in range(sample_count) + ] + + @staticmethod + def _sample_from_decoded_buffer( + decoded_buffer: DecodedBuffer, + ) -> SSDEncodedSample | None: + if ( + decoded_buffer.encoding not in {"tensor", "numpy"} + or decoded_buffer.buffer is None + or decoded_buffer.dtype is None + or decoded_buffer.shape is None + ): + return None + return SSDEncodedSample( + payload=memoryview(decoded_buffer.buffer).cast("B"), + codec=decoded_buffer.encoding, + dtype=decoded_buffer.dtype, + shape=decoded_buffer.shape, + ) + + @staticmethod + def _sample_from_value(value: Any) -> SSDEncodedSample | None: + if isinstance(value, torch.Tensor): + if value.is_nested or value.is_sparse: + return None + try: + tensor = value.detach() + if tensor.device.type != "cpu": + tensor = tensor.cpu() + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + payload = memoryview(tensor.flatten().view(torch.uint8).numpy()).cast("B") + except (RuntimeError, TypeError, ValueError): + return None + return SSDEncodedSample( + payload=payload, + codec="tensor", + dtype=str(tensor.dtype).removeprefix("torch."), + shape=tuple(tensor.shape), + ) + if isinstance(value, np.ndarray) and not value.dtype.hasobject: + try: + array = value if value.flags["C_CONTIGUOUS"] else np.ascontiguousarray(value) + payload = memoryview(array.view(np.uint8).ravel()).cast("B") + except (TypeError, ValueError): + return None + return SSDEncodedSample( + payload=payload, + codec="numpy", + dtype=str(array.dtype), + shape=tuple(array.shape), + ) + if isinstance(value, bytes): + return SSDEncodedSample(payload=memoryview(value), codec="bytes") + try: + pickled_payload = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + return None + return SSDEncodedSample(payload=memoryview(pickled_payload), codec="pickle") + + def _build_encoded_samples( + self, + values: Any, + sample_count: int, + storage_info: DecodeStorageInfo | None, + ) -> list[SSDEncodedSample] | None: + decoded_buffer = storage_info.get_buffer(values) if storage_info is not None else None + if decoded_buffer is not None: + if decoded_buffer.encoding in {"tensor", "numpy"}: + samples = self._split_batched_buffer(decoded_buffer, sample_count) + if samples is not None: + return samples + elif decoded_buffer.encoding == "nested_tensor" and len(decoded_buffer.children) == sample_count: + nested_samples: list[SSDEncodedSample] = [] + for child in decoded_buffer.children: + sample = self._sample_from_decoded_buffer(child) + if sample is None: + break + nested_samples.append(sample) + else: + return nested_samples + + if isinstance(values, torch.Tensor): + if values.is_nested: + value_list = list(values.unbind()) + else: + encoded_batch = self._sample_from_value(values) + if encoded_batch is None: + return None + batch_buffer = DecodedBuffer( + encoding=encoded_batch.codec, + buffer=encoded_batch.payload, + dtype=encoded_batch.dtype, + shape=encoded_batch.shape, + ) + return self._split_batched_buffer(batch_buffer, sample_count) + elif isinstance(values, np.ndarray): + encoded_batch = self._sample_from_value(values) + if encoded_batch is None: + return None + batch_buffer = DecodedBuffer( + encoding=encoded_batch.codec, + buffer=encoded_batch.payload, + dtype=encoded_batch.dtype, + shape=encoded_batch.shape, + ) + return self._split_batched_buffer(batch_buffer, sample_count) + else: + try: + value_list = list(values) + except TypeError: + return None + + if len(value_list) != sample_count: + return None + + samples = [] + for value in value_list: + decoded_value = storage_info.get_buffer(value) if storage_info is not None else None + sample = ( + self._sample_from_decoded_buffer(decoded_value) + if decoded_value is not None + else self._sample_from_value(value) + ) + if sample is None: + return None + samples.append(sample) + return samples + + @staticmethod + def _values_as_samples(values: Any, sample_count: int) -> list[Any]: + if isinstance(values, torch.Tensor): + samples = list(values.unbind()) + elif isinstance(values, np.ndarray): + samples = list(values) + else: + samples = list(values) + if len(samples) != sample_count: + raise ValueError(f"Expected {sample_count} samples, got {len(samples)}") + return samples + + def _validate_put(self, field_data: dict[str, Any], global_indexes: list[int]) -> None: + for field, values in field_data.items(): + if len(values) != len(global_indexes): + raise ValueError( + f"HybridStorageUnitData put_data: field '{field}' values length {len(values)} " + f"!= global_indexes length {len(global_indexes)}, length mismatch" + ) + if self._storage_size is not None: + resulting_keys = self._active_keys | set(global_indexes) + if len(resulting_keys) > self._storage_size: + raise ValueError( + f"Storage capacity exceeded: {len(self._active_keys)} existing + " + f"{len(resulting_keys - self._active_keys)} new > {self._storage_size}" + ) + + def _snapshot_memory( + self, memory_fields: dict[str, Any], global_indexes: list[int] + ) -> tuple[set[int], dict[str, tuple[bool, dict[int, Any]]]]: + active_keys = set(self._mem_store._active_keys) + field_snapshot: dict[str, tuple[bool, dict[int, Any]]] = {} + for field in memory_fields: + existed = field in self._mem_store.field_data + current = self._mem_store.field_data.get(field, {}) + field_snapshot[field] = ( + existed, + {key: current[key] for key in global_indexes if key in current}, + ) + return active_keys, field_snapshot + + def _restore_memory( + self, + snapshot: tuple[set[int], dict[str, tuple[bool, dict[int, Any]]]], + global_indexes: list[int], + ) -> None: + active_keys, fields = snapshot + for field, (existed, old_values) in fields.items(): + current = self._mem_store.field_data.get(field) + if current is None: + continue + for key in global_indexes: + if key in old_values: + current[key] = old_values[key] + else: + current.pop(key, None) + if not existed: + self._mem_store.field_data.pop(field, None) + self._mem_store._active_keys = active_keys + + def put_data( + self, + field_data: dict[str, Any], + global_indexes: list, + storage_info: DecodeStorageInfo | None = None, + ) -> None: + """Store each sample in memory or SSD according to its encoded size.""" + self._validate_put(field_data, global_indexes) + if not global_indexes: + return + + memory_writes: dict[str, tuple[list[int], list[Any]]] = {} + encoded_ssd_fields: dict[str, dict[int, SSDEncodedSample]] = {} + pending_locations: dict[str, dict[int, str]] = {} + for field, values in field_data.items(): + logical_samples = self._values_as_samples(values, len(global_indexes)) + encoded_samples = self._build_encoded_samples( + values, + len(global_indexes), + storage_info, + ) + field_locations: dict[int, str] = {} + memory_indexes: list[int] = [] + memory_values: list[Any] = [] + ssd_values: dict[int, SSDEncodedSample] = {} + for position, global_index in enumerate(global_indexes): + encoded = encoded_samples[position] if encoded_samples is not None else None + destination = "ssd" if encoded is not None and encoded.payload.nbytes >= self._threshold else "mem" + field_locations[global_index] = destination + if destination == "ssd": + assert encoded is not None + ssd_values[global_index] = encoded + else: + memory_indexes.append(global_index) + memory_values.append(logical_samples[position]) + + pending_locations[field] = field_locations + if memory_indexes: + memory_writes[field] = (memory_indexes, memory_values) + if ssd_values: + encoded_ssd_fields[field] = ssd_values + + memory_snapshot = self._snapshot_memory(field_data, global_indexes) + prepared: dict[str, dict[int, SSDIndexEntry]] = {} + try: + if encoded_ssd_fields: + prepared = self._ssd_store.prepare_encoded(encoded_ssd_fields) + for field, (indexes, values) in memory_writes.items(): + self._mem_store.put_data({field: values}, indexes) + except Exception: + self._ssd_store.discard_prepared(prepared) + self._restore_memory(memory_snapshot, global_indexes) + raise + + self._ssd_store.commit_prepared(prepared) + for field, locations in pending_locations.items(): + mem_field = self._mem_store.field_data.get(field) + for global_index, destination in locations.items(): + if destination == "ssd": + if mem_field is not None: + mem_field.pop(global_index, None) + else: + self._ssd_store.remove(field, global_index) + self._locations.setdefault(field, {}).update(locations) + self._active_keys.update(global_indexes) + + def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: + """Read mixed memory- and SSD-backed samples in request order.""" + result: dict[str, list] = {} + for field in fields: + if field not in self._locations: + raise ValueError( + f"HybridStorageUnitData get_data: field '{field}' not found. Available: {list(self._locations)}" + ) + field_locations = self._locations[field] + ssd_indexes = [ + global_index for global_index in global_indexes if field_locations.get(global_index) == "ssd" + ] + ssd_values = {} + if ssd_indexes: + decoded_values = self._ssd_store.get_data([field], ssd_indexes)[field] + ssd_values = dict(zip(ssd_indexes, decoded_values, strict=True)) + values = [] + for global_index in global_indexes: + location = field_locations.get(global_index) + if location == "ssd": + values.append(ssd_values[global_index]) + elif location == "mem": + try: + values.append(self._mem_store.field_data[field][global_index]) + except KeyError as e: + raise KeyError( + f"HybridStorageUnitData get_data: key {global_index} not found in field '{field}'" + ) from e + else: + raise KeyError(f"HybridStorageUnitData get_data: key {global_index} not found in field '{field}'") + result[field] = values + return result + + def clear(self, keys: list) -> None: + """Clear the requested keys from both tiers and the location index.""" + self._mem_store.clear(keys) + self._ssd_store.clear(keys) + for field, locations in list(self._locations.items()): + for key in keys: + locations.pop(key, None) + if not locations: + self._locations.pop(field, None) + self._active_keys -= set(keys) + + def get_state(self) -> tuple[dict, set]: + """Return ``(field_data, active_keys)`` for checkpoint serialisation. + + SSD-backed fields are loaded into memory temporarily so the caller can + write a single pickle file containing all data. + """ + field_data: dict[str, dict] = {field: dict(fd) for field, fd in self._mem_store.field_data.items()} + for field, values in self._ssd_store.get_state().items(): + field_data.setdefault(field, {}).update(values) + return field_data, set(self._active_keys) + + def load_state(self, field_data: dict, active_keys: set) -> None: + """Reset both stores and restore from checkpoint data. + + Each field is re-routed through the normal ``put_data`` path so the + threshold-based routing decision is re-applied to the restored data. + """ + if self._storage_size is not None and len(active_keys) > self._storage_size: + raise ValueError( + f"Checkpoint contains {len(active_keys)} active keys, exceeding storage capacity {self._storage_size}" + ) + + replacement = HybridStorageUnitData( + storage_size=self._storage_size, + threshold_bytes=self._threshold, + ssd_path=self._ssd_path, + run_id=self._run_id, + unit_id=f"{self._unit_id}.restore.{uuid4().hex}", + ) + try: + for field, field_dict in field_data.items(): + if not field_dict: + continue + indexes = sorted(field_dict) + replacement.put_data({field: [field_dict[key] for key in indexes]}, indexes) + replacement._active_keys = set(active_keys) + except Exception: + replacement.close() + raise + + old_ssd_store = self._ssd_store + self._mem_store = replacement._mem_store + self._ssd_store = replacement._ssd_store + self._locations = replacement._locations + self._active_keys = replacement._active_keys + old_ssd_store.close() + + def close(self) -> None: + """Release SSD resources owned by this hybrid store.""" + self._ssd_store.close() + + @ray.remote(num_cpus=1) class SimpleStorageUnit: """A storage unit that provides distributed data storage functionality. @@ -166,17 +973,44 @@ class SimpleStorageUnit: zmq_server_info: ZMQ connection information for clients. """ - def __init__(self, storage_unit_size: int | None = None): + def __init__( + self, + storage_unit_size: int | None = None, + ssd_config=None, + ssd_run_id: str | None = None, + ): """Initialize a SimpleStorageUnit with the specified size. Args: storage_unit_size: Maximum number of elements that can be stored in this storage unit. If None, the storage unit has unlimited capacity. + ssd_config: Optional OmegaConf DictConfig for SSD offload (the ``ssd_offload`` block + from config.yaml). When ``ssd_config.enabled`` is True a + ``HybridStorageUnitData`` is used in place of the default in-memory + ``StorageUnitData``. + ssd_run_id: Internal run identifier used to isolate SSD directories. """ self.storage_unit_id = f"TQ_STORAGE_UNIT_{uuid4().hex[:8]}" self.storage_unit_size = storage_unit_size - - self.storage_data = StorageUnitData(self.storage_unit_size) + self.storage_data: StorageUnitData | HybridStorageUnitData + + if ssd_config is not None and ssd_config.get("enabled", False): + ssd_path = ssd_config.get("path") + if not ssd_path: + raise ValueError("SimpleStorage SSD offload requires backend.SimpleStorage.ssd_offload.path") + self.storage_data = HybridStorageUnitData( + storage_size=self.storage_unit_size, + ssd_path=str(ssd_path), + run_id=ssd_run_id or uuid4().hex, + unit_id=self.storage_unit_id, + ) + logger.info( + f"[{self.storage_unit_id}]: SSD offload enabled — " + f"path={ssd_path}, " + f"threshold={DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES} B/sample" + ) + else: + self.storage_data = StorageUnitData(self.storage_unit_size) # Internal communication address for proxy and workers self._inproc_addr = f"inproc://simple_storage_workers_{self.storage_unit_id}" @@ -204,8 +1038,15 @@ def __init__(self, storage_unit_size: int | None = None): self.proxy_thread, self.zmq_context, self.put_get_socket, + self.worker_socket, + self.storage_data, ) + def shutdown(self) -> None: + """Stop request processing and release this storage unit's resources.""" + if self._finalizer.alive: + self._finalizer() + def _init_zmq_socket(self) -> None: """ Initialize ZMQ socket connections between storage unit and controller/clients: @@ -305,7 +1146,7 @@ def _worker_routine(self) -> None: identity = messages[0] serialized_msg = messages[1:] - request_msg = ZMQMessage.deserialize(serialized_msg) + request_msg, storage_info = ZMQMessage.deserialize_with_storage_info(serialized_msg) operation = request_msg.request_type try: @@ -314,7 +1155,7 @@ def _worker_routine(self) -> None: # Process request if operation == ZMQRequestType.PUT_DATA: # type: ignore[arg-type] with monitor.measure(op_type="PUT_DATA"): - response_msg = self._handle_put(request_msg) + response_msg = self._handle_put(request_msg, storage_info) elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] with monitor.measure(op_type="GET_DATA"): response_msg = self._handle_get(request_msg) @@ -357,7 +1198,11 @@ def _worker_routine(self) -> None: poller.unregister(worker_socket) worker_socket.close(linger=0) - def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: + def _handle_put( + self, + data_parts: ZMQMessage, + storage_info: DecodeStorageInfo | None = None, + ) -> ZMQMessage: """ Handle put request, add or update data into storage unit. @@ -417,7 +1262,16 @@ 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) + if isinstance(self.storage_data, HybridStorageUnitData): + # A parser may replace or mutate decoded values, invalidating wire-buffer layout. + effective_storage_info = None if data_parser is not None else storage_info + self.storage_data.put_data( + field_data, + global_indexes, + storage_info=effective_storage_info, + ) + else: + self.storage_data.put_data(field_data, global_indexes) # After put operation finish, send a message to the client response_msg = ZMQMessage.create( @@ -577,11 +1431,16 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: """ path = data_parts.body["path"] try: + if isinstance(self.storage_data, HybridStorageUnitData): + field_data, active_keys = self.storage_data.get_state() + else: + field_data = self.storage_data.field_data + active_keys = self.storage_data._active_keys state = { "storage_unit_id": self.storage_unit_id, "storage_unit_size": self.storage_unit_size, - "field_data": self.storage_data.field_data, - "active_keys": self.storage_data._active_keys, + "field_data": field_data, + "active_keys": active_keys, } with open(path, "wb") as f: pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) @@ -624,15 +1483,18 @@ def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: f"checkpoint={data['storage_unit_size']}, current={self.storage_unit_size}" ) - if self.storage_data._active_keys: - logger.warning( - f"[{self.storage_unit_id}]: overwriting {len(self.storage_data._active_keys)} " - f"existing keys with checkpoint data from {path}" - ) - self.storage_data.field_data.clear() - self.storage_data._active_keys.clear() - self.storage_data.field_data = data["field_data"] - self.storage_data._active_keys = data["active_keys"] + if isinstance(self.storage_data, HybridStorageUnitData): + self.storage_data.load_state(data["field_data"], data["active_keys"]) + else: + if self.storage_data._active_keys: + logger.warning( + f"[{self.storage_unit_id}]: overwriting {len(self.storage_data._active_keys)} " + f"existing keys with checkpoint data from {path}" + ) + self.storage_data.field_data.clear() + self.storage_data._active_keys.clear() + self.storage_data.field_data = data["field_data"] + self.storage_data._active_keys = data["active_keys"] logger.info( f"[{self.storage_unit_id}]: loaded checkpoint from {path} — " @@ -691,26 +1553,30 @@ def _shutdown_resources( proxy_thread: Thread | None, zmq_context: zmq.Context | None, put_get_socket: zmq.Socket | None, + worker_socket: zmq.Socket | None, + storage_data=None, ) -> None: """Clean up resources on garbage collection.""" logger.info("Shutting down SimpleStorageUnit resources...") - # Signal all threads to stop shutdown_event.set() - - # Terminate put_get_socket - if put_get_socket: - put_get_socket.close(linger=0) - - # Terminate ZMQ context to unblock proxy and workers - if zmq_context: - zmq_context.term() - - # Wait for threads to finish (with timeout) - if worker_thread and worker_thread.is_alive(): - worker_thread.join(timeout=5) - if proxy_thread and proxy_thread.is_alive(): - proxy_thread.join(timeout=5) + try: + if put_get_socket: + put_get_socket.close(linger=0) + if worker_socket: + worker_socket.close(linger=0) + if zmq_context: + zmq_context.term() + finally: + if worker_thread and worker_thread.is_alive(): + worker_thread.join(timeout=5) + if proxy_thread and proxy_thread.is_alive(): + proxy_thread.join(timeout=5) + if storage_data is not None and hasattr(storage_data, "close"): + try: + storage_data.close() + except Exception as e: + logger.warning(f"Error closing storage data on shutdown: {e}") logger.info("SimpleStorageUnit resources shutdown complete.") diff --git a/transfer_queue/utils/serial_utils.py b/transfer_queue/utils/serial_utils.py index 2ba8e7d1..244c5ddf 100644 --- a/transfer_queue/utils/serial_utils.py +++ b/transfer_queue/utils/serial_utils.py @@ -23,6 +23,7 @@ from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor from contextvars import ContextVar +from dataclasses import dataclass, field from typing import Any, TypeAlias import cloudpickle @@ -57,6 +58,36 @@ # This enables the global _encoder/_decoder instances to be safely used across threads _encoder_aux_buffers: ContextVar[list[bytestr] | None] = ContextVar("encoder_aux_buffers", default=None) _decoder_aux_buffers: ContextVar[Sequence[bytestr] | None] = ContextVar("decoder_aux_buffers", default=None) +_decoder_storage_info: ContextVar["DecodeStorageInfo | None"] = ContextVar("decoder_storage_info", default=None) + + +@dataclass(frozen=True) +class DecodedBuffer: + """Wire-buffer metadata retained while decoding a storage payload.""" + + encoding: str + buffer: bytestr | None = None + dtype: str | None = None + shape: tuple[int, ...] | None = None + children: tuple["DecodedBuffer", ...] = () + + +@dataclass +class DecodeStorageInfo: + """Associate decoded objects with the wire buffers from which they were built.""" + + _buffers_by_object_id: dict[int, DecodedBuffer] = field(default_factory=dict) + _object_owners: list[Any] = field(default_factory=list) + + def record(self, value: Any, decoded_buffer: DecodedBuffer) -> None: + """Associate a decoded object with its reusable wire-buffer metadata.""" + self._buffers_by_object_id[id(value)] = decoded_buffer + # Keep the object alive until the PUT completes so object ids cannot be reused. + self._object_owners.append(value) + + def get_buffer(self, value: Any) -> DecodedBuffer | None: + """Return reusable wire-buffer metadata for a decoded object.""" + return self._buffers_by_object_id.get(id(value)) class MsgpackEncoder: @@ -296,6 +327,12 @@ def _reconstruct_tensordict(self, obj: dict) -> Any: # If tensordict not available, return as dict return obj + @staticmethod + def _record_storage_buffer(value: Any, decoded_buffer: DecodedBuffer) -> None: + storage_info = _decoder_storage_info.get() + if storage_info is not None: + storage_info.record(value, decoded_buffer) + def _decode_tensor(self, meta: tuple) -> torch.Tensor: """Decode tensor from (dtype, shape, buffer_idx) tuple.""" dtype, shape, idx = meta @@ -303,12 +340,23 @@ def _decode_tensor(self, meta: tuple) -> torch.Tensor: torch_dtype = getattr(torch, dtype) if not buffer: # Handle empty tensors - return torch.empty(shape, dtype=torch_dtype) - - # Create uint8 tensor from buffer, then view as original dtype and reshape - arr = torch.frombuffer(buffer, dtype=torch.uint8) - # Convert back to proper shape & type - return arr.view(torch_dtype).view(shape) + result = torch.empty(shape, dtype=torch_dtype) + else: + # Create uint8 tensor from buffer, then view as original dtype and reshape + arr = torch.frombuffer(buffer, dtype=torch.uint8) + # Convert back to proper shape & type + result = arr.view(torch_dtype).view(shape) + + self._record_storage_buffer( + result, + DecodedBuffer( + encoding="tensor", + buffer=buffer, + dtype=dtype, + shape=tuple(shape), + ), + ) + return result def _decode_nested_tensor(self, nested_meta: dict) -> torch.Tensor: """Decode nested tensor from serialized sub-tensors.""" @@ -320,9 +368,24 @@ def _decode_nested_tensor(self, nested_meta: dict) -> torch.Tensor: # Reconstruct nested tensor with appropriate layout if layout == "jagged": - return torch.nested.as_nested_tensor(sub_tensors, layout=torch.jagged) + result = torch.nested.as_nested_tensor(sub_tensors, layout=torch.jagged) else: # strided - return torch.nested.as_nested_tensor(sub_tensors, layout=torch.strided) + result = torch.nested.as_nested_tensor(sub_tensors, layout=torch.strided) + + children = tuple( + DecodedBuffer( + encoding="tensor", + buffer=self.aux_buffers[idx], + dtype=dtype, + shape=tuple(shape), + ) + for dtype, shape, idx in tensor_metas + ) + self._record_storage_buffer( + result, + DecodedBuffer(encoding="nested_tensor", children=children), + ) + return result def _decode_numpy(self, meta: tuple) -> np.ndarray: """Decode numpy array from (dtype_str, shape, buffer_idx) tuple.""" @@ -331,11 +394,22 @@ def _decode_numpy(self, meta: tuple) -> np.ndarray: np_dtype = np.dtype(dtype_str) if not buffer: # empty array - return np.empty(shape, dtype=np_dtype) - - # Reconstruct from raw bytes: uint8 view → reinterpret as original dtype - arr = np.frombuffer(buffer, dtype=np.uint8) - return arr.view(np_dtype).reshape(shape) + result = np.empty(shape, dtype=np_dtype) + else: + # Reconstruct from raw bytes: uint8 view → reinterpret as original dtype + arr = np.frombuffer(buffer, dtype=np.uint8) + result = arr.view(np_dtype).reshape(shape) + + self._record_storage_buffer( + result, + DecodedBuffer( + encoding="numpy", + buffer=buffer, + dtype=dtype_str, + shape=tuple(shape), + ), + ) + return result def ext_hook(self, code: int, data: memoryview) -> Any: """Custom decoding hook for types msgspec doesn't natively support. @@ -421,6 +495,16 @@ def decode(frames: list) -> Any: return _decoder.decode(frames) +def decode_with_storage_info(frames: list) -> tuple[Any, DecodeStorageInfo]: + """Decode frames and retain reusable tensor/array wire-buffer metadata.""" + storage_info = DecodeStorageInfo() + token = _decoder_storage_info.set(storage_info) + try: + return decode(frames), storage_info + finally: + _decoder_storage_info.reset(token) + + # Packed buffer layout: # [item_count: uint32 LE] # [N × (payload_offset: uint32 LE, payload_size: uint32 LE)] diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 2ca63d7f..96a675e5 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -28,7 +28,7 @@ from transfer_queue.utils.enum_utils import ExplicitEnum, Role from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.serial_utils import decode, encode +from transfer_queue.utils.serial_utils import DecodeStorageInfo, decode, decode_with_storage_info, encode logger = get_logger(__name__) @@ -231,7 +231,29 @@ def deserialize(cls, frames: list) -> "ZMQMessage": result = decode(frames) except Exception as e: raise ZMQMessageDecodeError(f"{type(e).__name__}: {e}; {describe_frames(frames)}") from e + return cls._from_decoded(result) + @classmethod + def deserialize_with_storage_info( + cls, + frames: list, + ) -> tuple["ZMQMessage", DecodeStorageInfo]: + """Deserialize and retain buffer metadata needed by storage backends.""" + if not frames: + raise ValueError("Empty frames received") + + if frame_nbytes(frames[0]) == 0: + raise ZMQMessageDecodeError(f"leading frame is empty; {describe_frames(frames)}") + + try: + result, storage_info = decode_with_storage_info(frames) + except Exception as e: + raise ZMQMessageDecodeError(f"{type(e).__name__}: {e}; {describe_frames(frames)}") from e + return cls._from_decoded(result), storage_info + + @classmethod + def _from_decoded(cls, result: dict[str, Any]) -> "ZMQMessage": + """Build a message from the decoded wire dictionary.""" return cls( request_type=ZMQRequestType(result["request_type"]), sender_id=result["sender_id"],