EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[Router] remove duplicate char count (#2378)

Browse Source
This commit is contained in:
Byron Hsu
2024-12-06 11:26:07 -08:00
committed by GitHub
parent e5f227c0ee
commit 499c85f131
3 changed files with 41 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
rust/py_test/test_launch_server.py
View File

@@ -45,6 +45,8 @@ def popen_launch_router(
port,
"--dp",
str(dp_size), # Convert dp_size to string
"--router-eviction-interval",
"5", # frequent eviction for testing
]
# Use current environment

4
rust/src/server.rs
View File

@@ -1,8 +1,6 @@
use crate::router::PolicyConfig;
use crate::router::Router;
use actix_web::{
delete, get, post, put, web, App, HttpRequest, HttpResponse, HttpServer, Responder,
};
use actix_web::{get, post, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
use bytes::Bytes;
use env_logger::Builder;
use log::{info, LevelFilter};

90
rust/src/tree.rs
View File

@@ -24,7 +24,6 @@ struct Node {
#[derive(Debug)]
pub struct Tree {
root: NodeRef,
// TODO: Char Count per tenant
pub tenant_char_count: DashMap<String, usize>,
}
@@ -408,17 +407,9 @@ impl Tree {
pub fn evict_tenant_data(&self, max_size: usize) {
// Calculate used size and collect leaves
let mut stack = vec![Arc::clone(&self.root)];
let mut used_size_per_tenant: HashMap<String, usize> = HashMap::new();
let mut pq = BinaryHeap::new();
while let Some(curr) = stack.pop() {
for tenant in curr.tenant_last_access_time.iter() {
let size = used_size_per_tenant
.entry(tenant.key().clone())
.or_insert(0);
*size += curr.text.read().unwrap().chars().count();
}
for child in curr.children.iter() {
stack.push(Arc::clone(child.value()));
}
@@ -436,64 +427,59 @@ impl Tree {
}
info!("Before eviction - Used size per tenant:");
for (tenant, size) in &used_size_per_tenant {
info!("Tenant: {}, Size: {}", tenant, size);
for entry in self.tenant_char_count.iter() {
info!("Tenant: {}, Size: {}", entry.key(), entry.value());
}
// Process eviction
while let Some(Reverse(entry)) = pq.pop() {
let EvictionEntry { tenant, node, .. } = entry;
if let Some(&used_size) = used_size_per_tenant.get(&tenant) {
if used_size <= max_size {
if let Some(used_size) = self.tenant_char_count.get(&tenant) {
if *used_size <= max_size {
continue;
}
}
// Update used size
if let Some(size) = used_size_per_tenant.get_mut(&tenant) {
*size -= node.text.read().unwrap().chars().count();
}
// Decrement when removing tenant from node
if node.tenant_last_access_time.contains_key(&tenant) {
self.tenant_char_count
.entry(tenant.clone())
.and_modify(|count| {
if *count > 0 {
*count -= node.text.read().unwrap().chars().count();
}
});
}
// Remove tenant from node
node.tenant_last_access_time.remove(&tenant);
// Remove empty nodes
if node.children.is_empty() && node.tenant_last_access_time.is_empty() {
if let Some(parent) = node.parent.write().unwrap().as_ref() {
let first_char = node.text.read().unwrap().chars().next().unwrap();
parent.children.remove(&first_char);
}
}
// Add parent to queue if it becomes a leaf
if let Some(parent) = node.parent.read().unwrap().as_ref() {
if Tree::leaf_of(parent).contains(&tenant) {
if let Some(timestamp) = parent.tenant_last_access_time.get(&tenant) {
pq.push(Reverse(EvictionEntry {
timestamp: *timestamp,
tenant: tenant.clone(),
node: Arc::clone(parent),
}));
// Decrement when removing tenant from node
if node.tenant_last_access_time.contains_key(&tenant) {
self.tenant_char_count
.entry(tenant.clone())
.and_modify(|count| {
if *count > 0 {
*count -= node.text.read().unwrap().chars().count();
}
}
});
}
// Remove tenant from node
node.tenant_last_access_time.remove(&tenant);
// Remove empty nodes
if node.children.is_empty() && node.tenant_last_access_time.is_empty() {
if let Some(parent) = node.parent.write().unwrap().as_ref() {
let first_char = node.text.read().unwrap().chars().next().unwrap();
parent.children.remove(&first_char);
}
}
// Add parent to queue if it becomes a leaf
if let Some(parent) = node.parent.read().unwrap().as_ref() {
if Tree::leaf_of(parent).contains(&tenant) {
if let Some(timestamp) = parent.tenant_last_access_time.get(&tenant) {
pq.push(Reverse(EvictionEntry {
timestamp: *timestamp,
tenant: tenant.clone(),
node: Arc::clone(parent),
}));
}
}
};
}
info!("After eviction - Used size per tenant:");
for (tenant, size) in &used_size_per_tenant {
info!("Tenant: {}, Size: {}", tenant, size);
for entry in self.tenant_char_count.iter() {
info!("Tenant: {}, Size: {}", entry.key(), entry.value());
}
}