[DP] Fix Prometheus Logging

vllm/engine/arg_utils.py

@@ -1358,10 +1358,10 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool:
                 and not envs.is_set("VLLM_ATTENTION_BACKEND")
             ) or envs.VLLM_ATTENTION_BACKEND == "FLASH_ATTN_VLLM_V1"
             supported = False
-            if current_platform.is_rocm() or (
-                    current_platform.is_cuda()
-                    and current_platform.is_device_capability(100)
-            ):  # handle hpu also for OOT platform
+            if (current_platform.is_rocm()
+                    or (current_platform.is_cuda()
+                        and current_platform.is_device_capability(100))
+                    or current_platform.is_tpu()):
                 supported = True
             elif fp8_attention and will_use_fa:
                 from vllm.attention.utils.fa_utils import (

vllm/v1/engine/async_llm.py

@@ -36,10 +36,9 @@
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.engine.processor import Processor
 from vllm.v1.executor.abstract import Executor
-from vllm.v1.metrics.loggers import (StatLoggerBase, StatLoggerFactory,
-                                     setup_default_loggers)
+from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
 from vllm.v1.metrics.prometheus import shutdown_prometheus
-from vllm.v1.metrics.stats import IterationStats, SchedulerStats
+from vllm.v1.metrics.stats import IterationStats
 
 logger = init_logger(__name__)
 
@@ -96,12 +95,18 @@ def __init__(
         self.log_stats = log_stats
 
         # Set up stat loggers; independent set for each DP rank.
-        self.stat_loggers: list[list[StatLoggerBase]] = setup_default_loggers(
+        # HACK: asyncllm should not be aware of how many engines is it
+        # managing.
+        start_idx = vllm_config.parallel_config.data_parallel_rank
+        local_engines = vllm_config.parallel_config.data_parallel_size_local
+        engine_idxs = [
+            idx for idx in range(start_idx, start_idx + local_engines)
+        ]
+        self.logger_manager = StatLoggerManager(
             vllm_config=vllm_config,
-            log_stats=self.log_stats,
-            engine_num=vllm_config.parallel_config.data_parallel_size,
+            engine_idxs=engine_idxs,
             custom_stat_loggers=stat_loggers,
-        )
+        ) if self.log_stats else None
 
         # Tokenizer (+ ensure liveness if running in another process).
         self.tokenizer = init_tokenizer_from_configs(
@@ -129,9 +134,8 @@ def __init__(
             client_addresses=client_addresses,
             client_index=client_index,
         )
-        if self.stat_loggers:
-            for stat_logger in self.stat_loggers[0]:
-                stat_logger.log_engine_initialized()
+        if self.logger_manager:
+            self.logger_manager.log_engine_initialized()
         self.output_handler: Optional[asyncio.Task] = None
         try:
             # Start output handler eagerly if we are in the asyncio eventloop.
@@ -370,7 +374,7 @@ def _run_output_handler(self):
         engine_core = self.engine_core
         output_processor = self.output_processor
         log_stats = self.log_stats
-        stat_loggers = self.stat_loggers if log_stats else None
+        logger_manager = self.logger_manager
 
         async def output_handler():
             try:
@@ -410,11 +414,12 @@ async def output_handler():
                     # 4) Logging.
                     # TODO(rob): make into a coroutine and launch it in
                     # background thread once Prometheus overhead is non-trivial.
-                    if stat_loggers:
-                        AsyncLLM._record_stats(
-                            stat_loggers[outputs.engine_index],
+                    # NOTE: we do not use self.log
+                    if logger_manager:
+                        logger_manager.record(
                             scheduler_stats=outputs.scheduler_stats,
                             iteration_stats=iteration_stats,
+                            engine_idx=outputs.engine_index,
                         )
             except Exception as e:
                 logger.exception("AsyncLLM output_handler failed.")
@@ -431,18 +436,6 @@ async def abort(self, request_id: str) -> None:
         if self.log_requests:
             logger.info("Aborted request %s.", request_id)
 
-    @staticmethod
-    def _record_stats(
-        stat_loggers: list[StatLoggerBase],
-        scheduler_stats: Optional[SchedulerStats],
-        iteration_stats: Optional[IterationStats],
-    ):
-        """static so that it can be used from the output_handler task
-        without a circular ref to AsyncLLM."""
-        for stat_logger in stat_loggers:
-            stat_logger.record(scheduler_stats=scheduler_stats,
-                               iteration_stats=iteration_stats)
-
     async def encode(
         self,
         prompt: PromptType,
@@ -547,7 +540,11 @@ async def do_log_stats(
         scheduler_outputs=None,
         model_output=None,
     ) -> None:
-        for loggers in self.stat_loggers:
+        if self.stat_loggers is None:
+            return
+        # loggers, prom_logger
+        per_engine_loggers, _ = self.stat_loggers
+        for loggers in per_engine_loggers.values():
             for stat_logger in loggers:
                 stat_logger.log()
 
@@ -653,18 +650,18 @@ async def scale_elastic_ep(self,
             new_data_parallel_size
 
         # recreate stat loggers
-        if new_data_parallel_size > old_data_parallel_size:
-            stat_loggers: list[list[StatLoggerBase]] = setup_default_loggers(
-                vllm_config=self.vllm_config,
-                log_stats=self.log_stats,
-                engine_num=new_data_parallel_size,
-                custom_stat_loggers=None,
-            )
-            num_new_engines = len(stat_loggers) - len(self.stat_loggers)
-            self.stat_loggers.extend(stat_loggers[-num_new_engines:])
-        else:
-            for _ in range(old_data_parallel_size - new_data_parallel_size):
-                self.stat_loggers.pop()
+        # if new_data_parallel_size > old_data_parallel_size:
+        #     stat_loggers: list[list[StatLoggerBase]] = setup_default_loggers(
+        #         vllm_config=self.vllm_config,
+        #         log_stats=self.log_stats,
+        #         engine_num=new_data_parallel_size,
+        #         custom_stat_loggers=None,
+        #     )
+        #     num_new_engines = len(stat_loggers) - len(self.stat_loggers)
+        #     self.stat_loggers.extend(stat_loggers[-num_new_engines:])
+        # else:
+        #     for _ in range(old_data_parallel_size - new_data_parallel_size):
+        #         self.stat_loggers.pop()
 
     @property
     def is_running(self) -> bool: