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

fix(shrink): python-review findings (output-dir containment, exc breadth, type hints, epochs clamp) (v0.71.29)

This commit is contained in:
Alpamys 2026-07-05 10:51:07 +05:00
parent ee47ad3a3f
commit db58382506
3 changed files with 101 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
src/soup_cli/commands/shrink.py
View File

@ -17,7 +17,7 @@ from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Optional
from typing import Any, Optional, Sequence
import typer
from rich.console import Console
@ -25,24 +25,30 @@ from rich.markup import escape
from rich.panel import Panel
from rich.table import Table
from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink
from soup_cli.utils.paths import atomic_write_text, enforce_under_cwd_and_no_symlink
from soup_cli.utils.shrink import (
DECISION_SHIP,
MAX_TOLERANCE,
LayerImportance,
compute_layer_importance,
decide_shrink,
prune_model_layers,
render_shrink_panel,
resolve_drop_count,
select_drop_block,
shrink_arch_of,
shrink_verdict_to_dict,
)
console = Console()
# 64 MiB cap on the calibration JSONL (symlink-pointed-to-/dev/zero DoS guard).
_MAX_CALIB_BYTES = 64 * 1024 * 1024
# 64 MiB cap on the calibration / heal JSONL (symlink-to-/dev/zero DoS guard).
_MAX_INPUT_BYTES = 64 * 1024 * 1024
_MAX_CALIB_ROWS = 10_000
_MAX_HEAL_STEPS = 1_000_000
# Guard against an innocuous flag combo (huge --heal-steps, tiny heal set)
# expanding into millions of epochs; refuse rather than emit an absurd config.
_MAX_HEAL_EPOCHS = 100
_PPL_MAX_LENGTH = 512
@ -91,9 +97,9 @@ def _load_calib(path: str) -> list[str]:
except OSError as exc:
raise typer.BadParameter(f"calib path unreadable: {exc}") from exc
with os.fdopen(fd, "r", encoding="utf-8") as handle:
if os.fstat(handle.fileno()).st_size > _MAX_CALIB_BYTES:
if os.fstat(handle.fileno()).st_size > _MAX_INPUT_BYTES:
raise typer.BadParameter(
f"calib file exceeds {_MAX_CALIB_BYTES} bytes"
f"calib file exceeds {_MAX_INPUT_BYTES} bytes"
)
prompts: list[str] = []
for line in handle:
@ -122,11 +128,17 @@ def _count_params(model: object) -> int:
return sum(p.numel() for p in model.parameters()) # type: ignore[attr-defined]
def _perplexity(model: object, tokenizer: object, prompts: list[str], device: str) -> float:
"""Mean unconditional-LM perplexity of ``model`` over ``prompts``.
def _perplexity(
model: object, tokenizer: object, prompts: Sequence[str], device: str
) -> float:
"""Unweighted mean of per-example perplexities of ``model`` over ``prompts``.
``exp(mean per-example cross-entropy)`` with ``labels = input_ids`` (the
whole sequence is the target). Returns ``inf`` when no example is usable.
whole sequence is the target). Each prompt is weighted equally regardless of
length — this is NOT token-count-weighted corpus perplexity, so the absolute
numbers are not directly comparable to ``soup eval`` / lm-eval-harness; the
identical procedure is applied before and after, so the SHIP/DON'T-SHIP
*ratio* is valid. Returns ``inf`` when no example is usable.
"""
import math
@ -147,14 +159,16 @@ def _perplexity(model: object, tokenizer: object, prompts: list[str], device: st
continue
out = model(input_ids=input_ids, labels=input_ids) # type: ignore[operator]
loss = float(out.loss.item())
if loss == loss: # not NaN
if not math.isnan(loss):
losses.append(loss)
if not losses:
return float("inf")
return math.exp(sum(losses) / len(losses))
def _load_for_shrink(model_id: str, device: Optional[str], trust_remote_code: bool):
def _load_for_shrink(
model_id: str, device: Optional[str], trust_remote_code: bool
) -> tuple[Any, Any, str]:
"""Load a model + tokenizer for shrinking (trust_remote_code probe + warn)."""
from soup_cli.utils.live_eval import load_model_and_tokenizer
from soup_cli.utils.trust_remote import (
@ -169,7 +183,9 @@ def _load_for_shrink(model_id: str, device: Optional[str], trust_remote_code: bo
return load_model_and_tokenizer(model_id, device=device, trust_remote_code=trc)
def _render_importance_table(importances, chosen) -> None:
def _render_importance_table(
importances: Sequence[LayerImportance], chosen: LayerImportance
) -> None:
table = Table(title="soup shrink — layer importance (lower = safer to drop)")
table.add_column("Rank", justify="right")
table.add_column("Block (start..end)")
@ -246,7 +262,7 @@ def shrink(
)
except typer.Exit:
raise
except (typer.BadParameter, ValueError) as exc:
except (typer.BadParameter, ValueError, RuntimeError, OSError, ImportError) as exc:
console.print(f"[red]Error:[/] {escape(str(exc))}")
raise typer.Exit(1) from exc
@ -266,8 +282,11 @@ def _shrink_impl(
attach_to_registry: Optional[str],
plan_only: bool,
) -> None:
if not (0.0 <= tolerance <= 5.0):
raise typer.BadParameter("--tolerance must be in [0.0, 5.0]")
if not (0.0 <= tolerance <= MAX_TOLERANCE):
raise typer.BadParameter(f"--tolerance must be in [0.0, {MAX_TOLERANCE}]")
# Contain the output dir (arbitrary-write / symlink-redirect guard) BEFORE
# any mkdir / save_pretrained / report write.
enforce_under_cwd_and_no_symlink(output_dir, "--output-dir")
# Fail fast on the flag combination BEFORE loading a multi-GB model.
if (drop_ratio is None) == (drop_layers is None):
raise typer.BadParameter("set exactly one of --drop-ratio / --drop-layers")
@ -282,8 +301,6 @@ def _shrink_impl(
console.print(f"[dim]Loading {escape(model)} ...[/]")
mdl, tokenizer, dev = _load_for_shrink(model, device, trust_remote_code)
# Reject an unsupported architecture up front (before the importance scan).
from soup_cli.utils.shrink import shrink_arch_of
shrink_arch_of(mdl)
n_layers = int(mdl.config.num_hidden_layers)
count = resolve_drop_count(n_layers, drop_ratio=drop_ratio, drop_layers=drop_layers)
@ -359,7 +376,9 @@ def _shrink_impl(
report = shrink_verdict_to_dict(verdict)
report["model"] = model
report["dropped_block"] = [chosen.start, chosen.start + chosen.block_size - 1]
report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
atomic_write_text(
json.dumps(report, indent=2), str(report_path), field="shrink report"
)
console.print(f"[green]Wrote[/] {escape(str(report_path))}")
if attach_to_registry:
@ -372,7 +391,7 @@ def _attach_to_registry(registry_id: str, report_path: str) -> None:
"""Attach the shrink report JSON as a registry artifact (best-effort)."""
try:
from soup_cli.registry.attach import attach_artifact
except Exception as exc: # noqa: BLE001 — registry is optional
except ImportError as exc:
console.print(
f"[yellow]Warning:[/] could not import registry attach helper: {escape(str(exc))}"
)
@ -402,8 +421,8 @@ def _count_jsonl_rows(path: str) -> int:
except OSError as exc:
raise typer.BadParameter(f"heal path unreadable: {exc}") from exc
with os.fdopen(fd, "r", encoding="utf-8") as handle:
if os.fstat(handle.fileno()).st_size > _MAX_CALIB_BYTES:
raise typer.BadParameter(f"heal file exceeds {_MAX_CALIB_BYTES} bytes")
if os.fstat(handle.fileno()).st_size > _MAX_INPUT_BYTES:
raise typer.BadParameter(f"heal file exceeds {_MAX_INPUT_BYTES} bytes")
rows = sum(1 for line in handle if line.strip())
if rows == 0:
raise typer.BadParameter("heal file has no rows")
@ -429,6 +448,12 @@ def _build_heal_config_yaml(
import math
epochs = max(1, math.ceil(steps * _HEAL_BATCH_SIZE / max(1, heal_rows)))
if epochs > _MAX_HEAL_EPOCHS:
raise typer.BadParameter(
f"--heal-steps {steps} over {heal_rows} heal rows expands to "
f"{epochs} epochs (> {_MAX_HEAL_EPOCHS}); reduce --heal-steps or "
"grow the heal set."
)
return (
"base: {pruned}\n"
"task: distill\n"
@ -484,7 +509,7 @@ def _run_heal(
)
load_config_from_string(yaml_text) # validate before spending a subprocess
config_path = Path(pruned_dir).parent / "heal_config.yaml"
config_path.write_text(yaml_text, encoding="utf-8")
atomic_write_text(yaml_text, str(config_path), field="heal config")
argv = [
sys.executable,

16
src/soup_cli/utils/shrink.py
View File

@ -19,6 +19,7 @@ from __future__ import annotations
import math
import re
from dataclasses import asdict, dataclass
from typing import Any, Optional, Sequence
from rich.panel import Panel
@ -115,7 +116,7 @@ def decide_shrink(
)
def shrink_verdict_to_dict(verdict: ShrinkVerdict) -> dict:
def shrink_verdict_to_dict(verdict: ShrinkVerdict) -> dict[str, Any]:
"""Plain-dict view of a ``ShrinkVerdict`` (JSON-serialisable)."""
return asdict(verdict)
@ -137,6 +138,9 @@ def render_shrink_panel(verdict: ShrinkVerdict) -> Panel:
# ---------------------------------------------------------------------------
# Arch allowlist + prune (torch-lazy)
# ---------------------------------------------------------------------------
# Note: real SmolLM/SmolLM2 checkpoints report model_type="llama", so the
# "smollm" entry is aspirational (any HF model that literally reports a "smol"
# model_type) — "llama" is matched first and wins for the released checkpoints.
_ARCH_PATTERNS = {
"llama": re.compile(r"llama", re.I),
"qwen": re.compile(r"qwen", re.I),
@ -167,7 +171,7 @@ def shrink_arch_of(model: object) -> str:
)
def layer_list(model: object):
def layer_list(model: object) -> Any:
"""Return ``model.model.layers`` (the decoder ``ModuleList``), arch-guarded."""
shrink_arch_of(model) # raises on unsupported arch
try:
@ -215,7 +219,9 @@ _IMPORTANCE_MAX_LENGTH = 512
_DEFAULT_MAX_PROMPTS = 256
def resolve_drop_count(num_layers: int, *, drop_ratio, drop_layers) -> int:
def resolve_drop_count(
num_layers: int, *, drop_ratio: Optional[float], drop_layers: Optional[int]
) -> int:
"""Resolve the block **count** from exactly one of ratio / explicit count.
``drop_layers = round(drop_ratio * num_layers)`` when a ratio is given.
@ -246,7 +252,7 @@ def resolve_drop_count(num_layers: int, *, drop_ratio, drop_layers) -> int:
def compute_layer_importance(
model: object,
tokenizer: object,
prompts,
prompts: Sequence[str],
*,
block_size: int,
device: str,
@ -323,7 +329,7 @@ def compute_layer_importance(
return imps
def select_drop_block(importances) -> LayerImportance:
def select_drop_block(importances: Sequence[LayerImportance]) -> LayerImportance:
"""Return the least-important (min angular-distance) candidate block."""
if not importances:
raise ValueError("no importance scores to select from")

43
tests/test_v07129.py
View File

@ -645,3 +645,46 @@ class TestRegistryAttach:
report.write_text('{"decision":"SHIP"}', encoding="utf-8")
# Must not raise even for a nonexistent entry (best-effort warn).
_attach_to_registry("nonexistent-id", str(report))
# ---------------------------------------------------------------------------
# Review-fix regression guards (python-review CRITICAL + MEDIUM-2)
# ---------------------------------------------------------------------------
class TestReviewFixes:
def test_output_dir_outside_cwd_rejected(self, tmp_path, monkeypatch):
"""--output-dir must be cwd-contained (arbitrary-write guard)."""
from typer.testing import CliRunner
from soup_cli.cli import app
work = tmp_path / "work"
work.mkdir()
monkeypatch.chdir(work)
calib = work / "calib.jsonl"
calib.write_text('{"text":"hi there friend"}\n', encoding="utf-8")
model_dir = _write_tiny_model(work / "m", layers=6)
r = CliRunner().invoke(
app,
["shrink", "--model", model_dir, "--drop-layers", "2",
"--calib", "calib.jsonl", "--device", "cpu",
"--output-dir", str(tmp_path / "escape")],
)
assert r.exit_code != 0
def test_heal_epochs_clamp_rejects_absurd_combo(self):
"""Huge --heal-steps over a tiny heal set is refused, not silently run."""
import typer
from soup_cli.commands.shrink import _build_heal_config_yaml
with pytest.raises(typer.BadParameter):
_build_heal_config_yaml(
pruned_dir="./m", teacher="t", heal_data="./h.jsonl",
steps=1_000_000, out_dir="./o", heal_rows=1,
)
def test_perplexity_no_top_level_math_import_uses_isnan(self):
"""The NaN filter uses math.isnan, not the x == x self-compare idiom."""
src = pathlib.Path("src/soup_cli/commands/shrink.py").read_text(encoding="utf-8")
assert "loss == loss" not in src
assert "math.isnan(loss)" in src