mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

fix(preprocess): harden crash recovery integrity

Browse Source
This commit is contained in:
marcopan
2026-07-12 04:49:15 +02:00
parent 1d5f8c76a7
commit c964920f16
7 changed files with 440 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
harness/tht/jobs/runner.py
+106 -9
View File
@@ -23,6 +23,11 @@ class CorruptCheckpointError(RuntimeError):
pass
@dataclass(frozen=True)
class StageArtifacts:
required: tuple[str, ...] = ()
@dataclass(frozen=True)
class JobContext:
run_id: str
@@ -35,6 +40,84 @@ class JobContext:
Stage = Callable[[JobContext], Any]
def _artifact_digest(path: Path) -> dict[str, Any]:
payload = path.read_bytes()
return {"sha256": hashlib.sha256(payload).hexdigest(), "size": len(payload)}
def _seal_artifacts(context: JobContext, stage: str, result: Any, spec: JobSpec) -> None:
required = result.required if isinstance(result, StageArtifacts) else ()
root = context.run_dir / "artifacts"
root.mkdir(exist_ok=True)
manifest_path = root / "artifact-manifest.json"
manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {
"schema_version": 1,
"spec_fingerprint": _compatibility_fingerprint(spec, list(spec.stage_ids)),
"stages": {},
}
files = {}
for relative in required:
candidate = root / relative
if Path(relative).is_absolute() or ".." in Path(relative).parts or candidate.is_symlink():
raise CorruptCheckpointError("artifact path is unsafe")
if not candidate.is_file():
raise CorruptCheckpointError("required stage artifact is missing")
files[relative] = _artifact_digest(candidate)
for prior in manifest["stages"].values():
if relative in prior.get("required", []):
prior["required"].remove(relative)
prior["files"].pop(relative, None)
manifest["stages"][stage] = {"required": list(required), "files": files}
_atomic_write(manifest_path, json.dumps(manifest, sort_keys=True, separators=(",", ":")) + "\n")
def seal_stage_artifacts(
context: JobContext, stage: str, required: tuple[str, ...], spec: JobSpec,
) -> None:
"""Durably record external-effect intent before a stage performs that effect."""
_seal_artifacts(context, stage, StageArtifacts(required), spec)
def _validate_artifacts(run_dir: Path, spec: JobSpec, source: JobRun) -> set[str]:
root = run_dir / "artifacts"
manifest_path = root / "artifact-manifest.json"
successful = {stage.name for stage in source.stages if stage.status == "succeeded"}
if not successful and not manifest_path.exists():
return set()
try:
manifest = json.loads(manifest_path.read_text())
if manifest["spec_fingerprint"] != _compatibility_fingerprint(spec, list(spec.stage_ids)):
raise ValueError
sealed = set(manifest["stages"])
allowed = {"artifact-manifest.json"}
for stage, record in manifest["stages"].items():
for relative in record["required"]:
candidate = root / relative
if Path(relative).is_absolute() or ".." in Path(relative).parts or candidate.is_symlink():
raise ValueError
if not candidate.is_file() or _artifact_digest(candidate) != record["files"][relative]:
raise ValueError
allowed.add(relative)
if not successful.issubset(sealed):
raise ValueError
entries = list(root.iterdir())
if any(path.is_symlink() or not path.is_file() for path in entries):
raise ValueError
actual = {path.name for path in entries}
incomplete = next((stage for stage in source.stages if stage.status != "succeeded"), None)
marker = root / "compensated.json"
if incomplete is not None and incomplete.status in {"failed", "running"} and marker.is_file():
payload = json.loads(marker.read_text())
if not isinstance(payload.get("generation"), str):
raise ValueError
allowed.add("compensated.json")
if actual != allowed:
raise ValueError
return sealed
except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as error:
raise CorruptCheckpointError("resume artifact manifest is invalid") from error
def _atomic_write(path: Path, payload: str) -> None:
temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
@@ -91,7 +174,10 @@ def _new_run(spec: JobSpec, run_id: str, stages: Sequence[Stage]) -> JobRun:
)
def _resume_run(spec: JobSpec, run_id: str, stages: Sequence[Stage], source: JobRun) -> JobRun:
def _resume_run(
spec: JobSpec, run_id: str, stages: Sequence[Stage], source: JobRun,
effect_completed: set[str] | None = None,
) -> JobRun:
requested_names = list(spec.stage_ids)
if len(requested_names) != len(stages) or len(requested_names) != len(set(requested_names)):
raise CorruptCheckpointError("resume checkpoint is incompatible with requested stages")
@@ -107,11 +193,15 @@ def _resume_run(spec: JobSpec, run_id: str, stages: Sequence[Stage], source: Job
resumed_stages = []
for name in requested_names:
previous = source_by_name.get(name)
resumed_stages.append(
previous
if previous is not None and previous.status == "succeeded"
else StageRun(name=name)
)
if previous is not None and previous.status == "succeeded":
resumed_stages.append(previous)
elif previous is not None and previous.status == "running" and name in (effect_completed or set()):
resumed_stages.append(StageRun(
name=name, status="succeeded", started_at=previous.started_at or utc_now(),
finished_at=utc_now(),
))
else:
resumed_stages.append(StageRun(name=name))
return JobRun(
run_id=run_id,
compatibility_fingerprint=source.compatibility_fingerprint,
@@ -129,7 +219,10 @@ def _resume_run(spec: JobSpec, run_id: str, stages: Sequence[Stage], source: Job
)
def run_job(spec: JobSpec, stages: Sequence[Stage]) -> JobReport:
def run_job(
spec: JobSpec, stages: Sequence[Stage], *,
after_stage_return: Callable[[JobContext, str], Any] | None = None,
) -> JobReport:
"""Run stages once, returning a terminal report instead of leaking stage exceptions."""
with WorkspaceJobLock(spec.workspace_root, spec.workspace_id, spec.job_type):
jobs_root = spec.workspace_root / ".tht-jobs" / spec.job_type / "runs"
@@ -147,6 +240,7 @@ def run_job(spec: JobSpec, stages: Sequence[Stage]) -> JobReport:
source_path = matches[0]
source = _load_checkpoint(source_path)
_validate_resume_source(spec, stages, source)
effect_completed = _validate_artifacts(source_path.parent, spec, source)
run_id = uuid.uuid4().hex
run_dir = jobs_root / run_id
@@ -155,7 +249,7 @@ def run_job(spec: JobSpec, stages: Sequence[Stage]) -> JobReport:
if source is None:
run = _new_run(spec, run_id, stages)
else:
run = _resume_run(spec, run_id, stages, source)
run = _resume_run(spec, run_id, stages, source, effect_completed)
source_artifacts = jobs_root / source.run_id / "artifacts"
if source_artifacts.exists():
shutil.copytree(source_artifacts, run_dir / "artifacts")
@@ -173,7 +267,7 @@ def run_job(spec: JobSpec, stages: Sequence[Stage]) -> JobReport:
)
_persist(checkpoint_path, run)
try:
stage_callable(context)
stage_result = stage_callable(context)
except Exception:
failed = stage.model_copy(
update={
@@ -191,6 +285,9 @@ def run_job(spec: JobSpec, stages: Sequence[Stage]) -> JobReport:
)
_persist(checkpoint_path, run)
break
_seal_artifacts(context, stage.name, stage_result, spec)
if after_stage_return is not None:
after_stage_return(context, stage.name)
succeeded = stage.model_copy(update={"status": "succeeded", "finished_at": utc_now()})
run = run.model_copy(
update={"stages": run.stages[:index] + (succeeded,) + run.stages[index + 1 :]}