Clean up metrics code (#1972)

python/sglang/srt/server.py

@@ -25,20 +25,16 @@ import json
 import logging
 import multiprocessing as mp
 import os
-import re
-import tempfile
 import threading
 import time
 from http import HTTPStatus
 from typing import AsyncIterator, Dict, List, Optional, Union
 
-import orjson
-from starlette.routing import Mount
-
 # Fix a bug of Python threading
 setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
 
 import aiohttp
+import orjson
 import requests
 import uvicorn
 import uvloop
@@ -77,6 +73,7 @@ from sglang.srt.openai_api.protocol import ModelCard, ModelList
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import (
     add_api_key_middleware,
+    add_prometheus_middleware,
     assert_pkg_version,
     configure_logger,
     delete_directory,
@@ -84,16 +81,13 @@ from sglang.srt.utils import (
     kill_child_process,
     maybe_set_triton_cache_manager,
     prepare_model_and_tokenizer,
+    set_prometheus_multiproc_dir,
     set_ulimit,
 )
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
 
-# Temporary directory for prometheus multiprocess mode
-# Cleaned up automatically when this object is garbage collected
-prometheus_multiproc_dir: tempfile.TemporaryDirectory
-
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
@@ -445,10 +439,6 @@ def launch_server(
     1. The HTTP server and Tokenizer Manager both run in the main process.
     2. Inter-process communication is done through ICP (each process uses a different port) via the ZMQ library.
     """
-
-    if server_args.enable_metrics:
-        _set_prometheus_env()
-
     launch_engine(server_args=server_args)
 
     # Add api key authorization
@@ -487,36 +477,6 @@ def launch_server(
         t.join()
 
 
-def add_prometheus_middleware(app: FastAPI):
-    # Adapted from https://github.com/vllm-project/vllm/blob/v0.6.1/vllm/entrypoints/openai/api_server.py#L216
-    from prometheus_client import CollectorRegistry, make_asgi_app, multiprocess
-
-    registry = CollectorRegistry()
-    multiprocess.MultiProcessCollector(registry)
-    metrics_route = Mount("/metrics", make_asgi_app(registry=registry))
-
-    # Workaround for 307 Redirect for /metrics
-    metrics_route.path_regex = re.compile("^/metrics(?P<path>.*)$")
-    app.routes.append(metrics_route)
-
-
-def _set_prometheus_env():
-    # Set prometheus multiprocess directory
-    # sglang uses prometheus multiprocess mode
-    # we need to set this before importing prometheus_client
-    # https://prometheus.github.io/client_python/multiprocess/
-    global prometheus_multiproc_dir
-    if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
-        logger.debug(f"User set PROMETHEUS_MULTIPROC_DIR detected.")
-        prometheus_multiproc_dir = tempfile.TemporaryDirectory(
-            dir=os.environ["PROMETHEUS_MULTIPROC_DIR"]
-        )
-    else:
-        prometheus_multiproc_dir = tempfile.TemporaryDirectory()
-        os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name
-    logger.debug(f"PROMETHEUS_MULTIPROC_DIR: {os.environ['PROMETHEUS_MULTIPROC_DIR']}")
-
-
 def _set_envs_and_config(server_args: ServerArgs):
     # Set global environments
     os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -543,6 +503,10 @@ def _set_envs_and_config(server_args: ServerArgs):
             "at https://docs.flashinfer.ai/installation.html.",
         )
 
+    # Set prometheus env vars
+    if server_args.enable_metrics:
+        set_prometheus_multiproc_dir()
+
     mp.set_start_method("spawn", force=True)