dirigent-core 0.9.0__py3-none-any.whl

dirigent_core/documentschema.py

@@ -0,0 +1,92 @@
+"""The JSON Schema an editor checks a ``dirigent/v1`` document against.
+
+The document format is a pydantic model and every block's config is another one, so a document
+is only half-describable without the catalog: ``steps.push.config`` is whatever the block named
+in ``steps.push.block`` accepts. This module composes the two into one Draft 2020-12 schema, so
+an editor squiggles an unknown key in the place an apply would refuse it.
+"""
+
+from collections.abc import Mapping, Sequence
+from typing import Any, Final, cast
+
+from dirigent_client.schemas import Catalog
+from dirigent_common import HumaneJsonSchema, JsonMap
+from dirigent_core.engine.definition import PipelineDefinition, TriggersDefinition
+
+#: The dialect the composed schema declares, which is the one the engine validates parameters in.
+DIALECT: Final = "https://json-schema.org/draft/2020-12/schema"
+
+TITLE: Final = "dirigent/v1 document"
+
+
+def document_schema(catalog: Catalog) -> JsonMap:
+    """Compose both document kinds with this instance's own block config schemas.
+
+    The two branches share ``$defs``, and ``kind`` is what an editor picks between them by.
+    """
+    pipeline = _pipeline_schema(catalog)
+    triggers = TriggersDefinition.model_json_schema(ref_template="#/$defs/{model}", schema_generator=HumaneJsonSchema)
+    definitions = cast("dict[str, Any]", pipeline.pop("$defs", {}))
+    definitions.update(cast("dict[str, Any]", triggers.pop("$defs", {})))
+    return {
+        "$schema": DIALECT,
+        "title": TITLE,
+        "$defs": definitions,
+        "oneOf": [pipeline, triggers],
+    }
+
+
+def _pipeline_schema(catalog: Catalog) -> JsonMap:
+    """The pipeline branch: the document format with every block's config grafted onto its step."""
+    schema = PipelineDefinition.model_json_schema(schema_generator=HumaneJsonSchema)
+    definitions = cast("dict[str, Any]", schema.setdefault("$defs", {}))
+    step = cast("dict[str, Any]", definitions["StepDefinition"])
+    blocks = sorted(catalog.blocks, key=lambda entry: entry.id)
+
+    if blocks:
+        cast("dict[str, Any]", step["properties"])["block"]["enum"] = [entry.id for entry in blocks]
+
+    cases: list[JsonMap] = []
+    for entry in blocks:
+        config, lifted = _lift(entry.config_schema, entry.id)
+        definitions.update(lifted)
+        cases.append(
+            {
+                "if": {"properties": {"block": {"const": entry.id}}, "required": ["block"]},
+                "then": {"properties": {"config": config}},
+            }
+        )
+    if cases:
+        step["allOf"] = cases
+    return schema
+
+
+def _lift(config_schema: JsonMap, block_id: str) -> tuple[JsonMap, dict[str, Any]]:
+    """Move one block's ``$defs`` into the document's, under names only that block uses.
+
+    Two blocks contribute a ``Size`` each and a shared ``$defs`` has room for one, so every
+    definition is renamed for the block it came from and every ``$ref`` to it follows.
+    """
+    own = cast("dict[str, Any]", config_schema.get("$defs", {}))
+    body = {key: value for key, value in config_schema.items() if key != "$defs"}
+    lifted = {f"{block_id}.{name}": _repoint(value, block_id) for name, value in own.items()}
+    return cast("JsonMap", _repoint(body, block_id)), lifted
+
+
+def _repoint(value: object, block_id: str) -> Any:
+    """Point every local ``$ref`` at the name its definition was lifted under."""
+    if isinstance(value, dict):
+        mapping = cast("Mapping[str, object]", value)
+        return {key: _repoint_member(key, item, block_id) for key, item in mapping.items()}
+    if isinstance(value, list):
+        items = cast("Sequence[object]", value)
+        return [_repoint(item, block_id) for item in items]
+    return value
+
+
+def _repoint_member(key: str, value: object, block_id: str) -> Any:
+    """Rewrite one member, which is a rename only when it is the ``$ref`` itself."""
+    prefix = "#/$defs/"
+    if key == "$ref" and isinstance(value, str) and value.startswith(prefix):
+        return f"{prefix}{block_id}.{value.removeprefix(prefix)}"
+    return _repoint(value, block_id)

dirigent_core/engine/__init__.py

@@ -0,0 +1,120 @@
+"""The execution engine: definitions, references, claims, execution, and recovery."""
+
+from dirigent_core.engine.claim import ClaimedUnit, due_attempt_statement, select_due
+from dirigent_core.engine.context import (
+    BufferedLogger,
+    ConnectionRecord,
+    EngineStepContext,
+    load_connections,
+    load_schemas,
+)
+from dirigent_core.engine.definition import (
+    FORMAT_V1,
+    KIND_PIPELINE,
+    KIND_TRIGGERS,
+    ConcurrencyPolicy,
+    Document,
+    ItemPolicy,
+    ParameterError,
+    PipelineDefinition,
+    Requirements,
+    RetryPolicy,
+    ScheduleSpec,
+    StepDefinition,
+    TimeoutAction,
+    TriggerRule,
+    TriggersDefinition,
+    TriggerSpecs,
+    WebhookSpec,
+    canonical_document,
+    dump_definition,
+    load_definition,
+)
+from dirigent_core.engine.executor import Engine
+from dirigent_core.engine.failure import Failure, backoff_delay, classify, should_retry
+from dirigent_core.engine.recovery import detect_stuck_runs, sweep_leases
+from dirigent_core.engine.references import ReferenceScope, UnknownReference, resolve, resolve_config
+from dirigent_core.engine.runs import (
+    Attribution,
+    EngineRuns,
+    FanOutError,
+    Provenance,
+    RunCreationError,
+    cancel_run,
+    chain_depth,
+    create_run,
+    retry_step,
+    save_pipeline,
+)
+from dirigent_core.engine.services import EngineServices
+from dirigent_core.engine.state import (
+    Readiness,
+    StepOutcome,
+    StepState,
+    advance,
+    aggregate_items,
+    build_step_states,
+    derive_run_status,
+    evaluate_rule,
+)
+
+__all__ = [
+    "FORMAT_V1",
+    "KIND_PIPELINE",
+    "KIND_TRIGGERS",
+    "Attribution",
+    "BufferedLogger",
+    "ClaimedUnit",
+    "ConcurrencyPolicy",
+    "ConnectionRecord",
+    "Document",
+    "Engine",
+    "EngineRuns",
+    "EngineServices",
+    "EngineStepContext",
+    "Failure",
+    "FanOutError",
+    "ItemPolicy",
+    "ParameterError",
+    "PipelineDefinition",
+    "Provenance",
+    "Readiness",
+    "ReferenceScope",
+    "Requirements",
+    "RetryPolicy",
+    "RunCreationError",
+    "ScheduleSpec",
+    "StepDefinition",
+    "StepOutcome",
+    "StepState",
+    "TimeoutAction",
+    "TriggerRule",
+    "TriggerSpecs",
+    "TriggersDefinition",
+    "UnknownReference",
+    "WebhookSpec",
+    "advance",
+    "aggregate_items",
+    "backoff_delay",
+    "build_step_states",
+    "canonical_document",
+    "cancel_run",
+    "chain_depth",
+    "classify",
+    "create_run",
+    "derive_run_status",
+    "detect_stuck_runs",
+    "due_attempt_statement",
+    "dump_definition",
+    "evaluate_rule",
+    "load_connections",
+    "load_definition",
+    "load_schemas",
+    "resolve",
+    "resolve_config",
+    "retry_step",
+    "save_pipeline",
+    "select_due",
+    "should_retry",
+    "sweep_leases",
+]

dirigent_core/engine/claim.py

@@ -0,0 +1,156 @@
+"""The claim query: one statement that pulls the next due unit of work.
+
+On PostgreSQL the claim takes ``FOR UPDATE SKIP LOCKED``, so concurrent workers step over
+each other's rows. SQLite has no such clause, and the process-wide lock that stands in for it
+is only correct when exactly one process exists -- the guarantee the standalone mode makes.
+"""
+
+import json
+from collections.abc import Sequence
+from datetime import datetime
+from typing import cast
+from uuid import UUID
+
+import sqlalchemy as sa
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dirigent_client.enums import PRIORITY_RANK, AttemptStatus, RunPriority, RunStatus
+from dirigent_common import JsonMap
+from dirigent_core.engine.context import ConnectionRecord
+from dirigent_core.engine.definition import StepDefinition
+from dirigent_core.models import Run, StepAttempt
+from dirigent_plugin import RemoteHandle
+
+MAX_SETTLED_PER_CLAIM = 50
+
+CLAIMABLE_RUN_STATUSES = (RunStatus.QUEUED, RunStatus.RUNNING)
+
+
+class ClaimedUnit(BaseModel):
+    """Everything one block call needs, snapshotted by the transaction that claimed it."""
+
+    model_config = ConfigDict(frozen=True)
+
+    attempt_id: UUID
+    run_id: UUID
+    run_item_id: UUID | None
+    step_name: str
+    block_id: str
+    attempt_number: int
+    started_at: datetime
+    """When this attempt first started, which a wait measures itself from."""
+
+    previous_status: AttemptStatus
+    step: StepDefinition
+    config: JsonMap
+    params: JsonMap
+    log_levels: JsonMap | None = None
+    """The run's kept-level map, resolved against the block id when the buffer is built."""
+    scratch: str
+    item: JsonValue
+    has_item: bool
+    connections: dict[str, ConnectionRecord]
+    schemas: dict[str, JsonMap] = Field(default_factory=dict)
+    traceparent: str | None = None
+    """The run's trace context, so an attempt's span joins the trace the run started in."""
+
+    remote_handle: RemoteHandle | None = None
+    poke_cursor: JsonMap | None = None
+    """The cursor the last committed poke returned, handed to the next one."""
+
+    deadline_at: datetime | None = None
+    gone_probes: int = 0
+
+    @property
+    def is_probe(self) -> bool:
+        """Report whether this claim continues a wait rather than starting the work."""
+        return self.previous_status is AttemptStatus.WAITING
+
+
+def is_postgres(session: AsyncSession) -> bool:
+    """Report whether this session speaks PostgreSQL."""
+    return session.get_bind().dialect.name == "postgresql"
+
+
+def routable(tags: Sequence[str], *, postgres: bool) -> sa.ColumnElement[bool]:
+    """Build the test that a run's pinned tags are all carried by a worker holding ``tags``.
+
+    A run pinning no tags is claimable by anyone. On PostgreSQL the column is ``jsonb`` and
+    the test is containment, which a GIN index answers; on SQLite it is a correlated
+    ``json_each`` walk looking for one pinned tag the worker does not carry.
+    """
+    carried = sorted(set(tags))
+    if postgres:
+        contained = sa.type_coerce(Run.worker_tags, JSONB).contained_by(
+            sa.cast(sa.literal(json.dumps(carried), sa.Text), JSONB)
+        )
+        return cast("sa.ColumnElement[bool]", contained)
+    element = sa.func.json_each(Run.worker_tags).table_valued("value", joins_implicitly=True)
+    missing = sa.select(1).select_from(element).where(element.c.value.notin_(carried))
+    return ~sa.exists(missing)
+
+
+def priority_rank() -> sa.Case[int]:
+    """Score a run's priority so the claim can sort on it, highest first."""
+    return sa.case(
+        {priority.value: rank for priority, rank in PRIORITY_RANK.items()},
+        value=Run.priority,
+        else_=PRIORITY_RANK[RunPriority.NORMAL],
+    )
+
+
+def due_attempt_statement(
+    now: datetime, tags: Sequence[str] = (), *, postgres: bool = False
+) -> sa.Select[tuple[StepAttempt]]:
+    """Build the select that finds the next due unit of work this worker may run.
+
+    The order is priority first, then fairness, then due time. Fairness is round-robin
+    between runs: every attempt of a run still in flight is numbered within that run, and the
+    claim takes the lowest number among the due ones. The number counts the attempts already
+    taken, so a run that has been served falls behind one that has not, and a
+    four-hundred-item fan-out interleaves one attempt at a time with a two-step run rather
+    than holding every slot until it drains. Numbering only the due rows would not do it:
+    claiming one removes it, and the rest would renumber from one after every claim.
+
+    The subquery computes the ranking and nothing else. Every predicate that decides whether
+    an attempt may be claimed is repeated on the locked ``step_attempts`` relation in the
+    outer select, because ``SKIP LOCKED`` re-checks a row that changed under it against the
+    outer query's own predicates only: a status test living solely in the subquery would not
+    be re-evaluated, and two workers would take the same attempt.
+    """
+    due_at = sa.func.coalesce(StepAttempt.available_at, StepAttempt.next_poll_at)
+    ranked = (
+        sa.select(
+            StepAttempt.id.label("attempt_id"),
+            priority_rank().label("priority_rank"),
+            sa.func.row_number().over(partition_by=StepAttempt.run_id, order_by=(due_at, StepAttempt.id)).label("turn"),
+            due_at.label("due_at"),
+        )
+        .join(Run, Run.id == StepAttempt.run_id)
+        .where(Run.status.in_(CLAIMABLE_RUN_STATUSES), routable(tags, postgres=postgres))
+        .subquery()
+    )
+    due = sa.or_(
+        sa.and_(StepAttempt.status == AttemptStatus.QUEUED, StepAttempt.available_at <= now),
+        sa.and_(StepAttempt.status == AttemptStatus.WAITING, StepAttempt.next_poll_at <= now),
+    )
+    return (
+        sa.select(StepAttempt)
+        .join(Run, Run.id == StepAttempt.run_id)
+        .join(ranked, ranked.c.attempt_id == StepAttempt.id)
+        .where(due, Run.status.in_(CLAIMABLE_RUN_STATUSES), routable(tags, postgres=postgres))
+        .order_by(ranked.c.priority_rank.desc(), ranked.c.turn, ranked.c.due_at, ranked.c.attempt_id)
+        .limit(1)
+    )
+
+
+async def select_due(session: AsyncSession, now: datetime, tags: Sequence[str] = ()) -> StepAttempt | None:
+    """Run the claim select, locking the row on the dialect that can."""
+    postgres = is_postgres(session)
+    statement = due_attempt_statement(now, tags, postgres=postgres)
+    if postgres:
+        statement = statement.with_for_update(skip_locked=True, of=StepAttempt)
+    result = await session.execute(statement)
+    return result.scalars().first()