kvcache-ai/Mooncake

Fork 0

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

[PR #1212] [CLOSED] [Store] feat: introduce tiered backend #2075

New issue

Closed

opened 2026-10-29 21:32:36 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:32:36 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/1212
Author: @YiXR
Created: 12/15/2025
Status: ❌ Closed

Base: main ← Head: xinyi/tiered-backend

📝 Commits (1)

92e557e [Store] feat: introduce tired backend

📊 Changes

8 files changed (+986 additions, -0 deletions)

View changed files

➕ mooncake-store/include/tiered_cache/cache_tier.h (+86 -0)
➕ mooncake-store/include/tiered_cache/copier_registry.h (+86 -0)
➕ mooncake-store/include/tiered_cache/data_copier.h (+90 -0)
➕ mooncake-store/include/tiered_cache/tiered_backend.h (+197 -0)
📝 mooncake-store/src/CMakeLists.txt (+3 -0)
➕ mooncake-store/src/tiered_cache/copier_registry.cpp (+41 -0)
➕ mooncake-store/src/tiered_cache/data_copier.cpp (+116 -0)
➕ mooncake-store/src/tiered_cache/tiered_backend.cpp (+367 -0)

📄 Description

Description

This PR introduce the abstract base classes for cache tier management.

CacheTier

mooncake-store/include/tiered_cache/cache_tier.h

Responsibilities

This is an abstract base class that defines the unified interface a single cache tier (e.g., DRAM, SSD) must implement. It is responsible for managing the data lifecycle on a specific storage medium.

Core Interface:

Allocate: Finds free space of size bytes. Does NOT copy data.
Free: Releases space at offset. Used when writes fail or explicitly freeing anonymous blocks.

DataCopier

mooncake-store/include/tiered_cache/data_copier.h

Responsibilities:

Responsible for efficiently copying data between different memory types (MemoryType). This is a core utility class that decouples the data movement logic.

Core Functionality:

Built via DataCopierBuilder, allowing for the registration of direct copy functions between different memory types (e.g., DRAM -> VRAM).
Provides a unified Copy interface. When a direct copy path is unavailable, it automatically employs a fallback mechanism using DRAM as an intermediate buffer (e.g., VRAM -> DRAM -> SSD) (when implementing a new type, compilation requires copy functions between the new type and DRAM). This greatly simplifies the integration of new storage media.
- Interaction: Implementations of TieredBackend and CacheTier hold and use a DataCopier instance to execute all data copy operations, whether it's writing new data or moving data between tiers.

TieredBackend

mooncake-store/include/tiered_cache/tiered_backend.h

Responsibilities:

Serves as the Data Plane for the tiered cache. It does not contain any complex caching policies (e.g., LRU) and is solely responsible for faithfully executing data operation instructions.

Core Data Structures

`DataSource`

A generic descriptor for source data.

struct DataSource {
    const void* ptr;  // Pointer to source buffer
    size_t size;      // Data size
    MemoryType type;  // Source memory type (DRAM, NVME, etc.)
};

`TieredLocation`

Describes the physical location of a data segment.

struct TieredLocation {
    uint64_t tier_id;
    struct DataSource data;
};

`AllocationHandle`

/**
 * @struct AllocationEntry
 * @brief The internal state of an allocation.
 * acts as the "Control Block" for the resource.
 * When the last shared_ptr pointing to this entry dies, the destructor
 * triggers the physical release of the resource via the Backend.
 */
struct AllocationEntry {
    TieredBackend* backend;
    TieredLocation loc;

    AllocationEntry(TieredBackend* b, TieredLocation l) : backend(b), loc(l) {}

    // Destructor: Automatically releases the resource if valid
    ~AllocationEntry();
};

/**
 * @typedef AllocationHandle
 * @brief A reference-counted handle to a storage resource.
 */
using AllocationHandle = std::shared_ptr<AllocationEntry>;

Decomposed APIs

The traditional atomic Put operation is broken down into three granular steps:

Allocate
Reserves storage space locally without writing data.

AllocationHandle Allocate(
    size_t size, std::optional<uint64_t> preferred_tier);

Usage: Called by the AllocationStrategy to find free space.
Returns: true if space is reserved; false otherwise.

Write (local operation)
Writes data to a specific, reserved location.

bool Write(AllocationHandle handle, const DataSource& source);

Usage: Performs local write operation. If need remote write, please tranfer data into AllocationHandle by using TE.
Returns: true on success. If failed, the caller must free the allocated space.

Commit
Registers the key in the local metadata index.

bool Commit(const std::string& key, AllocationHandle handle);

Usage: Called after data is successfully written.
Behavior: It binds the key to the offset internally. This operation is designed to be infallible (must succeed). If an internal error occurs (e.g., OOM), it raises a fatal error.

Delete
Delete the key in the local metadata index.

bool Delete(const std::string& key,
                           std::optional<uint64_t> tier_id)

If tier_id is specified, removes only the replica on that tier. If nullopt, removes ALL replicas for this key.

Allocation Workflows

This decoupled design supports two distinct allocation patterns required by the new architecture.

Each backend can only use one allocation strategy after mooncake store start up. If the backend wants to use Master-Centric Allocation, it must realize all the func to register it's info and allocator to Master.

Pattern A: Client-Centric Allocation (Local Decision)

In this mode, the Client decides where to place data (e.g., "Try DRAM first, then SSD") and later informs the Master.

Workflow

Decision
The ClientCentricStrategy calls Allocate(size, preferred_tier=xxx).
TieredBackend checks local capacity and returns a AllocationHandle.
Execution
The strategy calls Write(data, AllocationHandle), this is a local write operation.
Please use TE to transfer data if you want to do remote write.
Local Commit
The strategy calls Commit(key, AllocationHandle).
The key is now visible to local Get() requests.
Global Sync
The Client sends an asynchronous RPC RegisterAllocation to the Master.
Master updates its directory: "Key X is located at Client Y, Tier Z".

Pattern B: Master-Centric Allocation (Global Decision)

In this mode, the Master manages all resources globally (Client should still call MountSegment() to register all the resources to Master) and tells the Client exactly where to write.
Client doesn't need to save metadata for this backend anyway.

Workflow:

Request
The Client sends an RPC PutStart to the Master.
Decision
Master checks its global resource view and returns {tier_id: 1, offset: 0x2000}.
Execution
Data can be write by TE into tier directly.

Data Migration & Consistency

Migrates a key between tiers using a "Make-Before-Break" protocol to ensure strong consistency.

bool CopyData(const std::string& key, const DataSource& source,
                             uint64_t dest_tier_id,
                             MetadataSyncCallback sync_cb);

If DataSource is Master-Centric Allocation, Master must 'tell' Client where the data is.

Workflow:

Allocate & Write
Data is copied to the destination tier.
Sync Callback
sync_cb is invoked.
Action
The Client sends an RPC to the Master: "Update Key location to Tier New".
Safety
If this RPC fails, CopyData returns false and rolls back the new copy. The old data remains valid.
Commit
Local metadata is updated to point to the new tier.
Delete
Only after the Master is updated and local commit succeeds, the old replica need to be deleted.
If DataSource is Master-Centric Allocation, Master will delete the key by itself.
If DataSource is Client-Centric Allocation, Client must call Delete(key) after CopyData.

sequenceDiagram
    participant Client
    participant TieredBackend
    participant CacheTier
    participant DataCopier
    participant MemoryPool

    Client->>TieredBackend: Allocate(size, preferred_tier)
    activate TieredBackend
    TieredBackend->>TieredBackend: GetSortedTiers()
    TieredBackend->>CacheTier: Allocate(size)
    activate CacheTier
    CacheTier->>MemoryPool: allocate
    MemoryPool-->>CacheTier: ptr
    CacheTier-->>TieredBackend: AllocationHandle(tier_id, DataSource)
    deactivate CacheTier
    TieredBackend-->>Client: AllocationHandle
    deactivate TieredBackend

    Client->>TieredBackend: Write(source_data, handle)
    activate TieredBackend
    TieredBackend->>DataCopier: Copy(source, dest_from_handle)
    activate DataCopier
    DataCopier->>DataCopier: FindCopier(src_type, dest_type)
    alt Direct Path Exists
        DataCopier->>DataCopier: Execute direct copy
    else No Direct Path
        note over DataCopier: Both non-DRAM?
        DataCopier->>TieredBackend: Allocate temp DRAM buffer
        DataCopier->>DataCopier: Copy src → DRAM
        DataCopier->>DataCopier: Copy DRAM → dest
    end
    DataCopier-->>TieredBackend: success/failure
    deactivate DataCopier
    TieredBackend-->>Client: bool
    deactivate TieredBackend

    Client->>TieredBackend: Commit(key, handle)
    activate TieredBackend
    TieredBackend->>TieredBackend: Update metadata_index<br/>(per-key replicas)
    TieredBackend-->>Client: bool
    deactivate TieredBackend

sequenceDiagram
    participant Client
    participant TieredBackend
    participant MetadataIndex
    participant CacheTier1
    participant CacheTier2

    Client->>TieredBackend: Get(key, optional_tier_id)
    activate TieredBackend
    TieredBackend->>MetadataIndex: Lookup replicas for key
    activate MetadataIndex
    MetadataIndex-->>TieredBackend: replica list sorted by priority
    deactivate MetadataIndex
    
    alt tier_id specified
        TieredBackend->>CacheTier1: Retrieve from specific tier
    else best available
        TieredBackend->>CacheTier1: Retrieve from highest-priority tier
    end
    activate CacheTier1
    CacheTier1-->>TieredBackend: AllocationHandle
    deactivate CacheTier1
    TieredBackend-->>Client: AllocationHandle
    deactivate TieredBackend

    Client->>TieredBackend: Delete(key, optional_tier_id)
    activate TieredBackend
    TieredBackend->>MetadataIndex: Remove replica from key
    activate MetadataIndex
    MetadataIndex->>CacheTier2: Free allocation
    activate CacheTier2
    CacheTier2-->>MetadataIndex: success
    deactivate CacheTier2
    MetadataIndex-->>TieredBackend: bool
    deactivate MetadataIndex
    TieredBackend-->>Client: bool
    deactivate TieredBackend

Type of Change

Types
- Bug fix
- New feature
  - Transfer Engine
  - Mooncake Store
  - Mooncake EP
  - Integration
  - P2P Store
  - Python Wheel
- Breaking change
- CI/CD
- Documentation update
- Other

Cohort / File(s)	Summary
Tier Abstraction & Metadatamooncake-store/include/tiered_cache/cache_tier.h, mooncake-store/include/tiered_cache/tiered_backend.h	Defines MemoryType enum (DRAM, NVME, UNKNOWN), DataSource struct, abstract CacheTier interface with lifecycle (Init, Allocate, Free) and introspection methods; introduces tiered storage data structures (TieredLocation, TierView, AllocationEntry, MetadataEntry) and TieredBackend orchestrator for multi-tier allocation, replication, metadata indexing, and tier management.
Copy Function Registrymooncake-store/include/tiered_cache/copier_registry.h, mooncake-store/src/tiered_cache/copier_registry.cpp	Implements singleton CopierRegistry for registering memory-type and direct-path copy operations; provides CopierRegistrar helper for static-initialization-time registration of copiers.
Data Copier with Fallbackmooncake-store/include/tiered_cache/data_copier.h, mooncake-store/src/tiered_cache/data_copier.cpp	Implements DataCopierBuilder (populates copy matrix from global registry, validates DRAM path availability) and DataCopier (executes direct or two-step DRAM-based fallback copy between memory types).
Backend Implementationmooncake-store/src/tiered_cache/tiered_backend.cpp	Implements TieredBackend lifecycle, allocation with tier selection, write/copy/commit operations with per-key metadata and replica tracking, concurrent Get/Delete with multi-layer locking, and tier introspection.

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have updated the documentation.
I have added tests to prove my changes are effective.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/Mooncake/pull/1212 **Author:** [@YiXR](https://github.com/YiXR) **Created:** 12/15/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `xinyi/tiered-backend` --- ### 📝 Commits (1) - [`92e557e`](https://github.com/kvcache-ai/Mooncake/commit/92e557e4eca5ac4d0e39e22e5d3d17b0edfba7e7) [Store] feat: introduce tired backend ### 📊 Changes **8 files changed** (+986 additions, -0 deletions) <details> <summary>View changed files</summary> ➕ `mooncake-store/include/tiered_cache/cache_tier.h` (+86 -0) ➕ `mooncake-store/include/tiered_cache/copier_registry.h` (+86 -0) ➕ `mooncake-store/include/tiered_cache/data_copier.h` (+90 -0) ➕ `mooncake-store/include/tiered_cache/tiered_backend.h` (+197 -0) 📝 `mooncake-store/src/CMakeLists.txt` (+3 -0) ➕ `mooncake-store/src/tiered_cache/copier_registry.cpp` (+41 -0) ➕ `mooncake-store/src/tiered_cache/data_copier.cpp` (+116 -0) ➕ `mooncake-store/src/tiered_cache/tiered_backend.cpp` (+367 -0) </details> ### 📄 Description ## Description This PR introduce the abstract base classes for cache tier management. # CacheTier `mooncake-store/include/tiered_cache/cache_tier.h` ## **Responsibilities** This is an abstract base class that defines the unified interface a single cache tier (e.g., DRAM, SSD) must implement. It is responsible for managing the data lifecycle on a specific storage medium. ## **Core Interface**: * `Allocate`: Finds free space of `size` bytes. Does NOT copy data. * `Free`: Releases space at offset. Used when writes fail or explicitly freeing anonymous blocks. # DataCopier `mooncake-store/include/tiered_cache/data_copier.h` ## **Responsibilities**: Responsible for efficiently copying data between different memory types (`MemoryType`). This is a core utility class that decouples the data movement logic. ## **Core Functionality**: * Built via `DataCopierBuilder`, allowing for the registration of direct copy functions between different memory types (e.g., DRAM -> VRAM). * Provides a unified `Copy` interface. When a direct copy path is unavailable, it automatically employs a fallback mechanism using **DRAM as an intermediate buffer** (e.g., VRAM -> DRAM -> SSD) (when implementing a new type, compilation requires copy functions between the new type and DRAM). This greatly simplifies the integration of new storage media. * **Interaction**: Implementations of `TieredBackend` and `CacheTier` hold and use a `DataCopier` instance to execute all data copy operations, whether it's writing new data or moving data between tiers. # TieredBackend `mooncake-store/include/tiered_cache/tiered_backend.h` ## **Responsibilities**: Serves as the **Data Plane** for the tiered cache. It does not contain any complex caching policies (e.g., LRU) and is solely responsible for faithfully executing data operation instructions. ## Core Data Structures ### `DataSource` A generic descriptor for source data. ```cpp struct DataSource { const void* ptr; // Pointer to source buffer size_t size; // Data size MemoryType type; // Source memory type (DRAM, NVME, etc.) }; ``` ### `TieredLocation` Describes the physical location of a data segment. ```cpp struct TieredLocation { uint64_t tier_id; struct DataSource data; }; ``` ### `AllocationHandle` ```cpp /** * @struct AllocationEntry * @brief The internal state of an allocation. * acts as the "Control Block" for the resource. * When the last shared_ptr pointing to this entry dies, the destructor * triggers the physical release of the resource via the Backend. */ struct AllocationEntry { TieredBackend* backend; TieredLocation loc; AllocationEntry(TieredBackend* b, TieredLocation l) : backend(b), loc(l) {} // Destructor: Automatically releases the resource if valid ~AllocationEntry(); }; /** * @typedef AllocationHandle * @brief A reference-counted handle to a storage resource. */ using AllocationHandle = std::shared_ptr<AllocationEntry>; ``` ## Decomposed APIs The traditional atomic Put operation is broken down into three granular steps: * `Allocate` Reserves storage space locally without writing data. ```cpp AllocationHandle Allocate( size_t size, std::optional<uint64_t> preferred_tier); ``` Usage: Called by the AllocationStrategy to find free space. Returns: true if space is reserved; false otherwise. * `Write (local operation)` Writes data to a specific, reserved location. ```cpp bool Write(AllocationHandle handle, const DataSource& source); ``` Usage: Performs local write operation. If need remote write, please tranfer data into AllocationHandle by using TE. Returns: true on success. If failed, the caller must free the allocated space. * `Commit` Registers the key in the local metadata index. ```cpp bool Commit(const std::string& key, AllocationHandle handle); ``` Usage: Called after data is successfully written. Behavior: It binds the key to the offset internally. This operation is designed to be infallible (must succeed). If an internal error occurs (e.g., OOM), it raises a fatal error. * `Delete` Delete the key in the local metadata index. ```cpp bool Delete(const std::string& key, std::optional<uint64_t> tier_id) ``` If `tier_id` is specified, removes only the replica on that tier. If nullopt, removes ALL replicas for this key. ## Allocation Workflows This decoupled design supports two distinct allocation patterns required by the new architecture. **Each backend can only use one allocation strategy after mooncake store start up. If the backend wants to use Master-Centric Allocation, it must realize all the func to register it's info and allocator to Master.** ### Pattern A: Client-Centric Allocation (Local Decision) In this mode, the Client decides where to place data (e.g., "Try DRAM first, then SSD") and later informs the Master. #### Workflow * Decision The ClientCentricStrategy calls `Allocate(size, preferred_tier=xxx)`. TieredBackend checks local capacity and returns a `AllocationHandle`. * Execution The strategy calls `Write(data, AllocationHandle)`, this is a local write operation. Please use TE to transfer data if you want to do remote write. * Local Commit The strategy calls `Commit(key, AllocationHandle)`. The key is now visible to local `Get()` requests. * Global Sync The Client sends an asynchronous RPC `RegisterAllocation` to the Master. Master updates its directory: "Key X is located at Client Y, Tier Z". ### Pattern B: Master-Centric Allocation (Global Decision) In this mode, the Master manages all resources globally (Client should still call `MountSegment()` to register all the resources to Master) and tells the Client exactly where to write. Client doesn't need to save metadata for this backend anyway. #### Workflow: * Request The Client sends an RPC `PutStart` to the Master. * Decision Master checks its global resource view and returns `{tier_id: 1, offset: 0x2000}`. * Execution Data can be write by TE into tier directly. ### Data Migration & Consistency Migrates a key between tiers using a "Make-Before-Break" protocol to ensure strong consistency. ```cpp bool CopyData(const std::string& key, const DataSource& source, uint64_t dest_tier_id, MetadataSyncCallback sync_cb); ``` If DataSource is Master-Centric Allocation, Master must 'tell' Client where the data is. #### Workflow: * Allocate & Write Data is copied to the destination tier. * Sync Callback sync_cb is invoked. * Action The Client sends an RPC to the Master: "Update Key location to Tier New". * Safety If this RPC fails, CopyData returns false and rolls back the new copy. The old data remains valid. * Commit Local metadata is updated to point to the new tier. * Delete Only after the Master is updated and local commit succeeds, the old replica need to be deleted. If DataSource is Master-Centric Allocation, Master will delete the key by itself. If DataSource is Client-Centric Allocation, Client must call `Delete(key)` after `CopyData`. ```mermaid sequenceDiagram participant Client participant TieredBackend participant CacheTier participant DataCopier participant MemoryPool Client->>TieredBackend: Allocate(size, preferred_tier) activate TieredBackend TieredBackend->>TieredBackend: GetSortedTiers() TieredBackend->>CacheTier: Allocate(size) activate CacheTier CacheTier->>MemoryPool: allocate MemoryPool-->>CacheTier: ptr CacheTier-->>TieredBackend: AllocationHandle(tier_id, DataSource) deactivate CacheTier TieredBackend-->>Client: AllocationHandle deactivate TieredBackend Client->>TieredBackend: Write(source_data, handle) activate TieredBackend TieredBackend->>DataCopier: Copy(source, dest_from_handle) activate DataCopier DataCopier->>DataCopier: FindCopier(src_type, dest_type) alt Direct Path Exists DataCopier->>DataCopier: Execute direct copy else No Direct Path note over DataCopier: Both non-DRAM? DataCopier->>TieredBackend: Allocate temp DRAM buffer DataCopier->>DataCopier: Copy src → DRAM DataCopier->>DataCopier: Copy DRAM → dest end DataCopier-->>TieredBackend: success/failure deactivate DataCopier TieredBackend-->>Client: bool deactivate TieredBackend Client->>TieredBackend: Commit(key, handle) activate TieredBackend TieredBackend->>TieredBackend: Update metadata_index<br/>(per-key replicas) TieredBackend-->>Client: bool deactivate TieredBackend ``` ```mermaid sequenceDiagram participant Client participant TieredBackend participant MetadataIndex participant CacheTier1 participant CacheTier2 Client->>TieredBackend: Get(key, optional_tier_id) activate TieredBackend TieredBackend->>MetadataIndex: Lookup replicas for key activate MetadataIndex MetadataIndex-->>TieredBackend: replica list sorted by priority deactivate MetadataIndex alt tier_id specified TieredBackend->>CacheTier1: Retrieve from specific tier else best available TieredBackend->>CacheTier1: Retrieve from highest-priority tier end activate CacheTier1 CacheTier1-->>TieredBackend: AllocationHandle deactivate CacheTier1 TieredBackend-->>Client: AllocationHandle deactivate TieredBackend Client->>TieredBackend: Delete(key, optional_tier_id) activate TieredBackend TieredBackend->>MetadataIndex: Remove replica from key activate MetadataIndex MetadataIndex->>CacheTier2: Free allocation activate CacheTier2 CacheTier2-->>MetadataIndex: success deactivate CacheTier2 MetadataIndex-->>TieredBackend: bool deactivate MetadataIndex TieredBackend-->>Client: bool deactivate TieredBackend ``` ## Type of Change * Types - [ ] Bug fix - [x] New feature - [ ] Transfer Engine - [x] Mooncake Store - [ ] Mooncake EP - [ ] Integration - [ ] P2P Store - [ ] Python Wheel - [ ] Breaking change - [ ] CI/CD - [ ] Documentation update - [ ] Other Cohort / File(s) | Summary -- | -- Tier Abstraction & Metadatamooncake-store/include/tiered_cache/cache_tier.h, mooncake-store/include/tiered_cache/tiered_backend.h | Defines MemoryType enum (DRAM, NVME, UNKNOWN), DataSource struct, abstract CacheTier interface with lifecycle (Init, Allocate, Free) and introspection methods; introduces tiered storage data structures (TieredLocation, TierView, AllocationEntry, MetadataEntry) and TieredBackend orchestrator for multi-tier allocation, replication, metadata indexing, and tier management. Copy Function Registrymooncake-store/include/tiered_cache/copier_registry.h, mooncake-store/src/tiered_cache/copier_registry.cpp | Implements singleton CopierRegistry for registering memory-type and direct-path copy operations; provides CopierRegistrar helper for static-initialization-time registration of copiers. Data Copier with Fallbackmooncake-store/include/tiered_cache/data_copier.h, mooncake-store/src/tiered_cache/data_copier.cpp | Implements DataCopierBuilder (populates copy matrix from global registry, validates DRAM path availability) and DataCopier (executes direct or two-step DRAM-based fallback copy between memory types). Backend Implementationmooncake-store/src/tiered_cache/tiered_backend.cpp | Implements TieredBackend lifecycle, allocation with tier selection, write/copy/commit operations with per-key metadata and replica tracking, concurrent Get/Delete with multi-layer locking, and tier introspection. ## How Has This Been Tested?  ## Checklist - [x] I have performed a self-review of my own code. - [ ] I have updated the documentation. - [ ] I have added tests to prove my changes are effective. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-29 21:32:36 +00:00