241 lines
8.3 KiB
Python
241 lines
8.3 KiB
Python
"""
|
|
Checkpoint Manager for Task Recovery
|
|
|
|
Provides partial checkpointing and retry strategies for agent tasks.
|
|
Enables recovery from failures without restarting entire workflows.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Dict, List, Optional, Any
|
|
from dataclasses import dataclass, asdict
|
|
from datetime import datetime
|
|
import redis
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class Checkpoint:
|
|
"""Represents a task checkpoint"""
|
|
task_id: str
|
|
agent_id: str
|
|
checkpoint_id: str
|
|
timestamp: float
|
|
phase: str # e.g., "analysis", "implementation", "testing"
|
|
state: Dict[str, Any] # Serializable state data
|
|
files_modified: List[str]
|
|
git_commit: Optional[str] = None
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
@dataclass
|
|
class RetryPolicy:
|
|
"""Retry strategy configuration"""
|
|
max_retries: int = 3
|
|
backoff_multiplier: float = 2.0
|
|
initial_delay_seconds: float = 1.0
|
|
max_delay_seconds: float = 60.0
|
|
retry_on_errors: List[str] = None # Error types to retry
|
|
|
|
|
|
class CheckpointManager:
|
|
"""Manages task checkpoints and recovery"""
|
|
|
|
def __init__(self, redis_client: redis.Redis):
|
|
self.redis = redis_client
|
|
self.checkpoint_ttl = 86400 * 7 # 7 days
|
|
|
|
def save_checkpoint(self, checkpoint: Checkpoint) -> bool:
|
|
"""Save a checkpoint to Redis"""
|
|
try:
|
|
key = f"checkpoint:{checkpoint.task_id}:{checkpoint.checkpoint_id}"
|
|
data = json.dumps(asdict(checkpoint))
|
|
self.redis.setex(key, self.checkpoint_ttl, data)
|
|
|
|
# Add to task's checkpoint list
|
|
list_key = f"checkpoints:{checkpoint.task_id}"
|
|
self.redis.lpush(list_key, checkpoint.checkpoint_id)
|
|
self.redis.expire(list_key, self.checkpoint_ttl)
|
|
|
|
logger.info(
|
|
f"Saved checkpoint {checkpoint.checkpoint_id} for task {checkpoint.task_id}"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to save checkpoint: {e}")
|
|
return False
|
|
|
|
def get_checkpoint(self, task_id: str, checkpoint_id: str) -> Optional[Checkpoint]:
|
|
"""Retrieve a specific checkpoint"""
|
|
try:
|
|
key = f"checkpoint:{task_id}:{checkpoint_id}"
|
|
data = self.redis.get(key)
|
|
if not data:
|
|
return None
|
|
|
|
checkpoint_dict = json.loads(data)
|
|
return Checkpoint(**checkpoint_dict)
|
|
except Exception as e:
|
|
logger.error(f"Failed to retrieve checkpoint: {e}")
|
|
return None
|
|
|
|
def get_latest_checkpoint(self, task_id: str) -> Optional[Checkpoint]:
|
|
"""Get the most recent checkpoint for a task"""
|
|
try:
|
|
list_key = f"checkpoints:{task_id}"
|
|
checkpoint_ids = self.redis.lrange(list_key, 0, 0)
|
|
|
|
if not checkpoint_ids:
|
|
return None
|
|
|
|
checkpoint_id = checkpoint_ids[0].decode('utf-8')
|
|
return self.get_checkpoint(task_id, checkpoint_id)
|
|
except Exception as e:
|
|
logger.error(f"Failed to get latest checkpoint: {e}")
|
|
return None
|
|
|
|
def list_checkpoints(self, task_id: str) -> List[str]:
|
|
"""List all checkpoint IDs for a task"""
|
|
try:
|
|
list_key = f"checkpoints:{task_id}"
|
|
checkpoint_ids = self.redis.lrange(list_key, 0, -1)
|
|
return [cid.decode('utf-8') for cid in checkpoint_ids]
|
|
except Exception as e:
|
|
logger.error(f"Failed to list checkpoints: {e}")
|
|
return []
|
|
|
|
def delete_checkpoint(self, task_id: str, checkpoint_id: str) -> bool:
|
|
"""Delete a specific checkpoint"""
|
|
try:
|
|
key = f"checkpoint:{task_id}:{checkpoint_id}"
|
|
self.redis.delete(key)
|
|
|
|
# Remove from list
|
|
list_key = f"checkpoints:{task_id}"
|
|
self.redis.lrem(list_key, 0, checkpoint_id)
|
|
|
|
logger.info(f"Deleted checkpoint {checkpoint_id} for task {task_id}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete checkpoint: {e}")
|
|
return False
|
|
|
|
def cleanup_task_checkpoints(self, task_id: str) -> bool:
|
|
"""Delete all checkpoints for a task"""
|
|
try:
|
|
checkpoint_ids = self.list_checkpoints(task_id)
|
|
for checkpoint_id in checkpoint_ids:
|
|
self.delete_checkpoint(task_id, checkpoint_id)
|
|
|
|
list_key = f"checkpoints:{task_id}"
|
|
self.redis.delete(list_key)
|
|
|
|
logger.info(f"Cleaned up all checkpoints for task {task_id}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to cleanup checkpoints: {e}")
|
|
return False
|
|
|
|
|
|
class RetryManager:
|
|
"""Manages retry logic with exponential backoff"""
|
|
|
|
def __init__(self, redis_client: redis.Redis):
|
|
self.redis = redis_client
|
|
|
|
def record_attempt(self, task_id: str, agent_id: str, error: Optional[str] = None):
|
|
"""Record a task attempt"""
|
|
key = f"retry:{task_id}"
|
|
attempt_data = {
|
|
"agent_id": agent_id,
|
|
"timestamp": time.time(),
|
|
"error": error
|
|
}
|
|
|
|
self.redis.lpush(key, json.dumps(attempt_data))
|
|
self.redis.expire(key, 86400) # 24 hour TTL
|
|
|
|
def get_attempt_count(self, task_id: str) -> int:
|
|
"""Get number of attempts for a task"""
|
|
key = f"retry:{task_id}"
|
|
return self.redis.llen(key)
|
|
|
|
def should_retry(self, task_id: str, policy: RetryPolicy) -> bool:
|
|
"""Determine if task should be retried"""
|
|
attempt_count = self.get_attempt_count(task_id)
|
|
return attempt_count < policy.max_retries
|
|
|
|
def get_retry_delay(self, task_id: str, policy: RetryPolicy) -> float:
|
|
"""Calculate delay before next retry (exponential backoff)"""
|
|
attempt_count = self.get_attempt_count(task_id)
|
|
delay = policy.initial_delay_seconds * (policy.backoff_multiplier ** attempt_count)
|
|
return min(delay, policy.max_delay_seconds)
|
|
|
|
def clear_attempts(self, task_id: str):
|
|
"""Clear retry history for a task"""
|
|
key = f"retry:{task_id}"
|
|
self.redis.delete(key)
|
|
|
|
|
|
class RecoveryCoordinator:
|
|
"""Coordinates task recovery from checkpoints"""
|
|
|
|
def __init__(self, checkpoint_manager: CheckpointManager, retry_manager: RetryManager):
|
|
self.checkpoint_manager = checkpoint_manager
|
|
self.retry_manager = retry_manager
|
|
|
|
def recover_task(
|
|
self,
|
|
task_id: str,
|
|
retry_policy: Optional[RetryPolicy] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Attempt to recover a failed task from its latest checkpoint
|
|
|
|
Returns recovery instructions or None if recovery not possible
|
|
"""
|
|
if retry_policy is None:
|
|
retry_policy = RetryPolicy()
|
|
|
|
# Check if we should retry
|
|
if not self.retry_manager.should_retry(task_id, retry_policy):
|
|
logger.warning(f"Task {task_id} exceeded max retries")
|
|
return None
|
|
|
|
# Get latest checkpoint
|
|
checkpoint = self.checkpoint_manager.get_latest_checkpoint(task_id)
|
|
if not checkpoint:
|
|
logger.warning(f"No checkpoint found for task {task_id}")
|
|
return None
|
|
|
|
# Calculate retry delay
|
|
delay = self.retry_manager.get_retry_delay(task_id, retry_policy)
|
|
|
|
recovery_plan = {
|
|
"task_id": task_id,
|
|
"checkpoint_id": checkpoint.checkpoint_id,
|
|
"resume_phase": checkpoint.phase,
|
|
"state": checkpoint.state,
|
|
"files_modified": checkpoint.files_modified,
|
|
"git_commit": checkpoint.git_commit,
|
|
"retry_delay_seconds": delay,
|
|
"attempt_number": self.retry_manager.get_attempt_count(task_id) + 1
|
|
}
|
|
|
|
logger.info(
|
|
f"Recovery plan created for task {task_id} from checkpoint {checkpoint.checkpoint_id}"
|
|
)
|
|
return recovery_plan
|
|
|
|
def mark_recovery_success(self, task_id: str):
|
|
"""Mark a task as successfully recovered"""
|
|
self.retry_manager.clear_attempts(task_id)
|
|
logger.info(f"Task {task_id} recovered successfully")
|
|
|
|
def mark_recovery_failure(self, task_id: str, agent_id: str, error: str):
|
|
"""Record a failed recovery attempt"""
|
|
self.retry_manager.record_attempt(task_id, agent_id, error)
|
|
logger.warning(f"Recovery attempt failed for task {task_id}: {error}")
|