♻️ refactor for using StatLoggerBase

Signed-off-by: Prashant Gupta <[email protected]>

src/vllm_tgis_adapter/grpc/grpc_server.py

@@ -194,7 +194,7 @@ async def post_init(self) -> None:
         assert hasattr(self.engine.engine, "stat_logger")
         assert self.engine.engine.stat_logger
 
-        vllm_stat_logger = self.engine.engine.stat_logger
+        vllm_stat_logger = self.engine.engine.stat_loggers["prometheus"]
         tgis_stats_logger = TGISStatLogger(
             vllm_stat_logger=vllm_stat_logger,
             max_sequence_len=self.config.max_model_len,

src/vllm_tgis_adapter/tgis_utils/metrics.py

@@ -5,7 +5,7 @@
 
 from prometheus_client import Counter, Gauge, Histogram
 from vllm import RequestOutput
-from vllm.engine.metrics import StatLogger, Stats
+from vllm.engine.metrics import StatLoggerBase, Stats
 
 from vllm_tgis_adapter.grpc.pb.generation_pb2 import (
     BatchedTokenizeRequest,
@@ -102,10 +102,10 @@ def observe_generation_success(self, start_time: float) -> None:
         self.tgi_request_duration.observe(duration)
 
 
-class TGISStatLogger(StatLogger):
-    """Wraps the vLLM StatLogger to report TGIS metric names for compatibility."""
+class TGISStatLogger(StatLoggerBase):
+    """Wraps the vLLM StatLoggerBase to report TGIS metric names for compatibility."""
 
-    def __init__(self, vllm_stat_logger: StatLogger, max_sequence_len: int):
+    def __init__(self, vllm_stat_logger: StatLoggerBase, max_sequence_len: int):
         # Not calling super-init because we're wrapping and delegating to
         # vllm_stat_logger
         self._vllm_stat_logger = vllm_stat_logger