PR feedback

tests/v1/engine/test_async_llm.py

@@ -14,9 +14,8 @@
 from vllm.platforms import current_platform
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.engine.async_llm import AsyncLLM
-from vllm.v1.metrics.loggers import (PROMETHEUS_LOGGING_LOGGER_NAME,
-                                     STANDARD_LOGGING_LOGGER_NAME,
-                                     LoggingStatLogger, StatLoggerBase)
+from vllm.v1.metrics.loggers import (BuiltinLoggerName, LoggingStatLogger,
+                                     StatLoggerBase)
 
 if not current_platform.is_cuda():
     pytest.skip(reason="V1 currently only supported on CUDA.",
@@ -263,8 +262,8 @@ async def test_customize_loggers(
         after.callback(engine.shutdown)
 
         if loggers is None:
-            engine.remove_logger(PROMETHEUS_LOGGING_LOGGER_NAME)
-            engine.remove_logger(STANDARD_LOGGING_LOGGER_NAME)
+            engine.remove_logger(BuiltinLoggerName.PROMETHEUS)
+            engine.remove_logger(BuiltinLoggerName.LOGGING)
             engine.add_logger(TEST_LOGGER_NAME, get_customized_logger_mock())
 
         await engine.do_log_stats()

vllm/third_party/pynvml.py

@@ -30,16 +30,16 @@
 # THE POSSIBILITY OF SUCH DAMAGE.
 #####
 
-import os
-import string
-import sys
-import threading
 ##
 # Python bindings for the NVML library
 ##
 from ctypes import *
 from ctypes.util import find_library
 from functools import wraps
+import sys
+import os
+import threading
+import string
 
 ## C Type mappings ##
 ## Enums

vllm/v1/engine/async_llm.py

@@ -56,14 +56,9 @@ def __init__(
 
         self.log_requests = log_requests
         self.log_stats = log_stats
-        self.stat_loggers: dict[str, StatLoggerBase] = dict()
-        if self.log_stats:
-            if stat_loggers is not None:
-                self.stat_loggers = stat_loggers
-            else:
-                setup_default_loggers(vllm_config,
-                                      logger.isEnabledFor(logging.INFO),
-                                      self.stat_loggers)
+        self.stat_loggers = setup_default_loggers(
+            vllm_config, self.log_stats, logger.isEnabledFor(logging.INFO),
+            stat_loggers)
 
         # Tokenizer (+ ensure liveness if running in another process).
         self.tokenizer = init_tokenizer_from_configs(

vllm/v1/metrics/loggers.py

@@ -2,6 +2,7 @@
 
 import time
 from abc import ABC, abstractmethod
+from enum import Enum
 from typing import Optional
 
 import numpy as np
@@ -17,11 +18,25 @@
 
 _LOCAL_LOGGING_INTERVAL_SEC = 5.0
 
-STANDARD_LOGGING_LOGGER_NAME = "logging"
-PROMETHEUS_LOGGING_LOGGER_NAME = "prometheus"
+
+class BuiltinLoggerName(str, Enum):
+    """Internal names for built-in metric loggers
+
+    These values are intended for internal use only. Accessing to
+    these built-in loggers is not considered to be part of the
+    public API.
+    """
+    LOGGING = "logging"
+    PROMETHEUS = "prometheus"
 
 
 class StatLoggerBase(ABC):
+    """Interface for logging metrics.
+
+    API users may define custom loggers that implement this interface.
+    However, note that the `SchedulerStats` and `IterationStats` classes
+    are not considered stable interfaces and may change in future versions.
+    """
 
     @abstractmethod
     def record(self, scheduler_stats: SchedulerStats,
@@ -433,11 +448,20 @@ def build_cudagraph_buckets(vllm_config: VllmConfig) -> list[int]:
         return [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8096]
 
 
-def setup_default_loggers(vllm_config: VllmConfig, enable_stats_log: bool,
-                          stat_loggers: dict[str, StatLoggerBase]):
+def setup_default_loggers(
+    vllm_config: VllmConfig, log_stats: bool, enable_logging: bool,
+    custom_stat_loggers: Optional[dict[str, StatLoggerBase]]
+) -> dict[str, StatLoggerBase]:
     """Setup logging and prometheus metrics."""
-    stat_loggers[PROMETHEUS_LOGGING_LOGGER_NAME] = (
-        PrometheusStatLogger(vllm_config))
-
-    if enable_stats_log:
-        stat_loggers[STANDARD_LOGGING_LOGGER_NAME] = (LoggingStatLogger())
+    if not log_stats:
+        return {}
+
+    if custom_stat_loggers is not None:
+        return custom_stat_loggers
+
+    stat_loggers: dict[str, StatLoggerBase] = {}
+    stat_loggers[BuiltinLoggerName.PROMETHEUS] = \
+        PrometheusStatLogger(vllm_config)
+    if enable_logging:
+        stat_loggers[BuiltinLoggerName.LOGGING] = LoggingStatLogger()
+    return stat_loggers