morphml 1.0.0__py3-none-any.whl

morphml/distributed/storage/cache.py

@@ -0,0 +1,366 @@
+"""Redis-based distributed cache for fast intermediate storage.
+
+Caches architecture evaluations, optimizer state, and temporary results.
+
+Author: Eshan Roy <eshanized@proton.me>
+Organization: TONMOY INFRASTRUCTURE & VISION
+"""
+
+import pickle
+from typing import Any, Dict, List, Optional
+
+try:
+    import redis
+
+    REDIS_AVAILABLE = True
+except ImportError:
+    REDIS_AVAILABLE = False
+
+from morphml.exceptions import DistributedError
+from morphml.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class DistributedCache:
+    """
+    Redis-based distributed cache.
+
+    Provides fast caching for:
+    - Architecture evaluation results
+    - Optimizer state
+    - Temporary computation results
+    - Worker metadata
+
+    Args:
+        redis_url: Redis connection URL (default: redis://localhost:6379)
+        prefix: Key prefix for namespacing (default: 'morphml')
+        default_ttl: Default time-to-live in seconds (default: None = no expiry)
+
+    Example:
+        >>> cache = DistributedCache('redis://localhost:6379')
+        >>> cache.set('key', {'value': 42}, ttl=3600)
+        >>> result = cache.get('key')
+        >>> cache.cache_architecture_result('abc123', {'fitness': 0.95})
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        prefix: str = "morphml",
+        default_ttl: Optional[int] = None,
+    ):
+        """Initialize distributed cache."""
+        if not REDIS_AVAILABLE:
+            raise DistributedError("Redis not available. Install with: pip install redis")
+
+        self.redis_url = redis_url
+        self.prefix = prefix
+        self.default_ttl = default_ttl
+
+        try:
+            self.client = redis.from_url(redis_url, decode_responses=False)
+            # Test connection
+            self.client.ping()
+            logger.info(f"Connected to Redis: {redis_url}")
+        except redis.ConnectionError as e:
+            raise DistributedError(f"Failed to connect to Redis: {e}")
+
+    def _make_key(self, key: str) -> str:
+        """Add prefix to key."""
+        return f"{self.prefix}:{key}"
+
+    def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
+        """
+        Set value in cache.
+
+        Args:
+            key: Cache key
+            value: Value to cache (will be pickled)
+            ttl: Time-to-live in seconds (None = use default)
+        """
+        full_key = self._make_key(key)
+        serialized = pickle.dumps(value)
+
+        ttl = ttl or self.default_ttl
+
+        if ttl:
+            self.client.setex(full_key, ttl, serialized)
+        else:
+            self.client.set(full_key, serialized)
+
+        logger.debug(f"Cached {key} (ttl={ttl})")
+
+    def get(self, key: str) -> Optional[Any]:
+        """
+        Get value from cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            Cached value or None if not found
+        """
+        full_key = self._make_key(key)
+        value = self.client.get(full_key)
+
+        if value is None:
+            logger.debug(f"Cache miss: {key}")
+            return None
+
+        logger.debug(f"Cache hit: {key}")
+        return pickle.loads(value)
+
+    def delete(self, key: str) -> bool:
+        """
+        Delete key from cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            True if key was deleted
+        """
+        full_key = self._make_key(key)
+        result = self.client.delete(full_key)
+        return result > 0
+
+    def exists(self, key: str) -> bool:
+        """
+        Check if key exists.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            True if key exists
+        """
+        full_key = self._make_key(key)
+        return self.client.exists(full_key) > 0
+
+    def cache_architecture_result(
+        self,
+        arch_hash: str,
+        result: Dict[str, Any],
+        ttl: int = 86400,  # 24 hours
+    ) -> None:
+        """
+        Cache architecture evaluation result.
+
+        Args:
+            arch_hash: Architecture hash
+            result: Evaluation result
+            ttl: Time-to-live (seconds)
+        """
+        key = f"arch:{arch_hash}"
+        self.set(key, result, ttl=ttl)
+
+    def get_architecture_result(self, arch_hash: str) -> Optional[Dict[str, Any]]:
+        """
+        Get cached architecture result.
+
+        Args:
+            arch_hash: Architecture hash
+
+        Returns:
+            Cached result or None
+        """
+        key = f"arch:{arch_hash}"
+        return self.get(key)
+
+    def cache_optimizer_state(
+        self,
+        experiment_id: str,
+        generation: int,
+        state: Dict[str, Any],
+        ttl: int = 3600,  # 1 hour
+    ) -> None:
+        """
+        Cache optimizer state for quick recovery.
+
+        Args:
+            experiment_id: Experiment ID
+            generation: Generation number
+            state: Optimizer state
+            ttl: Time-to-live (seconds)
+        """
+        key = f"optimizer:{experiment_id}:gen{generation}"
+        self.set(key, state, ttl=ttl)
+
+    def get_optimizer_state(self, experiment_id: str, generation: int) -> Optional[Dict[str, Any]]:
+        """
+        Get cached optimizer state.
+
+        Args:
+            experiment_id: Experiment ID
+            generation: Generation number
+
+        Returns:
+            Cached state or None
+        """
+        key = f"optimizer:{experiment_id}:gen{generation}"
+        return self.get(key)
+
+    def invalidate_pattern(self, pattern: str) -> int:
+        """
+        Invalidate all keys matching pattern.
+
+        Args:
+            pattern: Key pattern (supports * wildcard)
+
+        Returns:
+            Number of keys deleted
+        """
+        full_pattern = self._make_key(pattern)
+        keys = list(self.client.scan_iter(match=full_pattern))
+
+        if keys:
+            deleted = self.client.delete(*keys)
+            logger.info(f"Invalidated {deleted} keys matching {pattern}")
+            return deleted
+
+        return 0
+
+    def invalidate_experiment(self, experiment_id: str) -> int:
+        """
+        Invalidate all cache entries for experiment.
+
+        Args:
+            experiment_id: Experiment ID
+
+        Returns:
+            Number of keys deleted
+        """
+        return self.invalidate_pattern(f"*:{experiment_id}:*")
+
+    def get_statistics(self) -> Dict[str, Any]:
+        """
+        Get cache statistics.
+
+        Returns:
+            Statistics dictionary
+        """
+        info = self.client.info("stats")
+
+        return {
+            "total_connections": info.get("total_connections_received", 0),
+            "total_commands": info.get("total_commands_processed", 0),
+            "keyspace_hits": info.get("keyspace_hits", 0),
+            "keyspace_misses": info.get("keyspace_misses", 0),
+            "hit_rate": (
+                info.get("keyspace_hits", 0)
+                / max(
+                    info.get("keyspace_hits", 0) + info.get("keyspace_misses", 0),
+                    1,
+                )
+                * 100
+            ),
+        }
+
+    def clear_all(self) -> None:
+        """
+        Clear all cache entries with prefix.
+
+        Warning: This deletes all keys with the configured prefix!
+        """
+        pattern = f"{self.prefix}:*"
+        keys = list(self.client.scan_iter(match=pattern))
+
+        if keys:
+            self.client.delete(*keys)
+            logger.warning(f"Cleared {len(keys)} cache entries")
+
+    def set_multiple(self, mapping: Dict[str, Any], ttl: Optional[int] = None) -> None:
+        """
+        Set multiple key-value pairs.
+
+        Args:
+            mapping: Dictionary of key-value pairs
+            ttl: Time-to-live for all keys
+        """
+        pipe = self.client.pipeline()
+
+        for key, value in mapping.items():
+            full_key = self._make_key(key)
+            serialized = pickle.dumps(value)
+
+            if ttl:
+                pipe.setex(full_key, ttl, serialized)
+            else:
+                pipe.set(full_key, serialized)
+
+        pipe.execute()
+        logger.debug(f"Cached {len(mapping)} keys")
+
+    def get_multiple(self, keys: List[str]) -> Dict[str, Any]:
+        """
+        Get multiple values.
+
+        Args:
+            keys: List of cache keys
+
+        Returns:
+            Dictionary of found key-value pairs
+        """
+        full_keys = [self._make_key(k) for k in keys]
+        values = self.client.mget(full_keys)
+
+        result = {}
+        for key, value in zip(keys, values):
+            if value is not None:
+                result[key] = pickle.loads(value)
+
+        return result
+
+    def increment(self, key: str, amount: int = 1) -> int:
+        """
+        Increment counter.
+
+        Args:
+            key: Counter key
+            amount: Increment amount
+
+        Returns:
+            New counter value
+        """
+        full_key = self._make_key(key)
+        return self.client.incr(full_key, amount)
+
+    def decrement(self, key: str, amount: int = 1) -> int:
+        """
+        Decrement counter.
+
+        Args:
+            key: Counter key
+            amount: Decrement amount
+
+        Returns:
+            New counter value
+        """
+        full_key = self._make_key(key)
+        return self.client.decr(full_key, amount)
+
+    def get_ttl(self, key: str) -> Optional[int]:
+        """
+        Get remaining TTL for key.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            TTL in seconds, or None if no TTL set
+        """
+        full_key = self._make_key(key)
+        ttl = self.client.ttl(full_key)
+
+        if ttl == -1:  # No TTL
+            return None
+        elif ttl == -2:  # Key doesn't exist
+            return None
+        else:
+            return ttl
+
+    def close(self) -> None:
+        """Close Redis connection."""
+        self.client.close()
+        logger.info("Closed Redis connection")

morphml/distributed/storage/checkpointing.py

@@ -0,0 +1,329 @@
+"""Checkpoint management for experiment recovery.
+
+Enables saving and restoring experiment state for fault tolerance.
+
+Author: Eshan Roy <eshanized@proton.me>
+Organization: TONMOY INFRASTRUCTURE & VISION
+"""
+
+import pickle
+import time
+from typing import Any, Dict, List, Optional
+
+from morphml.core.search import Individual
+from morphml.distributed.storage.artifacts import ArtifactStore
+from morphml.distributed.storage.cache import DistributedCache
+from morphml.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class CheckpointManager:
+    """
+    Manage experiment checkpoints for recovery.
+
+    Uses both cache (fast) and artifact storage (persistent):
+    - Cache: Recent checkpoints for quick recovery
+    - Storage: All checkpoints for long-term persistence
+
+    Args:
+        artifact_store: ArtifactStore for persistent storage
+        cache: DistributedCache for fast access (optional)
+        checkpoint_interval: Save checkpoint every N generations
+        max_checkpoints: Maximum checkpoints to keep (0 = unlimited)
+
+    Example:
+        >>> manager = CheckpointManager(artifact_store, cache)
+        >>>
+        >>> # Save checkpoint
+        >>> manager.save_checkpoint(
+        ...     experiment_id='exp1',
+        ...     generation=10,
+        ...     optimizer_state={'population_size': 50},
+        ...     population=population
+        ... )
+        >>>
+        >>> # Load latest checkpoint
+        >>> checkpoint = manager.load_checkpoint('exp1')
+        >>> if checkpoint:
+        ...     generation = checkpoint['generation']
+        ...     optimizer_state = checkpoint['optimizer_state']
+    """
+
+    def __init__(
+        self,
+        artifact_store: ArtifactStore,
+        cache: Optional[DistributedCache] = None,
+        checkpoint_interval: int = 10,
+        max_checkpoints: int = 5,
+    ):
+        """Initialize checkpoint manager."""
+        self.store = artifact_store
+        self.cache = cache
+        self.checkpoint_interval = checkpoint_interval
+        self.max_checkpoints = max_checkpoints
+
+        logger.info(
+            f"Initialized CheckpointManager "
+            f"(interval={checkpoint_interval}, max={max_checkpoints})"
+        )
+
+    def save_checkpoint(
+        self,
+        experiment_id: str,
+        generation: int,
+        optimizer_state: Dict[str, Any],
+        population: Optional[List[Individual]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """
+        Save experiment checkpoint.
+
+        Args:
+            experiment_id: Experiment identifier
+            generation: Current generation number
+            optimizer_state: Optimizer state dictionary
+            population: Current population (optional)
+            metadata: Additional metadata (optional)
+        """
+        checkpoint = {
+            "experiment_id": experiment_id,
+            "generation": generation,
+            "optimizer_state": optimizer_state,
+            "population": [ind.to_dict() for ind in population] if population else [],
+            "metadata": metadata or {},
+            "timestamp": time.time(),
+        }
+
+        # Save to cache for fast recovery
+        if self.cache:
+            try:
+                self.cache.set(f"checkpoint:{experiment_id}:latest", checkpoint, ttl=3600)
+                self.cache.set(
+                    f"checkpoint:{experiment_id}:gen{generation}",
+                    checkpoint,
+                    ttl=7200,
+                )
+                logger.debug(f"Cached checkpoint for generation {generation}")
+            except Exception as e:
+                logger.warning(f"Failed to cache checkpoint: {e}")
+
+        # Save to persistent storage
+        try:
+            # Serialize checkpoint
+            checkpoint_bytes = pickle.dumps(checkpoint)
+
+            # Upload to S3
+            s3_key = f"checkpoints/{experiment_id}/gen_{generation:06d}.pkl"
+            self.store.upload_bytes(
+                checkpoint_bytes,
+                s3_key,
+                metadata={
+                    "experiment_id": experiment_id,
+                    "generation": str(generation),
+                    "timestamp": str(int(time.time())),
+                },
+            )
+
+            logger.info(f"Saved checkpoint for {experiment_id} generation {generation}")
+
+            # Clean up old checkpoints
+            if self.max_checkpoints > 0:
+                self._cleanup_old_checkpoints(experiment_id)
+
+        except Exception as e:
+            logger.error(f"Failed to save checkpoint: {e}")
+            raise
+
+    def load_checkpoint(
+        self, experiment_id: str, generation: Optional[int] = None
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Load checkpoint.
+
+        Args:
+            experiment_id: Experiment identifier
+            generation: Specific generation (None = latest)
+
+        Returns:
+            Checkpoint dictionary or None if not found
+        """
+        # Try cache first (latest only)
+        if generation is None and self.cache:
+            try:
+                checkpoint = self.cache.get(f"checkpoint:{experiment_id}:latest")
+
+                if checkpoint:
+                    logger.info("Loaded checkpoint from cache (latest)")
+                    return checkpoint
+            except Exception as e:
+                logger.warning(f"Failed to load from cache: {e}")
+
+        # Load from persistent storage
+        try:
+            if generation is not None:
+                # Load specific generation
+                s3_key = f"checkpoints/{experiment_id}/gen_{generation:06d}.pkl"
+
+                if not self.store.exists(s3_key):
+                    logger.warning(f"Checkpoint not found for generation {generation}")
+                    return None
+
+                checkpoint_bytes = self.store.download_bytes(s3_key)
+                checkpoint = pickle.loads(checkpoint_bytes)
+
+                logger.info(f"Loaded checkpoint for {experiment_id} generation {generation}")
+
+                return checkpoint
+
+            else:
+                # Load latest
+                checkpoints = self.list_checkpoints(experiment_id)
+
+                if not checkpoints:
+                    logger.info(f"No checkpoints found for {experiment_id}")
+                    return None
+
+                # Get latest
+                latest = checkpoints[-1]
+                s3_key = latest["key"]
+
+                checkpoint_bytes = self.store.download_bytes(s3_key)
+                checkpoint = pickle.loads(checkpoint_bytes)
+
+                logger.info(
+                    f"Loaded latest checkpoint for {experiment_id} "
+                    f"(generation {checkpoint['generation']})"
+                )
+
+                return checkpoint
+
+        except Exception as e:
+            logger.error(f"Failed to load checkpoint: {e}")
+            return None
+
+    def list_checkpoints(self, experiment_id: str) -> List[Dict[str, Any]]:
+        """
+        List all checkpoints for experiment.
+
+        Args:
+            experiment_id: Experiment identifier
+
+        Returns:
+            List of checkpoint metadata dictionaries
+        """
+        prefix = f"checkpoints/{experiment_id}/"
+
+        try:
+            checkpoints = self.store.list_objects(prefix)
+
+            # Sort by generation (extracted from key)
+            checkpoints.sort(key=lambda x: x["key"])
+
+            return checkpoints
+
+        except Exception as e:
+            logger.error(f"Failed to list checkpoints: {e}")
+            return []
+
+    def delete_checkpoint(self, experiment_id: str, generation: int) -> None:
+        """
+        Delete specific checkpoint.
+
+        Args:
+            experiment_id: Experiment identifier
+            generation: Generation number
+        """
+        s3_key = f"checkpoints/{experiment_id}/gen_{generation:06d}.pkl"
+
+        try:
+            self.store.delete(s3_key)
+            logger.info(f"Deleted checkpoint for generation {generation}")
+
+        except Exception as e:
+            logger.error(f"Failed to delete checkpoint: {e}")
+
+    def delete_all_checkpoints(self, experiment_id: str) -> int:
+        """
+        Delete all checkpoints for experiment.
+
+        Args:
+            experiment_id: Experiment identifier
+
+        Returns:
+            Number of checkpoints deleted
+        """
+        prefix = f"checkpoints/{experiment_id}/"
+
+        try:
+            deleted = self.store.delete_prefix(prefix)
+            logger.info(f"Deleted {deleted} checkpoints for {experiment_id}")
+
+            # Also clear cache
+            if self.cache:
+                self.cache.invalidate_pattern(f"checkpoint:{experiment_id}:*")
+
+            return deleted
+
+        except Exception as e:
+            logger.error(f"Failed to delete checkpoints: {e}")
+            return 0
+
+    def _cleanup_old_checkpoints(self, experiment_id: str) -> None:
+        """Remove old checkpoints keeping only max_checkpoints."""
+        checkpoints = self.list_checkpoints(experiment_id)
+
+        if len(checkpoints) <= self.max_checkpoints:
+            return
+
+        # Delete oldest checkpoints
+        num_to_delete = len(checkpoints) - self.max_checkpoints
+
+        for checkpoint in checkpoints[:num_to_delete]:
+            try:
+                self.store.delete(checkpoint["key"])
+                logger.debug(f"Cleaned up old checkpoint: {checkpoint['key']}")
+            except Exception as e:
+                logger.warning(f"Failed to cleanup checkpoint: {e}")
+
+    def should_checkpoint(self, generation: int) -> bool:
+        """
+        Check if should save checkpoint at this generation.
+
+        Args:
+            generation: Current generation
+
+        Returns:
+            True if should checkpoint
+        """
+        return generation % self.checkpoint_interval == 0
+
+    def get_statistics(self, experiment_id: str) -> Dict[str, Any]:
+        """
+        Get checkpoint statistics.
+
+        Args:
+            experiment_id: Experiment identifier
+
+        Returns:
+            Statistics dictionary
+        """
+        checkpoints = self.list_checkpoints(experiment_id)
+
+        if not checkpoints:
+            return {
+                "experiment_id": experiment_id,
+                "num_checkpoints": 0,
+                "total_size_bytes": 0,
+            }
+
+        total_size = sum(cp["size"] for cp in checkpoints)
+
+        return {
+            "experiment_id": experiment_id,
+            "num_checkpoints": len(checkpoints),
+            "total_size_bytes": total_size,
+            "total_size_mb": total_size / (1024 * 1024),
+            "oldest": checkpoints[0]["last_modified"],
+            "newest": checkpoints[-1]["last_modified"],
+        }