File size: 3,798 Bytes
5ef3875 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """End-to-end test of `ycode-lm autotrain` across two simulated GPU sessions (CPU, local hub)."""
import json
import pytest
torch = pytest.importorskip("torch")
pytest.importorskip("numpy")
from ycode.lm.autotrain import AutoTrainer, Plan # noqa: E402
from ycode.lm.hub import LocalHub # noqa: E402
from ycode.lm.train import TrainConfig, load_checkpoint, train # noqa: E402
@pytest.fixture
def base_model(data_dir, tmp_path):
cfg = TrainConfig(preset="v4-tiny", model_overrides={"block_size": 64}, batch_size=4, max_steps=10,
warmup_steps=2, eval_interval=10, eval_iters=1, log_interval=1000, device="cpu",
precision="fp32")
train(data_dir, tmp_path / "v4", cfg, log=lambda *_: None)
return tmp_path / "v4"
def _plan(base_model, src, **kw):
base = dict(base_model=str(base_model), preset="v5-tiny", sources=[str(src)], max_mb=None,
pretrain_steps=6, sft_steps=4, batch_size=2, nproc=1, precision="fp32",
grad_checkpoint=False, first_segment_steps=3, eval_interval=3, device="cpu",
session_hours=5, segment_minutes=60)
base.update(kw)
return Plan(**base)
def test_autotrain_resumes_across_sessions_and_publishes(base_model, tmp_path):
from tests.test_lm import CODE
src = tmp_path / "src"
src.mkdir()
for i in range(6):
(src / f"m{i}.py").write_text(CODE.replace("add", f"add{i}") * 3)
ckpt, release = LocalHub(tmp_path / "hub", "you/ckpt"), LocalHub(tmp_path / "hub", "you/ycode-lm-v5")
logs: list[str] = []
# Session 1: stops after one segment (as if the 12-hour session ended).
s1 = AutoTrainer(_plan(base_model, src, max_segments=1), tmp_path / "session1", ckpt, release,
log=logs.append).run()
assert (s1.phase, s1.pretrain_step) == ("pretrain", 3)
assert any("Grew" in line for line in logs)
hub_state = json.loads((ckpt.root / "state" / "state.json").read_text())
assert hub_state["pretrain_step"] == 3 and hub_state["seconds_per_step"] > 0
# Session 2: a brand-new machine/work dir, so it must resume purely from the hub.
logs.clear()
s2 = AutoTrainer(_plan(base_model, src), tmp_path / "session2", ckpt, release, log=logs.append).run()
assert s2.phase == "done" and s2.pretrain_step == 6 and s2.sft_step == 4 and s2.sessions == 2
assert not any("Grew" in line for line in logs) # resumed, did not start over
assert any("prepared in an earlier session" in line for line in logs) # data not rebuilt
assert any("training steps 3 -> 6" in line for line in logs)
# The published model loads, is version 5, and carries a model card + eval results.
model, _, payload = load_checkpoint(release.root)
assert payload["config"]["arch_version"] == 5 and model.cfg.n_embd == 128 and payload["stage"] == "sft"
card = (release.root / "README.md").read_text()
assert "YCode-LM v5" in card and "pass@1" in card and "fix@1" in card
assert json.loads((release.root / "eval.json").read_text())
# A third run is a no-op.
s3 = AutoTrainer(_plan(base_model, src), tmp_path / "session3", ckpt, release, log=logs.append).run()
assert s3.phase == "done"
def test_autotrain_cli_dry_run_requires_repo_or_local_hub(tmp_path, capsys):
from ycode.lm import cli
assert cli.main(["autotrain", "--base-model", str(tmp_path / "x")]) == 2
assert "--hf-repo" in capsys.readouterr().err
def test_hfhub_requires_token(monkeypatch):
pytest.importorskip("huggingface_hub")
from ycode.lm.hub import HFHub
import ycode.lm.hub as hub
monkeypatch.delenv("HF_TOKEN", raising=False)
monkeypatch.setattr(hub, "_saved_login", lambda: None)
with pytest.raises(RuntimeError, match="HF_TOKEN"):
HFHub("you/repo")
|