refactor(core): introduce backing-tier seam in prefetch pipeline#406
Open
FeathBow wants to merge 1 commit into
Open
refactor(core): introduce backing-tier seam in prefetch pipeline#406FeathBow wants to merge 1 commit into
FeathBow wants to merge 1 commit into
Conversation
Split storage/prefetch.rs into prefetch/{mod,scan,poll,task,state} and add
backing_tier.rs: RDMA and SSD now answer the same plan/fetch interface
(plan = metadata-only candidate discovery, fetch = data transfer), and
run_task walks the tier lineup instead of hardcoding RDMA-then-SSD branches.
Behavior-preserving: tier order, rdma_failed suppression, SSD budget
backpressure, RDMA shortfall terminality, attribution and lease paths are
unchanged and now pinned by fake-tier unit tests that need no io_uring or
IB hardware (first-tier short-circuit, plan fallthrough, budget cap and
exhaustion via real reservations, shortfall terminal + next-scan RDMA skip).
Also: drop the always-true emit_tier_metrics thread-through flag, merge the
isomorphic PrefetchSource/AttributionSource into TierSource, derive SSD
budget use from the tier source, and give skip_without_io_uring! a
PEGAFLOW_REQUIRE_IO_URING=1 hard-fail mode for merge gates.
There was a problem hiding this comment.
Pull request overview
Refactors pegaflow-core’s prefetch pipeline to introduce a backing-tier seam (RDMA/SSD via a shared plan() + fetch() interface), splitting the previous monolithic storage/prefetch.rs into focused submodules and adding unit-testable tier-orchestration semantics.
Changes:
- Added
storage/backing_tier.rsand rewired prefetch orchestration to walk an ordered tier lineup (RDMA-first, then SSD) instead of hardcoded branching. - Split prefetch into
prefetch/{mod,scan,poll,task,state}.rsand added unit tests using aFakeTierto validate tier selection, fallthrough, and SSD budget behavior. - Updated SSD-cache test gating (
PEGAFLOW_REQUIRE_IO_URING=1hard-fail mode), fixed a no-RDMA build import issue, and updatedCLAUDE.mddocumentation pointers.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| pegaflow-core/tests/ssd_cache.rs | Adds PEGAFLOW_REQUIRE_IO_URING=1 hard-fail behavior to avoid silently skipped SSD tests when io_uring is expected. |
| pegaflow-core/src/storage/tier_attribution.rs | Simplifies tier attribution plumbing to use TierSource and a TierDecision recorder. |
| pegaflow-core/src/storage/prefetch/tests/mod.rs | Introduces a shared test harness/utilities for the refactored prefetch modules. |
| pegaflow-core/src/storage/prefetch/tests/result.rs | Migrates/validates result assembly semantics (ordering + gap-stop). |
| pegaflow-core/src/storage/prefetch/tests/scheduler.rs | Adds scheduler-level tests for tier ordering and RDMA shortfall semantics. |
| pegaflow-core/src/storage/prefetch/tests/task.rs | Adds task-level tests for tier short-circuit, plan-miss fallthrough, and SSD budget reservation behavior. |
| pegaflow-core/src/storage/prefetch/mod.rs | New prefetch module entrypoint: Scheduler + shared types, hooks into StorageEngine. |
| pegaflow-core/src/storage/prefetch/scan.rs | Implements cache prefix scan + task spawning and tier lineup construction (RDMA slot + SSD slot). |
| pegaflow-core/src/storage/prefetch/poll.rs | Implements polling/draining task results and applying cache inserts + metaserver advertisement. |
| pegaflow-core/src/storage/prefetch/task.rs | Implements tier-walking task execution, SSD budget reservation, and ready-result assembly. |
| pegaflow-core/src/storage/prefetch/state.rs | Centralizes prefetch state (active tasks, SSD budget reservations, rdma-failed suppression). |
| pegaflow-core/src/storage/prefetch.rs | Removes the prior monolithic prefetch implementation (replaced by prefetch/ module). |
| pegaflow-core/src/storage/mod.rs | Wires StorageEngine to the new backing_tier + prefetch::Scheduler APIs and fixes a duplicate import. |
| pegaflow-core/src/storage/backing_tier.rs | Adds tier seam: BackingTier, TierPlan, TierQuery, TierSource, plus FakeTier for tests. |
| CLAUDE.md | Updates documentation references to reflect prefetch/ directory and new seam. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+73
to
+84
| async fn drain_task(req_id: &str, handle: JoinHandle<TaskResult>) -> TaskResult { | ||
| match handle.await { | ||
| Ok(result) => result, | ||
| Err(err) => { | ||
| warn!("Prefetch task failed for req_id={}: {}", req_id, err); | ||
| TaskResult { | ||
| source: None, | ||
| committed: 0, | ||
| inserts: Vec::new(), | ||
| ready_blocks: Vec::new(), | ||
| missing: 0, | ||
| } |
Author
There was a problem hiding this comment.
The underlying observation is real, but the failure mode is the opposite of a full cache hit. On JoinError the poll path returns Ready { blocks: [], missing: 0 }; the query response then reports num_hit_blocks = 0 and no lease is created, so the client recomputes everything — safe for correctness. The actual defect is observability only: missing: 0 suppresses the miss counter for that query. This PR is a behavior-preserving refactor, so I'd rather not change failure semantics here.
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
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.
Description
Refs #18
storage/prefetch.rs(686 lines) intostorage/prefetch/{mod,scan,poll,task,state}.rsand addstorage/backing_tier.rs: SSD and RDMA now answer the same two-phase interface —plan()(metadata-only candidate discovery: SSD index prefix / MetaServerquery_prefix) andfetch()(data transfer) — and the prefetch task walks a tier lineup instead of hardcoding RDMA-then-SSD branches.rdma_failedmarking suppresses RDMA on later scans for the request), SSD budget backpressure, tier attribution, and lease paths are unchanged.#[cfg(test)]FakeTier, and new unit tests pin those semantics: first-tier short-circuit, plan-miss fallthrough to the next tier, budget cap/exhaustion via real reservations, and shortfall terminality + next-scan RDMA skip. (The pure-function result-assembly tests already existed and are migrated, not new.) If feat(vllm): Add prefetch_threshold Parameter #18 adopts per-tier planned-length gating, the seam provides a single policy point for it.Changes
storage/backing_tier.rs:BackingTier(enum dispatch),TierPlan,TierQuery,TierSource;fetch()panics on a foreign tier's plan (unreachable!in the RDMA arm,assert!in the SSD/Fake arms).storage/prefetch/:mod.rs(scheduler entry),scan.rs(prefix scan + task start; one lock guard spans the check-then-spawn),poll.rs(poll and apply task results; guard released before anyawait),task.rs(tier walk),state.rs(active/rdma-failed/budget state;take_finishedchecks and removes a completed active task under one lock).emit_tier_metricsthread-through flag; merged the isomorphicPrefetchSource/AttributionSourceintoTierSource; SSD budget usage is now derived from the tier source.skip_without_io_uring!gains aPEGAFLOW_REQUIRE_IO_URING=1hard-fail mode: gates that set the variable fail loudly when io_uring is unavailable instead of silently skipping the SSD tests. Default behavior without the variable is unchanged (skip).use log::warnunder#[cfg(not(feature = "rdma"))]instorage/mod.rsthat failed every no-RDMA build with E0252.Test Env
Rust gates (rustc 1.97.0,
--no-default-features --features cuda-12,rdma):cargo clippy -p pegaflow-core --all-targets --no-default-features --features cuda-12,rdma -- -D warnings: cleanPEGAFLOW_REQUIRE_IO_URING=1 cargo test -p pegaflow-core --release --no-default-features --features cuda-12,rdma: lib 128 passed / 0 failed / 1 ignored; all 6 integration targets green (engine_basic 9, eviction 3, prefetch_lease 2, prefix_semantics 5, save_validation 2, ssd_cache 11)cargo test -p pegaflow-core --no-default-features --features cuda-12 --no-runpytest -m e2e tests/test_vllm_e2e_correctness.py --model <local Qwen3-4B> --max-model-len 4096on sm_89 x86_64, CUDA 12.8, vLLM 0.25.1, default sampler configuration (FlashInfer top-k/top-p engaged, confirmed in the server log).