fix(mocker): make G1 use allocation transactional [DYN-3299]

[PeaBrane]

Summary

• make each VLLM MoveBlock::Use allocation an all-or-nothing prepare/reserve/commit transaction, so failed attempts cannot publish phantom cache entries, partial active references, or Stored events
• propagate typed capacity, offload-dependency, and bounded immediate-retry outcomes through scheduling, ordinary decode growth, speculative decode reservation, destination reservation, and swap-in admission
• track the active G1 offload dependency that protects waiting work from capacity-triggered preemption while preserving finite mock-pipeline concurrency
• add a linearized per-pipeline snapshot and settlement token/target API so mocker advances only after drained batches and immediately runnable successor invocations are published
• document the balanced offline replay lifecycle contract at the zero-duration progress gate: prevent same-timestamp tight spins without mistaking unfinished work with no concrete wakeup for quiescence
• disable wall-clock request expiry only for offline replay, requiring explicit completion/cancellation while preserving the live router five-minute TTL

This implements the transactional allocation boundary described by #11016 and the semaphore-handoff settlement slice of #11018. Lifecycle watermarks and generalized auto-chain closure remain deferred.

Validation

Correctness and flake gate

Validated exact candidate 76e1d7afb7e4b7733e8e76bde6d3ca081c068094 on an SC-01 CPU-only x86_64 allocation:

• complete kvbm-engine: 125 passed
• complete kvbm-logical: 493 passed plus 3 doctests
• cargo test --locked -p dynamo-mocker --features kvbm-offload: 557 passed, 1 ignored
• full mocker feature suite: 20/20 clean repetitions
• idle_destination_reservation_wakes_at_offload_deadline: 200/200 repetitions
• live_destination_eviction_uses_command_application_time: 200/200 repetitions

Both known regressions fail on the exact baseline: the idle-destination case times out, and command-application-time eviction does not schedule G1→G2. Their candidate passes are intended fixes.

Deterministic tests also cover concurrency one, simultaneous completions, no successor, progress between subscription and initial check, block/object executors, cancellation, task abort, panic, executor error, shutdown, stale/foreign/unavailable tokens, overflow, G2-only release, partial chains, and speculative reservation.

White-box live path validation

Temporary tracing was applied only in the scratch candidate checkout, preserved as an artifact, then removed before rebuilding the exact candidate for A/B measurement.

• prefill completed 167 transfers / 333 blocks on each of G1→G2, G2→G3, and G2→G4
• decode completed 252 transfers / 420 blocks on each lane
• settlement targets exactly matched drained transfer counts and final completed-batch counters
• observed 1,257 permit releases and successor-worker invocations across block and object executors
• queued handoff was exercised (queued peaked at 4, starting at 2)
• released 753 G1 source slots; all final pipeline phase and pending-lane counts were zero
• router events: prefill 333 stored / 269 removed; decode 420 stored / 356 removed
• speculative decode: 476 bursts and 1,384 visible tokens

The traced path-validation load remained sub-saturation: frontend/prefill/decode CPU medians were 7.1%/10.6%/10.9%, with 39/39 successful measured requests.

Deterministic semantic A/B

Normalized traces remove wall timestamps, generated IDs, ports, paths, and absolute clock origins while preserving outputs, completion order, router decisions, modeled transfer durations, batch/block counts, and source-slot release.

• G2-only concurrency one: exact output token IDs/text, finish reasons, usage, completion order, transfer counts, modeled durations, and router metrics matched. The only event delta was intended: baseline released zero source slots for each of 101 drained one-block transfers; candidate released one.
• paced all-tier concurrency four: 24/24 exact replayed outputs in both arms, identical completion order, identical transfer/event multiset, and identical router metrics (147 stored / 22 removed)

Live before/after overhead ablation

Seven paired 60-second measured trials per workload after 15-second warmups, using closed-loop concurrency two, alternating arm order, paired seeds, and fresh services/discovery state per run:

• control: output-token throughput median -0.141% (bootstrap 95% CI -0.220% to +0.516%); p95 end-to-end latency +0.132% (CI -2.578% to +0.295%)
• offload-heavy: output-token throughput median -0.453% (CI -0.518% to -0.402%); p95 end-to-end latency +0.479% (CI +0.458% to +0.544%)
• all 28 counted runs had 100% successful requests, zero cancellations/errors, and zero output-length mismatches

All agreed performance gates passed.

Aggressive-tracing stress observation

With an intentionally tiny 12-block G1, concurrency four, and aggressive debug tracing, requests could not coexist once tracing stretched their overlap. The live decode scheduler entered 122,879 preempt/requeue passes, produced a 475 MB debug log, and seven clients timed out. The identical normal-logging workload completed 24/24 exact outputs in order.

The capacity preemption is expected for this artificial setup; the tight spin is a broader scheduler yield/backoff concern and is not attributed to pipeline settlement.

Summary by CodeRabbit

• New Features

  • Improved offload-aware scheduling and reservation handling, including offload dependency tracking so blocked work waits and resumes reliably.
  • Added richer transfer progress tracking for background transfers, including more precise completion timing and settle coordination.

• Bug Fixes

  • Reduced partial-commit behavior during capacity exhaustion and improved retry safety.
  • Prevented premature progress/possible lost-wakeup or livelock in same-timestamp replay progression.

• Documentation

  • Added/clarified offline replay liveness expectations and same-time progression guidance.
  • Enhanced sequence-expiry configuration for offline replay and explicit cleanup behavior.

[datadog-official]

⚠️ Warnings

🚦 1 Pipeline job failed

Docs link check | lychee     

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: f28bd54 | Docs | Give us feedback!}

[coderabbitai]

Walkthrough

The PR changes KV acquisition and reservation APIs to return G1Acquire outcomes, threads offload dependencies through the offload engine and vLLM scheduler, adds batch block registration and settlement tracking, and updates offline replay liveness documentation and comments.

Changes

G1 Acquire and Offload Flow

Layer / File(s)	Summary
Offload timing and worker notifications lib/mocker/src/kvbm_offload/bandwidth_sharing_model.rs, lib/mocker/src/kvbm_offload/worker.rs	The bandwidth model exposes transfer-aware deadlines and active checks, and the worker exposes those queries plus reservation-count watch updates and shared-G3 test helpers.
Settlement tracker and engine wiring lib/kvbm-engine/src/offload/settlement.rs, lib/kvbm-engine/src/offload/engine.rs, lib/kvbm-engine/src/offload/mod.rs, lib/kvbm-engine/src/offload/pipeline.rs	The offload engine adds settlement tokens, per-lane settlement tracking, released-slot accounting, and the settlement re-exports used by the pipeline code.
Acquire contract and use flow lib/mocker/src/kv_manager/kvbm_backend.rs	The KV manager returns typed acquire outcomes, tracks capacity generations, propagates offload dependencies, and stages use and eviction commits atomically.
Batch block registration lib/kvbm-logical/src/manager/mod.rs, lib/kvbm-logical/src/pools/store.rs, lib/kvbm-logical/src/registry/mod.rs, lib/kvbm-logical/src/metrics/pool_metrics.rs, lib/kvbm-logical/src/registry/tests.rs	The logical block manager, registry, store, and metrics add batched registration paths that preserve handle order and update registration counts per batch.
Scheduler offload gating lib/mocker/src/scheduler/source_holds.rs, lib/mocker/src/scheduler/vllm/core.rs, lib/mocker/src/scheduler/vllm/policy/tests.rs, lib/mocker/src/scheduler/vllm/tests.rs, lib/mocker/src/scheduler/vllm/live.rs	The vLLM scheduler stores offload dependencies per request, refreshes them during admission and preemption, and updates policy and scheduler tests to match typed acquire outcomes.

Offline Replay Liveness Docs

Layer / File(s)	Summary
Replay liveness guidance lib/mocker/src/replay/offline/CLAUDE.md, lib/mocker/src/replay/offline/components/engine.rs, lib/mocker/src/replay/offline/state.rs, lib/mocker/src/replay/router_shared.rs	The offline replay contract and engine comments describe queued-only passes, unfinished-work quiescence checks, same-time stopping conditions, and replay construction without expiry.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description covers the change and validation, but it misses the required Details, reviewer start, and Related Issues sections from the template.	Add the missing template sections, especially Details, Where should reviewer start?, and a completed Related Issues section with the issue link or no-issue checkbox.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: making G1 use allocation transactional in mocker.

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-4b297d6d-72c1-4f57-bde8-a0be39e588c0.docs.buildwithfern.com/dynamo/dev

[devin-ai-integration]

Devin Review found 2 potential issues.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/mocker/src/kv_manager/kvbm_backend.rs`:
- Around line 430-434: The offload dependency path in kvbm_backend.rs can return
Some(OffloadDependency) with both fields empty when
current_g1_offload_dependency() is None, which leaves refresh loops treating the
request as blocked without any wakeup. Update the logic around
has_nonterminal_g1_pipeline() and current_g1_offload_dependency() so that
offload blocking is only reported when there is a concrete future wakeup (real
transfer_id/deadline, scheduled completion, or similar notification), and
otherwise return no offload dependency. Use the existing OffloadDependency
construction in the current_g1_offload_dependency handling to locate and adjust
this branch.
- Around line 945-947: The eviction flow in
kvbm_backend::handle_evictions_with_source_slots_at is currently treated as
infallible by the caller that assigns to outcome and then unwraps with expect,
but it can legitimately return None when evicted PLHs do not map to
shadow-registry offload blocks. Update the allocation path at that call site to
handle the None case gracefully by returning a bounded retry/backoff instead of
panicking, and keep the existing dependency-outcome handling only for the
Some(outcome) path.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 12722e39-6679-4a54-8d6a-ed1e3c9d996b

📥 Commits

Reviewing files that changed from the base of the PR and between 4350517 and 95137a8.

📒 Files selected for processing (11)

• lib/mocker/src/kv_manager/kvbm_backend.rs
• lib/mocker/src/kvbm_offload/bandwidth_sharing_model.rs
• lib/mocker/src/kvbm_offload/engine.rs
• lib/mocker/src/kvbm_offload/mod.rs
• lib/mocker/src/kvbm_offload/worker.rs
• lib/mocker/src/replay/offline/CLAUDE.md
• lib/mocker/src/replay/offline/components/engine.rs
• lib/mocker/src/scheduler/source_holds.rs
• lib/mocker/src/scheduler/vllm/core.rs
• lib/mocker/src/scheduler/vllm/policy/tests.rs
• lib/mocker/src/scheduler/vllm/tests.rs

[dynamo-ops]

Architecture review: LGTM ✅

The core change — making G1 use allocation transactional via prepare_use/commit_use with an explicit G1Acquire<T> outcome enum — is well-designed. Key observations:

1. Atomicity is correctly preserved: The prepare phase collects existing handles (active/inactive hits) and reserves fresh slots in one batch via reserve_g1_slots, then commit applies all mutations. On any failure arm (CapacityExhausted, BlockedOnOffload), no partial state leaks — existing ImmutableBlock clones are dropped, unused MutableBlocks return to the reset pool via their RAII Drop.

2. RetryNow + validate_retry_witness: The bounded single-retry mechanism for presence-filtered evictions (where PLHs resolve to G2 shadow entries, releasing source slots without a DMA) is tight — asserts one retry max, positive slot delta, and generation advancement. The loop terminates.

3. Offload dependency propagation: The offload_dependency field on VllmRequestState + refresh_request_offload_dependency is clean. The refresh correctly clears the dependency when the specific transfer completes (checked by is_g1_offload_active) or when no nonterminal pipeline entries remain.

4. Speculative/decode gating: The early-exit on offload dependency before entering allocation loops (prefill admission, speculative batch, decode allocation) prevents wasted work. The sequence.pop() on BlockedOnOffload in the decode path correctly undoes the uncommitted speculative token.

5. CI test failure: The rust-tests failure should be investigated — likely a test not yet adapted to the new G1Acquire return type or a timing issue in the offload engine mock.

The two CodeRabbit suggestions (defensive handling for None dependency fields in refresh_offload_dependency, and graceful fallback on handle_evictions_with_source_slots_at returning None) are reasonable hardening but not blocking — the current code's expect() is a valid contract assertion given the offload-enabled feature gate.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[PeaBrane]

@coderabbitai full review

[coderabbitai]

✅ Action performed

Full review finished.

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

lib/kvbm-engine/src/offload/pipeline.rs (1)

1475-1507: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Finalize transfer states before returning block-transfer errors.

Line 1506 can return after states were marked Transferring/in_flight, but the affected TransferStates are never mark_failed/set_error, so TransferHandle::wait() can hang while settlement only records the batch failure. Apply the same terminal-state handling to the allocation and execute_local_transfer early-error paths too.

Proposed fix pattern

-            // Wait for transfer completion
-            notification.await?;
+            // Wait for transfer completion
+            if let Err(error) = notification.await {
+                let message = format!("block transfer failed: {error}");
+                for (_transfer_id, (state, block_ids)) in &transfer_states {
+                    let mut state_guard = state.lock().unwrap();
+                    state_guard.mark_failed(block_ids.iter().copied());
+                    state_guard.set_error(message.clone());
+                }
+                return Err(error);
+            }
             phase.mark_settling();

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/kvbm-engine/src/offload/pipeline.rs` around lines 1475 - 1507, The
transfer completion path in pipeline.rs leaves some TransferState entries stuck
in Transferring/in_flight when an error returns from notification.await or from
the earlier allocation/execute_local_transfer steps, which can make
TransferHandle::wait() hang. Update the transfer flow around
execute_local_transfer, notification.await, and the dst_blocks allocation path
to apply terminal failure handling before returning: mark each affected
TransferState as failed and set its error, and ensure phase/settlement state is
finalized consistently on all early-exit paths. Keep the existing successful
Transferring/in_flight transition, but add the same failure cleanup used
elsewhere in this transfer pipeline for both allocation failures and local
transfer execution failures.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/kvbm-logical/src/manager/mod.rs`:
- Around line 106-110: The registration counter in BlockPoolMetrics is being
used as a publication watermark, but inc_registrations_by() and snapshot() in
manager/mod.rs still rely on Ordering::Relaxed. Update the synchronization
around this watermark by using a release/acquire pair or a separate notifier so
that the batch’s store transitions and presence updates are fully visible before
another thread observes the increment. Keep the change centered on the
inc_registrations_by()/snapshot() flow and the check_presence()/swap-in planning
path.

---

Outside diff comments:
In `@lib/kvbm-engine/src/offload/pipeline.rs`:
- Around line 1475-1507: The transfer completion path in pipeline.rs leaves some
TransferState entries stuck in Transferring/in_flight when an error returns from
notification.await or from the earlier allocation/execute_local_transfer steps,
which can make TransferHandle::wait() hang. Update the transfer flow around
execute_local_transfer, notification.await, and the dst_blocks allocation path
to apply terminal failure handling before returning: mark each affected
TransferState as failed and set its error, and ensure phase/settlement state is
finalized consistently on all early-exit paths. Keep the existing successful
Transferring/in_flight transition, but add the same failure cleanup used
elsewhere in this transfer pipeline for both allocation failures and local
transfer execution failures.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f28e8ffa-12c1-4950-871a-97e4375d91f4

📥 Commits

Reviewing files that changed from the base of the PR and between f8a8591 and 76e1d7a.

📒 Files selected for processing (27)

• lib/kv-router/src/sequences/multi_worker.rs
• lib/kv-router/src/sequences/single.rs
• lib/kv-router/src/sequences/topology.rs
• lib/kvbm-engine/src/offload/engine.rs
• lib/kvbm-engine/src/offload/handle.rs
• lib/kvbm-engine/src/offload/mod.rs
• lib/kvbm-engine/src/offload/pipeline.rs
• lib/kvbm-engine/src/offload/settlement.rs
• lib/kvbm-logical/src/manager/mod.rs
• lib/kvbm-logical/src/metrics/pool_metrics.rs
• lib/kvbm-logical/src/pools/store.rs
• lib/kvbm-logical/src/registry/mod.rs
• lib/kvbm-logical/src/registry/tests.rs
• lib/mocker/src/kv_manager/kvbm_backend.rs
• lib/mocker/src/kvbm_offload/bandwidth_sharing_model.rs
• lib/mocker/src/kvbm_offload/engine.rs
• lib/mocker/src/kvbm_offload/mod.rs
• lib/mocker/src/kvbm_offload/worker.rs
• lib/mocker/src/replay/offline/CLAUDE.md
• lib/mocker/src/replay/offline/components/engine.rs
• lib/mocker/src/replay/offline/state.rs
• lib/mocker/src/replay/router_shared.rs
• lib/mocker/src/scheduler/source_holds.rs
• lib/mocker/src/scheduler/vllm/core.rs
• lib/mocker/src/scheduler/vllm/live.rs
• lib/mocker/src/scheduler/vllm/policy/tests.rs
• lib/mocker/src/scheduler/vllm/tests.rs

💤 Files with no reviewable changes (1)

• lib/mocker/src/replay/offline/state.rs

[PeaBrane]

@coderabbitai Addressed the block-transfer failure-path finding in e7d93b0. Destination-allocation, synchronous transfer-start, and asynchronous completion errors now mark every affected block failed and publish a terminal TransferState error before returning, so TransferHandle::wait() cannot remain stranded. Added focused coverage for both queued and in-flight handles. The separate relaxed registration-watermark comment remains intentionally deferred.

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

@PeaBrane Thanks — that addresses the block-transfer failure-path concern. Marking every affected block failed and publishing a terminal TransferState error on destination-allocation, sync start, and async completion failures is the right fix, and the added queued/in-flight coverage helps validate that TransferHandle::wait() won’t get stranded.

Keeping the relaxed registration-watermark comment deferred is fine. If you want, I can help open a follow-up issue for that separately.