smartify-ai 0.1.0__py3-none-any.whl

smartify/state/checkpoint.py

@@ -0,0 +1,558 @@
+"""Checkpoint and resume support for Smartify grid runs.
+
+Provides durable execution by:
+1. Checkpointing node outputs before marking complete
+2. Storing grid spec and scheduler state for resume
+3. Webhook retry queue for failed deliveries
+4. Resume incomplete runs on startup
+"""
+
+import json
+import logging
+import sqlite3
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from enum import Enum
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set
+from uuid import uuid4
+
+logger = logging.getLogger(__name__)
+
+
+class CheckpointStatus(str, Enum):
+    """Status of a checkpoint."""
+    ACTIVE = "active"      # Run in progress, checkpoints being written
+    COMPLETED = "completed" # Run finished successfully
+    FAILED = "failed"       # Run failed
+    ABANDONED = "abandoned" # Run stopped/crashed, can be resumed
+
+
+class WebhookDeliveryStatus(str, Enum):
+    """Status of a webhook delivery."""
+    PENDING = "pending"
+    DELIVERED = "delivered"
+    FAILED = "failed"
+    RETRYING = "retrying"
+
+
+@dataclass
+class Checkpoint:
+    """A checkpoint for a grid run."""
+    checkpoint_id: str
+    run_id: str
+    grid_id: str
+    grid_yaml: str  # Original grid YAML for reconstruction
+    status: CheckpointStatus
+    
+    # Execution state
+    completed_nodes: List[str]  # Node IDs that completed successfully
+    failed_nodes: List[str]     # Node IDs that failed
+    running_nodes: List[str]    # Node IDs that were running at checkpoint
+    
+    # Context state
+    inputs: Dict[str, Any]
+    outputs: Dict[str, Dict[str, Any]]  # node_id -> output
+    total_tokens: int
+    total_cost: float
+    
+    # Timestamps
+    created_at: datetime
+    updated_at: datetime
+    
+    # Resume info
+    resume_count: int = 0
+    last_error: Optional[str] = None
+
+
+@dataclass
+class WebhookRetryJob:
+    """A webhook delivery queued for retry."""
+    job_id: str
+    event_type: str
+    grid_id: str
+    webhook_url: str
+    payload: Dict[str, Any]
+    headers: Dict[str, str]
+    secret: Optional[str]
+    
+    status: WebhookDeliveryStatus
+    attempts: int
+    max_attempts: int
+    
+    created_at: datetime
+    next_retry_at: datetime
+    last_error: Optional[str] = None
+
+
+class CheckpointStore:
+    """SQLite-based checkpoint storage.
+    
+    Extends the base state store with checkpoint-specific tables
+    for durable execution and resume support.
+    """
+    
+    def __init__(self, db_path: str = "smartify_state.db"):
+        self.db_path = db_path
+        self._init_db()
+    
+    def _init_db(self) -> None:
+        """Initialize checkpoint tables."""
+        with sqlite3.connect(self.db_path) as conn:
+            # Checkpoints table - stores run state for resume
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS checkpoints (
+                    checkpoint_id TEXT PRIMARY KEY,
+                    run_id TEXT NOT NULL UNIQUE,
+                    grid_id TEXT NOT NULL,
+                    grid_yaml TEXT NOT NULL,
+                    status TEXT NOT NULL,
+                    completed_nodes TEXT NOT NULL,
+                    failed_nodes TEXT NOT NULL,
+                    running_nodes TEXT NOT NULL,
+                    inputs TEXT NOT NULL,
+                    outputs TEXT NOT NULL,
+                    total_tokens INTEGER DEFAULT 0,
+                    total_cost REAL DEFAULT 0.0,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL,
+                    resume_count INTEGER DEFAULT 0,
+                    last_error TEXT
+                )
+            """)
+            
+            # Webhook retry queue
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS webhook_retry_queue (
+                    job_id TEXT PRIMARY KEY,
+                    event_type TEXT NOT NULL,
+                    grid_id TEXT NOT NULL,
+                    webhook_url TEXT NOT NULL,
+                    payload TEXT NOT NULL,
+                    headers TEXT NOT NULL,
+                    secret TEXT,
+                    status TEXT NOT NULL,
+                    attempts INTEGER DEFAULT 0,
+                    max_attempts INTEGER DEFAULT 3,
+                    created_at TEXT NOT NULL,
+                    next_retry_at TEXT NOT NULL,
+                    last_error TEXT
+                )
+            """)
+            
+            # Indexes
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_checkpoints_status ON checkpoints(status)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_checkpoints_grid ON checkpoints(grid_id)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_webhook_queue_status ON webhook_retry_queue(status)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_webhook_queue_retry ON webhook_retry_queue(next_retry_at)")
+            
+            conn.commit()
+            logger.info(f"Checkpoint store initialized at {self.db_path}")
+    
+    # =========================================================================
+    # Checkpoint Operations
+    # =========================================================================
+    
+    def create_checkpoint(
+        self,
+        run_id: str,
+        grid_id: str,
+        grid_yaml: str,
+        inputs: Dict[str, Any],
+    ) -> Checkpoint:
+        """Create a new checkpoint for a run."""
+        checkpoint = Checkpoint(
+            checkpoint_id=f"ckpt-{uuid4().hex[:12]}",
+            run_id=run_id,
+            grid_id=grid_id,
+            grid_yaml=grid_yaml,
+            status=CheckpointStatus.ACTIVE,
+            completed_nodes=[],
+            failed_nodes=[],
+            running_nodes=[],
+            inputs=inputs,
+            outputs={},
+            total_tokens=0,
+            total_cost=0.0,
+            created_at=datetime.now(),
+            updated_at=datetime.now(),
+        )
+        
+        self._save_checkpoint(checkpoint)
+        logger.info(f"Created checkpoint {checkpoint.checkpoint_id} for run {run_id}")
+        return checkpoint
+    
+    def _save_checkpoint(self, checkpoint: Checkpoint) -> None:
+        """Save checkpoint to database."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute("""
+                INSERT OR REPLACE INTO checkpoints
+                (checkpoint_id, run_id, grid_id, grid_yaml, status,
+                 completed_nodes, failed_nodes, running_nodes,
+                 inputs, outputs, total_tokens, total_cost,
+                 created_at, updated_at, resume_count, last_error)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """, (
+                checkpoint.checkpoint_id,
+                checkpoint.run_id,
+                checkpoint.grid_id,
+                checkpoint.grid_yaml,
+                checkpoint.status.value,
+                json.dumps(checkpoint.completed_nodes),
+                json.dumps(checkpoint.failed_nodes),
+                json.dumps(checkpoint.running_nodes),
+                json.dumps(checkpoint.inputs),
+                json.dumps(checkpoint.outputs),
+                checkpoint.total_tokens,
+                checkpoint.total_cost,
+                checkpoint.created_at.isoformat(),
+                checkpoint.updated_at.isoformat(),
+                checkpoint.resume_count,
+                checkpoint.last_error,
+            ))
+            conn.commit()
+    
+    def get_checkpoint(self, run_id: str) -> Optional[Checkpoint]:
+        """Get checkpoint by run ID."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.execute(
+                "SELECT * FROM checkpoints WHERE run_id = ?", (run_id,)
+            )
+            row = cursor.fetchone()
+            
+            if not row:
+                return None
+            
+            return self._row_to_checkpoint(row)
+    
+    def _row_to_checkpoint(self, row: sqlite3.Row) -> Checkpoint:
+        """Convert database row to Checkpoint."""
+        return Checkpoint(
+            checkpoint_id=row['checkpoint_id'],
+            run_id=row['run_id'],
+            grid_id=row['grid_id'],
+            grid_yaml=row['grid_yaml'],
+            status=CheckpointStatus(row['status']),
+            completed_nodes=json.loads(row['completed_nodes']),
+            failed_nodes=json.loads(row['failed_nodes']),
+            running_nodes=json.loads(row['running_nodes']),
+            inputs=json.loads(row['inputs']),
+            outputs=json.loads(row['outputs']),
+            total_tokens=row['total_tokens'],
+            total_cost=row['total_cost'],
+            created_at=datetime.fromisoformat(row['created_at']),
+            updated_at=datetime.fromisoformat(row['updated_at']),
+            resume_count=row['resume_count'],
+            last_error=row['last_error'],
+        )
+    
+    def checkpoint_node_complete(
+        self,
+        run_id: str,
+        node_id: str,
+        output: Dict[str, Any],
+        tokens_used: int = 0,
+        cost: float = 0.0,
+    ) -> None:
+        """Record a node completion in the checkpoint."""
+        checkpoint = self.get_checkpoint(run_id)
+        if not checkpoint:
+            logger.warning(f"No checkpoint found for run {run_id}")
+            return
+        
+        # Update checkpoint state
+        if node_id not in checkpoint.completed_nodes:
+            checkpoint.completed_nodes.append(node_id)
+        if node_id in checkpoint.running_nodes:
+            checkpoint.running_nodes.remove(node_id)
+        
+        checkpoint.outputs[node_id] = output
+        checkpoint.total_tokens += tokens_used
+        checkpoint.total_cost += cost
+        checkpoint.updated_at = datetime.now()
+        
+        self._save_checkpoint(checkpoint)
+        logger.debug(f"Checkpointed node {node_id} for run {run_id}")
+    
+    def checkpoint_node_failed(
+        self,
+        run_id: str,
+        node_id: str,
+        error: str,
+    ) -> None:
+        """Record a node failure in the checkpoint."""
+        checkpoint = self.get_checkpoint(run_id)
+        if not checkpoint:
+            return
+        
+        if node_id not in checkpoint.failed_nodes:
+            checkpoint.failed_nodes.append(node_id)
+        if node_id in checkpoint.running_nodes:
+            checkpoint.running_nodes.remove(node_id)
+        
+        checkpoint.last_error = f"Node {node_id}: {error}"
+        checkpoint.updated_at = datetime.now()
+        
+        self._save_checkpoint(checkpoint)
+    
+    def checkpoint_node_started(self, run_id: str, node_id: str) -> None:
+        """Record that a node has started execution."""
+        checkpoint = self.get_checkpoint(run_id)
+        if not checkpoint:
+            return
+        
+        if node_id not in checkpoint.running_nodes:
+            checkpoint.running_nodes.append(node_id)
+        checkpoint.updated_at = datetime.now()
+        
+        self._save_checkpoint(checkpoint)
+    
+    def mark_completed(self, run_id: str) -> None:
+        """Mark a run as completed."""
+        checkpoint = self.get_checkpoint(run_id)
+        if checkpoint:
+            checkpoint.status = CheckpointStatus.COMPLETED
+            checkpoint.running_nodes = []
+            checkpoint.updated_at = datetime.now()
+            self._save_checkpoint(checkpoint)
+            logger.info(f"Marked run {run_id} as completed")
+    
+    def mark_failed(self, run_id: str, error: str) -> None:
+        """Mark a run as failed."""
+        checkpoint = self.get_checkpoint(run_id)
+        if checkpoint:
+            checkpoint.status = CheckpointStatus.FAILED
+            checkpoint.running_nodes = []
+            checkpoint.last_error = error
+            checkpoint.updated_at = datetime.now()
+            self._save_checkpoint(checkpoint)
+            logger.info(f"Marked run {run_id} as failed: {error}")
+    
+    def mark_abandoned(self, run_id: str) -> None:
+        """Mark a run as abandoned (can be resumed)."""
+        checkpoint = self.get_checkpoint(run_id)
+        if checkpoint:
+            checkpoint.status = CheckpointStatus.ABANDONED
+            checkpoint.updated_at = datetime.now()
+            self._save_checkpoint(checkpoint)
+    
+    def get_resumable_runs(self) -> List[Checkpoint]:
+        """Get all runs that can be resumed."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.execute(
+                "SELECT * FROM checkpoints WHERE status IN (?, ?)",
+                (CheckpointStatus.ACTIVE.value, CheckpointStatus.ABANDONED.value)
+            )
+            return [self._row_to_checkpoint(row) for row in cursor]
+    
+    def increment_resume_count(self, run_id: str) -> None:
+        """Increment the resume count for a run."""
+        checkpoint = self.get_checkpoint(run_id)
+        if checkpoint:
+            checkpoint.resume_count += 1
+            checkpoint.status = CheckpointStatus.ACTIVE
+            checkpoint.updated_at = datetime.now()
+            self._save_checkpoint(checkpoint)
+    
+    def delete_checkpoint(self, run_id: str) -> bool:
+        """Delete a checkpoint."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.execute(
+                "DELETE FROM checkpoints WHERE run_id = ?", (run_id,)
+            )
+            conn.commit()
+            return cursor.rowcount > 0
+    
+    # =========================================================================
+    # Webhook Retry Queue
+    # =========================================================================
+    
+    def queue_webhook_retry(
+        self,
+        event_type: str,
+        grid_id: str,
+        webhook_url: str,
+        payload: Dict[str, Any],
+        headers: Dict[str, str],
+        secret: Optional[str] = None,
+        max_attempts: int = 3,
+        retry_delay_seconds: float = 60.0,
+    ) -> WebhookRetryJob:
+        """Queue a webhook for retry."""
+        job = WebhookRetryJob(
+            job_id=f"whj-{uuid4().hex[:12]}",
+            event_type=event_type,
+            grid_id=grid_id,
+            webhook_url=webhook_url,
+            payload=payload,
+            headers=headers,
+            secret=secret,
+            status=WebhookDeliveryStatus.PENDING,
+            attempts=0,
+            max_attempts=max_attempts,
+            created_at=datetime.now(),
+            next_retry_at=datetime.now() + timedelta(seconds=retry_delay_seconds),
+        )
+        
+        self._save_webhook_job(job)
+        logger.info(f"Queued webhook retry job {job.job_id} for {webhook_url}")
+        return job
+    
+    def _save_webhook_job(self, job: WebhookRetryJob) -> None:
+        """Save webhook job to database."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute("""
+                INSERT OR REPLACE INTO webhook_retry_queue
+                (job_id, event_type, grid_id, webhook_url, payload, headers,
+                 secret, status, attempts, max_attempts, created_at, next_retry_at, last_error)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """, (
+                job.job_id,
+                job.event_type,
+                job.grid_id,
+                job.webhook_url,
+                json.dumps(job.payload),
+                json.dumps(job.headers),
+                job.secret,
+                job.status.value,
+                job.attempts,
+                job.max_attempts,
+                job.created_at.isoformat(),
+                job.next_retry_at.isoformat(),
+                job.last_error,
+            ))
+            conn.commit()
+    
+    def get_pending_webhook_jobs(self, limit: int = 100) -> List[WebhookRetryJob]:
+        """Get webhook jobs ready for retry."""
+        now = datetime.now().isoformat()
+        
+        with sqlite3.connect(self.db_path) as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.execute("""
+                SELECT * FROM webhook_retry_queue 
+                WHERE status IN (?, ?) AND next_retry_at <= ?
+                ORDER BY next_retry_at ASC
+                LIMIT ?
+            """, (
+                WebhookDeliveryStatus.PENDING.value,
+                WebhookDeliveryStatus.RETRYING.value,
+                now,
+                limit,
+            ))
+            
+            return [self._row_to_webhook_job(row) for row in cursor]
+    
+    def _row_to_webhook_job(self, row: sqlite3.Row) -> WebhookRetryJob:
+        """Convert database row to WebhookRetryJob."""
+        return WebhookRetryJob(
+            job_id=row['job_id'],
+            event_type=row['event_type'],
+            grid_id=row['grid_id'],
+            webhook_url=row['webhook_url'],
+            payload=json.loads(row['payload']),
+            headers=json.loads(row['headers']),
+            secret=row['secret'],
+            status=WebhookDeliveryStatus(row['status']),
+            attempts=row['attempts'],
+            max_attempts=row['max_attempts'],
+            created_at=datetime.fromisoformat(row['created_at']),
+            next_retry_at=datetime.fromisoformat(row['next_retry_at']),
+            last_error=row['last_error'],
+        )
+    
+    def mark_webhook_delivered(self, job_id: str) -> None:
+        """Mark a webhook job as delivered."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute("""
+                UPDATE webhook_retry_queue 
+                SET status = ?, attempts = attempts + 1
+                WHERE job_id = ?
+            """, (WebhookDeliveryStatus.DELIVERED.value, job_id))
+            conn.commit()
+    
+    def mark_webhook_failed(
+        self,
+        job_id: str,
+        error: str,
+        retry_delay_seconds: float = 60.0,
+        backoff_multiplier: float = 2.0,
+    ) -> None:
+        """Mark a webhook job as failed, schedule retry if attempts remain."""
+        job = self._get_webhook_job(job_id)
+        if not job:
+            return
+        
+        job.attempts += 1
+        job.last_error = error
+        
+        if job.attempts >= job.max_attempts:
+            job.status = WebhookDeliveryStatus.FAILED
+            logger.warning(f"Webhook job {job_id} failed permanently after {job.attempts} attempts")
+        else:
+            job.status = WebhookDeliveryStatus.RETRYING
+            # Exponential backoff
+            delay = retry_delay_seconds * (backoff_multiplier ** (job.attempts - 1))
+            job.next_retry_at = datetime.now() + timedelta(seconds=delay)
+            logger.info(f"Webhook job {job_id} retry scheduled in {delay}s")
+        
+        self._save_webhook_job(job)
+    
+    def _get_webhook_job(self, job_id: str) -> Optional[WebhookRetryJob]:
+        """Get a webhook job by ID."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.execute(
+                "SELECT * FROM webhook_retry_queue WHERE job_id = ?", (job_id,)
+            )
+            row = cursor.fetchone()
+            return self._row_to_webhook_job(row) if row else None
+    
+    def cleanup_old_jobs(self, days: int = 7) -> int:
+        """Clean up completed/failed jobs older than N days."""
+        cutoff = (datetime.now() - timedelta(days=days)).isoformat()
+        
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.execute("""
+                DELETE FROM webhook_retry_queue 
+                WHERE status IN (?, ?) AND created_at < ?
+            """, (
+                WebhookDeliveryStatus.DELIVERED.value,
+                WebhookDeliveryStatus.FAILED.value,
+                cutoff,
+            ))
+            conn.commit()
+            
+            deleted = cursor.rowcount
+            if deleted:
+                logger.info(f"Cleaned up {deleted} old webhook jobs")
+            return deleted
+    
+    def get_queue_stats(self) -> Dict[str, int]:
+        """Get webhook queue statistics."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.execute("""
+                SELECT status, COUNT(*) as count 
+                FROM webhook_retry_queue 
+                GROUP BY status
+            """)
+            
+            stats = {status.value: 0 for status in WebhookDeliveryStatus}
+            for row in cursor:
+                stats[row[0]] = row[1]
+            
+            return stats
+
+
+# Singleton instance
+_checkpoint_store: Optional[CheckpointStore] = None
+
+
+def get_checkpoint_store(db_path: str = "smartify_state.db") -> CheckpointStore:
+    """Get or create the checkpoint store singleton."""
+    global _checkpoint_store
+    if _checkpoint_store is None:
+        _checkpoint_store = CheckpointStore(db_path)
+    return _checkpoint_store