[Feat] Support record operations

flesher0813 · flesher0813 · commit 0e000f66ddf7 · 2025-11-27T15:58:41.000+08:00
diff --git a/ucm/integration/vllm/ucm_connector.py b/ucm/integration/vllm/ucm_connector.py
@@ -18,6 +18,7 @@
 from vllm.v1.request import Request
 
 from ucm.logger import init_logger
+from ucm.shared.metrics.async_op_db import AsyncUCMOperDB
 from ucm.store.factory import UcmConnectorFactory
 from ucm.store.ucmstore import Task, UcmKVStoreBase
 from ucm.utils import Config
@@ -159,6 +160,8 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
             config["kv_block_size"] / 1024 / 1024,
             config["io_size"] / 1024,
         )
+        if role == KVConnectorRole.WORKER and get_tp_group().is_first_rank:
+            self.oper_db = AsyncUCMOperDB()
 
     def generate_hash(self, block_size: int, request: "Request") -> list[str]:
         token_ids = request.all_token_ids
@@ -465,6 +468,14 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
                 request_to_task[request_id] = self.store.load(
                     ucm_total_block_ids, ucm_offsets, dst_tensor_addr
                 )
+                if self.oper_db:
+                    self.oper_db.log_operation(
+                        {
+                            "op_type": "load",
+                            "blocks": ucm_block_ids,
+                            "block_size": self.block_size,
+                        }
+                    )
             else:
                 request_to_task[request_id] = None
             req_broadcast_addr[request_id] = dst_tensor_addr
@@ -527,6 +538,14 @@ def wait_for_save(self) -> None:
             request_to_task[request_id] = self.store.dump(
                 ucm_total_block_ids, ucm_offsets, dst_tensor_addr
             )
+            if self.oper_db:
+                self.oper_db.log_operation(
+                    {
+                        "op_type": "dump",
+                        "blocks": ucm_block_ids,
+                        "block_size": self.block_size,
+                    }
+                )
             request_to_blocks[request_id] = ucm_block_ids
 
         for request_id, task in request_to_task.items():
diff --git a/ucm/shared/metrics/async_op_db.py b/ucm/shared/metrics/async_op_db.py
@@ -0,0 +1,297 @@
+import gzip
+import json
+import os
+import queue
+import shutil
+import threading
+import time
+from datetime import timedelta
+from pathlib import Path
+from typing import Any, Dict
+
+import yaml
+
+from ucm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class AsyncUCMOperDB:
+    """Asynchronous UCM operation log writer (supports log compression + auto-cleanup)"""
+
+    def _load_config(self, config_path: str) -> Dict[str, Any]:
+        """Load configuration from YAML file"""
+        try:
+            with open(config_path, "r") as f:
+                config = yaml.safe_load(f)
+                if config is None:
+                    logger.warning(
+                        f"Config file {config_path} is empty, using defaults"
+                    )
+                    return {}
+                return config
+        except FileNotFoundError:
+            logger.warning(f"Config file {config_path} not found, using defaults")
+            return {}
+        except yaml.YAMLError as e:
+            logger.error(f"Error parsing YAML config file {config_path}: {e}")
+            return {}
+
+    def __init__(
+        self,
+        config_path: str = "",
+    ):
+        if not config_path:
+            current_file_path = Path(__file__)
+            config_path = current_file_path.parent / "metrics_configs.yaml"
+        config = self._load_config(config_path)
+        oper_config = config.get("operation_db", {})
+        if not oper_config:
+            logger.warning(
+                f"Find no operation db configuration in config_path {config_path}, will use default configurations instead!"
+            )
+
+        # Basic configuration
+        self.log_dir = Path(oper_config.get("log_dir", "/vllm-workspace/ucm_logs"))
+        self.log_name = oper_config.get("log_name", "ucm_operation")
+        self.max_file_size = oper_config.get("max_file_size", 104857600)
+        self.batch_size = oper_config.get("batch_size", 100)
+        self.flush_interval = oper_config.get("flush_interval", 5.0)
+        self.encoding = oper_config.get("encoding", "utf-8")
+
+        # Compression configuration
+        self.compress_rotated = oper_config.get("compress_rotated", True)
+        self.compress_level = oper_config.get(
+            "compress_level", 6
+        )  # gzip compression level (1-9)
+
+        # Auto-cleanup configuration (cleanup when any condition is met)
+        self.max_log_files = oper_config.get("max_log_files", 30)
+        self.max_log_days = oper_config.get("max_log_days", 7)
+        self.max_log_total_size = oper_config.get("max_log_total_size", 1073741824)
+
+        # Initialize log directory
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+
+        # Thread-safe queue
+        self.log_queue = queue.Queue(maxsize=10000)  # Max cache: 10000 entries
+        self.is_running = oper_config.get("enabled", False)
+        self.current_file_path = self.log_dir / f"{self.log_name}.log"
+        self.latest_compress_file = self.log_dir / f"{self.log_name}.log.gz"
+        self.batch_buffer = []
+
+        # Start background thread (write + compression + cleanup)
+        self.write_thread = threading.Thread(target=self._async_write_loop, daemon=True)
+        self.write_thread.start()
+
+    def _get_compress_files(self) -> list[Path]:
+        """Get all compressed logs"""
+        compress_files = []
+
+        if self.latest_compress_file.exists():
+            compress_files.append((self.latest_compress_file, 0))
+
+        compress_pattern = f"{self.log_name}.log.gz.*"
+        for file in self.log_dir.glob(compress_pattern):
+            suffix = file.suffixes[-1].lstrip(".")
+            if suffix.isdigit():
+                compress_files.append((file, int(suffix)))
+        return compress_files
+
+    def _roll_old_compress(self) -> None:
+        """Handle old compressed logs"""
+        if not self.latest_compress_file.exists():
+            return
+        all_compress = self._get_compress_files()
+        current_count = len(all_compress)
+
+        if current_count >= self.max_log_files:
+            oldest_file = all_compress[-1]
+            oldest_file.unlink(missing_ok=True)
+            all_compress = self._get_compress_files()
+
+        numbered_files = [f for f, s in all_compress if s > 0]
+
+        for file in reversed(numbered_files):
+            current_suffix = int(file.suffixes[-1].lstrip("."))
+            new_suffix = current_suffix + 1
+            new_file = file.with_suffix(f".{new_suffix}")
+            try:
+                file.rename(new_file)
+
+            except PermissionError as e:
+                logger.error(
+                    f"Failed to rename compressed log file from {file} to {new_file}"
+                )
+                continue
+
+        old_compress_new_name = self.latest_compress_file.with_suffix(f".gz.1")
+        try:
+            self.latest_compress_file.rename(old_compress_new_name)
+        except PermissionError as e:
+            logger.error(
+                f"Failed to rename compressed log file from {self.latest_compress_file} to {old_compress_new_name}"
+            )
+            raise
+
+    def _compress_file(self, source_file: Path) -> None:
+        """Compress specified log file (gzip), delete original after successful compression"""
+        if not source_file.exists() or not self.compress_rotated:
+            return
+
+        try:
+            self._roll_old_compress()
+
+            with source_file.open("rb") as f_in, gzip.open(
+                self.latest_compress_file, "wb", compresslevel=self.compress_level
+            ) as f_out:
+                shutil.copyfileobj(f_in, f_out)
+
+            source_file.unlink(missing_ok=True)
+        except Exception as e:
+            logger.error(f"Failed to compress db log file! Exception {e}")
+
+    def _list_log_files(self) -> list[Path]:
+        """Get all log files (including uncompressed .log and compressed .log.gz), sorted by creation time ascending"""
+        log_files = list(self.log_dir.glob(f"{self.log_name}.log")) + list(
+            self.log_dir.glob(f"{self.log_name}.log.gz*")
+        )
+        # Sort by file creation time (oldest → newest)
+        log_files.sort(key=lambda x: x.stat().st_ctime)
+        return log_files
+
+    def _clean_old_logs(self) -> None:
+        """Auto-clean old logs (delete when any condition is met)"""
+        log_files = self._list_log_files()
+        if not log_files:
+            return
+
+        # Condition 1: Clean by file count (retain newest N files)
+        if self.max_log_files and len(log_files) > self.max_log_files:
+            files_to_delete = log_files[: -self.max_log_files]  # Delete oldest files
+            for file in files_to_delete:
+                self._delete_log_file(file)
+            return  # Stop after meeting one condition to avoid duplicate deletion
+
+        # Condition 2: Clean by time (retain logs from recent N days)
+        if self.max_log_days:
+            expire_time = (
+                time.time() - timedelta(days=self.max_log_days).total_seconds()
+            )
+            files_to_delete = [f for f in log_files if f.stat().st_ctime < expire_time]
+            for file in files_to_delete:
+                self._delete_log_file(file)
+            return
+
+        # Condition 3: Clean by total size (total size not exceeding M)
+        if self.max_log_total_size:
+            total_size = sum(f.stat().st_size for f in log_files)
+            if total_size > self.max_log_total_size:
+                # Delete oldest files first until total size meets requirement
+                current_total = total_size
+                for file in log_files:
+                    if current_total <= self.max_log_total_size:
+                        break
+                    current_total -= file.stat().st_size
+                    self._delete_log_file(file)
+
+    def _delete_log_file(self, file_path: Path) -> None:
+        """Delete log file (fault-tolerant handling)"""
+        try:
+            file_path.unlink()
+        except Exception as e:
+            logger.error(f"Failed to delete log {file_path.name}: {str(e)}")
+
+    def _check_file_rotate(self) -> None:
+        """Check file size, trigger rotation + compression + cleanup if needed"""
+        if (
+            not os.path.exists(self.current_file_path)
+            or os.path.getsize(self.current_file_path) < self.max_file_size
+        ):
+            return
+
+        # 1. Generate new file path (rotation)
+        rotate_file = self.current_file_path.with_suffix(f".1")
+        rotate_file.unlink(missing_ok=True)
+        self.current_file_path.rename(rotate_file)
+
+        # 3. Async compress old file (non-blocking write)
+        threading.Thread(
+            target=self._compress_file, args=(rotate_file,), daemon=True
+        ).start()
+        # 4. Auto-clean old logs
+        threading.Thread(target=self._clean_old_logs, daemon=True).start()
+
+    def _flush_batch(self) -> None:
+        """Batch write logs to file"""
+        if not self.batch_buffer:
+            return
+
+        self._check_file_rotate()  # Check if rotation is needed before writing
+
+        try:
+            with open(self.current_file_path, "a", encoding=self.encoding) as f:
+                for log_entry in self.batch_buffer:
+                    f.write(json.dumps(log_entry, ensure_ascii=False) + "\n")
+            self.batch_buffer.clear()
+        except Exception as e:
+            logger.error(f"Failed to write logs: {str(e)}")
+
+    def _async_write_loop(self) -> None:
+        """Background write loop"""
+        last_flush_time = time.time()
+        while self.is_running:
+            try:
+                # Get log from queue (1 second timeout)
+                log_entry = self.log_queue.get(timeout=1.0)
+                self.batch_buffer.append(log_entry)
+
+                # Flush if conditions are met
+                if (
+                    len(self.batch_buffer) >= self.batch_size
+                    or (last_flush_time - last_flush_time) >= self.flush_interval
+                ):
+                    self._flush_batch()
+                    last_flush_time = time.time()
+
+                self.log_queue.task_done()
+            except queue.Empty:
+                # Check forced flush and cleanup when queue is empty
+                if (time.time() - last_flush_time) >= self.flush_interval:
+                    self._flush_batch()
+                    last_flush_time = time.time()
+                # Periodic cleanup (check every 30 seconds)
+                if int(time.time() - last_flush_time) % 30 == 0:
+                    self._clean_old_logs()
+            except Exception as e:
+                logger.error(f"Log thread exception: {str(e)}")
+
+    def log_operation(self, operation_data: Dict[str, Any]) -> None:
+        """Record operation log (non-blocking)"""
+        if self.is_running:
+            return
+
+        default_data = {
+            "timestamp": time.time(),
+            "op_type": "None",
+        }
+        log_entry = {**default_data, **operation_data}
+
+        try:
+            self.log_queue.put_nowait(log_entry)
+        except queue.Full:
+            logger.error(
+                f"Log queue is full, dropping one log: {log_entry.get('request_id')}"
+            )
+
+    def stop(self) -> None:
+        """Stop logger (ensure data writing + compression)"""
+        self.is_running = False
+        self.write_thread.join(timeout=10.0)
+        self._flush_batch()  # Flush remaining data
+
+        # Final cleanup of old logs
+        self._clean_old_logs()
+
+    def __del__(self):
+        self.stop()
diff --git a/ucm/shared/metrics/metrics_configs.yaml b/ucm/shared/metrics/metrics_configs.yaml
@@ -0,0 +1,18 @@
+# Prometheus Metrics Configuration
+# This file defines which metrics should be enabled and their configurations
+operation_db:
+  enabled: True
+
+  log_dir: "/vllm-workspace/ucm_logs"
+  log_name: "ucm_operation"
+  max_file_size: 104857600          # Max single file size: 100MB (100 * 1024 * 1024 triggers rotation)
+  batch_size: 100                   # Batch write threshold
+  flush_interval: 5.0               # Force flush interval (5 seconds)
+  encoding: "utf-8"
+  # Compression related configuration
+  compress_rotated: True          # Enable rotated file compression (gzip by default)
+  compress_level: 6               # Compression level (1=fastest, 9=smallest, 6=balanced)
+  # Auto-cleanup related configuration
+  max_log_files: 30      # Max number of log files to retain (including compressed)
+  max_log_days: 7        # Max days to retain logs
+  max_log_total_size: 1073741824  # Max total log size: 1GB (1024 * 1024 * 1024)
