From f1d1d814b0f56561410860b4507fdeff88b5f59e Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sun, 4 Oct 2026 11:54:51 +0530 Subject: [PATCH] feat: make every adapter reproducible from a PEFT or Unsloth recipe Add a reviewed recipe per catalog adapter, check recipes against the catalog in CD, render the PEFT and Unsloth arguments a run would get, and register a finished adapter only after it passes the artifact scan. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/cd.yml | 2 + README.md | 26 ++ .../adapters/claims-extraction-lora-next.yaml | 24 ++ config/adapters/claims-extraction-lora.yaml | 25 ++ .../adapters/support-classification-lora.yaml | 24 ++ pyproject.toml | 19 + src/llm_router/adapters.py | 373 ++++++++++++++++++ tests/unit/test_adapters.py | 235 +++++++++++ 8 files changed, 728 insertions(+) create mode 100644 config/adapters/claims-extraction-lora-next.yaml create mode 100644 config/adapters/claims-extraction-lora.yaml create mode 100644 config/adapters/support-classification-lora.yaml create mode 100644 src/llm_router/adapters.py create mode 100644 tests/unit/test_adapters.py diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 244de26..90568b8 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -66,6 +66,8 @@ jobs: run: python -m llm_router.chart --check && helm lint deploy/helm/llm-routing - name: Package the Helm chart run: helm package deploy/helm/llm-routing --destination chart + - name: Verify every catalog adapter has a matching recipe + run: python -m llm_router.adapters check - name: Render the canary and rollback plans # One plan per track (model, adapter, policy), each naming what it # rolls back to and the criteria that trigger it. diff --git a/README.md b/README.md index a860cb9..79745d3 100644 --- a/README.md +++ b/README.md @@ -310,6 +310,32 @@ be served. A request can never introduce a model path, revision, or adapter. Send `routing.domain` to request a domain adapter; the router applies the promoted adapter with the largest measured quality gain for that base revision and task, or none at all. +### Adapter recipes + +Every adapter in the catalog has a recipe in [`config/adapters`](config/adapters): the base +model commit it was trained from, the dataset version, and the LoRA hyperparameters. A recipe +is what makes an adapter reproducible, so a catalog adapter without one fails the check. + +```bash +python -m pip install -e ".[training]" # PEFT; add ".[unsloth]" for Unsloth +python -m llm_router.adapters check # recipes agree with the catalog +python -m llm_router.adapters plan claims-extraction-lora # the arguments a run would get +python -m llm_router.adapters train claims-extraction-lora --output ./out/claims +python -m llm_router.adapters register claims-extraction-lora --output ./out/claims +``` + +- `method: lora` trains over the full-precision base; `method: qlora` loads the base in 4-bit. + `framework` selects PEFT or Unsloth, and both are given the same rank, alpha, target modules, + commit, and seed. +- The base is pinned to a 40-character commit. A branch or tag is rejected, because it can move. +- A rank above 32 is rejected: the serving configuration would not load it. +- `register` runs the [artifact scan](#artifact-scanning) on the finished adapter and prints its + catalog entry. The entry starts in `development` with no measured gain, and its revision is the + artifact's digest. Promotion is a separate, reviewed change once a benchmark exists. + +Training has not been run: it needs a GPU. The recipes' `hf_repo` and `hf_revision` are +placeholders, as the catalog's base models are mocks. + ### Governance in MLflow The catalog decides what is served; [MLflow](https://mlflow.org/docs/latest/) keeps the record. diff --git a/config/adapters/claims-extraction-lora-next.yaml b/config/adapters/claims-extraction-lora-next.yaml new file mode 100644 index 0000000..12fff1a --- /dev/null +++ b/config/adapters/claims-extraction-lora-next.yaml @@ -0,0 +1,24 @@ +# Recipe for the staged claims adapter: same base, the next dataset version, +# and full-precision LoRA. hf_repo and hf_revision are placeholders. +id: claims-extraction-lora-next +base_model_id: small-specialist +base_revision: mock-small@sha256:dev +hf_repo: REPLACE_ME/small-specialist +hf_revision: "0000000000000000000000000000000000000000" +domain: claims +intended_tasks: [extraction] +method: lora +framework: peft +rank: 16 +alpha: 32 +dropout: 0.05 +target_modules: [q_proj, k_proj, v_proj, o_proj] +dataset: + path: s3://llm-routing-artifacts/datasets/claims-2026-06.jsonl + version: claims-2026-06 +training: + epochs: 3 + learning_rate: 0.0002 + batch_size: 8 + max_seq_length: 2048 + seed: 7 diff --git a/config/adapters/claims-extraction-lora.yaml b/config/adapters/claims-extraction-lora.yaml new file mode 100644 index 0000000..916bf8e --- /dev/null +++ b/config/adapters/claims-extraction-lora.yaml @@ -0,0 +1,25 @@ +# Recipe for the claims extraction adapter. The catalog's base model is a mock, +# so hf_repo and hf_revision are placeholders to replace with the real base. +id: claims-extraction-lora +base_model_id: small-specialist +base_revision: mock-small@sha256:dev +hf_repo: REPLACE_ME/small-specialist +hf_revision: "0000000000000000000000000000000000000000" +domain: claims +intended_tasks: [extraction] +# QLoRA: trained over the base loaded in 4-bit, to fit an L4. +method: qlora +framework: unsloth +rank: 16 +alpha: 32 +dropout: 0.05 +target_modules: [q_proj, k_proj, v_proj, o_proj] +dataset: + path: s3://llm-routing-artifacts/datasets/claims-2026-05.jsonl + version: claims-2026-05 +training: + epochs: 3 + learning_rate: 0.0002 + batch_size: 8 + max_seq_length: 2048 + seed: 7 diff --git a/config/adapters/support-classification-lora.yaml b/config/adapters/support-classification-lora.yaml new file mode 100644 index 0000000..9f3c235 --- /dev/null +++ b/config/adapters/support-classification-lora.yaml @@ -0,0 +1,24 @@ +# Recipe for the support classification adapter. hf_repo and hf_revision are +# placeholders to replace with the real base. +id: support-classification-lora +base_model_id: small-specialist +base_revision: mock-small@sha256:dev +hf_repo: REPLACE_ME/small-specialist +hf_revision: "0000000000000000000000000000000000000000" +domain: support +intended_tasks: [classification] +method: lora +framework: peft +rank: 8 +alpha: 16 +dropout: 0.05 +target_modules: [q_proj, v_proj] +dataset: + path: s3://llm-routing-artifacts/datasets/support-2026-04.jsonl + version: support-2026-04 +training: + epochs: 2 + learning_rate: 0.0001 + batch_size: 16 + max_seq_length: 1024 + seed: 7 diff --git a/pyproject.toml b/pyproject.toml index 38a8fd0..46b32a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,18 @@ governance = [ "mlflow-skinny>=3.5,<4", "sqlalchemy>=2,<3", ] +# Adapter training. Needs a GPU; nothing in the verification environment +# installs it. Unsloth is an alternative trainer, installed separately. +training = [ + "accelerate>=1.0", + "bitsandbytes>=0.44", + "datasets>=3.0", + "peft>=0.13", + "transformers>=4.46", +] +unsloth = [ + "unsloth", +] tracing = [ "opentelemetry-exporter-otlp-proto-http>=1.30,<2", "opentelemetry-sdk>=1.30,<2", @@ -87,6 +99,13 @@ module = ["redis.*"] ignore_missing_imports = true follow_imports = "skip" +[[tool.mypy.overrides]] +# The training stack ships in optional extras and is imported only inside +# the function that trains an adapter. +module = ["datasets.*", "peft.*", "transformers.*", "unsloth.*"] +ignore_missing_imports = true +follow_imports = "skip" + [[tool.mypy.overrides]] # MLflow ships in the optional governance extra and is imported lazily, only # when a tracking URI is given; the client is used behind GovernanceStore. diff --git a/src/llm_router/adapters.py b/src/llm_router/adapters.py new file mode 100644 index 0000000..de9a06e --- /dev/null +++ b/src/llm_router/adapters.py @@ -0,0 +1,373 @@ +"""LoRA and QLoRA adapter recipes for PEFT and Unsloth (sections 6 and 7.5). + +An adapter in the catalog is the result of a training run. The recipe is the +reviewed record of that run: which immutable base revision it started from, +which dataset version it saw, and the hyperparameters used. This module +validates recipes against the catalog, turns one into the arguments PEFT or +Unsloth take, and turns a finished run into the catalog entry that describes +it. + +Training itself needs a GPU and the ``training`` extra. Everything else runs +anywhere, which is what lets a recipe be reviewed and checked in CI. +""" + +import json +import re +from collections.abc import Sequence +from enum import StrEnum +from pathlib import Path +from typing import Any + +import yaml +from pydantic import BaseModel, Field, model_validator + +from llm_router.artifact_scan import MAX_LORA_RANK, ArtifactKind, scan +from llm_router.models import TaskClass +from llm_router.registry import AdapterProfile, BenchmarkDelta, LifecycleStage, Registry + +RECIPE_DIRECTORY = Path("config/adapters") +COMMIT = re.compile(r"[0-9a-f]{40}") + + +class AdapterError(RuntimeError): + """Raised when a recipe or a finished adapter cannot be accepted.""" + + +class Method(StrEnum): + LORA = "lora" + # LoRA trained over a base model loaded in 4-bit. + QLORA = "qlora" + + +class Framework(StrEnum): + PEFT = "peft" + UNSLOTH = "unsloth" + + +class DatasetReference(BaseModel): + path: str + version: str + text_field: str = "text" + + +class TrainingSettings(BaseModel): + epochs: float = Field(default=3.0, gt=0) + learning_rate: float = Field(default=2e-4, gt=0) + batch_size: int = Field(default=8, ge=1) + gradient_accumulation_steps: int = Field(default=1, ge=1) + max_seq_length: int = Field(default=2048, ge=16) + warmup_ratio: float = Field(default=0.03, ge=0, le=1) + seed: int = 0 + + +class AdapterRecipe(BaseModel): + """Everything needed to reproduce one adapter.""" + + id: str + base_model_id: str + base_revision: str + # Where the base weights come from. The revision is a commit, never a + # branch or tag, so the base cannot move underneath the adapter. + hf_repo: str + hf_revision: str + domain: str + intended_tasks: frozenset[TaskClass] + method: Method = Method.LORA + framework: Framework = Framework.PEFT + rank: int = Field(default=16, ge=1) + alpha: int = Field(default=32, ge=1) + dropout: float = Field(default=0.05, ge=0, lt=1) + target_modules: tuple[str, ...] = ("q_proj", "k_proj", "v_proj", "o_proj") + dataset: DatasetReference + training: TrainingSettings = TrainingSettings() + + @model_validator(mode="after") + def keep_the_recipe_servable_and_reproducible(self) -> "AdapterRecipe": + if not COMMIT.fullmatch(self.hf_revision): + raise ValueError( + f"recipe {self.id}: hf_revision must be a 40-character commit, " + "not a branch or tag that can move" + ) + if self.rank > MAX_LORA_RANK: + raise ValueError( + f"recipe {self.id}: rank {self.rank} exceeds the {MAX_LORA_RANK} the engine loads" + ) + if not self.target_modules: + raise ValueError(f"recipe {self.id}: at least one target module is required") + return self + + +def load_recipes(directory: Path = RECIPE_DIRECTORY) -> dict[str, AdapterRecipe]: + recipes: dict[str, AdapterRecipe] = {} + for path in sorted(directory.glob("*.yaml")): + recipe = AdapterRecipe.model_validate(yaml.safe_load(path.read_text(encoding="utf-8"))) + if recipe.id != path.stem: + raise AdapterError(f"{path.name} holds recipe {recipe.id}; name the file after it") + recipes[recipe.id] = recipe + return recipes + + +def check_recipes(registry: Registry, recipes: dict[str, AdapterRecipe]) -> tuple[str, ...]: + """Every way the recipes and the catalog disagree. + + A catalog adapter with no recipe cannot be reproduced. A recipe that names + a different base, dataset or quantization than its catalog entry describes + some other adapter. + """ + + problems: list[str] = [] + bases = {card.id: card.revision for card in registry.models} + for adapter in registry.adapters: + recipe = recipes.get(adapter.id) + if recipe is None: + problems.append(f"{adapter.id}: in the catalog but has no recipe") + continue + if (recipe.base_model_id, recipe.base_revision) != ( + adapter.base_model_id, + adapter.base_revision, + ): + problems.append( + f"{adapter.id}: recipe trains on {recipe.base_model_id}@{recipe.base_revision}, " + f"the catalog says {adapter.base_model_id}@{adapter.base_revision}" + ) + if recipe.dataset.version != adapter.dataset_version: + problems.append( + f"{adapter.id}: recipe uses dataset {recipe.dataset.version}, " + f"the catalog says {adapter.dataset_version}" + ) + if (recipe.method is Method.QLORA) != adapter.quantized: + problems.append( + f"{adapter.id}: recipe method is {recipe.method.value} but the catalog " + f"records quantized={str(adapter.quantized).lower()}" + ) + if (recipe.domain, recipe.intended_tasks) != (adapter.domain, adapter.intended_tasks): + problems.append(f"{adapter.id}: recipe domain or tasks differ from the catalog") + for recipe in recipes.values(): + if bases.get(recipe.base_model_id) != recipe.base_revision: + problems.append( + f"{recipe.id}: base {recipe.base_model_id}@{recipe.base_revision} " + "is not a revision in the catalog" + ) + return tuple(dict.fromkeys(problems)) + + +def peft_arguments(recipe: AdapterRecipe) -> dict[str, Any]: + """Keyword arguments for ``peft.LoraConfig``.""" + + return { + "task_type": "CAUSAL_LM", + "r": recipe.rank, + "lora_alpha": recipe.alpha, + "lora_dropout": recipe.dropout, + "target_modules": list(recipe.target_modules), + "bias": "none", + } + + +def quantization_arguments(recipe: AdapterRecipe) -> dict[str, Any] | None: + """Keyword arguments for ``transformers.BitsAndBytesConfig``, for QLoRA only.""" + + if recipe.method is not Method.QLORA: + return None + return { + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_use_double_quant": True, + "bnb_4bit_compute_dtype": "bfloat16", + } + + +def unsloth_arguments(recipe: AdapterRecipe) -> tuple[dict[str, Any], dict[str, Any]]: + """Arguments for Unsloth's ``from_pretrained`` and ``get_peft_model``.""" + + load = { + "model_name": recipe.hf_repo, + "revision": recipe.hf_revision, + "max_seq_length": recipe.training.max_seq_length, + "load_in_4bit": recipe.method is Method.QLORA, + } + adapt = { + "r": recipe.rank, + "lora_alpha": recipe.alpha, + "lora_dropout": recipe.dropout, + "target_modules": list(recipe.target_modules), + "bias": "none", + "random_state": recipe.training.seed, + } + return load, adapt + + +def training_arguments(recipe: AdapterRecipe) -> dict[str, Any]: + """Keyword arguments for ``transformers.TrainingArguments``.""" + + settings = recipe.training + return { + "num_train_epochs": settings.epochs, + "learning_rate": settings.learning_rate, + "per_device_train_batch_size": settings.batch_size, + "gradient_accumulation_steps": settings.gradient_accumulation_steps, + "warmup_ratio": settings.warmup_ratio, + "seed": settings.seed, + # Checkpoints carry pickled optimizer state, which may not be released. + "save_strategy": "no", + "report_to": [], + } + + +def plan(recipe: AdapterRecipe) -> dict[str, Any]: + """What a training run for this recipe would be given, for review.""" + + document: dict[str, Any] = { + "recipe": recipe.id, + "framework": recipe.framework.value, + "method": recipe.method.value, + "base": {"repo": recipe.hf_repo, "revision": recipe.hf_revision}, + "dataset": recipe.dataset.model_dump(), + "training_arguments": training_arguments(recipe), + } + if recipe.framework is Framework.UNSLOTH: + load, adapt = unsloth_arguments(recipe) + document["unsloth"] = {"from_pretrained": load, "get_peft_model": adapt} + else: + document["peft"] = { + "lora_config": peft_arguments(recipe), + "quantization_config": quantization_arguments(recipe), + } + return document + + +def train(recipe: AdapterRecipe, output: Path) -> None: # pragma: no cover - needs a GPU + """Train the adapter and write it to ``output`` as safetensors. + + This has not been run: it needs a GPU and the training extra, neither of + which the verification environment has. + """ + + import tempfile + + from datasets import load_dataset + from transformers import ( + AutoTokenizer, + DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, + ) + + if recipe.framework is Framework.UNSLOTH: + from unsloth import FastLanguageModel + + load, adapt = unsloth_arguments(recipe) + model, tokenizer = FastLanguageModel.from_pretrained(**load) + model = FastLanguageModel.get_peft_model(model, **adapt) + else: + from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training + from transformers import AutoModelForCausalLM, BitsAndBytesConfig + + quantization = quantization_arguments(recipe) + model = AutoModelForCausalLM.from_pretrained( + recipe.hf_repo, + revision=recipe.hf_revision, + quantization_config=BitsAndBytesConfig(**quantization) if quantization else None, + device_map="auto", + ) + tokenizer = AutoTokenizer.from_pretrained(recipe.hf_repo, revision=recipe.hf_revision) + if quantization: + model = prepare_model_for_kbit_training(model) + model = get_peft_model(model, LoraConfig(**peft_arguments(recipe))) + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + dataset = load_dataset("json", data_files=recipe.dataset.path, split="train") + tokenized = dataset.map( + lambda row: tokenizer( + row[recipe.dataset.text_field], + truncation=True, + max_length=recipe.training.max_seq_length, + ), + remove_columns=dataset.column_names, + ) + with tempfile.TemporaryDirectory() as scratch: + Trainer( + model=model, + args=TrainingArguments(output_dir=scratch, **training_arguments(recipe)), + train_dataset=tokenized, + data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False), + ).train() + output.mkdir(parents=True, exist_ok=True) + model.save_pretrained(str(output), safe_serialization=True) + + +def catalog_entry(recipe: AdapterRecipe, output: Path) -> AdapterProfile: + """Scan a finished adapter and describe it as a catalog entry. + + The entry starts in development with no measured gain: promotion is a + separate, reviewed change made once a benchmark exists. The revision is + the artifact's own digest, so the entry cannot describe different bytes. + """ + + report = scan(output, ArtifactKind.ADAPTER) + if not report.passed: + reasons = "; ".join(f"{item.rule.value} at {item.path}" for item in report.findings) + raise AdapterError(f"adapter {recipe.id} failed its artifact scan: {reasons}") + return AdapterProfile( + id=recipe.id, + base_model_id=recipe.base_model_id, + base_revision=recipe.base_revision, + adapter_revision=f"{recipe.id}@{report.digest}", + domain=recipe.domain, + intended_tasks=recipe.intended_tasks, + dataset_version=recipe.dataset.version, + benchmark=BenchmarkDelta(quality_delta=0.0), + stage=LifecycleStage.DEVELOPMENT, + quantized=recipe.method is Method.QLORA, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """Check recipes against the catalog, or plan, train, or register one adapter. + + ``check`` exits 1 if the recipes and the catalog disagree. ``plan`` prints + what a run would be given. ``train`` runs it. ``register`` scans a + finished adapter directory and prints its catalog entry. + """ + + import argparse + + from llm_router.registry import load_registry + + parser = argparse.ArgumentParser(description=main.__doc__) + parser.add_argument("command", choices=["check", "plan", "train", "register"]) + parser.add_argument("recipe", nargs="?") + parser.add_argument("--recipes", default=str(RECIPE_DIRECTORY)) + parser.add_argument("--catalog", default="config/registry.yaml") + parser.add_argument("--output", help="adapter directory to write or to register") + arguments = parser.parse_args(argv) + + recipes = load_recipes(Path(arguments.recipes)) + if arguments.command == "check": + problems = check_recipes(load_registry(arguments.catalog), recipes) + print(json.dumps({"recipes": sorted(recipes), "problems": list(problems)}, indent=2)) + return 1 if problems else 0 + + recipe = recipes.get(arguments.recipe or "") + if recipe is None: + parser.error(f"{arguments.command} needs a recipe, one of: {', '.join(sorted(recipes))}") + if arguments.command == "plan": + print(json.dumps(plan(recipe), indent=2)) + return 0 + if not arguments.output: + parser.error(f"{arguments.command} needs --output") + output = Path(arguments.output) + if arguments.command == "train": # pragma: no cover - needs a GPU + train(recipe, output) + try: + entry = catalog_entry(recipe, output) + except AdapterError as error: + print(json.dumps({"error": str(error)}, indent=2)) + return 1 + print(yaml.safe_dump([entry.model_dump(mode="json")], sort_keys=False), end="") + return 0 + + +if __name__ == "__main__": # pragma: no cover - command-line entry point + raise SystemExit(main()) diff --git a/tests/unit/test_adapters.py b/tests/unit/test_adapters.py new file mode 100644 index 0000000..92848ec --- /dev/null +++ b/tests/unit/test_adapters.py @@ -0,0 +1,235 @@ +import json +import struct +from pathlib import Path +from typing import Any + +import pytest +import yaml +from pydantic import ValidationError + +from llm_router.adapters import ( + AdapterError, + AdapterRecipe, + Framework, + Method, + catalog_entry, + check_recipes, + load_recipes, + main, + peft_arguments, + plan, + quantization_arguments, + training_arguments, + unsloth_arguments, +) +from llm_router.artifact_scan import artifact_digest +from llm_router.registry import AdapterProfile, LifecycleStage, load_registry + +CATALOG = load_registry("config/registry.yaml") +RECIPES = load_recipes() +COMMIT = "a" * 40 + + +def recipe(**changes: Any) -> AdapterRecipe: + base = RECIPES["claims-extraction-lora"].model_dump(mode="json") + return AdapterRecipe.model_validate({**base, **changes}) + + +def write_adapter(root: Path, rank: int = 16) -> Path: + header = json.dumps({"lora_A": {"dtype": "F32", "shape": [8], "data_offsets": [0, 32]}}) + root.mkdir(parents=True, exist_ok=True) + (root / "adapter_model.safetensors").write_bytes( + struct.pack(" None: + assert set(RECIPES) == {adapter.id for adapter in CATALOG.adapters} + assert check_recipes(CATALOG, RECIPES) == () + + +def test_a_catalog_adapter_with_no_recipe_cannot_be_reproduced() -> None: + remaining = { + key: value for key, value in RECIPES.items() if key != "support-classification-lora" + } + + assert check_recipes(CATALOG, remaining) == ( + "support-classification-lora: in the catalog but has no recipe", + ) + + +@pytest.mark.parametrize( + ("changes", "problem"), + [ + ( + {"base_revision": "mock-small@sha256:old"}, + "the catalog says small-specialist@mock-small", + ), + ( + {"dataset": {"path": "x.jsonl", "version": "claims-2025-01"}}, + "uses dataset claims-2025-01", + ), + ({"method": "lora"}, "records quantized=true"), + ({"domain": "billing"}, "domain or tasks differ"), + ], +) +def test_a_recipe_that_describes_a_different_adapter_is_reported( + changes: dict[str, Any], problem: str +) -> None: + recipes = {**RECIPES, "claims-extraction-lora": recipe(**changes)} + + problems = check_recipes(CATALOG, recipes) + + assert any(problem in item for item in problems), problems + + +def test_a_recipe_for_a_base_the_catalog_does_not_hold_is_reported() -> None: + orphan = recipe(id="new-adapter", base_model_id="retired-model") + + assert check_recipes(CATALOG, {**RECIPES, "new-adapter": orphan}) == ( + "new-adapter: base retired-model@mock-small@sha256:dev is not a revision in the catalog", + ) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"hf_revision": "main"}, "must be a 40-character commit"), + ({"rank": 64}, "exceeds the 32 the engine loads"), + ({"target_modules": []}, "at least one target module"), + ], +) +def test_a_recipe_that_is_not_reproducible_or_servable_is_rejected( + changes: dict[str, Any], message: str +) -> None: + with pytest.raises(ValidationError, match=message): + recipe(**changes) + + +def test_a_recipe_file_must_be_named_after_its_adapter(tmp_path: Path) -> None: + (tmp_path / "misnamed.yaml").write_text( + yaml.safe_dump(recipe().model_dump(mode="json")), encoding="utf-8" + ) + + with pytest.raises(AdapterError, match="name the file after it"): + load_recipes(tmp_path) + + +def test_peft_gets_a_lora_config_and_four_bit_loading_only_for_qlora() -> None: + qlora = recipe(hf_revision=COMMIT) + lora = recipe(method="lora") + + assert peft_arguments(qlora) == { + "task_type": "CAUSAL_LM", + "r": 16, + "lora_alpha": 32, + "lora_dropout": 0.05, + "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj"], + "bias": "none", + } + assert quantization_arguments(qlora) == { + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_use_double_quant": True, + "bnb_4bit_compute_dtype": "bfloat16", + } + assert quantization_arguments(lora) is None + + +def test_unsloth_is_pinned_to_the_same_commit_and_seed() -> None: + load, adapt = unsloth_arguments(recipe(hf_revision=COMMIT)) + + assert load == { + "model_name": "REPLACE_ME/small-specialist", + "revision": COMMIT, + "max_seq_length": 2048, + "load_in_4bit": True, + } + assert adapt["r"] == 16 and adapt["random_state"] == 7 + assert unsloth_arguments(recipe(method="lora"))[0]["load_in_4bit"] is False + + +def test_training_never_writes_checkpoints_that_could_not_be_released() -> None: + arguments = training_arguments(recipe()) + + assert arguments["save_strategy"] == "no" + assert arguments["seed"] == 7 and arguments["num_train_epochs"] == 3 + + +def test_the_plan_shows_the_arguments_for_the_chosen_framework() -> None: + unsloth = plan(recipe()) + peft = plan(recipe(framework="peft")) + + assert unsloth["framework"] == Framework.UNSLOTH.value and "peft" not in unsloth + assert unsloth["unsloth"]["from_pretrained"]["load_in_4bit"] is True + assert peft["peft"]["lora_config"]["r"] == 16 and "unsloth" not in peft + assert peft["method"] == Method.QLORA.value + assert peft["dataset"]["version"] == "claims-2026-05" + + +def test_a_finished_adapter_becomes_a_development_entry_named_by_its_digest( + tmp_path: Path, +) -> None: + output = write_adapter(tmp_path) + + entry = catalog_entry(recipe(), output) + + assert entry.adapter_revision == f"claims-extraction-lora@{artifact_digest(output)}" + assert entry.stage is LifecycleStage.DEVELOPMENT + assert entry.benchmark.quality_delta == 0.0 + assert entry.quantized is True + assert (entry.base_model_id, entry.base_revision) == ( + "small-specialist", + "mock-small@sha256:dev", + ) + + +def test_an_adapter_that_fails_its_scan_is_never_described(tmp_path: Path) -> None: + output = write_adapter(tmp_path, rank=64) + (output / "optimizer.pt").write_bytes(b"\x80\x04N.") + + with pytest.raises(AdapterError, match=r"pickle-format at optimizer\.pt"): + catalog_entry(recipe(), output) + + +def test_the_command_line_checks_plans_and_registers( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + assert main(["check"]) == 0 + assert json.loads(capsys.readouterr().out)["problems"] == [] + + assert main(["plan", "support-classification-lora"]) == 0 + assert json.loads(capsys.readouterr().out)["peft"]["lora_config"]["r"] == 8 + + output = write_adapter(tmp_path / "adapter") + assert main(["register", "claims-extraction-lora", "--output", str(output)]) == 0 + (printed,) = yaml.safe_load(capsys.readouterr().out) + assert AdapterProfile.model_validate(printed).stage is LifecycleStage.DEVELOPMENT + + (output / "train.py").write_text("print('hello')\n", encoding="utf-8") + assert main(["register", "claims-extraction-lora", "--output", str(output)]) == 1 + assert "executable-code" in json.loads(capsys.readouterr().out)["error"] + + +def test_the_command_line_fails_when_recipes_and_catalog_disagree( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + (tmp_path / "claims-extraction-lora.yaml").write_text( + yaml.safe_dump(recipe().model_dump(mode="json")), encoding="utf-8" + ) + + assert main(["check", "--recipes", str(tmp_path)]) == 1 + assert len(json.loads(capsys.readouterr().out)["problems"]) == 2 + + +@pytest.mark.parametrize( + "arguments", [["plan"], ["plan", "nobody"], ["register", "claims-extraction-lora"]] +) +def test_the_command_line_says_what_it_is_missing(arguments: list[str]) -> None: + with pytest.raises(SystemExit) as raised: + main(arguments) + assert raised.value.code == 2