ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/deployment/base.py

@@ -0,0 +1,861 @@
+"""Core classes for pipeline deployments.
+
+Provides the PipelineDeployment base class and related types for
+creating unified, type-safe pipeline deployments with:
+- Per-flow resume (skip if outputs exist in DocumentStore)
+- Per-flow uploads (immediate, not just at end)
+- Prefect state hooks (on_running, on_completion, etc.)
+- Upload on failure (partial results saved)
+"""
+
+import asyncio
+import contextlib
+import hashlib
+import os
+import sys
+from abc import abstractmethod
+from collections.abc import Callable
+from contextlib import ExitStack
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
+from uuid import UUID, uuid4
+
+import httpx
+from lmnr import Laminar
+from opentelemetry import trace as otel_trace
+from prefect import flow, get_client, runtime
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic_settings import CliPositionalArg, SettingsConfigDict
+
+from ai_pipeline_core.document_store import SummaryGenerator, create_document_store, get_document_store, set_document_store
+from ai_pipeline_core.document_store.local import LocalDocumentStore
+from ai_pipeline_core.document_store.memory import MemoryDocumentStore
+from ai_pipeline_core.documents import Document
+from ai_pipeline_core.documents.context import RunContext, reset_run_context, set_run_context
+from ai_pipeline_core.logging import get_pipeline_logger, setup_logging
+from ai_pipeline_core.observability._debug import LocalDebugSpanProcessor, LocalTraceWriter, TraceDebugConfig
+from ai_pipeline_core.observability._initialization import get_tracking_service, initialize_observability
+from ai_pipeline_core.observability._tracking._models import RunStatus
+from ai_pipeline_core.pipeline.options import FlowOptions
+from ai_pipeline_core.settings import settings
+from ai_pipeline_core.testing import disable_run_logger, prefect_test_harness
+
+from .contract import CompletedRun, DeploymentResultData, FailedRun, ProgressRun
+from .helpers import (
+    StatusPayload,
+    class_name_to_deployment_name,
+    download_documents,
+    extract_generic_params,
+    send_webhook,
+    upload_documents,
+)
+
+logger = get_pipeline_logger(__name__)
+
+
+def _build_summary_generator() -> SummaryGenerator | None:
+    """Build a summary generator callable from settings, or None if disabled/unavailable."""
+    if not settings.doc_summary_enabled:
+        return None
+
+    from ai_pipeline_core.observability._summary import generate_document_summary
+
+    model = settings.doc_summary_model
+
+    async def _generator(name: str, excerpt: str) -> str:
+        return await generate_document_summary(name, excerpt, model=model)
+
+    return _generator
+
+
+# Fields added by run_cli()'s _CliOptions that should not affect the run scope fingerprint
+_CLI_FIELDS: set[str] = {"working_directory", "project_name", "start", "end", "no_trace"}
+
+
+def _compute_run_scope(project_name: str, documents: list[Document], options: FlowOptions) -> str:
+    """Compute a run scope that fingerprints inputs and options.
+
+    Different inputs or options produce a different scope, preventing
+    stale cache hits when re-running with the same project name.
+    Falls back to just project_name when no documents are provided
+    (e.g. --start N resume without initializer).
+    """
+    if not documents:
+        return project_name
+    sha256s = sorted(doc.sha256 for doc in documents)
+    exclude = _CLI_FIELDS & set(type(options).model_fields)
+    options_json = options.model_dump_json(exclude=exclude, exclude_none=True)
+    fingerprint = hashlib.sha256(f"{':'.join(sha256s)}|{options_json}".encode()).hexdigest()[:16]
+    return f"{project_name}:{fingerprint}"
+
+
+class DeploymentContext(BaseModel):
+    """Infrastructure configuration for deployments.
+
+    Webhooks are optional - provide URLs to enable:
+    - progress_webhook_url: Per-flow progress (started/completed/cached)
+    - status_webhook_url: Prefect state transitions (RUNNING/FAILED/etc)
+    - completion_webhook_url: Final result when deployment ends
+    """
+
+    input_documents_urls: tuple[str, ...] = Field(default_factory=tuple)
+    output_documents_urls: dict[str, str] = Field(default_factory=dict)  # nosemgrep: mutable-field-on-frozen-pydantic-model
+
+    progress_webhook_url: str = ""
+    status_webhook_url: str = ""
+    completion_webhook_url: str = ""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+
+class DeploymentResult(BaseModel):
+    """Base class for deployment results."""
+
+    success: bool
+    error: str | None = None
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+
+TOptions = TypeVar("TOptions", bound=FlowOptions)
+TResult = TypeVar("TResult", bound=DeploymentResult)
+
+
+class FlowCallable(Protocol):
+    """Protocol for @pipeline_flow decorated functions."""
+
+    name: str
+    __name__: str
+    input_document_types: list[type[Document]]
+    output_document_types: list[type[Document]]
+    estimated_minutes: int
+
+    def __call__(self, project_name: str, documents: list[Document], flow_options: FlowOptions) -> Any:  # type: ignore[type-arg]
+        """Execute the flow with standard pipeline signature."""
+        ...
+
+    def with_options(self, **kwargs: Any) -> "FlowCallable":
+        """Return a copy with overridden Prefect flow options (e.g., hooks)."""
+        ...
+
+
+def _reattach_flow_metadata(original: FlowCallable, target: Any) -> None:
+    """Reattach custom flow attributes that Prefect's with_options() may strip."""
+    for attr in ("input_document_types", "output_document_types", "estimated_minutes"):
+        if hasattr(original, attr) and not hasattr(target, attr):
+            setattr(target, attr, getattr(original, attr))
+
+
+@dataclass(slots=True)
+class _StatusWebhookHook:
+    """Prefect hook that sends status webhooks on state transitions."""
+
+    webhook_url: str
+    flow_run_id: str
+    project_name: str
+    step: int
+    total_steps: int
+    flow_name: str
+
+    async def __call__(self, flow: Any, flow_run: Any, state: Any) -> None:
+        payload: StatusPayload = {
+            "type": "status",
+            "flow_run_id": str(flow_run.id),
+            "project_name": self.project_name,
+            "step": self.step,
+            "total_steps": self.total_steps,
+            "flow_name": self.flow_name,
+            "state": state.type.value if hasattr(state.type, "value") else str(state.type),
+            "state_name": state.name or "",
+            "timestamp": datetime.now(UTC).isoformat(),
+        }
+        try:
+            async with httpx.AsyncClient(timeout=10) as client:
+                await client.post(self.webhook_url, json=payload)
+        except Exception as e:
+            logger.warning(f"Status webhook failed: {e}")
+
+
+def _validate_flow_chain(deployment_name: str, flows: list[Any]) -> None:
+    """Validate that each flow's input types are satisfiable by preceding flows' outputs.
+
+    Simulates a type pool: starts with the first flow's input types, adds each flow's
+    output types after processing. For subsequent flows, each required input type must
+    be satisfiable by at least one type in the pool (via issubclass).
+    """
+    type_pool: set[type[Document]] = set()
+
+    for i, flow_fn in enumerate(flows):
+        input_types: list[type[Document]] = getattr(flow_fn, "input_document_types", [])
+        output_types: list[type[Document]] = getattr(flow_fn, "output_document_types", [])
+        flow_name = getattr(flow_fn, "name", getattr(flow_fn, "__name__", f"flow[{i}]"))
+
+        if i == 0:
+            # First flow: its input types seed the pool
+            type_pool.update(input_types)
+        elif input_types:
+            # Subsequent flows: at least one declared input type must be satisfiable
+            # from the pool (union semantics — flow accepts any of the declared types)
+            any_satisfied = any(any(issubclass(available, t) for available in type_pool) for t in input_types)
+            if not any_satisfied:
+                input_names = sorted(t.__name__ for t in input_types)
+                pool_names = sorted(t.__name__ for t in type_pool) if type_pool else ["(empty)"]
+                raise TypeError(
+                    f"{deployment_name}: flow '{flow_name}' (step {i + 1}) requires input types "
+                    f"{input_names} but none are produced by preceding flows. "
+                    f"Available types: {pool_names}"
+                )
+
+        type_pool.update(output_types)
+
+
+class PipelineDeployment(Generic[TOptions, TResult]):
+    """Base class for pipeline deployments.
+
+    Features enabled by default:
+    - Per-flow resume: Skip flows if outputs exist in DocumentStore
+    - Per-flow uploads: Upload documents after each flow
+    - Prefect hooks: Attach state hooks if status_webhook_url provided
+    - Upload on failure: Save partial results if pipeline fails
+    """
+
+    flows: ClassVar[list[FlowCallable]]
+    name: ClassVar[str]
+    options_type: ClassVar[type[FlowOptions]]
+    result_type: ClassVar[type[DeploymentResult]]
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        if not hasattr(cls, "flows"):
+            return
+
+        if cls.__name__.startswith("Test"):
+            raise TypeError(f"Deployment class name cannot start with 'Test': {cls.__name__}")
+
+        cls.name = class_name_to_deployment_name(cls.__name__)
+
+        options_type, result_type = extract_generic_params(cls, PipelineDeployment)
+        if options_type is None or result_type is None:
+            raise TypeError(f"{cls.__name__} must specify Generic parameters: class {cls.__name__}(PipelineDeployment[MyOptions, MyResult])")
+
+        cls.options_type = options_type
+        cls.result_type = result_type
+
+        if not cls.flows:
+            raise TypeError(f"{cls.__name__}.flows cannot be empty")
+
+        # build_result must be implemented (not still abstract from PipelineDeployment)
+        build_result_fn = getattr(cls, "build_result", None)
+        if build_result_fn is None or getattr(build_result_fn, "__isabstractmethod__", False):
+            raise TypeError(f"{cls.__name__} must implement 'build_result' static method")
+
+        # No duplicate flows (by identity)
+        seen_ids: set[int] = set()
+        for flow_fn in cls.flows:
+            fid = id(flow_fn)
+            if fid in seen_ids:
+                flow_name = getattr(flow_fn, "name", getattr(flow_fn, "__name__", str(flow_fn)))
+                raise TypeError(f"{cls.__name__}.flows contains duplicate flow '{flow_name}'")
+            seen_ids.add(fid)
+
+        # Flow type chain validation: simulate a type pool
+        _validate_flow_chain(cls.__name__, cls.flows)
+
+    @staticmethod
+    @abstractmethod
+    def build_result(project_name: str, documents: list[Document], options: TOptions) -> TResult:
+        """Extract typed result from pipeline documents loaded from DocumentStore."""
+        ...
+
+    def _all_document_types(self) -> list[type[Document]]:
+        """Collect all document types from all flows (inputs + outputs), deduplicated."""
+        types: dict[str, type[Document]] = {}
+        for flow_fn in self.flows:
+            for t in getattr(flow_fn, "input_document_types", []):
+                types[t.__name__] = t
+            for t in getattr(flow_fn, "output_document_types", []):
+                types[t.__name__] = t
+        return list(types.values())
+
+    def _build_status_hooks(
+        self,
+        context: DeploymentContext,
+        flow_run_id: str,
+        project_name: str,
+        step: int,
+        total_steps: int,
+        flow_name: str,
+    ) -> dict[str, list[Callable[..., Any]]]:
+        """Build Prefect hooks for status webhooks."""
+        hook = _StatusWebhookHook(
+            webhook_url=context.status_webhook_url,
+            flow_run_id=flow_run_id,
+            project_name=project_name,
+            step=step,
+            total_steps=total_steps,
+            flow_name=flow_name,
+        )
+        return {
+            "on_running": [hook],
+            "on_completion": [hook],
+            "on_failure": [hook],
+            "on_crashed": [hook],
+            "on_cancellation": [hook],
+        }
+
+    async def _send_progress(
+        self,
+        context: DeploymentContext,
+        flow_run_id: str,
+        project_name: str,
+        step: int,
+        total_steps: int,
+        flow_name: str,
+        status: str,
+        step_progress: float = 0.0,
+        message: str = "",
+    ) -> None:
+        """Send progress webhook and update flow run labels."""
+        # Use estimated_minutes for weighted progress calculation
+        flow_minutes = [getattr(f, "estimated_minutes", 1) for f in self.flows]
+        total_minutes = sum(flow_minutes) or 1
+        completed_minutes = sum(flow_minutes[: max(step - 1, 0)])
+        current_flow_minutes = flow_minutes[step - 1] if step - 1 < len(flow_minutes) else 1
+        progress = round(max(0.0, min(1.0, (completed_minutes + current_flow_minutes * step_progress) / total_minutes)), 4)
+
+        if context.progress_webhook_url:
+            payload = ProgressRun(
+                flow_run_id=UUID(flow_run_id) if flow_run_id else UUID(int=0),
+                project_name=project_name,
+                state="RUNNING",
+                timestamp=datetime.now(UTC),
+                step=step,
+                total_steps=total_steps,
+                flow_name=flow_name,
+                status=status,
+                progress=progress,
+                step_progress=round(step_progress, 4),
+                message=message,
+            )
+            try:
+                await send_webhook(context.progress_webhook_url, payload)
+            except Exception as e:
+                logger.warning(f"Progress webhook failed: {e}")
+
+        if flow_run_id:
+            try:
+                async with get_client() as client:
+                    await client.update_flow_run_labels(
+                        flow_run_id=UUID(flow_run_id),
+                        labels={
+                            "progress.step": step,
+                            "progress.total_steps": total_steps,
+                            "progress.flow_name": flow_name,
+                            "progress.status": status,
+                            "progress.progress": progress,
+                            "progress.step_progress": round(step_progress, 4),
+                            "progress.message": message,
+                        },
+                    )
+            except Exception as e:
+                logger.warning(f"Progress label update failed: {e}")
+
+    async def _send_completion(
+        self,
+        context: DeploymentContext,
+        flow_run_id: str,
+        project_name: str,
+        result: TResult | None,
+        error: str | None,
+    ) -> None:
+        """Send completion webhook."""
+        if not context.completion_webhook_url:
+            return
+        try:
+            now = datetime.now(UTC)
+            frid = UUID(flow_run_id) if flow_run_id else UUID(int=0)
+            payload: CompletedRun | FailedRun
+            if result is not None:
+                payload = CompletedRun(
+                    flow_run_id=frid,
+                    project_name=project_name,
+                    timestamp=now,
+                    state="COMPLETED",
+                    result=DeploymentResultData.model_validate(result.model_dump()),
+                )
+            else:
+                payload = FailedRun(
+                    flow_run_id=frid,
+                    project_name=project_name,
+                    timestamp=now,
+                    state="FAILED",
+                    error=error or "Unknown error",
+                )
+            await send_webhook(context.completion_webhook_url, payload)
+        except Exception as e:
+            logger.warning(f"Completion webhook failed: {e}")
+
+    @final
+    async def run(
+        self,
+        project_name: str,
+        documents: list[Document],
+        options: TOptions,
+        context: DeploymentContext,
+    ) -> TResult:
+        """Execute all flows with resume, per-flow uploads, and webhooks.
+
+        Args:
+            project_name: Unique identifier for this pipeline run (used as run_scope).
+            documents: Initial input documents for the first flow.
+            options: Flow options passed to each flow.
+            context: Deployment context with webhook URLs and document upload config.
+
+        Returns:
+            Typed deployment result built from all pipeline documents.
+        """
+        store = get_document_store()
+        total_steps = len(self.flows)
+        flow_run_id: str = str(runtime.flow_run.get_id()) if runtime.flow_run else ""  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownArgumentType]
+
+        # Write identity labels for polling endpoint
+        if flow_run_id:
+            try:
+                async with get_client() as client:
+                    await client.update_flow_run_labels(
+                        flow_run_id=UUID(flow_run_id),
+                        labels={"pipeline.project_name": project_name},
+                    )
+            except Exception as e:
+                logger.warning(f"Identity label update failed: {e}")
+
+        # Download additional input documents
+        input_docs = list(documents)
+        if context.input_documents_urls:
+            downloaded = await download_documents(list(context.input_documents_urls))
+            input_docs.extend(downloaded)
+
+        # Compute run scope AFTER downloads so the fingerprint includes all inputs
+        run_scope = _compute_run_scope(project_name, input_docs, options)
+
+        if not store and total_steps > 1:
+            logger.warning("No DocumentStore configured for multi-step pipeline — intermediate outputs will not accumulate between flows")
+
+        completion_sent = False
+
+        # Tracking lifecycle
+        tracking_svc = None
+        run_uuid: UUID | None = None
+        run_failed = False
+        try:
+            tracking_svc = get_tracking_service()
+            if tracking_svc:
+                run_uuid = UUID(flow_run_id) if flow_run_id else uuid4()
+                tracking_svc.set_run_context(run_id=run_uuid, project_name=project_name, flow_name=self.name, run_scope=run_scope)
+                tracking_svc.track_run_start(run_id=run_uuid, project_name=project_name, flow_name=self.name, run_scope=run_scope)
+        except Exception:
+            tracking_svc = None
+
+        # Set RunContext for the entire pipeline run
+        run_token = set_run_context(RunContext(run_scope=run_scope))
+        try:
+            # Save initial input documents to store
+            if store and input_docs:
+                await store.save_batch(input_docs, run_scope)
+
+            for step, flow_fn in enumerate(self.flows, start=1):
+                flow_name = getattr(flow_fn, "name", flow_fn.__name__)
+                flow_run_id = str(runtime.flow_run.get_id()) if runtime.flow_run else ""  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownArgumentType]
+
+                # Resume check: skip if output documents already exist in store
+                output_types = getattr(flow_fn, "output_document_types", [])
+                if store and output_types:
+                    all_outputs_exist = all([await store.has_documents(run_scope, ot) for ot in output_types])
+                    if all_outputs_exist:
+                        logger.info(f"[{step}/{total_steps}] Resume: skipping {flow_name} (outputs exist)")
+                        await self._send_progress(
+                            context,
+                            flow_run_id,
+                            project_name,
+                            step,
+                            total_steps,
+                            flow_name,
+                            "cached",
+                            step_progress=1.0,
+                            message=f"Resumed from store: {flow_name}",
+                        )
+                        continue
+
+                # Prefect state hooks
+                active_flow = flow_fn
+                if context.status_webhook_url:
+                    hooks = self._build_status_hooks(context, flow_run_id, project_name, step, total_steps, flow_name)
+                    active_flow = flow_fn.with_options(**hooks)
+                    _reattach_flow_metadata(flow_fn, active_flow)
+
+                # Progress: started
+                await self._send_progress(
+                    context,
+                    flow_run_id,
+                    project_name,
+                    step,
+                    total_steps,
+                    flow_name,
+                    "started",
+                    step_progress=0.0,
+                    message=f"Starting: {flow_name}",
+                )
+
+                logger.info(f"[{step}/{total_steps}] Starting: {flow_name}")
+
+                # Load input documents from store
+                input_types = getattr(flow_fn, "input_document_types", [])
+                if store and input_types:
+                    current_docs = await store.load(run_scope, input_types)
+                else:
+                    current_docs = input_docs
+
+                try:
+                    await active_flow(project_name, current_docs, options)
+                except Exception as e:
+                    # Upload partial results on failure
+                    if context.output_documents_urls and store:
+                        all_docs = await store.load(run_scope, self._all_document_types())
+                        await upload_documents(all_docs, context.output_documents_urls)
+                    await self._send_completion(context, flow_run_id, project_name, result=None, error=str(e))
+                    completion_sent = True
+                    raise
+
+                # Per-flow upload (load from store since @pipeline_flow saves there)
+                if context.output_documents_urls and store and output_types:
+                    flow_docs = await store.load(run_scope, output_types)
+                    await upload_documents(flow_docs, context.output_documents_urls)
+
+                # Progress: completed
+                await self._send_progress(
+                    context,
+                    flow_run_id,
+                    project_name,
+                    step,
+                    total_steps,
+                    flow_name,
+                    "completed",
+                    step_progress=1.0,
+                    message=f"Completed: {flow_name}",
+                )
+
+                logger.info(f"[{step}/{total_steps}] Completed: {flow_name}")
+
+            # Build result from all documents in store
+            if store:
+                all_docs = await store.load(run_scope, self._all_document_types())
+            else:
+                all_docs = input_docs
+            result = self.build_result(project_name, all_docs, options)
+            await self._send_completion(context, flow_run_id, project_name, result=result, error=None)
+            return result
+
+        except Exception as e:
+            run_failed = True
+            if not completion_sent:
+                await self._send_completion(context, flow_run_id, project_name, result=None, error=str(e))
+            raise
+        finally:
+            reset_run_context(run_token)
+            store = get_document_store()
+            if store:
+                with contextlib.suppress(Exception):
+                    store.flush()
+            if (svc := tracking_svc) is not None and run_uuid is not None:
+                with contextlib.suppress(Exception):
+                    svc.track_run_end(run_id=run_uuid, status=RunStatus.FAILED if run_failed else RunStatus.COMPLETED)
+                    svc.flush()
+
+    @final
+    def run_local(
+        self,
+        project_name: str,
+        documents: list[Document],
+        options: TOptions,
+        context: DeploymentContext | None = None,
+        output_dir: Path | None = None,
+    ) -> TResult:
+        """Run locally with Prefect test harness and in-memory document store.
+
+        Args:
+            project_name: Pipeline run identifier.
+            documents: Initial input documents.
+            options: Flow options.
+            context: Optional deployment context (defaults to empty).
+            output_dir: Optional directory for writing result.json.
+
+        Returns:
+            Typed deployment result.
+        """
+        if context is None:
+            context = DeploymentContext()
+
+        if output_dir:
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+        store = MemoryDocumentStore()
+        set_document_store(store)
+        try:
+            with prefect_test_harness(), disable_run_logger():
+                result = asyncio.run(self.run(project_name, documents, options, context))
+        finally:
+            store.shutdown()
+            set_document_store(None)
+
+        if output_dir:
+            (output_dir / "result.json").write_text(result.model_dump_json(indent=2))
+
+        return result
+
+    @final
+    def run_cli(
+        self,
+        initializer: Callable[[TOptions], tuple[str, list[Document]]] | None = None,
+        trace_name: str | None = None,
+    ) -> None:
+        """Execute pipeline from CLI arguments with --start/--end step control.
+
+        Args:
+            initializer: Optional callback returning (project_name, documents) from options.
+            trace_name: Optional Laminar trace span name prefix.
+        """
+        if len(sys.argv) == 1:
+            sys.argv.append("--help")
+
+        setup_logging()
+        try:
+            initialize_observability()
+            logger.info("Observability initialized.")
+        except Exception as e:
+            logger.warning(f"Failed to initialize observability: {e}")
+            with contextlib.suppress(Exception):
+                Laminar.initialize(export_timeout_seconds=15)
+
+        deployment = self
+
+        class _CliOptions(
+            deployment.options_type,
+            cli_parse_args=True,
+            cli_kebab_case=True,
+            cli_exit_on_error=True,
+            cli_prog_name=deployment.name,
+            cli_use_class_docs_for_groups=True,
+        ):
+            working_directory: CliPositionalArg[Path]
+            project_name: str | None = None
+            start: int = 1
+            end: int | None = None
+            no_trace: bool = False
+
+            model_config = SettingsConfigDict(frozen=True, extra="ignore")
+
+        opts = cast(TOptions, _CliOptions())  # type: ignore[reportCallIssue]
+
+        wd = cast(Path, opts.working_directory)  # pyright: ignore[reportAttributeAccessIssue]
+        wd.mkdir(parents=True, exist_ok=True)
+
+        project_name = cast(str, opts.project_name or wd.name)  # pyright: ignore[reportAttributeAccessIssue]
+        start_step = getattr(opts, "start", 1)
+        end_step = getattr(opts, "end", None)
+        no_trace = getattr(opts, "no_trace", False)
+
+        # Set up local debug tracing (writes to <working_dir>/.trace)
+        debug_processor: LocalDebugSpanProcessor | None = None
+        if not no_trace:
+            try:
+                trace_path = wd / ".trace"
+                trace_path.mkdir(parents=True, exist_ok=True)
+                debug_config = TraceDebugConfig(path=trace_path, max_traces=20)
+                debug_writer = LocalTraceWriter(debug_config)
+                debug_processor = LocalDebugSpanProcessor(debug_writer)
+                provider: Any = otel_trace.get_tracer_provider()
+                if hasattr(provider, "add_span_processor"):
+                    provider.add_span_processor(debug_processor)
+                    logger.info(f"Local debug tracing enabled at {trace_path}")
+            except Exception as e:
+                logger.warning(f"Failed to set up local debug tracing: {e}")
+                debug_processor = None
+
+        # Initialize document store — ClickHouse when configured, local filesystem otherwise
+        summary_generator = _build_summary_generator()
+        if settings.clickhouse_host:
+            store = create_document_store(settings, summary_generator=summary_generator)
+        else:
+            store = LocalDocumentStore(base_path=wd, summary_generator=summary_generator)
+        set_document_store(store)
+
+        # Initialize documents (always run initializer for run scope fingerprinting,
+        # even when start_step > 1, so --start N resumes find the correct scope)
+        initial_documents: list[Document] = []
+        if initializer:
+            _, initial_documents = initializer(opts)
+
+        context = DeploymentContext()
+
+        with ExitStack() as stack:
+            if trace_name:
+                stack.enter_context(
+                    Laminar.start_as_current_span(
+                        name=f"{trace_name}-{project_name}",
+                        input=[opts.model_dump_json()],
+                    )
+                )
+
+            under_pytest = "PYTEST_CURRENT_TEST" in os.environ or "pytest" in sys.modules
+            if not settings.prefect_api_key and not under_pytest:
+                stack.enter_context(prefect_test_harness())
+                stack.enter_context(disable_run_logger())
+
+            result = asyncio.run(
+                self._run_with_steps(
+                    project_name=project_name,
+                    options=opts,
+                    context=context,
+                    start_step=start_step,
+                    end_step=end_step,
+                    initial_documents=initial_documents,
+                )
+            )
+
+        result_file = wd / "result.json"
+        result_file.write_text(result.model_dump_json(indent=2))
+        logger.info(f"Result saved to {result_file}")
+
+        # Shutdown background workers (debug tracing, document summaries, tracking)
+        if debug_processor is not None:
+            debug_processor.shutdown()
+        store = get_document_store()
+        if store:
+            store.shutdown()
+        tracking_svc = get_tracking_service()
+        if tracking_svc:
+            tracking_svc.shutdown()
+
+    async def _run_with_steps(
+        self,
+        project_name: str,
+        options: TOptions,
+        context: DeploymentContext,
+        start_step: int = 1,
+        end_step: int | None = None,
+        initial_documents: list[Document] | None = None,
+    ) -> TResult:
+        """Run pipeline with start/end step control and DocumentStore-based resume."""
+        store = get_document_store()
+        if end_step is None:
+            end_step = len(self.flows)
+
+        total_steps = len(self.flows)
+        run_scope = _compute_run_scope(project_name, initial_documents or [], options)
+
+        # Tracking lifecycle for CLI path
+        tracking_svc = None
+        run_uuid: UUID | None = None
+        run_failed = False
+        try:
+            tracking_svc = get_tracking_service()
+            if tracking_svc:
+                run_uuid = uuid4()
+                tracking_svc.set_run_context(run_id=run_uuid, project_name=project_name, flow_name=self.name, run_scope=run_scope)
+                tracking_svc.track_run_start(run_id=run_uuid, project_name=project_name, flow_name=self.name, run_scope=run_scope)
+        except Exception:
+            tracking_svc = None
+
+        # Set RunContext for the entire pipeline run
+        run_token = set_run_context(RunContext(run_scope=run_scope))
+        try:
+            # Save initial documents to store
+            if store and initial_documents:
+                await store.save_batch(initial_documents, run_scope)
+
+            for i in range(start_step - 1, end_step):
+                step = i + 1
+                flow_fn = self.flows[i]
+                flow_name = getattr(flow_fn, "name", flow_fn.__name__)
+                logger.info(f"--- [Step {step}/{total_steps}] {flow_name} ---")
+
+                # Resume check: skip if output documents already exist
+                output_types = getattr(flow_fn, "output_document_types", [])
+                if store and output_types:
+                    all_outputs_exist = all([await store.has_documents(run_scope, ot) for ot in output_types])
+                    if all_outputs_exist:
+                        logger.info(f"--- [Step {step}/{total_steps}] Skipping {flow_name} (outputs exist) ---")
+                        continue
+
+                # Load inputs from store
+                input_types = getattr(flow_fn, "input_document_types", [])
+                if store and input_types:
+                    current_docs = await store.load(run_scope, input_types)
+                else:
+                    current_docs = initial_documents or []
+
+                await flow_fn(project_name, current_docs, options)
+
+            # Build result from all documents in store
+            if store:
+                all_docs = await store.load(run_scope, self._all_document_types())
+            else:
+                all_docs = initial_documents or []
+            return self.build_result(project_name, all_docs, options)
+        except Exception:
+            run_failed = True
+            raise
+        finally:
+            reset_run_context(run_token)
+            store = get_document_store()
+            if store:
+                with contextlib.suppress(Exception):
+                    store.flush()
+            if (svc := tracking_svc) is not None and run_uuid is not None:
+                with contextlib.suppress(Exception):
+                    svc.track_run_end(run_id=run_uuid, status=RunStatus.FAILED if run_failed else RunStatus.COMPLETED)
+                    svc.flush()
+
+    @final
+    def as_prefect_flow(self) -> Callable[..., Any]:
+        """Generate a Prefect flow for production deployment.
+
+        Returns:
+            Async Prefect flow callable that initializes DocumentStore from settings.
+        """
+        deployment = self
+
+        @flow(
+            name=self.name,
+            flow_run_name=f"{self.name}-{{project_name}}",
+            persist_result=True,
+            result_serializer="json",
+        )
+        async def _deployment_flow(
+            project_name: str,
+            documents: list[Document],
+            options: FlowOptions,
+            context: DeploymentContext,
+        ) -> DeploymentResult:
+            store = create_document_store(
+                settings,
+                summary_generator=_build_summary_generator(),
+            )
+            set_document_store(store)
+            try:
+                return await deployment.run(project_name, documents, cast(Any, options), context)
+            finally:
+                store.shutdown()
+                set_document_store(None)
+
+        return _deployment_flow
+
+
+__all__ = [
+    "DeploymentContext",
+    "DeploymentResult",
+    "PipelineDeployment",
+]