quantumflow-sdk 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

quantumflow/pipeline/base_pipeline.py

@@ -0,0 +1,602 @@
+"""
+Base Pipeline with Checkpointing Support.
+
+Abstract base class for all quantum data pipelines with:
+- Configurable checkpointing (time/steps based)
+- State serialization and recovery
+- Integration with anomaly detection
+- Auto-rollback support
+
+Example:
+    class MyPipeline(BasePipeline):
+        def execute_step(self, step: int, state: PipelineState) -> PipelineState:
+            # Custom step logic
+            return state
+
+        def get_state_for_checkpoint(self, state: PipelineState) -> dict:
+            return state.to_dict()
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Callable
+from datetime import datetime
+import time
+import uuid
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PipelineConfig:
+    """Configuration for a pipeline."""
+
+    # Checkpointing
+    checkpoint_interval_steps: int = 10
+    checkpoint_interval_seconds: int = 300  # 5 minutes
+    max_checkpoints: int = 5
+    enable_quantum_compression: bool = False
+
+    # Anomaly detection
+    enable_anomaly_detection: bool = True
+    auto_rollback_on_critical: bool = True
+    gradient_explosion_threshold: float = 100.0
+    gradient_vanishing_threshold: float = 1e-7
+
+    # Execution
+    backend: str = "simulator"
+    max_retries: int = 3
+    retry_delay_seconds: float = 1.0
+
+    # Temporal memory
+    enable_temporal_memory: bool = True
+    memory_window_size: int = 100
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert config to dictionary."""
+        return {
+            "checkpoint_interval_steps": self.checkpoint_interval_steps,
+            "checkpoint_interval_seconds": self.checkpoint_interval_seconds,
+            "max_checkpoints": self.max_checkpoints,
+            "enable_quantum_compression": self.enable_quantum_compression,
+            "enable_anomaly_detection": self.enable_anomaly_detection,
+            "auto_rollback_on_critical": self.auto_rollback_on_critical,
+            "gradient_explosion_threshold": self.gradient_explosion_threshold,
+            "gradient_vanishing_threshold": self.gradient_vanishing_threshold,
+            "backend": self.backend,
+            "max_retries": self.max_retries,
+            "retry_delay_seconds": self.retry_delay_seconds,
+            "enable_temporal_memory": self.enable_temporal_memory,
+            "memory_window_size": self.memory_window_size,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "PipelineConfig":
+        """Create config from dictionary."""
+        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+
+@dataclass
+class PipelineState:
+    """Current state of a pipeline execution."""
+
+    # Core state
+    step: int = 0
+    data: Dict[str, Any] = field(default_factory=dict)
+
+    # Metrics
+    metrics: Dict[str, float] = field(default_factory=dict)
+    gradient_history: List[float] = field(default_factory=list)
+
+    # Algorithm-specific state
+    weights: Optional[List[float]] = None
+    parameters: Optional[Dict[str, Any]] = None
+
+    # History
+    step_history: List[Dict[str, Any]] = field(default_factory=list)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize state to dictionary."""
+        return {
+            "step": self.step,
+            "data": self.data,
+            "metrics": self.metrics,
+            "gradient_history": self.gradient_history,
+            "weights": self.weights,
+            "parameters": self.parameters,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "PipelineState":
+        """Deserialize state from dictionary."""
+        return cls(
+            step=data.get("step", 0),
+            data=data.get("data", {}),
+            metrics=data.get("metrics", {}),
+            gradient_history=data.get("gradient_history", []),
+            weights=data.get("weights"),
+            parameters=data.get("parameters"),
+        )
+
+    def update_metrics(self, **kwargs):
+        """Update metrics."""
+        self.metrics.update(kwargs)
+
+    def add_gradient(self, gradient: float):
+        """Add gradient to history."""
+        self.gradient_history.append(gradient)
+
+
+@dataclass
+class PipelineResult:
+    """Result of pipeline execution."""
+
+    pipeline_id: str
+    status: str  # completed, failed, rolled_back
+    final_state: PipelineState
+    total_steps: int
+    total_time_ms: float
+    checkpoints_created: int
+    anomalies_detected: int
+    rollbacks_performed: int
+    error: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary."""
+        return {
+            "pipeline_id": self.pipeline_id,
+            "status": self.status,
+            "final_state": self.final_state.to_dict(),
+            "total_steps": self.total_steps,
+            "total_time_ms": self.total_time_ms,
+            "checkpoints_created": self.checkpoints_created,
+            "anomalies_detected": self.anomalies_detected,
+            "rollbacks_performed": self.rollbacks_performed,
+            "error": self.error,
+        }
+
+
+class BasePipeline(ABC):
+    """
+    Abstract base class for quantum data pipelines.
+
+    Provides:
+    - Checkpointing at configurable intervals
+    - State serialization/deserialization
+    - Hooks for anomaly detection integration
+    - Auto-rollback on critical failures
+    """
+
+    def __init__(
+        self,
+        name: str,
+        config: Optional[PipelineConfig] = None,
+        pipeline_id: Optional[str] = None,
+    ):
+        """
+        Initialize pipeline.
+
+        Args:
+            name: Human-readable pipeline name
+            config: Pipeline configuration
+            pipeline_id: Optional existing pipeline ID for resumption
+        """
+        self.name = name
+        self.config = config or PipelineConfig()
+        self.pipeline_id = pipeline_id or str(uuid.uuid4())
+
+        # State
+        self._state = PipelineState()
+        self._is_running = False
+        self._is_paused = False
+
+        # Checkpointing
+        self._checkpoint_manager = None
+        self._last_checkpoint_time = 0.0
+        self._last_checkpoint_step = 0
+        self._checkpoints_created = 0
+
+        # Anomaly detection
+        self._anomaly_detector = None
+        self._anomalies_detected = 0
+
+        # Temporal memory
+        self._temporal_memory = None
+        self._run_id = str(uuid.uuid4())
+
+        # Rollback tracking
+        self._rollbacks_performed = 0
+
+        # Callbacks
+        self._on_step_complete: List[Callable] = []
+        self._on_checkpoint: List[Callable] = []
+        self._on_anomaly: List[Callable] = []
+        self._on_rollback: List[Callable] = []
+
+    @property
+    @abstractmethod
+    def pipeline_type(self) -> str:
+        """Return the pipeline type identifier."""
+        pass
+
+    @abstractmethod
+    def execute_step(self, step: int, state: PipelineState) -> PipelineState:
+        """
+        Execute a single pipeline step.
+
+        Args:
+            step: Current step number
+            state: Current pipeline state
+
+        Returns:
+            Updated pipeline state
+        """
+        pass
+
+    @abstractmethod
+    def get_state_for_checkpoint(self, state: PipelineState) -> Dict[str, Any]:
+        """
+        Get state data to save in checkpoint.
+
+        Override to include algorithm-specific state.
+
+        Args:
+            state: Current pipeline state
+
+        Returns:
+            Serializable state dictionary
+        """
+        pass
+
+    @abstractmethod
+    def restore_state_from_checkpoint(self, checkpoint_data: Dict[str, Any]) -> PipelineState:
+        """
+        Restore state from checkpoint data.
+
+        Args:
+            checkpoint_data: Saved checkpoint data
+
+        Returns:
+            Restored pipeline state
+        """
+        pass
+
+    def initialize(self) -> PipelineState:
+        """
+        Initialize pipeline state before execution.
+
+        Override for custom initialization.
+
+        Returns:
+            Initial pipeline state
+        """
+        return PipelineState()
+
+    def finalize(self, state: PipelineState) -> PipelineState:
+        """
+        Finalize pipeline after execution.
+
+        Override for custom finalization.
+
+        Args:
+            state: Final pipeline state
+
+        Returns:
+            Finalized state
+        """
+        return state
+
+    def should_stop(self, state: PipelineState) -> bool:
+        """
+        Check if pipeline should stop execution.
+
+        Override for custom stopping conditions.
+
+        Args:
+            state: Current pipeline state
+
+        Returns:
+            True if pipeline should stop
+        """
+        return False
+
+    def set_checkpoint_manager(self, manager: "CheckpointManager"):
+        """Set checkpoint manager for persistence."""
+        self._checkpoint_manager = manager
+
+    def set_anomaly_detector(self, detector: "AnomalyDetector"):
+        """Set anomaly detector."""
+        self._anomaly_detector = detector
+
+    def set_temporal_memory(self, memory: "TemporalMemoryStore"):
+        """Set temporal memory store."""
+        self._temporal_memory = memory
+
+    def on_step_complete(self, callback: Callable[[int, PipelineState], None]):
+        """Register step completion callback."""
+        self._on_step_complete.append(callback)
+
+    def on_checkpoint(self, callback: Callable[[int, Dict[str, Any]], None]):
+        """Register checkpoint callback."""
+        self._on_checkpoint.append(callback)
+
+    def on_anomaly(self, callback: Callable[[str, Dict[str, Any]], None]):
+        """Register anomaly callback."""
+        self._on_anomaly.append(callback)
+
+    def on_rollback(self, callback: Callable[[int, int], None]):
+        """Register rollback callback (from_step, to_step)."""
+        self._on_rollback.append(callback)
+
+    def _should_checkpoint(self, step: int) -> bool:
+        """Check if checkpoint should be created."""
+        # Step-based
+        if step - self._last_checkpoint_step >= self.config.checkpoint_interval_steps:
+            return True
+
+        # Time-based
+        current_time = time.time()
+        if current_time - self._last_checkpoint_time >= self.config.checkpoint_interval_seconds:
+            return True
+
+        return False
+
+    def _create_checkpoint(self, step: int, state: PipelineState):
+        """Create checkpoint at current step."""
+        if not self._checkpoint_manager:
+            return
+
+        checkpoint_data = self.get_state_for_checkpoint(state)
+
+        self._checkpoint_manager.save(
+            pipeline_id=self.pipeline_id,
+            step=step,
+            state_data=checkpoint_data,
+            metrics=state.metrics.copy(),
+            use_quantum_compression=self.config.enable_quantum_compression,
+        )
+
+        self._last_checkpoint_step = step
+        self._last_checkpoint_time = time.time()
+        self._checkpoints_created += 1
+
+        # Prune old checkpoints
+        self._checkpoint_manager.prune(self.pipeline_id, self.config.max_checkpoints)
+
+        # Notify callbacks
+        for callback in self._on_checkpoint:
+            callback(step, checkpoint_data)
+
+        logger.info(f"Checkpoint created at step {step}")
+
+    def _check_anomalies(self, state: PipelineState) -> Optional[Dict[str, Any]]:
+        """Check for anomalies in current state."""
+        if not self.config.enable_anomaly_detection:
+            return None
+
+        if not self._anomaly_detector:
+            return None
+
+        result = self._anomaly_detector.detect(
+            state=state,
+            step=state.step,
+            gradient_threshold=self.config.gradient_explosion_threshold,
+            vanishing_threshold=self.config.gradient_vanishing_threshold,
+        )
+
+        if result and result.is_anomaly:
+            self._anomalies_detected += 1
+
+            # Notify callbacks
+            for callback in self._on_anomaly:
+                callback(result.anomaly_type, result.to_dict())
+
+            logger.warning(f"Anomaly detected: {result.anomaly_type} at step {state.step}")
+
+            return result.to_dict()
+
+        return None
+
+    def _rollback_to_checkpoint(self, checkpoint_step: Optional[int] = None) -> Optional[PipelineState]:
+        """
+        Rollback to a previous checkpoint.
+
+        Args:
+            checkpoint_step: Specific step to rollback to (None = most recent valid)
+
+        Returns:
+            Restored state or None if rollback failed
+        """
+        if not self._checkpoint_manager:
+            logger.error("Cannot rollback: no checkpoint manager")
+            return None
+
+        # Find checkpoint
+        if checkpoint_step is not None:
+            checkpoint = self._checkpoint_manager.load(self.pipeline_id, checkpoint_step)
+        else:
+            checkpoint = self._checkpoint_manager.get_latest_valid(self.pipeline_id)
+
+        if not checkpoint:
+            logger.error("No valid checkpoint found for rollback")
+            return None
+
+        # Restore state
+        restored_state = self.restore_state_from_checkpoint(checkpoint["state_data"])
+
+        from_step = self._state.step
+        to_step = checkpoint["step_number"]
+
+        self._state = restored_state
+        self._rollbacks_performed += 1
+
+        # Notify callbacks
+        for callback in self._on_rollback:
+            callback(from_step, to_step)
+
+        logger.info(f"Rolled back from step {from_step} to step {to_step}")
+
+        return restored_state
+
+    def _store_temporal_state(self, state: PipelineState):
+        """Store state in temporal memory."""
+        if not self.config.enable_temporal_memory:
+            return
+
+        if not self._temporal_memory:
+            return
+
+        # Create state vector from metrics and gradients
+        state_vector = list(state.metrics.values())
+        if state.gradient_history:
+            state_vector.extend(state.gradient_history[-10:])  # Last 10 gradients
+
+        if state_vector:
+            self._temporal_memory.store(
+                pipeline_id=self.pipeline_id,
+                run_id=self._run_id,
+                sequence_number=state.step,
+                state_vector=state_vector,
+                metadata={"step": state.step, "metrics": state.metrics},
+            )
+
+    def run(self, total_steps: int, initial_state: Optional[PipelineState] = None) -> PipelineResult:
+        """
+        Execute the pipeline.
+
+        Args:
+            total_steps: Total number of steps to execute
+            initial_state: Optional initial state (for resumption)
+
+        Returns:
+            Pipeline execution result
+        """
+        start_time = time.time()
+        self._is_running = True
+        error = None
+        status = "completed"
+
+        try:
+            # Initialize
+            self._state = initial_state or self.initialize()
+            self._last_checkpoint_time = time.time()
+            self._last_checkpoint_step = self._state.step
+
+            logger.info(f"Starting pipeline {self.name} from step {self._state.step}")
+
+            # Execute steps
+            while self._state.step < total_steps:
+                if self._is_paused:
+                    time.sleep(0.1)
+                    continue
+
+                step = self._state.step
+
+                # Execute step
+                try:
+                    self._state = self.execute_step(step, self._state)
+                    self._state.step = step + 1
+                except Exception as e:
+                    logger.error(f"Step {step} failed: {e}")
+
+                    if self.config.auto_rollback_on_critical:
+                        restored = self._rollback_to_checkpoint()
+                        if restored:
+                            continue
+
+                    raise
+
+                # Notify step completion
+                for callback in self._on_step_complete:
+                    callback(step, self._state)
+
+                # Check for anomalies
+                anomaly = self._check_anomalies(self._state)
+                if anomaly and anomaly.get("severity") == "critical":
+                    if self.config.auto_rollback_on_critical:
+                        restored = self._rollback_to_checkpoint()
+                        if restored:
+                            continue
+                        else:
+                            status = "failed"
+                            error = f"Critical anomaly and rollback failed: {anomaly}"
+                            break
+
+                # Store temporal state
+                self._store_temporal_state(self._state)
+
+                # Create checkpoint if needed
+                if self._should_checkpoint(self._state.step):
+                    self._create_checkpoint(self._state.step, self._state)
+
+                # Check stopping condition
+                if self.should_stop(self._state):
+                    logger.info(f"Pipeline stopped early at step {self._state.step}")
+                    break
+
+            # Finalize
+            self._state = self.finalize(self._state)
+
+            # Final checkpoint
+            if self._checkpoint_manager:
+                self._create_checkpoint(self._state.step, self._state)
+
+        except Exception as e:
+            status = "failed"
+            error = str(e)
+            logger.error(f"Pipeline failed: {e}")
+
+            # Attempt rollback on failure
+            if self.config.auto_rollback_on_critical:
+                restored = self._rollback_to_checkpoint()
+                if restored:
+                    status = "rolled_back"
+
+        finally:
+            self._is_running = False
+
+        end_time = time.time()
+        total_time_ms = (end_time - start_time) * 1000
+
+        return PipelineResult(
+            pipeline_id=self.pipeline_id,
+            status=status,
+            final_state=self._state,
+            total_steps=self._state.step,
+            total_time_ms=total_time_ms,
+            checkpoints_created=self._checkpoints_created,
+            anomalies_detected=self._anomalies_detected,
+            rollbacks_performed=self._rollbacks_performed,
+            error=error,
+        )
+
+    def pause(self):
+        """Pause pipeline execution."""
+        self._is_paused = True
+        logger.info(f"Pipeline {self.name} paused")
+
+    def resume(self):
+        """Resume pipeline execution."""
+        self._is_paused = False
+        logger.info(f"Pipeline {self.name} resumed")
+
+    def stop(self):
+        """Stop pipeline execution."""
+        self._is_running = False
+        logger.info(f"Pipeline {self.name} stopped")
+
+    @property
+    def state(self) -> PipelineState:
+        """Get current pipeline state."""
+        return self._state
+
+    @property
+    def is_running(self) -> bool:
+        """Check if pipeline is running."""
+        return self._is_running
+
+    @property
+    def is_paused(self) -> bool:
+        """Check if pipeline is paused."""
+        return self._is_paused