kvcache-ai / kvcache-ai/Mooncake

[RFC]: LOCAL_DISK data migration during drain — Options (a) and (c)

Open
#2,885 5 comments 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

### Changes proposed

## Background

Issue [#2817](https://github.com/kvcache-ai/Mooncake/issues/2817) identified that `ScheduleDrainJobTasks` silently ignores LOCAL_DISK replicas because `Replica::get_segment_names()` returns an empty vector for `LOCAL_DISK`. PR [#2834](https://github.com/kvcache-ai/Mooncake/pull/2834) implements a fail-closed safety net (Step 1 in he-yufeng's two-step plan), but the actual migration capability — Step 2 — remains unimplemented.

This RFC follows up on [#2817](https://github.com/kvcache-ai/Mooncake/issues/2817) and the fail-closed fix in [#2834](https://github.com/kvcache-ai/Mooncake/pull/2834), proposing a concrete plan for Step 2 — actual LOCAL_DISK migration during drain.

It evaluates two viable approaches:
- **Option (a)**: Targeted file transfer — `BatchLoad` → `TransferWrite` → `BatchOffload`
- **Option (c)**: Model LOCAL_DISK as a `NoFSegment`-style allocator-managed resource

Option (b) (drop LOCAL_DISK and re-offload on target) is excluded because LOCAL_DISK data is **not recomputable** for some use cases.

---

## Current State: LOCAL_DISK in the Drain Flow

```
Drain Job
└─> ScheduleDrainJobTasks()
└─> GetReplicaSegmentNames()
└─> LOCAL_DISK → empty vector ← structurally invisible
└─> CreateMoveTask() — only for MEMORY replicas
└─> Client::Move()
└─> ExecuteReplicaTransfer()
└─> if (!source.is_memory_replica()) → INVALID_PARAMS ← hard rejection
```

Both the scheduling and execution paths structurally exclude LOCAL_DISK replicas.

---

## Option (a): Targeted File Transfer

### Approach

```
Data: source SSD ────→ source RAM ────→ target RAM ────→ target SSD
Action: BatchLoad TransferWrite BatchOffload
```

Three building blocks already exist in the codebase:
- **`BatchLoad`**: Reads data from local SSD files into staging buffers (all four storage backends).
- **`TransferWrite`** (TransferEngine): RDMA/TCP transfer from a local buffer to a remote MEMORY replica.
- **`BatchOffload`**: Writes data from memory buffers to local SSD on the target node.

The Promotion path (`ProcessPromotionTasks`) already proves the first two steps work together:
`AllocateBatch() → BatchLoad() → TransferWrite()`.

### What Already Exists

| Component | Location | Status |
| --- | --- | --- |
| BatchLoad (SSD → RAM) | `storage_backend.cpp` | ✅ 4 backends |
| TransferWrite (RDMA/TCP) | `transfer_task.cpp:981` | ✅ generic buffer support |
| BatchOffload (RAM → SSD) | `storage_backend.cpp` | ✅ 3 backends |
| Promotion: BatchLoad → TransferWrite | `file_storage.cpp:728` | ✅ proven pattern |

### What Needs to Be Built

1. **`Client::ExecuteReplicaTransfer()`** — lift the `is_memory_replica()` restriction (client_service.cpp:3182). Add a LOCAL_DISK path:
```
if (source.is_local_disk_replica()) {
allocate staging buffer;
BatchLoad(slices); // SSD → staging buffer
TransferWrite(target, slices); // staging buffer → target MEMORY replica
release staging buffer;
}
```

2. **`MasterService::ScheduleDrainJobTasks()`** — recognize LOCAL_DISK replicas as drainable (currently excluded via `get_segment_names()` returning empty). Generate move units for keys that have only LOCAL_DISK replicas on the draining segment.

3. **Staging buffer management** — reuse the existing `client_buffer_allocator_` (same pool used by Promotion). With `local_buffer_size` as hard cap, failures trigger GC retry, then skip. Concurrency is naturally bounded by `max_concurrency` (default 4).

4. **Source LOCAL_DISK cleanup** — after successful migration, delete the source `.bucket`/`.meta` entries. This requires a new cleanup path (MEMORY replicas use `discarded_replicas_` with delayed release; LOCAL_DISK has no equivalent).

5. **Deduplication** — a key may have both MEMORY and LOCAL_DISK replicas on the draining node. When both need to move, transferring both wastes bandwidth (LOCAL_DISK data is a superset of MEMORY data). Optimization: if a MEMORY replica exists for the same key on the same node, use the existing zero-copy TransferWrite path and skip the BatchLoad for LOCAL_DISK.

### Considerations

| Concern | Analysis |
| --- | --- |
| **Staging buffer OOM** | Buffer pool has O(n) hard cap. GC + retry on failure. Same risk profile as Promotion. |
| **Duplicate transfer** | Can be addressed by preferring MEMORY replica over LOCAL_DISK when both exist on same node. |
| **Source cleanup** | New path required. MEMORY replicas use `discarded_replicas_` (10-min delayed release). LOCAL_DISK cleanup is simpler — just delete the files. |
| **Concurrent drain pressure** | `max_concurrency=4` per drain job + 500ms scheduling interval provides natural throttling. |

### Estimation

- **Core change**: ~200–300 lines in `client_service.cpp` (new LOCAL_DISK branch in `ExecuteReplicaTransfer`, following the Promotion pattern)
- **Drain scheduling**: ~50–100 lines in `master_service.cpp` (recognize LOCAL_DISK replicas)
- **Cleanup**: ~50–100 lines (delete LOCAL_DISK entries after successful move)
- **Total**: moderate effort, low architectural risk

---

## Option (c): Model LOCAL_DISK as NoFSegment

### Approach

Treat LOCAL_DISK as a block-addressable resource managed by a `BufferAllocator`, similar to how NoF (NVMe-oF) segments work. This would allow LOCAL_DISK replicas to participate in the existing `REPLICA_MOVE` primitive (allocator-to-allocator copy).

### Why This Requires Architectural Changes

**1. Storage model mismatch — files vs blocks**

```
LOCAL_DISK (current): NoF (what option c requires):
┌─────────────────────┐ ┌─────────────────────┐
│ File system │ │ NVMe namespace │
│ bucket_1.bucket │ │ base=0 │
│ bucket_2.bucket │ │ offset=4096 (LBA 8) │
│ .meta │ │ offset=8192 │
│ Addressing: path │ │ Addressing: LBA │
│ I/O: open/read/write│ │ I/O: SPDK NVMe-oF │
└─────────────────────┘ └─────────────────────┘
```

NoF uses block-level LBA addressing via SPDK. LOCAL_DISK uses file-level path addressing. Bridging these requires a **file-backed block abstraction layer**.

**2. Replica data structure incompatibility**

| Aspect | NoF (`NoFReplicaData`) | LOCAL_DISK (`LocalDiskReplicaData`) |
| --- | --- | --- |
| Core field | `unique_ptr` (offset + endpoint) | `UUID client_id` + `transport_endpoint` |
| Allocator | CachelibBufferAllocator / OffsetBufferAllocator | None (client manages files) |
| get_segment_names() | Returns segment name | Returns empty vector |

To model LOCAL_DISK as NoF, `LocalDiskReplicaData` must carry an `AllocatedBuffer`, which means a new allocator type that manages file offsets rather than memory offsets.

**3. Requires a new `FileBackedBufferAllocator`**

Existing allocators (`CachelibBufferAllocator`, `OffsetBufferAllocator`) manage memory regions. LOCAL_DISK data is in files managed by `BucketStorageBackend` where multiple keys share a single bucket file. The allocator must:

- Track available space across bucket files
- Handle bucket fragmentation (keys of different sizes packed together)
- Coordinate with the bucket's LRU eviction policy
- Present a flat offset space that maps to `(file_path, file_offset)` pairs

This is effectively building a **file-system-backed block allocator** — a significant subsystem.

**4. Transfer path incompatibility**

```
NoF transfer path: LOCAL_DISK transfer path:
SPDK → NVMe-oF → remote NVMe open/read → user buffer → TransferEngine → remote
```

NoF's `submitSpdkNofOperation` issues NVMe commands directly. LOCAL_DISK reads must go through file I/O. To unify them, LOCAL_DISK would need an SPDK backend that emulates NVMe on top of files — unrealistic.

**5. Drain infrastructure doesn't support NoF segments**

| Drain step | NoF support | Needs |
| --- | --- | --- |
| `ValidateDrainRequest` | ❌ Only queries `segment_manager_` | Add `nof_segment_manager_` |
| `MaybeCompleteDrainJob` | ❌ Only sets `segment_manager_` status | Add NoF segment status handling |
| `ExecuteReplicaTransfer` | ❌ Hardcoded `is_memory_replica()` only | Add NoF source support |
| `CreateMoveTask` / `MoveStart` | ❌ Only allocates MEMORY replicas | Support NoF→NoF or NoF→MEMORY |

### Estimation

- **New allocator**: major new subsystem (thousands of lines)
- **SPDK/file bridge**: large effort, architecturally questionable
- **Drain full-chain NoF support**: hundreds of lines across validation, scheduling, and execution
- **Total**: high effort, high architectural risk, uncertain benefit

---

## Comparison

| | Option (a): File Transfer | Option (c): Model as NoF |
| --- | --- | --- |
| **Reuses existing code** | ✅ BatchLoad, TransferWrite, BatchOffload all exist | ❌ Requires new allocator, new transfer path |
| **New code required** | ~400–500 lines | Several thousand lines |
| **Architectural risk** | Low — pipelining existing primitives | High — abstracting files as blocks |
| **Drain scheduling change** | Recognize LOCAL_DISK in `ScheduleDrainJobTasks` | Full NoF drain support chain |
| **Staging buffer** | Reuse existing `client_buffer_allocator_` | N/A (block-level access) |
| **Transfer protocol** | File I/O + RDMA/TCP (existing) | Must bridge file I/O to SPDK NVMe-oF |
| **Source cleanup** | New delete path (simple) | Via allocator free (needs to be built) |
| **Deduplication** | Solvable — prefer MEMORY over LOCAL_DISK | Same |
| **Fits architecture** | ✅ Composes existing layers | ❌ Forces square peg into round hole |

---

## Open Questions

1. **Target replica type**: After moving LOCAL_DISK data, should the target store it as MEMORY (simpler, handles drain immediately) or trigger an offload to target's LOCAL_DISK (preserves storage tier semantics)?

2. **Partial drain semantics**: If some LOCAL_DISK objects are successfully moved and others fail, should the drain report PARTIAL rather than FAILED? This provides better visibility than the current all-or-nothing approach.

3. **Staging buffer sizing**: Should LOCAL_DISK drain use the same `local_buffer_size` pool as Promotion, or have a separate configurable limit?

### Before submitting a new issue...

- [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Contributor guide

Open the contributing guide

Research direction

Start by reading Client::ExecuteReplicaTransfer() in client_service.cpp:3182 and MasterService::ScheduleDrainJobTasks(), then compare the Promotion path in file_storage.cpp:728 and TransferWrite in transfer_task.cpp:981. Resolve the open questions about target replica type, partial drains, and staging limits before choosing an approach. Done means LOCAL_DISK replicas can be migrated during drain with source cleanup and defined failure semantics.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.