RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Introduce 'soup runs clean' for smart checkpoint space management (#9) 

* feat(cli): add 'soup runs clean' intelligent checkpoint cleanup to reclaim disk space

* update README

* feat(cli): add 'soup runs clean' intelligent checkpoint cleanup to reclaim disk space

* fixed whitespace trails

* style(cli): fix lints (line length and spacing) in runs.py

* style: fix all E501 line length lint errors

* fix test mismatch, improve deletion warnings, add path validation, and enforce argument exclusivity

* fix: break long message into multiple lines for Ruff compliance

* test: update runs clean test to use CWD-based output directory for security compliance
This commit is contained in:
Salil Mhatre 2026-04-06 23:05:16 +05:30 committed by GitHub
parent d748afd7f1
commit d134abb008
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 216 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
README.md
View File

@ -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:

133
soup_cli/commands/runs.py
View File

@ -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")

61
tests/test_runs.py
View File

@ -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)