forked from mooncake-track/Mooncake
Compare commits
12 Commits
main
...
dev/kv-ind
| Author | SHA1 | Date |
|---|---|---|
|
|
a6bfb2e195 | |
|
|
c837624931 | |
|
|
c67649c83a | |
|
|
6ab4e40ee5 | |
|
|
80c8b41cad | |
|
|
770031b101 | |
|
|
fb7313b87d | |
|
|
c15a911ae0 | |
|
|
45402ff9a1 | |
|
|
fbff581166 | |
|
|
0cc4fea731 | |
|
|
d0ec8173d3 |
|
|
@ -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: |
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
140
dependencies.sh
140
dependencies.sh
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
@ -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).
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
}'
|
||||
|
||||
```
|
||||
|
|
@ -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
|
||||
|
|
@ -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())
|
||||
}
|
||||
|
|
@ -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())
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Reference in New Issue