puffinflow 2.dev0__py3-none-any.whl

puffinflow/core/reliability/circuit_breaker.py

@@ -0,0 +1,149 @@
+"""Lightweight circuit breaker implementation."""
+
+import asyncio
+import time
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+
+class CircuitState(Enum):
+    CLOSED = "closed"  # Normal operation
+    OPEN = "open"  # Failing, blocking requests
+    HALF_OPEN = "half_open"  # Testing if service recovered
+
+
+@dataclass
+class CircuitBreakerConfig:
+    failure_threshold: int = 5
+    recovery_timeout: float = 60.0
+    success_threshold: int = 3  # For half-open state
+    timeout: float = 30.0
+    name: str = "default"
+
+
+class CircuitBreakerError(Exception):
+    """Raised when circuit breaker is open"""
+
+    pass
+
+
+class CircuitBreaker:
+    """Lightweight circuit breaker for state execution"""
+
+    def __init__(self, config: CircuitBreakerConfig):
+        self.config = config
+        self.state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._last_failure_time = 0
+        self._lock = asyncio.Lock()
+
+    @asynccontextmanager
+    async def protect(self) -> AsyncGenerator[None, None]:
+        """Context manager for protecting code blocks"""
+        async with self._lock:
+            await self._check_state()
+
+            if self.state == CircuitState.OPEN:
+                raise CircuitBreakerError(
+                    f"Circuit breaker '{self.config.name}' is OPEN"
+                )
+
+        time.time()
+        try:
+            yield
+            await self._record_success()
+        except Exception:
+            await self._record_failure()
+            raise
+
+    async def _check_state(self) -> None:
+        """Check and update circuit breaker state"""
+        now = time.time()
+
+        if (
+            self.state == CircuitState.OPEN
+            and now - self._last_failure_time >= self.config.recovery_timeout
+        ):
+            self.state = CircuitState.HALF_OPEN
+            self._success_count = 0
+
+        elif self.state == CircuitState.HALF_OPEN:
+            if self._success_count >= self.config.success_threshold:
+                self.state = CircuitState.CLOSED
+                self._failure_count = 0
+
+    async def _record_success(self) -> None:
+        """Record successful execution"""
+        async with self._lock:
+            if self.state == CircuitState.HALF_OPEN:
+                self._success_count += 1
+            elif self.state == CircuitState.CLOSED:
+                self._failure_count = max(
+                    0, self._failure_count - 1
+                )  # Gradually recover
+
+    async def _record_failure(self) -> None:
+        """Record failed execution"""
+        async with self._lock:
+            now = time.time()
+            self._failure_count += 1
+            self._last_failure_time = int(now)
+            self._success_count = 0
+
+            if (
+                self.state == CircuitState.CLOSED
+                and self._failure_count >= self.config.failure_threshold
+            ) or self.state == CircuitState.HALF_OPEN:
+                self.state = CircuitState.OPEN
+
+    def get_metrics(self) -> dict[str, Any]:
+        """Get current metrics"""
+        return {
+            "name": self.config.name,
+            "state": self.state.value,
+            "failure_count": self._failure_count,
+            "success_count": self._success_count,
+            "last_failure_time": self._last_failure_time,
+        }
+
+    async def force_open(self) -> None:
+        """Manually open the circuit"""
+        async with self._lock:
+            self.state = CircuitState.OPEN
+            self._last_failure_time = int(time.time())
+
+    async def force_close(self) -> None:
+        """Manually close the circuit"""
+        async with self._lock:
+            self.state = CircuitState.CLOSED
+            self._failure_count = 0
+
+
+# Global circuit breaker registry
+class CircuitBreakerRegistry:
+    """Simple registry for circuit breakers"""
+
+    def __init__(self) -> None:
+        self._breakers: dict[str, CircuitBreaker] = {}
+
+    def get_or_create(
+        self, name: str, config: Optional[CircuitBreakerConfig] = None
+    ) -> CircuitBreaker:
+        """Get existing or create new circuit breaker"""
+        if name not in self._breakers:
+            if config is None:
+                config = CircuitBreakerConfig(name=name)
+            self._breakers[name] = CircuitBreaker(config)
+        return self._breakers[name]
+
+    def get_all_metrics(self) -> dict[str, dict[str, Any]]:
+        """Get metrics for all circuit breakers"""
+        return {name: breaker.get_metrics() for name, breaker in self._breakers.items()}
+
+
+# Global registry instance
+circuit_registry = CircuitBreakerRegistry()

puffinflow/core/reliability/leak_detector.py

@@ -0,0 +1,122 @@
+"""Resource leak detection."""
+
+import time
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass
+class ResourceLeak:
+    state_name: str
+    agent_name: str
+    resources: dict[str, float]
+    allocated_at: float
+    held_for_seconds: float
+    leak_threshold_seconds: float
+
+
+@dataclass
+class ResourceAllocation:
+    state_name: str
+    agent_name: str
+    resources: dict[str, float]
+    allocated_at: float
+
+
+class ResourceLeakDetector:
+    """Detect and handle resource leaks"""
+
+    def __init__(self, leak_threshold_seconds: float = 300.0):  # 5 minutes default
+        self.leak_threshold = leak_threshold_seconds
+        self.allocations: dict[str, ResourceAllocation] = {}
+        self.detected_leaks: list[ResourceLeak] = []
+        self._max_leaks_history = 100
+
+    def track_allocation(
+        self, state_name: str, agent_name: str, resources: dict[str, float]
+    ) -> None:
+        """Track resource allocation"""
+        key = f"{agent_name}:{state_name}"
+        self.allocations[key] = ResourceAllocation(
+            state_name=state_name,
+            agent_name=agent_name,
+            resources=resources.copy(),
+            allocated_at=time.time(),
+        )
+
+    def track_release(self, state_name: str, agent_name: str) -> None:
+        """Track resource release"""
+        key = f"{agent_name}:{state_name}"
+        if key in self.allocations:
+            del self.allocations[key]
+
+    def detect_leaks(self) -> list[ResourceLeak]:
+        """Find resources held too long"""
+        current_leaks = []
+        now = time.time()
+
+        for _key, allocation in self.allocations.items():
+            held_for = now - allocation.allocated_at
+
+            if held_for > self.leak_threshold:
+                leak = ResourceLeak(
+                    state_name=allocation.state_name,
+                    agent_name=allocation.agent_name,
+                    resources=allocation.resources.copy(),
+                    allocated_at=allocation.allocated_at,
+                    held_for_seconds=held_for,
+                    leak_threshold_seconds=self.leak_threshold,
+                )
+                current_leaks.append(leak)
+
+                # Add to history if not already there
+                if not any(
+                    leak_item.state_name == leak.state_name
+                    and leak_item.agent_name == leak.agent_name
+                    and abs(leak_item.allocated_at - leak.allocated_at) < 1.0
+                    for leak_item in self.detected_leaks
+                ):
+                    self.detected_leaks.append(leak)
+
+        # Trim history
+        if len(self.detected_leaks) > self._max_leaks_history:
+            self.detected_leaks = self.detected_leaks[-self._max_leaks_history :]
+
+        return current_leaks
+
+    def get_metrics(self) -> dict[str, Any]:
+        """Get leak detection metrics"""
+        current_leaks = self.detect_leaks()
+
+        return {
+            "total_allocations": len(self.allocations),
+            "current_leaks": len(current_leaks),
+            "total_detected_leaks": len(self.detected_leaks),
+            "leak_threshold_seconds": self.leak_threshold,
+            "oldest_allocation_age": self._get_oldest_allocation_age(),
+            "leaks_by_agent": self._group_leaks_by_agent(current_leaks),
+        }
+
+    def _get_oldest_allocation_age(self) -> Optional[float]:
+        """Get age of oldest allocation"""
+        if not self.allocations:
+            return None
+
+        now = time.time()
+        oldest = min(alloc.allocated_at for alloc in self.allocations.values())
+        return now - oldest
+
+    def _group_leaks_by_agent(self, leaks: list[ResourceLeak]) -> dict[str, int]:
+        """Group leaks by agent"""
+        groups: dict[str, int] = {}
+        for leak in leaks:
+            groups[leak.agent_name] = groups.get(leak.agent_name, 0) + 1
+        return groups
+
+    def clear_leak_history(self) -> None:
+        """Clear leak detection history"""
+        self.detected_leaks.clear()
+
+
+# Global leak detector instance
+leak_detector = ResourceLeakDetector()

puffinflow/core/resources/__init__.py

@@ -0,0 +1,77 @@
+"""Resource management module for workflow orchestrator."""
+
+# Import submodules for import path tests
+from . import allocation, pool, quotas, requirements
+from .allocation import (
+    AllocationRequest,
+    AllocationResult,
+    AllocationStrategy,
+    BestFitAllocator,
+    FairShareAllocator,
+    FirstFitAllocator,
+    PriorityAllocator,
+    ResourceAllocator,
+    WorstFitAllocator,
+)
+from .pool import (
+    ResourceAllocationError,
+    ResourceOverflowError,
+    ResourcePool,
+    ResourceQuotaExceededError,
+    ResourceUsageStats,
+)
+from .quotas import (
+    QuotaExceededError,
+    QuotaLimit,
+    QuotaManager,
+    QuotaMetrics,
+    QuotaPolicy,
+    QuotaScope,
+)
+from .requirements import (
+    ResourceRequirements,
+    ResourceType,
+)
+
+__all__ = [
+    "AllocationRequest",
+    "AllocationResult",
+    # Allocation
+    "AllocationStrategy",
+    "BestFitAllocator",
+    "FairShareAllocator",
+    "FirstFitAllocator",
+    "PriorityAllocator",
+    "QuotaExceededError",
+    "QuotaLimit",
+    # Quotas
+    "QuotaManager",
+    "QuotaMetrics",
+    "QuotaPolicy",
+    "QuotaScope",
+    "ResourceAllocationError",
+    "ResourceAllocator",
+    "ResourceOverflowError",
+    # Pool
+    "ResourcePool",
+    "ResourceQuotaExceededError",
+    "ResourceRequirements",
+    # Requirements
+    "ResourceType",
+    "ResourceUsageStats",
+    "WorstFitAllocator",
+    "allocation",
+    # Submodules
+    "pool",
+    "quotas",
+    "requirements",
+]
+
+# Clean up module namespace
+import sys as _sys
+
+_current_module = _sys.modules[__name__]
+for _attr_name in dir(_current_module):
+    if not _attr_name.startswith("_") and _attr_name not in __all__:
+        delattr(_current_module, _attr_name)
+del _sys, _current_module, _attr_name