From 11892599f164d1c2da9c489e5f0b922c2ec6a3c4 Mon Sep 17 00:00:00 2001 From: Yinghai Lu Date: Sat, 31 Jan 2026 22:25:59 -0800 Subject: [PATCH] [metric] Optional extra metric labels (#18049) --- python/sglang/srt/managers/scheduler_metrics_mixin.py | 2 ++ python/sglang/srt/managers/tokenizer_manager.py | 2 ++ python/sglang/srt/mem_cache/base_prefix_cache.py | 10 +++++++--- python/sglang/srt/mem_cache/hiradix_cache.py | 9 ++++++++- python/sglang/srt/server_args.py | 8 ++++++++ 5 files changed, 27 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 1a3ba69a1..1c9c6a185 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -109,6 +109,8 @@ class SchedulerMetricsMixin: } if dp_rank is not None: labels["dp_rank"] = dp_rank + if self.server_args.extra_metric_labels: + labels.update(self.server_args.extra_metric_labels) self.metrics_collector = SchedulerMetricsCollector( labels=labels, enable_lora=self.enable_lora, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 69806de1d..8522da97c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -437,6 +437,8 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi if self.server_args.tokenizer_metrics_allowed_custom_labels: for label in self.server_args.tokenizer_metrics_allowed_custom_labels: labels[label] = "" + if self.server_args.extra_metric_labels: + labels.update(self.server_args.extra_metric_labels) self.metrics_collector = TokenizerMetricsCollector( server_args=self.server_args, labels=labels, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index e01f98dcb..98680d314 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -119,9 +119,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait): ) def init_metrics_collector(self): - self.metrics_collector = RadixCacheMetricsCollector( - labels={"cache_type": self.__class__.__name__} - ) + from sglang.srt.server_args import get_global_server_args + + server_args = get_global_server_args() + labels = {"cache_type": self.__class__.__name__} + if server_args.extra_metric_labels: + labels.update(server_args.extra_metric_labels) + self.metrics_collector = RadixCacheMetricsCollector(labels=labels) def update_eviction_metrics(self, num_evicted: int, start_time: float): if self.metrics_collector is not None and num_evicted > 0: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 268e4d72e..e8fa02cfd 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -8,7 +8,7 @@ import os import threading import time from queue import Empty -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional import torch @@ -88,6 +88,7 @@ class HiRadixCache(RadixCache): self.pp_size = params.pp_size self.enable_storage = server_args.hicache_storage_backend is not None self.enable_storage_metrics = self.enable_storage and params.enable_metrics + self.extra_metric_labels = server_args.extra_metric_labels ( extra_config, @@ -126,6 +127,7 @@ class HiRadixCache(RadixCache): hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys, enable_storage=self.enable_storage, enable_storage_metrics=self.enable_storage_metrics, + extra_metric_labels=self.extra_metric_labels, ) # record the nodes with ongoing write through @@ -168,6 +170,7 @@ class HiRadixCache(RadixCache): hicache_storage_pass_prefix_keys: bool, enable_storage: bool, enable_storage_metrics: bool, + extra_metric_labels: Optional[Dict[str, str]], ) -> None: prefetch_timeout_per_page = ( self.page_size / 1024 * prefetch_timeout_per_ki_token @@ -182,6 +185,9 @@ class HiRadixCache(RadixCache): "pp_rank": self.cache_controller.pp_rank, "pp_size": self.cache_controller.pp_size, } + if extra_metric_labels: + labels.update(extra_metric_labels) + self.storage_metrics_collector = StorageMetricsCollector(labels=labels) storage_metrics_collector = StorageMetricsCollector(labels=labels) self.enable_storage = enable_storage @@ -310,6 +316,7 @@ class HiRadixCache(RadixCache): hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys, enable_storage=True, enable_storage_metrics=self._enable_metrics_flag, + extra_metric_labels=self.extra_metric_labels, ) return True, "Attached HiCache storage backend successfully." diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1c7074715..b1fd7edf7 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -380,6 +380,7 @@ class ServerArgs: enable_metrics_for_all_schedulers: bool = False tokenizer_metrics_custom_labels_header: str = "x-custom-labels" tokenizer_metrics_allowed_custom_labels: Optional[List[str]] = None + extra_metric_labels: Optional[Dict[str, str]] = None bucket_time_to_first_token: Optional[List[float]] = None bucket_inter_token_latency: Optional[List[float]] = None bucket_e2e_request_latency: Optional[List[float]] = None @@ -3369,6 +3370,13 @@ class ServerArgs: "'--tokenizer-metrics-custom-labels-header' field in HTTP requests, e.g., {'label1': 'value1', 'label2': " "'value2'} is allowed if '--tokenizer-metrics-allowed-custom-labels label1 label2' is set.", ) + parser.add_argument( + "--extra-metric-labels", + type=json.loads, + default=ServerArgs.extra_metric_labels, + help="The custom labels for metrics. " + 'e.g. \'{"label1": "value1", "label2": "value2"}\'', + ) parser.add_argument( "--bucket-time-to-first-token", type=float,