hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/core/orchestration/components/lifecycle_manager.py

@@ -0,0 +1,583 @@
+"""Unified lifecycle manager for pipeline setup and cleanup.
+
+This module consolidates pre-DAG and post-DAG hook management into a single
+component that handles the complete pipeline lifecycle:
+
+- Health checks on adapters
+- Secret injection from KeyVault/SecretPort
+- Custom user hooks
+- Checkpoint saving
+- Secret cleanup (security)
+- Adapter cleanup (connections)
+
+The LifecycleManager replaces the separate PreDagHookManager, PostDagHookManager,
+HealthCheckManager, SecretManager, and AdapterLifecycleManager components.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import MappingProxyType
+
+    from hexdag.core.orchestration.models import NodeExecutionContext
+    from hexdag.core.ports.memory import Memory
+    from hexdag.core.ports.observer_manager import ObserverManagerPort
+    from hexdag.core.ports.secret import SecretPort
+
+from hexdag.core.logging import get_logger
+from hexdag.core.orchestration.events import HealthCheckCompleted
+from hexdag.core.ports.healthcheck import HealthStatus
+from hexdag.core.protocols import HealthCheckable
+
+logger = get_logger(__name__)
+
+# Constants
+DEFAULT_SECRET_PREFIX = "secret:"  # nosec B105 - Not a password, it's a key prefix
+MANAGER_PORT_NAMES = frozenset({"observer_manager"})
+LATENCY_PRECISION = 1
+CLEANUP_METHODS = ["aclose", "ashutdown", "cleanup"]
+
+__all__ = [
+    "HookConfig",
+    "LifecycleManager",
+    "PipelineStatus",
+    "PostDagHookConfig",
+]
+
+
+class PipelineStatus(StrEnum):
+    """Pipeline execution status enumeration."""
+
+    SUCCESS = "success"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+@dataclass(frozen=True, slots=True)
+class HookConfig:
+    """Configuration for pre-DAG hooks.
+
+    Attributes
+    ----------
+    enable_health_checks : bool
+        Run health checks on all adapters before pipeline execution
+    health_check_fail_fast : bool
+        If True, unhealthy adapters block pipeline execution
+    health_check_warn_only : bool
+        If True, log warnings for unhealthy adapters but don't block
+    enable_secret_injection : bool
+        Load secrets from SecretPort into Memory before execution
+    secret_keys : list[str] | None
+        Specific secret keys to load. If None, loads all available secrets.
+    secret_prefix : str
+        Prefix for secret keys in memory (default: "secret:")
+    custom_hooks : list[Callable]
+        User-defined pre-DAG hooks. Each receives (ports, context) and returns Any.
+    """
+
+    enable_health_checks: bool = True
+    health_check_fail_fast: bool = False
+    health_check_warn_only: bool = True
+    enable_secret_injection: bool = True
+    secret_keys: list[str] | None = None
+    secret_prefix: str = DEFAULT_SECRET_PREFIX
+    custom_hooks: list[Callable] = field(default_factory=list)
+
+
+@dataclass(frozen=True, slots=True)
+class PostDagHookConfig:
+    """Configuration for post-DAG hooks.
+
+    Attributes
+    ----------
+    enable_adapter_cleanup : bool
+        Call adapter.aclose() or adapter.ashutdown() if available
+    enable_secret_cleanup : bool
+        Remove secrets from Memory after pipeline execution
+    enable_checkpoint_save : bool
+        Save final checkpoint state
+    checkpoint_on_failure : bool
+        Save checkpoint even if pipeline fails (useful for debugging)
+    custom_hooks : list[Callable]
+        User-defined post-DAG hooks
+    run_on_success : bool
+        Run hooks when pipeline succeeds
+    run_on_failure : bool
+        Run hooks when pipeline fails
+    run_on_cancellation : bool
+        Run hooks when pipeline is cancelled
+    """
+
+    enable_adapter_cleanup: bool = True
+    enable_secret_cleanup: bool = True
+    enable_checkpoint_save: bool = False
+    checkpoint_on_failure: bool = True
+    custom_hooks: list[Callable] = field(default_factory=list)
+    run_on_success: bool = True
+    run_on_failure: bool = True
+    run_on_cancellation: bool = True
+
+
+class LifecycleManager:
+    """Unified manager for pipeline lifecycle: setup, health checks, secrets, cleanup.
+
+    This manager consolidates all pre-DAG and post-DAG operations:
+
+    Pre-execution (pre_execute):
+    1. Health checks on all adapters
+    2. Secret injection from SecretPort into Memory
+    3. Custom user-defined setup hooks
+
+    Post-execution (post_execute):
+    1. Checkpoint saving (if enabled)
+    2. Custom cleanup hooks
+    3. Secret cleanup (CRITICAL - always runs in finally)
+    4. Adapter cleanup (CRITICAL - always runs in finally)
+
+    Examples
+    --------
+    Basic usage::
+
+        manager = LifecycleManager(
+            pre_config=HookConfig(enable_health_checks=True),
+            post_config=PostDagHookConfig(enable_secret_cleanup=True)
+        )
+
+        # Before pipeline execution
+        pre_results = await manager.pre_execute(context, "my_pipeline")
+
+        # ... pipeline runs ...
+
+        # After pipeline execution (always call, even on failure)
+        post_results = await manager.post_execute(
+            context, "my_pipeline", "success", node_results
+        )
+    """
+
+    def __init__(
+        self,
+        pre_config: HookConfig | None = None,
+        post_config: PostDagHookConfig | None = None,
+    ):
+        """Initialize lifecycle manager.
+
+        Parameters
+        ----------
+        pre_config : HookConfig | None
+            Configuration for pre-DAG hooks. Uses defaults if None.
+        post_config : PostDagHookConfig | None
+            Configuration for post-DAG hooks. Uses defaults if None.
+        """
+        self.pre_config = pre_config or HookConfig()
+        self.post_config = post_config or PostDagHookConfig()
+        self._loaded_secret_keys: dict[str, list[str]] = {}  # dag_id -> memory_keys
+
+    # ========================================================================
+    # Pre-execution
+    # ========================================================================
+
+    async def pre_execute(
+        self,
+        context: NodeExecutionContext,
+        pipeline_name: str,
+    ) -> dict[str, Any]:
+        """Execute all pre-DAG lifecycle tasks.
+
+        Parameters
+        ----------
+        context : NodeExecutionContext
+            Execution context for this pipeline run
+        pipeline_name : str
+            Name of the pipeline being executed
+
+        Returns
+        -------
+        dict[str, Any]
+            Results from all pre-execution tasks
+        """
+        from hexdag.core.context import get_observer_manager, get_port, get_ports
+        from hexdag.core.exceptions import OrchestratorError
+
+        results: dict[str, Any] = {}
+        ports: MappingProxyType[str, Any] | dict[Any, Any] = get_ports() or {}
+        observer_manager = get_observer_manager()
+
+        # 1. Health checks
+        if self.pre_config.enable_health_checks:
+            logger.info(f"Running health checks for pipeline '{pipeline_name}'")
+            health_results = await self._check_all_adapters(
+                ports=dict(ports),
+                observer_manager=observer_manager,
+                pipeline_name=pipeline_name,
+            )
+            results["health_checks"] = health_results
+
+            # Check for critical failures
+            if unhealthy := self._get_unhealthy_adapters(health_results):
+                unhealthy_names = [h.adapter_name for h in unhealthy]
+                error_msg = f"Unhealthy adapters: {unhealthy_names}"
+
+                if self.pre_config.health_check_fail_fast:
+                    logger.error(f"Health check failed - blocking pipeline: {error_msg}")
+                    raise OrchestratorError(f"Health check failed: {error_msg}")
+                if self.pre_config.health_check_warn_only:
+                    logger.warning(f"Health check issues detected: {error_msg}")
+                else:
+                    logger.info(f"Health check issues: {error_msg}")
+
+        # 2. Secret injection
+        if self.pre_config.enable_secret_injection:
+            logger.info(f"Loading secrets for pipeline '{pipeline_name}'")
+            secret_port = get_port("secret")
+            memory = get_port("memory")
+            secret_results = await self._load_secrets(
+                secret_port=secret_port,
+                memory=memory,
+                dag_id=context.dag_id,
+            )
+            results["secrets_loaded"] = secret_results
+
+        # 3. Custom hooks
+        for hook in self.pre_config.custom_hooks:
+            hook_name = hook.__name__
+            logger.info(f"Running custom pre-DAG hook: {hook_name}")
+            try:
+                hook_result = await hook(ports, context)
+                results[hook_name] = hook_result
+            except (RuntimeError, ValueError, KeyError, TypeError) as e:
+                logger.error(f"Custom hook '{hook_name}' failed: {e}", exc_info=True)
+                results[hook_name] = {"error": str(e)}
+                raise
+
+        return results
+
+    # ========================================================================
+    # Post-execution
+    # ========================================================================
+
+    async def post_execute(
+        self,
+        context: NodeExecutionContext,
+        pipeline_name: str,
+        pipeline_status: Literal["success", "failed", "cancelled"],
+        node_results: dict[str, Any],
+        error: BaseException | None = None,
+    ) -> dict[str, Any]:
+        """Execute all post-DAG lifecycle tasks.
+
+        Parameters
+        ----------
+        context : NodeExecutionContext
+            Execution context
+        pipeline_name : str
+            Name of the pipeline
+        pipeline_status : Literal["success", "failed", "cancelled"]
+            Final pipeline status
+        node_results : dict[str, Any]
+            Results from all executed nodes
+        error : BaseException | None
+            Exception if pipeline failed
+
+        Returns
+        -------
+        dict[str, Any]
+            Results from all post-execution tasks
+        """
+        from hexdag.core.context import get_observer_manager, get_port, get_ports
+
+        results: dict[str, Any] = {}
+        ports: MappingProxyType[str, Any] | dict[Any, Any] = get_ports() or {}
+        observer_manager = get_observer_manager()
+
+        should_run = (
+            (pipeline_status == "success" and self.post_config.run_on_success)
+            or (pipeline_status == "failed" and self.post_config.run_on_failure)
+            or (pipeline_status == "cancelled" and self.post_config.run_on_cancellation)
+        )
+
+        if not should_run:
+            logger.debug(f"Skipping post-DAG hooks for status: {pipeline_status}")
+            return {"skipped": True, "reason": f"Not configured for {pipeline_status}"}
+
+        logger.info(f"Running post-DAG hooks for pipeline '{pipeline_name}' ({pipeline_status})")
+
+        try:
+            # 1. Save checkpoint (if enabled)
+            if self.post_config.enable_checkpoint_save and (
+                pipeline_status == "success" or self.post_config.checkpoint_on_failure
+            ):
+                try:
+                    checkpoint_result = await self._save_checkpoint(
+                        dict(ports), context, node_results, pipeline_status, observer_manager
+                    )
+                    results["checkpoint"] = checkpoint_result
+                except Exception as e:
+                    logger.error(f"Checkpoint save failed: {e}", exc_info=True)
+                    results["checkpoint"] = {"error": str(e)}
+
+            # 2. Custom hooks (user-defined)
+            for hook in self.post_config.custom_hooks:
+                hook_name = hook.__name__
+                try:
+                    logger.debug(f"Running custom post-DAG hook: {hook_name}")
+                    hook_result = await hook(ports, context, node_results, pipeline_status, error)
+                    results[hook_name] = hook_result
+                except Exception as e:
+                    logger.error(f"Custom hook '{hook_name}' failed: {e}", exc_info=True)
+                    results[hook_name] = {"error": str(e)}
+
+        finally:
+            # CRITICAL CLEANUP: Always run these, even if above hooks fail
+            # 3. Secret cleanup (security - do this before adapter cleanup)
+            if self.post_config.enable_secret_cleanup:
+                try:
+                    memory = get_port("memory")
+                    secret_cleanup = await self._cleanup_secrets(
+                        memory=memory, dag_id=context.dag_id
+                    )
+                    results["secret_cleanup"] = secret_cleanup
+                except Exception as e:
+                    logger.error(f"Secret cleanup failed: {e}", exc_info=True)
+                    results["secret_cleanup"] = {"error": str(e)}
+
+            # 4. Adapter cleanup (close connections - do this last)
+            if self.post_config.enable_adapter_cleanup:
+                try:
+                    adapter_cleanup = await self._cleanup_all_adapters(
+                        ports=dict(ports), observer_manager=observer_manager
+                    )
+                    results["adapter_cleanup"] = adapter_cleanup
+                except Exception as e:
+                    logger.error(f"Adapter cleanup failed: {e}", exc_info=True)
+                    results["adapter_cleanup"] = {"error": str(e)}
+
+        return results
+
+    # ========================================================================
+    # Health Checks (inlined from HealthCheckManager)
+    # ========================================================================
+
+    async def _check_all_adapters(
+        self,
+        ports: dict[str, Any],
+        observer_manager: ObserverManagerPort | None,
+        pipeline_name: str,
+    ) -> list[HealthStatus]:
+        """Run health checks on all adapters that implement ahealth_check()."""
+        health_results = []
+
+        for port_name, adapter in ports.items():
+            if port_name in MANAGER_PORT_NAMES:
+                continue
+
+            if isinstance(adapter, HealthCheckable):
+                status = await self._check_single_adapter(port_name, adapter, observer_manager)
+                health_results.append(status)
+
+        return health_results
+
+    async def _check_single_adapter(
+        self,
+        port_name: str,
+        adapter: Any,
+        observer_manager: ObserverManagerPort | None,
+    ) -> HealthStatus:
+        """Check health of a single adapter."""
+        try:
+            logger.debug(f"Running health check for {port_name}")
+            health_check = adapter.ahealth_check
+            status: HealthStatus = await health_check()
+            status.port_name = port_name
+
+            if observer_manager:
+                event = HealthCheckCompleted(
+                    adapter_name=status.adapter_name,
+                    port_name=port_name,
+                    status=status,
+                )
+                await observer_manager.notify(event)
+
+            self._log_health_result(port_name, status)
+            return status
+
+        except (RuntimeError, ConnectionError, TimeoutError, ValueError) as e:
+            logger.error(f"Health check failed for {port_name}: {e}", exc_info=True)
+            adapter_name = getattr(adapter, "_hexdag_name", port_name)
+            return HealthStatus(
+                status="unhealthy",
+                adapter_name=adapter_name,
+                port_name=port_name,
+                error=e,
+            )
+
+    def _log_health_result(self, port_name: str, status: HealthStatus) -> None:
+        """Log health check result."""
+        if status.status == "healthy":
+            latency_info = (
+                f" ({status.latency_ms:.{LATENCY_PRECISION}f}ms)" if status.latency_ms else ""
+            )
+            logger.info(f"✅ {port_name} health check: {status.status}{latency_info}")
+        else:
+            logger.warning(f"⚠️ {port_name} health check: {status.status} - {status.error}")
+
+    def _get_unhealthy_adapters(self, health_results: list[HealthStatus]) -> list[HealthStatus]:
+        """Filter health results to only unhealthy adapters."""
+        return [h for h in health_results if h.status == "unhealthy"]
+
+    # ========================================================================
+    # Secret Management (inlined from SecretManager)
+    # ========================================================================
+
+    async def _load_secrets(
+        self,
+        secret_port: SecretPort | None,
+        memory: Memory | None,
+        dag_id: str,
+    ) -> dict[str, str]:
+        """Load secrets from SecretPort into Memory."""
+        if not secret_port:
+            logger.debug("No secret port configured, skipping secret injection")
+            return {}
+
+        if not memory:
+            logger.warning("Secret port configured but no memory port available")
+            return {}
+
+        try:
+            mapping = await secret_port.aload_secrets_to_memory(
+                memory=memory,
+                prefix=self.pre_config.secret_prefix,
+                keys=self.pre_config.secret_keys,
+            )
+
+            memory_keys = list(mapping.values())
+            self._loaded_secret_keys[dag_id] = memory_keys
+
+            logger.info(
+                f"Loaded {len(mapping)} secrets into memory with prefix "
+                f"'{self.pre_config.secret_prefix}'"
+            )
+            logger.debug(f"Secret keys loaded: {list(mapping.keys())}")
+
+            return mapping
+
+        except (ValueError, KeyError, RuntimeError) as e:
+            logger.error(f"Failed to inject secrets: {e}", exc_info=True)
+            raise
+
+    async def _cleanup_secrets(
+        self,
+        memory: Memory | None,
+        dag_id: str,
+    ) -> dict[str, Any]:
+        """Remove secrets from Memory for security."""
+        if not memory:
+            logger.debug("No memory port available for secret cleanup")
+            return {"cleaned": False, "reason": "No memory port"}
+
+        secret_keys = self._loaded_secret_keys.get(dag_id, [])
+
+        if not secret_keys:
+            logger.debug("No secrets were loaded for this pipeline")
+            return {"cleaned": True, "keys_removed": 0}
+
+        removed_count = 0
+        for secret_key in secret_keys:
+            try:
+                await memory.aset(secret_key, None)
+                removed_count += 1
+                logger.debug(f"Removed secret from memory: {secret_key}")
+            except (RuntimeError, ValueError, KeyError) as e:
+                logger.warning(f"Failed to remove secret '{secret_key}': {e}")
+
+        if dag_id in self._loaded_secret_keys:
+            del self._loaded_secret_keys[dag_id]
+
+        logger.info(f"Secret cleanup: Removed {removed_count} secret(s) from memory")
+        return {"cleaned": True, "keys_removed": removed_count}
+
+    # ========================================================================
+    # Adapter Cleanup (inlined from AdapterLifecycleManager)
+    # ========================================================================
+
+    async def _cleanup_all_adapters(
+        self,
+        ports: dict[str, Any],
+        observer_manager: ObserverManagerPort | None,
+    ) -> dict[str, Any]:
+        """Close adapter connections and release resources."""
+        cleaned_adapters = []
+
+        for port_name, adapter in ports.items():
+            if port_name in MANAGER_PORT_NAMES:
+                continue
+
+            if await self._cleanup_single_adapter(port_name, adapter):
+                cleaned_adapters.append(port_name)
+
+        return {"cleaned_adapters": cleaned_adapters, "count": len(cleaned_adapters)}
+
+    async def _cleanup_single_adapter(self, port_name: str, adapter: Any) -> bool:
+        """Attempt to clean up a single adapter."""
+        for method_name in CLEANUP_METHODS:
+            if hasattr(adapter, method_name) and callable(getattr(adapter, method_name)):
+                cleanup_method = getattr(adapter, method_name)
+                try:
+                    logger.debug(f"Cleaning up adapter '{port_name}' via {method_name}()")
+                    await cleanup_method()
+                    logger.info(f"✅ Cleaned up adapter: {port_name}")
+                    return True
+                except (RuntimeError, ValueError, TypeError, ConnectionError, OSError) as e:
+                    logger.warning(f"Cleanup failed for {port_name}: {e}")
+                    return False
+
+        return False
+
+    # ========================================================================
+    # Checkpoint (inlined from PostDagHookManager)
+    # ========================================================================
+
+    async def _save_checkpoint(
+        self,
+        ports: dict[str, Any],
+        context: NodeExecutionContext,
+        node_results: dict[str, Any],
+        status: str,
+        observer_manager: ObserverManagerPort | None,
+    ) -> dict[str, Any]:
+        """Save final checkpoint state."""
+        from hexdag.core.context import get_port
+        from hexdag.core.orchestration.components import CheckpointManager
+        from hexdag.core.orchestration.models import CheckpointState
+
+        memory = get_port("memory")
+        if not memory:
+            logger.debug("No memory port available for checkpoint save")
+            return {"skipped": "No memory port available"}
+
+        checkpoint_mgr = CheckpointManager(storage=memory)
+
+        from datetime import UTC, datetime
+
+        state = CheckpointState(
+            run_id=context.dag_id,
+            dag_id=context.dag_id,
+            graph_snapshot={},
+            initial_input=None,
+            node_results=node_results,
+            completed_node_ids=list(node_results.keys()),
+            failed_node_ids=[],
+            created_at=datetime.now(UTC),
+            updated_at=datetime.now(UTC),
+            metadata={"pipeline_status": status},
+        )
+
+        await checkpoint_mgr.save(state)
+        logger.info(f"Saved checkpoint for run_id: {context.dag_id}")
+
+        return {"saved": True, "run_id": context.dag_id, "node_count": len(node_results)}