From ac78f96ec5c8c06f425514ae28be0769968c6988 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 29 Dec 2025 10:08:19 -0800 Subject: [PATCH] [model-gateway] Optimize radix tree timestamp updates for multi-tenant scaling (#16093) --- sgl-model-gateway/benches/tree_benchmark.rs | 9 +- sgl-model-gateway/src/policies/tree.rs | 121 +++++++------------- 2 files changed, 44 insertions(+), 86 deletions(-) diff --git a/sgl-model-gateway/benches/tree_benchmark.rs b/sgl-model-gateway/benches/tree_benchmark.rs index 8c908e720..5d86631f4 100644 --- a/sgl-model-gateway/benches/tree_benchmark.rs +++ b/sgl-model-gateway/benches/tree_benchmark.rs @@ -788,13 +788,12 @@ fn bench_summary(c: &mut Criterion) { // Configuration constants const TREE_SIZE: usize = 10_000; // Realistic cache size const INSERT_POOL_SIZE: usize = 10_000; // Unique requests for insert tests - const NUM_THREADS: usize = 8; - const OPS_PER_THREAD: usize = 100; + const NUM_THREADS: usize = 64; // Match GPU runner's 64 CPU cores + const OPS_PER_THREAD: usize = 200; // Worker scaling configurations to test - // Using 2 representative counts (low/high) to keep CI runtime reasonable - // while still demonstrating scaling behavior - const WORKER_COUNTS: [usize; 2] = [10, 100]; + // Full range to demonstrate scaling behavior on GPU runner + const WORKER_COUNTS: [usize; 4] = [10, 50, 100, 500]; // Pre-generate requests for tree population and queries let requests = generate_realistic_requests(TREE_SIZE); diff --git a/sgl-model-gateway/src/policies/tree.rs b/sgl-model-gateway/src/policies/tree.rs index a4c493c83..42c09b90c 100644 --- a/sgl-model-gateway/src/policies/tree.rs +++ b/sgl-model-gateway/src/policies/tree.rs @@ -6,7 +6,6 @@ use std::{ atomic::{AtomicU64, Ordering}, Arc, RwLock, }, - time::{SystemTime, UNIX_EPOCH}, }; use dashmap::{mapref::entry::Entry, DashMap}; @@ -176,35 +175,24 @@ impl Clone for NodeText { } } -/// Global timestamp that gets updated periodically to reduce syscalls. -/// Uses milliseconds since epoch. -static CURRENT_TIMESTAMP_MS: AtomicU64 = AtomicU64::new(0); +/// Global epoch counter for LRU ordering. +/// Uses a simple incrementing counter instead of wall clock time. +/// +/// Benefits: +/// - No syscall overhead (vs SystemTime::now()) +/// - Smaller memory footprint (u64 vs u128) +/// - Perfectly monotonic (no clock skew issues) +/// +/// For LRU eviction, relative ordering is all that matters. +static EPOCH_COUNTER: AtomicU64 = AtomicU64::new(0); -/// Staleness threshold in milliseconds for forced refresh. -/// If cached timestamp is older than this, always get fresh time. -const TIMESTAMP_STALENESS_MS: u64 = 5; - -/// Get current timestamp in milliseconds, using cached value when possible. -/// Refreshes if the cached value is stale (>TIMESTAMP_STALENESS_MS). -/// This provides ~99% syscall reduction under high load while maintaining accuracy. +/// Get the next epoch value for LRU timestamp ordering. +/// Uses fetch_add for lock-free, monotonically increasing values. +/// Relaxed ordering is sufficient since we only need eventual consistency +/// for approximate LRU behavior. #[inline] -fn get_timestamp_ms() -> u128 { - let cached = CURRENT_TIMESTAMP_MS.load(Ordering::Relaxed); - - // Always need syscall to check staleness, but it's cheap and necessary for correctness - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_millis() as u64; - - // Fast path: return cached if still fresh (within TIMESTAMP_STALENESS_MS) - if cached != 0 && now.saturating_sub(cached) < TIMESTAMP_STALENESS_MS { - return cached as u128; - } - - // Update cached value - CURRENT_TIMESTAMP_MS.store(now, Ordering::Relaxed); - now as u128 +fn get_epoch() -> u64 { + EPOCH_COUNTER.fetch_add(1, Ordering::Relaxed) } #[derive(Debug)] @@ -214,8 +202,8 @@ struct Node { children: DashMap, /// Node text with cached character count text: RwLock, - /// Per-tenant last access timestamps. Using TenantId (Arc) for cheap cloning. - tenant_last_access_time: DashMap, + /// Per-tenant last access epoch for LRU ordering. Using TenantId (Arc) for cheap cloning. + tenant_last_access_time: DashMap, /// Parent pointer for upward traversal during timestamp updates parent: RwLock>, } @@ -230,7 +218,7 @@ pub struct Tree { // For the heap struct EvictionEntry { - timestamp: u128, + timestamp: u64, tenant: TenantId, node: NodeRef, } @@ -314,8 +302,8 @@ impl Tree { // Insert text into tree with given tenant // Use slice-based traversal to avoid Vec allocation - // Use cached timestamp to reduce syscalls - let timestamp_ms = get_timestamp_ms(); + // Get epoch for LRU ordering + let epoch = get_epoch(); // Intern the tenant ID once for reuse let tenant_id = intern_tenant(tenant); @@ -325,7 +313,7 @@ impl Tree { self.root .tenant_last_access_time - .insert(Arc::clone(&tenant_id), timestamp_ms); + .insert(Arc::clone(&tenant_id), epoch); self.tenant_char_count .entry(Arc::clone(&tenant_id)) @@ -368,7 +356,7 @@ impl Tree { .or_insert(remaining_char_count); new_node .tenant_last_access_time - .insert(Arc::clone(&tenant_id), timestamp_ms); + .insert(Arc::clone(&tenant_id), epoch); entry.insert(new_node); InsertStep::Done @@ -420,10 +408,10 @@ impl Tree { .entry(Arc::clone(&tenant_id)) .and_modify(|count| *count += matched_text_count) .or_insert(matched_text_count); - v.insert(timestamp_ms); + v.insert(epoch); } Entry::Occupied(mut o) => { - o.insert(timestamp_ms); + o.insert(epoch); } } @@ -445,10 +433,10 @@ impl Tree { .entry(Arc::clone(&tenant_id)) .and_modify(|count| *count += matched_node_text_count) .or_insert(matched_node_text_count); - v.insert(timestamp_ms); + v.insert(epoch); } Entry::Occupied(mut o) => { - o.insert(timestamp_ms); + o.insert(epoch); } } @@ -524,17 +512,12 @@ impl Tree { .next() .map(|kv| Arc::clone(kv.key())); - // Use cached timestamp to reduce syscalls - let timestamp_ms = get_timestamp_ms(); - - // Traverse from the curr node to the root and update the timestamp + // Update timestamp on the matched node only (O(1)). + // Ancestor propagation is unnecessary - eviction only uses leaf timestamps. if let Some(ref tenant_id) = tenant { - let mut current_node = Some(Arc::clone(&curr)); - while let Some(node) = current_node { - node.tenant_last_access_time - .insert(Arc::clone(tenant_id), timestamp_ms); - current_node = node.parent.read().unwrap().clone(); - } + let epoch = get_epoch(); + curr.tenant_last_access_time + .insert(Arc::clone(tenant_id), epoch); } // Build matched text from original input using char count @@ -609,20 +592,15 @@ impl Tree { let curr = prev; - // Only update timestamp if we found a match for the specified tenant + // Only update timestamp if we found a match for the specified tenant. + // Update matched node only - ancestor propagation is unnecessary. if curr .tenant_last_access_time .contains_key(tenant_id.as_ref()) { - // Use cached timestamp to reduce syscalls - let timestamp_ms = get_timestamp_ms(); - - let mut current_node = Some(curr); - while let Some(node) = current_node { - node.tenant_last_access_time - .insert(Arc::clone(&tenant_id), timestamp_ms); - current_node = node.parent.read().unwrap().clone(); - } + let epoch = get_epoch(); + curr.tenant_last_access_time + .insert(Arc::clone(&tenant_id), epoch); } // Build result from original input using char count @@ -866,8 +844,6 @@ impl Tree { #[allow(dead_code)] fn node_to_string(node: &NodeRef, prefix: &str, is_last: bool) -> String { - use std::time::Duration; - let mut result = String::new(); // Add prefix and branch character @@ -878,29 +854,12 @@ impl Tree { let node_text = node.text.read().unwrap(); result.push_str(&format!("'{}' [", node_text.as_str())); - // Add tenant information with timestamps + // Add tenant information with epoch values let mut tenant_info = Vec::new(); for entry in node.tenant_last_access_time.iter() { let tenant_id = entry.key(); - let timestamp_ms = entry.value(); - - // Convert milliseconds to seconds and remaining milliseconds - let seconds = (timestamp_ms / 1000) as u64; - let millis = (timestamp_ms % 1000) as u32; - - // Create SystemTime from Unix timestamp - let system_time = UNIX_EPOCH + Duration::from_secs(seconds); - - // Format time as HH:MM:SS.mmm - let datetime = system_time.duration_since(UNIX_EPOCH).unwrap(); - let hours = (datetime.as_secs() % 86400) / 3600; - let minutes = (datetime.as_secs() % 3600) / 60; - let seconds = datetime.as_secs() % 60; - - tenant_info.push(format!( - "{} | {:02}:{:02}:{:02}.{:03}", - tenant_id, hours, minutes, seconds, millis - )); + let epoch = entry.value(); + tenant_info.push(format!("{} | epoch:{}", tenant_id, epoch)); } result.push_str(&tenant_info.join(", "));