[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__)
 
@@ -95,14 +94,6 @@ def __init__(
         self.log_requests = log_requests
         self.log_stats = log_stats
 
-        # Set up stat loggers; independent set for each DP rank.
-        self.stat_loggers: list[list[StatLoggerBase]] = setup_default_loggers(
-            vllm_config=vllm_config,
-            log_stats=self.log_stats,
-            engine_num=vllm_config.parallel_config.data_parallel_size,
-            custom_stat_loggers=stat_loggers,
-        )
-
         # Tokenizer (+ ensure liveness if running in another process).
         self.tokenizer = init_tokenizer_from_configs(
             model_config=vllm_config.model_config,
@@ -121,17 +112,24 @@ def __init__(
                                                 log_stats=self.log_stats)
 
         # EngineCore (starts the engine in background process).
-
         self.engine_core = EngineCoreClient.make_async_mp_client(
             vllm_config=vllm_config,
             executor_class=executor_class,
             log_stats=self.log_stats,
             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()
+
+        # Loggers.
+        self.logger_manager: Optional[StatLoggerManager] = None
+        if self.log_stats:
+            self.logger_manager = StatLoggerManager(
+                vllm_config=vllm_config,
+                engine_idxs=self.engine_core.engine_ranks,
+                custom_stat_loggers=stat_loggers,
+            )
+            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 +368,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 +408,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 +430,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 +534,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 +644,16 @@ 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(
+        if new_data_parallel_size > old_data_parallel_size and self.log_stats:
+            # TODO(rob): fix this after talking with Ray team.
+            # This resets all the prometheus metrics since we
+            # unregister during initialization. Need to understand
+            # the intended behavior here better.
+            self.logger_manager = StatLoggerManager(
                 vllm_config=self.vllm_config,
-                log_stats=self.log_stats,
-                engine_num=new_data_parallel_size,
+                engine_idxs=list(range(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:

vllm/v1/engine/core_client.py

@@ -432,14 +432,15 @@ def __init__(
             external_dp_lb = parallel_config.data_parallel_external_lb
 
             offline_mode = parallel_config.data_parallel_rank_local is not None
-            engine_ranks = [dp_rank] if (offline_mode
-                                         or external_dp_lb) else range(dp_size)
+            self.engine_ranks = ([dp_rank] if
+                                 (offline_mode or external_dp_lb) else list(
+                                     range(dp_size)))
             assert parallel_config.data_parallel_size_local <= len(
-                engine_ranks)
+                self.engine_ranks)
 
             # ZMQ identity of each engine that this client will talk to.
             self.core_engines: list[EngineIdentity] = [
-                index.to_bytes(2, "little") for index in engine_ranks
+                index.to_bytes(2, "little") for index in self.engine_ranks
             ]
 
             # Wait for ready messages from each engine on the input socket.