Skip to content

feat(load-monitor): per-group polling with rich load data #554

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slin1237 merged 3 commits into from
Feb 27, 2026
Merged

feat(load-monitor): per-group polling with rich load data #554

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9b5a802
feat(load-monitor): per-group polling with rich load data and token_u…
slin1237 Feb 27, 2026
952cf9d
fix(load-monitor): prevent zero-interval DoS and fix stale entry cleanup
slin1237 Feb 27, 2026
d9d106c
fix(load-monitor): fix cache merge, stale snapshots, and gRPC empty-l…
slin1237 Feb 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions grpc_client/src/sglang_scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -708,6 +708,38 @@ impl SglangSchedulerClient {
}
}

// ---------------------------------------------------------------------------
// Proto → protocol type conversions
// ---------------------------------------------------------------------------

impl From<proto::SchedulerLoad> for openai_protocol::worker::SchedulerLoadSnapshot {
fn from(load: proto::SchedulerLoad) -> Self {
Self {
dp_rank: load.dp_rank,
num_running_reqs: load.num_running_reqs,
num_waiting_reqs: load.num_waiting_reqs,
num_total_reqs: load.num_total_reqs,
num_used_tokens: load.num_used_tokens,
max_total_num_tokens: load.max_total_num_tokens,
token_usage: load.token_usage,
gen_throughput: load.gen_throughput,
cache_hit_rate: load.cache_hit_rate,
utilization: load.utilization,
max_running_requests: load.max_running_requests,
}
}
}

impl From<proto::GetLoadsResponse> for openai_protocol::worker::WorkerLoadResponse {
fn from(resp: proto::GetLoadsResponse) -> Self {
Self {
timestamp: resp.timestamp,
dp_rank_count: resp.dp_rank_count,
loads: resp.loads.into_iter().map(Into::into).collect(),
}
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down
2 changes: 1 addition & 1 deletion model_gateway/src/core/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ pub use job_queue::{Job, JobQueue, JobQueueConfig};
pub use openai_protocol::{
model_card::ModelCard,
model_type::{Endpoint, ModelType},
worker::ProviderType,
worker::{ProviderType, WorkerGroupKey},
};
pub use retry::{is_retryable_status, RetryExecutor};
pub use worker::{
Expand Down
57 changes: 42 additions & 15 deletions model_gateway/src/core/steps/worker/local/remove_from_worker_registry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
//! Step to remove workers from worker registry.

use std::collections::HashSet;
use std::collections::{HashMap, HashSet};

use async_trait::async_trait;
use tracing::{debug, warn};
Expand All @@ -10,6 +10,7 @@ use crate::{
core::{
steps::workflow_data::WorkerRemovalWorkflowData,
worker::{ConnectionModeExt, WorkerTypeExt},
WorkerGroupKey,
},
observability::metrics::Metrics,
};
Expand Down Expand Up @@ -37,19 +38,28 @@ impl StepExecutor<WorkerRemovalWorkflowData> for RemoveFromWorkerRegistryStep {
worker_urls.len()
);

// Collect unique worker configurations before removal for pool size updates
let unique_configs: HashSet<_> = worker_urls
.iter()
.filter_map(|url| app_context.worker_registry.get_by_url(url))
.map(|w| {
// Collect unique worker configurations and their URLs before removal.
// We need the URLs pre-removal so LoadMonitor can clean up stale load entries.
let mut urls_by_group: HashMap<
(
openai_protocol::worker::WorkerType,
openai_protocol::worker::ConnectionMode,
String,
),
Vec<String>,
> = HashMap::new();
for url in worker_urls {
if let Some(w) = app_context.worker_registry.get_by_url(url) {
let meta = w.metadata();
(
let key = (
meta.spec.worker_type,
meta.spec.connection_mode,
w.model_id().to_string(),
)
})
.collect();
);
urls_by_group.entry(key).or_default().push(url.clone());
}
}
let unique_configs: HashSet<_> = urls_by_group.keys().cloned().collect();

let mut removed_count = 0;
for worker_url in worker_urls {
Expand All @@ -74,17 +84,18 @@ impl StepExecutor<WorkerRemovalWorkflowData> for RemoveFromWorkerRegistryStep {
}

// Update Layer 3 worker pool size metrics for unique configurations
for (worker_type, connection_mode, model_id) in unique_configs {
// and notify LoadMonitor when groups become empty
for (worker_type, connection_mode, model_id) in &unique_configs {
// Get labels before moving values into get_workers_filtered
let worker_type_label = worker_type.as_metric_label();
let connection_mode_label = connection_mode.as_metric_label();

let pool_size = app_context
.worker_registry
.get_workers_filtered(
Some(&model_id),
Some(worker_type),
Some(connection_mode),
Some(model_id),
Some(*worker_type),
Some(*connection_mode),
None,
false,
)
Expand All @@ -93,9 +104,25 @@ impl StepExecutor<WorkerRemovalWorkflowData> for RemoveFromWorkerRegistryStep {
Metrics::set_worker_pool_size(
worker_type_label,
connection_mode_label,
&model_id,
model_id,
pool_size,
);

// If the group is now empty, stop its load monitor and clean up entries
if pool_size == 0 {

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clean removed worker loads even when group is non-empty

This cleanup path only runs when pool_size == 0, so removing a worker from a still-active group skips on_group_removed. The monitor loop later removes only URLs that still exist in that group, so departed URLs are never purged from the shared load snapshot for long-lived groups, causing stale-entry growth and allowing old load data to resurface if that URL is re-registered before the next successful poll.

Useful? React with 👍 / 👎.

@slin1237 slin1237 Feb 27, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicate of the earlier comment. Already addressed — the monitor loop clears and re-inserts on each tick (d9d106c), so departed URLs are cleaned up within one interval period.

if let Some(ref load_monitor) = app_context.load_monitor {
let key = WorkerGroupKey {
model_id: model_id.clone(),
worker_type: *worker_type,
connection_mode: *connection_mode,
};
let removed_urls = urls_by_group
.get(&(*worker_type, *connection_mode, model_id.clone()))
.map(|v| v.as_slice())
.unwrap_or_default();
load_monitor.on_group_removed(&key, removed_urls).await;
}
Comment on lines +111 to +124

@coderabbitai coderabbitai Bot Feb 27, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Cleanup is skipped for removed workers when the group still has members.

Line 112 gates cleanup behind pool_size == 0. When a subset is removed, those URLs are not pruned here, so stale load entries can persist for active groups.

Suggested direction 
-            // If the group is now empty, stop its load monitor and clean up entries
-            if pool_size == 0 {
-                if let Some(ref load_monitor) = app_context.load_monitor {
+            if let Some(ref load_monitor) = app_context.load_monitor {
+                let removed_urls = urls_by_group
+                    .get(&(*worker_type, *connection_mode, model_id.clone()))
+                    .map(|v| v.as_slice())
+                    .unwrap_or_default();
+
+                // New API: prune removed URLs even when group remains active.
+                load_monitor.on_workers_removed(removed_urls).await;
+
+                if pool_size == 0 {
                     let key = WorkerGroupKey {
                         model_id: model_id.clone(),
                         worker_type: *worker_type,
                         connection_mode: *connection_mode,
                     };
-                    let removed_urls = urls_by_group
-                        .get(&(*worker_type, *connection_mode, model_id.clone()))
-                        .map(|v| v.as_slice())
-                        .unwrap_or_default();
                     load_monitor.on_group_removed(&key, removed_urls).await;
                 }
             }
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@model_gateway/src/core/steps/worker/local/remove_from_worker_registry.rs`
around lines 111 - 124, The current cleanup only invokes
app_context.load_monitor.on_group_removed when pool_size == 0, so when a subset
of workers is removed from a non-empty group their URLs are never pruned; update
remove_from_worker_registry logic to always compute the removed URLs for the
group (using urls_by_group and the WorkerGroupKey built from model_id,
worker_type and connection_mode) and call the load monitor to prune those
specific URLs even when pool_size > 0 (e.g., invoke an existing or new load
monitor method to remove specific worker URLs such as
on_workers_removed/on_urls_removed with the key and removed_urls instead of only
calling on_group_removed for empty groups).

@slin1237 slin1237 Feb 27, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pushing back. Stale entries for removed workers are cleaned up on the next poll tick — group_monitor_loop clears all group URLs before inserting (d9d106c), and removed workers wont appear in get_workers_filtered, so their URLs wont be re-inserted. Cleanup window is bounded to one interval period.

@coderabbitai coderabbitai Bot Feb 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Find the group_monitor_loop implementation and the clearing logic
# Look for the section that clears group URLs before inserting new load data

# First, let's see the group_monitor_loop function
ast-grep --pattern $'async fn group_monitor_loop($$$) {
  $$$
}'

# Also search for the commit message mentioned by user
git log --oneline --all | grep -i "d9d106c5\|clear\|stale"

# Look for send_modify usage which should show the merge logic
rg -n -A10 -B5 "send_modify" model_gateway/src/core/worker_manager.rs

Length of output: 9054

@slin1237, you're absolutely right. I can see in lines 514-520 of model_gateway/src/core/worker_manager.rs that group_monitor_loop collects all_group_urls from the current get_workers_filtered result, then uses send_modify to remove all those URLs before extending with fresh group_loads. Since removed workers won't appear in the filtered list, their URLs are automatically pruned on the next tick. The comment at line 512 confirms this design. My concern was incorrect—apologies for the noise.

🧠 Learnings used 
Learnt from: slin1237
Repo: lightseekorg/smg PR: 489
File: protocols/src/worker.rs:338-343
Timestamp: 2026-02-21T02:36:06.251Z
Learning: Repo: lightseekorg/smg PR: 489
File: protocols/src/worker.rs (impl From<Vec<ModelCard>> for WorkerModels)
Learning: For the single-element case (models.len() == 1), keep the defensive Option guard with a non-panicking fallback to WorkerModels::Wildcard instead of using expect()/unwrap(). This is intentional to avoid introducing panic paths in production code for PR `#489` and similar clippy/lint-only efforts.

Learnt from: slin1237
Repo: lightseekorg/smg PR: 489
File: model_gateway/src/policies/bucket.rs:116-137
Timestamp: 2026-02-21T02:31:03.274Z
Learning: For clippy-focused PRs in this repository, avoid behavior changes (e.g., load-balancing logic). Defer such changes to a separate PR/issue.

Learnt from: XinyueZhang369
Repo: lightseekorg/smg PR: 399
File: protocols/src/interactions.rs:505-509
Timestamp: 2026-02-19T03:08:50.192Z
Learning: In code reviews for Rust projects using the validator crate (v0.20.0), ensure that custom validation functions for numeric primitive types (e.g., f32, i32, u32, i16, etc.) accept the value by value, not by reference. Example: fn validate(value: f32) { ... }. The validator derive macro has a hardcoded list of numeric types that are passed by value, while all other types are passed by reference. Apply this guideline whenever validating numeric fields to align with the derive macro behavior.

Learnt from: slin1237
Repo: lightseekorg/smg PR: 489
File: model_gateway/src/core/token_bucket.rs:58-63
Timestamp: 2026-02-21T02:30:51.443Z
Learning: For lint-only/Clippy enforcement PRs in this repository, avoid introducing behavioral changes (e.g., new input validation or logic changes). Treat such PRs as non-functional changes and plan a separate follow-up issue/PR for hardening or behavior changes. This applies broadly to Rust files across the repo; during review, focus on lint/style corrections and clearly note any intentional exceptions. 

Learnt from: slin1237
Repo: lightseekorg/smg PR: 489
File: protocols/src/responses.rs:928-931
Timestamp: 2026-02-21T02:36:00.882Z
Learning: In Rust code across the repository, use the marker INVARIANT: to document assumptions in safe code. Reserve SAFETY: for explaining why unsafe blocks are sound. This improves clarity of invariants and safety reasoning. Example reference: protocols/src/responses.rs near validate_tool_choice_with_tools().

Learnt from: slin1237
Repo: lightseekorg/smg PR: 489
File: mesh/src/sync.rs:83-83
Timestamp: 2026-02-21T02:37:01.416Z
Learning: General Rust formatting rule: format! with implicit captures only supports simple identifiers, not full expressions like {state.model_id}. For cases where you want to interpolate a field or expression, bind the value first and interpolate the binding, e.g., let model_id = &state.model_id; and then use format!("policy:{}", model_id). In the specific file mesh/src/sync.rs, prefer format!("policy:{}", state.model_id) or bind to a local variable if you need named interpolation, to keep clarity and avoid unintended captures.

}
}

Ok(StepResult::Success)
Expand Down
35 changes: 30 additions & 5 deletions model_gateway/src/core/steps/worker/shared/register.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ use crate::{
core::{
steps::workflow_data::WorkerRegistrationData,
worker::{ConnectionModeExt, WorkerTypeExt},
WorkerGroupKey,
},
observability::metrics::Metrics,
};
Expand Down Expand Up @@ -64,17 +65,18 @@ impl<D: WorkerRegistrationData + WorkflowData> StepExecutor<D> for RegisterWorke
.collect();

// Update Layer 3 worker pool size metrics per unique type/connection/model
for (worker_type, connection_mode, model_id) in unique_configs {
// and notify LoadMonitor of new groups
for (worker_type, connection_mode, model_id) in &unique_configs {
// Get labels before moving values into get_workers_filtered
let worker_type_label = worker_type.as_metric_label();
let connection_mode_label = connection_mode.as_metric_label();

let pool_size = app_context
.worker_registry
.get_workers_filtered(
Some(&model_id),
Some(worker_type),
Some(connection_mode),
Some(model_id),
Some(*worker_type),
Some(*connection_mode),
None,
false,
)
Expand All @@ -83,11 +85,34 @@ impl<D: WorkerRegistrationData + WorkflowData> StepExecutor<D> for RegisterWorke
Metrics::set_worker_pool_size(
worker_type_label,
connection_mode_label,
&model_id,
model_id,
pool_size,
);
}

// Notify LoadMonitor of new worker groups so it can start polling
if let Some(ref load_monitor) = app_context.load_monitor {
for (worker_type, connection_mode, model_id) in &unique_configs {
let key = WorkerGroupKey {
model_id: model_id.clone(),
worker_type: *worker_type,
connection_mode: *connection_mode,
};
// Use the minimum interval across all workers in this group.
// If workers disagree, the fastest polling rate wins (safest default).
let interval = workers
.iter()
.filter(|w| {
w.model_id() == model_id
&& w.worker_type() == worker_type
&& w.connection_mode() == connection_mode
})
.filter_map(|w| w.metadata().spec.load_monitor_interval_secs)

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Propagate per-worker monitor interval into worker metadata

This code derives the group polling interval from w.metadata().spec.load_monitor_interval_secs, but the worker creation paths still build workers via BasicWorkerBuilder::new(...) without ever copying WorkerSpec.load_monitor_interval_secs into metadata (e.g., CreateLocalWorkerStep / CreateExternalWorkersStep), so this filter always sees None and on_group_added always falls back to the global default. In practice, any user-specified per-worker override is silently ignored, so the new per-worker interval feature does not function.

Useful? React with 👍 / 👎.

@slin1237 slin1237 Feb 27, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Valid observation. The creation steps use individual builder setters and dont pass load_monitor_interval_secs through. The field exists on WorkerSpec and the monitoring code reads it, but the creation pipeline needs a builder setter or from_spec() usage. Will address in a follow-up.

.min();
Comment on lines +110 to +111

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Include default interval when computing group minimum

This minimum only considers explicit per-worker overrides, so workers with load_monitor_interval_secs = None are excluded from the calculation. In a mixed group (e.g., one worker with Some(30) and another with None while global is 10), this passes Some(30) to on_group_added, so the group polls slower than the global fallback and violates the documented “minimum effective interval wins” behavior.

Useful? React with 👍 / 👎.

@slin1237 slin1237 Feb 27, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pushing back. Per-worker interval is an explicit override. In a mixed group [Some(30), None] with global=10, we pass Some(30) — the user set 30s intentionally. Workers without an override defer to whatever the group picks. The global default only kicks in when no worker specifies an interval.

load_monitor.on_group_added(key, interval).await;
}
}
Comment thread
slin1237 marked this conversation as resolved.
Resolved

// Note: worker_ids are stored for potential future use but not persisted
// as they are internal registry identifiers
debug!(
Expand Down
Loading