diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 0a687e250..20c1a1671 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -386,6 +386,7 @@ class Envs: # Metrics SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False) + SGLANG_ENABLE_METRICS_DP_ATTENTION = EnvBool(False) # fmt: on diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4b41ad9a7..9bd1b62bf 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -74,7 +74,11 @@ from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache -from sglang.srt.metrics.collector import SchedulerMetricsCollector, TimeStats +from sglang.srt.metrics.collector import ( + DPCooperationInfo, + SchedulerMetricsCollector, + TimeStats, +) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -1249,6 +1253,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Diffusion LLM dllm_config: Optional[DllmConfig] = None + # Metrics + dp_cooperation_info: Optional[DPCooperationInfo] = None + @classmethod def init_new( cls, @@ -2161,6 +2168,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_track_indices=self.mamba_track_indices, mamba_track_mask=self.mamba_track_mask, mamba_track_seqlens=self.mamba_track_seqlens, + dp_cooperation_info=self.dp_cooperation_info, ) def _is_available_size_sufficient(self, num_tokens: int) -> bool: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c7f98abf1..391c15867 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -716,6 +716,7 @@ class Scheduler( self.last_batch: Optional[ScheduleBatch] = None self.forward_ct = 0 self.last_prefill_tokens = 0 + self.last_prefill_cache_tokens = 0 self.return_health_check_ct = 0 self.num_retracted_reqs: int = 0 self.num_paused_reqs: int = 0 @@ -1830,6 +1831,8 @@ class Scheduler( if ret: trace_event_batch("schedule", ret.reqs) + self.log_prefill_stats_late(ret) + return ret def get_num_allocatable_reqs(self, running_bs): diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index af239473b..d601b7f8c 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py @@ -1,13 +1,14 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING, Callable, Optional import torch from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.metrics.collector import DPCooperationInfo from sglang.srt.utils.common import require_mlp_tp_gather if TYPE_CHECKING: @@ -15,6 +16,9 @@ if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler +_ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get() + + @dataclass class MLPSyncBatchInfo: dp_size: int @@ -33,6 +37,7 @@ class MLPSyncBatchInfo: global_num_tokens_for_logprob: list[int] = None tbo_split_seq_index: torch.Tensor = None global_forward_mode: int = None + dp_cooperation_info: Optional[DPCooperationInfo] = None def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor: return torch.tensor( @@ -68,6 +73,8 @@ class MLPSyncBatchInfo: self.global_num_tokens_for_logprob = tp0_info[:, 1].tolist() self.can_cuda_graph = bool(tp0_info[:, 2].min().item()) self.is_extend_in_batch = bool(tp0_info[:, 3].max().item()) + if _ENABLE_METRICS_DP_ATTENTION: + self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist()) def _update_gather_batch( @@ -180,6 +187,9 @@ def prepare_mlp_sync_batch_raw( batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather ) + if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None: + local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info + return local_batch diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index fcce6bff1..26ace7f43 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -88,7 +88,7 @@ class SchedulerMetricsMixin: if ENABLE_METRICS_DEVICE_TIMER: self.forward_pass_device_timer = DeviceTimer( - reporter=self.metrics_collector.increment_gpu_execution_seconds + reporter=self.metrics_collector.increment_gpu_execution_seconds, ) if self.enable_kv_cache_events: @@ -124,6 +124,7 @@ class SchedulerMetricsMixin: self.last_prefill_stats_tic = time.perf_counter() self.last_input_throughput = self.last_prefill_tokens / gap_latency self.last_prefill_tokens = adder.log_input_tokens + self.last_prefill_cache_tokens = adder.log_hit_tokens # TODO: generalize this for various memory pools if self.is_hybrid_swa: @@ -231,23 +232,26 @@ class SchedulerMetricsMixin: self.disagg_decode_transfer_queue.queue ) - self.metrics_collector.increment_realtime_tokens( - prefill_compute_tokens=adder.log_input_tokens, - prefill_cache_tokens=adder.log_hit_tokens, - ) - # Others self.calculate_utilization() self.metrics_collector.log_stats(self.stats) self._emit_kv_metrics() self._publish_kv_events() + def log_prefill_stats_late(self: Scheduler, batch: Optional[ScheduleBatch]): + """This should be called after `batch` has gathered enough metadata.""" + if self.enable_metrics and batch is not None: + self.metrics_collector.increment_realtime_tokens( + prefill_compute_tokens=self.last_prefill_tokens, + prefill_cache_tokens=self.last_prefill_cache_tokens, + dp_cooperation_info=batch.dp_cooperation_info, + ) + def log_decode_stats( self: Scheduler, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None ): batch = running_batch or self.running_batch - last_num_generated_tokens = self.num_generated_tokens gap_latency = time.perf_counter() - self.last_decode_stats_tic self.last_decode_stats_tic = time.perf_counter() self.last_gen_throughput = self.num_generated_tokens / gap_latency @@ -388,16 +392,21 @@ class SchedulerMetricsMixin: self.disagg_decode_transfer_queue.queue ) - self.metrics_collector.increment_realtime_tokens( - decode_tokens=last_num_generated_tokens - ) - # Others self.calculate_utilization() self.metrics_collector.log_stats(self.stats) self._emit_kv_metrics() self._publish_kv_events() + def log_decode_stats_every_iteration( + self: Scheduler, batch: ScheduleBatch, num_accepted_tokens: int + ): + self.metrics_collector.increment_realtime_tokens( + # TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator + decode_tokens=batch.batch_size() + num_accepted_tokens, + dp_cooperation_info=batch.dp_cooperation_info, + ) + def log_batch_result_stats( self: Scheduler, batch: ScheduleBatch, @@ -491,5 +500,10 @@ class SchedulerMetricsMixin: return category = "forward_" + batch.forward_mode.name.lower() - with self.forward_pass_device_timer.wrap(category=category): + with self.forward_pass_device_timer.wrap( + metadata=dict( + category=category, + dp_cooperation_info=batch.dp_cooperation_info, + ), + ): yield diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 4228dc536..83050f76b 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -445,6 +445,10 @@ class SchedulerOutputProcessorMixin: and self.forward_ct_decode % self.server_args.decode_log_interval == 0 ): self.log_decode_stats(can_run_cuda_graph, running_batch=batch) + if self.enable_metrics: + self.log_decode_stats_every_iteration( + batch, num_accepted_tokens=result.num_accepted_tokens + ) def _mamba_prefix_cache_update( self, req: Req, batch: ScheduleBatch, result: GenerationBatchResult, i: int diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index 17623db2c..92736e894 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/python/sglang/srt/metrics/collector.py @@ -12,6 +12,7 @@ # limitations under the License. # ============================================================================== """Utilities for Prometheus Metrics Collection.""" +import dataclasses import logging import os import time @@ -21,6 +22,7 @@ from typing import Dict, List, Optional, Union from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs from sglang.srt.metrics.utils import exponential_buckets, generate_buckets +from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.server_args import ServerArgs from sglang.srt.utils import get_bool_env_var @@ -240,6 +242,24 @@ class SchedulerStats: is_cuda_graph: float = 0.0 +@dataclass +class DPCooperationInfo: + # Users can derive that, except for cases with idle, num_decode_ranks=world_size-num_prefill_ranks + # We do not provide `num_decode_ranks` to avoid cardinality explosion. + num_prefill_ranks: int + + @staticmethod + def create(forward_modes: List[int]): + return DPCooperationInfo( + num_prefill_ranks=sum( + 1 for mode in forward_modes if mode == ForwardMode.EXTEND.value + ), + ) + + def to_labels(self): + return dataclasses.asdict(self) + + class SchedulerMetricsCollector: def __init__( @@ -680,6 +700,17 @@ class SchedulerMetricsCollector: labelnames=list(labels.keys()) + ["category"], ) + self.dp_cooperation_realtime_tokens_total = Counter( + name="sglang:dp_cooperation_realtime_tokens_total", + documentation="Total number of tokens processed with labels about DP cooperation.", + labelnames=list(labels.keys()) + ["mode", "num_prefill_ranks"], + ) + self.dp_cooperation_gpu_execution_seconds_total = Counter( + name="sglang:dp_cooperation_gpu_execution_seconds_total", + documentation="Total time that GPU is busy executing a workload with labels about DP cooperation.", + labelnames=list(labels.keys()) + ["category", "num_prefill_ranks"], + ) + def _log_gauge(self, gauge, data: Union[int, float]) -> None: # Convenience function for logging to gauge. gauge.labels(**self.labels).set(data) @@ -716,7 +747,11 @@ class SchedulerMetricsCollector: ) def increment_realtime_tokens( - self, prefill_compute_tokens=0, prefill_cache_tokens=0, decode_tokens=0 + self, + dp_cooperation_info: Optional[DPCooperationInfo], + prefill_compute_tokens=0, + prefill_cache_tokens=0, + decode_tokens=0, ): for mode, delta in [ ("prefill_compute", prefill_compute_tokens), @@ -724,10 +759,27 @@ class SchedulerMetricsCollector: ("decode", decode_tokens), ]: self.realtime_tokens_total.labels(**self.labels, mode=mode).inc(delta) + if dp_cooperation_info is not None: + self.dp_cooperation_realtime_tokens_total.labels( + **self.labels, + mode=mode, + **dp_cooperation_info.to_labels(), + ).inc(delta) - def increment_gpu_execution_seconds(self, category: str, t: float): + def increment_gpu_execution_seconds( + self, + category: str, + t: float, + dp_cooperation_info: Optional[DPCooperationInfo], + ): logger.debug(f"GPU execution seconds: {category=} {t=:.3f}") self.gpu_execution_seconds_total.labels(**self.labels, category=category).inc(t) + if dp_cooperation_info is not None: + self.dp_cooperation_gpu_execution_seconds_total.labels( + **self.labels, + category=category, + **dp_cooperation_info.to_labels(), + ).inc(t) def log_stats(self, stats: SchedulerStats) -> None: self._log_gauge(self.num_running_reqs, stats.num_running_reqs) diff --git a/python/sglang/srt/utils/device_timer.py b/python/sglang/srt/utils/device_timer.py index e426918bb..686c8d3d9 100644 --- a/python/sglang/srt/utils/device_timer.py +++ b/python/sglang/srt/utils/device_timer.py @@ -1,23 +1,23 @@ from collections import deque from contextlib import contextmanager from dataclasses import dataclass -from typing import Callable, Deque, Optional +from typing import Callable, Deque, Dict, Optional import torch class DeviceTimer: - def __init__(self, reporter: Callable[[str, float], None]): + def __init__(self, reporter: Callable): self._intervals: Deque[_TimingInterval] = deque() self._reporter = reporter @contextmanager - def wrap(self, category: str): + def wrap(self, metadata: Dict): self._intervals.append(_TimingInterval.create()) try: yield finally: - self._intervals[-1].end(category=category) + self._intervals[-1].end(metadata=metadata) self._report() def _report(self): @@ -27,14 +27,14 @@ class DeviceTimer: break self._intervals.popleft() - self._reporter(interval.category, interval.elapsed_time() / 1000.0) + self._reporter(t=interval.elapsed_time() / 1000.0, **interval.metadata) @dataclass class _TimingInterval: start_event: torch.cuda.Event end_event: Optional[torch.cuda.Event] = None - category: Optional[str] = None + metadata: Optional[Dict] = None @staticmethod def create(): @@ -42,13 +42,13 @@ class _TimingInterval: start_event.record() return _TimingInterval(start_event=start_event) - def end(self, category: str): + def end(self, metadata: Dict): end_event = torch.cuda.Event(enable_timing=True) end_event.record() assert self.end_event is None self.end_event = end_event - self.category = category + self.metadata = metadata def elapsed_time(self) -> float: return self.start_event.elapsed_time(self.end_event) diff --git a/test/srt/test_metrics.py b/test/srt/test_metrics.py index ce85e0d8a..096a4e221 100644 --- a/test/srt/test_metrics.py +++ b/test/srt/test_metrics.py @@ -1,27 +1,79 @@ import unittest +from typing import Dict, List import requests +from prometheus_client.parser import text_string_to_metric_families +from prometheus_client.samples import Sample +from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree from sglang.test.test_utils import ( - DEFAULT_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, + is_in_ci, popen_launch_server, ) +_MODEL_NAME = "Qwen/Qwen3-0.6B" + class TestEnableMetrics(CustomTestCase): - def test_metrics_enabled(self): + def test_metrics_1gpu(self): """Test that metrics endpoint returns data when enabled""" - process = popen_launch_server( - DEFAULT_SMALL_MODEL_NAME_FOR_TEST, - DEFAULT_URL_FOR_TEST, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=["--enable-metrics", "--cuda-graph-max-bs", 2], + self._execute_core( + other_args=[], + verify_metrics_extra=None, ) + def test_metrics_2gpu(self): + # TODO enable when we have 2-gpu runner in nightly CI + if is_in_ci(): + print("Skip test_metrics_2gpu since in 1-gpu CI") + return + + def _verify_metrics_extra(metrics): + metrics_to_check = [ + ( + "sglang:dp_cooperation_realtime_tokens_total", + {"mode": "prefill_compute"}, + ), + ("sglang:dp_cooperation_realtime_tokens_total", {"mode": "decode"}), + ( + "sglang:dp_cooperation_gpu_execution_seconds_total", + {"category": "forward_prefill"}, + ), + ( + "sglang:dp_cooperation_gpu_execution_seconds_total", + {"category": "forward_decode"}, + ), + ] + _check_metrics_positive(self, metrics, metrics_to_check) + + num_prefill_ranks_values = { + s.labels["num_prefill_ranks"] + for s in metrics["sglang:dp_cooperation_realtime_tokens_total"] + } + self.assertIn("0", num_prefill_ranks_values) + self.assertIn("1", num_prefill_ranks_values) + + self._execute_core( + other_args=["--tp", "2", "--dp", "2", "--enable-dp-attention"], + verify_metrics_extra=_verify_metrics_extra, + ) + + def _execute_core(self, other_args, verify_metrics_extra): + with ( + envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.override(True), + envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.override(True), + ): + process = popen_launch_server( + _MODEL_NAME, + DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--enable-metrics", "--cuda-graph-max-bs", 2, *other_args], + ) + try: # Make some requests to generate some metrics response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate") @@ -45,44 +97,74 @@ class TestEnableMetrics(CustomTestCase): # Get metrics metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics") self.assertEqual(metrics_response.status_code, 200) - metrics_content = metrics_response.text + metrics_text = metrics_response.text - print(f"metrics_content=\n{metrics_content}") - - # Verify essential metrics are present - essential_metrics = [ - "sglang:num_running_reqs", - "sglang:num_used_tokens", - "sglang:token_usage", - "sglang:gen_throughput", - "sglang:num_queue_reqs", - "sglang:num_grammar_queue_reqs", - "sglang:cache_hit_rate", - "sglang:spec_accept_length", - "sglang:prompt_tokens_total", - "sglang:generation_tokens_total", - "sglang:cached_tokens_total", - "sglang:num_requests_total", - "sglang:time_to_first_token_seconds", - "sglang:inter_token_latency_seconds", - "sglang:e2e_request_latency_seconds", - ] - - for metric in essential_metrics: - self.assertIn(metric, metrics_content, f"Missing metric: {metric}") - - # Verify model name label is present and correct - expected_model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST - self.assertIn(f'model_name="{expected_model_name}"', metrics_content) - - # Verify metrics have values (not empty) - self.assertIn("_sum{", metrics_content) - self.assertIn("_count{", metrics_content) - self.assertIn("_bucket{", metrics_content) + print(f"metrics_text=\n{metrics_text}") + metrics = _parse_prometheus_metrics(metrics_text) + self._verify_metrics_common(metrics_text, metrics) + if verify_metrics_extra is not None: + verify_metrics_extra(metrics) finally: kill_process_tree(process.pid) + def _verify_metrics_common(self, metrics_text, metrics): + essential_metrics = [ + "sglang:num_running_reqs", + "sglang:num_used_tokens", + "sglang:token_usage", + "sglang:gen_throughput", + "sglang:num_queue_reqs", + "sglang:num_grammar_queue_reqs", + "sglang:cache_hit_rate", + "sglang:spec_accept_length", + "sglang:prompt_tokens_total", + "sglang:generation_tokens_total", + "sglang:cached_tokens_total", + "sglang:num_requests_total", + "sglang:time_to_first_token_seconds", + "sglang:inter_token_latency_seconds", + "sglang:e2e_request_latency_seconds", + ] + for metric in essential_metrics: + self.assertIn(metric, metrics_text, f"Missing metric: {metric}") + + self.assertIn(f'model_name="{_MODEL_NAME}"', metrics_text) + self.assertIn("_sum{", metrics_text) + self.assertIn("_count{", metrics_text) + self.assertIn("_bucket{", metrics_text) + + metrics_to_check = [ + ("sglang:realtime_tokens_total", {"mode": "prefill_compute"}), + ("sglang:realtime_tokens_total", {"mode": "decode"}), + ("sglang:gpu_execution_seconds_total", {"category": "forward_extend"}), + ("sglang:gpu_execution_seconds_total", {"category": "forward_decode"}), + ] + _check_metrics_positive(self, metrics, metrics_to_check) + + +def _parse_prometheus_metrics(metrics_text: str) -> Dict[str, List[Sample]]: + result = {} + for family in text_string_to_metric_families(metrics_text): + for sample in family.samples: + if sample.name not in result: + result[sample.name] = [] + result[sample.name].append(sample) + return result + + +def _get_sample_value_by_labels(samples: List[Sample], labels: Dict[str, str]) -> float: + for sample in samples: + if all(sample.labels.get(k) == v for k, v in labels.items()): + return sample.value + raise KeyError(f"No sample found with labels {labels}") + + +def _check_metrics_positive(test_case, metrics, metrics_to_check): + for metric_name, labels in metrics_to_check: + value = _get_sample_value_by_labels(metrics[metric_name], labels) + test_case.assertGreater(value, 0, f"{metric_name} {labels}") + if __name__ == "__main__": unittest.main()