[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. 👍 |
Signed-off-by: xupinjie <pinjiex@nvidia.com>
6b636de to
988f095
Compare
CLA Signature Passxupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍 |
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.