diff --git a/README.md b/README.md index e2395263..97d7096a 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ TransferQueue offers **fine-grained, sub-sample-level** data management and **lo

🔄 Updates

+ - **September 8, 2026**: 🔥 TransferQueue has been adopted in [Meshy](https://github.com/OpenBMB/Meshy). Meshy models each role as a dedicated service, and TransferQueue serves as the data flow manager that connects each role. + - **August 14, 2026**: 🔥 TransferQueue has been adopted in [Dressage](https://github.com/Accio-Lab/Dressage). By integrating TQ, peak memory usage of the data plane on the master node is reduced by 91% in a 32-node GLM-5.2 setup. Refer to the [report](https://github.com/Accio-Lab/Dressage/blob/main/docs/trajectory-storage-refactor-en.md) for details. - **June 18, 2026**: 🔥 TransferQueue has been adopted in [ROLL](https://github.com/alibaba/ROLL/pull/463). This integration introduces a [`RemoteBatch`](https://github.com/alibaba/ROLL/blob/main/docs_roll/docs/User%20Guides/Advanced%20Features/remote_batch_transfer.md) abstraction, enabling seamless compatibility with existing `DataProto` design. - **June 9, 2026**: 🔥 TransferQueue has been adopted in [UniRL](https://github.com/Tencent-Hunyuan/UniRL), a unified RL framework for multimodal models developed by Tencent Hunyuan. - **April 15, 2026**: 🔥 TransferQueue has been adopted in [Relax](https://github.com/redai-infra/Relax)! By leveraging the `StreamingDataLoader` abstraction, it schedules training data across the cluster at micro-batch granularity, reducing synchronization barriers in a single-controller setup. 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 05559dd5..8401ce68 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -60,6 +60,8 @@ Steps: |---------------------|---------|-------------| | `TQ_METRICS_COLLECT_INTERVAL` | `10` | Background collection interval (seconds) | | `TQ_METRICS_STORAGE_TIMEOUT` | `5` | ZMQ timeout for storage unit queries (seconds) | +| `TQ_ACCEPT_PROBE_INTERVAL` | `0` (off) | Accept-queue sampling period (seconds); enables the accept-queue metrics below | +| `TQ_STORAGE_ZMQ_BACKLOG` | `4096` | Accept-queue depth for the storage unit's listening socket | ## Architecture @@ -135,11 +137,54 @@ 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_request_ops` | Gauge | `storage_unit_id`, `op_type` | Total requests processed by storage unit | +| `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` | Gauge | `storage_unit_id` | Cumulative values retained in memory because SSD encoding was unavailable | +| `tq_storage_request_ops` | Gauge | `storage_unit_id`, `op_type` | Cumulative entries into measured PUT, GET, and CLEAR handlers | | `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) | | `tq_storage_request_latency_p99` | Gauge | `storage_unit_id`, `op_type` | P99 request latency (seconds) | +### Storage Request-Loss Diagnostics (collected via ZMQ, exposed on controller) + +Arrivals count decoded worker requests, including GET_METRICS and checkpoint operations. +For NIXL PUT/GET, PREPARE counts as a PUT_DATA/GET_DATA arrival; COMMIT and CANCEL are +continuations and are not counted again. + +`tq_storage_request_ops` counts entry into measured PUT, GET, and CLEAR handlers, not their +completion or success. For NIXL PUT/GET, that entry occurs at COMMIT. A failed or cancelled +PREPARE, including a key miss under a concurrent clear, therefore leaves an arrival without +a measured handler entry. GET_METRICS and checkpoint operations also have arrivals without +corresponding measured handler counts. A sustained gap alone does not prove request loss or +an unfinished request. These cumulative per-operation counters carry no request identity and +characterise a unit, not any individual request. + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `tq_storage_requests_arrived` | Gauge | `storage_unit_id` | Decoded worker requests, excluding NIXL COMMIT/CANCEL continuations | +| `tq_storage_arrivals_by_op` | Gauge | `storage_unit_id`, `op_type` | Same, by operation; NIXL PREPARE maps to PUT_DATA/GET_DATA | + +The accept-queue series below exist only when the unit runs with `TQ_ACCEPT_PROBE_INTERVAL` +set; the series are removed rather than reported as zero when the probe is off, so a missing +series means "not measured" rather than "no drops". `tq_storage_accept_queue_peak` and the two +drop counters are cumulative since the probe started, so use `rate()` on them. + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `tq_storage_accept_queue_backlog` | Gauge | `storage_unit_id` | Configured accept-queue depth (`TQ_STORAGE_ZMQ_BACKLOG`) | +| `tq_storage_accept_queue_peak` | Gauge | `storage_unit_id` | Deepest accept-queue occupancy seen by the probe | +| `tq_storage_accept_queue_peak_utilization_ratio` | Gauge | `storage_unit_id` | Peak occupancy as a fraction of the backlog | +| `tq_storage_socket_drops` | Gauge | `storage_unit_id` | Connections dropped on this listening socket since the probe started | +| `tq_storage_listen_overflows` | Gauge | `storage_unit_id` | Namespace-wide accept-queue overflows since the probe started | +| `tq_storage_listen_other_drops` | Gauge | `storage_unit_id` | Namespace-wide establishment drops that were **not** overflows | + +The kernel charges a listening socket's `sk_drops` on several connection-establishment +failures — a full accept queue, but also failures to allocate or route the new connection — so +`tq_storage_socket_drops` rising locates the socket, not the cause. Compare the last two +series: a non-zero `tq_storage_listen_other_drops` means raising the backlog would not have +prevented every drop in that window. + ### Storage Unit Native Metrics (exposed on each storage unit's own endpoint) | Metric | Type | Labels | Description | @@ -163,6 +208,7 @@ The dashboard ([`scripts/grafana_dashboard.json`](../scripts/grafana_dashboard.j | **Request Throughput & Latency** | Controller Request Rate (ops/s), Controller Request Latency (repeats per quantile) | | **Partition Status** | Samples per Partition, Production Progress, Consumption Progress | | **Storage Units** | Utilization Bar Gauge, Active Keys, Capacity vs Active Keys, RSS Memory, Storage Request Rate, Storage Request Latency (repeats per quantile), Produced vs Cleared Samples/s, Active Keys Delta | +| **Storage Request-Loss Diagnostics** | Arrived vs Completed/s, Accept-Queue Peak vs Backlog, Listening-Socket Drops/s, Overflow vs Other Establishment Drops/s | ### Template Variables diff --git a/docs/ssd_offload.md b/docs/ssd_offload.md new file mode 100644 index 00000000..b0e5b66a --- /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 | +| NumPy array with object dtype | Pickle payload | +| `bytes` | Original bytes | +| Other Python value that `pickle` can serialize | Pickle payload | + +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. + +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/grafana_dashboard.json b/scripts/grafana_dashboard.json index 0e6c5341..2a95c8a1 100644 --- a/scripts/grafana_dashboard.json +++ b/scripts/grafana_dashboard.json @@ -370,6 +370,68 @@ { "expr": "sum(tq_storage_active_keys_total)", "legendFormat": "Total Active Keys (all storage units)" } ], "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 86 }, + "id": 104, + "title": "Storage Request-Loss Diagnostics", + "type": "row" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "Arrivals count decoded requests, including GET_METRICS and checkpoint operations. NIXL PREPARE counts as PUT_DATA/GET_DATA; COMMIT/CANCEL are not counted again. Measured requests count entry into PUT/GET/CLEAR handlers, at COMMIT for NIXL, not completion or success. Unmeasured operations and failed or cancelled PREPAREs, including concurrent-clear key misses, can leave a gap. A gap alone does not prove request loss or unfinished requests. Neither series identifies an individual request.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 87 }, + "id": 40, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Arrived vs Measured Requests (per second)", + "targets": [ + { "expr": "sum by (op_type) (rate(tq_storage_arrivals_by_op{op_type=~\"$op_type\"}[$__rate_interval]))", "legendFormat": "arrived {{ op_type }}" }, + { "expr": "sum by (op_type) (rate(tq_storage_request_ops{op_type=~\"$op_type\"}[$__rate_interval]))", "legendFormat": "measured {{ op_type }}" } + ], + "type": "timeseries" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "Peak is a high-water mark since the probe started, so it never decreases. Only present when TQ_ACCEPT_PROBE_INTERVAL is set.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 87 }, + "id": 41, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Accept-Queue Peak vs Backlog", + "targets": [ + { "expr": "tq_storage_accept_queue_peak", "legendFormat": "peak {{ storage_unit_id }}" }, + { "expr": "tq_storage_accept_queue_backlog", "legendFormat": "backlog {{ storage_unit_id }}" } + ], + "type": "timeseries" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "The kernel charges sk_drops for several connection-establishment failures, of which a full accept queue is only one, so a rise locates the socket, not the cause.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 95 }, + "id": 42, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Listening-Socket Drops (per second)", + "targets": [ + { "expr": "sum by (storage_unit_id) (rate(tq_storage_socket_drops[$__rate_interval]))", "legendFormat": "{{ storage_unit_id }}" } + ], + "type": "timeseries" + }, + { + "datasource": { "type": "prometheus", "uid": "${datasource}" }, + "description": "Namespace-wide, not per port. A non-zero 'other' series means raising the backlog would not have prevented every drop in the window.", + "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "fillOpacity": 10, "lineWidth": 2 } } }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 95 }, + "id": 43, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi" } }, + "title": "Overflow vs Other Establishment Drops (per second)", + "targets": [ + { "expr": "sum (rate(tq_storage_listen_overflows[$__rate_interval]))", "legendFormat": "accept-queue overflow" }, + { "expr": "sum (rate(tq_storage_listen_other_drops[$__rate_interval]))", "legendFormat": "other establishment drops" } + ], + "type": "timeseries" } ], "refresh": "10s", 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_accept_probe.py b/tests/test_accept_probe.py new file mode 100644 index 00000000..a834955c --- /dev/null +++ b/tests/test_accept_probe.py @@ -0,0 +1,257 @@ +# 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. + +"""Tests for the accept-queue probe. + +A connection dropped during establishment is invisible to both ends, and the kernel +charges the listening socket's sk_drops on several such paths, of which a full accept +queue is only one. These tests pin the arithmetic the probe reports, the levels it logs +at, and that it never names a cause its counters cannot establish. +""" + +import threading + +from transfer_queue.utils.accept_probe import ( + AcceptQueueProbe, + AcceptQueueSample, + AcceptQueueStats, + _read_listen_overflows, + sample_accept_queue, +) + + +def _sample( + recv_q: int, backlog: int = 100, sk_drops: int = 0, overflows: int = 0, drops: int | None = None +) -> AcceptQueueSample: + return AcceptQueueSample( + timestamp=0.0, + recv_q=recv_q, + backlog=backlog, + sk_drops=sk_drops, + listen_overflows=overflows, + listen_drops=overflows if drops is None else drops, + ) + + +def _probe() -> AcceptQueueProbe: + return AcceptQueueProbe(port=34513, owner_id="TQ_STORAGE_UNIT_test", interval_s=0.01) + + +def test_utilization_is_depth_over_backlog(): + assert _sample(50, backlog=100).utilization == 0.5 + + +def test_zero_backlog_reports_zero_utilization_not_zero_division(): + """A socket read before bind reports backlog 0; that must not raise.""" + assert _sample(5, backlog=0).utilization == 0.0 + + +def test_full_queue_reports_full_utilization(): + assert _sample(100, backlog=100).utilization == 1.0 + + +def test_empty_stats_report_zero_deltas(): + stats = AcceptQueueStats(port=1234) + + assert stats.sk_drops_delta == 0 + assert stats.overflow_delta == 0 + + +def test_drop_delta_is_last_minus_first(): + """Deltas are what matter: the absolute counters carry the whole uptime's history.""" + stats = AcceptQueueStats(port=1234) + stats.first_sample = _sample(0, sk_drops=23, overflows=1775) + stats.last_sample = _sample(0, sk_drops=31, overflows=1790) + + assert stats.sk_drops_delta == 8 + assert stats.overflow_delta == 15 + + +def test_describe_names_port_and_deltas(): + stats = AcceptQueueStats(port=34513, backlog=100, peak_recv_q=97) + stats.first_sample = _sample(0, sk_drops=23) + stats.last_sample = _sample(0, sk_drops=24) + + text = stats.describe() + + assert "port=34513" in text + assert "peak_recv_q=97" in text + assert "sk_drops_delta=1" in text + + +def test_peak_tracks_highest_depth_not_latest(): + """The burst is the signal; sampling after it drains reads zero.""" + probe = _probe() + + probe._record(_sample(10)) + probe._record(_sample(97)) + probe._record(_sample(3)) + + assert probe.stats.peak_recv_q == 97 + + +def test_first_sample_is_retained_as_delta_baseline(): + probe = _probe() + + probe._record(_sample(1, sk_drops=23)) + probe._record(_sample(2, sk_drops=25)) + + assert probe.stats.first_sample.sk_drops == 23 + assert probe.stats.sk_drops_delta == 2 + + +def test_drop_increase_is_logged_at_error(caplog): + """A drop is the direct evidence, so it must not be buried at debug level.""" + probe = _probe() + probe._record(_sample(0, sk_drops=23)) + + with caplog.at_level("ERROR"): + probe._record(_sample(100, sk_drops=24)) + + assert "dropped an incoming connection" in caplog.text + + +def test_steady_drop_count_is_logged_once_not_every_sample(caplog): + """sk_drops is cumulative, so a past drop must not be re-reported forever. + + Measuring against the probe's first sample made the condition permanently true once + the socket had ever dropped a connection: at a 0.1s interval that is ten errors a + second for the life of the process, and it erases when the drop actually happened. + """ + probe = _probe() + probe._record(_sample(0, sk_drops=12)) + + with caplog.at_level("ERROR"): + probe._record(_sample(0, sk_drops=13)) # a new drop -- report it + for _ in range(20): + probe._record(_sample(0, sk_drops=13)) # unchanged -- stay quiet + + assert caplog.text.count("dropped an incoming connection") == 1 + + +def test_each_new_drop_is_reported(caplog): + """Quieting the repeat must not swallow genuinely new drops.""" + probe = _probe() + probe._record(_sample(0, sk_drops=12)) + + with caplog.at_level("ERROR"): + probe._record(_sample(0, sk_drops=13)) + probe._record(_sample(0, sk_drops=13)) + probe._record(_sample(0, sk_drops=14)) + + assert caplog.text.count("dropped an incoming connection") == 2 + + +def test_near_full_queue_warns_once(caplog): + """Repeating the warning every 100ms would flood the log of a 512-node job.""" + probe = _probe() + + with caplog.at_level("WARNING"): + probe._record(_sample(60)) + probe._record(_sample(70)) + + assert caplog.text.count("reached") == 1 + + +def test_no_warning_below_threshold(): + probe = _probe() + + probe._record(_sample(10)) + + assert probe._warned is False + + +def test_listen_overflows_returns_two_non_negative_ints(): + """Reads the real host, so assert shape rather than a specific value.""" + overflows, drops = _read_listen_overflows() + + assert overflows >= 0 + assert drops >= 0 + + +def test_sampling_an_unused_port_returns_none(): + assert sample_accept_queue(1) is None + + +def test_unit_shutdown_stops_the_probe(): + """Nothing else calls stop(), so the sampling thread would outlive the unit.""" + from unittest.mock import MagicMock + + from transfer_queue.storage.simple_storage import SimpleStorageUnit + + unit_class = SimpleStorageUnit.__ray_metadata__.modified_class + probe = MagicMock() + + unit_class._shutdown_resources( + shutdown_event=threading.Event(), + worker_thread=None, + proxy_thread=None, + zmq_context=None, + put_get_socket=None, + accept_probe=probe, + ) + + probe.stop.assert_called_once() + + +def test_shutdown_without_a_probe_is_a_no_op(): + """The probe is opt-in, so the default path must not require one.""" + from transfer_queue.storage.simple_storage import SimpleStorageUnit + + unit_class = SimpleStorageUnit.__ray_metadata__.modified_class + + unit_class._shutdown_resources( + shutdown_event=threading.Event(), + worker_thread=None, + proxy_thread=None, + zmq_context=None, + put_get_socket=None, + ) + + +def test_non_overflow_drops_are_separated_from_overflows(): + """ListenDrops counts every establishment failure; only some are queue overflows. + + The kernel charges a listening socket's sk_drops on several paths -- a full accept + queue, but also failures to allocate or route the new connection -- so this difference + is what says whether a bigger backlog could have helped. + """ + stats = AcceptQueueStats(port=1234) + stats.first_sample = _sample(0, overflows=10, drops=20) + stats.last_sample = _sample(0, overflows=11, drops=25) + + assert stats.overflow_delta == 1 + assert stats.non_overflow_drop_delta == 4 + + +def test_pure_overflow_window_reports_no_other_drops(): + stats = AcceptQueueStats(port=1234) + stats.first_sample = _sample(0, overflows=10, drops=10) + stats.last_sample = _sample(0, overflows=13, drops=13) + + assert stats.overflow_delta == 3 + assert stats.non_overflow_drop_delta == 0 + + +def test_drop_alert_does_not_assert_the_queue_overflowed(caplog): + """sk_drops locates the socket, not the cause, so the alert must not name one.""" + probe = _probe() + probe._record(_sample(0, sk_drops=1)) + + with caplog.at_level("ERROR"): + probe._record(_sample(0, sk_drops=2)) + + assert "dropped an incoming connection" in caplog.text + assert "accept queue dropped" not in caplog.text diff --git a/tests/test_client.py b/tests/test_client.py index d287108e..aa887656 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -68,6 +68,7 @@ def __init__(self, controller_id="controller_0"): ) self.running = True + self.caller_identities: set[bytes] = set() self.request_thread = Thread(target=self._handle_requests, daemon=True) self.request_thread.start() @@ -85,6 +86,9 @@ def _handle_requests(self): if self.request_socket in socks: messages = self.request_socket.recv_multipart(copy=False) identity = messages.pop(0) + # Each pooled socket dials with its own identity, so this records how + # many distinct connections the client opened. + self.caller_identities.add(bytes(identity)) serialized_msg = messages request_msg = ZMQMessage.deserialize(serialized_msg) @@ -1330,6 +1334,26 @@ def test_kv_retrieve_keys_type_validation(self, client_setup): ) +# ===================================================== +# Controller RPC Socket Pool Tests +# ===================================================== + + +def test_controller_rpc_reuses_one_socket(client_setup): + """Repeated calls must share a single connection to the controller. + + This is the client-side payoff of the socket pool: TransferQueueClient runs every call + on the one event loop it builds in __init__, so all of them land in the same pool + bucket. Counted from the controller, which sees one ZMQ identity per socket dialled. + """ + client, mock_controller, _ = client_setup + + for _ in range(5): + assert client.get_partition_list() == ["partition_0", "partition_1", "test_partition"] + + assert len(mock_controller.caller_identities) == 1, "each request opened its own socket" + + # ===================================================== # Checkpoint Interface Tests # ===================================================== diff --git a/tests/test_controller.py b/tests/test_controller.py index 8a735f6a..32c1603a 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -21,7 +21,7 @@ import torch import zmq -from transfer_queue.controller import TransferQueueController +from transfer_queue.controller import TQ_CONTROLLER_ZMQ_MAX_SOCKETS, TransferQueueController from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket # Set up logging @@ -1465,6 +1465,27 @@ def test_controller_checkpoint_load_nonexistent_file(self, ray_setup, tmp_path): print("✓ load_checkpoint raises on missing file") +class TestControllerSocketBudget: + """The controller's context must hold more sockets than libzmq's default 1023. + + The metrics exporter borrows this context to query storage units, so the budget has to + cover the fleet rather than just the controller's own two ROUTERs. Past the ceiling a + socket cannot be opened at all, which would take out the control plane too. + """ + + def test_context_ceiling_is_raised_above_the_libzmq_default(self, ray_setup): + controller = TransferQueueController.remote() + try: + # Read it from inside the actor: the ceiling is only meaningful on the context + # the controller actually built. + budget = ray.get(controller.__ray_call__.remote(lambda self: self.zmq_context.get(zmq.MAX_SOCKETS))) + finally: + ray.kill(controller) + + assert budget > 1023, "the controller kept libzmq's default socket ceiling" + assert budget == TQ_CONTROLLER_ZMQ_MAX_SOCKETS + + class TestTransferQueueControllerBadRequests: """The request loop must survive requests it cannot decode, does not handle, or fails on. diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index 7ac3744d..14980fab 100644 --- a/tests/test_kv_storage_manager.py +++ b/tests/test_kv_storage_manager.py @@ -23,6 +23,28 @@ from transfer_queue.storage.managers.base import KVStorageManager +def test_close_is_quiet_when_config_is_rejected_before_super_init(): + """A config rejected before super().__init__() leaves nothing for close() to release. + + Every KV manager validates its config first, so none of the base attributes close() + reaches for exist yet. __del__ calls close() regardless, and an AttributeError there + buries the real ValueError under a spurious teardown error. + """ + built = [] + + class Probe(KVStorageManager): + def __init__(self, *args, **kwargs): + built.append(self) + super().__init__(*args, **kwargs) + + with pytest.raises(ValueError, match="Missing client_name"): + Probe(None, {}) + + manager = built[0] + assert not hasattr(manager, "notify_pool"), "the test no longer exercises the pre-super path" + manager.close() + + def get_meta(data, global_indexes=None): if not global_indexes: global_indexes = list(range(data.batch_size[0])) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 15c43882..eea13564 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -13,15 +13,36 @@ # 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 -from unittest.mock import MagicMock +from concurrent.futures import Future +from threading import Thread +from unittest.mock import MagicMock, patch import pytest try: + import zmq + from transfer_queue.metrics import TQMetricsExporter + from transfer_queue.utils.enum_utils import Role + from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo _HAS_DEPS = True except (ImportError, OSError): @@ -93,6 +114,18 @@ 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_total", + "tq_storage_requests_arrived", + "tq_storage_arrivals_by_op", + "tq_storage_accept_queue_backlog", + "tq_storage_accept_queue_peak", + "tq_storage_accept_queue_peak_utilization_ratio", + "tq_storage_socket_drops", + "tq_storage_listen_overflows", + "tq_storage_listen_other_drops", ] registered = {m.name for m in exporter.registry.collect()} @@ -207,6 +240,77 @@ def test_multiple_ops_tracked_independently(self): # --------------------------------------------------------------------------- +class TestStorageQuerySocketPool: + def test_pool_borrows_the_owner_context(self): + """The exporter must query storage units over the context it was handed. + + The controller already holds a long-lived synchronous context. A second one would + add another native I/O thread and leave a context nobody closes, since the exporter + lives for the whole life of its Ray actor. + """ + ctx = zmq.Context() + try: + exporter = TQMetricsExporter(zmq_context=ctx) + with patch("zmq.Context") as minted: + exporter._get_socket_pool() + minted.assert_not_called() + finally: + ctx.destroy(linger=0) + + def test_missing_context_is_reported(self): + """Without a context there is nothing to query over, so say so rather than mint one.""" + exporter = TQMetricsExporter() + with pytest.raises(RuntimeError, match="without a ZMQ context"): + exporter._get_socket_pool() + + def test_collector_parks_no_socket_between_queries(self): + """A queried unit must leave nothing in the pool. + + Collection walks every unit once per cycle, so a parked socket is reused only a + cycle later while holding a slot in the controller context's budget for the whole + walk. At a few thousand units that budget is what runs out first. + """ + identities: set[bytes] = set() + ctx_peer = zmq.Context() + router = ctx_peer.socket(zmq.ROUTER) + port = router.bind_to_random_port("tcp://127.0.0.1") + running = True + + def serve(): + poller = zmq.Poller() + poller.register(router, zmq.POLLIN) + while running: + if not dict(poller.poll(50)): + continue + identity, _ = router.recv_multipart() + identities.add(bytes(identity)) + response = ZMQMessage.create( + request_type=ZMQRequestType.METRICS_RESPONSE, + sender_id="storage_0", + body={}, + ) + router.send_multipart([identity, *response.serialize()]) + + server = Thread(target=serve, daemon=True) + server.start() + + su_info = ZMQServerInfo(role=Role.STORAGE, id="storage_0", ip="127.0.0.1", ports={"put_get_socket": port}) + ctx = zmq.Context() + try: + exporter = TQMetricsExporter(zmq_context=ctx) + for _ in range(3): + assert exporter._query_storage_unit(su_info, "storage_0") == {} + # A parked socket would be reused, so a fresh identity per query is the + # externally visible proof that nothing was kept. + assert len(identities) == 3, "a queried unit left its socket in the pool" + finally: + running = False + server.join(timeout=2.0) + ctx.destroy(linger=0) + router.close(linger=0) + ctx_peer.term() + + class TestStorageMetricsCollection: def test_collect_with_no_storage_units(self): """No storage units registered — collect should be a no-op.""" @@ -229,6 +333,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, } ) @@ -238,6 +346,84 @@ 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._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() == 9 + + def test_arrival_counters_are_exported(self): + """Arrival counts reach Prometheus, so a dashboard can compare them with completions.""" + exporter = TQMetricsExporter() + fake_su_info = MagicMock() + fake_su_info.id = "SU_001" + exporter._storage_unit_infos = {"SU_001": fake_su_info} + exporter._query_storage_unit = MagicMock( + return_value={ + "storage_unit_id": "SU_001", + "capacity": 1000, + "active_keys": 1, + "requests_arrived": 42, + "arrivals_by_op": {"GET_DATA": 30, "PUT_DATA": 12}, + } + ) + + exporter.collect_storage_metrics() + + assert exporter.storage_requests_arrived.labels(storage_unit_id="SU_001")._value.get() == 42 + by_op = exporter.storage_arrivals_by_op + assert by_op.labels(storage_unit_id="SU_001", op_type="GET_DATA")._value.get() == 30 + assert by_op.labels(storage_unit_id="SU_001", op_type="PUT_DATA")._value.get() == 12 + + def test_accept_queue_metrics_are_exported(self): + """The overflow/non-overflow split is what tells a dashboard if backlog is the issue.""" + exporter = TQMetricsExporter() + fake_su_info = MagicMock() + fake_su_info.id = "SU_001" + exporter._storage_unit_infos = {"SU_001": fake_su_info} + exporter._query_storage_unit = MagicMock( + return_value={ + "storage_unit_id": "SU_001", + "capacity": 1000, + "active_keys": 1, + "accept_queue": { + "backlog": 4096, + "peak_recv_q": 97, + "peak_utilization": 0.02, + "sk_drops_delta": 5, + "listen_overflow_delta": 2, + "non_overflow_drop_delta": 3, + }, + } + ) + + exporter.collect_storage_metrics() + + label = {"storage_unit_id": "SU_001"} + assert exporter.storage_accept_queue_backlog.labels(**label)._value.get() == 4096 + assert exporter.storage_accept_queue_peak.labels(**label)._value.get() == 97 + assert exporter.storage_socket_drops.labels(**label)._value.get() == 5 + assert exporter.storage_listen_overflows.labels(**label)._value.get() == 2 + assert exporter.storage_listen_other_drops.labels(**label)._value.get() == 3 + + def test_accept_queue_series_absent_when_the_probe_is_off(self): + """The probe is opt-in; reporting 0 drops would read as 'measured, and none'.""" + exporter = TQMetricsExporter() + fake_su_info = MagicMock() + fake_su_info.id = "SU_001" + exporter._storage_unit_infos = {"SU_001": fake_su_info} + exporter._query_storage_unit = MagicMock( + return_value={"storage_unit_id": "SU_001", "capacity": 1000, "active_keys": 1} + ) + + exporter.collect_storage_metrics() + + exported = {sample.name for metric in exporter.registry.collect() for sample in metric.samples} + assert "tq_storage_socket_drops" not in exported + assert "tq_storage_accept_queue_backlog" not in exported def test_storage_metrics_handles_query_failure(self): """If a storage unit query fails, other units should still be collected.""" @@ -302,6 +488,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, @@ -354,3 +544,137 @@ def test_metrics_request_types_exist(self): assert ZMQRequestType.GET_METRICS.value == "GET_METRICS" assert ZMQRequestType.METRICS_RESPONSE.value == "METRICS_RESPONSE" + + +class TestStorageWorkerArrivals: + @pytest.fixture + def worker(self): + from transfer_queue.storage.payload_transfer.zmq import ZmqPayloadTransfer + from transfer_queue.storage.simple_storage import SimpleStorageUnit, StorageUnitData + from transfer_queue.utils.perf_utils import IntervalPerfMonitor + + unit_class = SimpleStorageUnit.__ray_metadata__.modified_class + unit = object.__new__(unit_class) + unit.storage_unit_id = "unit_a" + unit.storage_unit_size = 4 + unit.storage_data = StorageUnitData(4) + unit.payload_transfer = ZmqPayloadTransfer() + unit._requests_arrived = 0 + unit._arrivals_by_op = {} + unit._accept_probe = None + unit._metrics = None + return unit, IntervalPerfMonitor("unit_a") + + @staticmethod + def dispatch(worker, operation, body): + unit, monitor = worker + request = ZMQMessage.create(request_type=operation, sender_id="manager", body=body) + socket = MagicMock() + socket.recv_multipart.return_value = [b"manager"] + request.serialize() + unit._process_one_worker_request(socket, monitor) + sent = socket.send_multipart.call_args.args[0] + assert sent[0] == b"manager" + return ZMQMessage.deserialize(sent[1:]) + + def test_zmq_metrics_and_checkpoints_keep_their_arrival_labels(self, worker, tmp_path): + path = str(tmp_path / "storage.pkl") + requests = [ + (ZMQRequestType.PUT_DATA, {"global_indexes": [0], "data": {"value": [42]}}, "PUT_DATA_RESPONSE"), + (ZMQRequestType.GET_DATA, {"global_indexes": [0], "fields": ["value"]}, "GET_DATA_RESPONSE"), + (ZMQRequestType.SAVE_STORAGE_CHECKPOINT, {"path": path}, "SAVE_STORAGE_CHECKPOINT_RESPONSE"), + (ZMQRequestType.CLEAR_DATA, {"global_indexes": [0]}, "CLEAR_DATA_RESPONSE"), + (ZMQRequestType.LOAD_STORAGE_CHECKPOINT, {"path": path}, "LOAD_STORAGE_CHECKPOINT_RESPONSE"), + (ZMQRequestType.GET_METRICS, {}, "METRICS_RESPONSE"), + ] + for operation, body, response_type in requests: + reply = self.dispatch(worker, operation, body) + assert reply.request_type == ZMQRequestType[response_type] + if "CHECKPOINT" in operation.name: + assert reply.body["success"] + unit, monitor = worker + expected = {operation.name: 1 for operation, _, _ in requests} + assert reply.body["requests_arrived"] == len(requests) + assert reply.body["arrivals_by_op"] == expected + assert dict(monitor.success_counts) == {"PUT_DATA": 1, "GET_DATA": 1, "CLEAR_DATA": 1} + assert unit.storage_data.get_data(["value"], [0]) == {"value": [42]} + + exporter = TQMetricsExporter() + exporter.register_storage_units({"unit_a": MagicMock()}) + exporter._query_storage_unit = MagicMock(return_value=reply.body) + exporter.collect_storage_metrics() + for operation in expected: + assert exporter.storage_arrivals_by_op.labels(storage_unit_id="unit_a", op_type=operation)._value.get() == 1 + + @pytest.mark.parametrize("operation", ["put", "get"]) + @pytest.mark.parametrize("finish", ["COMMIT", "CANCEL"]) + def test_nixl_counts_prepare_once_and_skips_continuations(self, worker, operation, finish): + from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer, PayloadDescriptor + from transfer_queue.utils.serial_utils import calc_packed_size, encode, pack_into + + unit, monitor = worker + unit.storage_data.put_data({"value": [42]}, [0]) + transfer = object.__new__(NixlPayloadTransfer) + transfer._pending_puts = {} + transfer._pending_gets = {} + transfer._runtime = MagicMock() + unit.payload_transfer = transfer + frames = encode({"value": [42]}) + payload = bytearray(calc_packed_size(frames)) + pack_into(payload, frames) + descriptor = PayloadDescriptor("payload", len(payload), tuple(memoryview(frame).nbytes for frame in frames)) + received = Future() + received.set_result(memoryview(payload)) + transfer._runtime.receive.return_value = received + sent = Future() + sent.set_result(None) + transfer._runtime.send.return_value = sent + transfer._runtime.prepare_receive.return_value = {} + body = { + "global_indexes": [0], + "fields": ["value"], + "transfer_id": "payload", + "descriptor": descriptor.to_dict(), + } + + ready = self.dispatch(worker, ZMQRequestType[f"{operation.upper()}_DATA_PREPARE"], body) + assert ready.request_type == ZMQRequestType[f"{operation.upper()}_DATA_READY"] + assert unit._requests_arrived == 1 + assert dict(monitor.success_counts) == {} + + reply = self.dispatch( + worker, + ZMQRequestType[f"{operation.upper()}_DATA_{finish}"], + { + "transfer_id": "payload", + "receiver_endpoint": {"transport": "nixl-ucx", "data": {"agent_name": "client"}}, + "receive_token": { + "data": {"agent_name": "client", "agent_metadata": b"meta", "frame_remote_descs": b"desc"} + }, + }, + ) + assert reply.request_type == ZMQRequestType[f"{operation.upper()}_DATA_RESPONSE"] + assert unit._requests_arrived == 1 + assert unit._arrivals_by_op == {f"{operation.upper()}_DATA": 1} + assert dict(monitor.success_counts) == ({f"{operation.upper()}_DATA": 1} if finish == "COMMIT" else {}) + + def test_nixl_prepare_key_miss_is_an_arrival_without_measured_processing(self, worker): + from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer + from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER + + unit, monitor = worker + unit.storage_data.put_data({"value": [42]}, [0]) + unit.storage_data.clear([0]) + transfer = object.__new__(NixlPayloadTransfer) + transfer._pending_gets = {} + unit.payload_transfer = transfer + + reply = self.dispatch( + worker, + ZMQRequestType.GET_DATA_PREPARE, + {"transfer_id": "missing", "global_indexes": [0], "fields": ["value"]}, + ) + assert reply.request_type == ZMQRequestType.PUT_GET_ERROR + assert KEY_NOT_FOUND_MARKER in reply.body["message"] + assert unit._requests_arrived == 1 + assert unit._arrivals_by_op == {"GET_DATA": 1} + assert dict(monitor.success_counts) == {} diff --git a/tests/test_mooncake_buffer_pool.py b/tests/test_mooncake_buffer_pool.py new file mode 100644 index 00000000..400360ac --- /dev/null +++ b/tests/test_mooncake_buffer_pool.py @@ -0,0 +1,509 @@ +# 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. + +"""Pooled receive buffers on the MooncakeStore tensor read path. + +The store and the buffer pool are faked, so these tests need neither mooncake nor RDMA. +The fake pool keeps the upstream properties this path depends on: buffers are carved out +of one fixed arena and reused, exhaustion raises, a buffer cannot be returned while an +exported view of it is alive, and the pool cannot close while a buffer is out on loan. +""" + +import ctypes +import logging +import threading +import weakref + +import pytest +import torch + +from transfer_queue.storage.clients import mooncake_client as mcc + +LOGGER_NAME = "transfer_queue.storage.clients.mooncake_client" +ALIGN = 256 + +DTYPES = [torch.float32, torch.int64, torch.float32, torch.int16] +SHAPES = [(4, 3), (5,), (), (2, 8)] +KEYS = ["k0", "k1", "k2", "k3"] + + +def expected_tensors(): + """Deterministic payloads, one per key in KEYS.""" + out = [] + for seed, (dtype, shape) in enumerate(zip(DTYPES, SHAPES, strict=True)): + numel = torch.empty(shape).numel() + values = torch.arange(seed * 100, seed * 100 + numel) + out.append(values.to(dtype).reshape(shape)) + return out + + +class FakeStore: + """Writes the stored payload into whatever pointer batch_get_into is given.""" + + def __init__(self): + self.objects = { + key: bytes(t.contiguous().numpy().tobytes()) for key, t in zip(KEYS, expected_tensors(), strict=True) + } + self.registered: list[tuple[int, int]] = [] + self.unregistered: list[int] = [] + self.closed = False + self.transfer_error: Exception | None = None + self.on_transfer = None + + def setup(self, *args): + return 0 + + def register_buffer(self, ptr, size): + self.registered.append((ptr, size)) + return 0 + + def unregister_buffer(self, ptr): + self.unregistered.append(ptr) + return 0 + + def batch_get_into(self, keys, ptrs, sizes): + if self.transfer_error is not None: + raise self.transfer_error + if self.on_transfer is not None: + self.on_transfer() + for key, ptr, size in zip(keys, ptrs, sizes, strict=True): + ctypes.memmove(ptr, self.objects[key], size) + return list(sizes) + + def close(self): + self.closed = True + + +class FakeBuffer: + """One region carved out of the pool's arena.""" + + def __init__(self, pool, offset, nbytes): + self._pool = pool + self._offset = offset + self.nbytes = nbytes + self.ptr = pool.arena.data_ptr() + offset + self._views: list[weakref.ref] = [] + self.returned = False + + @property + def buffer(self): + # Mooncake hands out a fresh view per access and counts it as an export. + view = self._pool.arena_bytes[self._offset : self._offset + max(self.nbytes, 1)] + self._views.append(weakref.ref(view)) + return view + + def has_live_views(self): + return any(ref() is not None for ref in self._views) + + def release(self): + if self.has_live_views(): + raise RuntimeError("cannot release buffer while exported views exist") + self._pool._give_back(self._offset) + self.returned = True + + +class FakePool: + """Sub-allocates one arena, the way mooncake carves up its registered local buffer.""" + + def __init__(self, capacity, max_regions=None): + self.capacity = capacity + self.max_regions = max_regions + self.arena = torch.empty(max(capacity, 1), dtype=torch.uint8) + self.arena_bytes = self.arena.numpy() + self.acquired: list[int] = [] + self.offsets: list[int] = [] + self.returned = 0 + self.active: dict[int, int] = {} + self.peak_bytes = 0 + self.peak_regions = 0 + self.closed = False + self._lock = threading.Lock() + + def acquire(self, nbytes, block=True): + with self._lock: + if self.closed: + raise RuntimeError("buffer pool is closed") + size = max(nbytes, 1) + if size > self.capacity: + raise RuntimeError("requested buffer size exceeds pool capacity") + if self.max_regions is not None and len(self.active) >= self.max_regions: + raise RuntimeError("buffer pool is exhausted") + offset = self._first_fit(size) + if offset is None: + raise RuntimeError("buffer pool is exhausted") + self.active[offset] = size + self.acquired.append(nbytes) + self.offsets.append(offset) + self.peak_bytes = max(self.peak_bytes, sum(self.active.values())) + self.peak_regions = max(self.peak_regions, len(self.active)) + return FakeBuffer(self, offset, nbytes) + + def _first_fit(self, size): + cursor = 0 + for offset in sorted(self.active): + if offset - cursor >= size: + return cursor + cursor = -(-(offset + self.active[offset]) // ALIGN) * ALIGN + return cursor if self.capacity - cursor >= size else None + + def _give_back(self, offset): + with self._lock: + del self.active[offset] + self.returned += 1 + + def close(self): + with self._lock: + if self.active: + raise RuntimeError("cannot close buffer pool with active leases") + self.closed = True + + +class PoolFactory: + """Stands in for mooncake.store.BufferPool, recording how the client configured it.""" + + def __init__(self, capacity=None): + self._capacity = capacity + self.kwargs: dict | None = None + self.pool: FakePool | None = None + + def __call__(self, store, max_bytes=0, max_regions=None): + self.kwargs = {"max_bytes": max_bytes, "max_regions": max_regions} + self.pool = FakePool(max_bytes if self._capacity is None else self._capacity, max_regions) + return self.pool + + +@pytest.fixture +def store(monkeypatch): + fake = FakeStore() + monkeypatch.setattr(mcc, "MOONCAKE_STORE_IMPORTED", True) + # raising=False: these symbols are absent unless mooncake is installed. + monkeypatch.setattr(mcc, "MooncakeDistributedStore", lambda: fake, raising=False) + monkeypatch.setattr(mcc, "ReplicateConfig", type("ReplicateConfig", (), {}), raising=False) + return fake + + +def make_client(local_buffer_size=1 << 30, **extra): + config = { + "local_hostname": "127.0.0.1", + "metadata_server": "127.0.0.1:8080", + "master_server_address": "127.0.0.1:8081", + "local_buffer_size": local_buffer_size, + # The pool only runs on RDMA; default to it so the pooled-path tests exercise it. + "protocol": "rdma", + } + config.update(extra) + return mcc.MooncakeStoreClient(config) + + +def install_pool(monkeypatch, capacity=None): + """Make the client believe mooncake provides a buffer pool, and hand it a fake one.""" + factory = PoolFactory(capacity) + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", True) + monkeypatch.setattr(mcc, "BufferPool", factory, raising=False) + return factory + + +def make_client_and_pool(monkeypatch, local_buffer_size=1 << 30, capacity=None, **extra): + factory = install_pool(monkeypatch, capacity) + client = make_client(local_buffer_size, **extra) + return client, factory.pool + + +def read(client, keys, shapes, dtypes): + tensors, indexes = client._get_tensors_thread_worker(keys, shapes, dtypes, list(range(len(keys)))) + assert indexes == list(range(len(keys))) + return tensors + + +def read_all(client): + return read(client, KEYS, SHAPES, DTYPES) + + +def assert_payloads(tensors): + for got, want in zip(tensors, expected_tensors(), strict=True): + assert got.dtype == want.dtype + assert got.shape == want.shape + assert torch.equal(got, want) + + +def test_reads_land_in_pooled_buffer(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + tensors = read_all(client) + + # Mixed dtypes and a scalar all round-trip through the staged copy-out. + assert_payloads(tensors) + # The whole point: no registration on the data path, and no buffer left on loan. + assert store.registered == [] + assert len(pool.acquired) == 1 and pool.returned == 1 and pool.active == {} + + +def test_pool_budget_is_capped_by_the_registered_local_buffer(store, monkeypatch): + factory = install_pool(monkeypatch) + client = make_client(local_buffer_size=1 << 20) + + # max_bytes=0 would let mooncake register fresh memory per lease once the local + # buffer is full, which is exactly the cost this path removes. + assert factory.kwargs == {"max_bytes": 1 << 20, "max_regions": mcc.MAX_BATCH_WORKER_THREADS} + # Half of an equal per-thread share, leaving the store room for its own staging. + assert client._pool_share_bytes == (1 << 20) // (2 * mcc.MAX_BATCH_WORKER_THREADS) + + +def test_pool_is_disabled_on_tcp(store, monkeypatch): + # TCP has no registration cost to save, so the pool would only add the copy-out. + install_pool(monkeypatch) + client = make_client(protocol="tcp") + + assert client._buffer_pool is None + tensors = read_all(client) + + assert_payloads(tensors) + assert store.registered and len(store.unregistered) == len(store.registered) + + +def test_batch_larger_than_share_is_read_in_rounds(store, monkeypatch): + # Small local buffer: each thread's share holds only part of the batch. + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + + tensors = read_all(client) + + assert_payloads(tensors) + assert len(pool.acquired) > 1 + assert all(nbytes <= client._pool_share_bytes for nbytes in pool.acquired) + assert pool.returned == len(pool.acquired) + assert store.registered == [] + + +def test_tensor_larger_than_share_registers_its_own_buffer(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + keys, shapes, dtypes = ["big", "small"], [(256,), (4,)], [torch.float32, torch.float32] + payloads = [torch.arange(256, dtype=torch.float32), torch.arange(4, dtype=torch.float32)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, shapes, dtypes) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + # The oversized tensor must not lease the headroom the other readers need; the + # small one still takes the pooled path. + assert pool.acquired and all(nbytes <= client._pool_share_bytes for nbytes in pool.acquired) + assert store.registered and len(store.unregistered) == len(store.registered) + assert pool.active == {} + + +def test_fallback_registers_everything_in_one_pass(store, monkeypatch): + # Pool exhausted up front: every group falls back, but a single same-dtype batch must + # register one merged region, not one per group (the pre-pool cost, not N times it). + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + held = pool.acquire(pool.capacity) + + n = 8 + keys = [f"f{i}" for i in range(n)] + payloads = [torch.arange(i, i + 32, dtype=torch.float32) for i in range(n)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, [(32,)] * n, [torch.float32] * n) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + assert len(store.registered) == 1 and len(store.unregistered) == 1 + held.release() + + +def test_falls_back_to_own_buffers_when_pool_is_exhausted(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + # Mooncake's own staging holds the arena, so acquire() raises instead of returning None. + held = pool.acquire(pool.capacity) + + tensors = read_all(client) + + assert_payloads(tensors) + assert pool.acquired == [pool.capacity] + assert store.registered and len(store.unregistered) == len(store.registered) + held.release() + + +def test_exhaustion_warns_once(store, monkeypatch, caplog): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + held = pool.acquire(pool.capacity) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + read_all(client) + read_all(client) + + # Every group of every read falls back here; the hot path must not flood the log. + fallbacks = [r for r in caplog.records if "Falling back to registering" in r.message] + assert len(fallbacks) == 1 + held.release() + + +def test_registers_receive_regions_without_pool_support(store, monkeypatch): + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + + tensors = read_all(make_client()) + + assert_payloads(tensors) + assert store.registered and len(store.unregistered) == len(store.registered) + + +def test_uniform_group_copied_in_one_strided_pass(store, monkeypatch): + # Many identical small tensors (the fragmented-read case) take the single strided + # copy-out path instead of a per-tensor loop; the payloads must still round-trip. + client, pool = make_client_and_pool(monkeypatch) + n = 8 + keys = [f"u{i}" for i in range(n)] + payloads = [torch.arange(i, i + 16, dtype=torch.float32) for i in range(n)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, [(16,)] * n, [torch.float32] * n) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + assert pool.acquired and pool.returned == len(pool.acquired) + assert store.registered == [] + + +def test_returned_tensors_survive_buffer_reuse(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + + first = read_all(client) + rounds = len(pool.acquired) + second = read_all(client) + + # The second read gets the same arena offsets back, so the first read's tensors + # prove the copy-out really handed ownership to the caller. + assert pool.offsets[rounds:] == pool.offsets[:rounds] + assert pool.returned == len(pool.acquired) + assert_payloads(first) + assert_payloads(second) + + +def test_failed_transfer_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + store.transfer_error = RuntimeError("transport is down") + + with pytest.raises(RuntimeError, match="transport is down"): + read_all(client) + + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_failed_copy_out_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + def boom(self, source, *args, **kwargs): + # Drop the staged view the caller passed in: a real copy_ failure comes from C++ + # and leaves no Python frame holding it, so only the client's own reference counts. + del source + raise RuntimeError("copy-out failed") + + monkeypatch.setattr(torch.Tensor, "copy_", boom) + + # A staged view left alive by the failing copy would make release() raise and + # replace the real cause, so the view has to be dropped on the error path too. + with pytest.raises(RuntimeError, match="copy-out failed"): + read_all(client) + + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_concurrent_readers_stay_within_the_local_buffer(store, monkeypatch): + local_buffer_size = 8192 + readers = 8 + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=local_buffer_size) + barrier = threading.Barrier(readers) + # Hold every reader inside its transfer, so all of them need memory at the same time. + store.on_transfer = lambda: barrier.wait(timeout=30) + results: list[list[torch.Tensor]] = [None] * readers # type: ignore[list-item] + failures: list[Exception] = [] + + def reader(slot): + try: + results[slot] = read_all(client) + except Exception as e: # reported after the join below + failures.append(e) + + threads = [threading.Thread(target=reader, args=(i,)) for i in range(readers)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not failures + for tensors in results: + assert_payloads(tensors) + # Readers past the region cap fall back instead of eating into the store's headroom. + assert pool.peak_regions <= mcc.MAX_BATCH_WORKER_THREADS + assert pool.peak_bytes <= local_buffer_size // 2 + assert pool.active == {} + + +@pytest.mark.parametrize( + ("config", "warns"), + [ + ({"protocol": "rdma"}, True), + ({"protocol": "tcp"}, False), + ({"protocol": "rdma", "use_gdr": True}, False), + ], +) +def test_missing_pool_support_warns_only_for_cpu_rdma_reads(store, monkeypatch, caplog, config, warns): + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + # GDR staging defers cudaMalloc, so it can stand in without a CUDA device here. + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + make_client(**config) + + assert any("BufferPool is unavailable" in r.message for r in caplog.records) is warns + + +def test_close_returns_the_pool_before_the_store(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + read_all(client) + client.close() + + assert pool.closed and store.closed and client._buffer_pool is None + + +def test_close_keeps_the_store_open_while_a_buffer_is_out(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + held = pool.acquire(1024) + + # The store owns the memory the buffer points into, so a refused pool close must not + # tear it down underneath an active lease. + with pytest.raises(RuntimeError, match="cannot close buffer pool with active leases"): + client.close() + + assert not store.closed and client._buffer_pool is pool + held.release() + client.close() + assert pool.closed and store.closed + + +def test_empty_tensor_reads_like_the_register_path(store, monkeypatch): + # mooncake rejects a 0-byte put, so this is unreachable in practice; the guard only keeps + # the pooled path from raising where the register path returns, should it ever be hit. + keys, shapes, dtypes = ["empty", "normal"], [(0,), (4,)], [torch.float32, torch.float32] + payloads = [torch.empty(0), torch.arange(4, dtype=torch.float32)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + registered = read(make_client(), keys, shapes, dtypes) + client, pool = make_client_and_pool(monkeypatch) + pooled = read(client, keys, shapes, dtypes) + + for got, want in zip(pooled, registered, strict=True): + assert got.shape == want.shape and torch.equal(got, want) + assert pool.returned == len(pool.acquired) diff --git a/tests/test_mooncake_client_thread_context.py b/tests/test_mooncake_client_thread_context.py new file mode 100644 index 00000000..01ce3808 --- /dev/null +++ b/tests/test_mooncake_client_thread_context.py @@ -0,0 +1,105 @@ +# 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. + +"""Tests for the per-thread NPU device binding used by MooncakeStoreClient. + +Mooncake's Ascend transport is bound to a thread-local ACL context, but +register/unregister run on per-call ``ThreadPoolExecutor`` workers. These tests +stay hardware-free by faking the accelerator. +""" + +from transfer_queue.storage.clients import mooncake_client as mc + + +class _FakeNpu: + def __init__(self, current=5, initialized=True): + self._current = current + self._initialized = initialized + self.set_devices = [] + self.current_device_calls = 0 + + def is_initialized(self): + return self._initialized + + def current_device(self): + self.current_device_calls += 1 + return self._current + + def set_device(self, device): + self.set_devices.append(device) + + +def _new_client(npu_device): + client = object.__new__(mc.MooncakeStoreClient) + client._npu_device = npu_device + return client + + +def test_resolve_npu_device_captures_owning_thread_device(monkeypatch): + fake_npu = _FakeNpu(current=3) + monkeypatch.setattr(mc, "NPU_IMPORTED", True) + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + assert mc._resolve_npu_device() == 3 + assert fake_npu.current_device_calls == 1 + + +def test_resolve_npu_device_none_without_npu(monkeypatch): + monkeypatch.setattr(mc, "NPU_IMPORTED", False) + + assert mc._resolve_npu_device() is None + + +def test_resolve_npu_device_none_when_not_initialized(monkeypatch): + fake_npu = _FakeNpu(initialized=False) + monkeypatch.setattr(mc, "NPU_IMPORTED", True) + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + assert mc._resolve_npu_device() is None + assert fake_npu.current_device_calls == 0 + + +def test_worker_binds_device_captured_at_construction(monkeypatch): + # A context-less worker reports device 0; the client captured 3 in __init__. + fake_npu = _FakeNpu(current=0) + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + _new_client(npu_device=3)._bind_npu_device() + + assert fake_npu.set_devices == [3] + assert fake_npu.current_device_calls == 0 + + +def test_worker_bind_noop_without_npu_device(monkeypatch): + fake_npu = _FakeNpu() + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + + _new_client(npu_device=None)._bind_npu_device() + + assert fake_npu.set_devices == [] + + +def test_worker_bind_logs_and_swallows_failure(monkeypatch): + class _FailingNpu(_FakeNpu): + def set_device(self, device): + raise RuntimeError("set_device failed") + + fake_npu = _FailingNpu() + monkeypatch.setattr(mc.torch, "npu", fake_npu, raising=False) + warnings = [] + monkeypatch.setattr(mc.logger, "warning", lambda *a, **k: warnings.append((a, k))) + + _new_client(npu_device=3)._bind_npu_device() # must not raise + + assert warnings and warnings[0][1].get("exc_info") is True diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index e36063ee..a02a9b9c 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -13,24 +13,61 @@ # 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 gc +import pickle import time +import weakref +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.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket +from transfer_queue.storage.simple_storage import ( + HybridStorageUnitData, + SimpleStorageUnit, + SSDFileStore, + StorageKeyNotFoundError, + _SSDValueRef, +) +from transfer_queue.utils.serial_utils import calc_packed_size, decode, encode, pack_into, unpack_from +from transfer_queue.utils.zmq_utils import ( + STORAGE_MANAGER_IDENTITY_PREFIX, + ZMQMessage, + ZMQRequestType, + create_zmq_socket, +) class MockStorageClient: """Mock client for testing storage unit operations.""" - def __init__(self, storage_put_get_address, storage_ip): + def __init__(self, storage_put_get_address, storage_ip, identity=None): self.context = zmq.Context() - self.socket = create_zmq_socket(self.context, zmq.DEALER, storage_ip) + # The storage proxy drops identities outside STORAGE_CLIENT_IDENTITY_PREFIXES, so mirror + # the prefix a real storage manager uses instead of letting ZMQ auto-assign one. + if identity is None: + identity = f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}".encode() + self.socket = create_zmq_socket(self.context, zmq.DEALER, storage_ip, identity=identity) self.socket.setsockopt(zmq.RCVTIMEO, 5000) # 5 second timeout self.socket.connect(storage_put_get_address) @@ -102,7 +139,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()) @@ -457,6 +496,371 @@ 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 + assert storage.get_data(["unsupported"], [1])["unsupported"][0] is large_sparse + + 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() + + +@pytest.mark.parametrize("kind", ["tensor", "numpy"]) +@pytest.mark.parametrize("layout", ["packed", "frame-native"]) +@pytest.mark.parametrize("nested", [False, True]) +def test_hybrid_small_values_release_shared_receive_buffer(tmp_path, kind, layout, nested): + storage = HybridStorageUnitData(None, str(tmp_path), "test-run", "test-unit", threshold_bytes=1024) + large = np.arange(1024, dtype=np.float32).reshape(2, 512) + small = np.arange(4, dtype=np.float32).reshape(2, 2) + data = {"large": large, "small": small} + if kind == "tensor": + data = {field: torch.from_numpy(value) for field, value in data.items()} + if nested: + data["small"] = [{"value": [value]} for value in data["small"]] + encoded = encode(data) + if layout == "packed": + owner = np.empty(calc_packed_size(encoded), dtype=np.uint8) + pack_into(owner, encoded) + frames = unpack_from(owner) + else: + owner = np.concatenate([np.frombuffer(frame, dtype=np.uint8) for frame in encoded]) + frames = [] + offset = 0 + for frame in encoded: + size = memoryview(frame).nbytes + frames.append(memoryview(owner)[offset : offset + size]) + offset += size + owner_ref = weakref.ref(owner) + try: + storage.put_data(decode(frames), [0, 1]) + del frames, owner + gc.collect() + + assert storage.ssd_active_values == 2 + assert owner_ref() is None + result = storage.get_data(["large", "small"], [0, 1]) + for field in data: + for position, value in enumerate(result[field]): + expected = data[field][position] + if nested and field == "small": + value = value["value"][0] + expected = expected["value"][0] + if kind == "tensor": + torch.testing.assert_close(value, expected) + else: + np.testing.assert_array_equal(value, expected) + storage.clear([0, 1]) + assert storage.ssd_active_values == 0 + finally: + storage.close() + + +def test_hybrid_encoding_fallback_keeps_large_tensor_without_cloning(tmp_path, monkeypatch): + storage = HybridStorageUnitData(None, str(tmp_path), "test-run", "test-unit", threshold_bytes=64) + value = torch.arange(128, dtype=torch.float32) + monkeypatch.setattr(storage, "_sample_from_value", lambda sample: None) + try: + storage.put_data({"value": [value]}, [0]) + assert storage.get_data(["value"], [0])["value"][0] is value + assert storage.ssd_fallback_values_total == 1 + assert storage.ssd_active_values == 0 + finally: + storage.close() + + +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_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, + 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) + structured = np.zeros((2, 8), dtype=[("index", "= logging.ERROR] == [] + + +@pytest.mark.asyncio +async def test_key_not_found_reply_is_not_logged_as_error_by_caller(caplog): + manager = object.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "TQ_STORAGE_test" + manager.storage_unit_infos = {} + manager.payload_transfer = manager_module.create_payload_transfer(None) + + class GetErrorSocket: + async def send_multipart(self, *args, **kwargs): + pass + + async def recv_multipart(self, **kwargs): + return ZMQMessage.create( + request_type=ZMQRequestType.GET_ERROR, + sender_id="su0", + body={ + "message": f"Failed to get data from storage unit id #su0, detail error message: " + f"{KEY_NOT_FOUND_MARKER}: key 1 not found" + }, + ).serialize() + + with caplog.at_level(logging.DEBUG, logger="transfer_queue"): + with pytest.raises(StorageKeyNotFoundError): + await AsyncSimpleStorageManager._get_from_single_storage_unit.__wrapped__( + manager, + [1], + ["log_probs"], + target_storage_unit="su0", + socket=GetErrorSocket(), + ) + + errors = [ + record + for record in caplog.records + if record.name.startswith("transfer_queue") and record.levelno >= logging.ERROR + ] + assert errors == [], [record.getMessage() for record in errors] diff --git a/tests/test_storage_request_retry.py b/tests/test_storage_request_retry.py new file mode 100644 index 00000000..77fbb137 --- /dev/null +++ b/tests/test_storage_request_retry.py @@ -0,0 +1,373 @@ +# 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. + +"""Storage-unit retry and timeout-diagnosis tests.""" + +import logging +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest +import torch +import zmq +from tensordict import TensorDict + +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers import simple_storage_manager as ssm +from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager, StorageUnitTimeout +from transfer_queue.storage.payload_transfer import ( + PayloadTransferError, + PayloadTransferTimeout, + create_payload_transfer, +) +from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import NixlError +from transfer_queue.utils import common +from transfer_queue.utils.enum_utils import Role +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo + + +def _manager(with_unit: bool = True) -> AsyncSimpleStorageManager: + manager = AsyncSimpleStorageManager.__new__(AsyncSimpleStorageManager) + manager.storage_manager_id = "TQ_STORAGE_test" + manager.storage_unit_infos = {} + if with_unit: + manager.storage_unit_infos["unit_a"] = ZMQServerInfo( + role=Role.STORAGE, + id="unit_a", + ip="10.0.0.7", + ports={"put_get_socket": 5555}, + ) + manager.close = lambda: None + return manager + + +def _single_sample_batch() -> tuple[TensorDict, BatchMeta]: + metadata = BatchMeta( + global_indexes=[0], + partition_ids=["0"], + field_schema={"input_ids": {"dtype": torch.int64, "shape": (2,), "is_nested": False, "is_non_tensor": False}}, + production_status=np.ones(1, dtype=np.int8), + ) + return TensorDict({"input_ids": torch.zeros(1, 2, dtype=torch.int64)}, batch_size=1), metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data_parser, expected_attempts", [(None, 3), (lambda data: data, 1)]) +async def test_put_retries_only_without_a_parser(data_parser, expected_attempts): + manager = _manager() + manager.notify_data_update = AsyncMock() + manager._put_to_single_storage_unit = AsyncMock(side_effect=StorageUnitTimeout("no answer")) + data, metadata = _single_sample_batch() + + with ( + patch.object(ssm, "TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3), + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + pytest.raises(StorageUnitTimeout), + ): + await manager.put_data(data, metadata, data_parser=data_parser) + + assert manager._put_to_single_storage_unit.await_count == expected_attempts + manager.notify_data_update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lost_request_recovers_without_diagnosis(): + manager = _manager() + request = AsyncMock(side_effect=[StorageUnitTimeout("no answer"), "payload"]) + + with patch.object(manager, "_diagnose_storage_unit") as diagnose: + result = await manager._request_with_retry("get", "unit_a", "samples=4", request) + + assert result == "payload" + assert request.await_count == 2 + diagnose.assert_not_awaited() + + +class _ReplySocket: + def __init__(self, reply, timeout_on=None): + self.reply = reply + self.timeout_on = timeout_on + + async def send_multipart(self, *args, **kwargs): + if self.timeout_on == "send": + raise zmq.Again() + + async def recv_multipart(self, *args, **kwargs): + if self.timeout_on == "receive": + raise zmq.Again() + return self.reply.serialize() + + +async def _payload_request(manager, operation, socket): + if operation == "put": + return await manager._put_to_single_storage_unit.__wrapped__( + manager, [0], {"value": [42]}, "unit_a", socket=socket + ) + return await manager._get_from_single_storage_unit.__wrapped__(manager, [0], ["value"], "unit_a", socket=socket) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("timeout_on", ["send", "receive"]) +async def test_zmq_timeout_retry_recovers_without_error_log(operation, timeout_on, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + response = ZMQMessage.create( + request_type=ZMQRequestType[f"{operation.upper()}_DATA_RESPONSE"], + sender_id="unit_a", + body={"data": {"value": [42]}}, + ) + sockets = iter([_ReplySocket(None, timeout_on), _ReplySocket(response)]) + + async def attempt(): + return await _payload_request(manager, operation, next(sockets)) + + with patch.object(manager, "_diagnose_storage_unit") as diagnose, caplog.at_level(logging.WARNING): + result = await manager._request_with_retry(operation, "unit_a", "samples=1", attempt, max_attempts=2) + + assert result == (None if operation == "put" else (["value"], {"value": [42]})) + diagnose.assert_not_awaited() + assert [record for record in caplog.records if record.levelno >= logging.ERROR] == [] + assert sum("retry 2/2" in record.getMessage() for record in caplog.records) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("timeout_on", ["send", "receive"]) +async def test_zmq_exhausted_timeout_logs_once_and_preserves_cause(operation, timeout_on, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + attempts = 0 + + async def attempt(): + nonlocal attempts + attempts += 1 + return await _payload_request(manager, operation, _ReplySocket(None, timeout_on)) + + with ( + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis") as diagnose, + caplog.at_level(logging.WARNING), + pytest.raises(StorageUnitTimeout) as raised, + ): + await manager._request_with_retry(operation, "unit_a", "samples=1", attempt, max_attempts=2) + + assert attempts == 2 + diagnose.assert_awaited_once() + cause = raised.value.__cause__ + assert isinstance(cause, PayloadTransferTimeout) + assert isinstance(cause.__cause__, zmq.Again) + assert (cause.wire_bytes > 0) == (operation == "put") + errors = [record for record in caplog.records if record.levelno >= logging.ERROR] + assert len(errors) == 1 + assert "failed after 2 attempts" in errors[0].getMessage() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +async def test_zmq_peer_error_is_logged_once_without_retry(operation, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + response = ZMQMessage.create( + request_type=ZMQRequestType[f"{operation.upper()}_ERROR"], + sender_id="unit_a", + body={"message": "disk full"}, + ) + attempts = 0 + + async def attempt(): + nonlocal attempts + attempts += 1 + return await _payload_request(manager, operation, _ReplySocket(response)) + + with caplog.at_level(logging.ERROR), pytest.raises(RuntimeError, match="disk full") as raised: + await manager._request_with_retry(operation, "unit_a", "samples=1", attempt, max_attempts=2) + + assert attempts == 1 + assert len([record for record in caplog.records if record.levelno >= logging.ERROR]) == 1 + assert str(raised.value).count("Error in put to storage unit") == (1 if operation == "put" else 0) + assert str(raised.value).count("Error getting data from storage unit") == (1 if operation == "get" else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +async def test_zmq_socket_error_keeps_runtime_error_contract(operation, caplog): + manager = _manager() + manager.payload_transfer = create_payload_transfer(None) + error = zmq.ZMQError("connection failed") + socket = MagicMock() + socket.send_multipart = AsyncMock(side_effect=error) + socket.recv_multipart = AsyncMock() + + with caplog.at_level(logging.ERROR), pytest.raises(RuntimeError, match="connection failed") as raised: + await _payload_request(manager, operation, socket) + + assert raised.value.__cause__ is error + socket.send_multipart.assert_awaited_once() + socket.recv_multipart.assert_not_awaited() + assert len([record for record in caplog.records if record.levelno >= logging.ERROR]) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["put", "get"]) +@pytest.mark.parametrize("error_type", [PayloadTransferError, NixlError]) +async def test_payload_error_is_wrapped_once_without_retry(operation, error_type, caplog): + manager = _manager() + cause = zmq.ZMQError("connection failed") + error = error_type("transfer failed") + error.__cause__ = cause + manager.payload_transfer = MagicMock() + transfer_request = AsyncMock(side_effect=error) + setattr(manager.payload_transfer, operation, transfer_request) + + with ( + patch.object(manager, "_diagnose_storage_unit") as diagnose, + caplog.at_level(logging.ERROR), + pytest.raises(RuntimeError, match="transfer failed") as raised, + ): + await manager._request_with_retry( + operation, "unit_a", "samples=1", lambda: _payload_request(manager, operation, None), max_attempts=2 + ) + + assert type(raised.value) is RuntimeError + assert raised.value.__cause__ is error + assert raised.value.__cause__.__cause__ is cause + transfer_request.assert_awaited_once() + diagnose.assert_not_awaited() + assert len([record for record in caplog.records if record.levelno >= logging.ERROR]) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error, max_attempts, expected_attempts", + [ + (StorageUnitTimeout("no answer"), 3, 3), + (RuntimeError("unit rejected request"), 3, 1), + (StorageUnitTimeout("no answer"), 0, 1), + ], +) +async def test_retry_attempts_are_bounded_and_only_cover_timeouts(error, max_attempts, expected_attempts): + manager = _manager() + request = AsyncMock(side_effect=error) + + with ( + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + pytest.raises(type(error)), + ): + await manager._request_with_retry("get", "unit_a", "samples=4", request, max_attempts=max_attempts) + + assert request.await_count == expected_attempts + + +@pytest.mark.asyncio +async def test_failure_log_has_context_without_repeating_the_unit(caplog): + manager = _manager() + error = StorageUnitTimeout("no answer in 200s during get from storage unit unit_a at 10.0.0.7:5555") + + with ( + patch.object(manager, "_diagnose_storage_unit", return_value="diagnosis"), + caplog.at_level(logging.ERROR), + pytest.raises(StorageUnitTimeout), + ): + await manager._request_with_retry( + "get", "unit_a", "samples=4 fields=['input_ids']", AsyncMock(side_effect=error) + ) + + message = next(record.message for record in caplog.records if "failed after" in record.message) + assert "samples=4 fields=['input_ids']" in message + assert message.count("unit_a") == message.count("10.0.0.7:5555") == 1 + + +@pytest.mark.parametrize( + "elapsed, payload_bytes, should_log", + [(0.01, 2**20, False), (0.01, 512 * 2**20, True), (6.0, 2**20, True)], +) +def test_log_heavy_operation_thresholds(caplog, elapsed, payload_bytes, should_log): + with caplog.at_level(logging.WARNING): + common.log_heavy_operation("TQ_STORAGE_test", "put", elapsed, payload_bytes, "samples=4") + + assert bool(caplog.records) is should_log + if should_log: + assert "serialized_mb=" in caplog.records[0].message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tcp_result, probe_result, expected", + [ + # A probe that answers proves only that the unit is serving again, whatever the + # cumulative counters happen to say. + ((None, None), {"active_keys": 4, "arrivals_by_op": {"GET_DATA": 0}}, "unit_serving_again"), + ((None, None), {"active_keys": 4, "arrivals_by_op": {"GET_DATA": 9}}, "unit_serving_again"), + ((None, None), {"active_keys": 4}, "unit_serving_again"), + ((None, None), zmq.error.Again(), "unit_not_serving"), + (ConnectionRefusedError(), zmq.error.Again(), "tcp=down(ConnectionRefusedError)"), + ], +) +async def test_diagnosis_classifies_the_failure(tcp_result, probe_result, expected): + manager = _manager() + writer = MagicMock() + tcp = AsyncMock( + side_effect=tcp_result if isinstance(tcp_result, Exception) else None, + return_value=(None, writer), + ) + probe = AsyncMock( + side_effect=probe_result if isinstance(probe_result, Exception) else None, + return_value=probe_result, + ) + + with patch.object(ssm.asyncio, "open_connection", tcp), patch.object(manager, "_probe_storage_unit", probe): + diagnosis = await manager._diagnose_storage_unit("unit_a") + + assert expected in diagnosis + + +def test_diagnosis_does_not_report_empty_op_stats_as_zero_traffic(): + """op_stats is Prometheus-gated, so an empty dict must not read as 'served nothing'.""" + described = ssm._describe_unit_state({"requests_arrived": 7, "active_keys": 1}) + + assert "completed=unavailable(prometheus_disabled)" in described + assert "completed={}" not in described + + +@pytest.mark.parametrize( + "arrivals", + [ + # Nine historical GETs, all already completed: they say nothing about this request. + {"GET_DATA": 9}, + # The unit restarted during the timeout window, so its counters began again at zero. + {"GET_DATA": 0}, + # The request arrived and finished after the caller gave up waiting. + {"GET_DATA": 10}, + ], +) +def test_cumulative_counters_never_become_a_claim_about_this_request(arrivals): + """Counters are per-op and cumulative, so no value of them locates one request. + + They carry no request id, never decrease, and reset when the unit restarts, so both + directions of inference are unsound: a nonzero count may be entirely historical, and a + zero one may mean the request reached a previous process. + """ + described = ssm._describe_unit_state({"requests_arrived": sum(arrivals.values()), "arrivals_by_op": arrivals}) + + assert "verdict=unit_serving_again" in described + assert "arrived_but_unfinished" not in described + assert "request_lost_in_flight" not in described + # The raw counters stay in the line as triage input, just not as a verdict. + assert "arrivals_by_op=" in described + + +@pytest.mark.asyncio +async def test_diagnosis_handles_an_unknown_unit(): + assert "unit_not_registered" in await _manager(with_unit=False)._diagnose_storage_unit("missing") diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 2bd9a6c7..127e4d56 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -17,11 +17,13 @@ with_zmq_socket used to create and term() a context per RPC call, which churned libzmq signaler file descriptors and crashed the process under concurrency (Bad file descriptor --> SIGABRT). It now reuses the owner's context and only creates the socket per call, so +-> SIGABRT). It now reuses the owner's context and leases sockets from a shared pool, so these tests assert every concurrent call sees the SAME context, alive until close(). +Pool-specific behavior is covered in test_zmq_socket_pool.py. """ import asyncio +from contextlib import contextmanager from threading import Thread from unittest.mock import patch @@ -138,8 +140,10 @@ def _spy_create(ctx, *args, **kwargs): assert len(results) == num_calls assert all(isinstance(meta, BatchMeta) for meta in results) - # Every call must have used the SAME context, and it must be the client's context. - assert len(seen_contexts) == num_calls + # Every socket must come from the SAME context, and it must be the client's. The count + # is bounded by the pool's cap rather than by the call count, since concurrency past it + # waits for a socket instead of opening one. + assert seen_contexts, "no socket was created at all" assert all(ctx is client.zmq_context for ctx in seen_contexts) # The shared context must NOT have been terminated by any call. assert not client.zmq_context.closed @@ -168,6 +172,83 @@ def test_client_rejects_invalid_context_pool_size(echo_controller): ) +@contextmanager +def _no_context_left_open(): + """Assert every ZMQ context built inside the block is closed by the time it exits. + + __init__ arms the finalizer only on success, so an invalid argument must either be + rejected before the context is built or destroy it on the way out -- otherwise the + context and its native I/O threads leak for the process lifetime. + """ + created = [] + real_context = zmq.asyncio.Context + + def _spy(*args, **kwargs): + ctx = real_context(*args, **kwargs) + created.append(ctx) + return ctx + + with patch("zmq.asyncio.Context", side_effect=_spy): + try: + yield created + finally: + leaked = [ctx for ctx in created if not ctx.closed] + for ctx in leaked: + ctx.destroy(linger=0) + assert not leaked, f"{len(leaked)} ZMQ context(s) left open on the failure path" + + +def test_client_rejects_invalid_socket_pool_size(echo_controller): + """A bad TQ_SOCKET_POOL_SIZE must name the variable, not silently disable reuse. + + Below 1 nothing is ever parked, so every request pays a fresh connect while the client + still looks pooled. + """ + for bad in (-1, 0): + with patch("transfer_queue.client.TQ_SOCKET_POOL_SIZE", bad): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="TQ_SOCKET_POOL_SIZE must be at least 1"): + AsyncTransferQueueClient( + client_id="client_invalid_socket_pool", + controller_info=echo_controller.zmq_server_info, + ) + assert created == [], "checkable without a context, so none should be built" + + +def test_invalid_max_sockets_does_not_leak_a_context(echo_controller): + """No invalid max-sockets input may leave the context or its I/O threads behind.""" + # Rejectable without a live context. + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "not-a-number"): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer"): + AsyncTransferQueueClient( + client_id="client_garbage_max_sockets", + controller_info=echo_controller.zmq_server_info, + ) + assert created == [], "parsing needs no context, so none should be built" + + for bad in (0, -5): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="at least 1"): + AsyncTransferQueueClient( + client_id="client_low_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=bad, + ) + assert created == [], "the lower bound needs no context, so none should be built" + + # Above ZMQ_SOCKET_LIMIT: this one genuinely needs a live context, so it must be + # destroyed rather than hoisted. + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="ZMQ_SOCKET_LIMIT"): + AsyncTransferQueueClient( + client_id="client_huge_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=10**9, + ) + assert len(created) == 1, "the limit check requires a built context" + + def test_simple_storage_borrows_client_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_simple_storage_context", @@ -188,6 +269,105 @@ def test_simple_storage_borrows_client_context(echo_controller): client.close() +def test_each_scenario_gets_its_own_pool(echo_controller): + """Controller RPC, storage RPC and notify must each hold a separate pool. + + They share one context but never one pool: each dials a different peer or socket name, + so a shared pool reused nothing while letting one scenario's sockets be swept by + another's. Separate pools keep each scenario's failures to itself. + """ + client = AsyncTransferQueueClient( + client_id="client_scenario_pools", + controller_info=echo_controller.zmq_server_info, + ) + + with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"): + manager = AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, + zmq_context=client.zmq_context, + ) + + pools = [client.controller_rpc_pool, manager.storage_rpc_pool, manager.notify_pool] + assert len({id(pool) for pool in pools}) == 3, "scenarios must not share a pool" + + manager.close() + client.close() + + +@pytest.mark.parametrize( + "pool_size, units, warns", + [ + (8, 4, False), # 32 idle sockets against a 64-socket budget + (8, 32, True), # 256 would exceed it + ], +) +def test_warns_when_pool_size_times_units_exceeds_the_context(echo_controller, caplog, pool_size, units, warns): + """The per-address cap multiplies by storage-unit count; the context ceiling does not. + + At a few thousand units the product passes ZMQ_MAX_SOCKETS, where opening a socket + fails outright, so the mismatch is worth naming at construction rather than at a lease. + """ + client = AsyncTransferQueueClient( + client_id="client_pool_budget", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=64, + ) + zmq_info = { + f"storage_{i}": ZMQServerInfo( + role=Role.STORAGE, + id=f"storage_{i}", + ip="127.0.0.1", + ports={"put_get_socket": 5600 + i}, + ) + for i in range(units) + } + + with ( + patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"), + patch("transfer_queue.storage.managers.simple_storage_manager.TQ_SOCKET_POOL_SIZE", pool_size), + caplog.at_level("WARNING"), + ): + manager = AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": zmq_info}, + zmq_context=client.zmq_context, + ) + + assert ("above this context's ZMQ_MAX_SOCKETS" in caplog.text) is warns + + manager.close() + client.close() + + +def test_close_after_failed_handshake_still_releases_the_context(echo_controller): + """A manager whose base constructor raised must still tear down on close(). + + storage_rpc_pool is assigned after super().__init__(), so a handshake timeout leaves it + unset while the context and its native I/O threads are already allocated. __del__ calls + close() regardless, so dereferencing the pool there aborts the base teardown and leaks + the context. + """ + built = [] + + def _fail(self): + built.append(self) + raise TimeoutError("handshake failed") + + with patch.object(StorageManager, "_connect_to_controller", _fail): + with pytest.raises(TimeoutError): + AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, + ) + + manager = built[0] + assert not hasattr(manager, "storage_rpc_pool"), "the test no longer exercises the partial-build path" + + manager.close() + assert manager.zmq_context.closed, "the context outlived a failed construction" + + def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_borrowed_context_lifecycle", @@ -266,8 +446,11 @@ def test_close_skips_destroy_while_loop_thread_alive(echo_controller): context.destroy(linger=0) -def _make_borrowing_manager(zmq_context): - """A minimal manager that borrows a caller's context, like SimpleStorage does.""" +def _make_borrowing_manager(client=None): + """A minimal manager borrowing *client*'s context, like SimpleStorage does. + + With no client it creates its own context. Either way it builds its own notify pool. + """ class Borrower(StorageManager): def _connect_to_controller(self): @@ -282,7 +465,7 @@ async def get_data(self, *args, **kwargs): async def clear_data(self, *args, **kwargs): return None - return Borrower(None, {}, zmq_context=zmq_context) + return Borrower(None, {}) if client is None else Borrower(None, {}, zmq_context=client.zmq_context) def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller): @@ -295,7 +478,7 @@ def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller) client_id="client_notify_thread_veto", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(client.zmq_context) + client.storage_manager = _make_borrowing_manager(client) context = client.zmq_context with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True): @@ -314,7 +497,7 @@ def test_healthy_notify_thread_does_not_block_destroy(echo_controller): client_id="client_notify_thread_clean", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(client.zmq_context) + client.storage_manager = _make_borrowing_manager(client) client.close() assert client.zmq_context.closed @@ -326,7 +509,7 @@ def test_manager_with_own_context_does_not_veto(echo_controller): client_id="client_independent_manager", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(None) # creates its own context + client.storage_manager = _make_borrowing_manager() # creates its own context assert client.storage_manager.zmq_context is not client.zmq_context # Even a stuck notify thread on an unrelated context must not block the client. diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py new file mode 100644 index 00000000..6af2ebcf --- /dev/null +++ b/tests/test_zmq_socket_pool.py @@ -0,0 +1,390 @@ +# 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. + +"""Tests for ZMQSocketPool's concurrency invariants. + +Sockets used to be created and closed per request. The pool reuses them, which is only +safe because a lease is exclusive and a socket that did not complete a clean send/recv is +discarded: replies carry no request id (see ZMQMessage.create), so a reply left in flight +by a timed-out or cancelled request would be read by the next user of that socket as its +own. test_timed_out_socket_is_not_reused pins exactly that. + +Each role's own reuse is asserted where that role is tested -- the client's controller RPC +in test_client.py, the metrics collector in test_metrics.py, pool wiring and context +lifecycle in test_zmq_shared_context.py. What is left here needs control over reply timing +and loop lifetime that a caller-level test cannot reach. + +Reuse is asserted from the peer's side rather than from the pool's internals: every socket +dials with its own ZMQ identity, so one identity across many requests means the connection +was reused, and a fresh identity means the old socket was discarded. +""" + +import asyncio +import threading +from unittest.mock import patch + +import pytest +import zmq +import zmq.asyncio + +import transfer_queue.utils.zmq_utils as zmq_utils +from transfer_queue.utils.enum_utils import Role +from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool + + +class _Peer: + """A ROUTER that echoes one reply per request, recording who dialled it. + + ``tag`` distinguishes which endpoint answered, for the re-registration test. + """ + + def __init__(self, delay_first_reply: float = 0.0, peer_id: str = "peer_0", tag: bytes = b"reply-to-"): + self.context = zmq.Context() + self.socket = self.context.socket(zmq.ROUTER) + port = self.socket.bind_to_random_port("tcp://127.0.0.1") + self.info = ZMQServerInfo(role=Role.STORAGE, id=peer_id, ip="127.0.0.1", ports={"put_get_socket": port}) + self.identities: list[bytes] = [] + self._tag = tag + self._delay_first_reply = delay_first_reply + self._replies = 0 + self.running = True + self.thread = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + @property + def callers(self) -> int: + """How many distinct sockets have dialled this peer.""" + return len(set(self.identities)) + + def _serve(self): + poller = zmq.Poller() + poller.register(self.socket, zmq.POLLIN) + while self.running: + if not dict(poller.poll(50)): + continue + identity, request = self.socket.recv_multipart() + self.identities.append(identity) + if self._replies == 0 and self._delay_first_reply: + # Reply late enough that the requester has already timed out, leaving this + # reply in flight -- the poisoned-socket scenario. + time_left = self._delay_first_reply + while time_left > 0 and self.running: + sleep = min(0.05, time_left) + threading.Event().wait(sleep) + time_left -= sleep + self._replies += 1 + self.socket.send_multipart([identity, self._tag + request]) + + def stop(self): + self.running = False + self.thread.join(timeout=2.0) + self.socket.close(linger=0) + self.context.term() + + +@pytest.fixture +def peer(): + p = _Peer() + yield p + p.stop() + + +async def _round_trip(pool, peer_info, payload=b"req"): + with pool.lease(peer_info) as sock: + await sock.send_multipart([payload]) + return (await sock.recv_multipart())[0] + + +@pytest.mark.asyncio +async def test_timed_out_socket_is_not_reused(): + """A timed-out request must not leave its socket -- or its late reply -- in the pool. + + Regression guard for the core hazard: with the socket pooled, the *next* request would + receive the previous request's reply, silently attributing one response to another. + """ + # Replies to the first request only after the pool's 1s timeout has expired, so that + # reply is still in flight when the socket would otherwise go to the next caller. The + # peer serves serially, so the delay stays well inside the second request's own timeout. + peer = _Peer(delay_first_reply=1.4) + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", timeout=1) + try: + with pytest.raises(zmq.error.Again): + await _round_trip(pool, peer.info, b"first") + + # The late reply to "first" must not surface as the answer to "second". + assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" + assert peer.callers == 2, "the timed-out socket was handed to the next request" + finally: + pool.close() + ctx.destroy(linger=0) + peer.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["exception", "cancellation"]) +async def test_poisoned_lease_is_discarded(failure): + """A lease that did not complete cleanly must not return its socket to the pool. + + The request is on the wire before the failure and its reply is still in flight, so + parking the socket would hand the next caller a reply belonging to someone else. + Cancellation counts as well as a raise: asyncio.gather cancels its siblings on the + first failure, so this is the routine case rather than an exotic one. + """ + # Answers the first request only after it has been abandoned, which is what leaves the + # stale reply in flight. + peer = _Peer(delay_first_reply=0.6) + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + + async def poisoned(): + with pool.lease(peer.info) as sock: + await sock.send_multipart([b"doomed"]) + # Fail only once the peer has the request, so its reply really is outstanding. + while not peer.identities: + await asyncio.sleep(0.01) + if failure == "exception": + raise RuntimeError("handler blew up") + await asyncio.sleep(60) # cancelled here, after the lease was handed out + + try: + if failure == "exception": + with pytest.raises(RuntimeError): + await poisoned() + else: + task = asyncio.create_task(poisoned()) + while not peer.identities: + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + abandoned = set(peer.identities) + # The reply to "doomed" must not surface as the answer to "second". + assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" + assert set(peer.identities) - abandoned, "the poisoned socket was reused" + finally: + pool.close() + ctx.destroy(linger=0) + peer.stop() + + +def test_sockets_are_not_reused_across_event_loops(peer): + """A socket bound to a finished loop must never be handed to another one. + + pyzmq rebinds an async socket to whatever loop it next sees; one bound to a *closed* + loop is the "Bad file descriptor / SIGABRT" failure this keying exists to prevent. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + + async def lease_twice(tag): + # Twice per loop, so a socket IS reused within a loop -- which is what makes the + # cross-loop comparison meaningful rather than trivially true. + for i in range(2): + assert await _round_trip(pool, peer.info, f"{tag}{i}".encode()) == f"reply-to-{tag}{i}".encode() + + for tag in ("a", "b", "c"): + asyncio.run(lease_twice(tag)) + + assert peer.callers == 3, "one socket per loop, reused within it" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_reregistered_peer_is_not_served_a_stale_socket(): + """A peer that moves to a new address must not be answered by its old endpoint. + + A storage unit or controller restarted or re-registered keeps its id but gets a fresh + port, so keying reuse on the id alone would keep leasing a socket wired to the address + it no longer answers on -- silently talking to a dead or reassigned endpoint. All calls + share one event loop, as a real client or storage manager does; a loop per call would + discard the socket via owner eviction and hide the bug. + """ + old = _Peer(peer_id="su0", tag=b"OLD:") + new = _Peer(peer_id="su0", tag=b"NEW:") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + try: + assert await _round_trip(pool, old.info) == b"OLD:req" + # Same id, different port -- exactly what re-registration produces. + assert await _round_trip(pool, new.info) == b"NEW:req" + # And a peer that moves back is still reachable. + assert await _round_trip(pool, old.info) == b"OLD:req" + finally: + pool.close() + ctx.destroy(linger=0) + old.stop() + new.stop() + + +@pytest.mark.asyncio +async def test_burst_beyond_pool_size_is_served(peer): + """Concurrency above maxsize must not be refused or blocked, only left unparked.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) + + results = await asyncio.gather(*[_round_trip(pool, peer.info, f"c{i}".encode()) for i in range(8)]) + assert sorted(results) == sorted(f"reply-to-c{i}".encode() for i in range(8)) + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_alease_caps_sockets_in_flight(peer): + """alease waits for a socket rather than opening one past maxsize. + + This is what makes the context's socket budget a function of configuration: without it + the peak equals real concurrency, which at a few thousand peers exhausts ZMQ_MAX_SOCKETS + and fails a lease with EMFILE. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=4) + + async def one(i): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([f"c{i}".encode()]) + return (await sock.recv_multipart())[0] + + results = await asyncio.gather(*[one(i) for i in range(40)]) + + assert sorted(results) == sorted(f"reply-to-c{i}".encode() for i in range(40)), "a request was dropped" + assert peer.callers == 4, "concurrency past maxsize opened extra sockets instead of waiting" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_alease_permits_are_per_address(peer): + """A busy peer must not stall requests to a different one. + + Permits are keyed like the buckets, so a fan-out across N peers is not serialized by a + cap meant to bound one peer's concurrency. + """ + other = _Peer(peer_id="peer_1", tag=b"other-") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=1) + + async def hold(): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([b"slow"]) + await sock.recv_multipart() + await asyncio.sleep(0.3) # keeps peer's only permit + + async def other_peer(): + async with pool.alease(other.info) as sock: + await sock.send_multipart([b"req"]) + return (await sock.recv_multipart())[0] + + # Would time out if one peer's permit gated the other. + _, reply = await asyncio.wait_for(asyncio.gather(hold(), other_peer()), timeout=5) + assert reply == b"other-req" + + pool.close() + ctx.destroy(linger=0) + other.stop() + + +@pytest.mark.asyncio +async def test_nested_alease_raises_instead_of_hanging(peer): + """A lease inside another lease must fail loudly rather than wait on a held permit. + + A permit is released when its body ends, so nesting can wait on one the same task -- + or a sibling mid-cycle -- already holds. That hangs with no timeout and no error, the + failure mode hardest to diagnose in production. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=8) + + with pytest.raises(RuntimeError, match="already holding"): + async with pool.alease(peer.info): + async with pool.alease(peer.info): + pass + + # The outer permit is returned, so the pool still works. + async with pool.alease(peer.info) as sock: + await sock.send_multipart([b"after"]) + assert (await sock.recv_multipart())[0] == b"reply-to-after" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_nesting_is_rejected_across_pools(peer): + """Two pools do not make nesting safe: a cycle between them deadlocks just as well. + + Each task holds what the other waits for, and separate semaphores do not break that. + """ + other = _Peer(peer_id="peer_1", tag=b"other-") + ctx = zmq.asyncio.Context() + first = ZMQSocketPool(ctx, "first", "put_get_socket", maxsize=1) + second = ZMQSocketPool(ctx, "second", "put_get_socket", maxsize=1) + + with pytest.raises(RuntimeError, match="already holding"): + async with first.alease(peer.info): + async with second.alease(other.info): + pass + + first.close() + second.close() + ctx.destroy(linger=0) + other.stop() + + +@pytest.mark.asyncio +async def test_consecutive_leases_are_not_nesting(peer): + """Leasing again after the previous body ended is the supported shape.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) + + for i in range(3): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([f"s{i}".encode()]) + assert (await sock.recv_multipart())[0] == f"reply-to-s{i}".encode() + + pool.close() + ctx.destroy(linger=0) + + +def test_pool_size_comes_from_the_env_var_at_construction(): + """Every role's pool takes TQ_SOCKET_POOL_SIZE, resolved per pool rather than at import. + + A default bound in the signature would freeze whatever the environment held when this + module first loaded, which is what makes such a knob look settable but do nothing. + """ + ctx = zmq.Context() + try: + with patch.object(zmq_utils, "TQ_SOCKET_POOL_SIZE", 7): + assert ZMQSocketPool(ctx, "owner", "put_get_socket")._maxsize == 7 + # An explicit argument still wins, which is what lets a caller opt out. + assert ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=3)._maxsize == 3 + finally: + ctx.destroy(linger=0) + + +def test_pool_size_below_one_is_rejected(): + """Below 1 nothing is ever parked, so reuse is off while the pool still looks pooled.""" + ctx = zmq.Context() + try: + with patch.object(zmq_utils, "TQ_SOCKET_POOL_SIZE", 0): + with pytest.raises(ValueError, match="at least 1"): + ZMQSocketPool(ctx, "owner", "put_get_socket") + finally: + ctx.destroy(linger=0) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 01e50b85..0f7f13f1 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -29,9 +29,11 @@ from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + TQ_SOCKET_POOL_SIZE, ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, with_zmq_socket, ) @@ -49,10 +51,8 @@ # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( - "request_handle_socket", - get_identity=lambda self: self.client_id, get_peer=lambda self, target: self._controller, - get_context=lambda self: self.zmq_context, + get_pool=lambda self: self.controller_rpc_pool, ) @@ -61,6 +61,11 @@ class AsyncTransferQueueClient: This client provides async methods for data transfer operations including getting metadata, reading data from storage, writing data to storage, and clearing data. + + Await these methods from one long-lived event loop, as the examples below do. Pooled + sockets are keyed by their owning loop, so a fresh ``asyncio.run()`` per call gets a + fresh socket every time and pays a connect handshake it cannot amortize. Callers with + no loop of their own should use ``TransferQueueClient``, which keeps one internally. """ def __init__( @@ -91,12 +96,17 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info - # One long-lived context per client; sockets stay per-request because ZMQ sockets - # are not thread-safe. + # Check everything that does not need the context first: the finalizer is not armed + # until __init__ returns, so raising after allocation leaks it and its I/O threads. io_threads = TQ_CLIENT_ZMQ_IO_THREADS if zmq_io_threads is None else zmq_io_threads if io_threads < 1: raise ValueError(f"Client ZMQ I/O thread pool size must be at least 1, got {io_threads}") - self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) + if TQ_SOCKET_POOL_SIZE < 1: + # Name the variable: the pool's own error cannot say which knob supplied the value. + raise ValueError( + f"TQ_SOCKET_POOL_SIZE must be at least 1, got {TQ_SOCKET_POOL_SIZE}. " + f"The pool always reuses at least one socket per endpoint; it cannot be disabled." + ) max_sockets = zmq_max_sockets explicitly_requested = max_sockets is not None @@ -109,26 +119,44 @@ def __init__( f"TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer, got {TQ_CLIENT_ZMQ_MAX_SOCKETS!r}" ) from e explicitly_requested = True + if max_sockets is not None and max_sockets < 1: + # The upper bound needs ZMQ_SOCKET_LIMIT, hence a live context, but the lower one + # does not -- so reject it before allocating anything. + raise ValueError(f"Client ZMQ max sockets must be at least 1, got {max_sockets}") if max_sockets is None: max_sockets = DEFAULT_CLIENT_ZMQ_MAX_SOCKETS - socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) - if explicitly_requested: - # A value the caller asked for must not be silently reinterpreted. - if not 1 <= max_sockets <= socket_limit: - raise ValueError( - f"Client ZMQ max sockets must be between 1 and this build's " - f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" + self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) + try: + # ZMQ_SOCKET_LIMIT is a property of the built context, so this last check cannot + # be hoisted above it; destroy the context rather than leak it on failure. + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + if explicitly_requested: + # A value the caller asked for must not be silently reinterpreted. + if max_sockets > socket_limit: + raise ValueError( + f"Client ZMQ max sockets must be between 1 and this build's " + f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" + ) + elif max_sockets > socket_limit: + # Nobody asked for the default, so clamp instead of failing to construct on a + # build whose ZMQ_SOCKET_LIMIT is below it. + logger.debug( + f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " + f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." ) - elif max_sockets > socket_limit: - # Nobody asked for the default, so clamp instead of failing to construct on a - # build whose ZMQ_SOCKET_LIMIT is below it. - logger.debug( - f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " - f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." - ) - max_sockets = socket_limit - self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + max_sockets = socket_limit + self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + except BaseException: + self.zmq_context.destroy(linger=0) + raise + # Reused across requests, so the socket budget above is consumed by the concurrency + # high-water mark rather than by request count. + self.controller_rpc_pool = ZMQSocketPool( + self.zmq_context, + client_id, + "request_handle_socket", + ) # Backstop for a client that is never closed, so the context and its I/O threads do # not leak for the process lifetime. finalize() (not __del__) also runs at @@ -159,7 +187,8 @@ def initialize_storage_manager( The client's long-lived ZMQ context is offered to every backend uniformly; each registered manager decides whether to borrow it or keep its own, so the client - needs no knowledge of specific backend names. + needs no knowledge of specific backend names. Managers build their own pools over + whichever context they end up with, one per request scenario. Args: manager_type: Type of storage manager to create. Supported types include: @@ -236,33 +265,35 @@ async def async_get_meta( RuntimeError: If communication fails or controller returns error response Example: - >>> # Example 1: Basic fetch metadata - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["input_ids", "attention_mask"], - ... batch_size=4, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences" - ... )) - >>> print(batch_meta.is_ready) # True if all samples ready - >>> - >>> # Example 2: Fetch with self-defined samplers (using GRPOGroupNSampler as an example) - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["input_ids", "attention_mask"], - ... batch_size=8, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> print(batch_meta.is_ready) # True if all samples ready - >>> - >>> # Example 3: Force fetch metadata (bypass production status check and Sampler, - >>> # so may include unready and already-consumed samples. No filtering by consumption status is applied.) - >>> batch_meta = asyncio.run(client.async_get_meta( - ... partition_id="train_0", # optional - ... mode="force_fetch", - ... )) - >>> print(batch_meta.is_ready) # May be False if some samples not ready + >>> async def main(): + ... # Example 1: Basic fetch metadata + ... batch_meta = await client.async_get_meta( + ... data_fields=["input_ids", "attention_mask"], + ... batch_size=4, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences" + ... ) + ... print(batch_meta.is_ready) # True if all samples ready + ... + ... # Example 2: Fetch with self-defined samplers (using GRPOGroupNSampler as an example) + ... batch_meta = await client.async_get_meta( + ... data_fields=["input_ids", "attention_mask"], + ... batch_size=8, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... print(batch_meta.is_ready) # True if all samples ready + ... + ... # Example 3: Force fetch metadata (bypass production status check and Sampler, + ... # so may include unready and already-consumed samples. No filtering by + ... # consumption status is applied.) + ... batch_meta = await client.async_get_meta( + ... partition_id="train_0", # optional + ... mode="force_fetch", + ... ) + ... print(batch_meta.is_ready) # May be False if some samples not ready """ response_msg = await self._request_controller( socket=socket, @@ -302,10 +333,11 @@ async def async_set_custom_meta( RuntimeError: If communication fails or controller returns error response Example: - >>> # Create batch with custom metadata - >>> batch_meta = client.get_meta(data_fields=["input_ids"], batch_size=4, ...) - >>> batch_meta.update_custom_meta([{"score": 0.9}, {"score": 0.8}]) - >>> asyncio.run(client.async_set_custom_meta(batch_meta)) + >>> async def main(): + ... # Create batch with custom metadata + ... batch_meta = await client.async_get_meta(data_fields=["input_ids"], batch_size=4, ...) + ... batch_meta.update_custom_meta([{"score": 0.9}, {"score": 0.8}]) + ... await client.async_set_custom_meta(batch_meta) """ assert socket is not None @@ -387,30 +419,34 @@ async def async_put( >>> batch_size = 4 >>> seq_len = 16 >>> current_partition_id = "train_0" - >>> # Example 1: Normal usage with existing metadata - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["prompts", "attention_mask"], - ... batch_size=batch_size, - ... partition_id=current_partition_id, - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> batch = asyncio.run(client.async_get_data(batch_meta)) - >>> output = TensorDict({"response": torch.randn(batch_size, seq_len)}) - >>> asyncio.run(client.async_put(data=output, metadata=batch_meta)) - >>> - >>> # Example 2: Initial data insertion without pre-existing metadata - >>> # BE CAREFUL: this usage may overwrite any unconsumed data in the given partition_id! - >>> # Please make sure the corresponding partition_id is empty before calling the async_put() - >>> # without metadata. - >>> # Now we only support put all the data of the corresponding partition id in once. You should repeat with - >>> # interleave the initial data if n_sample > 1 before calling the async_put(). - >>> original_prompts = torch.randn(batch_size, seq_len) - >>> n_samples = 4 - >>> prompts_repeated = torch.repeat_interleave(original_prompts, n_samples, dim=0) - >>> prompts_repeated_batch = TensorDict({"prompts": prompts_repeated}) - >>> # This will create metadata in "insert" mode internally. - >>> metadata = asyncio.run(client.async_put(data=prompts_repeated_batch, partition_id=current_partition_id)) + >>> async def main(): + ... # Example 1: Normal usage with existing metadata + ... batch_meta = await client.async_get_meta( + ... data_fields=["prompts", "attention_mask"], + ... batch_size=batch_size, + ... partition_id=current_partition_id, + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... batch = await client.async_get_data(batch_meta) + ... output = TensorDict({"response": torch.randn(batch_size, seq_len)}) + ... await client.async_put(data=output, metadata=batch_meta) + ... + ... # Example 2: Initial data insertion without pre-existing metadata + ... # BE CAREFUL: this usage may overwrite any unconsumed data in the given + ... # partition_id! Please make sure the corresponding partition_id is empty + ... # before calling the async_put() without metadata. + ... # Now we only support put all the data of the corresponding partition id in + ... # once. You should repeat with interleave the initial data if n_sample > 1 + ... # before calling the async_put(). + ... original_prompts = torch.randn(batch_size, seq_len) + ... n_samples = 4 + ... prompts_repeated = torch.repeat_interleave(original_prompts, n_samples, dim=0) + ... prompts_repeated_batch = TensorDict({"prompts": prompts_repeated}) + ... # This will create metadata in "insert" mode internally. + ... metadata = await client.async_put( + ... data=prompts_repeated_batch, partition_id=current_partition_id + ... ) """ if not hasattr(self, "storage_manager") or self.storage_manager is None: @@ -467,16 +503,18 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict: - Requested data fields (e.g., "prompts", "attention_mask") Example: - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["prompts", "attention_mask"], - ... batch_size=4, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> batch = asyncio.run(client.async_get_data(batch_meta)) - >>> print(batch) - >>> # TensorDict with fields "prompts", "attention_mask", and sample order matching metadata global_indexes + >>> async def main(): + ... batch_meta = await client.async_get_meta( + ... data_fields=["prompts", "attention_mask"], + ... batch_size=4, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... batch = await client.async_get_data(batch_meta) + ... print(batch) + ... # TensorDict with fields "prompts", "attention_mask", and sample order + ... # matching metadata global_indexes """ if not hasattr(self, "storage_manager") or self.storage_manager is None: @@ -678,10 +716,10 @@ async def async_get_consumption_status( Example: >>> # Get consumption status - >>> global_index, consumption_status = asyncio.run(client.async_get_consumption_status( + >>> global_index, consumption_status = await client.async_get_consumption_status( ... task_name="generate_sequences", ... partition_id="train_0" - ... )) + ... ) >>> print(f"Global index: {global_index}, Consumption status: {consumption_status}") """ @@ -725,10 +763,10 @@ async def async_get_production_status( Example: >>> # Get production status - >>> global_index, production_status = asyncio.run(client.async_get_production_status( + >>> global_index, production_status = await client.async_get_production_status( ... data_fields=["input_ids", "attention_mask"], ... partition_id="train_0" - ... )) + ... ) >>> print(f"Global index: {global_index}, Production status: {production_status}") """ try: @@ -766,10 +804,10 @@ async def async_check_consumption_status( Example: >>> # Check if all samples have been consumed - >>> is_consumed = asyncio.run(client.async_check_consumption_status( + >>> is_consumed = await client.async_check_consumption_status( ... task_name="generate_sequences", ... partition_id="train_0" - ... )) + ... ) >>> print(f"All samples consumed: {is_consumed}") """ @@ -802,10 +840,10 @@ async def async_check_production_status( Example: >>> # Check if all samples are ready for consumption - >>> is_ready = asyncio.run(client.async_check_production_status( + >>> is_ready = await client.async_check_production_status( ... data_fields=["input_ids", "attention_mask"], ... partition_id="train_0" - ... )) + ... ) >>> print(f"All samples ready: {is_ready}") """ _, production_status = await self.async_get_production_status( @@ -842,10 +880,10 @@ async def async_reset_consumption( Example: >>> # Reset consumption for train task to re-train on same data - >>> success = asyncio.run(client.async_reset_consumption( + >>> success = await client.async_reset_consumption( ... partition_id="train_0", ... task_name="train" - ... )) + ... ) >>> print(f"Reset successful: {success}") """ body = {"partition_id": partition_id} @@ -879,7 +917,7 @@ async def async_get_partition_list( list[str]: List of partition ids managed by the controller Example: - >>> partition_ids = asyncio.run(client.get_partition_list()) + >>> partition_ids = await client.get_partition_list() >>> print(f"Available partitions: {partition_ids}") """ try: @@ -1053,6 +1091,8 @@ def close(self) -> None: ) return try: + # Close pooled sockets before the context that owns them. + self.controller_rpc_pool.close() if hasattr(self, "zmq_context") and self.zmq_context is not None: self.zmq_context.destroy(linger=0) except Exception as e: diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 9c721df3..53d2b18f 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -42,6 +42,18 @@ backend: # ucx_env_vars: {} # See docs/nixl_ucx_payload.md for UCX transport and device configuration. + # 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/controller.py b/transfer_queue/controller.py index ef71bb4c..b45a85a6 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -56,6 +56,9 @@ TQ_CONTROLLER_GET_METADATA_TIMEOUT = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_TIMEOUT", 1)) TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 5)) +# Above libzmq's default of 1023: the metrics exporter borrows this context to query +# storage units, so the budget covers a fleet rather than the two ROUTERs bound here. +TQ_CONTROLLER_ZMQ_MAX_SOCKETS = int(os.environ.get("TQ_CONTROLLER_ZMQ_MAX_SOCKETS", 4096)) # Sample pre-allocation for StreamingDataLoader compatibility. # By pre-allocating sample indices (typically global_batch_size), consumers can accurately @@ -1725,6 +1728,9 @@ def kv_retrieve_keys( def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" self.zmq_context = zmq.Context() + # Before any socket is opened on it: libzmq applies MAX_SOCKETS at socket creation. + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + self.zmq_context.set(zmq.MAX_SOCKETS, min(TQ_CONTROLLER_ZMQ_MAX_SOCKETS, socket_limit)) self._node_ip = get_node_ip_address() while True: @@ -2337,10 +2343,9 @@ def start_metrics(self, port: int = 0) -> str: return self._metrics_endpoint from transfer_queue.metrics import TQMetricsExporter - self._metrics = TQMetricsExporter() + # Lend the controller's context rather than let the exporter build a second one. + self._metrics = TQMetricsExporter(zmq_context=self.zmq_context) self._metrics_endpoint = self._metrics.start(node_ip=self._node_ip, port=port) - # Launch a daemon thread that periodically pushes controller state - # snapshots to the exporter, keeping them process-isolated. self._metrics_snapshot_thread = Thread( target=self._metrics_snapshot_loop, name="TQMetricsSnapshotThread", diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 98a12954..eca8762e 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 +# Idle workers and proxies observe shutdown within one second; leave time for +# Ray dispatch, in-flight work, and SSD cleanup before forcing termination. +_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 8d7d7b51..44ab0d16 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -13,12 +13,26 @@ # 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 from threading import Thread from typing import Any -from uuid import uuid4 import psutil import zmq @@ -26,11 +40,11 @@ from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + METRICS_COLLECTOR_IDENTITY_PREFIX, ZMQMessage, ZMQRequestType, ZMQServerInfo, - create_zmq_socket, - format_zmq_address, + ZMQSocketPool, ) logger = get_logger(__name__) @@ -67,13 +81,22 @@ class TQMetricsExporter: TQ_METRICS_STORAGE_TIMEOUT ZMQ timeout for storage queries (default 5s) """ - def __init__(self, role: str = "controller"): + def __init__(self, role: str = "controller", zmq_context: zmq.Context | None = None): + """ + Args: + role: Which process this exporter runs in; only "controller" collects from + storage units, so only that role needs a context. + zmq_context: The owner's long-lived synchronous context, borrowed for + storage-unit queries and never terminated here. Minting one instead would + add a second context and its native I/O thread with nobody to close them, + since the exporter lives as long as its Ray actor. + """ self._start_time = time.time() self._process = psutil.Process() self._role = role self._storage_unit_infos: dict[str, ZMQServerInfo] = {} - self._zmq_ctx: zmq.Context | None = None - self._zmq_sockets: dict[str, zmq.Socket] = {} + self._zmq_ctx = zmq_context + self._zmq_socket_pool: ZMQSocketPool | None = None 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() @@ -176,6 +199,86 @@ 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 = Gauge( + "tq_storage_ssd_fallback_values_total", + "Values retained in memory because SSD encoding was unavailable", + ["storage_unit_id"], + registry=r, + ) + + # ---- Storage-unit request-loss diagnostics ---- + # Arrival labels also include unmeasured operations and NIXL PREPARE failures, + # so a gap against measured handler entries alone cannot prove request loss. + self.storage_requests_arrived = Gauge( + "tq_storage_requests_arrived", + "Decoded storage worker requests, excluding NIXL COMMIT/CANCEL continuations", + ["storage_unit_id"], + registry=r, + ) + self.storage_arrivals_by_op = Gauge( + "tq_storage_arrivals_by_op", + "Decoded storage worker requests by operation; NIXL PREPARE maps to PUT_DATA/GET_DATA", + ["storage_unit_id", "op_type"], + registry=r, + ) + + # ---- Accept-queue probe (only populated when TQ_ACCEPT_PROBE_INTERVAL > 0) ---- + self.storage_accept_queue_backlog = Gauge( + "tq_storage_accept_queue_backlog", + "Configured accept-queue depth of the storage unit's listening socket", + ["storage_unit_id"], + registry=r, + ) + self.storage_accept_queue_peak = Gauge( + "tq_storage_accept_queue_peak", + "Deepest accept-queue occupancy seen by the probe", + ["storage_unit_id"], + registry=r, + ) + self.storage_accept_queue_peak_utilization = Gauge( + "tq_storage_accept_queue_peak_utilization_ratio", + "Peak accept-queue occupancy as a fraction of the backlog", + ["storage_unit_id"], + registry=r, + ) + self.storage_socket_drops = Gauge( + "tq_storage_socket_drops", + "Connections dropped on this listening socket since the probe started", + ["storage_unit_id"], + registry=r, + ) + # Split because sk_drops covers several establishment failures, not only a full + # queue; the difference is what says whether a bigger backlog would have helped. + self.storage_listen_overflows = Gauge( + "tq_storage_listen_overflows", + "Namespace-wide accept-queue overflows since the probe started", + ["storage_unit_id"], + registry=r, + ) + self.storage_listen_other_drops = Gauge( + "tq_storage_listen_other_drops", + "Namespace-wide establishment drops that were not accept-queue overflows", + ["storage_unit_id"], + registry=r, + ) # ---- Storage request metrics (collected via ZMQ, exposed as gauges) ---- # P50/P99 are pre-computed on the storage unit side and sent via ZMQ, @@ -183,7 +286,7 @@ def _define_controller_metrics(self) -> None: # conflict with Prometheus reserved suffixes and label semantics). self.storage_request_ops = Gauge( "tq_storage_request_ops", - "Total requests processed by storage unit", + "Cumulative entries into measured PUT, GET, and CLEAR handlers", ["storage_unit_id", "op_type"], registry=r, ) @@ -353,6 +456,38 @@ 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)) + self.storage_ssd_fallback_values.labels(storage_unit_id=label).set( + metrics.get("ssd_fallback_values_total", 0) + ) + + 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(): + self.storage_arrivals_by_op.labels(storage_unit_id=label, op_type=op_type).set(arrived) + + # Drop the series rather than report zero when the probe is off, so a + # disabled probe is not read as "measured, and no drops". + accept_queue = metrics.get("accept_queue") + accept_gauges = ( + (self.storage_accept_queue_backlog, "backlog"), + (self.storage_accept_queue_peak, "peak_recv_q"), + (self.storage_accept_queue_peak_utilization, "peak_utilization"), + (self.storage_socket_drops, "sk_drops_delta"), + (self.storage_listen_overflows, "listen_overflow_delta"), + (self.storage_listen_other_drops, "non_overflow_drop_delta"), + ) + for gauge, key in accept_gauges: + if accept_queue is None: + try: + gauge.remove(label) + except (KeyError, ValueError): + pass + else: + gauge.labels(storage_unit_id=label).set(accept_queue.get(key, 0)) # Per-operation request stats for op_type, op_data in metrics.get("op_stats", {}).items(): @@ -371,49 +506,50 @@ def collect_storage_metrics(self) -> None: except Exception as e: logger.warning(f"Failed to collect metrics from storage unit {su_id}: {e}") - def _get_or_create_socket(self, su_id: str, su_info: ZMQServerInfo) -> zmq.Socket: - """Return a cached ZMQ DEALER socket for *su_id*, creating one if needed.""" - if self._zmq_ctx is None: - self._zmq_ctx = zmq.Context() - - sock = self._zmq_sockets.get(su_id) - if sock is not None and not sock.closed: - return sock - - identity = f"metrics_collector_{uuid4().hex[:8]}".encode() - sock = create_zmq_socket(self._zmq_ctx, zmq.DEALER, su_info.ip, identity) - timeout_ms = TQ_METRICS_STORAGE_TIMEOUT * 1000 - address = format_zmq_address(su_info.ip, su_info.ports["put_get_socket"]) - sock.connect(address) - sock.setsockopt(zmq.RCVTIMEO, timeout_ms) - sock.setsockopt(zmq.SNDTIMEO, timeout_ms) - self._zmq_sockets[su_id] = sock - return sock + def _get_socket_pool(self) -> ZMQSocketPool: + """Return the lazily-created socket pool for storage-unit queries.""" + if self._zmq_socket_pool is None: + if self._zmq_ctx is None: + raise RuntimeError( + "TQMetricsExporter was built without a ZMQ context, so it cannot query " + "storage units; pass zmq_context= from the owning process." + ) + self._zmq_socket_pool = ZMQSocketPool( + self._zmq_ctx, + # The storage proxy drops identities without this prefix, and the pool + # builds each socket's identity from the owner id. + METRICS_COLLECTOR_IDENTITY_PREFIX.rstrip("_"), + "put_get_socket", + timeout=TQ_METRICS_STORAGE_TIMEOUT, + ) + return self._zmq_socket_pool def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, Any] | None: """Send a synchronous GET_METRICS request to a single storage unit.""" try: - sock = self._get_or_create_socket(su_id, su_info) - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_METRICS, - sender_id="metrics_collector", - body={}, - ) - sock.send_multipart(request_msg.serialize()) - response_frames = sock.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_frames) - if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: - return response_msg.body - return None + pool = self._get_socket_pool() + with pool.lease(su_info) as sock: + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_METRICS, + sender_id="metrics_collector", + body={}, + ) + sock.send_multipart(request_msg.serialize()) + response_frames = sock.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(response_frames) + # Closed rather than parked: collection walks every unit once per cycle, so + # a kept socket holds budget for the whole walk to save one handshake. + sock.close(linger=0) + if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: + return response_msg.body + return None except zmq.error.Again: + # The pool discarded the socket, so the next cycle starts clean. Reusing it + # would read this reply, if it lands late, as the answer to that cycle's query. logger.debug(f"Timeout querying metrics from {su_id}") return None except Exception as e: logger.warning(f"Error querying metrics from {su_id}: {e}") - # Close broken socket so it gets recreated next cycle - sock = self._zmq_sockets.pop(su_id, None) - if sock and not sock.closed: - sock.close(linger=0) return None def start(self, node_ip: str = "0.0.0.0", port: int = 0) -> str: diff --git a/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py b/transfer_queue/storage/bootstrap/simple_storage_bootstrap.py index 02db2ae9..ac2c1570 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 import ray from omegaconf import DictConfig @@ -28,33 +43,42 @@ 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) - payload_transfer_config = conf.backend.SimpleStorage.get("payload_transfer") - payload_transfer_backend, _ = parse_payload_transfer_config(payload_transfer_config) + 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") + # Validate the payload transfer block early so a bad backend fails before any unit starts. + payload_transfer_backend, _ = parse_payload_transfer_config(simple_storage_config.get("payload_transfer")) + 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, - payload_transfer=payload_transfer_config, + 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/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index d6914902..e765f3c8 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -33,6 +33,34 @@ logger = get_logger(__name__) +try: + import torch_npu # noqa: F401 + + NPU_IMPORTED: bool = True +except Exception: # pragma: no cover - accelerator optional + NPU_IMPORTED = False + + +def _resolve_npu_device() -> int | None: + """Return the NPU device this process is bound to, or ``None`` if it uses no NPU. + + Ascend ACL contexts are thread-local, so this must run on the thread that owns + the device (``__init__``). Resolving it on a context-less pool worker would + always report device 0. Gating on ``is_initialized()`` (rather than + ``is_available()``) keeps a process that never touches the NPU from allocating + device memory and contending for device 0. + """ + if not NPU_IMPORTED: + return None + try: + if not torch.npu.is_initialized(): + return None + return int(torch.npu.current_device()) + except Exception: + logger.warning("Failed to resolve the current NPU device.", exc_info=True) + return None + + MOONCAKE_STORE_IMPORTED: bool = True try: from mooncake.store import MooncakeDistributedStore, ReplicateConfig @@ -40,6 +68,13 @@ except ImportError: MOONCAKE_STORE_IMPORTED = False +MOONCAKE_BUFFER_POOL_IMPORTED: bool = True +try: + from mooncake.store import BufferPool +except ImportError: + # Older mooncake builds have no lease API; those fall back to registering per transfer. + MOONCAKE_BUFFER_POOL_IMPORTED = False + BATCH_SIZE_LIMIT: int = 400 MAX_BATCH_WORKER_THREADS = 4 MAX_SERIAL_WORKER_THREADS = 4 @@ -47,6 +82,57 @@ RETRY_DELAY_SECONDS = 1.0 +def _copy_buffer_content_into_tensors(lease, targets: list[Tensor], offsets: list[int]) -> None: + """Copy each staged region out of the pooled buffer into the caller's tensors. + + Views are dropped even when a copy raises: the pool refuses to return a buffer while an + exported view of it is alive, and an exception would otherwise keep it in its frame. + """ + # torch.frombuffer rejects count=0, and an empty tensor has nothing to copy. + staged = [(t, off) for t, off in zip(targets, offsets, strict=True) if t.numel()] + if not staged: + return + targets = [t for t, _ in staged] + offsets = [off for _, off in staged] + + t0 = targets[0] + numel = t0.numel() + # A uniform, tightly packed group copies in one strided pass; ragged groups copy per tensor. + uniform = len(targets) > 1 + if uniform: + base_off = t0.storage_offset() + base_ptr = t0.untyped_storage().data_ptr() + for j, t in enumerate(targets): + if ( + t.dtype != t0.dtype + or t.shape != t0.shape + or not t.is_contiguous() + or t.storage_offset() != base_off + j * numel + or t.untyped_storage().data_ptr() != base_ptr + ): + uniform = False + break + + if uniform: + k = len(targets) + stride_elems = (offsets[1] - offsets[0]) // t0.element_size() + src = torch.frombuffer(lease.buffer, dtype=t0.dtype, count=(k - 1) * stride_elems + numel).as_strided( + (k, numel), (stride_elems, 1) + ) + try: + t0.as_strided((k, numel), (numel, 1), t0.storage_offset()).copy_(src) + finally: + del src + return + + for target, off in zip(targets, offsets, strict=True): + staged = torch.frombuffer(lease.buffer, dtype=target.dtype, count=target.numel(), offset=off) + try: + target.copy_(staged.view(target.shape)) + finally: + del staged + + @StorageClientFactory.register("MooncakeStoreClient") class MooncakeStoreClient(StorageKVClient): """ @@ -121,6 +207,9 @@ def __init__(self, config: dict[str, Any]): hard_pin = not offload_enabled self.replica_config.with_hard_pin = bool(hard_pin) + # Capture on the owning thread; pool workers reuse it via _bind_npu_device. + self._npu_device = _resolve_npu_device() + self._store = MooncakeDistributedStore() ret = self._store.setup( self.local_hostname, @@ -134,6 +223,39 @@ def __init__(self, config: dict[str, Any]): if ret != 0: raise RuntimeError(f"Mooncake store setup failed with error code: {ret}") + # Only RDMA reads pay for register_buffer (a kernel op far slower than the transfer it + # enables); TCP reads would pay the extra copy-out for no gain, so leave them unpooled. + # See https://github.com/Ascend/TransferQueue/issues/169 + use_pool = MOONCAKE_BUFFER_POOL_IMPORTED and self.protocol == "rdma" + self._buffer_pool = ( + BufferPool(self._store, max_bytes=self.local_buffer_size, max_regions=MAX_BATCH_WORKER_THREADS) + if use_pool + else None + ) + self._pool_fallback_warned = False + if not MOONCAKE_BUFFER_POOL_IMPORTED and self.protocol == "rdma" and self._gdr_staging is None: + logger.warning( + "mooncake.store.BufferPool is unavailable, so every tensor read registers and " + "unregisters its own receive buffer. Upgrade to mooncake-transfer-engine >= 0.3.12 " + "to lease pre-registered buffers instead." + ) + # Claim at most half the registered buffer: mooncake stages its own reads from the + # same region, so splitting all of it across the reader threads would starve those reads. + self._pool_share_bytes = self.local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS) + + def _bind_npu_device(self) -> None: + """``ThreadPoolExecutor`` initializer: bind a worker to ``self._npu_device``.""" + if self._npu_device is None: + return + try: + torch.npu.set_device(self._npu_device) + except Exception: + logger.warning( + f"Failed to bind NPU device {self._npu_device} on a worker thread; " + "Mooncake buffer register/unregister may fail.", + exc_info=True, + ) + def put(self, keys: list[str], values: list[Any]) -> list[dict | None]: """Stores multiple key-value pairs to MooncakeStore. @@ -173,7 +295,10 @@ def put(self, keys: list[str], values: list[Any]) -> list[dict | None]: tensor_futures: list[Future[None]] = [] bytes_futures: list[Future[list[int]]] = [] - with ThreadPoolExecutor(max_workers=MAX_BATCH_WORKER_THREADS) as executor: + with ThreadPoolExecutor( + max_workers=MAX_BATCH_WORKER_THREADS, + initializer=self._bind_npu_device, + ) as executor: if not use_gdr_path: for i in range(0, len(tensor_keys), BATCH_SIZE_LIMIT): batch_keys = tensor_keys[i : i + BATCH_SIZE_LIMIT] @@ -373,7 +498,10 @@ def get( results[idx] = val futures = [] - with ThreadPoolExecutor(max_workers=MAX_BATCH_WORKER_THREADS) as executor: + with ThreadPoolExecutor( + max_workers=MAX_BATCH_WORKER_THREADS, + initializer=self._bind_npu_device, + ) as executor: for i in range(0, len(cpu_tensor_indices), BATCH_SIZE_LIMIT): batch_keys = cpu_tensor_keys[i : i + BATCH_SIZE_LIMIT] batch_shapes = cpu_tensor_shapes[i : i + BATCH_SIZE_LIMIT] @@ -404,17 +532,90 @@ def _get_tensors_thread_worker( self, batch_keys: list[str], batch_shapes: list[tuple], batch_dtypes: list[torch.dtype], indexes: list[int] ) -> tuple[list[Tensor], list[int]]: batch_nbytes = get_nbytes(batch_dtypes, batch_shapes) - batch_buffer_tensors, batch_buffer_ptrs, region_ptrs, region_sizes = allocate_empty_tensors( - batch_dtypes, batch_shapes - ) + batch_buffer_tensors, batch_buffer_ptrs, _, _ = allocate_empty_tensors(batch_dtypes, batch_shapes) + + if self._buffer_pool is None: + self._read_into_registered_tensors(batch_keys, batch_buffer_ptrs, batch_nbytes) + return batch_buffer_tensors, indexes + + # Read what the pool can serve; register the rest in one pass after the loop, so + # oversized tensors and pool exhaustion keep the pre-pool cost (one register + one + # transfer) instead of paying it per group. + fallback: list[int] = [] + for group in split_by_bytes(batch_nbytes, self._pool_share_bytes): + fallback += self._read_via_buffer_pool( + group, batch_keys, batch_buffer_ptrs, batch_nbytes, batch_buffer_tensors + ) + if fallback: + self._read_into_registered_tensors( + [batch_keys[i] for i in fallback], + [batch_buffer_ptrs[i] for i in fallback], + [batch_nbytes[i] for i in fallback], + ) + return batch_buffer_tensors, indexes + + def _read_into_registered_tensors(self, keys: list[str], ptrs: list[int], nbytes: list[int]) -> None: + """Register the receive regions for one transfer, read into them, then unregister.""" + region_ptrs, region_sizes = merge_contiguous_memory(ptrs, nbytes) self._register_all_buffers(region_ptrs, region_sizes) try: - self._batch_get_into_with_retry(batch_keys, batch_buffer_ptrs, batch_nbytes) + self._batch_get_into_with_retry(keys, ptrs, nbytes) finally: self._unregister_all_buffers(region_ptrs) - return batch_buffer_tensors, indexes + def _read_via_buffer_pool( + self, + group: list[int], + batch_keys: list[str], + batch_ptrs: list[int], + batch_nbytes: list[int], + batch_tensors: list[Tensor], + ) -> list[int]: + """Read one group into a pooled buffer and copy into the caller's tensors. + + Returns the group unread when the pool cannot serve it, so the caller can register + those keys in one batch. A group larger than one reader's share is never leased: + that headroom belongs to the other readers and to mooncake's own staging. + """ + nbytes = [batch_nbytes[i] for i in group] + offsets, total = _aligned_offsets(nbytes) + lease = self._acquire_buffer(total) if total <= self._pool_share_bytes else None + if lease is None: + return group + + keys = [batch_keys[i] for i in group] + try: + self._batch_get_into_with_retry(keys, [lease.ptr + off for off in offsets], nbytes) + _copy_buffer_content_into_tensors(lease, [batch_tensors[i] for i in group], offsets) + except Exception: + # A failed read can leave the buffer unreturnable; keep the original cause. + try: + lease.release() + except Exception as release_error: + logger.warning(f"Returning the pooled receive buffer failed: {release_error}") + raise + lease.release() + return [] + + def _acquire_buffer(self, nbytes: int): + """Lease ``nbytes`` of pre-registered memory, or None when the pool cannot serve it. + + Never block: a busy pool means registering a receive buffer beats waiting. mooncake + reports both the per-reader cap and an over-capacity request as an exception. + """ + assert self._buffer_pool is not None + try: + return self._buffer_pool.acquire(nbytes, block=False) + except Exception as e: + logger.debug(f"Leasing {nbytes} B of pre-registered memory failed ({e}); registering own buffer.") + if not self._pool_fallback_warned: + self._pool_fallback_warned = True + logger.warning( + f"Falling back to registering receive buffers per read ({e}). Raise local_buffer_size " + "or lower read concurrency to keep reads on the pre-registered fast path." + ) + return None def _get_tensors_gdr( self, @@ -534,7 +735,14 @@ def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) - logger.error(f"remove failed for key `{actual_keys[i]}` with error code: {ret}") def close(self): - """Closes MooncakeStore.""" + """Closes MooncakeStore. + + Callers must have finished their reads first: the pool refuses to close while a + buffer is still leased, and the store owns the memory those buffers point into. + """ + if self._buffer_pool is not None: + self._buffer_pool.close() + self._buffer_pool = None if self._gdr_staging is not None: self._gdr_staging.close(self._store) self._gdr_staging = None diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 9c32e213..369f800e 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -36,7 +36,14 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.clients.base import StorageClientFactory from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket +from transfer_queue.utils.zmq_utils import ( + STORAGE_MANAGER_IDENTITY_PREFIX, + ZMQMessage, + ZMQRequestType, + ZMQServerInfo, + ZMQSocketPool, + create_zmq_socket, +) logger = get_logger(__name__) @@ -70,7 +77,7 @@ def __init__( config: DictConfig, zmq_context: zmq.asyncio.Context | None = None, ): - self.storage_manager_id = f"TQ_STORAGE_{uuid4().hex[:8]}" + self.storage_manager_id = f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}" self.config = config self.controller_info = controller_info @@ -81,6 +88,9 @@ def __init__( # creates when handed nothing. Only an owner tears its context down (see close()). self._owns_zmq_context = zmq_context is None self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context + # Notify runs on a dedicated loop, so it could not share sockets with another + # scenario in any case. + self.notify_pool = ZMQSocketPool(self.zmq_context, self.storage_manager_id, "request_handle_socket") self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -265,56 +275,39 @@ async def notify_data_update( async def _notify_and_wait(self, request_msg: list) -> None: """Send a data status notification to the controller and block until ACK is received.""" - identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() - sock = create_zmq_socket( - ctx=self.zmq_context, socket_type=zmq.DEALER, ip=self.controller_info.ip, identity=identity - ) - sock.setsockopt(zmq.LINGER, 0) - sock.connect(self.controller_info.to_addr("request_handle_socket")) - - try: - await sock.send_multipart(request_msg) - logger.debug( - f"[{self.storage_manager_id}]: Sent data status update request " - f"to controller id #{self.controller_info.id} successfully." - ) - - response_received = False - timeout = TQ_DATA_UPDATE_RESPONSE_TIMEOUT + # Acquiring the lease sits outside the handler below: a missing socket name or a dead + # context is a configuration/lifecycle fault the caller must see, not a slow ACK. + async with self.notify_pool.alease(self.controller_info) as sock: + try: + await sock.send_multipart(request_msg) + logger.debug( + f"[{self.storage_manager_id}]: Sent data status update request " + f"to controller id #{self.controller_info.id} successfully." + ) - while not response_received and timeout > 0: - try: - poll_interval = min(TQ_STORAGE_POLLER_TIMEOUT, timeout) - messages = await asyncio.wait_for( - sock.recv_multipart(copy=False), - timeout=poll_interval, - ) + # One deadline for the whole wait, so unrelated traffic on this socket cannot + # extend it and a quiet controller still gets the full budget. + deadline = time.monotonic() + TQ_DATA_UPDATE_RESPONSE_TIMEOUT + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"no ACK from controller {self.controller_info.id} after {TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s" + ) + messages = await asyncio.wait_for(sock.recv_multipart(copy=False), timeout=remaining) response_msg = ZMQMessage.deserialize(messages) if response_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: # type: ignore[arg-type] - response_received = True logger.debug( f"[{self.storage_manager_id}]: Get data status update ACK response " f"from controller id #{response_msg.sender_id} successfully." ) - break - except asyncio.TimeoutError: - timeout -= poll_interval - except Exception as e: - logger.warning(f"[{self.storage_manager_id}]: Error receiving response: {e}") - break - - if not response_received: - logger.error( - f"[{self.storage_manager_id}]: Timeout waiting for data status update ACK " - f"from controller after {TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s." - ) - finally: - try: - if not sock.closed: - sock.close(linger=0) - except Exception: - pass + return + except Exception as e: + # Logged rather than raised, so a slow controller does not fail the put. Close + # the socket: a late ACK would otherwise be read as the next lessee's reply. + logger.error(f"[{self.storage_manager_id}]: Data status update failed: {type(e).__name__}: {e}") + sock.close(linger=0) @abstractmethod async def put_data( @@ -383,6 +376,11 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets/contexts and stop the notify loop.""" + # A subclass may reject its config before super().__init__() runs, and __del__ calls + # close() anyway; return rather than bury that error under an AttributeError. + if not hasattr(self, "notify_pool"): + return + if self.controller_handshake_socket: try: if not self.controller_handshake_socket.closed: @@ -402,13 +400,16 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") + # This manager always owns its notify pool, even over a borrowed context. Only once + # the notify thread is gone: Socket.close() is not thread-safe and it holds the leases. + if notify_thread_stopped: + self.notify_pool.close() + if self._owns_zmq_context: # destroy() calls Socket.close(), which is not thread-safe, so it must run only # after the notify thread holding sockets is gone. If that thread outlived its # join, leak the context rather than risk a crash on a terminating process. if notify_thread_stopped: - # linger=0 force-closes sockets left by an interrupted request, so this - # cannot hang on term(). self.zmq_context.destroy(linger=0) else: logger.warning( diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 598795d5..ba543db8 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -15,9 +15,11 @@ import asyncio import os +import time import warnings from collections import defaultdict from collections.abc import Mapping +from functools import partial from operator import itemgetter from pathlib import Path from typing import Any, Callable, NamedTuple @@ -30,13 +32,19 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.managers.base import StorageManager, StorageManagerFactory -from transfer_queue.storage.payload_transfer import create_payload_transfer +from transfer_queue.storage.payload_transfer import ( + PayloadTransferTimeout, + create_payload_transfer, +) from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError +from transfer_queue.utils.common import log_heavy_operation from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + TQ_SOCKET_POOL_SIZE, ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, with_zmq_socket, ) @@ -44,19 +52,57 @@ TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200)) # seconds +# Attempts per storage-unit request, including the first. Each one must use a new connection: a +# DEALER silently queues messages for a peer it never reached, so no timeout can catch that. +TQ_SIMPLE_STORAGE_MAX_ATTEMPTS = int(os.environ.get("TQ_SIMPLE_STORAGE_MAX_ATTEMPTS", 3)) + +# Timeout for the post-failure probe, which only has to answer whether the unit still serves. +TQ_SIMPLE_STORAGE_PROBE_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_PROBE_TIMEOUT", 10)) + + +class StorageUnitTimeout(RuntimeError): + """A storage unit did not answer within the send/recv timeout. + + Distinct from an error the unit reported: only a missing answer is worth a new connection. + """ + + +def _describe_unit_state(body: dict[str, Any]) -> str: + """Summarize a successful probe: the unit is serving again, plus its own counters. + + Draws no conclusion about the timed-out request: the counters are cumulative per + operation and carry no request identity, so no value of them locates one request. + """ + parts = [ + f"requests_arrived={body.get('requests_arrived')}", + f"arrivals_by_op={body.get('arrivals_by_op')}", + f"active_keys={body.get('active_keys')}", + f"rss_gb={body.get('process_rss_bytes', 0) / 2**30:.2f}", + ] + op_stats = body.get("op_stats") or {} + if op_stats: + parts.append(f"completed={ {op: stats.get('request_count') for op, stats in op_stats.items()} }") + else: + # Only with Prometheus enabled; an empty dict would read as "served nothing". + parts.append("completed=unavailable(prometheus_disabled)") + return f"verdict=unit_serving_again ({' '.join(parts)})" + + _SU_SUBDIR = "simple_storage" _SU_INFO_FILE = "storage_unit_info.json" # Pre-bound decorator for storage-unit socket operations. with_storage_unit_socket = with_zmq_socket( - "put_get_socket", - get_identity=lambda self: self.storage_manager_id, get_peer=lambda self, target: self.storage_unit_infos[target], - # Long-lived context from the base StorageManager, shared with the notify path. Safe - # because the context is loop-agnostic and each socket stays per-call. - get_context=lambda self: self.zmq_context, + get_pool=lambda self: self.storage_rpc_pool, + resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), +) + +# Same endpoint as above but with the short diagnostic timeout, used only after a failure. +with_storage_unit_probe_socket = with_zmq_socket( + get_peer=lambda self, target: self.storage_unit_infos[target], + get_pool=lambda self: self.storage_probe_pool, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), - timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, ) @@ -82,6 +128,20 @@ def __init__( zmq_context: zmq.asyncio.Context | None = None, ): super().__init__(controller_info, config, zmq_context=zmq_context) + # Storage-unit RPC, on whichever context the base class settled on. + self.storage_rpc_pool = ZMQSocketPool( + self.zmq_context, + self.storage_manager_id, + "put_get_socket", + timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, + ) + self.storage_probe_pool = ZMQSocketPool( + self.zmq_context, + f"{self.storage_manager_id}_probe", + "put_get_socket", + timeout=TQ_SIMPLE_STORAGE_PROBE_TIMEOUT, + maxsize=1, + ) self.config = config server_infos: ZMQServerInfo | dict[str, ZMQServerInfo] | None = config.get("zmq_info", None) @@ -104,6 +164,28 @@ def __init__( peer_infos=config.get("payload_transfer_infos", {}) or {}, control_peer_infos=self.storage_unit_infos, ) + self._warn_if_pool_can_exhaust_context(len(self.storage_unit_infos)) + + def _warn_if_pool_can_exhaust_context(self, num_units: int) -> None: + """Warn when the pool may park more sockets than the context can hold. + + The cap is per (owner, address), so it multiplies by storage-unit count while the + context ceiling does not. At a few thousand units the product passes ZMQ_MAX_SOCKETS, + where a lease fails with EMFILE instead of degrading. + """ + try: + budget = self.zmq_context.get(zmq.MAX_SOCKETS) + except zmq.ZMQError: # pragma: no cover - context already terminating + return + worst_case = (TQ_SOCKET_POOL_SIZE + 1) * num_units + if worst_case > budget: + logger.warning( + f"[{self.storage_manager_id}]: storage RPC pool may hold up to " + f"({TQ_SOCKET_POOL_SIZE} + 1 probe) x {num_units} = {worst_case} idle sockets, above this " + f"context's ZMQ_MAX_SOCKETS ({budget}). Concurrent requests can then fail to " + f"open a socket. Lower TQ_SOCKET_POOL_SIZE or raise TQ_CLIENT_ZMQ_MAX_SOCKETS " + f"(with enough file descriptors, see ulimit -n)." + ) def _register_servers(self, server_infos: "ZMQServerInfo | dict[Any, ZMQServerInfo]"): """Register and validate server information. @@ -151,6 +233,105 @@ def _group_by_hash(self, global_indexes: list[int]) -> dict[str, RoutingGroup]: pos_lists[key].append(pos) return {key: RoutingGroup(gi_lists[key], pos_lists[key]) for key in gi_lists} + def _describe_storage_unit(self, storage_unit_id: str) -> str: + """Return ``ip:port`` for a storage unit, for use in diagnostics. + + The unit id is a random uuid4 fragment that carries no location, so a bare id in an + error message cannot be traced back to a node. Never raises: it is only ever called + while reporting another failure, and must not mask it. + """ + info = self.storage_unit_infos.get(storage_unit_id) + if info is None: + return "endpoint unknown (unit not registered with this manager)" + return f"{info.ip}:{info.ports.get('put_get_socket')}" + + @with_storage_unit_probe_socket + async def _probe_storage_unit(self, target_storage_unit: str, socket: zmq.Socket = None) -> dict[str, Any]: + """Ask a storage unit for its own counters over a brand-new socket. + + Served by the same worker thread as put and get, so an answer proves the unit is serving. + """ + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_METRICS, # type: ignore[arg-type] + sender_id=f"{self.storage_manager_id}_probe", + receiver_id=target_storage_unit, + body={}, + ) + await socket.send_multipart(request_msg.serialize()) + messages = await socket.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(messages) + if response_msg.request_type != ZMQRequestType.METRICS_RESPONSE: + raise RuntimeError(f"unexpected probe response type {response_msg.request_type}") + return response_msg.body + + async def _diagnose_storage_unit(self, target_storage_unit: str) -> str: + """Report whether the unit is reachable and serving after a request to it timed out. + + Returns one log line and never raises: it runs while another failure is being reported. + Says nothing about where that request went; no per-request state exists to show it. + """ + info = self.storage_unit_infos.get(target_storage_unit) + if info is None: + return "verdict=unknown(unit_not_registered)" + + try: + _, writer = await asyncio.wait_for( + asyncio.open_connection(info.ip, info.ports.get("put_get_socket")), timeout=5 + ) + writer.close() + tcp = "tcp=up" + except Exception as e: + tcp = f"tcp=down({type(e).__name__})" + + try: + body = await self._probe_storage_unit(target_storage_unit=target_storage_unit) + except zmq.error.Again: + return f"{tcp} verdict=unit_not_serving (no probe answer in {TQ_SIMPLE_STORAGE_PROBE_TIMEOUT}s)" + except Exception as e: + return f"{tcp} verdict=unknown (probe failed: {type(e).__name__}: {e})" + + return f"{tcp} {_describe_unit_state(body)}" + + async def _request_with_retry( + self, + operation: str, + target_storage_unit: str, + request_context: str, + make_request: Callable[[], Any], + max_attempts: int | None = None, + ): + """Run one storage-unit request, retrying a missing answer on a fresh connection. + + Args: + operation: Operation name for logs, e.g. ``get`` or ``put``. + target_storage_unit: Unit this request is routed to. + request_context: Request shape, so both ends of a failure can be correlated. + make_request: Zero-arg callable returning a coroutine for one attempt. Must build a + new socket per call, which ``with_storage_unit_socket`` does. + max_attempts: Attempts allowed. Defaults to ``TQ_SIMPLE_STORAGE_MAX_ATTEMPTS``; pass 1 + for a request that must not be replayed. + """ + # Floored at one: a nonpositive count would skip the request and report success, which for + # put would publish metadata for data that was never sent. + attempts_allowed = max(1, TQ_SIMPLE_STORAGE_MAX_ATTEMPTS if max_attempts is None else max_attempts) + for attempt in range(1, attempts_allowed + 1): + try: + return await make_request() + except StorageUnitTimeout as e: + # The exception already names the unit, its endpoint and the timeout, so these + # lines only add what it cannot know: which attempt this was, and the shape. + if attempt < attempts_allowed: + logger.warning( + f"[{self.storage_manager_id}]: {operation} retry {attempt + 1}/{attempts_allowed} " + f"on a new connection. {request_context} {e}" + ) + continue + logger.error( + f"[{self.storage_manager_id}]: {operation} failed after {attempt} attempts. " + f"{request_context} {e} {await self._diagnose_storage_unit(target_storage_unit)}" + ) + raise + @staticmethod def _select_by_positions(field_data, positions: list[int]): """Slice a single field's data by non-contiguous batch positions. @@ -264,24 +445,37 @@ async def put_data( field_schema = extract_field_schema(data) routing = self._group_by_hash(metadata.global_indexes) - tasks = [ - self._put_to_single_storage_unit( - group.global_indexes, - {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()}, - target_storage_unit=su_id, - data_parser=data_parser, + # Parser-backed puts are not replayed: the public API does not constrain parser side effects. + max_attempts = 1 if data_parser is not None else None + tasks = [] + for su_id, group in routing.items(): + storage_data = {f: self._select_by_positions(data[f], group.batch_positions) for f in data.keys()} + tasks.append( + self._request_with_retry( + "put", + su_id, + f"samples={len(group.global_indexes)} fields={list(storage_data.keys())}", + partial( + self._put_to_single_storage_unit, + group.global_indexes, + storage_data, + target_storage_unit=su_id, + data_parser=data_parser, + ), + max_attempts=max_attempts, + ) ) - for su_id, group in routing.items() - ] try: await asyncio.gather(*tasks) except Exception as e: + # The offending unit is named in the error itself; the full routing list can run to + # hundreds of ids, which buries it. logger.error( f"[{self.storage_manager_id}]: put_data failed. " f"partition_id={metadata.partition_ids[0]}, " f"num_samples={metadata.size}, " - f"storage_units={list(routing.keys())}, " + f"num_storage_units={len(routing)}, " f"error={type(e).__name__}: {e}" ) raise @@ -306,14 +500,39 @@ async def _put_to_single_storage_unit( Send data to a specific storage unit. """ - await self.payload_transfer.put( - control_socket=socket, - sender_id=self.storage_manager_id, - target_id=target_storage_unit, - global_indexes=global_indexes, - data=storage_data, - data_parser=data_parser, - ) + started = time.perf_counter() + try: + # The payload transfer owns the wire protocol and reports the payload bytes it + # moved; this end waits out the round trip, so it holds both the true wire size + # and the end-to-end latency. + wire_bytes = await self.payload_transfer.put( + control_socket=socket, + sender_id=self.storage_manager_id, + target_id=target_storage_unit, + global_indexes=global_indexes, + data=storage_data, + data_parser=data_parser, + ) + log_heavy_operation( + self.storage_manager_id, + "put", + time.perf_counter() - started, + wire_bytes, + f"to {target_storage_unit} at {self._describe_storage_unit(target_storage_unit)} " + f"samples={len(global_indexes)} fields={list(storage_data.keys())}", + ) + except PayloadTransferTimeout as e: + raise StorageUnitTimeout( + f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during put to storage unit " + f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}; " + f"serialized_mb={e.wire_bytes / 2**20:.1f}" + ) from e + except Exception as e: + logger.error( + f"[{self.storage_manager_id}]: Unexpected error during put to storage unit " + f"{target_storage_unit}: {type(e).__name__}: {e}" + ) + raise RuntimeError(f"Error in put to storage unit {target_storage_unit}: {type(e).__name__}: {e}") from e @staticmethod def _pack_field_values(values: list) -> torch.Tensor | NonTensorStack: @@ -386,7 +605,17 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: routing = self._group_by_hash(metadata.global_indexes) tasks = [ - self._get_from_single_storage_unit(group.global_indexes, metadata.field_names, target_storage_unit=su_id) + self._request_with_retry( + "get", + su_id, + f"samples={len(group.global_indexes)} fields={list(metadata.field_names)}", + partial( + self._get_from_single_storage_unit, + group.global_indexes, + metadata.field_names, + target_storage_unit=su_id, + ), + ) for su_id, group in routing.items() ] try: @@ -434,12 +663,21 @@ async def _get_from_single_storage_unit( global_indexes=global_indexes, fields=fields, ) - except StorageKeyNotFoundError: - raise + except PayloadTransferTimeout as e: + raise StorageUnitTimeout( + f"no answer in {TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT}s during get from storage unit " + f"{target_storage_unit} at {self._describe_storage_unit(target_storage_unit)}" + ) from e except Exception as exc: if KEY_NOT_FOUND_MARKER in str(exc): raise StorageKeyNotFoundError(str(exc)) from exc - raise + logger.error( + f"[{self.storage_manager_id}]: Unexpected error from storage unit " + f"{target_storage_unit}: {type(exc).__name__}: {exc}" + ) + raise RuntimeError( + f"Error getting data from storage unit {target_storage_unit}: {type(exc).__name__}: {exc}" + ) from exc return fields, data async def clear_data(self, metadata: BatchMeta) -> None: @@ -617,8 +855,22 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: logger.info(f"[{self.storage_manager_id}]: restored {len(su_ids)} storage units from {su_dir}") def close(self) -> None: - """Close payload transfer resources before ZMQ ownership is released.""" + """Close payload transfer and pooled sockets before ZMQ ownership is released.""" + # Before super(), which may destroy the context these live on. getattr guards the + # base-constructor-raised case (a failed handshake does), and __del__ still calls this. + transfer = getattr(self, "payload_transfer", None) + pool = getattr(self, "storage_rpc_pool", None) + probe_pool = getattr(self, "storage_probe_pool", None) try: - self.payload_transfer.close() + if transfer is not None: + transfer.close() finally: - super().close() + try: + if pool is not None: + pool.close() + finally: + try: + if probe_pool is not None: + probe_pool.close() + finally: + super().close() diff --git a/transfer_queue/storage/payload_transfer/__init__.py b/transfer_queue/storage/payload_transfer/__init__.py index 8df2ac9c..621eae0b 100644 --- a/transfer_queue/storage/payload_transfer/__init__.py +++ b/transfer_queue/storage/payload_transfer/__init__.py @@ -18,6 +18,7 @@ from transfer_queue.storage.payload_transfer.base import ( PayloadTransfer, PayloadTransferError, + PayloadTransferTimeout, ) from transfer_queue.storage.payload_transfer.factory import ( create_payload_transfer, @@ -27,6 +28,7 @@ __all__ = [ "PayloadTransfer", "PayloadTransferError", + "PayloadTransferTimeout", "create_payload_transfer", "parse_payload_transfer_config", ] diff --git a/transfer_queue/storage/payload_transfer/base.py b/transfer_queue/storage/payload_transfer/base.py index fa2003e0..6c9c6d46 100644 --- a/transfer_queue/storage/payload_transfer/base.py +++ b/transfer_queue/storage/payload_transfer/base.py @@ -24,10 +24,27 @@ from transfer_queue.utils.zmq_utils import ZMQMessage +# Marks a GET error reply as an expected missing key so callers can distinguish it from a fault. +KEY_NOT_FOUND_MARKER = "TQKeyNotFound" + + class PayloadTransferError(RuntimeError): """A payload transfer could not be completed safely.""" +class PayloadTransferTimeout(RuntimeError): + """A payload transfer control message was not answered within the timeout. + + Distinct from an error the peer reported: only a missing answer is worth retrying on a + new connection. Carries the request's payload bytes so the caller can report the wire + size of the attempt that timed out. + """ + + def __init__(self, message: str, *, wire_bytes: int = 0): + super().__init__(message) + self.wire_bytes = wire_bytes + + class PayloadTransfer(ABC): """Complete SimpleStorage payload strategy, including its wire protocol.""" @@ -41,8 +58,12 @@ async def put( global_indexes: list[int], data: dict[str, Any], data_parser: Callable[[Any], Any] | None, - ) -> None: - """Put decoded storage data through this strategy.""" + ) -> int: + """Put decoded storage data through this strategy. + + Returns the payload bytes moved, so callers can report transfer sizes without + knowing the wire protocol. + """ @abstractmethod async def get( diff --git a/transfer_queue/storage/payload_transfer/nixl.py b/transfer_queue/storage/payload_transfer/nixl.py index 4350905a..d4cd070c 100644 --- a/transfer_queue/storage/payload_transfer/nixl.py +++ b/transfer_queue/storage/payload_transfer/nixl.py @@ -25,9 +25,10 @@ from typing import Any, Callable from uuid import uuid4 +import zmq import zmq.asyncio -from transfer_queue.storage.payload_transfer.base import PayloadTransfer, PayloadTransferError +from transfer_queue.storage.payload_transfer.base import PayloadTransfer, PayloadTransferError, PayloadTransferTimeout from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import NixlError, NixlRuntime from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger @@ -178,7 +179,7 @@ async def put( global_indexes: list[int], data: dict[str, Any], data_parser: Callable[[Any], Any] | None, - ) -> None: + ) -> int: frames = tuple(encode(data)) descriptor = PayloadDescriptor( transfer_id=uuid4().hex, @@ -217,9 +218,17 @@ async def put( response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) self._expect(response, ZMQRequestType.PUT_DATA_RESPONSE, target_id) remote_may_be_prepared = False - except BaseException: + return descriptor.payload_bytes + except BaseException as exc: if remote_may_be_prepared: await self._cancel(sender_id, target_id, ZMQRequestType.PUT_DATA_CANCEL, descriptor.transfer_id) + if isinstance(exc, zmq.error.Again): + # Only a missing answer is retried by the caller; the transfer_id above is + # fresh per attempt, so replaying the whole put cannot duplicate a commit. + raise PayloadTransferTimeout( + f"ZMQ recv timeout during put to storage unit {target_id}", + wire_bytes=descriptor.payload_bytes, + ) from exc raise async def get( @@ -267,11 +276,14 @@ async def get( remote_prepared = False payload = await asyncio.wrap_future(self.receive(descriptor)) return decode(unpack_from(payload)) - except BaseException: + except BaseException as exc: if receive_prepared and descriptor is not None: self.cancel_receive(descriptor.transfer_id) if remote_prepared: await self._cancel(sender_id, target_id, ZMQRequestType.GET_DATA_CANCEL, transfer_id) + if isinstance(exc, zmq.error.Again): + # A retried get prepares a fresh transfer_id, so replaying cannot duplicate. + raise PayloadTransferTimeout(f"ZMQ recv timeout during get from storage unit {target_id}") from exc raise def handle_request( diff --git a/transfer_queue/storage/payload_transfer/zmq.py b/transfer_queue/storage/payload_transfer/zmq.py index 50b082e2..ceaadc3e 100644 --- a/transfer_queue/storage/payload_transfer/zmq.py +++ b/transfer_queue/storage/payload_transfer/zmq.py @@ -18,20 +18,18 @@ from __future__ import annotations import os -from typing import Any, Callable, NoReturn +from typing import Any, Callable import zmq import zmq.asyncio -from transfer_queue.storage.payload_transfer.base import PayloadTransfer +from transfer_queue.storage.payload_transfer.base import KEY_NOT_FOUND_MARKER, PayloadTransfer, PayloadTransferTimeout from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, frame_nbytes logger = get_logger(__name__) TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) -TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT = int(os.environ.get("TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT", 200)) -KEY_NOT_FOUND_MARKER = "TQKeyNotFound" class ZmqPayloadTransfer(PayloadTransfer): @@ -46,27 +44,25 @@ async def put( global_indexes: list[int], data: dict[str, Any], data_parser: Callable[[Any], Any] | None, - ) -> None: + ) -> int: request = ZMQMessage.create( request_type=ZMQRequestType.PUT_DATA, sender_id=sender_id, receiver_id=target_id, body={"global_indexes": global_indexes, "data": data, "data_parser": data_parser}, ) + frames = request.serialize() + wire_bytes = sum(frame_nbytes(frame) or 0 for frame in frames) try: - await control_socket.send_multipart(request.serialize(), copy=False) + await control_socket.send_multipart(frames, copy=False) response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) if response.request_type != ZMQRequestType.PUT_DATA_RESPONSE: raise RuntimeError( f"Failed to put data to storage unit {target_id}: {response.body.get('message', 'Unknown error')}" ) + return wire_bytes except zmq.error.Again as exc: - self._raise_timeout(sender_id, target_id, "put", exc) - except Exception as exc: - logger.error( - f"[{sender_id}]: Unexpected error during put to storage unit {target_id}: {type(exc).__name__}: {exc}" - ) - raise RuntimeError(f"Error in put to storage unit {target_id}: {type(exc).__name__}: {exc}") from exc + raise PayloadTransferTimeout(f"no answer from storage unit {target_id}", wire_bytes=wire_bytes) from exc async def get( self, @@ -87,17 +83,11 @@ async def get( await control_socket.send_multipart(request.serialize()) response = ZMQMessage.deserialize(await control_socket.recv_multipart(copy=False)) if response.request_type != ZMQRequestType.GET_DATA_RESPONSE: - raise RuntimeError( - f"Failed to get data from storage unit {target_id}: {response.body.get('message', 'Unknown error')}" - ) + message = response.body.get("message", "Unknown error") + raise RuntimeError(f"Failed to get data from storage unit {target_id}: {message}") return response.body["data"] except zmq.error.Again as exc: - self._raise_timeout(sender_id, target_id, "get", exc) - except Exception as exc: - logger.error(f"[{sender_id}]: Unexpected error from storage unit {target_id}: {type(exc).__name__}: {exc}") - raise RuntimeError( - f"Error getting data from storage unit {target_id}: {type(exc).__name__}: {exc}" - ) from exc + raise PayloadTransferTimeout(f"no answer from storage unit {target_id}") from exc def handle_request( self, @@ -157,19 +147,3 @@ def handle_request( @staticmethod def _response(request_type: ZMQRequestType, sender_id: str, body: dict[str, Any] | None = None) -> ZMQMessage: return ZMQMessage.create(request_type=request_type, sender_id=sender_id, body=body or {}) - - @staticmethod - def _raise_timeout(sender_id: str, target_id: str, operation: str, error: Exception) -> NoReturn: - timeout = TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT - if operation == "put": - logger.error( - f"[{sender_id}]: ZMQ recv timeout ({timeout}s) during put to storage unit {target_id}. " - "The storage unit may be overloaded or crashed." - ) - raise RuntimeError(f"ZMQ recv timeout ({timeout}s) during put to storage unit {target_id}") from error - - logger.error( - f"[{sender_id}]: ZMQ recv timeout ({timeout}s) from storage unit {target_id}. " - "The storage unit may be overloaded or crashed." - ) - raise RuntimeError(f"ZMQ recv timeout ({timeout}s) from storage unit {target_id}") from error diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index eb1f4384..f3408f04 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -13,44 +13,116 @@ # 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 collections.abc import Mapping +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.storage.payload_transfer import PayloadTransfer, create_payload_transfer -from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads +from transfer_queue.storage.payload_transfer.base import KEY_NOT_FOUND_MARKER +from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.perf_utils import IntervalPerfMonitor from transfer_queue.utils.zmq_utils import ( + STORAGE_CLIENT_IDENTITY_PREFIXES, ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket, format_zmq_address, + frame_nbytes, get_free_port, get_node_ip_address, ) if TYPE_CHECKING: from transfer_queue.metrics import TQMetricsExporter + from transfer_queue.utils.accept_probe import AcceptQueueProbe logger = get_logger(__name__) + +def _storage_arrival_op_type(operation: ZMQRequestType) -> str | None: + """Count decoded operations while treating NIXL continuations as the same request.""" + if operation in ( + ZMQRequestType.PUT_DATA_COMMIT, + ZMQRequestType.GET_DATA_COMMIT, + ZMQRequestType.PUT_DATA_CANCEL, + ZMQRequestType.GET_DATA_CANCEL, + ): + return None + return { + ZMQRequestType.PUT_DATA_PREPARE: "PUT_DATA", + ZMQRequestType.GET_DATA_PREPARE: "GET_DATA", + }.get(operation, operation.name) + + 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 | np.dtype[Any] | 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 | np.dtype[Any] | 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" + +# Accept-queue depth for the client-facing ROUTER. A full queue loses connections silently. +TQ_STORAGE_ZMQ_BACKLOG = int(os.environ.get("TQ_STORAGE_ZMQ_BACKLOG", 4096)) + +# Accept-queue sampling period in seconds; 0 disables the probe. Keep sub-second. +TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0)) class StorageKeyNotFoundError(KeyError): @@ -88,6 +160,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. @@ -149,6 +241,507 @@ 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=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: + # Small views must not retain a shared receive buffer after SSD offload. + # Encoding failures may be large and retain their existing fallback behavior. + if encoded is not None: + if encoded.codec == "pickle": + value = pickle.loads(encoded.payload) + elif isinstance(value, torch.Tensor): + value = value.clone() + elif isinstance(value, np.ndarray): + value = value.copy() + 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 + + 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 unique_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 + + 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 obsolete_ssd_values + ) + self._ssd_fallback_values_total += fallback_values + 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.""" + 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: @@ -168,23 +761,49 @@ class SimpleStorageUnit: zmq_server_info: ZMQ connection information for clients. """ - def __init__( - self, - storage_unit_size: int | None = None, - payload_transfer: Mapping[str, object] | None = None, - ): - """Initialize a SimpleStorageUnit with the specified size. + _requests_arrived = 0 + _arrivals_by_op: dict[str, int] = {} + _accept_probe = None + + 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. - payload_transfer: Backend and optional backend-specific settings. + 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 + 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.storage_data = StorageUnitData(self.storage_unit_size) - self.payload_transfer = create_payload_transfer(payload_transfer) + self.payload_transfer = create_payload_transfer(config.get("payload_transfer")) + self._requests_arrived = 0 + self._arrivals_by_op = {} # Internal communication address for proxy and workers self._inproc_addr = f"inproc://simple_storage_workers_{self.storage_unit_id}" @@ -212,9 +831,17 @@ def __init__( self.proxy_thread, self.zmq_context, self.put_get_socket, + self._accept_probe, + self.worker_socket, + self.storage_data, self.payload_transfer, ) + 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: @@ -226,6 +853,7 @@ def _init_zmq_socket(self) -> None: # Frontend: ROUTER for receiving client requests self.put_get_socket = create_zmq_socket(self.zmq_context, zmq.ROUTER, self._node_ip) + self.put_get_socket.setsockopt(zmq.BACKLOG, TQ_STORAGE_ZMQ_BACKLOG) while True: try: @@ -236,6 +864,18 @@ def _init_zmq_socket(self) -> None: logger.warning(f"[{self.storage_unit_id}]: Try to bind ZMQ sockets failed, retrying...") continue + if TQ_ACCEPT_PROBE_INTERVAL > 0: + # Lazy: the probe shells out to ``ss`` on a timer, so keep it out of runs that + # have not enabled it. + from transfer_queue.utils.accept_probe import AcceptQueueProbe + + self._accept_probe = AcceptQueueProbe( + port=self._put_get_socket_port, + owner_id=str(self.storage_unit_id), + interval_s=TQ_ACCEPT_PROBE_INTERVAL, + ) + self._accept_probe.start() + # Backend: DEALER for worker communication (connected via zmq.proxy) self.worker_socket = create_zmq_socket(self.zmq_context, zmq.DEALER, self._node_ip) self.worker_socket.bind(self._inproc_addr) @@ -271,8 +911,33 @@ def _start_process_put_get(self) -> None: def _proxy_routine(self) -> None: """ZMQ proxy for message forwarding between frontend ROUTER and backend DEALER.""" logger.info(f"[{self.storage_unit_id}]: start ZMQ proxy...") + assert self.put_get_socket is not None, "put_get_socket is not properly initialized" + front, back = self.put_get_socket, self.worker_socket + poller = zmq.Poller() + poller.register(front, zmq.POLLIN) + poller.register(back, zmq.POLLIN) try: - zmq.proxy(self.put_get_socket, self.worker_socket) + # Forwarding by hand rather than via zmq.proxy() so a non-TransferQueue peer on + # this exposed TCP port cannot reach the worker and terminate its request loop. + while not self._shutdown_event.is_set(): + events = dict(poller.poll(1000)) + if front in events: + messages = front.recv_multipart(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) except zmq.ContextTerminated: logger.info(f"[{self.storage_unit_id}]: ZMQ Proxy stopped gracefully (Context Terminated)") except Exception as e: @@ -296,7 +961,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)") @@ -309,73 +975,9 @@ 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:] - - request_msg = ZMQMessage.deserialize(serialized_msg) - operation = request_msg.request_type - - try: - logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") - - metric_op = { - ZMQRequestType.PUT_DATA: "PUT_DATA", - ZMQRequestType.GET_DATA: "GET_DATA", - ZMQRequestType.PUT_DATA_COMMIT: "PUT_DATA", - ZMQRequestType.GET_DATA_COMMIT: "GET_DATA", - }.get(operation) - if metric_op is not None: - with monitor.measure(op_type=metric_op): - response_msg = self.payload_transfer.handle_request( - request_msg, - storage_id=self.storage_unit_id, - load_data=self._load_data, - store_data=self._put_decoded_data, - ) - else: - response_msg = self.payload_transfer.handle_request( - request_msg, - storage_id=self.storage_unit_id, - load_data=self._load_data, - store_data=self._put_decoded_data, - ) - - if response_msg is None and operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] - with monitor.measure(op_type="CLEAR_DATA"): - response_msg = self._handle_clear(request_msg) - elif response_msg is None and operation == ZMQRequestType.GET_METRICS: # type: ignore[arg-type] - response_msg = self._handle_get_metrics() - elif response_msg is None and operation == ZMQRequestType.SAVE_STORAGE_CHECKPOINT: # type: ignore[arg-type] - response_msg = self._handle_save_checkpoint(request_msg) - elif response_msg is None and operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] - response_msg = self._handle_load_checkpoint(request_msg) - elif response_msg is None: - 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 - worker_socket.send_multipart([identity] + response_msg.serialize(), 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) @@ -406,6 +1008,110 @@ def _put_decoded_data(self, global_indexes: list[int], field_data: dict[str, Any ) self.storage_data.put_data(field_data, global_indexes) + 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: + # Count each logical PUT/GET once at its first message; NIXL commit/cancel + # messages continue that request and must not inflate arrivals. + arrival_op = _storage_arrival_op_type(operation) + if arrival_op is not None: + self._requests_arrived += 1 + self._arrivals_by_op[arrival_op] = self._arrivals_by_op.get(arrival_op, 0) + 1 + + logger.debug(f"[{self.storage_unit_id}]: worker received operation: {operation}") + + # The payload transfer owns PUT/GET operations; everything else stays with the unit. + metric_op = { + ZMQRequestType.PUT_DATA: "PUT_DATA", + ZMQRequestType.GET_DATA: "GET_DATA", + ZMQRequestType.PUT_DATA_COMMIT: "PUT_DATA", + ZMQRequestType.GET_DATA_COMMIT: "GET_DATA", + }.get(operation) + if metric_op is not None: + with monitor.measure(op_type=metric_op): + response_msg = self.payload_transfer.handle_request( + request_msg, + storage_id=self.storage_unit_id, + load_data=self._load_data, + store_data=self._put_decoded_data, + ) + else: + response_msg = self.payload_transfer.handle_request( + request_msg, + storage_id=self.storage_unit_id, + load_data=self._load_data, + store_data=self._put_decoded_data, + ) + + if response_msg is None and operation == ZMQRequestType.CLEAR_DATA: # type: ignore[arg-type] + with monitor.measure(op_type="CLEAR_DATA"): + response_msg = self._handle_clear(request_msg) + elif response_msg is None and operation == ZMQRequestType.GET_METRICS: # type: ignore[arg-type] + response_msg = self._handle_get_metrics() + elif response_msg is None and operation == ZMQRequestType.SAVE_STORAGE_CHECKPOINT: # type: ignore[arg-type] + response_msg = self._handle_save_checkpoint(request_msg) + elif response_msg is None and operation == ZMQRequestType.LOAD_STORAGE_CHECKPOINT: # type: ignore[arg-type] + response_msg = self._handle_load_checkpoint(request_msg) + elif response_msg is None: + 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) + @staticmethod def _field_length(value: Any) -> int | None: if hasattr(value, "shape") and isinstance(value.shape, tuple | list) and len(value.shape) > 0: @@ -472,7 +1178,25 @@ 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, + # Keep arrivals independent of whether a request has a measured handler. + "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"] = { + "backlog": stats.backlog, + "peak_recv_q": stats.peak_recv_q, + "peak_utilization": stats.peak_utilization, + "sk_drops_delta": stats.sk_drops_delta, + "listen_overflow_delta": stats.overflow_delta, + "non_overflow_drop_delta": stats.non_overflow_drop_delta, + "samples": stats.samples, + } # Include per-operation stats if Prometheus metrics are enabled if self._metrics is not None: @@ -503,11 +1227,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). @@ -517,14 +1241,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] @@ -540,7 +1257,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: @@ -555,28 +1272,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] @@ -631,29 +1344,39 @@ def _shutdown_resources( proxy_thread: Thread | None, zmq_context: zmq.Context | None, put_get_socket: zmq.Socket | None, - payload_transfer: PayloadTransfer, + accept_probe: "AcceptQueueProbe | None" = None, + worker_socket: zmq.Socket | None = None, + storage_data: StorageUnitData | None = None, + payload_transfer: PayloadTransfer | None = None, ) -> None: """Clean up resources on garbage collection.""" logger.info("Shutting down SimpleStorageUnit resources...") - # Signal all threads to stop shutdown_event.set() - # Terminate put_get_socket + # 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() + + 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() + if storage_data is not None: + try: + storage_data.close() + except Exception as e: + logger.warning(f"Error closing storage data on shutdown: {e}") - # 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) - - payload_transfer.close() + if payload_transfer is not None: + payload_transfer.close() logger.info("SimpleStorageUnit resources shutdown complete.") diff --git a/transfer_queue/utils/accept_probe.py b/transfer_queue/utils/accept_probe.py new file mode 100644 index 00000000..916955c4 --- /dev/null +++ b/transfer_queue/utils/accept_probe.py @@ -0,0 +1,245 @@ +# 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. + +"""Accept-queue sampling for the storage-unit ROUTER socket. + +A connection dropped while being established is invisible to both ends, so the only +way to see it is to read the kernel's own counters. Two caveats govern how callers +may read what this reports: ``ListenOverflows`` and ``ListenDrops`` are per network +namespace rather than per port, and the kernel charges a socket's ``sk_drops`` for +several establishment failures, of which a full queue is only one. See +``docs/metrics.md`` for the exported metrics. +""" + +from __future__ import annotations + +import re +import subprocess +import threading +import time +from dataclasses import dataclass + +from transfer_queue.utils.logging_utils import get_logger + +logger = get_logger(__name__) + +_RECVQ_RE = re.compile(r"^LISTEN\s+(\d+)\s+(\d+)\s+\S*:(\d+)\s", re.M) +_DROPS_RE = re.compile(r"\bd(\d+)\b") + + +@dataclass +class AcceptQueueSample: + """One reading of the accept queue for a single listening port.""" + + timestamp: float + recv_q: int + backlog: int + sk_drops: int + listen_overflows: int + listen_drops: int + + @property + def utilization(self) -> float: + """Queue depth as a fraction of the configured backlog.""" + return self.recv_q / self.backlog if self.backlog else 0.0 + + +@dataclass +class AcceptQueueStats: + """Peak and delta summary over a sampling window.""" + + port: int + samples: int = 0 + peak_recv_q: int = 0 + peak_utilization: float = 0.0 + backlog: int = 0 + first_sample: AcceptQueueSample | None = None + last_sample: AcceptQueueSample | None = None + + @property + def sk_drops_delta(self) -> int: + """Connections this socket dropped during the window, 0 until two samples exist.""" + if self.first_sample is None or self.last_sample is None: + return 0 + return self.last_sample.sk_drops - self.first_sample.sk_drops + + @property + def overflow_delta(self) -> int: + """Namespace-wide accept-queue overflows during the window, 0 until two samples exist.""" + if self.first_sample is None or self.last_sample is None: + return 0 + return self.last_sample.listen_overflows - self.first_sample.listen_overflows + + @property + def non_overflow_drop_delta(self) -> int: + """Establishment drops in this window that were not accept-queue overflows.""" + if self.first_sample is None or self.last_sample is None: + return 0 + drops = self.last_sample.listen_drops - self.first_sample.listen_drops + return drops - self.overflow_delta + + def describe(self) -> str: + """Return a one-line summary of the window, for the probe's shutdown log.""" + return ( + f"port={self.port} samples={self.samples} backlog={self.backlog} " + f"peak_recv_q={self.peak_recv_q} peak_util={self.peak_utilization:.1%} " + f"sk_drops_delta={self.sk_drops_delta} listen_overflow_delta={self.overflow_delta}" + ) + + +def _read_listen_socket(port: int) -> tuple[int, int, int] | None: + """Return (recv_q, backlog, sk_drops) for the listening socket on ``port``.""" + try: + out = subprocess.run(["ss", "-lntm"], capture_output=True, text=True, timeout=5, check=False).stdout + except (OSError, subprocess.SubprocessError) as exc: + logger.debug(f"accept-probe: ss failed: {exc}") + return None + + for match in _RECVQ_RE.finditer(out): + recv_q, backlog, found_port = (int(g) for g in match.groups()) + if found_port != port: + continue + # skmem lives on the continuation line right after the match. + tail = out[match.end() : match.end() + 400] + drops_match = _DROPS_RE.search(tail.split("\n")[1] if "\n" in tail else "") + return recv_q, backlog, int(drops_match.group(1)) if drops_match else 0 + return None + + +def _read_listen_overflows() -> tuple[int, int]: + """Return this namespace's (ListenOverflows, ListenDrops) from /proc/net/netstat.""" + try: + with open("/proc/net/netstat") as handle: + lines = handle.read().splitlines() + except OSError: + return 0, 0 + + for i, line in enumerate(lines): + if not line.startswith("TcpExt:") or "ListenOverflows" not in line: + continue + keys = line.split() + values = lines[i + 1].split() + try: + return ( + int(values[keys.index("ListenOverflows")]), + int(values[keys.index("ListenDrops")]), + ) + except (ValueError, IndexError): + return 0, 0 + return 0, 0 + + +def sample_accept_queue(port: int) -> AcceptQueueSample | None: + """Take one accept-queue reading for ``port``, or None if it cannot be read.""" + listen = _read_listen_socket(port) + if listen is None: + return None + recv_q, backlog, sk_drops = listen + overflows, drops = _read_listen_overflows() + return AcceptQueueSample( + timestamp=time.time(), + recv_q=recv_q, + backlog=backlog, + sk_drops=sk_drops, + listen_overflows=overflows, + listen_drops=drops, + ) + + +class AcceptQueueProbe: + """Sample one listening port's accept queue from a background thread. + + Args: + port: Listening port to watch. + owner_id: Identifier used in log lines (the storage unit id). + interval_s: Seconds between samples. Keep sub-second; the queue drains in + milliseconds, so a slower cadence misses the burst entirely. + warn_utilization: Warn once when depth first reaches this fraction of the backlog. + """ + + def __init__( + self, + port: int, + owner_id: str, + interval_s: float = 0.1, + warn_utilization: float = 0.5, + ) -> None: + self.port = port + self.owner_id = owner_id + self.interval_s = interval_s + self.warn_utilization = warn_utilization + self.stats = AcceptQueueStats(port=port) + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._warned = False + + def start(self) -> None: + """Start the sampling thread; a second call is a no-op.""" + if self._thread is not None: + return + self._thread = threading.Thread(target=self._run, name=f"AcceptQueueProbe-{self.owner_id}", daemon=True) + self._thread.start() + logger.info(f"[{self.owner_id}]: accept-queue probe started on port {self.port} (interval={self.interval_s}s)") + + def stop(self) -> None: + """Stop the sampling thread and log the window summary.""" + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=5) + self._thread = None + logger.info(f"[{self.owner_id}]: accept-queue probe stopped. {self.stats.describe()}") + + def _run(self) -> None: + while not self._stop.is_set(): + sample = sample_accept_queue(self.port) + if sample is not None: + self._record(sample) + self._stop.wait(self.interval_s) + + def _record(self, sample: AcceptQueueSample) -> None: + stats = self.stats + stats.samples += 1 + stats.backlog = sample.backlog or stats.backlog + previous = stats.last_sample + if stats.first_sample is None: + stats.first_sample = sample + stats.last_sample = sample + + if sample.recv_q > stats.peak_recv_q: + stats.peak_recv_q = sample.recv_q + stats.peak_utilization = sample.utilization + + # Against the previous sample, not the first: sk_drops is cumulative, so measuring + # from probe start would re-report one old drop on every sample. + if previous is not None and sample.sk_drops > previous.sk_drops: + logger.error( + f"[{self.owner_id}]: listening socket on port {self.port} dropped an incoming " + f"connection. recv_q={sample.recv_q}/{sample.backlog} sk_drops={sample.sk_drops} " + f"(+{sample.sk_drops - previous.sk_drops} since the last sample, " + f"+{stats.sk_drops_delta} since probe start); netns since probe start: " + f"listen_overflows +{stats.overflow_delta}, other establishment drops " + f"+{stats.non_overflow_drop_delta}. The kernel charges sk_drops for several " + f"establishment failures, so a full queue is only one candidate: raising " + f"ZMQ_BACKLOG above {sample.backlog} helps only if recv_q above and the " + f"overflow delta point that way." + ) + + if not self._warned and sample.utilization >= self.warn_utilization: + self._warned = True + logger.warning( + f"[{self.owner_id}]: accept queue on port {self.port} reached " + f"{sample.recv_q}/{sample.backlog} ({sample.utilization:.0%}); " + f"connections are queueing faster than they are accepted." + ) diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py index 999e8183..1278dfd4 100644 --- a/transfer_queue/utils/common.py +++ b/transfer_queue/utils/common.py @@ -133,3 +133,27 @@ def get_env_bool(env_key: str, default: bool = False) -> bool: true_values = {"true", "1", "yes", "y", "on"} return env_value_lower in true_values + + +# Thresholds above which a single request earns a log line. Both sit far above the normal range +# of single-digit milliseconds, so tripping either one is itself the finding. +TQ_STORAGE_SLOW_REQUEST_SECONDS = float(os.environ.get("TQ_STORAGE_SLOW_REQUEST_SECONDS", 5.0)) +TQ_STORAGE_LARGE_PAYLOAD_MB = float(os.environ.get("TQ_STORAGE_LARGE_PAYLOAD_MB", 256)) + + +def log_heavy_operation(component_id: str, operation: str, elapsed: float, payload_bytes: int, detail: str) -> None: + """Warn about one slow or unusually large request; stay silent otherwise. + + Args: + component_id: Storage manager or storage unit reporting the request. + operation: Operation name, e.g. ``put`` or ``get``. + elapsed: Wall time spent on the request, in seconds. + payload_bytes: Serialized size of the payload on the wire. + detail: Request shape, appended to the message verbatim. + """ + payload_mb = payload_bytes / 2**20 + if elapsed < TQ_STORAGE_SLOW_REQUEST_SECONDS and payload_mb < TQ_STORAGE_LARGE_PAYLOAD_MB: + return + logger.warning( + f"[{component_id}]: heavy {operation} {detail} serialized_mb={payload_mb:.1f} elapsed={elapsed:.2f}s" + ) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index f63c39e7..26fe6f67 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -13,9 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import itertools +import os import socket +import threading import time -from collections.abc import Sequence +from collections.abc import AsyncIterator, Iterator, Sequence +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar from dataclasses import dataclass from functools import wraps from typing import Any, Callable, TypeAlias @@ -32,6 +38,18 @@ logger = get_logger(__name__) +# Identity prefixes of the peers allowed to reach a storage unit. The storage proxy drops +# anything else, so an identity built without these prefixes is silently unreachable. +STORAGE_MANAGER_IDENTITY_PREFIX = "TQ_STORAGE_" +METRICS_COLLECTOR_IDENTITY_PREFIX = "metrics_collector_" +STORAGE_CLIENT_IDENTITY_PREFIXES = ( + STORAGE_MANAGER_IDENTITY_PREFIX.encode(), + METRICS_COLLECTOR_IDENTITY_PREFIX.encode(), +) + +# Cap for every pool; see ZMQSocketPool for what it bounds. Small because it multiplies by +# peer count, and one socket per peer already avoids the repeated handshake. +TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 4)) bytestr: TypeAlias = bytes | bytearray | memoryview @@ -359,47 +377,263 @@ def create_zmq_socket( return socket +def _lease_owner() -> Any: + """The running event loop, or the current thread outside one. + + A ZMQ socket is safe on neither a second thread nor a second event loop, so leases + never cross either. pyzmq silently rebinds an async socket to whatever loop it next + sees, and one bound to a *closed* loop is the "Bad file descriptor / SIGABRT" failure + that per-call context churn used to cause. Returning the thread outside a loop lets + synchronous callers (the metrics collector) share this pool unchanged. + """ + try: + return asyncio.get_running_loop() + except RuntimeError: + return threading.current_thread() + + +def _owner_finished(owner: Any) -> bool: + """Whether a lease owner can no longer serve its sockets.""" + if isinstance(owner, threading.Thread): + return not owner.is_alive() + return owner.is_closed() + + +# Addresses whose permit the current task already holds. A ContextVar because asyncio +# copies the context per task, so siblings under one gather do not see each other's. +_held_permits: ContextVar[frozenset[str]] = ContextVar("_held_permits", default=frozenset()) + + +class ZMQSocketPool: + """Lends connected DEALER sockets for one request scenario, reusing them across calls. + + One pool serves one kind of request -- controller RPC, storage RPC, notify, metrics -- + so ``socket_name`` and ``timeout`` are fixed here and every socket it holds is + interchangeable but for the address it is connected to. Separate pools per scenario also + keep them from interfering: a pool shared across scenarios reused nothing anyway, since + scenarios differ by owning loop, socket name, or timeout. + + A lease is exclusive. Responses do not echo their request's ``request_id`` (see + ``ZMQMessage.create``), so a reply is matched to its request only by arrival order: + two concurrent users of one socket would read each other's replies. Each caller + therefore gets a socket to itself and returns it only after a clean send/recv. + + Reuse avoids a TCP + ZMTP handshake per request and stops the peer's ROUTER from + accreting a fresh identity every call. + """ + + def __init__( + self, + ctx: zmq.Context, + owner_id: str, + socket_name: str, + *, + timeout: int | None = None, + maxsize: int | None = None, + ): + """ + Args: + ctx: Long-lived context to open sockets on, sync or async. The pool borrows it + and never terminates it; the owner outlives the pool. + owner_id: Identity prefix for pooled sockets, for readable peer-side logs. + socket_name: Port key in ``ZMQServerInfo.ports`` that every lease dials. + timeout: Send/recv timeout in seconds applied to every socket, or None for none. + maxsize: Idle sockets kept per bucket, at least 1, defaulting to + ``TQ_SOCKET_POOL_SIZE``. A soft cap: a burst beyond it still gets sockets, + and the excess is closed on return rather than made to wait. Counted per + (owner, address), not per pool, so a pool dialling N peers may hold + N*maxsize idle sockets. + """ + # Resolved here rather than in the signature so the env var stays patchable; a + # default bound at import froze whatever the environment held when this module loaded. + maxsize = TQ_SOCKET_POOL_SIZE if maxsize is None else maxsize + if maxsize < 1: + # Below 1 nothing is ever parked, so every request pays a fresh connect while + # still looking pooled. Reject it rather than silently disable reuse. + raise ValueError(f"ZMQ socket pool size must be at least 1, got {maxsize}") + self._ctx = ctx + self._owner_id = owner_id + self._socket_name = socket_name + self._timeout = timeout + self._maxsize = maxsize + # Keyed by lease owner (see _lease_owner), then by address rather than peer id: a peer + # restarted under the same id at a new address must not get a socket wired to the old. + self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} + self._lock = threading.Lock() + # One semaphore per (owner, address), keyed like _idle because a semaphore belongs + # to the loop that awaits it. Used by alease only; see there. + self._permits: dict[Any, dict[str, asyncio.Semaphore]] = {} + # A ROUTER silently drops a second peer claiming an identity it already has, and + # owner_id alone repeats across nodes because client ids are pid-derived. + self._identity_prefix = f"{owner_id}_{uuid4().hex[:8]}" + self._counter = itertools.count() + + @contextmanager + def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: + """Yield a socket connected to ``peer``, returning it to the pool only on success. + + A plain (non-async) contextmanager on purpose: ``with`` still sees exceptions and + ``CancelledError`` raised across ``await``s in its body, so this one definition + serves both async and synchronous callers. + + Concurrency beyond ``maxsize`` opens extra sockets here rather than waiting; async + callers that want the cap enforced use ``alease``. + """ + address = self._address(peer) + sock = self._take(address) or self._connect(peer, address) + try: + yield sock + except BaseException: + # Poisoned: the request may already be on the wire, so its reply could still + # arrive and the next lessee would read it as its own. Cancellation counts too. + sock.close(linger=0) + raise + else: + self._release(address, sock) + + @asynccontextmanager + async def alease(self, peer: ZMQServerInfo) -> AsyncIterator[zmq.Socket]: + """Like ``lease``, but waits for a socket instead of opening one past ``maxsize``. + + This bounds sockets in flight at ``maxsize`` per (owner, address), so the context's + socket budget follows configuration rather than peak concurrency. + + A permit is held for the whole body, so a task that leases inside another lease can + wait on a permit it -- or a sibling mid-cycle -- already holds, and that hangs with + no timeout. Any nesting therefore raises, including across pools: an RPC's reply is + what releases its permit, so a second RPC belongs after the first returns. Notify + already works this way, running once the puts it reports have completed. + """ + address = self._address(peer) + held = _held_permits.get() + if held: + raise RuntimeError( + f"Lease on {address} from a task already holding {sorted(held)}. A lease " + f"keeps its permit until its body ends, so nesting one inside another can " + f"wait on a permit that is already held and hang. Complete the outer " + f"request first, then start this one." + ) + token = _held_permits.set(held | {address}) + permit = self._permit(address) + try: + async with permit: + with self.lease(peer) as sock: + yield sock + finally: + _held_permits.reset(token) + + def _address(self, peer: ZMQServerInfo) -> str: + port = peer.ports.get(self._socket_name) + if port is None: + raise RuntimeError(f"Socket '{self._socket_name}' not configured for server '{peer.id}'") + return format_zmq_address(peer.ip, port) + + def _permit(self, address: str) -> asyncio.Semaphore: + """The permit gating this (owner, address), created on first use by that owner.""" + with self._lock: + return self._permits.setdefault(_lease_owner(), {}).setdefault(address, asyncio.Semaphore(self._maxsize)) + + def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: + """Buckets for the current lease owner, first evicting any owner that has finished. + + Callers must hold ``self._lock``. A closed loop's sockets must not linger: a pooled + async socket keeps its loop referenced, so they would never be collected, and reusing + one is the SIGABRT hazard in _lease_owner. Sweeping here needs no background thread, + and the loop count stays tiny (one per client, plus one per notify thread). + """ + for owner in [o for o in self._idle if _owner_finished(o)]: + self._close_all(self._idle.pop(owner)) + self._permits.pop(owner, None) + return self._idle.setdefault(_lease_owner(), {}) + + def _take(self, address: str) -> zmq.Socket | None: + """Pop a live idle socket for *address*, discarding any found closed.""" + with self._lock: + bucket = self._owner_buckets().get(address) + while bucket: + sock = bucket.pop() + if not sock.closed: + return sock + return None + + def _release(self, address: str, sock: zmq.Socket) -> None: + """Return a socket, closing it if already closed or its bucket is full.""" + with self._lock: + # A caller may close the socket itself without raising, as the notify path does + # to discard a possibly-late ACK. + if not sock.closed: + bucket = self._owner_buckets().setdefault(address, []) + if len(bucket) < self._maxsize: + bucket.append(sock) + return + sock.close(linger=0) + + def _connect(self, peer: ZMQServerInfo, address: str) -> zmq.Socket: + """Open and connect a new DEALER socket to *address*.""" + identity = f"{self._identity_prefix}_to_{peer.id}_{next(self._counter)}".encode() + sock = create_zmq_socket(self._ctx, zmq.DEALER, peer.ip, identity=identity) + try: + if self._timeout is not None: + sock.setsockopt(zmq.RCVTIMEO, self._timeout * 1000) + sock.setsockopt(zmq.SNDTIMEO, self._timeout * 1000) + sock.connect(address) + except BaseException: + # Nothing owns the socket until it is handed to a lease, so close it here or it + # leaks. connect() raises on a malformed endpoint or a terminating context. + sock.close(linger=0) + raise + return sock + + def close(self) -> None: + """Close every idle socket. Safe to call twice, and after the context is gone. + + The pool stays usable afterwards: a lease still outstanding returns its socket to a + fresh bucket. Callers destroy the context right after, so nothing is reused. + """ + with self._lock: + owned = list(self._idle.values()) + self._idle = {} + for buckets in owned: + self._close_all(buckets) + + def _close_all(self, buckets: dict[str, list[zmq.Socket]]) -> None: + """Close every socket in *buckets*, tolerating an already-destroyed context.""" + for sock in itertools.chain.from_iterable(buckets.values()): + try: + if not sock.closed: + sock.close(linger=0) + except Exception as e: + logger.debug(f"[{self._owner_id}]: Error closing pooled socket: {e}") + + def with_zmq_socket( - socket_name: str, *, - get_identity: Callable[[Any], str], get_peer: Callable[[Any, str | None], ZMQServerInfo], - get_context: Callable[[Any], "zmq.asyncio.Context"], + get_pool: Callable[[Any], ZMQSocketPool], resolve_target: Callable[[tuple, dict], str | None] | None = None, - timeout: int | None = None, ): - """Create a reusable async decorator for request sockets. + """Create a reusable async decorator that injects a pooled request socket. - Lifecycle: get owner's shared context -> create/connect socket -> inject -> close socket. + Lifecycle: resolve peer -> lease a socket from ``self``'s pool -> inject as the + ``socket`` kwarg -> return it to the pool if the call succeeded, else discard it. - The context comes from ``self`` via ``get_context`` and is long-lived; only the DEALER - socket is per-call. Do NOT create or terminate a context here -- per-call churn corrupts - libzmq's signaler file descriptors under concurrency (Bad file descriptor / SIGABRT) and - can hang on term(). Contexts are thread-safe and loop-agnostic, so sharing one is safe. + The socket name and timeout come from the pool, which serves one request scenario. Args: - socket_name: Socket port key in ``ZMQServerInfo.ports``. - get_identity: Callable that extracts owner identity from ``self``. - Example: ``lambda self: self.client_id`` get_peer: Callable that returns ``ZMQServerInfo`` for the target. For single-target scenarios, ignore the target parameter. Example: ``lambda self, target: self.server_info`` Example: ``lambda self, target: self.storage_unit_infos[target]`` - get_context: Callable that returns the owner's long-lived ``zmq.asyncio.Context``. - Example: ``lambda self: self.zmq_context`` + get_pool: Callable that returns the pool for this scenario. + Example: ``lambda self: self.controller_rpc_pool`` resolve_target: Optional callable that extracts target identifier from function arguments. Receives (args, kwargs) and returns target name. Example: ``lambda args, kwargs: kwargs.get("target_storage_unit")`` - timeout: Optional timeout (seconds) for both send/recv operations. """ def decorator(func: Callable): @wraps(func) async def wrapper(self, *args, **kwargs): - owner_id = get_identity(self) - if owner_id is None: - raise RuntimeError("get_identity returned None") - target_name: str | None = None if resolve_target is not None: target_name = resolve_target(args, kwargs) @@ -408,31 +642,13 @@ async def wrapper(self, *args, **kwargs): if server_info is None: raise RuntimeError(f"get_peer returned None for target '{target_name}'") - port = server_info.ports.get(socket_name) - if port is None: - raise RuntimeError(f"Socket '{socket_name}' not configured for server '{server_info.id}'") + pool = get_pool(self) + if pool is None: + raise RuntimeError("get_pool returned None") - # Reuse the owner's long-lived context; only the socket is per-call. - context = get_context(self) - if context is None: - raise RuntimeError("get_context returned None") - - sock = None - try: - address = format_zmq_address(server_info.ip, port) - identity = f"{owner_id}_to_{server_info.id}_{uuid4().hex[:8]}".encode() - sock = create_zmq_socket(context, zmq.DEALER, server_info.ip, identity=identity) - sock.connect(address) - if timeout is not None: - sock.setsockopt(zmq.RCVTIMEO, timeout * 1000) - sock.setsockopt(zmq.SNDTIMEO, timeout * 1000) + async with pool.alease(server_info) as sock: kwargs["socket"] = sock return await func(self, *args, **kwargs) - finally: - # Close the per-call socket only; the context outlives the call. linger=0 - # drops unsent frames so close never blocks the event loop. - if sock is not None and not sock.closed: - sock.close(linger=0) return wrapper diff --git a/transfer_queue/version/version b/transfer_queue/version/version index ccd72abe..260e27af 100644 --- a/transfer_queue/version/version +++ b/transfer_queue/version/version @@ -1 +1 @@ -0.1.11.dev0 +0.1.12.dev0