Compare commits

...

12 Commits

Author SHA1 Message Date
yejj710 a6bfb2e195 do not use lru list defautly & add log for remove-event 2026-04-22 11:34:22 +08:00
yejj710 c837624931 Add LRU eviction policy to the index data of Mooncake-Store 2026-04-19 14:49:08 +08:00
yejj c67649c83a
support remove event & add some info-level log (#1881) 2026-04-13 21:09:18 +08:00
yejj 6ab4e40ee5
bugfix: Support multiple identical instances under a single tenant (#1840) 2026-04-08 15:43:51 +08:00
Asher Zhang 80c8b41cad
update Mooncake Store KVEvent Publisher (#1638) 2026-03-10 19:25:44 +08:00
Asher Zhang 770031b101
init publisher (#1635) 2026-03-09 20:25:49 +08:00
yejj710 fb7313b87d Merge remote-tracking branch 'center/main' into sync_code_from_master 2026-03-09 10:51:21 +08:00
yejj c15a911ae0
support mooncake_conductor global kv-indexer (#1614) 2026-03-05 20:26:51 +08:00
Ryan·Chang 45402ff9a1
[Misc] feat: support dynamic register and unregister API for KV Event Manager (#1569)
* Register/Unregister

* Register/Unregister

* add Register/Unregister interface

* fix: address review comments

* feat: replace LoraID with LoraName and optimize unregister
2026-02-28 21:21:57 +08:00
Ryan·Chang fbff581166
[Misc] Configure slog log level via environment variable (#1529)
* log level

* Delete .trae directory

Delete additional files.

---------

Co-authored-by: Ryanic-Chang <ryanchang1412@gmail.com>
2026-02-10 16:59:58 +08:00
yejj 0cc4fea731
Merge pull request #1478 from yejj710/kv-indexer-dev
[Misc][Dev] add indexer demo version & using example
2026-02-06 17:40:32 +08:00
Liziqi-77 d0ec8173d3 add indexer demo & using example 2026-02-03 10:23:48 +08:00
45 changed files with 10291 additions and 5 deletions

View File

@ -174,6 +174,13 @@ jobs:
echo "✅ Coverage collected successfully"
fi
- name: Test mooncake conductor
run: |
cd mooncake-conductor/conductor-ctrl
go mod tidy
go test ./... -v -cover
shell: bash
- name: Generate Python version tag
id: generate_tag_build
run: |

View File

@ -17,6 +17,7 @@ option(WITH_STORE "build mooncake store library and sample code" ON)
option(WITH_P2P_STORE "build p2p store library and sample code" OFF)
option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
option(WITH_EP "build mooncake with expert parallelism support" OFF)
option(WITH_CONDUCTOR "build mooncake conductor and sample code" OFF)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
set(PYTHON_EXECUTABLE "python3")
@ -71,3 +72,9 @@ if (WITH_P2P_STORE)
add_subdirectory(mooncake-p2p-store)
message(STATUS "P2P Store will be built")
endif()
if (WITH_CONDUCTOR)
add_subdirectory(mooncake-conductor)
message(STATUS "Conductor will be built")
endif()

View File

@ -127,6 +127,7 @@ SYSTEM_PACKAGES="build-essential \
libmsgpack-dev \
libzstd-dev \
libasio-dev \
libzmq3-dev \
pkg-config \
patchelf \
libc6-dev \
@ -195,6 +196,145 @@ check_success "Failed to install yalantinglibs"
print_success "yalantinglibs installed successfully"
# Install cppzmq (From source, similar to yalantinglibs)
print_section "Installing cppzmq from source"
CPPZMQ_DIR="${THIRDPARTIES_DIR}/cppzmq"
# Check if cppzmq directory already exists
if [ -d "$CPPZMQ_DIR" ]; then
echo -e "${YELLOW}cppzmq directory already exists. Removing for fresh install...${NC}"
rm -rf "$CPPZMQ_DIR"
check_success "Failed to remove existing cppzmq directory"
fi
# Clone cppzmq
echo "Cloning cppzmq from ${GITHUB_PROXY}/zeromq/cppzmq.git"
git clone ${GITHUB_PROXY}/zeromq/cppzmq.git "$CPPZMQ_DIR"
check_success "Failed to clone cppzmq"
# Build and install cppzmq
cd "$CPPZMQ_DIR"
check_success "Failed to change to cppzmq directory"
# Checkout a specific stable version v4.11.0
echo "Checking out cppzmq version v4.11.0..."
git checkout v4.11.0
check_success "Failed to checkout cppzmq version v4.11.0"
mkdir -p build
check_success "Failed to create build directory"
cd build
check_success "Failed to change to build directory"
echo "Configuring cppzmq..."
# Key configuration: Ensure it finds the system-installed libzmq
# and configure cppzmq to install to system directory /usr/local
cmake .. \
-DCMAKE_INSTALL_PREFIX=/usr/local \
-DCPPZMQ_BUILD_TESTS=OFF
check_success "Failed to configure cppzmq"
echo "Building cppzmq (using $(nproc) cores)..."
cmake --build . -j$(nproc)
check_success "Failed to build cppzmq"
echo "Installing cppzmq..."
cmake --install .
check_success "Failed to install cppzmq"
ldconfig
print_success "cppzmq installed successfully"
# Install msgpack-cxx from release source
print_section "Installing msgpack-cxx from release source (v7.0.0)"
# Change to thirdparties directory to ensure consistent starting point
cd "$THIRDPARTIES_DIR"
check_success "Failed to change to thirdparties directory"
MSGPACK_DIR="${THIRDPARTIES_DIR}/msgpack-cxx"
# Check if msgpack directory already exists
if [ -d "$MSGPACK_DIR" ]; then
echo -e "${YELLOW}msgpack-cxx directory already exists. Removing for fresh install...${NC}"
rm -rf "$MSGPACK_DIR"
check_success "Failed to remove existing msgpack-cxx directory"
fi
# Create msgpack directory
mkdir -p "$MSGPACK_DIR"
check_success "Failed to create msgpack-cxx directory"
cd "$MSGPACK_DIR"
check_success "Failed to change to msgpack-cxx directory"
# Download and extract msgpack-cxx 7.0.0 release
MSGPACK_VERSION="7.0.0"
MSGPACK_TARBALL="msgpack-cxx-${MSGPACK_VERSION}.tar.gz"
MSGPACK_URL="${GITHUB_PROXY}/msgpack/msgpack-c/releases/download/cpp-${MSGPACK_VERSION}/${MSGPACK_TARBALL}"
echo "Downloading msgpack-cxx v${MSGPACK_VERSION} from ${MSGPACK_URL}"
wget --show-progress -O "$MSGPACK_TARBALL" "$MSGPACK_URL"
check_success "Failed to download msgpack-cxx release"
# Verify the downloaded file is not empty
if [ ! -s "$MSGPACK_TARBALL" ]; then
print_error "Downloaded msgpack-cxx tarball is empty"
fi
# Extract the tarball
echo "Extracting msgpack-cxx source..."
tar -xzf "$MSGPACK_TARBALL"
check_success "Failed to extract msgpack-cxx tarball"
# The extracted directory name includes the version
EXTRACTED_DIR="msgpack-cxx-${MSGPACK_VERSION}"
if [ ! -d "$EXTRACTED_DIR" ]; then
print_error "Extracted directory '$EXTRACTED_DIR' not found"
fi
# Move into the extracted directory
cd "$EXTRACTED_DIR"
check_success "Failed to change to extracted msgpack-cxx directory"
# Build and install msgpack-cxx
mkdir -p build
check_success "Failed to create build directory"
cd build
check_success "Failed to change to build directory"
echo "Configuring msgpack-cxx..."
cmake .. \
-DCMAKE_INSTALL_PREFIX=/usr/local \
-DMSGPACK_BUILD_TESTS=OFF \
-DMSGPACK_BUILD_EXAMPLES=OFF \
-DMSGPACK_USE_BOOST=OFF
check_success "Failed to configure msgpack-cxx"
echo "Building msgpack-cxx (using $(nproc) cores)..."
cmake --build . -j$(nproc)
check_success "Failed to build msgpack-cxx"
echo "Installing msgpack-cxx..."
cmake --install .
check_success "Failed to install msgpack-cxx"
ldconfig # Update library cache
# Clean up: remove the downloaded tarball
rm -f "../${MSGPACK_TARBALL}"
print_success "Cleaned up downloaded tarball"
print_success "msgpack-cxx v${MSGPACK_VERSION} installed successfully"
# Initialize and update git submodules
print_section "Initializing Git Submodules"

View File

@ -0,0 +1,430 @@
# KV Event Subscriber Guide
This document provides guidance for developers who wish to subscribe to KV cache events in the Mooncake system. It focuses on the message schema and deserialization methods needed to properly handle events sent by the event publisher.
---
## Event Types & Meanings
Before diving into the technical schema, it's important to understand what each event type represents and when they are triggered in the Mooncake system.
### Event Overview Table
| Event Type | Trigger Condition | Purpose | Key Characteristics |
|------------|-------------------|---------|---------------------|
| **BlockStoreEvent** | First storage of KV cache block | Initial block creation with metadata | ✅ contains `StoreEventInfo` metadata<br>✅ Represents initial creation<br>✅ Generated during `Put` operations |
| **BlockUpdateEvent** | Replica management operations (copy/remove/migrate) | Track replica distribution changes | ❌ No `StoreEventInfo` metadata<br>✅ Focuses on replica locations<br>✅ Internal system operations |
| **RemoveAllEvent** | System-wide cache clearance | Signal cache invalidation | ❌ No additional fields<br>✅ System-wide invalidation<br>✅ Maintenance operations |
### Event Type Quick Reference
#### 🏗️ BlockStoreEvent - Storage Event
```mermaid
graph LR
A[Put Operation] --> B[BlockStoreEvent]
B --> C[Contains Metadata]
C --> D[Downstream Processing]
```
#### 🔄 BlockUpdateEvent - Replica Event
```mermaid
graph LR
A[Replica Operation] --> B[BlockUpdateEvent]
B --> C[Replica Location Changes]
C --> D[Downstream Processing]
```
#### 🧹 RemoveAllEvent - Clearance Event
```mermaid
graph LR
A[System Maintenance] --> B[RemoveAllEvent]
B --> C[Cache Invalidation]
C --> D[Downstream Processing]
```
### Field Definitions at a Glance
#### Core Event Fields
| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `mooncake_key` | `std::string` | ✅ All events | Unique identifier for cached object |
| `replicas` | Nested array | ✅ BlockStore/Update | Replica locations: `[type, location]` |
#### `StoreEventInfo` Fields (`BlockStoreEvent` only)
When processing `BlockStoreEvent`events, pay special attention to the StoreEventInfo fields which are appended to the event. **These fields have specific default values that indicate when they are not set**:
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `model_name` | `std::string` | `""` | Model identifier |
| `block_size` | `uint32_t` | `0` | Block size in bytes |
| `block_hash` | `std::string` | `""` | Current block hash |
| `parent_block_hash` | `std::string` | `""` | Parent block hash |
| `token_ids` | `std::vector<uint32_t>` | `[]` | Token ID sequence |
---
## Event Message Schema
The KV event system uses ZeroMQ (ZMQ) for message transport and MessagePack for serialization. Events are sent as multipart ZMQ messages containing three components:
### ZMQ Message Structure
Each event message consists of 3 ZMQ message parts:
```plainText
1. topic Part: Contains the topic string (default: "mooncake")
2. Sequence Number Part: 8-byte big-endian unsigned integer representing the sequence number
3. Payload Part: MessagePack-serialized event batch data
```
### Event Batch Schema
The payload part contains a serialized `EventBatch`object with the following structure:
```typescript
[
ts, // First element: Timestamp (double) - seconds since UNIX epoch
[ // Second element: Event list (array)
[event1_data], // Event 1
[event2_data], // Event 2
... // Additional events
]
]
```
Each event in the batch is a serialized event object that begins with an event type identifier string followed by its specific data fields.
### Event Types Schema
There are three main event types supported by the system:
#### Schema
```typescript
// Event triggered on the first storage occurrence - contains 8 fields
BlockStoreEvent {
"BlockStoreEvent", // Event type identifier, string type
std::string mooncake_key, // Mooncake key
[ // Replica location list (nested arrays)
["memory", "transport_endpoint"], // Memory replica
["disk", "file_path"], // Disk replica
["local_disk", "transport_endpoint"] // Local disk replica
],
std::string model_name, // Model name
uint32_t block_size, // Block size
std::string block_hash, // Current block hash
std::string parent_block_hash, // Parent block hash
std::vector<uint32_t> token_ids // Token ID sequence
}
// Contains all internal Mooncake system operations - contains 3 fields
BlockUpdateEvent {
"BlockUpdateEvent", // Event type identifier
std::string mooncake_key,
[ // Replica location list
["memory", "transport_endpoint"],
["disk", "file_path"],
["local_disk", "transport_endpoint"]
]
}
// Event for clearing all KV Cache in Mooncake - contains 1 field
RemoveAllEvent {
"RemoveAllEvent" // Event type identifier
}
```
#### Example
```typescript
// BlockStoreEvent
[
"BlockStoreEvent", // Event type identifier
"key_12345", // mooncake key
[ // Replica location list (nested arrays)
["memory", "tcp://192.168.1.10:6000"], // Memory type location
["disk", "/data/blocks/block_12345.bin"], // Disk type location
["local_disk", "tcp://192.168.1.10:7000"] // Local disk type location
],
"llama2-7b", // Model name
512, // block size
"0x41234125", // Current block hash
"0x51512342", // Parent block hash
[1,2,3,4,5] // Token id list
]
// BlockUpdateEvent
[
"BlockUpdateEvent", // Event type identifier
"key_12345", // mooncake key
[ // Replica location list (nested arrays)
["memory", "tcp://192.168.1.10:6000"], // Memory type location
],
]
// RemoveAllEvent
[
"RemoveAllEvent" // Event type identifier
]
```
---
## Deserialization Steps
To properly deserialize events from the KV event system, follow these steps:
1. **Receive the multipart message**: Use `zmq::recv_multipart`or equivalent to receive all 3 parts of the message.
2. **Extract the topic**: The first part contains the topic string.
3. **Extract and convert the sequence number**: The second part contains an 8-byte big-endian unsigned integer. On little-endian systems, convert it using `be64toh()`or equivalent function.
4. **Deserialize the payload**: The third part contains MessagePack-serialized data. Deserialize it to get the `EventBatch`.
5. **Process individual events**: Iterate through the events in the batch and handle each according to its type identifier (the first element in each event array).
## Required Libraries
- **ZeroMQ library**: For receiving multipart messages
- **MessagePack library**: For deserializing the payload
- **Byte order conversion functions**: For converting sequence numbers between big-endian and host byte order
---
## Event Batch Timestamp Handling
The event batch timestamp field `ts` is a double-precision floating-point number representing seconds since the UNIX epoch (January 1, 1970). When deserializing, convert this timestamp appropriately:
<details>
<summary>Click to expand: Python example</summary>
```python
import datetime
# Python example
timestamp = event_batch[0] # double type seconds
dt = datetime.datetime.fromtimestamp(timestamp)
print(f"Event batch timestamp: {dt}")
```
</details>
<details>
<summary>Click to expand: GoLang example</summary>
```go
// Go example
timestamp := eventBatch[0].(float64)
t := time.Unix(int64(timestamp), 0)
fmt.Printf("Event batch timestamp: %v\n", t)
```
</details>
---
## Python Example
Here's a Python example showing how to subscribe to and deserialize KV events:
<details>
<summary>Click to expand: Python example</summary>
```python
def deserialize_block_store_event(event_array):
if len(event_array) < 8:
raise ValueError("Invalid BlockStoreEvent array length")
# Extract fields by fixed position
event_type = event_array[0] # "BlockStoreEvent"
mooncake_key = event_array[1]
replicas = event_array[2] # Nested array of replica locations
model_name = event_array[3] # String, defaults to ""
block_size = event_array[4] # Integer, defaults to 0
block_hash = event_array[5] # String, defaults to ""
parent_block_hash = event_array[6] # String, defaults to ""
token_ids = event_array[7] # List of integers, defaults to []
# Handle default values according to StoreEventInfo specifications
if model_name == "":
# Model name not set, use default or skip processing
model_name = "unknown"
if block_size == 0:
# Invalid block size, may indicate an error
raise ValueError("Invalid block size: 0")
if block_hash == "":
# block_hash is not set, use None to indicate missing value
block_hash = None
if parent_block_hash == "":
# No parent block (root block), handle appropriately
parent_block_hash = None
return {
"type": event_type,
"mooncake_key": mooncake_key,
"replicas": replicas,
"model_name": model_name,
"block_size": block_size,
"block_hash": block_hash,
"parent_block_hash": parent_block_hash,
"token_ids": token_ids
}
def process_replica_locations(replicas):
"""Process nested replica location arrays"""
replica_info = []
for replica in replicas:
if len(replica) != 2:
continue
replica_type = replica[0] # "memory", "disk", or "local_disk"
location = replica[1] # Endpoint or file path
replica_info.append({"type": replica_type, "location": location})
return replica_info
```
</details>
## GoLang Example
Here's a Go example showing how to subscribe to and deserialize KV events:
<details>
<summary>Click to expand: GoLang example</summary>
```go
func deserializeBlockStoreEvent(eventSlice []interface{}) (map[string]interface{}, error) {
if len(eventSlice) < 8 {
return nil, fmt.Errorf("invalid BlockStoreEvent array length: %d", len(eventSlice))
}
result := make(map[string]interface{})
// Extract fields by fixed position
result["type"] = eventSlice[0].(string)
result["mooncake_key"] = eventSlice[1].(string)
result["replicas"] = eventSlice[2]
result["model_name"] = eventSlice[3].(string)
// Handle block_size with proper type assertion
if blockSize, ok := eventSlice[4].(uint32); ok {
result["block_size"] = blockSize
if blockSize == 0 {
return nil, fmt.Errorf("invalid block size: 0")
}
} else {
return nil, fmt.Errorf("invalid block_size type")
}
// block_hash field handling with default value check
blockHash := eventSlice[5].(string)
if blockHash == "" {
result["block_hash"] = nil // Indicates missing value
} else {
result["block_hash"] = blockHash
}
// parent_block_hash field handling
parentBlockHash := eventSlice[6].(string)
if parentBlockHash == "" {
result["parent_block_hash"] = nil // No parent block
} else {
result["parent_block_hash"] = parentBlockHash
}
result["token_ids"] = eventSlice[7]
return result, nil
}
func processReplicas(replicas interface{}) ([]map[string]string, error) {
// Process nested replica arrays
replicaSlice, ok := replicas.([]interface{})
if !ok {
return nil, fmt.Errorf("invalid replicas type")
}
var replicaInfo []map[string]string
for _, replica := range replicaSlice {
replicaArr, ok := replica.([]interface{})
if !ok || len(replicaArr) != 2 {
continue
}
replicaType, ok1 := replicaArr[0].(string)
location, ok2 := replicaArr[1].(string)
if ok1 && ok2 {
replicaInfo = append(replicaInfo, map[string]string{
"type": replicaType,
"location": location,
})
}
}
return replicaInfo, nil
}
```
</details>
---
## Replay Functionality
The KV event system supports replay functionality, allowing subscribers to request historical events. To implement replay:
1. **Send replay request**: Send a 3-part ZMQ message to the replay endpoint:
- Part 1: Client identifier (any data that identifies your client)
- Part 2: Empty frame
- Part 3: Starting sequence number (8-byte big-endian unsigned integer)
2. **Receive replay events**: The system will send historical events starting from the requested sequence number. Each replayed event is sent as a regular ZMQ multipart message.
3. **Replay end marker**: When all available historical events have been sent, the system sends a special end marker message containing the magic sequence `0xFFFFFFFFFFFFFFFF` (8 bytes of 0xFF).
4. **Handle replay completion**: Upon receiving the end marker, you know that the replay session has completed.
5. **Important considerations**:
- Replayed events may be older than real-time events already processed
- Implement duplicate detection if needed
- Replay buffer size is limited by the system configuration
Example replay request handling:
<details>
<summary>Click to expand: Python example</summary>
```python
def send_replay_request(socket, start_seq):
"""Send replay request to the replay endpoint"""
client_id = b"my_client_001" # Your client identifier
empty_frame = b"" # Empty frame
seq_be = start_seq.to_bytes(8, byteorder='big') # Big-endian sequence
messages = [client_id, empty_frame, seq_be]
socket.send_multipart(messages)
def handle_replay_message(messages):
"""Handle incoming replay message"""
if len(messages) == 4:
client_id = messages[0]
empty_frame = messages[1]
seq_be = messages[2]
payload = messages[3]
# Check for end marker
if payload == b'\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF':
print("Replay session completed")
return None
# Process regular replay event
seq = int.from_bytes(seq_be, byteorder='big')
return deserialize_payload(payload)
return None
```
</details>

View File

@ -0,0 +1,459 @@
# Mooncake Event System Developer Manual
## 1. System Overview
KVEventSystem is an asynchronous event processing system based on the publish-subscribe pattern. It encapsulates the complex logic of event publishing, consumption, and queue management using the Facade Pattern, providing a concise and unified interface for upper-layer applications.
### 1.1 Core Features
- **Asynchronous Event Publishing**: Supports publishing various types of events in a non-blocking manner.
- **Batch Processing**: Automatically batches and merges events to optimize network transmission efficiency.
- **Reliable Transmission**: Built-in retry mechanisms and queue buffering ensure no event loss.
- **Real-time Monitoring**: Provides detailed runtime statistical information.
- **Flexible Configuration**: Performance and reliability can be adjusted through configuration parameters.
### 1.2 Architecture Design
#### 1.2.1 Architecture Logic Diagram
```mermaid
graph TD
ApplicationLayer[API call site] -->|invoke interface<br>publish KVEvent| KVEventProducer
IR[indexer/router]
subgraph KVEventSystem
KVEventProducer -->|enqueue| KVEventQueue
subgraph KVEventConsumer
ZMQ
end
KVEventQueue -->|dequeue| KVEventConsumer
end
ZMQ -->|event batch serialization<br>ZMQ network transmission| IR
```
#### 1.2.2 Class Relationship Diagram
```mermaid
classDiagram
%% Event Class Hierarchy
class KVCacheEvent {
<<interface>>
+pack(msgpack::packer~msgpack::sbuffer~) void
+type_tag() string_view
}
KVCacheEvent <|-- BlockStoreEvent
KVCacheEvent <|-- BlockUpdateEvent
KVCacheEvent <|-- RemoveAllEvent
%% Core Components
class KVEventProducer {
-shared_ptr~KVEventQueue~ event_queue_
-unique_ptr~ThreadPool~ enqueue_pool_ : enqueue worker(1)
+publish~Event, Args...~(args) future~bool~
+publish_event_async(KVEventPtr) future~bool~
+shutdown() void
}
class KVEventConsumer {
-zmq::context_t context_
-shared_ptr~KVEventQueue~ event_queue_
-jthread publisher_thread_
+KVEventConsumer(shared_ptr~KVEventQueue~, Config)
+shutdown() void
+is_running() bool
}
class ThreadSafeQueue~T~ {
<<alias KVEventQueue>>
}
%% Facade Class
class KVEventSystem {
-shared_ptr~KVEventQueue~ event_queue_
+KVEventSystem(const KVEventPublisherConfig&)
+~KVEventSystem()
+shutdown() void
+is_running() bool
+publish~Event, Args...~(args) future~bool~
}
%% Configuration and Event Classes
class KVEventPublisherConfig {
}
class StoreEventInfo {
}
class EventBatch {
+serialize() msgpack::sbuffer
}
%% Relationships
KVEventSystem --> KVEventPublisherConfig : reads
KVEventSystem --> KVEventProducer : contains
KVEventSystem --> KVEventConsumer : contains
KVEventSystem --> ThreadSafeQueue~KVEventPtr~ : contains
KVEventProducer --> ThreadSafeQueue~KVEventPtr~ : enqueues
KVEventConsumer --> ThreadSafeQueue~KVEventPtr~ : dequeues
KVEventProducer --> KVCacheEvent : creates
BlockStoreEvent --> StoreEventInfo : contains
KVCacheEvent <.. EventBatch : aggregates
EventBatch <.. KVEventConsumer : serialize & publish
```
#### 1.2.3 Data Flow Sequence Diagram
```mermaid
sequenceDiagram
participant App as APICallSite
participant ES as KVEventSystem
participant EP as KVEventProducer
participant TSQ as ThreadSafeQueue
participant ZMQ as KVEventConsumer
participant Network as ZeroMQNetwork
App->>ES: 1. publish<KVEvent>(...)
ES->>EP: 2. publish<KVEvent>(...)
EP->>TSQ: 3. push(event)<br>(create KVCacheEvent object)
Note over TSQ: KVEvent Buffer
ZMQ->>TSQ: 4. peek_batch()
ZMQ->>ZMQ: 5. batch processing
ZMQ->>Network: 6. ZeroMQ Publish
Network-->>ZMQ: 7. Publish confirmation
ZMQ-->>TSQ: 8. pop_batch()
ES-->>App: 9. return future
Note over App,Network: Asynchronous processing flow
```
------
## 2. Key Class Interface Declarations
### 2.1 Event Classes (kv_event.hpp)
#### Event Class Relationships
| Class Name | Parent Class | Key Fields | Description |
| ------------------ | -------------- | ---------------------------------------------- | --------------------------------------------- |
| `KVCacheEvent` | - | No fields | Abstract base class, defines event interface |
| `StoreEventInfo` | - | All fields | Container for additional data of store events |
| `BlockStoreEvent` | `KVCacheEvent` | `mooncake_key`, `replicas`, `store_event_info` | Block storage event |
| `BlockUpdateEvent` | `KVCacheEvent` | `mooncake_key`, `replicas` | Block update event |
| `RemoveAllEvent` | `KVCacheEvent` | No fields | Clear all cache event |
**UML Class Diagram**:
```mermaid
classDiagram
%% Event Class Hierarchy
class KVCacheEvent {
<<abstract>>
+pack(msgpack::packer~msgpack::sbuffer~& pk) void
+type_tag() string_view
}
KVCacheEvent <|-- BlockStoreEvent
KVCacheEvent <|-- BlockUpdateEvent
KVCacheEvent <|-- RemoveAllEvent
%% StoreEventInfo Structure
class StoreEventInfo {
+std::string model_name
+uint32_t block_size
+std::string block_hash
+std::string parent_block_hash
+std::vector~uint32_t~ token_ids
}
%% BlockStoreEvent Class
class BlockStoreEvent {
+std::string mooncake_key
+std::vector~Replica::Descriptor~ replicas
+StoreEventInfo store_event_info
+BlockStoreEvent(key, replica_list, info)
+pack(pk) void
+type_tag() string_view
}
%% BlockUpdateEvent Class
class BlockUpdateEvent {
+std::string mooncake_key
+std::vector~Replica::Descriptor~ replicas
+BlockUpdateEvent(key, replica_list)
+pack(pk) void
+type_tag() string_view
}
%% RemoveAllEvent Class
class RemoveAllEvent {
+RemoveAllEvent()
+pack(pk) void
+type_tag() string_view
}
%% Relationships
BlockStoreEvent *-- StoreEventInfo : contains
```
------
#### Event Class Types
##### Event Base Class: `KVCacheEvent`
> The **abstract base class** for all KV cache events, defining the **unified interface** for event serialization and type identification. It uses the abstract base class pattern to provide a unified event processing interface. All concrete event types must inherit from this class and implement the interface methods.
```c++
struct KVCacheEvent {
virtual void pack(msgpack::packer<msgpack::sbuffer>& pk) const = 0; // Serialization method
virtual std::string_view type_tag() const = 0; // Type identification method
}
```
**Explanation**
| Method Name | Return Type | Parameters | Description |
| ----------- | ------------------ | --------------------------------------- | ------------------------------------------------------------ |
| `pack` | `void` | `msgpack::packer<msgpack::sbuffer>& pk` | Pure virtual function, serializes the event into MessagePack format. Derived classes must implement this method to provide type-specific serialization logic. |
| `type_tag` | `std::string_view` | None | Pure virtual function, returns a string view of the event type identifier, used for type identification during deserialization. |
------
##### Store Event Pass-Through Meta Data Structure: `StoreEventInfo`
> The `StoreEventInfo`struct is used to carry business metadata for KV cache storage events. All fields in this structure are optional; which fields to populate should be determined by the specific business scenario of the upper-layer cache-aware component (such as indexer or router). The Mooncake system acts only as a transparent pipeline for transmitting these fields and will not perform any business logic validation or interpretation of their content.
```c++
struct StoreEventInfo {
std::string model_name{""};
uint32_t block_size{0};
std::string block_hash{""};
std::string parent_block_hash{""};
std::vector token_ids{};
};
```
**Explanation**
| Field Name | Type | Default Value | Constraints/Explanation |
| ------------------- | ----------------------- | ------------- | --------------------------------------------------------- |
| `model_name` | `std::string` | `""` | Model name identifier, empty indicates not set |
| `block_size` | `uint32_t` | `0` | Block size (bytes), 0 indicates invalid value |
| `block_hash` | `std::string` | `""` | Hash of the current block, used for unique identification |
| `parent_block_hash` | `std::string` | `""` | Parent block hash, used for dependency chain |
| `token_ids` | `std::vector<uint32_t>` | Empty vector | Token ID sequence, can be empty |
------
##### Storage Event Structure: `BlockStoreEvent`
> Event type triggered when a KV cache block is stored for the first time, typically containing the additional storage information `StoreEventInfo`that needs to be passed.
```c++
struct BlockStoreEvent {
std::string mooncake_key;
std::vector<Replica::Descriptor> replicas;
StoreEventInfo store_event_info;
}
```
**Explanation**
| Field Name | Type | Default Value (set by constructor) | Constraints/Explanation |
| ------------------ | ---------------------------------- | ---------------------------------- | -------------------------------------------- |
| `mooncake_key` | `std::string` | N/A | Unique identifier of the cached `ObjectMetadata` in the `MasterService` |
| `replicas` | `std::vector<Replica::Descriptor>` | N/A | List of replica descriptors containing replica location and type information |
| `store_event_info` | `StoreEventInfo` | N/A | Additional information required for KVCache-aware algorithms |
> The `mooncake_key` field represents the primary key of the cached object in the distributed KV cache system. In the implementation, this value may originate from method parameters (e.g., the key parameter in `PutEnd` method) or from `it->first` when iterating through metadata maps (where it is an iterator over `std::unordered_map<std::string, ObjectMetadata>` in `MasterService::metadata_shards_`).
**Event Format Example**
```javascript
[
"BlockStoreEvent", // Event type identifier
"key_001", // mooncake key
[ // Replica location list
["memory", "tcp://192.168.1.10:6000"], // Memory type location
["disk", "/data/blocks/block_12345.bin"], // Disk type location
],
"llama2-7b", // Model name
512, // block size
"0x41234125", // Current block hash
"0x51512342", // Parent block hash
[1,2,3,4,5] // Token id list
]
```
------
##### Replica Change Event Structure: `BlockUpdateEvent`
> Event type triggered by any replica update operation other than storage events, including copy/eviction/migration.
```c++
struct BlockUpdateEvent {
std::string mooncake_key;
std::vector<Replica::Descriptor> replicas;
}
```
**Explanation**
| Field Name | Type | Default Value | Constraints/Explanation |
| -------------- | ---------------------------------- | ------------- | ----------------------------------------------- |
| `mooncake_key` | `std::string` | N/A | Unique identifier of the cached `ObjectMetadata` in the `MasterService` |
| `replicas` | `std::vector<Replica::Descriptor>` | N/A | List of replica descriptors containing replica location and type information |
**Event Format Example**
```javascript
[
"BlockUpdateEvent", // Event type identifier
"key_12345", // mooncake key
[ // Global replica location list
["memory", "tcp://192.168.1.10:6000"],
["memory", "tcp://192.168.1.11:5001"],
["disk", "/data/blocks/block_12345.bin"],
],
]
```
------
##### Clear All Cache Event Structure: `RemoveAllEvent`
> Event triggered to clear all caches in Mooncake-Store.
```c++
RemoveAllEvent {
}
```
**Explanation**: *This event has no fields, identified only by its type tag.*
**Event Format Example**
```javascript
[
"RemoveAllEvent"
]
```
------
#### Serialization-Related Types
##### Event Batch Structure: `EventBatch`
> Used to batch serialize events for transmission, thereby reducing network overhead and improving throughput, while timestamps ensure event ordering.
```c++
struct EventBatch {
double ts;
std::vector<std::shared_ptr<KVCacheEvent>> events;
}
```
**Explanation**
| Field Name | Type | Default Value (set by constructor) | Constraints/Explanation |
| ---------- | -------------------------------------------- | ---------------------------------- | ------------------------------------------------------------ |
| `ts` | `double` | N/A | Timestamp of batch creation, using Unix timestamp (seconds since 1970-01-01 00:00:00 UTC, floating-point), used by the receiver to process batches in chronological order. |
| `events` | `std::vector<std::shared_ptr<KVCacheEvent>>` | N/A | Array of event pointers, can contain any event object derived from `KVCacheEvent`. Events in the batch are serialized and processed in array order. |
**Event Batch Format Example**
The event batch is serialized as an array containing two elements:
- The first element is the timestamp (floating-point number)
- The second element is an array of events, where each event is its corresponding type's serialized array.
```javascript
[
170000000.123, // Element 1: Timestamp (double)
[ // Element 2: Event list (array)
[ // Event 1: BlockStoreEvent
"BlockStoreEvent",
"key_001",
[
["memory", "tcp://192.168.1.10:6000"],
["disk", "/data/blocks/block_12345.bin"]
],
"llama2-7b",
512,
"0x41234125",
"0x51512342",
[1, 2, 3, 4, 5]
],
[ // Event 2: BlockUpdateEvent
"BlockUpdateEvent",
"key_12345",
[
["memory", "tcp://192.168.1.10:6000"],
["memory", "tcp://192.168.1.11:5001"],
["disk", "/data/blocks/block_12345.bin"]
]
],
[ // Event 3: RemoveAllEvent
"RemoveAllEvent"
]
]
]
```
------
### 2.2 Event System Facade Class (kv_event_system.h)
#### 2.2.1 Configuration Item Descriptions
| Category | Parameter Name | Corresponding Field in `KVEventPublisherConfig` Structure | Type | Default Value | Description |
| ---------------------- | ------------------------------------- | --------------------------------------------------------- | ------------- | ----------------- | ------------------------------------------------------------ |
| **Basic Switch** | `enable_kv_event_publish` | None (system switch) | `bool` | `false` | Master switch for event publishing functionality. System starts event publishing only when set to `true`. |
| **Network Config** | `kv_event_publisher_endpoint` | `endpoint` | `std::string` | `"tcp://*:19997"` | ZeroMQ bind address, supports `tcp://`and `ipc://`protocols. |
| | `kv_event_publisher_replay_endpoint` | `replay_endpoint` | `std::string` | `""` | Replay endpoint address. Empty value disables replay functionality. |
| | `kv_event_publisher_topic` | `topic` | `std::string` | `"mooncake"` | Message topic. |
| **Performance Config** | `kv_event_publisher_hwm` | `hwm` | `int` | `100000` | ZeroMQ high-water mark, controls memory buffer size. |
| | `kv_event_publisher_send_interval_ms` | `send_interval` | `uint32_t` | `0` | Send interval (milliseconds), 0 means no delay. |
| | `kv_event_publisher_max_batch_size` | `max_batch_size` | `uint32_t` | `50` | Maximum batch size, affects throughput and latency. |
| **Advanced Config** | `kv_event_publisher_auto_port` | `auto_port` | `bool` | `true` | Automatic port switching, automatically tries other ports if the port is occupied. |
#### 2.2.2 Event Publishing Interface
##### Generic Publishing Method
```c++
// Template method, supports all event types derived from KVCacheEvent
template <DerivedFromKVCacheEvent Event, typename... Args>
std::future<bool> publish(Args&&... args);
```
##### Publishing Different Types of Events
```c++
// 1. Publish block storage event
system.publish<BlockStoreEvent>(
"key_001",
replica_list,
store_info
);
// 2. Publish block update event
system.publish<BlockUpdateEvent>(
"key_001",
updated_replicas
);
// 3. Publish clear all event
system.publish<RemoveAllEvent>();
// 4. Custom event (must inherit from KVCacheEvent)
class CustomEvent : public KVCacheEvent { /* ... */ };
system.publish<CustomEvent>(arg1, arg2);
```
------
## 3. Python API Integration
Since only the `BlockStoreEvent` requires the introduction of field information defined in `StoreEventInfo` passed from the inference engine side, a new default parameter (i.e., the optional parameter `store_event_infos`) has been added only to the interfaces related to `put` operations.
> Currently, only the `batch_put_from_multi_buffers` interface has been adapted accordingly. Other interfaces related to `put` operations will also need adaptation subsequently.
### `StoreEventInfo` Type
**Purpose**
The `StoreEventInfo` struct is used to carry business metadata for KV cache storage events. All fields in this structure are optional; which fields to populate should be determined by the specific business scenario of the upper-layer cache-aware component (such as indexer or router). The Mooncake system acts only as a transparent pipeline for transmitting these fields and will not perform any business logic validation or interpretation of their content.
For detailed API specification, please refer to the Python API Reference documentation: [mooncake-store.md](https://github.com/OpenMooncake/Mooncake/blob/main/docs/source/python-api-reference/mooncake-store.md).
### `batch_put_from_multi_buffers` Interface
The `batch_put_from_multi_buffers` interface has been enhanced to support passing `StoreEventInfo` objects via the optional `store_event_infos` parameter. This allows KV cache storage events to carry additional metadata that can be consumed by downstream components.
For detailed API specification and usage examples, please refer to the Python API Reference documentation: [mooncake-store.md](https://github.com/OpenMooncake/Mooncake/blob/main/docs/source/python-api-reference/mooncake-store.md).

View File

@ -0,0 +1,17 @@
find_program(GO_EXECUTABLE go)
if(NOT GO_EXECUTABLE)
message(FATAL_ERROR "Go compiler not found. Please install Golang first.")
endif()
add_custom_target(mooncake_conductor ALL
COMMAND ./build.sh ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_BINARY_DIR}
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
COMMENT "Building Go program: mooncake_conductor"
VERBATIM
)
set(GO_EXECUTABLE_PATH "${CMAKE_CURRENT_BINARY_DIR}/mooncake_conductor"
CACHE INTERNAL "Path to Go executable")
# Install target
install(PROGRAMS ${GO_EXECUTABLE_PATH} DESTINATION bin)

View File

@ -0,0 +1,167 @@
# vLLM V1 Disaggregated Serving with MooncakeConductor
## Overview
This is the latest version of the Mooncake Conductor integration doc with the vLLM project to support KVCache-Aware scheduling algorithm.
The conductor can be integrated as a plugin into any proxy to uniformly manage KV events from L1 to L3. We also provide a toy_proxy for those who want to try it out ([proxy](./cacheaware_disaggregated_proxy.py)). Benchmark results will be released soon.
- only vLLM and vLLM-Ascend are supported.
## Installation
The mooncake conductor will be compiled and installed together with the mooncake store. Refer to [Build Guide](https://github.com/kvcache-ai/Mooncake/blob/main/docs/source/getting_started/build.md).
- **WITH_CONDUCTOR must be set to ON in Mooncake/CMakeLists.txt.**
### Install the latest version of vLLM and vLLM-Ascend
#### 1. Clone vLLM from official repo
```bash
git clone git@github.com:vllm-project/vllm.git
```
#### 2. Build
##### 2.1 Build from source
```bash
cd vllm
pip3 install -e .
```
- If you encounter any problems that you cannot solve, please refer to the [vLLM official compilation guide](https://docs.vllm.ai/en/latest/getting_started/installation/index.html).
#### 3. Clone vLLM-Ascend from official repo
```bash
git clone git@github.com:vllm-project/vllm-ascend.git
```
#### 4. Build
##### 4.1 Build from source
```bash
cd vllm-ascend
pip install -e .
```
- If you encounter any problems that you cannot solve, please refer to the [vLLM-Ascend official compilation guide](https://docs.vllm.ai/projects/ascend/en/latest/).
## Configuration
### Prepare configuration file to Run Example
- Prepare a _**conductor_config.json**_ file for mooncake_conductor. Here is an example:
```json
{
"kvevent_instance":
{
"vllm-1":
{
"ip": "127.0.0.1",
"port": 5557,
"type": "vLLM",
"modelname": "qwen2.5_7B",
"lora_id": -1
},
"mooncake":
{
"ip": "127.0.0.1",
"port": 19997,
"type": "Mooncake",
"modelname": "qwen2.5_7B",
"lora_id": -1
}
},
"http_server_port": 13333
}
```
- `kvevent_instance`: Services capable of reporting KV events.
- `vllm-1/mooncake`: rename of a VLLM instance or Mooncake-master instance.You can modify it according to your own preferences.
- `ip`: zmq publisher IP.
- `port`: zmq publisher port.
- `type`: Mark the type of kv-event publisher. Generally, there are currently only two types: `vLLM` and `Mooncake`.
- `modelname`: Model name used for match the model.
- `lora_id`: LoRA Adapter ID.
- `http_server_port`: Conductor http server for querying cache hit rates, default use `13333`.
## Run Example
### 1. Start the mooncake_master server
```sh
# start mooncake_master without kv-event publish
mooncake_master --rpc_port 50051
# start moocake_master with kv-event
mooncake_master -enable_kv_event_publish -kv_event_publisher_endpoint tcp://*:19997 -rpc_port 50051
```
### 2. Run multiple vllm instances
```sh
# kv_producer role
vllm serve /qwen2.5_7B_instruct/ \
--enforce-eager \
--max-model-len 10000 \
--port 8100 \
--gpu-memory-utilization 0.8 \
--served-model-name "qwen2.5_7B" \
--trust-remote-code \
--kv-events-config \
'{
"publisher": "zmq",
"enable_kv_cache_events": true,
"endpoint": "tcp://*:5557",
"topic": "kv-events",
"replay_endpoint": "tcp://*:5558"
}' \
--kv-transfer-config \
'{
"kv_connector": "MooncakeConnectorStoreV1",
"kv_role":"kv_producer",
"kv_connector_extra_config":{"use_layerwise": false}
}'
```
```sh
# kv_consumer role
vllm serve /qwen2.5_7B_instruct/ \
--enforce-eager \
--max-model-len 10000 \
--port 8200 \
--gpu-memory-utilization 0.8 \
--served-model-name "qwen2.5_7B" \
--trust-remote-code \
--kv-transfer-config \
'{
"kv_connector": "MooncakeConnectorStoreV1",
"kv_role":"kv_consumer",
"kv_connector_extra_config":{"use_layerwise": false}
}'
```
### 3. Start the conductor server
```sh
export CONDUCTOR_CONFIG_PATH="./example/conductor_config.json"
mooncake_conductor
```
### 4. Run the proxy in the example
```sh
python cacheaware_disaggregated_proxy.py --prefiller-hosts 127.0.0.1 --prefiller-ports 8100 --decoder-host 127.0.0.1 --decoder-ports 8200 --conductor-address 127.0.0.1:13333
```
## Test with openai compatible request
```sh
curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
"model": "qwen2.5_7B",
"prompt": "What are the key architectural differences between vLLM and Mooncake when it comes to handling key-value (KV) cache events, and how can a centralized conductor component be designed in Go to normalize disparate event schemas from these systems, apply consistent metrics collection, and make dynamic scheduling decisions based on real-time KV cache hit rates without relying on Kubernetes-based autoscaling mechanisms?",
"max_tokens": 1000
}'
```

33
mooncake-conductor/build.sh Executable file
View File

@ -0,0 +1,33 @@
#!/bin/bash
if [ "$#" -ne 2 ]; then
echo "Usage: $0 TARGET_PATH BUILD_DIR"
exit 1
fi
TARGET=$1
BUILD_DIR=$2
cd conductor-ctrl
# Check if go.mod exists
if [ ! -f "go.mod" ]; then
echo "Error: go.mod file not found"
exit 1
fi
echo "Cleaning previous build..."
rm -f mooncake_conductor
go mod tidy
echo "Building Go program: mooncake_conductor"
go build -o "$TARGET/mooncake_conductor" main.go
if [ $? -eq 0 ] && [ -f "$TARGET/mooncake_conductor" ]; then
echo "mooncake_conductor built successfully"
else
echo "mooncake_conductor build failed"
exit 1
fi

View File

@ -0,0 +1,82 @@
package common
import (
"sync"
"sync/atomic"
)
type SyncMap[K any, V any] struct {
m sync.Map
len atomic.Int32
}
func (sm *SyncMap[K, V]) Delete(key K) {
sm.LoadAndDelete(key)
}
func (sm *SyncMap[K, V]) Load(key K) (typedVal V, ok bool) {
value, ok := sm.m.Load(key)
if ok {
typedVal = value.(V)
}
return
}
func (sm *SyncMap[K, V]) LoadAndDelete(key K) (typedVal V, loaded bool) {
value, loaded := sm.m.LoadAndDelete(key)
if loaded {
typedVal = value.(V)
sm.len.Add(-1)
}
return
}
func (sm *SyncMap[K, V]) LoadOrStore(key K, value V) (V, bool) {
actual, loaded := sm.m.LoadOrStore(key, value)
if !loaded {
sm.len.Add(1)
}
return actual.(V), loaded
}
func (sm *SyncMap[K, V]) Range(f func(key K, value V) bool) {
sm.m.Range(func(key, value any) bool {
return f(key.(K), value.(V))
})
}
func (sm *SyncMap[K, V]) Keys() []K {
k := make([]K, 0, sm.Len())
sm.m.Range(func(key, value any) bool {
k = append(k, key.(K))
return true
})
return k
}
func (sm *SyncMap[K, V]) Values() []V {
v := make([]V, 0, sm.Len())
sm.m.Range(func(key, value any) bool {
v = append(v, value.(V))
return true
})
return v
}
func (sm *SyncMap[K, V]) Store(key K, value V) {
sm.Swap(key, value)
}
func (sm *SyncMap[K, V]) Swap(key K, value V) (V, bool) {
old, loaded := sm.m.Swap(key, value)
if !loaded {
var ret V
sm.len.Add(1)
return ret, loaded
}
return old.(V), loaded
}
func (sm *SyncMap[K, V]) Len() int {
return int(sm.len.Load())
}

View File

@ -0,0 +1,145 @@
package common
import (
"fmt"
"sync"
"testing"
)
func TestSyncMap(t *testing.T) {
sm := &SyncMap[string, int]{}
// test Store and Load
sm.Store("key1", 1)
val, ok := sm.Load("key1")
if !ok || val != 1 {
t.Errorf("Expected Load('key1') to return (1, true), got (%d, %v)", val, ok)
}
// test Len
if sm.Len() != 1 {
t.Errorf("Expected Len() to return 1, got %d", sm.Len())
}
// test Load non-existent key
val, ok = sm.Load("key2")
if ok {
t.Errorf("Expected Load('key2') to return (0, false), got (%d, %v)", val, ok)
}
// test LoadOrStore with non-existent key
val, loaded := sm.LoadOrStore("key2", 2)
if loaded || val != 2 {
t.Errorf("Expected LoadOrStore('key2', 2) to return (2, false), got (%d, %v)", val, loaded)
}
if sm.Len() != 2 {
t.Errorf("Expected Len() to return 2, got %d", sm.Len())
}
// test LoadOrStore with existing key
val, loaded = sm.LoadOrStore("key1", 10)
if !loaded || val != 1 {
t.Errorf("Expected LoadOrStore('key1', 10) to return (1, true), got (%d, %v)", val, loaded)
}
// test Swap with existing key
val, loaded = sm.Swap("key1", 3)
if !loaded || val != 1 {
t.Errorf("Expected Swap('key1', 3) to return (1, true), got (%d, %v)", val, loaded)
}
val, ok = sm.Load("key1")
if !ok || val != 3 {
t.Errorf("Expected Load('key1') after Swap to return (3, true), got (%d, %v)", val, ok)
}
// test Swap with non-existent key
val, loaded = sm.Swap("key3", 4)
if loaded || val != 0 {
t.Errorf("Expected Swap('key3', 4) to return (0, false), got (%d, %v)", val, loaded)
}
if sm.Len() != 3 {
t.Errorf("Expected Len() to return 3, got %d", sm.Len())
}
// test Keys
keys := sm.Keys()
if len(keys) != 3 {
t.Errorf("Expected Keys() to return 3 keys, got %d", len(keys))
}
// test Keys is exists
keyMap := make(map[string]bool)
for _, k := range keys {
keyMap[k] = true
}
expectedKeys := []string{"key1", "key2", "key3"}
for _, k := range expectedKeys {
if !keyMap[k] {
t.Errorf("Expected key '%s' in Keys() result", k)
}
}
values := sm.Values()
if len(values) != 3 {
t.Errorf("Expected Values() to return 3 values, got %d", len(values))
}
// test Range
var rangeCount int
sm.Range(func(key string, value int) bool {
rangeCount++
return true
})
if rangeCount != 3 {
t.Errorf("Expected Range to iterate over 3 items, got %d", rangeCount)
}
// test Range with early termination
var earlyTerminateCount int
sm.Range(func(key string, value int) bool {
earlyTerminateCount++
return earlyTerminateCount < 2 // only iterate over the first two elements
})
if earlyTerminateCount != 2 {
t.Errorf("Expected Range with early termination to iterate over 2 items, got %d", earlyTerminateCount)
}
// test LoadAndDelete
val, loaded = sm.LoadAndDelete("key2")
if !loaded || val != 2 {
t.Errorf("Expected LoadAndDelete('key2') to return (2, true), got (%d, %v)", val, loaded)
}
if sm.Len() != 2 {
t.Errorf("Expected Len() after LoadAndDelete to return 2, got %d", sm.Len())
}
// test Delete
sm.Delete("key1")
if sm.Len() != 1 {
t.Errorf("Expected Len() after Delete to return 1, got %d", sm.Len())
}
val, ok = sm.Load("key1")
if ok {
t.Errorf("Expected Load('key1') after Delete to return (0, false), got (%d, %v)", val, ok)
}
// test concurrent safety
var wg sync.WaitGroup
concurrency := 100
wg.Add(concurrency)
for i := 0; i < concurrency; i++ {
go func(i int) {
defer wg.Done()
key := fmt.Sprintf("concurrent_key_%d", i)
sm.Store(key, i)
val, ok := sm.Load(key)
if !ok || val != i {
t.Errorf("Concurrent test failed: Expected Load('%s') to return (%d, true), got (%d, %v)", key, i, val, ok)
}
}(i)
}
wg.Wait()
if sm.Len() != concurrency+1 { // +1 for key3 still present
t.Errorf("Expected Len() after concurrent operations to return %d, got %d", concurrency+1, sm.Len())
}
}

View File

@ -0,0 +1,39 @@
package common
const (
ServiceTypeVLLM string = "vLLM"
ServiceTypeMooncake string = "Mooncake"
)
type ServiceConfig struct {
Endpoint string // kv publisher endpoint
ReplayEndpoint string // (optional)
Type string // kv publisher type, support: vLLM,Mooncake
ModelName string // Model name hosted by the service
LoraName string
TenantID string // (optional), default use 'default'
InstanceID string // required
BlockSize int64
DPRank int
AdditionalSalt string // (optional), default use empty string
}
type StoredEvent struct {
BlockHashes []uint64
BlockSize int64
ModelName string
LoraName string
InstanceID string
ParentBlockHash uint64
TokenIds []int32
Medium string
}
type RemovedEvent struct {
BlockHashes []uint64
ModelName string
LoraName string
InstanceID string
BlockSize int64
Medium string
}

View File

@ -0,0 +1,152 @@
package common
import (
"fmt"
"log/slog"
"os"
"strconv"
"strings"
)
func ParseLogLevel() slog.Level {
levelStr := os.Getenv("CONDUCTOR_LOG_LEVEL")
if levelStr == "" {
return slog.LevelInfo
}
switch strings.ToUpper(levelStr) {
case "DEBUG":
return slog.LevelDebug
case "INFO":
return slog.LevelInfo
case "WARN":
return slog.LevelWarn
case "ERROR":
return slog.LevelError
default:
// We use the default logger here to warn about the invalid config
slog.Warn("Invalid log level specified, defaulting to INFO", "level", levelStr)
return slog.LevelInfo
}
}
func LoadEnv(envName, defaultEnv string) string {
value := os.Getenv(envName)
if value == "" {
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
return value
}
func LoadIntEnv(envName string, defaultEnv int) int {
value := os.Getenv(envName)
trimmedValue := strings.TrimSpace(value)
if value != "" {
intValue, err := strconv.Atoi(value)
if err != nil {
slog.Error("invalid value for environment variable", "envName", envName, "value", trimmedValue)
} else {
return intValue
}
}
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
func LoadBoolEnv(envName string, defaultEnv bool) bool {
value := os.Getenv(envName)
trimmedValue := strings.TrimSpace(value)
if value != "" {
boolValue, err := strconv.ParseBool(value)
if err != nil {
slog.Error("invalid value for environment variable", "envName", envName, "value", trimmedValue)
} else {
return boolValue
}
}
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
func LoadFloatEnv(envName string, defaultEnv float64) float64 {
value := os.Getenv(envName)
trimmedValue := strings.TrimSpace(value)
if value != "" {
floatValue, err := strconv.ParseFloat(value, 64)
if err != nil {
slog.Error("invalid value for environment variable", "envName", envName, "value", trimmedValue)
} else {
return floatValue
}
}
slog.Warn("environment variable is not set, using default value", "envName", envName, "defaultValue", defaultEnv)
return defaultEnv
}
func ExtractTokenIdFromRequest(data map[string]interface{}, key string) ([]int32, error) {
raw, exists := data[key]
if !exists {
return nil, fmt.Errorf("missing key: %s", key)
}
arr, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("the value of %s is not an array", key)
}
result := make([]int32, len(arr))
for i, v := range arr {
switch val := v.(type) {
case float64:
result[i] = int32(val)
case int:
result[i] = int32(val)
default:
return nil, fmt.Errorf("unsupported value type of token_id at [%d], the type is %s", i, val)
}
}
return result, nil
}
func ExtractCandidateEngineFromRequest(data map[string]interface{}, key string) (map[string]struct{}, error) {
raw, exists := data[key]
if !exists {
return nil, fmt.Errorf("missing key: %s", key)
}
arr, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("the value of %s is not an array", key)
}
result := make(map[string]struct{}, len(arr))
for i, v := range arr {
str, ok := v.(string)
if !ok {
return nil, fmt.Errorf(`"instances[%d]" is not a string`, i)
}
result[str] = struct{}{}
}
return result, nil
}
func ExtractStringValueFromRequest(data map[string]interface{}, key string) (string, error) {
raw, exists := data[key]
if !exists {
return "", fmt.Errorf("missing key: %s", key)
}
str, ok := raw.(string)
if !ok {
return "", fmt.Errorf(`"the value of: %s" is not a string`, key)
}
return str, nil
}
func ExtractIntFromRequest(data map[string]interface{}, key string) (int64, error) {
raw, exists := data[key]
if !exists {
return -1, fmt.Errorf("missing key: %s", key)
}
result, ok := raw.(float64)
if !ok {
return -1, fmt.Errorf(`"the value of: %s" is not a number`, key)
}
return int64(result), nil
}

View File

@ -0,0 +1,265 @@
package common
import (
"log/slog"
"os"
"testing"
)
func TestParseLogLevel(t *testing.T) {
// test default log level
origLevel := os.Getenv("CONDUCTOR_LOG_LEVEL")
os.Unsetenv("CONDUCTOR_LOG_LEVEL")
defer os.Setenv("CONDUCTOR_LOG_LEVEL", origLevel)
level := ParseLogLevel()
if level != 0 { // slog.LevelInfo = 0
t.Errorf("Expected default log level to be Info, got %d", level)
}
// test various log levels
testCases := []struct {
name string
levelStr string
expected slog.Level
}{
{"Debug", "DEBUG", slog.LevelDebug},
{"Info", "INFO", slog.LevelInfo},
{"Warn", "WARN", slog.LevelWarn},
{"Error", "ERROR", slog.LevelError},
{"Lowercase", "debug", slog.LevelDebug},
{"MixedCase", "Debug", slog.LevelDebug},
{"Invalid", "INVALID", slog.LevelInfo},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
os.Setenv("CONDUCTOR_LOG_LEVEL", tc.levelStr)
level := ParseLogLevel()
if level != tc.expected {
t.Errorf("Expected log level %d for input %s, got %d", tc.expected, tc.levelStr, level)
}
})
}
}
func TestLoadEnv(t *testing.T) {
// test environment variable exists
origValue := os.Getenv("TEST_ENV_VAR")
os.Setenv("TEST_ENV_VAR", "test_value")
defer func() {
if origValue == "" {
os.Unsetenv("TEST_ENV_VAR")
} else {
os.Setenv("TEST_ENV_VAR", origValue)
}
}()
value := LoadEnv("TEST_ENV_VAR", "default_value")
if value != "test_value" {
t.Errorf("Expected LoadEnv to return 'test_value', got '%s'", value)
}
// test environment variable does not exist
origValue2 := os.Getenv("NON_EXISTENT_ENV_VAR")
os.Unsetenv("NON_EXISTENT_ENV_VAR")
defer func() {
if origValue2 != "" {
os.Setenv("NON_EXISTENT_ENV_VAR", origValue2)
}
}()
value = LoadEnv("NON_EXISTENT_ENV_VAR", "default_value")
if value != "default_value" {
t.Errorf("Expected LoadEnv to return 'default_value', got '%s'", value)
}
}
func TestLoadIntEnv(t *testing.T) {
// test environment variable exists and value is valid
origValue := os.Getenv("TEST_INT_ENV_VAR")
os.Setenv("TEST_INT_ENV_VAR", "42")
defer func() {
if origValue == "" {
os.Unsetenv("TEST_INT_ENV_VAR")
} else {
os.Setenv("TEST_INT_ENV_VAR", origValue)
}
}()
value := LoadIntEnv("TEST_INT_ENV_VAR", 100)
if value != 42 {
t.Errorf("Expected LoadIntEnv to return 42, got %d", value)
}
// test environment variable exists but value is invalid
os.Setenv("TEST_INT_ENV_VAR", "invalid")
value = LoadIntEnv("TEST_INT_ENV_VAR", 100)
if value != 100 {
t.Errorf("Expected LoadIntEnv to return 100 for invalid value, got %d", value)
}
// test environment variable does not exist
origValue2 := os.Getenv("NON_EXISTENT_INT_ENV_VAR")
os.Unsetenv("NON_EXISTENT_INT_ENV_VAR")
defer func() {
if origValue2 != "" {
os.Setenv("NON_EXISTENT_INT_ENV_VAR", origValue2)
}
}()
value = LoadIntEnv("NON_EXISTENT_INT_ENV_VAR", 100)
if value != 100 {
t.Errorf("Expected LoadIntEnv to return 100 for non-existent env, got %d", value)
}
}
func TestExtractTokenIdFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"token_ids": []interface{}{1.0, 2.0, 3.0},
}
result, err := ExtractTokenIdFromRequest(data, "token_ids")
if err != nil {
t.Errorf("Expected ExtractTokenIdFromRequest to succeed, got error: %v", err)
}
expected := []int32{1, 2, 3}
if len(result) != len(expected) {
t.Errorf("Expected length %d, got %d", len(expected), len(result))
}
for i, v := range result {
if v != expected[i] {
t.Errorf("Expected result[%d] = %d, got %d", i, expected[i], v)
}
}
// test mixed number types
data["token_ids"] = []interface{}{1, 2.0, 3}
result, err = ExtractTokenIdFromRequest(data, "token_ids")
if err != nil {
t.Errorf("Expected ExtractTokenIdFromRequest to succeed with mixed number types, got error: %v", err)
}
expected = []int32{1, 2, 3}
if len(result) != len(expected) {
t.Errorf("Expected length %d, got %d", len(expected), len(result))
}
for i, v := range result {
if v != expected[i] {
t.Errorf("Expected result[%d] = %d, got %d", i, expected[i], v)
}
}
// test missing key
result, err = ExtractTokenIdFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractTokenIdFromRequest to fail with missing key, got success")
}
// test non-array value
data["token_ids"] = "not an array"
result, err = ExtractTokenIdFromRequest(data, "token_ids")
if err == nil {
t.Errorf("Expected ExtractTokenIdFromRequest to fail with non-array value, got success")
}
// test unsupported types
data["token_ids"] = []interface{}{"string", true}
result, err = ExtractTokenIdFromRequest(data, "token_ids")
if err == nil {
t.Errorf("Expected ExtractTokenIdFromRequest to fail with unsupported types, got success")
}
}
func TestExtractCandidateEngineFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"instances": []interface{}{"engine1", "engine2", "engine3"},
}
result, err := ExtractCandidateEngineFromRequest(data, "instances")
if err != nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to succeed, got error: %v", err)
}
if len(result) != 3 {
t.Errorf("Expected 3 engines, got %d", len(result))
}
expectedEngines := []string{"engine1", "engine2", "engine3"}
for _, engine := range expectedEngines {
if _, ok := result[engine]; !ok {
t.Errorf("Expected engine '%s' not found in result", engine)
}
}
// test missing key
result, err = ExtractCandidateEngineFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to fail with missing key, got success")
}
// test non-array value
data["instances"] = "not an array"
result, err = ExtractCandidateEngineFromRequest(data, "instances")
if err == nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to fail with non-array value, got success")
}
// test non-string element
data["instances"] = []interface{}{"engine1", 2, "engine3"}
result, err = ExtractCandidateEngineFromRequest(data, "instances")
if err == nil {
t.Errorf("Expected ExtractCandidateEngineFromRequest to fail with non-string element, got success")
}
}
func TestExtractStringValueFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"str_key": "string_value",
}
result, err := ExtractStringValueFromRequest(data, "str_key")
if err != nil {
t.Errorf("Expected ExtractStringValueFromRequest to succeed, got error: %v", err)
}
if result != "string_value" {
t.Errorf("Expected 'string_value', got '%s'", result)
}
// test missing key
result, err = ExtractStringValueFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractStringValueFromRequest to fail with missing key, got success")
}
// test non-string value
data["str_key"] = 123
result, err = ExtractStringValueFromRequest(data, "str_key")
if err == nil {
t.Errorf("Expected ExtractStringValueFromRequest to fail with non-string value, got success")
}
}
func TestExtractIntFromRequest(t *testing.T) {
// test successful extraction
data := map[string]interface{}{
"int_key": 42.0,
}
result, err := ExtractIntFromRequest(data, "int_key")
if err != nil {
t.Errorf("Expected ExtractIntFromRequest to succeed, got error: %v", err)
}
if result != 42 {
t.Errorf("Expected 42, got %d", result)
}
// test missing key
result, err = ExtractIntFromRequest(data, "missing_key")
if err == nil {
t.Errorf("Expected ExtractIntFromRequest to fail with missing key, got success")
}
// test non-number value
data["int_key"] = "not a number"
result, err = ExtractIntFromRequest(data, "int_key")
if err == nil {
t.Errorf("Expected ExtractIntFromRequest to fail with non-number value, got success")
}
}

View File

@ -0,0 +1,9 @@
module conductor
go 1.23.8
require (
github.com/cespare/xxhash/v2 v2.3.0
github.com/pebbe/zmq4 v1.4.0
github.com/shamaton/msgpack/v2 v2.4.0
)

View File

@ -0,0 +1,112 @@
package kvevent
import (
"context"
"fmt"
"time"
"conductor/common"
"conductor/zmq"
"log/slog"
)
// KVEventHandler adapts the generic EventHandler interface for EventManager.
// It is instantiated in event_manager.go but implemented here to keep files clean.
type KVEventHandler struct {
manager *EventManager
tenant_id string
// svcName string
modelName string
loraName string
instanceID string
blockSize int64
additionalSalt string
}
func (h *KVEventHandler) HandleEvent(event zmq.KVEvent, dpRank int64) error {
h.manager.mu.RLock()
if h.manager.stopped {
h.manager.mu.RUnlock()
return fmt.Errorf("manager stopped")
}
h.manager.mu.RUnlock()
slog.Info("Handling KV event", "instance_id", h.instanceID, "dpRank", dpRank)
// Create context for processing
ctx, cancel := context.WithTimeout(h.manager.ctx, 10*time.Second)
defer cancel()
// Dispatch event
switch e := event.(type) {
case *zmq.BlockStoredEvent:
slog.Debug("BlockStored",
"instance_id", h.instanceID,
"dpRank", dpRank,
"blocks", len(e.BlockHashes),
)
slog.Info("Received BlockStoredEvent", "medium", e.Medium)
return h.handleBlockStored(ctx, e, dpRank)
case *zmq.BlockRemovedEvent:
slog.Debug("BlockRemoved",
"instance_id", h.instanceID,
"dpRank", dpRank,
"blocks", len(e.BlockHashes),
)
slog.Info("Received BlockRemovedEvent", "medium", e.Medium)
return h.handleBlockRemoved(ctx, e, dpRank)
default:
slog.Warn("Unknown event type",
"type", fmt.Sprintf("%T", event),
)
return nil
}
}
func (h *KVEventHandler) handleBlockStored(ctx context.Context, event *zmq.BlockStoredEvent, dpRank int64) error {
// Convert to kvindexer event
conductorEvent := common.StoredEvent{
BlockHashes: event.BlockHashes,
BlockSize: event.BlockSize,
ModelName: h.modelName,
LoraName: h.loraName,
InstanceID: h.instanceID,
ParentBlockHash: event.ParentBlockHash,
TokenIds: event.TokenIDs,
Medium: event.Medium,
}
indexer := h.manager.getIndexer()
er := indexer.ProcessStoreEvent(conductorEvent, dpRank)
// TODO support mooncake_key map
if er != nil {
slog.Error("process store event failed.", "error", er)
}
slog.Debug("in handleBlockStored", "conductorEvent", conductorEvent)
return nil
}
func (h *KVEventHandler) handleBlockRemoved(ctx context.Context, event *zmq.BlockRemovedEvent, dpRank int64) error {
// Convert to conductor event
conductorEvent := common.RemovedEvent{
BlockHashes: event.BlockHashes,
ModelName: h.modelName,
LoraName: h.loraName,
InstanceID: h.instanceID,
BlockSize: h.blockSize,
Medium: event.Medium,
}
indexer := h.manager.getIndexer()
er := indexer.ProcessRemoveEvent(conductorEvent, dpRank, h.instanceID)
if er != nil {
slog.Error("process remove event failed.")
}
slog.Debug("in handleBlockRemoved", "conductorEvent", conductorEvent)
return nil
}
// TODO support mooncake update kv event

View File

@ -0,0 +1,492 @@
package kvevent
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"net/http"
"sync"
"time"
"conductor/common"
"conductor/prefixindex"
"conductor/zmq"
)
// Dynamic register structure
type RegisterReq struct {
Endpoint string `json:"endpoint"`
ReplayEndpoint string `json:"replay_endpoint"`
Type string `json:"type"`
ModelName string `json:"modelname"`
LoraName *string `json:"lora_name"`
TenantID *string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
BlockSize int64 `json:"block_size"`
DPRank int `json:"dp_rank"`
AdditionalSalt *string `json:"additionalsalt"`
}
// Dynamic unregister structure
type UnregisterReq struct {
Type string `json:"type"`
ModelName string `json:"modelname"`
LoraName *string `json:"lora_name"`
TenantID *string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
BlockSize int `json:"block_size"`
DPRank int `json:"dp_rank"`
}
type QueryReq struct {
ModelName string `json:"model"`
LoraName *string `json:"lora_name"`
LoraID *int64 `json:"lora_id"`
TokenIDs []int32 `json:"token_ids"`
InstanceID *string `json:"instance_id"`
TenantID *string `json:"tenant_id"`
BlockSize int64 `json:"block_size"`
CacheSalt *string `json:"cache_salt"`
}
type EventManager struct {
indexer *prefixindex.PrefixCacheTable
services []common.ServiceConfig
httpserverport int
subscribers common.SyncMap[string, *zmq.ZMQClient]
// Map to store active configurations
activeConfigs common.SyncMap[string, common.ServiceConfig]
// Map to store tenant instance list
tenantInstanceMap map[string]map[string]struct{}
ctx context.Context
cancel context.CancelFunc
mu sync.RWMutex
tenantMutex sync.RWMutex
stopped bool
}
func NewEventManager(
services []common.ServiceConfig,
httpserverport int,
) *EventManager {
ctx, cancel := context.WithCancel(context.Background())
indexer := prefixindex.NewPrefixCacheTable()
// TODO 每个ModelContext创建一个独立的indexer
return &EventManager{
services: services,
indexer: indexer,
httpserverport: httpserverport,
ctx: ctx,
cancel: cancel,
tenantInstanceMap: make(map[string]map[string]struct{}),
}
}
func (m *EventManager) Start() error {
slog.Info("Starting KV Event Manager...")
// Subscribe to all services concurrently
var wg sync.WaitGroup
errCh := make(chan error, len(m.services))
for _, svc := range m.services {
wg.Add(1)
go func(service common.ServiceConfig) {
defer wg.Done()
if err := m.subscribeToService(service); err != nil {
slog.Error("Failed to initiate subscription",
"service_type", service.Type,
"instance_id", service.InstanceID,
"endpoint", service.Endpoint,
"error", err,
)
errCh <- fmt.Errorf("failed to subscribe to %s: %w", service.InstanceID, err)
}
}(svc)
}
wg.Wait()
close(errCh)
failureCount := len(errCh)
successCount := len(m.services) - failureCount
slog.Info("Static KV Event Manager started. Subscriptions",
"success", successCount,
"failed", failureCount,
)
return nil
}
func (m *EventManager) Stop() {
m.mu.Lock()
if m.stopped {
m.mu.Unlock()
return
}
m.stopped = true
m.mu.Unlock()
slog.Info("Stopping Conductor KV Event Manager.....")
// Cancel context
m.cancel()
// Stop all ZMQ clients
m.subscribers.Range(func(key string, client *zmq.ZMQClient) bool {
client.Stop()
slog.Info("Stopped all subscription",
"service_key", key,
)
return true
})
}
func makeServiceKey(instanceID string, tenantID string, dpRank int) string {
return fmt.Sprintf("%s|%s|%d", instanceID, tenantID, dpRank)
}
func (m *EventManager) subscribeToService(svc common.ServiceConfig) error {
// Use (instance_id, tenant_id) as composite key to support multi-tenant replicas
svcKey := makeServiceKey(svc.InstanceID, svc.TenantID, svc.DPRank)
if svc.InstanceID == "" {
svcKey = makeServiceKey(svc.Endpoint, svc.TenantID, svc.DPRank)
}
if _, exists := m.subscribers.Load(svcKey); exists {
return nil
}
// Validate endpoint
if svc.Endpoint == "" {
return fmt.Errorf("endpoint is required")
}
// Use ReplayEndpoint directly, fallback to empty if not provided
replayEndpoint := svc.ReplayEndpoint
handler := &KVEventHandler{
manager: m,
tenant_id: svc.TenantID,
modelName: svc.ModelName,
loraName: svc.LoraName,
instanceID: svc.InstanceID,
blockSize: svc.BlockSize,
additionalSalt: svc.AdditionalSalt,
}
// Configure ZMQ Client
zmqConfig := &zmq.ZMQClientConfig{
CachePoolKey: svcKey,
Endpoint: svc.Endpoint,
ReplayEndpoint: replayEndpoint,
ModelName: svc.ModelName,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 5 * time.Second,
ReconnectDelay: 1 * time.Second,
}
if err := zmq.ValidateConfig(zmqConfig); err != nil {
return fmt.Errorf("invalid ZMQ config: %w", err)
}
client := zmq.NewZMQClient(zmqConfig, handler)
if err := client.Start(); err != nil {
return fmt.Errorf("failed to start ZMQ client: %w", err)
}
m.subscribers.Store(svcKey, client)
m.activeConfigs.Store(svcKey, svc)
//Add instance to tenant's instance map
m.tenantMutex.Lock()
if _, exists := m.tenantInstanceMap[svc.TenantID]; !exists {
m.tenantInstanceMap[svc.TenantID] = make(map[string]struct{})
}
m.tenantInstanceMap[svc.TenantID][svc.InstanceID] = struct{}{}
m.tenantMutex.Unlock()
slog.Info("Successfully subscribed to service",
"service_type", svc.Type,
"service_key", svcKey,
"instance_id", svc.InstanceID,
"tenant_id", svc.TenantID,
"endpoint", svc.Endpoint,
"replay_endpoint", replayEndpoint,
)
return nil
}
func (m *EventManager) unsubscribeFromService(instanceID string, tenantID string, dpRank int) {
svcKey := makeServiceKey(instanceID, tenantID, dpRank)
if client, exists := m.subscribers.Load(svcKey); exists {
client.Stop()
m.subscribers.Delete(svcKey)
m.activeConfigs.Delete(svcKey)
// Remove engine_instance from tenant's instance set
m.tenantMutex.Lock()
if instanceSet, exists := m.tenantInstanceMap[tenantID]; exists {
delete(instanceSet, instanceID)
}
m.tenantMutex.Unlock()
slog.Info("Successfully unsubscribed from service",
"service_key", svcKey,
"instance_id", instanceID,
"tenant_id", tenantID,
)
}
}
func (m *EventManager) getIndexer() *prefixindex.PrefixCacheTable {
return m.indexer
}
func (m *EventManager) StartHTTPServer() error {
mux := http.NewServeMux()
mux.HandleFunc("/query", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req QueryReq
slog.Debug(
"receive req",
"method", r.Method,
"path", r.URL.Path,
"remote", r.RemoteAddr,
)
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
slog.Error("Failed to decode JSON", "err", err)
http.Error(w, "Invalid JSON", http.StatusBadRequest)
return
}
tenantID := "default"
if req.TenantID != nil && *req.TenantID != "" {
tenantID = *req.TenantID
}
loraName := ""
if req.LoraName != nil {
slog.Debug("LoraName is provided", "lora_name", *req.LoraName)
loraName = *req.LoraName
}
cacheSalt := ""
if req.CacheSalt != nil {
slog.Debug("cacheSalt is provided", "cacheSalt", *req.CacheSalt)
cacheSalt = *req.CacheSalt
}
response_result := make(map[string]map[string]prefixindex.CacheHitResult)
if req.InstanceID != nil {
slog.Info("search all engine instance for tenant. ", "instance_id", req.InstanceID)
modelContext := &prefixindex.ModelContext{
TenantID: tenantID,
ModelName: req.ModelName,
LoraName: loraName,
BlockSize: req.BlockSize,
AdditionalSalt: cacheSalt,
InstanceID: *req.InstanceID,
}
result := m.indexer.CacheHitCompute(modelContext, req.TokenIDs)
if result != nil {
if response_result[tenantID] == nil {
response_result[tenantID] = make(map[string]prefixindex.CacheHitResult)
}
response_result[tenantID][*req.InstanceID] = *result
}
} else {
if instanceSet, exists := m.tenantInstanceMap[tenantID]; exists {
for instanceID := range instanceSet {
modelContext := &prefixindex.ModelContext{
TenantID: tenantID,
ModelName: req.ModelName,
LoraName: loraName,
BlockSize: req.BlockSize,
AdditionalSalt: cacheSalt,
InstanceID: instanceID,
}
result := m.indexer.CacheHitCompute(modelContext, req.TokenIDs)
if result != nil {
if response_result[tenantID] == nil {
response_result[tenantID] = make(map[string]prefixindex.CacheHitResult)
}
response_result[tenantID][instanceID] = *result
}
}
} else {
slog.Warn("current tenant has no engine_instance. ", "tenant_id", tenantID)
}
}
slog.Debug("cache hit status", "hitresult", response_result)
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(response_result); err != nil {
slog.Error("Failed to encode response", "err", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
})
// Register interface
mux.HandleFunc("/register", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req RegisterReq
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
slog.Error("Failed to decode register JSON", "err", err)
http.Error(w, "Invalid JSON", http.StatusBadRequest)
return
}
// Handle Optional fields' default values
tenantID := "default"
if req.TenantID != nil && *req.TenantID != "" {
tenantID = *req.TenantID
}
loraName := ""
if req.LoraName != nil {
slog.Info("LoraName is provided", "lora_name", *req.LoraName)
loraName = *req.LoraName
}
additionalSalt := ""
if req.AdditionalSalt != nil {
additionalSalt = *req.AdditionalSalt
}
svc := common.ServiceConfig{
Endpoint: req.Endpoint,
ReplayEndpoint: req.ReplayEndpoint,
Type: req.Type,
ModelName: req.ModelName,
LoraName: loraName,
TenantID: tenantID,
InstanceID: req.InstanceID,
BlockSize: req.BlockSize,
DPRank: req.DPRank,
AdditionalSalt: additionalSalt,
}
if err := m.subscribeToService(svc); err != nil {
slog.Error("Dynamic register failed", "instance_id", req.InstanceID, "err", err)
http.Error(w, fmt.Sprintf("Failed to subscribe: %v", err), http.StatusInternalServerError)
return
}
m.services = append(m.services, svc)
modelContext := &prefixindex.ModelContext{
TenantID: tenantID,
ModelName: req.ModelName,
LoraName: loraName,
BlockSize: req.BlockSize,
AdditionalSalt: additionalSalt,
InstanceID: svc.InstanceID,
}
m.indexer.AddDpSize(modelContext, int64(svc.DPRank))
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]string{
"status": "registered successfully",
"instance_id": svc.InstanceID,
})
})
// Unregister interface
mux.HandleFunc("/unregister", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var req UnregisterReq
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
slog.Error("Failed to decode unregister JSON", "err", err)
http.Error(w, "Invalid JSON", http.StatusBadRequest)
return
}
// Build target service key from instance_id and tenant_id
targetTenant := "default"
if req.TenantID != nil && *req.TenantID != "" {
targetTenant = *req.TenantID
}
targetKey := makeServiceKey(req.InstanceID, targetTenant, req.DPRank)
// Direct lookup and removal
if _, exists := m.activeConfigs.Load(targetKey); !exists {
http.Error(w, fmt.Sprintf("service not found: %s", targetKey), http.StatusNotFound)
return
}
m.unsubscribeFromService(req.InstanceID, targetTenant, req.DPRank)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]interface{}{
"status": "unregistered successfully",
"removed_instances": []string{targetKey},
})
})
// Global view interface
mux.HandleFunc("/global_view", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
globalView := m.indexer.GetGlobalView()
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(globalView); err != nil {
slog.Error("Failed to encode global view response", "err", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
})
server := &http.Server{
Addr: fmt.Sprintf(":%d", m.httpserverport),
Handler: mux,
}
go func() {
slog.Info("HTTP server listening", "port", m.httpserverport)
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("HTTP server failed", "err", err)
}
}()
// Start a goroutine to listen for context cancellation, used for graceful shutdown.
go func() {
<-m.ctx.Done()
slog.Info("Shutting down HTTP server")
// 5-second timeout for forced shutdown
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := server.Shutdown(shutdownCtx); err != nil {
slog.Error("HTTP server shutdown error", "err", err)
server.Close()
}
}()
return nil
}

View File

@ -0,0 +1,131 @@
package main
import (
"encoding/json"
"errors"
"log/slog"
"os"
"os/signal"
"syscall"
"conductor/common"
"conductor/kvevent"
)
var (
// TODO change default config path
conductorConfigPath = common.LoadEnv("CONDUCTOR_CONFIG_PATH", "/root/conductor_config.json")
httpServerPort = 13333
)
type configStruct struct {
KVEventInstance map[string]serviceRaw `json:"kvevent_instance"`
HTTPPort int `json:"http_server_port"`
}
type serviceRaw struct {
Endpoint string `json:"endpoint"`
ReplayEndpoint string `json:"replay_endpoint"`
TypeStr string `json:"type"`
ModelName string `json:"modelname"`
LoraName string `json:"lora_name"`
TenantID string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
BlockSize int64 `json:"block_size"`
DPRank int `json:"dp_rank"`
AdditionalSalt string `json:"additionalsalt"`
}
func mapServiceType(s string) (string, bool) {
switch s {
case "vLLM":
return common.ServiceTypeVLLM, true
case "Mooncake":
return common.ServiceTypeMooncake, true
default:
return "None", false
}
}
func parseConfig() []common.ServiceConfig {
if _, err := os.Stat(conductorConfigPath); errors.Is(err, os.ErrNotExist) {
slog.Warn("Config file does not exist, exiting.", "path", conductorConfigPath)
// os.Exit(1)
return []common.ServiceConfig{}
} else if err != nil {
slog.Warn("Error accessing config file", "path", conductorConfigPath, "error", err)
// os.Exit(1)
return []common.ServiceConfig{}
}
data, err := os.ReadFile(conductorConfigPath)
if err != nil {
slog.Error("Failed to read config file", "path", conductorConfigPath, "error", err)
os.Exit(1)
}
var cfg configStruct
if err := json.Unmarshal(data, &cfg); err != nil {
slog.Error("Failed to parse JSON config", "error", err)
os.Exit(1)
}
httpServerPort = cfg.HTTPPort
services := make([]common.ServiceConfig, 0, len(cfg.KVEventInstance))
for name, raw := range cfg.KVEventInstance {
serviceType, ok := mapServiceType(raw.TypeStr)
if !ok {
slog.Error("Unknown service type", "type", raw.TypeStr)
continue
}
services = append(services, common.ServiceConfig{
Endpoint: raw.Endpoint,
ReplayEndpoint: raw.ReplayEndpoint,
Type: serviceType,
ModelName: raw.ModelName,
LoraName: raw.LoraName,
TenantID: raw.TenantID,
InstanceID: name,
BlockSize: raw.BlockSize,
DPRank: raw.DPRank,
AdditionalSalt: raw.AdditionalSalt,
})
}
return services
}
func main() {
// TODO support print metrics for conductor
logLevel := common.ParseLogLevel()
logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
Level: logLevel,
}))
slog.SetDefault(logger)
slog.Info("Starting Conductor KV Event Manager...", "logLevel", logLevel)
services := parseConfig()
manager := kvevent.NewEventManager(services, httpServerPort)
if err := manager.StartHTTPServer(); err != nil {
slog.Error("Failed to start HTTP server", "err", err)
}
if err := manager.Start(); err != nil {
slog.Error("Failed to start manager", "error", err)
os.Exit(1)
}
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
slog.Info("Manager is running. Press Ctrl+C to stop.")
<-sigChan
slog.Info("Shutting down...")
manager.Stop()
}

View File

@ -0,0 +1,561 @@
package prefixindex
import (
"encoding/binary"
"fmt"
"log/slog"
"math/rand"
"sync"
"sync/atomic"
"time"
"conductor/common"
"github.com/cespare/xxhash/v2"
)
var (
enableCpuEviction = common.LoadBoolEnv("ENABLE_CPU_EVICTION", false)
maxCpuKeyNum = int64(common.LoadIntEnv("MAX_CPU_KEY_NUM", 30000))
cpuKeyEvictionRatio = common.LoadFloatEnv("CPU_KEY_EVICTION_RATIO", 0.2)
// TODO
// The eviction parameter here is used to control that the number of CPU cache prefixes cannot grow without limit.
// The best way is to add an zmq publisher in mooncake-store to actively notify conductor of the block eviction.
)
type ModelContext struct {
ModelName string
LoraName string // None represents no LoRA adapter
BlockSize int64
// TODO @yejj710
// Confirm the difference between the previously discussed cache_salt and additionalSalt.
//The current understanding is that cache_salt is used to ensure data isolation between different customers,
// and it seems it can be directly added to additionalSalt.
AdditionalSalt string
TenantID string
InstanceID string // unique identifier for each API server
}
type CacheStoreInfo struct {
// TODO Currently, the KV cache at different levels is not distinguished.
// In the future, the caches of Mooncake and inference engines (vLLM, SGLang)
// should be handled separately.
engineLastAccessTime atomic.Int64
TotalReplicaNums atomic.Int64
mediumSet map[string]struct{}
dpRankSet map[int64]struct{} // indicate the dp_rank that the block is cached on
// LRU linked list pointers
lruPrev *CacheStoreInfo
lruNext *CacheStoreInfo
}
type HashMapStore struct {
// conductor prefixHash -> cachestore
prefixMap map[uint64]*CacheStoreInfo
createTime time.Time
lastAccess atomic.Int64
totalPrefixes int64
// LRU linked list: head is least recently used, tail is most recently used
lruHead *CacheStoreInfo
lruTail *CacheStoreInfo
}
type ContextData struct {
prefixMu sync.RWMutex
hashmapMu sync.RWMutex
prefixStore *HashMapStore
seed uint64
DpSize map[int64]struct{}
proxyHashMapping map[uint64]uint64 // engine block hash -> conductor prefix hash
}
type PrefixCacheTable struct {
// TODO use instance_id to distinguish different engine instances
contextMap sync.Map // ModelContext → *ContextData
contextCount atomic.Int32
}
type CacheHitResult struct {
LongestMatchTokens int64 `json:"longest_matched"`
DP map[int64]int64 `json:"DP"`
GPU int64 `json:"GPU"`
CPU int64 `json:"CPU"`
DISK int64 `json:"DISK"`
}
type ModelContextView struct {
ModelName string `json:"model_name"`
LoraName string `json:"lora_name"`
BlockSize int64 `json:"block_size"`
AdditionalSalt string `json:"additional_salt"`
TenantID string `json:"tenant_id"`
InstanceID string `json:"instance_id"`
}
type GlobalView struct {
ContextCount int32 `json:"context_count"`
ModelContexts []ModelContextView `json:"model_contexts"`
ProxyHashMap []map[uint64]uint64 `json:"hashmap"`
}
func GenerateSeedFromEnv() uint64 {
r := rand.New(rand.NewSource(time.Now().Unix()))
envSeed := common.LoadIntEnv("CONDUCTOR_SEED", -1)
var seed uint64
if envSeed != -1 {
seed = uint64(envSeed)
} else {
seed = r.Uint64()
}
return seed
}
func NewPrefixCacheTable() *PrefixCacheTable {
p := &PrefixCacheTable{}
return p
}
func (p *PrefixCacheTable) getContextData(modelcontext *ModelContext) *ContextData {
ctx_value := *modelcontext
value, exists := p.contextMap.Load(ctx_value)
if exists {
return value.(*ContextData)
}
seedValue := xxhash.Sum64String(modelcontext.AdditionalSalt)
newContextData := &ContextData{
prefixStore: &HashMapStore{
prefixMap: make(map[uint64]*CacheStoreInfo),
createTime: time.Now(),
totalPrefixes: 0,
},
proxyHashMapping: make(map[uint64]uint64),
seed: seedValue,
DpSize: make(map[int64]struct{}),
}
newContextData.prefixStore.lastAccess.Store(time.Now().Unix())
p.contextMap.Store(ctx_value, newContextData)
slog.Debug("in getContextData", "modelcontext", modelcontext)
p.contextCount.Add(1)
return newContextData
}
func (p *PrefixCacheTable) AddDpSize(modelcontext *ModelContext, dpRank int64) {
// value, exists := p.contextMap.Load(modelcontext)
contextData := p.getContextData(modelcontext)
contextData.DpSize[dpRank] = struct{}{}
}
func (p *PrefixCacheTable) ComputePrefixHash(modelcontext *ModelContext, tokenIds []int32, cacheSalt uint64) []uint64 {
// cacheSalt is used to seperate hash from different customers
numBlocks := len(tokenIds) / int(modelcontext.BlockSize)
prefixHashes := make([]uint64, 0, numBlocks)
var parentHash uint64 = cacheSalt
for i := 0; i < numBlocks; i++ {
start := i * int(modelcontext.BlockSize)
end := start + int(modelcontext.BlockSize)
if end > len(tokenIds) {
break
}
hashValue := p.computeHash(parentHash, tokenIds[start:end])
prefixHashes = append(prefixHashes, hashValue)
parentHash = hashValue
}
return prefixHashes
}
func (p *PrefixCacheTable) CacheHitCompute(modelcontext *ModelContext, tokenIds []int32) *CacheHitResult {
value, exists := p.contextMap.Load(*modelcontext)
prefixMatchResult := &CacheHitResult{
LongestMatchTokens: 0,
DP: map[int64]int64{},
GPU: 0,
CPU: 0,
DISK: 0,
}
if !exists {
slog.Error("In CacheHitCompute, contextData not found")
return prefixMatchResult
}
contextData := value.(*ContextData)
cacheSalt := xxhash.Sum64String(modelcontext.AdditionalSalt)
prefixHashes := p.ComputePrefixHash(modelcontext, tokenIds, cacheSalt)
// TODO @yejj710
// When there is no data in contextData, what information should be returned for the matched modelcontext
// This is related to function `AddDpSize`
slog.Debug("In CacheHitCompute", "prefixHashes", prefixHashes)
contextData.prefixMu.RLock()
defer contextData.prefixMu.RUnlock()
prefixStore := contextData.prefixStore
// reserve prefixHashes and then compute cache hit
for _, prefixHash := range prefixHashes {
cacheStoreInfo, exists := prefixStore.prefixMap[prefixHash]
slog.Debug("In CacheHitCompute", "cacheStoreInfo", cacheStoreInfo)
// chained hash, break if no replica exists
if !exists || cacheStoreInfo.TotalReplicaNums.Load() == 0 {
break
}
cacheHit := false
for key := range cacheStoreInfo.mediumSet {
if key == "cpu" {
prefixMatchResult.CPU += modelcontext.BlockSize
cacheHit = true
} else if key == "GPU" {
prefixMatchResult.GPU += modelcontext.BlockSize
cacheHit = true
} else {
slog.Warn("In CacheHitCompute, unknown medium type", "medium", key)
}
}
if cacheHit {
prefixMatchResult.LongestMatchTokens += modelcontext.BlockSize
for dpRank := range cacheStoreInfo.dpRankSet {
prefixMatchResult.DP[dpRank] += modelcontext.BlockSize
}
cacheStoreInfo.engineLastAccessTime.Store(time.Now().Unix())
// move to tail (most recently used)
if enableCpuEviction {
prefixStore.addToLRUTail(cacheStoreInfo)
// TODO update LRU list asynchronously
}
}
}
prefixStore.lastAccess.Store(time.Now().Unix())
return prefixMatchResult
}
func (p *PrefixCacheTable) ProcessStoreEvent(event common.StoredEvent, dpRank int64) error {
if len(event.BlockHashes) == 0 {
return nil
}
tenantID := "default"
slog.Debug("In ProcessStoreEvent", "modelName", event.ModelName, "instanceID", event.InstanceID, "dpRank", dpRank)
contextData := p.getContextData(&ModelContext{
ModelName: event.ModelName,
LoraName: event.LoraName,
BlockSize: event.BlockSize,
TenantID: tenantID,
AdditionalSalt: "",
InstanceID: event.InstanceID,
})
contextData.hashmapMu.Lock()
defer contextData.hashmapMu.Unlock()
proxyHashMap := contextData.proxyHashMapping
if len(event.BlockHashes)*int(event.BlockSize) != len(event.TokenIds) {
if len(event.BlockHashes) != 1 {
return fmt.Errorf("block hashes and tokens length mismatch")
}
}
newPrefixStore := make([]struct {
hashValue uint64
engineID string
}, 0)
var parentHash uint64 = contextData.seed
// TODO If the ParentBlockHash happens to be 0, a bug will occur here, because 0 is a valid hash value.
if event.ParentBlockHash != 0 {
slog.Debug("parent Block HASH is not None.")
if pbh, exists := proxyHashMap[event.ParentBlockHash]; exists {
parentHash = pbh
}
}
for i, blockHash := range event.BlockHashes {
// cache already exists, add engine info and continue
if existingHash, exists := proxyHashMap[blockHash]; exists {
newPrefixStore = append(newPrefixStore, struct {
hashValue uint64
engineID string
}{existingHash, event.InstanceID})
continue
}
// if not exists, compute hash
hashValue := p.computeHash(parentHash, event.TokenIds[i*int(event.BlockSize):(i+1)*int(event.BlockSize)])
parentHash = hashValue
proxyHashMap[blockHash] = hashValue
newPrefixStore = append(newPrefixStore, struct {
hashValue uint64
engineID string
}{
hashValue: hashValue,
engineID: event.InstanceID,
})
}
if len(newPrefixStore) > 0 {
contextData.prefixMu.Lock()
defer contextData.prefixMu.Unlock()
prefixStore := contextData.prefixStore
for _, newPrefix := range newPrefixStore {
slog.Debug("show new prefix data", "newPrefix", newPrefix)
p.addNewPrefixStore(prefixStore, newPrefix.hashValue, newPrefix.engineID, event.Medium, dpRank)
}
if enableCpuEviction && prefixStore.totalPrefixes > maxCpuKeyNum {
evictedCount := prefixStore.evictLRU(cpuKeyEvictionRatio)
slog.Info("LRU eviction triggered", "totalPrefixes", prefixStore.totalPrefixes, "evictedCount", evictedCount)
}
}
return nil
}
func (p *PrefixCacheTable) ProcessRemoveEvent(event common.RemovedEvent, dpRank int64, instanceID string) error {
if len(event.BlockHashes) == 0 {
return nil
}
contextData := p.getContextData(&ModelContext{
ModelName: event.ModelName,
LoraName: event.LoraName,
BlockSize: event.BlockSize,
TenantID: "default",
AdditionalSalt: "",
InstanceID: instanceID,
})
contextData.hashmapMu.Lock()
defer contextData.hashmapMu.Unlock()
proxyHashMap := contextData.proxyHashMapping
removeConductorHash := make([]uint64, 0, len(event.BlockHashes))
// delete proxyHashMapping
contextData.prefixMu.Lock()
defer contextData.prefixMu.Unlock()
prefixStore := contextData.prefixStore
for _, blockHash := range event.BlockHashes {
if conductorHash, exists := proxyHashMap[blockHash]; exists {
removeConductorHash = append(removeConductorHash, conductorHash)
// Only delete proxyHashMapping entry when all replicas are removed
cacheStoreInfo, exists := prefixStore.prefixMap[conductorHash]
if !exists {
continue
}
if cacheStoreInfo.TotalReplicaNums.Load() == 1 {
delete(proxyHashMap, blockHash)
}
}
}
// update prefixStore
for _, conductorHash := range removeConductorHash {
cacheStoreInfo, exists := prefixStore.prefixMap[conductorHash]
if !exists {
continue
}
cacheStoreInfo.TotalReplicaNums.Add(-1)
// Remove per-instance metadata
delete(cacheStoreInfo.dpRankSet, dpRank)
delete(cacheStoreInfo.mediumSet, event.Medium)
slog.Info("process remove event", "conductorHash", conductorHash, "dpRank", dpRank, "medium", event.Medium)
// Only delete entry when all replicas are removed
if cacheStoreInfo.TotalReplicaNums.Load() <= 0 {
delete(prefixStore.prefixMap, conductorHash)
prefixStore.totalPrefixes--
}
}
return nil
}
func (p *PrefixCacheTable) computeHash(parentHash uint64, blockTokenIDs []int32) uint64 {
// digest := xxhash.NewWithSeed()
digest := xxhash.New()
var parentHashBytes [8]byte
binary.LittleEndian.PutUint64(parentHashBytes[:], parentHash)
_, _ = digest.Write(parentHashBytes[:])
var tokenIDsBytes [8]byte
for _, tokenID := range blockTokenIDs {
binary.LittleEndian.PutUint32(tokenIDsBytes[:], uint32(tokenID))
_, _ = digest.Write(tokenIDsBytes[:])
}
return digest.Sum64()
}
func (p *PrefixCacheTable) addNewPrefixStore(prefixStore *HashMapStore, hashValue uint64, instanceID string, medium string, dpRank int64) {
now := time.Now().Unix()
if prefixStore.prefixMap[hashValue] == nil {
slog.Debug("in addNewPrefixStore, prefixStore.prefixMap[hashValue] is nil", "hashValue", hashValue)
prefixStore.prefixMap[hashValue] = &CacheStoreInfo{
mediumSet: make(map[string]struct{}),
dpRankSet: make(map[int64]struct{}),
}
prefixStore.totalPrefixes++
}
cacheStoreInfo := prefixStore.prefixMap[hashValue]
cacheStoreInfo.engineLastAccessTime.Store(now)
cacheStoreInfo.TotalReplicaNums.Add(1)
// TODO If using Mooncake-Store, you do not need to set dpRank, because it does not distinguish between kv-blocks of different dpRanks.
cacheStoreInfo.mediumSet[medium] = struct{}{}
cacheStoreInfo.dpRankSet[dpRank] = struct{}{}
slog.Debug("in addNewPrefixStore", "conductor_hash", hashValue, "current_mediumset", cacheStoreInfo.mediumSet[medium])
// move to tail (most recently used)
if enableCpuEviction {
prefixStore.addToLRUTail(cacheStoreInfo)
}
}
func (p *PrefixCacheTable) GetGlobalView() *GlobalView {
view := &GlobalView{
ContextCount: p.contextCount.Load(),
ModelContexts: make([]ModelContextView, 0),
ProxyHashMap: make([]map[uint64]uint64, 0),
}
p.contextMap.Range(func(key, value interface{}) bool {
ctx := key.(ModelContext)
contextData := value.(*ContextData)
ctxView := ModelContextView{
ModelName: ctx.ModelName,
LoraName: ctx.LoraName,
BlockSize: ctx.BlockSize,
AdditionalSalt: ctx.AdditionalSalt,
TenantID: ctx.TenantID,
InstanceID: ctx.InstanceID,
}
contextData.prefixMu.RLock()
defer contextData.prefixMu.RUnlock()
contextData.hashmapMu.RLock()
defer contextData.hashmapMu.RUnlock()
view.ProxyHashMap = append(view.ProxyHashMap, contextData.proxyHashMapping)
view.ModelContexts = append(view.ModelContexts, ctxView)
return true
})
return view
}
// move or add a CacheStoreInfo to the tail of the LRU list (most recently used)
func (h *HashMapStore) addToLRUTail(cacheNode *CacheStoreInfo) {
if cacheNode == nil {
return
}
// If already in list, remove it first
if cacheNode.lruPrev != nil || cacheNode.lruNext != nil || h.lruHead == cacheNode {
h.removeFromLRU(cacheNode)
}
if h.lruTail == nil {
// Empty list
h.lruHead = cacheNode
h.lruTail = cacheNode
cacheNode.lruPrev = nil
cacheNode.lruNext = nil
} else {
h.lruTail.lruNext = cacheNode
cacheNode.lruPrev = h.lruTail
cacheNode.lruNext = nil
h.lruTail = cacheNode
}
}
// remove a CacheStoreInfo from the LRU list
func (h *HashMapStore) removeFromLRU(cacheNode *CacheStoreInfo) {
if cacheNode == nil {
return
}
if cacheNode.lruPrev != nil {
cacheNode.lruPrev.lruNext = cacheNode.lruNext
} else {
// cacheNode is head
h.lruHead = cacheNode.lruNext
}
if cacheNode.lruNext != nil {
cacheNode.lruNext.lruPrev = cacheNode.lruPrev
} else {
// cacheNode is tail
h.lruTail = cacheNode.lruPrev
}
cacheNode.lruPrev = nil
cacheNode.lruNext = nil
}
// evict the least recently used cpu mediums based on eviction ratio
func (h *HashMapStore) evictLRU(ratio float64) int {
if h.lruHead == nil || ratio <= 0 {
return 0
}
// TODO Here it is assumed that each CacheStoreInfo contains a cpu medium, but in practice such assumptions should not be made.
// Instead, a separate variable should be used to record the number of cpu mediums.
targetCount := int(float64(h.totalPrefixes) * ratio)
evictedCount := 0
// Traverse from head (least recently used)
for h.lruHead != nil && evictedCount < targetCount {
cacheNode := h.lruHead
// Only evict `cpu` medium
if _, exist := cacheNode.mediumSet["cpu"]; exist {
delete(cacheNode.mediumSet, "cpu")
// If no mediums left, remove the entry entirely
if len(cacheNode.mediumSet) == 0 {
h.removeFromLRU(cacheNode)
h.deleteCacheStoreInfo(cacheNode)
// TODO remove proxyHashMapping in ContextData,
// currently we do not maintain reverse mapping from conductor hash to proxy hash,
// which makes it hard to delete the proxyHashMapping entry.
h.totalPrefixes--
}
} else {
// Move to LRU-List tail
h.removeFromLRU(cacheNode)
h.addToLRUTail(cacheNode)
}
evictedCount++
}
return evictedCount
}
func (h *HashMapStore) deleteCacheStoreInfo(cacheNode *CacheStoreInfo) {
for k, v := range h.prefixMap {
if v == cacheNode {
delete(h.prefixMap, k)
return
}
}
}

View File

@ -0,0 +1,106 @@
package zmq
import "time"
type EventType string
const (
EventTypeBlockStored EventType = "BlockStored"
EventTypeBlockRemoved EventType = "BlockRemoved"
// EventTypeBlockUpdate indicates that blocks have been updated from the KV cache
EventTypeBlockUpdate EventType = "BlockUpdate"
EventTypeAllCleared EventType = "AllBlocksCleared"
)
type KVEvent interface {
GetType() EventType
GetTimestamp() time.Time
}
type BlockStoredEvent struct {
Type EventType
Timestamp time.Time
BlockHashes []uint64
TokenIDs []int32
ParentBlockHash uint64
BlockSize int64
MooncakeKey string
ReplicaList [][]string
ModelName string
LoraID int64
LoraName string
PodName string
Medium string
}
func (e *BlockStoredEvent) GetType() EventType {
return e.Type
}
func (e *BlockStoredEvent) GetTimestamp() time.Time {
return e.Timestamp
}
type BlockRemovedEvent struct {
Type EventType
Timestamp time.Time
BlockHashes []uint64
ModelName string
PodName string
Medium string
}
func (e *BlockRemovedEvent) GetType() EventType {
return e.Type
}
func (e *BlockRemovedEvent) GetTimestamp() time.Time {
return e.Timestamp
}
type AllBlocksClearedEvent struct {
Type EventType
Timestamp time.Time
ModelName string
PodName string
}
func (e *AllBlocksClearedEvent) GetType() EventType {
return e.Type
}
func (e *AllBlocksClearedEvent) GetTimestamp() time.Time {
return e.Timestamp
}
type BlockUpdateEvent struct {
Type EventType
Timestamp time.Time
BlockHashes []uint64
TokenIDs []int32
ParentBlockHash uint64
ModelName string
PodName string
BlockSize int64
}
// GetType returns the event type
func (e *BlockUpdateEvent) GetType() EventType {
return e.Type
}
func (e *BlockUpdateEvent) GetTimestamp() time.Time {
return e.Timestamp
}
const (
SourceMooncake string = "mooncake"
SourceVLLM string = "vllm"
)
type EventBatch struct {
Source string // indicates the origin of the event batch
Events []KVEvent
DataParallelRank int64
}

View File

@ -0,0 +1,676 @@
package zmq
import (
"fmt"
"log/slog"
"strconv"
"strings"
"time"
msgpack "github.com/shamaton/msgpack/v2"
)
type EventParser interface {
ParseEvent(raw []interface{}, timestamp interface{}) (KVEvent, error)
EventMappings() map[string]EventType
Source() string
}
type mooncakeParser struct{}
func (p *mooncakeParser) Source() string { return SourceMooncake }
func (p *mooncakeParser) EventMappings() map[string]EventType {
return map[string]EventType{
"BlockStoreEvent": EventTypeBlockStored,
"BlockUpdateEvent": EventTypeBlockUpdate,
"RemoveAllEvent": EventTypeAllCleared,
}
}
func (p *mooncakeParser) ParseEvent(raw []interface{}, timestamp interface{}) (KVEvent, error) {
eventTypeStr, ok := raw[0].(string)
if !ok {
return nil, fmt.Errorf("invalid event type format: %T", raw[0])
}
eventType, exists := p.EventMappings()[eventTypeStr]
if !exists {
return nil, fmt.Errorf("unknown mooncake event type: %s", eventTypeStr)
}
switch eventType {
case EventTypeBlockStored:
return parseMooncakeBlockStored(raw, timestamp)
default:
return nil, fmt.Errorf("unhandled event: %s", eventType)
}
}
func decodeCommonEventBatch(
data []byte,
expectedLength int,
extractEvents func([]interface{}) ([]interface{}, interface{}, error),
parser EventParser,
) (*EventBatch, error) {
if len(data) > 0 {
slog.Debug("First byte of payload", "hex", fmt.Sprintf("%02x", data[0]))
}
var arr []interface{}
if err := msgpack.Unmarshal(data, &arr); err != nil {
return nil, fmt.Errorf("failed to unmarshal event batch: %w", err)
}
if len(arr) != expectedLength {
return nil, fmt.Errorf("expected %d-element array, got %d", expectedLength, len(arr))
}
events, timestamp, err := extractEvents(arr)
if err != nil {
return nil, err
}
if len(events) == 0 {
slog.Warn("Received empty event list")
}
var dpRank int64 = -1
if expectedLength == 3 {
dpRank, err = parseInt64(arr[2])
if err != nil {
return nil, fmt.Errorf("failed to parse dpRank: %w", err)
}
}
batch := &EventBatch{
Source: parser.Source(),
Events: make([]KVEvent, 0, len(events)),
DataParallelRank: dpRank,
}
slog.Info("Receive batched kv-event", "source", batch.Source, "dpRank", dpRank)
for i, rawEvent := range events {
eventSlice, ok := rawEvent.([]interface{})
if !ok {
return nil, fmt.Errorf("event at index %d is not a slice: %T", i, rawEvent)
}
event, err := parser.ParseEvent(eventSlice, timestamp)
if err != nil {
return nil, fmt.Errorf("failed to parse event at index %d: %w", i, err)
}
batch.Events = append(batch.Events, event)
}
return batch, nil
}
func newMooncakeParser() EventParser {
return &mooncakeParser{}
}
func DecodeMooncakeEventBatch(data []byte) (*EventBatch, error) {
return decodeCommonEventBatch(
data,
2,
func(arr []interface{}) ([]interface{}, interface{}, error) {
events, ok := arr[1].([]interface{})
if !ok {
return nil, nil, fmt.Errorf("invalid events type: %T", arr[1])
}
return events, arr[0], nil
},
newMooncakeParser(),
)
}
type vllmParser struct{}
func (p *vllmParser) Source() string { return SourceVLLM }
func (p *vllmParser) EventMappings() map[string]EventType {
return map[string]EventType{
"BlockStored": EventTypeBlockStored,
"BlockRemoved": EventTypeBlockRemoved,
"AllBlocksCleared": EventTypeAllCleared,
}
}
func (p *vllmParser) ParseEvent(raw []interface{}, timestamp interface{}) (KVEvent, error) {
eventTypeStr, ok := raw[0].(string)
if !ok {
return nil, fmt.Errorf("invalid event type format: %T", raw[0])
}
eventType, exists := p.EventMappings()[eventTypeStr]
if !exists {
return nil, fmt.Errorf("unknown vllm event type: %s", eventTypeStr)
}
switch eventType {
case EventTypeBlockStored:
return parseVllmBlockStored(raw, timestamp)
case EventTypeBlockRemoved:
return parseVllmBlockRemoved(raw, timestamp)
default:
return nil, fmt.Errorf("unhandled event: %s", eventType)
}
}
func newVLLMParser() EventParser {
return &vllmParser{}
}
func DecodeVllmEventBatch(data []byte) (*EventBatch, error) {
return decodeCommonEventBatch(
data,
3,
func(arr []interface{}) ([]interface{}, interface{}, error) {
events, ok := arr[1].([]interface{})
if !ok {
return nil, nil, fmt.Errorf("invalid events type: %T", arr[1])
}
return events, arr[0], nil
},
newVLLMParser(),
)
}
func parseMooncakeBlockStored(data []interface{}, timestamp interface{}) (*BlockStoredEvent, error) {
event := &BlockStoredEvent{
Type: EventTypeBlockStored,
}
if mooncakekey, err := safeGetString(data[1]); err == nil {
event.MooncakeKey = mooncakekey
} else {
return nil, fmt.Errorf("failed to parse MooncakeKey from field at index 1: %w", err)
}
if replicalist, err := convertToReplicaList(data[2]); err == nil {
event.ReplicaList = replicalist
slog.Debug("ReplicaList:", "ReplicaList", event.ReplicaList)
} else {
return nil, fmt.Errorf("failed to parse ReplicaList from field at index 2: %w", err)
}
if blocksize, err := parseInt64(data[4]); err == nil {
event.BlockSize = blocksize
slog.Debug("BlockSize:", "BlockSize", event.BlockSize)
} else {
return nil, fmt.Errorf("failed to parse BlockSize from field at index 4: %w", err)
}
if blockhash, err := parseMooncakeParentUint64(data[5]); err == nil {
event.BlockHashes = blockhash
slog.Debug("BlockHashes:", "BlockHashes", event.BlockHashes)
} else {
return nil, fmt.Errorf("failed to parse BlockHashes from field at index 5: %w", err)
}
if parentblockhash, err := parseMooncakeUint64(data[6]); err == nil {
event.ParentBlockHash = parentblockhash
slog.Debug("ParentBlockHash:", "ParentBlockHash", event.ParentBlockHash)
} else {
return nil, fmt.Errorf("failed to parse ParentBlockHash from field at index 6: %w", err)
}
if tokenIDsRaw, ok := data[7].([]interface{}); ok {
tokens, err := parseInt32Array(tokenIDsRaw)
if err != nil {
return nil, fmt.Errorf("failed to parse TokenIDs from field at index 7: %w", err)
}
event.TokenIDs = tokens
slog.Debug("TokenIDs:", "TokenIDs", event.TokenIDs)
} else {
return nil, fmt.Errorf("missing or invalid token_ids")
}
return event, nil
}
// parseBlockStoredEvent parses a BlockStoredEvent from raw data
func parseVllmBlockStored(data []interface{}, timestamp interface{}) (*BlockStoredEvent, error) {
event := &BlockStoredEvent{
Type: EventTypeBlockStored,
}
for i, elem := range data {
slog.Debug("in parseVllmBlockStored:", "index", i, "type", fmt.Sprintf("%T", elem), "value", elem)
}
slog.Debug("in parseVllmBlockStored:", "timestamp", timestamp)
// Parse timestamp
if ts, err := parseTimestamp(timestamp); err == nil {
event.Timestamp = ts
} else {
return nil, fmt.Errorf("failed to parse timestamp: %w", err)
}
// Parse block hashes
if hashes, err := parseUint64Array(data[1]); err == nil {
event.BlockHashes = hashes
} else {
return nil, fmt.Errorf("failed to parse block_hashes: %w", err)
}
// Parse token IDs (array of arrays)
if tokenIDsRaw, ok := data[3].([]interface{}); ok {
tokens, err := parseInt32Array(tokenIDsRaw)
if err != nil {
return nil, fmt.Errorf("failed to parse token_ids at index %w", err)
}
event.TokenIDs = tokens
} else {
return nil, fmt.Errorf("missing or invalid token_ids")
}
var parentHash uint64
if data[2] == nil {
parentHash = uint64(0)
} else {
hash := data[2]
// fmt.Printf("Type of hash>>>>: %T\n", hash)
if h, ok := hash.(uint64); ok {
parentHash = h
} else {
return nil, fmt.Errorf("expected uint64, got %T", hash)
}
}
event.ParentBlockHash = parentHash
if blocksize, err := parseInt64(data[4]); err == nil {
event.BlockSize = blocksize
} else {
return nil, fmt.Errorf("failed to parse field at index 4 as 'block_size': %w", err)
}
if medium, err := safeGetString(data[6]); err == nil {
event.Medium = medium
} else {
return nil, fmt.Errorf("failed to parse 'medium' from field at index 6: %w", err)
}
return event, nil
}
func parseVllmBlockRemoved(data []interface{}, timestamp interface{}) (*BlockRemovedEvent, error) {
event := &BlockRemovedEvent{
Type: EventTypeBlockRemoved,
}
for i, elem := range data {
slog.Debug("in parseVllmBlockRemoved:", "index", i, "type", fmt.Sprintf("%T", elem), "value", elem)
}
// Parse block hashes
if hashes, err := parseUint64Array(data[1]); err == nil {
event.BlockHashes = hashes
} else {
return nil, fmt.Errorf("failed to parse block_hashes: %w", err)
}
// parse medium
if medium, err := safeGetString(data[2]); err == nil {
event.Medium = medium
} else {
return nil, fmt.Errorf("failed to parse 'medium' from field at index 6: %w", err)
}
return event, nil
}
func convertToReplicaList(raw interface{}) ([][]string, error) {
list, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("expected []interface{}, got %T", raw)
}
result := make([][]string, len(list))
for i, item := range list {
subList, ok := item.([]interface{})
if !ok {
return nil, fmt.Errorf("item %d is not []interface{}, got %T", i, item)
}
result[i] = make([]string, len(subList))
for j, v := range subList {
str, ok := v.(string)
if !ok {
return nil, fmt.Errorf("element [%d][%d] is not string, got %T", i, j, v)
}
result[i][j] = str
}
}
return result, nil
}
func safeGetString(val interface{}) (string, error) {
switch v := val.(type) {
case string:
return v, nil
case []byte:
return string(v), nil
case fmt.Stringer:
return v.String(), nil
case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64:
return fmt.Sprintf("%v", v), nil
case nil:
return "", nil
default:
slog.Warn("Unexpected type in string field",
"type", fmt.Sprintf("%T", v),
"value", v)
return fmt.Sprintf("%v", v), nil
}
}
// Helper functions for parsing common types
func parseTimestamp(v interface{}) (time.Time, error) {
switch t := v.(type) {
case time.Time:
return t, nil
case int64:
return time.Unix(t, 0).UTC(), nil
case int:
return time.Unix(int64(t), 0).UTC(), nil
case int32:
return time.Unix(int64(t), 0).UTC(), nil
case uint32:
return time.Unix(int64(t), 0).UTC(), nil
case uint64:
return time.Unix(int64(t), 0).UTC(), nil
case float64:
sec := int64(t)
nsec := int64((t - float64(sec)) * 1e9)
return time.Unix(sec, nsec).UTC().Truncate(time.Microsecond), nil
case float32:
f64 := float64(t)
sec := int64(f64)
nsec := int64((f64 - float64(sec)) * 1e9)
return time.Unix(sec, nsec).UTC().Truncate(time.Microsecond), nil
case string:
// Try to parse RFC3339 format
return time.Parse(time.RFC3339, t)
default:
return time.Time{}, fmt.Errorf("unsupported timestamp type: %T", v)
}
}
func parseInt64(v interface{}) (int64, error) {
switch n := v.(type) {
case int64:
return n, nil
case int:
return int64(n), nil
case int32:
return int64(n), nil
case int16:
return int64(n), nil
case int8:
return int64(n), nil
case uint:
return int64(n), nil
case uint64:
return int64(n), nil
case uint32:
return int64(n), nil
case uint16:
return int64(n), nil
case uint8:
return int64(n), nil
case float64:
return int64(n), nil
case float32:
return int64(n), nil
default:
return 0, fmt.Errorf("unsupported int64 type: %T", v)
}
}
func parseUint64(v interface{}) (uint64, error) {
switch n := v.(type) {
case uint64:
return n, nil
case int:
return uint64(n), nil
case int32:
return uint64(n), nil
case int16:
return uint64(n), nil
case int8:
return uint64(n), nil
case uint:
return uint64(n), nil
case int64:
return uint64(n), nil
case uint32:
return uint64(n), nil
case uint16:
return uint64(n), nil
case uint8:
return uint64(n), nil
case float64:
return uint64(n), nil
case float32:
return uint64(n), nil
default:
return 0, fmt.Errorf("unsupported int64 type: %T", v)
}
}
func parseUint64Array(v interface{}) ([]uint64, error) {
arr, ok := v.([]interface{})
if !ok {
return nil, fmt.Errorf("expected array, got %T", v)
}
result := make([]uint64, 0, len(arr))
for i, item := range arr {
val, err := parseUint64(item)
if err != nil {
return nil, fmt.Errorf("failed to parse element at index %d: %w", i, err)
}
result = append(result, val)
}
return result, nil
}
func parseInt32Array(v interface{}) ([]int32, error) {
arr, ok := v.([]interface{})
if !ok {
return nil, fmt.Errorf("expected array, got %T", v)
}
result := make([]int32, 0, len(arr))
for i, item := range arr {
switch n := item.(type) {
case int32:
result = append(result, n)
case int:
result = append(result, int32(n))
case int64:
result = append(result, int32(n))
case int16:
result = append(result, int32(n))
case int8:
result = append(result, int32(n))
case uint:
result = append(result, int32(n))
case uint64:
result = append(result, int32(n))
case uint32:
result = append(result, int32(n))
case uint16:
result = append(result, int32(n))
case uint8:
result = append(result, int32(n))
case float64:
result = append(result, int32(n))
case float32:
result = append(result, int32(n))
default:
return nil, fmt.Errorf("unsupported int32 type at index %d: %T", i, item)
}
}
return result, nil
}
func parseMooncakeUint64(v interface{}) (uint64, error) {
var s string
switch val := v.(type) {
case string:
s = val
case nil:
s = ""
case uint64:
return val, nil
case int:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int8:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int16:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int32:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case int64:
if val < 0 {
return 0, fmt.Errorf("negative value %d", val)
}
return uint64(val), nil
case uint:
return uint64(val), nil
case uint8:
return uint64(val), nil
case uint16:
return uint64(val), nil
case uint32:
return uint64(val), nil
default:
s = fmt.Sprint(v)
}
if s == "" {
return 0, nil
}
return strconv.ParseUint(s, 10, 64)
}
func parseMooncakeParentUint64(v interface{}) ([]uint64, error) {
switch val := v.(type) {
case nil:
return []uint64{}, nil
case string:
if val == "" {
return []uint64{}, nil
}
parts := strings.FieldsFunc(val, func(r rune) bool {
return r == ',' || r == ' ' || r == '\t' || r == '\n'
})
result := make([]uint64, 0, len(parts))
for _, part := range parts {
if part == "" {
continue
}
u, err := strconv.ParseUint(part, 10, 64)
if err != nil {
return nil, fmt.Errorf("failed to parse uint64 from string %q: %w", part, err)
}
result = append(result, u)
}
return result, nil
case []interface{}:
result := make([]uint64, 0, len(val))
for _, item := range val {
u, err := parseSingleUint64(item)
if err != nil {
return nil, fmt.Errorf("failed to parse element %v: %w", item, err)
}
result = append(result, u)
}
return result, nil
case []uint64:
// already correct type
return val, nil
default:
// try parse as single uint64
u, err := parseSingleUint64(val)
if err != nil {
return nil, err
}
return []uint64{u}, nil
}
}
func parseSingleUint64(v interface{}) (uint64, error) {
switch val := v.(type) {
case uint64:
return val, nil
case int:
if val < 0 {
return 0, fmt.Errorf("negative int %d cannot convert to uint64", val)
}
return uint64(val), nil
case int8:
if val < 0 {
return 0, fmt.Errorf("negative int8 %d cannot convert to uint64", val)
}
return uint64(val), nil
case int16:
if val < 0 {
return 0, fmt.Errorf("negative int16 %d cannot convert to uint64", val)
}
return uint64(val), nil
case int32:
if val < 0 {
return 0, fmt.Errorf("negative int32 %d cannot convert to uint64", val)
}
return uint64(val), nil
case int64:
if val < 0 {
return 0, fmt.Errorf("negative int64 %d cannot convert to uint64", val)
}
return uint64(val), nil
case uint:
return uint64(val), nil
case uint8:
return uint64(val), nil
case uint16:
return uint64(val), nil
case uint32:
return uint64(val), nil
case float32:
f := float64(val)
if f < 0 || f != float64(uint64(f)) {
return 0, fmt.Errorf("float32 %v invalid for uint64", val)
}
return uint64(f), nil
case float64:
if val < 0 || val != float64(uint64(val)) {
return 0, fmt.Errorf("float64 %v invalid for uint64", val)
}
return uint64(val), nil
case string:
if val == "" {
return 0, fmt.Errorf("empty string cannot be parsed as uint64")
}
return strconv.ParseUint(val, 10, 64)
case nil:
return 0, fmt.Errorf("nil cannot be parsed as uint64")
default:
return 0, fmt.Errorf("unsupported type %T for uint64 conversion", v)
}
}

View File

@ -0,0 +1,178 @@
package zmq_test
import (
"testing"
"time"
"conductor/zmq"
msgpack "github.com/shamaton/msgpack/v2"
)
func TestDecodeMooncakeEventBatch(t *testing.T) {
timestamp := int64(1700000000)
event := []interface{}{
"BlockStoreEvent",
"mooncake-key-123",
[][]interface{}{
[]interface{}{"replica1", "replica2"},
[]interface{}{"replica3"},
},
nil, // index 3 is not used
int64(1024), // BlockSize at index 4
[]interface{}{uint64(100), uint64(200)}, // BlockHashes at index 5
uint64(50), // ParentBlockHash at index 6
[]interface{}{int32(1), int32(2), int32(3)}, // TokenIDs at index 7
}
events := []interface{}{event}
batch := []interface{}{timestamp, events}
data, err := msgpack.Marshal(batch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
result, err := zmq.DecodeMooncakeEventBatch(data)
if err != nil {
t.Fatalf("DecodeMooncakeEventBatch failed: %v", err)
}
if result.Source != zmq.SourceMooncake {
t.Errorf("Expected source %v, got %v", zmq.SourceMooncake, result.Source)
}
if len(result.Events) != 1 {
t.Fatalf("Expected 1 event, got %d", len(result.Events))
}
blockEvent, ok := result.Events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", result.Events[0])
}
if blockEvent.Type != zmq.EventTypeBlockStored {
t.Errorf("Expected type %v, got %v", zmq.EventTypeBlockStored, blockEvent.Type)
}
if blockEvent.MooncakeKey != "mooncake-key-123" {
t.Errorf("Expected MooncakeKey 'mooncake-key-123', got '%s'", blockEvent.MooncakeKey)
}
if blockEvent.BlockSize != 1024 {
t.Errorf("Expected BlockSize 1024, got %d", blockEvent.BlockSize)
}
if len(blockEvent.BlockHashes) != 2 {
t.Fatalf("Expected 2 block hashes, got %d", len(blockEvent.BlockHashes))
}
if blockEvent.BlockHashes[0] != 100 || blockEvent.BlockHashes[1] != 200 {
t.Errorf("Expected block hashes [100, 200], got %v", blockEvent.BlockHashes)
}
if blockEvent.ParentBlockHash != 50 {
t.Errorf("Expected ParentBlockHash 50, got %d", blockEvent.ParentBlockHash)
}
if len(blockEvent.TokenIDs) != 3 {
t.Fatalf("Expected 3 token IDs, got %d", len(blockEvent.TokenIDs))
}
if len(blockEvent.ReplicaList) != 2 {
t.Fatalf("Expected 2 replica lists, got %d", len(blockEvent.ReplicaList))
}
}
func TestDecodeVllmEventBatch(t *testing.T) {
timestamp := int64(1700000000)
event := []interface{}{
"BlockStored",
[]interface{}{uint64(100), uint64(200)}, // BlockHashes
uint64(5000000000), // ParentBlockHash
[]interface{}{int32(10000000), int32(2), int32(3)}, // TokenIDs
int64(1024), // BlockSize
}
events := []interface{}{event}
status := "ok"
batch := []interface{}{timestamp, events, status}
data, err := msgpack.Marshal(batch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
result, err := zmq.DecodeVllmEventBatch(data)
if err != nil {
t.Fatalf("DecodeVllmEventBatch failed: %v", err)
}
if result.Source != zmq.SourceVLLM {
t.Errorf("Expected source %v, got %v", zmq.SourceVLLM, result.Source)
}
if len(result.Events) != 1 {
t.Fatalf("Expected 1 event, got %d", len(result.Events))
}
blockEvent, ok := result.Events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", result.Events[0])
}
if blockEvent.Type != zmq.EventTypeBlockStored {
t.Errorf("Expected type %v, got %v", zmq.EventTypeBlockStored, blockEvent.Type)
}
expectedTime := time.Unix(1700000000, 0).UTC()
if !blockEvent.Timestamp.Equal(expectedTime) {
t.Errorf("Expected timestamp %v, got %v", expectedTime, blockEvent.Timestamp)
}
if len(blockEvent.BlockHashes) != 2 {
t.Fatalf("Expected 2 block hashes, got %d", len(blockEvent.BlockHashes))
}
if blockEvent.BlockHashes[0] != 100 || blockEvent.BlockHashes[1] != 200 {
t.Errorf("Expected block hashes [100, 200], got %v", blockEvent.BlockHashes)
}
if blockEvent.ParentBlockHash != 5000000000 {
t.Errorf("Expected ParentBlockHash 50, got %d", blockEvent.ParentBlockHash)
}
if blockEvent.BlockSize != 1024 {
t.Errorf("Expected BlockSize 1024, got %d", blockEvent.BlockSize)
}
if len(blockEvent.TokenIDs) != 3 {
t.Fatalf("Expected 3 token IDs, got %d", len(blockEvent.TokenIDs))
}
}
func TestDecodeMooncakeEventBatch_InvalidData(t *testing.T) {
// Test with invalid array length
invalidBatch := []interface{}{int64(1700000000)} // Missing events
data, err := msgpack.Marshal(invalidBatch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
_, err = zmq.DecodeMooncakeEventBatch(data)
if err == nil {
t.Error("Expected error for invalid array length, got nil")
}
}
func TestDecodeVllmEventBatch_InvalidData(t *testing.T) {
// Test with invalid array length
invalidBatch := []interface{}{int64(1700000000)} // Missing events and status
data, err := msgpack.Marshal(invalidBatch)
if err != nil {
t.Fatalf("Failed to marshal test data: %v", err)
}
_, err = zmq.DecodeVllmEventBatch(data)
if err == nil {
t.Error("Expected error for invalid array length, got nil")
}
}

View File

@ -0,0 +1,41 @@
package zmq
import (
"fmt"
"time"
)
// EventHandler processes received KV events
type EventHandler interface {
HandleEvent(event KVEvent, dpRank int64) error
}
// ZMQClientConfig contains configuration for the ZMQ client
type ZMQClientConfig struct {
CachePoolKey string
Endpoint string
ReplayEndpoint string
ModelName string
PollTimeout time.Duration
ReplayTimeout time.Duration
ReconnectDelay time.Duration
}
const (
// Timeouts and intervals
DefaultPollTimeout = 100 * time.Millisecond
DefaultReplayTimeout = 5 * time.Second
DefaultReconnectInterval = 1 * time.Second
MaxReconnectInterval = 30 * time.Second
ReconnectBackoffFactor = 2.0
EventChannelBufferSize = 1000
)
func ValidateConfig(config *ZMQClientConfig) error {
if config.Endpoint == "" {
return fmt.Errorf("endpoint is required")
}
return nil
}

View File

@ -0,0 +1,355 @@
package zmq
import (
"context"
"encoding/binary"
"fmt"
"log/slog"
"sync"
"time"
zmq "github.com/pebbe/zmq4"
)
type ZMQClient struct {
config *ZMQClientConfig
subSocket *zmq.Socket
replaySocket *zmq.Socket
eventHandler EventHandler
// State management
mu sync.RWMutex
connected bool
lastSeq int64
reconnectDelay time.Duration
// Lifecycle
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
func NewZMQClient(config *ZMQClientConfig, handler EventHandler) *ZMQClient {
ctx, cancel := context.WithCancel(context.Background())
return &ZMQClient{
config: config,
eventHandler: handler,
lastSeq: -1,
reconnectDelay: config.ReconnectDelay,
ctx: ctx,
cancel: cancel,
}
}
// Start initiates the connection and background event consumption loop.
func (c *ZMQClient) Start() error {
// Attempt initial connection
if err := c.Connect(); err != nil {
return fmt.Errorf("initial connection failed: %w", err)
}
c.wg.Add(1)
go c.loop()
slog.Info("ZMQ client started", "service", c.config.CachePoolKey)
return nil
}
func (c *ZMQClient) Stop() {
c.cancel()
c.wg.Wait()
c.mu.Lock()
c.cleanupSockets()
c.mu.Unlock()
slog.Info("ZMQ client stopped", "service", c.config.CachePoolKey)
}
// loop is the main background loop handling events and reconnections.
// Simplified: Fixed reconnect interval, single loop structure.
func (c *ZMQClient) loop() {
defer c.wg.Done()
for {
// Check if we should stop
select {
case <-c.ctx.Done():
return
default:
}
// 1. If disconnected, wait for ticker then try to reconnect
if !c.isConnected() {
c.handleReconnect()
continue
}
// 2. If connected, consume events
if err := c.consume(); err != nil {
slog.Error("Consumption error", "service", c.config.CachePoolKey, "error", err)
c.markDisconnected()
}
}
}
func (c *ZMQClient) handleReconnect() {
slog.Info("Attempting to reconnect to the service.", "service", c.config.CachePoolKey, "reconnectDelay", c.reconnectDelay)
ticker := time.NewTicker(c.config.ReconnectDelay)
defer ticker.Stop()
select {
case <-c.ctx.Done():
return
case <-ticker.C:
}
if err := c.Connect(); err != nil {
slog.Error("Reconnect failed", "service", c.config.CachePoolKey, "error", err)
}
// Reconnected! Request replay from last known sequence
lastSeq := c.getLastSequence()
if lastSeq >= 0 {
slog.Info("Reconnected", "service", c.config.CachePoolKey, "resuming_from", lastSeq+1)
if err := c.requestReplay(lastSeq + 1); err != nil {
slog.Warn("Failed to request replay after reconnect", "service", c.config.CachePoolKey, "error", err)
}
}
}
// Connect establishes the ZMQ SUB and DEALER sockets.
func (c *ZMQClient) Connect() error {
c.mu.Lock()
defer c.mu.Unlock()
if c.connected {
return nil
}
// Ensure clean state
c.cleanupSockets()
sock, err := zmq.NewSocket(zmq.SUB)
if err != nil {
return fmt.Errorf("create socket failed: %w", err)
}
if err := sock.SetIpv6(true); err != nil {
_ = sock.Close()
return fmt.Errorf("failed to enable IPv6 on socket: %w", err)
}
if err := sock.Connect(c.config.Endpoint); err != nil {
_ = sock.Close()
return fmt.Errorf("failed to connect to %s: %w", c.config.Endpoint, err)
}
// Important: Subscribe to all topics
if err := sock.SetSubscribe(""); err != nil {
_ = sock.Close()
return fmt.Errorf("failed to subscribe: %w", err)
}
replaySocket, err := zmq.NewSocket(zmq.DEALER)
if err != nil {
sock.Close()
return fmt.Errorf("failed to create DEALER socket: %w", err)
}
// Enable IPv6 for dual-stack support
if err := replaySocket.SetIpv6(true); err != nil {
_ = sock.Close()
_ = replaySocket.Close()
return fmt.Errorf("failed to enable IPv6 on DEALER socket: %w", err)
}
if err := replaySocket.Connect(c.config.ReplayEndpoint); err != nil {
_ = sock.Close()
_ = replaySocket.Close()
return fmt.Errorf("failed to connect to replay endpoint %s: %w", c.config.ReplayEndpoint, err)
}
c.subSocket = sock
c.replaySocket = replaySocket
c.connected = true
c.reconnectDelay = c.config.ReconnectDelay
slog.Info("Successfully connected to vLLM publisher", "service", c.config.CachePoolKey, "endpoint", c.config.Endpoint)
return nil
}
// consume reads and processes messages from the SUB socket.
func (c *ZMQClient) consume() error {
c.mu.RLock()
socket := c.subSocket
c.mu.RUnlock()
if socket == nil {
return fmt.Errorf("socket is nil")
}
poller := zmq.NewPoller()
poller.Add(socket, zmq.POLLIN)
// Poll for data
polled, err := poller.Poll(c.config.PollTimeout)
if err != nil {
return fmt.Errorf("poll error: %w", err)
}
if len(polled) == 0 {
return nil // No data, continue loop
}
if err := c.processMessage(socket); err != nil {
return fmt.Errorf("failed to process message: %w", err)
}
return nil
}
func (c *ZMQClient) processMessage(socket *zmq.Socket) error {
if socket == nil {
return fmt.Errorf("socket is nil")
}
// Read Frames: [Topic, Seq, Payload]
topic, err := socket.RecvBytes(0)
if err != nil {
return err
}
seqBytes, err := socket.RecvBytes(0)
if err != nil {
return err
}
payload, err := socket.RecvBytes(0)
if err != nil {
return err
}
if len(seqBytes) != 8 {
return fmt.Errorf("invalid sequence length")
}
seq := int64(binary.BigEndian.Uint64(seqBytes))
c.mu.RLock()
lastSeq := c.lastSeq
c.mu.RUnlock()
if lastSeq != -1 && seq > lastSeq+1 {
slog.Warn("Event gap detected",
"service", c.config.CachePoolKey,
"missed", seq-lastSeq-1,
"last", lastSeq,
"current", seq,
)
// Trigger replay for missed events?
// Usually we just log warning here, or could auto-trigger requestReplay
}
// Update Sequence immediately to keep state fresh
c.mu.Lock()
c.lastSeq = seq
c.mu.Unlock()
slog.Debug("enter deal topic", "topic", topic)
var batch *EventBatch
switch string(topic) {
case "mooncake":
batch, err = DecodeMooncakeEventBatch(payload)
default:
batch, err = DecodeVllmEventBatch(payload)
}
if err != nil {
return fmt.Errorf("decode failed: %w", err)
}
for _, event := range batch.Events {
// Inject Source Name
switch e := event.(type) {
case *BlockStoredEvent:
e.PodName = c.config.CachePoolKey
case *BlockRemovedEvent:
e.PodName = c.config.CachePoolKey
}
if err := c.eventHandler.HandleEvent(event, batch.DataParallelRank); err != nil {
slog.Error("Handler error", "service", c.config.CachePoolKey, "error", err)
}
}
slog.Debug("Processed batch", "service", c.config.CachePoolKey, "seq", seq, "topic", string(topic))
return nil
}
func (c *ZMQClient) requestReplay(fromSeq int64) error {
c.mu.RLock()
socket := c.replaySocket
c.mu.RUnlock()
if socket == nil {
return fmt.Errorf("replay socket is nil")
}
req := make([]byte, 8)
binary.BigEndian.PutUint64(req, uint64(fromSeq))
if _, err := socket.SendBytes(req, 0); err != nil {
return fmt.Errorf("failed to send replay request: %w", err)
}
// Ideally, we should wait for an ACK here if the protocol supports it
// For simplicity in static client, we fire and forget the request,
// assuming the server will send the replayed events via the SUB channel (or DEALER response)
// Original code read response from DEALER, let's keep that.
_ = socket.SetRcvtimeo(c.config.ReplayTimeout)
resp, err := socket.RecvBytes(0)
if err != nil {
return fmt.Errorf("failed to receive replay response: %w", err)
}
slog.Info("Replay requested", "service", c.config.CachePoolKey, "from", fromSeq, "resp_len", len(resp))
return nil
}
func (c *ZMQClient) cleanupSockets() {
if c.subSocket != nil {
c.subSocket.Close()
c.subSocket = nil
}
if c.replaySocket != nil {
c.replaySocket.Close()
c.replaySocket = nil
}
c.connected = false
}
func (c *ZMQClient) markDisconnected() {
c.mu.Lock()
defer c.mu.Unlock()
c.connected = false
}
func (c *ZMQClient) isConnected() bool {
c.mu.RLock()
defer c.mu.RUnlock()
return c.connected
}
func (c *ZMQClient) getLastSequence() int64 {
c.mu.RLock()
defer c.mu.RUnlock()
return c.lastSeq
}

View File

@ -0,0 +1,663 @@
package zmq_test
import (
"context"
"encoding/binary"
"errors"
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"conductor/zmq"
zmq4 "github.com/pebbe/zmq4"
msgpack "github.com/shamaton/msgpack/v2"
)
// MockEventHandler implements EventHandler for testing
type MockEventHandler struct {
mu sync.Mutex
events []zmq.KVEvent
handleError error
callCount int64
}
func NewMockEventHandler() *MockEventHandler {
return &MockEventHandler{
events: make([]zmq.KVEvent, 0),
}
}
func (m *MockEventHandler) HandleEvent(event zmq.KVEvent) error {
m.mu.Lock()
defer m.mu.Unlock()
atomic.AddInt64(&m.callCount, 1)
if m.handleError != nil {
return m.handleError
}
m.events = append(m.events, event)
return nil
}
func (m *MockEventHandler) GetEvents() []zmq.KVEvent {
m.mu.Lock()
defer m.mu.Unlock()
events := make([]zmq.KVEvent, len(m.events))
copy(events, m.events)
return events
}
func (m *MockEventHandler) GetCallCount() int64 {
return atomic.LoadInt64(&m.callCount)
}
func (m *MockEventHandler) SetHandleError(err error) {
m.mu.Lock()
defer m.mu.Unlock()
m.handleError = err
}
func (m *MockEventHandler) Clear() {
m.mu.Lock()
defer m.mu.Unlock()
m.events = m.events[:0]
m.handleError = nil
atomic.StoreInt64(&m.callCount, 0)
}
// MockPublisher simulates a ZMQ publisher for testing
type MockPublisher struct {
pubSocket *zmq4.Socket
routerSocket *zmq4.Socket
sequence int64
mu sync.Mutex
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
func createMockPublisher(t *testing.T, pubPort, routerPort int) *MockPublisher {
ctx, cancel := context.WithCancel(context.Background())
// Create PUB socket
pubSocket, err := zmq4.NewSocket(zmq4.PUB)
if err != nil {
t.Fatalf("Failed to create PUB socket: %v", err)
}
err = pubSocket.SetIpv6(true)
if err != nil {
pubSocket.Close()
t.Fatalf("Failed to enable IPv6 on PUB socket: %v", err)
}
err = pubSocket.Bind("tcp://127.0.0.1:*")
if err != nil {
pubSocket.Close()
t.Fatalf("Failed to bind PUB socket: %v", err)
}
// Create ROUTER socket for replay
routerSocket, err := zmq4.NewSocket(zmq4.ROUTER)
if err != nil {
pubSocket.Close()
t.Fatalf("Failed to create ROUTER socket: %v", err)
}
err = routerSocket.SetIpv6(true)
if err != nil {
pubSocket.Close()
routerSocket.Close()
t.Fatalf("Failed to enable IPv6 on ROUTER socket: %v", err)
}
err = routerSocket.Bind("tcp://127.0.0.1:*")
if err != nil {
pubSocket.Close()
routerSocket.Close()
t.Fatalf("Failed to bind ROUTER socket: %v", err)
}
mp := &MockPublisher{
pubSocket: pubSocket,
routerSocket: routerSocket,
ctx: ctx,
cancel: cancel,
}
// Start replay handler
mp.wg.Add(1)
go mp.handleReplay()
// Wait for sockets to bind
time.Sleep(100 * time.Millisecond)
return mp
}
func (mp *MockPublisher) PublishEvent(topic string, event zmq.KVEvent) error {
// Encode event based on topic
var payload []byte
var err error
if topic == "mooncake" {
payload, err = encodeMooncakeEvent(event)
} else {
payload, err = encodeVllmEvent(event)
}
if err != nil {
return err
}
mp.mu.Lock()
mp.sequence++
seq := mp.sequence
mp.mu.Unlock()
seqBytes := make([]byte, 8)
binary.BigEndian.PutUint64(seqBytes, uint64(seq))
_, err = mp.pubSocket.SendMessage(topic, seqBytes, payload)
return err
}
func (mp *MockPublisher) handleReplay() {
defer mp.wg.Done()
for {
select {
case <-mp.ctx.Done():
return
default:
}
// Set receive timeout
_ = mp.routerSocket.SetRcvtimeo(100 * time.Millisecond)
msg, err := mp.routerSocket.RecvBytes(0)
if err != nil {
continue
}
if len(msg) == 8 {
// Send ACK
_, _ = mp.routerSocket.SendBytes([]byte("OK"), 0)
}
}
}
func (mp *MockPublisher) Close() {
mp.cancel()
mp.wg.Wait()
_ = mp.pubSocket.Close()
_ = mp.routerSocket.Close()
}
func encodeMooncakeEvent(event zmq.KVEvent) ([]byte, error) {
switch e := event.(type) {
case *zmq.BlockStoredEvent:
timestamp := e.Timestamp.Unix()
eventData := []interface{}{
"BlockStoreEvent",
e.MooncakeKey,
e.ReplicaList,
nil, // index 3 not used
e.BlockSize,
convertUint64Slice(e.BlockHashes),
e.ParentBlockHash,
convertInt32Slice(e.TokenIDs),
}
batch := []interface{}{timestamp, []interface{}{eventData}}
return msgpack.Marshal(batch)
default:
return nil, errors.New("unsupported event type for mooncake")
}
}
func encodeVllmEvent(event zmq.KVEvent) ([]byte, error) {
switch e := event.(type) {
case *zmq.BlockStoredEvent:
timestamp := e.Timestamp.Unix()
eventData := []interface{}{
"BlockStored",
convertUint64Slice(e.BlockHashes),
e.ParentBlockHash,
convertInt32Slice(e.TokenIDs),
e.BlockSize,
}
batch := []interface{}{timestamp, []interface{}{eventData}, "ok"}
return msgpack.Marshal(batch)
case *zmq.BlockRemovedEvent:
timestamp := e.Timestamp.Unix()
eventData := []interface{}{
"BlockRemoved",
convertUint64Slice(e.BlockHashes),
}
batch := []interface{}{timestamp, []interface{}{eventData}, "ok"}
return msgpack.Marshal(batch)
default:
return nil, errors.New("unsupported event type for vllm")
}
}
func convertUint64Slice(slice []uint64) []interface{} {
result := make([]interface{}, len(slice))
for i, v := range slice {
result[i] = uint64(v)
}
return result
}
func convertInt32Slice(slice []int32) []interface{} {
result := make([]interface{}, len(slice))
for i, v := range slice {
result[i] = int32(v)
}
return result
}
func skipIfZMQUnavailable(t *testing.T) {
ctx, err := zmq4.NewContext()
if err != nil {
t.Skip("ZMQ not available:", err)
}
_ = ctx.Term()
}
func TestZMQClient_Connect_Success(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5547, 5548)
defer publisher.Close()
// Get actual bound ports
pubEndpoint, err := publisher.pubSocket.GetLastEndpoint()
if err != nil {
t.Fatalf("Failed to get PUB endpoint: %v", err)
}
routerEndpoint, err := publisher.routerSocket.GetLastEndpoint()
if err != nil {
t.Fatalf("Failed to get ROUTER endpoint: %v", err)
}
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err = client.Connect()
if err != nil {
t.Fatalf("Connect failed: %v", err)
}
client.Stop()
}
func TestZMQClient_Connect_AlreadyConnected(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Connect()
if err != nil {
t.Fatalf("First Connect failed: %v", err)
}
// Connect again should not error
err = client.Connect()
if err != nil {
t.Fatalf("Second Connect failed: %v", err)
}
client.Stop()
}
func TestZMQClient_Start_Stop(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
// Wait a bit for loop to start
time.Sleep(50 * time.Millisecond)
// Stop should work gracefully
client.Stop()
}
func TestZMQClient_ProcessMessage_MooncakeTopic(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Create and publish mooncake event
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{100, 200},
TokenIDs: []int32{1, 2, 3},
ParentBlockHash: 50,
BlockSize: 1024,
MooncakeKey: "mooncake-key-123",
ReplicaList: [][]string{{"replica1", "replica2"}},
ModelName: "test-model",
}
err = publisher.PublishEvent("mooncake", event)
if err != nil {
t.Fatalf("PublishEvent failed: %v", err)
}
// Wait for processing
time.Sleep(200 * time.Millisecond)
events := handler.GetEvents()
if len(events) < 1 {
t.Fatalf("Expected at least 1 event, got %d", len(events))
}
blockEvent, ok := events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", events[0])
}
if blockEvent.PodName != "test-pod" {
t.Errorf("Expected PodName 'test-pod', got '%s'", blockEvent.PodName)
}
if blockEvent.MooncakeKey != "mooncake-key-123" {
t.Errorf("Expected MooncakeKey 'mooncake-key-123', got '%s'", blockEvent.MooncakeKey)
}
if len(blockEvent.BlockHashes) == 0 || blockEvent.BlockHashes[0] != 100 {
t.Errorf("Expected first BlockHash 100, got %v", blockEvent.BlockHashes)
}
}
func TestZMQClient_ProcessMessage_VllmTopic(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Create and publish vllm event
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{300, 400},
TokenIDs: []int32{4, 5, 6},
ParentBlockHash: 1500000000000000,
BlockSize: 2048,
ModelName: "test-model",
}
err = publisher.PublishEvent("vllm", event)
if err != nil {
t.Fatalf("PublishEvent failed: %v", err)
}
// Wait for processing
time.Sleep(200 * time.Millisecond)
events := handler.GetEvents()
if len(events) < 1 {
t.Fatalf("Expected at least 1 event, got %d", len(events))
}
blockEvent, ok := events[0].(*zmq.BlockStoredEvent)
if !ok {
t.Fatalf("Expected BlockStoredEvent, got %T", events[0])
}
if blockEvent.PodName != "test-pod" {
t.Errorf("Expected PodName 'test-pod', got '%s'", blockEvent.PodName)
}
if len(blockEvent.BlockHashes) == 0 || blockEvent.BlockHashes[0] != 300 {
t.Errorf("Expected first BlockHash 300, got %v", blockEvent.BlockHashes)
}
}
func TestZMQClient_SequenceTracking(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Publish multiple events
for i := 0; i < 5; i++ {
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{uint64(i)},
TokenIDs: []int32{int32(i)},
ParentBlockHash: 1000000000000000000,
BlockSize: 128,
ModelName: "test-model",
}
_ = publisher.PublishEvent("vllm", event)
time.Sleep(50 * time.Millisecond)
}
// Wait for processing
time.Sleep(500 * time.Millisecond)
// Verify events were processed
events := handler.GetEvents()
if len(events) < 5 {
t.Errorf("Expected at least 5 events, got %d", len(events))
}
}
func TestZMQClient_ProcessMessage_EventGap(t *testing.T) {
skipIfZMQUnavailable(t)
publisher := createMockPublisher(t, 5557, 5558)
defer publisher.Close()
pubEndpoint, _ := publisher.pubSocket.GetLastEndpoint()
routerEndpoint, _ := publisher.routerSocket.GetLastEndpoint()
pubPort := extractPortFromEndpoint(pubEndpoint)
routerPort := extractPortFromEndpoint(routerEndpoint)
config := &zmq.ZMQClientConfig{
CachePoolKey: "test-pod",
ServiceIP: "127.0.0.1",
ModelName: "test-model",
Port: pubPort,
RouterPort: routerPort,
PollTimeout: 100 * time.Millisecond,
ReplayTimeout: 1 * time.Second,
ReconnectDelay: 100 * time.Millisecond,
}
handler := NewMockEventHandler()
client := zmq.NewZMQClient(config, handler)
err := client.Start()
if err != nil {
t.Fatalf("Start failed: %v", err)
}
defer client.Stop()
// Wait for connection
time.Sleep(100 * time.Millisecond)
// Manually set last sequence to simulate gap
// This is a bit tricky since we need to access internal state
// For now, we publish events and verify gap detection works
// The actual gap detection is logged, so we verify the code path exists
event := &zmq.BlockStoredEvent{
Type: zmq.EventTypeBlockStored,
Timestamp: time.Now().UTC(),
BlockHashes: []uint64{100},
TokenIDs: []int32{1},
ModelName: "test-model",
}
_ = publisher.PublishEvent("vllm", event)
time.Sleep(100 * time.Millisecond)
// Publish another event - gap detection should work for subsequent events
_ = publisher.PublishEvent("vllm", event)
time.Sleep(100 * time.Millisecond)
}
// Helper function to extract port from endpoint string like "tcp://127.0.0.1:5557"
func extractPortFromEndpoint(endpoint string) int {
// Parse "tcp://127.0.0.1:5557" to get 5557
parts := strings.Split(endpoint, ":")
if len(parts) < 3 {
return 5557 // Default fallback
}
portStr := parts[len(parts)-1]
port, err := strconv.Atoi(portStr)
if err != nil {
return 5557 // Default fallback
}
return port
}

View File

@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from the vLLM repository's toy_proxy_server.py in tests/v1/kv_connector/nixl_integration/.
import argparse
import itertools
import logging
import os
import uuid
from contextlib import asynccontextmanager
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
class CacheAwareRouter():
def __init__(self, address, endpoint):
self.address = address
self.endpoint = endpoint
self.client = httpx.AsyncClient(timeout=None, base_url=f'http://{address}')
async def get_best_prefiller(self, token_ids: list, ready_instances, req_data):
# call conductor restful api to get cache hit situation
model_name = req_data.get("model", "ds")
lora_id = req_data.get("lora_id", -1)
request_data = {
"instances": ready_instances,
"token_ids": token_ids,
"model_name": model_name,
"lora_id": lora_id
}
headers = {
"Content-Type": "application/json",
}
logger.debug(f"conductor request_data: {request_data}")
response = await self.client.post(self.endpoint, json=request_data, headers=headers)
response.raise_for_status()
return response.json()["HitStatus"]
async def close(self) -> None:
if self.client:
await self.client.aclose()
@asynccontextmanager
async def lifespan(app: FastAPI):
"""
Lifespan context manager to handle startup and shutdown events.
"""
# Startup: Initialize client pools for prefiller and decoder services
app.state.prefill_clients = []
app.state.decode_clients = []
# Create prefill clients
for i, (host, port) in enumerate(global_args.prefiller_instances):
prefiller_base_url = f'http://{host}:{port}'
app.state.prefill_clients.append({
'client':
httpx.AsyncClient(timeout=None, base_url=prefiller_base_url),
'host':
host,
'port':
port,
'id':
i
})
# Create decode clients
for i, (host, port) in enumerate(global_args.decoder_instances):
decoder_base_url = f'http://{host}:{port}'
app.state.decode_clients.append({
'client':
httpx.AsyncClient(timeout=None, base_url=decoder_base_url),
'host':
host,
'port':
port,
'id':
i
})
# Create conductor client
app.state.conductor_client = CacheAwareRouter(global_args.conductor_address, "/cache")
# Initialize round-robin iterators
app.state.prefill_iterator = itertools.cycle(
range(len(app.state.prefill_clients)))
app.state.decode_iterator = itertools.cycle(
range(len(app.state.decode_clients)))
logger.info(f"Initialized {len(app.state.prefill_clients)} prefill clients "
f"and {len(app.state.decode_clients)} decode clients.")
yield
# Shutdown: Close all clients
for client_info in app.state.prefill_clients:
await client_info['client'].aclose()
for client_info in app.state.decode_clients:
await client_info['client'].aclose()
# Close conductor client
await app.state.conductor_client.close()
# Update FastAPI app initialization to use lifespan
app = FastAPI(lifespan=lifespan)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", type=str, default="localhost")
# For prefiller instances
parser.add_argument("--prefiller-hosts",
"--prefiller-host",
type=str,
nargs="+",
default=["localhost"])
parser.add_argument("--prefiller-ports",
"--prefiller-port",
type=int,
nargs="+",
default=[8100])
# For decoder instances
parser.add_argument("--decoder-hosts",
"--decoder-host",
type=str,
nargs="+",
default=["localhost"])
parser.add_argument("--decoder-ports",
"--decoder-port",
type=int,
nargs="+",
default=[8200])
parser.add_argument("--conductor-address", type=str, default="127.0.0.1:13333")
args = parser.parse_args()
# Validate and pair hosts with ports
if len(args.prefiller_hosts) != len(args.prefiller_ports):
raise ValueError(
"Number of prefiller hosts must match number of prefiller ports")
if len(args.decoder_hosts) != len(args.decoder_ports):
raise ValueError(
"Number of decoder hosts must match number of decoder ports")
# Create tuples of (host, port) for each service type
args.prefiller_instances = list(
zip(args.prefiller_hosts, args.prefiller_ports))
args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
return args
def get_next_client(app, service_type: str):
"""
Get the next client in round-robin fashion.
Args:
app: The FastAPI app instance
service_type: Either 'prefill' or 'decode'
Returns:
The next client to use
"""
if service_type == 'prefill':
client_idx = next(app.state.prefill_iterator)
return app.state.prefill_clients[client_idx]
elif service_type == 'decode':
client_idx = next(app.state.decode_iterator)
return app.state.decode_clients[client_idx]
else:
raise ValueError(f"Unknown service type: {service_type}")
async def get_best_prefiller(app, token_ids: list, round_robin_prefill, req_data):
# Get all prefill instances
ready_instances = []
index_map = {}
for index, client_info in enumerate(app.state.prefill_clients):
if client_info['client'].is_closed:
continue
ready_instances.append(client_info['host'])
index_map[client_info['host']] = index
cache_hit_status = await app.state.conductor_client.get_best_prefiller(token_ids, ready_instances, req_data)
if not cache_hit_status:
return round_robin_prefill
best_prefiller_index = None
max_hit_value = -1
for k, v in cache_hit_status.items():
if v > max_hit_value:
best_prefiller_index = index_map[k]
max_hit_value = v
return app.state.prefill_clients[best_prefiller_index]
async def get_tokenid(client_info: dict, req_data: dict, request_id: str):
req_data = req_data.copy()
req_data["stream"] = False
req_data["max_tokens"] = 1
if "stream_options" in req_data:
del req_data["stream_options"]
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id
}
response = await client_info['client'].post("/tokenize",
json=req_data,
headers=headers)
response.raise_for_status()
token_id = response.json()["tokens"]
return token_id
async def send_request_to_service(client_info: dict, endpoint: str,
req_data: dict, request_id: str):
"""
Send a request to a service using a client from the pool.
"""
req_data = req_data.copy()
req_data['kv_transfer_params'] = {
"do_remote_decode": True,
"do_remote_prefill": False,
"remote_engine_id": None,
"remote_block_ids": None,
"remote_host": None,
"remote_port": None
}
req_data["stream"] = False
req_data["max_tokens"] = 1
if "max_completion_tokens" in req_data:
req_data["max_completion_tokens"] = 1
if "stream_options" in req_data:
del req_data["stream_options"]
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id
}
logger.debug(f"req_data: {req_data}")
response = await client_info['client'].post(endpoint,
json=req_data,
headers=headers)
response.raise_for_status()
return response
async def stream_service_response(client_info: dict, endpoint: str,
req_data: dict, request_id: str):
"""
Asynchronously stream response from a service using a client from the pool.
"""
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id
}
async with client_info['client'].stream("POST",
endpoint,
json=req_data,
headers=headers) as response:
response.raise_for_status()
async for chunk in response.aiter_bytes():
yield chunk
async def _handle_completions(api: str, request: Request):
try:
req_data = await request.json()
request_id = str(uuid.uuid4())
# select tokenizer client in round-robin fashion
remote_tokenizer_client_info = get_next_client(request.app, 'prefill')
token_ids = await get_tokenid(remote_tokenizer_client_info, req_data, request_id)
# choice best cache hit prefill instance
prefill_client_info = await get_best_prefiller(request.app, token_ids, remote_tokenizer_client_info, req_data)
response = await send_request_to_service(prefill_client_info, api,
req_data, request_id)
# Extract the needed fields
response_json = response.json()
kv_transfer_params = response_json.get('kv_transfer_params', {})
if kv_transfer_params:
req_data["kv_transfer_params"] = kv_transfer_params
# Get the next decode client in round-robin fashion
decode_client_info = get_next_client(request.app, 'decode')
logger.debug("Using %s %s", prefill_client_info, decode_client_info)
# Stream response from decode service
async def generate_stream():
async for chunk in stream_service_response(decode_client_info,
api,
req_data,
request_id=request_id):
yield chunk
return StreamingResponse(generate_stream(),
media_type="application/json")
except Exception as e:
import sys
import traceback
exc_info = sys.exc_info()
logger.error("Error occurred in disagg prefill proxy server"
f" - {api} endpoint")
logger.error(e)
logger.error("".join(traceback.format_exception(*exc_info)))
raise
@app.post("/v1/completions")
async def handle_completions(request: Request):
return await _handle_completions("/v1/completions", request)
@app.post("/v1/chat/completions")
async def handle_chat_completions(request: Request):
return await _handle_completions("/v1/chat/completions", request)
@app.get("/healthcheck")
async def healthcheck():
"""Simple endpoint to check if the server is running."""
return {
"status": "ok",
"prefill_instances": len(app.state.prefill_clients),
"decode_instances": len(app.state.decode_clients)
}
if __name__ == '__main__':
global global_args
global_args = parse_args()
import uvicorn
uvicorn.run(app, host=global_args.host, port=global_args.port)

View File

@ -0,0 +1,32 @@
{
"kvevent_instance":
{
"vllm-prefill-node1":
{
"endpoint": "tcp://127.0.0.1:5557",
"replay_endpoint": "tcp://127.0.0.1:5558",
"type": "vLLM",
"modelname": "qwen2.5",
"lora_name": "xx-adapter",
"tenant_id": "default",
"instance_id": "vllm-prefill-node1",
"block_size": 128,
"dp_rank": 0,
"additionalsalt": ""
},
"mooncake":
{
"endpoint": "tcp://127.0.0.1:6667",
"replay_endpoint": "tcp://127.0.0.1:6668",
"type": "Mooncake",
"modelname": "qwen2.5",
"lora_name": "xx-adapter",
"tenant_id": "default",
"instance_id": "vllm-prefill-node1",
"block_size": 128,
"dp_rank": 0,
"additionalsalt": ""
}
},
"http_server_port": 13333
}

View File

@ -0,0 +1,150 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_H
#define MOONCAKE_KV_EVENT_H
#include "replica.h"
#include <msgpack.hpp>
#include <atomic>
#include <vector>
#include <memory>
#include <string>
#include <chrono>
namespace mooncake {
/**
* @brief Base class for KV cache events
* Defines the interface that all KV cache events must implement for msgpack
* serialization Uses the Abstract Base Class pattern to provide a uniform event
* handling interface
*/
class KVCacheEvent {
public:
virtual ~KVCacheEvent() = default;
/**
* @brief Serialize the event to msgpack format
* @param pk Reference to msgpack packer
* Derived classes must implement this method to provide type-specific
* serialization logic
*/
virtual void pack(msgpack::packer<msgpack::sbuffer>& pk) const = 0;
/**
* @brief Get the event type identifier
* @return String view of the event type
* Used for event type identification during deserialization
*/
virtual std::string_view type_tag() const = 0;
};
/**
* @brief Block update event
* Used for replica management operations within Mooncake-Store:
* 1. Replica increment
* 2. Replica removal
* 3. Replica migration
*/
class BlockUpdateEvent : public KVCacheEvent {
public:
std::string mooncake_key;
std::vector<Replica::Descriptor> replicas;
static constexpr size_t kFieldCount{3};
BlockUpdateEvent(std::string key,
std::vector<Replica::Descriptor> replica_list)
: mooncake_key(std::move(key)), replicas(std::move(replica_list)) {}
void pack(msgpack::packer<msgpack::sbuffer>& pk) const override {
pk.pack_array(kFieldCount);
pk.pack(type_tag());
pk.pack(mooncake_key);
pk.pack_array(replicas.size());
for (const auto& replica : replicas) {
if (replica.is_memory_replica()) {
pk.pack_array(2);
pk.pack("memory");
pk.pack(replica.get_memory_descriptor()
.buffer_descriptor.transport_endpoint_);
} else if (replica.is_disk_replica()) {
pk.pack_array(2);
pk.pack("disk");
pk.pack(replica.get_disk_descriptor().file_path);
} else if (replica.is_local_disk_replica()) {
pk.pack_array(2);
pk.pack("local_disk");
pk.pack(replica.get_local_disk_descriptor().transport_endpoint);
} else {
throw std::runtime_error(
"Unknown replica type in BlockUpdateEvent");
}
}
}
std::string_view type_tag() const override { return "BlockUpdateEvent"; }
};
/**
* @brief Remove all event
* Special event that instructs receivers to clear all cached data
*/
class RemoveAllEvent : public KVCacheEvent {
public:
static constexpr size_t kFieldCount{1};
RemoveAllEvent() = default;
void pack(msgpack::packer<msgpack::sbuffer>& pk) const override {
pk.pack_array(kFieldCount);
pk.pack(type_tag());
}
std::string_view type_tag() const override { return "RemoveAllEvent"; }
};
/**
* @brief Event batch
* Packages multiple events into a batch for transmission to improve network
* efficiency Includes timestamp for receivers to process events in
* chronological order
*/
class EventBatch {
public:
double ts;
std::vector<std::shared_ptr<KVCacheEvent>> events;
static constexpr size_t kFieldCount{2};
EventBatch(std::vector<std::shared_ptr<KVCacheEvent>> evts)
: ts(get_current_time()), events(std::move(evts)) {}
// [ts, events]
msgpack::sbuffer serialize() const {
msgpack::sbuffer buffer;
msgpack::packer<msgpack::sbuffer> pk(buffer);
pk.pack_array(kFieldCount);
pk.pack(ts);
pk.pack_array(events.size());
for (const auto& event : events) {
event->pack(pk);
}
return buffer;
}
private:
static double get_current_time() {
auto now = std::chrono::system_clock::now();
return std::chrono::duration<double>(now.time_since_epoch()).count();
}
};
} // namespace mooncake
#endif // MOONCAKE_KV_EVENT_H

View File

@ -0,0 +1,220 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_CONSUMER_H
#define MOONCAKE_KV_EVENT_CONSUMER_H
#include "kv_event/kv_event.hpp"
#include "kv_event/kv_event_types.h"
#include <zmq.hpp>
#include <zmq_addon.hpp>
#include <future>
#include <memory>
#include <optional>
#include <algorithm>
namespace mooncake {
class KVEventSystem;
/**
* @brief ZeroMQ event consumer
* Consumes events from the event queue and publishes them via ZeroMQ
* Supports event batch processing and replay functionality
*/
class KVEventConsumer {
public:
/**
* @brief Consumer configuration structure
*/
struct Config {
// Network endpoint configuration
std::string endpoint{"tcp://*:19997"};
std::optional<std::string> replay_endpoint = std::nullopt;
// Performance configuration
size_t buffer_steps{10000}; // Replay buffer size
int hwm{100000}; // ZeroMQ high-water mark
std::chrono::milliseconds send_interval{0};
// Batching configuration
size_t max_batch_size{50}; // Maximum event batch size
std::chrono::milliseconds pop_timeout{100};
// Message configuration
std::string topic{"mooncake"};
// Auto-port switching configuration
bool auto_port{true}; // Whether to enable auto-port switching
size_t max_port_attempts{10};
};
/**
* @brief Consumer statistics
*/
struct Stats {
size_t total_events{0};
size_t total_batches{0};
size_t failed_events{0};
size_t replay_requests{0};
double success_rate{0.0};
double events_per_batch{0.0};
/**
* @brief Check if any events have been processed
* @return true if at least one event has been received, false otherwise
*/
bool has_data() const { return total_events > 0; }
void calculate_derived_metrics();
friend std::ostream& operator<<(std::ostream& os, const Stats& stats);
};
/**
* @brief Constructor
* @param event_queue Shared event queue
* @param config Consumer configuration
* Initializes ZeroMQ context, publishing thread and configures the shared
* queue
*/
explicit KVEventConsumer(const std::shared_ptr<KVEventQueue> event_queue,
const Config& config);
/**
* @brief Destructor
* Automatically calls shutdown() to ensure resource release
*/
~KVEventConsumer();
// Prohibit copy and move operations
KVEventConsumer(const KVEventConsumer&) = delete;
KVEventConsumer& operator=(const KVEventConsumer&) = delete;
KVEventConsumer(KVEventConsumer&&) = delete;
KVEventConsumer& operator=(KVEventConsumer&&) = delete;
/**
* @brief Gracefully shutdown the consumer
* 1. Stop accepting new events
* 2. Stop publishing thread
* 3. Close all network connections
*/
void shutdown();
/**
* @brief Get current statistics
* @return Stats object containing all performance metrics
* Note: Frequent calls may impact performance
*/
Stats get_stats() const;
/**
* @brief Check if consumer is running
* @return true if consumer is active, false if shutdown
*/
bool is_running() const noexcept {
return running_.load(std::memory_order_acquire);
};
private:
/**
* @brief Replay buffer entry
* Stores recently transmitted events with their sequence numbers to support
* replay for late subscribers
*/
struct ReplayEntry {
uint64_t seq; // Sequence number for ordering
msgpack::sbuffer payload; // Serialized event data
ReplayEntry(uint64_t s, msgpack::sbuffer p)
: seq(s), payload(std::move(p)) {}
};
/**
* @brief Thread-specific resources
* Contains all resources required by the publishing thread to ensure thread
* safety
*/
struct ThreadResources {
std::unique_ptr<zmq::socket_t> pub_socket;
std::unique_ptr<zmq::socket_t> replay_socket;
std::deque<ReplayEntry> replay_buffer;
uint64_t next_seq = 0;
/**
* @brief Constructor
* @param ctx ZeroMQ context
* @param config Consumer configuration
*/
explicit ThreadResources(zmq::context_t& ctx, const Config& config);
};
/**
* @brief Publishing thread main function
* @param stop_token Stop token for cooperative thread cancellation
* Continuous event processing workflow:
* 1. Check and service replay requests
* 2. Wait for configured send interval
* 3. Batch events (up to max_batch_size)
* 4. Serialize batches and send via ZeroMQ
* 5. Update replay buffer
* 6. Fulfill promises for successfully transmitted events
*/
void publisher_thread(std::stop_token stop_token);
/**
* @brief Set up ZeroMQ sockets
* @param resources Thread resources reference
* @throws std::runtime_error if socket setup fails
* Performs automatic port selection based on configuration
*/
void setup_sockets(ThreadResources& resources);
/**
* @brief Service replay requests
* @param resources Thread resources
* Handles replay requests:
* 1. Reads requested starting sequence number
* 2. Sends all buffered events with equal or higher sequence numbers
* 3. Sends an end marker
* Processes only one request per call to avoid blocking the main publishing
* loop
*/
void service_replay(ThreadResources& resources);
/**
* @brief Handle errors during publishing operations
* @param e Exception that was caught
* @param context Description of where the error occurred
* Logs error details and implements brief sleep to prevent tight error
* loops
*/
void handle_error(const std::exception& e, const std::string& context);
// Graceful shutdown timeout in seconds
static constexpr double SHUTDOWN_TIMEOUT = 2.0;
// Magic sequence number indicating end of replay transmission
static constexpr std::array<uint8_t, 8> END_SEQ = {0xFF, 0xFF, 0xFF, 0xFF,
0xFF, 0xFF, 0xFF, 0xFF};
Config config_;
zmq::context_t context_;
std::shared_ptr<KVEventQueue> event_queue_;
std::jthread publisher_thread_; // Publishing thread
std::atomic<bool> running_{false};
std::atomic<size_t> total_events_{0};
std::atomic<size_t> total_batches_{0};
std::atomic<size_t> failed_events_{0};
std::atomic<size_t> replay_requests_{0};
}; // class KVEventConsumer
} // namespace mooncake
#endif // MOONCAKE_KV_EVENT_CONSUMER_H

View File

@ -0,0 +1,159 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_PRODUCER_H
#define MOONCAKE_KV_EVENT_PRODUCER_H
#include "kv_event/kv_event.hpp"
#include "kv_event/kv_event_types.h"
#include "thread_pool.h"
#include <future>
#include <memory>
namespace mooncake {
class KVEventSystem;
/**
* @brief Event producer
* Responsible for creating events and asynchronously pushing them to the event
* queue
*/
class KVEventProducer {
public:
/**
* @brief Producer configuration structure
*/
struct Config {
size_t enqueue_thread_pool_size{
1}; // Number of enqueue thread pool workers
std::chrono::milliseconds enqueue_timeout{
10}; // Enqueue timeout duration
size_t enqueue_max_retries{10}; // Maximum enqueue retry attempts
};
/**
* @brief Producer statistics
*/
struct Stats {
size_t events_created{0};
size_t enqueue_failed{0};
double success_rate{0.0};
size_t update_event{0};
size_t remove_all_event{0};
bool has_data() const { return events_created > 0; };
void calculate_derived_metrics();
friend std::ostream& operator<<(std::ostream& os, const Stats& stats);
};
/**
* @brief Constructor
* @param event_queue Shared event queue
* @param config Producer configuration
* Initializes thread pool and configures the shared queue
*/
explicit KVEventProducer(const std::shared_ptr<KVEventQueue> event_queue,
const Config& config);
/**
* @brief Destructor
* Automatically calls shutdown() to ensure thread pool is properly closed
*/
~KVEventProducer();
// Prohibit copy and move operations
KVEventProducer(const KVEventProducer&) = delete;
KVEventProducer& operator=(const KVEventProducer&) = delete;
KVEventProducer(KVEventProducer&&) = delete;
KVEventProducer& operator=(KVEventProducer&&) = delete;
bool is_running() const noexcept {
return running_.load(std::memory_order_acquire);
};
/**
* @brief Increment counter based on event type
* @tparam Event Event type
*/
template <DerivedFromKVCacheEvent Event>
void increment_event_count() {
if constexpr (std::is_same_v<Event, BlockUpdateEvent>) {
update_event_.fetch_add(1, std::memory_order_relaxed);
} else if constexpr (std::is_same_v<Event, RemoveAllEvent>) {
remove_all_event_.fetch_add(1, std::memory_order_relaxed);
} else {
static_assert(std::is_same_v<Event, void>,
"Unknown event type! Please add handling for this "
"event type in Stats::increment_event_count()");
}
events_created_.fetch_add(1, std::memory_order_relaxed);
};
/**
* @brief Publish an event
* @tparam Event Event type, must derive from KVCacheEvent
* @tparam Args Event constructor parameter types
* @param args Arguments passed to the event constructor
* @return std::future<bool> Future containing the publication result
* Creates and enqueues events asynchronously
*/
template <DerivedFromKVCacheEvent Event, typename... Args>
std::future<bool> publish(Args&&... args) {
if (!is_running()) {
auto promise = std::make_shared<std::promise<bool>>();
promise->set_value(false);
return promise->get_future();
}
auto event = std::make_shared<Event>(std::forward<Args>(args)...);
increment_event_count<Event>();
return publish_event_async(std::move(event));
};
/**
* @brief Asynchronously publish a pre-created event object
* @param event Shared pointer to event
* @return std::future<bool> Future containing the publication result
* Submits the event to the thread pool for enqueuing operation
*/
std::future<bool> publish_event_async(KVEventPtr event);
/**
* @brief Get the event queue
* @return Shared pointer to the event queue
*/
std::shared_ptr<KVEventQueue> get_queue() const { return event_queue_; };
/**
* @brief Get statistics
* @return Producer statistics
*/
Stats get_stats() const;
/**
* @brief Stop the producer
* Shuts down thread pool and stops accepting new events
*/
void shutdown();
private:
Config config_;
std::atomic<bool> running_{false};
std::shared_ptr<KVEventQueue> event_queue_;
std::unique_ptr<ThreadPool> enqueue_pool_; // Enqueuing thread pool
std::atomic<size_t> events_created_{0};
std::atomic<size_t> enqueue_failed_{0};
std::atomic<size_t> update_event_{0};
std::atomic<size_t> remove_all_event_{0};
};
} // namespace mooncake
#endif // MOONCAKE_KV_EVENT_PRODUCER_H

View File

@ -0,0 +1,47 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_PUBLISHER_CONFIG_H
#define MOONCAKE_KV_EVENT_PUBLISHER_CONFIG_H
#include <string>
#include <atomic>
#include <chrono>
#include <optional>
namespace mooncake {
struct KVEventPublisherConfig {
// Network endpoint configuration
std::string endpoint{"tcp://*:19997"};
std::optional<std::string> replay_endpoint = std::nullopt;
// Performance configuration
size_t buffer_steps{10000}; // Replay buffer size
int hwm{100000}; // ZeroMQ high-water mark
size_t max_queue_size{100000};
std::chrono::milliseconds send_interval{0};
// Batching configuration
size_t max_batch_size{50}; // Maximum event batch size
std::chrono::milliseconds batch_timeout{200};
std::chrono::milliseconds pop_timeout{100};
// Thread pool configuration
size_t enqueue_thread_pool_size{1};
std::chrono::milliseconds enqueue_timeout{10};
size_t enqueue_max_retries = 10;
// Message configuration
std::string topic{"mooncake"};
// Auto-port switching configuration
bool auto_port{true}; // Whether to enable auto-port switching
size_t max_port_attempts{10};
// Validate configuration validity
bool validate() const noexcept;
};
} // namespace mooncake
#endif // MOONCAKE_KV_EVENT_PUBLISHER_CONFIG_H

View File

@ -0,0 +1,159 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_SYSTEM_H
#define MOONCAKE_KV_EVENT_SYSTEM_H
#include "kv_event/kv_event.hpp"
#include "kv_event/kv_event_publisher_config.h"
#include "kv_event/kv_event_types.h"
#include <future>
#include <atomic>
#include <memory>
namespace mooncake {
class KVEventProducer;
class KVEventConsumer;
/**
* @brief Event system facade class
* Provides a simplified interface for using the event system's
* publish/subscribe functionality Implements the Facade pattern to hide complex
* interactions between producers and consumers
*/
class KVEventSystem {
public:
/**
* @brief Constructor
* @param config System configuration parameters
* Initializes event queue, producer, and consumer, and starts the event
* processing pipeline
*/
explicit KVEventSystem(
const KVEventPublisherConfig& config = KVEventPublisherConfig{});
/**
* @brief Destructor
* Automatically calls shutdown() to ensure proper resource release
*/
~KVEventSystem();
// Disable copy and move operations to ensure singleton-like behavior
KVEventSystem(const KVEventSystem&) = delete;
KVEventSystem& operator=(const KVEventSystem&) = delete;
KVEventSystem(KVEventSystem&&) = delete;
KVEventSystem& operator=(KVEventSystem&&) = delete;
/**
* @brief Stop the event system
* Gracefully shuts down producer and consumer, ensuring all events are
* processed
*/
void shutdown();
/**
* @brief Check if the system is running
* @return System running status
*/
bool is_running() const noexcept {
return running_.load(std::memory_order_acquire);
};
/**
* @brief Generic event publishing interface
* @tparam Event Event type, must derive from KVCacheEvent
* @tparam Args Event constructor parameter types
* @param args Arguments passed to the event constructor
* @return Future containing the publication result (true for success, false
* for failure)
*/
template <DerivedFromKVCacheEvent Event, typename... Args>
std::future<bool> publish(Args&&... args) {
return producer_->publish<Event>(std::forward<Args>(args)...);
};
/**
* @brief Queue statistics structure
* Monitors event queue status for performance tuning and capacity planning
*/
struct QueueStats {
size_t queue_remain_events =
0; // Number of pending events in the queue
size_t queue_capacity = 0; // Total queue capacity
friend std::ostream& operator<<(std::ostream& os,
const QueueStats& stats);
};
/**
* @brief Complete system statistics
* Aggregates statistics from producer, consumer, and queue
*/
struct Stats {
KVEventProducer::Stats producer_stats;
KVEventConsumer::Stats consumer_stats;
QueueStats event_queue_stats;
double success_rate{0.0};
bool has_data() const;
void calculate_derived_metrics();
friend std::ostream& operator<<(std::ostream& os, const Stats& stats);
};
/**
* @brief Get complete system statistics
* @return Stats object containing all statistical information
* Note: Frequent calls may impact performance
*/
Stats get_stats() const;
/**
* @brief Get producer statistics
* @return Producer statistics
*/
KVEventProducer::Stats get_producer_stats() const {
return producer_->get_stats();
};
/**
* @brief Get consumer statistics
* @return Consumer statistics
*/
KVEventConsumer::Stats get_consumer_stats() const {
return consumer_->get_stats();
};
/**
* @brief Get queue statistics
* @return Queue statistics
*/
QueueStats get_queue_stats() const;
/**
* @brief Get internal producer instance
* @return Shared pointer to event producer
* Used for advanced usage to directly access producer interface
*/
std::shared_ptr<KVEventProducer> get_producer() const { return producer_; };
/**
* @brief Get internal consumer instance
* @return Shared pointer to event consumer
* Used for advanced usage to directly access consumer interface
*/
std::shared_ptr<KVEventConsumer> get_consumer() const { return consumer_; };
private:
KVEventPublisherConfig config_;
std::shared_ptr<KVEventQueue> event_queue_;
std::shared_ptr<KVEventProducer> producer_;
std::shared_ptr<KVEventConsumer> consumer_;
std::atomic<bool> running_{false};
};
} // namespace mooncake
#endif // MOONCAKE_KV_EVENT_SYSTEM_H

View File

@ -0,0 +1,50 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_TYPES_H
#define MOONCAKE_KV_EVENT_TYPES_H
#include "kv_event/kv_event.hpp"
#include "thread_safe_queue.h"
#include <future>
#include <memory>
#include <optional>
#include <type_traits>
namespace mooncake {
/**
* @brief Type trait concept for events derived from KVCacheEvent
* Ensures compile-time type checking for event types used in templates
*/
template <typename Event>
concept DerivedFromKVCacheEvent = std::is_base_of_v<KVCacheEvent, Event>;
/**
* @brief Type alias for KVCacheEvent shared pointer
* Standardized pointer type for event objects throughout the system
*/
using KVEventPtr = std::shared_ptr<KVCacheEvent>;
/**
* @brief Structure representing an item in the event queue
* Combines an event with its associated promise for notification of
* transmission completion The promise is fulfilled by the consumer when the
* event is successfully published
*/
struct QueuedItem {
KVEventPtr event; // Shared pointer to the event
std::shared_ptr<std::promise<bool>>
promise; // Promise to notify publishing result
};
/**
* @brief Type alias for thread-safe event queue
* Queue element is optional to support sentinel values for shutdown signaling
* Uses ThreadSafeQueue implementation with configurable capacity
*/
using KVEventQueue = ThreadSafeQueue<std::optional<QueuedItem>>;
} // namespace mooncake
#endif // MOONCAKE_KV_EVENT_TYPES_H

View File

@ -0,0 +1,13 @@
#pragma once
#ifndef MOONCAKE_KV_EVENT_PUBLISHER_H
#define MOONCAKE_KV_EVENT_PUBLISHER_H
#include "kv_event/kv_event.hpp"
#include "kv_event/kv_event_types.h"
#include "kv_event/kv_event_publisher_config.h"
#include "kv_event/kv_event_producer.h"
#include "kv_event/kv_event_consumer.h"
#include "kv_event/kv_event_system.h"
#endif // MOONCAKE_KV_EVENT_PUBLISHER_H

View File

@ -7,6 +7,7 @@
#include "config_helper.h"
#include "types.h"
#include "kv_event/kv_event_publisher_config.h"
namespace mooncake {
@ -70,6 +71,10 @@ struct MasterConfig {
std::string cxl_path;
size_t cxl_size;
bool enable_cxl = false;
// KV Event Publisher configuration
bool enable_kv_event_publish = false;
KVEventPublisherConfig kv_event_publisher_config{};
};
class MasterServiceSupervisorConfig {
@ -125,6 +130,11 @@ class MasterServiceSupervisorConfig {
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
// KV Event Publisher configuration
bool enable_kv_event_publish = false;
KVEventPublisherConfig kv_event_publisher_config{};
MasterServiceSupervisorConfig() = default;
// From MasterConfig
@ -183,6 +193,13 @@ class MasterServiceSupervisorConfig {
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
// KV Event Publisher configuration
enable_kv_event_publish = config.enable_kv_event_publish;
if (config.enable_kv_event_publish) {
kv_event_publisher_config = config.kv_event_publisher_config;
}
validate();
}
@ -223,6 +240,9 @@ class MasterServiceSupervisorConfig {
if (!rpc_thread_num.IsSet()) {
throw std::runtime_error("rpc_thread_num is not set");
}
if (enable_kv_event_publish && !kv_event_publisher_config.validate()) {
throw std::runtime_error("Invalid KVEventPublisher configuration");
}
}
};
@ -275,6 +295,11 @@ class WrappedMasterServiceConfig {
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
// KV Event Publisher configuration
bool enable_kv_event_publish = false;
KVEventPublisherConfig kv_event_publisher_config{};
WrappedMasterServiceConfig() = default;
// From MasterConfig
@ -343,6 +368,12 @@ class WrappedMasterServiceConfig {
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
// KV Event Publisher configuration
enable_kv_event_publish = config.enable_kv_event_publish;
if (config.enable_kv_event_publish) {
kv_event_publisher_config = config.kv_event_publisher_config;
}
}
// From MasterServiceSupervisorConfig, enable_ha is set to true
@ -391,6 +422,12 @@ class WrappedMasterServiceConfig {
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
// KV Event Publisher configuration
enable_kv_event_publish = config.enable_kv_event_publish;
if (config.enable_kv_event_publish) {
kv_event_publisher_config = config.kv_event_publisher_config;
}
}
};
@ -440,6 +477,10 @@ class MasterServiceConfigBuilder {
size_t cxl_size_ = DEFAULT_CXL_SIZE;
bool enable_cxl_ = false;
// KV Event Publisher configuration
bool enable_kv_event_publish_ = false;
KVEventPublisherConfig kv_event_publisher_config_{};
public:
MasterServiceConfigBuilder() = default;
@ -617,6 +658,18 @@ class MasterServiceConfigBuilder {
return *this;
}
// KV Event Publisher configuration
MasterServiceConfigBuilder& set_enable_kv_event_publish(bool enable) {
enable_kv_event_publish_ = enable;
return *this;
}
MasterServiceConfigBuilder& set_kv_event_publisher_config(
const KVEventPublisherConfig& config) {
kv_event_publisher_config_ = config;
return *this;
}
MasterServiceConfig build() const;
};
@ -674,6 +727,11 @@ class MasterServiceConfig {
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
// KV Event Publisher configuration
bool enable_kv_event_publish = false;
KVEventPublisherConfig kv_event_publisher_config{};
MasterServiceConfig() = default;
// From WrappedMasterServiceConfig
@ -723,6 +781,9 @@ class MasterServiceConfig {
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
// KV Event Publisher configuration
enable_kv_event_publish = config.enable_kv_event_publish;
kv_event_publisher_config = config.kv_event_publisher_config;
}
// Static factory method to create a builder
@ -771,6 +832,9 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const {
config.cxl_path = cxl_path_;
config.cxl_size = cxl_size_;
config.enable_cxl = enable_cxl_;
// KV Event Publisher configuration
config.enable_kv_event_publish = enable_kv_event_publish_;
config.kv_event_publisher_config = kv_event_publisher_config_;
return config;
}
@ -792,6 +856,9 @@ struct InProcMasterConfig {
std::optional<std::string> root_fs_dir;
std::optional<bool> enable_disk_eviction;
std::optional<uint64_t> quota_bytes;
// KV Event Publisher configuration
std::optional<bool> enable_kv_event_publish;
std::optional<KVEventPublisherConfig> kv_event_publisher_config;
};
// Builder class for InProcMasterConfig
@ -808,6 +875,10 @@ class InProcMasterConfigBuilder {
std::optional<std::string> root_fs_dir_ = std::nullopt;
std::optional<bool> enable_disk_eviction_ = std::nullopt;
std::optional<uint64_t> quota_bytes_ = std::nullopt;
// KV Event Publisher configuration
std::optional<bool> enable_kv_event_publish_ = std::nullopt;
std::optional<KVEventPublisherConfig> kv_event_publisher_config_ =
std::nullopt;
public:
InProcMasterConfigBuilder() = default;
@ -871,6 +942,18 @@ class InProcMasterConfigBuilder {
return *this;
}
// KV Event Publisher configuration
InProcMasterConfigBuilder& set_enable_kv_event_publish(bool enable) {
enable_kv_event_publish_ = enable;
return *this;
}
InProcMasterConfigBuilder& set_kv_event_publisher_config(
const KVEventPublisherConfig& config) {
kv_event_publisher_config_ = config;
return *this;
}
InProcMasterConfig build() const;
};
@ -888,6 +971,9 @@ inline InProcMasterConfig InProcMasterConfigBuilder::build() const {
config.root_fs_dir = root_fs_dir_;
config.enable_disk_eviction = enable_disk_eviction_;
config.quota_bytes = quota_bytes_;
// KV Event Publisher configuration
config.enable_kv_event_publish = enable_kv_event_publish_;
config.kv_event_publisher_config = kv_event_publisher_config_;
return config;
}

View File

@ -28,6 +28,7 @@
#include "replica.h"
#include "serialize/serializer_backend.h"
#include "task_manager.h"
#include "kv_event_publisher.h"
namespace mooncake {
// Forward declarations
@ -326,6 +327,24 @@ class MasterService {
*/
size_t GetKeyCount() const;
/**
* @brief Check if the event publisher is enabled.
* @return true if the publisher is enabled, false if disabled.
*/
bool IsPublisherEnabled() const { return enable_kv_event_publish; }
/**
* @brief Retrieve statistics from the event publisher.
* @return On success, returns KVEventSystem::Stats.
* On failure, returns an ErrorCode indicating the reason:
* - ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the publisher is
* not enabled
* - ErrorCode::INTERNAL_ERROR if publisher exists but statistics
* cannot be retrieved
*/
auto GetPublisherStats() const
-> tl::expected<KVEventSystem::Stats, ErrorCode>;
/**
* @brief Heartbeat from client
* @param client_id The uuid of the client
@ -696,6 +715,21 @@ class MasterService {
return segment_names;
}
std::vector<Replica::Descriptor> GetReplicasDescriptorList() {
if (!IsValid()) {
return {};
}
std::vector<Replica::Descriptor> replica_list;
replica_list.reserve(replicas_.size());
metadata.VisitReplicas(
&Replica::fn_is_completed, [&replica_list](const Replica& replica) {
replica_list.emplace_back(replica.get_descriptor());
});
return replica_list;
}
private:
// Use the accessors to visit and modify the replicas.
std::vector<Replica> replicas_;
@ -1084,6 +1118,10 @@ class MasterService {
// Task manager
ClientTaskManager task_manager_;
// KV Event Publisher configuration
bool enable_kv_event_publish;
std::unique_ptr<KVEventSystem> publisher;
};
} // namespace mooncake

View File

@ -1,6 +1,8 @@
// ThreadPool.h
#pragma once
#ifndef MOONCAKE_THREAD_POOL_H
#define MOONCAKE_THREAD_POOL_H
#include <vector>
#include <queue>
#include <functional>
@ -8,6 +10,7 @@
#include <mutex>
#include <condition_variable>
#include <atomic>
namespace mooncake {
/**
* @class ThreadPool
@ -68,3 +71,5 @@ void ThreadPool::enqueue(F&& f, Args&&... args) {
condition.notify_one(); ///< Wake one waiting worker
}
} // namespace mooncake
#endif // MOONCAKE_THREAD_POOL_H

View File

@ -0,0 +1,564 @@
#pragma once
#ifndef MOONCAKE_THREAD_SAFE_QUEUE_H
#define MOONCAKE_THREAD_SAFE_QUEUE_H
#include <deque>
#include <mutex>
#include <condition_variable>
#include <chrono>
#include <optional>
#include <atomic>
#include <vector>
#include <type_traits>
namespace mooncake {
/**
* @brief Thread-safe queue implementation with bounded capacity and shutdown
* mechanism
*
* This template class provides a producer-consumer queue that can be safely
* accessed from multiple threads. It supports both blocking and timeout-based
* operations, and includes a shutdown mechanism to gracefully terminate waiting
* threads.
*
* Optimized for MPSC (Multiple Producers, Single Consumer) scenarios with
* atomic size tracking to reduce lock contention.
*
* @note Type T must be nothrow move constructible or copy constructible
* for exception safety in batch operations. Uses std::deque as the
* underlying container to support efficient random access for
* peek operations.
*
* @tparam T Type of elements stored in the queue
*/
template <typename T>
class ThreadSafeQueue {
// Static assertion for exception safety requirements
static_assert(
std::is_nothrow_move_constructible_v<T> ||
std::is_copy_constructible_v<T>,
"Type T must be nothrow move constructible or copy constructible "
"for exception safety in batch operations");
public:
/**
* @brief Construct a new ThreadSafeQueue with specified maximum capacity
* @param max_size Maximum number of elements the queue can hold
*/
explicit ThreadSafeQueue(size_t max_size = 100000) : max_size_(max_size) {}
/**
* @brief Push an item into the queue (blocking)
* @param item Item to be pushed
* @return true if item was successfully pushed, false if queue is shutdown
*/
bool push(T item);
/**
* @brief Push an item into the queue with timeout
* @param item Item to be pushed
* @param timeout Maximum time to wait for push operation
* @return true if item was successfully pushed, false on timeout or
* shutdown
*/
bool push(T item, std::chrono::milliseconds timeout);
/**
* @brief Attempt to push an item into the queue without blocking
*
* This method attempts to push an item immediately without waiting.
* It returns immediately with the result of the attempt.
*
* @param item Item to be pushed
* @return true if item was successfully pushed, false if queue is full or
* shutdown
*/
bool try_push(T item);
/**
* @brief Pop an item from the queue (blocking)
* @return std::optional containing the popped item, or std::nullopt if
* queue is shutdown
*/
std::optional<T> pop();
/**
* @brief Pop an item from the queue with timeout
* @param timeout Maximum time to wait for pop operation
* @return std::optional containing the popped item, or std::nullopt on
* timeout or shutdown
*/
std::optional<T> pop(std::chrono::milliseconds timeout);
/**
* @brief Batch pop items from the queue
* @param max_batch_size Maximum number of items to pop
* @param timeout Maximum time to wait for at least one item
* @return Vector of popped items, std::nullopt if timeout or shutdown
*/
std::optional<std::vector<T>> pop_batch(size_t max_batch_size,
std::chrono::milliseconds timeout);
/**
* @brief Non-blocking batch pop
* @param max_batch_size Maximum number of items to pop
* @return Vector of popped items, empty if queue is empty
*/
std::optional<std::vector<T>> try_pop_batch(size_t max_batch_size);
/**
* @brief Get the first max_batch_size elements from the queue without
* removing them
*
* This method returns a copy of the first max_batch_size elements in the
* queue without popping. The operation is thread-safe and efficient due to
* the use of std::deque as the underlying container, which supports random
* access.
*
* In MPSC scenarios, producers are only blocked for the minimal time
* needed to create a copy of the required elements.
*
* @param max_batch_size Number of elements to retrieve. If
* max_batch_size=0, returns std::nullopt.
* @return std::optional<std::vector<T>> containing the first max_batch_size
* elements, or std::nullopt if the queue is empty
*/
std::optional<std::vector<T>> peek_batch(size_t max_batch_size) const;
/**
* @brief Get elements at specific positions in the queue without removing
* them
*
* This method returns copies of elements at the specified positions
* without removing them. Positions are 0-based from the front of the queue.
* The operation is thread-safe and efficient with std::deque.
*
* @param positions Vector of 0-based positions to peek
* @return std::vector<std::optional<T>> containing the elements at the
* requested positions, with nullopt for invalid positions
*/
std::vector<std::optional<T>> peek_at(
const std::vector<size_t>& positions) const;
/**
* @brief Check if the queue is empty
* @return true if queue contains no elements, false otherwise
*/
bool empty() const;
/**
* @brief Get approximate queue size without locking
*
* This method provides fast size approximation using atomic operations.
* In MPSC scenarios, this provides good accuracy with minimal contention.
*
* @return Approximate number of elements in the queue
*/
size_t size_approx() const {
// Use relaxed memory order as this is an approximate value
return approximate_size_.load(std::memory_order_acquire);
}
/**
* @brief Get the exact number of elements in the queue
*
* This method acquires a lock to get the exact size of the queue.
* It is safe to call this method after shutdown.
*
* @return Exact number of elements in the queue
*/
size_t size() const;
/**
* @brief Get the maximum capacity of the queue
*
* This method returns the maximum number of elements the queue can hold.
* The value is set during construction and remains constant throughout
* the lifetime of the queue.
*
* @return Maximum capacity of the queue
*/
size_t capacity() const { return max_size_; }
/**
* @brief Shutdown the queue and wake up all waiting threads
*
* After shutdown, all push operations will fail immediately, while pop
* operations will continue to return remaining elements until the queue is
* empty.
*/
void shutdown();
/**
* @brief Check if the queue has been shutdown
* @return true if queue is shutdown, false otherwise
*/
bool is_shutdown() const;
/**
* @brief Check if the queue is full
* @return true if queue is at maximum capacity, false otherwise
*/
bool is_full() const;
private:
// Helper methods for batch operations
template <bool UseCopy>
std::vector<T> batch_pop_impl(size_t max_batch_size, size_t original_size,
bool& was_full);
// Helper to check if we should proceed with push operation
bool should_proceed_with_push() const {
// For push operations, we should not proceed if shutdown
return !shutdown_.load(std::memory_order_acquire);
}
private:
mutable std::mutex mutex_;
std::condition_variable not_empty_;
std::condition_variable not_full_;
std::deque<T> deque_;
size_t max_size_;
std::atomic<bool> shutdown_{false};
std::atomic<size_t> approximate_size_{0};
};
// Implementation
template <typename T>
bool ThreadSafeQueue<T>::push(T item) {
if (!should_proceed_with_push()) {
return false;
}
std::unique_lock lock(mutex_);
not_full_.wait(lock, [this] {
return deque_.size() < max_size_ ||
shutdown_.load(std::memory_order_acquire);
});
if (shutdown_.load(std::memory_order_acquire)) {
return false;
}
bool was_empty = deque_.empty();
deque_.push_back(std::move(item));
approximate_size_.fetch_add(1, std::memory_order_release);
lock.unlock();
if (was_empty) {
not_empty_.notify_one();
}
return true;
}
template <typename T>
bool ThreadSafeQueue<T>::push(T item, std::chrono::milliseconds timeout) {
if (!should_proceed_with_push()) {
return false;
}
std::unique_lock lock(mutex_);
if (!not_full_.wait_for(lock, timeout, [this] {
return deque_.size() < max_size_ ||
shutdown_.load(std::memory_order_acquire);
})) {
return false;
}
if (shutdown_.load(std::memory_order_acquire)) {
return false;
}
bool was_empty = deque_.empty();
deque_.push_back(std::move(item));
approximate_size_.fetch_add(1, std::memory_order_release);
lock.unlock();
if (was_empty) {
not_empty_.notify_one();
}
return true;
}
template <typename T>
bool ThreadSafeQueue<T>::try_push(T item) {
if (!should_proceed_with_push()) {
return false;
}
if (approximate_size_.load(std::memory_order_acquire) >= max_size_) {
return false;
}
std::unique_lock lock(mutex_, std::try_to_lock);
if (!lock.owns_lock() || shutdown_.load(std::memory_order_acquire)) {
return false;
}
if (deque_.size() >= max_size_) {
return false;
}
bool was_empty = deque_.empty();
deque_.push_back(std::move(item));
approximate_size_.fetch_add(1, std::memory_order_release);
lock.unlock();
if (was_empty) {
not_empty_.notify_one();
}
return true;
}
template <typename T>
std::optional<T> ThreadSafeQueue<T>::pop() {
std::unique_lock lock(mutex_);
not_empty_.wait(lock, [this] {
return !deque_.empty() || shutdown_.load(std::memory_order_acquire);
});
if (deque_.empty()) {
return std::nullopt;
}
bool was_full = (deque_.size() == max_size_);
T item = std::move(deque_.front());
deque_.pop_front();
approximate_size_.fetch_sub(1, std::memory_order_release);
lock.unlock();
if (was_full) {
not_full_.notify_one();
}
return item;
}
template <typename T>
std::optional<T> ThreadSafeQueue<T>::pop(std::chrono::milliseconds timeout) {
std::unique_lock lock(mutex_);
if (!not_empty_.wait_for(lock, timeout, [this] {
return !deque_.empty() || shutdown_.load(std::memory_order_acquire);
})) {
return std::nullopt;
}
if (deque_.empty()) {
return std::nullopt;
}
bool was_full = (deque_.size() == max_size_);
T item = std::move(deque_.front());
deque_.pop_front();
approximate_size_.fetch_sub(1, std::memory_order_release);
lock.unlock();
if (was_full) {
not_full_.notify_one();
}
return item;
}
template <typename T>
template <bool UseCopy>
std::vector<T> ThreadSafeQueue<T>::batch_pop_impl(size_t max_batch_size,
size_t original_size,
bool& was_full) {
const size_t count = std::min(max_batch_size, original_size);
was_full = (original_size == max_size_);
if (count == 0) {
return {};
}
std::vector<T> batch;
batch.reserve(count);
auto begin = deque_.begin();
auto end = begin;
std::advance(end, count);
if constexpr (UseCopy) {
batch.insert(batch.end(), begin, end);
} else {
batch.insert(batch.end(), std::make_move_iterator(begin),
std::make_move_iterator(end));
}
approximate_size_.fetch_sub(count, std::memory_order_release);
deque_.erase(begin, end);
return batch;
}
template <typename T>
std::optional<std::vector<T>> ThreadSafeQueue<T>::pop_batch(
size_t max_batch_size, std::chrono::milliseconds timeout) {
if (max_batch_size == 0) {
return std::nullopt;
}
std::unique_lock lock(mutex_);
if (!not_empty_.wait_for(lock, timeout, [this] {
return !deque_.empty() || shutdown_.load(std::memory_order_acquire);
})) {
return std::nullopt;
}
if (deque_.empty()) {
return std::nullopt;
}
const size_t original_size = deque_.size();
bool was_full = false;
std::vector<T> batch;
if constexpr (std::is_nothrow_move_constructible_v<T>) {
batch = batch_pop_impl<false>(max_batch_size, original_size, was_full);
} else {
batch = batch_pop_impl<true>(max_batch_size, original_size, was_full);
}
lock.unlock();
if (was_full) {
not_full_.notify_all();
}
if (!batch.empty()) {
return batch;
}
return std::nullopt;
}
template <typename T>
std::optional<std::vector<T>> ThreadSafeQueue<T>::try_pop_batch(
size_t max_batch_size) {
if (max_batch_size == 0) {
return std::nullopt;
}
std::unique_lock lock(mutex_);
if (deque_.empty()) {
return std::nullopt;
}
const size_t original_size = deque_.size();
bool was_full = false;
std::vector<T> batch;
if constexpr (std::is_nothrow_move_constructible_v<T>) {
batch = batch_pop_impl<false>(max_batch_size, original_size, was_full);
} else {
batch = batch_pop_impl<true>(max_batch_size, original_size, was_full);
}
lock.unlock();
if (was_full) {
not_full_.notify_all();
}
if (!batch.empty()) {
return batch;
}
return std::nullopt;
}
template <typename T>
std::optional<std::vector<T>> ThreadSafeQueue<T>::peek_batch(
size_t max_batch_size) const {
if (max_batch_size == 0) {
return std::nullopt;
}
if (approximate_size_.load(std::memory_order_acquire) == 0) {
return std::nullopt;
}
std::unique_lock lock(mutex_);
if (deque_.empty()) {
return std::nullopt;
}
const size_t count = std::min(max_batch_size, deque_.size());
return std::make_optional<std::vector<T>>(deque_.begin(),
deque_.begin() + count);
}
template <typename T>
std::vector<std::optional<T>> ThreadSafeQueue<T>::peek_at(
const std::vector<size_t>& positions) const {
std::vector<std::optional<T>> result;
std::unique_lock lock(mutex_);
if (deque_.empty()) {
result.resize(positions.size());
return result;
}
const size_t deque_size = deque_.size();
result.reserve(positions.size());
for (size_t pos : positions) {
if (pos < deque_size) {
result.emplace_back(deque_[pos]);
} else {
result.emplace_back(std::nullopt);
}
}
return result;
}
template <typename T>
size_t ThreadSafeQueue<T>::size() const {
std::lock_guard lock(mutex_);
return deque_.size();
}
template <typename T>
bool ThreadSafeQueue<T>::empty() const {
if (approximate_size_.load(std::memory_order_acquire) > 0) {
return false;
}
std::lock_guard lock(mutex_);
return deque_.empty();
}
template <typename T>
void ThreadSafeQueue<T>::shutdown() {
shutdown_.store(true, std::memory_order_release);
{
std::lock_guard lock(mutex_);
not_empty_.notify_all();
not_full_.notify_all();
}
}
template <typename T>
bool ThreadSafeQueue<T>::is_shutdown() const {
return shutdown_.load(std::memory_order_acquire);
}
template <typename T>
bool ThreadSafeQueue<T>::is_full() const {
std::lock_guard lock(mutex_);
return deque_.size() >= max_size_;
}
} // namespace mooncake
#endif // MOONCAKE_THREAD_SAFE_QUEUE_H

View File

@ -1,6 +1,9 @@
# Find Python package
add_subdirectory(cachelib_memory_allocator)
find_package(cppzmq REQUIRED)
find_package(msgpack-cxx REQUIRED)
set(MOONCAKE_STORE_SOURCES
allocator.cpp
master_service.cpp
@ -34,6 +37,7 @@ set(MOONCAKE_STORE_SOURCES
utils/file_util.cpp
task_manager.cpp
local_hot_cache.cpp
kv_event_publisher.cpp
)
set(EXTRA_LIBS "")
@ -103,6 +107,8 @@ target_link_libraries(mooncake_store
gflags::gflags
${EXTRA_LIBS}
asio_shared
cppzmq
msgpack-cxx
PRIVATE
transfer_engine
)

View File

@ -0,0 +1,875 @@
#include "kv_event_publisher.h"
#include <zmq_addon.hpp>
#include <glog/logging.h>
#include <regex>
#include <random>
#include <chrono>
#include <algorithm>
#include <iostream>
namespace mooncake {
namespace {
enum class EndpointType { TCP, IPC, INPROC, UNKNOWN };
struct EndpointInfo {
EndpointType type = EndpointType::UNKNOWN;
std::string protocol;
std::string host;
int port = 0;
std::string path;
static EndpointInfo parse(const std::string& endpoint) {
EndpointInfo info;
if (endpoint.find("tcp://") == 0) {
info.type = EndpointType::TCP;
info.protocol = "tcp";
std::string rest = endpoint.substr(6);
size_t colon_pos = rest.find_last_of(':');
if (colon_pos != std::string::npos) {
info.host = rest.substr(0, colon_pos);
std::string port_str = rest.substr(colon_pos + 1);
try {
info.port = std::stoi(port_str);
} catch (...) {
info.port = 0;
}
}
} else if (endpoint.find("ipc://") == 0) {
info.type = EndpointType::IPC;
info.protocol = "ipc";
info.path = endpoint.substr(6);
} else if (endpoint.find("inproc://") == 0) {
info.type = EndpointType::INPROC;
info.protocol = "inproc";
info.path = endpoint.substr(9);
} else {
info.type = EndpointType::UNKNOWN;
}
return info;
}
std::string to_string() const {
switch (type) {
case EndpointType::TCP:
return protocol + "://" + host + ":" + std::to_string(port);
case EndpointType::IPC:
return protocol + "://" + path;
case EndpointType::INPROC:
return protocol + "://" + path;
default:
return "";
}
}
std::string to_string_with_port(int new_port) const {
if (type != EndpointType::TCP) {
return to_string();
}
return protocol + "://" + host + ":" + std::to_string(new_port);
}
};
std::pair<bool, std::string> smart_bind(zmq::socket_t& socket,
const std::string& endpoint,
const KVEventConsumer::Config& config) {
EndpointInfo info = EndpointInfo::parse(endpoint);
if (info.type == EndpointType::UNKNOWN) {
LOG(ERROR) << "Unknown endpoint type: " << endpoint;
return {false, endpoint};
}
try {
socket.bind(endpoint);
LOG(INFO) << "Bound to: " << endpoint;
return {true, endpoint};
} catch (const zmq::error_t& e) {
if (e.num() != EADDRINUSE) {
LOG(ERROR) << "Failed to bind to " << endpoint << ": " << e.what();
return {false, endpoint};
}
if (info.type != EndpointType::TCP || !config.auto_port) {
LOG(ERROR) << "Port in use and auto_port not enabled: " << endpoint;
return {false, endpoint};
}
LOG(WARNING) << "Port {" << info.port
<< "} was in use, trying random ports";
}
static constexpr std::array<int, 4> PORT_ERRORS = {
EADDRINUSE, // Address already in use
EACCES, // Permission denied
EADDRNOTAVAIL, // Address not available
EINVAL // Invalid argument
};
std::random_device rd;
std::mt19937_64 rng(rd());
std::uniform_int_distribution<int> dist(1024, 65535);
std::string random_endpoint;
random_endpoint.reserve(endpoint.size() + 10);
for (size_t attempts = 0; attempts < config.max_port_attempts; ++attempts) {
int random_port = dist(rng);
random_endpoint = info.to_string_with_port(random_port);
try {
socket.bind(random_endpoint);
LOG(WARNING) << "WARNING: Failed to bind to in-use port {"
<< info.port << "}"
<< ", successfully switched to port {" << random_port
<< "} (attempt " << (attempts + 1) << ")";
return {true, random_endpoint};
} catch (const zmq::error_t& e) {
bool is_port_unavailable = false;
for (int error_code : PORT_ERRORS) {
if (e.num() == error_code) {
is_port_unavailable = true;
break;
}
}
if (!is_port_unavailable) {
LOG(ERROR) << "Failed to bind to " << random_endpoint << ": "
<< e.what();
return {false, endpoint};
}
}
}
LOG(ERROR) << "Failed to find an available port after "
<< config.max_port_attempts << " random attempts";
return {false, endpoint};
}
} // namespace
// KVEventPublisherConfig Implementation
bool KVEventPublisherConfig::validate() const noexcept {
if (endpoint.empty() || topic.empty()) {
LOG(ERROR) << "Endpoint and topic cannot be empty";
return false;
}
auto is_valid_endpoint = [](const std::string& ep) -> bool {
if (ep.find("://") == std::string::npos) {
LOG(ERROR) << "Endpoint missing protocol: " << ep;
return false;
}
if (ep.find("tcp://") != 0 && ep.find("ipc://") != 0 &&
ep.find("inproc://") != 0) {
LOG(ERROR) << "Unsupported protocol in endpoint: " << ep;
return false;
}
return true;
};
if (!is_valid_endpoint(endpoint)) {
return false;
}
if (replay_endpoint.has_value()) {
if (!is_valid_endpoint(*replay_endpoint)) {
return false;
}
}
if (max_queue_size == 0 || max_queue_size > 10000000) {
LOG(ERROR) << "max_queue_size out of range: " << max_queue_size;
return false;
}
if (max_batch_size == 0 || max_batch_size > 100) {
LOG(ERROR) << "max_batch_size out of range: " << max_batch_size;
return false;
}
if (hwm < 0) {
LOG(ERROR) << "hwm cannot be negative: " << hwm;
return false;
}
if (buffer_steps == 0 || buffer_steps > 1000000) {
LOG(ERROR) << "buffer_steps out of range: " << buffer_steps;
return false;
}
if (max_port_attempts == 0 || max_port_attempts > 1000) {
LOG(ERROR) << "max_port_attempts out of range: " << max_port_attempts;
return false;
}
if (enqueue_max_retries == 0 || enqueue_max_retries > 1000) {
LOG(ERROR) << "enqueue_max_retries out of range: "
<< enqueue_max_retries;
return false;
}
if (batch_timeout.count() < 0 || pop_timeout.count() < 0 ||
enqueue_timeout.count() < 0) {
LOG(ERROR) << "Timeout values cannot be negative";
return false;
}
if (topic.size() > 255) {
LOG(ERROR) << "Topic too long: " << topic.size();
return false;
}
return true;
}
// KVEventSystem Implementation
KVEventSystem::KVEventSystem(const KVEventPublisherConfig& config)
: config_(config) {
if (!config_.validate()) {
throw std::runtime_error("Invalid KV Event System configuration");
}
if (config_.max_batch_size > 100) {
config_.max_batch_size = 100;
LOG(WARNING)
<< "KV Event System Config"
<< " max_batch_size cannot be negative or greater than 100;\n"
<< "KV Event System Config max_batch_size has been reset to: "
<< config_.max_batch_size;
}
if (config_.send_interval.count() < 0) {
config_.send_interval = std::chrono::milliseconds(0);
;
LOG(WARNING) << "KV Event System Config"
<< " send_interval has been reset to: "
<< config_.send_interval.count();
}
event_queue_ = std::make_shared<KVEventQueue>(config_.max_queue_size);
KVEventProducer::Config producer_config{
.enqueue_thread_pool_size = config_.enqueue_thread_pool_size,
.enqueue_timeout = config_.enqueue_timeout,
.enqueue_max_retries = config_.enqueue_max_retries};
KVEventConsumer::Config consumer_config{
.endpoint = config_.endpoint,
.replay_endpoint = config_.replay_endpoint,
.buffer_steps = config_.buffer_steps,
.hwm = config_.hwm,
.send_interval = config_.send_interval,
.max_batch_size = config_.max_batch_size,
.pop_timeout = config_.pop_timeout,
.topic = config_.topic,
.auto_port = config_.auto_port,
.max_port_attempts = config_.max_port_attempts,
};
producer_ =
std::make_shared<KVEventProducer>(event_queue_, producer_config);
consumer_ =
std::make_shared<KVEventConsumer>(event_queue_, consumer_config);
running_.store(true, std::memory_order_release);
LOG(INFO) << "KV Event Publish System started.";
LOG(INFO) << " Endpoint" << (config_.auto_port ? "(mutable)" : "") << ": "
<< config_.endpoint;
LOG(INFO) << " Max batch size: " << config_.max_batch_size;
LOG(INFO) << " Send interval: " << config_.send_interval.count() << "ms";
LOG(INFO) << " " << get_stats();
}
KVEventSystem::~KVEventSystem() {
if (!is_running()) {
return;
}
shutdown();
}
void KVEventSystem::shutdown() {
bool expected = true;
if (!running_.compare_exchange_strong(expected, false,
std::memory_order_release,
std::memory_order_relaxed)) {
return;
}
LOG(INFO) << "Shutting down KV Event System...";
producer_->shutdown();
event_queue_->shutdown();
consumer_->shutdown();
LOG(INFO) << "KV Event System shutdown complete: " << get_stats();
}
bool KVEventSystem::Stats::has_data() const {
return producer_stats.has_data();
}
void KVEventSystem::Stats::calculate_derived_metrics() {
auto succeed_events =
consumer_stats.total_events - consumer_stats.failed_events;
auto total_events = producer_stats.events_created;
success_rate =
total_events > 0 ? succeed_events * 100.0 / total_events : 0.0;
}
KVEventSystem::Stats KVEventSystem::get_stats() const {
Stats stats{.producer_stats = get_producer_stats(),
.consumer_stats = get_consumer_stats(),
.event_queue_stats = get_queue_stats()};
stats.calculate_derived_metrics();
return stats;
}
KVEventSystem::QueueStats KVEventSystem::get_queue_stats() const {
return QueueStats{
.queue_remain_events =
event_queue_->size(), // Expensive operation, avoid frequent calls
.queue_capacity = event_queue_->capacity()};
}
std::ostream& operator<<(std::ostream& os,
const KVEventSystem::QueueStats& stats) {
os << "Queue(pending/cap): " << stats.queue_remain_events << "/"
<< stats.queue_capacity;
return os;
}
std::ostream& operator<<(std::ostream& os, const KVEventSystem::Stats& stats) {
os << "KV Event System: Succ=";
if (stats.has_data()) {
os << std::fixed << std::setprecision(1) << stats.success_rate << "%";
} else {
os << "--/--";
}
os << " | " << stats.event_queue_stats << " | " << stats.producer_stats
<< " | " << stats.consumer_stats;
return os;
}
// KVEventProducer Implementation
KVEventProducer::KVEventProducer(
const std::shared_ptr<KVEventQueue> event_queue,
const Config& config = Config{})
: config_(config), event_queue_(event_queue) {
enqueue_pool_ =
std::make_unique<ThreadPool>(config_.enqueue_thread_pool_size);
running_.store(true, std::memory_order_release);
}
KVEventProducer::~KVEventProducer() {
if (is_running()) {
shutdown();
}
}
void KVEventProducer::shutdown() {
bool expected = true;
if (!running_.compare_exchange_strong(expected, false,
std::memory_order_release,
std::memory_order_relaxed)) {
return;
}
if (enqueue_pool_) {
LOG(INFO) << "Stopping enqueue thread pool...";
enqueue_pool_->stop();
}
}
KVEventProducer::Stats KVEventProducer::get_stats() const {
Stats stats{
.events_created = events_created_.load(std::memory_order_relaxed),
.enqueue_failed = enqueue_failed_.load(std::memory_order_relaxed),
.update_event = update_event_.load(std::memory_order_relaxed),
.remove_all_event = remove_all_event_.load(std::memory_order_relaxed),
};
stats.calculate_derived_metrics();
return stats;
}
void KVEventProducer::Stats::calculate_derived_metrics() {
success_rate =
has_data() ? (events_created - enqueue_failed) * 100.0 / events_created
: 0.0;
}
std::ostream& operator<<(std::ostream& os,
const KVEventProducer::Stats& stats) {
os << "SuccEqueue: ";
if (stats.has_data()) {
os << std::fixed << std::setprecision(1) << stats.success_rate << "%";
} else {
os << "--/--";
}
os << "(Total=" << stats.events_created << ", Fail=" << stats.enqueue_failed
<< ")";
os << " | Evt Types: Update=" << stats.update_event
<< ", RemoveAll=" << stats.remove_all_event;
return os;
}
std::future<bool> KVEventProducer::publish_event_async(
std::shared_ptr<KVCacheEvent> event) {
auto promise = std::make_shared<std::promise<bool>>();
auto future = promise->get_future();
QueuedItem item{std::move(event), promise};
auto task = [this, item = std::move(item), promise]() mutable {
if (!this->is_running()) {
promise->set_value(false);
enqueue_failed_.fetch_add(1, std::memory_order_relaxed);
return;
}
// Retry logic
for (size_t attempt = 0; attempt < config_.enqueue_max_retries;
++attempt) {
if (!this->is_running()) {
break;
}
// Attempt fast enqueue
if (event_queue_->try_push(item)) {
return; // Success, promise is set by the consumer
}
// Attempt enqueue with timeout
if (event_queue_->push(item, config_.enqueue_timeout)) {
return; // Success, promise is set by the consumer
}
// Retry interval
if (attempt < config_.enqueue_max_retries - 1) {
std::this_thread::sleep_for(std::chrono::microseconds(5));
}
}
LOG(ERROR) << "Failed to enqueue event after "
<< config_.enqueue_max_retries << " retries";
promise->set_value(false);
enqueue_failed_.fetch_add(1, std::memory_order_relaxed);
};
try {
enqueue_pool_->enqueue(std::move(task));
} catch (const std::exception& e) {
LOG(ERROR) << "Failed to enqueue task: " << e.what();
promise->set_value(false);
enqueue_failed_.fetch_add(1, std::memory_order_relaxed);
} catch (...) {
LOG(ERROR) << "Failed to enqueue task: unknown exception";
promise->set_exception(std::current_exception());
enqueue_failed_.fetch_add(1, std::memory_order_relaxed);
}
return future;
}
// KVEventConsumer Implementation
KVEventConsumer::KVEventConsumer(
const std::shared_ptr<KVEventQueue> event_queue,
const Config& config = Config{})
: config_(config), event_queue_(event_queue) {
context_ = zmq::context_t(1);
publisher_thread_ = std::jthread([this](std::stop_token token) {
this->publisher_thread(std::move(token));
});
running_.store(true, std::memory_order_release);
}
KVEventConsumer::~KVEventConsumer() {
if (is_running()) {
shutdown();
}
}
// Shutdown
void KVEventConsumer::shutdown() {
bool expected = true;
if (!running_.compare_exchange_strong(expected, false,
std::memory_order_release,
std::memory_order_relaxed)) {
return;
}
// Wait for consumer thread to finish
if (publisher_thread_.joinable()) {
publisher_thread_.request_stop();
publisher_thread_.join();
}
}
// Publisher thread
void KVEventConsumer::publisher_thread(std::stop_token stop_token) {
ThreadResources resources(context_, config_);
setup_sockets(resources);
auto last_send_time = std::chrono::steady_clock::now();
while (is_running() && !stop_token.stop_requested()) {
try {
// Check replay requests
if (resources.replay_socket) {
zmq::pollitem_t items[] = {
{*resources.replay_socket, 0, ZMQ_POLLIN, 0}};
if (zmq::poll(items, 1, std::chrono::milliseconds(0)) > 0) {
service_replay(resources);
}
}
if (config_.send_interval.count() > 0) {
auto now = std::chrono::steady_clock::now();
auto time_since_last_send =
std::chrono::duration_cast<std::chrono::milliseconds>(
now - last_send_time);
if (time_since_last_send < config_.send_interval) {
// Wait until next send window
auto wait_time =
config_.send_interval - time_since_last_send;
// Handle replay requests while waiting
if (resources.replay_socket) {
zmq::pollitem_t items[] = {
{*resources.replay_socket, 0, ZMQ_POLLIN, 0}};
int poll_timeout = static_cast<int>(wait_time.count());
if (zmq::poll(items, 1,
std::chrono::milliseconds(poll_timeout)) >
0) {
service_replay(resources);
}
} else {
std::this_thread::sleep_for(wait_time);
}
// Update last send time
last_send_time = std::chrono::steady_clock::now();
}
}
// Batch peek events from queue
auto batch_opt = event_queue_->peek_batch(config_.max_batch_size);
if (!batch_opt.has_value()) {
last_send_time = std::chrono::steady_clock::now();
continue;
}
auto batch_items = *batch_opt;
if (batch_items.empty()) {
// Timeout or empty queue, continue
last_send_time = std::chrono::steady_clock::now();
continue;
}
// Prepare events and promises
std::vector<std::shared_ptr<KVCacheEvent>> events;
std::vector<std::shared_ptr<std::promise<bool>>> promises;
events.reserve(batch_items.size());
promises.reserve(batch_items.size());
for (auto& item : batch_items) {
if (!item) continue; // Sentinel value
events.push_back(std::move(item->event));
promises.push_back(std::move(item->promise));
}
if (events.empty()) {
last_send_time = std::chrono::steady_clock::now();
continue;
}
total_events_.fetch_add(events.size(), std::memory_order_relaxed);
// Create batch and serialize
auto event_batch = std::make_shared<EventBatch>(std::move(events));
uint64_t seq = resources.next_seq++;
auto payload = event_batch->serialize();
total_batches_.fetch_add(1, std::memory_order_relaxed);
// Prepare ZeroMQ messages
std::vector<zmq::message_t> messages;
// Topic part
if (!config_.topic.empty()) {
messages.emplace_back(config_.topic.data(),
config_.topic.size());
} else {
messages.emplace_back();
}
// Sequence number part
uint64_t seq_be = htobe64(seq);
messages.emplace_back(&seq_be, sizeof(seq_be));
// Payload part
messages.emplace_back(payload.data(), payload.size());
// Send messages
try {
zmq::send_multipart(*resources.pub_socket, messages,
zmq::send_flags::dontwait);
// Store in replay buffer
if (resources.replay_buffer.size() >= config_.buffer_steps) {
resources.replay_buffer.pop_front();
}
resources.replay_buffer.emplace_back(seq, std::move(payload));
// Set all promises to success
for (auto& promise : promises) {
promise->set_value(true);
}
// consume the events of this batch from event_queue_
auto effective_timeout = config_.pop_timeout;
if (config_.send_interval.count() > 0) {
effective_timeout = std::min(config_.pop_timeout,
std::chrono::milliseconds(1));
}
event_queue_->pop_batch(batch_items.size(), effective_timeout);
} catch (const zmq::error_t& e) {
handle_error(e, "send_multipart");
for (auto& promise : promises) {
promise->set_value(false);
}
failed_events_.fetch_add(promises.size(),
std::memory_order_relaxed);
}
last_send_time = std::chrono::steady_clock::now();
} catch (const std::exception& e) {
handle_error(e, "publisher_thread");
}
}
}
// ThreadResources constructor
KVEventConsumer::ThreadResources::ThreadResources(zmq::context_t& ctx,
const Config& config) {
pub_socket = std::make_unique<zmq::socket_t>(ctx, zmq::socket_type::pub);
pub_socket->set(zmq::sockopt::sndhwm, config.hwm);
if (config.replay_endpoint) {
replay_socket =
std::make_unique<zmq::socket_t>(ctx, zmq::socket_type::router);
}
}
// Statistics
KVEventConsumer::Stats KVEventConsumer::get_stats() const {
Stats stats{
.total_events = total_events_.load(std::memory_order_relaxed),
.total_batches = total_batches_.load(std::memory_order_relaxed),
.failed_events = failed_events_.load(std::memory_order_relaxed),
.replay_requests = replay_requests_.load(std::memory_order_relaxed),
};
stats.calculate_derived_metrics();
return stats;
}
void KVEventConsumer::Stats::calculate_derived_metrics() {
if (has_data()) {
events_per_batch =
total_batches > 0
? static_cast<double>(total_events) / total_batches
: 0.0;
success_rate = (total_events - failed_events) * 100.0 / total_events;
} else {
events_per_batch = 0.0;
success_rate = 0.0;
}
}
std::ostream& operator<<(std::ostream& os,
const KVEventConsumer::Stats& stats) {
os << "Publish Evts: " << stats.total_events
<< " (Batch=" << stats.total_batches << ", Avg/Batch=";
if (stats.has_data() && stats.total_batches > 0) {
os << std::fixed << std::setprecision(1) << stats.events_per_batch;
} else {
os << "--/--";
}
os << ")"
<< " | Succ: ";
if (stats.has_data()) {
os << std::fixed << std::setprecision(1) << stats.success_rate << "%";
} else {
os << "--/--";
}
os << " (Fail: " << stats.failed_events << ")"
<< " | Replay: " << stats.replay_requests;
return os;
}
// Socket setup
void KVEventConsumer::setup_sockets(ThreadResources& resources) {
bool should_bind = (config_.endpoint.find('*') != std::string::npos ||
config_.endpoint.find("::") != std::string::npos ||
config_.endpoint.find("ipc://") == 0 ||
config_.endpoint.find("inproc://") == 0);
try {
if (should_bind) {
auto [success, bound_endpoint] =
smart_bind(*resources.pub_socket, config_.endpoint, config_);
if (!success) {
throw std::runtime_error("Failed to bind PUB socket");
}
// Update configuration
if (bound_endpoint != config_.endpoint) {
LOG(WARNING) << "Updated PUB endpoint from " << config_.endpoint
<< " to " << bound_endpoint;
config_.endpoint = bound_endpoint;
}
} else {
resources.pub_socket->connect(config_.endpoint);
LOG(INFO) << "Connected PUB socket to: " << config_.endpoint;
}
} catch (const zmq::error_t& e) {
throw std::runtime_error("Failed to setup PUB socket: " +
std::string(e.what()));
}
// Setup replay ROUTER socket
if (resources.replay_socket && config_.replay_endpoint) {
try {
auto [success, bound_endpoint] = smart_bind(
*resources.replay_socket, *config_.replay_endpoint, config_);
if (!success) {
throw std::runtime_error("Failed to bind ROUTER socket");
}
// Update configuration
if (bound_endpoint != *config_.replay_endpoint) {
LOG(WARNING)
<< "Updated replay endpoint from "
<< *config_.replay_endpoint << " to " << bound_endpoint;
config_.replay_endpoint = bound_endpoint;
}
} catch (const zmq::error_t& e) {
throw std::runtime_error("Failed to setup ROUTER socket: " +
std::string(e.what()));
}
}
}
// Replay service
void KVEventConsumer::service_replay(ThreadResources& resources) {
if (!resources.replay_socket) return;
try {
std::vector<zmq::message_t> frames;
if (!zmq::recv_multipart(*resources.replay_socket,
std::back_inserter(frames))) {
return;
}
if (frames.size() != 3) {
LOG(ERROR) << "Invalid replay request: " << frames.size()
<< " frames";
return;
}
zmq::message_t client_id_frame = std::move(frames[0]);
zmq::message_t empty_frame = std::move(frames[1]);
zmq::message_t start_seq_frame = std::move(frames[2]);
if (start_seq_frame.size() != 8) {
LOG(ERROR) << "Invalid replay sequence number size";
return;
}
replay_requests_.fetch_add(1, std::memory_order_relaxed);
uint64_t start_seq =
be64toh(*reinterpret_cast<const uint64_t*>(frames[2].data()));
for (const auto& entry : resources.replay_buffer) {
if (entry.seq >= start_seq) {
std::vector<zmq::message_t> reply;
zmq::message_t reply_client_id(client_id_frame.data(),
client_id_frame.size());
reply.push_back(std::move(reply_client_id));
reply.emplace_back();
uint64_t seq_be = htobe64(entry.seq);
reply.emplace_back(&seq_be, sizeof(seq_be));
reply.emplace_back(entry.payload.data(), entry.payload.size());
try {
zmq::send_multipart(*resources.replay_socket, reply,
zmq::send_flags::dontwait);
} catch (const std::exception& e) {
LOG(ERROR) << "Error sending replay event: " << e.what();
break;
}
}
}
std::vector<zmq::message_t> end_reply;
zmq::message_t end_client_id(client_id_frame.data(),
client_id_frame.size());
end_reply.push_back(std::move(end_client_id));
end_reply.emplace_back(); // ZMQ ROUTER-DEALER delimiter
end_reply.emplace_back(); // End marker's empty sequence part
end_reply.emplace_back(END_SEQ.data(),
END_SEQ.size()); // Payload is the end marker
zmq::send_multipart(*resources.replay_socket, end_reply,
zmq::send_flags::dontwait);
} catch (const std::exception& e) {
LOG(ERROR) << "Error in replay service: " << e.what();
}
}
// Error handling
void KVEventConsumer::handle_error(const std::exception& e,
const std::string& context) {
LOG(ERROR) << "Error in " << context << ": " << e.what();
std::this_thread::sleep_for(std::chrono::milliseconds(100));
}
} // namespace mooncake

View File

@ -142,6 +142,73 @@ DEFINE_string(cxl_path, mooncake::DEFAULT_CXL_PATH,
"DAX device path for CXL memory");
DEFINE_uint64(cxl_size, mooncake::DEFAULT_CXL_SIZE, "CXL memory size in bytes");
DEFINE_bool(enable_cxl, false, "Whether to enable CXL memory support");
// KV Event Publisher
DEFINE_bool(enable_kv_event_publish, false,
"Enable KV event publishing functionality");
DEFINE_string(kv_event_publisher_endpoint, "tcp://*:19997",
"Endpoint for KV event publisher (ZeroMQ address)");
DEFINE_string(kv_event_publisher_replay_endpoint, "",
"Optional replay endpoint for KV event publisher (ZeroMQ "
"address, empty means disabled)");
DEFINE_int32(kv_event_publisher_hwm, 100000,
"ZeroMQ high water mark for event publisher");
DEFINE_uint32(
kv_event_publisher_send_interval_ms, 0,
"Send interval in milliseconds for event publisher (0 = no delay)");
DEFINE_uint32(kv_event_publisher_max_batch_size, 50,
"Maximum batch size for event publishing");
DEFINE_bool(kv_event_publisher_auto_port, true,
"Enable automatic port switching for event publisher");
DEFINE_string(kv_event_publisher_topic, "mooncake",
"Topic for published events");
void InitKVEventPublisherConf(const mooncake::DefaultConfig& default_config,
mooncake::MasterConfig& master_config) {
default_config.GetString("kv_event_publisher_endpoint",
&master_config.kv_event_publisher_config.endpoint,
FLAGS_kv_event_publisher_endpoint);
std::string replay_endpoint = FLAGS_kv_event_publisher_replay_endpoint;
default_config.GetString("kv_event_publisher_replay_endpoint",
&replay_endpoint,
FLAGS_kv_event_publisher_replay_endpoint);
if (!replay_endpoint.empty()) {
master_config.kv_event_publisher_config.replay_endpoint =
replay_endpoint;
} else {
master_config.kv_event_publisher_config.replay_endpoint = std::nullopt;
}
int32_t hwm_value = FLAGS_kv_event_publisher_hwm;
default_config.GetInt32("kv_event_publisher_hwm", &hwm_value,
FLAGS_kv_event_publisher_hwm);
master_config.kv_event_publisher_config.hwm = hwm_value;
uint32_t send_interval_ms = FLAGS_kv_event_publisher_send_interval_ms;
default_config.GetUInt32("kv_event_publisher_send_interval_ms",
&send_interval_ms,
FLAGS_kv_event_publisher_send_interval_ms);
master_config.kv_event_publisher_config.send_interval =
std::chrono::milliseconds(static_cast<int64_t>(send_interval_ms));
uint32_t max_batch_size_temp = FLAGS_kv_event_publisher_max_batch_size;
default_config.GetUInt32("kv_event_publisher_max_batch_size",
&max_batch_size_temp,
FLAGS_kv_event_publisher_max_batch_size);
master_config.kv_event_publisher_config.max_batch_size =
static_cast<size_t>(max_batch_size_temp);
default_config.GetBool("kv_event_publisher_auto_port",
&master_config.kv_event_publisher_config.auto_port,
FLAGS_kv_event_publisher_auto_port);
default_config.GetString("kv_event_publisher_topic",
&master_config.kv_event_publisher_config.topic,
FLAGS_kv_event_publisher_topic);
}
void InitMasterConf(const mooncake::DefaultConfig& default_config,
mooncake::MasterConfig& master_config) {
// Initialize the master service configuration from the default config
@ -264,6 +331,77 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config,
default_config.GetUInt32("max_retry_attempts",
&master_config.max_retry_attempts,
FLAGS_max_retry_attempts);
default_config.GetBool("enable_kv_event_publish",
&master_config.enable_kv_event_publish,
FLAGS_enable_kv_event_publish);
if (master_config.enable_kv_event_publish) {
InitKVEventPublisherConf(default_config, master_config);
}
}
void LoadKVEventPublisherConfigFromCmdline(
mooncake::MasterConfig& master_config, bool conf_set,
google::CommandLineFlagInfo& info) {
if ((google::GetCommandLineFlagInfo("kv_event_publisher_endpoint", &info) &&
!info.is_default) ||
!conf_set) {
master_config.kv_event_publisher_config.endpoint =
FLAGS_kv_event_publisher_endpoint;
}
if ((google::GetCommandLineFlagInfo("kv_event_publisher_replay_endpoint",
&info) &&
!info.is_default) ||
!conf_set) {
if (!FLAGS_kv_event_publisher_replay_endpoint.empty()) {
master_config.kv_event_publisher_config.replay_endpoint =
FLAGS_kv_event_publisher_replay_endpoint;
} else {
master_config.kv_event_publisher_config.replay_endpoint =
std::nullopt;
}
}
if ((google::GetCommandLineFlagInfo("kv_event_publisher_hwm", &info) &&
!info.is_default) ||
!conf_set) {
master_config.kv_event_publisher_config.hwm =
FLAGS_kv_event_publisher_hwm;
}
if ((google::GetCommandLineFlagInfo("kv_event_publisher_send_interval_ms",
&info) &&
!info.is_default) ||
!conf_set) {
master_config.kv_event_publisher_config.send_interval =
std::chrono::milliseconds(
FLAGS_kv_event_publisher_send_interval_ms);
}
if ((google::GetCommandLineFlagInfo("kv_event_publisher_max_batch_size",
&info) &&
!info.is_default) ||
!conf_set) {
master_config.kv_event_publisher_config.max_batch_size =
FLAGS_kv_event_publisher_max_batch_size;
}
if ((google::GetCommandLineFlagInfo("kv_event_publisher_auto_port",
&info) &&
!info.is_default) ||
!conf_set) {
master_config.kv_event_publisher_config.auto_port =
FLAGS_kv_event_publisher_auto_port;
}
if ((google::GetCommandLineFlagInfo("kv_event_publisher_topic", &info) &&
!info.is_default) ||
!conf_set) {
master_config.kv_event_publisher_config.topic =
FLAGS_kv_event_publisher_topic;
}
}
void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
@ -539,6 +677,16 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
!conf_set) {
master_config.snapshot_backend_type = FLAGS_snapshot_backend_type;
}
if ((google::GetCommandLineFlagInfo("enable_kv_event_publish", &info) &&
!info.is_default) ||
!conf_set) {
master_config.enable_kv_event_publish = FLAGS_enable_kv_event_publish;
}
if (master_config.enable_kv_event_publish) {
LoadKVEventPublisherConfigFromCmdline(master_config, conf_set, info);
}
}
// Function to start HTTP metadata server
@ -606,6 +754,13 @@ int main(int argc, char* argv[]) {
return 1;
}
if (master_config.enable_kv_event_publish) {
if (!master_config.kv_event_publisher_config.validate()) {
LOG(FATAL) << "Invalid KVEventPublisher configuration";
return 1;
}
}
const char* value = std::getenv("MC_RPC_PROTOCOL");
std::string protocol = "tcp";
if (value && std::string_view(value) == "rdma") {
@ -671,6 +826,49 @@ int main(int argc, char* argv[]) {
<< ", cxl_path=" << master_config.cxl_path
<< ", cxl_size=" << master_config.cxl_size;
if (master_config.enable_kv_event_publish) {
LOG(INFO)
<< "KV event publisher configuration: "
<< "endpoint=" << master_config.kv_event_publisher_config.endpoint
<< ", replay_endpoint="
<< (master_config.kv_event_publisher_config.replay_endpoint
.has_value()
? master_config.kv_event_publisher_config.replay_endpoint
.value()
: "[disabled]")
<< ", hwm=" << master_config.kv_event_publisher_config.hwm
<< ", send_interval="
<< master_config.kv_event_publisher_config.send_interval.count()
<< "ms"
<< ", max_batch_size="
<< master_config.kv_event_publisher_config.max_batch_size
<< ", auto_port="
<< (master_config.kv_event_publisher_config.auto_port ? "true"
: "false")
<< ", topic=" << master_config.kv_event_publisher_config.topic
<< ", buffer_steps="
<< master_config.kv_event_publisher_config.buffer_steps
<< ", max_queue_size="
<< master_config.kv_event_publisher_config.max_queue_size
<< ", batch_timeout="
<< master_config.kv_event_publisher_config.batch_timeout.count()
<< "ms"
<< ", pop_timeout="
<< master_config.kv_event_publisher_config.pop_timeout.count()
<< "ms"
<< ", enqueue_thread_pool_size="
<< master_config.kv_event_publisher_config.enqueue_thread_pool_size
<< ", enqueue_timeout="
<< master_config.kv_event_publisher_config.enqueue_timeout.count()
<< "ms"
<< ", enqueue_max_retries="
<< master_config.kv_event_publisher_config.enqueue_max_retries
<< ", max_port_attempts="
<< master_config.kv_event_publisher_config.max_port_attempts;
} else {
LOG(INFO) << "KV event publish: [disabled]";
}
// Start HTTP metadata server if enabled
std::unique_ptr<mooncake::HttpMetadataServer> http_metadata_server;
if (master_config.enable_http_metadata_server) {

View File

@ -68,7 +68,8 @@ MasterService::MasterService(const MasterServiceConfig& config)
task_manager_(config.task_manager_config),
cxl_path_(config.cxl_path),
cxl_size_(config.cxl_size),
enable_cxl_(config.enable_cxl) {
enable_cxl_(config.enable_cxl),
enable_kv_event_publish(config.enable_kv_event_publish) {
if (enable_snapshot_ || enable_snapshot_restore_) {
try {
auto backend_type =
@ -149,6 +150,11 @@ MasterService::MasterService(const MasterServiceConfig& config)
segment_manager_.initializeCxlAllocator(cxl_path_, cxl_size_);
VLOG(1) << "action=start_cxl_global_allocator";
}
if (enable_kv_event_publish) {
publisher =
std::make_unique<KVEventSystem>(config.kv_event_publisher_config);
}
}
MasterService::~MasterService() {
@ -167,6 +173,9 @@ MasterService::~MasterService() {
if (client_monitor_thread_.joinable()) {
client_monitor_thread_.join();
}
if (publisher) {
publisher->shutdown();
}
if (snapshot_thread_.joinable()) {
snapshot_thread_.join();
}
@ -262,18 +271,30 @@ auto MasterService::ReMountSegment(const std::vector<Segment>& segments,
}
void MasterService::ClearInvalidHandles() {
size_t replica_count{0};
for (size_t i = 0; i < kNumShards; i++) {
MetadataShardAccessorRW shard(this, i);
auto it = shard->metadata.begin();
while (it != shard->metadata.end()) {
if (publisher) {
replica_count = it->second.CountReplicas();
}
if (CleanupStaleHandles(it->second)) {
// If the object is empty, we need to erase the iterator and
// also erase the key from processing_keys and
// replication_tasks.
if (publisher) {
publisher->publish<BlockUpdateEvent>(
it->first, std::vector<Replica::Descriptor>{});
}
shard->processing_keys.erase(it->first);
shard->replication_tasks.erase(it->first);
it = shard->metadata.erase(it);
} else {
if (publisher && replica_count != it->second.CountReplicas()) {
publisher->publish<BlockUpdateEvent>(
it->first, it->second.GetReplicasDescriptorList());
}
++it;
}
}
@ -523,6 +544,12 @@ auto MasterService::BatchReplicaClear(
VLOG(1) << "BatchReplicaClear: successfully cleared all replicas "
"for key="
<< key << " for client_id=" << client_id;
if (publisher) {
publisher->publish<BlockUpdateEvent>(
key, std::vector<Replica::Descriptor>{});
}
} else {
// Clear only replicas on the specified segment_name
bool has_replica_on_segment = false;
@ -571,6 +598,11 @@ auto MasterService::BatchReplicaClear(
"segment_name="
<< segment_name << " for key=" << key
<< " for client_id=" << client_id;
if (publisher) {
publisher->publish<BlockUpdateEvent>(
key, metadata.GetReplicasDescriptorList());
}
}
}
@ -808,6 +840,11 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key,
accessor.EraseFromProcessing();
}
if (publisher) {
publisher->publish<BlockUpdateEvent>(
key, metadata.GetReplicasDescriptorList());
}
if (replica_type == ReplicaType::MEMORY) {
MasterMetricManager::instance().inc_mem_cache_nums();
} else if (replica_type == ReplicaType::DISK) {
@ -1115,6 +1152,11 @@ tl::expected<void, ErrorCode> MasterService::CopyEnd(const UUID& client_id,
accessor.EraseReplicationTask();
if (publisher) {
publisher->publish<BlockUpdateEvent>(
key, metadata.GetReplicasDescriptorList());
}
return all_complete ? tl::expected<void, ErrorCode>()
: tl::make_unexpected(ErrorCode::REPLICA_IS_GONE);
}
@ -1331,6 +1373,11 @@ tl::expected<void, ErrorCode> MasterService::MoveEnd(const UUID& client_id,
accessor.EraseReplicationTask();
if (publisher) {
publisher->publish<BlockUpdateEvent>(
key, metadata.GetReplicasDescriptorList());
}
return {};
}
@ -1421,6 +1468,12 @@ auto MasterService::Remove(const std::string& key, bool force)
// Remove object metadata
accessor.Erase();
if (publisher) {
publisher->publish<BlockUpdateEvent>(
key, std::vector<Replica::Descriptor>{});
}
return {};
}
@ -1474,6 +1527,12 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, bool force)
VLOG(1) << "key=" << it->first
<< " matched by regex. Removing.";
if (publisher) {
publisher->publish<BlockUpdateEvent>(
it->first, std::vector<Replica::Descriptor>{});
}
it = shard->metadata.erase(it);
removed_count++;
} else {
@ -1524,6 +1583,10 @@ long MasterService::RemoveAll(bool force) {
}
}
if (publisher) {
publisher->publish<RemoveAllEvent>();
}
VLOG(1) << "action=remove_all_objects"
<< ", removed_count=" << removed_count
<< ", total_freed_size=" << total_freed_size;
@ -1549,6 +1612,19 @@ size_t MasterService::GetKeyCount() const {
return total;
}
auto MasterService::GetPublisherStats() const
-> tl::expected<KVEventSystem::Stats, ErrorCode> {
if (!enable_kv_event_publish || !publisher) {
return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS);
}
try {
return publisher->get_stats();
} catch (const std::exception& e) {
return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
}
}
auto MasterService::Ping(const UUID& client_id)
-> tl::expected<PingResponse, ErrorCode> {
std::shared_lock<std::shared_mutex> lock(client_mutex_);
@ -1645,6 +1721,17 @@ auto MasterService::NotifyOffloadSuccess(
Replica replica(client_id, metadata.data_size,
metadata.transport_endpoint, ReplicaStatus::COMPLETE);
auto res = AddReplica(client_id, key, replica);
if (publisher) {
MetadataAccessor accessor(this, key);
if (!accessor.Exists()) {
LOG(ERROR) << "key=" << key << ", error=object_not_found";
return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
}
publisher->publish<BlockUpdateEvent>(
key, accessor.Get().GetReplicasDescriptorList());
}
if (!res && res.error() != ErrorCode::OBJECT_NOT_FOUND) {
LOG(ERROR) << "Failed to add replica: error=" << res.error()
<< ", client_id=" << client_id << ", key=" << key;
@ -2867,8 +2954,13 @@ void MasterService::BatchEvict(double evict_ratio_target,
if (it->second.lease_timeout <= target_timeout) {
// Evict this object
total_freed_size +=
it->second.size *
evict_replicas(it->second); // Erase memory replicas
it->second.size * evict_replicas(it->second);
if (publisher) {
publisher->publish<BlockUpdateEvent>(
it->first, it->second.GetReplicasDescriptorList());
}
// Erase memory replicas
if (it->second.IsValid() == false) {
it = shard->metadata.erase(it);
} else {
@ -2929,6 +3021,13 @@ void MasterService::BatchEvict(double evict_ratio_target,
it->second.size *
evict_replicas(
it->second); // Erase memory replicas
if (publisher) {
publisher->publish<BlockUpdateEvent>(
it->first,
it->second.GetReplicasDescriptorList());
}
if (it->second.IsValid() == false) {
it = shard->metadata.erase(it);
} else {
@ -2978,6 +3077,13 @@ void MasterService::BatchEvict(double evict_ratio_target,
it->second.size *
evict_replicas(
it->second); // Erase memory replicas
if (publisher) {
publisher->publish<BlockUpdateEvent>(
it->first,
it->second.GetReplicasDescriptorList());
}
if (it->second.IsValid() == false) {
it = shard->metadata.erase(it);
} else {
@ -3015,7 +3121,8 @@ void MasterService::BatchEvict(double evict_ratio_target,
}
MasterMetricManager::instance().inc_eviction_fail();
}
VLOG(1) << "action=evict_objects" << ", evicted_count=" << evicted_count
VLOG(1) << "action=evict_objects"
<< ", evicted_count=" << evicted_count
<< ", total_freed_size=" << total_freed_size;
}

View File

@ -39,6 +39,15 @@ WrappedMasterService::WrappedMasterService(
std::string metrics_summary =
MasterMetricManager::instance().get_summary_string();
LOG(INFO) << "Master Metrics: " << metrics_summary;
if (master_service_.IsPublisherEnabled()) {
auto result = master_service_.GetPublisherStats();
if (result) {
LOG(INFO)
<< "KV Event Publisher Metrics: " << result.value();
}
}
std::this_thread::sleep_for(
std::chrono::seconds(kMetricReportIntervalSeconds));
}

View File

@ -25,6 +25,7 @@ add_store_test(cxl_client_integration_test cxl_client_integration_test.cpp)
add_store_test(master_metrics_test master_metrics_test.cpp)
add_store_test(posix_file_test posix_file_test.cpp)
add_store_test(thread_pool_test thread_pool_test.cpp)
add_store_test(thread_safe_queue_test thread_safe_queue_test.cpp)
add_store_test(transfer_task_test transfer_task_test.cpp)
add_store_test(segment_test segment_test.cpp)
add_store_test(offset_allocator_test offset_allocator_test.cpp)
@ -68,3 +69,21 @@ target_link_libraries(stress_workload_test PUBLIC
gflags
pthread
)
find_package(cppzmq REQUIRED)
find_package(msgpack-cxx REQUIRED)
add_executable(kv_event_publisher_test kv_event_publisher_test.cpp)
target_link_libraries(kv_event_publisher_test PUBLIC
mooncake_store
cachelib_memory_allocator
${ETCD_WRAPPER_LIB}
glog
ibverbs
gtest
gtest_main
pthread
cppzmq
msgpack-cxx
)

View File

@ -0,0 +1,366 @@
#include <gtest/gtest.h>
#include <glog/logging.h>
#include <msgpack.hpp>
#include <zmq.hpp>
#include <chrono>
#include <thread>
#include <atomic>
#include <future>
#include <type_traits>
#include "kv_event_publisher.h"
namespace mooncake {
class KVEventSystemTest : public ::testing::Test {
protected:
static void SetUpTestSuite() {
google::InitGoogleLogging("KVEventSystemTest");
FLAGS_logtostderr = 1;
}
static void TearDownTestSuite() { google::ShutdownGoogleLogging(); }
void SetUp() override { create_test_descriptors(); }
Replica::Descriptor create_memory_replica() {
Replica::Descriptor desc;
uintptr_t dummy_address = 0xDEADBEEF;
desc.descriptor_variant = MemoryDescriptor{
.buffer_descriptor = {
.size_ = 1024,
.buffer_address_ = dummy_address,
.transport_endpoint_ = "tcp://192.168.1.100:5555"}};
desc.status = ReplicaStatus::COMPLETE;
return desc;
}
Replica::Descriptor create_disk_replica() {
Replica::Descriptor desc;
desc.descriptor_variant = DiskDescriptor{
.file_path = "/data/blocks/block_123.bin", .object_size = 4096};
desc.status = ReplicaStatus::COMPLETE;
return desc;
}
void create_test_descriptors() {
memory_replica_ = create_memory_replica();
disk_replica_ = create_disk_replica();
replicas_mixed_ = {memory_replica_, disk_replica_};
replicas_memory_only_ = {memory_replica_, memory_replica_};
replicas_disk_only_ = {disk_replica_, disk_replica_};
}
template <typename ConfigType>
std::unique_ptr<KVEventSystem> create_publish_system(ConfigType&& config) {
static_assert(
std::is_same_v<std::decay_t<ConfigType>, KVEventPublisherConfig>,
"ConfigType must be KVEventPublisherConfig");
return std::make_unique<KVEventSystem>(
std::forward<ConfigType>(config));
}
protected:
Replica::Descriptor memory_replica_;
Replica::Descriptor disk_replica_;
std::vector<Replica::Descriptor> replicas_mixed_;
std::vector<Replica::Descriptor> replicas_memory_only_;
std::vector<Replica::Descriptor> replicas_disk_only_;
};
TEST_F(KVEventSystemTest, BlockUpdateEventSerialization) {
BlockUpdateEvent event("key123", replicas_mixed_);
EXPECT_EQ(event.type_tag(), "BlockUpdateEvent");
msgpack::sbuffer buffer;
msgpack::packer<msgpack::sbuffer> pk(buffer);
event.pack(pk);
msgpack::object_handle oh = msgpack::unpack(buffer.data(), buffer.size());
msgpack::object obj = oh.get();
EXPECT_EQ(obj.type, msgpack::type::ARRAY);
EXPECT_EQ(obj.via.array.size, 3);
}
TEST_F(KVEventSystemTest, EventBatchSerialization) {
auto event1 =
std::make_shared<BlockUpdateEvent>("key1", replicas_memory_only_);
auto event2 =
std::make_shared<BlockUpdateEvent>("key2", replicas_disk_only_);
EventBatch batch({event1, event2});
EXPECT_GT(batch.ts, 0.0);
msgpack::sbuffer serialized = batch.serialize();
EXPECT_GT(serialized.size(), 0);
}
TEST_F(KVEventSystemTest, KVEventSystemConstruction) {
EXPECT_NO_THROW({
KVEventPublisherConfig config;
auto publisher = create_publish_system(config);
EXPECT_TRUE(publisher->is_running());
auto stats = publisher->get_stats();
EXPECT_EQ(stats.producer_stats.events_created, 0);
EXPECT_EQ(stats.event_queue_stats.queue_remain_events, 0);
EXPECT_EQ(stats.event_queue_stats.queue_capacity,
config.max_queue_size);
publisher->shutdown();
EXPECT_FALSE(publisher->is_running());
});
}
TEST_F(KVEventSystemTest, KVEventSystemMultipleConstruction) {
auto publisher1 = create_publish_system(KVEventPublisherConfig{});
auto publisher2 = create_publish_system(KVEventPublisherConfig{});
EXPECT_TRUE(publisher1->is_running());
EXPECT_TRUE(publisher2->is_running());
auto future1 =
publisher1->publish<BlockUpdateEvent>("key1", replicas_mixed_);
auto status1 = future1.wait_for(std::chrono::seconds(2));
EXPECT_EQ(status1, std::future_status::ready);
auto success1 = future1.get();
EXPECT_TRUE(success1);
publisher1->shutdown();
publisher2->shutdown();
}
TEST_F(KVEventSystemTest, KVEventSystemBasicPublish) {
auto publisher_ = create_publish_system(KVEventPublisherConfig{});
auto future_store = publisher_->publish<BlockUpdateEvent>("test_key_store",
replicas_mixed_);
auto future_update = publisher_->publish<BlockUpdateEvent>(
"test_key_update", replicas_mixed_);
auto future_remove_all = publisher_->publish<RemoveAllEvent>();
auto status_store = future_store.wait_for(std::chrono::seconds(2));
auto status_update = future_update.wait_for(std::chrono::seconds(2));
auto status_remove_all =
future_remove_all.wait_for(std::chrono::seconds(2));
EXPECT_EQ(status_store, std::future_status::ready);
EXPECT_EQ(status_update, std::future_status::ready);
EXPECT_EQ(status_remove_all, std::future_status::ready);
EXPECT_TRUE(future_store.get());
EXPECT_TRUE(future_update.get());
EXPECT_TRUE(future_remove_all.get());
auto stats = publisher_->get_stats();
EXPECT_EQ(stats.producer_stats.events_created, 3);
publisher_->shutdown();
EXPECT_FALSE(publisher_->is_running());
}
TEST_F(KVEventSystemTest, KVEventSystemConcurrentPublishing) {
auto publisher_ =
create_publish_system(KVEventPublisherConfig{.max_queue_size = 1000});
const int num_threads = 4;
const int events_per_thread = 100;
std::atomic<int> events_published{0};
std::atomic<int> failed_events{0};
std::vector<std::thread> threads;
std::vector<std::future<bool>> futures;
std::mutex futures_mutex;
futures.reserve(num_threads * events_per_thread);
for (int i = 0; i < num_threads; ++i) {
threads.emplace_back([this, i, &publisher_, &events_published, &futures,
&futures_mutex]() {
for (int j = 0; j < events_per_thread; ++j) {
std::string key =
"key_" + std::to_string(i) + "_" + std::to_string(j);
try {
auto future = publisher_->publish<BlockUpdateEvent>(
key, replicas_mixed_);
{
std::lock_guard<std::mutex> lock(futures_mutex);
futures.push_back(std::move(future));
}
events_published++;
} catch (const std::exception& e) {
FAIL() << "Publish failed: " << e.what();
}
std::this_thread::sleep_for(std::chrono::microseconds(5));
}
});
}
for (auto& t : threads) {
t.join();
}
int successful_events = 0;
for (auto& future : futures) {
if (future.wait_for(std::chrono::seconds(2)) ==
std::future_status::ready) {
if (future.get()) {
successful_events++;
}
}
}
EXPECT_EQ(events_published.load(), num_threads * events_per_thread);
auto stats = publisher_->get_stats();
EXPECT_GE(stats.producer_stats.events_created, events_published.load());
publisher_->shutdown();
EXPECT_FALSE(publisher_->is_running());
}
TEST_F(KVEventSystemTest, KVEventSystemGracefulShutdown) {
auto publisher_ =
create_publish_system(KVEventPublisherConfig{.max_queue_size = 100});
std::vector<std::future<bool>> futures;
for (int i = 0; i < 50; ++i) {
futures.push_back(publisher_->publish<BlockUpdateEvent>(
"key_" + std::to_string(i), replicas_mixed_));
}
for (auto& future : futures) {
future.wait_for(std::chrono::milliseconds(100));
}
EXPECT_NO_THROW(publisher_->shutdown());
auto future_after_shutdown =
publisher_->publish<BlockUpdateEvent>("should_fail", replicas_mixed_);
EXPECT_EQ(future_after_shutdown.wait_for(std::chrono::milliseconds(100)),
std::future_status::ready);
EXPECT_FALSE(future_after_shutdown.get());
}
TEST_F(KVEventSystemTest, KVEventSystemEdgeCases) {
auto publisher_ = create_publish_system(KVEventPublisherConfig{});
std::vector<Replica::Descriptor> empty_replicas;
auto future1 = publisher_->publish<BlockUpdateEvent>("empty_replicas_key",
empty_replicas);
EXPECT_EQ(future1.wait_for(std::chrono::seconds(2)),
std::future_status::ready);
EXPECT_TRUE(future1.get());
std::vector<uint32_t> large_tokens(10000);
for (size_t i = 0; i < large_tokens.size(); ++i) {
large_tokens[i] = static_cast<uint32_t>(i);
}
auto future2 = publisher_->publish<BlockUpdateEvent>("large_tokens_key",
replicas_mixed_);
EXPECT_EQ(future2.wait_for(std::chrono::seconds(2)),
std::future_status::ready);
EXPECT_TRUE(future2.get());
publisher_->shutdown();
EXPECT_FALSE(publisher_->is_running());
}
TEST_F(KVEventSystemTest, KVEventSystemStats) {
auto publisher_ = create_publish_system(KVEventPublisherConfig{});
auto initial_stats = publisher_->get_stats();
EXPECT_EQ(initial_stats.producer_stats.events_created, 0);
EXPECT_EQ(initial_stats.consumer_stats.total_batches, 0);
EXPECT_EQ(initial_stats.consumer_stats.failed_events, 0);
std::vector<std::future<bool>> futures;
for (int i = 0; i < 10; ++i) {
futures.push_back(publisher_->publish<BlockUpdateEvent>(
"key_" + std::to_string(i), replicas_mixed_));
}
for (auto& future : futures) {
future.wait_for(std::chrono::seconds(1));
}
auto final_stats = publisher_->get_stats();
EXPECT_GE(final_stats.producer_stats.events_created, 10);
EXPECT_GE(final_stats.consumer_stats.total_batches, 0);
EXPECT_EQ(final_stats.consumer_stats.failed_events, 0);
publisher_->shutdown();
EXPECT_FALSE(publisher_->is_running());
}
TEST_F(KVEventSystemTest, PerformanceTest) {
auto publisher_ = create_publish_system(KVEventPublisherConfig{
.max_queue_size = 10000,
.max_batch_size = 100,
.batch_timeout = std::chrono::milliseconds(500)});
const int num_events = 1000;
auto start_time = std::chrono::high_resolution_clock::now();
std::vector<std::future<bool>> futures;
futures.reserve(num_events);
for (int i = 0; i < num_events; ++i) {
futures.push_back(publisher_->publish<BlockUpdateEvent>(
"perf_key_" + std::to_string(i), replicas_mixed_));
}
int successful_events = 0;
for (auto& future : futures) {
if (future.wait_for(std::chrono::seconds(5)) ==
std::future_status::ready) {
if (future.get()) {
successful_events++;
}
}
}
auto end_time = std::chrono::high_resolution_clock::now();
auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
end_time - start_time);
LOG(INFO) << "Published " << num_events << " events in " << duration.count()
<< "ms";
LOG(INFO) << "Successful events: " << successful_events << "/"
<< num_events;
LOG(INFO) << "Throughput: " << (num_events * 1000.0 / duration.count())
<< " events/sec";
auto stats = publisher_->get_stats();
LOG(INFO) << "Stats - Total events: " << stats.producer_stats.events_created
<< ", Total batches: " << stats.consumer_stats.total_batches
<< ", Failed events: " << stats.consumer_stats.failed_events;
// Allow up to 10 failed events out of 1000 for transient issues
const int MAX_ALLOWED_FAILURES = 10;
EXPECT_GE(successful_events, num_events - MAX_ALLOWED_FAILURES);
publisher_->shutdown();
EXPECT_FALSE(publisher_->is_running());
}
} // namespace mooncake
int main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}

File diff suppressed because it is too large Load Diff