pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/pipeline/service.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import Any
+
+from pulse_engine.pipeline.models import PipelineRunModel
+from pulse_engine.pipeline.repositories import (
+    ModuleRegistryRepository,
+    PipelineRunRepository,
+    PipelineStepRunRepository,
+)
+from pulse_engine.pipeline.schemas import PipelineConfig, StepStatus
+from pulse_engine.pipeline.translators.base import (
+    BaseStatusProvider,
+    BaseTranslator,
+    PipelineStatus,
+)
+
+
+class PipelineServiceError(Exception):
+    pass
+
+
+class PipelineSubmissionError(PipelineServiceError):
+    pass
+
+
+class PipelineService:
+    def __init__(
+        self,
+        module_repo: ModuleRegistryRepository,
+        run_repo: PipelineRunRepository,
+        translators: dict[str, BaseTranslator],
+        status_providers: dict[str, BaseStatusProvider],
+        settings: Any,
+        token_issuer: Any,
+        step_repo: PipelineStepRunRepository | None = None,
+    ) -> None:
+        self._module_repo = module_repo
+        self._run_repo = run_repo
+        self._translators = translators
+        self._status_providers = status_providers
+        self._settings = settings
+        self._token_issuer = token_issuer
+        self._step_repo = step_repo
+
+    async def trigger(
+        self,
+        tenant_id: str,
+        product: str,
+        orchestrator: str,
+        config: PipelineConfig,
+        global_config: dict[str, Any],
+    ) -> str:
+        """Trigger a pipeline. Returns the pipeline run ID."""
+        # Validate orchestrator
+        if orchestrator not in self._translators:
+            raise PipelineServiceError(f"Unsupported orchestrator: '{orchestrator}'")
+
+        # Resolve module images
+        images = await self._module_repo.get_images_map(tenant_id, product)
+        missing = [m.name for m in config.modules if m.module not in images]
+        if missing:
+            raise PipelineServiceError(
+                f"Modules not registered for product '{product}': {missing}"
+            )
+
+        # Create pipeline run record
+        run = await self._run_repo.create(
+            tenant_id=tenant_id,
+            product=product,
+            orchestrator=orchestrator,
+            config_snapshot=config.model_dump(mode="json", by_alias=True),
+            global_config=global_config,
+        )
+
+        # Issue run-scoped token
+        token = self._token_issuer.issue_token(
+            pipeline_run_id=run.id,
+            tenant_id=tenant_id,
+        )
+
+        # Submit to orchestrator
+        translator = self._translators[orchestrator]
+        try:
+            orchestrator_run_id = await translator.submit(
+                pipeline_run_id=run.id,
+                parsed_config=config,
+                module_images=images,
+                global_config=global_config,
+                tenant_id=tenant_id,
+                pulse_engine_url=self._settings.pulse_engine_url,
+                pulse_api_token=token,
+            )
+        except Exception as e:
+            await self._run_repo.update_status(run.id, "submission_failed")
+            raise PipelineSubmissionError(
+                f"Failed to submit pipeline to {orchestrator}: {e}"
+            ) from e
+
+        await self._run_repo.set_orchestrator_run_id(run.id, orchestrator_run_id)
+        await self._run_repo.update_status(run.id, "running")
+        return run.id
+
+    async def get_run(self, tenant_id: str, run_id: str) -> PipelineRunModel | None:
+        """Get a pipeline run record."""
+        return await self._run_repo.get(run_id, tenant_id)
+
+    async def get_status(self, tenant_id: str, run_id: str) -> PipelineStatus | None:
+        """Get normalized pipeline status, merging DB step records as fallback."""
+        run = await self._run_repo.get(run_id, tenant_id)
+        if run is None:
+            return None
+
+        # Load DB-persisted step records for merge / fallback
+        db_steps: list[StepStatus] = []
+        if self._step_repo is not None:
+            rows = await self._step_repo.list_by_run(run_id, tenant_id)
+            db_steps = [
+                StepStatus(
+                    step=r.step_name,
+                    module=r.module_type,
+                    status=r.status,
+                    started_at=r.started_at,
+                    completed_at=r.completed_at,
+                    error_message=r.error_message,
+                    output_ref=r.output_ref,
+                )
+                for r in rows
+            ]
+
+        if run.orchestrator_run_id is None:
+            return PipelineStatus(status=run.status, steps=db_steps)
+
+        provider = self._status_providers.get(run.orchestrator)
+        if provider is None:
+            return PipelineStatus(status=run.status, steps=db_steps)
+
+        try:
+            status = await provider.get_status(run.orchestrator_run_id)
+        except Exception:
+            # Orchestrator unreachable — serve from DB
+            return PipelineStatus(status=run.status, steps=db_steps)
+
+        # Merge: enrich orchestrator steps with DB fields (error_message, output_ref)
+        db_map = {s.step: s for s in db_steps}
+        merged: list[StepStatus] = []
+        for s in status.steps:
+            db = db_map.get(s.step)
+            if db is not None:
+                merged.append(
+                    StepStatus(
+                        step=s.step,
+                        module=s.module,
+                        status=s.status,
+                        fan_out_count=s.fan_out_count,
+                        completed_count=s.completed_count,
+                        failed_count=s.failed_count,
+                        started_at=s.started_at or db.started_at,
+                        completed_at=db.completed_at,
+                        error_message=db.error_message,
+                        output_ref=db.output_ref,
+                    )
+                )
+            else:
+                merged.append(s)
+
+        # Add DB-only steps not returned by orchestrator (e.g. after purge)
+        orchestrator_step_names = {s.step for s in status.steps}
+        for s in db_steps:
+            if s.step not in orchestrator_step_names:
+                merged.append(s)
+
+        status.steps = merged
+
+        # Cache top-level status
+        if status.status != run.status:
+            await self._run_repo.update_status(run_id, status.status)
+
+        return status
+
+    async def record_step_status(
+        self,
+        tenant_id: str,
+        run_id: str,
+        step_name: str,
+        status: str,
+        output_ref: dict[str, Any] | None = None,
+        error_message: str | None = None,
+    ) -> None:
+        """Persist a step status callback from the orchestrator runner."""
+        run = await self._run_repo.get(run_id, tenant_id)
+        if run is None or self._step_repo is None:
+            return
+
+        now = datetime.now(tz=UTC)
+        started_at = now if status == "running" else None
+        completed_at = now if status in {"completed", "failed"} else None
+
+        # Derive module_type from config snapshot if possible
+        module_type = ""
+        dag = run.config_snapshot.get("dag", [])
+        for step in dag:
+            base_name = step_name.split("[")[0]
+            if step.get("step") == base_name:
+                module_type = step.get("module_type") or step.get("module", "")
+                break
+
+        await self._step_repo.upsert(
+            pipeline_run_id=run_id,
+            tenant_id=tenant_id,
+            step_name=step_name,
+            module_type=module_type,
+            status=status,
+            started_at=started_at,
+            completed_at=completed_at,
+            error_message=error_message,
+            output_ref=output_ref,
+        )
+
+    async def cancel(self, tenant_id: str, run_id: str) -> bool:
+        """Cancel a running pipeline."""
+        run = await self._run_repo.get(run_id, tenant_id)
+        if run is None:
+            return False
+
+        if run.orchestrator_run_id is None:
+            await self._run_repo.update_status(run_id, "cancelled")
+            return True
+
+        provider = self._status_providers.get(run.orchestrator)
+        if provider is None:
+            return False
+
+        result = await provider.cancel(run.orchestrator_run_id)
+        if result:
+            await self._run_repo.update_status(run_id, "cancelled")
+        return result
+
+    async def list_runs(
+        self,
+        tenant_id: str,
+        product: str | None = None,
+        status: str | None = None,
+        limit: int = 20,
+        offset: int = 0,
+    ) -> tuple[list[PipelineRunModel], int]:
+        return await self._run_repo.list_runs(
+            tenant_id, product=product, status=status, limit=limit, offset=offset
+        )

pulse_engine/pipeline/translators/__init__.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pulse_engine.pipeline.translators.base import BaseStatusProvider, BaseTranslator
+
+__all__ = [
+    "BaseTranslator",
+    "BaseStatusProvider",
+    "get_translator",
+    "get_status_provider",
+]
+
+
+def get_translator(orchestrator: str, **kwargs: Any) -> BaseTranslator:
+    if orchestrator == "prefect":
+        from pulse_engine.pipeline.translators.prefect_translator import (
+            PrefectTranslator,
+        )
+
+        return PrefectTranslator(**kwargs)
+    elif orchestrator == "airflow":
+        from pulse_engine.pipeline.translators.airflow_translator import (
+            AirflowTranslator,
+        )
+
+        return AirflowTranslator()
+    raise ValueError(f"Unknown orchestrator: {orchestrator}")
+
+
+def get_status_provider(orchestrator: str, **kwargs: Any) -> BaseStatusProvider:
+    if orchestrator == "prefect":
+        from pulse_engine.pipeline.translators.prefect_status import (
+            PrefectStatusProvider,
+        )
+
+        return PrefectStatusProvider(**kwargs)
+    elif orchestrator == "airflow":
+        from pulse_engine.pipeline.translators.airflow_status import (
+            AirflowStatusProvider,
+        )
+
+        return AirflowStatusProvider()
+    raise ValueError(f"Unknown orchestrator: {orchestrator}")

pulse_engine/pipeline/translators/airflow_status.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from pulse_engine.pipeline.translators.base import BaseStatusProvider, PipelineStatus
+
+
+class AirflowStatusProvider(BaseStatusProvider):
+    async def get_status(self, orchestrator_run_id: str) -> PipelineStatus:
+        raise NotImplementedError("Airflow status provider is not yet implemented.")
+
+    async def cancel(self, orchestrator_run_id: str) -> bool:
+        raise NotImplementedError("Airflow cancel is not yet implemented.")

pulse_engine/pipeline/translators/airflow_translator.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pulse_engine.pipeline.schemas import PipelineConfig
+from pulse_engine.pipeline.translators.base import BaseTranslator
+
+
+class AirflowTranslator(BaseTranslator):
+    async def submit(
+        self,
+        pipeline_run_id: str,
+        parsed_config: PipelineConfig,
+        module_images: dict[str, str],
+        global_config: dict[str, Any],
+        tenant_id: str,
+        pulse_engine_url: str,
+        pulse_api_token: str,
+    ) -> str:
+        raise NotImplementedError(
+            "Airflow translator is not yet implemented. Use orchestrator='prefect'."
+        )

pulse_engine/pipeline/translators/base.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+from pulse_engine.pipeline.schemas import PipelineConfig, StepStatus
+
+
+@dataclass
+class PipelineStatus:
+    status: str  # pending, running, completed, failed, cancelled
+    started_at: str | None = None
+    steps: list[StepStatus] = field(default_factory=list)
+
+
+class BaseTranslator(ABC):
+    @abstractmethod
+    async def submit(
+        self,
+        pipeline_run_id: str,
+        parsed_config: PipelineConfig,
+        module_images: dict[str, str],
+        global_config: dict[str, Any],
+        tenant_id: str,
+        pulse_engine_url: str,
+        pulse_api_token: str,
+    ) -> str:
+        """Submit pipeline to orchestrator. Returns orchestrator's native run ID."""
+        ...
+
+
+class BaseStatusProvider(ABC):
+    @abstractmethod
+    async def get_status(self, orchestrator_run_id: str) -> PipelineStatus:
+        """Query orchestrator and return normalized status."""
+        ...
+
+    @abstractmethod
+    async def cancel(self, orchestrator_run_id: str) -> bool:
+        """Cancel a running pipeline. Returns True if successfully cancelled."""
+        ...

pulse_engine/pipeline/translators/prefect_status.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import base64
+
+import httpx
+
+from pulse_engine.pipeline.schemas import StepStatus
+from pulse_engine.pipeline.translators.base import BaseStatusProvider, PipelineStatus
+
+_PREFECT_STATE_MAP: dict[str, str] = {
+    "COMPLETED": "completed",
+    "FAILED": "failed",
+    "CANCELLED": "cancelled",
+    "CANCELLING": "cancelled",
+    "RUNNING": "running",
+    "PENDING": "pending",
+    "SCHEDULED": "pending",
+}
+
+
+class PrefectStatusProvider(BaseStatusProvider):
+    def __init__(
+        self,
+        prefect_api_url: str,
+        prefect_api_key: str,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self._api_url = prefect_api_url.rstrip("/")
+        self._api_key = prefect_api_key
+        self._client = http_client
+
+    def _get_client(self) -> httpx.AsyncClient:
+        if self._client is not None:
+            return self._client
+        auth_header = base64.b64encode(f":{self._api_key}".encode()).decode()
+        return httpx.AsyncClient(
+            headers={"Authorization": f"Basic {auth_header}"},
+            timeout=30.0,
+        )
+
+    async def get_status(self, orchestrator_run_id: str) -> PipelineStatus:
+        client = self._get_client()
+        try:
+            resp = await client.get(f"{self._api_url}/flow_runs/{orchestrator_run_id}")
+            resp.raise_for_status()
+            flow_run = resp.json()
+
+            prefect_state = flow_run.get("state", {}).get("type", "PENDING")
+            status = _PREFECT_STATE_MAP.get(prefect_state, "pending")
+            started_at = flow_run.get("start_time")
+
+            task_resp = await client.post(
+                f"{self._api_url}/task_runs/filter",
+                json={"flow_runs": {"id": {"any_": [orchestrator_run_id]}}},
+            )
+            task_resp.raise_for_status()
+            task_runs = task_resp.json()
+
+            steps: list[StepStatus] = []
+            for task in task_runs:
+                task_state = task.get("state", {}).get("type", "PENDING")
+                steps.append(
+                    StepStatus(
+                        step=task.get("name", "unknown"),
+                        module=task.get("name", "unknown"),
+                        status=_PREFECT_STATE_MAP.get(task_state, "pending"),
+                        started_at=task.get("start_time"),
+                    )
+                )
+
+            return PipelineStatus(
+                status=status,
+                started_at=started_at,
+                steps=steps,
+            )
+        finally:
+            if self._client is None:
+                await client.aclose()
+
+    async def cancel(self, orchestrator_run_id: str) -> bool:
+        client = self._get_client()
+        try:
+            resp = await client.post(
+                f"{self._api_url}/flow_runs/{orchestrator_run_id}/set_state",
+                json={"state": {"type": "CANCELLED"}},
+            )
+            resp.raise_for_status()
+            return True
+        except httpx.HTTPStatusError:
+            return False
+        finally:
+            if self._client is None:
+                await client.aclose()

pulse_engine/pipeline/translators/prefect_translator.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import httpx
+
+from pulse_engine.pipeline.schemas import PipelineConfig, SecretRef
+from pulse_engine.pipeline.translators.base import BaseTranslator
+
+
+class PrefectTranslator(BaseTranslator):
+    """Translates a PipelineConfig into a Prefect flow run submission."""
+
+    FLOW_ENTRYPOINT = "pulse_engine.runners.prefect_pipeline_flow:pipeline_flow"
+
+    def __init__(
+        self,
+        prefect_api_url: str,
+        prefect_api_key: str,
+        work_pool_name: str,
+        engine_image: str = "",
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self._api_url = prefect_api_url.rstrip("/")
+        self._api_key = prefect_api_key
+        self._work_pool_name = work_pool_name
+        self._engine_image = engine_image
+        self._client = http_client
+
+    def _get_client(self) -> httpx.AsyncClient:
+        if self._client is not None:
+            return self._client
+        auth_header = base64.b64encode(f":{self._api_key}".encode()).decode()
+        return httpx.AsyncClient(
+            headers={"Authorization": f"Basic {auth_header}"},
+            timeout=30.0,
+        )
+
+    def _serialize_dag(self, config: PipelineConfig) -> list[dict[str, Any]]:
+        """Serialize v2 DAG into a JSON-friendly list for Prefect flow parameters."""
+        module_map = {m.name: m for m in config.modules}
+        steps = []
+        for s in config.dag:
+            module = module_map[s.module]
+            step_dict: dict[str, Any] = {
+                "step": s.step,
+                "module": s.module,
+                "module_type": module.module,
+                "retries": module.retries,
+                "retry_delay": module.retry_delay,
+                "timeout": module.timeout,
+                "max_concurrency": module.max_concurrency,
+                "args": module.args,
+                "env": {
+                    k: v.model_dump() if isinstance(v, SecretRef) else v
+                    for k, v in module.env.items()
+                },
+                "resources": {
+                    **module.resources.model_dump(),
+                    "compute": module.resources.compute
+                    or config.infra.default_compute
+                    or "ecs",
+                    "cpu_units": module.resources.ecs_cpu_units,
+                },
+            }
+            if s.for_each is not None:
+                step_dict["for_each"] = s.for_each
+            if s.collect_from is not None:
+                step_dict["collect_from"] = (
+                    [s.collect_from]
+                    if isinstance(s.collect_from, str)
+                    else list(s.collect_from)
+                )
+            if s.when is not None:
+                step_dict["when"] = s.when
+            if s.trigger_rule is not None:
+                step_dict["trigger_rule"] = s.trigger_rule
+            if s.depends_on:
+                step_dict["depends_on"] = [{"step": d.step} for d in s.depends_on]
+            steps.append(step_dict)
+        return steps
+
+    async def _ensure_deployment(
+        self,
+        client: httpx.AsyncClient,
+        pipeline_name: str,
+        cron_schedule: str | None = None,
+    ) -> str:
+        """Create or update the Prefect deployment for this pipeline.
+
+        Returns the deployment ID.
+        """
+        deployment_name = self._get_deployment_name(pipeline_name)
+
+        # Get or create the flow (idempotent — Prefect returns existing if name matches)
+        flow_resp = await client.post(
+            f"{self._api_url}/flows/",
+            json={"name": deployment_name},
+        )
+        flow_resp.raise_for_status()
+        flow_id = str(flow_resp.json()["id"])
+
+        deploy_body: dict[str, Any] = {
+            "name": deployment_name,
+            "flow_id": flow_id,
+            "entrypoint": self.FLOW_ENTRYPOINT,
+            "work_pool_name": self._work_pool_name,
+            "job_variables": {"image": self._engine_image},
+            "path": "/app",
+        }
+        if cron_schedule:
+            deploy_body["schedules"] = [
+                {"schedule": {"cron": cron_schedule, "timezone": "UTC"}, "active": True}
+            ]
+
+        # Upsert deployment
+        # path="/app" required for Prefect process worker to locate the flow
+        # when the entrypoint is a package import (not a relative file path).
+        deploy_resp = await client.post(
+            f"{self._api_url}/deployments/",
+            json=deploy_body,
+        )
+        deploy_resp.raise_for_status()
+        return str(deploy_resp.json()["id"])
+
+    async def submit(
+        self,
+        pipeline_run_id: str,
+        parsed_config: PipelineConfig,
+        module_images: dict[str, str],
+        global_config: dict[str, Any],
+        tenant_id: str,
+        pulse_engine_url: str,
+        pulse_api_token: str,
+    ) -> str:
+        """Submit a pipeline as a Prefect flow run, upserting the deployment first."""
+        client = self._get_client()
+
+        parameters = {
+            "pipeline_run_id": pipeline_run_id,
+            "tenant_id": tenant_id,
+            "dag": self._serialize_dag(parsed_config),
+            "module_images": module_images,
+            "global_config": global_config,
+            "pulse_engine_url": pulse_engine_url,
+            "pulse_api_token": pulse_api_token,
+            "results_backend": parsed_config.infra.results_backend.model_dump(),
+        }
+
+        body = {
+            "parameters": parameters,
+            "state": {"type": "SCHEDULED"},
+            "tags": [f"product:{parsed_config.name}", f"pipeline:{pipeline_run_id}"],
+        }
+
+        try:
+            deployment_id = await self._ensure_deployment(
+                client, parsed_config.deployment_name, parsed_config.infra.schedule
+            )
+
+            response = await client.post(
+                f"{self._api_url}/deployments/{deployment_id}/create_flow_run",
+                json=body,
+            )
+            response.raise_for_status()
+            data = response.json()
+            return str(data["id"])
+        finally:
+            if self._client is None:
+                await client.aclose()
+
+    async def _resolve_deployment_id(
+        self, client: httpx.AsyncClient, deployment_name: str
+    ) -> str:
+        """Look up a Prefect deployment ID by name."""
+        response = await client.post(
+            f"{self._api_url}/deployments/filter",
+            json={
+                "deployments": {"name": {"any_": [deployment_name]}},
+                "limit": 1,
+            },
+        )
+        response.raise_for_status()
+        deployments = response.json()
+        if not deployments:
+            raise RuntimeError(
+                f"Prefect deployment '{deployment_name}' not found. "
+                f"Create it first with: prefect deployment build ..."
+            )
+        return str(deployments[0]["id"])
+
+    def _get_deployment_name(self, pipeline_name: str) -> str:
+        """Deployment name for the generic pipeline flow."""
+        return f"pulse-pipeline-{pipeline_name}"