add example

flesher0813 · flesher0813 · commit a2b04379b890 · 2025-12-18T10:29:42.000+08:00
diff --git a/examples/metrics/metrics_configs.yaml b/examples/metrics/metrics_configs.yaml
@@ -19,10 +19,10 @@ metric_prefix: "ucm:"
 
 # Histogram metrics configuration
 histogram:
-  - name: "load_requests_num"
+  - name: "load_requests_total"
     documentation: "Number of requests loaded from ucm"
     buckets: [1, 5, 10, 20, 50, 100, 200, 500, 1000]
-  - name: "load_blocks_num"
+  - name: "load_blocks_total"
     documentation: "Number of blocks loaded from ucm"
     buckets: [0, 50, 100, 150, 200, 250, 300, 350, 400, 550, 600, 750, 800, 850, 900, 950, 1000]
   - name: "load_duration"
@@ -31,10 +31,10 @@ histogram:
   - name: "load_speed"
     documentation: "Speed of loading from ucm (GB/s)"
     buckets: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 50, 60, 70, 80, 90, 100]
-  - name: "save_requests_num"
+  - name: "save_requests_total"
     documentation: "Number of requests saved to ucm"
     buckets: [1, 5, 10, 20, 50, 100, 200, 500, 1000]
-  - name: "save_blocks_num"
+  - name: "save_blocks_total"
     documentation: "Number of blocks saved to ucm"
     buckets: [0, 50, 100, 150, 200, 250, 300, 350, 400, 550, 600, 750, 800, 850, 900, 950, 1000]
   - name: "save_duration"
@@ -46,4 +46,9 @@ histogram:
   - name: "interval_lookup_hit_rates"
     documentation: "Hit rates of ucm lookup requests"
     buckets: [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
-
+  - name: "s2d_bandwidth"
+    documentation: "Band width of uc store task s2d, copy tensors from storage to device"
+    buckets: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+  - name: "d2s_bandwidth"
+    documentation: "Band width of uc store task d2s, copy tensors from device to storage"
+    buckets: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
diff --git a/ucm/integration/vllm/ucm_connector.py b/ucm/integration/vllm/ucm_connector.py
@@ -18,8 +18,8 @@
 from vllm.v1.core.sched.output import SchedulerOutput
 
 from ucm.logger import init_logger
-from ucm.shared.metrics import ucmmonitor
-from ucm.shared.metrics.observability import UCMStatsLogger
+from ucm.observability import PrometheusStatsLogger
+from ucm.shared.metrics import ucmmetrics
 from ucm.store.factory_v1 import UcmConnectorFactoryV1
 from ucm.store.ucmstore_v1 import Task, UcmKVStoreBaseV1
 from ucm.utils import Config
@@ -160,12 +160,11 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
 
         self.metrics_config = self.launch_config.get("metrics_config_path", "")
         if self.metrics_config:
-            self.stats_logger = UCMStatsLogger(
+            self.stats_logger = PrometheusStatsLogger(
                 vllm_config.model_config.served_model_name,
                 self.global_rank,
                 self.metrics_config,
             )
-            self.monitor = ucmmonitor.StatsMonitor.get_instance()
 
         self.synchronize = (
             torch.cuda.synchronize
@@ -295,8 +294,7 @@ def get_num_new_matched_tokens(
             f"hit external: {external_hit_blocks}"
         )
         if self.metrics_config:
-            self.monitor.update_stats(
-                "ConnStats",
+            ucmmetrics.update_stats(
                 {"interval_lookup_hit_rates": external_hit_blocks / len(ucm_block_ids)},
             )
 
@@ -539,14 +537,13 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
             / 1024
         )  # GB/s
         if self.metrics_config and is_load:
-            self.monitor.update_stats(
-                "ConnStats",
+            ucmmetrics.update_stats(
                 {
-                    "load_requests_num": num_loaded_request,
-                    "load_blocks_num": num_loaded_block,
+                    "load_requests_total": num_loaded_request,
+                    "load_blocks_total": num_loaded_block,
                     "load_duration": load_end_time - load_start_time,
                     "load_speed": load_speed,
-                },
+                }
             )
 
     def wait_for_layer_load(self, layer_name: str) -> None:
@@ -615,11 +612,10 @@ def wait_for_save(self) -> None:
             / 1024
         )  # GB/s
         if self.metrics_config and is_save:
-            self.monitor.update_stats(
-                "ConnStats",
+            ucmmetrics.update_stats(
                 {
-                    "save_requests_num": num_saved_request,
-                    "save_blocks_num": num_saved_block,
+                    "save_requests_total": num_saved_request,
+                    "save_blocks_total": num_saved_block,
                     "save_duration": save_end_time - save_start_time,
                     "save_speed": save_speed,
                 },
diff --git a/ucm/observability.py b/ucm/observability.py
@@ -28,16 +28,17 @@
 import time
 from typing import Any, Union
 
-from prometheus_client import Gauge, Counter, Histogram
 import yaml
+from prometheus_client import Counter, Gauge, Histogram
 
 from ucm.logger import init_logger
 from ucm.shared.metrics import ucmmetrics
 
 logger = init_logger(__name__)
 
+
 class PrometheusStatsLogger:
-    
+
     def _load_config(self, config_path: str) -> dict[str, Any]:
         """Load configuration from YAML file"""
         try:
@@ -57,9 +58,14 @@ def _load_config(self, config_path: str) -> dict[str, Any]:
             return {}
 
     def __init__(self, model_name, worker_id, config_path):
-        # Ensure PROMETHEUS_MULTIPROC_DIR is set before any metric registration
+        """
+        Load metrics config from YAML file (config_path),
+        register metrics using prometheus_client, and start a thread to get updated metrics.
+        """
+        # Load metrics config
         self.config = self._load_config(config_path)
         self.log_interval = self.config.get("log_interval", 10)
+
         multiproc_dir = self.config.get("multiproc_dir", "/vllm-workspace")
         if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:
             os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir
@@ -71,60 +77,56 @@ def __init__(self, model_name, worker_id, config_path):
             "worker_id": worker_id,
         }
         self.labelnames = list(self.labels.keys())
-        
-        # Initialize metrics based on configuration
+
         self.metric_type_config = {
-            "counter": (
-                Counter,
-                {}
-            ),
-            "gauge": (
-                Gauge,
-                {"multiprocess_mode": "all"}
-            ),
-            "histogram": (
-                Histogram,
-                {"buckets": []}
-            )
+            "counter": (Counter, {}),
+            "gauge": (Gauge, {"multiprocess_mode": "all"}),
+            "histogram": (Histogram, {"buckets": []}),
         }
+        # Initialize metrics based on config
         self._init_metrics_from_config()
+
+        # Start thread to update metrics
         self.is_running = True
-        self.thread = threading.Thread(target=self.obtain_stats_thread, daemon=True)
+        self.thread = threading.Thread(target=self.update_stats_loop, daemon=True)
         self.thread.start()
-        
-    def _process_metric_group(self, group_name):
-        metric_cls, default_kwargs = self.metric_type_config[group_name]
-        cfg_list = self.config.get(group_name, [])
-        
+
+    def _register_metrics_by_type(self, metric_type):
+        """
+        Register metrics by different metric types.
+        """
+        metric_cls, default_kwargs = self.metric_type_config[metric_type]
+        cfg_list = self.config.get(metric_type, [])
+
         for cfg in cfg_list:
             name = cfg.get("name")
             doc = cfg.get("documentation", "")
             # Prometheus metric name with prefix
             prometheus_name = f"{self.metric_prefix}{name}"
-            ucmmetrics.create_stats(name, group_name)
-            
+            ucmmetrics.create_stats(name, metric_type)
+
             metric_kwargs = {
                 "name": prometheus_name,
                 "documentation": doc,
                 "labelnames": self.labelnames,
                 **default_kwargs,
-                **{k: v for k, v in cfg.items() if k in default_kwargs}
+                **{k: v for k, v in cfg.items() if k in default_kwargs},
             }
-            
+
             self.metric_mappings[name] = metric_cls(**metric_kwargs)
 
     def _init_metrics_from_config(self):
-        """Initialize metrics based on configuration"""
+        """Initialize metrics based on config"""
         # Get metric name prefix from config (e.g., "ucm:")
         self.metric_prefix = self.config.get("metric_prefix", "ucm:")
 
         # Store metric mapping: metric_name -> Union[Counter, Gauge, Histogram]
         # This mapping will be used in update_stats to dynamically log metrics
         self.metric_mappings: dict[str, Union[Counter, Gauge, Histogram]] = {}
 
-        for group_name in self.metric_type_config.keys():
-            self._process_metric_group(group_name)
-    
+        for metric_type in self.metric_type_config.keys():
+            self._register_metrics_by_type(metric_type)
+
     def _update_counter(self, metric, value):
         if value < 0:
             return
@@ -136,13 +138,17 @@ def _update_gauge(self, metric, value):
     def _update_histogram(self, metric, value):
         for data in value:
             metric.observe(data)
-        
+
     def _update_with_func(self, update_func, stats: dict[str, Any], op_desc: str):
+        """
+        Generic update for Prometheus metrics: match metrics by name, bind labels,
+        and update values via the specified function (update_func).
+        """
         for stat_name, value in stats.items():
             if stat_name not in self.metric_mappings:
                 logger.error(f"Metric {stat_name} not found")
                 continue
-            
+
             metric = self.metric_mappings[stat_name]
             try:
                 metric_with_labels = metric.labels(**self.labels)
@@ -153,8 +159,9 @@ def _update_with_func(self, update_func, stats: dict[str, Any], op_desc: str):
                 logger.debug(f"Failed to {op_desc} {stat_name}: {e}")
 
     def update_stats(self, counter_stats, gauge_stats, histogram_stats):
-        """Log metrics to Prometheus based on configuration file"""
-        # Dynamically log metrics based on what's configured in YAML
+        """
+        Update all Prometheus metrics (Counter/Gauge/Histogram) with given stats.
+        """
         update_tasks = [
             (self._update_counter, counter_stats, "increment"),
             (self._update_gauge, gauge_stats, "set"),
@@ -163,9 +170,14 @@ def update_stats(self, counter_stats, gauge_stats, histogram_stats):
         for update_func, stats, op_desc in update_tasks:
             self._update_with_func(update_func, stats, op_desc)
 
-    def obtain_stats_thread(self):
+    def update_stats_loop(self):
+        """
+        Periodically update Prometheus metrics in a loop until stopped.
+        """
         while self.is_running:
-            counter_stats, gauge_stats, histogram_stats = ucmmetrics.get_all_stats_and_clear()
+            counter_stats, gauge_stats, histogram_stats = (
+                ucmmetrics.get_all_stats_and_clear()
+            )
             self.update_stats(counter_stats, gauge_stats, histogram_stats)
             time.sleep(self.log_interval)
 
diff --git a/ucm/shared/metrics/cc/domain/metrics.h b/ucm/shared/metrics/cc/domain/metrics.h
@@ -21,8 +21,8 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  * */
-#ifndef UNIFIEDCACHE_MONITOR_H
-#define UNIFIEDCACHE_MONITOR_H
+#ifndef UNIFIEDCACHE_METRICS_H
+#define UNIFIEDCACHE_METRICS_H
 
 #include <memory>
 #include <mutex>
diff --git a/ucm/shared/test/example/metrics/monitor_stats_example.py b/ucm/shared/test/example/metrics/monitor_stats_example.py
@@ -42,15 +42,17 @@ def wrapper(*args, **kwargs):
 def metrics_with_update_stats():
     ucmmetrics.create_stats("counter_1", "counter")
     ucmmetrics.update_stats("counter_1", 1.2)
-    
+
     ucmmetrics.create_stats("gauge_1", "gauge")
     ucmmetrics.update_stats("gauge_1", 2.2)
-    
+
     ucmmetrics.create_stats("histogram_1", "histogram")
     ucmmetrics.update_stats("histogram_1", 1)
 
     counters, gauges, histograms = ucmmetrics.get_all_stats_and_clear()
-    print(f"After clear then get counters: {counters}, gauges: {gauges}, histograms: {histograms}")
+    print(
+        f"After clear then get counters: {counters}, gauges: {gauges}, histograms: {histograms}"
+    )
     ucmmetrics.update_stats(
         {
             "counter_1": 5,
@@ -59,7 +61,9 @@ def metrics_with_update_stats():
         }
     )
     counters, gauges, histograms = ucmmetrics.get_all_stats_and_clear()
-    print(f"After clear then get counters: {counters}, gauges: {gauges}, histograms: {histograms}")
+    print(
+        f"After clear then get counters: {counters}, gauges: {gauges}, histograms: {histograms}"
+    )
 
 
 if __name__ == "__main__":
diff --git a/ucm/store/detail/task/task_shard.h b/ucm/store/detail/task/task_shard.h
@@ -35,6 +35,7 @@
 #include <vector>
 #include "logger/logger.h"
 #include "task_waiter.h"
+#include "metrics_api.h"
 
 namespace UC {
 
@@ -129,6 +130,17 @@ class Task {
         auto wait = execTp_ - startTp_;
         auto exec = NowTp() - execTp_;
         auto bw = size_ / exec / 1024 / 1024 / 1024;
+        switch (type_)
+        {
+        case Type::DUMP:
+            UC::Metrics::UpdateStats("d2s_bandwidth", bw);
+            break;
+        case Type::LOAD:
+            UC::Metrics::UpdateStats("s2d_bandwidth", bw);
+            break;
+        default:
+            break;
+        }
         return fmt::format("wait={:.06f}s, exec={:.06f}s, bw={:.06f}GB/s", wait, exec, bw);
     }
 
diff --git a/ucm/store/nfsstore/CMakeLists.txt b/ucm/store/nfsstore/CMakeLists.txt
@@ -6,7 +6,7 @@ target_include_directories(nfsstore PUBLIC
     ${CMAKE_CURRENT_SOURCE_DIR}/cc/api
     ${CMAKE_CURRENT_SOURCE_DIR}/cc/domain
 )
-target_link_libraries(nfsstore PUBLIC storeintf storedevice infra_logger)
+target_link_libraries(nfsstore PUBLIC storeintf storedevice infra_logger metrics)
 
 file(GLOB_RECURSE UCMSTORE_NFS_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmnfsstore ${UCMSTORE_NFS_CPY_SOURCE_FILES})

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@ target_include_directories(nfsstore PUBLIC`
`6`	`6`	`${CMAKE_CURRENT_SOURCE_DIR}/cc/api`
`7`	`7`	`${CMAKE_CURRENT_SOURCE_DIR}/cc/domain`
`8`	`8`	`)`
`9`		`-target_link_libraries(nfsstore PUBLIC storeintf storedevice infra_logger)`
	`9`	`+target_link_libraries(nfsstore PUBLIC storeintf storedevice infra_logger metrics)`
`10`	`10`
`11`	`11`	`file(GLOB_RECURSE UCMSTORE_NFS_CPY_SOURCE_FILES "./cpy/*.cc")`
`12`	`12`	`pybind11_add_module(ucmnfsstore ${UCMSTORE_NFS_CPY_SOURCE_FILES})`