Repository navigation
[feat] simple storage support ssd offload - #162
Conversation
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ed9768257d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| done | ||
| done | ||
|
|
||
| python "${SCRIPT_DIR}/draw_figure.py" |
There was a problem hiding this comment.
When this documented SSD benchmark runs with a clean results/ directory, every generated filename ends in sample1mib, sample2mib, or sample4mib, but draw_figure.py accepts only small, medium, and large; it skips every new CSV and then fails at pd.concat([]). If older standard benchmark CSVs exist, it instead produces a chart from those stale, unrelated results. Thus the expensive benchmark either exits unsuccessfully or emits a misleading comparison rather than the advertised SSD chart; extend the plotter for these names/scenarios or invoke a dedicated SSD plotter.
AGENTS.md reference: AGENTS.md:L17-L18
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
I delete the draw command.
| @@ -0,0 +1,85 @@ | |||
| #!/bin/bash | |||
There was a problem hiding this comment.
Can we integrate the SSD test in existing test shell scripts?
There was a problem hiding this comment.
The current implementation uses a 1 MB threshold per sample, but the data sizes in the existing benchmark scripts are too small to trigger the offloading mechanism.
| storage.close() | ||
|
|
||
|
|
||
| def test_hybrid_storage_round_trips_raw_tensor_samples(tmp_path): |
There was a problem hiding this comment.
We may not need some of the tests here. Consider simplify the test.
| enabled: false | ||
| # Directory on a local NVMe SSD where offload files are written. | ||
| # Each storage unit creates its own sub-directory; files are removed on shutdown. | ||
| path: /tmp/tq_ssd_offload |
There was a problem hiding this comment.
/tmp usually is not a SSD path, but a tmpfs RAM space. This may explain the high performance in PR description. If this is the case, we have to rerun it.
There was a problem hiding this comment.
I have changed the default value to null, so users must explicitly provide a path. I also removed the code that handled the previous default value and added validation to ensure that at least the parent directory of the user-provided path exists.
ed97682 to
6b636de
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
6b636de to
988f095
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
988f095 to
b1d5412
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
| storage_unit_size: int | None = None, | ||
| ssd_config=None, | ||
| ssd_run_id: str | None = None, |
There was a problem hiding this comment.
It might be better to pass the simple storage related config in a single config param, e.g. https://github.com/Ascend/TransferQueue/blob/main/transfer_queue/storage/clients/mooncake_client.py#L56
| ssd_offload: | ||
| # Master switch. Set to true to enable SSD offload. | ||
| enabled: false | ||
| # Required when enabled. Its parent must exist on a local NVMe SSD; |
There was a problem hiding this comment.
we should allow the case when parent dir not exist.
| fd = os.open(temp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) | ||
| try: | ||
| self._write_many(fd, [sample.payload]) | ||
| finally: | ||
| os.close(fd) |
There was a problem hiding this comment.
Maybe we can try to find some existing library that can help us to enable io_uring
There was a problem hiding this comment.
Thanks for the suggestion. For this PR, we intentionally use a small thread pool with regular POSIX file I/O to keep the initial implementation simple and avoid introducing additional dependencies. I plan to evaluate existing io_uring libraries in a follow-up PR, together with dedicated performance benchmarks.
|
Great work! I've already verified the effectiveness and stability of offloading in long runs locally. That said, there's one point worth raising: Symptom: After enabling SSD offload, storage unit RSS actually increases (+75% on the data plane in our 4-node test), even though payloads are correctly written to disk and everything functions normally. Diagnosis: Call Mechanism: The contiguous payload buffers generated by Fix: Set Suggestion: Consider adding an RSS measurement to |
b1d5412 to
0aec5e7
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
| # Required when enabled. TransferQueue creates this directory and any | ||
| # missing parents on every storage node, stores data under its | ||
| # transfer_queue_ssd_offload/ child, and leaves this directory intact on close. | ||
| # Example: /path/to/local/ssd |
There was a problem hiding this comment.
We can simplify the comments. Best fit in one line
| # SSD offload reduces the long-lived host-memory footprint by keeping large | ||
| # fields on local SSD. Small fields remain in memory. | ||
| # Each offloaded (field, global_index) sample owns one temporary file so | ||
| # CLEAR can reclaim its disk space immediately. |
There was a problem hiding this comment.
I suggest we provide a dedicate doc to describe the implementation details
There was a problem hiding this comment.
I add a new file docs/ssd_offload.md.
|
|
||
| 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 |
There was a problem hiding this comment.
The default value is inconsistent with config.yaml
There was a problem hiding this comment.
Python：1024 * 1024 = 1048576
YAML：threshold_bytes: 1048576
They are both 1 MiB (1,048,576 bytes).
1f436c2 to
5f2c7eb
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
|
scripts/put_benchmark.py is broken: SimpleStorageUnit.init now only accepts config. Please help to fix this benchmark |
| @staticmethod | ||
| def _sample_from_value(value: Any) -> SSDEncodedSample | None: | ||
| if isinstance(value, torch.Tensor): | ||
| if value.is_nested or value.is_sparse: |
There was a problem hiding this comment.
Could we add support for sparse tensors now? Also, data larger than threshold_bytes but not eligible for SSD offload currently falls back to DRAM silently — could we add some metrics for this so that we can observe it?
There was a problem hiding this comment.
Thanks for the suggestion. I've added tq_storage_ssd_fallback_values_total, a per-storage-unit counter that tracks values that cannot be encoded for SSD offload and therefore remain in DRAM. This makes silent fallbacks observable through the metrics endpoint, and I've documented its semantics as well. For sparse tensors and broader data type support, I'd prefer to address those in a follow-up PR so that we can define and test the support matrix separately.
Signed-off-by: pinjie <pinjiex@nvidia.com>
5f2c7eb to
874934b
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
Impact of
|
| glibc_mmap_threshold_bytes | RSS before kv_clear | RSS after kv_clear | Matched in-memory run: RSS before kv_clear | kv_batch_put throughput | kv_batch_get throughput |
|---|---|---|---|---|---|
| unset (null) | 11.858 GiB | 11.787 GiB | 13.205 GiB | 44.14 Gbit/s | 36.69 Gbit/s |
| 128 KiB | 4.782 GiB | 4.045 GiB | 13.570 GiB | 44.53 Gbit/s | 38.24 Gbit/s |
| 1 MiB | 4.747 GiB | 4.081 GiB | 13.376 GiB | 44.90 Gbit/s | 37.99 Gbit/s |
| 4 MiB | 4.519 GiB | 3.799 GiB | 12.902 GiB | 44.72 Gbit/s | 38.12 Gbit/s |
The benchmark ran on two nodes with Arm Neoverse V2 CPUs and NVIDIA/Mellanox ConnectX-6 Lx networking. Each iteration processed 5 GiB of data, with 4 GiB (80%) offloaded to local SSD. Each configuration ran 10 consecutive PUT–GET–CLEAR iterations, and the table reports results from the final iteration.
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
Sync 11 main PRs (Ascend#162, Ascend#167, Ascend#168, Ascend#170-Ascend#172, Ascend#174, Ascend#175, Ascend#177-Ascend#179) into the payload-transfer branch. Main's PR history is preserved; the payload transfer abstraction stays on top. Conflict resolutions: - SimpleStorageUnit: keep main's config-dict constructor, hybrid memory/SSD storage, proxy identity filter, arrival counters, accept probe, and checkpoint delegation; the worker dispatches PUT/GET families to the configured PayloadTransfer as before. - AsyncSimpleStorageManager: keep main's socket pools, retry with diagnosis, and heavy-request logging; per-unit put/get delegate to the configured PayloadTransfer. - Bootstrap: keep main's config-dict actor construction plus payload transfer validation and NIXL endpoint collection. Semantic adaptations (not textual conflicts): - Add PayloadTransferTimeout (base.py) carrying wire_bytes; both backends raise it on an unanswered control message, and the manager translates it to StorageUnitTimeout so Ascend#171 retry/diagnosis covers the ZMQ and NIXL-UCX paths uniformly. - PayloadTransfer.put returns payload bytes moved, so Ascend#171's manager side heavy-put logging keeps true wire sizes per backend. Verified: python -m compileall -q transfer_queue tutorial tests clean; ruff check clean; pytest 622 passed, 8 skipped, 8 errors (the 8 errors are the pre-existing optional-Yuanrong-backend baseline).
## Summary
Adds two operations on existing KV keys:
- **`kv_update` / `kv_batch_update`** merge new values into fields that
are already stored, using `merge_fn(old, new)`. The merge runs on the
SimpleStorage unit that holds each sample, so the stored payload never
makes a round trip through the client.
- **`kv_clear(keys, partition_id, fields=...)`** clears only the given
fields and keeps the keys, for example to release `pixel_values` after
rollout while training still reads the rest of the sample. It works on
every storage backend. Without `fields`, `kv_clear` behaves as before.
To make multi-field updates all-or-nothing on a unit, SSD-offloaded
`put_data` is now two-phase, which also affects plain multi-field `put`
(see "SSD offload").
## Usage
### `kv_update`: append the response to the prompt
```python
import torch
import transfer_queue as tq
tq.kv_put(key="s0", partition_id="train", fields={
"input_ids": torch.tensor([10, 11, 12]),
"attention_mask": torch.tensor([1, 1, 1]),
})
tq.kv_update(
key="s0",
partition_id="train",
fields={"input_ids": torch.tensor([20, 21]), "attention_mask": torch.tensor([1, 1])},
merge_fn=lambda old, new: torch.cat([old, new]),
)
# input_ids -> tensor([10, 11, 12, 20, 21])
# attention_mask -> tensor([1, 1, 1, 1, 1])
```
For a batch, `kv_batch_update(keys, partition_id, fields: TensorDict,
merge_fn)` takes a TensorDict whose batch size matches `keys` (nested
tensors are allowed here). Both have `async_` variants.
### `kv_clear(fields=...)`: release pixel values after rollout
In a multimodal RL step, the rollout engine needs the raw
`pixel_values`, while training only needs the processed tensors. Once
rollout is done, the pixel values only take up storage:
```python
tq.kv_batch_put(keys=keys, partition_id="train", fields=TensorDict({
"input_ids": input_ids,
"pixel_values": pixel_values, # large; only rollout reads it
"mm_tensors": mm_tensors, # what training reads
}, batch_size=len(keys)))
# ... rollout reads pixel_values and writes the responses back ...
tq.kv_clear(keys=keys, partition_id="train", fields="pixel_values")
# Training reads input_ids / responses / mm_tensors exactly as before.
```
After a field-level clear:
- The storage backend no longer holds the values. With SSD offload on,
their SSD files are deleted.
- The field reads as never written for those keys: it is not ready, and
the controller no longer describes its old shape or dtype for them.
- The keys, their other fields, tags and consumption status are
unchanged, so samples are not consumed again.
- A later `kv_put` to the field writes it normally and marks it ready
again.
- Missing keys and fields are ignored, matching whole-key `kv_clear`.
Clearing every field still keeps the key in `kv_list`; removing the key
itself is still `kv_clear` without `fields`.
## Design
### `kv_update` (data plane, SimpleStorage only)
`interface.async_kv_batch_update` → `client.async_update` →
`AsyncSimpleStorageManager.update_data` → new ZMQ request `UPDATE_DATA`
→ `SimpleStorageUnit._handle_update` → `StorageUnitData.apply_update`.
- **Preconditions.** Every key must exist and every field must already
be produced for every key; otherwise `ValueError`. Nothing is created.
- **Compute, validate, then write once.** Each unit computes every merge
result and checks that each field keeps its tensor/non-tensor type and
tensor dtype, then commits with a single `put_data`. A raising
`merge_fn` or an incompatible result leaves that unit unchanged.
- **Schema.** Each unit reports the shapes it stored; the manager
assembles one batch-level `field_schema` in `global_indexes` order (the
controller reads `per_sample_shapes` by position) and notifies the
controller. The client merges it into its `BatchMeta` with the new
`BatchMeta.apply_field_schema` (`add_fields` now uses the same helper).
- **Controller ACK is checked.** `notify_data_update` now returns the
ACK's `success`. On the update path, a rejected or missing ACK raises
`RuntimeError` stating that storage may already have committed and the
merge must not be retried blindly.
- **No replays.** Updates use `max_attempts=1`, because a replay after a
lost reply would fold the same value in twice. A timeout therefore has
an unknown outcome: the unit may commit after the caller stops waiting.
This is documented on the API.
- **Atomicity is per storage-unit request.** A batch spanning several
units is not a distributed transaction.
- **Other backends.** The base `StorageManager.update_data` raises
`NotImplementedError("kv_update is not supported by <ManagerClass>; it
requires the SimpleStorage backend")`. KV backends store each
sample-field under its own key and have no hook to run `merge_fn` where
the value lives.
### Field-level `kv_clear` (control plane first, then storage)
`interface.async_kv_clear(fields=...)` → `client.async_kv_clear_fields`
→ new controller request `KV_CLEAR_FIELDS` →
`storage_manager.clear_data(metadata, fields=...)`.
- **Controller first.** `TransferQueueController.kv_clear_fields`
resolves the keys and captures the metadata of the affected cells,
including custom backend metadata that KV backends need for deletion.
Then `DataPartitionStatus.clear_fields` zeroes those production-status
cells, calls the existing `FieldMeta.remove_samples` (which already
falls back from nested to regular), and drops the per-field backend
metadata. It does not touch key mappings, tags, consumption status or
indexes, so there is no index-reuse race and no `MARK_CLEARING` step is
needed.
- **Then storage.** `clear_data` gains an optional `fields` filter. KV
backends (Mooncake, Yuanrong, Ray) narrow the metadata with
`select_fields`. SimpleStorage forwards `fields` in `CLEAR_DATA`, and
units drop only those fields, unlinking their SSD files; the sample
stays counted against unit capacity.
- The controller resolves fields per cell rather than relying on
`kv_retrieve_meta`, which only reports fields that every requested key
has, so a field present on only some keys is still cleared on those
keys.
### SSD offload: all-or-nothing multi-field put (#162)
`HybridStorageUnitData.put_data` used to write, commit and clean up one
field at a time, so a failure on a later field (for example, a full
disk) left earlier fields replaced. It is now two-phase: write every
field's SSD files first (on failure, unlink only the files this call
wrote and re-raise), then swap all references and unlink superseded
files. The base `StorageUnitData.put_data` checks every field's length
before writing any.
**This also affects plain multi-field `put` with SSD offload on.** The
trade-off is a higher temporary peak: the old files plus all replacement
files exist at once. `docs/ssd_offload.md` documents it.
### Shared backing memory
SimpleStorage slices dense batches without copying, so samples written
by one batched put share an allocation. Replacing or clearing some of
them does not return memory until every sibling is replaced or cleared.
`docs/ssd_offload.md` documents this alongside the existing
`data_parser` note.
## Docs and tutorial
- **README:** `(async_)kv_update` / `(async_)kv_batch_update` added,
`kv_clear` notes the `fields` option, and the
`StorageManager.clear_data` signature is updated.
- **`tutorial/02_kv_interface.py`:** a new step that concatenates a
field with `kv_update` and then clears it with `kv_clear(fields=...)`.
- **`docs/ssd_offload.md`:** shared backing memory and SSD headroom for
atomic multi-field writes.
## Testing
- `python -m compileall -q transfer_queue tutorial tests` and pre-commit
(ruff, ruff-format, mypy) are clean.
- `tests/e2e/test_kv_interface_e2e.py` on SimpleStorage with 2 units,
sync and async: **64 passed, 2 skipped** (the skips are KV-backend-only
rejection tests).
- Full local suite, `python -m pytest -q`: **699 passed, 12 skipped, 8
errors**. The errors are all in
`tests/test_yuanrong_storage_client_e2e.py` at fixture setup, because
the Yuanrong `datasystem` package is not installed locally.
- Not run: the Mooncake and Yuanrong e2e backends.
New and updated tests:
- `tests/test_kv_update.py`: merge, atomicity on merge error or
incompatible result, unproduced field rejection, decoded SSD-offloaded
old values, atomicity when a later field fails to reach SSD, cross-unit
schema assembly, KV-backend rejection naming the class.
- `tests/test_storage_request_retry.py`: the controller ACK status is
returned; a rejected ACK after a storage commit raises; an update
timeout is neither retried nor published.
- `tests/test_controller_data_partitions.py`: field-level clear keeps
keys, tags, consumption status and other fields; the cleared field's
schema for the remaining samples is unchanged; a later nested write to
the same field succeeds; `FieldMeta` is dropped once no sample holds the
field.
- `tests/test_simple_storage_unit.py`: field-level clear in memory and
on SSD (only that field's file is unlinked); all-or-nothing multi-field
SSD put.
- `tests/test_kv_storage_manager.py`: `clear_data(fields=...)` deletes
only those keys with their backend metadata.
- `tests/e2e/test_kv_interface_e2e.py`: `TestKVUpdateE2E` (update then
clear a field, then nested writes; batch update across units; rejected
arguments) and
`TestKVClearE2E::test_kv_clear_fields_keeps_keys_and_other_fields`.
## Out of scope
- `kv_batch_get(select_fields=...)` on a field that is not ready returns
an empty result rather than raising. This is existing behavior.
---------
Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Motivation
Large training samples can make the long-lived host-memory footprint of
SimpleStoragea capacity bottleneck as training scale and sequence length increase.Design
This PR adds an optional hybrid storage layer inside
SimpleStorageUnitwithout changing the public API, controller, or client behavior. Each(field, global index)sample is routed independently according to its actual stored payload size: samples at or above the internal 1 MiB threshold are written to node-local SSD, while smaller samples remain in host memory. The hybrid store preserves existing PUT, GET, CLEAR, overwrite, capacity, and checkpoint semantics; it reuses decoded contiguous buffers when possible and falls back to pickle for composite values. SSD files are temporary, isolated by run and storage unit, and removed during clear or graceful shutdown. The feature is disabled by default and configured throughssd_offload.enabledandssd_offload.path.E2E Performance
The following results are averages over four iterations with a batch size of 512 and eight fields. Throughput is reported in Gbit/s.
Test environment: 2 × Intel Xeon Platinum 8480C CPUs (224 cores), 10 × NVIDIA/Mellanox ConnectX-7 NDR 400 Gb/s InfiniBand adapters plus an Intel E810-C 100 GbE adapter, and 8 × KIOXIA CM6-R 3.84 TB NVMe SSDs (
KCM6DRUL3T84). 2 Node.