add

Author: qimcis

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -25,6 +25,10 @@
 from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils import (
+    add_prometheus_middleware,
+    add_prometheus_track_response_middleware,
+)
 
 if TYPE_CHECKING:
     from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
@@ -207,6 +211,10 @@ def create_app(server_args: ServerArgs):
     """
     app = FastAPI(lifespan=lifespan)
 
+    if server_args.enable_metrics:
+        add_prometheus_middleware(app)
+        add_prometheus_track_response_middleware(app)
+
     app.include_router(health_router)
     app.include_router(vertex_router)
 

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -17,6 +17,7 @@
     set_global_server_args,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
+from sglang.srt.utils import set_prometheus_multiproc_dir
 
 
 def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
@@ -68,6 +69,9 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True):
     # Start a new server with multiple worker processes
     logger.info("Starting server...")
 
+    if server_args.enable_metrics:
+        set_prometheus_multiproc_dir()
+
     num_gpus = server_args.num_gpus
     processes = []
 

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -43,7 +43,10 @@
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
-from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch
+from sglang.multimodal_gen.runtime.utils.common import (
+    get_diffusion_metrics_collector,
+    set_cuda_arch,
+)
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
     OffloadableDiTMixin,
     iter_materialized_weights,
@@ -79,6 +82,11 @@ def __init__(
         # FIXME: should we use tcp as distribute init method?
         self.server_args = server_args
         self.pipeline: ComposedPipelineBase = None
+        self.metrics_collector = (
+            get_diffusion_metrics_collector(server_args)
+            if server_args.enable_metrics and rank == 0
+            else None
+        )
 
         self.init_device_and_model()
         self.sp_group = get_sp_group()
@@ -89,6 +97,8 @@ def __init__(
         self.cfg_group = get_cfg_group()
         self.cfg_cpu_group = self.cfg_group.cpu_group
 
+        self._update_lora_metrics()
+
     def init_device_and_model(self) -> None:
         """Initialize the device and load the model."""
         torch.get_device_module().set_device(self.local_rank)
@@ -199,7 +209,7 @@ def do_mem_analysis(self, output_batch: OutputBatch):
         logger.info(
             f"Peak GPU memory: {peak_reserved_gb:.2f} GB, "
             f"Peak allocated: {peak_allocated_gb:.2f} GB, "
-            f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb/peak_reserved_gb*100:.1f}%), "
+            f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb / peak_reserved_gb * 100:.1f}%), "
             f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
             f"Components that could stay resident (based on the last request workload): {can_stay_resident}. "
             f"Related offload server args to disable: {suggested_args_str}"
@@ -212,12 +222,12 @@ def execute_forward(self, batch: List[Req]) -> OutputBatch:
         assert self.pipeline is not None
         req = batch[0]
         output_batch = None
+        status = "success"
+        start_time = time.monotonic()
         try:
             if self.rank == 0:
                 torch.get_device_module().reset_peak_memory_stats()
 
-            start_time = time.monotonic()
-
             # capture memory baseline before forward
             if self.rank == 0 and req.metrics:
                 baseline_snapshot = capture_memory_snapshot()
@@ -274,7 +284,10 @@ def execute_forward(self, batch: List[Req]) -> OutputBatch:
                 # Avoid logging warmup perf records that share the same request_id.
                 if not req.is_warmup:
                     PerformanceLogger.log_request_summary(metrics=output_batch.metrics)
+            if output_batch is not None and output_batch.error is not None:
+                status = "error"
         except Exception as e:
+            status = "error"
             logger.error(
                 f"Error executing request {req.request_id}: {e}", exc_info=True
             )
@@ -283,8 +296,25 @@ def execute_forward(self, batch: List[Req]) -> OutputBatch:
             if output_batch is None:
                 output_batch = OutputBatch()
             output_batch.error = f"Error executing request {req.request_id}: {e}"
+        finally:
+            if self.metrics_collector is not None:
+                self.metrics_collector.observe_request(
+                    status=status,
+                    is_warmup=req.is_warmup,
+                    latency_s=time.monotonic() - start_time,
+                )
         return output_batch
 
+    def _update_lora_metrics(self):
+        if self.metrics_collector is None:
+            return
+
+        if not isinstance(self.pipeline, LoRAPipeline):
+            self.metrics_collector.clear_lora_status()
+            return
+
+        self.metrics_collector.update_lora_status(self.pipeline.get_lora_status())
+
     def get_can_stay_resident_components(
         self, remaining_gpu_mem_gb: float
     ) -> List[str]:
@@ -339,6 +369,7 @@ def set_lora(
         if not isinstance(self.pipeline, LoRAPipeline):
             return OutputBatch(error="Lora is not enabled")
         self.pipeline.set_lora(lora_nickname, lora_path, target, strength)
+        self._update_lora_metrics()
         return OutputBatch()
 
     def merge_lora_weights(
@@ -354,6 +385,7 @@ def merge_lora_weights(
         if not isinstance(self.pipeline, LoRAPipeline):
             return OutputBatch(error="Lora is not enabled")
         self.pipeline.merge_lora_weights(target, strength)
+        self._update_lora_metrics()
         return OutputBatch()
 
     def unmerge_lora_weights(self, target: str = "all") -> OutputBatch:
@@ -366,6 +398,7 @@ def unmerge_lora_weights(self, target: str = "all") -> OutputBatch:
         if not isinstance(self.pipeline, LoRAPipeline):
             return OutputBatch(error="Lora is not enabled")
         self.pipeline.unmerge_lora_weights(target)
+        self._update_lora_metrics()
         return OutputBatch()
 
     def list_loras(self) -> OutputBatch:

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -4,6 +4,7 @@
 import asyncio
 import os
 import pickle
+import time
 from collections import deque
 from copy import deepcopy
 from typing import Any, List
@@ -34,7 +35,10 @@
     ServerArgs,
     set_global_server_args,
 )
-from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket
+from sglang.multimodal_gen.runtime.utils.common import (
+    get_diffusion_metrics_collector,
+    get_zmq_socket,
+)
 from sglang.multimodal_gen.runtime.utils.distributed import broadcast_pyobj
 from sglang.multimodal_gen.runtime.utils.logging_utils import GREEN, RESET, init_logger
 
@@ -101,6 +105,8 @@ def __init__(
 
         # FIFO, new reqs are appended
         self.waiting_queue: deque[tuple[bytes, Req]] = deque()
+        self._generation_waiting_count = 0
+        self._generation_enqueue_timestamps: dict[int, float] = {}
 
         # whether we've send the necessary warmup reqs
         self.warmed_up = False
@@ -110,6 +116,15 @@ def __init__(
 
         self.prepare_server_warmup_reqs()
 
+        self.metrics_collector = (
+            get_diffusion_metrics_collector(server_args)
+            if server_args.enable_metrics and gpu_id == 0
+            else None
+        )
+        if self.metrics_collector is not None:
+            self.metrics_collector.set_queue_depth(self._generation_waiting_count)
+            self.metrics_collector.set_running_reqs(0)
+
         # Maximum consecutive errors before terminating the event loop
         self._max_consecutive_errors = 3
         self._consecutive_error_count = 0
@@ -187,9 +202,32 @@ def get_next_batch_to_run(self) -> list[tuple[bytes, Req]] | None:
 
         # pop the first (earliest)
         item = self.waiting_queue.popleft()
+        self._on_req_dequeued(item[1])
+        if self.metrics_collector is not None:
+            self.metrics_collector.set_queue_depth(self._generation_waiting_count)
 
         return [item]
 
+    def _on_req_enqueued(self, req: Any) -> None:
+        if not isinstance(req, Req):
+            return
+        self._generation_waiting_count += 1
+        self._generation_enqueue_timestamps[id(req)] = time.monotonic()
+
+    def _on_req_dequeued(self, req: Any) -> None:
+        if not isinstance(req, Req):
+            return
+        if self._generation_waiting_count > 0:
+            self._generation_waiting_count -= 1
+        enqueue_ts = self._generation_enqueue_timestamps.pop(id(req), None)
+        if enqueue_ts is not None and self.metrics_collector is not None:
+            self.metrics_collector.observe_queue_time(time.monotonic() - enqueue_ts)
+
+    def _enqueue_received_reqs(self, new_reqs: list[tuple[bytes, Any]]) -> None:
+        self.waiting_queue.extend(new_reqs)
+        for _, req in new_reqs:
+            self._on_req_enqueued(req)
+
     def prepare_server_warmup_reqs(self):
         if (
             self.server_args.warmup
@@ -235,6 +273,7 @@ def prepare_server_warmup_reqs(self):
                     )
                 req.set_as_warmup()
                 self.waiting_queue.append((None, req))
+                self._on_req_enqueued(req)
             # if server is warmed-up, set this flag to avoid req-based warmup
             self.warmed_up = True
 
@@ -334,7 +373,11 @@ def event_loop(self) -> None:
             try:
                 new_reqs = self.recv_reqs()
                 new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs)
-                self.waiting_queue.extend(new_reqs)
+                self._enqueue_received_reqs(new_reqs)
+                if self.metrics_collector is not None:
+                    self.metrics_collector.set_queue_depth(
+                        self._generation_waiting_count
+                    )
                 # Reset error count on success
                 self._consecutive_error_count = 0
             except Exception as e:
@@ -362,60 +405,72 @@ def event_loop(self) -> None:
 
             identities = [item[0] for item in items]
             reqs = [item[1] for item in items]
+            generation_running_reqs = sum(1 for req in reqs if isinstance(req, Req))
+            if self.metrics_collector is not None:
+                self.metrics_collector.set_running_reqs(generation_running_reqs)
 
             try:
-                processed_req = reqs[0]
-                handler = self.request_handlers.get(type(processed_req))
-                if handler:
-                    output_batch = handler(reqs)
-                else:
-                    output_batch = OutputBatch(
-                        error=f"Unknown request type: {type(processed_req)}"
+                try:
+                    processed_req = reqs[0]
+                    handler = self.request_handlers.get(type(processed_req))
+                    if handler:
+                        output_batch = handler(reqs)
+                    else:
+                        output_batch = OutputBatch(
+                            error=f"Unknown request type: {type(processed_req)}"
+                        )
+                except Exception as e:
+                    logger.error(
+                        f"Error executing request in scheduler event loop: {e}",
+                        exc_info=True,
+                    )
+                    # Determine appropriate error response format
+                    output_batch = (
+                        OutputBatch(error=str(e))
+                        if reqs and isinstance(reqs[0], Req)
+                        else OutputBatch(error=str(e))
                     )
-            except Exception as e:
-                logger.error(
-                    f"Error executing request in scheduler event loop: {e}",
-                    exc_info=True,
-                )
-                # Determine appropriate error response format
-                output_batch = (
-                    OutputBatch(error=str(e))
-                    if reqs and isinstance(reqs[0], Req)
-                    else OutputBatch(error=str(e))
-                )
 
-            # 3. return results
-            try:
-                # log warmup info
-                is_warmup = (
-                    processed_req.is_warmup if isinstance(processed_req, Req) else False
-                )
-                if is_warmup:
-                    if output_batch.error is None:
-                        if self._warmup_total > 0:
-                            logger.info(
-                                f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processed in {GREEN}%.2f{RESET} seconds",
-                                output_batch.metrics.total_duration_s,
-                            )
-                        else:
-                            logger.info(
-                                f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
-                                output_batch.metrics.total_duration_s,
-                            )
-                    else:
-                        if self._warmup_total > 0:
-                            logger.info(
-                                f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processing failed"
-                            )
+                # 3. return results
+                try:
+                    # log warmup info
+                    is_warmup = (
+                        processed_req.is_warmup
+                        if isinstance(processed_req, Req)
+                        else False
+                    )
+                    if is_warmup:
+                        if output_batch.error is None:
+                            if self._warmup_total > 0:
+                                logger.info(
+                                    f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processed in {GREEN}%.2f{RESET} seconds",
+                                    output_batch.metrics.total_duration_s,
+                                )
+                            else:
+                                logger.info(
+                                    f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
+                                    output_batch.metrics.total_duration_s,
+                                )
                         else:
-                            logger.info(f"Warmup req processing failed")
-
-                # TODO: Support sending back to multiple identities if batched
-                self.return_result(output_batch, identities[0], is_warmup=is_warmup)
-            except zmq.ZMQError as e:
-                # Reply failed; log and keep loop alive to accept future requests
-                logger.error(f"ZMQ error sending reply: {e}")
-                continue
+                            if self._warmup_total > 0:
+                                logger.info(
+                                    f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processing failed"
+                                )
+                            else:
+                                logger.info(f"Warmup req processing failed")
+
+                    # TODO: Support sending back to multiple identities if batched
+                    self.return_result(output_batch, identities[0], is_warmup=is_warmup)
+                except zmq.ZMQError as e:
+                    # Reply failed; log and keep loop alive to accept future requests
+                    logger.error(f"ZMQ error sending reply: {e}")
+                    continue
+            finally:
+                if self.metrics_collector is not None:
+                    self.metrics_collector.set_running_reqs(0)
+                    self.metrics_collector.set_queue_depth(
+                        self._generation_waiting_count
+                    )
 
         if self.receiver is not None:
             self.receiver.close()

python/sglang/multimodal_gen/runtime/server_args.py

@@ -327,6 +327,7 @@ class ServerArgs:
     # http server endpoint config
     host: str | None = "127.0.0.1"
     port: int | None = 30000
+    enable_metrics: bool = False
 
     # TODO: webui and their endpoint, check if webui_port is available.
     webui: bool = False
@@ -858,6 +859,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=ServerArgs.port,
             help="Port for the HTTP API server.",
         )
+        parser.add_argument(
+            "--enable-metrics",
+            action=StoreBoolean,
+            default=ServerArgs.enable_metrics,
+            help="Enable Prometheus metrics endpoint at /metrics.",
+        )
         parser.add_argument(
             "--webui",
             action=StoreBoolean,