diff --git a/src/kev_laya/cli.py b/src/kev_laya/cli.py index 9782a74..a83afa5 100644 --- a/src/kev_laya/cli.py +++ b/src/kev_laya/cli.py @@ -223,11 +223,13 @@ def main(argv=None) -> int: if any(supplied_wandb) and not all(supplied_wandb): raise ValueError("all W&B identity components are required") wandb_ref = dict(zip(("entity", "project", "run_id"), supplied_wandb)) if all(supplied_wandb) else None + parent_sha256 = point.get("checkpoint_sha256") if not args.resume else None + del point # release deserialized weights and optimizer state before training result = train(model, tokenizer, suite["train"], manifest, args.out, settings, loss_cfg, Limits(**cfg.get("limits", {})), resume=args.resume, stop_after=args.stop_after, experiment_id=args.experiment_id, run_id=args.run_id, scheduler_ref=scheduler_ref, wandb_ref=wandb_ref, publish_wandb=args.wandb_publish, - parent_sha256=point.get("checkpoint_sha256") if not args.resume else None, + parent_sha256=parent_sha256, quality_budget=quality_budget, quality_allocation_id=f"{arm}:{settings.seed}" if quality_protocol is not None else None) print(json.dumps({k: v for k, v in result.items() if k != "metrics"}, indent=2)) diff --git a/tests/test_cli_checkpoint_lifetime.py b/tests/test_cli_checkpoint_lifetime.py new file mode 100644 index 0000000..fc90d41 --- /dev/null +++ b/tests/test_cli_checkpoint_lifetime.py @@ -0,0 +1,37 @@ +"""CPU-only regression for CLI checkpoint payload ownership.""" +import gc +import json +import weakref + +import pytest +import torch + +from kev_laya import cli + + +@pytest.mark.parametrize("command", ["train", "reward-train"]) +@pytest.mark.parametrize("initial_arg", ["--init", "--resume"]) +def test_training_cli_releases_checkpoint_payload(tmp_path, monkeypatch, command, initial_arg): + config = tmp_path / "config.json" + config.write_text(json.dumps({"training": {"steps": 2}, "objective": {"reinforce": 1.0}})) + observed = {} + digest = "a" * 64 + + def load_checkpoint(path, device): + tensor = torch.empty(1, device="cpu") + observed["payload_tensor"] = weakref.ref(tensor) + return object(), object(), {"checkpoint_sha256": digest, "model": {"weight": tensor}} + + def train(model, tokenizer, data, manifest, output, settings, objective, limits, **kwargs): + assert settings.steps == 2 + gc.collect() + assert observed["payload_tensor"]() is None, "CLI retained checkpoint tensor during train" + assert kwargs["parent_sha256"] == (digest if initial_arg == "--init" else None) + assert kwargs["resume"] == (tmp_path / "checkpoint.pt" if initial_arg == "--resume" else None) + return {"state": {"step": 1}} + + monkeypatch.setattr(cli, "load_checkpoint", load_checkpoint) + monkeypatch.setattr(cli, "load_suite", lambda path: ({"train": []}, {})) + monkeypatch.setattr(cli, "train", train) + assert cli.main([command, "--suite", str(tmp_path), "--out", str(tmp_path / "run"), + "--config", str(config), initial_arg, str(tmp_path / "checkpoint.pt")]) == 0