[sgl-router] refactor - move policy-required states under src/state (#40272)
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
acd20a516e
commit
4a9dc5c4af
@@ -90,7 +90,8 @@ async fn static_urls_pd_role_resolved_end_to_end() {
|
||||
use axum::{routing::get, Json, Router};
|
||||
use serde_json::json;
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ObservabilityConfig, ProxyConfig, ServerConfig,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ObservabilityConfig, ProxyConfig,
|
||||
ServerConfig,
|
||||
};
|
||||
use sgl_router::discovery::{spawn_discovery, WorkerId};
|
||||
use sgl_router::workers::{manager, WorkerRegistry};
|
||||
@@ -146,7 +147,7 @@ async fn static_urls_pd_role_resolved_end_to_end() {
|
||||
urls: vec![url.clone()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
|
||||
@@ -5,10 +5,8 @@ use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use sgl_kv_indexer::PrefixOutcome;
|
||||
use sgl_router::policies::kv_events::{
|
||||
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
|
||||
};
|
||||
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
|
||||
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
|
||||
|
||||
#[test]
|
||||
fn radix_tree_reports_contiguous_prefix_depth_per_worker() {
|
||||
|
||||
@@ -12,8 +12,10 @@ use sgl_router::policies::decode::{
|
||||
resolve_decode_with_capacity_fallback, DecodePolicy, DecodePowerOfTwoPolicy,
|
||||
DecodeSelectionContext, LegacyHostAffinityDecodePolicy,
|
||||
};
|
||||
use sgl_router::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
|
||||
use sgl_router::policies::SelectionProposal;
|
||||
use sgl_router::state::load_monitor::engine_reported_load::{
|
||||
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad,
|
||||
};
|
||||
use sgl_router::workers::Worker;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::Ordering;
|
||||
@@ -30,15 +32,15 @@ fn worker(id: &str) -> Arc<Worker> {
|
||||
}))
|
||||
}
|
||||
|
||||
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineLoadSnapshot {
|
||||
EngineLoadSnapshot::from_native_cache_workers(
|
||||
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineReportedLoadSnapshot {
|
||||
EngineReportedLoadSnapshot::from_native_cache_workers(
|
||||
7,
|
||||
entries
|
||||
.iter()
|
||||
.map(|(worker, running, waiting, used, capacity)| {
|
||||
(
|
||||
worker.url.clone(),
|
||||
NativeCacheWorkerLoad {
|
||||
EngineReportedSchedulingLoad {
|
||||
num_running_reqs: *running,
|
||||
num_waiting_reqs: *waiting,
|
||||
num_waiting_uncached_tokens: *waiting,
|
||||
|
||||
@@ -12,15 +12,15 @@
|
||||
//! and the idlest at 1.0, so every assertion below holds either way.
|
||||
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad};
|
||||
use sgl_router::policies::kv_events::{
|
||||
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
|
||||
};
|
||||
use sgl_router::policies::load_based::LoadBasedPolicy;
|
||||
use sgl_router::policies::scoring::{
|
||||
prefix_cache::PrefixCachePolicy, FusedScorePolicy, ScorePolicy,
|
||||
};
|
||||
use sgl_router::policies::{Policy, SelectionContext};
|
||||
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
|
||||
use sgl_router::state::load_monitor::engine_reported_load::{
|
||||
EngineReportedLoadSnapshot, EngineReportedWorkerLoad,
|
||||
};
|
||||
use sgl_router::workers::Worker;
|
||||
use std::{collections::HashMap, sync::Arc, time::Instant};
|
||||
|
||||
@@ -82,12 +82,12 @@ fn fused_load_based_term_uses_the_request_snapshot() {
|
||||
let ws = vec![worker("w0"), worker("w1")];
|
||||
// Local counters changed after the request snapshot and prefer w0.
|
||||
let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect();
|
||||
let snapshot = EngineLoadSnapshot::from_workers(
|
||||
let snapshot = EngineReportedLoadSnapshot::from_workers(
|
||||
29,
|
||||
HashMap::from([
|
||||
(
|
||||
ws[0].url.clone(),
|
||||
EngineWorkerLoad {
|
||||
EngineReportedWorkerLoad {
|
||||
num_running_reqs: 50,
|
||||
num_waiting_reqs: 0,
|
||||
num_tokens: 0,
|
||||
@@ -97,7 +97,7 @@ fn fused_load_based_term_uses_the_request_snapshot() {
|
||||
),
|
||||
(
|
||||
ws[1].url.clone(),
|
||||
EngineWorkerLoad {
|
||||
EngineReportedWorkerLoad {
|
||||
num_running_reqs: 1,
|
||||
num_waiting_reqs: 0,
|
||||
num_tokens: 0,
|
||||
@@ -123,12 +123,12 @@ fn fused_load_based_term_uses_the_request_snapshot() {
|
||||
fn score_policy_forwards_the_request_snapshot_to_load_based() {
|
||||
let ws = vec![worker("w0"), worker("w1")];
|
||||
let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect();
|
||||
let snapshot = EngineLoadSnapshot::from_workers(
|
||||
let snapshot = EngineReportedLoadSnapshot::from_workers(
|
||||
31,
|
||||
HashMap::from([
|
||||
(
|
||||
ws[0].url.clone(),
|
||||
EngineWorkerLoad {
|
||||
EngineReportedWorkerLoad {
|
||||
num_running_reqs: 50,
|
||||
num_waiting_reqs: 0,
|
||||
num_tokens: 0,
|
||||
@@ -138,7 +138,7 @@ fn score_policy_forwards_the_request_snapshot_to_load_based() {
|
||||
),
|
||||
(
|
||||
ws[1].url.clone(),
|
||||
EngineWorkerLoad {
|
||||
EngineReportedWorkerLoad {
|
||||
num_running_reqs: 1,
|
||||
num_waiting_reqs: 0,
|
||||
num_tokens: 0,
|
||||
|
||||
@@ -19,7 +19,7 @@
|
||||
//! whatever fixture is checked in.
|
||||
|
||||
use serde::Deserialize;
|
||||
use sgl_router::policies::kv_events::compute_block_hashes;
|
||||
use sgl_router::state::kv_events::compute_block_hashes;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
|
||||
//! Concurrent-mutation stress test for `HashTree`.
|
||||
//!
|
||||
//! The inline tests in `policies::kv_events::tree` are all
|
||||
//! The inline tests in `state::kv_events::tree` are all
|
||||
//! single-threaded. Under production load, multiple worker subscribers
|
||||
//! drive `insert` / `remove` / `clear_worker` against the same tree from
|
||||
//! tokio worker threads while the chat handler simultaneously calls
|
||||
@@ -25,7 +25,7 @@
|
||||
use std::sync::Arc;
|
||||
use std::thread;
|
||||
|
||||
use sgl_router::policies::kv_events::{HashTree, KvWorkerId};
|
||||
use sgl_router::state::kv_events::{HashTree, KvWorkerId};
|
||||
|
||||
fn worker(i: usize) -> KvWorkerId {
|
||||
KvWorkerId {
|
||||
|
||||
@@ -22,8 +22,8 @@ use std::time::Duration;
|
||||
|
||||
use zeromq::SocketSend;
|
||||
|
||||
use sgl_router::policies::kv_events::discovery::EventConfig;
|
||||
use sgl_router::policies::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
|
||||
use sgl_router::state::kv_events::discovery::EventConfig;
|
||||
use sgl_router::state::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
|
||||
|
||||
use super::zmq_helpers::{
|
||||
build_multipart, encode_block_stored_event, encode_event_batch, make_pub_bound,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! Shared ZMQ wire-format helpers for the `policies::kv_events` component
|
||||
//! Shared ZMQ wire-format helpers for the `state::kv_events` component
|
||||
//! tests. Encodes events in the same msgspec layout SGLang emits, builds
|
||||
//! the two-frame `[seq, payload]` ZMQ message a real publisher sends, and
|
||||
//! binds a loopback PUB socket on an OS-assigned port.
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
use serde::Deserialize;
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::ModelId;
|
||||
@@ -128,7 +128,7 @@ fn registry(model_id: &str, tokenizer_path: PathBuf) -> TokenizerRegistry {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
TokenizerRegistry::load_from_config(&cfg).unwrap()
|
||||
}
|
||||
|
||||
@@ -131,17 +131,17 @@ fn load_guard_decrements_on_panic_unwind() {
|
||||
model_ids: vec![ModelId("m".into())],
|
||||
bootstrap_port: None,
|
||||
}));
|
||||
assert_eq!(w.active_load(), 0);
|
||||
assert_eq!(w.router_inflight_load(), 0);
|
||||
|
||||
let w_inner = w.clone();
|
||||
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
|
||||
let _g = w_inner.load_guard();
|
||||
assert_eq!(w_inner.active_load(), 1);
|
||||
assert_eq!(w_inner.router_inflight_load(), 1);
|
||||
panic!("synthetic panic to exercise Drop on unwind");
|
||||
}));
|
||||
assert!(result.is_err(), "the closure must have panicked");
|
||||
assert_eq!(
|
||||
w.active_load(),
|
||||
w.router_inflight_load(),
|
||||
0,
|
||||
"LoadGuard's Drop must decrement even when the holder panics",
|
||||
);
|
||||
|
||||
@@ -740,7 +740,7 @@ async fn removed_awaits_pending_added() {
|
||||
/// independently — 2N round-trips for N workers.
|
||||
#[tokio::test]
|
||||
async fn manager_emits_single_server_info_fetch_per_worker() {
|
||||
use sgl_router::policies::kv_events::KvEventIndex;
|
||||
use sgl_router::state::kv_events::KvEventIndex;
|
||||
|
||||
let body = json!({
|
||||
"served_model_name": "m",
|
||||
|
||||
@@ -10,17 +10,17 @@ use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use sgl_kv_indexer::{PrefixIndex, PrefixIndexError, PrefixMatch, PrefixOutcome};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, AffinityConfig, BucketConfig, BucketSpec, BucketStage, CacheAwareConfig,
|
||||
CachePrefixProvider, Config, DiscoveryBackend, KvIndexerEndpointConfig, ModelConfig,
|
||||
AffinityConfig, BucketConfig, BucketSpec, BucketStage, CacheAwareConfig, CachePrefixProvider,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, KvIndexerEndpointConfig, ModelConfig,
|
||||
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode,
|
||||
SloBucketPolicy, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::engine_load::{LoadStat, NativeCacheRankLoad};
|
||||
use sgl_router::policies::factory::build_registry_with_defaults;
|
||||
use sgl_router::proxy::Proxy;
|
||||
use sgl_router::server::app::build_router;
|
||||
use sgl_router::server::app_context::AppContext;
|
||||
use sgl_router::state::load_monitor::engine_reported_load::{LoadStat, NativeCacheRankLoad};
|
||||
use sgl_router::tokenizer::TokenizerRegistry;
|
||||
use sgl_router::workers::WorkerRegistry;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
@@ -77,7 +77,7 @@ fn build_app_context(
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&config).unwrap());
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
@@ -236,7 +236,7 @@ fn set_native_load_with_waiting(
|
||||
max_total_num_tokens: u64,
|
||||
num_waiting_uncached_tokens: u64,
|
||||
) {
|
||||
ctx.engine_load.set(
|
||||
ctx.engine_reported_load.set(
|
||||
worker_url,
|
||||
0,
|
||||
LoadStat {
|
||||
|
||||
@@ -18,10 +18,10 @@ use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry;
|
||||
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
||||
use sgl_router::proxy::Proxy;
|
||||
use sgl_router::server::app::build_router;
|
||||
use sgl_router::server::app_context::AppContext;
|
||||
use sgl_router::state::kv_events::{BlockSizeOracle, HashTree};
|
||||
use sgl_router::tokenizer::TokenizerRegistry;
|
||||
use sgl_router::workers::WorkerRegistry;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -50,7 +50,7 @@ fn config_for(_worker_url: &str) -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1322,7 +1322,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br
|
||||
/// streaming response, not just for the handler lifetime.
|
||||
///
|
||||
/// Before the fix, the handler dropped `_guard` as soon as it returned
|
||||
/// (which happens when headers arrive), so `active_load()` was 0 while
|
||||
/// (which happens when headers arrive), so `router_inflight_load()` was 0 while
|
||||
/// the SSE pump was still relaying bytes. This test catches that bug.
|
||||
#[tokio::test]
|
||||
async fn streaming_load_guard_persists_for_body_lifetime() {
|
||||
@@ -1360,7 +1360,7 @@ async fn streaming_load_guard_persists_for_body_lifetime() {
|
||||
));
|
||||
let app = build_router(ctx);
|
||||
|
||||
// Grab the Worker handle so we can assert active_load().
|
||||
// Grab the Worker handle so we can assert router_inflight_load().
|
||||
let w_handle: Arc<Worker> = registry
|
||||
.workers_for(&ModelId("tiny".into()))
|
||||
.into_iter()
|
||||
@@ -1386,9 +1386,9 @@ async fn streaming_load_guard_persists_for_body_lifetime() {
|
||||
// first chunk's delay to pass, then assert load is still held.
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
assert!(
|
||||
w_handle.active_load() >= 1,
|
||||
w_handle.router_inflight_load() >= 1,
|
||||
"load should be >= 1 mid-stream, got {}",
|
||||
w_handle.active_load()
|
||||
w_handle.router_inflight_load()
|
||||
);
|
||||
|
||||
// Drain the entire body — this drives the SSE pump to completion.
|
||||
@@ -1398,25 +1398,25 @@ async fn streaming_load_guard_persists_for_body_lifetime() {
|
||||
// released. Give the spawned task a brief moment to clean up.
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
assert_eq!(
|
||||
w_handle.active_load(),
|
||||
w_handle.router_inflight_load(),
|
||||
0,
|
||||
"load should be 0 after stream completes"
|
||||
);
|
||||
}
|
||||
|
||||
/// Task A: the chat handler mints an `ActiveLoadGuard` from the shared
|
||||
/// `ActiveLoadRegistry` and drops it when the request completes. The
|
||||
/// Task A: the chat handler mints an `RouterInflightLoadGuard` from the shared
|
||||
/// `RouterInflightLoadRegistry` and drops it when the request completes. The
|
||||
/// non-streaming path drops the guard on handler exit; this test
|
||||
/// asserts the round-trip increment → 0 across a single request.
|
||||
#[tokio::test]
|
||||
async fn non_streaming_active_load_increments_then_returns_to_zero() {
|
||||
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx_with_worker(&worker.url);
|
||||
let active_load = Arc::clone(&ctx.active_load);
|
||||
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
|
||||
let app = build_router(ctx);
|
||||
|
||||
assert_eq!(
|
||||
active_load.inflight_count(),
|
||||
router_inflight_load.inflight_count(),
|
||||
0,
|
||||
"registry must start with no in-flight requests",
|
||||
);
|
||||
@@ -1442,19 +1442,19 @@ async fn non_streaming_active_load_increments_then_returns_to_zero() {
|
||||
// The handler has returned, so the active-load guard must have
|
||||
// dropped — counters are back to zero.
|
||||
assert_eq!(
|
||||
active_load.inflight_count(),
|
||||
router_inflight_load.inflight_count(),
|
||||
0,
|
||||
"active-load registry must be empty after non-streaming handler returns",
|
||||
);
|
||||
let w_id = WorkerId("w1".into());
|
||||
assert_eq!(
|
||||
active_load.prefill_load(&w_id),
|
||||
router_inflight_load.prefill_load(&w_id),
|
||||
0,
|
||||
"prefill_load must decrement on response end",
|
||||
);
|
||||
}
|
||||
|
||||
/// Task A: the streaming path holds the `ActiveLoadGuard` until the
|
||||
/// Task A: the streaming path holds the `RouterInflightLoadGuard` until the
|
||||
/// SSE pump finishes. Mid-stream the registry shows `inflight_count >= 1`;
|
||||
/// after the body drains it returns to 0. Counterpart to
|
||||
/// `streaming_load_guard_persists_for_body_lifetime` — both guards must
|
||||
@@ -1486,7 +1486,7 @@ async fn streaming_active_load_persists_for_body_lifetime() {
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
|
||||
let ctx = Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies));
|
||||
let active_load = Arc::clone(&ctx.active_load);
|
||||
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
|
||||
let app = build_router(ctx);
|
||||
|
||||
let req = Request::builder()
|
||||
@@ -1508,15 +1508,15 @@ async fn streaming_active_load_persists_for_body_lifetime() {
|
||||
// still running, so the registry's per-request entry must remain.
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
assert!(
|
||||
active_load.inflight_count() >= 1,
|
||||
router_inflight_load.inflight_count() >= 1,
|
||||
"registry inflight must be >= 1 mid-stream, got {}",
|
||||
active_load.inflight_count(),
|
||||
router_inflight_load.inflight_count(),
|
||||
);
|
||||
let w_id = WorkerId("w1".into());
|
||||
assert!(
|
||||
active_load.prefill_load(&w_id) >= 1,
|
||||
router_inflight_load.prefill_load(&w_id) >= 1,
|
||||
"prefill_load must be > 0 mid-stream, got {}",
|
||||
active_load.prefill_load(&w_id),
|
||||
router_inflight_load.prefill_load(&w_id),
|
||||
);
|
||||
|
||||
// Drain the body — drives the SSE pump to completion.
|
||||
@@ -1524,12 +1524,12 @@ async fn streaming_active_load_persists_for_body_lifetime() {
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
|
||||
assert_eq!(
|
||||
active_load.inflight_count(),
|
||||
router_inflight_load.inflight_count(),
|
||||
0,
|
||||
"registry must be empty after stream drains",
|
||||
);
|
||||
assert_eq!(
|
||||
active_load.prefill_load(&w_id),
|
||||
router_inflight_load.prefill_load(&w_id),
|
||||
0,
|
||||
"prefill_load must be 0 after stream drains",
|
||||
);
|
||||
@@ -1555,7 +1555,7 @@ async fn streaming_active_load_drops_on_client_disconnect() {
|
||||
)
|
||||
.await;
|
||||
let (ctx, body) = stream_chat(&worker.url).await;
|
||||
let active_load = Arc::clone(&ctx.active_load);
|
||||
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
|
||||
|
||||
// Read one chunk to confirm the stream is live, then drop the body.
|
||||
use futures::StreamExt;
|
||||
@@ -1570,7 +1570,7 @@ async fn streaming_active_load_drops_on_client_disconnect() {
|
||||
wait_for_metric(&ctx, &expected).await;
|
||||
|
||||
assert_eq!(
|
||||
active_load.inflight_count(),
|
||||
router_inflight_load.inflight_count(),
|
||||
0,
|
||||
"client disconnect must drop the streaming pump's guards within one tick",
|
||||
);
|
||||
@@ -1583,13 +1583,15 @@ async fn streaming_active_load_drops_on_client_disconnect() {
|
||||
/// `ApiError::StaleRequestExpired`.
|
||||
///
|
||||
/// Wiring: build an `AppContext` with a short
|
||||
/// `stale_request_timeout` `ActiveLoadRegistry` + spawn a janitor
|
||||
/// `stale_request_timeout` `RouterInflightLoadRegistry` + spawn a janitor
|
||||
/// with sub-second cadence + dispatch to a slow upstream that takes
|
||||
/// longer than the timeout. The janitor sweeps before the upstream
|
||||
/// returns; cancellation fires; handler returns 504.
|
||||
#[tokio::test]
|
||||
async fn janitor_expiry_returns_504_stale_request_expired() {
|
||||
use sgl_router::policies::active_load::{spawn_janitor, ActiveLoadRegistry};
|
||||
use sgl_router::state::load_monitor::router_inflight_load::{
|
||||
spawn_janitor, RouterInflightLoadRegistry,
|
||||
};
|
||||
// Upstream that takes 2s to respond — longer than our 50ms
|
||||
// stale_request_timeout.
|
||||
let worker =
|
||||
@@ -1610,18 +1612,18 @@ async fn janitor_expiry_returns_504_stale_request_expired() {
|
||||
// Aggressive 50ms timeout: the janitor will sweep on the next
|
||||
// tick (every 20ms) and fire the cancellation token before the
|
||||
// upstream returns.
|
||||
let active_load = ActiveLoadRegistry::new(
|
||||
Arc::new(sgl_router::policies::active_load::SystemTimeClock),
|
||||
let router_inflight_load = RouterInflightLoadRegistry::new(
|
||||
Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock),
|
||||
Duration::from_millis(50),
|
||||
);
|
||||
let _janitor = spawn_janitor(Arc::clone(&active_load), Duration::from_millis(20));
|
||||
let ctx = Arc::new(AppContext::with_active_load(
|
||||
let _janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20));
|
||||
let ctx = Arc::new(AppContext::with_router_inflight_load(
|
||||
cfg,
|
||||
tokenizers,
|
||||
proxy,
|
||||
registry,
|
||||
policies,
|
||||
active_load,
|
||||
router_inflight_load,
|
||||
));
|
||||
let app = build_router(ctx);
|
||||
|
||||
@@ -1671,7 +1673,7 @@ async fn non_streaming_error_path_drops_active_load_guard() {
|
||||
drop(listener);
|
||||
|
||||
let ctx = build_ctx_with_worker(&dead_url);
|
||||
let active_load = Arc::clone(&ctx.active_load);
|
||||
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
|
||||
let app = build_router(ctx);
|
||||
|
||||
let req = Request::builder()
|
||||
@@ -1692,7 +1694,7 @@ async fn non_streaming_error_path_drops_active_load_guard() {
|
||||
// Drain so any drop-on-body-end work runs.
|
||||
let _ = res.into_body().collect().await.unwrap().to_bytes();
|
||||
assert_eq!(
|
||||
active_load.inflight_count(),
|
||||
router_inflight_load.inflight_count(),
|
||||
0,
|
||||
"error path must drop the active-load guard",
|
||||
);
|
||||
|
||||
@@ -7,8 +7,8 @@
|
||||
//! built-in V4 chat formatter — the engine-equivalent path — with no template fixture.
|
||||
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
|
||||
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
CacheAwareConfig, Config, DiscoveryBackend, InflightLoadConfig, ModelConfig,
|
||||
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
|
||||
pub const MODEL: &str = "deepseek-v4-tiny";
|
||||
@@ -42,6 +42,6 @@ pub fn config() -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,11 +19,11 @@ use sgl_kv_indexer::{
|
||||
use sgl_router::config::{AffinityConfig, CachePrefixProvider, PolicyKind};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry;
|
||||
use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
|
||||
use sgl_router::policies::request_tokens_for;
|
||||
use sgl_router::proxy::Proxy;
|
||||
use sgl_router::server::app::build_router;
|
||||
use sgl_router::server::app_context::AppContext;
|
||||
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
|
||||
use sgl_router::tokenizer::TokenizerRegistry;
|
||||
use sgl_router::workers::WorkerRegistry;
|
||||
use tokio_stream::wrappers::TcpListenerStream;
|
||||
|
||||
@@ -54,7 +54,7 @@ async fn failover_when_one_worker_dies() {
|
||||
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
|
||||
@@ -25,7 +25,7 @@
|
||||
|
||||
use futures::future::join_all;
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -69,7 +69,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
@@ -554,7 +554,7 @@ async fn wait_for_inflight_http(ctx: &Arc<AppContext>, want: usize) {
|
||||
/// response BODY finishing, not the handler returning. A streaming completion
|
||||
/// hands back its headers immediately, so a count released at handler exit
|
||||
/// would read 0 for the entire window the heartbeat exists to explain — the
|
||||
/// same blind spot `active_load.inflight_count()` has, reproduced in the
|
||||
/// same blind spot `router_inflight_load.inflight_count()` has, reproduced in the
|
||||
/// replacement.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn inflight_http_counts_a_streaming_response_until_its_body_finishes() {
|
||||
@@ -623,7 +623,7 @@ async fn inflight_http_counts_a_streaming_response_until_its_body_finishes() {
|
||||
|
||||
/// Every route is instrumented, not only the proxied ones. `/metrics`,
|
||||
/// `/readyz` and a 404 are exchanges axum's drain waits on too, and they are
|
||||
/// exactly the traffic `active_load` cannot see — so a guard that leaked on a
|
||||
/// exactly the traffic `router_inflight_load` cannot see — so a guard that leaked on a
|
||||
/// non-proxied route would leave the heartbeat permanently busy and turn the
|
||||
/// drain report back into noise.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -47,7 +47,7 @@ async fn forwards_whitelisted_headers_strips_others() {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
|
||||
@@ -21,7 +21,7 @@ use axum::http::{Request, StatusCode};
|
||||
use bytes::Bytes;
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -62,7 +62,7 @@ fn config() -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -225,7 +225,7 @@ async fn round_robin_pd_prefill_does_not_track_dispatch_timestamps() {
|
||||
let request = tokio::spawn(build_router(Arc::clone(&ctx)).oneshot(chat_request()));
|
||||
|
||||
await_captured_body(&prefill, Duration::from_secs(2), "prefill").await;
|
||||
assert_eq!(prefill_worker.active_load(), 1);
|
||||
assert_eq!(prefill_worker.router_inflight_load(), 1);
|
||||
assert_eq!(prefill_worker.slots_acquired_since(cutoff), 0);
|
||||
|
||||
assert_eq!(request.await.unwrap().unwrap().status(), StatusCode::OK);
|
||||
|
||||
@@ -20,7 +20,7 @@ use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use http_body_util::BodyExt;
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -61,7 +61,7 @@ fn config() -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -32,7 +32,7 @@ use hyper::service::service_fn;
|
||||
use hyper::Response as HyperResponse;
|
||||
use hyper_util::rt::{TokioExecutor, TokioIo};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -149,7 +149,7 @@ fn config() -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -10,14 +10,12 @@ use serde_json::json;
|
||||
use sgl_router::config::{AffinityConfig, CachePrefixProvider, PolicyKind};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry;
|
||||
use sgl_router::policies::kv_events::{
|
||||
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
|
||||
};
|
||||
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
|
||||
use sgl_router::policies::request_tokens_for;
|
||||
use sgl_router::proxy::Proxy;
|
||||
use sgl_router::server::app::build_router;
|
||||
use sgl_router::server::app_context::AppContext;
|
||||
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
|
||||
use sgl_router::tokenizer::TokenizerRegistry;
|
||||
use sgl_router::workers::WorkerRegistry;
|
||||
use tower::ServiceExt;
|
||||
|
||||
@@ -11,7 +11,7 @@ use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -59,7 +59,7 @@ fn config() -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -7,11 +7,10 @@ use std::time::{Duration, Instant};
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::engine_load::{LoadStat, NativeCacheRankLoad};
|
||||
use sgl_router::policies::{
|
||||
CacheCandidate, CacheCandidateProposal, Policy, PolicyRegistry, PrefillProposal, ProposalKind,
|
||||
SelectionContext, SelectionProposal,
|
||||
@@ -19,6 +18,7 @@ use sgl_router::policies::{
|
||||
use sgl_router::proxy::Proxy;
|
||||
use sgl_router::server::app::build_router;
|
||||
use sgl_router::server::app_context::AppContext;
|
||||
use sgl_router::state::load_monitor::engine_reported_load::{LoadStat, NativeCacheRankLoad};
|
||||
use sgl_router::tokenizer::TokenizerRegistry;
|
||||
use sgl_router::workers::{Worker, WorkerRegistry};
|
||||
use tower::ServiceExt;
|
||||
@@ -165,7 +165,7 @@ fn config(policy: PolicyKind) -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -307,7 +307,7 @@ async fn chat_commits_the_admitted_prefill_backup() {
|
||||
total_prefill_busy_us,
|
||||
}),
|
||||
};
|
||||
fixture.ctx.engine_load.set(
|
||||
fixture.ctx.engine_reported_load.set(
|
||||
&fixture.workers[0].url,
|
||||
0,
|
||||
native_load(1, 1),
|
||||
@@ -315,7 +315,7 @@ async fn chat_commits_the_admitted_prefill_backup() {
|
||||
);
|
||||
fixture
|
||||
.ctx
|
||||
.engine_load
|
||||
.engine_reported_load
|
||||
.set(&fixture.workers[0].url, 0, native_load(2, 2), now);
|
||||
|
||||
assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK);
|
||||
@@ -351,7 +351,7 @@ async fn capacity_exhaustion_does_not_return_503() {
|
||||
})
|
||||
.await;
|
||||
for worker in &fixture.workers {
|
||||
fixture.ctx.engine_load.set(
|
||||
fixture.ctx.engine_reported_load.set(
|
||||
&worker.url,
|
||||
0,
|
||||
LoadStat {
|
||||
@@ -434,7 +434,7 @@ async fn chat_records_cache_candidates_exhausted() {
|
||||
})
|
||||
})
|
||||
.await;
|
||||
fixture.ctx.engine_load.set(
|
||||
fixture.ctx.engine_reported_load.set(
|
||||
&fixture.workers[0].url,
|
||||
0,
|
||||
LoadStat {
|
||||
|
||||
@@ -24,7 +24,7 @@ use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -77,7 +77,7 @@ fn config() -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
//! `MockWorker` backends (CPU-only, no GPU).
|
||||
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -63,7 +63,7 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc<AppContext
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
};
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
|
||||
@@ -14,7 +14,7 @@ use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use http_body_util::BodyExt;
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
@@ -55,7 +55,7 @@ fn config(_worker_url: &str) -> Config {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
router_inflight_load: InflightLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user