From 17b38f88bc2da5d701250a4893095ce3561795d0 Mon Sep 17 00:00:00 2001 From: Chenxi Li <41864925+ConnorLi96@users.noreply.github.com> Date: Wed, 24 Dec 2025 12:13:52 -0800 Subject: [PATCH] Add LoRA metrics for potential auto scaling (#15149) --- .../srt/managers/scheduler_metrics_mixin.py | 50 ++++++++++++++++++- python/sglang/srt/metrics/collector.py | 34 +++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 26ace7f43..22437b4b3 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -84,7 +84,9 @@ class SchedulerMetricsMixin: } if dp_rank is not None: labels["dp_rank"] = dp_rank - self.metrics_collector = SchedulerMetricsCollector(labels=labels) + self.metrics_collector = SchedulerMetricsCollector( + labels=labels, enable_lora=self.enable_lora + ) if ENABLE_METRICS_DEVICE_TIMER: self.forward_pass_device_timer = DeviceTimer( @@ -234,6 +236,7 @@ class SchedulerMetricsMixin: # Others self.calculate_utilization() + self.update_lora_metrics() self.metrics_collector.log_stats(self.stats) self._emit_kv_metrics() self._publish_kv_events() @@ -394,6 +397,7 @@ class SchedulerMetricsMixin: # Others self.calculate_utilization() + self.update_lora_metrics() self.metrics_collector.log_stats(self.stats) self._emit_kv_metrics() self._publish_kv_events() @@ -451,6 +455,50 @@ class SchedulerMetricsMixin: batch = KVEventBatch(ts=time.time(), events=events) self.kv_event_publisher.publish(batch) + def update_lora_metrics(self: Scheduler): + """Update LoRA pool metrics for monitoring and autoscaling.""" + if not self.enable_lora: + return + + try: + # Get LoRA memory pool stats + lora_manager = self.tp_worker.model_runner.lora_manager + if lora_manager is None or lora_manager.memory_pool is None: + return + + mem_pool = lora_manager.memory_pool + slots_total = mem_pool.max_loras_per_batch + + # Calculate active adapters from running batch + # This gives a true measure of current load for autoscaling purposes + active_lora_ids = set() + + # For PP mode, check all running micro batches + if hasattr(self, "running_mbs") and self.running_mbs: + for batch in self.running_mbs: + if batch and hasattr(batch, "reqs"): + for req in batch.reqs: + if hasattr(req, "lora_id") and req.lora_id is not None: + active_lora_ids.add(req.lora_id) + # For normal mode, check running_batch + elif hasattr(self, "running_batch") and self.running_batch: + if hasattr(self.running_batch, "reqs"): + for req in self.running_batch.reqs: + if hasattr(req, "lora_id") and req.lora_id is not None: + active_lora_ids.add(req.lora_id) + + # Count active adapters (excluding None for base model) + slots_used = len(active_lora_ids) + utilization = slots_used / slots_total if slots_total > 0 else 0.0 + + # Update stats + self.stats.lora_pool_slots_used = slots_used + self.stats.lora_pool_slots_total = slots_total + self.stats.lora_pool_utilization = utilization + + except Exception as e: + logger.warning(f"Failed to update LoRA metrics: {e}") + def calculate_utilization(self): if self.disaggregation_mode == DisaggregationMode.PREFILL: self.stats.utilization = -1 diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index edaada3b4..c85fc07b7 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/python/sglang/srt/metrics/collector.py @@ -241,6 +241,11 @@ class SchedulerStats: # CUDA graph is_cuda_graph: float = 0.0 + # LoRA pool metrics + lora_pool_slots_used: int = 0 + lora_pool_slots_total: int = 0 + lora_pool_utilization: float = 0.0 + @dataclass class DPCooperationInfo: @@ -265,11 +270,13 @@ class SchedulerMetricsCollector: def __init__( self, labels: Dict[str, str], + enable_lora: bool = False, ) -> None: # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` from prometheus_client import Counter, Gauge, Histogram, Summary self.labels = labels + self.enable_lora = enable_lora self.last_log_time = time.perf_counter() self.num_running_reqs = Gauge( @@ -692,6 +699,27 @@ class SchedulerMetricsCollector: labelnames=list(labels.keys()) + ["forward_mode"], ) + # LoRA pool metrics (only created when LoRA is enabled) + if self.enable_lora: + self.lora_pool_slots_used = Gauge( + name="sglang:lora_pool_slots_used", + documentation="Number of LoRA adapter slots currently occupied in GPU memory.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.lora_pool_slots_total = Gauge( + name="sglang:lora_pool_slots_total", + documentation="Total number of LoRA adapter slots available (max_loras_per_batch).", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.lora_pool_utilization = Gauge( + name="sglang:lora_pool_utilization", + documentation="LoRA pool utilization ratio (used/total). 1.0 means pool is full.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.new_token_ratio = Gauge( name="sglang:new_token_ratio", documentation="The new token ratio.", @@ -868,6 +896,12 @@ class SchedulerMetricsCollector: # CUDA graph self._log_gauge(self.is_cuda_graph, stats.is_cuda_graph) + # LoRA pool metrics (only logged if LoRA is enabled) + if self.enable_lora: + self._log_gauge(self.lora_pool_slots_used, stats.lora_pool_slots_used) + self._log_gauge(self.lora_pool_slots_total, stats.lora_pool_slots_total) + self._log_gauge(self.lora_pool_utilization, stats.lora_pool_utilization) + self.last_log_time = time.perf_counter() def log_grammar_stats(self, grammar_stats) -> None: