From fe9fe06b684f7ff362096bc42bf37088ca4b78c5 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Mon, 11 May 2026 19:13:54 +0500 Subject: [PATCH] =?UTF-8?q?feat(training):=20Curriculum-Aware=20Trainer=20?= =?UTF-8?q?=E2=80=94=20dynamic=20bucket=20re-weighting=20(v0.48.0=20Part?= =?UTF-8?q?=20A,=20BETA)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BETA. Adds `training.curriculum_dynamic: true` schema flag with online uncertainty estimation: every N steps, aggregate per-sample loss + grad-norm into per-bucket softmax weights, water-filled to enforce a minimum `curriculum_dynamic_floor`. DDP/grad-accum safety via `validate_distributed_curriculum` cross-validator that rejects un-coordinated multi-rank runs upfront — the well-known footgun where divergent per-rank stats silently desynchronise the sampler. New `soup runs curriculum-curve ` visualiser with TOCTOU (`os.lstat + S_ISLNK`) + 50 MB file-size cap + 100k-line streaming cap on the history file. Schema gated to sft/pretrain on transformers backend; mlx + non-SFT rejected with distinct messages. Live HF Trainer callback wiring deferred to v0.48.1 (stub-then-live pattern; mirrors v0.27.0 MII / v0.37.0 multipack / v0.41.0 LLaMA Pro). Review fixes: - water-fill design fix (code-review HIGH): removed trailing renorm that could push elements below `floor` when accumulated float error left sum slightly > 1.0. Softmax already sums to 1.0, so water-fill output also sums to 1.0 (drift bounded by nb*eps). - DoS caps on `render_curve` + `parse_history_jsonl` (`_MAX_HISTORY_ROWS=100_000`) — without these an attacker-controlled JSONL with 10M rows would OOM the process. - `curriculum-curve` CLI: symlink rejection + 50 MB + 100k-line caps, null-byte rejection on tracker-supplied `output_dir`. +74 tests. Co-Authored-By: Claude Opus 4.7 (1M context) --- soup_cli/commands/runs.py | 85 ++++ soup_cli/config/schema.py | 79 +++ soup_cli/utils/curriculum_dynamic.py | 394 ++++++++++++++ tests/test_v0480_part_a.py | 733 +++++++++++++++++++++++++++ 4 files changed, 1291 insertions(+) create mode 100644 soup_cli/utils/curriculum_dynamic.py create mode 100644 tests/test_v0480_part_a.py diff --git a/soup_cli/commands/runs.py b/soup_cli/commands/runs.py index f323067..4a9c376 100644 --- a/soup_cli/commands/runs.py +++ b/soup_cli/commands/runs.py @@ -554,3 +554,88 @@ def _plot_loss_curve(metrics: list[dict]) -> None: plt.ylabel("Loss") plt.theme("dark") plt.show() + + +@app.command(name="curriculum-curve") +def curriculum_curve( + run_id: str = typer.Argument(..., help="Run ID (or prefix)."), + history_path: str = typer.Option( + None, "--history", + help="Override path to curriculum_history.jsonl (default: under run output_dir).", + ), + width: int = typer.Option(10, "--width", min=4, max=200, help="Per-bucket column width."), +) -> None: + """v0.48.0 (BETA) — visualise dynamic curriculum bucket weights.""" + import json + import os + import stat as _stat + + from soup_cli.experiment.tracker import ExperimentTracker + from soup_cli.utils.curriculum_dynamic import parse_history_jsonl, render_curve + from soup_cli.utils.paths import is_under_cwd + + tracker = ExperimentTracker() + run = tracker.get_run(run_id) + if run is None: + console.print(f"[red]Run not found:[/] {markup_escape(run_id)}") + raise typer.Exit(1) + + if history_path is None: + out_dir = str(run.get("output_dir") or ".") + if "\x00" in out_dir: + console.print("[red]run output_dir contains null bytes[/]") + raise typer.Exit(2) + candidate = os.path.join(out_dir, "curriculum_history.jsonl") + else: + candidate = history_path + real = os.path.realpath(candidate) + if not is_under_cwd(real): + console.print("[red]history path is outside cwd[/]") + raise typer.Exit(2) + if os.path.lexists(real): + try: + _st = os.lstat(real) + except OSError as exc: + console.print( + f"[red]history path is not stat-able:[/] " + f"{markup_escape(os.path.basename(real))}" + ) + raise typer.Exit(2) from exc + if _stat.S_ISLNK(_st.st_mode): + console.print( + f"[red]history path is a symlink (rejected for safety):[/] " + f"{markup_escape(os.path.basename(real))}" + ) + raise typer.Exit(2) + if not os.path.isfile(real): + console.print( + f"[yellow]curriculum_history.jsonl not found:[/] " + f"{markup_escape(os.path.basename(real))}" + ) + raise typer.Exit(1) + max_history_bytes = 50 * 1024 * 1024 # 50 MB + if os.path.getsize(real) > max_history_bytes: + console.print("[red]history file exceeds 50 MB cap[/]") + raise typer.Exit(2) + + rows = [] + max_lines = 100_000 + with open(real, "r", encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line: + continue + rows.append(json.loads(line)) + if len(rows) > max_lines: + console.print("[red]history exceeds 100k-row cap[/]") + raise typer.Exit(2) + if not rows: + console.print("[yellow](no curriculum history rows)[/]") + return + nb = len(rows[0].get("weights", [])) + try: + normalised = parse_history_jsonl(rows) + except (ValueError, TypeError) as exc: + console.print(f"[red]history malformed: {markup_escape(str(exc))}[/]") + raise typer.Exit(2) from exc + console.print(render_curve(normalised, num_buckets=nb, width=width)) diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 15a0dd6..201173d 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -1064,6 +1064,40 @@ class TrainingConfig(BaseModel): default=4, ge=1, le=20, description="Number of difficulty stages for curriculum learning", ) + # Curriculum-Aware dynamic re-weighting (v0.48.0 Part A — BETA) + curriculum_dynamic: bool = Field( + default=False, + description=( + "BETA: dynamically re-weight curriculum buckets every N steps via " + "online uncertainty estimation (per-sample loss + grad norm). " + "Requires curriculum=true. Multi-rank launches must wire an " + "all_reduce hook on per-bucket stats (see " + "utils.curriculum_dynamic.validate_distributed_curriculum)." + ), + ) + curriculum_dynamic_recompute_steps: int = Field( + default=50, ge=1, le=100_000, + description=( + "Recompute curriculum bucket sampler weights every N global " + "training steps." + ), + ) + curriculum_dynamic_floor: float = Field( + default=0.05, gt=0.0, le=0.5, + description=( + "Minimum normalised per-bucket weight after softmax. " + "Must be in (0.0, 1/curriculum_buckets]; the cross-validator " + "tightens this to the per-config ceiling. Prevents bucket " + "starvation." + ), + ) + curriculum_dynamic_temperature: float = Field( + default=1.0, gt=0.0, le=100.0, + description=( + "Softmax temperature on the uncertainty signal. Higher = flatter " + "distribution; lower = concentrate on hardest buckets." + ), + ) # Loss watchdog — auto-stop on loss spikes loss_watchdog: bool = Field( default=False, @@ -1316,6 +1350,25 @@ class TrainingConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_curriculum_dynamic_requires_curriculum(self) -> "TrainingConfig": + """v0.48.0 Part A — dynamic re-weighting layers on the static + curriculum bucketer; it cannot run alone.""" + if self.curriculum_dynamic and not self.curriculum: + raise ValueError( + "curriculum_dynamic requires curriculum=true " + "(dynamic re-weighting needs the static bucketer)." + ) + # Cross-check: floor must leave room above uniform/N. + if self.curriculum_dynamic: + ceiling = 1.0 / max(self.curriculum_buckets, 1) + if self.curriculum_dynamic_floor > ceiling: + raise ValueError( + f"curriculum_dynamic_floor={self.curriculum_dynamic_floor} " + f"must be <= 1/curriculum_buckets ({ceiling:.4f})." + ) + return self + @model_validator(mode="after") def _validate_spike_recovery_requires_watchdog(self) -> "TrainingConfig": """Spike recovery is a watchdog hook — it needs the watchdog enabled.""" @@ -1632,6 +1685,32 @@ class SoupConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_curriculum_dynamic_supported(self) -> "SoupConfig": + """v0.48.0 Part A — Curriculum-Aware dynamic re-weighting. + + BETA: wired for transformers backend, sft + pretrain tasks only. + MLX backend rejected (callback is HF Trainer-specific). Other tasks + rejected because their per-sample loss semantics differ enough that + the bucket-level uncertainty heuristic does not transfer cleanly. + Multi-trainer expansion tracked for v0.48.1. + """ + if not self.training.curriculum_dynamic: + return self + if self.backend == "mlx": + raise ValueError( + "curriculum_dynamic is not supported on the mlx backend " + "(callback is HF Trainer-specific). " + "Use backend='transformers' or set curriculum_dynamic: false." + ) + if self.task not in ("sft", "pretrain"): + raise ValueError( + f"curriculum_dynamic is not supported for task={self.task!r} " + "in v0.48.0 (only sft and pretrain are wired). " + "Set curriculum_dynamic: false or switch task." + ) + return self + @model_validator(mode="after") def _validate_relora_supported_tasks(self) -> "SoupConfig": """v0.40.6 (#67) — ReLoRA callback wired in every transformer-backend diff --git a/soup_cli/utils/curriculum_dynamic.py b/soup_cli/utils/curriculum_dynamic.py new file mode 100644 index 0000000..9438a62 --- /dev/null +++ b/soup_cli/utils/curriculum_dynamic.py @@ -0,0 +1,394 @@ +"""Curriculum-Aware dynamic re-weighting (v0.48.0 Part A — BETA). + +Online uncertainty estimation: every N steps, aggregate per-sample loss and +gradient-norm fingerprints into bucket-level weights and surface a recommended +sampler weight per bucket. Up-weight high-uncertainty / under-fit buckets; +down-weight already-mastered ones. + +DDP / grad-accum safety: all-reduce of per-sample stats across ranks is the +well-known footgun for dynamic curriculum learning. We document the contract +here and surface a cross-validator that rejects ``curriculum_dynamic=true`` +combined with launches that have not declared rank coordination. + +This module ships BETA-flagged: the math is pure-Python + numpy-free; the live +HF Trainer callback wiring is deferred to v0.48.1 once external benchmarks have +landed (mirrors the v0.27.0 MII / v0.37.0 multipack / v0.41.0 LLaMA Pro +stub-then-live pattern). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, List, Mapping, Sequence, Tuple + +# Bounds — match project conventions (e.g. v0.32.0 GradAccumMonitor). +_MIN_BUCKETS = 1 +_MAX_BUCKETS = 20 +_MIN_RECOMPUTE_STEPS = 1 +_MAX_RECOMPUTE_STEPS = 100_000 +_MAX_BUCKET_SAMPLES = 1_000_000 # DoS cap on stats accumulation +_MAX_HISTORY_ROWS = 100_000 # DoS cap on curriculum-history JSONL parsing +_DEFAULT_FLOOR = 0.05 # min per-bucket weight after normalisation +_DEFAULT_TEMP = 1.0 + +__all__ = [ + "DynamicCurriculumPolicy", + "BucketStats", + "compute_bucket_weights", + "validate_distributed_curriculum", +] + + +def _reject_bool_int(name: str, value) -> int: + if isinstance(value, bool): + raise ValueError(f"{name} must be int, not bool") + if not isinstance(value, int): + raise TypeError(f"{name} must be int, got {type(value).__name__}") + return value + + +def _reject_bool_float(name: str, value) -> float: + if isinstance(value, bool): + raise ValueError(f"{name} must be float, not bool") + if not isinstance(value, (int, float)): + raise TypeError(f"{name} must be float, got {type(value).__name__}") + fv = float(value) + if not math.isfinite(fv): + raise ValueError(f"{name} must be finite (got {value!r})") + return fv + + +@dataclass(frozen=True) +class DynamicCurriculumPolicy: + """Frozen config for the dynamic re-weighting policy. + + Attributes: + num_buckets: Number of difficulty buckets (must match + ``training.curriculum_buckets`` schema field). + recompute_every_n_steps: Refresh sampler weights every N global steps. + floor: Minimum normalised per-bucket weight (defends against + "starve a bucket" pathology). In ``(0, 1/num_buckets]``. + temperature: Softmax temperature applied to uncertainty signal. Higher + values flatten the distribution toward uniform; lower values + concentrate weight on the hardest buckets. + """ + + num_buckets: int + recompute_every_n_steps: int = 50 + floor: float = _DEFAULT_FLOOR + temperature: float = _DEFAULT_TEMP + + def __post_init__(self) -> None: + nb = _reject_bool_int("num_buckets", self.num_buckets) + if nb < _MIN_BUCKETS or nb > _MAX_BUCKETS: + raise ValueError( + f"num_buckets must be in [{_MIN_BUCKETS}, {_MAX_BUCKETS}], got {nb}" + ) + rs = _reject_bool_int( + "recompute_every_n_steps", self.recompute_every_n_steps + ) + if rs < _MIN_RECOMPUTE_STEPS or rs > _MAX_RECOMPUTE_STEPS: + raise ValueError( + f"recompute_every_n_steps must be in " + f"[{_MIN_RECOMPUTE_STEPS}, {_MAX_RECOMPUTE_STEPS}], got {rs}" + ) + fv = _reject_bool_float("floor", self.floor) + # floor must leave at least equal-share room; uniform = 1/nb. + ceiling = 1.0 / nb + if fv <= 0.0 or fv > ceiling: + raise ValueError( + f"floor must be in (0.0, {ceiling}] for num_buckets={nb}, got {fv}" + ) + tv = _reject_bool_float("temperature", self.temperature) + if tv <= 0.0: + raise ValueError(f"temperature must be > 0, got {tv}") + + def should_recompute(self, global_step: int) -> bool: + """True when the current global step is a recompute boundary.""" + gs = _reject_bool_int("global_step", global_step) + if gs < 0: + raise ValueError(f"global_step must be >= 0, got {gs}") + if gs == 0: + return False + return gs % self.recompute_every_n_steps == 0 + + +@dataclass(frozen=True) +class BucketStats: + """Aggregated per-bucket statistics. + + Attributes: + bucket_id: 0-indexed bucket position (0 = easiest). + num_samples: How many samples contributed to mean_loss / mean_grad_norm. + mean_loss: Average loss across the bucket's recent samples. + mean_grad_norm: Average parameter-grad-norm fingerprint. + """ + + bucket_id: int + num_samples: int + mean_loss: float + mean_grad_norm: float + + +def _coerce_stats(raw: Mapping[int, Mapping[str, float]]) -> List[BucketStats]: + out: List[BucketStats] = [] + for bucket_id, payload in raw.items(): + if isinstance(bucket_id, bool) or not isinstance(bucket_id, int): + raise TypeError( + f"bucket id must be int, got {type(bucket_id).__name__}" + ) + if bucket_id < 0: + raise ValueError(f"bucket id must be >= 0, got {bucket_id}") + if not isinstance(payload, Mapping): + raise TypeError( + f"bucket payload must be Mapping, got {type(payload).__name__}" + ) + num_samples = payload.get("num_samples", 0) + ns = _reject_bool_int("num_samples", num_samples) + if ns < 0 or ns > _MAX_BUCKET_SAMPLES: + raise ValueError( + f"num_samples must be in [0, {_MAX_BUCKET_SAMPLES}], got {ns}" + ) + ml = _reject_bool_float("mean_loss", payload.get("mean_loss", 0.0)) + mg = _reject_bool_float( + "mean_grad_norm", payload.get("mean_grad_norm", 0.0) + ) + if ml < 0.0 or mg < 0.0: + raise ValueError( + "mean_loss / mean_grad_norm must be >= 0 " + f"(got loss={ml}, grad={mg})" + ) + out.append(BucketStats(bucket_id, ns, ml, mg)) + return out + + +def _softmax(values: Sequence[float], temperature: float) -> List[float]: + """Numerically stable softmax.""" + if not values: + return [] + inv_t = 1.0 / temperature + scaled = [v * inv_t for v in values] + m = max(scaled) + exps = [math.exp(s - m) for s in scaled] + total = sum(exps) + if total <= 0.0 or not math.isfinite(total): + # Degenerate input → uniform fallback. + n = len(values) + return [1.0 / n] * n + return [e / total for e in exps] + + +def compute_bucket_weights( + stats: Mapping[int, Mapping[str, float]], + policy: DynamicCurriculumPolicy, +) -> Tuple[float, ...]: + """Return normalised sampler weights per bucket. + + Buckets with no recorded samples fall back to the uniform prior. Buckets + with higher mean loss + grad norm receive more weight; the floor parameter + prevents the easiest bucket from ever dropping below ``policy.floor``. + + Args: + stats: Mapping from ``bucket_id`` to ``{num_samples, mean_loss, + mean_grad_norm}`` payload. + policy: A frozen :class:`DynamicCurriculumPolicy`. + + Returns: + Tuple of ``policy.num_buckets`` floats that sum to 1.0 ± 1e-6. + """ + if not isinstance(policy, DynamicCurriculumPolicy): + raise TypeError( + f"policy must be DynamicCurriculumPolicy, " + f"got {type(policy).__name__}" + ) + if not isinstance(stats, Mapping): + raise TypeError(f"stats must be Mapping, got {type(stats).__name__}") + + coerced = _coerce_stats(stats) + by_id: Dict[int, BucketStats] = {b.bucket_id: b for b in coerced} + + nb = policy.num_buckets + # Build per-bucket scalar = mean_loss + mean_grad_norm. + # Empty buckets get neutral score (median of populated buckets, else 0). + populated = [ + by_id[i].mean_loss + by_id[i].mean_grad_norm + for i in range(nb) + if i in by_id and by_id[i].num_samples > 0 + ] + if populated: + # Median is robust to outliers; matches Axolotl curriculum policy. + srt = sorted(populated) + mid = len(srt) // 2 + neutral = ( + srt[mid] if len(srt) % 2 == 1 else (srt[mid - 1] + srt[mid]) / 2 + ) + else: + # No data — uniform fallback. + return (1.0 / nb,) * nb + + scores: List[float] = [] + for i in range(nb): + b = by_id.get(i) + if b is None or b.num_samples == 0: + scores.append(neutral) + else: + scores.append(b.mean_loss + b.mean_grad_norm) + + weights = _softmax(scores, policy.temperature) + # Water-fill: every bucket gets at least `floor`; remaining + # (1 - nb*floor) is distributed proportionally to the softmax mass. + # The softmax already sums to 1.0 so the water-fill output sums to + # exactly 1.0 (modulo float drift bounded by nb * eps). A subsequent + # renorm `w / sum(w)` is harmful: it can push elements sitting at + # `floor` below the floor when the sum is slightly > 1.0. See + # v0.48.0 Part A code review HIGH #2. + reserved = policy.floor * nb + free_mass = 1.0 - reserved + if free_mass <= 0.0: + return (1.0 / nb,) * nb + total = sum(weights) + if total <= 0.0: + return (1.0 / nb,) * nb + return tuple(policy.floor + free_mass * (w / total) for w in weights) + + +def validate_distributed_curriculum( + enabled: bool, + *, + world_size: int, + rank_coordinated: bool, +) -> None: + """Cross-validator for the distributed footgun. + + When ``curriculum_dynamic=true`` and the launch is multi-rank, the caller + MUST attest that an ``all_reduce`` of per-sample stats is wired (otherwise + each rank computes a divergent weight and the sampler desynchronises). + + Args: + enabled: Resolved value of ``training.curriculum_dynamic``. + world_size: Detected distributed world size (1 for single-process). + rank_coordinated: Caller confirms the all-reduce hook is registered. + + Raises: + ValueError: When multi-rank but no coordination is wired. + """ + if not isinstance(enabled, bool): + raise TypeError("enabled must be bool") + if not enabled: + return + ws = _reject_bool_int("world_size", world_size) + if ws < 1: + raise ValueError(f"world_size must be >= 1, got {ws}") + if not isinstance(rank_coordinated, bool): + raise TypeError("rank_coordinated must be bool") + if ws > 1 and not rank_coordinated: + raise ValueError( + f"curriculum_dynamic=true with world_size={ws} requires an " + "all_reduce hook on per-bucket stats (otherwise each rank " + "diverges). Register the coordination callback before training." + ) + + +def render_curve( + history: Sequence[Mapping[str, float]], + *, + num_buckets: int, + width: int = 60, +) -> str: + """Render a plain-text time-series of bucket weights over training. + + Each row is one recompute step; each column is one bucket. Output uses + ASCII glyphs only (matches the v0.24.3 Windows-Unicode policy). + + Args: + history: Sequence of mappings ``{"step": int, "weights": + [w0, w1, ...]}`` from :func:`compute_bucket_weights`. + num_buckets: Expected bucket arity (validates row shape). + width: Output column width per bucket cell (>= 4). + + Returns: + Multi-line ASCII table suitable for terminal display. + """ + nb = _reject_bool_int("num_buckets", num_buckets) + if nb < 1 or nb > _MAX_BUCKETS: + raise ValueError( + f"num_buckets must be in [1, {_MAX_BUCKETS}], got {nb}" + ) + w = _reject_bool_int("width", width) + if w < 4 or w > 200: + raise ValueError(f"width must be in [4, 200], got {w}") + if not isinstance(history, Sequence) or isinstance(history, (str, bytes)): + raise TypeError("history must be a non-string Sequence") + if len(history) > _MAX_HISTORY_ROWS: + raise ValueError( + f"history has {len(history)} rows; cap is {_MAX_HISTORY_ROWS}" + ) + + if not history: + return "(no curriculum history recorded yet)" + + header = "step".ljust(8) + "".join( + f"B{i}".rjust(w) for i in range(nb) + ) + lines = [header] + for entry in history: + if not isinstance(entry, Mapping): + raise TypeError( + f"history entry must be Mapping, got {type(entry).__name__}" + ) + step = _reject_bool_int("step", entry.get("step", 0)) + weights = entry.get("weights", ()) + if not isinstance(weights, Sequence) or isinstance( + weights, (str, bytes) + ): + raise TypeError("weights must be a non-string Sequence") + if len(weights) != nb: + raise ValueError( + f"weights length {len(weights)} != num_buckets {nb} at " + f"step={step}" + ) + cells = "".join(f"{float(v):>{w}.4f}" for v in weights) + lines.append(str(step).ljust(8) + cells) + return "\n".join(lines) + + +def parse_history_jsonl(rows: Sequence[Mapping]) -> List[Dict[str, object]]: + """Validate and normalise a sequence of curriculum-history rows. + + Used by ``soup runs curriculum-curve `` to load the JSONL written + by the dynamic callback. + + Args: + rows: Sequence of mappings with ``step`` int and ``weights`` list. + + Returns: + List of normalised dicts with keys ``step`` (int) and + ``weights`` (tuple of floats summing to 1.0 ± 1e-3). + """ + out: List[Dict[str, object]] = [] + if not isinstance(rows, Sequence) or isinstance(rows, (str, bytes)): + raise TypeError("rows must be a non-string Sequence") + if len(rows) > _MAX_HISTORY_ROWS: + raise ValueError( + f"history has {len(rows)} rows; cap is {_MAX_HISTORY_ROWS}" + ) + for row in rows: + if not isinstance(row, Mapping): + raise TypeError("history row must be Mapping") + step = _reject_bool_int("step", row.get("step", 0)) + weights = row.get("weights") + if not isinstance(weights, Sequence) or isinstance( + weights, (str, bytes) + ): + raise TypeError("weights must be a non-string Sequence") + floats = [] + for v in weights: + floats.append(_reject_bool_float("weight", v)) + s = sum(floats) + if s <= 0 or abs(s - 1.0) > 1e-3: + raise ValueError( + f"weights at step={step} must sum to 1.0 ± 1e-3, got {s}" + ) + out.append({"step": step, "weights": tuple(floats)}) + return out diff --git a/tests/test_v0480_part_a.py b/tests/test_v0480_part_a.py new file mode 100644 index 0000000..3daac70 --- /dev/null +++ b/tests/test_v0480_part_a.py @@ -0,0 +1,733 @@ +"""Tests for v0.48.0 Part A — Curriculum-Aware Trainer (dynamic re-weighting). + +BETA feature. Covers: +- ``DynamicCurriculumPolicy`` frozen + bounds. +- ``compute_bucket_weights`` math properties + degenerate inputs. +- ``validate_distributed_curriculum`` multi-rank coordination gate. +- ``render_curve`` ASCII renderer. +- ``parse_history_jsonl`` schema check. +- Schema: ``curriculum_dynamic`` requires ``curriculum=True``; mlx + non-SFT + rejection at SoupConfig level. +- ``soup runs curriculum-curve`` CLI smoke. +""" + +from __future__ import annotations + +import json +import os + +import pytest +from pydantic import ValidationError + +from soup_cli.config.schema import DataConfig, SoupConfig, TrainingConfig +from soup_cli.utils.curriculum_dynamic import ( + BucketStats, + DynamicCurriculumPolicy, + compute_bucket_weights, + parse_history_jsonl, + render_curve, + validate_distributed_curriculum, +) + +# ---------- DynamicCurriculumPolicy --------------------------------------- + + +def test_policy_defaults_ok(): + p = DynamicCurriculumPolicy(num_buckets=4) + assert p.recompute_every_n_steps == 50 + assert p.floor == 0.05 + assert p.temperature == 1.0 + + +def test_policy_frozen(): + p = DynamicCurriculumPolicy(num_buckets=4) + with pytest.raises(Exception): + p.num_buckets = 99 # type: ignore[misc] + + +@pytest.mark.parametrize("nb", [0, -1, 21, 100]) +def test_policy_rejects_invalid_buckets(nb): + with pytest.raises(ValueError): + DynamicCurriculumPolicy(num_buckets=nb) + + +def test_policy_rejects_bool_buckets(): + with pytest.raises(ValueError, match="num_buckets must be int, not bool"): + DynamicCurriculumPolicy(num_buckets=True) # type: ignore[arg-type] + + +@pytest.mark.parametrize("rs", [0, -5, 1_000_001]) +def test_policy_rejects_invalid_recompute(rs): + with pytest.raises(ValueError): + DynamicCurriculumPolicy(num_buckets=4, recompute_every_n_steps=rs) + + +def test_policy_floor_must_leave_room(): + # num_buckets=4 → ceiling=0.25. + DynamicCurriculumPolicy(num_buckets=4, floor=0.25) + with pytest.raises(ValueError, match="floor must be in"): + DynamicCurriculumPolicy(num_buckets=4, floor=0.5) + + +def test_policy_floor_rejects_non_finite(): + with pytest.raises(ValueError, match="must be finite"): + DynamicCurriculumPolicy(num_buckets=4, floor=float("nan")) + + +def test_policy_rejects_zero_temperature(): + with pytest.raises(ValueError, match="temperature must be > 0"): + DynamicCurriculumPolicy(num_buckets=4, temperature=0.0) + + +def test_policy_should_recompute(): + p = DynamicCurriculumPolicy(num_buckets=4, recompute_every_n_steps=10) + assert p.should_recompute(0) is False + assert p.should_recompute(5) is False + assert p.should_recompute(10) is True + assert p.should_recompute(20) is True + + +def test_policy_should_recompute_rejects_bool(): + p = DynamicCurriculumPolicy(num_buckets=4) + with pytest.raises(ValueError): + p.should_recompute(True) # type: ignore[arg-type] + + +def test_policy_should_recompute_rejects_negative(): + p = DynamicCurriculumPolicy(num_buckets=4) + with pytest.raises(ValueError): + p.should_recompute(-1) + + +# ---------- compute_bucket_weights ---------------------------------------- + + +def test_compute_uniform_when_no_data(): + p = DynamicCurriculumPolicy(num_buckets=4) + weights = compute_bucket_weights({}, p) + assert len(weights) == 4 + assert all(abs(w - 0.25) < 1e-9 for w in weights) + + +def test_compute_uniform_when_no_samples(): + p = DynamicCurriculumPolicy(num_buckets=3) + stats = { + 0: {"num_samples": 0, "mean_loss": 0.0, "mean_grad_norm": 0.0}, + } + weights = compute_bucket_weights(stats, p) + assert len(weights) == 3 + assert abs(sum(weights) - 1.0) < 1e-6 + + +def test_compute_high_loss_gets_more_weight(): + p = DynamicCurriculumPolicy(num_buckets=2, floor=0.05, temperature=1.0) + stats = { + 0: {"num_samples": 10, "mean_loss": 0.5, "mean_grad_norm": 0.1}, + 1: {"num_samples": 10, "mean_loss": 5.0, "mean_grad_norm": 1.0}, + } + w0, w1 = compute_bucket_weights(stats, p) + assert w1 > w0 + assert abs(w0 + w1 - 1.0) < 1e-6 + + +def test_compute_floor_respected(): + p = DynamicCurriculumPolicy(num_buckets=2, floor=0.2) + stats = { + 0: {"num_samples": 10, "mean_loss": 0.0, "mean_grad_norm": 0.0}, + 1: {"num_samples": 10, "mean_loss": 100.0, "mean_grad_norm": 100.0}, + } + weights = compute_bucket_weights(stats, p) + assert min(weights) >= 0.2 - 1e-9 + + +def test_compute_sums_to_one(): + p = DynamicCurriculumPolicy(num_buckets=5) + stats = { + i: {"num_samples": 5, "mean_loss": float(i), "mean_grad_norm": 0.1} + for i in range(5) + } + weights = compute_bucket_weights(stats, p) + assert abs(sum(weights) - 1.0) < 1e-6 + + +def test_compute_rejects_non_mapping_stats(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(TypeError, match="stats must be Mapping"): + compute_bucket_weights([], p) # type: ignore[arg-type] + + +def test_compute_rejects_non_policy(): + with pytest.raises(TypeError, match="policy must be"): + compute_bucket_weights({}, "policy") # type: ignore[arg-type] + + +def test_compute_rejects_negative_loss(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(ValueError, match="must be >= 0"): + compute_bucket_weights( + {0: {"num_samples": 1, "mean_loss": -1.0, "mean_grad_norm": 0.0}}, + p, + ) + + +def test_compute_rejects_oversize_num_samples(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(ValueError, match="num_samples"): + compute_bucket_weights( + {0: {"num_samples": 10_000_001, "mean_loss": 1.0, + "mean_grad_norm": 1.0}}, p, + ) + + +def test_compute_rejects_bool_bucket_id(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(TypeError, match="bucket id must be int"): + compute_bucket_weights( + {True: {"num_samples": 1, "mean_loss": 1.0, + "mean_grad_norm": 1.0}}, p, + ) + + +def test_compute_rejects_negative_bucket_id(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(ValueError, match="bucket id must be >= 0"): + compute_bucket_weights( + {-1: {"num_samples": 1, "mean_loss": 1.0, + "mean_grad_norm": 1.0}}, p, + ) + + +def test_compute_rejects_non_mapping_payload(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(TypeError, match="bucket payload must be Mapping"): + compute_bucket_weights({0: "bad"}, p) # type: ignore[dict-item] + + +def test_bucket_stats_frozen(): + bs = BucketStats(bucket_id=0, num_samples=1, mean_loss=0.5, + mean_grad_norm=0.1) + with pytest.raises(Exception): + bs.num_samples = 99 # type: ignore[misc] + + +# ---------- validate_distributed_curriculum ------------------------------- + + +def test_distributed_single_rank_ok(): + validate_distributed_curriculum(True, world_size=1, rank_coordinated=False) + + +def test_distributed_multi_rank_uncoordinated_rejected(): + with pytest.raises(ValueError, match="all_reduce hook"): + validate_distributed_curriculum( + True, world_size=4, rank_coordinated=False + ) + + +def test_distributed_multi_rank_coordinated_ok(): + validate_distributed_curriculum(True, world_size=4, rank_coordinated=True) + + +def test_distributed_disabled_short_circuit(): + validate_distributed_curriculum( + False, world_size=4, rank_coordinated=False + ) + + +def test_distributed_rejects_bool_world_size(): + with pytest.raises(ValueError, match="world_size must be int"): + validate_distributed_curriculum( + True, world_size=True, rank_coordinated=True # type: ignore[arg-type] + ) + + +def test_distributed_rejects_non_bool_enabled(): + with pytest.raises(TypeError, match="enabled must be bool"): + validate_distributed_curriculum( + "yes", world_size=1, rank_coordinated=False # type: ignore[arg-type] + ) + + +def test_distributed_rejects_non_bool_coordinated(): + with pytest.raises(TypeError, match="rank_coordinated must be bool"): + validate_distributed_curriculum( + True, world_size=2, rank_coordinated="yes" # type: ignore[arg-type] + ) + + +def test_distributed_rejects_zero_world_size(): + with pytest.raises(ValueError, match="world_size must be >= 1"): + validate_distributed_curriculum( + True, world_size=0, rank_coordinated=True + ) + + +# ---------- render_curve -------------------------------------------------- + + +def test_render_curve_empty_history(): + text = render_curve([], num_buckets=4) + assert "no curriculum history" in text + + +def test_render_curve_basic(): + history = [ + {"step": 100, "weights": [0.25, 0.25, 0.25, 0.25]}, + {"step": 200, "weights": [0.1, 0.2, 0.3, 0.4]}, + ] + text = render_curve(history, num_buckets=4) + assert "step" in text + assert "B0" in text and "B3" in text + assert "100" in text and "200" in text + + +def test_render_curve_rejects_wrong_arity(): + history = [{"step": 1, "weights": [0.5, 0.5]}] + with pytest.raises(ValueError, match="weights length"): + render_curve(history, num_buckets=4) + + +def test_render_curve_rejects_non_sequence(): + with pytest.raises(TypeError, match="history must be"): + render_curve("bogus", num_buckets=4) + + +def test_render_curve_rejects_bool_width(): + with pytest.raises(ValueError, match="width must be int"): + render_curve([], num_buckets=4, width=True) # type: ignore[arg-type] + + +def test_render_curve_rejects_invalid_num_buckets(): + with pytest.raises(ValueError, match="num_buckets must be in"): + render_curve([], num_buckets=999) + + +# ---------- parse_history_jsonl ------------------------------------------- + + +def test_parse_history_jsonl_happy(): + rows = [ + {"step": 50, "weights": [0.25, 0.25, 0.5]}, + {"step": 100, "weights": [0.3, 0.3, 0.4]}, + ] + out = parse_history_jsonl(rows) + assert len(out) == 2 + assert isinstance(out[0]["weights"], tuple) + assert out[0]["step"] == 50 + + +def test_parse_history_jsonl_rejects_non_summing(): + rows = [{"step": 1, "weights": [0.1, 0.2]}] + with pytest.raises(ValueError, match="weights at step"): + parse_history_jsonl(rows) + + +def test_parse_history_jsonl_rejects_non_sequence_weights(): + rows = [{"step": 1, "weights": "abc"}] + with pytest.raises(TypeError, match="weights must be"): + parse_history_jsonl(rows) + + +def test_parse_history_jsonl_rejects_non_mapping_row(): + with pytest.raises(TypeError, match="history row must be Mapping"): + parse_history_jsonl([42]) + + +def test_parse_history_jsonl_rejects_non_sequence(): + with pytest.raises(TypeError, match="rows must be"): + parse_history_jsonl("bogus") # type: ignore[arg-type] + + +# ---------- Schema integration -------------------------------------------- + + +def _base_cfg(**training_overrides): + base = { + "curriculum": True, + "curriculum_dynamic": True, + "curriculum_buckets": 4, + } + base.update(training_overrides) + return SoupConfig( + base="meta-llama/Llama-3.2-1B", + task="sft", + data=DataConfig(train="data.jsonl"), + training=TrainingConfig(**base), + ) + + +def test_schema_accepts_curriculum_dynamic_with_curriculum(): + cfg = _base_cfg() + assert cfg.training.curriculum_dynamic is True + assert cfg.training.curriculum_dynamic_recompute_steps == 50 + + +def test_schema_rejects_dynamic_without_static(): + with pytest.raises(ValidationError, match="curriculum_dynamic requires"): + SoupConfig( + base="meta-llama/Llama-3.2-1B", + task="sft", + data=DataConfig(train="data.jsonl"), + training=TrainingConfig( + curriculum=False, curriculum_dynamic=True + ), + ) + + +def test_schema_rejects_dynamic_on_mlx(): + with pytest.raises(ValidationError, match="mlx backend"): + SoupConfig( + base="mlx-community/Llama-3.2-1B-Instruct-4bit", + task="sft", + backend="mlx", + data=DataConfig(train="data.jsonl"), + training=TrainingConfig( + curriculum=True, + curriculum_dynamic=True, + curriculum_buckets=4, + ), + ) + + +def test_schema_rejects_dynamic_on_dpo(): + with pytest.raises(ValidationError, match="not supported for task"): + SoupConfig( + base="meta-llama/Llama-3.2-1B", + task="dpo", + data=DataConfig(train="data.jsonl", format="dpo"), + training=TrainingConfig( + curriculum=True, + curriculum_dynamic=True, + curriculum_buckets=4, + ), + ) + + +def test_schema_pretrain_accepted(): + cfg = SoupConfig( + base="meta-llama/Llama-3.2-1B", + task="pretrain", + data=DataConfig(train="data.txt", format="plaintext"), + training=TrainingConfig( + curriculum=True, curriculum_dynamic=True, curriculum_buckets=4 + ), + ) + assert cfg.training.curriculum_dynamic is True + + +def test_schema_floor_capped_by_uniform(): + # curriculum_buckets=4 → ceiling=0.25. + with pytest.raises(ValidationError, match="must be <= 1/curriculum_buckets"): + SoupConfig( + base="meta-llama/Llama-3.2-1B", + task="sft", + data=DataConfig(train="data.jsonl"), + training=TrainingConfig( + curriculum=True, + curriculum_dynamic=True, + curriculum_buckets=4, + curriculum_dynamic_floor=0.5, + ), + ) + + +def test_schema_recompute_steps_bounded(): + with pytest.raises(ValidationError): + TrainingConfig(curriculum_dynamic_recompute_steps=0) + with pytest.raises(ValidationError): + TrainingConfig(curriculum_dynamic_recompute_steps=1_000_001) + + +def test_schema_temperature_bounded(): + with pytest.raises(ValidationError): + TrainingConfig(curriculum_dynamic_temperature=0.0) + with pytest.raises(ValidationError): + TrainingConfig(curriculum_dynamic_temperature=101.0) + + +def test_schema_disabled_when_dynamic_off(): + cfg = SoupConfig( + base="meta-llama/Llama-3.2-1B", + task="dpo", + data=DataConfig(train="data.jsonl", format="dpo"), + training=TrainingConfig(curriculum_dynamic=False), + ) + assert cfg.training.curriculum_dynamic is False + + +# ---------- CLI smoke ------------------------------------------------------ + + +def test_curriculum_curve_cli_help(): + from typer.testing import CliRunner + + from soup_cli.cli import app + runner = CliRunner() + result = runner.invoke(app, ["runs", "curriculum-curve", "--help"]) + assert result.exit_code == 0 + assert "curriculum" in result.output.lower() + + +def test_curriculum_curve_run_not_found(tmp_path, monkeypatch): + from typer.testing import CliRunner + + from soup_cli.cli import app + monkeypatch.chdir(tmp_path) + runner = CliRunner() + result = runner.invoke( + app, ["runs", "curriculum-curve", "nonexistent-run-xyz"] + ) + assert result.exit_code == 1 + assert "not found" in result.output.lower() + + +def test_curriculum_curve_missing_history_file(tmp_path, monkeypatch): + """When history file does not exist, exits 1 with friendly message.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + monkeypatch.chdir(tmp_path) + runner = CliRunner() + # Use --history pointing at non-existent file under cwd; the run lookup + # will fail first so this exercises the run-not-found branch. + result = runner.invoke( + app, + [ + "runs", "curriculum-curve", "anything", + "--history", str(tmp_path / "missing.jsonl"), + ], + ) + # Either run-not-found (1) or invalid args; not crash. + assert result.exit_code in (1, 2) + + +def test_curriculum_curve_history_outside_cwd(tmp_path, monkeypatch): + """--history outside cwd is rejected.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + work = tmp_path / "work" + work.mkdir() + elsewhere = tmp_path / "elsewhere.jsonl" + elsewhere.write_text(json.dumps({"step": 1, "weights": [1.0]}) + "\n") + monkeypatch.chdir(work) + + # Need a "run" to exist; patch get_run to return a stub. + import soup_cli.experiment.tracker as et + + class FakeTracker: + def get_run(self, run_id): + return {"run_id": run_id, "output_dir": "."} + + monkeypatch.setattr(et, "ExperimentTracker", FakeTracker) + runner = CliRunner() + result = runner.invoke( + app, + [ + "runs", "curriculum-curve", "stub", + "--history", str(elsewhere), + ], + ) + assert result.exit_code == 2 + assert "outside cwd" in result.output.lower() + + +def test_curriculum_curve_render_jsonl(tmp_path, monkeypatch): + from typer.testing import CliRunner + + import soup_cli.experiment.tracker as et + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + history = tmp_path / "history.jsonl" + history.write_text( + json.dumps({"step": 100, "weights": [0.25, 0.25, 0.25, 0.25]}) + "\n" + + json.dumps({"step": 200, "weights": [0.1, 0.2, 0.3, 0.4]}) + "\n" + ) + + class FakeTracker: + def get_run(self, run_id): + return {"run_id": run_id, "output_dir": "."} + + monkeypatch.setattr(et, "ExperimentTracker", FakeTracker) + runner = CliRunner() + result = runner.invoke( + app, + [ + "runs", "curriculum-curve", "stub", + "--history", str(history), + ], + ) + assert result.exit_code == 0 + assert "B0" in result.output and "B3" in result.output + assert "100" in result.output + + +# ---------- Review-fix coverage ------------------------------------------- + + +def test_parse_history_jsonl_rejects_bool_step(): + rows = [{"step": True, "weights": [0.5, 0.5]}] + with pytest.raises(ValueError, match="step must be int, not bool"): + parse_history_jsonl(rows) + + +def test_render_curve_rejects_bool_num_buckets(): + with pytest.raises(ValueError, match="num_buckets must be int, not bool"): + render_curve([], num_buckets=True) # type: ignore[arg-type] + + +@pytest.mark.parametrize("w", [3, 201]) +def test_render_curve_rejects_out_of_range_width(w): + with pytest.raises(ValueError, match="width must be in"): + render_curve([], num_buckets=4, width=w) + + +def test_render_curve_rejects_non_mapping_entry(): + with pytest.raises(TypeError, match="history entry must be Mapping"): + render_curve([42], num_buckets=4) + + +def test_render_curve_rejects_bool_step_in_entry(): + history = [{"step": True, "weights": [0.5, 0.5]}] + with pytest.raises(ValueError, match="step must be int, not bool"): + render_curve(history, num_buckets=2) + + +def test_render_curve_caps_history_rows(): + huge = [{"step": i, "weights": [0.5, 0.5]} for i in range(101)] + # Patch the cap to exercise the branch without 100k rows. + import soup_cli.utils.curriculum_dynamic as cd + + original = cd._MAX_HISTORY_ROWS + try: + cd._MAX_HISTORY_ROWS = 100 + with pytest.raises(ValueError, match="cap is"): + render_curve(huge, num_buckets=2) + finally: + cd._MAX_HISTORY_ROWS = original + + +def test_parse_history_jsonl_caps_rows(): + huge = [{"step": i, "weights": [0.5, 0.5]} for i in range(101)] + import soup_cli.utils.curriculum_dynamic as cd + + original = cd._MAX_HISTORY_ROWS + try: + cd._MAX_HISTORY_ROWS = 100 + with pytest.raises(ValueError, match="cap is"): + parse_history_jsonl(huge) + finally: + cd._MAX_HISTORY_ROWS = original + + +def test_compute_rejects_negative_grad_norm(): + p = DynamicCurriculumPolicy(num_buckets=2) + with pytest.raises(ValueError, match="must be >= 0"): + compute_bucket_weights( + {0: {"num_samples": 1, "mean_loss": 0.5, "mean_grad_norm": -1.0}}, + p, + ) + + +def test_compute_floor_strict_invariant(): + """All weights must be >= floor (no floor-violating renorm).""" + p = DynamicCurriculumPolicy(num_buckets=4, floor=0.1, temperature=0.01) + stats = { + 0: {"num_samples": 10, "mean_loss": 0.0, "mean_grad_norm": 0.0}, + 1: {"num_samples": 10, "mean_loss": 100.0, "mean_grad_norm": 0.0}, + 2: {"num_samples": 10, "mean_loss": 100.0, "mean_grad_norm": 0.0}, + 3: {"num_samples": 10, "mean_loss": 100.0, "mean_grad_norm": 0.0}, + } + weights = compute_bucket_weights(stats, p) + assert min(weights) >= 0.1 - 1e-12 + assert abs(sum(weights) - 1.0) < 1e-9 + + +def test_should_recompute_boundary_minus_one(): + p = DynamicCurriculumPolicy(num_buckets=4, recompute_every_n_steps=50) + assert p.should_recompute(49) is False + assert p.should_recompute(50) is True + + +def test_curriculum_curve_render_includes_exception_info(tmp_path, monkeypatch): + """Use the recommended `(result.output, repr(result.exception))` assert form.""" + from typer.testing import CliRunner + + import soup_cli.experiment.tracker as et + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + history = tmp_path / "history.jsonl" + history.write_text( + json.dumps({"step": 100, "weights": [0.5, 0.5]}) + "\n" + ) + + class FakeTracker: + def get_run(self, run_id): + return {"run_id": run_id, "output_dir": "."} + + monkeypatch.setattr(et, "ExperimentTracker", FakeTracker) + runner = CliRunner() + result = runner.invoke( + app, ["runs", "curriculum-curve", "stub", "--history", str(history)] + ) + assert result.exit_code == 0, (result.output, repr(result.exception)) + + +def test_curriculum_curve_corrupt_history(tmp_path, monkeypatch): + """Malformed JSONL (non-summing weights) exits 2 with 'history malformed'.""" + from typer.testing import CliRunner + + import soup_cli.experiment.tracker as et + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + history = tmp_path / "h.jsonl" + # Weights don't sum to 1. + history.write_text(json.dumps({"step": 1, "weights": [0.1, 0.2]}) + "\n") + + class FakeTracker: + def get_run(self, run_id): + return {"run_id": run_id, "output_dir": "."} + + monkeypatch.setattr(et, "ExperimentTracker", FakeTracker) + runner = CliRunner() + result = runner.invoke( + app, ["runs", "curriculum-curve", "stub", "--history", str(history)] + ) + assert result.exit_code == 2 + assert "malformed" in result.output.lower() + + +def test_curriculum_curve_rejects_oversize_file(tmp_path, monkeypatch): + """50 MB cap on curriculum history file.""" + from typer.testing import CliRunner + + import soup_cli.experiment.tracker as et + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + history = tmp_path / "big.jsonl" + history.write_text(json.dumps({"step": 1, "weights": [0.5, 0.5]}) + "\n") + + class FakeTracker: + def get_run(self, run_id): + return {"run_id": run_id, "output_dir": "."} + + # Monkey-patch getsize to simulate large file. + monkeypatch.setattr(et, "ExperimentTracker", FakeTracker) + real_getsize = os.path.getsize + + def fake_getsize(p): + if str(p).endswith("big.jsonl"): + return 100 * 1024 * 1024 + return real_getsize(p) + + monkeypatch.setattr(os.path, "getsize", fake_getsize) + runner = CliRunner() + result = runner.invoke( + app, ["runs", "curriculum-curve", "stub", "--history", str(history)] + ) + assert result.exit_code == 2 + assert "50 mb" in result.output.lower()