"""Shared validation, checkpoint selection, and output conversion for teachers.""" from __future__ import annotations from dataclasses import dataclass import json import math from pathlib import Path import sys from typing import ClassVar, Iterable import yaml from dlb.checkpoints import load_checkpoint_manifest 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_checkpoint_provenance from dlb.schema import ConditionalSampleRecord, SampleRecord class AdapterError(ValueError): """A teacher request or upstream artifact cannot satisfy the canonical contract.""" @dataclass(frozen=True) class CheckpointSelection: path: Path teacher_family: str source: str required_files: tuple[str, ...] = () config_path: Path | None = None class BaseTeacherAdapter: """Strict adapter base used by the FLM-lineage and discrete teachers.""" identity: ClassVar[str] upstream: ClassVar[str] supported_models: ClassVar[frozenset[str]] teacher_families: ClassVar[dict[str, str]] batch_sizes: ClassVar[dict[tuple[str, str], int]] _DATA_CONFIGS: ClassVar[dict[str, str]] = { "lm1b": "lm1b-wrap", "owt": "openwebtext-split", } _TOKENIZER_VOCAB_SIZES: ClassVar[dict[str, int]] = { "bert-base-uncased": 30_522, "gpt2": 50_257, } _CHECKPOINT_SUFFIXES: ClassVar[frozenset[str]] = frozenset( {".ckpt", ".pt", ".pth", ".bin", ".safetensors"} ) def build_command(self, request: RunRequest, run_dir: Path) -> list[str]: """Build a real argv, rejecting absent canonical checkpoint bytes.""" return self.render_command(request, run_dir, dry_run=False) def benchmark_hook(self, request: RunRequest) -> str: """Name the concrete loaded-model sampling boundary used on the server.""" del request return { "mdlm": "mdlm._sample", "langflow": "langflow.generate_samples", "rdlm": "rdlm.sampling_fn", "sdtt": "distilled.model.sample", "di4c": "distilled.model.sample", }.get(self.upstream, "teacher.generate_samples") def author_precision_policy(self, request: RunRequest) -> dict[str, str]: """Return the audited inference policy of this pinned upstream sampler.""" policies = { "flm": "flm:checkpoint_config_and_loaded_runtime_code_bound", "duo": "duo:pinned_internal_bf16_autocast_with_fp32_sensitive_ops", "mdlm": "mdlm:pinned_internal_bf16_autocast_with_fp32_sensitive_ops", "candi": "candi:pinned_internal_bf16_autocast_with_fp32_sensitive_ops", "langflow": "langflow:pinned_internal_bf16_autocast_with_fp32_sensitive_ops", "rdlm": "rdlm:pinned_internal_bf16_autocast_with_fp32_sensitive_ops", "sdtt": "sdtt:pinned_bf16_mixed_config_and_internal_autocast", "di4c": "di4c:pinned_bf16_mixed_config_and_internal_autocast", } precision = ( "resolved-from-checkpoint-config-at-execution" if request.model_id in {"flm", "fmlm"} else "bf16-mixed" ) return { "precision": precision, "precision_policy": policies[self.upstream], "precision_evidence": ( "static_policy_bound_to_checkpoint_and_runtime_code_" "not_runtime_autocast_observation" ), } def render_benchmark_command( self, request: RunRequest, run_dir: Path, *, output: Path, metadata_path: Path, precision: str, dry_run: bool, ) -> list[str]: """Inject timing into this adapter's real loaded-model sampler command.""" if request.generation_mode == "conditional_prefix": if request.sample_count != 2048: raise AdapterError( "conditional latency benchmark must bind the full 2048-sample schedule" ) elif request.sample_count != 1: raise AdapterError("primary latency benchmark requires sample_count=1") if precision != "author": raise AdapterError("benchmark precision must use the pinned author policy") output = output.resolve() metadata_path = metadata_path.resolve() if output.is_symlink() or metadata_path.is_symlink(): raise AdapterError("benchmark output paths must not be symlinks") arguments = self.render_command(request, run_dir, dry_run=dry_run) marker = "-m" capture_command = marker in arguments and "dlb.adapters.capture" in arguments if capture_command: separator = arguments.index("--") arguments[separator:separator] = [ f"--benchmark-output={output}", f"--benchmark-metadata={metadata_path}", f"--benchmark-precision={precision}", ] replacements = { "loader.eval_batch_size": "1", "sampling.num_sample_batches": "1", "sampling.batch_per_gpu": "1", "--expected-samples": "1", } for key, value in replacements.items(): prefix = key + "=" for index, argument in enumerate(arguments): if argument.startswith(prefix): arguments[index] = prefix + value for key in ("--num_samples", "--batch_size"): if key in arguments: arguments[arguments.index(key) + 1] = "1" else: arguments.extend( [ "--benchmark-output", str(output), "--benchmark-metadata", str(metadata_path), "--benchmark-precision", precision, ] ) for key in ("--sample-count", "--batch-size"): if key in arguments: arguments[arguments.index(key) + 1] = "1" self._validate_argv(arguments) return arguments def _conditional_capture_flags(self, request: RunRequest) -> list[str]: """Return capture-wrapper key=value flags for conditional requests.""" if request.generation_mode == "unconditional": return [] required = { "--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, } missing = [name for name, value in required.items() if value is None] if missing: raise AdapterError( "conditional request is missing adapter fields: " + ", ".join(missing) ) return [f"{name}={value}" for name, value in required.items()] def _conditional_script_flags(self, request: RunRequest) -> list[str]: """Return argparse-style flag/value pairs for conditional script wrappers.""" capture_flags = self._conditional_capture_flags(request) arguments: list[str] = [] for flag in capture_flags: name, _, value = flag.partition("=") arguments.extend([name, value]) return arguments def render_command( self, request: RunRequest, run_dir: Path, *, dry_run: bool ) -> list[str]: root, length, batch_size = self._validate_request(request, run_dir) checkpoint = self._resolve_checkpoint(root, request, dry_run=dry_run) entrypoint = root / "upstreams" / self.upstream / "main.py" if not entrypoint.is_file() or entrypoint.is_symlink(): raise AdapterError(f"pinned upstream entrypoint is missing or unsafe: {entrypoint}") generated_count = math.ceil(request.sample_count / batch_size) * batch_size batches = generated_count // batch_size capture_path = run_dir.resolve() / "upstream_token_ids.json" upstream_output = run_dir.resolve() / "upstream_samples.json" hydra_dir = run_dir.resolve() / "upstream_hydra" arguments = [ sys.executable, "-u", "-m", "dlb.adapters.capture", f"--upstream-entrypoint={entrypoint}", f"--capture-path={capture_path}", *self._conditional_capture_flags(request), "--", "mode=sample_eval", f"seed={request.seed}", f"data={self._data_config(request)}", "model=small", f"model.length={length}", f"loader.batch_size={self._loader_batch_size(request)}", f"loader.eval_batch_size={batch_size}", f"sampling.num_sample_batches={batches}", f"sampling.steps={request.step_count}", f"eval.checkpoint_path={checkpoint.path}", "trainer.devices=1", f"hydra.run.dir={hydra_dir}", f"checkpointing.save_dir={hydra_dir}", "+wandb.offline=true", ] if self.upstream != "mdlm": arguments.append(f"eval.generated_samples_path={upstream_output}") arguments.extend(self._sampling_overrides(request, checkpoint)) self._validate_argv(arguments) return arguments def convert_outputs( self, request: RunRequest, run_dir: Path ) -> Iterable[SampleRecord | ConditionalSampleRecord]: if request.generation_mode == "conditional_prefix": return self._convert_conditional_capture_outputs(request, run_dir) root, sequence_length, batch_size = self._validate_conversion_request( request, run_dir ) self._require_conversion_provenance(root, request) expected_generated = math.ceil(request.sample_count / batch_size) * batch_size capture_path = run_dir.resolve() / "upstream_token_ids.json" capture = self._read_capture(capture_path, expected_generated) if self.upstream == "mdlm": texts = [item["text"] for item in capture] token_ids = [item["token_ids"] for item in capture] output_format = "dlb-upstream-token-capture-v1" token_ids_source = "upstream" token_ids_transformation = { "max_length": sequence_length, "operation": "validated_exact_length", } upstream_path = capture_path else: upstream_path = run_dir.resolve() / "upstream_samples.json" texts = self._read_standard_output(upstream_path, expected_generated) self._validate_texts(texts) output_format = "upstream-generated-seqs-v1" if capture: captured_texts = [item["text"] for item in capture] if captured_texts != texts: raise AdapterError("captured token batches do not match upstream generated_seqs") token_ids = [item["token_ids"] for item in capture] token_ids_source = "upstream" token_ids_transformation = { "max_length": sequence_length, "operation": "validated_exact_length", } else: token_ids, token_ids_transformation = self._retokenize( root, request.dataset_id, texts, sequence_length ) token_ids_source = "retokenized" dataset = self._load_data_contract(root, request.dataset_id) tokenizer_name = dataset["tokenizer"] vocab_size = self._TOKENIZER_VOCAB_SIZES[tokenizer_name] self._validate_samples( texts, token_ids, expected_generated, vocab_size, sequence_length ) records = [ SampleRecord( sample_id=index, text=texts[index], token_ids=token_ids[index], seed=request.seed, generation_seconds=0.0, ) for index in range(request.sample_count) ] atomic_json_write( run_dir.resolve() / "conversion_metadata.json", { "format": output_format, "upstream_output": str(upstream_path), "generated_samples": expected_generated, "requested_samples": request.sample_count, "trimmed_samples": expected_generated - request.sample_count, "trim_policy": "stable_prefix", "token_ids_source": token_ids_source, "token_ids_transformation": token_ids_transformation, "tokenizer": tokenizer_name, "tokenizer_revision": self._tokenizer_revision(root, tokenizer_name), "checkpoint_sha256": request.checkpoint_sha256, "checkpoint_lock_id": request.checkpoint_lock_id, "checkpoint_selection": request.checkpoint_selection, "teacher_family": request.checkpoint_teacher_family, "generation_seconds_source": "unavailable_excluded_sentinel", "generation_seconds_sentinel": 0.0, "exclude_from_latency": True, }, ) return records def _convert_capture_outputs( self, request: RunRequest, run_dir: Path, *, retokenize_inexact_rows: bool = False, ) -> Iterable[SampleRecord | ConditionalSampleRecord]: """Convert an exact project capture produced around a pinned sampler.""" if request.generation_mode == "conditional_prefix": return self._convert_conditional_capture_outputs( request, run_dir, retokenize_inexact_rows=retokenize_inexact_rows, ) root, sequence_length, _ = self._validate_conversion_request(request, run_dir) self._require_conversion_provenance(root, request) capture_path = run_dir.resolve() / "upstream_token_ids.json" capture = self._read_capture(capture_path, request.sample_count) if not capture: raise AdapterError(f"token capture is missing or empty: {capture_path}") texts = [item["text"] for item in capture] upstream_token_ids = [item["token_ids"] for item in capture] token_ids = upstream_token_ids token_ids_source = "upstream" dataset = self._load_data_contract(root, request.dataset_id) tokenizer_name = dataset["tokenizer"] vocab_size = self._TOKENIZER_VOCAB_SIZES[tokenizer_name] token_ids_transformation: dict[str, object] = { "max_length": sequence_length, "operation": "validated_exact_length", } retokenize_reason = ( self._retokenize_reason(upstream_token_ids, sequence_length, vocab_size) if retokenize_inexact_rows else None ) if retokenize_reason is not None: token_ids, token_ids_transformation = self._retokenize( root, request.dataset_id, texts, sequence_length ) token_ids_transformation = { **token_ids_transformation, "reason": retokenize_reason, } token_ids_source = "retokenized" self._validate_samples( texts, token_ids, request.sample_count, vocab_size, sequence_length, ) records = [ SampleRecord( sample_id=index, text=texts[index], token_ids=token_ids[index], seed=request.seed, generation_seconds=0.0, ) for index in range(request.sample_count) ] atomic_json_write( run_dir.resolve() / "conversion_metadata.json", { "format": "dlb-upstream-token-capture-v1", "upstream_output": str(capture_path), "generated_samples": request.sample_count, "requested_samples": request.sample_count, "trimmed_samples": 0, "trim_policy": "none_exact_count", "token_ids_source": token_ids_source, "token_ids_transformation": token_ids_transformation, "tokenizer": tokenizer_name, "tokenizer_revision": self._tokenizer_revision(root, tokenizer_name), "checkpoint_sha256": request.checkpoint_sha256, "checkpoint_lock_id": request.checkpoint_lock_id, "checkpoint_selection": request.checkpoint_selection, "teacher_family": request.checkpoint_teacher_family, "generation_seconds_source": "unavailable_excluded_sentinel", "generation_seconds_sentinel": 0.0, "exclude_from_latency": True, }, ) return records def _convert_conditional_capture_outputs( self, request: RunRequest, run_dir: Path, *, retokenize_inexact_rows: bool = False, ) -> Iterable[ConditionalSampleRecord]: """Convert a prompt-enriched capture into conditional sample records.""" root, sequence_length, _ = self._validate_conversion_request(request, run_dir) self._require_conversion_provenance(root, request) capture_path = run_dir.resolve() / "upstream_token_ids.json" capture = self._read_conditional_capture(capture_path, request.sample_count) dataset = self._load_data_contract(root, request.dataset_id) tokenizer_name = dataset["tokenizer"] vocab_size = self._TOKENIZER_VOCAB_SIZES[tokenizer_name] prefix_lengths: set[int] = set() retokenize_reason = None for index, item in enumerate(capture): full = item["full_token_ids"] prefix = item["prefix_token_ids"] reference = item["reference_token_ids"] if not isinstance(full, list) or not isinstance(prefix, list) or not isinstance(reference, list): raise AdapterError(f"conditional capture record {index} has invalid token fields") prefix_lengths.add(len(prefix)) if len(prefix) >= sequence_length: raise AdapterError(f"conditional capture record {index} has invalid prefix length") if len(full) >= len(prefix) and full[: len(prefix)] != prefix: raise AdapterError(f"conditional capture record {index} has prefix mismatch") if any( type(token) is not int or token < 0 or token >= vocab_size for token in [*prefix, *reference] ): raise AdapterError(f"conditional capture record {index} has token outside vocabulary") reason = ( self._retokenize_reason([full], sequence_length, vocab_size) if retokenize_inexact_rows else None ) if reason is not None and retokenize_reason is None: retokenize_reason = reason if len(prefix_lengths) != 1: raise AdapterError("conditional capture records have inconsistent prefix lengths") prefix_length = next(iter(prefix_lengths), 0) continuation_length = sequence_length - prefix_length retokenized_continuations: list[list[int]] | None = None token_ids_source = "upstream_conditional_capture" token_ids_transformation: dict[str, object] = { "max_length": sequence_length, "operation": "validated_exact_length", } if retokenize_reason is not None: texts: list[str] = [] for index, item in enumerate(capture): text = item.get("continuation_text") if not isinstance(text, str) or not text.strip(): raise AdapterError( f"conditional capture record {index} has empty continuation text" ) texts.append(text) retokenized_continuations, suffix_transformation = self._retokenize( root, request.dataset_id, texts, continuation_length, ) token_ids_source = "retokenized_conditional_suffix" token_ids_transformation = { **suffix_transformation, "operation": "retokenized_continuation_only", "reason": retokenize_reason, "preserved_prefix_tokens": prefix_length, } records: list[ConditionalSampleRecord] = [] for index, item in enumerate(capture): full = item["full_token_ids"] prefix = item["prefix_token_ids"] reference = item["reference_token_ids"] if not isinstance(full, list) or not isinstance(prefix, list) or not isinstance(reference, list): raise AdapterError(f"conditional capture record {index} has invalid token fields") if retokenized_continuations is None: if len(full) != sequence_length: raise AdapterError( f"conditional capture record {index} has {len(full)} full tokens, " f"expected {sequence_length}" ) if any( type(token) is not int or token < 0 or token >= vocab_size for token in full ): raise AdapterError(f"conditional capture record {index} has token outside vocabulary") continuation = full[len(prefix) :] else: continuation = retokenized_continuations[index] if any( type(token) is not int or token < 0 or token >= vocab_size for token in continuation ): raise AdapterError( f"conditional capture record {index} retokenized outside vocabulary" ) full = prefix + continuation continuation = full[len(prefix) :] records.append( ConditionalSampleRecord( sample_id=index, prompt_id=item["prompt_id"], completion_id=item["completion_id"], source_index=item["source_index"], prefix_token_ids=prefix, continuation_token_ids=continuation, reference_token_ids=reference, full_token_ids=full, prefix_text=str(item.get("prefix_text") or " ".join(map(str, prefix))), continuation_text=str(item.get("continuation_text") or " ".join(map(str, continuation))), reference_text=str(item.get("reference_text") or " ".join(map(str, reference))), full_text=str(item.get("full_text") or item.get("text") or " ".join(map(str, full))), seed=request.seed, generation_seconds=0.0, prefix_exact_match=True, ) ) atomic_json_write( run_dir.resolve() / "conversion_metadata.json", { "format": "dlb-conditional-token-capture-v1", "upstream_output": str(capture_path), "generated_samples": len(records), "requested_samples": request.sample_count, "trimmed_samples": 0, "trim_policy": "none_exact_conditional_schedule", "token_ids_source": token_ids_source, "token_ids_transformation": token_ids_transformation, "tokenizer": tokenizer_name, "tokenizer_revision": self._tokenizer_revision(root, tokenizer_name), "checkpoint_sha256": request.checkpoint_sha256, "checkpoint_lock_id": request.checkpoint_lock_id, "checkpoint_selection": request.checkpoint_selection, "teacher_family": request.checkpoint_teacher_family, "conditioning_implementation": ( "native_projection" if request.model_id in {"candi", "rdlm", "mdlm_sdtt", "mdlm_di4c", "duo_di4c"} else "zero_shot_runtime_projection" ), "paper_conditional_training_reproduced": False, "generation_seconds_source": "unavailable_excluded_sentinel", "generation_seconds_sentinel": 0.0, "exclude_from_latency": True, }, ) return records def _require_conversion_provenance(self, root: Path, request: RunRequest) -> None: if ( not isinstance(request.checkpoint_sha256, str) or len(request.checkpoint_sha256) != 64 or any( character not in "0123456789abcdef" for character in request.checkpoint_sha256 ) or not isinstance(request.checkpoint_lock_id, str) or not request.checkpoint_lock_id or not isinstance(request.checkpoint_selection, dict) or not request.checkpoint_selection or not isinstance(request.checkpoint_teacher_family, str) or not request.checkpoint_teacher_family ): raise AdapterError("runner-resolved checkpoint provenance is required") registry = load_registry(root / "configs" / "experiments.yaml") try: support = registry.models[request.model_id].datasets[request.dataset_id] canonical = _resolve_checkpoint_provenance( root, request, support.train_recipe ) except (KeyError, OSError, ValueError) as error: raise AdapterError(f"canonical checkpoint provenance is invalid: {error}") from error expected_family = self.teacher_families[request.model_id] if canonical.teacher_family != expected_family: raise AdapterError( f"canonical checkpoint teacher family {canonical.teacher_family!r} " f"does not match {expected_family!r}" ) if request.checkpoint_sha256 != canonical.sha256: raise AdapterError("checkpoint SHA differs from canonical checkpoint provenance") if request.checkpoint_lock_id != canonical.lock_id: raise AdapterError("checkpoint lock ID differs from canonical checkpoint provenance") if request.checkpoint_selection != canonical.selection: raise AdapterError("checkpoint selection differs from canonical checkpoint provenance") if request.checkpoint_teacher_family != canonical.teacher_family: raise AdapterError( "checkpoint teacher family differs from canonical checkpoint provenance" ) def _runtime_asset_digests( self, root: Path, request: RunRequest, checkpoint_path: Path, config_path: Path, *, dry_run: bool, ) -> tuple[str, str]: """Authenticate selected bytes before a server wrapper may deserialize them.""" if dry_run: return "dry-run-unverified", "dry-run-unverified" self._require_conversion_provenance(root, request) checkpoint_digest = sha256_file(checkpoint_path) config_digest = sha256_file(config_path) selection = request.checkpoint_selection or {} resource_id = selection.get("resource") if isinstance(resource_id, str): lock_path = root / "artifacts" / "checkpoint_lock.json" try: lock = json.loads(lock_path.read_text(encoding="utf-8")) record = lock["resources"][resource_id] files = record["files"] except (OSError, UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as error: raise AdapterError("canonical checkpoint lock is invalid") from error if not isinstance(files, list): raise AdapterError("canonical checkpoint lock file inventory is invalid") expected = { str(item.get("path")): item.get("sha256") for item in files if isinstance(item, dict) } for path, digest, label in ( (checkpoint_path, checkpoint_digest, "student checkpoint"), (config_path, config_digest, "sampling config"), ): relative = path.relative_to(root).as_posix() if selection.get("sampling_config_source") == "project" and path == config_path: locked_digest = selection.get("sampling_config_sha256") else: locked_digest = expected.get(relative) if locked_digest != digest: raise AdapterError( f"{label} SHA differs from canonical checkpoint provenance: {path}" ) return checkpoint_digest, config_digest def _validate_request( self, request: RunRequest, run_dir: Path ) -> tuple[Path, int, int]: root, length, batch_size = self._validate_conversion_request(request, run_dir) registry_path = root / "configs" / "experiments.yaml" if not registry_path.is_file(): raise AdapterError(f"canonical registry is missing: {registry_path}") registry = load_registry(registry_path) try: model = registry.models[request.model_id] support = model.datasets[request.dataset_id] except KeyError as error: raise AdapterError( f"unknown model/dataset cell: {request.model_id}/{request.dataset_id}" ) from error if support.status != "supported": raise AdapterError( 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 AdapterError( f"invalid step count {request.step_count} for " f"{request.model_id}/{request.dataset_id}; allowed: {joined_steps}" ) if model.adapter != self.upstream: raise AdapterError( f"registry adapter {model.adapter!r} does not match {self.upstream!r}" ) return root, length, batch_size def _validate_conversion_request( self, request: RunRequest, run_dir: Path ) -> tuple[Path, int, int]: if request.model_id not in self.supported_models: raise AdapterError( f"{type(self).__name__} does not support model {request.model_id!r}" ) if request.dataset_id not in self._DATA_CONFIGS: raise AdapterError(f"unsupported dataset: {request.dataset_id!r}") for name, value, positive in ( ("step_count", request.step_count, True), ("sample_count", request.sample_count, True), ("seed", request.seed, False), ): if type(value) is not int or (positive and value <= 0): qualifier = "a positive integer" if positive else "an integer" raise AdapterError(f"{name} must be {qualifier}") run_dir = run_dir.resolve() sample_suffix = ( Path("samples") / request.dataset_id / request.model_id / f"steps_{request.step_count}" ) if request.results_root is not None: results_root = Path(request.results_root).resolve() if results_root.is_symlink(): raise AdapterError(f"results root is unsafe: {results_root}") root = self._project_root_for_results_root(results_root) if results_root / sample_suffix != run_dir: raise AdapterError(f"run directory does not match the request: {run_dir}") else: suffix = Path("results") / sample_suffix if len(run_dir.parents) < 5: raise AdapterError(f"run directory is not rooted at a request root: {run_dir}") root = run_dir.parents[4] if root / suffix != run_dir: raise AdapterError(f"run directory does not match the request: {run_dir}") dataset = self._load_data_contract(root, request.dataset_id) length = dataset.get("sequence_length") if type(length) is not int or length <= 0: raise AdapterError(f"invalid canonical sequence length for {request.dataset_id}") batch_size = self.batch_sizes.get((request.model_id, request.dataset_id)) if type(batch_size) is not int or batch_size <= 0: raise AdapterError( f"missing eval batch size for {request.model_id}/{request.dataset_id}" ) return root, length, batch_size def _project_root_for_results_root(self, results_root: Path) -> Path: for candidate in results_root.parents: if ( (candidate / "configs" / "experiments.yaml").is_file() and (candidate / "artifacts" / "data.yaml").is_file() ): return candidate.resolve() raise AdapterError( f"results root is not under a canonical project root: {results_root}" ) def _resolve_checkpoint( self, root: Path, request: RunRequest, *, dry_run: bool ) -> CheckpointSelection: manifest_path = root / "artifacts" / "checkpoints.yaml" if not manifest_path.is_file(): raise AdapterError(f"canonical checkpoint manifest is missing: {manifest_path}") manifest = load_checkpoint_manifest(manifest_path) registry = load_registry(root / "configs" / "experiments.yaml") support = registry.models[request.model_id].datasets[request.dataset_id] expected_family = self.teacher_families[request.model_id] coverage = manifest.coverage.get((request.model_id, request.dataset_id)) if coverage is not None: resource = manifest.resources[coverage.resource] family = coverage.teacher_family or resource.teacher_family if family != expected_family: raise AdapterError( f"checkpoint teacher family {family!r} does not match {expected_family!r}" ) base = root / "checkpoints" / resource.destination path = base / coverage.path if coverage.path else base config_path: Path | None = None config_sha256: str | None = None if coverage.sampling_config is not None: if coverage.sampling_config_source == "resource": config_path = (base / coverage.sampling_config).absolute() else: config_path = (root / coverage.sampling_config).absolute() config_sha256 = coverage.sampling_config_sha256 if ( config_path.is_symlink() or not config_path.is_file() or sha256_file(config_path) != config_sha256 ): raise AdapterError( f"project sampling config is missing or differs from manifest: {config_path}" ) selection = CheckpointSelection( path=path.resolve(), teacher_family=family, source=f"resource:{coverage.resource}", required_files=tuple(resource.required_files), config_path=config_path, ) if not dry_run: self._require_resource_checkpoint(selection, base, coverage.path) return selection recipe_id = support.train_recipe if recipe_id is None or recipe_id not in manifest.recipes: raise AdapterError( f"canonical checkpoint selection is missing for {request.model_id}/{request.dataset_id}" ) recipe = manifest.recipes[recipe_id] if (recipe.model, recipe.dataset) != (request.model_id, request.dataset_id): raise AdapterError(f"checkpoint recipe {recipe_id} describes the wrong cell") if recipe.teacher_family != expected_family: raise AdapterError( f"checkpoint teacher family {recipe.teacher_family!r} does not match {expected_family!r}" ) output = (root / recipe.output).absolute() if recipe.sampling_checkpoint is not None: path = (output / recipe.sampling_checkpoint).absolute() if not dry_run: relative = Path(recipe.sampling_checkpoint) ancestors = [output] current = output for part in relative.parts[:-1]: current = current / part ancestors.append(current) if ( any(ancestor.is_symlink() for ancestor in ancestors) or path.is_symlink() or not path.is_file() or path.stat().st_size <= 0 ): raise AdapterError( f"recipe sampling checkpoint is missing or unsafe: {path}" ) else: path = output if dry_run else self._select_recipe_checkpoint(output) config_path = ( (output / recipe.sampling_config).absolute() if recipe.sampling_config is not None else None ) if not dry_run and config_path is not None: relative = Path(recipe.sampling_config) ancestors = [output] current = output for part in relative.parts[:-1]: current = current / part ancestors.append(current) if ( any(ancestor.is_symlink() for ancestor in ancestors) or config_path.is_symlink() or not config_path.is_file() or config_path.stat().st_size <= 0 ): raise AdapterError( f"recipe sampling config is missing or unsafe: {config_path}" ) return CheckpointSelection( path=path, teacher_family=recipe.teacher_family, source=f"recipe:{recipe_id}", config_path=config_path, ) def _require_resource_checkpoint( self, selection: CheckpointSelection, base: Path, selected_name: str | None ) -> None: path = selection.path if selected_name is not None: if path.is_symlink() or not path.is_file() or path.stat().st_size <= 0: raise AdapterError(f"required checkpoint file is missing or unsafe: {path}") return if path.is_symlink() or not path.is_dir(): raise AdapterError(f"required checkpoint directory is missing or unsafe: {path}") for relative in selection.required_files: required = base / relative if required.is_symlink() or not required.is_file() or required.stat().st_size <= 0: raise AdapterError(f"required checkpoint file is missing or unsafe: {required}") def _select_recipe_checkpoint(self, output: Path) -> Path: if output.is_symlink() or not output.exists(): raise AdapterError(f"recipe checkpoint output is missing or unsafe: {output}") if output.is_file(): if output.stat().st_size <= 0: raise AdapterError(f"recipe checkpoint output is empty: {output}") return output candidates = sorted( path for path in output.rglob("*") if path.is_file() and not path.is_symlink() and path.suffix.lower() in self._CHECKPOINT_SUFFIXES and path.stat().st_size > 0 ) preferred = [path for path in candidates if path.name == "last.ckpt"] if len(preferred) == 1: return preferred[0] if len(candidates) != 1: raise AdapterError( f"recipe output must contain exactly one selectable checkpoint, found {len(candidates)}" ) return candidates[0] def _read_standard_output(self, path: Path, expected: int) -> list[str]: value = self._read_json(path, "upstream sample output") if not isinstance(value, dict) or set(value) != { "generative_ppl", "entropy", "generated_seqs", }: raise AdapterError("unexpected upstream output format") texts = value["generated_seqs"] if not isinstance(texts, list): raise AdapterError("unexpected upstream generated_seqs format") if len(texts) != expected: raise AdapterError(f"expected {expected} generated samples, found {len(texts)}") return texts def _read_capture(self, path: Path, expected: int) -> list[dict[str, object]]: if not path.exists() and self.upstream != "mdlm": return [] value = self._read_json(path, "token capture") if ( not isinstance(value, dict) or set(value) != {"schema", "samples"} or value.get("schema") != "dlb-upstream-token-capture-v1" or not isinstance(value.get("samples"), list) ): raise AdapterError("unexpected capture format") samples = value["samples"] if len(samples) != expected: raise AdapterError(f"expected {expected} generated samples, found {len(samples)}") seen: set[int] = set() normalized: list[dict[str, object]] = [] for index, item in enumerate(samples): if not isinstance(item, dict) or set(item) != {"sample_id", "text", "token_ids"}: raise AdapterError(f"unexpected capture record format at index {index}") sample_id = item["sample_id"] if type(sample_id) is not int: raise AdapterError(f"invalid sample_id at index {index}") if sample_id in seen: raise AdapterError(f"duplicate sample_id {sample_id}") seen.add(sample_id) if sample_id != index: raise AdapterError(f"expected sample_id {index}, found {sample_id}") normalized.append(item) return normalized def _read_conditional_capture( self, path: Path, expected: int ) -> list[dict[str, object]]: value = self._read_json(path, "conditional token capture") required = { "sample_id", "text", "token_ids", "prompt_id", "completion_id", "source_index", "prefix_token_ids", "reference_token_ids", "full_token_ids", "prefix_exact_match", } optional = { "prefix_text", "continuation_text", "reference_text", "full_text", } if ( not isinstance(value, dict) or set(value) != {"schema", "samples"} or value.get("schema") != "dlb-upstream-token-capture-v1" or not isinstance(value.get("samples"), list) ): raise AdapterError("unexpected conditional capture format") samples = value["samples"] if len(samples) != expected: raise AdapterError( f"expected {expected} conditional samples, found {len(samples)}" ) seen: set[int] = set() normalized: list[dict[str, object]] = [] for index, item in enumerate(samples): if not isinstance(item, dict): raise AdapterError(f"unexpected conditional capture record at index {index}") keys = set(item) if not required.issubset(keys) or keys - required - optional: raise AdapterError( f"unexpected conditional capture record format at index {index}" ) sample_id = item["sample_id"] if type(sample_id) is not int: raise AdapterError(f"invalid sample_id at index {index}") if sample_id in seen: raise AdapterError(f"duplicate sample_id {sample_id}") seen.add(sample_id) if sample_id != index: raise AdapterError(f"expected sample_id {index}, found {sample_id}") if item["prefix_exact_match"] is not True: raise AdapterError(f"conditional prefix mismatch at sample {index}") for field in ("prompt_id", "completion_id", "source_index"): if type(item[field]) is not int or item[field] < 0: raise AdapterError( f"conditional capture record {index} has invalid {field}" ) normalized.append(item) return normalized def _validate_samples( self, texts: list[object], token_ids: list[object], expected: int, vocab_size: int, sequence_length: int, ) -> None: if len(texts) != expected or len(token_ids) != expected: raise AdapterError( f"expected {expected} generated samples, found {min(len(texts), len(token_ids))}" ) self._validate_texts(texts) for index, (_, tokens) in enumerate(zip(texts, token_ids)): if not isinstance(tokens, list) or not tokens: raise AdapterError(f"empty token_ids at sample {index}") if len(tokens) != sequence_length: raise AdapterError( f"expected {sequence_length} tokens at sample {index}, found {len(tokens)}" ) for token in tokens: if type(token) is not int or token < 0 or token >= vocab_size: raise AdapterError(f"invalid token at sample {index}: {token!r}") @staticmethod def _retokenize_reason( token_ids: list[object], sequence_length: int, vocab_size: int ) -> str | None: for tokens in token_ids: if not isinstance(tokens, list) or len(tokens) != sequence_length: return "upstream_ids_not_canonical_length" for token in tokens: if type(token) is not int or token < 0 or token >= vocab_size: return "upstream_ids_outside_canonical_vocab" return None @staticmethod def _validate_texts(texts: list[object]) -> None: for index, text in enumerate(texts): if not isinstance(text, str) or not text.strip(): raise AdapterError(f"empty text at sample {index}") def _retokenize( self, root: Path, dataset_id: str, texts: list[str], sequence_length: int ) -> tuple[list[list[int]], dict[str, object]]: contract = self._load_data_contract(root, dataset_id) tokenizer_name = contract["tokenizer"] revision = self._tokenizer_revision(root, tokenizer_name) try: from transformers import AutoTokenizer except ImportError as error: raise AdapterError( "transformers is required to retokenize upstream text without captured IDs" ) from error try: tokenizer = AutoTokenizer.from_pretrained( tokenizer_name, revision=revision, local_files_only=True ) except Exception as error: raise AdapterError( f"pinned tokenizer {tokenizer_name}@{revision} is not available locally" ) from error if tokenizer.pad_token_id is None: if tokenizer.eos_token is None: raise AdapterError(f"canonical tokenizer {tokenizer_name} has no padding token") tokenizer.pad_token = tokenizer.eos_token settings = { "add_special_tokens": False, "padding": "max_length", "truncation": True, "max_length": sequence_length, "return_attention_mask": False, } encoded = tokenizer(texts, **settings) try: token_ids = [list(tokens) for tokens in encoded["input_ids"]] except (KeyError, TypeError) as error: raise AdapterError("pinned tokenizer returned invalid input_ids") from error return token_ids, { "add_special_tokens": False, "max_length": sequence_length, "padding": "max_length", "pad_token_id": tokenizer.pad_token_id, "truncation": True, } def _load_data_contract(self, root: Path, dataset_id: str) -> dict[str, object]: path = root / "artifacts" / "data.yaml" if not path.is_file(): raise AdapterError(f"canonical data manifest is missing: {path}") document = yaml.safe_load(path.read_text(encoding="utf-8")) try: dataset = document["datasets"][dataset_id] except (KeyError, TypeError) as error: raise AdapterError(f"canonical data contract is missing for {dataset_id}") from error if not isinstance(dataset, dict): raise AdapterError(f"invalid canonical data contract for {dataset_id}") tokenizer = dataset.get("tokenizer") if tokenizer not in self._TOKENIZER_VOCAB_SIZES: raise AdapterError(f"unsupported canonical tokenizer: {tokenizer!r}") return dataset def _tokenizer_revision(self, root: Path, tokenizer_name: str) -> str: document = yaml.safe_load( (root / "artifacts" / "data.yaml").read_text(encoding="utf-8") ) try: revision = document["models"][tokenizer_name] except (KeyError, TypeError) as error: raise AdapterError(f"missing pinned tokenizer revision for {tokenizer_name}") from error if not isinstance(revision, str) or len(revision) != 40: raise AdapterError(f"invalid pinned tokenizer revision for {tokenizer_name}") return revision @staticmethod def _read_json(path: Path, label: str) -> object: if path.is_symlink() or not path.is_file(): raise AdapterError(f"{label} is missing or unsafe: {path}") try: return json.loads(path.read_text(encoding="utf-8")) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: raise AdapterError(f"{label} is invalid JSON: {path}") from error @staticmethod def _validate_argv(arguments: list[str]) -> None: if not arguments: raise AdapterError("adapter command is empty") for argument in arguments: if not isinstance(argument, str) or not argument: raise AdapterError("adapter command contains an empty argument") if "${" in argument or "{" in argument or "}" in argument or "None" in argument: raise AdapterError(f"adapter command contains an unresolved argument: {argument}") for argument in arguments: if argument.startswith("eval.checkpoint_path=") and not argument.split("=", 1)[1]: raise AdapterError("checkpoint path must not be empty") def _loader_batch_size(self, request: RunRequest) -> int: del request return 2 def _data_config(self, request: RunRequest) -> str: return self._DATA_CONFIGS[request.dataset_id] def _sampling_overrides( self, request: RunRequest, checkpoint: CheckpointSelection ) -> list[str]: raise NotImplementedError