Unverified Commit f0365820 authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[Misc] add structure logging, write to file and log tracing for SGL Router

parent 86317c09
......@@ -525,6 +525,15 @@ dependencies = [
"cfg-if",
]
[[package]]
name = "crossbeam-channel"
version = "0.5.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.20"
......@@ -1204,6 +1213,12 @@ version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388"
[[package]]
name = "lazy_static"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "libc"
version = "0.2.161"
......@@ -1255,6 +1270,15 @@ version = "0.4.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24"
[[package]]
name = "matchers"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558"
dependencies = [
"regex-automata 0.1.10",
]
[[package]]
name = "memchr"
version = "2.7.4"
......@@ -1315,6 +1339,16 @@ dependencies = [
"tempfile",
]
[[package]]
name = "nu-ansi-term"
version = "0.46.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84"
dependencies = [
"overload",
"winapi",
]
[[package]]
name = "num-conv"
version = "0.1.0"
......@@ -1389,6 +1423,12 @@ dependencies = [
"vcpkg",
]
[[package]]
name = "overload"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
[[package]]
name = "parking_lot"
version = "0.12.3"
......@@ -1591,8 +1631,17 @@ checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191"
dependencies = [
"aho-corasick",
"memchr",
"regex-automata",
"regex-syntax",
"regex-automata 0.4.8",
"regex-syntax 0.8.5",
]
[[package]]
name = "regex-automata"
version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132"
dependencies = [
"regex-syntax 0.6.29",
]
[[package]]
......@@ -1603,7 +1652,7 @@ checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3"
dependencies = [
"aho-corasick",
"memchr",
"regex-syntax",
"regex-syntax 0.8.5",
]
[[package]]
......@@ -1612,6 +1661,12 @@ version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "53a49587ad06b26609c52e423de037e7f57f20d53535d66e08c695f347df952a"
[[package]]
name = "regex-syntax"
version = "0.6.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1"
[[package]]
name = "regex-syntax"
version = "0.8.5"
......@@ -1859,6 +1914,10 @@ dependencies = [
"serde",
"serde_json",
"tokio",
"tracing",
"tracing-appender",
"tracing-log",
"tracing-subscriber",
]
[[package]]
......@@ -1872,6 +1931,15 @@ dependencies = [
"digest",
]
[[package]]
name = "sharded-slab"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
dependencies = [
"lazy_static",
]
[[package]]
name = "shlex"
version = "1.3.0"
......@@ -2007,6 +2075,36 @@ dependencies = [
"windows-sys 0.59.0",
]
[[package]]
name = "thiserror"
version = "1.0.69"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
dependencies = [
"thiserror-impl",
]
[[package]]
name = "thiserror-impl"
version = "1.0.69"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "thread_local"
version = "1.1.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c"
dependencies = [
"cfg-if",
"once_cell",
]
[[package]]
name = "time"
version = "0.3.36"
......@@ -2107,22 +2205,90 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3"
[[package]]
name = "tracing"
version = "0.1.40"
version = "0.1.41"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef"
checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0"
dependencies = [
"log",
"pin-project-lite",
"tracing-attributes",
"tracing-core",
]
[[package]]
name = "tracing-appender"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3566e8ce28cc0a3fe42519fc80e6b4c943cc4c8cef275620eb8dac2d3d4e06cf"
dependencies = [
"crossbeam-channel",
"thiserror",
"time",
"tracing-subscriber",
]
[[package]]
name = "tracing-attributes"
version = "0.1.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "tracing-core"
version = "0.1.32"
version = "0.1.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c"
dependencies = [
"once_cell",
"valuable",
]
[[package]]
name = "tracing-log"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54"
checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3"
dependencies = [
"log",
"once_cell",
"tracing-core",
]
[[package]]
name = "tracing-serde"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "704b1aeb7be0d0a84fc9828cae51dab5970fee5088f83d1dd7ee6f6246fc6ff1"
dependencies = [
"serde",
"tracing-core",
]
[[package]]
name = "tracing-subscriber"
version = "0.3.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008"
dependencies = [
"chrono",
"matchers",
"nu-ansi-term",
"once_cell",
"regex",
"serde",
"serde_json",
"sharded-slab",
"smallvec",
"thread_local",
"tracing",
"tracing-core",
"tracing-log",
"tracing-serde",
]
[[package]]
......@@ -2184,6 +2350,12 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "valuable"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
[[package]]
name = "vcpkg"
version = "0.2.15"
......@@ -2301,6 +2473,28 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "winapi"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
dependencies = [
"winapi-i686-pc-windows-gnu",
"winapi-x86_64-pc-windows-gnu",
]
[[package]]
name = "winapi-i686-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
[[package]]
name = "winapi-x86_64-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows-core"
version = "0.52.0"
......
......@@ -25,6 +25,11 @@ env_logger = "0.11.5"
log = "0.4.22"
chrono = "0.4.38"
tokio = "1.42.0"
# Added for enhanced logging system
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "chrono"] }
tracing-log = "0.2"
tracing-appender = "0.2.3"
[profile.release]
lto = "thin"
......
......@@ -67,6 +67,20 @@ $ pip install -e .
**Note:** When modifying Rust code, you must rebuild the wheel for changes to take effect.
### Logging
The SGL Router includes structured logging with console output by default. To enable log files:
```python
# Enable file logging when creating a router
router = Router(
worker_urls=["http://worker1:8000", "http://worker2:8000"],
log_dir="./logs" # Daily log files will be created here
)
```
Use the `--verbose` flag with the CLI for more detailed logs.
### Troubleshooting
1. If rust analyzer is not working in VSCode, set `rust-analyzer.linkedProjects` to the absolute path of `Cargo.toml` in your repo. For example:
......
......@@ -42,6 +42,7 @@ class RouterArgs:
max_tree_size: int = 2**24
max_payload_size: int = 4 * 1024 * 1024 # 4MB
verbose: bool = False
log_dir: Optional[str] = None
@staticmethod
def add_cli_args(
......@@ -142,6 +143,12 @@ class RouterArgs:
action="store_true",
help="Enable verbose logging",
)
parser.add_argument(
f"--{prefix}log-dir",
type=str,
default=None,
help="Directory to store log files. If not specified, logs are only output to console.",
)
@classmethod
def from_cli_args(
......@@ -174,6 +181,7 @@ class RouterArgs:
max_tree_size=getattr(args, f"{prefix}max_tree_size"),
max_payload_size=getattr(args, f"{prefix}max_payload_size"),
verbose=getattr(args, f"{prefix}verbose", False),
log_dir=getattr(args, f"{prefix}log_dir", None),
)
......@@ -220,6 +228,7 @@ def launch_router(args: argparse.Namespace) -> Optional[Router]:
max_tree_size=router_args.max_tree_size,
max_payload_size=router_args.max_payload_size,
verbose=router_args.verbose,
log_dir=router_args.log_dir,
)
router.start()
......
......@@ -31,6 +31,7 @@ class Router:
max_payload_size: Maximum payload size in bytes. Default: 4MB
max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24
verbose: Enable verbose logging. Default: False
log_dir: Directory to store log files. If None, logs are only output to console. Default: None
"""
def __init__(
......@@ -48,6 +49,7 @@ class Router:
max_tree_size: int = 2**24,
max_payload_size: int = 4 * 1024 * 1024, # 4MB
verbose: bool = False,
log_dir: Optional[str] = None,
):
self._router = _Router(
worker_urls=worker_urls,
......@@ -63,6 +65,7 @@ class Router:
max_tree_size=max_tree_size,
max_payload_size=max_payload_size,
verbose=verbose,
log_dir=log_dir,
)
def start(self) -> None:
......
......@@ -37,6 +37,7 @@ class TestLaunchRouter(unittest.TestCase):
max_tree_size=2**24,
max_payload_size=4 * 1024 * 1024, # 4MB
verbose=False,
log_dir=None,
)
def create_router_args(self, **kwargs):
......
......@@ -23,6 +23,7 @@ def popen_launch_router(
policy: str = "cache_aware",
max_payload_size: int = None,
api_key: str = None,
log_dir: str = None,
):
"""
Launch the router server process.
......@@ -35,6 +36,7 @@ def popen_launch_router(
policy: Router policy, one of "cache_aware", "round_robin", "random"
max_payload_size: Maximum payload size in bytes
api_key: API key for the router
log_dir: Directory to store log files. If None, logs are only output to console.
"""
_, host, port = base_url.split(":")
host = host[2:]
......@@ -63,6 +65,9 @@ def popen_launch_router(
if max_payload_size is not None:
command.extend(["--router-max-payload-size", str(max_payload_size)])
if log_dir is not None:
command.extend(["--log-dir", log_dir])
process = subprocess.Popen(command, stdout=None, stderr=None)
start_time = time.time()
......
use pyo3::prelude::*;
pub mod logging;
pub mod router;
pub mod server;
pub mod tree;
#[pyclass(eq)]
#[derive(Clone, PartialEq)]
#[derive(Clone, PartialEq, Debug)]
pub enum PolicyType {
Random,
RoundRobin,
......@@ -12,6 +13,7 @@ pub enum PolicyType {
}
#[pyclass]
#[derive(Debug, Clone, PartialEq)]
struct Router {
host: String,
port: u16,
......@@ -26,6 +28,7 @@ struct Router {
max_tree_size: usize,
max_payload_size: usize,
verbose: bool,
log_dir: Option<String>,
}
#[pymethods]
......@@ -44,7 +47,8 @@ impl Router {
eviction_interval_secs = 60,
max_tree_size = 2usize.pow(24),
max_payload_size = 4 * 1024 * 1024,
verbose = false
verbose = false,
log_dir = None,
))]
fn new(
worker_urls: Vec<String>,
......@@ -60,6 +64,7 @@ impl Router {
max_tree_size: usize,
max_payload_size: usize,
verbose: bool,
log_dir: Option<String>,
) -> PyResult<Self> {
Ok(Router {
host,
......@@ -75,6 +80,7 @@ impl Router {
max_tree_size,
max_payload_size,
verbose,
log_dir,
})
}
......@@ -107,6 +113,7 @@ impl Router {
policy_config,
verbose: self.verbose,
max_payload_size: self.max_payload_size,
log_dir: self.log_dir.clone(),
})
.await
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
......
use std::path::PathBuf;
use tracing::Level;
use tracing_appender::non_blocking::WorkerGuard;
use tracing_appender::rolling::{RollingFileAppender, Rotation};
use tracing_log::LogTracer;
use tracing_subscriber::fmt::time::ChronoUtc;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{EnvFilter, Layer};
/// Configuration for the logging system
#[derive(Debug, Clone)]
pub struct LoggingConfig {
/// Log level for the application (default: INFO)
pub level: Level,
/// Whether to use json format for logs (default: false)
pub json_format: bool,
/// Path to store log files. If None, logs will only go to stdout/stderr
pub log_dir: Option<String>,
/// Whether to colorize logs when output is a terminal (default: true)
pub colorize: bool,
/// Log file name to use if log_dir is specified (default: "sgl-router")
pub log_file_name: String,
/// Custom log targets to filter (default: "sglang_router_rs")
pub log_targets: Option<Vec<String>>,
}
impl Default for LoggingConfig {
fn default() -> Self {
Self {
level: Level::INFO,
json_format: false,
log_dir: None,
colorize: true,
log_file_name: "sgl-router".to_string(),
log_targets: Some(vec!["sglang_router_rs".to_string()]),
}
}
}
/// Guard that keeps the file appender worker thread alive
///
/// This must be kept in scope for the duration of the program
/// to ensure logs are properly written to files
#[allow(dead_code)]
pub struct LogGuard {
_file_guard: Option<WorkerGuard>,
}
/// Initialize the logging system with the given configuration
///
/// # Arguments
/// * `config` - Configuration for the logging system
///
/// # Returns
/// A LogGuard that must be kept alive for the duration of the program
///
/// # Panics
/// Will not panic, as initialization errors are handled gracefully
pub fn init_logging(config: LoggingConfig) -> LogGuard {
// Forward logs to tracing - ignore errors to allow for multiple initialization
let _ = LogTracer::init();
// Convert log level to filter string
let level_filter = match config.level {
Level::TRACE => "trace",
Level::DEBUG => "debug",
Level::INFO => "info",
Level::WARN => "warn",
Level::ERROR => "error",
};
// Create env filter
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
// Format: <target>=<level>,<target2>=<level2>,...
let filter_string = if let Some(targets) = &config.log_targets {
targets
.iter()
.enumerate()
.map(|(i, target)| {
if i > 0 {
format!(",{}={}", target, level_filter)
} else {
format!("{}={}", target, level_filter)
}
})
.collect::<String>()
} else {
format!("sglang_router_rs={}", level_filter)
};
EnvFilter::new(filter_string)
});
// Setup stdout/stderr layer
let mut layers = Vec::new();
// Standard timestamp format: YYYY-MM-DD HH:MM:SS
let time_format = "%Y-%m-%d %H:%M:%S".to_string();
// Configure the console stdout layer
let stdout_layer = tracing_subscriber::fmt::layer()
.with_ansi(config.colorize)
.with_file(true)
.with_line_number(true)
.with_timer(ChronoUtc::new(time_format.clone()));
let stdout_layer = if config.json_format {
stdout_layer.json().flatten_event(true).boxed()
} else {
stdout_layer.boxed()
};
layers.push(stdout_layer);
// Create a file appender if log_dir is specified
let mut file_guard = None;
if let Some(log_dir) = &config.log_dir {
let file_name = config.log_file_name.clone();
let log_dir = PathBuf::from(log_dir);
// Create log directory if it doesn't exist
if !log_dir.exists() {
if let Err(e) = std::fs::create_dir_all(&log_dir) {
eprintln!("Failed to create log directory: {}", e);
return LogGuard { _file_guard: None };
}
}
let file_appender = RollingFileAppender::new(Rotation::DAILY, log_dir, file_name);
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
file_guard = Some(guard);
let file_layer = tracing_subscriber::fmt::layer()
.with_ansi(false) // Never use ANSI colors in log files
.with_file(true)
.with_line_number(true)
.with_timer(ChronoUtc::new(time_format))
.with_writer(non_blocking);
let file_layer = if config.json_format {
file_layer.json().flatten_event(true).boxed()
} else {
file_layer.boxed()
};
layers.push(file_layer);
}
// Initialize the subscriber with all layers
// Use try_init to handle errors gracefully in case another subscriber is already set
let _ = tracing_subscriber::registry()
.with(env_filter)
.with(layers)
.try_init();
// Return the guard to keep the file appender worker thread alive
LogGuard {
_file_guard: file_guard,
}
}
use crate::logging::{self, LoggingConfig};
use crate::router::PolicyConfig;
use crate::router::Router;
use actix_web::{
error, get, post, web, App, Error, HttpRequest, HttpResponse, HttpServer, Responder,
};
use bytes::Bytes;
use env_logger::Builder;
use futures_util::StreamExt;
use log::{info, LevelFilter};
use std::collections::HashMap;
use std::io::Write;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tracing::{info, Level};
#[derive(Debug)]
pub struct AppState {
......@@ -148,30 +148,29 @@ pub struct ServerConfig {
pub policy_config: PolicyConfig,
pub verbose: bool,
pub max_payload_size: usize,
pub log_dir: Option<String>,
}
pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
// Initialize logger
Builder::new()
.format(|buf, record| {
use chrono::Local;
writeln!(
buf,
"[Router (Rust)] {} - {} - {}",
Local::now().format("%Y-%m-%d %H:%M:%S"),
record.level(),
record.args()
)
})
.filter(
None,
if config.verbose {
LevelFilter::Debug
// Only initialize logging if not already done (for Python bindings support)
static LOGGING_INITIALIZED: AtomicBool = AtomicBool::new(false);
let _log_guard = if !LOGGING_INITIALIZED.swap(true, Ordering::SeqCst) {
Some(logging::init_logging(LoggingConfig {
level: if config.verbose {
Level::DEBUG
} else {
LevelFilter::Info
Level::INFO
},
)
.init();
json_format: false,
log_dir: config.log_dir.clone(),
colorize: true,
log_file_name: "sgl-router".to_string(),
log_targets: None,
}))
} else {
None
};
info!("🚧 Initializing router on {}:{}", config.host, config.port);
info!("🚧 Initializing workers on {:?}", config.worker_urls);
......@@ -189,7 +188,7 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
let app_state = web::Data::new(
AppState::new(
config.worker_urls.clone(),
client,
client.clone(), // Clone the client here
config.policy_config.clone(),
)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment