gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
third_party/sglang/sgl-model-gateway/src/observability/events.rs vendored Normal file
View File

@@ -0,0 +1,90 @@
//! Request events for observability and monitoring.
//!
//! Events use DEBUG level when OTEL is disabled, INFO when enabled.
use tracing::{debug, event, Level};
use super::otel_trace::is_otel_enabled;
/// Module path used by CustomOtelFilter to identify events for OTEL export.
#[inline]
pub const fn get_module_path() -> &'static str {
"smg::observability::events"
}
pub trait Event {
fn emit(&self);
}
/// Event emitted when a prefill-decode request pair is sent.
#[derive(Debug, Clone, Copy)]
pub struct RequestPDSentEvent<'a> {
pub prefill_url: &'a str,
pub decode_url: &'a str,
}
impl Event for RequestPDSentEvent<'_> {
#[inline]
fn emit(&self) {
if is_otel_enabled() {
event!(
Level::INFO,
prefill_url = %self.prefill_url,
decode_url = %self.decode_url,
"Sending concurrent requests"
);
} else {
debug!(
prefill_url = %self.prefill_url,
decode_url = %self.decode_url,
"Sending concurrent requests"
);
}
}
}
/// Event emitted when a request is sent to a worker.
#[derive(Debug, Clone, Copy)]
pub struct RequestSentEvent<'a> {
pub url: &'a str,
}
impl Event for RequestSentEvent<'_> {
#[inline]
fn emit(&self) {
if is_otel_enabled() {
event!(Level::INFO, url = %self.url, "Sending request");
} else {
debug!(url = %self.url, "Sending request");
}
}
}
/// Event emitted when concurrent requests are received.
#[derive(Debug, Clone, Copy)]
pub struct RequestReceivedEvent;
impl Event for RequestReceivedEvent {
#[inline]
fn emit(&self) {
if is_otel_enabled() {
event!(Level::INFO, "Received concurrent requests");
} else {
debug!("Received concurrent requests");
}
}
}
#[cfg(test)]
mod tests {
use std::mem::size_of;
use super::*;
#[test]
fn test_event_sizes() {
assert_eq!(size_of::<RequestReceivedEvent>(), 0);
assert_eq!(size_of::<RequestSentEvent>(), 16);
assert_eq!(size_of::<RequestPDSentEvent>(), 32);
}
}

583
third_party/sglang/sgl-model-gateway/src/observability/gauge_histogram.rs vendored Normal file
View File

@@ -0,0 +1,583 @@
//! Non-cumulative gauge histogram for Grafana heatmap visualization.
//!
//! Unlike Prometheus Histogram which uses cumulative `le` buckets, this emits
//! non-cumulative bucket counts with `(gt, le]` ranges suitable for heatmaps.
//!
//! # Design: True Zero-Allocation Hot Path
//!
//! The key insight is that `gauge!` returns a `Gauge` handle that can be stored.
//! By pre-registering all gauge handles at startup, the hot path becomes just
//! N+1 atomic `gauge.set()` calls with zero allocations.
//!
//! # Performance Characteristics
//!
//! Setup (once per label combination):
//! - `register()`: N+1 gauge registrations, N+1 String allocations for gt/le
//!
//! Hot path (`set_counts()`):
//! - **Zero heap allocations**
//! - **Zero key lookups** (handles are pre-registered)
//! - N+1 atomic `gauge.set()` calls
//!
//! # Example
//!
//! ```ignore
//! use crate::observability::gauge_histogram::{BucketBounds, GaugeHistogramVec};
//!
//! // Define at module level
//! static BOUNDS: BucketBounds<10> = BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
//! static HISTOGRAM: GaugeHistogramVec<10> = GaugeHistogramVec::new("smg_request_dist", &BOUNDS);
//!
//! // At startup: register for each label combination
//! let handle = HISTOGRAM.register(&[("router", "round_robin"), ("model", "llama")]);
//!
//! // Pre-allocate counts buffer
//! let mut counts = vec![0usize; BOUNDS.bucket_count()];
//!
//! // Hot path: TRUE zero allocation
//! fn update(handle: &GaugeHistogramHandle, counts: &mut [usize], observations: &[u64]) {
//! BOUNDS.compute_counts_into(counts, observations);
//! handle.set_counts(counts); // Just N+1 atomic gauge.set() calls!
//! }
//! ```
use std::sync::Arc;
use dashmap::DashMap;
use metrics::{gauge, Label};
// =============================================================================
// BUCKET BOUNDS
// =============================================================================
/// Static bucket boundary configuration.
///
/// Uses const generics to define bucket bounds at compile time with validation.
/// The bounds define `N` upper limits, creating `N + 1` buckets:
/// `(0, b[0]], (b[0], b[1]], ..., (b[N-1], +Inf]`.
#[derive(Debug)]
pub struct BucketBounds<const N: usize> {
bounds: [u64; N],
}
impl<const N: usize> BucketBounds<N> {
/// Create new bucket bounds from a sorted array of upper limits.
///
/// # Panics
///
/// Panics at compile time (in const context) or runtime if bounds are not
/// strictly ascending.
#[must_use]
pub const fn new(bounds: [u64; N]) -> Self {
let mut i = 1;
while i < N {
assert!(
bounds[i] > bounds[i - 1],
"bucket bounds must be strictly ascending"
);
i += 1;
}
Self { bounds }
}
/// Returns the number of buckets (one more than the number of bounds).
#[inline]
#[must_use]
pub const fn bucket_count(&self) -> usize {
N + 1
}
/// Returns the number of bounds.
#[inline]
#[must_use]
pub const fn bound_count(&self) -> usize {
N
}
/// Get the bounds array.
#[inline]
#[must_use]
pub const fn bounds(&self) -> &[u64; N] {
&self.bounds
}
/// Find the bucket index for a value. O(log N).
#[inline]
#[must_use]
pub fn bucket_index(&self, value: u64) -> usize {
self.bounds.partition_point(|&bound| bound < value)
}
/// Get the upper bound for a bucket index, or None for the +Inf bucket.
#[inline]
#[must_use]
pub const fn upper_bound(&self, idx: usize) -> Option<u64> {
if idx < N {
Some(self.bounds[idx])
} else {
None
}
}
/// Get the lower bound for a bucket index (0 for the first bucket).
#[inline]
#[must_use]
pub const fn lower_bound(&self, idx: usize) -> u64 {
if idx == 0 {
0
} else {
self.bounds[idx - 1]
}
}
/// Compute bucket counts into a pre-allocated buffer. **Zero allocation.**
///
/// # Panics
///
/// Panics if `counts.len() < bucket_count()`.
#[inline]
pub fn compute_counts_into(&self, counts: &mut [usize], observations: &[u64]) {
debug_assert!(
counts.len() >= self.bucket_count(),
"counts buffer too small"
);
counts[..self.bucket_count()].fill(0);
for &value in observations {
let idx = self.bucket_index(value);
counts[idx] += 1;
}
}
/// Compute bucket counts, allocating a new Vec.
///
/// Prefer `compute_counts_into` in hot paths to avoid allocation.
#[must_use]
pub fn compute_counts(&self, observations: &[u64]) -> Vec<usize> {
let mut counts = vec![0usize; self.bucket_count()];
self.compute_counts_into(&mut counts, observations);
counts
}
}
// =============================================================================
// GAUGE HISTOGRAM HANDLE (pre-registered, zero-alloc hot path)
// =============================================================================
/// Pre-registered gauge handles for a histogram with specific labels.
///
/// This is what you use in the hot path. Calling `set_counts()` does only
/// N+1 atomic `gauge.set()` operations - zero allocations, zero lookups.
#[derive(Clone)]
pub struct GaugeHistogramHandle {
gauges: Vec<metrics::Gauge>,
}
impl GaugeHistogramHandle {
/// Set bucket counts. **TRUE zero allocation.**
///
/// Just N+1 atomic `gauge.set()` calls - no key lookup, no allocation.
#[inline]
pub fn set_counts(&self, counts: &[usize]) {
debug_assert_eq!(
counts.len(),
self.gauges.len(),
"counts length must match bucket count"
);
for (gauge, &count) in self.gauges.iter().zip(counts.iter()) {
gauge.set(count as f64);
}
}
/// Number of buckets.
#[inline]
pub fn bucket_count(&self) -> usize {
self.gauges.len()
}
/// Zero out all gauges. **Zero allocation.**
#[inline]
pub fn zero_counts(&self) {
for gauge in &self.gauges {
gauge.set(0.0);
}
}
}
// =============================================================================
// GAUGE HISTOGRAM VEC (factory for registering handles)
// =============================================================================
/// Factory for creating pre-registered histogram handles.
///
/// Define as a static constant, then call `register()` for each label combination
/// you need. The returned `GaugeHistogramHandle` provides zero-allocation updates.
#[derive(Debug)]
pub struct GaugeHistogramVec<const N: usize> {
name: &'static str,
bounds: &'static BucketBounds<N>,
}
impl<const N: usize> GaugeHistogramVec<N> {
/// Create a new gauge histogram factory.
///
/// This just stores the name and bounds - no allocation or registration yet.
#[must_use]
pub const fn new(name: &'static str, bounds: &'static BucketBounds<N>) -> Self {
Self { name, bounds }
}
#[inline]
pub const fn name(&self) -> &'static str {
self.name
}
#[inline]
pub const fn bounds(&self) -> &BucketBounds<N> {
self.bounds
}
/// Register gauges for a specific label combination.
///
/// Call this once per unique label combination (at startup or when first seen).
/// The returned handle can be cloned cheaply (just Arc clones internally).
///
/// # Arguments
///
/// - `labels`: Static key-value label pairs for this histogram instance
///
/// # Example
///
/// ```ignore
/// let handle = HISTOGRAM.register(&[("router", "round_robin"), ("model", "llama")]);
/// ```
pub fn register(&self, labels: &[(&'static str, &str)]) -> GaugeHistogramHandle {
let bucket_count = self.bounds.bucket_count();
let mut gauges = Vec::with_capacity(bucket_count);
for i in 0..bucket_count {
// Build gt/le labels for this bucket
let gt_str = if i == 0 {
"0".to_string()
} else {
self.bounds.bounds[i - 1].to_string()
};
let le_str = if i < N {
self.bounds.bounds[i].to_string()
} else {
"+Inf".to_string()
};
// Build complete label set
let mut all_labels: Vec<Label> = Vec::with_capacity(labels.len() + 2);
for &(k, v) in labels {
all_labels.push(Label::new(k, v.to_string()));
}
all_labels.push(Label::new("gt", gt_str));
all_labels.push(Label::new("le", le_str));
// Register and store the gauge handle
let g = gauge!(self.name, all_labels);
gauges.push(g);
}
GaugeHistogramHandle { gauges }
}
/// Register with no additional labels (just gt/le).
pub fn register_no_labels(&self) -> GaugeHistogramHandle {
self.register(&[])
}
}
// =============================================================================
// CACHED GAUGE HISTOGRAM (for dynamic labels discovered at runtime)
// =============================================================================
/// A gauge histogram with automatic handle caching for dynamic labels.
///
/// Use this when label values (like worker names) are discovered at runtime.
/// Handles are registered on first use and cached for subsequent calls.
///
/// # Example
///
/// ```ignore
/// static BOUNDS: BucketBounds<10> = BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
/// static HISTOGRAM: GaugeHistogramVec<10> = GaugeHistogramVec::new("smg_worker_dist", &BOUNDS);
///
/// // Create cached wrapper (do this once, store in your router state)
/// let cached = CachedGaugeHistogram::new(&HISTOGRAM);
///
/// // Hot path - first call registers, subsequent calls use cached handle
/// cached.observe("worker-1", &request_counts);
/// cached.observe("worker-2", &request_counts);
/// cached.observe("worker-1", &request_counts); // Uses cached handle
/// ```
pub struct CachedGaugeHistogram<const N: usize> {
histogram: &'static GaugeHistogramVec<N>,
/// Cache of label value -> (handle, counts_buffer)
cache: DashMap<Arc<str>, (GaugeHistogramHandle, Vec<usize>)>,
/// Static label key (e.g., "worker", "model")
label_key: &'static str,
}
impl<const N: usize> CachedGaugeHistogram<N> {
/// Create a new cached histogram for a single dynamic label.
///
/// # Arguments
///
/// - `histogram`: The static histogram factory
/// - `label_key`: The label key for the dynamic value (e.g., "worker")
pub fn new(histogram: &'static GaugeHistogramVec<N>, label_key: &'static str) -> Self {
Self {
histogram,
cache: DashMap::new(),
label_key,
}
}
/// Get or create a handle for the given label value.
///
/// First call for a label value registers the gauges (allocates).
/// Subsequent calls return the cached handle (no allocation).
/// Thread-safe: uses entry API to avoid race conditions.
pub fn get_or_register(
&self,
label_value: &str,
) -> dashmap::mapref::one::Ref<'_, Arc<str>, (GaugeHistogramHandle, Vec<usize>)> {
// Fast path: already cached
if let Some(entry) = self.cache.get(label_value) {
return entry;
}
// Slow path: use entry API to handle concurrent inserts atomically
self.cache.entry(Arc::from(label_value)).or_insert_with(|| {
let handle = self.histogram.register(&[(self.label_key, label_value)]);
let counts_buf = vec![0usize; self.histogram.bounds.bucket_count()];
(handle, counts_buf)
});
self.cache.get(label_value).unwrap()
}
/// Observe a distribution for a label value. **Zero allocation after first call.**
///
/// First call for a new label value registers gauges (allocates).
/// All subsequent calls are zero-allocation.
/// Thread-safe: uses entry API to avoid race conditions.
pub fn observe(&self, label_value: &str, observations: &[u64]) {
// Fast path: existing entry
if let Some(mut entry) = self.cache.get_mut(label_value) {
let (ref handle, ref mut counts_buf) = entry.value_mut();
self.histogram
.bounds
.compute_counts_into(counts_buf, observations);
handle.set_counts(counts_buf);
return;
}
// Slow path: use entry API to handle concurrent inserts atomically
let mut entry = self.cache.entry(Arc::from(label_value)).or_insert_with(|| {
let handle = self.histogram.register(&[(self.label_key, label_value)]);
let counts_buf = vec![0usize; self.histogram.bounds.bucket_count()];
(handle, counts_buf)
});
let (ref handle, ref mut counts_buf) = entry.value_mut();
self.histogram
.bounds
.compute_counts_into(counts_buf, observations);
handle.set_counts(counts_buf);
}
/// Number of cached label combinations.
pub fn cache_size(&self) -> usize {
self.cache.len()
}
/// Remove a worker and zero out its metrics. **Zero allocation.**
///
/// Call this when a worker is removed from the pool.
/// Sets all bucket counts to 0 (so Grafana shows it as empty).
///
/// Note: The gauge handles remain in the Prometheus registry (the `metrics`
/// crate doesn't support unregistering). But memory in our cache is freed.
pub fn remove(&self, label_value: &str) {
if let Some((_, (handle, _))) = self.cache.remove(label_value) {
handle.zero_counts();
}
}
/// Remove workers not in the provided set.
///
/// Call this periodically with your current active workers to clean up stale entries.
/// Uses `DashMap::retain` for atomic operation without intermediate allocation.
///
/// # Example
///
/// ```ignore
/// let active: HashSet<&str> = workers.iter().map(|w| w.name.as_str()).collect();
/// cached.retain_only(&active);
/// ```
pub fn retain_only<S: std::borrow::Borrow<str> + std::hash::Hash + Eq>(
&self,
active_labels: &std::collections::HashSet<S>,
) {
self.cache.retain(|key, (handle, _)| {
if active_labels.contains(key.as_ref()) {
true
} else {
handle.zero_counts();
false
}
});
}
/// Get all currently tracked label values.
pub fn tracked_labels(&self) -> Vec<Arc<str>> {
self.cache.iter().map(|e| Arc::clone(e.key())).collect()
}
}
// =============================================================================
// CONVENIENCE CONSTANTS
// =============================================================================
/// Common bucket bounds for request counts.
pub static REQUEST_COUNT_BOUNDS: BucketBounds<10> =
BucketBounds::new([1, 2, 3, 5, 7, 10, 20, 50, 100, 200]);
// =============================================================================
// TESTS
// =============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bucket_bounds_creation() {
let bounds = BucketBounds::new([10, 30, 60]);
assert_eq!(bounds.bucket_count(), 4);
assert_eq!(bounds.bound_count(), 3);
}
#[test]
fn test_bucket_bounds_const_creation() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
assert_eq!(BOUNDS.bucket_count(), 4);
}
#[test]
#[should_panic(expected = "bucket bounds must be strictly ascending")]
fn test_bucket_bounds_not_ascending_panics() {
let _ = BucketBounds::new([10, 5, 60]);
}
#[test]
fn test_bucket_index() {
let bounds = BucketBounds::new([10, 30, 60]);
assert_eq!(bounds.bucket_index(0), 0);
assert_eq!(bounds.bucket_index(10), 0);
assert_eq!(bounds.bucket_index(11), 1);
assert_eq!(bounds.bucket_index(30), 1);
assert_eq!(bounds.bucket_index(31), 2);
assert_eq!(bounds.bucket_index(60), 2);
assert_eq!(bounds.bucket_index(61), 3);
}
#[test]
fn test_compute_counts() {
let bounds = BucketBounds::new([10, 30, 60]);
assert_eq!(
bounds.compute_counts(&[5, 10, 15, 40, 100]),
vec![2, 1, 1, 1]
);
}
#[test]
fn test_compute_counts_into() {
let bounds = BucketBounds::new([10, 30, 60]);
let mut counts = [0usize; 4];
bounds.compute_counts_into(&mut counts, &[5, 10, 15, 40, 100]);
assert_eq!(counts, [2, 1, 1, 1]);
}
#[test]
fn test_gauge_histogram_vec_creation() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_metric", &BOUNDS);
assert_eq!(HISTOGRAM.name(), "test_metric");
assert_eq!(HISTOGRAM.bounds().bucket_count(), 4);
}
#[test]
fn test_gauge_histogram_handle_registration() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_hist", &BOUNDS);
let handle = HISTOGRAM.register(&[("router", "rr")]);
assert_eq!(handle.bucket_count(), 4);
// This should be zero-allocation
handle.set_counts(&[1, 2, 3, 4]);
}
#[test]
fn test_request_count_bounds() {
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_count(), 11);
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(1), 0);
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(2), 1);
assert_eq!(REQUEST_COUNT_BOUNDS.bucket_index(201), 10);
}
#[test]
fn test_cached_histogram() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> = GaugeHistogramVec::new("test_cached", &BOUNDS);
let cached = CachedGaugeHistogram::new(&HISTOGRAM, "worker");
// First call registers
cached.observe("worker-1", &[5, 15, 45, 100]);
assert_eq!(cached.cache_size(), 1);
// Second call uses cache
cached.observe("worker-1", &[1, 2, 3]);
assert_eq!(cached.cache_size(), 1);
// New worker registers
cached.observe("worker-2", &[10, 20, 30]);
assert_eq!(cached.cache_size(), 2);
}
#[test]
fn test_cached_histogram_removal() {
static BOUNDS: BucketBounds<3> = BucketBounds::new([10, 30, 60]);
static HISTOGRAM: GaugeHistogramVec<3> =
GaugeHistogramVec::new("test_cached_remove", &BOUNDS);
let cached = CachedGaugeHistogram::new(&HISTOGRAM, "worker");
// Add some workers
cached.observe("worker-1", &[5, 15]);
cached.observe("worker-2", &[10, 20]);
cached.observe("worker-3", &[1, 2]);
assert_eq!(cached.cache_size(), 3);
// Remove one
cached.remove("worker-2");
assert_eq!(cached.cache_size(), 2);
// retain_only
let active: std::collections::HashSet<&str> = ["worker-1"].into_iter().collect();
cached.retain_only(&active);
assert_eq!(cached.cache_size(), 1);
// Check tracked labels
let labels = cached.tracked_labels();
assert_eq!(labels.len(), 1);
assert_eq!(&*labels[0], "worker-1");
}
}

195
third_party/sglang/sgl-model-gateway/src/observability/inflight_tracker.rs vendored Normal file
View File

@@ -0,0 +1,195 @@
use std::{
sync::{
atomic::{AtomicU64, Ordering},
Arc, LazyLock, OnceLock,
},
time::Instant,
};
use dashmap::DashMap;
use super::gauge_histogram::{BucketBounds, GaugeHistogramHandle, GaugeHistogramVec};
use crate::policies::utils::PeriodicTask;
static INFLIGHT_AGE_BOUNDS: BucketBounds<11> =
BucketBounds::new([30, 60, 180, 300, 600, 1200, 3600, 7200, 14400, 28800, 86400]);
static INFLIGHT_AGE_HISTOGRAM: GaugeHistogramVec<11> =
GaugeHistogramVec::new("smg_http_inflight_request_age_count", &INFLIGHT_AGE_BOUNDS);
static INFLIGHT_AGE_HANDLE: LazyLock<GaugeHistogramHandle> =
LazyLock::new(|| INFLIGHT_AGE_HISTOGRAM.register_no_labels());
pub struct InFlightRequestTracker {
requests: DashMap<u64, Instant>,
next_id: AtomicU64,
sampler: OnceLock<PeriodicTask>,
}
impl InFlightRequestTracker {
pub fn new() -> Arc<Self> {
Arc::new(Self {
requests: DashMap::new(),
next_id: AtomicU64::new(0),
sampler: OnceLock::new(),
})
}
pub fn start_sampler(self: &Arc<Self>, interval_secs: u64) {
let tracker = self.clone();
let task = PeriodicTask::spawn(interval_secs, "InFlightRequestSampler", move || {
tracker.sample_and_record();
});
self.sampler.set(task).unwrap();
}
pub fn track(self: &Arc<Self>) -> InFlightGuard {
let request_id = self.next_id.fetch_add(1, Ordering::Relaxed);
self.requests.insert(request_id, Instant::now());
InFlightGuard {
tracker: self.clone(),
request_id,
}
}
pub fn len(&self) -> usize {
self.requests.len()
}
pub fn is_empty(&self) -> bool {
self.requests.is_empty()
}
pub fn compute_bucket_counts(&self) -> Vec<usize> {
let ages = self.collect_ages();
INFLIGHT_AGE_BOUNDS.compute_counts(&ages)
}
fn collect_ages(&self) -> Vec<u64> {
let now = Instant::now();
self.requests
.iter()
.map(|entry| now.duration_since(*entry.value()).as_secs())
.collect()
}
fn sample_and_record(&self) {
let counts = self.compute_bucket_counts();
INFLIGHT_AGE_HANDLE.set_counts(&counts);
}
}
pub struct InFlightGuard {
tracker: Arc<InFlightRequestTracker>,
request_id: u64,
}
impl Drop for InFlightGuard {
fn drop(&mut self) {
self.tracker.requests.remove(&self.request_id);
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
impl InFlightRequestTracker {
fn insert_with_time(&self, request_id: u64, start_time: Instant) {
self.requests.insert(request_id, start_time);
}
}
#[test]
fn test_track_and_drop() {
let tracker = InFlightRequestTracker::new();
{
let _guard1 = tracker.track();
let _guard2 = tracker.track();
assert_eq!(tracker.len(), 2);
}
assert_eq!(tracker.len(), 0);
}
#[test]
fn test_guard_auto_deregister() {
let tracker = InFlightRequestTracker::new();
let guard = tracker.track();
assert_eq!(tracker.len(), 1);
drop(guard);
assert_eq!(tracker.len(), 0);
}
#[test]
fn test_request_age_tracking() {
let tracker = InFlightRequestTracker::new();
let _guard = tracker.track();
std::thread::sleep(Duration::from_millis(100));
let entry = tracker.requests.iter().next().unwrap();
let age = entry.value().elapsed();
assert!(age >= Duration::from_millis(100));
}
#[test]
fn test_collect_ages_empty() {
let tracker = InFlightRequestTracker::new();
let ages = tracker.collect_ages();
assert!(ages.is_empty());
}
#[test]
fn test_collect_ages() {
let tracker = InFlightRequestTracker::new();
let now = Instant::now();
tracker.insert_with_time(1, now);
tracker.insert_with_time(2, now - Duration::from_secs(45));
tracker.insert_with_time(3, now - Duration::from_secs(100));
let ages = tracker.collect_ages();
assert_eq!(ages.len(), 3);
// Ages should be approximately 0, 45, 100 (order may vary due to DashMap)
let mut sorted_ages = ages.clone();
sorted_ages.sort();
assert!(sorted_ages[0] <= 1); // ~0s
assert!((44..=46).contains(&sorted_ages[1])); // ~45s
assert!((99..=101).contains(&sorted_ages[2])); // ~100s
}
#[test]
fn test_concurrent_tracking() {
use std::thread;
let tracker = InFlightRequestTracker::new();
let mut handles = vec![];
for _ in 0..10 {
let t = tracker.clone();
handles.push(thread::spawn(move || {
(0..100).map(|_| t.track()).collect::<Vec<_>>()
}));
}
let all_guards: Vec<_> = handles
.into_iter()
.flat_map(|h| h.join().unwrap())
.collect();
assert_eq!(tracker.len(), 1000);
drop(all_guards);
assert_eq!(tracker.len(), 0);
}
#[test]
fn test_unique_ids() {
let tracker = InFlightRequestTracker::new();
let g1 = tracker.track();
let g2 = tracker.track();
let g3 = tracker.track();
assert_ne!(g1.request_id, g2.request_id);
assert_ne!(g2.request_id, g3.request_id);
assert_eq!(tracker.len(), 3);
}
}

174
third_party/sglang/sgl-model-gateway/src/observability/logging.rs vendored Normal file
View File

@@ -0,0 +1,174 @@
//! Logging infrastructure with non-blocking file I/O.
use std::path::PathBuf;
use tracing::Level;
use tracing_appender::{
non_blocking::WorkerGuard,
rolling::{RollingFileAppender, Rotation},
};
use tracing_log::LogTracer;
use tracing_subscriber::{
fmt::time::ChronoUtc, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer,
};
use super::otel_trace::get_otel_layer;
use crate::config::TraceConfig;
const TIME_FORMAT: &str = "%Y-%m-%d %H:%M:%S";
const DEFAULT_LOG_TARGET: &str = "smg";
#[derive(Debug, Clone)]
pub struct LoggingConfig {
pub level: Level,
pub json_format: bool,
pub log_dir: Option<String>,
pub colorize: bool,
pub log_file_name: String,
pub log_targets: Option<Vec<String>>,
}
impl Default for LoggingConfig {
#[inline]
fn default() -> Self {
Self {
level: Level::INFO,
json_format: false,
log_dir: None,
colorize: true,
log_file_name: "smg".to_string(),
log_targets: Some(vec![DEFAULT_LOG_TARGET.to_string()]),
}
}
}
/// Guard that keeps the file appender thread alive.
#[allow(dead_code)]
pub struct LogGuard {
_file_guard: Option<WorkerGuard>,
}
#[inline]
const fn level_to_str(level: Level) -> &'static str {
match level {
Level::TRACE => "trace",
Level::DEBUG => "debug",
Level::INFO => "info",
Level::WARN => "warn",
Level::ERROR => "error",
}
}
#[inline]
fn build_filter_string(targets: &[String], level_filter: &str) -> String {
// Exact capacity: sum of target lengths + "=" and level per target + commas between
let capacity = targets.iter().map(String::len).sum::<usize>()
+ targets.len() * (1 + level_filter.len())
+ targets.len().saturating_sub(1);
let mut filter_string = String::with_capacity(capacity);
for (i, target) in targets.iter().enumerate() {
if i > 0 {
filter_string.push(',');
}
filter_string.push_str(target);
filter_string.push('=');
filter_string.push_str(level_filter);
}
filter_string
}
pub fn init_logging(config: LoggingConfig, otel_layer_config: Option<TraceConfig>) -> LogGuard {
let _ = LogTracer::init();
let level_filter = level_to_str(config.level);
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
let filter_string = match &config.log_targets {
Some(targets) if !targets.is_empty() => build_filter_string(targets, level_filter),
_ => {
let mut s =
String::with_capacity(DEFAULT_LOG_TARGET.len() + 1 + level_filter.len());
s.push_str(DEFAULT_LOG_TARGET);
s.push('=');
s.push_str(level_filter);
s
}
};
EnvFilter::new(filter_string)
});
let mut layers = Vec::with_capacity(3);
let stdout_layer = tracing_subscriber::fmt::layer()
.with_ansi(config.colorize)
.with_file(true)
.with_line_number(true)
.with_timer(ChronoUtc::new(TIME_FORMAT.to_string()));
let stdout_layer = if config.json_format {
stdout_layer.json().flatten_event(true).boxed()
} else {
stdout_layer.boxed()
};
layers.push(stdout_layer);
let mut file_guard = None;
if let Some(log_dir) = &config.log_dir {
let log_dir = PathBuf::from(log_dir);
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, &config.log_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)
.with_file(true)
.with_line_number(true)
.with_timer(ChronoUtc::new(TIME_FORMAT.to_string()))
.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);
}
if let Some(otel_layer_config) = &otel_layer_config {
if otel_layer_config.enable_trace {
match get_otel_layer() {
Ok(otel_layer) => {
layers.push(otel_layer);
}
Err(e) => {
eprintln!("Failed to initialize OpenTelemetry: {}", e);
}
}
}
}
let _ = tracing_subscriber::registry()
.with(env_filter)
.with(layers)
.try_init();
LogGuard {
_file_guard: file_guard,
}
}

1516
third_party/sglang/sgl-model-gateway/src/observability/metrics.rs vendored Normal file
View File

File diff suppressed because it is too large Load Diff

8
third_party/sglang/sgl-model-gateway/src/observability/mod.rs vendored Normal file
View File

@@ -0,0 +1,8 @@
//! Observability utilities for logging, metrics, and tracing.
pub mod events;
pub mod gauge_histogram;
pub mod inflight_tracker;
pub mod logging;
pub mod metrics;
pub mod otel_trace;

279
third_party/sglang/sgl-model-gateway/src/observability/otel_trace.rs vendored Normal file
View File

@@ -0,0 +1,279 @@
//! OpenTelemetry tracing integration.
use std::{
sync::{
atomic::{AtomicBool, Ordering},
OnceLock,
},
time::Duration,
};
use anyhow::Result;
use axum::http::{HeaderMap, HeaderName, HeaderValue};
use opentelemetry::{global, trace::TracerProvider as _, KeyValue};
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{
propagation::TraceContextPropagator,
runtime,
trace::{BatchConfigBuilder, BatchSpanProcessor, Tracer as SdkTracer, TracerProvider},
Resource,
};
use tokio::task::spawn_blocking;
use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
use tracing::{Metadata, Subscriber};
use tracing_opentelemetry::{self, OpenTelemetrySpanExt};
use tracing_subscriber::{
layer::{Context, Filter},
Layer,
};
use super::events::get_module_path as events_module_path;
/// Whether OpenTelemetry tracing is enabled.
///
/// This flag guards access to TRACER and PROVIDER. We use Release/Acquire
/// ordering to ensure proper synchronization: writes to TRACER/PROVIDER
/// happen-before the Release store, and Acquire loads happen-before reads.
static ENABLED: AtomicBool = AtomicBool::new(false);
static TRACER: OnceLock<SdkTracer> = OnceLock::new();
static PROVIDER: OnceLock<TracerProvider> = OnceLock::new();
static ALLOWED_TARGETS: OnceLock<[&'static str; 3]> = OnceLock::new();
#[inline]
fn get_allowed_targets() -> &'static [&'static str; 3] {
ALLOWED_TARGETS.get_or_init(|| {
[
"smg::otel-trace",
"smg::observability::otel_trace",
events_module_path(),
]
})
}
/// Filter that only allows specific module targets to be exported to OTEL.
#[derive(Clone, Copy, Default)]
pub(crate) struct CustomOtelFilter;
impl CustomOtelFilter {
#[inline]
pub const fn new() -> Self {
Self
}
#[inline]
fn is_allowed(target: &str) -> bool {
get_allowed_targets()
.iter()
.any(|allowed| target.starts_with(allowed))
}
}
impl<S> Filter<S> for CustomOtelFilter
where
S: Subscriber,
{
#[inline]
fn enabled(&self, meta: &Metadata<'_>, _cx: &Context<'_, S>) -> bool {
Self::is_allowed(meta.target())
}
#[inline]
fn callsite_enabled(&self, meta: &'static Metadata<'static>) -> tracing::subscriber::Interest {
if Self::is_allowed(meta.target()) {
tracing::subscriber::Interest::always()
} else {
tracing::subscriber::Interest::never()
}
}
}
pub fn otel_tracing_init(enable: bool, otlp_endpoint: Option<&str>) -> Result<()> {
if !enable {
// Use Release to ensure any prior OTEL state changes are visible
ENABLED.store(false, Ordering::Release);
return Ok(());
}
let endpoint = otlp_endpoint.unwrap_or("localhost:4317");
let endpoint = if !endpoint.starts_with("http://") && !endpoint.starts_with("https://") {
format!("http://{}", endpoint)
} else {
endpoint.to_string()
};
global::set_text_map_propagator(TraceContextPropagator::new());
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(&endpoint)
.with_protocol(opentelemetry_otlp::Protocol::Grpc)
.build()
.map_err(|e| {
eprintln!("[tracing] Failed to create OTLP exporter: {}", e);
anyhow::anyhow!("Failed to create OTLP exporter: {}", e)
})?;
let batch_config = BatchConfigBuilder::default()
.with_scheduled_delay(Duration::from_millis(500))
.with_max_export_batch_size(64)
.build();
let span_processor = BatchSpanProcessor::builder(exporter, runtime::Tokio)
.with_batch_config(batch_config)
.build();
let resource =
Resource::default().merge(&Resource::new(vec![KeyValue::new("service.name", "smg")]));
let provider = TracerProvider::builder()
.with_span_processor(span_processor)
.with_resource(resource)
.build();
PROVIDER
.set(provider.clone())
.map_err(|_| anyhow::anyhow!("Provider already initialized"))?;
let tracer = provider.tracer("smg");
TRACER
.set(tracer)
.map_err(|_| anyhow::anyhow!("Tracer already initialized"))?;
let _ = global::set_tracer_provider(provider);
// Use Release ordering: all writes to TRACER/PROVIDER happen-before this store,
// so any thread that loads ENABLED with Acquire will see the initialized state.
ENABLED.store(true, Ordering::Release);
eprintln!("[tracing] OpenTelemetry initialized successfully");
Ok(())
}
/// Get the OpenTelemetry tracing layer. Must be called after `otel_tracing_init`.
pub fn get_otel_layer<S>() -> Result<Box<dyn Layer<S> + Send + Sync + 'static>>
where
S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a> + Send + Sync,
{
if !is_otel_enabled() {
anyhow::bail!("OpenTelemetry is not enabled");
}
let tracer = TRACER
.get()
.ok_or_else(|| anyhow::anyhow!("Tracer not initialized. Call otel_tracing_init first."))?
.clone();
let layer = tracing_opentelemetry::layer()
.with_tracer(tracer)
.with_filter(CustomOtelFilter::new());
Ok(Box::new(layer))
}
/// Check if OpenTelemetry tracing is enabled.
///
/// Uses Acquire ordering to synchronize with the Release store in `otel_tracing_init`,
/// ensuring that if this returns true, TRACER and PROVIDER are fully initialized.
#[inline]
pub fn is_otel_enabled() -> bool {
ENABLED.load(Ordering::Acquire)
}
pub async fn flush_spans_async() -> Result<()> {
if !is_otel_enabled() {
return Ok(());
}
let provider = PROVIDER
.get()
.ok_or_else(|| anyhow::anyhow!("Provider not initialized"))?
.clone();
spawn_blocking(move || provider.force_flush())
.await
.map_err(|e| anyhow::anyhow!("Failed to flush spans: {}", e))?;
Ok(())
}
pub fn shutdown_otel() {
// Use Acquire to ensure we see any prior OTEL operations
if ENABLED.load(Ordering::Acquire) {
global::shutdown_tracer_provider();
// Use Release to ensure shutdown completes before flag is cleared
ENABLED.store(false, Ordering::Release);
eprintln!("[tracing] OpenTelemetry shut down");
}
}
/// Inject W3C trace context headers into an HTTP request.
#[inline]
pub fn inject_trace_context_http(headers: &mut HeaderMap) {
if !is_otel_enabled() {
return;
}
let context = tracing::Span::current().context();
struct HeaderInjector<'a>(&'a mut HeaderMap);
impl opentelemetry::propagation::Injector for HeaderInjector<'_> {
#[inline]
fn set(&mut self, key: &str, value: String) {
if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
if let Ok(header_value) = HeaderValue::from_str(&value) {
self.0.insert(header_name, header_value);
}
}
}
}
global::get_text_map_propagator(|propagator| {
propagator.inject_context(&context, &mut HeaderInjector(headers));
});
}
/// Inject W3C trace context into gRPC metadata.
#[inline]
pub fn inject_trace_context_grpc(metadata: &mut MetadataMap) {
if !is_otel_enabled() {
return;
}
let context = tracing::Span::current().context();
struct MetadataInjector<'a>(&'a mut MetadataMap);
impl opentelemetry::propagation::Injector for MetadataInjector<'_> {
#[inline]
fn set(&mut self, key: &str, value: String) {
if let Ok(metadata_key) = MetadataKey::from_bytes(key.as_bytes()) {
if let Ok(metadata_value) = MetadataValue::try_from(&value) {
self.0.insert(metadata_key, metadata_value);
}
}
}
}
global::get_text_map_propagator(|propagator| {
propagator.inject_context(&context, &mut MetadataInjector(metadata));
});
}
/// OpenTelemetry trace injector implementing the `smg_grpc_client::TraceInjector` trait.
///
/// This bridges sglang's OTel integration with the `smg-grpc-client` crate's
/// trace injection interface, enabling distributed tracing across gRPC calls.
#[derive(Clone, Default)]
pub struct OtelTraceInjector;
impl smg_grpc_client::TraceInjector for OtelTraceInjector {
fn inject(
&self,
metadata: &mut MetadataMap,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
inject_trace_context_grpc(metadata);
Ok(())
}
}