diff --git a/README.md b/README.md index e73a970..8ec0f76 100644 --- a/README.md +++ b/README.md @@ -810,6 +810,28 @@ soup train --config soup.yaml --resume auto soup train --config soup.yaml --resume ./output/checkpoint-500 ``` +## Run Management & Cleanup + +LLM training generates massive checkpoint files. Soup automatically manages an SQLite database of your training loss and metrics, empowering you to safely reclaim disk space once training is complete. + +```bash +# List all historical training runs +soup runs list + +# Compare two differing experiments side-by-side +soup runs compare run_202611... run_202612... + +# Intelligently clean up redundant checkpoints +# (Preserves the final model and the checkpoint with the lowest loss) +soup runs clean run_202611... + +# Preview space that would be reclaimed across ALL experiments +soup runs clean --all --dry-run +``` + +By default, the `clean` command operates in "surgical mode" (`--keep-weights`), deleting huge optimizer state files (`optimizer.pt`) from lesser checkpoints to save gigabytes, but keeping their lightweight evaluation weights just in case you want to load them later. + + ## Batch Inference Run a model on a list of prompts and save results: diff --git a/soup_cli/commands/runs.py b/soup_cli/commands/runs.py index 3042166..cb58259 100644 --- a/soup_cli/commands/runs.py +++ b/soup_cli/commands/runs.py @@ -271,6 +271,139 @@ def delete( console.print(f"[green]Deleted run: {run['run_id']}[/]") +@app.command() +def clean( + run_id: Optional[str] = typer.Argument( + None, help="Run ID (or prefix) to clean. Omit if --all is used." + ), + all_runs: bool = typer.Option(False, "--all", help="Cleanup all historical runs."), + dry_run: bool = typer.Option( + False, "--dry-run", help="Estimate space savings without deleting." + ), + force: bool = typer.Option(False, "--force", "-f", help="Skip confirmation"), + keep_weights: bool = typer.Option( + True, "--keep-weights", + help="Keep intermediate model weights but delete optimizer states." + ), +): + """Intelligently clean up redundant checkpoint files to reclaim disk space.""" + import shutil + from pathlib import Path + + from soup_cli.experiment.tracker import ExperimentTracker + + tracker = ExperimentTracker() + + if all_runs and run_id: + raise typer.BadParameter("Cannot use --all and RUN_ID together.") + + if all_runs: + runs_to_clean = tracker.list_runs(limit=100000) + elif run_id: + run = tracker.get_run(run_id) + if not run: + console.print(f"[red]Run not found: {run_id}[/]") + raise typer.Exit(1) + runs_to_clean = [run] + else: + console.print("[red]Must specify RUN_ID or use --all[/]") + console.print("[dim]Example: soup runs clean run_2026...[/]") + raise typer.Exit(1) + + if not runs_to_clean: + console.print("[yellow]No runs to clean.[/]") + raise typer.Exit() + + total_bytes_to_reclaim = 0 + files_to_delete = [] + dirs_to_delete = [] + + for run in runs_to_clean: + out_dir_str = run.get("output_dir") + if not out_dir_str: + continue + try: + output_dir = Path(out_dir_str).resolve() + output_dir.relative_to(Path.cwd().resolve()) + except ValueError: + continue # skip paths outside cwd + if not output_dir.exists(): + continue + + metrics = tracker.get_metrics(run["run_id"]) + + best_step = -1 + valid_metrics = [m for m in metrics if m.get("loss") is not None] + if valid_metrics: + best_metric = min(valid_metrics, key=lambda x: x["loss"]) + best_step = best_metric["step"] + + checkpoints = [d for d in output_dir.glob("checkpoint-*") if d.is_dir()] + for ckpt in checkpoints: + is_best = False + try: + step = int(ckpt.name.split("-")[-1]) + if step == best_step: + is_best = True + except ValueError: + pass + + if is_best: + continue + + if keep_weights: + for opt_file in ckpt.glob("optimizer*.pt"): + total_bytes_to_reclaim += opt_file.stat().st_size + files_to_delete.append(opt_file) + for opt_file in ckpt.glob("optimizer*.safetensors"): + total_bytes_to_reclaim += opt_file.stat().st_size + files_to_delete.append(opt_file) + for sch_file in ckpt.glob("scheduler.pt"): + total_bytes_to_reclaim += sch_file.stat().st_size + files_to_delete.append(sch_file) + else: + size = sum(f.stat().st_size for f in ckpt.rglob('*') if f.is_file()) + total_bytes_to_reclaim += size + dirs_to_delete.append(ckpt) + + if total_bytes_to_reclaim == 0: + console.print( + "[green]No disposable checkpoint files found. " + "Storage is already optimized.[/]" + ) + raise typer.Exit() + + gb_to_reclaim = total_bytes_to_reclaim / (1024 ** 3) + + if dry_run: + console.print(f"[bold]Dry Run:[/] Would reclaim [green]{gb_to_reclaim:.2f} GB[/] " + f"from {len(files_to_delete)} files " + f"and {len(dirs_to_delete)} directories.") + for d in dirs_to_delete: + console.print(f" [red]Delete dir:[/]\t{d}") + for f in files_to_delete: + console.print(f" [red]Delete file:[/]\t{f}") + raise typer.Exit() + + if not force: + console.print(f"Ready to reclaim [green]{gb_to_reclaim:.2f} GB[/] by pruning checkpoints.") + if not typer.confirm("Do you want to proceed?"): + raise typer.Exit() + + for f in files_to_delete: + try: + f.unlink() + except OSError as e: + console.print(f"[yellow]Warning:[/] Failed to delete file {f}: {e}") + for d in dirs_to_delete: + try: + shutil.rmtree(d) + except OSError as e: + console.print(f"[yellow]Warning:[/] Failed to delete directory {d}: {e}") + + console.print(f"[green]Successfully reclaimed {gb_to_reclaim:.2f} GB.[/]") + + def _fmt_loss(run: dict) -> str: """Format loss as 'initial -> final'.""" init = run.get("initial_loss") diff --git a/tests/test_runs.py b/tests/test_runs.py index 09861e9..d3a969c 100644 --- a/tests/test_runs.py +++ b/tests/test_runs.py @@ -119,3 +119,64 @@ def test_runs_delete_with_data(tracker): assert "Deleted" in result.output # Verify it's gone assert tracker.get_run(run_id) is None + + +def test_runs_clean_not_found(): + """soup runs clean with bad IDs should fail.""" + result = runner.invoke(app, ["runs", "clean", "nonexistent"]) + assert result.exit_code == 1 + + +def test_runs_clean_with_data(tracker, tmp_path): + """soup runs clean should reclaim space natively.""" + # Note: Using Path.cwd() / "test_output" to satisfy security check that paths stay under CWD + out_dir = Path.cwd() / "test_output_clean" + if out_dir.exists(): + import shutil + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True) + + try: + # create some dummy checkpoints + ckpt1 = out_dir / "checkpoint-100" + ckpt1.mkdir() + (ckpt1 / "optimizer.pt").write_text("dummy") + (ckpt1 / "adapter_model.bin").write_text("model") + + ckpt2 = out_dir / "checkpoint-200" + ckpt2.mkdir() + (ckpt2 / "optimizer.pt").write_text("dummy") + (ckpt2 / "adapter_model.bin").write_text("model") + + run_id = tracker.start_run( + config_dict={}, device="cpu", device_name="CPU", gpu_info={}, + ) + tracker.finish_run( + run_id=run_id, + initial_loss=2.0, + final_loss=1.0, # let's say step 200 has lower loss + total_steps=200, + duration_secs=100.0, + output_dir=str(out_dir), + ) + # mock metrics: step 200 is best + tracker.log_metrics(run_id, step=100, loss=2.0) + tracker.log_metrics(run_id, step=200, loss=1.0) + + result = runner.invoke(app, ["runs", "clean", run_id, "--force"]) + assert result.exit_code == 0 + assert "Successfully reclaimed" in result.output + + # verify optimizer in ckpt1 is gone, but in ckpt2 it's potentially kept or not + # based on keep-weights + assert not (ckpt1 / "optimizer.pt").exists() + assert (ckpt2 / "optimizer.pt").exists() + + # test --all + result = runner.invoke(app, ["runs", "clean", "--all", "--force"]) + assert "No disposable checkpoint files found" in result.output + finally: + import shutil + if out_dir.exists(): + shutil.rmtree(out_dir) +