"""Determined Core snapshots and a deterministic CPU pause/resume fixture.""" from __future__ import annotations import argparse from contextlib import contextmanager import hashlib import json import math import os from pathlib import Path, PurePosixPath import shutil import tempfile import time from typing import Any, Callable import uuid FORMAT = "hessian-managed-snapshot-v1" SEAL = "checkpoint-complete.json" CANONICAL_ROOT = Path("/pfss/mlde/workspaces/mlde_wsp_JAIF/hessian_recipes") ROOT_FILES = {"run-plan.json": "run_plan_sha256", "runtime-binding.json": "runtime_binding_sha256", "initial-adapter.pt": "initial_adapter_sha256"} PREFIX_FILES = {"rollout-evidence.jsonl": "evidence_prefix", "native-tool-diagnostics.jsonl": "diagnostics_prefix"} NATIVE_STATE = {"optimizer.pt", "scheduler.pt", "rng_state.pth", "trainer_state.json", "adapter_config.json", "adapter_model.safetensors", "trainable-state.pt", "learning-evidence.json"} def _json_bytes(value: Any) -> bytes: return (json.dumps(value, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n").encode() def _read_json(path: Path) -> dict: _regular(path) def unique(pairs): result = {} for key, value in pairs: if key in result: raise ValueError("Duplicate key in managed state") result[key] = value return result value = json.loads(path.read_text(), object_pairs_hook=unique) if not isinstance(value, dict): raise ValueError("Managed state must be a JSON object") return value def _write_json(path: Path, value: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile(dir=path.parent, prefix=path.name + ".", delete=False) as stream: temporary = Path(stream.name) stream.write(_json_bytes(value)) stream.flush() os.fsync(stream.fileno()) try: os.replace(temporary, path) finally: temporary.unlink(missing_ok=True) def _relative(name: str) -> str: if (not isinstance(name, str) or not name or "\\" in name or PurePosixPath(name).is_absolute() or any(part in {"", ".", ".."} for part in name.split("/"))): raise ValueError("Unsafe snapshot filename") return name def _regular(path: Path) -> None: if path.is_symlink() or any(parent.is_symlink() for parent in path.parents): raise ValueError("Managed state may not contain symlinks") if not path.is_file(): raise ValueError("Managed snapshot source is not a regular file") def _copy(source: Path, destination: Path, *, limit: int | None = None) -> dict: _regular(source) if limit is not None and (type(limit) is not int or limit < 0): raise ValueError("Invalid evidence prefix length") digest, consumed = hashlib.sha256(), 0 destination.parent.mkdir(parents=True, exist_ok=True) with source.open("rb") as reader, destination.open("wb") as writer: while limit is None or consumed < limit: block = reader.read(1024 * 1024 if limit is None else min(1024 * 1024, limit - consumed)) if not block: break writer.write(block) digest.update(block) consumed += len(block) writer.flush() os.fsync(writer.fileno()) if limit is not None and consumed != limit: raise ValueError("Evidence is shorter than its sealed prefix") return {"bytes": consumed, "sha256": digest.hexdigest()} def _digest(path: Path) -> str: _regular(path) digest = hashlib.sha256() with path.open("rb") as stream: for block in iter(lambda: stream.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def _inventory(path: Path) -> dict: result = {} if path.is_symlink(): raise ValueError("Managed snapshot directory is a symlink") for entry in sorted(path.rglob("*")): if entry.is_symlink(): raise ValueError("Managed snapshot may not contain symlinks") if entry.is_file(): result[_relative(entry.relative_to(path).as_posix())] = { "bytes": entry.stat().st_size, "sha256": _digest(entry)} return result class ManagedTrial: """Bind one worker to frozen input identity and a deadline persistent across allocations. The binding file stays outside native output so interruption before the first scheduler checkpoint cannot reset the wall budget. """ def __init__(self, context: Any, info: Any, *, identity: dict, binding_path: str | Path, max_wall_seconds: int, clock: Callable[[], float] = time.time, native_validator: Callable[[Path, Path], Any] | None = None): if info is None or info.task_type != "TRIAL": raise ValueError("Managed execution requires an actual Determined trial") if (len(info.container_addrs) > 1 or len(info.slot_ids) > 1 or type(max_wall_seconds) is not int or max_wall_seconds <= 0): raise ValueError("Managed adapter supports one bounded worker") if not isinstance(identity, dict) or not identity: raise ValueError("Managed trial requires source/config identity") self.identity = json.loads(_json_bytes(identity)) self.context, self.info, self.clock = context, info, clock self.native_validator = native_validator self.binding_path = Path(binding_path).absolute() self.max_wall_seconds = max_wall_seconds self.binding: dict | None = None self.last_uuid: str | None = None self.last_step: int | None = None self.preempt_step: int | None = None if self.binding_path.exists(): self._adopt_binding(_read_json(self.binding_path)) elif not info.latest_checkpoint: started = clock() self._adopt_binding({"format_version": FORMAT, "experiment_id": info.trial.experiment_id, "trial_id": info.trial.trial_id, "identity": self.identity, "started_at": started, "deadline": started + max_wall_seconds, "max_wall_seconds": max_wall_seconds}) def _adopt_binding(self, binding: dict) -> None: expected = {"format_version": FORMAT, "experiment_id": self.info.trial.experiment_id, "trial_id": self.info.trial.trial_id, "identity": self.identity, "max_wall_seconds": self.max_wall_seconds} if any(binding.get(key) != value for key, value in expected.items()): raise ValueError("Managed trial binding or source/config identity changed") started, deadline = binding.get("started_at"), binding.get("deadline") if (type(started) not in (int, float) or type(deadline) not in (int, float) or not math.isfinite(started) or not math.isfinite(deadline) or deadline != started + self.max_wall_seconds): raise ValueError("Invalid persisted managed deadline") if self.binding is not None and binding != self.binding: raise ValueError("Scheduler snapshot differs from persistent trial binding") if self.clock() >= deadline: raise TimeoutError("Managed trial's original wall-clock deadline expired") self.binding = binding if not self.binding_path.exists(): self.binding_path.parent.mkdir(parents=True, exist_ok=True) try: descriptor = os.open(self.binding_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) except FileExistsError: if _read_json(self.binding_path) != binding: raise ValueError("Concurrent managed binding differs") from None else: with os.fdopen(descriptor, "wb") as stream: stream.write(_json_bytes(binding)) stream.flush() os.fsync(stream.fileno()) def check_deadline(self) -> None: if self.binding is None: raise ValueError("Restore the scheduler checkpoint before managed execution") if self.clock() >= self.binding["deadline"]: raise TimeoutError("Managed trial's original wall-clock deadline expired") def report(self, step: int, metrics: dict, operation: Any = None) -> None: self.check_deadline() if (type(step) is not int or step < 0 or not metrics or any(type(value) not in (int, float) or not math.isfinite(value) for value in metrics.values())): raise ValueError("Managed metrics require finite numbers and an optimizer step") self.context.train.report_training_metrics(steps_completed=step, metrics=metrics) if operation is not None: operation.report_progress(step) def should_preempt(self, step: int) -> bool: self.check_deadline() requested = self.context.preempt.should_preempt(auto_ack=False) if requested: self.preempt_step = step return requested def preempt_requested(self, step: int) -> bool: return self.should_preempt(step) def acknowledge_preemption(self) -> None: if self.preempt_step is None or self.last_step != self.preempt_step or not self.last_uuid: raise ValueError("Publish the current optimizer state before acknowledging preemption") self.context.preempt.acknowledge_preemption_signal() def _publish(self, files: dict, *, step: int, kind: str, output: Path | None = None, checkpoint_name: str | None = None, retained_checkpoints: list[str] | None = None) -> str: self.check_deadline() if type(step) is not int or step < 1: raise ValueError("Managed checkpoint requires completed steps") if self.last_step is not None and step <= self.last_step: raise ValueError("Managed checkpoint must advance beyond the restored or last published step") metadata = {"steps_completed": step, "kind": kind, "experiment_id": self.info.trial.experiment_id, "trial_id": self.info.trial.trial_id} with self.context.checkpoint.store_path(metadata) as (snapshot, storage_id): snapshot = Path(snapshot) payload = snapshot / "payload" payload.mkdir() for name, (source, limit, expected) in files.items(): actual = _copy(source, payload / _relative(name), limit=limit) if expected is not None and actual != expected: raise ValueError("Copied state differs from its native seal") _write_json(snapshot / "snapshot.json", { "format_version": FORMAT, "binding": self.binding, "kind": kind, "steps_completed": step, "scheduler_uuid": storage_id, "publishing_allocation_id": self.info.allocation_id, "canonical_output": str(output) if output is not None else None, "checkpoint_name": checkpoint_name, "retained_checkpoints": retained_checkpoints, "files": _inventory(payload)}) self.last_uuid, self.last_step = storage_id, step return storage_id @contextmanager def _restored(self, kind: str): storage_id = self.info.latest_checkpoint if not storage_id: yield None, None return with self.context.checkpoint.restore_path(storage_id) as snapshot: snapshot = Path(snapshot) envelope = _read_json(snapshot / "snapshot.json") if (envelope.get("format_version") != FORMAT or envelope.get("kind") != kind or envelope.get("scheduler_uuid") != storage_id or type(envelope.get("steps_completed")) is not int): raise ValueError("Scheduler checkpoint snapshot identity differs") self._adopt_binding(envelope.get("binding", {})) payload = snapshot / "payload" if _inventory(payload) != envelope.get("files"): raise ValueError("Scheduler checkpoint payload hashes differ") self.last_uuid, self.last_step = storage_id, envelope["steps_completed"] yield payload, envelope def publish_trainer(self, output: str | Path, checkpoint: str | Path) -> str: """Retain native recovery state and all earlier checkpoints needed by DEV selection.""" output, checkpoint = Path(output).absolute(), Path(checkpoint).absolute() if checkpoint.parent != output or checkpoint.is_symlink(): raise ValueError("Managed checkpoint must belong to its canonical run") current = _read_json(checkpoint / SEAL) step = current.get("optimizer_step") if type(step) is not int or checkpoint.name != "checkpoint-" + str(step): raise ValueError("Native seal and checkpoint step differ") candidates = sorted((path for path in output.glob("checkpoint-*") if path.name.removeprefix("checkpoint-").isdigit() and int(path.name.removeprefix("checkpoint-")) <= step), key=lambda path: int(path.name.removeprefix("checkpoint-"))) if checkpoint not in candidates: raise ValueError("Current native checkpoint disappeared") if len(candidates) > 1 and self.native_validator is None: raise ValueError("Retaining earlier checkpoints requires the native validator") files = {} for candidate in candidates: if self.native_validator is not None: self.native_validator(output, candidate) seal = _read_json(candidate / SEAL) if seal.get("optimizer_step") != int(candidate.name.removeprefix("checkpoint-")): raise ValueError("Retained checkpoint step differs from its seal") native_files = _inventory(candidate) expected = seal.get("files_sha256") if (not isinstance(expected, dict) or not NATIVE_STATE <= set(expected) or {name: value["sha256"] for name, value in native_files.items() if name != SEAL} != expected): raise ValueError("Native checkpoint is incomplete or changed") files.update({candidate.name + "/" + name: (candidate / name, None, metadata) for name, metadata in native_files.items()}) for name, seal_key in ROOT_FILES.items(): source = output / name if _digest(source) != seal.get(seal_key): raise ValueError("Native root provenance differs from checkpoint seal") files[name] = (source, None, {"bytes": source.stat().st_size, "sha256": seal[seal_key]}) for name, seal_key in PREFIX_FILES.items(): prefix = seal.get(seal_key) if (not isinstance(prefix, dict) or set(prefix) != {"bytes", "sha256"} or prefix["bytes"] > current[seal_key]["bytes"]): raise ValueError("Native seal lacks an ordered evidence prefix") # Backend validation checks each earlier prefix. The newest prefix # contains all committed evidence needed to verify earlier seals. files[name] = (output / name, current[seal_key]["bytes"], current[seal_key]) if "managed_binding_sha256" in seal: source = output / "managed-binding.json" if (_read_json(source) != self.binding or _digest(source) != seal["managed_binding_sha256"]): raise ValueError("Native managed binding differs from this trial") files["managed-binding.json"] = (source, None, { "bytes": source.stat().st_size, "sha256": seal["managed_binding_sha256"]}) tokenizer = output / "policy-tokenizer" if tokenizer.exists(): for name, metadata in _inventory(tokenizer).items(): files["policy-tokenizer/" + name] = (tokenizer / name, None, metadata) manifest = output / "training-manifest.json" if manifest.exists(): files["training-manifest.json"] = (manifest, None, None) return self._publish(files, step=step, kind="native-trainer", output=output, checkpoint_name=checkpoint.name, retained_checkpoints=[path.name for path in candidates]) def restore_trainer(self, output: str | Path, validator: Callable[[Path, Path, str], Any]) -> Path | None: """Validate (staged_root, canonical_output, checkpoint_name) before replacement. The validator must check native seals without rewriting absolute provenance. Existing local output is retained in an adjacent recovery directory. """ output = Path(output).absolute() with self._restored("native-trainer") as (payload, envelope): if payload is None: return None name = _relative(envelope["checkpoint_name"]) if ("/" in name or name != "checkpoint-" + str(envelope["steps_completed"]) or envelope.get("canonical_output") != str(output)): raise ValueError("Scheduler snapshot belongs to another canonical output") if output.is_symlink() or any(parent.is_symlink() for parent in output.parents): raise ValueError("Canonical output may not be a symlink") output.parent.mkdir(parents=True, exist_ok=True) staging = Path(tempfile.mkdtemp(prefix=output.name + ".managed-restore-", dir=output.parent)) try: shutil.copytree(payload, staging, dirs_exist_ok=True) retained = envelope.get("retained_checkpoints") if (not isinstance(retained, list) or name not in retained or len(set(retained)) != len(retained) or any(not isinstance(item, str) or "/" in item or not item.startswith("checkpoint-") or not item.removeprefix("checkpoint-").isdigit() or not 1 <= int(item.removeprefix("checkpoint-")) <= envelope["steps_completed"] for item in retained) or {entry.name for entry in staging.glob("checkpoint-*")} != set(retained)): raise ValueError("Scheduler snapshot retained checkpoint list differs") for retained_name in retained: validator(staging, output, retained_name) self.check_deadline() quarantine = None if output.exists(): history = output.parent / (output.name + ".managed-recovery") if history.is_symlink(): raise ValueError("Managed recovery directory may not be a symlink") history.mkdir(exist_ok=True) quarantine = history / ("before-restore-" + uuid.uuid4().hex) os.replace(output, quarantine) try: os.replace(staging, output) except BaseException: if quarantine is not None: os.replace(quarantine, output) raise finally: if staging.exists(): shutil.rmtree(staging) return output / name publish = publish_trainer restore = restore_trainer def publish_cpu(self, state: dict, output: Path) -> str: with tempfile.TemporaryDirectory(prefix="managed-cpu-state-", dir=output) as temporary: source = Path(temporary) / "cpu-state.json" _write_json(source, state) return self._publish({"cpu-state.json": (source, None, None)}, step=state["steps_completed"], kind="cpu-lifecycle") def restore_cpu(self) -> dict | None: with self._restored("cpu-lifecycle") as (payload, envelope): if payload is None: return None if set(envelope["files"]) != {"cpu-state.json"}: raise ValueError("Unexpected CPU snapshot files") state = _read_json(payload / "cpu-state.json") if state.get("steps_completed") != envelope["steps_completed"]: raise ValueError("CPU state and scheduler step disagree") return state def _advance(state: dict) -> dict: step = state["steps_completed"] + 1 value = (state["value"] * 1664525 + 1013904223) % (2 ** 32) chain = hashlib.sha256((state["chain_sha256"] + ":" + str(step) + ":" + str(value)).encode()).hexdigest() return state | {"steps_completed": step, "value": value, "chain_sha256": chain} def cpu_lifecycle(context: Any, info: Any, output: str | Path, *, steps: int = 8, delay: float = 3, max_wall_seconds: int = 900, clock: Callable[[], float] = time.time, sleeper: Callable[[float], None] = time.sleep) -> dict: """CPU scheduler lifecycle fixture; does not qualify GPU training recovery.""" if type(steps) is not int or not 1 <= steps <= 100 or delay < 0 or not math.isfinite(delay): raise ValueError("CPU lifecycle requires bounded steps and delay") if info is None or info.task_type != "TRIAL" or info.slot_ids: raise ValueError("CPU lifecycle requires a zero-slot Determined trial") output = Path(output).absolute() / ("trial-" + str(info.trial.trial_id)) output.mkdir(parents=True, exist_ok=True) identity = {"protocol": "hessian-core-cpu-lifecycle-v1", "steps": steps, "source_sha256": _digest(Path(__file__))} managed = ManagedTrial(context, info, identity=identity, binding_path=output / "binding.json", max_wall_seconds=max_wall_seconds, clock=clock) state = managed.restore_cpu() restored = state is not None if state is None: state = {"steps_completed": 0, "value": 17, "chain_sha256": hashlib.sha256(b"hessian-cpu").hexdigest(), "allocations": []} expected = {"steps_completed": 0, "value": 17, "chain_sha256": hashlib.sha256(b"hessian-cpu").hexdigest()} for _ in range(state["steps_completed"]): expected = _advance(expected) if any(state.get(key) != value for key, value in expected.items()): raise ValueError("Restored CPU state differs from uninterrupted reference") if not isinstance(state.get("allocations"), list): raise ValueError("CPU allocation history is invalid") if info.allocation_id not in state["allocations"]: state["allocations"].append(info.allocation_id) result = {"status": "running", "experiment_id": info.trial.experiment_id, "trial_id": info.trial.trial_id, "allocation_id": info.allocation_id, "restored_from_uuid": info.latest_checkpoint, "restored": restored} for operation in context.searcher.operations(): if operation.length != steps: raise ValueError("CPU searcher horizon differs from fixture protocol") while state["steps_completed"] < operation.length: managed.check_deadline() state = _advance(state) sleeper(delay) managed.publish_cpu(state, output) managed.report(state["steps_completed"], {"lifecycle_value": state["value"]}, operation) result.update(steps_completed=state["steps_completed"], scheduler_uuid=managed.last_uuid, state=state, deadline=managed.binding["deadline"]) _write_json(output / "status.json", result) if managed.should_preempt(state["steps_completed"]): result["status"] = "preempted" _write_json(output / "status.json", result) managed.acknowledge_preemption() return result reference = {"steps_completed": 0, "value": 17, "chain_sha256": hashlib.sha256(b"hessian-cpu").hexdigest()} for _ in range(steps): reference = _advance(reference) if any(state[key] != value for key, value in reference.items()): raise ValueError("Final CPU state differs from uninterrupted reference") context.train.report_validation_metrics(steps_completed=state["steps_completed"], metrics={"lifecycle_value": state["value"], "reference_matches": 1}) operation.report_completed(float(state["value"])) result["status"] = "completed" _write_json(output / "status.json", result) _write_json(output / "result.json", result) return result def main() -> None: parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) cpu = commands.add_parser("cpu") cpu.add_argument("--output", type=Path, default=CANONICAL_ROOT / "runs/managed-cpu") cpu.add_argument("--steps", type=int, default=8) cpu.add_argument("--delay", type=float, default=3) cpu.add_argument("--max-wall-seconds", type=int, default=900) args = parser.parse_args() import determined as det from determined import core info = det.get_cluster_info() if info is None or info.task_type != "TRIAL": parser.error("Run in a managed Determined experiment") with core.init(tensorboard_mode=core.TensorboardMode.MANUAL) as context: result = cpu_lifecycle(context, info, args.output, steps=args.steps, delay=args.delay, max_wall_seconds=args.max_wall_seconds) print(json.dumps(result, sort_keys=True)) if __name__ == "__main__": main()