[Feat] UCM supports metrics display online via Grafana and Promethues

Author: sumingZero

CMakeLists.txt

@@ -8,7 +8,6 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 option(BUILD_UCM_STORE "build ucm store module." ON)
 option(BUILD_UCM_SPARSE "build ucm sparse module." ON)
-option(BUILD_UCM_METRICS "build ucm metrics module." ON)
 option(BUILD_UNIT_TESTS "build all unit test suits." OFF)
 option(BUILD_NUMA "build numactl library." OFF)
 option(DOWNLOAD_DEPENDENCE "download dependence by cmake." ON)

ucm/CMakeLists.txt

@@ -5,6 +5,3 @@ endif()
 if(BUILD_UCM_SPARSE)
     add_subdirectory(sparse)
 endif()
-if(BUILD_UCM_METRICS)
-    add_subdirectory(metrics)
-endif()

ucm/integration/vllm/ucm_connector.py

@@ -1,3 +1,4 @@
+import ctypes
 import hashlib
 import itertools
 import os
@@ -19,8 +20,8 @@
 from vllm.v1.request import Request
 
 from ucm.logger import init_logger
-from ucm.metrics.ucm_obser import UCMStatsLogger
-from ucm.metrics.ucmmonitor import UCMStatsMonitor
+from ucm.shared.metrics import monitor
+from ucm.shared.metrics.observability import UCMStatsLogger
 from ucm.store.factory import UcmConnectorFactory
 from ucm.store.ucmstore import Task, UcmKVStoreBase
 from ucm.utils import Config
@@ -129,10 +130,9 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
                 self.broadcast_fn = self.group_coordinator.broadcast
                 self.broadcast_stream = torch.cuda.Stream()
 
+        logger.info(f"self.launch_config: {self.launch_config}")
         connector_configs = self.launch_config.get("ucm_connectors", [])
         assert len(connector_configs) > 0, "no storage connector name in config."
-            self.io_size = config["io_size"]
-            self.num_layers = num_layers
 
         name = connector_configs[0].get("ucm_connector_name")
         config = connector_configs[0].get("ucm_connector_config") or {}
@@ -157,6 +157,7 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
             1 if self.is_mla else num_head_per_tp
         )
         self.store = UcmConnectorFactory.create_connector(name, config)
+        self.block_data_size = config["kv_block_size"]
 
         logger.info("init UCConnectorImpl, connector: %s", name)
         logger.info(
@@ -165,6 +166,16 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
             config["io_size"] / 1024,
         )
 
+        self.stats_logger = UCMStatsLogger(
+            vllm_config.model_config.served_model_name, self.rank
+        )
+        self.monitor = monitor.StatsMonitor.get_instance()
+        self.synchronize = (
+            torch.cuda.synchronize
+            if current_platform.is_cuda_alike()
+            else torch.npu.synchronize
+        )
+
     def generate_hash(self, block_size: int, request: "Request") -> list[str]:
         token_ids = request.all_token_ids
 
@@ -213,6 +224,10 @@ def get_num_new_matched_tokens(
             f"hit hbm: {hbm_hit_block_num}, "
             f"hit external: {external_hit_blocks}"
         )
+        self.monitor.update_stats(
+            "ConnStats",
+            {"interval_lookup_hit_rates": external_hit_blocks / len(ucm_block_ids)},
+        )
 
         total_hit_block_num = hbm_hit_block_num + external_hit_blocks
 
@@ -456,12 +471,14 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         req_broadcast_addr = {}
         is_load = False
         num_loaded_block = 0
+        num_loaded_request = 0
         load_start_time = time.perf_counter() * 1000
         for request_id, request in metadata.request_meta.items():
             if len(request.load_block_ids[0]) == 0:
                 continue
             is_load = True
             num_loaded_block += len(request.load_block_ids[0])
+            num_loaded_request += 1
 
             ucm_block_ids, vllm_block_ids = request.load_block_ids
             if self.rank != 0 and not self.is_mla:
@@ -486,17 +503,21 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
             if self.load_only_first_rank:
                 self._broadcast(req_broadcast_addr[request_id])
         load_end_time = time.perf_counter() * 1000
+        load_speed = (
+            num_loaded_block
+            * self.block_data_size
+            / (load_end_time - load_start_time)
+            / 1024
+            / 1024
+        )  # GB/s
         if is_load:
-            UCMStatsMonitor.get_instance().update_stats(
-                "UCMStats",
+            self.monitor.update_stats(
+                "ConnStats",
                 {
+                    "load_requests_num": num_loaded_request,
+                    "load_blocks_num": num_loaded_block,
                     "load_duration": load_end_time - load_start_time,
-                    "load_speed": num_loaded_block
-                    * self.io_size
-                    * self.num_layers
-                    / (load_end_time - load_start_time)
-                    / 1024
-                    / 1024,  # GB/s
+                    "load_speed": load_speed,
                 },
             )
 
@@ -515,20 +536,24 @@ def save_kv_layer(
     def wait_for_save(self) -> None:
         if self.is_mla and self.rank != 0:
             return
+        if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
+            self.synchronize()
 
         metadata = self._get_connector_metadata()
         assert isinstance(metadata, UCMConnectorMetadata)
 
         request_to_task: dict[str, Task] = {}
         request_to_blocks: dict[str, list[str]] = {}
-        is_dump = False
-        dump_start_time = time.perf_counter() * 1000
-        num_dumped_block = 0
+        is_save = False
+        num_saved_block = 0
+        num_saved_request = 0
+        save_start_time = time.perf_counter() * 1000
         for request_id, request in metadata.request_meta.items():
             if len(request.dump_block_ids[0]) == 0:
                 continue
-            is_dump = True
-            num_dumped_block += len(request.dump_block_ids[0])
+            is_save = True
+            num_saved_block += len(request.dump_block_ids[0])
+            num_saved_request += 1
 
             ucm_block_ids, vllm_block_ids = request.dump_block_ids
             if self.rank != 0:
@@ -563,18 +588,22 @@ def wait_for_save(self) -> None:
             else:
                 logger.error(f"request {request_id} dump kv cache failed.")
                 self.store.commit(ucm_block_ids, False)
-        dump_end_time = time.perf_counter() * 1000
-        if is_dump:
-            UCMStatsMonitor.get_instance().update_stats(
-                "UCMStats",
+        save_end_time = time.perf_counter() * 1000
+        save_speed = (
+            num_saved_block
+            * self.block_data_size
+            / (save_end_time - save_start_time)
+            / 1024
+            / 1024
+        )  # GB/s
+        if is_save:
+            self.monitor.update_stats(
+                "ConnStats",
                 {
-                    "save_duration": dump_end_time - dump_start_time,
-                    "save_speed": num_dumped_block
-                    * self.io_size
-                    * self.num_layers
-                    / (dump_end_time - dump_start_time)
-                    / 1024
-                    / 1024,  # GB/s
+                    "save_requests_num": num_saved_request,
+                    "save_blocks_num": num_saved_block,
+                    "save_duration": save_end_time - save_start_time,
+                    "save_speed": save_speed,
                 },
             )
 
@@ -699,10 +728,6 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         else:
             self.connector = UCMDirectConnector(vllm_config, role)
 
-        if role == KVConnectorRole.WORKER:
-            self.stats_logger = UCMStatsLogger(vllm_config, 10)
-            self.monitor = UCMStatsMonitor.get_instance()
-
     def get_num_new_matched_tokens(
         self,
         request: "Request",