114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
|
"""Sequential, resumable execution of an immutable paired trial plan."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Callable, Mapping
|
||
|
|
|
||
|
|
from .io import (
|
||
|
|
AtomicTrialWriter,
|
||
|
|
TrialPayload,
|
||
|
|
atomic_write_json,
|
||
|
|
record_trial_failure,
|
||
|
|
utc_now,
|
||
|
|
)
|
||
|
|
from .manifest import ensure_batch
|
||
|
|
from .plan import validate_trial_plan
|
||
|
|
from .validate import validate_trial_directory
|
||
|
|
|
||
|
|
|
||
|
|
TrialExecutor = Callable[[Mapping[str, Any]], TrialPayload | Mapping[str, Any]]
|
||
|
|
|
||
|
|
|
||
|
|
def _normalize_payload(value: TrialPayload | Mapping[str, Any]) -> TrialPayload:
|
||
|
|
if isinstance(value, TrialPayload):
|
||
|
|
return value
|
||
|
|
if not isinstance(value, Mapping) or "samples" not in value:
|
||
|
|
raise TypeError("executor must return TrialPayload or a mapping with samples")
|
||
|
|
return TrialPayload(
|
||
|
|
samples=value["samples"],
|
||
|
|
events=value.get("events", ()),
|
||
|
|
metadata=value.get("metadata", {}),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def run_trial_plan(
|
||
|
|
plan: Mapping[str, Any],
|
||
|
|
batch_dir: Path,
|
||
|
|
executor: TrialExecutor | None = None,
|
||
|
|
*,
|
||
|
|
resume: bool = True,
|
||
|
|
dry_run: bool = False,
|
||
|
|
continue_on_error: bool = True,
|
||
|
|
max_trials: int | None = None,
|
||
|
|
source_root: Path | None = None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Run trials in plan order and atomically commit successful artifacts."""
|
||
|
|
validate_trial_plan(plan)
|
||
|
|
selected = list(plan["trials"])
|
||
|
|
if max_trials is not None:
|
||
|
|
if max_trials < 0:
|
||
|
|
raise ValueError("max_trials must be non-negative")
|
||
|
|
selected = selected[:max_trials]
|
||
|
|
if dry_run:
|
||
|
|
return {
|
||
|
|
"dry_run": True,
|
||
|
|
"study_id": plan["study_id"],
|
||
|
|
"split": plan["split"],
|
||
|
|
"pair_count": len({trial["pair_id"] for trial in selected}),
|
||
|
|
"trial_count": len(selected),
|
||
|
|
"plan_hash": plan["plan_hash"],
|
||
|
|
}
|
||
|
|
if executor is None:
|
||
|
|
raise ValueError("executor is required unless dry_run=True")
|
||
|
|
|
||
|
|
batch_dir = Path(batch_dir)
|
||
|
|
ensure_batch(batch_dir, plan, source_root=source_root)
|
||
|
|
writer = AtomicTrialWriter(batch_dir)
|
||
|
|
completed = 0
|
||
|
|
skipped = 0
|
||
|
|
failed = 0
|
||
|
|
failures: list[dict[str, str]] = []
|
||
|
|
for trial in selected:
|
||
|
|
destination = writer.trial_dir(trial["trial_id"])
|
||
|
|
if destination.exists():
|
||
|
|
validation_errors = validate_trial_directory(destination, trial)
|
||
|
|
if resume and not validation_errors:
|
||
|
|
skipped += 1
|
||
|
|
continue
|
||
|
|
if validation_errors:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"existing trial {trial['trial_id']} is invalid: "
|
||
|
|
+ "; ".join(validation_errors)
|
||
|
|
)
|
||
|
|
raise FileExistsError(f"trial already completed: {trial['trial_id']}")
|
||
|
|
try:
|
||
|
|
payload = _normalize_payload(executor(trial))
|
||
|
|
writer.write(trial, payload)
|
||
|
|
completed += 1
|
||
|
|
except Exception as exc:
|
||
|
|
failed += 1
|
||
|
|
record_trial_failure(batch_dir, trial, exc)
|
||
|
|
failures.append(
|
||
|
|
{
|
||
|
|
"trial_id": trial["trial_id"],
|
||
|
|
"error_type": type(exc).__name__,
|
||
|
|
"error_message": str(exc),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
if not continue_on_error:
|
||
|
|
raise
|
||
|
|
summary = {
|
||
|
|
"kind": "run_summary",
|
||
|
|
"schema_version": plan["schema_version"],
|
||
|
|
"plan_hash": plan["plan_hash"],
|
||
|
|
"recorded_utc": utc_now(),
|
||
|
|
"selected_trials": len(selected),
|
||
|
|
"completed_trials": completed,
|
||
|
|
"skipped_trials": skipped,
|
||
|
|
"failed_trials": failed,
|
||
|
|
"failures": failures,
|
||
|
|
}
|
||
|
|
atomic_write_json(batch_dir / "run_summary.json", summary)
|
||
|
|
return summary
|