275 lines
8.4 KiB
Python
275 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from uuid import uuid4
|
|
|
|
from wf_api.saved_subgraphs import SavedSubgraphTree
|
|
from wf_artifacts import (
|
|
AvailableSource,
|
|
CheckpointReason,
|
|
DependencyDiagnostic,
|
|
DiagnosticSeverity,
|
|
PinnedRunEnvironment,
|
|
ResumeReadiness,
|
|
RunAdmission,
|
|
RunCheckpoint,
|
|
RunStore,
|
|
StoredRunStatus,
|
|
WorkflowArtifact,
|
|
WorkflowDeployment,
|
|
WorkflowRunRecord,
|
|
validate_deployment_dependencies,
|
|
)
|
|
from wf_core import (
|
|
PersistedRunState,
|
|
RunState,
|
|
RunStatus,
|
|
dump_run_state,
|
|
load_run_state,
|
|
load_run_state_with_upgrade,
|
|
)
|
|
|
|
|
|
def create_pinned_environment(
|
|
*,
|
|
deployment: WorkflowDeployment,
|
|
artifact: WorkflowArtifact,
|
|
tree: SavedSubgraphTree,
|
|
) -> PinnedRunEnvironment:
|
|
"""Capture exact root, deployment, and child definitions for one run."""
|
|
return PinnedRunEnvironment(
|
|
deployment=deployment,
|
|
root_artifact=artifact,
|
|
child_artifacts=list(tree.artifacts_by_ref.values()),
|
|
)
|
|
|
|
|
|
def persist_admission(
|
|
*,
|
|
store: RunStore,
|
|
run_id: str,
|
|
environment: PinnedRunEnvironment,
|
|
resolved_input: dict[str, object],
|
|
max_steps: int | None,
|
|
scheduled_at: datetime | None = None,
|
|
schedule_id: str | None = None,
|
|
schedule_revision: int | None = None,
|
|
deployment_revision: int | None = None,
|
|
) -> RunAdmission:
|
|
"""Persist the authoritative admission record before dispatch.
|
|
|
|
The occurrence is decided here: the caller must treat this persist as the
|
|
single authority for the run identity and frozen invocation. A failed
|
|
durable admission must never dispatch.
|
|
"""
|
|
from wf_core.models.json_values import validate_strict_json_value
|
|
|
|
frozen = validate_strict_json_value(dict(resolved_input))
|
|
if not isinstance(frozen, dict):
|
|
raise ValueError("resolved workflow input must be a JSON object")
|
|
admission = RunAdmission(
|
|
id=run_id,
|
|
environment=environment,
|
|
resolved_input=frozen,
|
|
max_steps=max_steps,
|
|
scheduled_at=scheduled_at,
|
|
schedule_id=schedule_id,
|
|
schedule_revision=schedule_revision,
|
|
deployment_revision=deployment_revision,
|
|
created_at=datetime.now(UTC),
|
|
)
|
|
store.save_admission(admission)
|
|
return admission
|
|
|
|
|
|
def materialize_admitted_view(
|
|
*,
|
|
store: RunStore,
|
|
admission: RunAdmission,
|
|
) -> WorkflowRunRecord:
|
|
"""Materialize the run admission view using the admission identity.
|
|
|
|
The view carries no checkpoint, trace, output, or step counts: the outcome
|
|
is unknown until dispatch completes and persists a stopped checkpoint.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
try:
|
|
existing = store.get_run(admission.id)
|
|
except KeyError:
|
|
existing = None
|
|
if existing is not None:
|
|
return existing
|
|
record = WorkflowRunRecord(
|
|
id=admission.id,
|
|
status=StoredRunStatus.ADMITTED,
|
|
resume_readiness=ResumeReadiness.NOT_APPLICABLE,
|
|
environment=admission.environment,
|
|
latest_checkpoint_id=None,
|
|
created_at=admission.created_at,
|
|
updated_at=now,
|
|
)
|
|
store.save_run(record)
|
|
return record
|
|
|
|
|
|
def recover_admission_view(*, store: RunStore, run_id: str) -> WorkflowRunRecord:
|
|
"""Reconcile a missing run view from its admission record.
|
|
|
|
Recovery never executes work: it only completes the missing view so a
|
|
later poll can dispatch the captured invocation exactly once.
|
|
"""
|
|
try:
|
|
return store.get_run(run_id)
|
|
except KeyError:
|
|
pass
|
|
admission = store.get_admission(run_id)
|
|
return materialize_admitted_view(store=store, admission=admission)
|
|
|
|
|
|
def persist_stopped_run(
|
|
*,
|
|
store: RunStore,
|
|
environment: PinnedRunEnvironment,
|
|
run: RunState,
|
|
run_id: str | None = None,
|
|
) -> WorkflowRunRecord:
|
|
"""Persist one externally visible stopped state and its typed checkpoint."""
|
|
if run.status not in {
|
|
RunStatus.INTERRUPTED,
|
|
RunStatus.COMPLETED,
|
|
RunStatus.FAILED,
|
|
}:
|
|
raise ValueError(
|
|
f"cannot persist active workflow run with status {run.status!s}"
|
|
)
|
|
|
|
key = run_id or f"run_{uuid4().hex}"
|
|
now = datetime.now(UTC)
|
|
sequence = 1
|
|
created_at = now
|
|
if run_id is not None:
|
|
try:
|
|
existing = store.get_run(run_id)
|
|
except KeyError:
|
|
existing = None
|
|
if existing is not None:
|
|
created_at = existing.created_at
|
|
try:
|
|
sequence = store.get_latest_checkpoint(run_id).sequence + 1
|
|
except KeyError:
|
|
sequence = 1
|
|
|
|
status = StoredRunStatus(run.status.value)
|
|
readiness = (
|
|
ResumeReadiness.READY
|
|
if status is StoredRunStatus.INTERRUPTED
|
|
else ResumeReadiness.NOT_APPLICABLE
|
|
)
|
|
checkpoint_id = f"{key}.{sequence:06d}"
|
|
checkpoint = RunCheckpoint(
|
|
id=checkpoint_id,
|
|
run_id=key,
|
|
sequence=sequence,
|
|
reason=CheckpointReason(status.value),
|
|
state=PersistedRunState.model_validate(dump_run_state(run)),
|
|
created_at=now,
|
|
)
|
|
record = WorkflowRunRecord(
|
|
id=key,
|
|
status=status,
|
|
resume_readiness=readiness,
|
|
environment=environment,
|
|
latest_checkpoint_id=checkpoint_id,
|
|
created_at=created_at,
|
|
updated_at=now,
|
|
)
|
|
store.save_checkpoint(checkpoint)
|
|
store.save_run(record)
|
|
return record
|
|
|
|
|
|
def restore_interrupted_run(
|
|
store: RunStore, run_id: str
|
|
) -> tuple[WorkflowRunRecord, RunState]:
|
|
"""Load a persisted interrupted run, persisting a v1 upgrade first.
|
|
|
|
A pre-budget (v1) checkpoint receives its one-time defaults and is
|
|
rewritten as a new v2 interrupted checkpoint under the same run id and
|
|
pinned environment *before* the run is returned, so resume dispatch
|
|
never runs on unmigrated state and a failed upgrade fails resume before
|
|
any handler runs. Ordinary inspection uses :func:`load_stored_run`,
|
|
which decodes v1 prospectively without mutating the store.
|
|
"""
|
|
record = store.get_run(run_id)
|
|
if record.status is not StoredRunStatus.INTERRUPTED:
|
|
raise ValueError(f"workflow run {run_id!r} is not interrupted")
|
|
checkpoint = store.get_latest_checkpoint(run_id)
|
|
run, upgraded = load_run_state_with_upgrade(
|
|
checkpoint.state.model_dump(mode="json")
|
|
)
|
|
if upgraded:
|
|
record = persist_stopped_run(
|
|
store=store,
|
|
environment=record.environment,
|
|
run=run,
|
|
run_id=run_id,
|
|
)
|
|
return record, run
|
|
|
|
|
|
def load_stored_run(store: RunStore, run_id: str) -> tuple[WorkflowRunRecord, RunState]:
|
|
"""Load any stopped run record together with its latest typed checkpoint."""
|
|
record = store.get_run(run_id)
|
|
if record.status is StoredRunStatus.ADMITTED:
|
|
raise ValueError(
|
|
f"workflow run {run_id!r} is admitted but has no stopped checkpoint yet"
|
|
)
|
|
checkpoint = store.get_latest_checkpoint(run_id)
|
|
return record, load_run_state(checkpoint.state.model_dump(mode="json"))
|
|
|
|
|
|
def validate_pinned_resume_environment(
|
|
*,
|
|
record: WorkflowRunRecord,
|
|
sources: list[AvailableSource],
|
|
) -> list[DependencyDiagnostic]:
|
|
"""Revalidate exact stored graph definitions before a resume mutates state."""
|
|
environment = record.environment
|
|
diagnostics = validate_deployment_dependencies(
|
|
artifact=environment.root_artifact,
|
|
deployment=environment.deployment,
|
|
sources=sources,
|
|
)
|
|
for child in environment.child_artifacts:
|
|
diagnostics.extend(
|
|
validate_deployment_dependencies(
|
|
artifact=child,
|
|
deployment=environment.deployment,
|
|
sources=sources,
|
|
)
|
|
)
|
|
return diagnostics
|
|
|
|
|
|
def has_blocking_diagnostics(diagnostics: list[DependencyDiagnostic]) -> bool:
|
|
"""Return whether dependency diagnostics prohibit executing a resume."""
|
|
return any(item.severity is DiagnosticSeverity.ERROR for item in diagnostics)
|
|
|
|
|
|
def mark_resume_blocked(
|
|
*,
|
|
store: RunStore,
|
|
record: WorkflowRunRecord,
|
|
diagnostics: list[DependencyDiagnostic],
|
|
) -> WorkflowRunRecord:
|
|
"""Record blocked readiness without writing a new execution checkpoint."""
|
|
blocked = record.model_copy(
|
|
update={
|
|
"resume_readiness": ResumeReadiness.BLOCKED,
|
|
"diagnostics": diagnostics,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
)
|
|
store.save_run(blocked)
|
|
return blocked
|