429 lines
16 KiB
Python
429 lines
16 KiB
Python
|
"""Private, atomic checkpoints for host-executed GEWU research runs."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import importlib
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
import os
|
||
|
|
import re
|
||
|
|
import secrets
|
||
|
|
import stat
|
||
|
|
import sys
|
||
|
|
from datetime import UTC, datetime
|
||
|
|
from pathlib import Path
|
||
|
|
from types import ModuleType
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
if __package__:
|
||
|
|
manifest_module: ModuleType = importlib.import_module(".manifest_tool", __package__)
|
||
|
|
else:
|
||
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||
|
|
manifest_module = importlib.import_module("manifest_tool")
|
||
|
|
|
||
|
|
load_manifest = manifest_module.load_manifest
|
||
|
|
verify = manifest_module.verify
|
||
|
|
|
||
|
|
_RUN_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{1,63}$")
|
||
|
|
_SHA256 = re.compile(r"^[a-f0-9]{64}$")
|
||
|
|
_MUTABLE = {
|
||
|
|
"stage",
|
||
|
|
"sources",
|
||
|
|
"coverage",
|
||
|
|
"budgets",
|
||
|
|
"pending_operations",
|
||
|
|
"proposals",
|
||
|
|
"outputs",
|
||
|
|
"last_failure",
|
||
|
|
"runtime",
|
||
|
|
}
|
||
|
|
_FORBIDDEN_KEYS = {
|
||
|
|
"token",
|
||
|
|
"secret",
|
||
|
|
"password",
|
||
|
|
"credential",
|
||
|
|
"source_body",
|
||
|
|
"content",
|
||
|
|
"cookie",
|
||
|
|
"authorization",
|
||
|
|
"auth_header",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
|
||
|
|
result: dict[str, Any] = {}
|
||
|
|
for key, value in pairs:
|
||
|
|
if key in result:
|
||
|
|
raise ValueError("checkpoint JSON contains a duplicate key")
|
||
|
|
result[key] = value
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _safe_payload(value: Any) -> None:
|
||
|
|
if isinstance(value, dict):
|
||
|
|
for key, item in value.items():
|
||
|
|
normalized = str(key).lower().replace("-", "_")
|
||
|
|
if (
|
||
|
|
normalized in _FORBIDDEN_KEYS
|
||
|
|
or any(
|
||
|
|
part in normalized
|
||
|
|
for part in ("secret", "password", "credential", "bearer", "api_key")
|
||
|
|
)
|
||
|
|
or normalized.endswith("_token")
|
||
|
|
):
|
||
|
|
raise ValueError("checkpoint cannot store credentials or source bodies")
|
||
|
|
_safe_payload(item)
|
||
|
|
elif isinstance(value, list):
|
||
|
|
for item in value:
|
||
|
|
_safe_payload(item)
|
||
|
|
elif isinstance(value, float) and not math.isfinite(value):
|
||
|
|
raise ValueError("checkpoint contains a non-finite number")
|
||
|
|
elif not isinstance(value, (str, int, float, bool, type(None))):
|
||
|
|
raise ValueError("checkpoint contains a non-JSON value")
|
||
|
|
|
||
|
|
|
||
|
|
def _runs_directory(workspace: Path) -> Path:
|
||
|
|
root = workspace.expanduser().absolute()
|
||
|
|
if root.is_symlink():
|
||
|
|
raise ValueError("workspace root must not be a symbolic link")
|
||
|
|
runs = root / "Runs"
|
||
|
|
if runs.is_symlink():
|
||
|
|
raise ValueError("Runs directory must not be a symbolic link")
|
||
|
|
runs.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||
|
|
if not runs.is_dir() or runs.is_symlink():
|
||
|
|
raise ValueError("Runs path is unavailable")
|
||
|
|
if os.name != "nt":
|
||
|
|
runs.chmod(stat.S_IRWXU)
|
||
|
|
return runs
|
||
|
|
|
||
|
|
|
||
|
|
def _checkpoint_path(workspace: Path, run_id: str) -> Path:
|
||
|
|
if not _RUN_ID.fullmatch(run_id):
|
||
|
|
raise ValueError("run ID is invalid")
|
||
|
|
return _runs_directory(workspace) / f"{run_id}.json"
|
||
|
|
|
||
|
|
|
||
|
|
def _write_atomic(path: Path, value: dict[str, Any]) -> None:
|
||
|
|
_safe_payload(value)
|
||
|
|
payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
||
|
|
if len(payload.encode("utf-8")) > 1_000_000:
|
||
|
|
raise ValueError("checkpoint exceeds the one-megabyte limit")
|
||
|
|
temporary = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
|
||
|
|
try:
|
||
|
|
descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||
|
|
with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as stream:
|
||
|
|
stream.write(payload + "\n")
|
||
|
|
stream.flush()
|
||
|
|
os.fsync(stream.fileno())
|
||
|
|
os.replace(temporary, path)
|
||
|
|
if os.name != "nt":
|
||
|
|
path.chmod(stat.S_IRUSR | stat.S_IWUSR)
|
||
|
|
directory_fd = os.open(path.parent, os.O_RDONLY)
|
||
|
|
try:
|
||
|
|
os.fsync(directory_fd)
|
||
|
|
finally:
|
||
|
|
os.close(directory_fd)
|
||
|
|
finally:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
def _module_pins(manifest: dict[str, Any]) -> list[dict[str, str]]:
|
||
|
|
return [
|
||
|
|
{
|
||
|
|
"module_id": module["module_id"],
|
||
|
|
"entry": module["entry"],
|
||
|
|
"contract_id": module["provided_contract"]["id"],
|
||
|
|
"contract_version": module["provided_contract"]["version"],
|
||
|
|
"implementation_version": module["implementation_version"],
|
||
|
|
"sha256": module["sha256"],
|
||
|
|
}
|
||
|
|
for module in manifest["modules"]
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def create_checkpoint(
|
||
|
|
workspace: Path,
|
||
|
|
*,
|
||
|
|
run_id: str,
|
||
|
|
objective: str,
|
||
|
|
projects: list[str],
|
||
|
|
include_shared: bool,
|
||
|
|
lens: str,
|
||
|
|
budgets: dict[str, int] | None = None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
manifest = load_manifest()
|
||
|
|
verify(manifest)
|
||
|
|
if not objective.strip() or len(objective) > 4_000:
|
||
|
|
raise ValueError("objective must contain 1 to 4000 characters")
|
||
|
|
if (
|
||
|
|
not projects
|
||
|
|
or len(projects) > 100
|
||
|
|
or len(set(projects)) != len(projects)
|
||
|
|
or any(
|
||
|
|
not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", project) for project in projects
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise ValueError("scope must contain distinct explicit project IDs")
|
||
|
|
if not lens.strip() or len(lens) > 2_000:
|
||
|
|
raise ValueError("lens must contain 1 to 2000 characters")
|
||
|
|
path = _checkpoint_path(workspace, run_id)
|
||
|
|
if path.is_symlink() or path.exists():
|
||
|
|
raise ValueError("run checkpoint already exists")
|
||
|
|
now = datetime.now(UTC).isoformat()
|
||
|
|
checkpoint: dict[str, Any] = {
|
||
|
|
"checkpoint_version": "0.1.0",
|
||
|
|
"run_id": run_id,
|
||
|
|
"created_at": now,
|
||
|
|
"updated_at": now,
|
||
|
|
"pack": {
|
||
|
|
"pack_id": manifest["pack_id"],
|
||
|
|
"pack_version": manifest["pack_version"],
|
||
|
|
"snapshot_digest": manifest["snapshot"]["digest"],
|
||
|
|
"modules": _module_pins(manifest),
|
||
|
|
"service_contracts": [
|
||
|
|
{"contract_id": item["contract_id"], "sha256": item["sha256"]}
|
||
|
|
for item in manifest["verified_service_contracts"]
|
||
|
|
],
|
||
|
|
},
|
||
|
|
"objective": objective,
|
||
|
|
"scope": {"projects": projects, "include_shared": include_shared},
|
||
|
|
"lens": lens,
|
||
|
|
"stage": "context",
|
||
|
|
"sources": [],
|
||
|
|
"coverage": {"selected": [], "included": [], "excluded": [], "skipped": []},
|
||
|
|
"budgets": budgets or {},
|
||
|
|
"runtime": {"host": None, "model": None},
|
||
|
|
"pending_operations": [],
|
||
|
|
"proposals": [],
|
||
|
|
"outputs": [],
|
||
|
|
"last_failure": None,
|
||
|
|
}
|
||
|
|
_validate_checkpoint(checkpoint)
|
||
|
|
_write_atomic(path, checkpoint)
|
||
|
|
return checkpoint
|
||
|
|
|
||
|
|
|
||
|
|
def load_checkpoint(workspace: Path, run_id: str) -> dict[str, Any]:
|
||
|
|
manifest = load_manifest()
|
||
|
|
verify(manifest)
|
||
|
|
path = _checkpoint_path(workspace, run_id)
|
||
|
|
if path.is_symlink() or not path.is_file() or path.stat().st_size > 1_000_000:
|
||
|
|
raise ValueError("run checkpoint is unavailable")
|
||
|
|
checkpoint = json.loads(
|
||
|
|
path.read_text(encoding="utf-8"), object_pairs_hook=_reject_duplicate_keys
|
||
|
|
)
|
||
|
|
if not isinstance(checkpoint, dict) or checkpoint.get("run_id") != run_id:
|
||
|
|
raise ValueError("run checkpoint identity is invalid")
|
||
|
|
pinned_pack = checkpoint.get("pack")
|
||
|
|
if not isinstance(pinned_pack, dict) or (
|
||
|
|
pinned_pack.get("snapshot_digest") != manifest["snapshot"]["digest"]
|
||
|
|
or pinned_pack.get("modules") != _module_pins(manifest)
|
||
|
|
):
|
||
|
|
raise ValueError("pack changed since this run; revalidate before resuming")
|
||
|
|
_safe_payload(checkpoint)
|
||
|
|
_validate_checkpoint(checkpoint)
|
||
|
|
return checkpoint
|
||
|
|
|
||
|
|
|
||
|
|
def update_checkpoint(workspace: Path, run_id: str, patch: dict[str, Any]) -> dict[str, Any]:
|
||
|
|
checkpoint = load_checkpoint(workspace, run_id)
|
||
|
|
if not patch or not set(patch).issubset(_MUTABLE):
|
||
|
|
raise ValueError("checkpoint update contains unsupported fields")
|
||
|
|
_safe_payload(patch)
|
||
|
|
updated = {**checkpoint, **patch, "updated_at": datetime.now(UTC).isoformat()}
|
||
|
|
_validate_checkpoint(updated)
|
||
|
|
_write_atomic(_checkpoint_path(workspace, run_id), updated)
|
||
|
|
return updated
|
||
|
|
|
||
|
|
|
||
|
|
def _validate_checkpoint(checkpoint: dict[str, Any]) -> None:
|
||
|
|
required = {
|
||
|
|
"checkpoint_version",
|
||
|
|
"run_id",
|
||
|
|
"created_at",
|
||
|
|
"updated_at",
|
||
|
|
"pack",
|
||
|
|
"objective",
|
||
|
|
"scope",
|
||
|
|
"lens",
|
||
|
|
"stage",
|
||
|
|
"sources",
|
||
|
|
"coverage",
|
||
|
|
"budgets",
|
||
|
|
"runtime",
|
||
|
|
"pending_operations",
|
||
|
|
"proposals",
|
||
|
|
"outputs",
|
||
|
|
"last_failure",
|
||
|
|
}
|
||
|
|
if set(checkpoint) != required or checkpoint.get("checkpoint_version") != "0.1.0":
|
||
|
|
raise ValueError("checkpoint fields or version are invalid")
|
||
|
|
if not isinstance(checkpoint["objective"], str) or not checkpoint["objective"].strip():
|
||
|
|
raise ValueError("checkpoint objective is required")
|
||
|
|
if not isinstance(checkpoint["lens"], str) or not checkpoint["lens"].strip():
|
||
|
|
raise ValueError("checkpoint lens is required")
|
||
|
|
scope = checkpoint["scope"]
|
||
|
|
if (
|
||
|
|
not isinstance(scope, dict)
|
||
|
|
or set(scope) != {"projects", "include_shared"}
|
||
|
|
or not isinstance(scope["projects"], list)
|
||
|
|
or not scope["projects"]
|
||
|
|
or len(scope["projects"]) > 100
|
||
|
|
or any(
|
||
|
|
not isinstance(project, str)
|
||
|
|
or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", project)
|
||
|
|
for project in scope["projects"]
|
||
|
|
)
|
||
|
|
or len(set(scope["projects"])) != len(scope["projects"])
|
||
|
|
or type(scope["include_shared"]) is not bool
|
||
|
|
):
|
||
|
|
raise ValueError("checkpoint scope is invalid")
|
||
|
|
for field in ("sources", "pending_operations", "proposals", "outputs"):
|
||
|
|
if not isinstance(checkpoint[field], list) or len(checkpoint[field]) > 10_000:
|
||
|
|
raise ValueError(f"{field} must be a bounded list")
|
||
|
|
for source in checkpoint["sources"]:
|
||
|
|
if (
|
||
|
|
not isinstance(source, dict)
|
||
|
|
or not isinstance(source.get("source_revision_id"), str)
|
||
|
|
or not isinstance(source.get("status"), str)
|
||
|
|
):
|
||
|
|
raise ValueError("source records require a revision ID and capture status")
|
||
|
|
extraction_id = source.get("extraction_id")
|
||
|
|
if extraction_id is not None and not _SHA256.fullmatch(str(extraction_id)):
|
||
|
|
raise ValueError("source extraction references must use a SHA-256 ID")
|
||
|
|
if not isinstance(checkpoint["coverage"], dict) or set(checkpoint["coverage"]) != {
|
||
|
|
"selected",
|
||
|
|
"included",
|
||
|
|
"excluded",
|
||
|
|
"skipped",
|
||
|
|
}:
|
||
|
|
raise ValueError("coverage must record selected, included, excluded and skipped items")
|
||
|
|
for key, items in checkpoint["coverage"].items():
|
||
|
|
if not isinstance(items, list) or len(items) > 10_000:
|
||
|
|
raise ValueError("coverage categories must be bounded lists")
|
||
|
|
if key in {"selected", "included"} and any(not isinstance(item, str) for item in items):
|
||
|
|
raise ValueError("selected and included coverage entries must be strings")
|
||
|
|
if key in {"excluded", "skipped"} and any(
|
||
|
|
not isinstance(item, dict)
|
||
|
|
or not isinstance(item.get("item"), str)
|
||
|
|
or not isinstance(item.get("reason"), str)
|
||
|
|
or not item["reason"].strip()
|
||
|
|
for item in items
|
||
|
|
):
|
||
|
|
raise ValueError("excluded and skipped coverage items require a reason")
|
||
|
|
if not isinstance(checkpoint["budgets"], dict) or len(checkpoint["budgets"]) > 100:
|
||
|
|
raise ValueError("budgets must be a bounded object")
|
||
|
|
if any(type(value) is not int or value < 0 for value in checkpoint["budgets"].values()):
|
||
|
|
raise ValueError("budget values must be non-negative integers")
|
||
|
|
if (
|
||
|
|
not isinstance(checkpoint["runtime"], dict)
|
||
|
|
or set(checkpoint["runtime"]) != {"host", "model"}
|
||
|
|
or any(
|
||
|
|
value is not None and not isinstance(value, str)
|
||
|
|
for value in checkpoint["runtime"].values()
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise ValueError(
|
||
|
|
"runtime metadata must identify the host and model or explicitly omit them"
|
||
|
|
)
|
||
|
|
if not isinstance(checkpoint["stage"], str) or not checkpoint["stage"]:
|
||
|
|
raise ValueError("stage must be explicit")
|
||
|
|
for operation in checkpoint["pending_operations"]:
|
||
|
|
if (
|
||
|
|
not isinstance(operation, dict)
|
||
|
|
or not isinstance(operation.get("operation_id"), str)
|
||
|
|
or not isinstance(operation.get("tool_name"), str)
|
||
|
|
or not _SHA256.fullmatch(str(operation.get("input_sha256", "")))
|
||
|
|
or not isinstance(operation.get("status"), str)
|
||
|
|
or operation["status"] not in {"pending", "uncertain", "complete", "failed"}
|
||
|
|
):
|
||
|
|
raise ValueError("pending operation identity, input hash and status are required")
|
||
|
|
for proposal in checkpoint["proposals"]:
|
||
|
|
if (
|
||
|
|
not isinstance(proposal, dict)
|
||
|
|
or not isinstance(proposal.get("proposal_id"), str)
|
||
|
|
or type(proposal.get("revision")) is not int
|
||
|
|
or proposal["revision"] < 1
|
||
|
|
or not isinstance(proposal.get("status"), str)
|
||
|
|
or proposal["status"]
|
||
|
|
not in {
|
||
|
|
"staged",
|
||
|
|
"evaluated",
|
||
|
|
"eligible",
|
||
|
|
"needs-review",
|
||
|
|
"invalid",
|
||
|
|
"conflict",
|
||
|
|
"publishing",
|
||
|
|
"accepted",
|
||
|
|
"rejected",
|
||
|
|
"withdrawn",
|
||
|
|
"uncertain",
|
||
|
|
}
|
||
|
|
or (
|
||
|
|
proposal.get("status") == "accepted"
|
||
|
|
and not isinstance(proposal.get("receipt"), dict)
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise ValueError(
|
||
|
|
"proposal records require an ID, revision, status and accepted receipt"
|
||
|
|
)
|
||
|
|
if proposal["status"] == "accepted":
|
||
|
|
receipt = proposal["receipt"]
|
||
|
|
if (
|
||
|
|
receipt.get("proposal_id") != proposal["proposal_id"]
|
||
|
|
or receipt.get("revision") != proposal["revision"]
|
||
|
|
or receipt.get("state") != "accepted"
|
||
|
|
or not re.fullmatch(
|
||
|
|
r"(?:[a-f0-9]{40}|[a-f0-9]{64})", str(receipt.get("commit_id", ""))
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise ValueError("accepted proposal receipt does not match its revision")
|
||
|
|
for output in checkpoint["outputs"]:
|
||
|
|
if not isinstance(output, dict) or not _SHA256.fullmatch(str(output.get("sha256", ""))):
|
||
|
|
raise ValueError("output records require a SHA-256")
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
||
|
|
init = subparsers.add_parser("init")
|
||
|
|
init.add_argument("--workspace", type=Path, default=Path.cwd())
|
||
|
|
init.add_argument("--run-id", required=True)
|
||
|
|
init.add_argument("--objective", required=True)
|
||
|
|
init.add_argument("--projects", required=True)
|
||
|
|
init.add_argument("--include-shared", action="store_true")
|
||
|
|
init.add_argument("--lens", required=True)
|
||
|
|
init.add_argument("--budgets-json", default="{}")
|
||
|
|
update = subparsers.add_parser("update")
|
||
|
|
update.add_argument("--workspace", type=Path, default=Path.cwd())
|
||
|
|
update.add_argument("--run-id", required=True)
|
||
|
|
show = subparsers.add_parser("show")
|
||
|
|
show.add_argument("--workspace", type=Path, default=Path.cwd())
|
||
|
|
show.add_argument("--run-id", required=True)
|
||
|
|
arguments = parser.parse_args()
|
||
|
|
if arguments.command == "init":
|
||
|
|
checkpoint = create_checkpoint(
|
||
|
|
arguments.workspace,
|
||
|
|
run_id=arguments.run_id,
|
||
|
|
objective=arguments.objective,
|
||
|
|
projects=[item for item in arguments.projects.split(",") if item],
|
||
|
|
include_shared=arguments.include_shared,
|
||
|
|
lens=arguments.lens,
|
||
|
|
budgets=json.loads(arguments.budgets_json),
|
||
|
|
)
|
||
|
|
elif arguments.command == "show":
|
||
|
|
checkpoint = load_checkpoint(arguments.workspace, arguments.run_id)
|
||
|
|
else:
|
||
|
|
patch_text = sys.stdin.read(64_001)
|
||
|
|
if len(patch_text.encode("utf-8")) > 64_000:
|
||
|
|
raise SystemExit("checkpoint update exceeds 64 KiB")
|
||
|
|
patch = json.loads(patch_text, object_pairs_hook=_reject_duplicate_keys)
|
||
|
|
if not isinstance(patch, dict):
|
||
|
|
raise SystemExit("checkpoint update must be one JSON object on stdin")
|
||
|
|
checkpoint = update_checkpoint(arguments.workspace, arguments.run_id, patch)
|
||
|
|
print(json.dumps(checkpoint, ensure_ascii=False, sort_keys=True))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|