kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #4431] [RFC]: [Store] Completing the cross-client move/migration control plane (DrainJob HA recovery, task lifecycle APIs, surface parity) #1118

New issue

Open

opened 2026-10-29 21:23:03 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 21:23:03 +00:00

Owner

Originally created by @stmatengss on GitHub (Sep 30, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4431

Changes proposed

Summary

Mooncake Store already has a solid data plane for cross-client replica movement:

the synchronous lease RPCs CopyStart/CopyEnd/CopyRevoke and MoveStart/MoveEnd/MoveRevoke (plus the DynamicReplicaCopy* variants),
the asynchronous REPLICA_COPY / REPLICA_MOVE task queue (CreateCopyTask / CreateMoveTask / QueryTask / FetchTasks / MarkTaskToComplete), and
the segment-level DrainJob that fans out into per-key move tasks.

The control plane around that data plane is incomplete, however. This RFC inventories the existing interfaces, lists the gaps, and proposes a phased completion plan:

DrainJob HA persistence and resume (correctness)
Task/job lifecycle APIs: register drain RPCs, add enumeration and per-task cancel (operability)
Batching and granularity: batch task creation, drain-by-client (ergonomics)
Surface parity and documentation: C API / Python / Admin HTTP / docs (adoption)

Current state (inventory)

Synchronous lease RPCs (executor client ↔ master)

RPC	Shape	Notes
`CopyStart` / `CopyEnd` / `CopyRevoke`	1 source → N targets	plus `DynamicReplicaCopyStart/End/Revoke` for dynamic replication
`MoveStart` / `MoveEnd` / `MoveRevoke`	1 source → 1 target	`MoveStartResponse.target` is `optional` (degenerates to source-only removal when the target replica already exists); target allocation is memory-segment only

All registered in RegisterRpcService (rpc_service.cpp).

Asynchronous task queue

CreateCopyTask(key, targets) / CreateMoveTask(key, source, target) / QueryTask / FetchTasks / MarkTaskToComplete, all registered RPCs.
A task is assigned to the client owning the source segment and executed by its background loop (Client::PollAndDispatchTasks → ExecuteTask → Client::Copy / Client::Move).
Task state survives master failover: task_manager is one of the three snapshot payloads (MasterSnapshotPayloads{metadata, segments, task_manager}, see master_snapshot_codec.h).

Segment-level migration (DrainJob)

CreateDrainJob / QueryDrainJob / CancelDrainJob are declared on WrappedMasterService but are not registered as coro_rpc handlers; they are exposed only through the admin HTTP server (POST /api/v1/drain_jobs, GET /api/v1/drain_jobs/query, POST /api/v1/drain_jobs/cancel, GET /api/v1/segments/status, see master_admin_service.cpp).
The scheduler (ScheduleDrainJobTasks in master_service.cpp) expands a job into per-(tenant, key, segment) units, each executed as a REPLICA_MOVE task, bounded by max_concurrency (default 4) and retried up to kMaxDrainUnitRetries.
Segment lifecycle OK → DRAINING → DRAINED (or back to OK on cancel/failure) is persisted: MountedSegment.status is covered by the segment snapshot payload.

Cross-media migration (separate mechanism)

memory→SSD offload and SSD→memory promotion are heartbeat-driven (OffloadObjectHeartbeat / NotifyOffloadSuccess, PromotionObjectHeartbeat / PromotionAllocStart / NotifyPromotionSuccess/Failure), with EvictDiskReplica / BatchEvictDiskReplica for disk-replica cleanup. They are not task-based "moves" and are only listed here for completeness.

Client surfaces

Surface	move/copy/task	drain
`MasterClient` (C++)	full	none
`RealClient` public API	`create_copy_task`, `create_move_task`, `query_task`	none
Python `MooncakeDistributedStore`	same three	none
C API (store_c.h)	none	none
Admin HTTP	none	drain jobs + segment status

Gaps

G1. DrainJob state is lost on master failover (correctness)

drain_jobs_ is an in-memory map; the snapshot payloads are only {metadata, segments, task_manager}. After a failover:

The job itself, its per-unit progress (completed_unit_keys, terminal_failed_unit_keys, retry_counts) and migrated_bytes are gone.
The per-key move tasks it spawned do survive (they live in the task-manager payload) and will run to completion, but no job remains to consume their results.
Segments the job marked DRAINING keep that persisted status, yet no job exists to drive them to DRAINED or back to OK. They stay non-allocatable indefinitely — and there is currently no admin endpoint to force a segment status change either.
CancelDrainJob / QueryDrainJob on the new leader return JOB_NOT_FOUND; operators cannot even tell that a drain was in progress.

G2. DrainJob has no RPC surface

The three drain methods exist on WrappedMasterService but are absent from RegisterRpcService, and MasterClient has no corresponding methods. Anything drain-related requires raw HTTP access to the admin server, bypassing the RPC path that gives clients service discovery, leader redirection, and retries.

G3. No enumeration or per-task lifecycle

No ListDrainJobs: a lost job_id makes the job invisible.
No ListTasks / query-by-key: QueryTask is by UUID only.
No CancelTask: an individual copy/move task cannot be canceled. Only a whole drain job can be, and only while it has no active tasks. A wrongly targeted CreateMoveTask can only be waited out.

G4. No batch move primitive

CreateCopyTask / CreateMoveTask take a single key; MoveStart takes a single target. Migrating a key set (tenant compaction, rebalancer output) costs one RPC per key with no grouping handle. DrainJob covers the segment-level case only.

G5. Drain granularity is segment-name only

CreateDrainJobRequest.segments takes segment names. Evacuating a client — the common decommissioning case — requires enumerating its segments out-of-band first; there is no client_ids field or resolution helper.

G6. Surface parity

C API has no move/copy/task functions at all.
Python has no drain client (create/query/cancel).
Admin HTTP has no endpoint to trigger a single-object move/copy or to query a task — only whole-segment drain.
The drain admin HTTP API itself is not documented under docs/.

Proposal

Phase 1 — DrainJob recovery (correctness)

Serialize drain_jobs_ (job id, request, status, per-unit progress, retry counts, counters) as a fourth snapshot payload and rebuild the bookkeeping on restore. The tasks a job owns already persist via the task-manager payload, so a restored leader can re-attach to them by id and resume scheduling the remaining units.
Fail-safe on restore: a segment found in persisted DRAINING status whose owning job was not restored (snapshot skew, job created after the last snapshot) must not be left stuck — either recreate a job for it or reset it to OK, but always with an explicit, observable log + metric.
Longer term, align with the HA roadmap (#3808): job lifecycle events (JOB_CREATE / JOB_UPDATE / JOB_TERMINAL) as OpLog entries so standby replay keeps job state fresh instead of snapshot-granular.

Phase 2 — Task/job lifecycle APIs (operability)

Register CreateDrainJob / QueryDrainJob / CancelDrainJob / QuerySegmentStatus(ById) in RegisterRpcService and add matching MasterClient methods; keep the admin HTTP endpoints as thin wrappers.
ListDrainJobs(filter) and ListTasks(filter by status/type/client/key) with bounded page size.
CancelTask(task_id): PENDING tasks are removed immediately; PROCESSING tasks get a best-effort cancel flag the executor checks between transfer slices (final state still reported via MarkTaskToComplete).
Enrich QueryJobResponse with bounded terminal-failed unit details (key, source segment, last error).

Phase 3 — Batching and granularity (ergonomics)

BatchCreateCopyTask / BatchCreateMoveTask: one RPC, one task per key (preserves per-key retry/cancel granularity), returning per-key task_id | error.
CreateDrainJobRequest.client_ids: the master resolves them to the client's mounted segments at validation time (snapshot semantics: later unmounts do not silently shrink the job's scope).

Phase 4 — Surface parity and docs

C API: mooncake_store_create_move_task / mooncake_store_create_copy_task / mooncake_store_query_task.
Python: create_drain_job / query_drain_job / cancel_drain_job, plus the list/cancel APIs from Phase 2.
Admin HTTP: single-object POST /api/v1/move_tasks, GET /api/v1/tasks/{id}, and a docs/ page for the full drain/task admin API.

Non-goals / relation to existing work

LOCAL_DISK replica migration during drain is the data-plane design of #2817 / #2885 and intentionally out of scope here; Phases 1–4 are replica-type agnostic and should compose with whatever lands there.
Drain bandwidth/rate limiting is #2875; pacing controls should attach to the persisted job state from Phase 1.
Durable Copy/Move lifecycle events (RL03/RL05 in the HA roadmap #3808) harden the per-replica data path; Phase 1's job events would follow the same OpLog conventions.
Client liveness/offboarding (#2953) handles dead-client cleanup; drain-by-client (Phase 3) targets live evacuation and complements it.
No change to the MoveStart single-target invariant is proposed; multi-destination "move" remains expressible as copy + remove.

Compatibility

Phases 2–4 are purely additive: new RPC methods, new optional request fields, new endpoints.
Phase 1 adds a snapshot payload; restore must treat its absence in older snapshots as "no jobs" and then apply the orphaned-DRAINING fail-safe.
No on-wire changes to existing task payloads.

Open questions

Phase 1 persistence model: snapshot-only (simpler; job progress may roll back to the last snapshot) vs. OpLog job events (standby-fresh; more machinery). Is snapshot-only acceptable as a first step, given that the underlying tasks already resume?
Orphaned DRAINING segments on restore: auto-reset to OK, auto-recreate a drain job, or require an explicit admin decision (with a new force-status endpoint)?
CancelTask on PROCESSING tasks: is a best-effort flag checked between transfer slices enough, or do we want preemption at the Transfer Engine batch level?
Batch creation granularity: per-key tasks (proposed) vs. a single multi-key task with internal progress — which unit do operators actually want to retry/cancel at?
Drain-by-client resolution: resolve segments once at creation (snapshot semantics, proposed) or track the client's segment set dynamically?

Before submitting a new issue...

Make sure you already searched for relevant issues and read the documentation

Originally created by @stmatengss on GitHub (Sep 30, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4431 ### Changes proposed ## Summary Mooncake Store already has a solid data plane for cross-client replica movement: - the synchronous lease RPCs `CopyStart/CopyEnd/CopyRevoke` and `MoveStart/MoveEnd/MoveRevoke` (plus the `DynamicReplicaCopy*` variants), - the asynchronous `REPLICA_COPY` / `REPLICA_MOVE` task queue (`CreateCopyTask` / `CreateMoveTask` / `QueryTask` / `FetchTasks` / `MarkTaskToComplete`), and - the segment-level `DrainJob` that fans out into per-key move tasks. The **control plane** around that data plane is incomplete, however. This RFC inventories the existing interfaces, lists the gaps, and proposes a phased completion plan: 1. **DrainJob HA persistence and resume** (correctness) 2. **Task/job lifecycle APIs**: register drain RPCs, add enumeration and per-task cancel (operability) 3. **Batching and granularity**: batch task creation, drain-by-client (ergonomics) 4. **Surface parity and documentation**: C API / Python / Admin HTTP / docs (adoption) ## Current state (inventory) ### Synchronous lease RPCs (executor client ↔ master) | RPC | Shape | Notes | | --- | --- | --- | | `CopyStart` / `CopyEnd` / `CopyRevoke` | 1 source → N targets | plus `DynamicReplicaCopyStart/End/Revoke` for dynamic replication | | `MoveStart` / `MoveEnd` / `MoveRevoke` | 1 source → 1 target | `MoveStartResponse.target` is `optional` (degenerates to source-only removal when the target replica already exists); target allocation is memory-segment only | All registered in `RegisterRpcService` ([rpc_service.cpp](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/src/rpc_service.cpp)). ### Asynchronous task queue - `CreateCopyTask(key, targets)` / `CreateMoveTask(key, source, target)` / `QueryTask` / `FetchTasks` / `MarkTaskToComplete`, all registered RPCs. - A task is assigned to the client owning the source segment and executed by its background loop (`Client::PollAndDispatchTasks` → `ExecuteTask` → `Client::Copy` / `Client::Move`). - Task state survives master failover: `task_manager` is one of the three snapshot payloads (`MasterSnapshotPayloads{metadata, segments, task_manager}`, see [master_snapshot_codec.h](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/include/ha/snapshot/master_snapshot_codec.h)). ### Segment-level migration (DrainJob) - `CreateDrainJob` / `QueryDrainJob` / `CancelDrainJob` are declared on `WrappedMasterService` but are **not registered as coro_rpc handlers**; they are exposed only through the admin HTTP server (`POST /api/v1/drain_jobs`, `GET /api/v1/drain_jobs/query`, `POST /api/v1/drain_jobs/cancel`, `GET /api/v1/segments/status`, see [master_admin_service.cpp](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/src/master_admin_service.cpp)). - The scheduler (`ScheduleDrainJobTasks` in [master_service.cpp](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/src/master_service.cpp)) expands a job into per-(tenant, key, segment) units, each executed as a `REPLICA_MOVE` task, bounded by `max_concurrency` (default 4) and retried up to `kMaxDrainUnitRetries`. - Segment lifecycle `OK → DRAINING → DRAINED` (or back to `OK` on cancel/failure) **is** persisted: `MountedSegment.status` is covered by the segment snapshot payload. ### Cross-media migration (separate mechanism) memory→SSD offload and SSD→memory promotion are heartbeat-driven (`OffloadObjectHeartbeat` / `NotifyOffloadSuccess`, `PromotionObjectHeartbeat` / `PromotionAllocStart` / `NotifyPromotionSuccess/Failure`), with `EvictDiskReplica` / `BatchEvictDiskReplica` for disk-replica cleanup. They are not task-based "moves" and are only listed here for completeness. ### Client surfaces | Surface | move/copy/task | drain | | --- | --- | --- | | `MasterClient` (C++) | full | none | | `RealClient` public API | `create_copy_task`, `create_move_task`, `query_task` | none | | Python `MooncakeDistributedStore` | same three | none | | C API ([store_c.h](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/include/store_c.h)) | none | none | | Admin HTTP | none | drain jobs + segment status | ## Gaps ### G1. DrainJob state is lost on master failover (correctness) `drain_jobs_` is an in-memory map; the snapshot payloads are only `{metadata, segments, task_manager}`. After a failover: - The job itself, its per-unit progress (`completed_unit_keys`, `terminal_failed_unit_keys`, `retry_counts`) and `migrated_bytes` are gone. - The per-key move tasks it spawned **do** survive (they live in the task-manager payload) and will run to completion, but no job remains to consume their results. - Segments the job marked `DRAINING` keep that persisted status, yet no job exists to drive them to `DRAINED` or back to `OK`. They stay non-allocatable indefinitely — and there is currently no admin endpoint to force a segment status change either. - `CancelDrainJob` / `QueryDrainJob` on the new leader return `JOB_NOT_FOUND`; operators cannot even tell that a drain was in progress. ### G2. DrainJob has no RPC surface The three drain methods exist on `WrappedMasterService` but are absent from `RegisterRpcService`, and `MasterClient` has no corresponding methods. Anything drain-related requires raw HTTP access to the admin server, bypassing the RPC path that gives clients service discovery, leader redirection, and retries. ### G3. No enumeration or per-task lifecycle - No `ListDrainJobs`: a lost `job_id` makes the job invisible. - No `ListTasks` / query-by-key: `QueryTask` is by UUID only. - No `CancelTask`: an individual copy/move task cannot be canceled. Only a whole drain job can be, and only while it has no active tasks. A wrongly targeted `CreateMoveTask` can only be waited out. ### G4. No batch move primitive `CreateCopyTask` / `CreateMoveTask` take a single key; `MoveStart` takes a single target. Migrating a key set (tenant compaction, rebalancer output) costs one RPC per key with no grouping handle. DrainJob covers the segment-level case only. ### G5. Drain granularity is segment-name only `CreateDrainJobRequest.segments` takes segment names. Evacuating a *client* — the common decommissioning case — requires enumerating its segments out-of-band first; there is no `client_ids` field or resolution helper. ### G6. Surface parity - C API has no move/copy/task functions at all. - Python has no drain client (create/query/cancel). - Admin HTTP has no endpoint to trigger a single-object move/copy or to query a task — only whole-segment drain. - The drain admin HTTP API itself is not documented under `docs/`. ## Proposal ### Phase 1 — DrainJob recovery (correctness) - Serialize `drain_jobs_` (job id, request, status, per-unit progress, retry counts, counters) as a fourth snapshot payload and rebuild the bookkeeping on restore. The tasks a job owns already persist via the task-manager payload, so a restored leader can re-attach to them by id and resume scheduling the remaining units. - Fail-safe on restore: a segment found in persisted `DRAINING` status whose owning job was not restored (snapshot skew, job created after the last snapshot) must not be left stuck — either recreate a job for it or reset it to `OK`, but always with an explicit, observable log + metric. - Longer term, align with the HA roadmap (#3808): job lifecycle events (`JOB_CREATE` / `JOB_UPDATE` / `JOB_TERMINAL`) as OpLog entries so standby replay keeps job state fresh instead of snapshot-granular. ### Phase 2 — Task/job lifecycle APIs (operability) - Register `CreateDrainJob` / `QueryDrainJob` / `CancelDrainJob` / `QuerySegmentStatus(ById)` in `RegisterRpcService` and add matching `MasterClient` methods; keep the admin HTTP endpoints as thin wrappers. - `ListDrainJobs(filter)` and `ListTasks(filter by status/type/client/key)` with bounded page size. - `CancelTask(task_id)`: PENDING tasks are removed immediately; PROCESSING tasks get a best-effort cancel flag the executor checks between transfer slices (final state still reported via `MarkTaskToComplete`). - Enrich `QueryJobResponse` with bounded terminal-failed unit details (key, source segment, last error). ### Phase 3 — Batching and granularity (ergonomics) - `BatchCreateCopyTask` / `BatchCreateMoveTask`: one RPC, one task per key (preserves per-key retry/cancel granularity), returning per-key `task_id | error`. - `CreateDrainJobRequest.client_ids`: the master resolves them to the client's mounted segments at validation time (snapshot semantics: later unmounts do not silently shrink the job's scope). ### Phase 4 — Surface parity and docs - C API: `mooncake_store_create_move_task` / `mooncake_store_create_copy_task` / `mooncake_store_query_task`. - Python: `create_drain_job` / `query_drain_job` / `cancel_drain_job`, plus the list/cancel APIs from Phase 2. - Admin HTTP: single-object `POST /api/v1/move_tasks`, `GET /api/v1/tasks/{id}`, and a `docs/` page for the full drain/task admin API. ## Non-goals / relation to existing work - **LOCAL_DISK replica migration during drain** is the data-plane design of #2817 / #2885 and intentionally out of scope here; Phases 1–4 are replica-type agnostic and should compose with whatever lands there. - **Drain bandwidth/rate limiting** is #2875; pacing controls should attach to the persisted job state from Phase 1. - **Durable Copy/Move lifecycle events** (RL03/RL05 in the HA roadmap #3808) harden the per-replica data path; Phase 1's job events would follow the same OpLog conventions. - **Client liveness/offboarding** (#2953) handles dead-client cleanup; drain-by-client (Phase 3) targets *live* evacuation and complements it. - No change to the `MoveStart` single-target invariant is proposed; multi-destination "move" remains expressible as copy + remove. ## Compatibility - Phases 2–4 are purely additive: new RPC methods, new optional request fields, new endpoints. - Phase 1 adds a snapshot payload; restore must treat its absence in older snapshots as "no jobs" and then apply the orphaned-`DRAINING` fail-safe. - No on-wire changes to existing task payloads. ## Open questions 1. **Phase 1 persistence model**: snapshot-only (simpler; job progress may roll back to the last snapshot) vs. OpLog job events (standby-fresh; more machinery). Is snapshot-only acceptable as a first step, given that the underlying tasks already resume? 2. **Orphaned `DRAINING` segments on restore**: auto-reset to `OK`, auto-recreate a drain job, or require an explicit admin decision (with a new force-status endpoint)? 3. **`CancelTask` on PROCESSING tasks**: is a best-effort flag checked between transfer slices enough, or do we want preemption at the Transfer Engine batch level? 4. **Batch creation granularity**: per-key tasks (proposed) vs. a single multi-key task with internal progress — which unit do operators actually want to retry/cancel at? 5. **Drain-by-client resolution**: resolve segments once at creation (snapshot semantics, proposed) or track the client's segment set dynamically? ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

bsakaguchi added the

RFC

label

2026-10-29 21:23:03 +00:00

bsakaguchi commented

2026-10-29 21:23:04 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Sep 30, 2026):

Thanks for opening this issue, @stmatengss!

Field	Value
Issue	#4431
GitHub user ID	`11641725`
Reporter	@stmatengss

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Sep 30, 2026): Thanks for opening this issue, @stmatengss! | Field | Value | |-------|-------| | **Issue** | #4431 | | **GitHub user ID** | `11641725` | | **Reporter** | @stmatengss | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi commented

2026-10-29 21:23:05 +00:00

Author

Owner

@SuperGoodGame commented on GitHub (Oct 3, 2026):

Hi @stmatengss, I'd like to take the snapshot-only DrainJob recovery part of Phase 1. I've started a local implementation.

The plan is to reuse the existing master snapshot/storage path to persist job state, progress, retry budgets and task associations, then resume after cold restart or snapshot-only standby promotion. Initial investigation also found that DRAINING/DRAINED allocator recovery, client remount and in-flight COPY/MOVE bookkeeping need to remain consistent for safe continuation.

I'll keep legacy snapshots readable and add on-disk restart and failure regressions. For orphaned DRAINING segments, I propose preserving readable data, keeping allocation blocked, and reporting logs/metrics for explicit operator recovery through the existing segment-status API.

Job OpLog and batch-generated standby snapshot recovery would follow separately. This slice is separate from the native RPC work in #4436. Please let me know if someone is already working on it.

@SuperGoodGame commented on GitHub (Oct 3, 2026): Hi @stmatengss, I'd like to take the snapshot-only DrainJob recovery part of Phase 1. I've started a local implementation. The plan is to reuse the existing master snapshot/storage path to persist job state, progress, retry budgets and task associations, then resume after cold restart or snapshot-only standby promotion. Initial investigation also found that DRAINING/DRAINED allocator recovery, client remount and in-flight COPY/MOVE bookkeeping need to remain consistent for safe continuation. I'll keep legacy snapshots readable and add on-disk restart and failure regressions. For orphaned DRAINING segments, I propose preserving readable data, keeping allocation blocked, and reporting logs/metrics for explicit operator recovery through the existing segment-status API. Job OpLog and batch-generated standby snapshot recovery would follow separately. This slice is separate from the native RPC work in #4436. Please let me know if someone is already working on it.