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

quantumflow/pipeline/checkpoint_manager.py

@@ -0,0 +1,587 @@
+"""
+Checkpoint Manager for Pipeline State Persistence.
+
+Provides:
+- Save/load checkpoints with SHA-256 integrity verification
+- Optional quantum compression for large state vectors
+- Automatic pruning of old checkpoints
+- Database persistence
+
+Example:
+    manager = CheckpointManager(backend="simulator")
+
+    # Save checkpoint
+    manager.save(
+        pipeline_id="...",
+        step=100,
+        state_data={"weights": [...], "energy": -1.5},
+        metrics={"loss": 0.01},
+    )
+
+    # Load checkpoint
+    checkpoint = manager.load(pipeline_id, step=100)
+
+    # Get latest valid checkpoint
+    latest = manager.get_latest_valid(pipeline_id)
+"""
+
+import hashlib
+import json
+import logging
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import uuid
+
+logger = logging.getLogger(__name__)
+
+
+class CheckpointManager:
+    """
+    Manages checkpoint persistence for pipelines.
+
+    Supports:
+    - State serialization with SHA-256 integrity hashes
+    - Optional quantum compression for large states
+    - Automatic pruning to keep N most recent
+    - In-memory cache for fast access
+    """
+
+    def __init__(
+        self,
+        backend: str = "simulator",
+        use_database: bool = True,
+        cache_size: int = 10,
+    ):
+        """
+        Initialize checkpoint manager.
+
+        Args:
+            backend: Quantum backend for compression
+            use_database: Whether to persist to database
+            cache_size: Number of checkpoints to cache in memory
+        """
+        self.backend = backend
+        self.use_database = use_database
+        self.cache_size = cache_size
+
+        # In-memory storage (for non-DB mode or caching)
+        self._checkpoints: Dict[str, Dict[int, Dict[str, Any]]] = {}
+        self._cache: Dict[str, Dict[int, Dict[str, Any]]] = {}
+
+        # Quantum compressor (lazy loaded)
+        self._compressor = None
+
+    def _get_compressor(self):
+        """Get or create quantum compressor."""
+        if self._compressor is None:
+            try:
+                from quantumflow.core.quantum_compressor import QuantumCompressor
+                self._compressor = QuantumCompressor(backend=self.backend)
+            except ImportError:
+                logger.warning("QuantumCompressor not available, compression disabled")
+        return self._compressor
+
+    def _compute_hash(self, data: Dict[str, Any]) -> str:
+        """Compute SHA-256 hash of state data."""
+        serialized = json.dumps(data, sort_keys=True, default=str)
+        return hashlib.sha256(serialized.encode()).hexdigest()
+
+    def _serialize_state(self, state_data: Dict[str, Any]) -> str:
+        """Serialize state data to JSON string."""
+        return json.dumps(state_data, default=self._json_serializer)
+
+    def _deserialize_state(self, state_json: str) -> Dict[str, Any]:
+        """Deserialize JSON string to state data."""
+        return json.loads(state_json)
+
+    @staticmethod
+    def _json_serializer(obj):
+        """Custom JSON serializer for complex types."""
+        if hasattr(obj, "tolist"):  # numpy arrays
+            return obj.tolist()
+        if hasattr(obj, "isoformat"):  # datetime
+            return obj.isoformat()
+        if hasattr(obj, "__dict__"):
+            return obj.__dict__
+        return str(obj)
+
+    def _compress_state(self, state_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """
+        Quantum-compress state data.
+
+        Args:
+            state_data: State to compress
+
+        Returns:
+            Compressed state or None if compression not available
+        """
+        compressor = self._get_compressor()
+        if not compressor:
+            return None
+
+        try:
+            # Extract numeric data for compression
+            vectors = []
+            vector_keys = []
+
+            for key, value in state_data.items():
+                if isinstance(value, (list, tuple)) and all(
+                    isinstance(v, (int, float)) for v in value
+                ):
+                    vectors.extend(value)
+                    vector_keys.append((key, len(value)))
+
+            if not vectors:
+                return None
+
+            # Compress
+            compressed = compressor.compress(vectors)
+
+            return {
+                "amplitudes": compressed.amplitudes.tolist() if hasattr(compressed.amplitudes, "tolist") else list(compressed.amplitudes),
+                "n_qubits": compressed.n_qubits,
+                "vector_keys": vector_keys,
+                "compression_ratio": compressed.compression_percentage / 100,
+            }
+
+        except Exception as e:
+            logger.warning(f"Compression failed: {e}")
+            return None
+
+    def _decompress_state(
+        self, compressed: Dict[str, Any], original_keys: List[tuple]
+    ) -> Dict[str, Any]:
+        """
+        Decompress quantum-compressed state.
+
+        Args:
+            compressed: Compressed state data
+            original_keys: List of (key, length) tuples
+
+        Returns:
+            Decompressed state data
+        """
+        compressor = self._get_compressor()
+        if not compressor:
+            return {}
+
+        try:
+            from quantumflow.core.quantum_compressor import CompressedState
+
+            cs = CompressedState(
+                amplitudes=compressed["amplitudes"],
+                n_qubits=compressed["n_qubits"],
+                original_length=sum(length for _, length in original_keys),
+                compression_percentage=compressed.get("compression_ratio", 0) * 100,
+                input_token_count=0,
+            )
+
+            decompressed = compressor.decompress(cs)
+
+            # Reconstruct dictionary
+            result = {}
+            offset = 0
+            for key, length in original_keys:
+                result[key] = decompressed[offset : offset + length]
+                offset += length
+
+            return result
+
+        except Exception as e:
+            logger.warning(f"Decompression failed: {e}")
+            return {}
+
+    def save(
+        self,
+        pipeline_id: str,
+        step: int,
+        state_data: Dict[str, Any],
+        metrics: Optional[Dict[str, float]] = None,
+        name: Optional[str] = None,
+        use_quantum_compression: bool = False,
+    ) -> str:
+        """
+        Save a checkpoint.
+
+        Args:
+            pipeline_id: Pipeline identifier
+            step: Step number
+            state_data: State data to save
+            metrics: Optional metrics at checkpoint
+            name: Optional checkpoint name
+            use_quantum_compression: Whether to use quantum compression
+
+        Returns:
+            Checkpoint ID
+        """
+        checkpoint_id = str(uuid.uuid4())
+
+        # Compute integrity hash
+        state_hash = self._compute_hash(state_data)
+
+        # Serialize state
+        state_json = self._serialize_state(state_data)
+        state_size = len(state_json.encode())
+
+        # Quantum compression (optional)
+        compressed_state = None
+        compression_ratio = None
+        if use_quantum_compression:
+            compressed = self._compress_state(state_data)
+            if compressed:
+                compressed_state = compressed
+                compression_ratio = compressed.get("compression_ratio")
+
+        checkpoint = {
+            "id": checkpoint_id,
+            "pipeline_id": pipeline_id,
+            "step_number": step,
+            "checkpoint_name": name,
+            "state_data": state_data,
+            "state_hash": state_hash,
+            "compressed_state": compressed_state,
+            "compression_ratio": compression_ratio,
+            "metrics": metrics or {},
+            "is_valid": True,
+            "validation_error": None,
+            "created_at": datetime.utcnow().isoformat(),
+            "state_size_bytes": state_size,
+        }
+
+        # Store in memory
+        if pipeline_id not in self._checkpoints:
+            self._checkpoints[pipeline_id] = {}
+        self._checkpoints[pipeline_id][step] = checkpoint
+
+        # Update cache
+        if pipeline_id not in self._cache:
+            self._cache[pipeline_id] = {}
+        self._cache[pipeline_id][step] = checkpoint
+
+        # Database persistence
+        if self.use_database:
+            self._save_to_database(checkpoint)
+
+        logger.debug(f"Checkpoint saved: pipeline={pipeline_id}, step={step}")
+
+        return checkpoint_id
+
+    def _save_to_database(self, checkpoint: Dict[str, Any]):
+        """Save checkpoint to database."""
+        try:
+            from db.database import get_session
+            from db.models import Checkpoint as CheckpointModel
+
+            with get_session() as session:
+                db_checkpoint = CheckpointModel(
+                    id=uuid.UUID(checkpoint["id"]),
+                    pipeline_id=uuid.UUID(checkpoint["pipeline_id"]),
+                    step_number=checkpoint["step_number"],
+                    checkpoint_name=checkpoint["checkpoint_name"],
+                    state_data=checkpoint["state_data"],
+                    state_hash=checkpoint["state_hash"],
+                    compressed_state=checkpoint["compressed_state"],
+                    compression_ratio=checkpoint["compression_ratio"],
+                    metrics=checkpoint["metrics"],
+                    is_valid=checkpoint["is_valid"],
+                    state_size_bytes=checkpoint["state_size_bytes"],
+                )
+                session.add(db_checkpoint)
+                session.commit()
+
+        except Exception as e:
+            logger.warning(f"Database save failed, using in-memory only: {e}")
+
+    def load(
+        self, pipeline_id: str, step: int, verify_integrity: bool = True
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Load a checkpoint.
+
+        Args:
+            pipeline_id: Pipeline identifier
+            step: Step number
+            verify_integrity: Whether to verify hash integrity
+
+        Returns:
+            Checkpoint data or None if not found
+        """
+        # Check cache first
+        if pipeline_id in self._cache and step in self._cache[pipeline_id]:
+            checkpoint = self._cache[pipeline_id][step]
+            if verify_integrity and not self._verify_integrity(checkpoint):
+                return None
+            return checkpoint
+
+        # Check in-memory storage
+        if pipeline_id in self._checkpoints and step in self._checkpoints[pipeline_id]:
+            checkpoint = self._checkpoints[pipeline_id][step]
+            if verify_integrity and not self._verify_integrity(checkpoint):
+                return None
+            return checkpoint
+
+        # Load from database
+        if self.use_database:
+            checkpoint = self._load_from_database(pipeline_id, step)
+            if checkpoint:
+                if verify_integrity and not self._verify_integrity(checkpoint):
+                    return None
+
+                # Update cache
+                if pipeline_id not in self._cache:
+                    self._cache[pipeline_id] = {}
+                self._cache[pipeline_id][step] = checkpoint
+
+                return checkpoint
+
+        return None
+
+    def _load_from_database(
+        self, pipeline_id: str, step: int
+    ) -> Optional[Dict[str, Any]]:
+        """Load checkpoint from database."""
+        try:
+            from db.database import get_session
+            from db.models import Checkpoint as CheckpointModel
+
+            with get_session() as session:
+                checkpoint = (
+                    session.query(CheckpointModel)
+                    .filter(
+                        CheckpointModel.pipeline_id == uuid.UUID(pipeline_id),
+                        CheckpointModel.step_number == step,
+                    )
+                    .first()
+                )
+
+                if checkpoint:
+                    return {
+                        "id": str(checkpoint.id),
+                        "pipeline_id": str(checkpoint.pipeline_id),
+                        "step_number": checkpoint.step_number,
+                        "checkpoint_name": checkpoint.checkpoint_name,
+                        "state_data": checkpoint.state_data,
+                        "state_hash": checkpoint.state_hash,
+                        "compressed_state": checkpoint.compressed_state,
+                        "compression_ratio": checkpoint.compression_ratio,
+                        "metrics": checkpoint.metrics,
+                        "is_valid": checkpoint.is_valid,
+                        "validation_error": checkpoint.validation_error,
+                        "created_at": checkpoint.created_at.isoformat(),
+                        "state_size_bytes": checkpoint.state_size_bytes,
+                    }
+
+        except Exception as e:
+            logger.warning(f"Database load failed: {e}")
+
+        return None
+
+    def _verify_integrity(self, checkpoint: Dict[str, Any]) -> bool:
+        """Verify checkpoint integrity using hash."""
+        computed_hash = self._compute_hash(checkpoint["state_data"])
+        is_valid = computed_hash == checkpoint["state_hash"]
+
+        if not is_valid:
+            logger.warning(
+                f"Checkpoint integrity check failed: step={checkpoint['step_number']}"
+            )
+
+        return is_valid
+
+    def get_latest_valid(self, pipeline_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get the most recent valid checkpoint.
+
+        Args:
+            pipeline_id: Pipeline identifier
+
+        Returns:
+            Latest valid checkpoint or None
+        """
+        # Get all checkpoints for pipeline
+        checkpoints = self.list_checkpoints(pipeline_id)
+
+        if not checkpoints:
+            return None
+
+        # Sort by step descending
+        sorted_checkpoints = sorted(checkpoints, key=lambda c: c["step_number"], reverse=True)
+
+        # Find first valid one
+        for cp in sorted_checkpoints:
+            if cp["is_valid"] and self._verify_integrity(cp):
+                return cp
+
+        return None
+
+    def list_checkpoints(
+        self, pipeline_id: str, limit: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        List all checkpoints for a pipeline.
+
+        Args:
+            pipeline_id: Pipeline identifier
+            limit: Maximum number to return
+
+        Returns:
+            List of checkpoints sorted by step descending
+        """
+        checkpoints = []
+
+        # From in-memory
+        if pipeline_id in self._checkpoints:
+            checkpoints = list(self._checkpoints[pipeline_id].values())
+
+        # From database if empty
+        if not checkpoints and self.use_database:
+            checkpoints = self._list_from_database(pipeline_id, limit)
+
+        # Sort by step descending
+        checkpoints.sort(key=lambda c: c["step_number"], reverse=True)
+
+        if limit:
+            checkpoints = checkpoints[:limit]
+
+        return checkpoints
+
+    def _list_from_database(
+        self, pipeline_id: str, limit: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
+        """List checkpoints from database."""
+        try:
+            from db.database import get_session
+            from db.models import Checkpoint as CheckpointModel
+
+            with get_session() as session:
+                query = (
+                    session.query(CheckpointModel)
+                    .filter(CheckpointModel.pipeline_id == uuid.UUID(pipeline_id))
+                    .order_by(CheckpointModel.step_number.desc())
+                )
+
+                if limit:
+                    query = query.limit(limit)
+
+                return [
+                    {
+                        "id": str(cp.id),
+                        "pipeline_id": str(cp.pipeline_id),
+                        "step_number": cp.step_number,
+                        "checkpoint_name": cp.checkpoint_name,
+                        "state_data": cp.state_data,
+                        "state_hash": cp.state_hash,
+                        "compressed_state": cp.compressed_state,
+                        "compression_ratio": cp.compression_ratio,
+                        "metrics": cp.metrics,
+                        "is_valid": cp.is_valid,
+                        "validation_error": cp.validation_error,
+                        "created_at": cp.created_at.isoformat(),
+                        "state_size_bytes": cp.state_size_bytes,
+                    }
+                    for cp in query.all()
+                ]
+
+        except Exception as e:
+            logger.warning(f"Database list failed: {e}")
+            return []
+
+    def prune(self, pipeline_id: str, keep_count: int):
+        """
+        Prune old checkpoints, keeping only the most recent N.
+
+        Args:
+            pipeline_id: Pipeline identifier
+            keep_count: Number of checkpoints to keep
+        """
+        checkpoints = self.list_checkpoints(pipeline_id)
+
+        if len(checkpoints) <= keep_count:
+            return
+
+        # Checkpoints to delete (oldest ones)
+        to_delete = checkpoints[keep_count:]
+
+        for cp in to_delete:
+            step = cp["step_number"]
+
+            # Remove from in-memory
+            if pipeline_id in self._checkpoints and step in self._checkpoints[pipeline_id]:
+                del self._checkpoints[pipeline_id][step]
+
+            # Remove from cache
+            if pipeline_id in self._cache and step in self._cache[pipeline_id]:
+                del self._cache[pipeline_id][step]
+
+            # Remove from database
+            if self.use_database:
+                self._delete_from_database(pipeline_id, step)
+
+        logger.debug(f"Pruned {len(to_delete)} old checkpoints for pipeline {pipeline_id}")
+
+    def _delete_from_database(self, pipeline_id: str, step: int):
+        """Delete checkpoint from database."""
+        try:
+            from db.database import get_session
+            from db.models import Checkpoint as CheckpointModel
+
+            with get_session() as session:
+                session.query(CheckpointModel).filter(
+                    CheckpointModel.pipeline_id == uuid.UUID(pipeline_id),
+                    CheckpointModel.step_number == step,
+                ).delete()
+                session.commit()
+
+        except Exception as e:
+            logger.warning(f"Database delete failed: {e}")
+
+    def invalidate(self, pipeline_id: str, step: int, reason: str):
+        """
+        Mark a checkpoint as invalid.
+
+        Args:
+            pipeline_id: Pipeline identifier
+            step: Step number
+            reason: Reason for invalidation
+        """
+        checkpoint = self.load(pipeline_id, step, verify_integrity=False)
+
+        if checkpoint:
+            checkpoint["is_valid"] = False
+            checkpoint["validation_error"] = reason
+
+            # Update in-memory
+            if pipeline_id in self._checkpoints and step in self._checkpoints[pipeline_id]:
+                self._checkpoints[pipeline_id][step] = checkpoint
+
+            # Update database
+            if self.use_database:
+                self._update_validity_in_database(pipeline_id, step, False, reason)
+
+            logger.info(f"Checkpoint invalidated: step={step}, reason={reason}")
+
+    def _update_validity_in_database(
+        self, pipeline_id: str, step: int, is_valid: bool, error: Optional[str]
+    ):
+        """Update checkpoint validity in database."""
+        try:
+            from db.database import get_session
+            from db.models import Checkpoint as CheckpointModel
+
+            with get_session() as session:
+                session.query(CheckpointModel).filter(
+                    CheckpointModel.pipeline_id == uuid.UUID(pipeline_id),
+                    CheckpointModel.step_number == step,
+                ).update({"is_valid": is_valid, "validation_error": error})
+                session.commit()
+
+        except Exception as e:
+            logger.warning(f"Database validity update failed: {e}")
+
+    def clear_cache(self, pipeline_id: Optional[str] = None):
+        """Clear checkpoint cache."""
+        if pipeline_id:
+            if pipeline_id in self._cache:
+                del self._cache[pipeline_id]
+        else:
+            self._cache.clear()

quantumflow/pipeline/finance/__init__.py

@@ -0,0 +1,5 @@
+"""Finance domain pipelines."""
+
+from quantumflow.pipeline.finance.portfolio_optimization import PortfolioOptimizationPipeline
+
+__all__ = ["PortfolioOptimizationPipeline"]