"""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())