Skip to content

[feat] simple storage support ssd offload - #162

Draft
xupinjie wants to merge 1 commit into
Ascend:mainfrom
xupinjie:pinjie/simple_storage_offload_ssd
Draft

xupinjie wants to merge 1 commit into
Ascend:mainfrom
xupinjie:pinjie/simple_storage_offload_ssd

Conversation

@xupinjie

@xupinjie xupinjie commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Motivation

Large training samples can make the long-lived host-memory footprint of SimpleStorage a capacity bottleneck as training scale and sequence length increase.

Design

This PR adds an optional hybrid storage layer inside SimpleStorageUnit without 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 through ssd_offload.enabled and ssd_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.

Per-field sample size Total payload Memory PUT SSD PUT Memory GET SSD GET Memory E2E SSD E2E E2E change
1 MiB 4 GiB 55.35 51.33 29.58 28.42 38.34 36.44 -5.0%
2 MiB 8 GiB 57.17 53.53 25.37 24.65 35.04 33.73 -3.7%
4 MiB 16 GiB 58.18 50.80 27.60 26.06 37.39 34.41 -8.0%

@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

xupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 1, 2026

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review Completed 2026-09-01T04:01:12.713288Z ed97682 PR opened
ℹ️ 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" or "@codex security review".

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 chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Use an SSD-compatible result plotter

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 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I delete the draw command.

@0oshowero0
0oshowero0 requested a balanced review from Copilot September 1, 2026 07:35

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@@ -0,0 +1,85 @@
#!/bin/bash

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we integrate the SSD test in existing test shell scripts?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread tests/test_simple_storage_unit.py Outdated
storage.close()


def test_hybrid_storage_round_trips_raw_tensor_samples(tmp_path):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We may not need some of the tests here. Consider simplify the test.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deleted some test.

Comment thread transfer_queue/config.yaml Outdated
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

@0oshowero0 0oshowero0 Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@xupinjie
xupinjie force-pushed the pinjie/simple_storage_offload_ssd branch from ed97682 to 6b636de Compare September 1, 2026 15:33
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

xupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍

@xupinjie
xupinjie force-pushed the pinjie/simple_storage_offload_ssd branch from 6b636de to 988f095 Compare September 2, 2026 02:10
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

xupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍

@xupinjie
xupinjie force-pushed the pinjie/simple_storage_offload_ssd branch from 988f095 to b1d5412 Compare September 6, 2026 11:08
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

xupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Comment on lines +978 to +980
storage_unit_size: int | None = None,
ssd_config=None,
ssd_run_id: str | None = None,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread transfer_queue/config.yaml Outdated
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;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should allow the case when parent dir not exist.

Comment on lines +368 to +372
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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can try to find some existing library that can help us to enable io_uring

@huniu20

huniu20 commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

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 malloc_trim(0) after CLEAR and check whether RSS drops. If it does, the memory is unreturned free heap space (high-water mark), not data still being held. Our test confirms it drops, so this isn't a reference leak in the code.

Mechanism: The contiguous payload buffers generated by _build_encoded_samples for disk writes trigger glibc's dynamic mmap threshold adjustment — after the first large mmap'd block is freed, the threshold is raised to that block's size, and subsequent same-sized allocations are served from the heap and no longer returned to the OS. RSS stays at the high-water mark as a result.

Fix: Set MALLOC_MMAP_THRESHOLD_ on the storage actor below the typical payload size. Scope it narrowly — setting it globally imposes unnecessary mmap/munmap overhead on training/inference hot paths that benefit from heap reuse.

Suggestion: Consider adding an RSS measurement to run_ssd_offload_perf_test.sh. This PR's motivation is reducing host memory footprint, but the current self-test only covers throughput.

@xupinjie
xupinjie force-pushed the pinjie/simple_storage_offload_ssd branch from b1d5412 to 0aec5e7 Compare September 17, 2026 10:13
@ascend-robot

Copy link
Copy Markdown

CLA Signature Pass

xupinjie, thanks for your pull request. All authors of the commits have signed the CLA. 👍

@xupinjie
xupinjie marked this pull request as draft September 17, 2026 10:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants