lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/v1/pin_monitor.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Optional
+import threading
+import time
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.observability import LMCStatsMonitor, PrometheusLogger
+from lmcache.v1.periodic_thread import (
+    PeriodicThread,
+    PeriodicThreadRegistry,
+    ThreadLevel,
+    ThreadRunSummary,
+)
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.config import LMCacheEngineConfig
+    from lmcache.v1.memory_management import MemoryObj
+
+
+logger = init_logger(__name__)
+
+
+class PinMonitor(PeriodicThread):
+    """
+    Global monitor (singleton per process, shared across all cache engines)
+    for pinned TensorMemoryObj instances to handle timeout detection.
+    This class runs a background thread that periodically checks for pinned objects
+    that have exceeded their timeout duration.
+    """
+
+    _instance = None
+    _instance_lock = threading.Lock()  # Class-level lock for singleton pattern
+
+    def __init__(self, config: "LMCacheEngineConfig"):
+        # Initialize PeriodicThread base class
+        super().__init__(
+            name="PinMonitor-thread",
+            interval=config.pin_check_interval_sec,
+            level=ThreadLevel.CRITICAL,
+            init_wait=0.0,
+        )
+
+        # obj_id is the virtual memory address given by Python's id() function
+        self._pinned_objects: dict[
+            int, tuple["MemoryObj", float]
+        ] = {}  # {obj_id: (memory_obj, register_time)}
+        self._objects_lock = threading.Lock()
+        self._check_interval = config.pin_check_interval_sec
+        self._pin_timeout_sec = config.pin_timeout_sec
+
+        # Register with the global registry
+        PeriodicThreadRegistry.get_instance().register(self)
+
+        # Auto-start the monitor on first instance creation
+        self.start_monitoring()
+
+    @staticmethod
+    def GetOrCreate(config: Optional["LMCacheEngineConfig"] = None) -> "PinMonitor":
+        """Get or create the singleton instance.
+
+        Args:
+            config: Required for first-time initialization.
+                Optional for subsequent calls.
+
+        Raises:
+            ValueError: If config is None when creating the instance
+                for the first time.
+        """
+        if PinMonitor._instance is None:
+            with PinMonitor._instance_lock:
+                if PinMonitor._instance is None:
+                    assert config is not None, "config is required"
+                    PinMonitor._instance = PinMonitor(config)
+        return PinMonitor._instance
+
+    def on_pin(self, memory_obj: "MemoryObj"):
+        """Register a pinned memory object for timeout monitoring.
+
+        Note: The same memory_obj can be pinned multiple times, so this
+        function may be called multiple times with the same object.
+        Each call updates the register time, effectively resetting the
+        timeout countdown.
+        """
+        obj_id = id(memory_obj)
+        with self._objects_lock:
+            current_time = time.time()
+            self._pinned_objects[obj_id] = (memory_obj, current_time)
+            logger.debug(
+                "Registered pinned object %s for timeout monitoring at time %.2f",
+                obj_id,
+                current_time,
+            )
+
+    def on_unpin(self, memory_obj: "MemoryObj"):
+        """Unregister a memory object from timeout monitoring."""
+        obj_id = id(memory_obj)
+        with self._objects_lock:
+            if obj_id in self._pinned_objects:
+                del self._pinned_objects[obj_id]
+                logger.debug(
+                    "Unregistered pinned object %s from timeout monitoring",
+                    obj_id,
+                )
+
+    def _check_timeouts(self) -> tuple[int, int, int]:
+        """Check all registered pinned objects for timeout.
+
+        Returns:
+            tuple: (pinned_count, timeout_count, force_unpin_success_count)
+        """
+        current_time = time.time()
+        timeout_objects = []
+
+        with self._objects_lock:
+            pinned_count = len(self._pinned_objects)
+            for obj_id, (memory_obj, register_time) in list(
+                self._pinned_objects.items()
+            ):
+                # Check if object is still pinned and has exceeded timeout
+                if memory_obj.meta.pin_count > 0:
+                    elapsed_time = current_time - register_time
+                    if elapsed_time > self._pin_timeout_sec:
+                        timeout_objects.append((memory_obj, elapsed_time))
+
+        # Force unpin timeout objects outside the lock to avoid deadlocks
+        force_unpin_success_count = 0
+        for memory_obj, elapsed_time in timeout_objects:
+            try:
+                self._force_unpin_timeout_object(memory_obj, elapsed_time)
+                force_unpin_success_count += 1
+            except Exception as e:
+                logger.error(
+                    "Error forcing unpin for timeout object %s: %s", id(memory_obj), e
+                )
+        if force_unpin_success_count > 0:
+            logger.warning(
+                "Force unpinned %d timeout objects in %d pinned_objects "
+                "within %d seconds",
+                force_unpin_success_count,
+                pinned_count,
+                self._pin_timeout_sec,
+            )
+        else:
+            logger.debug(
+                "PinMonitor check: pinned_objects=%d, timeout_objects=%d, "
+                "force_unpin_success=%d",
+                pinned_count,
+                len(timeout_objects),
+                force_unpin_success_count,
+            )
+
+        return pinned_count, len(timeout_objects), force_unpin_success_count
+
+    def _force_unpin_timeout_object(self, memory_obj: "MemoryObj", elapsed_time: float):
+        """Force unpin a timeout object and log the event."""
+        # Get current pin_count without holding the lock for unpin calls
+        # Use nullcontext if memory_obj doesn't have a lock attribute
+        obj_lock = getattr(memory_obj, "lock", None) or nullcontext()
+        with obj_lock:
+            current_pin_count = memory_obj.meta.pin_count
+            if current_pin_count <= 0:
+                return
+
+            logger.warning(
+                "Pin timeout detected for MemoryObj %s. "
+                "Pin count: %s, Elapsed time: %.2fs. Forcing unpin to 0.",
+                memory_obj.meta.address,
+                current_pin_count,
+                elapsed_time,
+            )
+
+        # Update forced unpin statistics
+        LMCStatsMonitor.GetOrCreate().update_forced_unpin_count(1)
+
+        # Call unpin() while pin_count > 0 to properly release resources
+        while memory_obj.meta.pin_count > 0:
+            memory_obj.unpin()
+
+    def _execute(self) -> ThreadRunSummary:
+        """
+        Execute one pin monitor check cycle.
+
+        This method is called by the PeriodicThread base class.
+
+        Returns:
+            ThreadRunSummary: Summary of the check cycle
+        """
+        pinned_count, timeout_count, force_unpin_count = self._check_timeouts()
+
+        return ThreadRunSummary(
+            success=True,
+            message=f"Checked {pinned_count} objects, {timeout_count} timeouts, "
+            f"{force_unpin_count} force unpinned",
+            extra_info={
+                "pinned_count": str(pinned_count),
+                "timeout_count": str(timeout_count),
+                "force_unpin_count": str(force_unpin_count),
+            },
+        )
+
+    def start_monitoring(self):
+        """Start the background monitoring thread."""
+        if self.is_running:
+            return
+
+        # Use base class start method
+        self.start()
+        logger.info("PinMonitor started")
+
+        # Setup metrics callback
+        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
+        if prometheus_logger is not None:
+            prometheus_logger.pin_monitor_pinned_objects_count.set_function(
+                lambda: len(self._pinned_objects)
+            )
+
+    def stop_monitoring(self):
+        """Stop the background monitoring thread."""
+        if not self.is_running:
+            return
+
+        # Unregister from the global registry
+        PeriodicThreadRegistry.get_instance().unregister(self.name)
+
+        # Use base class stop method
+        self.stop()
+        logger.info("PinMonitor stopped")
+
+    def get_monitored_count(self) -> int:
+        """Get the number of currently monitored pinned objects."""
+        with self._objects_lock:
+            return len(self._pinned_objects)
+
+    @staticmethod
+    def DestroyInstance():
+        """Destroy the singleton instance and stop monitoring.
+        This is mainly used for testing to ensure clean state between tests.
+        """
+        with PinMonitor._instance_lock:
+            if PinMonitor._instance is not None:
+                PinMonitor._instance.stop_monitoring()
+                PinMonitor._instance = None

lmcache/v1/plugin/runtime_plugin_launcher.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from pathlib import Path
+import atexit
+import os
+import shutil
+import subprocess
+import threading
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class RuntimePluginLauncher:
+    def __init__(
+        self,
+        config,
+        role: str | None,
+        worker_count: int,
+        worker_id: int,
+    ):
+        self.config = config
+        self.role = role
+        self.worker_count = worker_count
+        self.worker_id = worker_id
+        self.plugin_processes: list[subprocess.Popen] = []
+        # Register cleanup handler
+        atexit.register(self.stop_plugins)
+
+    def launch_plugins(self):
+        """Launch all configured plugins"""
+        if not self.config.runtime_plugin_locations:
+            return
+
+        for loc in self.config.runtime_plugin_locations:
+            self._launch_plugins(loc)
+
+    def _launch_plugins(self, loc: str):
+        """Launch plugins from specified location"""
+        path = Path(loc)
+        if not path.exists():
+            logger.warning(f"Runtime plugin location {loc} does not exist")
+            return
+
+        files = []
+        if path.is_file():
+            files = [path]
+        elif path.is_dir():
+            # Recursively find all .py and .sh files
+            for ext in ["*.py", "*.sh"]:
+                files.extend(path.rglob(ext))
+
+        for file in files:
+            self._launch_plugin(file)
+
+    def _should_skip_plugin(self, file: Path, parts: list[str]) -> bool:
+        """Determine if plugin should be skipped based on role/worker ID.
+
+        The naming convention is ROLE_WORKERID_NAME (e.g. server_0_foo.py).
+        When role is None (MP mode), this convention does not apply,
+        so both role and worker ID filtering are skipped entirely.
+        """
+        if len(parts) < 2:
+            return False
+
+        # When role is None, skip all role/worker-based filtering
+        # (e.g. MP mode has no role concept)
+        if self.role is None:
+            return False
+
+        plugin_role = parts[0].upper()
+        if plugin_role != "ALL" and plugin_role != self.role.upper():
+            logger.info(
+                "Skipping %s: requires role %s",
+                file,
+                plugin_role,
+            )
+            return True
+
+        # Check worker ID match (parts[1] in ROLE_WORKERID_NAME)
+        if len(parts) > 2 and parts[1].isdigit():
+            plugin_worker_id = int(parts[1])
+            if plugin_worker_id != self.worker_id:
+                logger.info(
+                    "worker %d is skipping plugin %s, which is only for worker ID %d",
+                    self.worker_id,
+                    file,
+                    plugin_worker_id,
+                )
+                return True
+
+        return False
+
+    def _launch_plugin(self, file: Path):
+        """Launch a plugin"""
+        try:
+            filename = file.stem.lower()
+            parts = filename.split("_")
+
+            if self._should_skip_plugin(file, parts):
+                return
+
+            # Get interpreter from first line (shebang)
+            interpreter = self._get_interpreter(file)
+
+            # Pass role and config as environment variables
+            env = os.environ.copy()
+            role_str = self.role or ""
+            env["LMCACHE_RUNTIME_PLUGIN_ROLE"] = role_str
+            env["LMCACHE_RUNTIME_PLUGIN_CONFIG"] = self.config.to_json()
+            env["LMCACHE_RUNTIME_PLUGIN_WORKER_COUNT"] = str(self.worker_count)
+            env["LMCACHE_RUNTIME_PLUGIN_WORKER_ID"] = str(self.worker_id)
+
+            # TODO: For backwards compatibility, remove when applicable
+            env["LMCACHE_PLUGIN_ROLE"] = role_str
+            env["LMCACHE_PLUGIN_CONFIG"] = self.config.to_json()
+            env["LMCACHE_PLUGIN_WORKER_COUNT"] = str(self.worker_count)
+            env["LMCACHE_PLUGIN_WORKER_ID"] = str(self.worker_id)
+
+            # Force line-buffered stdout for Python sub-processes
+            # so that output is captured in real-time rather than
+            # being held in the block buffer until exit.
+            env["PYTHONUNBUFFERED"] = "1"
+
+            proc = subprocess.Popen(
+                [interpreter, str(file)],
+                env=env,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+            )
+            self.plugin_processes.append(proc)
+            logger.info(f"Launched runtime plugin: {file} with {interpreter}")
+
+            # Start thread to capture output continuously
+            threading.Thread(
+                target=self._capture_plugin_output, args=(proc, str(file)), daemon=True
+            ).start()
+        except Exception as e:
+            logger.error(f"Failed to launch plugin {file}: {e}")
+
+    def _get_interpreter(self, file: Path) -> str:
+        """Get interpreter from first line comment"""
+        interpreters = []
+        try:
+            with open(file, "r", encoding="utf-8") as f:
+                first_line = f.readline().strip()
+                if first_line.startswith("#!"):
+                    # Extract interpreter
+                    interpreter_str = first_line[2:].strip()
+                    interpreters.append(interpreter_str)
+        except Exception as e:
+            logger.error(
+                f"Error reading interpreter from runtime plugin file {file} - "
+                f"using default interpreters: {e}"
+            )
+            pass
+
+        # Fallback to default interpreters
+        if file.suffix == ".py":
+            interpreters.append("python")
+            interpreters.append("python3")
+        elif file.suffix == ".sh":
+            interpreters.append("bash")
+        else:
+            raise ValueError(f"Plugin type {file.suffix} not supported ")
+
+        # Try each interpreter until we find one that exists
+        for interpreter in interpreters:
+            interpreter = interpreter.strip()
+            resolved_interpreter = shutil.which(interpreter)
+            if resolved_interpreter:
+                return resolved_interpreter
+
+        raise ValueError(f"No valid interpreter found for {file} from {interpreters}")
+
+    def _capture_plugin_output(self, proc: subprocess.Popen, plugin_name: str):
+        """Continuously capture and log plugin output"""
+        try:
+            assert proc.stdout is not None, (
+                "The runtime plugin subprocess does not have stdout"
+            )
+            while True:
+                line = proc.stdout.readline()
+                if not line:
+                    break
+                logger.info(f"[{plugin_name}] {line.strip()}")
+
+            proc.wait()
+            logger.info(
+                f"Runtime plugin {plugin_name} exited with code {proc.returncode}"
+            )
+        except Exception as e:
+            logger.error(f"Error capturing output for {plugin_name}: {e}")
+
+    def stop_plugins(self):
+        """Terminate all plugin processes.
+
+        Note: This method may be called via atexit during interpreter shutdown,
+        when the logging system may already be closed. We avoid using logger
+        here to prevent "I/O operation on closed file" errors.
+        """
+        for proc in self.plugin_processes:
+            try:
+                if proc.poll() is None:
+                    proc.terminate()
+            except Exception:
+                # Silently ignore errors during shutdown
+                pass