pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/deployment/backends/native_lambda.py

@@ -0,0 +1,125 @@
+"""Native Lambda runner backend — invokes Lambda directly via boto3.
+
+When InfraProvisioner is configured, Lambda functions are auto-provisioned
+from container images at registration time. Otherwise falls back to
+expecting pre-provisioned functions (backward compatible).
+"""
+
+from __future__ import annotations
+
+import json
+import uuid
+from typing import TYPE_CHECKING, Any
+
+import boto3
+import structlog
+
+from pulse_engine.deployment.backends.base import BaseRunnerBackend
+from pulse_engine.extractor.orchestrator.base import OrchestratorRunStatus
+
+if TYPE_CHECKING:
+    from pulse_engine.config import Settings
+
+logger = structlog.get_logger(__name__)
+
+
+class NativeLambdaBackend(BaseRunnerBackend):
+    """Invokes Lambda functions directly via boto3 (no Prefect involved).
+
+    Lambda functions follow the naming convention: {product}-{stage}.
+    If pipeline infrastructure settings are configured, functions are
+    auto-provisioned from container images. Otherwise they must be
+    pre-provisioned by Terraform.
+    """
+
+    def __init__(self, settings: Settings) -> None:
+        self._settings = settings
+        self._region = settings.aws_region
+        self._provisioner: Any = None
+
+    def _get_provisioner(self) -> Any:
+        if self._provisioner is not None:
+            return self._provisioner
+        role_arn = getattr(self._settings, "lambda_execution_role_arn", "")
+        if not isinstance(role_arn, str) or not role_arn:
+            return None
+        try:
+            from pulse_engine.deployment.infra_provisioner import InfraProvisioner
+
+            self._provisioner = InfraProvisioner(
+                region=self._settings.aws_region,
+                pipeline_cluster_name=self._settings.pipeline_cluster_name,
+                pipeline_execution_role_arn=self._settings.pipeline_execution_role_arn,
+                pipeline_task_role_arn=self._settings.pipeline_task_role_arn,
+                pipeline_log_group=self._settings.pipeline_log_group,
+                pipeline_subnets=self._settings.pipeline_subnet_list,
+                pipeline_security_groups=self._settings.pipeline_sg_list,
+                lambda_execution_role_arn=self._settings.lambda_execution_role_arn,
+                lambda_subnets=self._settings.lambda_subnet_list,
+                lambda_security_groups=self._settings.lambda_sg_list,
+                lambda_log_group=self._settings.lambda_log_group,
+            )
+            return self._provisioner
+        except Exception:
+            logger.warning("Failed to create InfraProvisioner", exc_info=True)
+            return None
+
+    async def prepare(self) -> None:
+        pass
+
+    async def register(
+        self,
+        product: str,
+        stage: str,
+        image: str,
+        entrypoint: str | None = None,
+    ) -> str:
+        func_name = f"{product}-{stage}"
+
+        # Auto-provision Lambda function from container image
+        provisioner = self._get_provisioner()
+        if provisioner and image:
+            provisioner.ensure_lambda_function(func_name, image)
+            logger.info(
+                "lambda_auto_provisioned",
+                function=func_name,
+                image=image,
+            )
+
+        return func_name
+
+    async def trigger(self, handle: str, parameters: dict[str, Any]) -> str:
+        from botocore.exceptions import ClientError
+
+        from pulse_engine.core.exceptions import NotFoundError
+
+        run_id = str(uuid.uuid4())
+        payload = {**parameters, "run_id": run_id}
+        client = boto3.client("lambda", region_name=self._region)
+        try:
+            client.invoke(
+                FunctionName=handle,
+                InvocationType="Event",  # async fire-and-forget
+                Payload=json.dumps(payload),
+            )
+        except ClientError as exc:
+            code = exc.response["Error"]["Code"]
+            if code == "ResourceNotFoundException":
+                raise NotFoundError(
+                    f"Lambda function not found: {handle}. "
+                    "Ensure pipeline infrastructure settings are configured "
+                    "or the function is deployed manually.",
+                    function=handle,
+                ) from exc
+            raise
+        logger.info("lambda_triggered", function=handle, run_id=run_id)
+        return run_id
+
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        # Status arrives via Lambda callback to POST /jobs/{id}/status
+        # The engine's job record is the source of truth — not Lambda
+        return OrchestratorRunStatus(run_id=run_id, status="unknown")
+
+    async def cancel_run(self, run_id: str) -> bool:
+        # Async Lambda invocations cannot be cancelled post-fire
+        return False

pulse_engine/deployment/backends/prefect_ecs.py

@@ -0,0 +1,116 @@
+"""Prefect + ECS runner backend."""
+
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING, Any
+
+import httpx
+import structlog
+
+from pulse_engine.deployment.backends.base import BaseRunnerBackend
+from pulse_engine.extractor.orchestrator.base import OrchestratorRunStatus
+
+if TYPE_CHECKING:
+    from pulse_engine.config import Settings
+
+logger = structlog.get_logger(__name__)
+
+_PREFECT_STATE_MAP: dict[str, str] = {
+    "COMPLETED": "completed",
+    "FAILED": "failed",
+    "CRASHED": "failed",
+    "CANCELLED": "cancelled",
+    "CANCELLING": "cancelled",
+    "RUNNING": "running",
+    "PENDING": "pending",
+    "SCHEDULED": "pending",
+}
+
+
+class PrefectECSBackend(BaseRunnerBackend):
+    """Runs Prefect flow runs on an ECS work pool."""
+
+    def __init__(self, settings: Settings) -> None:
+        self._work_pool_name = settings.prefect_ecs_work_pool_name
+        headers: dict[str, str] = {"Content-Type": "application/json"}
+        if settings.prefect_api_key:
+            encoded = base64.b64encode(settings.prefect_api_key.encode()).decode()
+            headers["Authorization"] = f"Basic {encoded}"
+        self._client = httpx.AsyncClient(
+            base_url=settings.prefect_api_url.rstrip("/"),
+            headers=headers,
+            timeout=10.0,
+        )
+
+    async def prepare(self) -> None:
+        resp = await self._client.get(f"/work_pools/{self._work_pool_name}")
+        if resp.is_success:
+            return
+        payload: dict[str, Any] = {
+            "name": self._work_pool_name,
+            "type": "ecs",
+        }
+        resp = await self._client.post("/work_pools/", json=payload)
+        resp.raise_for_status()
+
+    async def register(
+        self,
+        product: str,
+        stage: str,
+        image: str,
+        entrypoint: str | None = None,
+    ) -> str:
+        name = f"{product}-{stage}"
+        flow_id = await self._get_or_create_flow_id(name)
+        payload: dict[str, Any] = {
+            "name": name,
+            "flow_id": flow_id,
+            "entrypoint": entrypoint,
+            "work_pool_name": self._work_pool_name,
+            "job_variables": {"image": image},
+        }
+        resp = await self._client.post("/deployments/", json=payload)
+        resp.raise_for_status()
+        data: dict[str, Any] = resp.json()
+        return str(data["id"])
+
+    async def trigger(self, handle: str, parameters: dict[str, Any]) -> str:
+        resp = await self._client.post(
+            f"/deployments/{handle}/create_flow_run",
+            json={"parameters": parameters},
+        )
+        resp.raise_for_status()
+        data: dict[str, Any] = resp.json()
+        return str(data["id"])
+
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        try:
+            resp = await self._client.get(f"/flow_runs/{run_id}")
+            resp.raise_for_status()
+            data = resp.json()
+            raw_state = data.get("state", {}).get("type", "UNKNOWN")
+            canonical = _PREFECT_STATE_MAP.get(raw_state.upper(), "unknown")
+            return OrchestratorRunStatus(
+                run_id=run_id, status=canonical, raw_state=raw_state
+            )
+        except Exception:
+            logger.warning("ecs_backend_status_failed", run_id=run_id, exc_info=True)
+            return OrchestratorRunStatus(run_id=run_id, status="unknown")
+
+    async def cancel_run(self, run_id: str) -> bool:
+        try:
+            resp = await self._client.post(
+                f"/flow_runs/{run_id}/set_state",
+                json={"state": {"type": "CANCELLED"}},
+            )
+            return resp.is_success
+        except Exception:
+            logger.warning("ecs_backend_cancel_failed", run_id=run_id, exc_info=True)
+            return False
+
+    async def _get_or_create_flow_id(self, flow_name: str) -> str:
+        resp = await self._client.post("/flows/", json={"name": flow_name})
+        resp.raise_for_status()
+        data: dict[str, Any] = resp.json()
+        return str(data["id"])

pulse_engine/deployment/backends/prefect_k8s.py

@@ -0,0 +1,131 @@
+"""Prefect + Kubernetes runner backend."""
+
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING, Any
+
+import httpx
+import structlog
+
+from pulse_engine.deployment.backends.base import BaseRunnerBackend
+from pulse_engine.deployment.backends.prefect_ecs import _PREFECT_STATE_MAP
+from pulse_engine.extractor.orchestrator.base import OrchestratorRunStatus
+
+if TYPE_CHECKING:
+    from pulse_engine.config import Settings
+
+logger = structlog.get_logger(__name__)
+
+
+class PrefectK8sBackend(BaseRunnerBackend):
+    """Runs Prefect flow runs on a Kubernetes work pool.
+
+    A Prefect Kubernetes worker runs inside the cluster (deployed via Terraform/Helm).
+    S3 access uses IRSA — no credentials are injected here.
+    CPU/memory defaults come from settings; overrides can be passed in
+    trigger parameters via config["cpu"] and config["memory"].
+    """
+
+    def __init__(self, settings: Settings) -> None:
+        self._work_pool_name = settings.prefect_k8s_work_pool_name
+        self._namespace = settings.prefect_k8s_namespace
+        self._default_cpu = settings.prefect_k8s_default_cpu
+        self._default_memory = settings.prefect_k8s_default_memory
+        headers: dict[str, str] = {"Content-Type": "application/json"}
+        if settings.prefect_api_key:
+            encoded = base64.b64encode(settings.prefect_api_key.encode()).decode()
+            headers["Authorization"] = f"Basic {encoded}"
+        self._client = httpx.AsyncClient(
+            base_url=settings.prefect_api_url.rstrip("/"),
+            headers=headers,
+            timeout=10.0,
+        )
+
+    async def prepare(self) -> None:
+        resp = await self._client.get(f"/work_pools/{self._work_pool_name}")
+        if resp.is_success:
+            return
+        payload: dict[str, Any] = {
+            "name": self._work_pool_name,
+            "type": "kubernetes",
+        }
+        resp = await self._client.post("/work_pools/", json=payload)
+        resp.raise_for_status()
+
+    async def register(
+        self,
+        product: str,
+        stage: str,
+        image: str,
+        entrypoint: str | None = None,
+    ) -> str:
+        name = f"{product}-{stage}"
+        flow_id = await self._get_or_create_flow_id(name)
+        payload: dict[str, Any] = {
+            "name": name,
+            "flow_id": flow_id,
+            "entrypoint": entrypoint,
+            "work_pool_name": self._work_pool_name,
+            "job_variables": {
+                "image": image,
+                "namespace": self._namespace,
+                "cpu": self._default_cpu,
+                "memory": self._default_memory,
+            },
+        }
+        resp = await self._client.post("/deployments/", json=payload)
+        resp.raise_for_status()
+        dep_data: dict[str, Any] = resp.json()
+        return str(dep_data["id"])
+
+    async def trigger(self, handle: str, parameters: dict[str, Any]) -> str:
+        config: dict[str, Any] = parameters.get("config", {})
+        job_variables: dict[str, Any] = {}
+        if "cpu" in config:
+            job_variables["cpu"] = config["cpu"]
+        if "memory" in config:
+            job_variables["memory"] = config["memory"]
+
+        payload: dict[str, Any] = {"parameters": parameters}
+        if job_variables:
+            payload["job_variables"] = job_variables
+
+        resp = await self._client.post(
+            f"/deployments/{handle}/create_flow_run",
+            json=payload,
+        )
+        resp.raise_for_status()
+        run_data: dict[str, Any] = resp.json()
+        return str(run_data["id"])
+
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        try:
+            resp = await self._client.get(f"/flow_runs/{run_id}")
+            resp.raise_for_status()
+            data = resp.json()
+            raw_state = data.get("state", {}).get("type", "UNKNOWN")
+            canonical = _PREFECT_STATE_MAP.get(raw_state.upper(), "unknown")
+            return OrchestratorRunStatus(
+                run_id=run_id, status=canonical, raw_state=raw_state
+            )
+        except Exception:
+            logger.warning("k8s_backend_status_failed", run_id=run_id, exc_info=True)
+            return OrchestratorRunStatus(run_id=run_id, status="unknown")
+
+    async def cancel_run(self, run_id: str) -> bool:
+        try:
+            resp = await self._client.post(
+                f"/flow_runs/{run_id}/set_state",
+                json={"state": {"type": "CANCELLED"}},
+            )
+            return resp.is_success
+        except Exception:
+            logger.warning("k8s_backend_cancel_failed", run_id=run_id, exc_info=True)
+            return False
+
+    async def _get_or_create_flow_id(self, flow_name: str) -> str:
+        resp = await self._client.post("/flows/", json={"name": flow_name})
+        resp.raise_for_status()
+        flow_data: dict[str, Any] = resp.json()
+        return str(flow_data["id"])

pulse_engine/deployment/backends/registry.py

@@ -0,0 +1,50 @@
+"""Registry mapping (orchestrator, compute) pairs to BaseRunnerBackend instances."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from pulse_engine.deployment.backends.base import BaseRunnerBackend
+from pulse_engine.deployment.backends.exceptions import BackendNotAvailableError
+
+if TYPE_CHECKING:
+    from pulse_engine.config import Settings
+
+
+class BackendRegistry:
+    """Lazy-initializing registry. Backends are instantiated on first get()."""
+
+    def __init__(self, settings: Settings) -> None:
+        self._settings = settings
+        self._instances: dict[tuple[str, str], BaseRunnerBackend] = {}
+
+    def get(self, orchestrator: str, compute: str) -> BaseRunnerBackend:
+        key = (orchestrator, compute)
+        if key not in self._instances:
+            self._instances[key] = self._create(orchestrator, compute)
+        return self._instances[key]
+
+    def available_backends(self) -> list[tuple[str, str]]:
+        return [
+            ("prefect", "ecs"),
+            ("prefect", "kubernetes"),
+            ("native", "lambda"),
+        ]
+
+    def _create(self, orchestrator: str, compute: str) -> BaseRunnerBackend:
+        from pulse_engine.deployment.backends.native_lambda import NativeLambdaBackend
+        from pulse_engine.deployment.backends.prefect_ecs import PrefectECSBackend
+        from pulse_engine.deployment.backends.prefect_k8s import PrefectK8sBackend
+
+        # IMPORTANT: use lambdas so only the matched backend is instantiated.
+        # A plain dict literal would call all three constructors on every _create().
+        factories: dict[tuple[str, str], Callable[[], BaseRunnerBackend]] = {
+            ("prefect", "ecs"): lambda: PrefectECSBackend(self._settings),
+            ("prefect", "kubernetes"): lambda: PrefectK8sBackend(self._settings),
+            ("native", "lambda"): lambda: NativeLambdaBackend(self._settings),
+        }
+        if (orchestrator, compute) not in factories:
+            available = [f"{o}×{c}" for o, c in self.available_backends()]
+            raise BackendNotAvailableError(orchestrator, compute, available)
+        return factories[(orchestrator, compute)]()