| """Render and execute registry-bound, concrete sampler timing commands.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import re |
| import subprocess |
| from typing import Iterable, Literal |
| import uuid |
|
|
| from dlb.adapters.base import AdapterError, BaseTeacherAdapter |
| from dlb.checkpoints import load_checkpoint_manifest |
| from dlb.command import ADAPTERS |
| from dlb.io import atomic_json_write, sha256_file |
| from dlb.registry import load_registry, step_grid_for_model |
| from dlb.runner import RunRequest, _resolve_request |
| from dlb.timing import _validate_metadata |
|
|
|
|
| def _run_dir(root: Path, request: RunRequest) -> Path: |
| result_root = ( |
| root / "results/conditional" |
| if request.generation_mode == "conditional_prefix" |
| else root / "results" |
| ) |
| return ( |
| result_root |
| / "samples" |
| / request.dataset_id |
| / request.model_id |
| / f"steps_{request.step_count}" |
| ) |
|
|
|
|
| def _output_path(output_root: Path, request: RunRequest) -> Path: |
| return ( |
| output_root |
| / request.dataset_id |
| / request.model_id |
| / f"steps_{request.step_count}" |
| / "timing.json" |
| ) |
|
|
|
|
| def _request( |
| model: str, |
| dataset: str, |
| steps: int, |
| seed: int, |
| generation_mode: Literal["unconditional", "conditional_prefix"] = "unconditional", |
| ) -> RunRequest: |
| return RunRequest( |
| run_id=f"benchmark-{model}-{dataset}-steps-{steps}", |
| model_id=model, |
| dataset_id=dataset, |
| step_count=steps, |
| seed=seed, |
| sample_count=2048 if generation_mode == "conditional_prefix" else 1, |
| generation_mode=generation_mode, |
| ) |
|
|
|
|
| def render_benchmark_matrix( |
| *, |
| root: Path, |
| models: Iterable[str] | None, |
| datasets: Iterable[str], |
| steps: int, |
| seed: int, |
| precision: str, |
| output_root: Path, |
| dry_run: bool, |
| generation_mode: Literal["unconditional", "conditional_prefix"] = "unconditional", |
| ) -> list[dict[str, object]]: |
| """Render every selected cell, preserving explicit unsupported records.""" |
|
|
| root = root.resolve() |
| registry = load_registry(root / "configs/experiments.yaml") |
| selected_models = tuple(registry.models) if models is None else tuple(models) |
| records: list[dict[str, object]] = [] |
| for model_id in selected_models: |
| for dataset_id in datasets: |
| model = registry.models.get(model_id) |
| support = model.datasets.get(dataset_id) if model is not None else None |
| if model is None or support is None or support.status != "supported": |
| records.append( |
| { |
| "status": "unsupported", |
| "model": model_id, |
| "dataset": dataset_id, |
| "reason": getattr(support, "reason", None) or "unknown model/dataset cell", |
| } |
| ) |
| continue |
| allowed_steps = step_grid_for_model(registry, model_id) |
| if steps not in allowed_steps: |
| joined_steps = ",".join(str(step) for step in allowed_steps) |
| records.append( |
| { |
| "status": "error", |
| "model": model_id, |
| "dataset": dataset_id, |
| "reason": ( |
| f"invalid step count {steps} for {model_id}/{dataset_id}; " |
| f"allowed: {joined_steps}" |
| ), |
| } |
| ) |
| continue |
| adapter = ADAPTERS[model_id] |
| request = _request(model_id, dataset_id, steps, seed, generation_mode) |
| if generation_mode == "conditional_prefix": |
| try: |
| request, _ = _resolve_request(request, root, adapter) |
| except (OSError, ValueError) as error: |
| records.append( |
| { |
| "status": "error", |
| "model": model_id, |
| "dataset": dataset_id, |
| "reason": str(error), |
| } |
| ) |
| continue |
| output = _output_path(output_root.resolve(), request) |
| dry_attempt = "0" * 32 |
| staged_output = output.with_name( |
| f".{output.name}.{dry_attempt}.staged.json" |
| ) |
| metadata_path = output.with_name( |
| f".benchmark_metadata.{dry_attempt}.json" |
| ) |
| try: |
| command = adapter.render_benchmark_command( |
| request, |
| _run_dir(root, request), |
| output=staged_output, |
| metadata_path=metadata_path, |
| precision=precision, |
| dry_run=dry_run, |
| ) |
| except (AdapterError, OSError, ValueError) as error: |
| records.append( |
| { |
| "status": "error", |
| "model": model_id, |
| "dataset": dataset_id, |
| "reason": str(error), |
| } |
| ) |
| continue |
| records.append( |
| { |
| "status": "supported", |
| "model": model_id, |
| "dataset": dataset_id, |
| "steps": steps, |
| "environment": model.environment, |
| "hook": adapter.benchmark_hook(request), |
| **adapter.author_precision_policy(request), |
| "batch_size": 1, |
| "warmups": 5, |
| "repeats": 32, |
| "output": str(output), |
| "staged_output": str(staged_output), |
| "metadata_path": str(metadata_path), |
| "generation_mode": generation_mode, |
| "command": command, |
| } |
| ) |
| return records |
|
|
|
|
| def _precision_policy_binding( |
| root: Path, request: RunRequest, adapter: BaseTeacherAdapter |
| ) -> dict[str, object]: |
| """Bind an audited static policy to exact run, checkpoint and config bytes.""" |
|
|
| binding: dict[str, object] = { |
| "model": request.model_id, |
| "dataset": request.dataset_id, |
| "upstream": adapter.upstream, |
| "source_commit": request.source_sha256, |
| "experiment_config_sha256": request.config_sha256, |
| "checkpoint_sha256": request.checkpoint_sha256, |
| "checkpoint_lock_id": request.checkpoint_lock_id, |
| "checkpoint_selection": request.checkpoint_selection, |
| "adapter_identity": request.adapter_identity, |
| } |
| if request.generation_mode == "conditional_prefix": |
| binding.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, |
| } |
| ) |
| selection = request.checkpoint_selection |
| resource_id = selection.get("resource") if isinstance(selection, dict) else None |
| if isinstance(resource_id, str): |
| manifest = load_checkpoint_manifest(root / "artifacts/checkpoints.yaml") |
| resource = manifest.resources.get(resource_id) |
| if resource is None: |
| raise ValueError(f"precision policy references unknown checkpoint {resource_id}") |
| binding.update( |
| { |
| "checkpoint_resource": resource_id, |
| "checkpoint_repository": getattr(resource.source, "repo_id", None), |
| "checkpoint_revision": getattr(resource.source, "revision", None), |
| } |
| ) |
| if request.model_id in {"flm", "fmlm"}: |
| config_path = root / "checkpoints" / resource.destination / "config.json" |
| if config_path.is_symlink() or not config_path.is_file(): |
| raise ValueError( |
| "FLM/FMLM precision policy requires the downloaded checkpoint config.json" |
| ) |
| try: |
| config = json.loads(config_path.read_text(encoding="utf-8")) |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: |
| raise ValueError("FLM/FMLM checkpoint config.json is invalid") from error |
| if not isinstance(config, dict): |
| raise ValueError("FLM/FMLM checkpoint config.json must be a mapping") |
| dtype = config.get("torch_dtype") |
| if not isinstance(dtype, str) or not dtype: |
| raise ValueError( |
| "FLM/FMLM checkpoint config does not declare torch_dtype" |
| ) |
| binding.update( |
| { |
| "checkpoint_config_path": str(config_path.resolve()), |
| "checkpoint_config_sha256": sha256_file(config_path), |
| "checkpoint_config_torch_dtype": dtype, |
| "checkpoint_config_architectures": config.get("architectures"), |
| "checkpoint_config_auto_map": config.get("auto_map"), |
| } |
| ) |
| return binding |
|
|
|
|
| def _metadata( |
| root: Path, request: RunRequest, adapter: BaseTeacherAdapter, attempt_id: str |
| ) -> dict[str, object]: |
| binding = _precision_policy_binding(root, request, adapter) |
| policy: dict[str, object] = adapter.author_precision_policy(request) |
| if policy.get("precision") == "resolved-from-checkpoint-config-at-execution": |
| policy["precision"] = "bf16-mixed-static-author-policy" |
| return { |
| "seed": request.seed, |
| "sample_count": request.sample_count, |
| "dataset": request.dataset_id, |
| "model": request.model_id, |
| "steps": request.step_count, |
| "environment": request.environment, |
| "source_commit": request.source_sha256, |
| "config_sha256": request.config_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, |
| "adapter_identity": request.adapter_identity, |
| "requested_precision": "author", |
| "attempt_id": attempt_id, |
| "precision_policy_binding": binding, |
| **policy, |
| **( |
| { |
| "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, |
| } |
| if request.generation_mode == "conditional_prefix" |
| else {} |
| ), |
| } |
|
|
|
|
| def _load_staged_timing( |
| path: Path, expected_metadata: dict[str, object], attempt_id: str |
| ) -> None: |
| if path.is_symlink() or not path.is_file(): |
| raise RuntimeError("benchmark sampler completed without a fresh staged timing result") |
| try: |
| payload = json.loads(path.read_text(encoding="utf-8")) |
| timing = payload["timing"] |
| metadata = payload["metadata"] |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as error: |
| raise RuntimeError("fresh staged timing result is invalid") from error |
| if payload.get("schema") != "dlb-generation-timing-v1": |
| raise RuntimeError("fresh staged timing result has the wrong schema") |
| try: |
| _validate_metadata(metadata) |
| except (TypeError, ValueError) as error: |
| raise RuntimeError(f"fresh staged timing metadata is invalid: {error}") from error |
| if metadata.get("attempt_id") != attempt_id: |
| raise RuntimeError("fresh staged timing attempt provenance does not match") |
| for key, expected in expected_metadata.items(): |
| if metadata.get(key) != expected: |
| raise RuntimeError(f"fresh staged timing provenance differs at {key}") |
| if not isinstance(timing, dict) or any( |
| timing.get(key) != expected |
| for key, expected in { |
| "mode": "primary_latency", |
| "warmups": 5, |
| "repeats": 32, |
| "batch_size": 1, |
| "num_timed_samples": 32, |
| }.items() |
| ): |
| raise RuntimeError("fresh staged timing protocol is invalid") |
| raw = timing.get("raw_durations_seconds") |
| seconds = timing.get("seconds_per_sample") |
| if ( |
| not isinstance(raw, list) |
| or len(raw) != 32 |
| or any(type(value) not in {int, float} or not math.isfinite(value) or value < 0 for value in raw) |
| or type(seconds) not in {int, float} |
| or not math.isfinite(seconds) |
| or seconds < 0 |
| ): |
| raise RuntimeError("fresh staged timing durations are invalid") |
|
|
|
|
| def run_timing_attempt( |
| command: list[str], |
| *, |
| cwd: Path, |
| final_output: Path, |
| staged_output: Path, |
| expected_metadata: dict[str, object], |
| attempt_id: str, |
| ) -> Path: |
| """Publish a staged result only after the complete sampler process succeeds.""" |
|
|
| if re.fullmatch(r"[0-9a-f]{32}", attempt_id) is None: |
| raise ValueError("attempt_id must be 32 lowercase hexadecimal characters") |
| final_output = final_output.absolute() |
| staged_output = staged_output.absolute() |
| expected_stage = final_output.with_name( |
| f".{final_output.name}.{attempt_id}.staged.json" |
| ) |
| if staged_output != expected_stage: |
| raise ValueError("staged timing path does not match this attempt") |
| final_output.parent.mkdir(parents=True, exist_ok=True) |
| if staged_output.exists() or staged_output.is_symlink(): |
| raise RuntimeError("staged timing attempt path already exists") |
| superseded = final_output.with_name( |
| f".{final_output.name}.{attempt_id}.superseded.json" |
| ) |
| if superseded.exists() or superseded.is_symlink(): |
| raise RuntimeError("superseded timing path already exists") |
| if final_output.exists() or final_output.is_symlink(): |
| if final_output.is_symlink() or not final_output.is_file(): |
| raise RuntimeError("existing timing result is unsafe") |
| os.replace(final_output, superseded) |
| try: |
| completed = subprocess.run(command, cwd=cwd, check=False) |
| if completed.returncode: |
| raise RuntimeError( |
| f"benchmark sampler exited with status {completed.returncode}" |
| ) |
| _load_staged_timing(staged_output, expected_metadata, attempt_id) |
| os.replace(staged_output, final_output) |
| except BaseException: |
| staged_output.unlink(missing_ok=True) |
| raise |
| return final_output |
|
|
|
|
| def execute_one( |
| *, |
| root: Path, |
| model: str, |
| dataset: str, |
| steps: int, |
| seed: int, |
| precision: str, |
| generation_mode: Literal["unconditional", "conditional_prefix"] = "unconditional", |
| ) -> Path: |
| """Resolve provenance, then execute one real sampler wrapper in this environment.""" |
|
|
| root = root.resolve() |
| request = _request(model, dataset, steps, seed, generation_mode) |
| adapter: BaseTeacherAdapter = ADAPTERS[model] |
| resolved, adapter = _resolve_request(request, root, adapter) |
| timing_root = ( |
| root / "results/conditional/timing" |
| if generation_mode == "conditional_prefix" |
| else root / "results/timing" |
| ) |
| output = _output_path(timing_root, resolved) |
| attempt_id = uuid.uuid4().hex |
| staged_output = output.with_name(f".{output.name}.{attempt_id}.staged.json") |
| metadata_path = output.with_name(f".benchmark_metadata.{attempt_id}.json") |
| metadata = _metadata(root, resolved, adapter, attempt_id) |
| atomic_json_write(metadata_path, metadata) |
| try: |
| command = adapter.render_benchmark_command( |
| resolved, |
| _run_dir(root, resolved), |
| output=staged_output, |
| metadata_path=metadata_path, |
| precision=precision, |
| dry_run=False, |
| ) |
| return run_timing_attempt( |
| command, |
| cwd=root, |
| final_output=output, |
| staged_output=staged_output, |
| expected_metadata=metadata, |
| attempt_id=attempt_id, |
| ) |
| finally: |
| metadata_path.unlink(missing_ok=True) |
|
|
|
|
| def _csv(value: str) -> tuple[str, ...]: |
| items = tuple(value.split(",")) |
| if not items or any(not item for item in items): |
| raise argparse.ArgumentTypeError("comma-separated values must not be empty") |
| return items |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--root", type=Path, required=True) |
| parser.add_argument("--models", type=_csv, required=True) |
| parser.add_argument("--datasets", type=_csv, required=True) |
| parser.add_argument("--steps", type=int, required=True) |
| parser.add_argument("--seed", type=int, required=True) |
| parser.add_argument("--precision", choices=("author",), required=True) |
| parser.add_argument( |
| "--generation-mode", |
| choices=("unconditional", "conditional_prefix"), |
| default="unconditional", |
| ) |
| parser.add_argument("--dry-run", action="store_true") |
| arguments = parser.parse_args(argv) |
| if arguments.dry_run: |
| records = render_benchmark_matrix( |
| root=arguments.root, |
| models=arguments.models, |
| datasets=arguments.datasets, |
| steps=arguments.steps, |
| seed=arguments.seed, |
| precision=arguments.precision, |
| output_root=( |
| arguments.root / "results/conditional/timing" |
| if arguments.generation_mode == "conditional_prefix" |
| else arguments.root / "results/timing" |
| ), |
| dry_run=True, |
| generation_mode=arguments.generation_mode, |
| ) |
| for record in records: |
| print(json.dumps(record, sort_keys=True)) |
| return 1 if any(record["status"] == "error" for record in records) else 0 |
| if len(arguments.models) != 1 or len(arguments.datasets) != 1: |
| parser.error("execution requires exactly one model and one dataset") |
| try: |
| output = execute_one( |
| root=arguments.root, |
| model=arguments.models[0], |
| dataset=arguments.datasets[0], |
| steps=arguments.steps, |
| seed=arguments.seed, |
| precision=arguments.precision, |
| generation_mode=arguments.generation_mode, |
| ) |
| except (KeyError, OSError, RuntimeError, ValueError) as error: |
| parser.error(str(error)) |
| print(json.dumps({"status": "succeeded", "output": str(output)})) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
