pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/deployment/infra_provisioner.py

@@ -0,0 +1,285 @@
+"""Dynamic AWS infrastructure provisioner for pipeline step execution.
+
+The engine auto-provisions ECS task definitions, Lambda functions, and ECR
+repositories so that products can execute immediately after registration
+without requiring manual Terraform runs for each new product/module.
+
+Prerequisites (created by Terraform once):
+  - ECS cluster, execution role, task role, security group (pipeline_cluster module)
+  - Lambda execution role, security group (lambda_pipeline module)
+  - VPC with private subnets (networking module)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import boto3
+from botocore.exceptions import ClientError
+
+logger = logging.getLogger(__name__)
+
+
+class InfraProvisioner:
+    """Auto-provisions AWS infrastructure for pipeline step containers.
+
+    Uses IAM roles and networking created by Terraform. Creates the dynamic
+    resources (task definitions, Lambda functions, ECR repos) on demand.
+    """
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        pipeline_cluster_name: str,
+        pipeline_execution_role_arn: str,
+        pipeline_task_role_arn: str,
+        pipeline_log_group: str,
+        pipeline_subnets: list[str],
+        pipeline_security_groups: list[str],
+        lambda_execution_role_arn: str,
+        lambda_subnets: list[str],
+        lambda_security_groups: list[str],
+        lambda_log_group: str,
+    ) -> None:
+        self._region = region
+        self._cluster = pipeline_cluster_name
+        self._exec_role_arn = pipeline_execution_role_arn
+        self._task_role_arn = pipeline_task_role_arn
+        self._log_group = pipeline_log_group
+        self._subnets = pipeline_subnets
+        self._security_groups = pipeline_security_groups
+        self._lambda_role_arn = lambda_execution_role_arn
+        self._lambda_subnets = lambda_subnets
+        self._lambda_sgs = lambda_security_groups
+        self._lambda_log_group = lambda_log_group
+
+        # Cache to avoid redundant API calls within same process lifetime
+        self._ecs_task_def_cache: dict[str, str] = {}
+        self._lambda_fn_cache: set[str] = set()
+        self._ecr_repo_cache: dict[str, str] = {}
+
+    # ------------------------------------------------------------------
+    # ECS Task Definition
+    # ------------------------------------------------------------------
+
+    def ensure_ecs_task_definition(
+        self,
+        family: str = "pulse-pipeline-step",
+        cpu: str = "256",
+        memory: str = "512",
+        image: str | None = None,
+    ) -> str:
+        """Create or update a Fargate task definition for pipeline steps.
+
+        When image is provided, registers a new revision with that image.
+        When image is None, reuses existing revision or creates one with
+        a placeholder image.
+
+        Returns the task definition ARN.
+        """
+        cache_key = f"{family}:{cpu}:{memory}:{image or 'default'}"
+        if cache_key in self._ecs_task_def_cache:
+            return self._ecs_task_def_cache[cache_key]
+
+        ecs = boto3.client("ecs", region_name=self._region)
+
+        # If no specific image requested, reuse existing task def
+        if image is None:
+            try:
+                resp = ecs.describe_task_definition(taskDefinition=family)
+                existing = resp["taskDefinition"]
+                result_arn = str(existing["taskDefinitionArn"])
+                logger.info(
+                    "ECS task definition already exists: %s (rev %s)",
+                    family,
+                    existing["revision"],
+                )
+                self._ecs_task_def_cache[cache_key] = result_arn
+                return result_arn
+            except ClientError as e:
+                if e.response["Error"]["Code"] != "ClientException":
+                    raise
+
+        # Register new task definition with the specified (or placeholder) image
+        container_def: dict[str, Any] = {
+            "name": "step",
+            "image": image or "amazon/amazon-ecs-sample",
+            "essential": True,
+            # Set entryPoint to "python -c" so that the ECS
+            # containerOverrides.command becomes the script body.
+            # This overrides any image ENTRYPOINT (e.g. Lambda images
+            # with awslambdaric) ensuring the pipeline runner script
+            # executes instead of the Lambda runtime.
+            "entryPoint": ["python", "-c"],
+            "command": ["print('task-definition placeholder')"],
+            "logConfiguration": {
+                "logDriver": "awslogs",
+                "options": {
+                    "awslogs-group": self._log_group,
+                    "awslogs-region": self._region,
+                    "awslogs-stream-prefix": "step",
+                },
+            },
+        }
+
+        resp = ecs.register_task_definition(
+            family=family,
+            networkMode="awsvpc",
+            requiresCompatibilities=["FARGATE"],
+            cpu=cpu,
+            memory=memory,
+            executionRoleArn=self._exec_role_arn,
+            taskRoleArn=self._task_role_arn,
+            containerDefinitions=[container_def],
+        )
+
+        result_arn = str(resp["taskDefinition"]["taskDefinitionArn"])
+        logger.info("Registered ECS task definition: %s (image=%s)", result_arn, image)
+        self._ecs_task_def_cache[cache_key] = result_arn
+        return result_arn
+
+    # ------------------------------------------------------------------
+    # Lambda Function
+    # ------------------------------------------------------------------
+
+    def ensure_lambda_function(
+        self,
+        function_name: str,
+        image_uri: str,
+        timeout: int = 900,
+        memory_size: int = 3008,
+    ) -> str:
+        """Create or update a Lambda function backed by a container image.
+
+        Always waits for the function to be Active before returning.
+        Returns the function name.
+        """
+        if function_name in self._lambda_fn_cache:
+            return function_name
+
+        lam = boto3.client("lambda", region_name=self._region)
+
+        try:
+            lam.get_function(FunctionName=function_name)
+            # Function exists — update code to latest image
+            self._update_lambda_code(function_name, image_uri)
+            logger.info("Lambda function exists, updated code: %s", function_name)
+        except ClientError as e:
+            if e.response["Error"]["Code"] != "ResourceNotFoundException":
+                raise
+            # Create the function
+            lam.create_function(
+                FunctionName=function_name,
+                Role=self._lambda_role_arn,
+                Code={"ImageUri": image_uri},
+                PackageType="Image",
+                Timeout=timeout,
+                MemorySize=memory_size,
+                VpcConfig={
+                    "SubnetIds": self._lambda_subnets,
+                    "SecurityGroupIds": self._lambda_sgs,
+                },
+                Environment={"Variables": {}},
+            )
+            logger.info("Created Lambda function: %s", function_name)
+
+        # Always wait for function to be Active (handles both create and update)
+        waiter = lam.get_waiter("function_active_v2")
+        waiter.wait(FunctionName=function_name, WaiterConfig={"MaxAttempts": 60})
+        logger.info("Lambda function active: %s", function_name)
+
+        self._lambda_fn_cache.add(function_name)
+        return function_name
+
+    def _update_lambda_code(self, function_name: str, image_uri: str) -> None:
+        """Update Lambda function code to a new container image."""
+        lam = boto3.client("lambda", region_name=self._region)
+        try:
+            lam.update_function_code(FunctionName=function_name, ImageUri=image_uri)
+        except ClientError:
+            logger.warning(
+                "Failed to update Lambda code for %s", function_name, exc_info=True
+            )
+
+    # ------------------------------------------------------------------
+    # ECR Repository
+    # ------------------------------------------------------------------
+
+    def ensure_ecr_repository(self, repo_name: str) -> str:
+        """Create an ECR repository if it doesn't exist.
+
+        Returns the repository URI.
+        """
+        if repo_name in self._ecr_repo_cache:
+            return self._ecr_repo_cache[repo_name]
+
+        ecr = boto3.client("ecr", region_name=self._region)
+
+        lifecycle_policy = (
+            '{"rules":[{"rulePriority":1,"description":"Keep last 30",'
+            '"selection":{"tagStatus":"any","countType":"imageCountMoreThan",'
+            '"countNumber":30},"action":{"type":"expire"}}]}'
+        )
+
+        try:
+            resp = ecr.describe_repositories(repositoryNames=[repo_name])
+            uri = str(resp["repositories"][0]["repositoryUri"])
+        except ClientError as e:
+            if e.response["Error"]["Code"] != "RepositoryNotFoundException":
+                raise
+            resp = ecr.create_repository(
+                repositoryName=repo_name,
+                imageScanningConfiguration={"scanOnPush": True},
+                imageTagMutability="MUTABLE",
+            )
+            uri = str(resp["repository"]["repositoryUri"])
+            ecr.put_lifecycle_policy(
+                repositoryName=repo_name,
+                lifecyclePolicyText=lifecycle_policy,
+            )
+            logger.info("Created ECR repository: %s (%s)", repo_name, uri)
+
+        self._ecr_repo_cache[repo_name] = uri
+        return uri
+
+    # ------------------------------------------------------------------
+    # Convenience: provision everything a module needs
+    # ------------------------------------------------------------------
+
+    def provision_for_ecs_module(
+        self,
+        product: str,
+        module: str,
+        task_def_family: str = "pulse-pipeline-step",
+    ) -> dict[str, str]:
+        """Ensure all ECS infrastructure exists for a pipeline module.
+
+        Returns dict with cluster, task_definition, subnets, security_groups.
+        """
+        task_def_arn = self.ensure_ecs_task_definition(family=task_def_family)
+        ecr_repo = self.ensure_ecr_repository(f"{product}-{module}")
+
+        return {
+            "cluster": self._cluster,
+            "task_definition": task_def_arn,
+            "subnets": ",".join(self._subnets),
+            "security_groups": ",".join(self._security_groups),
+            "ecr_repository": ecr_repo,
+        }
+
+    def provision_for_lambda_module(
+        self,
+        product: str,
+        module: str,
+        image_uri: str,
+    ) -> dict[str, str]:
+        """Ensure Lambda function exists for a pipeline module.
+
+        Returns dict with function_name.
+        """
+        func_name = f"{product}-{module}"
+        self.ensure_lambda_function(func_name, image_uri)
+
+        return {"function_name": func_name}

pulse_engine/deployment/job_launcher.py

@@ -0,0 +1,178 @@
+"""Central dispatch for job triggering across all runner backends."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from pulse_engine.core.exceptions import (
+    BadRequestError,
+    NotFoundError,
+    UnprocessableEntityError,
+)
+from pulse_engine.deployment.backends.exceptions import BackendNotAvailableError
+
+if TYPE_CHECKING:
+    from pulse_engine.config import Settings
+    from pulse_engine.core.job_token import JobTokenIssuer
+    from pulse_engine.deployment.backend_deployment_repository import (
+        BackendDeploymentRepository,
+    )
+    from pulse_engine.deployment.backends.registry import BackendRegistry
+    from pulse_engine.deployment.repository import RegistrationRepository
+
+logger = structlog.get_logger(__name__)
+
+# The engine-owned Prefect flow entrypoint. All product containers must expose
+# entrypoint:run as a plain function; this wrapper (shipped with the engine)
+# decorates it with @flow so product code stays Prefect-free.
+PULSE_FLOW_ENTRYPOINT = "pulse_engine.runners.prefect_runner:main"
+
+_STAGE_NAME_MAP = {
+    "extractor": "extraction",
+    "processor": "processing",
+    "storage": "storage",
+}
+
+
+@dataclass
+class LaunchResult:
+    flow_run_id: str
+    token: str
+
+
+class JobLauncher:
+    """Orchestrates the full job trigger sequence.
+
+    Steps:
+    1. Look up product registration -> image
+    2. Resolve backend from registry
+    3. prepare() — work pool setup or no-op
+    4. register() — upsert backend run unit (cached)
+    5. Issue job-scoped JWT
+    6. trigger() — fire the run, return run ID
+    7. Return LaunchResult
+    """
+
+    def __init__(
+        self,
+        registration_repo: RegistrationRepository,
+        backend_deployment_repo: BackendDeploymentRepository,
+        registry: BackendRegistry,
+        token_issuer: JobTokenIssuer,
+        settings: Settings,
+    ) -> None:
+        self._reg_repo = registration_repo
+        self._backend_repo = backend_deployment_repo
+        self._registry = registry
+        self._token_issuer = token_issuer
+        self._settings = settings
+
+    async def launch(
+        self,
+        job_id: str,
+        tenant_id: str,
+        product: str,
+        stage: str,
+        orchestrator: str,
+        compute: str,
+        chain: bool,
+        config: dict[str, Any],
+    ) -> LaunchResult:
+        # 1. Look up registration
+        registration = await self._reg_repo.get(product, stage)
+        if registration is None:
+            raise NotFoundError(
+                f"No deployment registered for {product}/{stage}",
+                product=product,
+                stage=stage,
+            )
+
+        # 2. Resolve backend (raises 422 if unsupported)
+        try:
+            backend = self._registry.get(orchestrator, compute)
+        except BackendNotAvailableError as exc:
+            logger.warning(
+                "backend_not_available",
+                orchestrator=orchestrator,
+                compute=compute,
+                exc_info=True,
+            )
+            raise UnprocessableEntityError(
+                "Unsupported orchestrator/compute combination"
+            ) from exc
+
+        # 3. prepare() — idempotent setup (work pool creation, no-op for native)
+        await backend.prepare()
+
+        # 4. register() — cached; Prefect gets image + entrypoint, native Lambda
+        # derives function name from product+stage convention.
+        # NOTE: Concurrent callers may both see a cache miss and both call register().
+        # This is safe: upsert() uses ON CONFLICT DO UPDATE, and register()
+        # is idempotent.
+        cached = await self._backend_repo.get(product, stage, orchestrator, compute)
+        if cached:
+            handle = cached.deployment_id
+        else:
+            handle = await backend.register(
+                product=product,
+                stage=stage,
+                image=registration.image,
+                entrypoint=PULSE_FLOW_ENTRYPOINT,
+            )
+            await self._backend_repo.upsert(
+                product=product,
+                stage=stage,
+                orchestrator=orchestrator,
+                compute=compute,
+                deployment_id=handle,
+            )
+
+        # 5. Issue job-scoped JWT
+        stage_name = _STAGE_NAME_MAP.get(stage)
+        if stage_name is None:
+            raise BadRequestError(
+                f"Unknown stage: {stage!r}. Must be one of: {list(_STAGE_NAME_MAP)}"
+            )
+        scope = ["jobs:status"]
+        if stage == "storage":
+            scope.append("kb:write")
+        if chain:
+            scope.append("jobs:trigger_next")
+
+        token = self._token_issuer.issue(
+            job_id=job_id,
+            tenant_id=tenant_id,
+            product=product,
+            stage=stage_name,
+            scope=scope,
+            orchestrator=orchestrator,
+            compute=compute,
+        )
+
+        # 6. trigger() — returns run ID (UUID for native Lambda, Prefect flow
+        # run UUID for Prefect)
+        run_id = await backend.trigger(
+            handle,
+            {
+                "job_id": job_id,
+                "chain": chain,
+                "config": config,
+                "pulse_api_token": token,
+                "pulse_engine_url": self._settings.pulse_engine_url,
+            },
+        )
+
+        logger.info(
+            "job_launched",
+            job_id=job_id,
+            product=product,
+            stage=stage_name,
+            orchestrator=orchestrator,
+            compute=compute,
+            run_id=run_id,
+        )
+
+        return LaunchResult(flow_run_id=run_id, token=token)

pulse_engine/deployment/models.py

@@ -0,0 +1,48 @@
+"""ORM models for product registrations and backend deployment cache."""
+
+from datetime import datetime
+
+import sqlalchemy as sa
+from sqlalchemy.orm import Mapped, mapped_column
+
+from pulse_engine.database import Base
+
+
+class ProductRegistrationModel(Base):
+    """Stores product stage images registered via POST /api/v1/deployments."""
+
+    __tablename__ = "product_registrations"
+
+    product: Mapped[str] = mapped_column(sa.String(100), primary_key=True)
+    stage: Mapped[str] = mapped_column(sa.String(50), primary_key=True)
+    image: Mapped[str] = mapped_column(sa.String(500), nullable=False)
+    registered_at: Mapped[datetime] = mapped_column(
+        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        sa.DateTime(timezone=True),
+        server_default=sa.func.now(),
+        onupdate=sa.func.now(),
+        nullable=False,
+    )
+
+
+class ProductBackendDeploymentModel(Base):
+    """Lazy cache: Prefect deployment IDs per deployment dimension tuple."""
+
+    __tablename__ = "product_deployments"
+
+    product: Mapped[str] = mapped_column(sa.String(100), primary_key=True)
+    stage: Mapped[str] = mapped_column(sa.String(50), primary_key=True)
+    orchestrator: Mapped[str] = mapped_column(sa.String(50), primary_key=True)
+    compute: Mapped[str] = mapped_column(sa.String(50), primary_key=True)
+    deployment_id: Mapped[str] = mapped_column(sa.String(255), nullable=False)
+    created_at: Mapped[datetime] = mapped_column(
+        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        sa.DateTime(timezone=True),
+        server_default=sa.func.now(),
+        onupdate=sa.func.now(),
+        nullable=False,
+    )

pulse_engine/deployment/repository.py

@@ -0,0 +1,54 @@
+"""Data access layer for product registrations."""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from pulse_engine.deployment.models import ProductRegistrationModel
+
+
+class RegistrationRepository:
+    """CRUD for product_registrations table."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def upsert(
+        self,
+        product: str,
+        stage: str,
+        image: str,
+    ) -> ProductRegistrationModel:
+        existing = await self.get(product, stage)
+        if existing:
+            existing.image = image
+            await self._session.commit()
+            await self._session.refresh(existing)
+            return existing
+        record = ProductRegistrationModel(
+            product=product,
+            stage=stage,
+            image=image,
+        )
+        self._session.add(record)
+        await self._session.commit()
+        await self._session.refresh(record)
+        return record
+
+    async def get(self, product: str, stage: str) -> ProductRegistrationModel | None:
+        stmt = sa.select(ProductRegistrationModel).where(
+            ProductRegistrationModel.product == product,
+            ProductRegistrationModel.stage == stage,
+        )
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def list_by_product(self, product: str) -> list[ProductRegistrationModel]:
+        stmt = (
+            sa.select(ProductRegistrationModel)
+            .where(ProductRegistrationModel.product == product)
+            .order_by(ProductRegistrationModel.stage)
+        )
+        result = await self._session.execute(stmt)
+        return list(result.scalars().all())

pulse_engine/deployment/router.py

@@ -0,0 +1,22 @@
+"""Deployments API router."""
+
+from fastapi import APIRouter, Depends
+
+from pulse_engine.dependencies import get_deployment_service
+from pulse_engine.deployment.schemas import (
+    RegisterDeploymentRequest,
+    RegisterDeploymentResponse,
+)
+from pulse_engine.deployment.service import DeploymentService
+from pulse_engine.middleware.tenant import get_tenant_id
+
+router = APIRouter(prefix="/deployments", tags=["Deployments"])
+
+
+@router.post("/", response_model=RegisterDeploymentResponse, status_code=201)
+async def register_deployment(
+    body: RegisterDeploymentRequest,
+    tenant_id: str = Depends(get_tenant_id),
+    service: DeploymentService = Depends(get_deployment_service),
+) -> RegisterDeploymentResponse:
+    return await service.register(body)

pulse_engine/deployment/schemas.py

@@ -0,0 +1,18 @@
+# src/pulse_engine/deployment/schemas.py
+"""Pydantic schemas for the Deployments API."""
+
+from pydantic import BaseModel
+
+
+class StageDefinition(BaseModel):
+    image: str  # ECR URI
+
+
+class RegisterDeploymentRequest(BaseModel):
+    product: str
+    stages: dict[str, StageDefinition]  # key: "extractor" | "processor" | "storage"
+
+
+class RegisterDeploymentResponse(BaseModel):
+    product: str
+    stages: dict[str, dict[str, str]]  # {"extractor": {"status": "registered"}, ...}

pulse_engine/deployment/service.py

@@ -0,0 +1,65 @@
+# src/pulse_engine/deployment/service.py
+"""Business logic for registering product stage deployments."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import structlog
+
+from pulse_engine.deployment.schemas import (
+    RegisterDeploymentRequest,
+    RegisterDeploymentResponse,
+)
+
+if TYPE_CHECKING:
+    from pulse_engine.deployment.backend_deployment_repository import (
+        BackendDeploymentRepository,
+    )
+    from pulse_engine.deployment.repository import RegistrationRepository
+
+logger = structlog.get_logger(__name__)
+
+
+class DeploymentService:
+    """Stores product stage image registrations. Does not create Prefect deployments.
+
+    Prefect deployments are created lazily by JobLauncher on first trigger.
+    When a product re-registers, all cached backend deployment IDs for that
+    product+stage are invalidated so JobLauncher recreates them with the new image.
+    """
+
+    def __init__(
+        self,
+        registration_repo: RegistrationRepository,
+        backend_deployment_repo: BackendDeploymentRepository,
+    ) -> None:
+        self._repo = registration_repo
+        self._backend_repo = backend_deployment_repo
+
+    async def register(
+        self, request: RegisterDeploymentRequest
+    ) -> RegisterDeploymentResponse:
+        result_stages: dict[str, dict[str, str]] = {}
+
+        for stage_name, stage_def in request.stages.items():
+            await self._repo.upsert(
+                product=request.product,
+                stage=stage_name,
+                image=stage_def.image,
+            )
+            # Invalidate cached Prefect deployment IDs — image may have changed
+            await self._backend_repo.delete_by_product_stage(
+                request.product, stage_name
+            )
+            result_stages[stage_name] = {"status": "registered"}
+            logger.info(
+                "stage_registration_stored",
+                product=request.product,
+                stage=stage_name,
+            )
+
+        return RegisterDeploymentResponse(
+            product=request.product,
+            stages=result_stages,
+        )