diff --git a/src/soup_cli/commands/shrink.py b/src/soup_cli/commands/shrink.py index 77832ff..0911d26 100644 --- a/src/soup_cli/commands/shrink.py +++ b/src/soup_cli/commands/shrink.py @@ -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, diff --git a/src/soup_cli/utils/shrink.py b/src/soup_cli/utils/shrink.py index a07b2a5..9f7231f 100644 --- a/src/soup_cli/utils/shrink.py +++ b/src/soup_cli/utils/shrink.py @@ -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") diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 02fe487..683c290 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -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