kvcache-ai / kvcache-ai/Mooncake

[RFC]: Best-fit segment placement for mixed-size workloads (RL data-plane offload)

Open
#3,965 1 comment 0 reactions 0 assignees View on GitHub
RFC
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

#### Motivation

In an RL data-plane offload deployment a 1.25 GiB `PutStart` failed with `NO_AVAILABLE_HANDLE` while the cluster still reported about 20 GB of free memory. KV-cache deployments with a fixed object size never show this.

The first explanation was allocator precision (`MANTISSA_BITS`, see #3787) combined with eviction behavior. We reproduced the incident on a real `mooncake_master` and traced it to segment placement instead:

- The master allocates every object as one contiguous region inside one segment (`put_parts`, `put_tensor`, or the structured object API with a raised `chunk_bytes` all take this path). A 1.25 GiB object therefore needs one segment with a 1.25 GiB hole.
- `random` and `free_ratio_first` both spread allocations evenly. Under an RL workload, where object sizes span several orders of magnitude and every object is removed explicitly one or two steps after it was written, this leaves every segment with a set of medium-sized holes and no segment with a large one.
- Eviction is not involved. The failures occur with zero evictions, and in this workload eviction would also be the wrong remedy because it destroys rollout data the trainer has not consumed yet.
- Raising `MANTISSA_BITS` reduces rounding waste but does not change the hole structure; it is a mitigation, not a fix.

#### Evidence

All experiments use the current `main`, 16 segments x 4 GiB, and a driver that plays 8 data-parallel ranks writing a GRPO-style step (5 stages, 12 to 16 keys, plus one whole object of 640 to 1920 MiB per step) and removing each step two steps later, i.e. the same interface (`put_parts`, `get_buffer`, `remove`) and lifetime pattern as the production data plane. `mooncake_master --v=1` logged every `PutStart`, allocation failure, `Remove`, and eviction.

Reproduction on the real master (`--eviction_ratio=0`, watermark 0.95):

| Run | Strategy | Puts | Removes | Evict | Fails | Free at failure | Largest hole |
| --- | --- | --- | --- | --- | --- | --- | --- |
| runF | random | 7680 | 7673 | 0 | 7 | 19 to 29 GB | 1.38 to 1.75 GiB |
| runH | random | 7680 | 7675 | 0 | 5 | 19 to 25 GB | 1.5 to 1.75 GiB |
| runG, runI, runJ | best_fit | 7680 | 7680 | 0 | 0 | n/a | n/a |

Every failure is a request between 1.4 and 1.9 GiB that does not fit the largest hole while total free space is an order of magnitude larger. With a constant large-object size (runD) there are no failures because the freed hole always matches the next request; variable large sizes are the trigger, and RL batches vary every step.

Offline replay of the recorded event logs (same put/remove order, no eviction retry) in `allocation_strategy_bench`. Failure counts are for the runF and runE event logs; the last column is the p50 largest hole at probe time on runC, a workload without large objects.

| Strategy | Fails, runF | Fails, runE | Largest hole p50, runC |
| --- | --- | --- | --- |
| random | 11 to 15 | 15 | 1792 MiB |
| free_ratio_first | 12 to 15 | 19 | 1280 MiB |
| largest_hole_first (bench only) | n/a | 34 | 1280 MiB |
| reserved segments (bench only) | n/a | 44 | 4096 MiB (1) |
| best_fit | 0 to 1 (2) | 5 (2) | 4096 MiB |

(1) 9 of 16 segments reserved for large objects; on runC the reservation causes 1165 capacity failures instead.
(2) Every best_fit failure happened with less than 14 GB free, i.e. near capacity rather than from fragmentation.

Two real RL framework traces (miles with the Mooncake object-store backend, 10 and 20 rollouts) confirmed the lifetime model: puts and removes are strictly paired, zero evictions, 34 keys per rollout per DP shard, lifetimes of 7 to 160 seconds depending on step length.

#### Design

Add `AllocationStrategyType::BEST_FIT`, selectable with `--allocation_strategy=best_fit`. For each replica the segment whose largest contiguous free region is the smallest one that still fits the request is chosen. Small objects fill segments that are already partially used instead of carving holes into the emptiest ones, so large contiguous regions survive for large objects.

- `BestFitAllocationStrategy` reuses `RankedAllocationStrategy::AllocateRanked` with a score of `-(largest_free_region - request)`. Segments whose allocator cannot report a largest free region (cachelib) rank last; segments that cannot fit the request rank after those and are only tried as a fallback. Preferred segments, distinct segments per replica, and best-effort semantics are inherited unchanged.
- `AllocateRanked` gains a `rank_all_segments` flag. The default keeps today's sampling of `6 x remaining` candidates; best-fit ranks every segment because the choice has to be the global tightest fit. Ties are now sorted with `stable_sort` so the random start index actually spreads ties instead of favoring low segment indices; this also applies to `free_ratio_first`.
- `OffsetBufferAllocator::getLargestFreeRegion()` already returns exactly the largest request the segment can satisfy (highest non-empty bin floor under the allocator mutex, consistent with the fast-fail check inside `allocate()`). A property test pins this contract because best-fit depends on it.
- Cost: one mutex-protected bin lookup per segment per allocation. In the bench p99 allocation latency goes from 0.7 us (random) to 1.4 us (best_fit) with 16 segments and scales linearly with the segment count.

No change to the default strategy, the wire protocol, or the allocator layout.

#### Trade-off: traffic concentration

Best-fit deliberately concentrates data. Measured on the real master over 60 steps with identical workloads, using per-segment placement logged by the master:

| Metric | random | best_fit |
| --- | --- | --- |
| Total bytes written, busiest segment / mean | 1.11 | 1.43 |
| Total bytes written, idlest segment / mean | 0.92 | 0.54 |
| Write skew inside a 200-put window (instantaneous) | 2.24 | 2.32 |
| Stddev of live bytes across segments | 0.90 GiB | 1.38 GiB |

Instantaneous hot-spotting is unchanged because one step's objects are too many to fit in one segment anyway. The long-run imbalance is real: over 60 steps the busiest segment wrote 2.6x the idlest one. Deployments whose per-segment bandwidth is already the bottleneck should measure this before switching. A bucketed variant (random choice among segments whose slack falls in the same 512 MiB bucket) trades only a little skew (1.43 to 1.36) for more failures, so it is not proposed; a bandwidth-aware secondary sort is the natural follow-up if the concentration matters in practice.

#### Alternatives considered

- Client-side chunking of large objects. The structured object API already chunks at 64 MiB and never failed in these runs. Effective, but it changes the `get_buffer` contract for callers that expect one contiguous buffer, so it cannot be applied from the master side.
- `fragmentation_aware` (#2797) ranks sampled candidates by largest free region. Its "largest hole first" rule fixes the case where a hole exists but was not picked; in these traces the dominant failure is that no hole exists, and sending small objects to the largest hole makes that worse (34 failures vs 15 for random). The two are complementary: best-fit preserves holes, fragmentation-aware finds them.
- Reserved segments for large objects. Zero failures when the reservation matches the size mix and capacity failures when it does not; brittle.
- Higher `MANTISSA_BITS` (#3787). Raises the fill ceiling by a few percent; orthogonal and still worth having.
- Eviction changes. Not applicable: the failures happen with zero evictions, and RL data must not be evicted before it is consumed.

#### Tooling added to reproduce and evaluate

To be submitted as a separate PR:

- `allocation_strategy_bench`: `rl` size pattern, `--probe_mib` large-object probe, `--rl_trace_file` size replay, `--rl_trace_events` put/remove replay with real lifetimes and a per-failure free-space report, `--size_class_strategies` comparison including bench-only placement experiments, per-segment traffic skew and utilization spread.
- `mooncake_master` `VLOG(1)` records: `put_start_allocated` with the chosen segments, `put_start_alloc_failed` with total free space and largest free region, per-key `remove_object` and `evict_object`.
- `extract_alloc_trace.py`: turns a master log into a size file, an ordered event log, and a per-octave histogram.
- `mooncake-rl/examples/rl_dataproto_trace_driver.py`: multi-rank RL data-plane load generator with both the structured object API and the `put_parts` and `get_buffer` write path.

#### Rollout

1. PR 1, `[Store] Add best_fit allocation strategy`: strategy, `AllocateRanked` flag and stable sort, unit tests (parameterized suite plus tightest-fit and largest-free-region contract tests), documentation. About 100 lines of non-test code, off by default.
2. PR 2, `[Store] RL data-plane trace tooling`: bench, master logs, extraction script, driver, benchmark README.

#### Open questions

- Should best-fit become the default for the `offset` allocator, or stay opt-in until the traffic-concentration trade-off is measured on a bandwidth-bound cluster?
- Is a bandwidth-aware or utilization-aware secondary sort wanted inside the same strategy, or better as a separate ranking policy?
- The largest-free-region query takes the allocator mutex once per segment per allocation. Is an atomic hint (the allocator already keeps one as an upper bound) preferable for clusters with thousands of segments?

Contributor guide

Open the contributing guide

Research direction

Start by locating AllocationStrategyType, RankedAllocationStrategy::AllocateRanked, BestFitAllocationStrategy, and OffsetBufferAllocator::getLargestFreeRegion. Read the parameterized strategy tests and the largest-free-region property test, then use allocation_strategy_bench to compare placement and failures. Done means an opt-in best_fit strategy, stable tie handling, passing tightest-fit and allocator contract tests, and documentation without changing the default strategy or wire protocol.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, distributed-systems, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.