tests/test_train_smoke.py

"""Smoke #6: the shared train.py wires up the REAL trl 1.7.1 SFT API.

Trains a tiny random Qwen2 for 1 step on a handful of dummy-teacher rows, no
4-bit, on CPU. Proves SFTConfig(max_length, completion_only_loss) +
SFTTrainer(processing_class, peft_config) + conversational prompt/completion
data all fit together. Requires network (downloads the tiny test model).
"""
import argparse
from pathlib import Path

import pytest

from mathcompose.common.io import write_jsonl
from mathcompose.datagen.prm800k_loader import iter_prm800k
from mathcompose.datagen.gen_verifier_data import generate_verifier_dataset
from mathcompose.teachers import get_teacher

TINY = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
FIX = Path(__file__).parent / "fixtures" / "prm800k_sample.jsonl"


@pytest.mark.slow
@pytest.mark.network
def test_train_one_step(tmp_path):
    pytest.importorskip("torch")
    from mathcompose.train.train import build_and_train

    rows = list(generate_verifier_dataset(iter_prm800k(FIX), get_teacher("dummy"), banned=set()))
    train_file = tmp_path / "train.jsonl"
    write_jsonl(rows, train_file)

    out_dir = tmp_path / "vsmoke"
    args = argparse.Namespace(
        task="v", config="configs/verifier_v.yaml", base_id=TINY,
        train_file=str(train_file), val_file=None, output_dir=str(out_dir),
        no_4bit=True, max_steps=1, wandb=False, push=False, hub_model_id=None,
    )
    build_and_train(args)
    # a LoRA adapter should have been saved
    assert (out_dir / "adapter_config.json").exists() or (out_dir / "adapter_model.safetensors").exists()