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

lmcache/v1/cache_controller/controllers/full_sync_tracker.py

@@ -0,0 +1,473 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Tuple
+import time
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.utils import (
+    FullSyncState,
+    RegistryTree,
+    WorkerNode,
+    WorkerSyncInfo,
+)
+
+logger = init_logger(__name__)
+
+
+class FullSyncTracker:
+    """
+    Tracks full sync state for all workers.
+
+    This class manages the state of full sync operations, including:
+    - Tracking which workers need full sync
+    - Monitoring sync progress
+    - Handling sync timeout
+    - Determining when freeze mode can be exited
+    """
+
+    def __init__(
+        self,
+        registry_tree: RegistryTree,
+        completion_threshold: float = 0.8,
+        sync_timeout_s: float = 300.0,
+    ):
+        """
+        Initialize the FullSyncTracker.
+
+        Args:
+            registry_tree: The registry tree containing worker nodes
+            completion_threshold: Percentage of workers that need to complete
+                                  sync before others can exit freeze mode (default: 80%)
+            sync_timeout_s: Timeout in seconds for a single worker's sync
+                                  (default: 300s)
+        """
+        self.registry_tree = registry_tree
+        self.completion_threshold = completion_threshold
+        self.sync_timeout_s = sync_timeout_s
+
+        # Flag to indicate if controller just restarted and needs full sync
+        self._need_full_sync_all = True
+
+    def _get_sync_info(
+        self, instance_id: str, worker_id: int
+    ) -> Optional[WorkerSyncInfo]:
+        """Get sync info for a worker from the registry tree."""
+        worker_node = self.registry_tree.get_worker(instance_id, worker_id)
+        if worker_node is None:
+            return None
+        return worker_node.sync_info
+
+    def _set_sync_info(
+        self, instance_id: str, worker_id: int, sync_info: Optional[WorkerSyncInfo]
+    ) -> bool:
+        """Set sync info for a worker. Returns True if successful."""
+        worker_node = self.registry_tree.get_worker(instance_id, worker_id)
+        if worker_node is None:
+            return False
+        worker_node.sync_info = sync_info
+        return True
+
+    def set_need_full_sync_all(self, need: bool) -> None:
+        """Set whether all workers need full sync (e.g., after controller restart)"""
+        self._need_full_sync_all = need
+        logger.info("Set need_full_sync_all to %s", need)
+
+    def _get_all_workers_cached(
+        self, timeout_seconds: Optional[float] = None
+    ) -> list[tuple[str, WorkerNode]]:
+        """Get all registered workers from the registry tree."""
+        return self.registry_tree.get_all_worker_nodes_cached(timeout_seconds)
+
+    def should_request_full_sync(
+        self, instance_id: str, worker_id: int
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Check if a worker should perform full sync.
+
+        Returns:
+            Tuple of (need_sync, reason)
+        """
+        # Case 1: Controller just restarted, all workers need sync
+        if self._need_full_sync_all:
+            sync_info = self._get_sync_info(instance_id, worker_id)
+            if sync_info is None or sync_info.state not in (
+                FullSyncState.SYNCING,
+                FullSyncState.COMPLETED,
+            ):
+                return True, "controller_restart"
+
+        # Case 2: Worker sync failed/timeout, needs retry
+        sync_info = self._get_sync_info(instance_id, worker_id)
+        if sync_info is not None and sync_info.state == FullSyncState.FAILED:
+            return True, "sync_failed_retry"
+
+        return False, None
+
+    def is_worker_syncing(self, instance_id: str, worker_id: int) -> bool:
+        """
+        Check if a worker is currently in sync state.
+
+        When a worker is syncing, incremental events should be discarded.
+        """
+        sync_info = self._get_sync_info(instance_id, worker_id)
+        if sync_info is None:
+            return False
+        return sync_info.state == FullSyncState.SYNCING
+
+    def get_sync_id(self, instance_id: str, worker_id: int) -> Optional[str]:
+        """
+        Get sync ID for a worker.
+
+        Returns sync ID if worker is syncing, None otherwise.
+        """
+        sync_info = self._get_sync_info(instance_id, worker_id)
+        if sync_info is None:
+            return None
+        return sync_info.sync_id
+
+    def start_sync(
+        self,
+        instance_id: str,
+        worker_id: int,
+        sync_id: str,
+        total_keys: int,
+        batch_count: int,
+    ) -> bool:
+        """
+        Start sync for a worker.
+
+        Returns:
+            True if sync started successfully, False otherwise
+        """
+        report_id = (instance_id, worker_id)
+        current_time = time.time()
+
+        # Check if already syncing with different sync_id
+        existing_sync = self._get_sync_info(instance_id, worker_id)
+        if existing_sync is not None and existing_sync.state == FullSyncState.SYNCING:
+            if existing_sync.sync_id != sync_id:
+                logger.warning(
+                    "Worker %s already syncing with different sync_id: "
+                    "existing=%s, new=%s",
+                    report_id,
+                    existing_sync.sync_id,
+                    sync_id,
+                )
+                return False
+
+        new_sync_info = WorkerSyncInfo(
+            sync_id=sync_id,
+            state=FullSyncState.SYNCING,
+            start_time=current_time,
+            expected_total_keys=total_keys,
+            expected_batch_count=batch_count,
+            last_activity_time=current_time,
+        )
+        if not self._set_sync_info(instance_id, worker_id, new_sync_info):
+            logger.warning(
+                "Failed to start sync for worker %s: worker not found", report_id
+            )
+            return False
+
+        logger.info(
+            "Started full sync for worker %s: sync_id=%s, "
+            "expected_keys=%d, expected_batches=%d",
+            report_id,
+            sync_id,
+            total_keys,
+            batch_count,
+        )
+        return True
+
+    def receive_batch(
+        self,
+        instance_id: str,
+        worker_id: int,
+        sync_id: str,
+        batch_id: int,
+        keys_count: int,
+    ) -> bool:
+        """
+        Record receipt of a sync batch.
+
+        Returns:
+            True if batch was recorded, False if invalid
+        """
+        report_id = (instance_id, worker_id)
+        sync_info = self._get_sync_info(instance_id, worker_id)
+
+        if sync_info is None:
+            logger.warning(
+                "Received batch for unknown sync session: worker=%s, sync_id=%s",
+                report_id,
+                sync_id,
+            )
+            return False
+
+        if sync_info.sync_id != sync_id:
+            logger.warning(
+                "Sync ID mismatch: expected=%s, received=%s",
+                sync_info.sync_id,
+                sync_id,
+            )
+            return False
+
+        if sync_info.state != FullSyncState.SYNCING:
+            logger.warning(
+                "Received batch for non-syncing worker: worker=%s, state=%s",
+                report_id,
+                sync_info.state,
+            )
+            return False
+
+        sync_info.received_batches.add(batch_id)
+        sync_info.received_keys_count += keys_count
+        sync_info.last_activity_time = time.time()
+
+        logger.debug(
+            "Received batch %d for worker %s: keys=%d, total_received=%d",
+            batch_id,
+            report_id,
+            keys_count,
+            sync_info.received_keys_count,
+        )
+        return True
+
+    def complete_sync(
+        self,
+        instance_id: str,
+        worker_id: int,
+        sync_id: str,
+        actual_total_keys: int,
+    ) -> bool:
+        """
+        Mark sync as completed for a worker.
+
+        Returns:
+            True if completion was successful, False otherwise
+        """
+        report_id = (instance_id, worker_id)
+        sync_info = self._get_sync_info(instance_id, worker_id)
+
+        if sync_info is None:
+            logger.warning(
+                "Received sync end for unknown session: worker=%s, sync_id=%s",
+                report_id,
+                sync_id,
+            )
+            return False
+
+        if sync_info.sync_id != sync_id:
+            logger.warning(
+                "Sync ID mismatch on completion: expected=%s, received=%s",
+                sync_info.sync_id,
+                sync_id,
+            )
+            return False
+
+        # Verify key count
+        if sync_info.received_keys_count != actual_total_keys:
+            logger.warning(
+                "Key count mismatch on completion: received=%d, reported=%d",
+                sync_info.received_keys_count,
+                actual_total_keys,
+            )
+            # Still mark as completed but log the discrepancy
+
+        sync_info.state = FullSyncState.COMPLETED
+        sync_info.last_activity_time = time.time()
+
+        logger.info(
+            "Completed full sync for worker %s: sync_id=%s, "
+            "received_keys=%d, batches=%d",
+            report_id,
+            sync_id,
+            sync_info.received_keys_count,
+            len(sync_info.received_batches),
+        )
+        return True
+
+    def mark_failed(self, instance_id: str, worker_id: int, reason: str) -> None:
+        """Mark a worker's sync as failed"""
+        report_id = (instance_id, worker_id)
+        sync_info = self._get_sync_info(instance_id, worker_id)
+
+        if sync_info is not None:
+            sync_info.state = FullSyncState.FAILED
+            logger.warning(
+                "Marked sync as failed for worker %s: reason=%s", report_id, reason
+            )
+
+    # Only for testing
+    def check_sync_timeout(self) -> None:
+        """
+        Check for sync timeouts and mark failed workers.
+
+        This should be called periodically (e.g., in health check loop).
+        """
+        current_time = time.time()
+        for instance_id, worker_node in self._get_all_workers_cached():
+            sync_info = worker_node.sync_info
+            if sync_info is not None and sync_info.state == FullSyncState.SYNCING:
+                if current_time - sync_info.last_activity_time > self.sync_timeout_s:
+                    self.mark_failed(
+                        instance_id,
+                        worker_node.worker_id,
+                        f"timeout after {self.sync_timeout_s}s",
+                    )
+
+    def get_global_progress(self) -> float:
+        """
+        Get the global sync progress.
+
+        Returns:
+            Progress as a float between 0.0 and 1.0
+
+        Note: Uses cached worker list for Prometheus metrics efficiency.
+
+        Progress calculation:
+        - Denominator: total number of all workers
+        - Numerator: workers that are ready to serve (COMPLETED or no sync info needed)
+        - Workers in SYNCING or FAILED state are NOT considered ready
+        """
+        all_workers = self._get_all_workers_cached()
+        if not all_workers:
+            return 0.0
+
+        total = len(all_workers)
+        ready_count = sum(
+            1
+            for _, worker_node in all_workers
+            if worker_node.sync_info is None
+            or worker_node.sync_info.state == FullSyncState.COMPLETED
+        )
+
+        # Progress = ready / total
+        return ready_count / total
+
+    def get_completed_count(self) -> int:
+        """Get count of workers that have completed sync.
+
+        Note: Uses cached worker list for Prometheus metrics efficiency.
+        """
+        return sum(
+            1
+            for _, worker_node in self._get_all_workers_cached()
+            if worker_node.sync_info is not None
+            and worker_node.sync_info.state == FullSyncState.COMPLETED
+        )
+
+    def get_syncing_count(self) -> int:
+        """Get count of workers currently syncing.
+
+        Note: Uses cached worker list for Prometheus metrics efficiency.
+        """
+        return sum(
+            1
+            for _, worker_node in self._get_all_workers_cached()
+            if worker_node.sync_info is not None
+            and worker_node.sync_info.state == FullSyncState.SYNCING
+        )
+
+    def can_exit_freeze(self, progress: Optional[float] = None) -> bool:
+        """
+        Check if the completion threshold is reached and freeze mode can be exited.
+
+        Args:
+            progress: Pre-computed global progress. If None, will be computed.
+
+        Returns:
+            True if enough workers have completed sync
+        """
+        if progress is None:
+            progress = self.get_global_progress()
+        can_exit = progress >= self.completion_threshold
+
+        if can_exit and self._need_full_sync_all:
+            # Once threshold is reached, disable the global full sync flag
+            logger.info(
+                "Full sync completion threshold reached (%.1f%%), "
+                "disabling need_full_sync_all",
+                progress * 100,
+            )
+            self._need_full_sync_all = False
+
+        return can_exit
+
+    def get_total_missing_batches_count(self) -> int:
+        """
+        Get total count of missing batches across all syncing workers.
+
+        Returns:
+            Total number of missing batches
+
+        Note: Uses cached worker list for Prometheus metrics efficiency.
+        """
+        total = 0
+        for instance_id, worker_node in self._get_all_workers_cached():
+            sync_info = worker_node.sync_info
+            if sync_info is not None and sync_info.state == FullSyncState.SYNCING:
+                expected_batches = set(range(sync_info.expected_batch_count))
+                missing = expected_batches - sync_info.received_batches
+                total += len(missing)
+        return total
+
+    def get_missing_batches(
+        self, instance_id: str, worker_id: int, sync_id: str
+    ) -> list[int]:
+        """
+        Get list of missing batch IDs that need to be resent.
+
+        Args:
+            instance_id: The instance ID
+            worker_id: The worker ID
+            sync_id: The sync session ID
+
+        Returns:
+            List of missing batch IDs, empty if sync is complete or invalid
+        """
+        sync_info = self._get_sync_info(instance_id, worker_id)
+
+        if sync_info is None:
+            return []
+
+        # Check sync_id matches
+        if sync_info.sync_id != sync_id:
+            return []
+
+        # If already completed, no missing batches
+        if sync_info.state == FullSyncState.COMPLETED:
+            return []
+
+        # If not syncing, no missing batches
+        if sync_info.state != FullSyncState.SYNCING:
+            return []
+
+        # Calculate missing batches
+        expected_batches = set(range(sync_info.expected_batch_count))
+        missing = expected_batches - sync_info.received_batches
+
+        return sorted(missing)
+
+    def get_sync_status(
+        self, instance_id: str, worker_id: int, sync_id: str
+    ) -> Tuple[bool, float, bool, list[int]]:
+        """
+        Get sync status for a specific worker.
+
+        Returns:
+            Tuple of (is_complete, global_progress, can_exit_freeze, missing_batches)
+        """
+        sync_info = self._get_sync_info(instance_id, worker_id)
+
+        is_complete = (
+            sync_info is not None
+            and sync_info.sync_id == sync_id
+            and sync_info.state == FullSyncState.COMPLETED
+        )
+        global_progress = self.get_global_progress()
+        can_exit = self.can_exit_freeze(global_progress)
+        missing_batches = self.get_missing_batches(instance_id, worker_id, sync_id)
+
+        return is_complete, global_progress, can_exit, missing_batches