code/fastfill/tests/test_training_repairs.py

"""Bounded offline regressions for evaluation, merge artifacts, and download exits."""
import importlib.util
import json
import runpy
import socket
import sys
import types
from pathlib import Path
from unittest.mock import Mock

import numpy as np
import pytest

ROOT = Path(__file__).resolve().parents[2]


@pytest.fixture(autouse=True)
def offline(monkeypatch):
    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
    monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")

    def denied(*args, **kwargs):
        raise AssertionError("These tests must not access the network")

    monkeypatch.setattr(socket.socket, "connect", denied)


def load_merge():
    spec = importlib.util.spec_from_file_location("merge_under_test", ROOT / "fastfill/merge_lora.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


@pytest.fixture
def hf(monkeypatch):
    calls = []

    class Batch(dict):
        def to(self, device):
            return self

    class Tokenizer:
        @classmethod
        def from_pretrained(cls, *args, **kwargs):
            return cls()

        def __call__(self, prompts, **kwargs):
            assert kwargs.get("add_special_tokens") is False
            if isinstance(prompts, str):
                return {"input_ids": [ord(c) for c in prompts]}
            width = max(map(len, prompts))
            ids = np.array([[0] * (width - len(p)) + [ord(c) for c in p] for p in prompts])
            return Batch(input_ids=ids, attention_mask=(ids != 0).astype(int))

        def batch_decode(self, rows, **kwargs):
            return ["".join(chr(int(c)) for c in row) for row in rows]

    class Model:
        device = "cpu"

        @classmethod
        def from_pretrained(cls, *args, **kwargs):
            return cls()

        def generate(self, input_ids, attention_mask, max_new_tokens, **kwargs):
            calls.append((attention_mask.sum(axis=1).tolist(), max_new_tokens))
            assert kwargs["do_sample"] is False
            assert input_ids.shape[1] + max_new_tokens == 10
            return np.concatenate((input_ids, np.repeat(input_ids[:, -1:], max_new_tokens, axis=1)), axis=1)

    monkeypatch.setitem(sys.modules, "vllm", None)
    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bfloat16"))
    monkeypatch.setitem(sys.modules, "transformers", types.SimpleNamespace(
        AutoTokenizer=Tokenizer, AutoModelForCausalLM=Model))
    return calls


@pytest.mark.parametrize("batch", [1, 2, 3, 20])
def test_hf_preserves_each_prompt_budget_and_order(hf, batch):
    from fastfill.evaluate import generate

    prompts = ["a", "bbbbbbbb", "cc", "d", "e" * 10, "f" * 11, "gg"]
    assert generate("mock", prompts, 10, batch) == ["a" * 9, "bb", "c" * 8, "d" * 9, None, None, "g" * 8]
    assert all(len(lengths) <= batch for lengths, _ in hf)
    if batch > 1:
        assert any(len(lengths) == 2 for lengths, _ in hf)


@pytest.mark.parametrize("prompts", [[], ["x" * 10, "y" * 11]])
def test_hf_does_not_generate_over_limit_prompts(hf, prompts):
    from fastfill.evaluate import generate

    assert generate("mock", prompts, 10, 2) == [None] * len(prompts)
    assert hf == []


@pytest.fixture
def merge_mock(monkeypatch):
    class Base:
        def save_pretrained(self, path):
            Path(path).mkdir(parents=True, exist_ok=True)
            (Path(path) / "config.json").write_text(json.dumps({"rope_parameters": {"rope_theta": 1000000}}))

    class Tokenizer:
        def __init__(self, source):
            self.config = json.loads((Path(source) / "tokenizer_config.json").read_text())

        def save_pretrained(self, path):
            (Path(path) / "tokenizer_config.json").write_text(json.dumps(self.config))

    loader = Mock(side_effect=lambda path, **kwargs: Tokenizer(path))
    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bfloat16"))
    monkeypatch.setitem(sys.modules, "transformers", types.SimpleNamespace(
        AutoTokenizer=types.SimpleNamespace(from_pretrained=loader),
        AutoModelForCausalLM=types.SimpleNamespace(from_pretrained=Mock(return_value=Base()))))
    monkeypatch.setitem(sys.modules, "peft", types.SimpleNamespace(
        PeftModel=types.SimpleNamespace(from_pretrained=Mock(return_value=types.SimpleNamespace(
            merge_and_unload=lambda: Base())))))
    return load_merge(), loader


def merge_paths(tmp_path):
    base, adapter, merged = tmp_path / "base", tmp_path / "run/final", tmp_path / "merged"
    base.mkdir()
    adapter.mkdir(parents=True)
    (base / "tokenizer_config.json").write_text(json.dumps({"pad_token": None, "extra_special_tokens": ["base"]}))
    return base, adapter, merged


def test_merge_preserves_adapter_tokenizer_and_compatibility(merge_mock, tmp_path):
    merge, loader = merge_mock
    base, adapter, merged = merge_paths(tmp_path)
    (adapter / "tokenizer_config.json").write_text(json.dumps({
        "pad_token": "<eos>", "chat_template": "adapter template", "extra_special_tokens": ["trained"]}))
    merge.apply_lora(str(base), str(merged), str(adapter))
    tc = json.loads((merged / "tokenizer_config.json").read_text())
    assert tc == {"pad_token": "<eos>", "chat_template": "adapter template", "additional_special_tokens": ["trained"]}
    assert Path(loader.call_args.args[0]) == adapter
    cfg = json.loads((merged / "config.json").read_text())
    assert cfg["rope_theta"] == cfg["rope_parameters"]["rope_theta"] == 1000000


def test_merge_falls_back_only_when_adapter_tokenizer_absent(merge_mock, tmp_path):
    merge, loader = merge_mock
    base, adapter, merged = merge_paths(tmp_path)
    merge.apply_lora(str(base), str(merged), str(adapter))
    assert Path(loader.call_args.args[0]) == base
    (adapter / "tokenizer_config.json").write_text("invalid json")
    with pytest.raises(json.JSONDecodeError):
        merge.apply_lora(str(base), str(tmp_path / "bad"), str(adapter))
    assert not (tmp_path / "bad").exists()


@pytest.mark.parametrize("local", [False, True])
def test_merge_preserves_one_manifest_family(merge_mock, tmp_path, local):
    merge, _ = merge_mock
    base, adapter, merged = merge_paths(tmp_path)
    run = adapter if local else adapter.parent
    manifests = {"run_manifest.json": '{"data":"v3"}\n',
                 "run_manifest.resume1.json": '{"data":"v3.1"}\n',
                 "run_manifest.resume10.json": '{"args":{"resume":true}}\n'}
    for name, content in manifests.items():
        (run / name).write_text(content)
    (run / "run_manifest.resume-backup.json").write_text("not a run")
    if local:
        (adapter.parent / "run_manifest.json").write_text('{"unrelated":true}')
        (adapter.parent / "run_manifest.resume2.json").write_text('{"unrelated":true}')
    merge.apply_lora(str(base), str(merged), str(adapter))
    actual = {p.name: p.read_text() for p in merged.glob("run_manifest*.json")}
    assert actual == manifests


@pytest.fixture
def download_mock(monkeypatch, tmp_path):
    from fastfill import download
    from huggingface_hub.hf_api import RepoFile

    sources = {"first": ("owner/one", "first"), "second": ("owner/one", "second"),
               "third": ("owner/two", "third")}
    api = Mock()
    api.list_repo_tree.side_effect = lambda repo, **kwargs: [
        RepoFile(path=f"{folder}/data.json", size=2, oid="0" * 40)
        for owner, folder in sources.values() if owner == repo]
    monkeypatch.setattr(download, "SOURCES", sources)
    monkeypatch.setattr(download, "HfApi", Mock(return_value=api))
    monkeypatch.setattr(sys, "argv", ["download", "--out", str(tmp_path), "--workers", "1"])
    return download, api


@pytest.mark.parametrize("failures", [{"first/data.json"}, {"first/data.json", "second/data.json", "third/data.json"}])
def test_download_failures_exit_nonzero_after_all_sources(download_mock, monkeypatch, failures):
    download, api = download_mock
    fetch = Mock(side_effect=lambda repo, path, local: path not in failures)
    monkeypatch.setattr(download, "fetch", fetch)
    with pytest.raises(SystemExit) as error:
        download.main()
    assert error.value.code == 1
    assert fetch.call_count == 3
    assert api.list_repo_tree.call_count == 2


def test_download_success_and_size_resume(download_mock, monkeypatch, tmp_path):
    download, _ = download_mock
    cached = tmp_path / "owner__one/first/data.json"
    cached.parent.mkdir(parents=True)
    cached.write_text("{}")
    fetch = Mock(return_value=True)
    monkeypatch.setattr(download, "fetch", fetch)
    assert download.main() in (None, 0)
    assert [c.args[1] for c in fetch.call_args_list] == ["second/data.json", "third/data.json"]


def test_download_module_entrypoint_exits_after_retries(monkeypatch, tmp_path):
    import huggingface_hub
    from huggingface_hub.hf_api import RepoFile

    api = Mock()
    api.list_repo_tree.return_value = [RepoFile(path="SAGE10k/data.json", size=2, oid="0" * 40)]
    monkeypatch.setattr(huggingface_hub, "HfApi", Mock(return_value=api))
    fetch = Mock(side_effect=OSError("offline fixture failure"))
    monkeypatch.setattr(huggingface_hub, "hf_hub_download", fetch)
    monkeypatch.setattr("time.sleep", lambda _: None)
    monkeypatch.setattr(sys, "argv", ["download", "--out", str(tmp_path), "--only", "SAGE-10k", "--workers", "1"])
    with pytest.raises(SystemExit) as error:
        runpy.run_path(str(ROOT / "fastfill/download.py"), run_name="__main__")
    assert error.value.code == 1
    assert fetch.call_count == 8


def test_tiny_real_cpu_merge(tmp_path):
    import torch
    from peft import LoraConfig, PeftModel, get_peft_model
    from tokenizers import Tokenizer
    from tokenizers.models import WordLevel
    from tokenizers.pre_tokenizers import Whitespace
    from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, Qwen3Config, Qwen3ForCausalLM

    torch.set_num_threads(1)
    torch.manual_seed(42)
    backend = Tokenizer(WordLevel({"<unk>": 0, "<eos>": 1, "hello": 2, "world": 3, "<user>": 4}, unk_token="<unk>"))
    backend.pre_tokenizer = Whitespace()
    tok = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="<unk>", eos_token="<eos>")
    tok.add_special_tokens({"additional_special_tokens": ["<user>"]})
    tok.chat_template = "{% for m in messages %}{{ m['content'] }}<eos>{% endfor %}"
    cfg = Qwen3Config(vocab_size=len(tok), hidden_size=16, intermediate_size=32, num_hidden_layers=1,
                      num_attention_heads=2, num_key_value_heads=1, head_dim=8, max_position_embeddings=32,
                      eos_token_id=1, rope_theta=1000000.)
    base, adapter, merged = tmp_path / "base", tmp_path / "run/final", tmp_path / "merged"
    model = Qwen3ForCausalLM(cfg).eval()
    model.save_pretrained(base)
    tok.save_pretrained(base)
    lora = get_peft_model(model, LoraConfig(task_type="CAUSAL_LM", r=2, lora_alpha=4, target_modules=["q_proj"]))
    with torch.no_grad():
        for name, parameter in lora.named_parameters():
            if "lora_" in name:
                parameter.fill_(0.125)
    lora.save_pretrained(adapter)
    tok.pad_token = tok.eos_token
    tok.save_pretrained(adapter)
    for name in ("run_manifest.json", "run_manifest.resume1.json"):
        (adapter.parent / name).write_text('{"fixture":true}\n')
    reference = PeftModel.from_pretrained(
        AutoModelForCausalLM.from_pretrained(base, dtype=torch.bfloat16, local_files_only=True),
        adapter, torch_dtype=torch.bfloat16).eval()
    ids = torch.tensor([[2, 3, 1]])
    with torch.no_grad():
        before = reference(input_ids=ids).logits.float()
    load_merge().apply_lora(str(base), str(merged), str(adapter))
    reloaded = AutoModelForCausalLM.from_pretrained(merged, dtype=torch.bfloat16, local_files_only=True).eval()
    merged_tok = AutoTokenizer.from_pretrained(merged, local_files_only=True)
    with torch.no_grad():
        after = reloaded(input_ids=ids).logits.float()
    assert (before - after).abs().max().item() < 0.01
    assert merged_tok.pad_token_id == tok.pad_token_id == 1
    assert merged_tok.get_vocab() == tok.get_vocab()
    assert merged_tok.chat_template == tok.chat_template
    assert merged_tok(["hello", "hello world"], padding=True, return_tensors="pt")["input_ids"].shape == (2, 2)
    saved = json.loads((merged / "config.json").read_text())
    assert saved["rope_theta"] == saved["rope_parameters"]["rope_theta"] == 1000000.
    tc = json.loads((merged / "tokenizer_config.json").read_text())
    assert "additional_special_tokens" in tc and not isinstance(tc.get("extra_special_tokens"), list)
    assert (merged / "run_manifest.resume1.json").read_text() == '{"fixture":true}\n'