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

lmcache/v1/cache_controller/utils.py

@@ -0,0 +1,679 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Callable, NamedTuple, Optional, Set
+import threading
+
+# Third Party
+import zmq.asyncio
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.locks import FastLockWithTimeout, RWLockWithTimeout
+from lmcache.v1.cache_controller.message import BatchedKVOperationMsg, WorkerInfo
+from lmcache.v1.utils.cache_utils import TTLListCache
+
+logger = init_logger(__name__)
+
+
+class KVChunkInfo(NamedTuple):
+    """
+    Represents the location information of a KV chunk in the cluster.
+    This class is immutable and can be used as a dictionary key.
+    """
+
+    instance_id: str
+    worker_id: int
+    location: str
+
+
+class FullSyncState(Enum):
+    """State of full sync for a worker"""
+
+    IDLE = "idle"  # Not in sync
+    SYNCING = "syncing"  # Sync in progress
+    COMPLETED = "completed"  # Sync completed
+    FAILED = "failed"  # Sync failed/timeout
+
+
+@dataclass
+class WorkerSyncInfo:
+    """Information about a worker's sync state"""
+
+    sync_id: str
+    state: FullSyncState
+    start_time: float
+    expected_total_keys: int
+    expected_batch_count: int
+    received_batches: Set[int] = field(default_factory=set)
+    received_keys_count: int = 0
+    last_activity_time: float = 0.0
+
+    def __post_init__(self):
+        if self.last_activity_time == 0.0:
+            self.last_activity_time = self.start_time
+
+
+@dataclass
+class WorkerNode:
+    """
+    Represents a single worker with all its associated metadata.
+    """
+
+    worker_id: int
+    ip: str
+    port: int
+    peer_init_url: Optional[str]
+    socket: Optional[zmq.asyncio.Socket]
+    registration_time: float
+    last_heartbeat_time: float
+    # Guarded by _lock
+    seq_tracker: dict[str, int] = field(default_factory=dict)  # location -> seq_num
+    kv_store: dict[str, set[int]] = field(
+        default_factory=dict
+    )  # location -> set[chunk_hash]
+    sync_info: Optional[WorkerSyncInfo] = None  # Full sync state
+
+    def __post_init__(self):
+        # Fast lock with timeout for WorkerNode operations
+        self._lock = FastLockWithTimeout()
+
+    def handle_batched_kv_operations(
+        self,
+        msg: BatchedKVOperationMsg,
+        on_seq_num_out_of_order: Optional[Callable[[], None]] = None,
+        is_full_sync: bool = False,
+    ) -> bool:
+        """
+        Handle batched KV operations with single lock acquisition.
+        Logs warning and calls callback if sequence out of order is detected.
+
+        Args:
+            msg: The batched KV operation message.
+            on_seq_num_out_of_order: Callback when sequence out of order detected.
+            is_full_sync: If True, this is a full sync batch operation.
+                         During full sync, sequence check is skipped.
+                         When worker is syncing, non-full-sync operations are rejected.
+
+        Returns:
+            True if operations were processed, False if rejected
+            (e.g., during syncing when is_full_sync=False).
+        """
+        seq_warning: Optional[tuple[int, int, int]] = None
+        with self._lock:
+            # During syncing, reject non-full-sync operations
+            if (
+                self.sync_info is not None
+                and self.sync_info.state == FullSyncState.SYNCING
+                and not is_full_sync
+            ):
+                return False
+
+            location = msg.location
+
+            if location not in self.kv_store:
+                self.kv_store[location] = set()
+
+            for op in msg.operations:
+                # Apply operation first, then check sequence
+                # (sequence check is skipped during full sync)
+                if op.op_type.value == "admit":
+                    self.kv_store[location].add(op.key)
+                elif op.op_type.value == "evict":
+                    self.kv_store[location].discard(op.key)
+                else:
+                    logger.error(f"Unknown op_type: {op.op_type}")
+
+                # Skip sequence check during full sync
+                if is_full_sync:
+                    continue
+
+                # Sequence check
+                last_seq_num = self.seq_tracker.get(location)
+                if last_seq_num is not None:
+                    expected_seq = last_seq_num + 1
+                    if op.seq_num != expected_seq and seq_warning is None:
+                        seq_warning = (
+                            expected_seq,
+                            op.seq_num,
+                            op.seq_num - expected_seq,
+                        )
+                self.seq_tracker[location] = op.seq_num
+
+            # Clean up empty set
+            if not self.kv_store[location]:
+                self.kv_store.pop(location, None)
+
+        # Log warning and call callback outside lock
+        if seq_warning is not None:
+            if on_seq_num_out_of_order is not None:
+                on_seq_num_out_of_order()
+            logger.warning(
+                "KV batch sequence out of order detected: "
+                "key=%s, expected_seq=%s, actual_seq=%s, gap=%s",
+                (msg.instance_id, msg.worker_id, msg.location),
+                seq_warning[0],
+                seq_warning[1],
+                seq_warning[2],
+            )
+        return True
+
+    def has_kv(self, location: str, key: int) -> bool:
+        """Check if a KV chunk exists in this worker."""
+        with self._lock:
+            return location in self.kv_store and key in self.kv_store[location]
+
+    def get_kv_keys(self, location: str) -> set[int]:
+        """Get all keys for a location."""
+        with self._lock:
+            keys = self.kv_store.get(location)
+            if keys is None:
+                return set()
+            # Return a shallow copy for thread safety
+            return set(keys)
+
+    def clear_kv_store(self) -> None:
+        """Clear all KV data for this worker."""
+        with self._lock:
+            self.kv_store.clear()
+            self.seq_tracker.clear()
+
+    def get_kv_count(self) -> int:
+        """Get total count of KV chunks."""
+        with self._lock:
+            return sum(len(keys) for keys in self.kv_store.values())
+
+    def update_seq_num(self, location: str, seq_num: int) -> None:
+        """Update sequence number for a location."""
+        with self._lock:
+            self.seq_tracker[location] = seq_num
+
+    def get_seq_num(self, location: str) -> Optional[int]:
+        """Get sequence number for a location."""
+        with self._lock:
+            return self.seq_tracker.get(location)
+
+    def find_key(
+        self, key: int
+    ) -> Optional[tuple[KVChunkInfo, Optional[str], set[int]]]:
+        """
+        Find a key in this worker's kv_store.
+        Returns: (KVChunkInfo, peer_init_url, keys) if found, None otherwise.
+        KVChunkInfo will have instance_id="" since WorkerNode doesn't know
+        its instance.
+        """
+        with self._lock:
+            for location, keys in self.kv_store.items():
+                if key in keys:
+                    # WorkerNode doesn't know its instance_id, so we leave it empty
+                    # The caller should fill in the instance_id
+                    return (
+                        KVChunkInfo("", self.worker_id, location),
+                        self.peer_init_url,
+                        keys,
+                    )
+        return None
+
+    def find_key_simple(self, key: int) -> Optional[KVChunkInfo]:
+        """
+        Find a key in this worker's kv_store, returning only KVChunkInfo.
+        KVChunkInfo will have instance_id="" since WorkerNode doesn't know
+        its instance.
+        """
+        with self._lock:
+            for location, keys in self.kv_store.items():
+                if key in keys:
+                    return KVChunkInfo("", self.worker_id, location)
+        return None
+
+    def to_worker_info(self, instance_id: str) -> WorkerInfo:
+        """Convert to WorkerInfo for backward compatibility."""
+        # No need to lock here
+        return WorkerInfo(
+            instance_id=instance_id,
+            worker_id=self.worker_id,
+            ip=self.ip,
+            port=self.port,
+            peer_init_url=self.peer_init_url,
+            registration_time=self.registration_time,
+            last_heartbeat_time=self.last_heartbeat_time,
+        )
+
+
+@dataclass
+class InstanceNode:
+    """
+    Represents an instance with all its workers.
+    Tree structure: InstanceNode -> WorkerNode
+    Each InstanceNode has its own lock for thread-safe worker operations.
+    """
+
+    instance_id: str
+    # Guarded by _rwlock
+    workers: dict[int, WorkerNode] = field(
+        default_factory=dict
+    )  # worker_id -> WorkerNode
+
+    def __post_init__(self):
+        # RW lock for protecting workers dict access
+        self._rwlock = RWLockWithTimeout()
+
+    def add_worker(self, worker_node: WorkerNode) -> None:
+        """Add a worker to this instance."""
+        with self._rwlock.write_lock(timeout=10):
+            self.workers[worker_node.worker_id] = worker_node
+
+    def remove_worker(self, worker_id: int) -> Optional[WorkerNode]:
+        """Remove and return a worker from this instance."""
+        with self._rwlock.write_lock(timeout=10):
+            return self.workers.pop(worker_id, None)
+
+    def get_worker(self, worker_id: int) -> Optional[WorkerNode]:
+        """Get a worker by worker_id. Uses optimistic locking (lock-free read)."""
+        # Optimistic read: workers dict rarely changes
+        return self.workers.get(worker_id)
+
+    def get_worker_ids(self) -> list[int]:
+        """Get sorted list of worker IDs."""
+        # Snapshot keys to avoid RuntimeError if dict changes during iteration
+        return sorted(list(self.workers.keys()))
+
+    def has_workers(self) -> bool:
+        """Check if instance has any workers."""
+        # Optimistic read: workers dict rarely changes
+        return len(self.workers) > 0
+
+    def get_all_worker_infos(self) -> list[WorkerInfo]:
+        """Get WorkerInfo for all workers in this instance."""
+        # Snapshot values to avoid RuntimeError if dict changes during iteration
+        workers_snapshot = list(self.workers.values())
+        return [worker.to_worker_info(self.instance_id) for worker in workers_snapshot]
+
+    def find_key(
+        self, key: int
+    ) -> Optional[tuple[KVChunkInfo, Optional[str], set[int]]]:
+        """
+        Find a key in any worker within this instance.
+        Returns: (KVChunkInfo, peer_init_url, keys) if found, None otherwise.
+        """
+        # Snapshot workers to avoid RuntimeError if dict changes during iteration
+        workers_snapshot = list(self.workers.items())
+        for worker_id, worker_node in workers_snapshot:
+            if result := worker_node.find_key(key):
+                # Fill in the instance_id in KVChunkInfo
+                kv_info, peer_init_url, keys = result
+                return (
+                    KVChunkInfo(self.instance_id, worker_id, kv_info.location),
+                    peer_init_url,
+                    keys,
+                )
+        return None
+
+    def find_key_simple(self, key: int) -> Optional[KVChunkInfo]:
+        """
+        Find a key in any worker within this instance, returning only KVChunkInfo.
+        """
+        # Snapshot workers to avoid RuntimeError if dict changes during iteration
+        workers_snapshot = list(self.workers.items())
+        for worker_id, worker_node in workers_snapshot:
+            if kv_info := worker_node.find_key_simple(key):
+                return KVChunkInfo(self.instance_id, worker_id, kv_info.location)
+        return None
+
+    def has_worker_with_ip(self, ip: str) -> bool:
+        """
+        Check if any worker in this instance has the specified IP address.
+        """
+        # Snapshot workers to avoid RuntimeError if dict changes during iteration
+        return any(worker.ip == ip for worker in list(self.workers.values()))
+
+    def get_total_kv_count(self) -> int:
+        """
+        Get total count of KV chunks across all workers in this instance.
+        """
+        # Snapshot workers to avoid RuntimeError if dict changes during iteration
+        return sum(worker.get_kv_count() for worker in list(self.workers.values()))
+
+
+class RegistryTree:
+    """
+    Central registry managing the tree structure of instances and workers.
+    Structure: instance_id -> InstanceNode -> WorkerNode
+
+    Lock hierarchy (from coarse to fine):
+    1. RegistryTree._rwlock: protects instances dict access
+    2. InstanceNode._rwlock: protects workers dict access
+    3. WorkerNode._lock: protects kv_store and seq_tracker access
+
+    This fine-grained locking allows concurrent operations on different
+    instances/workers, improving throughput significantly.
+    """
+
+    def __init__(self):
+        # Guarded by _rwlock
+        # instance_id -> InstanceNode
+        self.instances: dict[str, InstanceNode] = {}
+        # RW lock only for protecting instances dict access
+        self._rwlock = RWLockWithTimeout()
+        # Atomic counter for sequence discontinuity (protected by _counter_lock)
+        self._seq_discontinuity_count = 0
+        self._counter_lock = threading.Lock()
+        # Cache for get_all_worker_infos (for Prometheus metrics)
+        self._worker_info_cache: TTLListCache[WorkerInfo] = TTLListCache[WorkerInfo]()
+        # Cache for get_all_worker_nodes (for FullSyncTracker and other use cases)
+        self._worker_node_cache: TTLListCache[tuple[str, WorkerNode]] = TTLListCache[
+            tuple[str, WorkerNode]
+        ]()
+
+    def get_seq_discontinuity_count(self) -> int:
+        """Get the count of sequence discontinuities (thread-safe)."""
+        # Lock-free read, no need for lock
+        return self._seq_discontinuity_count
+
+    def _incr_seq_discontinuity_count(self) -> None:
+        """Increment the sequence discontinuity counter (thread-safe)."""
+        with self._counter_lock:
+            self._seq_discontinuity_count += 1
+
+    def _get_or_create_instance(self, instance_id: str) -> InstanceNode:
+        """Get or create an instance node. Internal use only."""
+        # Optimistic read first: instances dict rarely changes
+        instance_node = self.instances.get(instance_id)
+        if instance_node is not None:
+            return instance_node
+
+        # Need to create, use write lock
+        with self._rwlock.write_lock(timeout=10):
+            # Double-check after acquiring write lock
+            instance_node = self.instances.get(instance_id)
+            if instance_node is None:
+                instance_node = InstanceNode(instance_id=instance_id)
+                self.instances[instance_id] = instance_node
+            return instance_node
+
+    def _get_instance(self, instance_id: str) -> Optional[InstanceNode]:
+        """Get an instance node. Uses optimistic locking (lock-free read)."""
+        # Optimistic read: instances dict rarely changes
+        return self.instances.get(instance_id)
+
+    def register_worker(
+        self,
+        instance_id: str,
+        worker_id: int,
+        ip: str,
+        port: int,
+        peer_init_url: Optional[str],
+        socket: zmq.asyncio.Socket,
+        registration_time: float,
+    ) -> WorkerNode:
+        """Register a new worker, creating instance if needed."""
+        # Get or create instance (locks instances dict)
+        instance_node = self._get_or_create_instance(instance_id)
+
+        # Create worker node
+        worker_node = WorkerNode(
+            worker_id=worker_id,
+            ip=ip,
+            port=port,
+            peer_init_url=peer_init_url,
+            socket=socket,
+            registration_time=registration_time,
+            last_heartbeat_time=registration_time,
+        )
+        # Add worker (locks workers dict in instance_node)
+        instance_node.add_worker(worker_node)
+        return worker_node
+
+    def deregister_worker(
+        self, instance_id: str, worker_id: int
+    ) -> Optional[WorkerNode]:
+        """Deregister a worker and clean up empty instances."""
+        instance_node = self._get_instance(instance_id)
+        if instance_node is None:
+            return None
+
+        # Remove worker (locks workers dict in instance_node)
+        worker_node = instance_node.remove_worker(worker_id)
+
+        # Clean up empty instance (need write lock on instances)
+        if not instance_node.has_workers():
+            # TODO(baoloongmao): Move timeout values to configuration
+            with self._rwlock.write_lock(timeout=100):
+                # Double-check after acquiring write lock
+                # Use pop to avoid KeyError if another thread already removed it
+                if not instance_node.has_workers() and instance_id in self.instances:
+                    self.instances.pop(instance_id, None)
+
+        return worker_node
+
+    def get_worker(self, instance_id: str, worker_id: int) -> Optional[WorkerNode]:
+        """Get a specific worker."""
+        instance_node = self._get_instance(instance_id)
+        if instance_node is None:
+            return None
+        return instance_node.get_worker(worker_id)
+
+    def get_instance(self, instance_id: str) -> Optional[InstanceNode]:
+        """Get an instance by instance_id."""
+        return self._get_instance(instance_id)
+
+    def get_instance_by_ip(self, ip: str) -> Optional[InstanceNode]:
+        """Get an instance by IP address. Returns first instance if multiple exist."""
+        # Snapshot to avoid RuntimeError if dict changes during iteration
+        for instance_node in list(self.instances.values()):
+            if instance_node.has_worker_with_ip(ip):
+                return instance_node
+        return None
+
+    def get_instances_by_ip(self, ip: str) -> list[InstanceNode]:
+        """Get all instances by IP address."""
+        # Snapshot to avoid RuntimeError if dict changes during iteration
+        result = []
+        for instance_node in list(self.instances.values()):
+            if instance_node.has_worker_with_ip(ip):
+                result.append(instance_node)
+        return result
+
+    def get_worker_ids(self, instance_id: str) -> list[int]:
+        """Get sorted list of worker IDs for an instance."""
+        instance_node = self._get_instance(instance_id)
+        if instance_node is None:
+            return []
+        return instance_node.get_worker_ids()
+
+    def get_all_worker_infos_cached(
+        self, timeout_seconds: Optional[float] = None
+    ) -> list[WorkerInfo]:
+        """Get WorkerInfo for all workers across all instances.
+
+        Args:
+            timeout_seconds: Cache timeout in seconds.
+                - If None (default): use the cache's default timeout.
+                - If 0: cache immediately expires, always get fresh data
+                  and update cache.
+                - If > 0: use cache with specified timeout.
+
+        Returns:
+            List of WorkerInfo, either from cache or fresh data
+        """
+        return self._worker_info_cache.get_cached(
+            lambda: self._get_all_worker_infos_fresh(), timeout_override=timeout_seconds
+        )
+
+    def _get_all_worker_infos_fresh(self) -> list[WorkerInfo]:
+        """Internal method to get fresh WorkerInfo without cache."""
+        # Snapshot to avoid RuntimeError if dict changes during iteration
+        result = []
+        for instance_node in list(self.instances.values()):
+            result.extend(instance_node.get_all_worker_infos())
+        return result
+
+    def get_all_worker_nodes_cached(
+        self, timeout_seconds: Optional[float] = None
+    ) -> list[tuple[str, WorkerNode]]:
+        """Get all (instance_id, WorkerNode) tuples across all instances, with caching.
+
+        Args:
+            timeout_seconds: Cache timeout in seconds.
+                - If None (default): use the cache's default timeout.
+                - If 0: cache immediately expires, always get fresh data
+                  and update cache.
+                - If > 0: use cache with specified timeout.
+
+        Returns:
+            List of (instance_id, WorkerNode) tuples, either from cache or fresh data
+        """
+        return self._worker_node_cache.get_cached(
+            lambda: self._get_all_worker_nodes_fresh(), timeout_override=timeout_seconds
+        )
+
+    def _get_all_worker_nodes_fresh(self) -> list[tuple[str, WorkerNode]]:
+        """Internal method to get fresh (instance_id, WorkerNode) tuples
+        without cache."""
+        # Snapshot to avoid RuntimeError if dict changes during iteration
+        result = []
+        for instance_id, instance_node in self.instances.items():
+            for worker_node in instance_node.workers.values():
+                result.append((instance_id, worker_node))
+        return result
+
+    def update_heartbeat(
+        self, instance_id: str, worker_id: int, timestamp: float
+    ) -> bool:
+        """Update worker heartbeat timestamp. Returns True if successful."""
+        instance_node = self._get_instance(instance_id)
+        if instance_node is None:
+            return False
+        worker_node = instance_node.get_worker(worker_id)
+        if worker_node is None:
+            return False
+        worker_node.last_heartbeat_time = timestamp
+        return True
+
+    def handle_batched_kv_operations(
+        self, msg: BatchedKVOperationMsg, is_full_sync: bool = False
+    ) -> bool:
+        """
+        Handle batched KV operations by forwarding to WorkerNode.
+
+        Args:
+            msg: The batched KV operation message.
+            is_full_sync: If True, this is a full sync batch operation.
+
+        Returns:
+            True if operations were processed successfully.
+            False if worker not found or operations were rejected.
+        """
+        instance_node = self._get_instance(msg.instance_id)
+        if instance_node is None:
+            return False
+        worker_node = instance_node.get_worker(msg.worker_id)
+        if worker_node is None:
+            return False
+        return worker_node.handle_batched_kv_operations(
+            msg,
+            on_seq_num_out_of_order=self._incr_seq_discontinuity_count,
+            is_full_sync=is_full_sync,
+        )
+
+    def find_kv(
+        self,
+        key: int,
+        exclude_instance_id: Optional[str] = None,
+    ) -> Optional[KVChunkInfo]:
+        """
+        Find a KV chunk across all workers.
+
+        Args:
+            key: The KV chunk key to find.
+            exclude_instance_id: Instance ID to exclude
+            (all workers in this instance will be excluded).
+
+        Returns: KVChunkInfo if found, None otherwise.
+        """
+        with self._rwlock.read_lock(timeout=1):
+            for instance_id, instance_node in self.instances.items():
+                # Exclude all workers in the specified instance
+                if (
+                    exclude_instance_id is not None
+                    and instance_id == exclude_instance_id
+                ):
+                    continue
+                result = instance_node.find_key_simple(key)
+                if result is not None:
+                    return result
+            return None
+
+    def find_kv_with_worker_info(
+        self,
+        key: int,
+        exclude_instance_id: Optional[str] = None,
+    ) -> Optional[tuple[KVChunkInfo, Optional[str], set[int]]]:
+        """
+        Find a KV chunk and return worker info in one lookup.
+        Optimized for batched_p2p_lookup to avoid multiple lookups.
+
+        Returns: (KVChunkInfo, peer_init_url, kv_keys) if found, None otherwise.
+        """
+        # Snapshot instances to avoid RuntimeError if dict changes during iteration
+        for instance_id, instance_node in list(self.instances.items()):
+            if exclude_instance_id is not None and instance_id == exclude_instance_id:
+                continue
+            result = instance_node.find_key(key)
+            if result is not None:
+                return result
+        return None
+
+    def get_total_kv_count(self) -> int:
+        """Get total count of KV chunks across all workers."""
+        # Snapshot to avoid RuntimeError if dict changes during iteration
+        total = 0
+        for instance_node in list(self.instances.values()):
+            total += instance_node.get_total_kv_count()
+        return total
+
+    def get_worker_kv_keys(
+        self, instance_id: str, worker_id: int, location: str
+    ) -> set[int]:
+        """Get all KV keys for a specific worker and location."""
+        instance_node = self._get_instance(instance_id)
+        if instance_node is None:
+            return set()
+        worker_node = instance_node.get_worker(worker_id)
+        if worker_node is None:
+            return set()
+        return worker_node.get_kv_keys(location)
+
+    def clear_worker_kv(
+        self, instance_id: str, worker_id: int, location: Optional[str] = None
+    ) -> bool:
+        """
+        Clear all KV chunks for a specific worker and location.
+        Returns True if successful, False if worker not found.
+        """
+        worker_node = self.get_worker(instance_id, worker_id)
+        if worker_node is None:
+            return False
+        with worker_node._lock:
+            if location is None:
+                worker_node.kv_store.clear()
+                worker_node.seq_tracker.clear()
+                return True
+            if location in worker_node.kv_store:
+                del worker_node.kv_store[location]
+            if location in worker_node.seq_tracker:
+                del worker_node.seq_tracker[location]
+        return True
+
+    def clear_all_worker_kv(self, instance_id: str, worker_id: int) -> bool:
+        """
+        Clear all KV chunks for a worker across all locations.
+        Returns True if successful, False if worker not found.
+        """
+        worker_node = self.get_worker(instance_id, worker_id)
+        if worker_node is None:
+            return False
+
+        worker_node.clear_kv_store()
+        return True