feat: Priority-based scheduling optimization (including default priority, preemption toggle, priority-based metrics, etc.)

python/sglang/srt/managers/scheduler.py

@@ -814,7 +814,10 @@ def init_schedule_policy(self):
                 ),
             )
         # Enable preemption for priority scheduling.
-        self.try_preemption = self.enable_priority_scheduling
+        self.try_preemption = (
+            self.enable_priority_scheduling and not self.server_args.disable_preemption
+        )
+
         self.init_new_token_ratio = min(
             envs.SGLANG_INIT_NEW_TOKEN_RATIO.get()
             * self.server_args.schedule_conservativeness,
@@ -2004,6 +2007,7 @@ def _get_new_batch_prefill_raw(
             return None
 
         running_bs = len(self.running_batch.reqs)
+
         # Ignore the check if self.chunked_req is not None.
         # In the non-PP case, when self.chunked_req is not None, num_allocatable_reqs should always be greater than 0,
         # as the space for the chunked requests has just been released.
@@ -2174,12 +2178,16 @@ def _get_new_batch_prefill_raw(
         new_batch.prepare_for_extend()
 
         # Record prefill stats for logging after forward
+        num_running_reqs_by_priority = self._count_reqs_by_priority(
+            self.running_batch.reqs
+        )
         new_batch.prefill_stats = PrefillStats(
             log_input_tokens=adder.log_input_tokens,
             log_hit_tokens=adder.log_hit_tokens,
             new_token_ratio=adder.new_token_ratio,
             running_bs=len(self.running_batch.reqs),
             num_new_seqs=len(can_run_list),
+            num_running_reqs_by_priority=num_running_reqs_by_priority,
         )
 
         # Mixed-style chunked prefill

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -303,25 +303,49 @@ def check_memory(self: Scheduler):
                 num_used, token_usage, _, _ = self._get_token_info()
             num_running_reqs = len(self.running_batch.reqs)
             self.stats.num_running_reqs = num_running_reqs
+            self.stats.num_running_reqs_by_priority = self._count_reqs_by_priority(
+                self.running_batch.reqs
+            )
             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 = len(self.waiting_queue)
+            self.stats.num_queue_reqs_by_priority = self._count_reqs_by_priority(
+                self.waiting_queue
+            )
             self.stats.num_grammar_queue_reqs = len(self.grammar_manager)
             if self.disaggregation_mode == DisaggregationMode.PREFILL:
                 self.stats.num_prefill_prealloc_queue_reqs = len(
                     self.disagg_prefill_bootstrap_queue.queue
                 )
+                self.stats.num_prefill_prealloc_queue_reqs_by_priority = (
+                    self._count_reqs_by_priority(
+                        self.disagg_prefill_bootstrap_queue.queue
+                    )
+                )
                 self.stats.num_prefill_inflight_queue_reqs = len(
                     self.disagg_prefill_inflight_queue
                 )
+                self.stats.num_prefill_inflight_queue_reqs_by_priority = (
+                    self._count_reqs_by_priority(self.disagg_prefill_inflight_queue)
+                )
             if self.disaggregation_mode == DisaggregationMode.DECODE:
                 self.stats.num_decode_prealloc_queue_reqs = len(
                     self.disagg_decode_prealloc_queue.queue
                 )
+                self.stats.num_decode_prealloc_queue_reqs_by_priority = (
+                    self._count_reqs_by_priority(
+                        self.disagg_decode_prealloc_queue.queue
+                    )
+                )
                 self.stats.num_decode_transfer_queue_reqs = len(
                     self.disagg_decode_transfer_queue.queue
                 )
+                self.stats.num_decode_transfer_queue_reqs_by_priority = (
+                    self._count_reqs_by_priority(
+                        self.disagg_decode_transfer_queue.queue
+                    )
+                )
             self.metrics_collector.log_stats(self.stats)
         self._publish_kv_events()
 

python/sglang/srt/managers/tokenizer_manager.py

@@ -233,6 +233,8 @@ def init_model_config(self):
         self.context_len = self.model_config.context_len
         self.image_token_id = self.model_config.image_token_id
         self.max_req_input_len = None  # Will be set later in engine.py
+        self.enable_priority_scheduling = server_args.enable_priority_scheduling
+        self.default_priority_value = server_args.default_priority_value
         speculative_algorithm = SpeculativeAlgorithm.from_string(
             server_args.speculative_algorithm
         )
@@ -421,6 +423,8 @@ def init_metric_collector_watchdog(self):
                 "model_name": self.server_args.served_model_name,
                 # TODO: Add lora name/path in the future,
             }
+            if self.enable_priority_scheduling:
+                labels["priority"] = ""
             if self.server_args.tokenizer_metrics_allowed_custom_labels:
                 for label in self.server_args.tokenizer_metrics_allowed_custom_labels:
                     labels[label] = ""
@@ -482,6 +486,7 @@ async def generate_request(
 
         # Normalize the request
         obj.normalize_batch_and_arguments()
+        self._set_default_priority(obj)
         self._validate_rid(obj)
 
         if isinstance(obj, GenerateReqInput) and obj.routed_dp_rank is not None:
@@ -1899,6 +1904,10 @@ def collect_metrics(self, state: ReqState, recv_obj: BatchStrOutput, i: int):
             if custom_labels
             else self.metrics_collector.labels
         )
+        if self.enable_priority_scheduling:
+            priority = getattr(state.obj, "priority", None)
+            if priority is not None:
+                labels["priority"] = priority
         if (
             state.time_stats.first_token_time == 0.0
             and self.disaggregation_mode != DisaggregationMode.PREFILL
@@ -2367,6 +2376,15 @@ def convert_to_span_attrs(
 
         return span_attrs
 
+    def _set_default_priority(self, obj: Union[GenerateReqInput, EmbeddingReqInput]):
+        """Set the default priority value."""
+        if (
+            self.enable_priority_scheduling
+            and obj.priority is None
+            and self.default_priority_value is not None
+        ):
+            obj.priority = self.default_priority_value
+
 
 class ServerStatus(Enum):
     Up = "Up"

python/sglang/srt/observability/metrics_collector.py

@@ -18,7 +18,7 @@
 import os
 import time
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Set, Union
 
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
@@ -29,7 +29,6 @@
 
 SGLANG_TEST_REQUEST_TIME_STATS = get_bool_env_var("SGLANG_TEST_REQUEST_TIME_STATS")
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -51,6 +50,7 @@ def get_histogram_conf_from_env(env_var_name: str) -> Optional[List[float]]:
 class SchedulerStats:
     # Basics
     num_running_reqs: int = 0
+    num_running_reqs_by_priority: Optional[Dict[int, int]] = None
     num_used_tokens: int = 0
     token_usage: float = 0.0
     pending_prealloc_token_usage: float = 0.0
@@ -59,6 +59,7 @@ class SchedulerStats:
     decode_sum_seq_lens: int = 0
     gen_throughput: float = 0.0
     num_queue_reqs: int = 0
+    num_queue_reqs_by_priority: Optional[Dict[int, int]] = None
     num_grammar_queue_reqs: int = 0
     num_running_reqs_offline_batch: int = 0
     cache_hit_rate: float = 0.0
@@ -75,9 +76,13 @@ class SchedulerStats:
 
     # PD disaggregation
     num_prefill_prealloc_queue_reqs: int = 0
+    num_prefill_prealloc_queue_reqs_by_priority: Optional[Dict[int, int]] = None
     num_prefill_inflight_queue_reqs: int = 0
+    num_prefill_inflight_queue_reqs_by_priority: Optional[Dict[int, int]] = None
     num_decode_prealloc_queue_reqs: int = 0
+    num_decode_prealloc_queue_reqs_by_priority: Optional[Dict[int, int]] = None
     num_decode_transfer_queue_reqs: int = 0
+    num_decode_transfer_queue_reqs_by_priority: Optional[Dict[int, int]] = None
     kv_transfer_speed_gb_s: float = 0.0
     kv_transfer_latency_ms: float = 0.0
     kv_transfer_bootstrap_ms: float = 0.0
@@ -158,6 +163,7 @@ def __init__(
         self.enable_lora = enable_lora
         self.enable_hierarchical_cache = enable_hierarchical_cache
         self.last_log_time = time.perf_counter()
+        self._known_priorities: Set[int] = set()
 
         self.num_running_reqs = Gauge(
             name="sglang:num_running_reqs",
@@ -733,9 +739,22 @@ def __init__(
             multiprocess_mode="mostrecent",
         )
 
-    def _log_gauge(self, gauge, data: Union[int, float]) -> None:
+    def _log_gauge(
+        self,
+        gauge,
+        data: Union[int, float],
+        by_priority_dict: Optional[Dict[int, int]] = None,
+    ) -> None:
         # Convenience function for logging to gauge.
-        gauge.labels(**self.labels).set(data)
+        if by_priority_dict is not None:
+            self._known_priorities.update(by_priority_dict.keys())
+            for priority in self._known_priorities:
+                value = by_priority_dict.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_histogram(self, histogram, data: Union[int, float]) -> None:
         histogram.labels(**self.labels).observe(data)
@@ -844,7 +863,11 @@ def increment_gpu_execution_seconds(
             ).inc(t)
 
     def log_stats(self, stats: SchedulerStats) -> None:
-        self._log_gauge(self.num_running_reqs, stats.num_running_reqs)
+        self._log_gauge(
+            self.num_running_reqs,
+            stats.num_running_reqs,
+            stats.num_running_reqs_by_priority,
+        )
         self._log_gauge(self.num_used_tokens, stats.num_used_tokens)
         self._log_gauge(self.token_usage, stats.token_usage)
         self._log_gauge(
@@ -854,7 +877,9 @@ def log_stats(self, stats: SchedulerStats) -> None:
         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(
+            self.num_queue_reqs, stats.num_queue_reqs, stats.num_queue_reqs_by_priority
+        )
         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
@@ -869,16 +894,24 @@ def log_stats(self, stats: SchedulerStats) -> None:
 
         # PD disaggregation
         self._log_gauge(
-            self.num_prefill_prealloc_queue_reqs, stats.num_prefill_prealloc_queue_reqs
+            self.num_prefill_prealloc_queue_reqs,
+            stats.num_prefill_prealloc_queue_reqs,
+            stats.num_prefill_prealloc_queue_reqs_by_priority,
         )
         self._log_gauge(
-            self.num_prefill_inflight_queue_reqs, stats.num_prefill_inflight_queue_reqs
+            self.num_prefill_inflight_queue_reqs,
+            stats.num_prefill_inflight_queue_reqs,
+            stats.num_prefill_inflight_queue_reqs_by_priority,
         )
         self._log_gauge(
-            self.num_decode_prealloc_queue_reqs, stats.num_decode_prealloc_queue_reqs
+            self.num_decode_prealloc_queue_reqs,
+            stats.num_decode_prealloc_queue_reqs,
+            stats.num_decode_prealloc_queue_reqs_by_priority,
         )
         self._log_gauge(
-            self.num_decode_transfer_queue_reqs, stats.num_decode_transfer_queue_reqs
+            self.num_decode_transfer_queue_reqs,
+            stats.num_decode_transfer_queue_reqs,
+            stats.num_decode_transfer_queue_reqs_by_priority,
         )
         self._log_gauge(self.kv_transfer_speed_gb_s, stats.kv_transfer_speed_gb_s)
         self._log_gauge(self.kv_transfer_latency_ms, stats.kv_transfer_latency_ms)