Files
lda-wf/src/wf_artifacts/runs/models.py
T

103 lines
2.9 KiB
Python

from __future__ import annotations
import re
from datetime import datetime
from enum import StrEnum
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field
from wf_core import PersistedRunState
from ..models import DependencyDiagnostic, WorkflowArtifact, WorkflowDeployment
RUN_ID_PATTERN = r"^[A-Za-z0-9_][A-Za-z0-9_.-]*$"
def ensure_run_id(run_id: str) -> str:
"""Reject ids that cannot safely identify one local run directory."""
if not re.fullmatch(RUN_ID_PATTERN, run_id):
raise ValueError(
"run_id must start with alphanumeric or underscore and contain only "
"[A-Za-z0-9_.-]"
)
return run_id
class StoredRunStatus(StrEnum):
"""Stopped runtime statuses supported by durable run persistence."""
INTERRUPTED = "interrupted"
COMPLETED = "completed"
FAILED = "failed"
class ResumeReadiness(StrEnum):
"""Whether an interrupted stored run may currently continue."""
READY = "ready"
BLOCKED = "blocked"
NOT_APPLICABLE = "not_applicable"
class CheckpointReason(StrEnum):
"""Why a stopped-state checkpoint was written."""
INTERRUPTED = "interrupted"
COMPLETED = "completed"
FAILED = "failed"
class PinnedRunEnvironment(BaseModel):
"""Exact execution definitions captured when a run starts."""
model_config = ConfigDict(extra="forbid")
deployment: WorkflowDeployment
root_artifact: WorkflowArtifact
child_artifacts: list[WorkflowArtifact] = Field(default_factory=list)
class WorkflowRunRecord(BaseModel):
"""Durable summary and pinned environment for one started workflow run."""
model_config = ConfigDict(extra="forbid")
id: str = Field(pattern=RUN_ID_PATTERN)
status: StoredRunStatus
resume_readiness: ResumeReadiness
environment: PinnedRunEnvironment
latest_checkpoint_id: str = Field(pattern=RUN_ID_PATTERN)
diagnostics: list[DependencyDiagnostic] = Field(default_factory=list)
created_at: datetime
updated_at: datetime
class VersionedCheckpointState(BaseModel):
"""Lenient read envelope for stopped-run checkpoints.
Writes always produce version 2 via ``wf_core.dump_run_state``; reads
accept version 1 so pre-budget checkpoints reach
``load_run_state_with_upgrade`` instead of failing checkpoint validation
with a ``version == 2`` literal error first. The inner state stays an
untyped dict because core owns strict budget validation there.
"""
model_config = ConfigDict(extra="forbid")
version: Literal[1, 2] = 2
state: dict[str, Any]
class RunCheckpoint(BaseModel):
"""One stopped-state snapshot persisted at an external run boundary."""
model_config = ConfigDict(extra="forbid")
id: str = Field(pattern=RUN_ID_PATTERN)
run_id: str = Field(pattern=RUN_ID_PATTERN)
sequence: int = Field(ge=1)
reason: CheckpointReason
state: PersistedRunState | VersionedCheckpointState
created_at: datetime