comet-ml--opik/apps/opik-python-backend/src/opik_backend/studio/cancellation.py

"""
Cancellation monitor for Optimization Studio jobs.

Uses a single background thread to monitor cancellation signals for all
running optimizations via Redis MGET (multi-get). This is more efficient
than having a separate polling thread per optimization.

Architecture:
    - CancellationMonitor: Singleton that polls Redis for all registered optimizations
    - CancellationHandle: Per-optimization handle for registration and status checking

Environment Variables (for self-hosted deployments):
    OPTSTUDIO_CANCEL_POLL_INTERVAL_SECS: Polling interval in seconds (default: 2)
"""

import logging
import os
import threading
from typing import Any, Callable, Dict, Optional, Type

from opik_backend.utils.redis_utils import get_redis_client

logger = logging.getLogger(__name__)

# Redis key pattern for cancellation signals
CANCEL_KEY_PATTERN = "opik:cancel:{}"

# Environment variable for polling interval
ENV_CANCEL_POLL_INTERVAL_SECS = "OPTSTUDIO_CANCEL_POLL_INTERVAL_SECS"

# Default polling interval
DEFAULT_POLL_INTERVAL_SECS = 2


class CancellationMonitor:
    """
    Singleton monitor that checks cancellation status for all running optimizations.
    
    Uses a single thread with Redis MGET to efficiently check multiple keys at once.
    Optimizations register themselves and provide a callback for when cancelled.
    """
    
    _instance: Optional["CancellationMonitor"] = None
    _lock: threading.Lock = threading.Lock()
    
    def __new__(cls) -> "CancellationMonitor":
        with cls._lock:
            if cls._instance is None:
                cls._instance = super().__new__(cls)
                cls._instance._initialized = False
            return cls._instance
    
    def __init__(self) -> None:
        if self._initialized:
            return
        
        self._registrations: Dict[str, Callable[[], None]] = {}
        self._registrations_lock = threading.Lock()
        self._stop_event = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._poll_interval = float(os.getenv(
            ENV_CANCEL_POLL_INTERVAL_SECS,
            str(DEFAULT_POLL_INTERVAL_SECS)
        ))
        self._initialized = True
        
        logger.debug(f"CancellationMonitor initialized (poll interval: {self._poll_interval}s)")
    
    def register(self, optimization_id: str, on_cancelled: Callable[[], None]) -> None:
        """
        Register an optimization for cancellation monitoring.
        
        Args:
            optimization_id: The optimization ID to monitor
            on_cancelled: Callback to invoke when cancellation is detected
        """
        with self._registrations_lock:
            self._registrations[optimization_id] = on_cancelled
            logger.debug(f"Registered optimization '{optimization_id}' for cancellation monitoring")
            
            # Start monitor thread if not running
            if self._thread is None or not self._thread.is_alive():
                self._start_monitor()
    
    def unregister(self, optimization_id: str) -> None:
        """
        Unregister an optimization from cancellation monitoring.
        
        Args:
            optimization_id: The optimization ID to stop monitoring
        """
        with self._registrations_lock:
            if optimization_id in self._registrations:
                del self._registrations[optimization_id]
                logger.debug(f"Unregistered optimization '{optimization_id}' from cancellation monitoring")
            
            # Stop monitor thread if no more registrations
            if not self._registrations and self._thread and self._thread.is_alive():
                self._stop_monitor()
    
    def _start_monitor(self) -> None:
        """Start the background monitoring thread."""
        self._stop_event.clear()
        self._thread = threading.Thread(target=self._monitor_loop, daemon=True, name="CancellationMonitor")
        self._thread.start()
        logger.debug("CancellationMonitor thread started")
    
    def _stop_monitor(self) -> None:
        """Stop the background monitoring thread."""
        self._stop_event.set()
        if self._thread:
            self._thread.join(timeout=5)
            if self._thread.is_alive():
                logger.warning("CancellationMonitor thread did not stop gracefully")
            self._thread = None
        logger.debug("CancellationMonitor thread stopped")
    
    def _monitor_loop(self) -> None:
        """Main monitoring loop - checks all registered optimizations via MGET."""
        redis_client = get_redis_client()
        
        while not self._stop_event.is_set():
            try:
                with self._registrations_lock:
                    if not self._registrations:
                        # No registrations, exit loop
                        break
                    
                    # Get all optimization IDs and their callbacks
                    opt_ids = list(self._registrations.keys())
                    callbacks = dict(self._registrations)
                
                if opt_ids:
                    # Build keys for MGET
                    keys = [CANCEL_KEY_PATTERN.format(opt_id) for opt_id in opt_ids]
                    
                    # Single MGET call for all keys
                    results = redis_client.mget(keys)
                    
                    # Process results
                    for opt_id, key, value in zip(opt_ids, keys, results):
                        if value is not None:
                            # Cancellation detected
                            logger.info(f"Cancellation signal detected for optimization '{opt_id}'")
                            
                            # Get and invoke callback
                            callback = callbacks.get(opt_id)
                            if callback:
                                try:
                                    callback()
                                except Exception as e:
                                    logger.error(f"Error in cancellation callback for '{opt_id}': {e}")
                            
                            # Clean up the key
                            try:
                                redis_client.delete(key)
                                logger.debug(f"Cleaned up cancellation key for '{opt_id}'")
                            except Exception as e:
                                logger.warning(f"Error cleaning up cancel key for '{opt_id}': {e}")
                            
                            # Remove from registrations (don't call unregister to avoid potential deadlock)
                            with self._registrations_lock:
                                self._registrations.pop(opt_id, None)
                
            except Exception as e:
                logger.warning(f"Error in cancellation monitor loop: {e}", exc_info=True)
            
            # Wait for next poll interval
            self._stop_event.wait(self._poll_interval)
        
        logger.debug("CancellationMonitor loop exited")


def get_cancellation_monitor() -> CancellationMonitor:
    """Get the global CancellationMonitor singleton instance."""
    return CancellationMonitor()


class CancellationHandle:
    """
    Handle for a single optimization's cancellation status.
    
    Registers with the global CancellationMonitor and provides a simple
    interface for checking cancellation status.
    
    Usage:
        with CancellationHandle(optimization_id, on_cancelled=my_callback) as handle:
            # ... do work ...
            if handle.was_cancelled:
                # Handle cancellation
    """
    
    def __init__(self, optimization_id: str, on_cancelled: Optional[Callable[[], None]] = None) -> None:
        self.optimization_id: str = optimization_id
        self._cancelled_event: threading.Event = threading.Event()
        self._on_cancelled: Optional[Callable[[], None]] = on_cancelled
        self._registered: bool = False
    
    def __enter__(self) -> "CancellationHandle":
        """Context manager entry - registers with monitor."""
        self.register()
        return self
    
    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[Any],
    ) -> bool:
        """Context manager exit - unregisters from monitor."""
        self.unregister()
        return False  # Don't suppress exceptions
    
    @property
    def was_cancelled(self) -> bool:
        """Thread-safe check if cancellation was triggered."""
        return self._cancelled_event.is_set()
    
    def register(self) -> None:
        """Register this optimization for cancellation monitoring."""
        if self._registered:
            return
        
        def on_cancelled_wrapper():
            self._cancelled_event.set()
            if self._on_cancelled:
                self._on_cancelled()
        
        get_cancellation_monitor().register(self.optimization_id, on_cancelled_wrapper)
        self._registered = True
        logger.debug(f"CancellationHandle registered for '{self.optimization_id}'")
    
    def unregister(self) -> None:
        """Unregister this optimization from cancellation monitoring."""
        if not self._registered:
            return
        
        get_cancellation_monitor().unregister(self.optimization_id)
        self._registered = False
        logger.debug(f"CancellationHandle unregistered for '{self.optimization_id}'")