reader-workbench 1.0.0__py3-none-any.whl

reader_workbench/workbench/engine/contracts.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+from reader_workbench.contracts import ContractCatalog
+from reader_workbench.errors import ContractError, ExecutionError
+from reader_workbench.workbench.graph import OutputRef
+from reader_workbench.workbench.ports import InputPortSpec, OutputPortSpec
+from reader_workbench.workbench.records import PathDescription, SourceRecordCollection
+from reader_workbench.workbench.registry import Plugin
+
+
+def _assert_input_ports(
+    plugin: Plugin,
+    inputs: dict[str, Any],
+    *,
+    contracts: ContractCatalog,
+    where: str,
+) -> None:
+    declared = plugin.input_ports()
+    allowed = set(declared)
+    for name, port in declared.items():
+        if name not in inputs:
+            if port.optional:
+                continue
+            raise ExecutionError(f"[{where}] input '{name}' is required by plugin but not provided in 'reads'")
+        value = inputs[name]
+        _assert_input_port_value(
+            name=name,
+            port=port,
+            value=value,
+            contracts=contracts,
+            where=where,
+        )
+    extra = sorted(set(inputs) - allowed)
+    if extra:
+        raise ExecutionError(f"[{where}] unexpected inputs provided: {extra} (allowed: {sorted(allowed)})")
+
+
+def _assert_output_ports(
+    output_ports: Mapping[str, OutputPortSpec],
+    outputs: dict[str, Any],
+    *,
+    contracts: ContractCatalog,
+    where: str,
+) -> None:
+    if set(outputs) != set(output_ports):
+        raise ExecutionError(f"[{where}] plugin must emit outputs {sorted(output_ports)} but emitted {sorted(outputs)}")
+    for name, port in output_ports.items():
+        value = outputs[name]
+        if port.kind == "dataframe":
+            if not isinstance(value, pd.DataFrame):
+                raise ExecutionError(
+                    f"[{where}] output '{name}' must be a DataFrame for dataframe port kind, got {type(value).__name__}"
+                )
+            try:
+                contracts.validate(value, contract_id=port.contract or "", where=where)
+            except ContractError as err:
+                raise ExecutionError(str(err)) from err
+            continue
+        if isinstance(value, pd.DataFrame):
+            msg = f"[{where}] output '{name}' is declared as {port.kind} but returned a DataFrame"
+            raise ExecutionError(msg)
+        _coerce_file_output(port=port, value=value, where=where, name=name)
+
+
+def _resolve_runtime_output_ports(
+    plugin: Plugin,
+    *,
+    inputs: Mapping[str, Any],
+    outputs: Mapping[str, Any],
+    cfg: Any,
+    contracts: ContractCatalog,
+    where: str,
+) -> dict[str, OutputPortSpec]:
+    declared = dict(plugin.output_ports())
+    resolved = dict(plugin.resolve_output_ports(inputs=inputs, outputs=outputs, cfg=cfg, where=where))
+    if set(resolved) != set(declared):
+        raise ExecutionError(
+            f"[{where}] runtime output ports must match declared outputs: "
+            f"declared={sorted(declared)} resolved={sorted(resolved)}"
+        )
+    for name, declared_port in declared.items():
+        resolved_port = resolved[name]
+        if resolved_port.kind != declared_port.kind:
+            raise ExecutionError(
+                f"[{where}] runtime output port '{name}' changed kind from {declared_port.kind!r} "
+                f"to {resolved_port.kind!r}"
+            )
+        if declared_port.kind != "dataframe":
+            if resolved_port.contract is not None:
+                raise ExecutionError(
+                    f"[{where}] runtime output port '{name}' of kind {declared_port.kind!r} "
+                    "must not resolve a dataframe contract"
+                )
+            continue
+        if declared_port.contract is None or resolved_port.contract is None:
+            raise ExecutionError(f"[{where}] dataframe output port '{name}' must resolve a non-empty contract id")
+        try:
+            if contracts.satisfies(actual=resolved_port.contract, expected=declared_port.contract):
+                continue
+        except ContractError as err:
+            raise ExecutionError(str(err)) from err
+        raise ExecutionError(
+            f"[{where}] runtime contract for output '{name}' must satisfy declared contract "
+            f"{declared_port.contract!r}, got {resolved_port.contract!r}"
+        )
+    return resolved
+
+
+def _resolve_output_labels(
+    *,
+    step_id: str,
+    output_ports: Mapping[str, OutputPortSpec],
+    writes: dict[str, OutputRef],
+) -> dict[str, OutputRef]:
+    unknown = sorted(set(writes) - set(output_ports))
+    if unknown:
+        raise ExecutionError(
+            f"[{step_id}] writes includes unknown outputs: {unknown} (expected: {sorted(output_ports)})"
+        )
+    labels: dict[str, OutputRef] = {}
+    for out_name, port in output_ports.items():
+        if port.kind != "dataframe":
+            if out_name in writes:
+                raise ExecutionError(
+                    f"[{step_id}] writes cannot target output '{out_name}' (port kind is {port.kind!r})."
+                )
+            continue
+        label = writes[out_name] if out_name in writes else OutputRef(record_id=f"{step_id}/{out_name}")
+        if not isinstance(label, OutputRef) or not label.record_id.strip():
+            raise ExecutionError(f"[{step_id}] writes for '{out_name}' must be a non-empty record ref.")
+        labels[out_name] = label
+    if len({ref.record_id for ref in labels.values()}) != len(labels):
+        rendered = {name: ref.record_id for name, ref in labels.items()}
+        raise ExecutionError(f"[{step_id}] writes produce duplicate output labels: {rendered}")
+    return labels
+
+
+def collect_file_output_paths(
+    *,
+    output_ports: Mapping[str, OutputPortSpec],
+    outputs: Mapping[str, Any],
+    where: str,
+) -> list[Path]:
+    collected: list[Path] = []
+    for name, port in output_ports.items():
+        if port.kind == "dataframe":
+            continue
+        collected.extend(_coerce_file_output(port=port, value=outputs[name], where=where, name=name))
+    return collected
+
+
+def collect_file_output_descriptions(
+    *,
+    output_ports: Mapping[str, OutputPortSpec],
+    outputs: Mapping[str, Any],
+) -> list[PathDescription]:
+    descriptions: list[PathDescription] = []
+    for name, port in output_ports.items():
+        if port.kind == "dataframe":
+            continue
+        value = outputs[name]
+        if isinstance(value, PathDescription):
+            descriptions.append(value)
+        elif isinstance(value, Sequence) and not isinstance(value, str | bytes):
+            descriptions.extend(item for item in value if isinstance(item, PathDescription))
+    return descriptions
+
+
+def _assert_input_port_value(
+    *,
+    name: str,
+    port: InputPortSpec,
+    value: Any,
+    contracts: ContractCatalog,
+    where: str,
+) -> None:
+    if port.kind == "dataframe":
+        actual_contract = getattr(value, "contract_id", None)
+        if actual_contract is None:
+            raise ExecutionError(
+                f"[{where}] input '{name}' expects a dataframe artifact but got {type(value).__name__}"
+            )
+        if port.contract is None:
+            return
+        try:
+            if contracts.satisfies(actual=actual_contract, expected=port.contract):
+                return
+        except ContractError as err:
+            msg = str(err)
+        else:
+            msg = f"[{where}] input '{name}' must be contract {port.contract} but got {actual_contract}"
+        raise ExecutionError(msg)
+    if port.kind == "record_collection":
+        if not isinstance(value, SourceRecordCollection):
+            raise ExecutionError(
+                f"[{where}] input '{name}' expects a source record collection but got {type(value).__name__}"
+            )
+        for item in value:
+            actual_contract = getattr(item.record, "contract_id", None)
+            if actual_contract is None:
+                raise ExecutionError(
+                    f"[{where}] input '{name}' source {item.ref.resource_id!r} is not a dataframe artifact"
+                )
+            if port.contract is None:
+                continue
+            try:
+                compatible = contracts.satisfies(actual=actual_contract, expected=port.contract)
+            except ContractError as err:
+                raise ExecutionError(str(err)) from err
+            if not compatible:
+                raise ExecutionError(
+                    f"[{where}] input '{name}' source {item.ref.resource_id!r} must satisfy "
+                    f"contract {port.contract!r} but got {actual_contract!r}"
+                )
+        return
+    if port.kind == "file_path":
+        if isinstance(value, Path):
+            return
+        raise ExecutionError(f"[{where}] input '{name}' expects a file path but got {type(value).__name__}")
+    if port.kind == "file_set":
+        if isinstance(value, tuple) and value and all(isinstance(item, Path) for item in value):
+            return
+        raise ExecutionError(
+            f"[{where}] input '{name}' expects a non-empty tuple of file paths but got {type(value).__name__}"
+        )
+    if port.kind == "file_bundle":
+        files = getattr(value, "files", None)
+        if files is None:
+            raise ExecutionError(f"[{where}] input '{name}' expects a file bundle but got {type(value).__name__}")
+        return
+    raise ExecutionError(f"[{where}] input '{name}' uses unknown port kind {port.kind!r}")
+
+
+def _coerce_file_output(
+    *,
+    port: OutputPortSpec,
+    value: Any,
+    where: str,
+    name: str,
+) -> list[Path]:
+    if port.kind == "file_path":
+        if isinstance(value, PathDescription):
+            return [value.path]
+        if isinstance(value, str | Path):
+            return [_validated_output_path(value, where=where, name=name)]
+        raise ExecutionError(f"[{where}] output '{name}' must be a path-like value for file_path ports")
+    if port.kind == "file_bundle":
+        if isinstance(value, PathDescription):
+            return [value.path]
+        if isinstance(value, str | Path):
+            return [_validated_output_path(value, where=where, name=name)]
+        if isinstance(value, Sequence) and not isinstance(value, str | bytes):
+            if not value:
+                raise ExecutionError(f"[{where}] output '{name}' file bundle must contain at least one file")
+            paths: list[Path] = []
+            for item in value:
+                if isinstance(item, PathDescription):
+                    paths.append(item.path)
+                    continue
+                if not isinstance(item, str | Path):
+                    raise ExecutionError(
+                        f"[{where}] output '{name}' file bundle entries must be path-like or PathDescription, "
+                        f"got {type(item).__name__}"
+                    )
+                paths.append(_validated_output_path(item, where=where, name=name))
+            return paths
+        raise ExecutionError(f"[{where}] output '{name}' must be a list of path-like values for file_bundle ports")
+    raise ExecutionError(f"[{where}] output '{name}' uses unknown non-dataframe port kind {port.kind!r}")
+
+
+def _validated_output_path(value: str | Path, *, where: str, name: str) -> Path:
+    if isinstance(value, str) and not value.strip():
+        raise ExecutionError(f"[{where}] output '{name}' paths must be non-empty")
+    path = Path(value)
+    if path == Path("."):
+        raise ExecutionError(f"[{where}] output '{name}' paths must identify files")
+    return path

reader_workbench/workbench/engine/execution.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from rich.console import Console
+from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+
+from reader_workbench.errors import ExecutionError, ReaderError
+from reader_workbench.workbench.context import RunContext
+from reader_workbench.workbench.graph import (
+    FileRef,
+    ProvenanceInput,
+    RecordCollectionRef,
+    RecordRef,
+    ResourceRef,
+    SourceRecordRef,
+)
+from reader_workbench.workbench.records import PathDescription, RecordInputEvidence, RecordStore
+
+from ._shared import digest_cfg
+from .contracts import (
+    _assert_input_ports,
+    _assert_output_ports,
+    _resolve_output_labels,
+    _resolve_runtime_output_ports,
+    collect_file_output_descriptions,
+    collect_file_output_paths,
+)
+from .file_outputs import FileOutputTransaction
+from .inputs import _resolve_inputs, resolve_missing_file_inputs
+
+
+def _runtime_provenance_inputs(*, step: Any, inputs: dict[str, Any]) -> list[ProvenanceInput]:
+    provenance: list[ProvenanceInput] = []
+    declared_reads = dict(step.reads or {})
+    for label in sorted(inputs):
+        ref = declared_reads.get(label)
+        value = inputs[label]
+        if isinstance(ref, RecordCollectionRef):
+            provenance.extend(
+                ProvenanceInput(
+                    label=f"{label}[{item.resource_id}]",
+                    ref=item,
+                    discovery_policy="source_record",
+                )
+                for item in ref.records
+            )
+            continue
+        if isinstance(value, tuple) and value and all(isinstance(item, Path) for item in value):
+            for index, path in enumerate(value):
+                item_ref = ref if ref is not None and len(value) == 1 else FileRef(path=path)
+                if isinstance(item_ref, ResourceRef):
+                    policy = "declared_resource"
+                elif ref is not None:
+                    policy = "declared_file"
+                else:
+                    policy = "plugin_discovery"
+                provenance.append(
+                    ProvenanceInput(
+                        label=f"{label}[{index}]",
+                        ref=item_ref,
+                        discovery_policy=policy,
+                    )
+                )
+            continue
+        if ref is None:
+            if not isinstance(value, Path):
+                raise ExecutionError(
+                    f"{step.id}: resolved input {label!r} has no declared reference and is not a file path"
+                )
+            ref = FileRef(path=value)
+            discovery_policy = "plugin_discovery"
+        elif isinstance(ref, RecordRef):
+            discovery_policy = "record"
+        elif isinstance(ref, SourceRecordRef):
+            discovery_policy = "source_record"
+        elif isinstance(ref, ResourceRef):
+            discovery_policy = "declared_resource"
+        else:
+            discovery_policy = "declared_file"
+        provenance.append(ProvenanceInput(label=label, ref=ref, discovery_policy=discovery_policy))
+    return provenance
+
+
+def _persist_dataframe_outputs(
+    *,
+    store: RecordStore,
+    ctx: RunContext,
+    step: Any,
+    cfg: Any,
+    input_evidence: tuple[RecordInputEvidence, ...],
+    outputs: dict[str, Any],
+    resolved_output_ports: dict[str, Any],
+    output_labels: dict[str, Any],
+    phase: str,
+) -> None:
+    for out_name, obj in outputs.items():
+        port = resolved_output_ports[out_name]
+        if port.kind != "dataframe":
+            continue
+        if isinstance(obj, pd.DataFrame):
+            store.persist_dataframe(
+                producer_id=step.id,
+                producer_kind="pipeline" if phase == "pipeline" else ("plot" if phase == "plots" else "export"),
+                producer_plugin=step.plugin,
+                out_name=out_name,
+                record_id=output_labels[out_name].record_id,
+                df=obj,
+                contract_id=port.contract or "",
+                inputs=input_evidence,
+                config_digest=ctx.config_digest,
+                producer_config_digest=digest_cfg(cfg),
+                source_recipe=step.source_recipe,
+            )
+            continue
+        raise ExecutionError(f"{phase} {step.id}: unsupported output type for {out_name}")
+
+
+def _persist_file_bundle_record(
+    *,
+    store: RecordStore,
+    ctx: RunContext,
+    step: Any,
+    cfg: Any,
+    input_evidence: tuple[RecordInputEvidence, ...],
+    outputs: dict[str, Any],
+    resolved_output_ports: dict[str, Any],
+    phase: str,
+    description: str,
+    protocol_figure_description: str | None,
+) -> None:
+    explicit_files = collect_file_output_paths(
+        output_ports=resolved_output_ports,
+        outputs=outputs,
+        where=f"{phase}:{step.id}",
+    )
+    if not explicit_files:
+        raise ExecutionError(f"{phase} {step.id}: must emit at least one explicit file output")
+    producer_kind = "plot" if phase == "plots" else "export"
+    record_files = sorted(
+        {_record_output_path(path, outputs_dir=ctx.outputs_dir) for path in explicit_files},
+        key=str,
+    )
+    path_descriptions: tuple[PathDescription, ...] = ()
+    if phase == "plots":
+        explicit_descriptions = collect_file_output_descriptions(
+            output_ports=resolved_output_ports,
+            outputs=outputs,
+        )
+        explicit_by_path: dict[Path, str] = {}
+        for item in explicit_descriptions:
+            record_path = _record_output_path(item.path, outputs_dir=ctx.outputs_dir)
+            if record_path in explicit_by_path:
+                raise ExecutionError(f"plots {step.id}: duplicate descriptions for {record_path}")
+            explicit_by_path[record_path] = item.description
+        fallback_description = protocol_figure_description or description
+        path_descriptions = tuple(
+            PathDescription(path=path, description=explicit_by_path.get(path, fallback_description))
+            for path in record_files
+        )
+    store.append_file_bundle(
+        producer_kind=producer_kind,
+        producer_id=step.id,
+        producer_plugin=step.plugin,
+        record_id=f"{producer_kind}:{step.id}",
+        inputs=input_evidence,
+        config_digest=ctx.config_digest,
+        producer_config_digest=digest_cfg(cfg),
+        files=record_files,
+        description=description,
+        path_descriptions=path_descriptions,
+        source_recipe=step.source_recipe,
+    )
+
+
+def _record_output_path(path: Path, *, outputs_dir: Path) -> Path:
+    try:
+        return path.relative_to(outputs_dir)
+    except ValueError:
+        return path
+
+
+def _protocol_figure_description(*, ctx: RunContext, step: Any) -> str | None:
+    for figure in ctx.protocol.descriptor.figures:
+        if figure.id == step.id:
+            return figure.summary
+    return None
+
+
+def execute_step(*, step: Any, phase: str, store: RecordStore, ctx: RunContext, registry: Any) -> None:
+    descriptor = registry.resolve_descriptor(step.plugin)
+    plugin_cls = descriptor.cls
+    if ctx.protocol is None:
+        raise ExecutionError(f"{phase} {step.id}: run context is missing a bound protocol")
+    effective_with = ctx.protocol.effective_plugin_config(plugin_id=step.plugin, step_with=(step.with_ or {}))
+    cfg = plugin_cls.ConfigModel.model_validate(effective_with)
+    plugin = plugin_cls()
+    plugin.bind_runtime(descriptor=descriptor, contracts=registry.contracts)
+    input_ports = plugin.input_ports()
+    output_ports = plugin.output_ports()
+    output_labels = _resolve_output_labels(
+        step_id=step.id,
+        output_ports=output_ports,
+        writes=(step.writes or {}),
+    )
+    inputs = _resolve_inputs(store, step.reads or {}, input_ports=input_ports, exp_dir=ctx.exp_dir)
+    inputs = resolve_missing_file_inputs(
+        plugin=plugin,
+        exp_dir=ctx.exp_dir,
+        cfg=cfg,
+        inputs=inputs,
+        input_ports=input_ports,
+    )
+    _assert_input_ports(
+        plugin,
+        inputs,
+        contracts=registry.contracts,
+        where=step.id,
+    )
+    provenance_inputs = _runtime_provenance_inputs(step=step, inputs=inputs)
+    input_evidence = store.capture_inputs(provenance_inputs, resolved_inputs=inputs)
+
+    plug_inputs: dict[str, Any] = {}
+    for key, value in inputs.items():
+        if hasattr(value, "load_dataframe"):
+            plug_inputs[key] = value.load_dataframe()
+        else:
+            plug_inputs[key] = value
+
+    transaction_context = (
+        FileOutputTransaction(context=ctx, step_id=step.id, phase=phase)
+        if phase in {"plots", "exports"}
+        else nullcontext()
+    )
+    with transaction_context as file_transaction:
+        plugin_context = file_transaction.context if file_transaction is not None else ctx
+        try:
+            outputs = plugin.run(plugin_context, plug_inputs, cfg)
+        except ReaderError:
+            raise
+        except Exception as err:
+            raise ExecutionError(f"{phase} {step.id} crashed: {err}") from err
+
+        resolved_output_ports = _resolve_runtime_output_ports(
+            plugin,
+            inputs=inputs,
+            outputs=outputs,
+            cfg=cfg,
+            contracts=registry.contracts,
+            where=step.id,
+        )
+        _assert_output_ports(
+            resolved_output_ports,
+            outputs,
+            contracts=registry.contracts,
+            where=step.id,
+        )
+        _persist_dataframe_outputs(
+            store=store,
+            ctx=ctx,
+            step=step,
+            cfg=cfg,
+            input_evidence=input_evidence,
+            outputs=outputs,
+            resolved_output_ports=resolved_output_ports,
+            output_labels=output_labels,
+            phase=phase,
+        )
+        if phase in {"plots", "exports"}:
+            outputs = file_transaction.promote(
+                outputs=outputs,
+                output_ports=resolved_output_ports,
+                where=f"{phase}:{step.id}",
+            )
+            protocol_figure_description = _protocol_figure_description(ctx=ctx, step=step) if phase == "plots" else None
+            description = protocol_figure_description or descriptor.summary
+            _persist_file_bundle_record(
+                store=store,
+                ctx=ctx,
+                step=step,
+                cfg=cfg,
+                input_evidence=input_evidence,
+                outputs=outputs,
+                resolved_output_ports=resolved_output_ports,
+                phase=phase,
+                description=description,
+                protocol_figure_description=protocol_figure_description,
+            )
+            file_transaction.commit()
+
+
+def run_steps(
+    *,
+    items: list[Any],
+    phase: str,
+    verbose: bool,
+    console: Console,
+    ctx: RunContext,
+    store: RecordStore,
+    registry: Any,
+) -> None:
+    if not items:
+        return
+    if verbose:
+        for ordinal, step in enumerate(items, 1):
+            ctx.logger.info("→ %s %s [%d/%d] plugin=%s", phase, step.id, ordinal, len(items), step.plugin)
+            execute_step(step=step, phase=phase, store=store, ctx=ctx, registry=registry)
+        return
+
+    with Progress(
+        SpinnerColumn(style="accent"),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.completed}/{task.total}"),
+        TimeElapsedColumn(),
+        console=console,
+        transient=True,
+    ) as progress:
+        task = progress.add_task(f"{phase.title()} ({len(items)} steps)", total=len(items))
+        for step in items:
+            progress.update(task, description=f"{phase}: {step.id}")
+            execute_step(step=step, phase=phase, store=store, ctx=ctx, registry=registry)
+            progress.advance(task)