[Feature] Add process-local in-memory KV indexer and Router integration (#33370)
Co-authored-by: Wu, Yutong <yutong.wu@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
co-authored by
Wu, Yutong
TianDi101
Zhangheng
parent
238ba40c27
commit
360d10d6bc
@@ -0,0 +1,11 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
pub fn nanos() -> u128 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
use sgl_kv_indexer::pb::{
|
||||
ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType,
|
||||
};
|
||||
|
||||
pub fn hbm() -> i32 {
|
||||
TierType::TierHbm as i32
|
||||
}
|
||||
|
||||
pub fn dram() -> i32 {
|
||||
TierType::TierDram as i32
|
||||
}
|
||||
|
||||
pub fn action(kind: ExternalKvActionType, tier: i32, hashes: &[i64]) -> ExternalKvAction {
|
||||
ExternalKvAction {
|
||||
r#type: kind as i32,
|
||||
tier,
|
||||
hashes: hashes.to_vec(),
|
||||
component_masks: Vec::new(),
|
||||
block_sizes: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// A component-aware REPORT action: each hash carries its component bitmask and
|
||||
/// token count, index-aligned with `hashes`.
|
||||
#[allow(dead_code)] // used by memory_integration, not grpc_contract
|
||||
pub fn component_report(
|
||||
tier: i32,
|
||||
hashes: &[i64],
|
||||
masks: &[u32],
|
||||
block_sizes: &[u32],
|
||||
) -> ExternalKvAction {
|
||||
ExternalKvAction {
|
||||
r#type: ExternalKvActionType::ActionReport as i32,
|
||||
tier,
|
||||
hashes: hashes.to_vec(),
|
||||
component_masks: masks.to_vec(),
|
||||
block_sizes: block_sizes.to_vec(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn apply_request(
|
||||
worker: &str,
|
||||
address: &str,
|
||||
seq: u64,
|
||||
actions: Vec<ExternalKvAction>,
|
||||
) -> ApplyExternalKvBatchRequest {
|
||||
ApplyExternalKvBatchRequest {
|
||||
worker_id: worker.to_string(),
|
||||
seq,
|
||||
actions,
|
||||
worker_address: address.to_string(),
|
||||
cache_spec: None,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
use std::net::SocketAddr;
|
||||
|
||||
/// Reserves an ephemeral loopback port. The listener is dropped immediately,
|
||||
/// so callers that spawn a server should retain their connect-retry loop.
|
||||
pub fn free_addr() -> SocketAddr {
|
||||
std::net::TcpListener::bind("127.0.0.1:0")
|
||||
.unwrap()
|
||||
.local_addr()
|
||||
.unwrap()
|
||||
}
|
||||
@@ -0,0 +1,517 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! gRPC contract tests: exercise all four RPCs of the `KVIndexer` service
|
||||
//! over the wire (real tonic server + client), not just the backend trait.
|
||||
|
||||
#[path = "common/id.rs"]
|
||||
mod test_id;
|
||||
#[allow(dead_code)]
|
||||
#[path = "common/kv.rs"]
|
||||
mod test_kv;
|
||||
#[path = "common/net.rs"]
|
||||
mod test_net;
|
||||
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use prost::Message;
|
||||
use tokio::sync::Semaphore;
|
||||
use tonic::transport::Server;
|
||||
use tonic::{Code, Status};
|
||||
|
||||
use sgl_kv_indexer::pb::kv_indexer_client::KvIndexerClient;
|
||||
use sgl_kv_indexer::pb::kv_indexer_server::KvIndexerServer;
|
||||
use sgl_kv_indexer::pb::{
|
||||
ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvAction,
|
||||
ExternalKvActionType, GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse,
|
||||
MatchExternalKvPrefixRequest, MatchExternalKvPrefixResponse, MatchExternalKvRequest,
|
||||
MatchExternalKvResponse,
|
||||
};
|
||||
use sgl_kv_indexer::{
|
||||
server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerBackend, KvIndexerService,
|
||||
PrefixIndex, PrefixIndexConfig, MAX_GRPC_DECODING_MESSAGE_SIZE,
|
||||
};
|
||||
use test_id::nanos;
|
||||
use test_kv::{action, apply_request, hbm};
|
||||
use test_net::free_addr;
|
||||
|
||||
async fn start_backend(
|
||||
backend: InMemoryKvIndexerBackend,
|
||||
) -> KvIndexerClient<tonic::transport::Channel> {
|
||||
let svc = KvIndexerService::new(backend).into_server();
|
||||
let addr = free_addr();
|
||||
tokio::spawn(async move {
|
||||
server_builder()
|
||||
.add_service(svc)
|
||||
.serve(addr)
|
||||
.await
|
||||
.expect("server serve");
|
||||
});
|
||||
|
||||
let endpoint = format!("http://{addr}");
|
||||
for _ in 0..50 {
|
||||
if let Ok(c) = KvIndexerClient::connect(endpoint.clone()).await {
|
||||
return c;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
panic!("client failed to connect to {endpoint}");
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct BlockingPrefixBackend {
|
||||
entered: Arc<AtomicUsize>,
|
||||
release: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl KvIndexerBackend for BlockingPrefixBackend {
|
||||
async fn apply_external_kv_batch(
|
||||
&self,
|
||||
_request: ApplyExternalKvBatchRequest,
|
||||
) -> Result<ApplyExternalKvBatchResponse, Status> {
|
||||
Ok(ApplyExternalKvBatchResponse::default())
|
||||
}
|
||||
|
||||
async fn match_external_kv(
|
||||
&self,
|
||||
_request: MatchExternalKvRequest,
|
||||
) -> Result<MatchExternalKvResponse, Status> {
|
||||
Ok(MatchExternalKvResponse::default())
|
||||
}
|
||||
|
||||
async fn match_external_kv_prefix(
|
||||
&self,
|
||||
_request: MatchExternalKvPrefixRequest,
|
||||
) -> Result<MatchExternalKvPrefixResponse, Status> {
|
||||
self.entered.fetch_add(1, Ordering::SeqCst);
|
||||
let _permit = self.release.acquire().await.expect("semaphore open");
|
||||
Ok(MatchExternalKvPrefixResponse::default())
|
||||
}
|
||||
|
||||
async fn get_external_kv_hit_counts(
|
||||
&self,
|
||||
_request: GetExternalKvHitCountsRequest,
|
||||
) -> Result<GetExternalKvHitCountsResponse, Status> {
|
||||
Ok(GetExternalKvHitCountsResponse::default())
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_blocking_backend(
|
||||
backend: BlockingPrefixBackend,
|
||||
) -> KvIndexerClient<tonic::transport::Channel> {
|
||||
let svc = KvIndexerService::with_prefix_query_max_inflight(backend, 2).into_server();
|
||||
let addr = free_addr();
|
||||
tokio::spawn(async move {
|
||||
server_builder()
|
||||
.add_service(svc)
|
||||
.serve(addr)
|
||||
.await
|
||||
.expect("server serve");
|
||||
});
|
||||
|
||||
let endpoint = format!("http://{addr}");
|
||||
for _ in 0..50 {
|
||||
if let Ok(client) = KvIndexerClient::connect(endpoint.clone()).await {
|
||||
return client;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
panic!("client failed to connect to {endpoint}");
|
||||
}
|
||||
|
||||
/// Starts a real gRPC server with isolated process-local state.
|
||||
async fn start() -> KvIndexerClient<tonic::transport::Channel> {
|
||||
start_backend(InMemoryKvIndexerBackend::new()).await
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn prefix_limit_rejects_over_real_grpc_without_blocking_writes() {
|
||||
let entered = Arc::new(AtomicUsize::new(0));
|
||||
let release = Arc::new(Semaphore::new(0));
|
||||
let backend = BlockingPrefixBackend {
|
||||
entered: Arc::clone(&entered),
|
||||
release: Arc::clone(&release),
|
||||
};
|
||||
let client = start_blocking_backend(backend).await;
|
||||
let request = || MatchExternalKvPrefixRequest {
|
||||
hashes: vec![-1],
|
||||
max_blocks: 0,
|
||||
};
|
||||
|
||||
let mut first_client = client.clone();
|
||||
let first = tokio::spawn(async move { first_client.match_external_kv_prefix(request()).await });
|
||||
let mut second_client = client.clone();
|
||||
let second =
|
||||
tokio::spawn(async move { second_client.match_external_kv_prefix(request()).await });
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while entered.load(Ordering::SeqCst) != 2 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("two prefix queries should enter the backend");
|
||||
|
||||
let mut rejected_client = client.clone();
|
||||
let rejected = tokio::time::timeout(
|
||||
Duration::from_secs(1),
|
||||
rejected_client.match_external_kv_prefix(request()),
|
||||
)
|
||||
.await
|
||||
.expect("overload response should be immediate")
|
||||
.expect_err("third prefix query should be rejected");
|
||||
assert_eq!(rejected.code(), Code::ResourceExhausted);
|
||||
assert_eq!(entered.load(Ordering::SeqCst), 2);
|
||||
|
||||
let mut write_client = client.clone();
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(1),
|
||||
write_client.apply_external_kv_batch(ApplyExternalKvBatchRequest {
|
||||
worker_id: "worker".into(),
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.expect("writes should not share the prefix-query limit")
|
||||
.expect("write should succeed");
|
||||
|
||||
release.add_permits(2);
|
||||
first.await.expect("first task").expect("first response");
|
||||
second.await.expect("second task").expect("second response");
|
||||
}
|
||||
|
||||
fn apply(
|
||||
worker: &str,
|
||||
addr: &str,
|
||||
seq: u64,
|
||||
action_type: ExternalKvActionType,
|
||||
tier: i32,
|
||||
hashes: &[i64],
|
||||
) -> ApplyExternalKvBatchRequest {
|
||||
apply_request(worker, addr, seq, vec![action(action_type, tier, hashes)])
|
||||
}
|
||||
|
||||
fn apply_report(
|
||||
worker: &str,
|
||||
addr: &str,
|
||||
seq: u64,
|
||||
tier: i32,
|
||||
hashes: &[i64],
|
||||
) -> ApplyExternalKvBatchRequest {
|
||||
apply(
|
||||
worker,
|
||||
addr,
|
||||
seq,
|
||||
ExternalKvActionType::ActionReport,
|
||||
tier,
|
||||
hashes,
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multiple_workers_share_one_indexer_server() {
|
||||
let mut indexer = start().await;
|
||||
let suffix = nanos();
|
||||
let worker_0 = format!("worker-0-{suffix}");
|
||||
let worker_1 = format!("worker-1-{suffix}");
|
||||
let (hash_0, hash_1, shared_hash) = (1, 2, 3);
|
||||
|
||||
indexer
|
||||
.apply_external_kv_batch(apply_report(
|
||||
&worker_0,
|
||||
"10.0.0.1:9000",
|
||||
1,
|
||||
hbm(),
|
||||
&[hash_0, shared_hash],
|
||||
))
|
||||
.await
|
||||
.expect("apply worker-0");
|
||||
indexer
|
||||
.apply_external_kv_batch(apply_report(
|
||||
&worker_1,
|
||||
"10.0.0.2:9000",
|
||||
1,
|
||||
hbm(),
|
||||
&[hash_1, shared_hash],
|
||||
))
|
||||
.await
|
||||
.expect("apply worker-1");
|
||||
|
||||
let response = indexer
|
||||
.match_external_kv(MatchExternalKvRequest {
|
||||
hashes: vec![hash_0, hash_1, shared_hash],
|
||||
count_as_hit: false,
|
||||
})
|
||||
.await
|
||||
.expect("query indexer")
|
||||
.into_inner();
|
||||
assert!(response
|
||||
.matches
|
||||
.iter()
|
||||
.any(|entry| entry.worker_id == worker_0));
|
||||
assert!(response
|
||||
.matches
|
||||
.iter()
|
||||
.any(|entry| entry.worker_id == worker_1));
|
||||
|
||||
// Keep one wire-level smoke check for hit counting; detailed counter
|
||||
// semantics live in memory_integration.rs.
|
||||
indexer
|
||||
.match_external_kv(MatchExternalKvRequest {
|
||||
hashes: vec![hash_0],
|
||||
count_as_hit: true,
|
||||
})
|
||||
.await
|
||||
.expect("counting match over gRPC");
|
||||
let miss = 4;
|
||||
let counts = indexer
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest {
|
||||
hashes: vec![hash_0, miss],
|
||||
})
|
||||
.await
|
||||
.expect("hit counts over gRPC")
|
||||
.into_inner();
|
||||
let count = |hash: i64| {
|
||||
counts
|
||||
.entries
|
||||
.iter()
|
||||
.find(|entry| entry.hash == hash)
|
||||
.map(|entry| entry.hit_count_total)
|
||||
.unwrap_or(0)
|
||||
};
|
||||
assert!(count(hash_0) >= 1, "matched hash should have a hit");
|
||||
assert_eq!(count(miss), 0, "unmatched hash must not be counted");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn validation_errors_map_to_invalid_argument_over_grpc() {
|
||||
let mut c = start().await;
|
||||
|
||||
let err = c
|
||||
.apply_external_kv_batch(apply_report("", "addr", 1, hbm(), &[1]))
|
||||
.await
|
||||
.expect_err("empty worker_id must be rejected");
|
||||
assert_eq!(err.code(), Code::InvalidArgument);
|
||||
|
||||
// An action type outside the enum can only arrive over the wire; the
|
||||
// in-process tests cover the mapped `ActionUnknown` variant instead.
|
||||
let unmapped_action_type = ApplyExternalKvBatchRequest {
|
||||
worker_id: "w".into(),
|
||||
seq: 1,
|
||||
worker_address: String::new(),
|
||||
cache_spec: None,
|
||||
actions: vec![ExternalKvAction {
|
||||
r#type: 999,
|
||||
tier: hbm(),
|
||||
hashes: vec![1],
|
||||
component_masks: Vec::new(),
|
||||
block_sizes: Vec::new(),
|
||||
}],
|
||||
};
|
||||
let err = c
|
||||
.apply_external_kv_batch(unmapped_action_type)
|
||||
.await
|
||||
.expect_err("unknown action type must be rejected");
|
||||
assert_eq!(err.code(), Code::InvalidArgument);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn match_prefix_over_grpc() {
|
||||
let mut c = start().await;
|
||||
let (w_long, w_short) = (format!("long-{}", nanos()), format!("short-{}", nanos()));
|
||||
let (a, b, d) = (1, 2, 3);
|
||||
|
||||
c.apply_external_kv_batch(apply_report(&w_long, "10.0.0.1:9000", 1, hbm(), &[a, b, d]))
|
||||
.await
|
||||
.expect("apply long");
|
||||
c.apply_external_kv_batch(apply_report(&w_short, "10.0.0.2:9000", 1, hbm(), &[a]))
|
||||
.await
|
||||
.expect("apply short");
|
||||
|
||||
let resp = c
|
||||
.match_external_kv_prefix(MatchExternalKvPrefixRequest {
|
||||
hashes: vec![a, b, d],
|
||||
max_blocks: 0,
|
||||
})
|
||||
.await
|
||||
.expect("prefix ok")
|
||||
.into_inner();
|
||||
|
||||
assert_eq!(resp.best_prefix_blocks, 3);
|
||||
assert_eq!(resp.blocks_read, 3);
|
||||
// Descending by prefix length: long (3) before short (1).
|
||||
assert_eq!(resp.matches.len(), 2);
|
||||
assert_eq!(resp.matches[0].worker_id, w_long);
|
||||
assert_eq!(resp.matches[0].matched_prefix_blocks, 3);
|
||||
assert_eq!(resp.matches[0].worker_address, "10.0.0.1:9000");
|
||||
assert_eq!(resp.matches[1].worker_id, w_short);
|
||||
assert_eq!(resp.matches[1].matched_prefix_blocks, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prefix_query_scans_more_than_one_apply_chunk_over_grpc() {
|
||||
const APPLY_CHUNK_SIZE: usize = 16_384;
|
||||
|
||||
let mut indexer = start().await;
|
||||
let hashes: Vec<i64> = (0..=APPLY_CHUNK_SIZE as i64).collect();
|
||||
for (seq, chunk) in hashes.chunks(APPLY_CHUNK_SIZE).enumerate() {
|
||||
indexer
|
||||
.apply_external_kv_batch(apply_report(
|
||||
"large-prefix-worker",
|
||||
"10.0.0.1:9000",
|
||||
seq as u64,
|
||||
hbm(),
|
||||
chunk,
|
||||
))
|
||||
.await
|
||||
.expect("bounded apply chunk");
|
||||
}
|
||||
|
||||
let response = indexer
|
||||
.match_external_kv_prefix(MatchExternalKvPrefixRequest {
|
||||
hashes,
|
||||
max_blocks: 0,
|
||||
})
|
||||
.await
|
||||
.expect("prefix request larger than one apply chunk")
|
||||
.into_inner();
|
||||
|
||||
assert_eq!(response.best_prefix_blocks as usize, APPLY_CHUNK_SIZE + 1);
|
||||
assert_eq!(response.blocks_read as usize, APPLY_CHUNK_SIZE + 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn packed_signed_hash_query_can_exceed_tonics_default_receive_limit() {
|
||||
const TONIC_DEFAULT_RECEIVE_LIMIT: usize = 4 * 1024 * 1024;
|
||||
const HASH_COUNT: usize = 600_000;
|
||||
|
||||
let mut indexer = start().await;
|
||||
indexer
|
||||
.apply_external_kv_batch(apply_report(
|
||||
"large-wire-worker",
|
||||
"10.0.0.1:9000",
|
||||
1,
|
||||
hbm(),
|
||||
&[-1],
|
||||
))
|
||||
.await
|
||||
.expect("store the signed first hash");
|
||||
|
||||
let mut hashes = Vec::with_capacity(HASH_COUNT);
|
||||
hashes.push(-1);
|
||||
hashes.extend((1..HASH_COUNT).map(|value| value as i64));
|
||||
let request = MatchExternalKvPrefixRequest {
|
||||
hashes,
|
||||
max_blocks: 0,
|
||||
};
|
||||
assert!(request.encoded_len() > TONIC_DEFAULT_RECEIVE_LIMIT);
|
||||
assert!(request.encoded_len() < MAX_GRPC_DECODING_MESSAGE_SIZE);
|
||||
|
||||
let response = indexer
|
||||
.match_external_kv_prefix(request)
|
||||
.await
|
||||
.expect("configured server accepts a packed request larger than 4 MiB")
|
||||
.into_inner();
|
||||
assert_eq!(response.best_prefix_blocks, 1);
|
||||
}
|
||||
|
||||
/// Past the configured ceiling the server must answer OUT_OF_RANGE, because that
|
||||
/// is the code the router maps to a degraded (cache-affinity-free) route rather
|
||||
/// than to a failed request. A different code there would fail the request.
|
||||
#[tokio::test]
|
||||
async fn query_past_the_configured_limit_is_refused_as_out_of_range() {
|
||||
let hash_count = MAX_GRPC_DECODING_MESSAGE_SIZE / std::mem::size_of::<i64>() + 1_024;
|
||||
let request = MatchExternalKvPrefixRequest {
|
||||
hashes: (0..hash_count).map(|value| value as i64).collect(),
|
||||
max_blocks: 0,
|
||||
};
|
||||
assert!(request.encoded_len() > MAX_GRPC_DECODING_MESSAGE_SIZE);
|
||||
|
||||
let status = start()
|
||||
.await
|
||||
.match_external_kv_prefix(request)
|
||||
.await
|
||||
.expect_err("a request past the ceiling must be refused");
|
||||
assert_eq!(status.code(), Code::OutOfRange);
|
||||
}
|
||||
|
||||
/// Serves an empty backend behind an interceptor that records the `grpc-timeout`
|
||||
/// of every request, and returns the router-facing client alongside the capture.
|
||||
async fn start_recording_deadlines(
|
||||
query_deadline: Duration,
|
||||
) -> (GrpcPrefixIndex, Arc<Mutex<Vec<String>>>) {
|
||||
let seen = Arc::new(Mutex::new(Vec::new()));
|
||||
let recorder = Arc::clone(&seen);
|
||||
let svc = KvIndexerServer::with_interceptor(
|
||||
KvIndexerService::new(InMemoryKvIndexerBackend::new()),
|
||||
move |request: tonic::Request<()>| {
|
||||
if let Some(timeout) = request.metadata().get("grpc-timeout") {
|
||||
recorder
|
||||
.lock()
|
||||
.expect("deadline recorder")
|
||||
.push(timeout.to_str().expect("ascii timeout").to_string());
|
||||
}
|
||||
Ok(request)
|
||||
},
|
||||
);
|
||||
let addr = free_addr();
|
||||
tokio::spawn(async move {
|
||||
Server::builder()
|
||||
.add_service(svc)
|
||||
.serve(addr)
|
||||
.await
|
||||
.expect("server serve");
|
||||
});
|
||||
|
||||
let endpoint = format!("http://{addr}");
|
||||
for _ in 0..50 {
|
||||
if KvIndexerClient::connect(endpoint.clone()).await.is_ok() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
|
||||
let index = GrpcPrefixIndex::new(PrefixIndexConfig {
|
||||
endpoint,
|
||||
query_deadline,
|
||||
max_inflight: sgl_kv_indexer::DEFAULT_QUERY_MAX_INFLIGHT,
|
||||
})
|
||||
.expect("test endpoint is valid");
|
||||
(index, seen)
|
||||
}
|
||||
|
||||
/// The router-facing client must publish its deadline on the wire: that header is
|
||||
/// the only thing letting the indexer shed a query whose caller gave up.
|
||||
#[tokio::test]
|
||||
async fn router_client_publishes_its_deadline_on_the_wire() {
|
||||
let (index, seen) = start_recording_deadlines(Duration::from_millis(250)).await;
|
||||
|
||||
index
|
||||
.match_prefix(vec![1, 2, 3])
|
||||
.await
|
||||
.expect("query reaches the indexer");
|
||||
|
||||
let seen = seen.lock().expect("deadline recorder").clone();
|
||||
assert_eq!(
|
||||
seen.len(),
|
||||
1,
|
||||
"exactly one query reached the server: {seen:?}"
|
||||
);
|
||||
let raw = &seen[0];
|
||||
// Asserted structurally, not byte-for-byte: the wire spec lets the sender
|
||||
// pick any unit that fits, so pinning tonic's choice would fail on a
|
||||
// change that is still correct.
|
||||
let (digits, unit) = raw.split_at(raw.len() - 1);
|
||||
assert!(
|
||||
matches!(unit, "H" | "M" | "S" | "m" | "u" | "n"),
|
||||
"unit is one the wire spec defines: {raw:?}"
|
||||
);
|
||||
let value: u64 = digits.parse().expect("timeout value is numeric");
|
||||
assert!(
|
||||
value > 0,
|
||||
"a budget of zero would shed every query: {raw:?}"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,886 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! Integration tests for the process-local in-memory backend.
|
||||
|
||||
#[path = "common/kv.rs"]
|
||||
mod test_kv;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use sgl_kv_indexer::pb::{
|
||||
ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvActionType,
|
||||
GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse, MatchExternalKvPrefixRequest,
|
||||
MatchExternalKvPrefixResponse, MatchExternalKvRequest, MatchExternalKvResponse,
|
||||
WorkerCacheSpec,
|
||||
};
|
||||
use sgl_kv_indexer::{
|
||||
InMemoryKvIndexerBackend, KvIndexerBackend, WorkerPrefixInput, COMPONENT_FULL, COMPONENT_SWA,
|
||||
};
|
||||
use test_kv::{action, apply_request as apply_req, component_report, dram, hbm};
|
||||
use tonic::Status;
|
||||
|
||||
fn backend() -> InMemoryKvIndexerBackend {
|
||||
InMemoryKvIndexerBackend::new()
|
||||
}
|
||||
|
||||
fn match_req(hs: &[i64], count_as_hit: bool) -> MatchExternalKvRequest {
|
||||
MatchExternalKvRequest {
|
||||
hashes: hs.to_vec(),
|
||||
count_as_hit,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the tiers a worker holds a hash at, per the match response.
|
||||
fn tiers_for(resp: &MatchExternalKvResponse, worker: &str, hash: i64) -> Vec<i32> {
|
||||
let mut tiers = Vec::new();
|
||||
for m in &resp.matches {
|
||||
if m.worker_id != worker {
|
||||
continue;
|
||||
}
|
||||
for th in &m.hashes_by_tier {
|
||||
if th.hashes.contains(&hash) {
|
||||
tiers.push(th.tier);
|
||||
}
|
||||
}
|
||||
}
|
||||
tiers.sort_unstable();
|
||||
tiers
|
||||
}
|
||||
|
||||
macro_rules! itest {
|
||||
($name:ident, $b:ident, $body:block) => {
|
||||
#[tokio::test]
|
||||
async fn $name() {
|
||||
let $b = backend();
|
||||
$body
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
itest!(report_then_match_returns_worker_and_address, b, {
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"10.0.0.1:9000",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let resp = b
|
||||
.match_external_kv(match_req(&[1, 2, 3], false))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.matches.len(), 1);
|
||||
let m = &resp.matches[0];
|
||||
assert_eq!(m.worker_id, "w1");
|
||||
assert_eq!(m.address, "10.0.0.1:9000");
|
||||
assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]);
|
||||
assert_eq!(tiers_for(&resp, "w1", 2), vec![hbm()]);
|
||||
assert!(tiers_for(&resp, "w1", 3).is_empty());
|
||||
});
|
||||
|
||||
itest!(large_request_preserves_complete_ordered_results, b, {
|
||||
// Exercise a large write and read while preserving complete ordered results.
|
||||
let expected_hashes: Vec<i64> = (0..300).collect();
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(
|
||||
ExternalKvActionType::ActionReport,
|
||||
hbm(),
|
||||
&expected_hashes,
|
||||
)],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let resp = b
|
||||
.match_external_kv(match_req(&expected_hashes, false))
|
||||
.await
|
||||
.unwrap();
|
||||
let worker = resp
|
||||
.matches
|
||||
.iter()
|
||||
.find(|m| m.worker_id == "w1")
|
||||
.expect("worker must match");
|
||||
let tier = worker
|
||||
.hashes_by_tier
|
||||
.iter()
|
||||
.find(|t| t.tier == hbm())
|
||||
.expect("HBM tier must match");
|
||||
assert_eq!(tier.hashes, expected_hashes);
|
||||
});
|
||||
|
||||
itest!(duplicate_report_is_idempotent, b, {
|
||||
for _ in 0..3 {
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]);
|
||||
});
|
||||
|
||||
itest!(identical_batch_replay_is_idempotent, b, {
|
||||
// Stores, removes, then stores the same hash again; the net state is
|
||||
// "stored". Re-delivering the identical batch must not change it.
|
||||
let batch = apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
7,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[9]),
|
||||
action(ExternalKvActionType::ActionRevoke, hbm(), &[9]),
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[9]),
|
||||
],
|
||||
);
|
||||
b.apply_external_kv_batch(batch.clone()).await.unwrap();
|
||||
let first = b.match_external_kv(match_req(&[9], false)).await.unwrap();
|
||||
b.apply_external_kv_batch(batch).await.unwrap();
|
||||
let second = b.match_external_kv(match_req(&[9], false)).await.unwrap();
|
||||
assert_eq!(tiers_for(&first, "w1", 9), vec![hbm()]);
|
||||
assert_eq!(tiers_for(&second, "w1", 9), vec![hbm()]);
|
||||
});
|
||||
|
||||
itest!(recomputed_full_node_restores_hbm_placement, b, {
|
||||
// HiRadixCache lifecycle for an exact-match recomputation:
|
||||
// BlockStored(GPU) -> BlockRemoved(GPU) -> BlockStored(GPU).
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let evicted = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert!(tiers_for(&evicted, "w1", 1).is_empty());
|
||||
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
3,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let restored = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert_eq!(tiers_for(&restored, "w1", 1), vec![hbm()]);
|
||||
});
|
||||
|
||||
itest!(recomputed_split_reports_only_materialized_hashes, b, {
|
||||
// An evicted [prefix -> old suffix] is partially recomputed as
|
||||
// [prefix -> new suffix]. The old suffix must remain absent.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1, 2])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
3,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[1]),
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[3]),
|
||||
],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = b
|
||||
.match_external_kv(match_req(&[1, 2, 3], false))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(tiers_for(&result, "w1", 1), vec![hbm()]);
|
||||
assert!(tiers_for(&result, "w1", 2).is_empty());
|
||||
assert_eq!(tiers_for(&result, "w1", 3), vec![hbm()]);
|
||||
});
|
||||
|
||||
itest!(recomputed_batch_replay_keeps_cpu_copy, b, {
|
||||
// Re-materializing on GPU must not revoke the existing host backup, and
|
||||
// re-delivering the same batch must leave both tiers unchanged.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, dram(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let recomputed = apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
);
|
||||
b.apply_external_kv_batch(recomputed.clone()).await.unwrap();
|
||||
b.apply_external_kv_batch(recomputed).await.unwrap();
|
||||
|
||||
let result = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert_eq!(tiers_for(&result, "w1", 1), vec![hbm(), dram()]);
|
||||
});
|
||||
|
||||
itest!(revoke_partial_tier_keeps_other_tier, b, {
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[1]),
|
||||
action(ExternalKvActionType::ActionReport, dram(), &[1]),
|
||||
action(ExternalKvActionType::ActionRevoke, hbm(), &[1]),
|
||||
],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert_eq!(tiers_for(&resp, "w1", 1), vec![dram()]);
|
||||
});
|
||||
|
||||
itest!(revoke_missing_hash_is_idempotent, b, {
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[404])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b.match_external_kv(match_req(&[404], false)).await.unwrap();
|
||||
assert!(resp.matches.is_empty());
|
||||
});
|
||||
|
||||
itest!(multi_worker_multi_tier, b, {
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a1",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w2",
|
||||
"a2",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, dram(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert_eq!(resp.matches.len(), 2);
|
||||
assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]);
|
||||
assert_eq!(tiers_for(&resp, "w2", 1), vec![dram()]);
|
||||
});
|
||||
|
||||
itest!(clear_all_at_tier_removes_only_that_tier, b, {
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[1, 2, 3]),
|
||||
action(ExternalKvActionType::ActionReport, dram(), &[1]),
|
||||
],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![action(
|
||||
ExternalKvActionType::ActionClearAllAtTier,
|
||||
hbm(),
|
||||
&[],
|
||||
)],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b
|
||||
.match_external_kv(match_req(&[1, 2, 3], false))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(tiers_for(&resp, "w1", 1), vec![dram()]);
|
||||
assert!(tiers_for(&resp, "w1", 2).is_empty());
|
||||
assert!(tiers_for(&resp, "w1", 3).is_empty());
|
||||
});
|
||||
|
||||
itest!(
|
||||
count_as_hit_only_counts_matched_and_replay_does_not_double,
|
||||
b,
|
||||
{
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Diagnostic match (count_as_hit=false) must not count.
|
||||
b.match_external_kv(match_req(&[1, 2], false))
|
||||
.await
|
||||
.unwrap();
|
||||
let counts = b
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1, 2] })
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(counts.entries.is_empty());
|
||||
|
||||
// Counting match: only the matched hash "1" is counted, "2" (a miss) is not.
|
||||
b.match_external_kv(match_req(&[1, 2], true)).await.unwrap();
|
||||
let counts = b
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1, 2] })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts.entries.len(), 1);
|
||||
assert_eq!(counts.entries[0].hash, 1);
|
||||
assert_eq!(counts.entries[0].hit_count_total, 1);
|
||||
|
||||
// Replaying the apply batch must not touch hit counts.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let counts = b
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts.entries[0].hit_count_total, 1);
|
||||
}
|
||||
);
|
||||
|
||||
itest!(full_revoke_drops_hit_key, b, {
|
||||
// Report a block, count a hit (creates the co-located :h key), then fully
|
||||
// revoke it. The hit key must go with the placement, or a
|
||||
// matched-then-evicted block leaks its counter forever.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Counting match creates the hit key with c=1.
|
||||
b.match_external_kv(match_req(&[1], true)).await.unwrap();
|
||||
let counts = b
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts.entries.len(), 1);
|
||||
assert_eq!(counts.entries[0].hit_count_total, 1);
|
||||
|
||||
// Fully revoke the block: placement empties, so the hit key must go too.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
|
||||
assert!(resp.matches.is_empty());
|
||||
|
||||
// Hit key is gone too: a leaked :h would still report a count here.
|
||||
let counts = b
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
counts.entries.is_empty(),
|
||||
"hit key leaked after full revoke: {:?}",
|
||||
counts.entries
|
||||
);
|
||||
});
|
||||
|
||||
itest!(partial_revoke_keeps_hit_key, b, {
|
||||
// Block present at two tiers; count a hit, then revoke only one tier. Placement
|
||||
// is still non-empty, so the hit key must survive (guard against over-deletion).
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[1]),
|
||||
action(ExternalKvActionType::ActionReport, dram(), &[1]),
|
||||
],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
b.match_external_kv(match_req(&[1], true)).await.unwrap();
|
||||
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let counts = b
|
||||
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(counts.entries.len(), 1);
|
||||
assert_eq!(counts.entries[0].hit_count_total, 1);
|
||||
});
|
||||
|
||||
itest!(batch_action_order_is_preserved, b, {
|
||||
// revoke-then-report on the same hash within one batch must net to "stored".
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
1,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionRevoke, hbm(), &[5]),
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[5]),
|
||||
],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b.match_external_kv(match_req(&[5], false)).await.unwrap();
|
||||
assert_eq!(tiers_for(&resp, "w1", 5), vec![hbm()]);
|
||||
|
||||
// report-then-revoke on the same hash must net to "absent".
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"a",
|
||||
2,
|
||||
vec![
|
||||
action(ExternalKvActionType::ActionReport, hbm(), &[6]),
|
||||
action(ExternalKvActionType::ActionRevoke, hbm(), &[6]),
|
||||
],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b.match_external_kv(match_req(&[6], false)).await.unwrap();
|
||||
assert!(tiers_for(&resp, "w1", 6).is_empty());
|
||||
});
|
||||
|
||||
// --- prefix query: backend override vs. the trait's default implementation ---
|
||||
//
|
||||
// The trait default is the written semantics and the backend override is a read
|
||||
// optimization, so they must agree field-for-field on the parts that ARE the
|
||||
// contract (per-worker prefix set and best_prefix_blocks). `blocks_read` is
|
||||
// observability and legitimately differs, so it is not compared.
|
||||
|
||||
/// Delegates every RPC to an in-memory backend EXCEPT `match_external_kv_prefix`,
|
||||
/// which it leaves to the trait default — giving a reference answer computed from
|
||||
/// the same state the optimized path reads.
|
||||
struct DefaultViaMemory(Arc<InMemoryKvIndexerBackend>);
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl KvIndexerBackend for DefaultViaMemory {
|
||||
async fn apply_external_kv_batch(
|
||||
&self,
|
||||
request: ApplyExternalKvBatchRequest,
|
||||
) -> Result<ApplyExternalKvBatchResponse, Status> {
|
||||
self.0.apply_external_kv_batch(request).await
|
||||
}
|
||||
|
||||
async fn match_external_kv(
|
||||
&self,
|
||||
request: MatchExternalKvRequest,
|
||||
) -> Result<MatchExternalKvResponse, Status> {
|
||||
self.0.match_external_kv(request).await
|
||||
}
|
||||
|
||||
// Delegate the component-aware read to the same backend so the trait
|
||||
// default computes over the same placement and specs the fast path sees.
|
||||
async fn collect_worker_prefix_inputs(
|
||||
&self,
|
||||
hashes: &[i64],
|
||||
) -> Result<Vec<WorkerPrefixInput>, Status> {
|
||||
self.0.collect_worker_prefix_inputs(hashes).await
|
||||
}
|
||||
|
||||
async fn get_external_kv_hit_counts(
|
||||
&self,
|
||||
request: GetExternalKvHitCountsRequest,
|
||||
) -> Result<GetExternalKvHitCountsResponse, Status> {
|
||||
self.0.get_external_kv_hit_counts(request).await
|
||||
}
|
||||
}
|
||||
|
||||
fn shared_state_pair() -> (Arc<InMemoryKvIndexerBackend>, DefaultViaMemory) {
|
||||
let backend = Arc::new(InMemoryKvIndexerBackend::new());
|
||||
let reference = DefaultViaMemory(Arc::clone(&backend));
|
||||
(backend, reference)
|
||||
}
|
||||
|
||||
fn prefix_req(hs: &[i64]) -> MatchExternalKvPrefixRequest {
|
||||
MatchExternalKvPrefixRequest {
|
||||
hashes: hs.to_vec(),
|
||||
max_blocks: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sorted `(worker_id, matched_prefix_blocks)` — the semantic content of a
|
||||
/// prefix response, independent of `blocks_read`.
|
||||
fn prefix_pairs(resp: &MatchExternalKvPrefixResponse) -> Vec<(String, u32)> {
|
||||
let mut pairs: Vec<(String, u32)> = resp
|
||||
.matches
|
||||
.iter()
|
||||
.map(|m| (m.worker_id.clone(), m.matched_prefix_blocks))
|
||||
.collect();
|
||||
pairs.sort();
|
||||
pairs
|
||||
}
|
||||
|
||||
fn report(worker: &str, addr: &str, seq: u64, hs: &[i64]) -> ApplyExternalKvBatchRequest {
|
||||
apply_req(
|
||||
worker,
|
||||
addr,
|
||||
seq,
|
||||
vec![action(ExternalKvActionType::ActionReport, hbm(), hs)],
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prefix_fast_path_matches_default_impl() {
|
||||
let (fast, reference) = shared_state_pair();
|
||||
|
||||
// Nested prefixes (hole-free), a diverging branch, and a hole.
|
||||
fast.apply_external_kv_batch(report("w-long", "10.0.0.1:1", 1, &[1, 2, 3, 4]))
|
||||
.await
|
||||
.unwrap();
|
||||
fast.apply_external_kv_batch(report("w-short", "10.0.0.2:1", 1, &[1, 2]))
|
||||
.await
|
||||
.unwrap();
|
||||
// w-hole holds 1, 3, 4 but not 2: strict prefix must be 1.
|
||||
fast.apply_external_kv_batch(report("w-hole", "10.0.0.3:1", 1, &[1, 3, 4]))
|
||||
.await
|
||||
.unwrap();
|
||||
// w-noaddr is unroutable and must be excluded by both paths.
|
||||
fast.apply_external_kv_batch(report("w-noaddr", "", 1, &[1, 2]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let query = [1, 2, 3, 4];
|
||||
let fast_resp = fast
|
||||
.match_external_kv_prefix(prefix_req(&query))
|
||||
.await
|
||||
.unwrap();
|
||||
let ref_resp = reference
|
||||
.match_external_kv_prefix(prefix_req(&query))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(prefix_pairs(&fast_resp), prefix_pairs(&ref_resp));
|
||||
assert_eq!(fast_resp.best_prefix_blocks, ref_resp.best_prefix_blocks);
|
||||
assert_eq!(fast_resp.best_prefix_blocks, 4);
|
||||
assert_eq!(
|
||||
prefix_pairs(&fast_resp),
|
||||
vec![
|
||||
("w-hole".to_string(), 1),
|
||||
("w-long".to_string(), 4),
|
||||
("w-short".to_string(), 2),
|
||||
]
|
||||
);
|
||||
assert!(fast_resp
|
||||
.matches
|
||||
.iter()
|
||||
.all(|m| !m.worker_address.is_empty()));
|
||||
// Descending order and first-block read are part of the response contract.
|
||||
assert_eq!(fast_resp.matches[0].matched_prefix_blocks, 4);
|
||||
assert!(fast_resp.blocks_read >= 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prefix_first_block_miss_reads_one_block() {
|
||||
let b = backend();
|
||||
// No worker holds the first queried block; the scan stops after one read.
|
||||
b.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[2, 3]))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b
|
||||
.match_external_kv_prefix(prefix_req(&[1, 2, 3]))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.matches.is_empty());
|
||||
assert_eq!(resp.best_prefix_blocks, 0);
|
||||
assert_eq!(resp.blocks_read, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prefix_max_blocks_caps_the_scan() {
|
||||
let b = backend();
|
||||
b.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[1, 2, 3, 4]))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b
|
||||
.match_external_kv_prefix(MatchExternalKvPrefixRequest {
|
||||
hashes: vec![1, 2, 3, 4],
|
||||
max_blocks: 2,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
// Capped at 2 even though the worker holds all four.
|
||||
assert_eq!(resp.best_prefix_blocks, 2);
|
||||
assert_eq!(resp.blocks_read, 2);
|
||||
assert_eq!(resp.matches.len(), 1);
|
||||
assert_eq!(resp.matches[0].matched_prefix_blocks, 2);
|
||||
}
|
||||
|
||||
// --- component-aware placement & prefix -------------------------------------
|
||||
|
||||
/// A hybrid-SWA spec: full servable from HBM+DRAM, swa a 100-token trailing
|
||||
/// window servable from HBM.
|
||||
fn swa_spec() -> WorkerCacheSpec {
|
||||
WorkerCacheSpec {
|
||||
version: 1,
|
||||
components: COMPONENT_FULL | COMPONENT_SWA,
|
||||
swa_window_tokens: 100,
|
||||
full_tier_mask: (1 << hbm()) | (1 << dram()),
|
||||
swa_tier_mask: 1 << hbm(),
|
||||
mamba_tier_mask: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_with_spec(
|
||||
worker: &str,
|
||||
addr: &str,
|
||||
seq: u64,
|
||||
spec: WorkerCacheSpec,
|
||||
actions: Vec<sgl_kv_indexer::pb::ExternalKvAction>,
|
||||
) -> ApplyExternalKvBatchRequest {
|
||||
let mut req = apply_req(worker, addr, seq, actions);
|
||||
req.cache_spec = Some(spec);
|
||||
req
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn component_prefix_matches_default_impl() {
|
||||
let (fast, reference) = shared_state_pair();
|
||||
|
||||
// Four full blocks (50 tokens each); swa present on all but the 4th, so the
|
||||
// largest boundary with an unbroken 100-token swa window is 3.
|
||||
let report = component_report(
|
||||
hbm(),
|
||||
&[1, 2, 3, 4],
|
||||
&[
|
||||
COMPONENT_FULL | COMPONENT_SWA,
|
||||
COMPONENT_FULL | COMPONENT_SWA,
|
||||
COMPONENT_FULL | COMPONENT_SWA,
|
||||
COMPONENT_FULL,
|
||||
],
|
||||
&[50, 50, 50, 50],
|
||||
);
|
||||
fast.apply_external_kv_batch(apply_with_spec(
|
||||
"w-swa",
|
||||
"10.0.0.1:1",
|
||||
1,
|
||||
swa_spec(),
|
||||
vec![report],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let query = [1, 2, 3, 4];
|
||||
let fast_resp = fast
|
||||
.match_external_kv_prefix(prefix_req(&query))
|
||||
.await
|
||||
.unwrap();
|
||||
let ref_resp = reference
|
||||
.match_external_kv_prefix(prefix_req(&query))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(prefix_pairs(&fast_resp), prefix_pairs(&ref_resp));
|
||||
assert_eq!(fast_resp.best_prefix_blocks, ref_resp.best_prefix_blocks);
|
||||
assert_eq!(prefix_pairs(&fast_resp), vec![("w-swa".to_string(), 3)]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn partial_eviction_replace_shrinks_component_set() {
|
||||
let b = backend();
|
||||
// Store full+swa, then restate to full-only (partial swa eviction) via a
|
||||
// REPLACE snapshot for the same (hash, tier). No BlockRemoved is involved.
|
||||
b.apply_external_kv_batch(apply_with_spec(
|
||||
"w1",
|
||||
"10.0.0.1:1",
|
||||
1,
|
||||
swa_spec(),
|
||||
vec![component_report(
|
||||
hbm(),
|
||||
&[1, 2],
|
||||
&[
|
||||
COMPONENT_FULL | COMPONENT_SWA,
|
||||
COMPONENT_FULL | COMPONENT_SWA,
|
||||
],
|
||||
&[80, 80],
|
||||
)],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
// Both blocks reusable (window 100 met by 2x80 tokens; head rule also holds).
|
||||
let before = b
|
||||
.match_external_kv_prefix(prefix_req(&[1, 2]))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(before.best_prefix_blocks, 2);
|
||||
|
||||
// Restate the second block to full only: swa gone there.
|
||||
b.apply_external_kv_batch(apply_with_spec(
|
||||
"w1",
|
||||
"10.0.0.1:1",
|
||||
2,
|
||||
swa_spec(),
|
||||
vec![component_report(hbm(), &[2], &[COMPONENT_FULL], &[80])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
// That block has no swa now, and its trailing window (only 80 < 100) is not
|
||||
// headed, so the largest valid boundary drops to 1.
|
||||
let after = b
|
||||
.match_external_kv_prefix(prefix_req(&[1, 2]))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(after.best_prefix_blocks, 1);
|
||||
|
||||
let snapshot = b
|
||||
.match_external_kv(match_req(&[1, 2], false))
|
||||
.await
|
||||
.unwrap();
|
||||
let tier = &snapshot.matches[0].hashes_by_tier[0];
|
||||
assert_eq!(tier.hashes, vec![1, 2]);
|
||||
assert_eq!(
|
||||
tier.component_masks,
|
||||
vec![COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL]
|
||||
);
|
||||
assert_eq!(tier.block_sizes, vec![80, 80]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn component_aware_worker_without_spec_is_excluded() {
|
||||
let b = backend();
|
||||
// Report component-aware placement but never send a spec: the worker cannot
|
||||
// be interpreted and must be excluded (NoSignal-safe), never over-reported.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"10.0.0.1:1",
|
||||
1,
|
||||
vec![component_report(
|
||||
hbm(),
|
||||
&[1, 2],
|
||||
&[COMPONENT_FULL, COMPONENT_FULL],
|
||||
&[16, 16],
|
||||
)],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b
|
||||
.match_external_kv_prefix(prefix_req(&[1, 2]))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.matches.is_empty());
|
||||
assert_eq!(resp.best_prefix_blocks, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_hash_in_one_report_keeps_last_snapshot() {
|
||||
let b = backend();
|
||||
// A single REPORT action naming the same hash twice (a coalesced
|
||||
// store+restate): the LAST snapshot must win deterministically, never a race.
|
||||
b.apply_external_kv_batch(apply_with_spec(
|
||||
"w1",
|
||||
"10.0.0.1:1",
|
||||
1,
|
||||
swa_spec(),
|
||||
vec![component_report(
|
||||
hbm(),
|
||||
&[1, 1],
|
||||
&[COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL],
|
||||
&[80, 80],
|
||||
)],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
// The hash ends as full-only (last snapshot); with swa required and a lone 80-token
|
||||
// block that is not a full head window, the boundary requiring swa fails,
|
||||
// so no reusable prefix.
|
||||
let resp = b.match_external_kv_prefix(prefix_req(&[1])).await.unwrap();
|
||||
assert_eq!(resp.best_prefix_blocks, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn absent_spec_batch_clears_stored_spec() {
|
||||
let b = backend();
|
||||
// First a component-aware batch establishes a spec + a reusable block.
|
||||
b.apply_external_kv_batch(apply_with_spec(
|
||||
"w1",
|
||||
"10.0.0.1:1",
|
||||
1,
|
||||
swa_spec(),
|
||||
vec![component_report(
|
||||
hbm(),
|
||||
&[1],
|
||||
&[COMPONENT_FULL | COMPONENT_SWA],
|
||||
&[200],
|
||||
)],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
b.match_external_kv_prefix(prefix_req(&[1]))
|
||||
.await
|
||||
.unwrap()
|
||||
.best_prefix_blocks,
|
||||
1
|
||||
);
|
||||
|
||||
// A later batch with NO spec (worker reverted to legacy) must clear the old
|
||||
// spec. The still-component-aware placement can then no longer be interpreted
|
||||
// (component data but no spec) -> fail closed, never scored on stale rules.
|
||||
b.apply_external_kv_batch(apply_req(
|
||||
"w1",
|
||||
"10.0.0.1:1",
|
||||
2,
|
||||
vec![action(ExternalKvActionType::ActionReport, dram(), &[2])],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let resp = b.match_external_kv_prefix(prefix_req(&[1])).await.unwrap();
|
||||
assert!(resp.matches.is_empty());
|
||||
assert_eq!(resp.best_prefix_blocks, 0);
|
||||
}
|
||||
Reference in New Issue
Block a user