Skip to content

fix(ublk): decouple proactive prewarm from idle cache capacity - #234

Open
ztting01 wants to merge 6 commits into
kvcache-ai:mainfrom
ztting01:fix/ublk-decouple-prewarm-capacity
Open

fix(ublk): decouple proactive prewarm from idle cache capacity#234
ztting01 wants to merge 6 commits into
kvcache-ai:mainfrom
ztting01:fix/ublk-decouple-prewarm-capacity

Conversation

@ztting01

@ztting01 ztting01 commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

What

Add an optional [pool.block].prewarm_high_watermark that limits only proactive OverlayBD ublk-device refill, independently from the normal idle-device cache capacity in [pool].high_watermark.

Returned workload devices may still populate the idle cache up to high_watermark. Proactive refill stops at prewarm_high_watermark.

This change also:

  • adds --pool-prewarm-high-watermark to uvm-ublk-daemon;
  • defaults an omitted prewarm watermark to high_watermark, preserving the previous behavior;
  • rejects prewarm_high_watermark > high_watermark;
  • logs both limits when enabling the OverlayBD warm pool;
  • documents the new optional setting and includes a commented prewarm_high_watermark = 8 example in config/default.toml; operators must enable it explicitly;
  • adds configuration, refill-policy, and live ublk tests.

Why

PR #205 made block prewarm capability-aware and explicitly left concurrent refill-target amplification out of scope.

On kernels that support UBLK_F_UPDATE_SIZE, or when operators explicitly enable startup_prewarm, a concurrent sandbox burst can still grow the refill target to the shared pool high watermark. With high_watermark = 64, proactive refill can create up to 64 idle ublk devices even when the workload also has many instance-owned rootfs devices. This produces a large startup device spike and leaves little room for real devices returned during pause.

Lowering the shared high_watermark is not equivalent: it also reduces the useful cache capacity for devices returned by completed or paused workloads. Operators need to retain a larger reuse cache while independently bounding speculative prewarm.

Related issue

Follow-up to #205 and #203. PR #205 fixed the no-resize default behavior and identified concurrent refill-target amplification as a separate issue.

Scope and non-goals

Included:

  • block-pool config and CLI plumbing for prewarm_high_watermark;
  • proactive refill bounded by the block-specific limit;
  • backward-compatible fallback to high_watermark when the setting is omitted;
  • config validation, startup logging, documentation, unit tests, and a live ublk test.

Not included:

  • changes to the normal returned-device cache limit;
  • changes to pause, resume, release, or device-deletion synchronization;
  • changes to the generic warm-pool growth policy or the Firecracker pool;
  • kernel backports or alternative implementations of ublk dynamic resize.

Design and behavior changes

The daemon now loads two related limits:

  • PoolConfig::high_watermark: maximum idle cache capacity, including devices returned by real workloads;
  • prewarm_high_watermark: maximum idle count that proactive refill may reach.

prewarm_refill_count() first asks the existing warm-pool policy how many entries it wants to create, then caps that count by the remaining block-specific prewarm capacity.

The refill loop checks the prewarm limit both before creating another device and before inserting the newly created device. The second check handles concurrent releases that may increase the idle count while an asynchronous device creation is in progress; an excess device is stopped through the existing cleanup path.

The existing refill_inflight guard still ensures that only one refill task is elected at a time. Rescheduling now occurs only while the prewarm-specific limit has remaining capacity.

An omitted prewarm_high_watermark resolves to high_watermark, so existing configurations retain the previous coupled behavior. It is intentionally valid for the prewarm watermark to be below the shared low watermark: the block-specific limit is an operator cap on speculative creation, while returned devices can still warm the larger cache naturally.

Compatibility and operations

  • Public API or generated protocol: No change.
  • Configuration or defaults: Adds optional [pool.block].prewarm_high_watermark. Omission preserves the old effective limit. config/default.toml leaves prewarm_high_watermark = 8 commented; operators can uncomment it to opt in while [pool].high_watermark remains 64.
  • Snapshot manifest, artifact layout, or storage format: No change.
  • Upgrade and rollback: Safe with no data migration. Remove the new key to restore the historical coupled behavior. Values above high_watermark fail startup with an actionable error.
  • Host requirements, permissions, ports, or dependencies: No change.

Validation

  • make fmt
  • make clippy
  • make test-unit
  • Relevant Rust integration tests
  • make -C services test (required when services/ changes)
  • Generated clients/server regenerated with the documented make target
  • Documentation updated
  • Benchmarks or performance comparison completed

Commands and results:

cargo fmt --all -- --check
# passed

git diff --check origin/main...HEAD
# passed

cargo clippy -p agentenv -p uvm-ublk-daemon --all-targets -- -D warnings
# passed

cargo test -p agentenv \
  block_pool_prewarm_high_watermark_rejects_values_above_cache_capacity
# 1 passed

cargo test -p uvm-ublk-daemon prewarm_high_watermark
# 3 passed

cargo test -p uvm-ublk-daemon \
  prewarm_refill_count_is_capped_below_idle_cache_capacity
# 1 passed

cargo test -p uvm-ublk-daemon \
  live_prewarm_refill_creates_only_the_block_specific_limit -- --ignored
# 1 passed on a host with real /dev/ublk-control access

API_ADDR=127.0.0.1:8000 AENV_RUN_USER=root make start-server-release
# passed; the release server was used for the replay benchmark below

Replay benchmark

Host and workload:

  • openEuler 24.03 LTS-SP3, aarch64, bare metal;
  • real uvm-ublk-daemon and /dev/ublk-control;
  • ublk feature mask 0x1fe; UBLK_F_UPDATE_SIZE unavailable;
  • template: django-money-task-v2;
  • 128 trajectories submitted concurrently, with at most 64 sandboxes running;
  • zero launch interval, control-plane QPS limit 1000, and a 600-second per-trajectory timeout;
  • the same trajectory set and server build were used for every round;
  • each round started from a clean server restart with zero existing sandboxes;
  • host, ublk-device, CPU, memory, disk, and sandbox-state metrics were sampled every 100 ms;
  • rounds were interleaved as A1, B1, A2, B2 to reduce ordering bias;
  • startup_prewarm = true and [pool].high_watermark = 64 in both variants;
  • two retained repetitions per variant.

Compared configurations:

# A: historical coupled behavior
[pool]
high_watermark = 64

[pool.block]
startup_prewarm = true
# prewarm_high_watermark omitted -> 64
# B: decoupled proactive refill
[pool]
high_watermark = 64

[pool.block]
startup_prewarm = true
prewarm_high_watermark = 8
Metric Coupled limit 64, run 1 / run 2 Prewarm limit 8, run 1 / run 2 Mean change
ublk device maximum 177 / 168 136 / 136 -21.2%
ublk device p95 132 / 126 92 / 84 -31.8%
first-60s ublk maximum 158 / 165 136 / 136 -15.8%
ublk-daemon CPU p95 2.870 / 2.754 cores 2.086 / 1.995 cores -27.4%
elapsed time 623.94 / 621.08 s 613.32 / 609.55 s -1.8%
resume API p95 3.390 / 2.697 s 2.793 / 2.579 s -11.7%
pause API p95 2.551 / 2.176 s 1.949 / 1.937 s -17.8%
successful trajectories 121/128 / 121/128 121/128 / 121/128 no change

Both runs with the new limit peaked at 136 devices, consistent with up to 128 workload-owned devices plus 8 proactively prewarmed devices. Both coupled-limit runs had higher peaks. The primary acceptance criterion was therefore met in both repetitions, with no success-rate regression.

Pause/resume latency moved in a favorable direction, but two repetitions are not enough to claim a statistically significant latency improvement. These measurements are treated as a non-regression signal; the primary result is the lower ublk device footprint.

The benchmark explicitly set startup_prewarm = true, so it exercises the operator opt-in path retained by #205 on this no-resize kernel.

An unrelated TensorFlow/Bazel build started during the tail of the second coupled-limit run, after that run's device peak had already occurred. Two prewarm-limit attempts that overlapped external builds were discarded. The retained second prewarm-limit run started after a 61-second quiet window and was checked throughout for absence of the external CI process. This interference does not affect the two-run device-peak comparison, but it is another reason not to overstate the secondary latency numbers.

Skipped checks and reasons:

  • Full make test-unit was not run; the tests added or directly affected by this change were run explicitly, including the ignored live ublk test.
  • Full repository make clippy was not run; Clippy passed with warnings denied for the two changed Rust packages and all their targets.
  • make -C services test is not applicable because services/ is unchanged.
  • Generated-code checks are not applicable because no generated API or protocol changed.

Risks and reviewer notes

  • The repository default keeps the setting commented, so omitted configurations retain the historical effective limit of 64. Operators can explicitly set prewarm_high_watermark = 8 to reduce proactive refill while retaining a returned-device cache capacity of 64.
  • This setting does not cap total ublk devices. Active and paused sandboxes still own devices; it only caps additional proactively created idle devices.
  • Returned devices may still grow the idle cache above prewarm_high_watermark, up to the existing high_watermark. This is intentional and preserves reuse capacity.
  • Review the second prewarm-limit check in refill_idle_pool: it prevents an asynchronous create from overshooting when another task releases a device while creation is in flight.
  • Key files: storage/ublk-daemon/src/server.rs, storage/ublk-daemon/src/main.rs, and src/cfg.rs.

Checklist

  • The PR contains one coherent change and no unrelated formatting or refactoring.
  • New behavior is covered by tests, or I explained why testing is impractical.
  • Logs and examples contain no credentials, tokens, or private registry information.
  • I did not manually edit generated code without updating its source and regenerating it.

@ztting01
ztting01 marked this pull request as ready for review September 2, 2026 16:24
@github-actions

github-actions Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

🔍 OpenCodeReview found 2 issue(s) in this PR.

  • ✅ Successfully posted inline: 2 comment(s)

Comment on lines +135 to +141
let pool_prewarm_high_watermark = pool_config.as_ref().map(|pool| {
config
.pool
.block
.prewarm_high_watermark
.unwrap_or(pool.high_watermark)
});

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.

bug · medium
Preserve an omitted setting as None here instead of materializing the default. When app_config is unavailable, append_pool_cli_args treats this as explicit and always passes the newly added --pool-prewarm-high-watermark option. A configured/PATH-resolved older daemon will then reject an option the user never set and fail to start. The daemon already defaults this value to its resolved high_watermark, so forwarding only the explicit option retains the same behavior and permits mixed-version startup.

Suggestion:

Suggested change
let pool_prewarm_high_watermark = pool_config.as_ref().map(|pool| {
config
.pool
.block
.prewarm_high_watermark
.unwrap_or(pool.high_watermark)
});
let pool_prewarm_high_watermark = pool_config
.as_ref()
.and(config.pool.block.prewarm_high_watermark);

Comment on lines +424 to +426
#[cfg(test)]
mod tests {
use super::*;

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.

test · medium
These new configuration tests are in the binary target, but the repository's test-unit target runs cargo test -p uvm-ublk-daemon --lib, so CI will not compile or execute them. Move load_pool_config and these tests into the library, or update test-unit to include --bin uvm-ublk-daemon so the validation/defaulting behavior remains covered.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant