Blame for src/dlb/runner.py

"""Resumable, provenance-bound execution of one baseline sample request."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime, timezone
import argparse
import hashlib
import importlib
import inspect
from importlib import metadata as importlib_metadata
import json
import os
from pathlib import Path
import platform
import socket
import stat
import subprocess
import sys
from typing import Iterable, Literal, Protocol, Sequence

from dlb.checkpoints import load_checkpoint_manifest
from dlb.conditional_prompts import load_protocol, verify_prompts
from dlb.io import (
    SampleValidationError,
    atomic_json_write,
    ensure_safe_directory,
    expected_conditional_schedule,
    open_safe_output,
    remove_safe_file,
    sha256_file,
    validate_conditional_samples,
    validate_samples,
    write_conditional_samples_atomic,
    write_samples_atomic,
)
from dlb.registry import load_registry, step_grid_for_model
from dlb.schema import SampleRecord


class SampleAdapter(Protocol):
    """Small adapter boundary implemented by the model-specific Task 7 modules."""

    identity: str

    def build_command(self, request: "RunRequest", run_dir: Path) -> Sequence[str]: ...

    def convert_outputs(
        self, request: "RunRequest", run_dir: Path
    ) -> Iterable[SampleRecord | dict[str, object]]: ...


@dataclass(frozen=True)
class RunRequest:
    run_id: str
    model_id: str
    dataset_id: str
    step_count: int
    seed: int
    sample_count: int = 1024
    command: list[str] | None = None
    config_sha256: str | None = None
    source_sha256: str | None = None
    checkpoint_sha256: str | None = None
    checkpoint_lock_id: str | None = None
    checkpoint_selection: dict[str, object] | None = None
    checkpoint_teacher_family: str | None = None
    adapter_identity: str | None = None
    environment: str | None = None
    device: str | None = None
    results_root: str | None = None
    generation_mode: Literal["unconditional", "conditional_prefix"] = "unconditional"
    conditioning_manifest: str | None = None
    conditioning_manifest_sha256: str | None = None
    conditioning_config_sha256: str | None = None
    prefix_length: int | None = None
    evaluation_continuation_length: int | None = None
    prompt_count: int | None = None
    diversity_prompt_count: int | None = None
    completions_per_diversity_prompt: int | None = None
    completion_schedule: str | None = None


@dataclass(frozen=True)
class RunResult:
    status: str
    run_dir: Path
    returncode: int


@dataclass(frozen=True)
class CheckpointProvenance:
    sha256: str
    lock_id: str
    selection: dict[str, object]
    teacher_family: str


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _canonical_sha256(value: object) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def _tail(path: Path, limit: int = 8192) -> str:
    with path.open("rb") as log_file:
        log_file.seek(0, os.SEEK_END)
        size = log_file.tell()
        log_file.seek(max(0, size - limit))
        return log_file.read().decode("utf-8", errors="replace")


def _safe_json_load(path: Path) -> dict[str, object] | None:
    try:
        if path.is_symlink() or not path.is_file():
            return None
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
        return None
    return value if isinstance(value, dict) else None


def _environment_evidence() -> dict[str, object]:
    torch_version: str | None
    try:
        torch_version = importlib_metadata.version("torch")
    except importlib_metadata.PackageNotFoundError:
        torch_version = None
    return {
        "hostname": socket.gethostname(),
        "platform": platform.platform(),
        "python": platform.python_version(),
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "cuda_version": os.environ.get("CUDA_VERSION"),
        "torch_version": torch_version,
    }


def _adapter_identity(adapter: SampleAdapter, request: RunRequest) -> str:
    del request
    declared_identity = getattr(adapter, "identity", None)
    if not isinstance(declared_identity, str) or not declared_identity:
        raise ValueError("adapter must declare a non-empty identity")
    try:
        source_path = inspect.getsourcefile(type(adapter))
    except (OSError, TypeError):
        source_path = None
    class_identity = f"{type(adapter).__module__}:{type(adapter).__qualname__}"
    if source_path is not None and Path(source_path).is_file():
        implementation = sha256_file(Path(source_path))
    else:
        implementation = "source-unavailable"
    return f"{declared_identity}|{class_identity}|{implementation}"


def _resolve_checkpoint_provenance(
    root: Path, request: RunRequest, recipe_id: str | None
) -> CheckpointProvenance:
    """Resolve a coverage resource to the exact verified lock-file inventory hash."""

    manifest_path = root / "artifacts" / "checkpoints.yaml"
    if not manifest_path.is_file():
        raise ValueError("canonical checkpoint manifest is missing")
    manifest = load_checkpoint_manifest(manifest_path)
    coverage = manifest.coverage.get((request.model_id, request.dataset_id))
    manifest_sha256 = sha256_file(manifest_path)
    if coverage is None and recipe_id is not None:
        recipe = manifest.recipes.get(recipe_id)
        if recipe is None:
            raise ValueError(f"unknown training recipe: {recipe_id}")
        output = root / recipe.output
        inventory = _checkpoint_inventory(root, output)
        selector = {
            "recipe": recipe_id,
            "output": recipe.output,
            "source": recipe.source,
            "source_commit": recipe.source_commit,
            "teacher_family": recipe.teacher_family,
            "teacher_adapter": recipe.teacher_adapter,
            "sampling_checkpoint": recipe.sampling_checkpoint,
            "sampling_config": recipe.sampling_config,
        }
        digest = _canonical_sha256(
            {"manifest_sha256": manifest_sha256, "selector": selector, "files": inventory}
        )
        return CheckpointProvenance(
            sha256=digest,
            lock_id=f"recipe:{recipe_id}:{manifest_sha256}",
            selection=selector,
            teacher_family=recipe.teacher_family,
        )
    if coverage is None:
        raise ValueError(
            f"canonical checkpoint selection is missing for {request.model_id}/{request.dataset_id}"
        )
    lock = _safe_json_load(root / "artifacts" / "checkpoint_lock.json")
    if lock is None:
        raise ValueError("canonical checkpoint lock is missing or invalid")
    if lock.get("manifest_sha256") != manifest_sha256:
        raise ValueError("checkpoint lock does not match the current checkpoint manifest")
    resources = lock.get("resources")
    record = resources.get(coverage.resource) if isinstance(resources, dict) else None
    files = record.get("files") if isinstance(record, dict) else None
    if record is None or record.get("status") != "downloaded" or not isinstance(files, list):
        raise ValueError(f"checkpoint lock is not verified for {coverage.resource}")
    inventory: list[dict[str, object]] = []
    for item in files:
        if not isinstance(item, dict) or not isinstance(item.get("sha256"), str):
            raise ValueError(f"checkpoint lock has invalid file inventory for {coverage.resource}")
        if coverage.path is not None and not str(item.get("path", "")).endswith(
            "/" + coverage.path
        ):
            continue
        inventory.append(
            {
                "path": item.get("path"),
                "size_bytes": item.get("size_bytes"),
                "sha256": item["sha256"],
            }
        )
    if coverage.sampling_config_source == "project":
        config_path = root / str(coverage.sampling_config)
        if config_path.is_symlink() or not config_path.is_file():
            raise ValueError("project sampling config is missing or unsafe")
        observed_config_sha256 = sha256_file(config_path)
        if observed_config_sha256 != coverage.sampling_config_sha256:
            raise ValueError("project sampling config differs from checkpoint manifest")
        inventory.append(
            {
                "path": config_path.relative_to(root).as_posix(),
                "size_bytes": config_path.stat().st_size,
                "sha256": observed_config_sha256,
            }
        )
    if not inventory:
        raise ValueError(f"checkpoint lock lacks selected file for {coverage.resource}")
    selector = {
        "resource": coverage.resource,
        "path": coverage.path,
        "teacher_family": coverage.teacher_family,
    }
    if coverage.sampling_config is not None:
        selector.update(
            {
                "sampling_config": coverage.sampling_config,
                "sampling_config_source": coverage.sampling_config_source,
                "sampling_config_sha256": coverage.sampling_config_sha256,
                "sampling_config_source_commit": coverage.sampling_config_source_commit,
            }
        )
    digest = _canonical_sha256(
        {
            "manifest_sha256": manifest_sha256,
            "selector": selector,
            "files": sorted(inventory, key=lambda item: str(item["path"])),
        }
    )
    resource = manifest.resources[coverage.resource]
    return CheckpointProvenance(
        sha256=digest,
        lock_id=f"{coverage.resource}:{manifest_sha256}:{coverage.path or 'all'}",
        selection=selector,
        teacher_family=coverage.teacher_family or resource.teacher_family,
    )


def _resolve_checkpoint_identity(
    root: Path, request: RunRequest, recipe_id: str | None
) -> tuple[str, str]:
    """Compatibility wrapper for callers that need only the content and lock identities."""

    provenance = _resolve_checkpoint_provenance(root, request, recipe_id)
    return provenance.sha256, provenance.lock_id


def _checkpoint_inventory(root: Path, output: Path) -> list[dict[str, object]]:
    """Hash a recipe output tree while refusing symlinks and empty outputs."""

    if not output.exists() or output.is_symlink():
        raise ValueError(f"recipe checkpoint output is missing or unsafe: {output}")
    paths = [output] if output.is_file() else sorted(output.rglob("*"))
    inventory: list[dict[str, object]] = []
    for path in paths:
        metadata = path.lstat()
        if stat.S_ISLNK(metadata.st_mode):
            raise ValueError(f"recipe checkpoint output contains a symlink: {path}")
        if not stat.S_ISREG(metadata.st_mode):
            continue
        if metadata.st_size <= 0:
            if path.suffix == ".log":
                continue
            raise ValueError(f"recipe checkpoint output contains an empty file: {path}")
        inventory.append(
            {
                "path": path.relative_to(root).as_posix(),
                "size_bytes": metadata.st_size,
                "sha256": sha256_file(path),
            }
        )
    if not inventory:
        raise ValueError(f"recipe checkpoint output contains no files: {output}")
    return inventory


def load_adapter(adapter_id: str) -> SampleAdapter:
    """Import a Task 7 adapter without encoding model implementation details here."""

    module_name, separator, attribute = adapter_id.partition(":")
    module = importlib.import_module(module_name if separator else f"dlb.adapters.{adapter_id}")
    candidate = getattr(module, attribute, None) if separator else getattr(module, "adapter", None)
    candidate = candidate if candidate is not None else getattr(module, "ADAPTER", None)
    if candidate is None:
        candidate = getattr(module, "Adapter", None)
    if candidate is None:
        raise ValueError(f"adapter {adapter_id!r} has no adapter, ADAPTER, or Adapter export")
    return candidate() if isinstance(candidate, type) else candidate


_CONDITIONAL_FIELDS = (
    "conditioning_manifest",
    "conditioning_manifest_sha256",
    "conditioning_config_sha256",
    "prefix_length",
    "evaluation_continuation_length",
    "prompt_count",
    "diversity_prompt_count",
    "completions_per_diversity_prompt",
    "completion_schedule",
)


def _completion_schedule_text(prompt_count: int, diversity_prompt_count: int, completions: int) -> str:
    """Serialize the canonical schedule in the adapter-facing compact form."""

    return f"c0:p0-{prompt_count - 1};c1-{completions - 1}:p0-{diversity_prompt_count - 1}"


def _conditioning_sidecar_path(root: Path, dataset_id: str) -> Path:
    return root / "data" / "manifests" / f"conditional-{dataset_id}-c64.json"


def _resolve_conditional_request(request: RunRequest, root: Path) -> dict[str, object]:
    """Verify prompt/data bindings and derive the only publishable C64 contract."""

    config_path = root / "configs" / "conditional.yaml"
    if config_path.is_symlink() or not config_path.is_file():
        raise ValueError("canonical conditional protocol is missing: configs/conditional.yaml")
    protocol = load_protocol(config_path)
    if request.seed != protocol.sampling_seed:
        raise ValueError("conditional seed differs from the verified protocol")
    manifest = verify_prompts(root, request.dataset_id, protocol)
    if manifest.dataset != request.dataset_id or manifest.protocol != protocol.protocol:
        raise ValueError("verified conditional prompt manifest differs from request")
    if (
        manifest.prompt_count != protocol.prompt_count
        or manifest.prefix_length != protocol.prefix_length
        or manifest.evaluation_continuation_length != protocol.evaluation_continuation_length
        or manifest.model_length != protocol.datasets[request.dataset_id].model_length
    ):
        raise ValueError("verified conditional prompt manifest differs from protocol")
    sidecar_path = _conditioning_sidecar_path(root, request.dataset_id)
    if sidecar_path.is_symlink() or not sidecar_path.is_file():
        raise ValueError(f"conditional prompt manifest is missing or unsafe: {sidecar_path}")
    config_sha256 = sha256_file(config_path)
    schedule = expected_conditional_schedule(
        protocol.prompt_count,
        protocol.diversity_prompt_count,
        protocol.completions_per_diversity_prompt,
    )
    completion_schedule = _completion_schedule_text(
        protocol.prompt_count,
        protocol.diversity_prompt_count,
        protocol.completions_per_diversity_prompt,
    )
    expected = {
        "conditioning_manifest": str(sidecar_path),
        "conditioning_manifest_sha256": sha256_file(sidecar_path),
        "conditioning_config_sha256": config_sha256,
        "prefix_length": protocol.prefix_length,
        "evaluation_continuation_length": protocol.evaluation_continuation_length,
        "prompt_count": protocol.prompt_count,
        "diversity_prompt_count": protocol.diversity_prompt_count,
        "completions_per_diversity_prompt": protocol.completions_per_diversity_prompt,
        "completion_schedule": completion_schedule,
    }
    for name, value in expected.items():
        asserted = getattr(request, name)
        if asserted is not None and asserted != value:
            raise ValueError(f"conditional {name} assertion differs from verified contract")
    if request.sample_count != len(schedule):
        raise ValueError(f"conditional sample_count must be {len(schedule)} for the verified schedule")
    return expected


def _conditional_results_root(request: RunRequest, root: Path) -> str:
    """Return an isolated conditional result root, preserving smoke descendants."""

    isolated_root = (root / "results" / "conditional").resolve()
    if request.results_root is None:
        return str(isolated_root)
    requested = Path(request.results_root)
    if requested.is_symlink():
        raise ValueError("conditional results_root is a symlink")
    resolved = requested.resolve()
    try:
        resolved.relative_to(isolated_root)
    except ValueError as error:
        raise ValueError(
            "conditional results_root must be root/results/conditional or a descendant"
        ) from error
    return str(resolved)


def _conditional_io_contract(request: RunRequest, root: Path) -> tuple[list[tuple[int, int]], int, int]:
    """Re-read the verified prompt contract needed by conditional publication."""

    protocol = load_protocol(root / "configs" / "conditional.yaml")
    manifest = verify_prompts(root, request.dataset_id, protocol)
    if (
        request.prompt_count != protocol.prompt_count
        or request.diversity_prompt_count != protocol.diversity_prompt_count
        or request.completions_per_diversity_prompt != protocol.completions_per_diversity_prompt
        or request.prefix_length != manifest.prefix_length
        or request.evaluation_continuation_length != manifest.evaluation_continuation_length
    ):
        raise ValueError("conditional request differs from verified prompt contract")
    return (
        expected_conditional_schedule(
            protocol.prompt_count,
            protocol.diversity_prompt_count,
            protocol.completions_per_diversity_prompt,
        ),
        manifest.model_length,
        manifest.vocabulary_size,
    )


def _resolve_request(request: RunRequest, root: Path, adapter: SampleAdapter | None) -> tuple[RunRequest, SampleAdapter]:
    registry_path = root / "configs" / "experiments.yaml"
    if not registry_path.is_file():
        raise ValueError("canonical registry is missing: configs/experiments.yaml")
    registry = load_registry(registry_path)
    try:
        model = registry.models[request.model_id]
        support = model.datasets[request.dataset_id]
    except KeyError as error:
        raise ValueError(f"unknown model/dataset cell: {request.model_id}/{request.dataset_id}") from error
    if support.status != "supported":
        raise ValueError(f"unsupported model/dataset cell: {request.model_id}/{request.dataset_id}")
    allowed_steps = step_grid_for_model(registry, request.model_id)
    if request.step_count not in allowed_steps:
        joined_steps = ",".join(str(step) for step in allowed_steps)
        raise ValueError(
            f"invalid step count {request.step_count} for "
            f"{request.model_id}/{request.dataset_id}; allowed: {joined_steps}"
        )
    if request.environment is not None and request.environment != model.environment:
        raise ValueError("environment assertion differs from canonical registry")
    if request.sample_count <= 0:
        raise ValueError("sample_count must be positive")
    if request.step_count <= 0:
        raise ValueError("step_count must be positive")
    if request.generation_mode not in {"unconditional", "conditional_prefix"}:
        raise ValueError(f"unsupported generation mode: {request.generation_mode}")
    if request.generation_mode == "unconditional":
        if any(getattr(request, name) is not None for name in _CONDITIONAL_FIELDS):
            raise ValueError("unconditional requests must not include conditional fields")
        conditional = {}
    else:
        conditional = {
            **_resolve_conditional_request(request, root),
            "results_root": _conditional_results_root(request, root),
        }

    source_lock = _safe_json_load(root / "artifacts" / "source_lock.json")
    sources = source_lock.get("sources") if source_lock is not None else None
    source = sources.get(model.source) if isinstance(sources, dict) else None
    source_sha256 = source.get("commit") if isinstance(source, dict) else None
    if not isinstance(source_sha256, str) or not source_sha256:
        raise ValueError(f"canonical source lock lacks SHA for {model.source}")
    if request.source_sha256 is not None and request.source_sha256 != source_sha256:
        raise ValueError("source SHA assertion differs from canonical source lock")

    config_sha256 = sha256_file(registry_path)
    if request.config_sha256 is not None and request.config_sha256 != config_sha256:
        raise ValueError("config SHA assertion differs from canonical registry")

    checkpoint = _resolve_checkpoint_provenance(
        root, request, support.train_recipe
    )
    if request.checkpoint_sha256 is not None and request.checkpoint_sha256 != checkpoint.sha256:
        raise ValueError("checkpoint SHA assertion differs from canonical checkpoint identity")
    if request.checkpoint_lock_id is not None and request.checkpoint_lock_id != checkpoint.lock_id:
        raise ValueError("checkpoint lock assertion differs from canonical checkpoint identity")
    if (
        request.checkpoint_selection is not None
        and request.checkpoint_selection != checkpoint.selection
    ):
        raise ValueError("checkpoint selection assertion differs from canonical selection")
    if (
        request.checkpoint_teacher_family is not None
        and request.checkpoint_teacher_family != checkpoint.teacher_family
    ):
        raise ValueError("checkpoint teacher family differs from canonical selection")
    resolved_adapter = adapter or load_adapter(model.adapter)
    resolved = RunRequest(
        **{
            **asdict(request),
            "config_sha256": config_sha256,
            "source_sha256": source_sha256,
            "adapter_identity": _adapter_identity(resolved_adapter, request),
            "environment": model.environment,
            "checkpoint_sha256": checkpoint.sha256,
            "checkpoint_lock_id": checkpoint.lock_id,
            "checkpoint_selection": checkpoint.selection,
            "checkpoint_teacher_family": checkpoint.teacher_family,
            **conditional,
        }
    )
    if resolved.checkpoint_sha256 is None or resolved.checkpoint_lock_id is None:
        raise ValueError("checkpoint SHA and checkpoint lock ID are required")
    return resolved, resolved_adapter


def _identity(request: RunRequest, command: list[str]) -> dict[str, object]:
    identity = {
        "run_id": request.run_id,
        "model_id": request.model_id,
        "dataset_id": request.dataset_id,
        "step_count": request.step_count,
        "seed": request.seed,
        "sample_count": request.sample_count,
        "device": request.device,
        "environment": request.environment,
        "adapter_identity": request.adapter_identity,
        "config_sha256": request.config_sha256,
        "source_sha256": request.source_sha256,
        "checkpoint_sha256": request.checkpoint_sha256,
        "checkpoint_lock_id": request.checkpoint_lock_id,
        "checkpoint_selection": request.checkpoint_selection,
        "checkpoint_teacher_family": request.checkpoint_teacher_family,
        "results_root": request.results_root,
        "command_sha256": _canonical_sha256(command),
    }
    if request.generation_mode == "conditional_prefix":
        identity.update(
            {
                "generation_mode": request.generation_mode,
                "conditioning_manifest": request.conditioning_manifest,
                "conditioning_manifest_sha256": request.conditioning_manifest_sha256,
                "conditioning_config_sha256": request.conditioning_config_sha256,
                "prefix_length": request.prefix_length,
                "evaluation_continuation_length": request.evaluation_continuation_length,
                "prompt_count": request.prompt_count,
                "diversity_prompt_count": request.diversity_prompt_count,
                "completions_per_diversity_prompt": request.completions_per_diversity_prompt,
                "completion_schedule": request.completion_schedule,
            }
        )
    return identity


def _write_failure(
    run_dir: Path,
    request: RunRequest,
    *,
    stage: str,
    message: str,
    returncode: int,
    command: list[str],
    stdout_path: Path,
    stderr_path: Path,
    started_at: str,
) -> None:
    atomic_json_write(
        run_dir / "failure.json",
        {
            "run_id": request.run_id,
            "stage": stage,
            "message": message,
            "exit_code": returncode if returncode >= 0 else None,
            "signal": -returncode if returncode < 0 else None,
            "command": command,
            "stdout_log": str(stdout_path),
            "stderr_log": str(stderr_path),
            "stdout_tail": _tail(stdout_path) if stdout_path.exists() else "",
            "stderr_tail": _tail(stderr_path) if stderr_path.exists() else "",
            "started_at": started_at,
            "finished_at": _utc_now(),
        },
    )


def run_experiment(
    request: RunRequest, root: Path | None = None, *, adapter: SampleAdapter | None = None
) -> RunResult:
    """Run, convert, validate, and atomically publish one baseline request."""

    root = (root or Path.cwd()).resolve()
    request, adapter = _resolve_request(request, root, adapter)
    results_root = (
        Path(request.results_root).absolute()
        if request.results_root is not None
        else root / "results"
    )
    if results_root.is_symlink() or not results_root.is_absolute():
        raise ValueError("results_root must be an absolute, non-symlink path")
    run_dir = (
        results_root
        / "samples"
        / request.dataset_id
        / request.model_id
        / f"steps_{request.step_count}"
    )
    ensure_safe_directory(run_dir)
    command = list(request.command) if request.command is not None else list(adapter.build_command(request, run_dir))
    if not command or any(not isinstance(argument, str) or not argument for argument in command):
        raise ValueError("adapter command must be a non-empty argv array of non-empty strings")
    identity = _identity(request, command)
    samples_path = run_dir / "samples.jsonl"
    metadata_path = run_dir / "run_metadata.json"
    metadata = _safe_json_load(metadata_path)
    if metadata is not None and metadata.get("status") == "succeeded" and metadata.get("identity") == identity:
        try:
            if request.generation_mode == "conditional_prefix":
                schedule, sequence_length, vocab_size = _conditional_io_contract(request, root)
                validate_conditional_samples(
                    samples_path,
                    expected=request.sample_count,
                    schedule=schedule,
                    sequence_length=sequence_length,
                    vocab_size=vocab_size,
                )
            else:
                validate_samples(samples_path, expected=request.sample_count)
        except SampleValidationError:
            pass
        else:
            remove_safe_file(run_dir / "failure.json")
            return RunResult("skipped", run_dir, 0)
    if metadata_path.exists():
        if metadata_path.is_symlink() or not metadata_path.is_file():
            raise ValueError(f"run metadata path is unsafe: {metadata_path}")
        metadata_path.unlink()

    stdout_path = run_dir / "stdout.log"
    stderr_path = run_dir / "stderr.log"
    started_at = _utc_now()
    atomic_json_write(
        run_dir / "request.json",
        {"request": asdict(request), "command": command, "identity": identity, "started_at": started_at},
    )
    try:
        with open_safe_output(stdout_path) as stdout_file, open_safe_output(stderr_path) as stderr_file:
            completed = subprocess.run(command, check=False, stdout=stdout_file, stderr=stderr_file)
    except OSError as error:
        _write_failure(
            run_dir, request, stage="command", message=str(error), returncode=1, command=command,
            stdout_path=stdout_path, stderr_path=stderr_path, started_at=started_at,
        )
        return RunResult("failed", run_dir, 1)
    if completed.returncode != 0:
        _write_failure(
            run_dir, request, stage="command", message="child command failed", returncode=completed.returncode,
            command=command, stdout_path=stdout_path, stderr_path=stderr_path, started_at=started_at,
        )
        return RunResult("failed", run_dir, completed.returncode)
    try:
        records = adapter.convert_outputs(request, run_dir)
        if request.generation_mode == "conditional_prefix":
            schedule, sequence_length, vocab_size = _conditional_io_contract(request, root)
            write_conditional_samples_atomic(
                samples_path,
                records,
                expected=request.sample_count,
                schedule=schedule,
                sequence_length=sequence_length,
                vocab_size=vocab_size,
            )
            validate_conditional_samples(
                samples_path,
                expected=request.sample_count,
                schedule=schedule,
                sequence_length=sequence_length,
                vocab_size=vocab_size,
            )
        else:
            write_samples_atomic(samples_path, records, expected=request.sample_count)
            validate_samples(samples_path, expected=request.sample_count)
    except Exception as error:
        _write_failure(
            run_dir, request, stage="conversion", message=str(error), returncode=1, command=command,
            stdout_path=stdout_path, stderr_path=stderr_path, started_at=started_at,
        )
        return RunResult("failed", run_dir, 1)
    atomic_json_write(
        metadata_path,
        {
            "status": "succeeded",
            "identity": identity,
            "command": command,
            "command_sha256": identity["command_sha256"],
            "config_sha256": request.config_sha256,
            "source_sha256": request.source_sha256,
            "checkpoint_sha256": request.checkpoint_sha256,
            "checkpoint_lock_id": request.checkpoint_lock_id,
            "checkpoint_selection": request.checkpoint_selection,
            "checkpoint_teacher_family": request.checkpoint_teacher_family,
            "environment": _environment_evidence(),
            "stdout_log": str(stdout_path),
            "stderr_log": str(stderr_path),
            "started_at": started_at,
            "finished_at": _utc_now(),
        },
    )
    remove_safe_file(run_dir / "failure.json")
    return RunResult("succeeded", run_dir, 0)


def _validate_registry_request(root: Path, model_id: str, dataset_id: str, step_count: int) -> str:
    registry_path = root / "configs" / "experiments.yaml"
    if not registry_path.is_file():
        raise ValueError("canonical registry is missing: configs/experiments.yaml")
    registry = load_registry(registry_path)
    try:
        model = registry.models[model_id]
        support = model.datasets[dataset_id]
    except KeyError as error:
        raise ValueError(f"unknown model/dataset cell: {model_id}/{dataset_id}") from error
    if support.status != "supported":
        raise ValueError(f"unsupported model/dataset cell: {model_id}/{dataset_id}")
    allowed_steps = step_grid_for_model(registry, model_id)
    if step_count not in allowed_steps:
        joined_steps = ",".join(str(step) for step in allowed_steps)
        raise ValueError(
            f"invalid step count {step_count} for {model_id}/{dataset_id}; "
            f"allowed: {joined_steps}"
        )
    return model.environment


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--model", required=True)
    parser.add_argument("--dataset", required=True)
    parser.add_argument("--steps", type=int, required=True)
    parser.add_argument("--num-samples", type=int)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--device")
    parser.add_argument("--results-root", type=Path)
    parser.add_argument(
        "--generation-mode",
        choices=("unconditional", "conditional_prefix"),
        default="unconditional",
    )
    parser.add_argument("--conditioning-manifest", type=Path)
    parser.add_argument("--conditioning-manifest-sha256")
    parser.add_argument("--conditioning-config", type=Path)
    parser.add_argument("--prefix-length", type=int)
    parser.add_argument("--evaluation-continuation-length", type=int)
    parser.add_argument("--prompt-count", type=int)
    parser.add_argument("--diversity-prompt-count", type=int)
    parser.add_argument("--completions-per-diversity-prompt", type=int)
    parser.add_argument("--completion-schedule")
    parser.add_argument("--validate-only", action="store_true")
    arguments = parser.parse_args(argv)
    if arguments.validate_only:
        try:
            print(_validate_registry_request(arguments.root, arguments.model, arguments.dataset, arguments.steps))
        except ValueError as error:
            parser.error(str(error))
        return 0
    request = RunRequest(
        run_id=f"{arguments.model}-{arguments.dataset}-steps-{arguments.steps}",
        model_id=arguments.model,
        dataset_id=arguments.dataset,
        step_count=arguments.steps,
        seed=arguments.seed,
        sample_count=(
            arguments.num_samples
            if arguments.num_samples is not None
            else (2048 if arguments.generation_mode == "conditional_prefix" else 1024)
        ),
        device=arguments.device,
        results_root=str(arguments.results_root.resolve()) if arguments.results_root else None,
        generation_mode=arguments.generation_mode,
        conditioning_manifest=(
            str(arguments.conditioning_manifest.resolve())
            if arguments.conditioning_manifest is not None
            else None
        ),
        conditioning_manifest_sha256=arguments.conditioning_manifest_sha256,
        conditioning_config_sha256=(
            sha256_file(arguments.conditioning_config.resolve())
            if arguments.conditioning_config is not None
            else None
        ),
        prefix_length=arguments.prefix_length,
        evaluation_continuation_length=arguments.evaluation_continuation_length,
        prompt_count=arguments.prompt_count,
        diversity_prompt_count=arguments.diversity_prompt_count,
        completions_per_diversity_prompt=arguments.completions_per_diversity_prompt,
        completion_schedule=arguments.completion_schedule,
    )
    try:
        result = run_experiment(request, arguments.root)
    except (ValueError, OSError) as error:
        parser.error(str(error))
    print(json.dumps({"status": result.status, "run_dir": str(result.run_dir)}))
    if result.returncode < 0:
        signum = -result.returncode
        if os.name == "posix":
            import signal

            sys.stdout.flush()
            sys.stderr.flush()
            if signum not in {signal.SIGKILL, signal.SIGSTOP}:
                signal.signal(signum, signal.SIG_DFL)
            os.kill(os.getpid(), signum)
        return 128 + signum
    return result.returncode


if __name__ == "__main__":
    raise SystemExit(main())