diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 2a6a0566e8e3..29feca9135bf 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/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() diff --git a/vllm/third_party/pynvml.py b/vllm/third_party/pynvml.py index 2b6031ba9dc0..0a4be23a0936 100644 --- a/vllm/third_party/pynvml.py +++ b/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 diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 8b722f160471..734df74a5517 100644 --- a/vllm/v1/engine/async_llm.py +++ b/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( diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 2462ebbd759f..5bca7e6af5cd 100644 --- a/vllm/v1/metrics/loggers.py +++ b/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