From 874934bf9c8d06d01127b7cd896e0687919cc04b Mon Sep 17 00:00:00 2001 From: pinjie Date: Wed, 23 Sep 2026 03:29:27 +0000 Subject: [PATCH 1/3] feat: add ssd offload for simplestorage Signed-off-by: pinjie --- docs/checkpoint.md | 4 +- docs/metrics.md | 4 + docs/ssd_offload.md | 222 +++++ .../README_PERFTEST_SSD_OFFLOAD.md | 157 ++++ .../performance_test_ssd_offload/perftest.py | 501 +++++++++++ .../perftest_config.yaml | 47 + .../run_perf_test.sh | 24 + scripts/put_benchmark.py | 9 +- tests/e2e/test_ssd_offload_e2e.py | 143 +++ tests/test_metrics.py | 34 + tests/test_simple_storage_unit.py | 280 +++++- transfer_queue/config.yaml | 12 + transfer_queue/interface.py | 34 +- transfer_queue/metrics.py | 57 ++ .../bootstrap/simple_storage_bootstrap.py | 53 +- transfer_queue/storage/simple_storage.py | 845 +++++++++++++++--- 16 files changed, 2271 insertions(+), 155 deletions(-) create mode 100644 docs/ssd_offload.md create mode 100644 scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md create mode 100644 scripts/performance_test_ssd_offload/perftest.py create mode 100644 scripts/performance_test_ssd_offload/perftest_config.yaml create mode 100644 scripts/performance_test_ssd_offload/run_perf_test.sh create mode 100644 tests/e2e/test_ssd_offload_e2e.py diff --git a/docs/checkpoint.md b/docs/checkpoint.md index a5b36803..8c0a2829 100644 --- a/docs/checkpoint.md +++ b/docs/checkpoint.md @@ -77,7 +77,9 @@ checkpoint_dir/ └── simple_storage/ ├── storage_unit_info.json # Position-to-ID manifest ├── su_0_.pkl # StorageUnit at position 0 + ├── su_0_.pkl.blobs/ # Its SSD values, when offload is enabled ├── su_1_.pkl # StorageUnit at position 1 + ├── su_1_.pkl.blobs/ # Its SSD values, when offload is enabled └── ... ``` @@ -123,7 +125,7 @@ tq.save_checkpoint / tq.load_checkpoint └── ... ``` -Both the controller and each storage unit write their data directly to disk from within their own processes. The ZMQ RPC carries only the target file path and an ACK, not the payload — this avoids routing large tensors through the Ray object store. +Both the controller and each storage unit write their data directly to disk from within their own processes. The ZMQ RPC carries only the target file path and an ACK, not the payload — this avoids routing large tensors through the Ray object store. With SSD offload enabled, each `su_*.pkl` stores in-memory values plus an `ssd_index` of plain metadata dictionaries; SSD-backed values are copied unchanged into the adjacent `.blobs` directory. Save and load therefore do not materialize the full SSD tier in host memory. ## Save Order and Consistency diff --git a/docs/metrics.md b/docs/metrics.md index 327f6079..4f59db79 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -137,6 +137,10 @@ Steps: | `tq_storage_active_keys_total` | Gauge | `storage_unit_id` | Active keys in storage | | `tq_storage_utilization_ratio` | Gauge | `storage_unit_id` | Utilization (active/capacity) | | `tq_storage_memory_rss_bytes` | Gauge | `storage_unit_id` | Storage process RSS memory | +| `tq_storage_ssd_offload_enabled` | Gauge | `storage_unit_id` | `1` when SSD offload is enabled, otherwise `0` | +| `tq_storage_ssd_active_values` | Gauge | `storage_unit_id` | Active field values stored on SSD | +| `tq_storage_ssd_active_bytes` | Gauge | `storage_unit_id` | Logical bytes held by active SSD-backed values | +| `tq_storage_ssd_fallback_values_total` | Counter | `storage_unit_id` | Values retained in memory because SSD encoding was unavailable | | `tq_storage_request_ops` | Gauge | `storage_unit_id`, `op_type` | Total requests processed by storage unit | | `tq_storage_request_latency_avg` | Gauge | `storage_unit_id`, `op_type` | Average request latency (seconds) | | `tq_storage_request_latency_p50` | Gauge | `storage_unit_id`, `op_type` | P50 request latency (seconds) | diff --git a/docs/ssd_offload.md b/docs/ssd_offload.md new file mode 100644 index 00000000..63111065 --- /dev/null +++ b/docs/ssd_offload.md @@ -0,0 +1,222 @@ +# SimpleStorage SSD Offload + +> Last updated: 09/20/2026 + +## Overview + +SimpleStorage can keep large field values on local SSD instead of keeping them +in host memory for their full lifetime. This reduces long-lived host-memory +use while preserving the existing TransferQueue APIs. + +SSD offload is available only with the `SimpleStorage` backend. Placement is +decided for each field value in each sample. Values whose encoded size is at +or above the configured threshold are stored on SSD; smaller values remain in +memory. + +SSD offload is temporary storage, not durable storage. `CLEAR` operations and +`tq.close()` remove offloaded data. Use a checkpoint when data must survive a +TransferQueue restart. + +## Quick Start + +Enable SSD offload under `backend.SimpleStorage`: + +```yaml +backend: + storage_backend: SimpleStorage + SimpleStorage: + ssd_offload: + enabled: true + path: /local/ssd/transfer-queue-job-a + threshold_bytes: 1048576 + glibc_mmap_threshold_bytes: 1048576 +``` + +## Configuration + +| Config key | Default | Description | +|------------|---------|-------------| +| `backend.SimpleStorage.ssd_offload.enabled` | `false` | Enables SSD offload for SimpleStorage. | +| `backend.SimpleStorage.ssd_offload.path` | `null` | Target directory for SSD offload. TransferQueue creates it, including missing parent directories, if needed. | +| `backend.SimpleStorage.ssd_offload.threshold_bytes` | `1048576` | Minimum encoded size, in bytes, for storing one field value on SSD. The value must be greater than zero. | +| `backend.SimpleStorage.ssd_offload.glibc_mmap_threshold_bytes` | `1048576` | glibc allocation size at which storage actors prefer independently releasable mappings. Set to `null` to leave the process environment unchanged. | + +`glibc_mmap_threshold_bytes` is independent of the SSD placement threshold. It +sets `MALLOC_MMAP_THRESHOLD_` only for storage actors and only when SSD offload +is enabled. Lower values can reduce RSS retained after transient allocations +are freed, but can increase mapping system calls and page faults. The setting +is effective only with glibc and disables its dynamic mmap-threshold adjustment. + +### Path ownership + +TransferQueue stores each run under +`/transfer_queue_ssd_offload//`. + +The configured `path` must be a directory, or a path that TransferQueue can +create. The `transfer_queue_ssd_offload` child must be a real directory, not a +symbolic link. + +Each storage unit owns only its `/` directory. A +successful `tq.close()` removes the current run's unit directories without +deleting sibling runs or unrelated files under `transfer_queue_ssd_offload`. + +## Runtime Directory Layout + +The current implementation stores active SSD data under this layout: + +```text +/ +└── transfer_queue_ssd_offload/ + └── / + └── / + └── <00..ff>/ + ├── .bin + └── .tmp- # Present only while a value is being written +``` + +One `run_id` is created for each TransferQueue initialization and shared by its +storage units. Each storage unit has its own directory and creates `00` through +`ff` subdirectories. + +Each active SSD-backed `(field, global_index)` value is stored in one `.bin` +file. The file contains only the encoded value. Its name does not contain the +field or global index; the storage unit keeps that mapping and the decode +metadata in memory. + +TransferQueue first writes a `.tmp-` file, then renames it to +`.bin` after the full value has been written. If a write fails, +TransferQueue attempts to remove its temporary file. + +This directory layout is an internal implementation detail and may change. +Applications should access offloaded values through TransferQueue APIs rather +than reading these files directly. Checkpoint files use the separate layout +described in [Checkpoint Integration](#checkpoint-integration). + +## Data Placement + +Placement is based on the encoded size of one field value for one sample. The +threshold is inclusive: + +```text +encoded size < threshold_bytes -> host memory +encoded size >= threshold_bytes -> local SSD +``` + +For a batched tensor, each item along the first dimension is considered one +sample. Different fields in the same sample can be placed in different tiers. +Overwriting a key can also move its values between memory and SSD. + +The following value types can be stored on SSD: + +| Value type | SSD representation | +|------------|--------------------| +| Dense PyTorch tensor | Raw tensor bytes, dtype, and shape | +| NumPy array without object dtype | Raw array bytes, dtype, and shape | +| `bytes` | Original bytes | +| Other Python value that `pickle` can serialize | Pickle payload | + +Nested or sparse PyTorch tensors, NumPy arrays with object dtype, and values +that cannot be encoded remain in memory. A non-CPU tensor is copied to CPU +before it is written to SSD. + +GET reads SSD-backed values and reconstructs their original types. SSD file +references are internal and are not returned to the application. + +The storage unit encodes values, performs SSD I/O, and keeps the file +references. The controller continues to manage metadata and scheduling; it +does not move SSD payloads. + +## Data Cleanup + +TransferQueue removes SSD files at these points: + +- `CLEAR` deletes the files for the cleared samples. +- A successful overwrite deletes files that are no longer referenced. +- A successful `tq.close()` stops each SimpleStorage unit and deletes that + unit's `/` directory. + +TransferQueue does not sweep sibling or stale run directories. A process that +exits without completing `tq.close()` can leave its run directory behind for +manual cleanup. The configured `path` and `transfer_queue_ssd_offload` child +directory remain in place. + +## Multi-Node Requirements + +All storage units receive the same `ssd_offload.path`. Each unit uses that path +on its own node, so SSD offload does not require a shared filesystem. + +Checkpoint directories must be shared across nodes; see [Checkpoint: Save and +Restore System State](checkpoint.md#multi-node-requirements). + +## Checkpoint Integration + +SSD offload works with `tq.save_checkpoint` and `tq.load_checkpoint`. Each +storage unit has a `.pkl` file. Its `.pkl.blobs/` directory exists only when +that unit has SSD-backed values: + +```text +simple_storage/ +├── su__.pkl # In-memory values and SSD metadata +└── su__.pkl.blobs/ # SSD-backed values, when present +``` + +Saving a checkpoint copies SSD files directly into the `.blobs` directory. It +does not read and decode the full SSD tier into host memory. Loading the +checkpoint copies the blobs into the current SSD offload directory. + +An SSD checkpoint must be loaded with SSD offload enabled. A checkpoint made +without SSD offload can be loaded with SSD offload enabled; its values are +placed using the current `threshold_bytes` setting. + +See [Checkpoint: Save and Restore System State](checkpoint.md) for checkpoint +consistency, replacement, and multi-node rules. + +## Metrics + +When metrics are enabled, the controller endpoint exposes these metrics for +each storage unit: + +| Metric | Description | +|--------|-------------| +| `tq_storage_ssd_offload_enabled` | `1` when SSD offload is enabled, otherwise `0` | +| `tq_storage_ssd_active_values` | Number of active field values stored on SSD | +| `tq_storage_ssd_active_bytes` | Logical encoded bytes held by active SSD-backed values | +| `tq_storage_ssd_fallback_values_total` | Cumulative number of field values retained in memory because SSD encoding was unavailable | + +`tq_storage_ssd_active_bytes` is logical payload size, not filesystem usage. +Filesystem blocks, directories, and temporary files are not included. See +[Prometheus Metrics & Grafana Dashboard](metrics.md) for metrics setup. + +`tq_storage_ssd_fallback_values_total` is monotonic for the lifetime of a +storage unit. A non-zero or increasing value means SSD offload is enabled but +some values cannot use any supported SSD encoding and remain in DRAM. Clearing +or overwriting those values does not decrease the counter. + +## Choosing a Threshold + +A lower threshold moves more values to SSD and can reduce long-lived host +memory use. It also creates more files, consumes more filesystem inodes, and +increases SSD I/O. A higher threshold keeps more values in memory and reduces +SSD work. + +Choose the threshold using the real field sizes and access pattern of the +workload. The benchmark under +[`scripts/performance_test_ssd_offload`](../scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md) +compares host-memory mode with SSD offload and reports throughput, storage-unit +RSS, and active SSD bytes. + +## Known Limitations + +### GET loads SSD-backed values into memory + +GET reads each requested SSD file and reconstructs the value in host memory. +SSD offload reduces long-lived memory use, but it does not remove temporary +memory use while values are read or encoded. + +### `data_parser` outputs should not share backing storage across samples + +If a `data_parser` returns differently sized tensor or NumPy views backed by +the same allocation, an in-memory view can keep the entire allocation resident +after another view is offloaded. When parsing URLs or file paths, return an +independently owned value for each sample, or copy retained views before +returning them. diff --git a/scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md b/scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md new file mode 100644 index 00000000..ae5b6bd7 --- /dev/null +++ b/scripts/performance_test_ssd_offload/README_PERFTEST_SSD_OFFLOAD.md @@ -0,0 +1,157 @@ +# TransferQueue SimpleStorage SSD Offload Performance Benchmark + +This benchmark is separate from the general backend throughput benchmarks under +`scripts/performance_test`. It runs the same mixed workload with SimpleStorage +host-memory storage and SSD offload, and records throughput, SSD byte accounting, +and storage-unit RSS. + +## Prerequisites + +1. Start a Ray cluster with exactly two live nodes. Each node must advertise + the `node:` resource used by the benchmark for actor placement: + + ```bash + # On the head node + ray start --head --resources='{"node:10.0.0.1":1}' + + # On the worker node + ray start --address=10.0.0.1:6379 --resources='{"node:10.0.0.2":1}' + ``` + +2. Create the same node-local SSD path on both nodes. The path must also be + visible to the benchmark driver and storage actors. When they run in + containers, bind-mount the path at the same absolute location. +3. Run the benchmark on the Ray head node in an environment containing the + TransferQueue runtime dependencies. A containerized driver that joins a + node-local Ray runtime must also bind-mount that runtime's `/tmp` directory + so it can reach the Ray socket. + +## Usage + +Use the wrapper to run the default 5 GiB mixed workload: + +```bash +HEAD_NODE_IP=10.0.0.1 \ +WORKER_NODE_IP=10.0.0.2 \ +SSD_OFFLOAD_PATH=/path/to/local/ssd \ +RESULTS_DIR=/path/to/results \ +bash scripts/performance_test_ssd_offload/run_perf_test.sh +``` + +### Wrapper configuration + +| Variable | Description | Default | Required | +|----------|-------------|---------|----------| +| `HEAD_NODE_IP` | Ray head-node IP | - | Yes | +| `WORKER_NODE_IP` | Ray worker-node IP | - | Yes | +| `SSD_OFFLOAD_PATH` | Node-local SSD path on both nodes | - | Yes | +| `RESULTS_DIR` | Output directory | `scripts/performance_test_ssd_offload/results` | No | +| `NUM_TEST_ITERATIONS` | Iterations per storage mode | `4` | No | + +## Arguments + +Run `perftest.py` directly to change the workload: + +```bash +python scripts/performance_test_ssd_offload/perftest.py \ + --backend_config=scripts/performance_test_ssd_offload/perftest_config.yaml \ + --head_node_ip=10.0.0.1 \ + --worker_node_ip=10.0.0.2 \ + --ssd_path=/path/to/local/ssd \ + --output_dir=/path/to/results \ + --global_batch_size=512 \ + --small_fields=4 \ + --small_sample_bytes=524288 \ + --large_fields=4 \ + --large_sample_bytes=2097152 \ + --num_test_iterations=4 \ + --warmup_iterations=1 +``` + +| Argument | Description | Default | Required | +|----------|-------------|---------|----------| +| `--backend_config` | Path to the benchmark configuration | - | Yes | +| `--head_node_ip` | Ray head-node IP used for writer placement | - | Yes | +| `--worker_node_ip` | Ray worker-node IP used for reader placement | - | Yes | +| `--ssd_path` | Node-local SSD path | - | Yes | +| `--output_dir` | Output directory | `scripts/performance_test_ssd_offload/results` | No | +| `--num_test_iterations` | Iterations per storage mode | `4` | No | +| `--warmup_iterations` | Leading iterations excluded from the summary | `1` | No | +| `--global_batch_size` | Samples per iteration | `512` | No | +| `--small_fields` | Tensor fields below the offload threshold | `4` | No | +| `--large_fields` | Tensor fields above the offload threshold | `4` | No | +| `--small_sample_bytes` | Bytes per sample in each small field | `524288` | No | +| `--large_sample_bytes` | Bytes per sample in each large field | `2097152` | No | + +Sample sizes are bytes per tensor field per sample. Small samples must be +strictly below the configured `ssd_offload.threshold_bytes`, and large samples +must be strictly above it. The SSD threshold and number of storage units come +from `perftest_config.yaml`. The number of iterations must exceed the number of +warm-up iterations. + +## Benchmark Scenario + +The default batch contains eight float32 tensor fields whose per-sample sizes +straddle the 1 MiB offload threshold: + +| Fields | Bytes per sample | Placement | +|--------|-----------------:|-----------| +| 4 small fields | 512 KiB | Host memory | +| 4 large fields | 2 MiB | Local SSD | + +With a batch size of 512, the logical payload is 5 GiB: 1 GiB remains in host +memory and 4 GiB is eligible for SSD offload. The benchmark first runs this +workload in pure host-memory mode and then repeats it with SSD offload. + +## Benchmark Flow + +Each iteration performs a PUT -> LIST -> GET -> CLEAR cycle through the +TransferQueue KV API. Storage-unit Prometheus metrics are sampled before PUT, +after PUT, after GET, and after CLEAR. RSS is accepted only after the configured +number of samples agree within the stability tolerance. + +After excluding warm-up iterations, the benchmark reports the theoretical and +actual SSD-active bytes, verifies disk reclamation after CLEAR, and summarizes +SSD-mode RSS before PUT and after PUT, GET, and CLEAR. The post-CLEAR RSS delta +and its growth across iterations show whether memory is progressively retained. + +## Distributed Smoke Test + +Use a smaller workload to verify node placement, mixed-tier routing, metrics, +and cleanup: + +```bash +python scripts/performance_test_ssd_offload/perftest.py \ + --backend_config=scripts/performance_test_ssd_offload/perftest_config.yaml \ + --head_node_ip="$HEAD_NODE_IP" \ + --worker_node_ip="$WORKER_NODE_IP" \ + --ssd_path="$SSD_OFFLOAD_PATH" \ + --output_dir=/tmp/tq-ssd-smoke \ + --global_batch_size=8 \ + --num_test_iterations=2 \ + --warmup_iterations=1 +``` + +This workload keeps 16 MiB in memory and offloads 64 MiB for functional +validation. Use the default 5 GiB workload when collecting post-CLEAR RSS +trends. + +## Output Format + +The output directory contains: + +- `mixed_memory.csv`: memory-mode throughput and diagnostic RSS samples; +- `mixed_ssd.csv`: SSD-mode throughput, RSS samples, and active SSD bytes; +- `summary.json`: warm-up-filtered SSD-byte accounting and post-CLEAR RSS + retention measurements. + +Important `summary.json` fields are: + +| Field | Description | +|-------|-------------| +| `theoretical_offload_bytes` | Logical payload expected to be stored on SSD | +| `ssd_median_active_bytes_after_put` | Median bytes reported as active on SSD after PUT | +| `ssd_max_active_bytes_after_clear` | Maximum active SSD bytes after CLEAR; expected to be zero | +| `ssd_median_rss_retained_after_clear_bytes` | Median SSD-mode RSS increase remaining after CLEAR | +| `ssd_max_rss_retained_after_clear_bytes` | Maximum SSD-mode RSS increase remaining after CLEAR | +| `ssd_clear_rss_growth_bytes` | Last post-CLEAR RSS minus the first post-CLEAR RSS | diff --git a/scripts/performance_test_ssd_offload/perftest.py b/scripts/performance_test_ssd_offload/perftest.py new file mode 100644 index 00000000..81e39bdc --- /dev/null +++ b/scripts/performance_test_ssd_offload/perftest.py @@ -0,0 +1,501 @@ +#!/usr/bin/env python3 +# 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. + +# 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. + +"""Benchmark mixed in-memory and SSD-backed SimpleStorage samples.""" + +import argparse +import csv +import json +import logging +import statistics +import time +import urllib.request +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import ray +import torch +from omegaconf import DictConfig, OmegaConf +from prometheus_client.parser import text_string_to_metric_families +from tensordict import TensorDict + +import transfer_queue as tq + +LOGGER = logging.getLogger("ssd_offload_benchmark") +DEFAULT_SMALL_SAMPLE_BYTES = 512 * 1024 +DEFAULT_LARGE_SAMPLE_BYTES = 2 * 1024 * 1024 +BYTES_PER_FLOAT32 = 4 +RAY_ADDRESS = "auto" +METRICS_TIMEOUT_SECONDS = 180 +RSS_SAMPLE_INTERVAL_SECONDS = 11 +RSS_STABLE_SAMPLES = 3 +RSS_STABILITY_TOLERANCE_BYTES = 32 * 2**20 +MODE_GAP_SECONDS = 10 + + +@dataclass(frozen=True) +class Workload: + batch_size: int + small_fields: int + large_fields: int + small_sample_bytes: int + large_sample_bytes: int + offload_threshold_bytes: int + + def validate(self) -> None: + sizes = { + "batch_size": self.batch_size, + "small_sample_bytes": self.small_sample_bytes, + "large_sample_bytes": self.large_sample_bytes, + "offload_threshold_bytes": self.offload_threshold_bytes, + } + if any(value <= 0 for value in sizes.values()): + raise ValueError(f"Workload sizes must be positive: {sizes}") + if self.small_fields < 0 or self.large_fields < 0: + raise ValueError("Field counts must be non-negative") + if self.small_fields + self.large_fields == 0: + raise ValueError("At least one tensor field is required") + if self.small_sample_bytes % BYTES_PER_FLOAT32 or self.large_sample_bytes % BYTES_PER_FLOAT32: + raise ValueError("Sample sizes must be divisible by four for float32 tensors") + if self.small_sample_bytes >= self.offload_threshold_bytes: + raise ValueError("small_sample_bytes must be below the SSD offload threshold") + if self.large_sample_bytes <= self.offload_threshold_bytes: + raise ValueError("large_sample_bytes must be above the SSD offload threshold") + + @property + def inline_bytes(self) -> int: + return self.batch_size * self.small_fields * self.small_sample_bytes + + @property + def offloaded_bytes(self) -> int: + return self.batch_size * self.large_fields * self.large_sample_bytes + + @property + def total_bytes(self) -> int: + return self.inline_bytes + self.offloaded_bytes + + +@dataclass(frozen=True) +class StorageMetrics: + active_keys: int + rss_bytes: int + ssd_active_bytes: int + + +def create_mixed_tensors(workload: Workload) -> TensorDict: + """Create fields whose individual samples straddle the offload threshold.""" + workload.validate() + torch.manual_seed(0) + fields = { + f"small_{index}": torch.randn( + workload.batch_size, + workload.small_sample_bytes // BYTES_PER_FLOAT32, + dtype=torch.float32, + ) + for index in range(workload.small_fields) + } + fields.update( + { + f"large_{index}": torch.randn( + workload.batch_size, + workload.large_sample_bytes // BYTES_PER_FLOAT32, + dtype=torch.float32, + ) + for index in range(workload.large_fields) + } + ) + return TensorDict(fields, batch_size=[workload.batch_size]) + + +def read_storage_metrics(endpoint: str, expected_storage_units: int) -> StorageMetrics: + """Read one aggregate snapshot from the controller's Prometheus endpoint.""" + with urllib.request.urlopen(f"http://{endpoint}/metrics", timeout=5) as response: + payload = response.read().decode("utf-8") + + metric_names = { + "tq_storage_active_keys_total": "active_keys", + "tq_storage_memory_rss_bytes": "rss_bytes", + "tq_storage_ssd_active_bytes": "ssd_active_bytes", + } + units: dict[str, dict[str, int]] = {} + for family in text_string_to_metric_families(payload): + for sample in family.samples: + field = metric_names.get(sample.name) + storage_unit_id = sample.labels.get("storage_unit_id") + if field is not None and storage_unit_id is not None: + units.setdefault(storage_unit_id, {})[field] = int(sample.value) + + complete = len(units) == expected_storage_units and all( + set(values) >= {"active_keys", "rss_bytes", "ssd_active_bytes"} for values in units.values() + ) + if not complete: + fields = sorted({tuple(sorted(values)) for values in units.values()}) + raise RuntimeError(f"Incomplete storage metrics: units={len(units)}/{expected_storage_units}, fields={fields}") + return StorageMetrics( + active_keys=sum(values["active_keys"] for values in units.values()), + rss_bytes=sum(values["rss_bytes"] for values in units.values()), + ssd_active_bytes=sum(values["ssd_active_bytes"] for values in units.values()), + ) + + +def wait_for_storage_state( + endpoint: str, + expected_storage_units: int, + expected_active_keys: int, + timeout_seconds: float, +) -> StorageMetrics: + """Wait until a fresh metrics collection reports the requested key count.""" + deadline = time.monotonic() + timeout_seconds + last_metrics = None + last_error = None + while time.monotonic() < deadline: + try: + last_metrics = read_storage_metrics(endpoint, expected_storage_units) + last_error = None + if last_metrics.active_keys == expected_active_keys: + return last_metrics + except (OSError, RuntimeError) as error: + last_error = str(error) + time.sleep(0.5) + raise TimeoutError( + f"Storage metrics did not reach active_keys={expected_active_keys}: " + f"last_metrics={last_metrics}, last_error={last_error!r}" + ) + + +def wait_for_stable_rss( + endpoint: str, + expected_storage_units: int, + expected_active_keys: int, + timeout_seconds: float, + sample_interval_seconds: float, + stable_samples: int, + tolerance_bytes: int, +) -> StorageMetrics: + """Return a median RSS after consecutive full metrics collection periods agree.""" + deadline = time.monotonic() + timeout_seconds + samples: list[StorageMetrics] = [] + while time.monotonic() < deadline: + remaining = max(0.5, deadline - time.monotonic()) + sample = wait_for_storage_state( + endpoint, + expected_storage_units, + expected_active_keys, + remaining, + ) + samples.append(sample) + window = samples[-stable_samples:] + if len(window) == stable_samples: + rss_values = [item.rss_bytes for item in window] + if max(rss_values) - min(rss_values) <= tolerance_bytes: + return StorageMetrics( + active_keys=expected_active_keys, + rss_bytes=int(statistics.median(rss_values)), + ssd_active_bytes=int(statistics.median(item.ssd_active_bytes for item in window)), + ) + time.sleep(sample_interval_seconds) + rss_values = [sample.rss_bytes for sample in samples[-stable_samples:]] + raise TimeoutError( + f"Storage RSS did not stabilize within {timeout_seconds}s: " + f"last_rss_values={rss_values}, tolerance_bytes={tolerance_bytes}" + ) + + +@ray.remote +class BenchmarkClient: + """Own one public TransferQueue client and, for the writer, the test payload.""" + + def __init__(self, config: dict[str, Any]): + self._config = config + self._data: TensorDict | None = None + self._keys: list[str] = [] + + def initialize(self) -> None: + tq.init(OmegaConf.create(self._config)) + + def create_data(self, workload: Workload) -> None: + self._data = create_mixed_tensors(workload) + self._keys = [f"mixed_{index}" for index in range(workload.batch_size)] + + def put(self, partition_id: str) -> None: + if self._data is None: + raise RuntimeError("Benchmark data has not been created") + tq.kv_batch_put(keys=self._keys, partition_id=partition_id, fields=self._data) + + def list_keys(self, partition_id: str) -> list[str]: + partitions = tq.kv_list(partition_id=partition_id) + return list(partitions.get(partition_id, {})) + + def get(self, partition_id: str, keys: list[str]) -> None: + tq.kv_batch_get(keys=keys, partition_id=partition_id) + + def clear(self, partition_id: str, keys: list[str]) -> None: + tq.kv_clear(keys=keys, partition_id=partition_id) + + def metrics_endpoint(self) -> str | None: + return tq.get_metrics_endpoint() + + def close(self) -> None: + tq.close() + + +def build_config( + base_config: DictConfig, + workload: Workload, + ssd_enabled: bool, + ssd_path: str, +) -> dict[str, Any]: + """Specialize the benchmark config for one SimpleStorage mode.""" + config = OmegaConf.create(OmegaConf.to_container(base_config, resolve=True)) + config.metrics.enabled = True + config.metrics.port = 0 + config.backend.storage_backend = "SimpleStorage" + config.backend.SimpleStorage.total_storage_size = workload.batch_size + config.backend.SimpleStorage.required_node_resource = None + config.backend.SimpleStorage.ssd_offload.enabled = ssd_enabled + config.backend.SimpleStorage.ssd_offload.path = ssd_path if ssd_enabled else None + return OmegaConf.to_container(config, resolve=True) + + +def validate_cluster(head_node_ip: str, worker_node_ip: str) -> None: + alive_addresses = {node["NodeManagerAddress"] for node in ray.nodes() if node["Alive"]} + expected = {head_node_ip, worker_node_ip} + if alive_addresses != expected: + raise RuntimeError(f"Expected exactly two Ray nodes {sorted(expected)}, got {sorted(alive_addresses)}") + + +class BenchmarkRun: + def __init__( + self, + mode: str, + base_config: DictConfig, + workload: Workload, + head_node_ip: str, + worker_node_ip: str, + ssd_path: str, + ) -> None: + self.mode = mode + self.workload = workload + + config = build_config(base_config, workload, mode == "ssd", ssd_path) + self.expected_storage_units = int(config["backend"]["SimpleStorage"]["num_data_storage_units"]) + self.writer = BenchmarkClient.options( + num_cpus=0.001, + resources={f"node:{head_node_ip}": 0.001}, + ).remote(config) + self.reader = BenchmarkClient.options( + num_cpus=0.001, + resources={f"node:{worker_node_ip}": 0.001}, + ).remote(config) + ray.get(self.writer.initialize.remote()) + ray.get(self.reader.initialize.remote()) + ray.get(self.writer.create_data.remote(workload)) + self.endpoint = ray.get(self.writer.metrics_endpoint.remote()) + if not self.endpoint: + raise RuntimeError("SimpleStorage metrics endpoint is unavailable") + + def stable_metrics(self, active_keys: int) -> StorageMetrics: + return wait_for_stable_rss( + self.endpoint, + self.expected_storage_units, + active_keys, + METRICS_TIMEOUT_SECONDS, + RSS_SAMPLE_INTERVAL_SECONDS, + RSS_STABLE_SAMPLES, + RSS_STABILITY_TOLERANCE_BYTES, + ) + + def run(self, iterations: int) -> list[dict[str, Any]]: + partition_id = "mixed" + baseline = self.stable_metrics(active_keys=0) + rows = [] + + for iteration in range(1, iterations + 1): + LOGGER.info("%s iteration %d/%d", self.mode, iteration, iterations) + put_start = time.perf_counter() + ray.get(self.writer.put.remote(partition_id)) + put_seconds = time.perf_counter() - put_start + after_put = self.stable_metrics(active_keys=self.workload.batch_size) + + keys = ray.get(self.reader.list_keys.remote(partition_id)) + if len(keys) != self.workload.batch_size: + raise AssertionError(f"Listed {len(keys)} keys, expected {self.workload.batch_size}") + get_start = time.perf_counter() + ray.get(self.reader.get.remote(partition_id, keys)) + get_seconds = time.perf_counter() - get_start + time.sleep(RSS_SAMPLE_INTERVAL_SECONDS) + after_get = self.stable_metrics(active_keys=self.workload.batch_size) + + ray.get(self.writer.clear.remote(partition_id, keys)) + after_clear = self.stable_metrics(active_keys=0) + + retained_bytes = after_clear.rss_bytes - baseline.rss_bytes + row = { + "mode": self.mode, + "iteration": iteration, + **asdict(self.workload), + "inline_payload_bytes": self.workload.inline_bytes, + "offload_candidate_bytes": self.workload.offloaded_bytes, + "total_payload_bytes": self.workload.total_bytes, + "put_seconds": put_seconds, + "get_seconds": get_seconds, + "put_gbit_per_second": self.workload.total_bytes * 8 / put_seconds / 1e9, + "get_gbit_per_second": self.workload.total_bytes * 8 / get_seconds / 1e9, + "storage_rss_before_put_bytes": baseline.rss_bytes, + "storage_rss_after_put_bytes": after_put.rss_bytes, + "storage_rss_after_get_bytes": after_get.rss_bytes, + "storage_rss_after_clear_bytes": after_clear.rss_bytes, + "storage_rss_retained_after_clear_bytes": retained_bytes, + "storage_ssd_active_bytes_after_put": after_put.ssd_active_bytes, + "storage_ssd_active_bytes_after_clear": after_clear.ssd_active_bytes, + } + rows.append(row) + LOGGER.info( + "%s RSS before/put/get/clear: %.3f/%.3f/%.3f/%.3f GiB; retained: %.1f MiB", + self.mode, + baseline.rss_bytes / 2**30, + after_put.rss_bytes / 2**30, + after_get.rss_bytes / 2**30, + after_clear.rss_bytes / 2**30, + retained_bytes / 2**20, + ) + baseline = after_clear + return rows + + def close(self) -> None: + try: + ray.get([self.writer.close.remote(), self.reader.close.remote()]) + finally: + ray.kill(self.writer, no_restart=True) + ray.kill(self.reader, no_restart=True) + + +def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def summarize( + ssd_rows: list[dict[str, Any]], + workload: Workload, + warmup_iterations: int, +) -> dict[str, Any]: + ssd_measured = ssd_rows[warmup_iterations:] + retained_values = [row["storage_rss_retained_after_clear_bytes"] for row in ssd_measured] + clear_rss_values = [row["storage_rss_after_clear_bytes"] for row in ssd_measured] + return { + "workload": asdict(workload), + "analyzed_iterations": len(ssd_measured), + "theoretical_offload_bytes": workload.offloaded_bytes, + "ssd_median_active_bytes_after_put": int( + statistics.median(row["storage_ssd_active_bytes_after_put"] for row in ssd_measured) + ), + "ssd_max_active_bytes_after_clear": max(row["storage_ssd_active_bytes_after_clear"] for row in ssd_measured), + "ssd_median_rss_retained_after_clear_bytes": int(statistics.median(retained_values)), + "ssd_max_rss_retained_after_clear_bytes": max(retained_values), + "ssd_clear_rss_growth_bytes": clear_rss_values[-1] - clear_rss_values[0], + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend_config", type=Path, required=True) + parser.add_argument("--head_node_ip", required=True) + parser.add_argument("--worker_node_ip", required=True) + parser.add_argument("--ssd_path", type=Path, required=True) + parser.add_argument("--output_dir", type=Path, default=Path(__file__).resolve().parent / "results") + parser.add_argument("--num_test_iterations", type=int, default=4) + parser.add_argument("--warmup_iterations", type=int, default=1) + parser.add_argument("--global_batch_size", type=int, default=512) + parser.add_argument("--small_fields", type=int, default=4) + parser.add_argument("--large_fields", type=int, default=4) + parser.add_argument("--small_sample_bytes", type=int, default=DEFAULT_SMALL_SAMPLE_BYTES) + parser.add_argument("--large_sample_bytes", type=int, default=DEFAULT_LARGE_SAMPLE_BYTES) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + base_config = OmegaConf.load(args.backend_config) + workload = Workload( + batch_size=args.global_batch_size, + small_fields=args.small_fields, + large_fields=args.large_fields, + small_sample_bytes=args.small_sample_bytes, + large_sample_bytes=args.large_sample_bytes, + offload_threshold_bytes=int(base_config.backend.SimpleStorage.ssd_offload.threshold_bytes), + ) + workload.validate() + if args.num_test_iterations <= args.warmup_iterations: + raise ValueError("num_test_iterations must be greater than warmup_iterations") + if not args.ssd_path.is_dir(): + raise ValueError(f"SSD path does not exist on the benchmark driver node: {args.ssd_path}") + + ray.init(address=RAY_ADDRESS) + validate_cluster(args.head_node_ip, args.worker_node_ip) + results: dict[str, list[dict[str, Any]]] = {} + try: + for mode in ("memory", "ssd"): + benchmark = BenchmarkRun( + mode=mode, + base_config=base_config, + workload=workload, + head_node_ip=args.head_node_ip, + worker_node_ip=args.worker_node_ip, + ssd_path=str(args.ssd_path), + ) + try: + results[mode] = benchmark.run(args.num_test_iterations) + finally: + benchmark.close() + write_csv(args.output_dir / f"mixed_{mode}.csv", results[mode]) + time.sleep(MODE_GAP_SECONDS) + + summary = summarize( + results["ssd"], + workload, + args.warmup_iterations, + ) + args.output_dir.mkdir(parents=True, exist_ok=True) + summary_path = args.output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2) + "\n") + LOGGER.info("Mixed SSD offload summary:\n%s", json.dumps(summary, indent=2)) + finally: + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/scripts/performance_test_ssd_offload/perftest_config.yaml b/scripts/performance_test_ssd_offload/perftest_config.yaml new file mode 100644 index 00000000..381c6b0d --- /dev/null +++ b/scripts/performance_test_ssd_offload/perftest_config.yaml @@ -0,0 +1,47 @@ +# This is the default configuration of TransferQueue. Users may modify the default value +# and use transfer_queue.init(conf) to overwrite the config entries. + +# Prometheus metrics exporter. +metrics: + enabled: true + # HTTP port for /metrics endpoint (0 = auto-assign free port) + port: 0 + +controller: + # User-defined sampler. User can pass sampler instance to overwrite this string config. + sampler: SequentialSampler + # Whether return an empty BatchMeta to prevent request blocking when no enough data is available + polling_mode: False + # ZMQ Server IP & Ports (automatically generated during init) + zmq_info: null + + +backend: + # Pluggable storage/transport backend of TransferQueue. Choose from: + # SimpleStorage, Yuanrong, MooncakeStore, ... + storage_backend: SimpleStorage + + # SimpleStorage: ZMQ-based in-memory storage for out-of-the-box usage + SimpleStorage: + # Maximum number of experience samples to hold across all storage units + total_storage_size: 100000 + # Number of distributed storage units. Units are round-robin scheduled across eligible + # alive Ray nodes, guaranteeing an even split of memory/bandwidth usage per node. + # Recommended: >= 2 x number of nodes so each node hosts multiple units. + num_data_storage_units: 16 + # Optional Ray custom resource required on storage nodes. null keeps all alive nodes eligible. + required_node_resource: null + # ZMQ Server IP & Ports (automatically generated during init) + zmq_info: null + + # SSD offload stores large field values on local SSD to reduce long-lived host-memory use. + # See docs/ssd_offload.md for configuration and cleanup details. + ssd_offload: + # Master switch. Set to true to enable SSD offload. + enabled: false + # SSD offload target directory; created automatically if it does not exist. + path: null + # Samples at or above this encoded-size threshold (in bytes) are stored on SSD. + threshold_bytes: 1048576 + # glibc-only allocator threshold for reclaiming large transient allocations. + glibc_mmap_threshold_bytes: 1048576 diff --git a/scripts/performance_test_ssd_offload/run_perf_test.sh b/scripts/performance_test_ssd_offload/run_perf_test.sh new file mode 100644 index 00000000..6b447149 --- /dev/null +++ b/scripts/performance_test_ssd_offload/run_perf_test.sh @@ -0,0 +1,24 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +RESULTS_DIR="${RESULTS_DIR:-${SCRIPT_DIR}/results}" +PERFTEST_PY="${SCRIPT_DIR}/perftest.py" +CONFIG_YAML="${SCRIPT_DIR}/perftest_config.yaml" + +: "${HEAD_NODE_IP:?Set HEAD_NODE_IP to the Ray head node IP}" +: "${WORKER_NODE_IP:?Set WORKER_NODE_IP to the Ray worker node IP}" +: "${SSD_OFFLOAD_PATH:?Set SSD_OFFLOAD_PATH to node-local SSD storage}" + +NUM_TEST_ITERATIONS="${NUM_TEST_ITERATIONS:-4}" + +mkdir -p "${RESULTS_DIR}" + +python "${PERFTEST_PY}" \ + --backend_config="${CONFIG_YAML}" \ + --head_node_ip="${HEAD_NODE_IP}" \ + --worker_node_ip="${WORKER_NODE_IP}" \ + --ssd_path="${SSD_OFFLOAD_PATH}" \ + --output_dir="${RESULTS_DIR}" \ + --num_test_iterations="${NUM_TEST_ITERATIONS}" diff --git a/scripts/put_benchmark.py b/scripts/put_benchmark.py index c67bb54c..822f8dd7 100644 --- a/scripts/put_benchmark.py +++ b/scripts/put_benchmark.py @@ -17,7 +17,6 @@ import asyncio import json import logging -import math import os import time @@ -267,6 +266,10 @@ def initialize_system(self, config_dict): ) total_storage_size = self.tq_config.global_batch_size * 2 + storage_config = { + "total_storage_size": total_storage_size, + "num_data_storage_units": self.num_storage_units, + } logger.info(f"Initializing Storage Units (Remote={self.remote_mode}, Target={self.target_ip})...") @@ -277,7 +280,7 @@ def initialize_system(self, config_dict): num_cpus=1, resources={f"node:{self.target_ip}": 0.001}, runtime_env={"env_vars": {"OMP_NUM_THREADS": "2"}}, - ).remote(storage_unit_size=math.ceil(total_storage_size / self.num_storage_units)) + ).remote(config=storage_config) else: # Local Mode: Use placement group self.storage_placement_group = get_placement_group(self.num_storage_units, num_cpus_per_actor=2) @@ -286,7 +289,7 @@ def initialize_system(self, config_dict): placement_group=self.storage_placement_group, placement_group_bundle_index=rank, runtime_env={"env_vars": {"OMP_NUM_THREADS": "2"}}, - ).remote(storage_unit_size=math.ceil(total_storage_size / self.num_storage_units)) + ).remote(config=storage_config) # Controller Init self.data_system_controller = TransferQueueController.remote() diff --git a/tests/e2e/test_ssd_offload_e2e.py b/tests/e2e/test_ssd_offload_e2e.py new file mode 100644 index 00000000..eb9060b6 --- /dev/null +++ b/tests/e2e/test_ssd_offload_e2e.py @@ -0,0 +1,143 @@ +# 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. + +# 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. + +"""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 ssd_root.is_dir() + assert not list((ssd_root / "transfer_queue_ssd_offload").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) + assert list((checkpoint_dir / "simple_storage").glob("*.pkl.blobs/*.bin")) + + 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/test_metrics.py b/tests/test_metrics.py index 66e82a3f..b3b077bb 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -13,6 +13,21 @@ # 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. + """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time @@ -98,6 +113,10 @@ def test_all_metrics_are_registered(self): "tq_storage_active_keys_total", "tq_storage_utilization_ratio", "tq_storage_memory_rss_bytes", + "tq_storage_ssd_offload_enabled", + "tq_storage_ssd_active_values", + "tq_storage_ssd_active_bytes", + "tq_storage_ssd_fallback_values", "tq_storage_requests_arrived", "tq_storage_arrivals_by_op", "tq_storage_accept_queue_backlog", @@ -313,6 +332,10 @@ def test_storage_metrics_populated_on_success(self): "capacity": 1000, "active_keys": 250, "process_rss_bytes": 512 * 1024 * 1024, + "ssd_offload_enabled": 1, + "ssd_active_values": 120, + "ssd_active_bytes": 4 * 1024 * 1024 * 1024, + "ssd_fallback_values_total": 7, } ) @@ -322,6 +345,13 @@ def test_storage_metrics_populated_on_success(self): assert exporter.storage_active_keys.labels(storage_unit_id="SU_001")._value.get() == 250 assert exporter.storage_utilization.labels(storage_unit_id="SU_001")._value.get() == 0.25 assert exporter.storage_memory_rss.labels(storage_unit_id="SU_001")._value.get() == 512 * 1024 * 1024 + assert exporter.storage_ssd_offload_enabled.labels(storage_unit_id="SU_001")._value.get() == 1 + assert exporter.storage_ssd_active_values.labels(storage_unit_id="SU_001")._value.get() == 120 + assert exporter.storage_ssd_active_bytes.labels(storage_unit_id="SU_001")._value.get() == 4 * 1024 * 1024 * 1024 + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 7 + + exporter.collect_storage_metrics() + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 7 def test_arrival_counters_are_exported(self): """Arrival counts reach Prometheus, so a dashboard can compare them with completions.""" @@ -456,6 +486,10 @@ def test_storage_metrics_skips_capacity_when_none(self): # active_keys and memory_rss are still reported assert exporter.storage_active_keys.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 42 assert exporter.storage_memory_rss.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 128 * 1024 * 1024 + assert exporter.storage_ssd_offload_enabled.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 + assert exporter.storage_ssd_active_values.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 + assert exporter.storage_ssd_active_bytes.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_UNLIMITED")._value.get() == 0 def test_storage_metrics_prunes_stale_capacity_on_switch_to_unlimited(self): """If a storage unit transitions from a numeric capacity to unlimited, diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 1145d963..5747c6a3 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -13,16 +13,40 @@ # 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 pickle import time +from pathlib import Path from uuid import uuid4 +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, + SSDFileStore, + StorageKeyNotFoundError, + _SSDValueRef, +) from transfer_queue.utils.zmq_utils import ( STORAGE_MANAGER_IDENTITY_PREFIX, ZMQMessage, @@ -112,7 +136,9 @@ def storage_setup(ray_setup): tensordict.set_list_to_stack(True).set() # Start Ray actor for SimpleStorageUnit - storage_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote(storage_unit_size=storage_size) + storage_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote( + config={"total_storage_size": storage_size, "num_data_storage_units": 1} + ) # Get ZMQ server info from storage unit zmq_info = ray.get(storage_actor.get_zmq_server_info.remote()) @@ -467,6 +493,248 @@ 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): + """Large samples use SSD while small samples preserve the in-memory contract.""" + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + large_sparse = torch.sparse_coo_tensor( + torch.stack((torch.arange(16), torch.arange(16))), + torch.ones(16), + size=(16, 16), + ) + try: + storage.put_data( + { + "mixed": [b"x" * 100, b"y" * 10], + "small": [1, 2], + "unsupported": [large_sparse, b"small"], + }, + [1, 2], + ) + + assert storage.get_data(["mixed", "small"], [1, 2]) == { + "small": [1, 2], + "mixed": [b"x" * 100, b"y" * 10], + } + assert len(list(tmp_path.rglob("*.bin"))) == 1 + assert storage.ssd_active_values == 1 + assert storage.ssd_active_bytes == 100 + assert storage.ssd_fallback_values_total == 1 + + storage.clear([1]) + assert storage.active_key_count == 1 + assert storage.ssd_active_values == 0 + assert storage.ssd_active_bytes == 0 + assert storage.ssd_fallback_values_total == 1 + assert not list(tmp_path.rglob("*.bin")) + with pytest.raises(StorageKeyNotFoundError): + storage.get_data(["mixed"], [1]) + + storage.clear([2]) + with pytest.raises(StorageKeyNotFoundError): + storage.get_data(["mixed"], [2]) + finally: + storage.close() + + assert not (tmp_path / "transfer_queue_ssd_offload" / "test-run").exists() + + +def test_hybrid_storage_cleans_completed_ssd_writes_after_batch_failure(tmp_path, monkeypatch): + """A failed SSD batch removes files written by its successful tasks.""" + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + + write_value_file = storage._ssd_store._write_value_file + + def fail_one_write(encoded_value): + if encoded_value.payload[0] == ord("y"): + raise OSError("injected SSD failure") + return write_value_file(encoded_value) + + monkeypatch.setattr(storage._ssd_store, "_write_value_file", fail_one_write) + try: + with pytest.raises(OSError, match="injected SSD failure"): + storage.put_data( + {"small": [1, 2], "large": [b"x" * 100, b"y" * 100]}, + [1, 2], + ) + + assert not list(tmp_path.rglob("*.bin")) + 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_all_bytes", 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 storage._ssd_active_values == 1 + assert storage._ssd_active_bytes == len(old_value) + assert old_path.exists() + assert list(tmp_path.rglob("*.bin")) == [old_path] + finally: + storage.close() + + +@pytest.mark.parametrize("threshold_bytes", [0, -1]) +def test_hybrid_storage_requires_positive_threshold(tmp_path, threshold_bytes): + with pytest.raises(ValueError, match="threshold must be greater than zero"): + HybridStorageUnitData( + storage_size=10, + threshold_bytes=threshold_bytes, + ssd_path=str(tmp_path), + run_id="test-run", + unit_id="test-unit", + ) + + +def test_hybrid_checkpoint_copies_ssd_values_without_materializing(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path / "ssd"), + run_id="test-run", + unit_id="test-unit", + ) + checkpoint_path = tmp_path / "storage_unit.pkl" + try: + storage.put_data({"value": [b"small", b"x" * 100]}, [1, 2]) + + def fail_materialization(*_args): + raise AssertionError("checkpoint must not materialize or re-encode SSD values") + + with monkeypatch.context() as patch: + patch.setattr(storage._ssd_store, "read_values", fail_materialization) + storage.save_checkpoint(checkpoint_path, "test-unit") + + with open(checkpoint_path, "rb") as f: + manifest = pickle.load(f) + assert manifest["field_data"] == {"value": {1: b"small"}} + checkpoint_ref = manifest["ssd_index"]["value"][2] + blob_path = Path(f"{checkpoint_path}.blobs") / checkpoint_ref["filename"] + assert blob_path.read_bytes() == b"x" * 100 + + storage.clear([1, 2]) + with monkeypatch.context() as patch: + patch.setattr( + HybridStorageUnitData, + "_sample_from_value", + staticmethod(fail_materialization), + ) + storage.load_checkpoint(checkpoint_path) + + assert storage.get_data(["value"], [1, 2]) == {"value": [b"small", b"x" * 100]} + assert storage._ssd_active_values == 1 + assert storage._ssd_active_bytes == 100 + finally: + storage.close() + + +def test_hybrid_storage_loads_legacy_logical_checkpoint(tmp_path): + storage = HybridStorageUnitData( + storage_size=1, + threshold_bytes=64, + ssd_path=str(tmp_path / "ssd"), + run_id="test-run", + unit_id="test-unit", + ) + checkpoint_path = tmp_path / "legacy.pkl" + with open(checkpoint_path, "wb") as f: + pickle.dump( + { + "storage_unit_id": "old-unit", + "storage_unit_size": 10, + "field_data": {"value": {1: b"x" * 100, 2: b"y" * 100}}, + "active_keys": {1, 2}, + }, + f, + ) + + try: + storage.load_checkpoint(checkpoint_path) + assert all(isinstance(value, _SSDValueRef) for value in storage.field_data["value"].values()) + assert storage._ssd_active_values == 2 + assert storage._ssd_active_bytes == 200 + assert storage.get_data(["value"], [1, 2]) == {"value": [b"x" * 100, b"y" * 100]} + 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 all( + isinstance(value, _SSDValueRef) + for stored_values in storage.field_data.values() + for value in stored_values.values() + ) + assert {entry.codec for entry in storage.field_data["tensor"].values()} == {"tensor"} + assert {entry.codec for entry in storage.field_data["array"].values()} == {"numpy"} + assert {entry.codec for entry in storage.field_data["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_file_store_creates_missing_parent_directories(tmp_path): + ssd_path = tmp_path / "missing" / "nested" + store = SSDFileStore(str(ssd_path), "test-run", "test-unit") + try: + assert (ssd_path / "transfer_queue_ssd_offload").is_dir() + finally: + store.close() + + def test_storage_unit_data_parser(storage_setup): """Test data_parser functionality in SimpleStorageUnit. @@ -663,7 +931,9 @@ def test_storage_unit_checkpoint_round_trip(storage_setup, tmp_path): assert (tmp_path / "storage_unit.pkl").exists() # 3. Create a fresh storage unit and load the checkpoint into it - fresh_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote(storage_unit_size=10000) + fresh_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote( + config={"total_storage_size": 10000, "num_data_storage_units": 1} + ) fresh_zmq_info = ray.get(fresh_actor.get_zmq_server_info.remote()) import time as _time @@ -701,7 +971,9 @@ def test_storage_unit_checkpoint_overwrites_existing_data(storage_setup, tmp_pat assert response.body["success"] is True # 2. Create a second unit, pre-populate it with different data, then load the checkpoint - second_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote(storage_unit_size=10000) + second_actor = SimpleStorageUnit.options(max_concurrency=50, num_cpus=1).remote( + config={"total_storage_size": 10000, "num_data_storage_units": 1} + ) second_zmq_info = ray.get(second_actor.get_zmq_server_info.remote()) import time as _time diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index bd83a599..c4da5fba 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -35,6 +35,18 @@ backend: # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null + # SSD offload stores large field values on local SSD to reduce long-lived host-memory use. + # See docs/ssd_offload.md for configuration and cleanup details. + ssd_offload: + # Master switch. Set to true to enable SSD offload. + enabled: false + # SSD offload target directory; created automatically if it does not exist. + path: null + # Samples at or above this encoded-size threshold (in bytes) are stored on SSD. + threshold_bytes: 1048576 + # glibc-only allocator threshold for reclaiming large transient allocations. + glibc_mmap_threshold_bytes: 1048576 + # 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..cffc4cc2 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -13,6 +13,21 @@ # 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 json import os import shutil @@ -45,6 +60,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 @@ -238,9 +257,18 @@ def close(): if _TQ_STORAGE: 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) + # Only the initializing process owns and shuts down these storage actors. + 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/metrics.py b/transfer_queue/metrics.py index e75792b5..fc62f7ac 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -13,6 +13,21 @@ # 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 os import time from contextlib import contextmanager @@ -85,6 +100,7 @@ def __init__(self, role: str = "controller", zmq_context: zmq.Context | None = N self._known_partition_ids: set[str] = set() self._known_production_labels: set[tuple[str, str]] = set() self._known_consumption_labels: set[tuple[str, str]] = set() + self._storage_ssd_fallback_values_seen: dict[str, int] = {} self._metrics_endpoint: str = "" # Plain-dict snapshot pushed by the controller via update_controller_snapshot(). @@ -184,6 +200,30 @@ def _define_controller_metrics(self) -> None: self.storage_memory_rss = Gauge( "tq_storage_memory_rss_bytes", "Storage unit process RSS memory", ["storage_unit_id"], registry=r ) + self.storage_ssd_offload_enabled = Gauge( + "tq_storage_ssd_offload_enabled", + "Whether SSD offload is enabled for the storage unit", + ["storage_unit_id"], + registry=r, + ) + self.storage_ssd_active_values = Gauge( + "tq_storage_ssd_active_values", + "Active field values stored on SSD", + ["storage_unit_id"], + registry=r, + ) + self.storage_ssd_active_bytes = Gauge( + "tq_storage_ssd_active_bytes", + "Logical bytes held by active SSD-backed values", + ["storage_unit_id"], + registry=r, + ) + self.storage_ssd_fallback_values = Counter( + "tq_storage_ssd_fallback_values", + "Values retained in memory because SSD encoding was unavailable", + ["storage_unit_id"], + registry=r, + ) # ---- Storage-unit request-loss diagnostics ---- # Read against tq_storage_request_ops, which advances only on completion: a gap @@ -417,6 +457,23 @@ def collect_storage_metrics(self) -> None: pass self.storage_active_keys.labels(storage_unit_id=label).set(active) self.storage_memory_rss.labels(storage_unit_id=label).set(metrics.get("process_rss_bytes", 0)) + self.storage_ssd_offload_enabled.labels(storage_unit_id=label).set( + metrics.get("ssd_offload_enabled", 0) + ) + self.storage_ssd_active_values.labels(storage_unit_id=label).set(metrics.get("ssd_active_values", 0)) + self.storage_ssd_active_bytes.labels(storage_unit_id=label).set(metrics.get("ssd_active_bytes", 0)) + fallback_values = metrics.get("ssd_fallback_values_total", 0) + previous_fallback_values = self._storage_ssd_fallback_values_seen.get(label, 0) + # Storage units report lifetime totals. Advance the exporter counter only + # by unseen events, treating a lower value as a storage-unit restart. + fallback_delta = ( + fallback_values - previous_fallback_values + if fallback_values >= previous_fallback_values + else fallback_values + ) + if fallback_delta: + self.storage_ssd_fallback_values.labels(storage_unit_id=label).inc(fallback_delta) + self._storage_ssd_fallback_values_seen[label] = fallback_values self.storage_requests_arrived.labels(storage_unit_id=label).set(metrics.get("requests_arrived", 0)) for op_type, arrived in (metrics.get("arrivals_by_op") or {}).items(): diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 381341ec..f0de7a0b 100644 --- a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py +++ b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py @@ -13,8 +13,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math +# 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. + from typing import Any +from uuid import uuid4 from omegaconf import DictConfig @@ -26,30 +41,40 @@ logger = get_logger(__name__) +DEFAULT_GLIBC_MMAP_THRESHOLD_BYTES = 1024 * 1024 + @StorageBootstrapProvider.register_provider("SimpleStorage") def initialize_simple_storage(conf: DictConfig) -> dict[str, Any]: """Initialize Simple storage with metastore mode.""" simple_storage_handles = {} - num_data_storage_units = conf.backend.SimpleStorage.num_data_storage_units - total_storage_size = conf.backend.SimpleStorage.get("total_storage_size", None) - required_node_resource = conf.backend.SimpleStorage.get("required_node_resource", None) + simple_storage_config = dict(conf.backend.SimpleStorage) + simple_storage_config["_run_id"] = uuid4().hex + num_data_storage_units = simple_storage_config["num_data_storage_units"] + required_node_resource = simple_storage_config.get("required_node_resource") + ssd_config = simple_storage_config.get("ssd_offload") or {} + storage_actor_runtime_env = None + if ssd_config.get("enabled", False): + mmap_threshold = ssd_config.get("glibc_mmap_threshold_bytes", DEFAULT_GLIBC_MMAP_THRESHOLD_BYTES) + if mmap_threshold is not None: + mmap_threshold = int(mmap_threshold) + if mmap_threshold <= 0: + raise ValueError("glibc_mmap_threshold_bytes must be greater than zero or null") + storage_actor_runtime_env = {"env_vars": {"MALLOC_MMAP_THRESHOLD_": str(mmap_threshold)}} scheduling_strategies = get_node_round_robin_scheduling_strategies( num_data_storage_units, required_node_resource=required_node_resource ) - # Compute per-unit capacity: None means unlimited - storage_unit_size = ( - math.ceil(total_storage_size / num_data_storage_units) if total_storage_size is not None 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, + actor_options = { + "scheduling_strategy": scheduling_strategies[storage_unit_rank], + "name": f"TransferQueueStorageUnit#{storage_unit_rank}", + } + if storage_actor_runtime_env is not None: + actor_options["runtime_env"] = storage_actor_runtime_env + storage_node = SimpleStorageUnit.options(**actor_options).remote( # type: ignore[attr-defined] + config=simple_storage_config ) 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 61c5e9a1..bc900b3b 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -13,16 +13,38 @@ # 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 math import os import pickle +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, log_heavy_operation @@ -49,6 +71,34 @@ 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)) +DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES = 1024 * 1024 +DEFAULT_SSD_READ_THREADS = 32 +DEFAULT_SSD_WRITE_THREADS = 8 +SSD_OFFLOAD_DIRECTORY_NAME = "transfer_queue_ssd_offload" + +_HYBRID_CHECKPOINT_FORMAT = "transfer_queue_hybrid_storage_v1" + + +@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 _SSDValueRef: + """Internal reference to one SSD-backed value.""" + + path: Path + size_bytes: 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" @@ -95,6 +145,26 @@ def active_key_count(self) -> int: """Number of active keys currently stored.""" return len(self._active_keys) + @property + def ssd_offload_enabled(self) -> bool: + """Whether this storage data offloads values to SSD.""" + return False + + @property + def ssd_active_values(self) -> int: + """Number of active field values stored on SSD.""" + return 0 + + @property + def ssd_active_bytes(self) -> int: + """Logical bytes held by active SSD-backed values.""" + return 0 + + @property + def ssd_fallback_values_total(self) -> int: + """Values retained in memory because SSD encoding was unavailable.""" + return 0 + def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: """Get data by global index keys. @@ -156,6 +226,486 @@ def clear(self, keys: list[int]) -> None: self.field_data[f].pop(key, None) self._active_keys -= set(keys) + def save_checkpoint(self, path: str | Path, storage_unit_id: str) -> None: + """Write in-memory storage state to a checkpoint.""" + state = { + "storage_unit_id": storage_unit_id, + "storage_unit_size": self.storage_size, + "field_data": self.field_data, + "active_keys": self._active_keys, + } + with open(path, "wb") as f: + pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + + def load_checkpoint(self, path: str | Path) -> int | None: + """Replace in-memory storage state from a checkpoint.""" + with open(path, "rb") as f: + state = pickle.load(f) + if state.get("format") == _HYBRID_CHECKPOINT_FORMAT: + raise ValueError("An SSD checkpoint requires SSD offload to be enabled") + + checkpoint_size = state["storage_unit_size"] + restored_field_data = state["field_data"] + restored_active_keys = set(state["active_keys"]) + self.field_data = restored_field_data + self._active_keys = restored_active_keys + return checkpoint_size + + def close(self) -> None: + """Release resources owned by this storage data.""" + return None + + +class SSDFileStore: + """Own and read/write one SSD file per offloaded value.""" + + def __init__( + self, + ssd_path: str, + run_id: str, + unit_id: str, + ) -> None: + self._closed = False + configured_path = Path(ssd_path).resolve() + if configured_path.exists() and not configured_path.is_dir(): + raise ValueError(f"SSD offload path is not a directory: {configured_path}") + configured_path.mkdir(parents=True, exist_ok=True) + self._ssd_root = configured_path / SSD_OFFLOAD_DIRECTORY_NAME + self._ssd_root.mkdir(exist_ok=True) + if self._ssd_root.is_symlink() or not self._ssd_root.is_dir(): + raise ValueError(f"SSD offload working path must be a directory, not a symlink: {self._ssd_root}") + self._base_path = self._ssd_root / run_id / unit_id + self._base_path.mkdir(parents=True) + for prefix in range(256): + (self._base_path / f"{prefix:02x}").mkdir() + self._read_pool = ThreadPoolExecutor( + max_workers=DEFAULT_SSD_READ_THREADS, + thread_name_prefix="tq-ssd-read", + ) + self._write_pool = ThreadPoolExecutor( + max_workers=DEFAULT_SSD_WRITE_THREADS, + thread_name_prefix="tq-ssd-write", + ) + + @staticmethod + def _write_all_bytes(fd: int, payload: memoryview) -> None: + """Write the entire payload to an open file descriptor, handling partial writes.""" + view = payload.cast("B") + while view: + written = os.write(fd, view) + if written <= 0: + raise OSError("SSDFileStore write returned no progress") + view = view[written:] + + def _write_value_file( + self, + encoded_value: SSDEncodedSample, + ) -> _SSDValueRef: + """Atomically write one encoded value to a new SSD file and return its reference.""" + token = uuid4().hex + directory = self._base_path / token[:2] + 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_all_bytes(fd, encoded_value.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 _SSDValueRef( + path=final_path, + size_bytes=encoded_value.payload.nbytes, + codec=encoded_value.codec, + dtype=encoded_value.dtype, + shape=encoded_value.shape, + ) + + def write_values( + self, + encoded_values: list[SSDEncodedSample], + ) -> list[_SSDValueRef]: + """Write encoded values concurrently in input order and remove completed files if the batch fails.""" + refs: list[_SSDValueRef | None] = [None] * len(encoded_values) + futures = { + self._write_pool.submit(self._write_value_file, encoded_value): position + for position, encoded_value in enumerate(encoded_values) + } + first_error: Exception | None = None + for future, position in futures.items(): + try: + refs[position] = future.result() + except Exception as e: + if first_error is None: + first_error = e + if first_error is not None: + for ref in refs: + if ref is not None: + self.unlink(ref) + raise first_error + return [ref for ref in refs if ref is not None] + + def import_file(self, source: Path, metadata: dict[str, Any]) -> _SSDValueRef: + """Copy one checkpoint blob into this store without materializing it.""" + token = uuid4().hex + directory = self._base_path / token[:2] + temp_path = directory / f".tmp-{token}" + final_path = directory / f"{token}.bin" + try: + shutil.copyfile(source, temp_path) + actual_size = temp_path.stat().st_size + if actual_size != metadata["size_bytes"]: + raise OSError( + f"SSD checkpoint blob has {actual_size} bytes, expected {metadata['size_bytes']}: {source}" + ) + temp_path.rename(final_path) + except Exception: + temp_path.unlink(missing_ok=True) + raise + return _SSDValueRef( + path=final_path, + size_bytes=metadata["size_bytes"], + codec=metadata["codec"], + dtype=metadata["dtype"], + shape=metadata["shape"], + ) + + @staticmethod + def unlink(entry: _SSDValueRef) -> None: + """Delete one offloaded value, tolerating cleanup failures.""" + try: + entry.path.unlink(missing_ok=True) + except OSError as e: + logger.warning(f"Failed to delete superseded SSD sample {entry.path}: {e}") + + def read_values(self, refs: list[_SSDValueRef]) -> list[bytes]: + """Read SSD-backed values concurrently and return their payloads in input order.""" + return list(self._read_pool.map(self._read_value_file, refs)) + + @staticmethod + def _read_value_file(ref: _SSDValueRef) -> bytes: + """Read one SSD value file and verify its size against the stored reference.""" + payload = ref.path.read_bytes() + if len(payload) != ref.size_bytes: + raise OSError( + f"SSDFileStore short read from {ref.path}: expected {ref.size_bytes} bytes, got {len(payload)}" + ) + return payload + + 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) + shutil.rmtree(self._base_path, ignore_errors=True) + try: + self._base_path.parent.rmdir() + except OSError: + pass + + +class HybridStorageUnitData(StorageUnitData): + """Store each value inline or as an SSD file reference in one field map.""" + + 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: + if threshold_bytes <= 0: + raise ValueError("SSD offload threshold must be greater than zero") + super().__init__(storage_size) + self._ssd_store = SSDFileStore(ssd_path, run_id, unit_id) + self._threshold = threshold_bytes + self._ssd_active_values = 0 + self._ssd_active_bytes = 0 + self._ssd_fallback_values_total = 0 + + @property + def ssd_offload_enabled(self) -> bool: + """Whether this storage data offloads values to SSD.""" + return True + + @property + def ssd_active_values(self) -> int: + """Number of active field values stored on SSD.""" + return self._ssd_active_values + + @property + def ssd_active_bytes(self) -> int: + """Logical bytes held by active SSD-backed values.""" + return self._ssd_active_bytes + + @property + def ssd_fallback_values_total(self) -> int: + """Values retained in memory because SSD encoding was unavailable.""" + return self._ssd_fallback_values_total + + @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") + + @staticmethod + def _decode_sample(raw: bytes, entry: _SSDValueRef) -> 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) + 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") + 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 _prepare_field_values( + self, + values: Any, + ) -> tuple[list[Any], list[_SSDValueRef], int]: + """Encode one field and write its offloaded values to the active SSD directory.""" + logical_samples = list(values.unbind()) if isinstance(values, torch.Tensor) else list(values) + encoded_samples = [self._sample_from_value(sample) for sample in logical_samples] + prepared_values: list[Any] = [] + ssd_positions: list[int] = [] + ssd_samples: list[SSDEncodedSample] = [] + fallback_values = 0 + for position, (value, encoded) in enumerate(zip(logical_samples, encoded_samples, strict=True)): + if encoded is not None and encoded.payload.nbytes >= self._threshold: + ssd_positions.append(position) + ssd_samples.append(encoded) + prepared_values.append(None) + else: + prepared_values.append(value) + if encoded is None: + fallback_values += 1 + + entries = self._ssd_store.write_values(ssd_samples) + for position, entry in zip(ssd_positions, entries, strict=True): + prepared_values[position] = entry + return prepared_values, entries, fallback_values + + def put_data( + self, + field_data: dict[str, Any], + global_indexes: list, + ) -> None: + """Store each sample in memory or SSD according to its encoded size.""" + if not global_indexes or not field_data: + super().put_data(field_data, global_indexes) + return + + for field, values in field_data.items(): + prepared_values, entries, fallback_values = self._prepare_field_values(values) + + stored_field = self.field_data.get(field, {}) + old_ssd_values = [] + for global_index in set(global_indexes): + old_value = stored_field.get(global_index) + if isinstance(old_value, _SSDValueRef): + old_ssd_values.append(old_value) + try: + super().put_data({field: prepared_values}, global_indexes) + except Exception: + for entry in entries: + self._ssd_store.unlink(entry) + raise + + self._ssd_active_values += len(entries) - len(old_ssd_values) + self._ssd_active_bytes += sum(entry.size_bytes for entry in entries) - sum( + value.size_bytes for value in old_ssd_values + ) + self._ssd_fallback_values_total += fallback_values + for old_value in old_ssd_values: + self._ssd_store.unlink(old_value) + + def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: + """Read mixed memory- and SSD-backed samples in request order.""" + result = super().get_data(fields, global_indexes) + for values in result.values(): + ssd_values = [(position, value) for position, value in enumerate(values) if isinstance(value, _SSDValueRef)] + raw_values = self._ssd_store.read_values([value for _, value in ssd_values]) + for (position, entry), raw in zip(ssd_values, raw_values, strict=True): + values[position] = self._decode_sample(raw, entry) + return result + + def clear(self, keys: list) -> None: + """Remove values and unlink any files they reference.""" + ssd_values: list[_SSDValueRef] = [] + for values in self.field_data.values(): + for key in set(keys): + value = values.get(key) + if isinstance(value, _SSDValueRef): + ssd_values.append(value) + super().clear(keys) + self._ssd_active_values -= len(ssd_values) + self._ssd_active_bytes -= sum(value.size_bytes for value in ssd_values) + for value in ssd_values: + self._ssd_store.unlink(value) + + def save_checkpoint(self, path: str | Path, storage_unit_id: str) -> None: + """Write memory values to a manifest and copy SSD values beside it.""" + manifest_path = Path(path) + blob_dir = Path(f"{manifest_path}.blobs") + shutil.rmtree(blob_dir, ignore_errors=True) + checkpoint_fields: dict[str, dict[int, Any]] = {} + ssd_index: dict[str, dict[int, dict[str, Any]]] = {} + try: + for field, values in self.field_data.items(): + checkpoint_values = {} + field_ssd_index = {} + for global_index, value in values.items(): + if not isinstance(value, _SSDValueRef): + checkpoint_values[global_index] = value + continue + + blob_dir.mkdir(parents=True, exist_ok=True) + filename = value.path.name + destination = blob_dir / filename + shutil.copyfile(value.path, destination) + actual_size = destination.stat().st_size + if actual_size != value.size_bytes: + raise OSError(f"SSD value has {actual_size} bytes, expected {value.size_bytes}: {value.path}") + field_ssd_index[global_index] = { + "filename": filename, + "size_bytes": value.size_bytes, + "codec": value.codec, + "dtype": value.dtype, + "shape": value.shape, + } + checkpoint_fields[field] = checkpoint_values + if field_ssd_index: + ssd_index[field] = field_ssd_index + + state = { + "format": _HYBRID_CHECKPOINT_FORMAT, + "storage_unit_id": storage_unit_id, + "storage_unit_size": self.storage_size, + "field_data": checkpoint_fields, + "ssd_index": ssd_index, + "active_keys": set(self._active_keys), + } + with open(manifest_path, "wb") as f: + pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + manifest_path.unlink(missing_ok=True) + shutil.rmtree(blob_dir, ignore_errors=True) + raise + + def load_checkpoint(self, path: str | Path) -> int | None: + """Replace memory and SSD state after all checkpoint data is ready.""" + manifest_path = Path(path) + with open(manifest_path, "rb") as f: + state = pickle.load(f) + checkpoint_format = state.get("format") + if checkpoint_format not in (None, _HYBRID_CHECKPOINT_FORMAT): + raise ValueError(f"Unsupported HybridStorageUnitData checkpoint format: {checkpoint_format}") + checkpoint_size = state["storage_unit_size"] + field_data = state["field_data"] + restored_active_keys = set(state["active_keys"]) + old_ssd_values = [ + value for values in self.field_data.values() for value in values.values() if isinstance(value, _SSDValueRef) + ] + restored_ssd_values: list[_SSDValueRef] = [] + restored_fallback_values = 0 + + try: + if checkpoint_format == _HYBRID_CHECKPOINT_FORMAT: + blob_dir = Path(f"{manifest_path}.blobs") + restored_field_data = {field: dict(values) for field, values in field_data.items()} + for field, entries in state["ssd_index"].items(): + restored_values = restored_field_data.setdefault(field, {}) + for global_index, metadata in entries.items(): + filename = metadata["filename"] + if Path(filename).name != filename: + raise ValueError(f"Invalid SSD checkpoint blob name: {filename}") + restored_value = self._ssd_store.import_file(blob_dir / filename, metadata) + restored_ssd_values.append(restored_value) + restored_values[global_index] = restored_value + else: + restored_field_data = {} + for field, values in field_data.items(): + if not values: + restored_field_data[field] = {} + continue + indexes = list(values) + prepared_values, entries, fallback_values = self._prepare_field_values( + [values[key] for key in indexes] + ) + restored_ssd_values.extend(entries) + restored_fallback_values += fallback_values + restored_field_data[field] = dict(zip(indexes, prepared_values, strict=True)) + restored_ssd_active_values = len(restored_ssd_values) + restored_ssd_active_bytes = sum(value.size_bytes for value in restored_ssd_values) + except Exception: + for value in restored_ssd_values: + self._ssd_store.unlink(value) + raise + + self.field_data = restored_field_data + self._active_keys = restored_active_keys + self._ssd_active_values = restored_ssd_active_values + self._ssd_active_bytes = restored_ssd_active_bytes + self._ssd_fallback_values_total += restored_fallback_values + for value in old_ssd_values: + self._ssd_store.unlink(value) + return checkpoint_size + + def close(self) -> None: + """Release SSD resources owned by this hybrid store.""" + self._ssd_store.close() + @ray.remote(num_cpus=1) class SimpleStorageUnit: @@ -179,17 +729,41 @@ class SimpleStorageUnit: _arrivals_by_op: dict[str, int] = {} _accept_probe = None - def __init__(self, storage_unit_size: int | None = None): - """Initialize a SimpleStorageUnit with the specified size. + def __init__(self, config: dict[str, Any]): + """Initialize a SimpleStorageUnit from the SimpleStorage config. Args: - storage_unit_size: Maximum number of elements that can be stored in this storage unit. - If None, the storage unit has unlimited capacity. + config: The ``backend.SimpleStorage`` configuration. Bootstrap adds one + shared internal run ID. """ 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) + total_storage_size = config.get("total_storage_size") + num_data_storage_units = int(config.get("num_data_storage_units", 1)) + self.storage_unit_size = ( + math.ceil(total_storage_size / num_data_storage_units) if total_storage_size is not None else None + ) + self.storage_data: StorageUnitData + + ssd_config = config.get("ssd_offload") + 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") + threshold_bytes = int(ssd_config.get("threshold_bytes", DEFAULT_SSD_OFFLOAD_THRESHOLD_BYTES)) + self.storage_data = HybridStorageUnitData( + storage_size=self.storage_unit_size, + ssd_path=str(ssd_path), + run_id=str(config.get("_run_id") or uuid4().hex), + unit_id=self.storage_unit_id, + threshold_bytes=threshold_bytes, + ) + logger.info( + f"[{self.storage_unit_id}]: SSD offload enabled — " + f"path={Path(ssd_path).resolve() / SSD_OFFLOAD_DIRECTORY_NAME}, " + f"threshold={threshold_bytes} B/sample" + ) + else: + self.storage_data = StorageUnitData(self.storage_unit_size) self._requests_arrived = 0 self._arrivals_by_op = {} @@ -221,8 +795,15 @@ def __init__(self, storage_unit_size: int | None = None): self.zmq_context, self.put_get_socket, self._accept_probe, + 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: @@ -304,14 +885,18 @@ def _proxy_routine(self) -> None: events = dict(poller.poll(1000)) if front in events: messages = front.recv_multipart(copy=False) - identity = bytes(messages[0]) if messages else b"" - if not identity.startswith(STORAGE_CLIENT_IDENTITY_PREFIXES): - logger.warning( - "[%s]: dropping request with unrecognized ZMQ identity", - self.storage_unit_id, - ) - continue - back.send_multipart(messages, copy=False) + try: + identity = bytes(messages[0]) if messages else b"" + if not identity.startswith(STORAGE_CLIENT_IDENTITY_PREFIXES): + logger.warning( + "[%s]: dropping request with unrecognized ZMQ identity", + self.storage_unit_id, + ) + continue + back.send_multipart(messages, copy=False) + finally: + # This thread outlives each request, so drop forwarded payload frames while idle. + del messages if back in events: front.send_multipart(back.recv_multipart(copy=False), copy=False) @@ -351,93 +936,97 @@ def _worker_routine(self) -> None: break if worker_socket in socks: - # Messages received from proxy: [identity, serialized_msg_frame1, ...] - messages = worker_socket.recv_multipart(copy=False) - identity = messages[0] - serialized_msg = messages[1:] - - try: - request_msg = ZMQMessage.deserialize(serialized_msg) - except Exception as e: - # The identity filter cannot cover this: an allowed peer can still send - # frames that fail to decode, and decoding here used to kill the thread. - logger.error( - f"[{self.storage_unit_id}]: undecodable request from " - f"identity={bytes(identity)!r}: {type(e).__name__}: {e}" - ) - error_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={"message": f"undecodable request: {type(e).__name__}: {e}"}, - ) - worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) - continue - operation = request_msg.request_type - started = time.perf_counter() - - try: - self._requests_arrived += 1 - self._arrivals_by_op[operation.name] = self._arrivals_by_op.get(operation.name, 0) + 1 - - logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") - - # 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) - elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="GET_DATA"): - response_msg = self._handle_get(request_msg) - elif operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] - 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() - 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] - response_msg = self._handle_load_checkpoint(request_msg) - else: - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_GET_OPERATION_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"Storage unit id #{self.storage_unit_id} " - f"receive invalid operation: {operation}." - }, - ) - except Exception as e: - logger.error( - f"[{self.storage_unit_id}]: worker error during {operation} " - f"from sender={request_msg.sender_id}: {type(e).__name__}: {e}" - ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] - sender_id=self.storage_unit_id, - body={ - "message": f"{self.storage_unit_id}, worker encountered error " - f"during operation {operation}: {str(e)}." - }, - ) - - # Send response back with identity for routing - response_frames = response_msg.serialize() - if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] - # This end serializes the get response, so its frames give the true wire size. - log_heavy_operation( - self.storage_unit_id, - "get", - time.perf_counter() - started, - sum(frame_nbytes(frame) or 0 for frame in response_frames), - f"samples={len(request_msg.body.get('global_indexes', []))} " - f"fields={list(request_msg.body.get('fields', []))}", - ) - worker_socket.send_multipart([identity] + response_frames, copy=False) + # Keep request handling in a separate scope so payload references are + # released before the worker waits for the next request. + self._process_one_worker_request(worker_socket, monitor) logger.info(f"[{self.storage_unit_id}]: worker stopped.") poller.unregister(worker_socket) worker_socket.close(linger=0) + def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) -> None: + """Process one storage request and send its response.""" + # Messages received from proxy: [identity, serialized_msg_frame1, ...] + messages = worker_socket.recv_multipart(copy=False) + identity = messages[0] + serialized_msg = messages[1:] + + try: + request_msg = ZMQMessage.deserialize(serialized_msg) + except Exception as e: + # The identity filter cannot cover this: an allowed peer can still send + # frames that fail to decode, and decoding here used to kill the thread. + logger.error( + f"[{self.storage_unit_id}]: undecodable request from " + f"identity={bytes(identity)!r}: {type(e).__name__}: {e}" + ) + error_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={"message": f"undecodable request: {type(e).__name__}: {e}"}, + ) + worker_socket.send_multipart([identity] + error_msg.serialize(), copy=False) + return + + operation = request_msg.request_type + started = time.perf_counter() + + try: + self._requests_arrived += 1 + self._arrivals_by_op[operation.name] = self._arrivals_by_op.get(operation.name, 0) + 1 + + logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") + + if operation == ZMQRequestType.PUT_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="PUT_DATA"): + response_msg = self._handle_put(request_msg) + elif operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="GET_DATA"): + response_msg = self._handle_get(request_msg) + elif operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] + 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() + 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] + response_msg = self._handle_load_checkpoint(request_msg) + else: + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_OPERATION_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={ + "message": f"Storage unit id #{self.storage_unit_id} receive invalid operation: {operation}." + }, + ) + except Exception as e: + logger.error( + f"[{self.storage_unit_id}]: worker error during {operation} " + f"from sender={request_msg.sender_id}: {type(e).__name__}: {e}" + ) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.PUT_GET_ERROR, # type: ignore[arg-type] + sender_id=self.storage_unit_id, + body={ + "message": f"{self.storage_unit_id}, worker encountered error " + f"during operation {operation}: {str(e)}." + }, + ) + + response_frames = response_msg.serialize() + if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type] + # This end serializes the get response, so its frames give the true wire size. + log_heavy_operation( + self.storage_unit_id, + "get", + time.perf_counter() - started, + sum(frame_nbytes(frame) or 0 for frame in response_frames), + f"samples={len(request_msg.body.get('global_indexes', []))} " + f"fields={list(request_msg.body.get('fields', []))}", + ) + worker_socket.send_multipart([identity] + response_frames, copy=False) + def _handle_put(self, data_parts: ZMQMessage) -> ZMQMessage: """ Handle put request, add or update data into storage unit. @@ -613,11 +1202,14 @@ def _handle_get_metrics(self) -> ZMQMessage: "capacity": self.storage_unit_size, "active_keys": self.storage_data.active_key_count, "process_rss_bytes": process_rss, + "ssd_offload_enabled": int(self.storage_data.ssd_offload_enabled), + "ssd_active_values": self.storage_data.ssd_active_values, + "ssd_active_bytes": self.storage_data.ssd_active_bytes, + "ssd_fallback_values_total": self.storage_data.ssd_fallback_values_total, # Counted on arrival; op_stats below only advances on completion. "requests_arrived": self._requests_arrived, "arrivals_by_op": dict(self._arrivals_by_op), } - if self._accept_probe is not None: stats = self._accept_probe.stats metrics["accept_queue"] = { @@ -659,11 +1251,11 @@ def _handle_get_metrics(self) -> ZMQMessage: ) def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: - """Serialize storage unit data directly to a file. + """Write storage unit data directly to its checkpoint path. Args: data_parts: ZMQMessage from client, containing ``path`` in body: - absolute path for the output .pkl file. The caller must ensure + absolute path for the output manifest. The caller must ensure this path is reachable from the node running this actor (shared filesystem required for multi-node setups). @@ -673,14 +1265,7 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: """ path = data_parts.body["path"] try: - 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, - } - with open(path, "wb") as f: - pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL) + self.storage_data.save_checkpoint(path, self.storage_unit_id) logger.info(f"[{self.storage_unit_id}]: saved checkpoint to {path}") return ZMQMessage.create( request_type=ZMQRequestType.SAVE_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] @@ -696,7 +1281,7 @@ def _handle_save_checkpoint(self, data_parts) -> ZMQMessage: ) def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: - """Restore storage unit data directly from a file. + """Restore storage unit data directly from its checkpoint path. Args: data_parts: ZMQMessage from client, containing ``path`` in body: @@ -711,28 +1296,24 @@ def _handle_load_checkpoint(self, data_parts) -> ZMQMessage: """ path = data_parts.body["path"] try: - with open(path, "rb") as f: - data = pickle.load(f) + previous_key_count = self.storage_data.active_key_count + checkpoint_size = self.storage_data.load_checkpoint(path) - if data["storage_unit_size"] != self.storage_unit_size: + if checkpoint_size != self.storage_unit_size: logger.warning( f"[{self.storage_unit_id}]: storage_unit_size mismatch — " - f"checkpoint={data['storage_unit_size']}, current={self.storage_unit_size}" + f"checkpoint={checkpoint_size}, current={self.storage_unit_size}" ) - if self.storage_data._active_keys: + if previous_key_count: logger.warning( - f"[{self.storage_unit_id}]: overwriting {len(self.storage_data._active_keys)} " + f"[{self.storage_unit_id}]: overwriting {previous_key_count} " 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} — " - f"{len(data['active_keys'])} keys, {len(data['field_data'])} fields" + f"{self.storage_data.active_key_count} keys, {len(self.storage_data.field_data)} fields" ) return ZMQMessage.create( request_type=ZMQRequestType.LOAD_STORAGE_CHECKPOINT_RESPONSE, # type: ignore[arg-type] @@ -788,30 +1369,34 @@ def _shutdown_resources( zmq_context: zmq.Context | None, put_get_socket: zmq.Socket | None, accept_probe: "AcceptQueueProbe | None" = None, + worker_socket: zmq.Socket | None = None, + storage_data: StorageUnitData | None = None, ) -> None: """Clean up resources on garbage collection.""" logger.info("Shutting down SimpleStorageUnit resources...") - # Signal all threads to stop shutdown_event.set() # Before the ZMQ teardown: the probe runs on its own timer and would outlive the unit. if accept_probe is not None: accept_probe.stop() - # Terminate put_get_socket + if worker_thread and worker_thread.is_alive(): + worker_thread.join() + if proxy_thread and proxy_thread.is_alive(): + proxy_thread.join() + if put_get_socket: put_get_socket.close(linger=0) - - # Terminate ZMQ context to unblock proxy and workers + if worker_socket: + worker_socket.close(linger=0) 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) + if storage_data is not None: + try: + storage_data.close() + except Exception as e: + logger.warning(f"Error closing storage data on shutdown: {e}") logger.info("SimpleStorageUnit resources shutdown complete.") From f5e9f55248231fd802bb93455b328029aa46285f Mon Sep 17 00:00:00 2001 From: pinjie Date: Wed, 23 Sep 2026 06:29:08 +0000 Subject: [PATCH 2/3] small fix Signed-off-by: pinjie --- transfer_queue/storage/simple_storage.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index bc900b3b..c5d18a63 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -85,7 +85,7 @@ class SSDEncodedSample: payload: memoryview codec: str - dtype: str | None = None + dtype: str | np.dtype[Any] | None = None shape: tuple[int, ...] | None = None @@ -96,7 +96,7 @@ class _SSDValueRef: path: Path size_bytes: int codec: str - dtype: str | None = None + dtype: str | np.dtype[Any] | None = None shape: tuple[int, ...] | None = None @@ -482,7 +482,7 @@ def _sample_from_value(value: Any) -> SSDEncodedSample | None: return SSDEncodedSample( payload=payload, codec="numpy", - dtype=str(array.dtype), + dtype=array.dtype, shape=tuple(array.shape), ) if isinstance(value, bytes): From 42e33f6bde8e711cb5d025a38880f0664ba2d3ec Mon Sep 17 00:00:00 2001 From: pinjie Date: Wed, 23 Sep 2026 13:52:11 +0000 Subject: [PATCH 3/3] fix Signed-off-by: pinjie --- docs/metrics.md | 2 +- docs/ssd_offload.md | 6 +-- tests/test_metrics.py | 5 +- tests/test_simple_storage_unit.py | 61 +++++++++++++++++++++++- transfer_queue/interface.py | 4 +- transfer_queue/metrics.py | 18 ++----- transfer_queue/storage/simple_storage.py | 25 +++++++--- 7 files changed, 92 insertions(+), 29 deletions(-) diff --git a/docs/metrics.md b/docs/metrics.md index 4f59db79..70ba670f 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -140,7 +140,7 @@ Steps: | `tq_storage_ssd_offload_enabled` | Gauge | `storage_unit_id` | `1` when SSD offload is enabled, otherwise `0` | | `tq_storage_ssd_active_values` | Gauge | `storage_unit_id` | Active field values stored on SSD | | `tq_storage_ssd_active_bytes` | Gauge | `storage_unit_id` | Logical bytes held by active SSD-backed values | -| `tq_storage_ssd_fallback_values_total` | Counter | `storage_unit_id` | Values retained in memory because SSD encoding was unavailable | +| `tq_storage_ssd_fallback_values_total` | Gauge | `storage_unit_id` | Cumulative values retained in memory because SSD encoding was unavailable | | `tq_storage_request_ops` | Gauge | `storage_unit_id`, `op_type` | Total requests processed by storage unit | | `tq_storage_request_latency_avg` | Gauge | `storage_unit_id`, `op_type` | Average request latency (seconds) | | `tq_storage_request_latency_p50` | Gauge | `storage_unit_id`, `op_type` | P50 request latency (seconds) | diff --git a/docs/ssd_offload.md b/docs/ssd_offload.md index 63111065..b0e5b66a 100644 --- a/docs/ssd_offload.md +++ b/docs/ssd_offload.md @@ -112,12 +112,12 @@ The following value types can be stored on SSD: |------------|--------------------| | Dense PyTorch tensor | Raw tensor bytes, dtype, and shape | | NumPy array without object dtype | Raw array bytes, dtype, and shape | +| NumPy array with object dtype | Pickle payload | | `bytes` | Original bytes | | Other Python value that `pickle` can serialize | Pickle payload | -Nested or sparse PyTorch tensors, NumPy arrays with object dtype, and values -that cannot be encoded remain in memory. A non-CPU tensor is copied to CPU -before it is written to SSD. +Nested or sparse PyTorch tensors and values that cannot be encoded remain in +memory. A non-CPU tensor is copied to CPU before it is written to SSD. GET reads SSD-backed values and reconstructs their original types. SSD file references are internal and are not returned to the application. diff --git a/tests/test_metrics.py b/tests/test_metrics.py index b3b077bb..f3e338b0 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -116,7 +116,7 @@ def test_all_metrics_are_registered(self): "tq_storage_ssd_offload_enabled", "tq_storage_ssd_active_values", "tq_storage_ssd_active_bytes", - "tq_storage_ssd_fallback_values", + "tq_storage_ssd_fallback_values_total", "tq_storage_requests_arrived", "tq_storage_arrivals_by_op", "tq_storage_accept_queue_backlog", @@ -350,8 +350,9 @@ def test_storage_metrics_populated_on_success(self): assert exporter.storage_ssd_active_bytes.labels(storage_unit_id="SU_001")._value.get() == 4 * 1024 * 1024 * 1024 assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 7 + exporter._query_storage_unit.return_value["ssd_fallback_values_total"] = 9 exporter.collect_storage_metrics() - assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 7 + assert exporter.storage_ssd_fallback_values.labels(storage_unit_id="SU_001")._value.get() == 9 def test_arrival_counters_are_exported(self): """Arrival counts reach Prometheus, so a dashboard can compare them with completions.""" diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 5747c6a3..99ceb4ec 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -657,6 +657,52 @@ def fail_materialization(*_args): storage.close() +def test_hybrid_checkpoint_load_failure_preserves_existing_state(tmp_path, monkeypatch): + storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path / "target-ssd"), + run_id="target-run", + unit_id="target-unit", + ) + checkpoint_storage = HybridStorageUnitData( + storage_size=10, + threshold_bytes=64, + ssd_path=str(tmp_path / "checkpoint-ssd"), + run_id="checkpoint-run", + unit_id="checkpoint-unit", + ) + checkpoint_path = tmp_path / "storage_unit.pkl" + try: + storage.put_data({"value": [b"original" * 16]}, [1]) + original_files = set((tmp_path / "target-ssd").rglob("*.bin")) + + checkpoint_storage.put_data({"value": [b"x" * 100, b"y" * 100]}, [2, 3]) + checkpoint_storage.save_checkpoint(checkpoint_path, "checkpoint-unit") + + import_file = storage._ssd_store.import_file + import_count = 0 + + def fail_on_second_import(source, metadata): + nonlocal import_count + import_count += 1 + if import_count == 2: + raise OSError("injected checkpoint import failure") + return import_file(source, metadata) + + monkeypatch.setattr(storage._ssd_store, "import_file", fail_on_second_import) + + with pytest.raises(OSError, match="injected checkpoint import failure"): + storage.load_checkpoint(checkpoint_path) + + assert storage.get_data(["value"], [1]) == {"value": [b"original" * 16]} + assert storage.active_key_count == 1 + assert set((tmp_path / "target-ssd").rglob("*.bin")) == original_files + finally: + storage.close() + checkpoint_storage.close() + + def test_hybrid_storage_loads_legacy_logical_checkpoint(tmp_path): storage = HybridStorageUnitData( storage_size=1, @@ -701,9 +747,18 @@ def test_hybrid_storage_round_trips_supported_codecs(tmp_path): torch.arange(40, dtype=torch.float32), ] arrays = np.arange(64, dtype=np.float32).reshape(2, 32) + structured = np.zeros((2, 8), dtype=[("index", " None: ["storage_unit_id"], registry=r, ) - self.storage_ssd_fallback_values = Counter( - "tq_storage_ssd_fallback_values", + self.storage_ssd_fallback_values = Gauge( + "tq_storage_ssd_fallback_values_total", "Values retained in memory because SSD encoding was unavailable", ["storage_unit_id"], registry=r, @@ -462,18 +461,9 @@ def collect_storage_metrics(self) -> None: ) self.storage_ssd_active_values.labels(storage_unit_id=label).set(metrics.get("ssd_active_values", 0)) self.storage_ssd_active_bytes.labels(storage_unit_id=label).set(metrics.get("ssd_active_bytes", 0)) - fallback_values = metrics.get("ssd_fallback_values_total", 0) - previous_fallback_values = self._storage_ssd_fallback_values_seen.get(label, 0) - # Storage units report lifetime totals. Advance the exporter counter only - # by unseen events, treating a lower value as a storage-unit restart. - fallback_delta = ( - fallback_values - previous_fallback_values - if fallback_values >= previous_fallback_values - else fallback_values + self.storage_ssd_fallback_values.labels(storage_unit_id=label).set( + metrics.get("ssd_fallback_values_total", 0) ) - if fallback_delta: - self.storage_ssd_fallback_values.labels(storage_unit_id=label).inc(fallback_delta) - self._storage_ssd_fallback_values_seen[label] = fallback_values self.storage_requests_arrived.labels(storage_unit_id=label).set(metrics.get("requests_arrived", 0)) for op_type, arrived in (metrics.get("arrivals_by_op") or {}).items(): diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index c5d18a63..df28521e 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -546,12 +546,15 @@ def put_data( super().put_data(field_data, global_indexes) return + unique_global_indexes = set(global_indexes) + has_duplicate_indexes = len(unique_global_indexes) != len(global_indexes) + for field, values in field_data.items(): prepared_values, entries, fallback_values = self._prepare_field_values(values) stored_field = self.field_data.get(field, {}) old_ssd_values = [] - for global_index in set(global_indexes): + for global_index in unique_global_indexes: old_value = stored_field.get(global_index) if isinstance(old_value, _SSDValueRef): old_ssd_values.append(old_value) @@ -562,13 +565,22 @@ def put_data( self._ssd_store.unlink(entry) raise - self._ssd_active_values += len(entries) - len(old_ssd_values) + obsolete_ssd_values = old_ssd_values + if has_duplicate_indexes: + retained_ssd_paths = set() + for global_index in unique_global_indexes: + retained_value = self.field_data[field][global_index] + if isinstance(retained_value, _SSDValueRef): + retained_ssd_paths.add(retained_value.path) + obsolete_ssd_values.extend(entry for entry in entries if entry.path not in retained_ssd_paths) + + self._ssd_active_values += len(entries) - len(obsolete_ssd_values) self._ssd_active_bytes += sum(entry.size_bytes for entry in entries) - sum( - value.size_bytes for value in old_ssd_values + value.size_bytes for value in obsolete_ssd_values ) self._ssd_fallback_values_total += fallback_values - for old_value in old_ssd_values: - self._ssd_store.unlink(old_value) + for obsolete_value in obsolete_ssd_values: + self._ssd_store.unlink(obsolete_value) def get_data(self, fields: list[str], global_indexes: list) -> dict[str, list]: """Read mixed memory- and SSD-backed samples in request order.""" @@ -923,7 +935,8 @@ def _worker_routine(self) -> None: while not self._shutdown_event.is_set(): monitor = self._metrics if self._metrics is not None else perf_monitor try: - socks = dict(poller.poll(TQ_STORAGE_POLLER_TIMEOUT * 1000)) + # The event cannot wake a ZMQ poll, so bound idle shutdown latency. + socks = dict(poller.poll(min(TQ_STORAGE_POLLER_TIMEOUT, 1) * 1000)) except zmq.error.ContextTerminated: # ZMQ context was terminated, exit gracefully logger.info(f"[{self.storage_unit_id}]: worker stopped gracefully (Context Terminated)")