[misc] Priority scheduling metrics cleanup (#19927)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
sglang-bot
2026-03-05 12:42:42 -08:00
committed by GitHub
parent ff6048fb9c
commit ebb66cc1de
5 changed files with 116 additions and 54 deletions

View File

@@ -303,30 +303,30 @@ class SchedulerRuntimeCheckerMixin:
else:
num_used, token_usage, _, _ = self._get_token_info()
_enable_ps = self.enable_priority_scheduling
priority_enabled = self.enable_priority_scheduling
self.stats.num_running_reqs = QueueCount.from_reqs(
self.running_batch.reqs, _enable_ps
self.running_batch.reqs, priority_enabled
)
self.stats.num_used_tokens = num_used
self.stats.token_usage = round(token_usage, 2)
self.stats.gen_throughput = 0
self.stats.num_queue_reqs = QueueCount.from_reqs(
self.waiting_queue, _enable_ps
self.waiting_queue, priority_enabled
)
self.stats.num_grammar_queue_reqs = len(self.grammar_manager)
if self.disaggregation_mode == DisaggregationMode.PREFILL:
self.stats.num_prefill_prealloc_queue_reqs = QueueCount.from_reqs(
self.disagg_prefill_bootstrap_queue.queue, _enable_ps
self.disagg_prefill_bootstrap_queue.queue, priority_enabled
)
self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
self.disagg_prefill_inflight_queue, _enable_ps
self.disagg_prefill_inflight_queue, priority_enabled
)
if self.disaggregation_mode == DisaggregationMode.DECODE:
self.stats.num_decode_prealloc_queue_reqs = QueueCount.from_reqs(
self.disagg_decode_prealloc_queue.queue, _enable_ps
self.disagg_decode_prealloc_queue.queue, priority_enabled
)
self.stats.num_decode_transfer_queue_reqs = QueueCount.from_reqs(
self.disagg_decode_transfer_queue.queue, _enable_ps
self.disagg_decode_transfer_queue.queue, priority_enabled
)
self.metrics_collector.log_stats(self.stats)
self._publish_kv_events()

View File

@@ -758,22 +758,23 @@ class SchedulerMetricsCollector:
multiprocess_mode="mostrecent",
)
def _log_gauge(self, gauge: Gauge, data: Union[int, float, QueueCount]) -> None:
# Convenience function for logging to gauge.
if isinstance(data, QueueCount):
# NOTE: When priority scheduling is enabled, the total is recorded under
# priority="" (the default label value). Per-priority breakdowns are recorded
# with priority="<int>". Grafana queries should use priority="" for totals.
gauge.labels(**self.labels).set(data.total)
if data.by_priority is not None:
self._known_priorities.update(data.by_priority.keys())
for priority in self._known_priorities:
value = data.by_priority.get(priority, 0)
labels = dict(self.labels)
labels["priority"] = str(priority)
gauge.labels(**labels).set(value)
else:
gauge.labels(**self.labels).set(data)
def _log_gauge(self, gauge: Gauge, data: Union[int, float]) -> None:
# Convenience function for logging a scalar to gauge.
gauge.labels(**self.labels).set(data)
def _log_gauge_queue_count(self, gauge: Gauge, data: QueueCount) -> None:
# Log a QueueCount to gauge: total under default labels, per-priority breakdown under priority="<int>".
# NOTE: When priority scheduling is enabled, the total is recorded under
# priority="" (the default label value). Per-priority breakdowns are recorded
# with priority="<int>". Grafana queries should use priority="" for totals.
gauge.labels(**self.labels).set(data.total)
if data.by_priority is not None:
self._known_priorities.update(data.by_priority.keys())
for priority in self._known_priorities:
value = data.by_priority.get(priority, 0)
labels = dict(self.labels)
labels["priority"] = str(priority)
gauge.labels(**labels).set(value)
def _log_histogram(self, histogram, data: Union[int, float]) -> None:
histogram.labels(**self.labels).observe(data)
@@ -900,7 +901,7 @@ class SchedulerMetricsCollector:
).inc(t)
def log_stats(self, stats: SchedulerStats) -> None:
self._log_gauge(self.num_running_reqs, stats.num_running_reqs)
self._log_gauge_queue_count(self.num_running_reqs, stats.num_running_reqs)
self._log_gauge(self.num_used_tokens, stats.num_used_tokens)
self._log_gauge(self.token_usage, stats.token_usage)
self._log_gauge(
@@ -910,7 +911,7 @@ class SchedulerMetricsCollector:
self._log_gauge(self.mamba_usage, stats.mamba_usage)
self._log_gauge(self.decode_sum_seq_lens, stats.decode_sum_seq_lens)
self._log_gauge(self.gen_throughput, stats.gen_throughput)
self._log_gauge(self.num_queue_reqs, stats.num_queue_reqs)
self._log_gauge_queue_count(self.num_queue_reqs, stats.num_queue_reqs)
self._log_gauge(self.num_grammar_queue_reqs, stats.num_grammar_queue_reqs)
self._log_gauge(
self.num_running_reqs_offline_batch, stats.num_running_reqs_offline_batch
@@ -924,16 +925,16 @@ class SchedulerMetricsCollector:
self._log_gauge(self.spec_accept_rate, stats.spec_accept_rate)
# PD disaggregation
self._log_gauge(
self._log_gauge_queue_count(
self.num_prefill_prealloc_queue_reqs, stats.num_prefill_prealloc_queue_reqs
)
self._log_gauge(
self._log_gauge_queue_count(
self.num_prefill_inflight_queue_reqs, stats.num_prefill_inflight_queue_reqs
)
self._log_gauge(
self._log_gauge_queue_count(
self.num_decode_prealloc_queue_reqs, stats.num_decode_prealloc_queue_reqs
)
self._log_gauge(
self._log_gauge_queue_count(
self.num_decode_transfer_queue_reqs, stats.num_decode_transfer_queue_reqs
)
# Retract

View File

@@ -283,9 +283,9 @@ class SchedulerMetricsMixin:
if self.is_hybrid_ssm:
self.stats.mamba_usage = mamba_usage
_enable_ps = self.enable_priority_scheduling
priority_enabled = self.enable_priority_scheduling
self.stats.num_queue_reqs = QueueCount.from_reqs(
self.waiting_queue, _enable_ps
self.waiting_queue, priority_enabled
)
self.stats.num_grammar_queue_reqs = len(self.grammar_manager)
self.stats.cache_hit_rate = cache_hit_rate
@@ -300,19 +300,19 @@ class SchedulerMetricsMixin:
# PD disaggregation
if self.disaggregation_mode == DisaggregationMode.PREFILL:
self.stats.num_prefill_prealloc_queue_reqs = QueueCount.from_reqs(
self.disagg_prefill_bootstrap_queue.queue, _enable_ps
self.disagg_prefill_bootstrap_queue.queue, priority_enabled
)
self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
self.disagg_prefill_inflight_queue, _enable_ps
self.disagg_prefill_inflight_queue, priority_enabled
)
self.stats.kv_transfer_speed_gb_s = self.kv_transfer_speed_gb_s
self.stats.kv_transfer_latency_ms = self.kv_transfer_latency_ms
elif self.disaggregation_mode == DisaggregationMode.DECODE:
self.stats.num_decode_prealloc_queue_reqs = QueueCount.from_reqs(
self.disagg_decode_prealloc_queue.queue, _enable_ps
self.disagg_decode_prealloc_queue.queue, priority_enabled
)
self.stats.num_decode_transfer_queue_reqs = QueueCount.from_reqs(
self.disagg_decode_transfer_queue.queue, _enable_ps
self.disagg_decode_transfer_queue.queue, priority_enabled
)
# Others
@@ -436,9 +436,11 @@ class SchedulerMetricsMixin:
logger.info(msg)
if self.enable_metrics:
_enable_ps = self.enable_priority_scheduling
priority_enabled = self.enable_priority_scheduling
# Basics
self.stats.num_running_reqs = QueueCount.from_reqs(batch.reqs, _enable_ps)
self.stats.num_running_reqs = QueueCount.from_reqs(
batch.reqs, priority_enabled
)
self.stats.num_running_reqs_offline_batch = num_running_reqs_offline_batch
self.stats.num_used_tokens = num_used
self.stats.token_usage = token_usage
@@ -449,7 +451,7 @@ class SchedulerMetricsMixin:
self.stats.decode_sum_seq_lens = batch.seq_lens_cpu.sum().item()
self.stats.gen_throughput = self.last_gen_throughput
self.stats.num_queue_reqs = QueueCount.from_reqs(
self.waiting_queue, _enable_ps
self.waiting_queue, priority_enabled
)
self.stats.num_grammar_queue_reqs = len(self.grammar_manager)
self.stats.cache_hit_rate = cache_hit_rate
@@ -468,17 +470,17 @@ class SchedulerMetricsMixin:
# PD disaggregation
if self.disaggregation_mode == DisaggregationMode.PREFILL:
self.stats.num_prefill_prealloc_queue_reqs = QueueCount.from_reqs(
self.disagg_prefill_bootstrap_queue.queue, _enable_ps
self.disagg_prefill_bootstrap_queue.queue, priority_enabled
)
self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
self.disagg_prefill_inflight_queue, _enable_ps
self.disagg_prefill_inflight_queue, priority_enabled
)
elif self.disaggregation_mode == DisaggregationMode.DECODE:
self.stats.num_decode_prealloc_queue_reqs = QueueCount.from_reqs(
self.disagg_decode_prealloc_queue.queue, _enable_ps
self.disagg_decode_prealloc_queue.queue, priority_enabled
)
self.stats.num_decode_transfer_queue_reqs = QueueCount.from_reqs(
self.disagg_decode_transfer_queue.queue, _enable_ps
self.disagg_decode_transfer_queue.queue, priority_enabled
)
running_routing_keys = [r.routing_key for r in batch.reqs]
waiting_routing_keys = [r.routing_key for r in self.waiting_queue]

View File

@@ -5619,6 +5619,21 @@ class ServerArgs:
"fcfs",
"lof",
], f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{self.schedule_policy}' is not supported."
if self.default_priority_value is None:
logger.warning(
"--default-priority-value is not set while --enable-priority-scheduling is enabled. "
"Requests without explicit priority will have priority=None, "
"resulting in priority='None' string labels in Prometheus metrics."
)
else:
if self.disable_priority_preemption:
logger.warning(
"--disable-priority-preemption has no effect without --enable-priority-scheduling"
)
if self.default_priority_value is not None:
logger.warning(
"--default-priority-value has no effect without --enable-priority-scheduling"
)
# Check multi-item scoring
if self.multi_item_scoring_delimiter is not None:

View File

@@ -1,10 +1,12 @@
import unittest
from typing import Dict, List
from unittest.mock import Mock
import requests
from prometheus_client.parser import text_string_to_metric_families
from prometheus_client.samples import Sample
from sglang.srt.observability.metrics_collector import QueueCount
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import (
@@ -41,6 +43,36 @@ def _get_sample_value_by_labels(samples: List[Sample], labels: Dict[str, str]) -
raise KeyError(f"No sample found with labels {labels}")
class TestQueueCount(CustomTestCase):
"""Unit tests for QueueCount (no server needed)."""
def test_queue_count_from_reqs(self):
"""QueueCount correctly counts per-priority breakdown."""
reqs = [
Mock(priority=1),
Mock(priority=1),
Mock(priority=5),
Mock(priority=5),
Mock(priority=10),
]
qc = QueueCount.from_reqs(reqs, enable_priority_scheduling=True)
self.assertEqual(qc.total, 5)
self.assertEqual(qc.by_priority, {1: 2, 5: 2, 10: 1})
def test_queue_count_from_reqs_disabled(self):
"""Priority scheduling disabled → no breakdown."""
reqs = [Mock(priority=1), Mock(priority=5)]
qc = QueueCount.from_reqs(reqs, enable_priority_scheduling=False)
self.assertEqual(qc.total, 2)
self.assertIsNone(qc.by_priority)
def test_queue_count_empty(self):
"""Empty request list."""
qc = QueueCount.from_reqs([], enable_priority_scheduling=True)
self.assertEqual(qc.total, 0)
self.assertEqual(qc.by_priority, {})
class TestPriorityMetrics(CustomTestCase):
"""Test that priority-based metrics are correctly emitted when
--enable-priority-scheduling is enabled."""
@@ -118,27 +150,39 @@ class TestPriorityMetrics(CustomTestCase):
self.assertEqual(metrics_response.status_code, 200)
metrics = _parse_prometheus_metrics(metrics_response.text)
# Check histogram metrics have priority label
# Check histogram metrics have priority label with per-priority breakdown
histogram_metrics = [
"sglang:time_to_first_token_seconds",
"sglang:e2e_request_latency_seconds",
]
for metric_name in histogram_metrics:
# Histogram metrics are emitted as _sum, _count, _bucket
sum_name = f"{metric_name}_sum"
count_name = f"{metric_name}_count"
for suffix_name in [sum_name, count_name]:
samples = _get_samples_by_name(metrics, suffix_name)
if not samples:
continue
# At least one sample should have a non-empty priority label
priority_values = {s.labels.get("priority", "") for s in samples}
non_empty = priority_values - {""}
samples = _get_samples_by_name(metrics, count_name)
self.assertGreater(len(samples), 0, f"No samples found for {count_name}")
# At least one sample should have a non-empty priority label
priority_values = {s.labels.get("priority", "") for s in samples}
non_empty = priority_values - {""}
self.assertGreater(
len(non_empty),
0,
f"{count_name}: expected per-priority samples, "
f"got priority labels: {priority_values}",
)
# Verify that both priority="1" and priority="5" have count > 0
for expected_priority in ["1", "5"]:
matching = [
s for s in samples if s.labels.get("priority") == expected_priority
]
self.assertGreater(
len(non_empty),
len(matching),
0,
f"{suffix_name}: expected per-priority samples, "
f"got priority labels: {priority_values}",
f"{count_name}: no sample with priority='{expected_priority}'",
)
self.assertGreater(
matching[0].value,
0,
f"{count_name}: priority='{expected_priority}' count should be > 0",
)
def test_default_priority_value(self):