feat: add CPU staging buffer to access GPU memory via unified buffer - #37
Open
WillowWang0216 wants to merge 94 commits into
Open
feat: add CPU staging buffer to access GPU memory via unified buffer#37WillowWang0216 wants to merge 94 commits into
WillowWang0216 wants to merge 94 commits into
Conversation
Design for reclaiming SSD bucket space only via Remove/BatchRemove: - tombstone + background copy-on-write compaction in BucketStorageBackend - no LRU/FIFO eviction (must keep eviction_policy=none) - RemoveByRegex/RemoveAll unchanged, isolated from bucket GC - concurrency via existing mutex_ + BucketReadGuard + inflight_reads_
- Config: eviction_policy=LRU + disable_ssd_eviction=true * disable_ssd_eviction makes PrepareEviction a no-op (no bucket deletion) * eviction_policy=LRU keeps last_access_ns_ updated and lru_index_ maintained - Drop new gc_last_read_ns_ field; reuse last_access_ns_/lru_index_ - IsEnableOffloading quota check handles space-exhausted offload failure - GC reuses SelectEvictionCandidate LRU logic + deleted_bytes_>0 filter
Four diagrams covering: tombstone marking, compaction with concurrent read + two-phase validation, space-exhausted offload failure, and GC vs concurrent Remove non-loss guarantee.
Remove ASCII commas/semicolons from mermaid message text (parsed as participant separator / statement terminator) and split multi-action statements. Use full-width punctuation or rewrite sentences.
13 TDD tasks covering: MarkRemoved interface, BucketMetadata GC fields, config + env parsing, candidate selection, copy-on-write compaction, background GC thread, concurrency tests, FileStorage forwarding, RealClient wiring, deployment docs, and final verification.
… selection Tasks 1-6 of the SSD explicit-delete-only GC plan: - Add MarkRemoved/BatchMarkRemoved virtual no-op to StorageBackendInterface - Add deleted_bytes_/compacting_ runtime fields to BucketMetadata (with copy/move ctor reset handling) - Add GC config fields + env var parsing to BucketBackendConfig - Declare GC methods/members on BucketStorageBackend (incl. <thread>/<condition_variable> includes) - Implement MarkRemoved/BatchMarkRemoved tombstone marking (delete key from object_bucket_map_, bump deleted_bytes_, no disk IO) - Implement SelectGCCandidate (deleted_ratio + LRU coldness ordering) Design: docs/superpowers/specs/2026-06-25-ssd-explicit-delete-gc-design.md Plan: docs/superpowers/plans/2026-06-25-ssd-explicit-delete-gc.md Build not verified on this host (no C++ toolchain; Linux devcontainer only).
Task 7: CompactBucket rewrites surviving live keys into a new bucket via copy-on-write, atomically swaps mappings under mutex_ with re-validation, then deletes old bucket files after in-flight reads drain. Includes WaitForInflightReads (reuse FinalizeEviction pattern) and DeleteBucketFiles helpers. Reuses BuildBucket/WriteBucket/StoreBucketMetadata and vector_read.
Task 8: GCThreadFunc loop sleeps gc_interval_ms, checks space pressure (shared lock on mutex_ for total_size_), selects candidates via SelectGCCandidate + deleted_ratio threshold (or forced under space pressure), compacts up to gc_max_buckets_per_round. Thread started in Init (when gc_enable), stopped+joined in destructor. Uses std::mutex for gc_mutex_ to be compatible with std::condition_variable (matches existing master_service TimerLoop pattern).
…y tests Task 9: tests for MarkRemoved (hides key, idempotent), CompactBucket (reclaims deleted keys, preserves live key data integrity), concurrent MarkRemoved+BatchLoad, and disable_ssd_eviction no-op under space pressure.
…ient Tasks 10-11: - FileStorage::MarkRemoved/BatchMarkRemoved forward to storage_backend_ (dispatches to BucketStorageBackend tombstone; no-op for file-per-key) - RealClient::remove_internal calls file_storage_->MarkRemoved(key) after successful master Remove - RealClient::batchRemove_internal collects successfully removed keys and calls file_storage_->BatchMarkRemoved(removed) - Remove/BatchRemove interface signatures unchanged
Task 12: document the required eviction_policy=LRU + disable_ssd_eviction=true config and the GC tuning environment variables in ssd-offload.md.
GetEnvOr<T> uses std::stoll internally which cannot parse fractional values like '0.25'. Parse gc_deleted_ratio and gc_high_watermark_ratio manually with std::stod so the env config actually takes effect.
The methods were declared after IsEnableOffloading which is in the private section, making them inaccessible to RealClient. Move them to the public section (before 'private:').
CompactBucket was private, but tests (and explicit compaction) need to call it directly. Move it to the public section alongside DeleteBucket, which serves the same role. Internal helpers SelectGCCandidate/ GCThreadFunc/WaitForInflightReads/DeleteBucketFiles stay private.
…uota Three test failures fixed: 1. Self-deadlock (inflight_reads timeout): BucketReadGuard was held across WaitForInflightReads, so the bucket waited for its own inflight read to drain. Scope the guard to the Step 2 read block so it releases before Step 5 waits. The empty-bucket path needs no guard (it doesn't read). 2. Data corruption (b1 got v3's data): Step 4 re-mapped live_keys[i] to new_metas[i], but BuildBucket iterates an unordered_map in its own order, so new_metas is aligned with build_result->keys, NOT live_keys. Re-validate using new_bucket->keys instead. 3. DisableEvictionNoopUnderPressure: set total_size_limit (checked by IsEnableOffloading) instead of max_total_size (only used by the no-op'd PrepareEviction) so the quota check actually rejects.
Plan A: gc_e2e_test.cpp exercises the full pipeline through RealClient with enable_ssd_offload=true (BucketStorageBackend), unlike the existing storage_backend_e2e_test which uses the Client base class + file-per-key. Three tests: - RemoveReclaimsSSDSpace: put 2 keys, remove 1, verify survivor stays readable and GC compaction reclaims bucket files - RemoveMiddleKeyPreservesSurvivors: remove middle of 3 keys, verify copy-on-write preserves both survivors with correct data - BatchRemoveMixedExistingAndAbsent: batch remove with absent key, verify existing key is gone after GC Sets env vars (eviction_policy=lru, disable_ssd_eviction=true, gc_interval_ms=200, gc_deleted_ratio=0.1) per the GC config requirement.
Plan B: run_gc_e2e.sh launches real production binaries (mooncake_master + mooncake_client/real_client_main) as separate OS processes and drives put/get/remove via Python MooncakeDistributedStore, which connects to the real_client RPC server — the only path that triggers RealClient::remove_internal -> MarkRemoved -> GC. Sets GC env vars (eviction_policy=lru, disable_ssd_eviction=true, gc_interval_ms, gc_deleted_ratio) before launching the client so BucketBackendConfig picks them up. Verifies: 3 keys offloaded, middle key removed, GC compacts bucket, both survivors stay readable with correct data, removed key stays gone.
Existing e2e tests link only 'gtest' (not 'gtest_main') and provide their own main(). Add the same main() entry point to gc_e2e_test.cpp to resolve 'undefined reference to main' linker error.
getenv() returns nullptr when the env var is unset; assigning that to std::optional<std::string> constructs std::string(nullptr) which throws 'basic_string: construction from null is not valid'. Add GetEnvOpt() helper that returns std::nullopt for unset vars.
PutEnd sets an object lease; Remove with force=false is rejected with OBJECT_HAS_LEASE (-706) while the lease is active. Use force=true for all remove/batchRemove calls in gc_e2e_test so the tombstone marking proceeds immediately after put.
The tests were passing vacuously: PutAndWaitReadable returned as soon as the key was readable from MEMORY, but offload to BucketStorageBackend (async via heartbeat) hadn't completed yet. MarkRemoved was a no-op (key not in object_bucket_map_), so GC never ran — 'buckets_before=0'. Replace PutAndWaitReadable with PutAndWaitOffloaded which waits for a .bucket file to appear in the SSD dir AND get_buffer to return correct data. Test 2 now polls for bucket-file-count change to detect compaction rather than a fixed sleep.
Increase wait to 40s (heartbeat interval is 10s). On timeout, log the watched dir path, list all files in ssd_dir and tmp_dir_ tree to find where bucket files actually land.
Setting root_fs_dir made the master return a non-empty fsdir, causing the client to set up its own file-per-key StorageBackend and write DISK replicas to master-side disk cache (path .../mooncake_cluster/n/l/key). Master then sees a DISK replica exists and never pushes an offload task, so the client's FileStorage/BucketStorageBackend never gets a task via heartbeat and no .bucket file is written. With enable_offload=true but no root_fs_dir, master pushes offload tasks that the client's FileStorage drains via OffloadObjectHeartbeat, writing .bucket files to the BucketStorageBackend's storage_path_.
Set default_kv_lease_ttl=300000 (matches production). Set MOONCAKE_OFFLOAD_FILE_STORAGE_PATH and MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR env vars before client init (same as production deployment) so FileStorageConfig::FromEnvironment picks them up.
Log whether PushOffloadingQueue is called, succeeds, or fails (with error code) at PutEnd, plus whether the offload path is even entered. Temporary diagnostic to find why offload tasks aren't reaching the client heartbeat.
Log OffloadObjectHeartbeat result (ok/error, task count) to find why offload tasks aren't being drained even though PushOffloadingQueue OK.
…_object Log BatchQuerySlices result (ok/error, got count) and when batch_object is empty (skipping BatchOffload). This pinpoints where the offload data preparation fails.
…op entry Pinpoint whether AllocateOffloadingBuckets returns empty buckets_keys or the for-loop is never entered.
Root cause: GroupOffloadingKeysByBucket puts keys into the ungrouped pool until bucket_keys_limit (default 500) keys accumulate, then writes a .bucket file. With only 1-3 test keys, no bucket was ever written. Set MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT=1 so each offloaded key immediately fills a bucket and a .bucket file is written on the next heartbeat.
在 invoke_rpc 和 invoke_batch_rpc 中添加 MC_RPC_BEGIN/MC_RPC_END 日志, 打印 Mooncake 的 trace_id 和函数名, 通过时间戳与 yalanting 的 CTIMING/TIMING 日志对齐, 实现端到端 trace 关联。
This reverts commit 736a052.
在Mooncake线程构造lambda时用CurrentTraceId()捕获trace_id, 通过request_config_t.trace_id传入yalanting,使CTIMING日志能 关联Mooncake的trace_id,解决并发请求无法对齐的问题。
This reverts commit 87078a9.
Supercache->supercache-snapshot
fix(tracing): 添加缺失的 span 结束调用
…失;3. 为 MarkRemoved 增加持久化 tombstone；4.修复 compaction 与删除并发时的版本校验；
…失;3. 为 MarkRemoved 增加持久化 tombstone；4.修复 compaction 与删除并发时的版本校验；
…失;3. 为 MarkRemoved 增加持久化 tombstone；4.修复 compaction 与删除并发时的版本校验；
ssd remove problem fix
依据 vllm_spdiag_logging_plan.md，按 Q1b/Q2A/Q3b 方案为 Mooncake Store vLLM 调用链添加 SpDiag 性能打点与日志输出： - mooncake_perf_points.def: 追加 25 个新 PerfKey，覆盖 store_py 入口 层（STORE_PY_*）、RealClient 核心层（RC_*）、Client 服务层 （CLIENT_BATCH_QUERY）。 - store_py.cpp: 为 setup/register_buffer/batch_put_from_multi_buffers/ batch_get_into_multi_buffers/batchIsExist/remove_all/close 等 8 个 Python 绑定方法添加 PerfPoint（KEY_MODULE），remove_all 与 close 额外输出 MC_LOG 汇总行。 - real_client.cpp: 为 setup_real/batchIsExist/batch_put_from_multi_ buffers/batch_get_into_multi_buffers 等 9 个方法添加 PerfPoint； 下沉层 _internal 方法用 MODULE 级别，并按 Q1b 输出 MC_LOG 汇总行 + per-key 多行（包含 success/key/size/replica/endpoint 字段， 缺失字段不输出）。 - client_service.cpp: 在 L1117 BatchQuery 实际实现中添加 CLIENT_BATCH_QUERY PerfPoint（Q3b，仅实际实现，避免重复统计）。
SpDiag::PerfPoint v1.0.0 仅暴露 Start()/End()/Abandon() 三个方法， 没有 ElapsedMicros()。服务器构建报错： real_client.cpp:1218: error: 'class SpDiag::PerfPoint' has no member named 'ElapsedMicros' 修复方式：在 setup_real / batchIsExist / register_buffer 以及 store_py.cpp 的 remove_all / close 入口层用 std::chrono::steady_clock 手动测量 elapsed_us，与下沉层 batch_get_into_multi_buffers_internal 已有的 t0/t1/total_us 风格保持一致。
现有 stress_cluster_bench 只调用老接口 batch_get_into，无法触发 vllm 调用路径的新打点（batch_put_from_multi_buffers / batchIsExist / batch_get_into_multi_buffers）。 按方案文档第 6 章实现 store_connector_bench.cpp： - 4 个 scenario: write / is_exist / get / all - 模拟 vllm 多 layer KV cache 场景（每请求 num_layers 个 key + buffer） - 调用 RealClient::batch_put_from_multi_buffers / batchIsExist / batch_get_into_multi_buffers，触发 STORE_PY_* + RC_* + CLIENT_BATCH_QUERY SpDiag 打点 + MC_LOG 汇总日志 - 输出带宽 + 时延分位数（P50/P90/P99） - 同步追加 CMakeLists.txt 编译目标，链接库参考 stress_cluster_bench
malloc 返回的内存不保证页对齐和 NUMA 本地性，UB/RDMA driver 会拒绝注册为 DMA 内存，导致 register_buffer 失败： Failed to register segment ... : Success [0] UbTransport: cannot register LocalMemory 参考 stress_cluster_bench 的做法，改用 numa_alloc_local (mmap + page-aligned + NUMA-aware)，匹配 UB driver 的要求。 析构同步改用 numa_free。
batchIsExist 返回 1=存在(成功),0=不存在(失败) batch_get_into_multi_buffers 返回 >0=字节数(成功),<0=错误码(失败) 原代码统一用 v!=0 判断失败,导致 IS_EXIST 和 GET 全部误判为失败 修复:按 phase 分别判断 - IS_EXIST: v != 1 为失败 - GET: v <= 0 为失败 - WRITE: v != 0 为失败(保持不变)
feat: add CPU staging buffer to access GPU memory via unified buffer
feat(spdiag): 为 vLLM 调用路径添加 PerfPoint 打点与 MC_LOG 日志
github-actions
Bot
added
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When integrating vllm with mooncake, GPU memory addresses need to be passed to mooncake for data transfer. However, the existing ubtransport implementation does not support registering GPU memory directly, making it impossible to fetch data from the GPU side.
To address this, a CPU staging buffer mechanism is introduced:
Add a CPU staging area as an intermediate buffer between GPU memory and ubtransport
Before data transfer, copy data from GPU memory to the CPU staging buffer (D2H)
mooncake reads data from the CPU staging buffer via ubtransport
For reverse transfers, data is first written to the CPU staging buffer, then copied back to GPU memory (H2D)
This approach enables indirect access and transfer of GPU memory data without modifying the ubtransport registration logic.