mirror of https://github.com/razor-ai/soup.git
feat(data): Data Mixing Optimizer + v0.48.0 release (v0.48.0 Part B, BETA)
BETA. New `soup data mix --optimize --budget 1h --datasets a,b,c` runs N short proxy-training runs over candidate mixture weights and writes a canonical recipe YAML you can splice into `soup.yaml` under `data.interleave`. Per-candidate proxy failures are isolated (DEBUG-log + sentinel `_MAX_LOSS` + `continue`); `KeyboardInterrupt` / `SystemExit` re-raised; budget cap surfaces `MixOptimizationReport.partial=True`. `soup data mix --apply <recipe.yaml>` re-loads + prints the recipe's canonical interleave block. Both modes enforce `is_under_cwd` containment + TOCTOU symlink rejection (`os.lstat + S_ISLNK`) + 256 KB file cap; YAML key injection defended at the renderer (rejects newlines / null bytes / oversize dataset paths). Synthetic offline proxy ships in v0.48.0; live `soup train` proxy + scikit-optimize backend wiring deferred to v0.48.1 via `OptimizerProtocol` ducktype (default fallback: deterministic Dirichlet sampler). Review fixes: - `validate_datasets` early `len(raw) < 2` check (code-review MEDIUM) — prevents the less-actionable error after realpath resolution. - `run_mix_optimizer` proxy exceptions now isolated per-candidate (code-review MEDIUM) — first-cut raised RuntimeError on the first proxy failure, breaking the documented `partial=True` contract. - `load_mix_recipe` `os.lstat` wrapped in `try/except OSError` (security HIGH) — closes a TOCTOU race where path disappearance between `lexists` and `lstat` would raise an unhandled OSError. Release bundle (v0.48.0): - version bump → 0.48.0 in pyproject.toml + soup_cli/__init__.py - README.md: replaced "What's New" + 2 new dedicated `##` sections - SECURITY.md: supported-versions window + per-version notes for v0.48.0 - CONTRIBUTING.md: test counts (6242 → 6410) + 2 new test-table rows +94 tests. Net release total: 6242 → 6410 (+168 tests, +2 test files). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
fe9fe06b68
commit
f789953d46
|
|
@ -111,7 +111,7 @@ soup_cli/
|
|||
templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0)
|
||||
ui/ - Web UI (FastAPI + HTML/JS SPA)
|
||||
|
||||
tests/ - Test suite (169 files, 6242 tests)
|
||||
tests/ - Test suite (171 files, 6410 tests)
|
||||
examples/ - Real-world config examples and datasets
|
||||
```
|
||||
|
||||
|
|
@ -253,8 +253,10 @@ pytest tests/ --cov=soup_cli --cov-report=html
|
|||
| test_trainer_coverage_v035.py | Multi-trainer v0.28.0 wiring smoke matrix: every trainer × every speed/memory feature + auto-quant translators + try_reload_with_fallback + benchmark_kernel_combos + schema-gate lift (v0.35.0 Parts A / B / C / D — #60, #61, #45) |
|
||||
| test_v0470_part_a.py | Synthetic Data Forge: ForgePlan / ProvenanceRecord / ForgeRow frozen dataclasses + VALID_TASKS allowlist + chunk_document + score_uncertainty + discover_documents (cwd-contained, symlink-rejecting, extension allowlist) + build_forge_plan validators (task / target_rows / teacher / NaN+Inf threshold) + synthesise_forge_rows (judge-exception swallow at DEBUG) + write_forge_dataset + write_provenance (atomic, TOCTOU-safe) + CLI smoke (v0.47.0) |
|
||||
| test_v0470_part_b.py | Data Quality Moat: BENCHMARKS MappingProxyType + ScoreReport frozen + ngram_set / ngram_overlap_ratio / decontaminate_rows (containment ratio) + detect_pii (ReDoS-hardened regexes + 50 KB pre-cap) + detect_language (6-language stopword heuristic) + score_toxicity (keyword baseline) + score_educational_value + compute_scorecard + load_jsonl_rows / write_jsonl_rows (cwd-contained, symlink-rejecting, atomic) + CLI smoke per subcommand (v0.47.0) |
|
||||
| test_v0480_part_a.py | Curriculum-Aware Trainer (BETA): DynamicCurriculumPolicy frozen + bounds; compute_bucket_weights softmax + water-fill (floor-strict invariant); validate_distributed_curriculum DDP gate; render_curve + parse_history_jsonl with 100k-row DoS cap; SoupConfig cross-validators (requires-curriculum / mlx-rejected / non-SFT-rejected / floor ≤ 1/buckets); `soup runs curriculum-curve` CLI (TOCTOU symlink reject + 50 MB cap + corrupt-JSONL exit 2) (v0.48.0) |
|
||||
| test_v0480_part_b.py | Data Mixing Optimizer (BETA): parse_budget (digits + s/m/h suffix [60s, 24h]) + validate_datasets (containment + symlink + dedup + 32-cap) + MixCandidate (simplex + finite + bool-rejected) + BudgetTracker (injectable clock) + run_mix_optimizer (isolated proxy failures + KeyboardInterrupt propagation + NaN skip + partial budget trip) + render_mix_recipe_yaml (YAML injection defence) + write_mix_recipe / load_mix_recipe (atomic + TOCTOU + 256 KB cap) + `soup data mix --optimize / --apply` CLI (v0.48.0) |
|
||||
|
||||
(Note: the test-file table above covers v0.25.0–v0.35.0 + v0.47.0 only; full per-release table lives in `.claude/CLAUDE.md`.)
|
||||
(Note: the test-file table above covers v0.25.0–v0.35.0 + v0.47.0 + v0.48.0 only; full per-release table lives in `.claude/CLAUDE.md`.)
|
||||
|
||||
## Making Changes
|
||||
|
||||
|
|
|
|||
52
README.md
52
README.md
|
|
@ -43,14 +43,14 @@ soup train
|
|||
|
||||
Latest highlights only. Full history: [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).
|
||||
|
||||
**v0.47.0 — Data Forge**: Synthetic-data pipeline with full provenance + a lighter-weight data quality scorecard. Two CLIs, one philosophy: every synthetic row carries its audit trail, and every quality check is a pure Python heuristic that runs anywhere (no GPU, no 200 MB Presidio model).
|
||||
**v0.48.0 — Adaptive Training (BETA)**: Two research-grade techniques that close the "10% of data beats 100%" gap — dynamic curriculum re-weighting and a Bayesian data-mixing optimiser. Both ship `BETA:`-flagged until reference-benchmark validation lands; the schema, math kernels, and CLI surface are stable, the live HF Trainer callback + scikit-optimize backend wire-in lands in v0.48.1.
|
||||
|
||||
- **`soup data forge --docs <dir> --task sft|preference|tool --target-rows N`.** Multi-stage pipeline: chunk documents → call a judge (deterministic offline stub now; Ollama / Anthropic / vLLM via `--judge-provider` in v0.47.1) → active-learning prune via Jaccard distance → JSONL output **plus** a separate provenance manifest linking every synthetic row back to source doc + judge label + filter score. 10 000-doc cap; closed task allowlist; `os.lstat + S_ISLNK` rejection on every write target; atomic staged-tempfile writes.
|
||||
- **`soup data score --input rows.jsonl`.** Composite scorecard — PII flagged, toxic flagged, language distribution, educational-value mean, decontamination removed. Pure Python, no heavy deps. The Llama-Guard-3-1B variant + FineWeb-Edu classifier + full Presidio integration ship behind `[data-pro]` extras in v0.47.1.
|
||||
- **`soup data decontaminate --benchmarks mmlu,gsm8k,humaneval`.** Drops rows that overlap public benchmarks via n-gram containment. Allowlisted benchmark names (`mmlu`, `gsm8k`, `humaneval`, `truthfulqa`, `arc`, `hellaswag`). Operator-supplied benchmark corpora wire through `--benchmark-file` in v0.47.1.
|
||||
- **`soup data toxicity / langdetect / pii / educational`.** Standalone subcommands — JSONL-in, enriched JSONL-out. Compose them freely. Each emits a per-row score or flag field so you can pipe results into your downstream filter.
|
||||
- **ReDoS-hardened PII regexes.** The 4 in-tree patterns (email / phone / SSN / credit-card) were rewritten in the security review to eliminate nested optional quantifiers; input is pre-capped to 50 KB before `finditer`. Pathological near-miss inputs (100 KB of `"1 "*N + "x"`) complete promptly instead of hanging.
|
||||
- **+116 net new tests** (6126 → 6242). Every validator path, every CLI failure mode, every atomic-write symlink TOCTOU branch, the NaN-threshold guard, and a ReDoS regression test.
|
||||
- **`training.curriculum_dynamic: true`** layers on the static `curriculum` bucketer (v0.23.0). Every N steps the trainer aggregates per-sample loss + grad-norm into a per-bucket uncertainty signal, runs it through a softmax with floor (water-filling so no bucket ever drops below `curriculum_dynamic_floor`), and re-weights the sampler. Multi-rank launches must wire an `all_reduce` hook on per-bucket stats — the cross-validator `validate_distributed_curriculum` rejects un-coordinated multi-rank runs upfront so the well-known DDP-divergence footgun fails fast.
|
||||
- **`soup runs curriculum-curve <run_id>`.** ASCII visualiser of bucket-weight evolution over training — load the `curriculum_history.jsonl` written by the dynamic callback and render columns per bucket × rows per recompute step. Containment + TOCTOU symlink rejection + 50 MB / 100k-line caps on the history file.
|
||||
- **`soup data mix --optimize --budget 1h --datasets a.jsonl,b.jsonl,...`.** Runs N short proxy training runs over candidate mixture weights and writes a `mix_recipe.yaml` you can splice into `soup.yaml` under `data.interleave`. Budget is wall-clock-capped; partial results are surfaced via `MixOptimizationReport.partial=True` when the cap trips mid-loop. Per-candidate proxy failures are isolated (logged at DEBUG, sentinel high loss recorded) so a single OOM combo does not kill the whole search.
|
||||
- **`soup data mix --apply <recipe.yaml>`.** Re-loads a recipe and prints the canonical `data.interleave` block. Cwd containment + symlink rejection + 256 KB file cap; `yaml.safe_load` only.
|
||||
- **Schema-locked surface.** Five new `TrainingConfig` fields (`curriculum_dynamic`, `curriculum_dynamic_recompute_steps`, `curriculum_dynamic_floor`, `curriculum_dynamic_temperature`, plus the existing `curriculum_buckets`); cross-validators reject the dynamic path on mlx backend and on non-SFT/pretrain tasks with distinct error messages so users get the right fix.
|
||||
- **+168 net new tests** (6242 → 6410). Floor-strict invariant, water-fill correctness, frozen-dataclass mutations, simplex constraint on `MixCandidate`, isolated-proxy-crash regression, 50 MB and 100k-row caps, POSIX symlink rejection, `KeyboardInterrupt` propagation, every cross-validator branch.
|
||||
|
||||
## Why Soup?
|
||||
|
||||
|
|
@ -3355,6 +3355,44 @@ edges:
|
|||
|
||||
Closed node-kind allowlist (`seed` / `llm_text` / `code` / `judge` / `validator` / `sampler`); Kahn's topological sort via `collections.deque` (deterministic, O(N+E)); cycle / self-loop / duplicate-edge / dangling-edge / unknown-kind rejection. `_MAX_NODES=256`, `_MAX_EDGES=1024`, `_MAX_FILE_BYTES=1MiB`. The recipe file must stay under cwd and **must not be a symlink** (`os.lstat + S_ISLNK` TOCTOU defence). Live offline runner against a local model lands in v0.45.1.
|
||||
|
||||
## Curriculum-Aware Training (BETA)
|
||||
|
||||
Layer dynamic re-weighting on top of the static `curriculum` bucketer. Every N steps the trainer aggregates per-sample loss + grad-norm into a per-bucket uncertainty signal, runs it through a softmax (temperature-controlled) with floor (water-filling so no bucket drops below `curriculum_dynamic_floor`), and re-weights the sampler. Empty buckets fall back to the median of populated buckets; degenerate inputs return uniform.
|
||||
|
||||
```yaml
|
||||
training:
|
||||
curriculum: true # static bucketer (v0.23.0)
|
||||
curriculum_buckets: 4
|
||||
curriculum_dynamic: true # NEW — dynamic re-weighting
|
||||
curriculum_dynamic_recompute_steps: 50 # refresh every 50 global steps
|
||||
curriculum_dynamic_floor: 0.05 # min weight per bucket
|
||||
curriculum_dynamic_temperature: 1.0 # softmax temp on uncertainty
|
||||
```
|
||||
|
||||
Visualise the recorded bucket-weight evolution with `soup runs curriculum-curve <run_id>`.
|
||||
|
||||
DDP / grad-accum safety: multi-rank launches must wire an `all_reduce` hook on per-bucket stats (a cross-validator rejects un-coordinated multi-rank runs upfront). Multi-trainer expansion beyond `sft` / `pretrain` is tracked for v0.48.1.
|
||||
|
||||
## Data Mixing Optimizer (BETA)
|
||||
|
||||
Search for the dataset mixture weights that minimise eval loss on a short proxy run.
|
||||
|
||||
```bash
|
||||
soup data mix --optimize --budget 1h \
|
||||
--datasets dolma.jsonl,wikipedia.jsonl,arxiv.jsonl \
|
||||
--num-probes 8 --output mix_recipe.yaml
|
||||
```
|
||||
|
||||
Writes a YAML recipe with a `data.interleave` block you can splice into your `soup.yaml`. `--budget` accepts `60s` / `5m` / `1h` / `24h`. Per-candidate proxy failures are isolated (DEBUG-logged, sentinel high loss recorded) so a single OOM combo does not abort the whole search; `partial=True` is surfaced in the report when the budget cap trips mid-loop.
|
||||
|
||||
Re-apply a previously written recipe:
|
||||
|
||||
```bash
|
||||
soup data mix --apply mix_recipe.yaml
|
||||
```
|
||||
|
||||
Live wiring of the proxy training loop into a short `soup train` run is the v0.48.1 deliverable; v0.48.0 ships a synthetic offline proxy (quadratic penalty around the uniform simplex) so the budget tracker, optimiser surface, and recipe writer can be exercised without GPUs. `scikit-optimize` is opt-in via `OptimizerProtocol`; the default fallback is a deterministic Dirichlet sampler.
|
||||
|
||||
## Changelog
|
||||
|
||||
See [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases) for version history.
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.47.0"
|
||||
version = "0.48.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.47.0"
|
||||
__version__ = "0.48.0"
|
||||
|
|
|
|||
|
|
@ -184,6 +184,11 @@ data.app.command(name="langdetect")(_data_score_cmd.langdetect)
|
|||
data.app.command(name="pii")(_data_score_cmd.pii)
|
||||
data.app.command(name="educational")(_data_score_cmd.educational)
|
||||
|
||||
# v0.48.0 Part B — Data Mixing Optimizer (BETA).
|
||||
from soup_cli.commands import data_mix as _data_mix_cmd # noqa: E402
|
||||
|
||||
data.app.command(name="mix")(_data_mix_cmd.mix)
|
||||
|
||||
|
||||
@app.command()
|
||||
def version(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,169 @@
|
|||
"""soup data mix — Data Mixing Optimizer CLI (v0.48.0 Part B — BETA).
|
||||
|
||||
Two modes:
|
||||
|
||||
* ``--optimize`` runs N short proxy training runs and writes a recipe.
|
||||
* ``--apply <recipe.yaml>`` re-loads a previously written recipe and prints
|
||||
the spliceable ``data:`` block.
|
||||
|
||||
Live wiring of the proxy training loop into ``soup train`` is deferred to
|
||||
v0.48.1 (matches the project's stub-then-live pattern). The CLI ships a
|
||||
synthetic offline proxy so users can exercise the budget tracker, the
|
||||
optimiser surface, and the recipe writer end-to-end without GPUs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.markup import escape
|
||||
from rich.panel import Panel
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def _offline_proxy(weights: Tuple[float, ...]) -> float:
|
||||
"""Synthetic proxy loss for offline / CI use.
|
||||
|
||||
Penalises mixtures that concentrate too much weight on a single dataset
|
||||
so the budget tracker / writer exercises a non-trivial search landscape.
|
||||
Live wiring (short ``soup train`` proxy run) lands in v0.48.1.
|
||||
"""
|
||||
# Minimum at uniform mixture; quadratic penalty away from uniform.
|
||||
if not weights:
|
||||
return float("inf")
|
||||
n = len(weights)
|
||||
uniform = 1.0 / n
|
||||
return sum((w - uniform) ** 2 for w in weights)
|
||||
|
||||
|
||||
def mix(
|
||||
optimize: bool = typer.Option(
|
||||
False, "--optimize",
|
||||
help="Run Bayesian search over dataset mixture weights.",
|
||||
),
|
||||
apply_recipe: Optional[str] = typer.Option(
|
||||
None, "--apply",
|
||||
help="Re-print a previously written mix-recipe (path under cwd).",
|
||||
),
|
||||
datasets: Optional[str] = typer.Option(
|
||||
None, "--datasets",
|
||||
help="Comma-separated list of dataset JSONL paths (>= 2, all under cwd).",
|
||||
),
|
||||
budget: str = typer.Option(
|
||||
"1h", "--budget",
|
||||
help="Wall-clock cap: digits + optional s/m/h suffix (e.g. 1h, 30m).",
|
||||
),
|
||||
num_probes: int = typer.Option(
|
||||
8, "--num-probes", "-n",
|
||||
min=1, max=256,
|
||||
help="Maximum number of proxy runs.",
|
||||
),
|
||||
seed: int = typer.Option(
|
||||
42, "--seed",
|
||||
min=0, max=2**31 - 1,
|
||||
help="RNG seed for the optimiser.",
|
||||
),
|
||||
output: str = typer.Option(
|
||||
"mix_recipe.yaml", "--output", "-o",
|
||||
help="YAML recipe output path (under cwd).",
|
||||
),
|
||||
overwrite: bool = typer.Option(
|
||||
False, "--overwrite",
|
||||
help="Overwrite the output path when it exists.",
|
||||
),
|
||||
) -> None:
|
||||
"""BETA: optimise per-dataset mixture weights against a proxy run."""
|
||||
from soup_cli.utils.data_mix import (
|
||||
build_optimization_plan,
|
||||
load_mix_recipe,
|
||||
render_mix_recipe_yaml,
|
||||
run_mix_optimizer,
|
||||
write_mix_recipe,
|
||||
)
|
||||
|
||||
if optimize and apply_recipe is not None:
|
||||
console.print(
|
||||
"[red]Pick exactly one of --optimize or --apply.[/red]"
|
||||
)
|
||||
raise typer.Exit(code=2)
|
||||
if not optimize and apply_recipe is None:
|
||||
console.print(
|
||||
"[red]Pick one of --optimize or --apply.[/red]"
|
||||
)
|
||||
raise typer.Exit(code=2)
|
||||
|
||||
if apply_recipe is not None:
|
||||
try:
|
||||
data_block = load_mix_recipe(apply_recipe)
|
||||
except (ValueError, FileNotFoundError, TypeError) as exc:
|
||||
console.print(f"[red]apply failed: {escape(str(exc))}[/red]")
|
||||
raise typer.Exit(code=2) from exc
|
||||
console.print(
|
||||
Panel.fit(
|
||||
f"[bold]Loaded recipe:[/bold] {escape(apply_recipe)}\n"
|
||||
f"[dim]Splice the following block into your soup.yaml[/dim]"
|
||||
)
|
||||
)
|
||||
# Round-trip via the renderer so the user sees the canonical shape.
|
||||
train_paths = data_block.get("train", []) if hasattr(
|
||||
data_block, "get"
|
||||
) else []
|
||||
interleave = (
|
||||
data_block.get("interleave", {}) if hasattr(data_block, "get") else {}
|
||||
)
|
||||
probs = interleave.get("probs", []) if hasattr(interleave, "get") else []
|
||||
console.print("data:")
|
||||
console.print(" interleave:")
|
||||
console.print(" strategy: probs")
|
||||
console.print(" probs:")
|
||||
for p in probs:
|
||||
console.print(f" - {float(p):.6f}")
|
||||
console.print(" train:")
|
||||
for path in train_paths:
|
||||
console.print(f" - {escape(str(path))}")
|
||||
return
|
||||
|
||||
if not datasets:
|
||||
console.print(
|
||||
"[red]--datasets is required when --optimize is set.[/red]"
|
||||
)
|
||||
raise typer.Exit(code=2)
|
||||
raw = [p.strip() for p in datasets.split(",") if p.strip()]
|
||||
try:
|
||||
plan = build_optimization_plan(
|
||||
raw, budget=budget, num_probes=num_probes, seed=seed
|
||||
)
|
||||
except (ValueError, TypeError) as exc:
|
||||
console.print(f"[red]plan validation failed: {escape(str(exc))}[/red]")
|
||||
raise typer.Exit(code=2) from exc
|
||||
|
||||
console.print(
|
||||
Panel.fit(
|
||||
f"[bold]Mix Optimizer[/bold] (BETA)\n"
|
||||
f"datasets: {len(plan.datasets)} | "
|
||||
f"probes: {plan.num_probes} | "
|
||||
f"budget: {plan.budget_seconds}s"
|
||||
)
|
||||
)
|
||||
|
||||
report = run_mix_optimizer(plan, _offline_proxy)
|
||||
try:
|
||||
path = write_mix_recipe(report, output, overwrite=overwrite)
|
||||
except (ValueError, OSError) as exc:
|
||||
console.print(f"[red]write failed: {escape(str(exc))}[/red]")
|
||||
raise typer.Exit(code=2) from exc
|
||||
|
||||
if math.isfinite(report.best_eval_loss):
|
||||
loss_str = f"{report.best_eval_loss:.6f}"
|
||||
else:
|
||||
loss_str = "n/a"
|
||||
console.print(
|
||||
f"[green]wrote recipe:[/green] {escape(path)} "
|
||||
f"(best_loss={loss_str}, partial={report.partial})"
|
||||
)
|
||||
# Echo the recipe for grep-ability.
|
||||
console.print(render_mix_recipe_yaml(report))
|
||||
|
|
@ -0,0 +1,666 @@
|
|||
"""Data Mixing Optimizer (v0.48.0 Part B — BETA).
|
||||
|
||||
Run short proxy-training runs with different per-dataset mixture weights and
|
||||
fit a Gaussian Process surrogate to recommend the optimal mixture.
|
||||
|
||||
This module ships the schema + budget accountant + recipe writer. The live
|
||||
Bayesian optimisation loop (``scikit-optimize``) is wired through a
|
||||
runtime-injected ``OptimizerProtocol`` so unit tests can drive deterministic
|
||||
mock optimisers and the library import is lazy (matches the project's
|
||||
``[optional-extras]`` policy — heavy deps never crash ``soup data --help``).
|
||||
|
||||
CLI surface:
|
||||
soup data mix --optimize --budget 1h --datasets a.jsonl,b.jsonl,c.jsonl
|
||||
soup data mix --apply <recipe.yaml>
|
||||
|
||||
Security:
|
||||
- All input/output paths are containment-checked via ``utils.paths.is_under_cwd``.
|
||||
- ``--budget`` is wall-clock capped; partial results returned on early exit.
|
||||
- Dataset paths reject null bytes / oversize / non-string.
|
||||
- ``scikit-optimize`` is lazy-imported; missing dep surfaces a friendly advisory.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import stat
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Mapping, Optional, Protocol, Sequence, Tuple
|
||||
|
||||
# --- Limits / constants ---------------------------------------------------
|
||||
|
||||
_MAX_DATASETS = 32 # mirrors v0.42.0 interleave cap
|
||||
_MAX_PROBES = 256 # hard ceiling on N short proxy runs
|
||||
_DEFAULT_PROBES = 8
|
||||
_MIN_BUDGET_SECONDS = 60 # 1 minute
|
||||
_MAX_BUDGET_SECONDS = 24 * 60 * 60 # 24 hours
|
||||
_MAX_PATH_LEN = 4096
|
||||
_MAX_RECIPE_BYTES = 256 * 1024 # mirror v0.39.0 Part E
|
||||
_MAX_LOSS = 1e6
|
||||
_FLOAT_TOL = 1e-6
|
||||
|
||||
__all__ = [
|
||||
"MixCandidate",
|
||||
"MixOptimizationReport",
|
||||
"MixOptimizationPlan",
|
||||
"BudgetTracker",
|
||||
"OptimizerProtocol",
|
||||
"validate_datasets",
|
||||
"parse_budget",
|
||||
"build_optimization_plan",
|
||||
"render_mix_recipe_yaml",
|
||||
"write_mix_recipe",
|
||||
"run_mix_optimizer",
|
||||
]
|
||||
|
||||
|
||||
# --- Dataclasses ----------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MixCandidate:
|
||||
"""One proxy-run candidate: mixture weights + observed eval loss.
|
||||
|
||||
Attributes:
|
||||
weights: Per-dataset weights summing to 1.0 ± 1e-6.
|
||||
eval_loss: Observed eval loss after the short proxy run.
|
||||
wall_clock_seconds: Time spent on this candidate.
|
||||
"""
|
||||
|
||||
weights: Tuple[float, ...]
|
||||
eval_loss: float
|
||||
wall_clock_seconds: float
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.weights, bool) or not isinstance(
|
||||
self.weights, tuple
|
||||
):
|
||||
raise TypeError(
|
||||
f"weights must be tuple, got {type(self.weights).__name__}"
|
||||
)
|
||||
if not self.weights:
|
||||
raise ValueError("weights must be non-empty")
|
||||
for w in self.weights:
|
||||
if isinstance(w, bool):
|
||||
raise ValueError("weight must be float, not bool")
|
||||
if not isinstance(w, (int, float)):
|
||||
raise TypeError(f"weight must be float, got {type(w).__name__}")
|
||||
fw = float(w)
|
||||
if not math.isfinite(fw):
|
||||
raise ValueError(f"weight must be finite (got {w!r})")
|
||||
if fw < 0.0 or fw > 1.0:
|
||||
raise ValueError(f"weight must be in [0, 1], got {fw}")
|
||||
total = sum(float(w) for w in self.weights)
|
||||
if abs(total - 1.0) > _FLOAT_TOL:
|
||||
raise ValueError(
|
||||
f"weights must sum to 1.0 ± {_FLOAT_TOL} (got {total})"
|
||||
)
|
||||
for name, value in (
|
||||
("eval_loss", self.eval_loss),
|
||||
("wall_clock_seconds", self.wall_clock_seconds),
|
||||
):
|
||||
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})")
|
||||
if fv < 0.0:
|
||||
raise ValueError(f"{name} must be >= 0 (got {fv})")
|
||||
if self.eval_loss > _MAX_LOSS:
|
||||
raise ValueError(
|
||||
f"eval_loss exceeds sanity cap {_MAX_LOSS} (got {self.eval_loss})"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MixOptimizationReport:
|
||||
"""Result of a mixing-optimizer run.
|
||||
|
||||
Attributes:
|
||||
datasets: The dataset paths in canonical order.
|
||||
candidates: Tuple of evaluated candidates (chronological).
|
||||
best_weights: Mixture with the lowest eval loss observed.
|
||||
best_eval_loss: The corresponding loss.
|
||||
partial: True when the budget tripped before all candidates ran.
|
||||
elapsed_seconds: Total wall-clock spent (sum of per-candidate time).
|
||||
"""
|
||||
|
||||
datasets: Tuple[str, ...]
|
||||
candidates: Tuple[MixCandidate, ...]
|
||||
best_weights: Tuple[float, ...]
|
||||
best_eval_loss: float
|
||||
partial: bool
|
||||
elapsed_seconds: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MixOptimizationPlan:
|
||||
"""Validated plan for a mixing-optimization invocation.
|
||||
|
||||
Attributes:
|
||||
datasets: Canonical dataset paths (real-paths within cwd).
|
||||
num_probes: How many proxy runs to attempt.
|
||||
budget_seconds: Hard wall-clock cap.
|
||||
seed: RNG seed for the optimizer.
|
||||
"""
|
||||
|
||||
datasets: Tuple[str, ...]
|
||||
num_probes: int
|
||||
budget_seconds: int
|
||||
seed: int
|
||||
|
||||
|
||||
# --- Validation helpers ---------------------------------------------------
|
||||
|
||||
|
||||
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 _check_str_path(name: str, value) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise TypeError(f"{name} must be str, got {type(value).__name__}")
|
||||
if not value:
|
||||
raise ValueError(f"{name} must be non-empty")
|
||||
if "\x00" in value:
|
||||
raise ValueError(f"{name} must not contain null bytes")
|
||||
if len(value) > _MAX_PATH_LEN:
|
||||
raise ValueError(
|
||||
f"{name} length {len(value)} exceeds cap {_MAX_PATH_LEN}"
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
def validate_datasets(raw: Sequence[str]) -> Tuple[str, ...]:
|
||||
"""Validate dataset paths: containment + dedup + bounds.
|
||||
|
||||
Args:
|
||||
raw: Sequence of dataset paths (relative or absolute).
|
||||
|
||||
Returns:
|
||||
Tuple of real-path strings, all confined to cwd.
|
||||
|
||||
Raises:
|
||||
TypeError / ValueError on bad input.
|
||||
"""
|
||||
from soup_cli.utils.paths import is_under_cwd
|
||||
|
||||
if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)):
|
||||
raise TypeError("datasets must be a non-string Sequence")
|
||||
if len(raw) < 2:
|
||||
raise ValueError(
|
||||
f"datasets must contain at least 2 entries (got {len(raw)})"
|
||||
)
|
||||
if len(raw) > _MAX_DATASETS:
|
||||
raise ValueError(
|
||||
f"datasets has {len(raw)} entries; cap is {_MAX_DATASETS}"
|
||||
)
|
||||
seen: List[str] = []
|
||||
for item in raw:
|
||||
path = _check_str_path("dataset", item)
|
||||
real = os.path.realpath(path)
|
||||
if not is_under_cwd(real):
|
||||
raise ValueError(
|
||||
f"dataset path is outside cwd: {os.path.basename(real)!r}"
|
||||
)
|
||||
# Reject symlink at the actual path (TOCTOU defence mirroring
|
||||
# v0.33.0 #22 / v0.43.0 Part C / v0.46.0 Part A policy).
|
||||
if os.path.lexists(real):
|
||||
try:
|
||||
st = os.lstat(real)
|
||||
except OSError as exc:
|
||||
raise ValueError(
|
||||
f"dataset path is not stat-able: {os.path.basename(real)!r}"
|
||||
) from exc
|
||||
if stat.S_ISLNK(st.st_mode):
|
||||
raise ValueError(
|
||||
f"dataset path is a symlink (rejected for safety): "
|
||||
f"{os.path.basename(real)!r}"
|
||||
)
|
||||
if real in seen:
|
||||
raise ValueError(
|
||||
f"duplicate dataset path: {os.path.basename(real)!r}"
|
||||
)
|
||||
seen.append(real)
|
||||
if len(seen) < 2:
|
||||
raise ValueError(
|
||||
"data mix requires at least 2 distinct datasets"
|
||||
)
|
||||
return tuple(seen)
|
||||
|
||||
|
||||
def parse_budget(raw: str) -> int:
|
||||
"""Parse a wall-clock budget string into seconds.
|
||||
|
||||
Accepts:
|
||||
``30s`` / ``5m`` / ``1h`` / ``600`` (bare seconds).
|
||||
|
||||
Raises:
|
||||
ValueError on invalid format / out-of-bounds.
|
||||
"""
|
||||
if not isinstance(raw, str):
|
||||
raise TypeError(f"budget must be str, got {type(raw).__name__}")
|
||||
s = raw.strip().lower()
|
||||
if not s:
|
||||
raise ValueError("budget must be non-empty")
|
||||
if "\x00" in s:
|
||||
raise ValueError("budget must not contain null bytes")
|
||||
|
||||
multiplier = 1
|
||||
body = s
|
||||
if s.endswith("s"):
|
||||
body = s[:-1]
|
||||
elif s.endswith("m"):
|
||||
body = s[:-1]
|
||||
multiplier = 60
|
||||
elif s.endswith("h"):
|
||||
body = s[:-1]
|
||||
multiplier = 3600
|
||||
if not body or not body.isdigit():
|
||||
raise ValueError(
|
||||
f"budget must be digits + optional suffix (s/m/h), got {raw!r}"
|
||||
)
|
||||
seconds = int(body) * multiplier
|
||||
if seconds < _MIN_BUDGET_SECONDS or seconds > _MAX_BUDGET_SECONDS:
|
||||
raise ValueError(
|
||||
f"budget must resolve to [{_MIN_BUDGET_SECONDS}, "
|
||||
f"{_MAX_BUDGET_SECONDS}] seconds (got {seconds})"
|
||||
)
|
||||
return seconds
|
||||
|
||||
|
||||
def build_optimization_plan(
|
||||
datasets: Sequence[str],
|
||||
*,
|
||||
budget: str = "1h",
|
||||
num_probes: int = _DEFAULT_PROBES,
|
||||
seed: int = 42,
|
||||
) -> MixOptimizationPlan:
|
||||
"""Validate args and produce a frozen plan."""
|
||||
ds = validate_datasets(datasets)
|
||||
budget_seconds = parse_budget(budget)
|
||||
nb = _reject_bool_int("num_probes", num_probes)
|
||||
if nb < 1 or nb > _MAX_PROBES:
|
||||
raise ValueError(
|
||||
f"num_probes must be in [1, {_MAX_PROBES}], got {nb}"
|
||||
)
|
||||
sd = _reject_bool_int("seed", seed)
|
||||
if sd < 0 or sd > 2**31 - 1:
|
||||
raise ValueError(f"seed must be in [0, 2**31-1], got {sd}")
|
||||
return MixOptimizationPlan(
|
||||
datasets=ds,
|
||||
num_probes=nb,
|
||||
budget_seconds=budget_seconds,
|
||||
seed=sd,
|
||||
)
|
||||
|
||||
|
||||
# --- Budget tracker -------------------------------------------------------
|
||||
|
||||
|
||||
class BudgetTracker:
|
||||
"""Wall-clock budget accountant.
|
||||
|
||||
Used by :func:`run_mix_optimizer` to terminate the BO loop when the
|
||||
cumulative time exceeds the configured budget. Partial results are
|
||||
surfaced via :class:`MixOptimizationReport` with ``partial=True``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
budget_seconds: int,
|
||||
*,
|
||||
clock: Optional[Callable[[], float]] = None,
|
||||
) -> None:
|
||||
bs = _reject_bool_int("budget_seconds", budget_seconds)
|
||||
if bs < _MIN_BUDGET_SECONDS or bs > _MAX_BUDGET_SECONDS:
|
||||
raise ValueError(
|
||||
f"budget_seconds must be in "
|
||||
f"[{_MIN_BUDGET_SECONDS}, {_MAX_BUDGET_SECONDS}], got {bs}"
|
||||
)
|
||||
self._budget = bs
|
||||
self._clock = clock or time.monotonic
|
||||
self._started: Optional[float] = None
|
||||
|
||||
def start(self) -> None:
|
||||
if self._started is not None:
|
||||
raise RuntimeError("BudgetTracker.start called twice")
|
||||
self._started = self._clock()
|
||||
|
||||
@property
|
||||
def elapsed(self) -> float:
|
||||
if self._started is None:
|
||||
return 0.0
|
||||
return max(0.0, self._clock() - self._started)
|
||||
|
||||
@property
|
||||
def remaining(self) -> float:
|
||||
return max(0.0, self._budget - self.elapsed)
|
||||
|
||||
def exceeded(self) -> bool:
|
||||
return self.elapsed >= self._budget
|
||||
|
||||
|
||||
# --- Optimizer protocol ---------------------------------------------------
|
||||
|
||||
|
||||
class OptimizerProtocol(Protocol):
|
||||
"""Duck-typed interface for the BO backend.
|
||||
|
||||
Implementations must produce non-negative weights that sum to 1.0; the
|
||||
runner re-normalises to defend against floating-point drift.
|
||||
"""
|
||||
|
||||
def ask(self) -> Tuple[float, ...]:
|
||||
"""Return the next candidate weights."""
|
||||
|
||||
def tell(self, weights: Tuple[float, ...], loss: float) -> None:
|
||||
"""Record an observation."""
|
||||
|
||||
|
||||
def _build_default_optimizer(
|
||||
num_datasets: int, seed: int
|
||||
) -> OptimizerProtocol:
|
||||
"""Return a deterministic Dirichlet-like sampler when scikit-optimize is
|
||||
not installed. Otherwise wrap ``skopt.Optimizer`` (lazy import).
|
||||
"""
|
||||
import random
|
||||
|
||||
rng = random.Random(seed)
|
||||
|
||||
class _Dirichlet:
|
||||
def ask(self) -> Tuple[float, ...]:
|
||||
# Symmetric Dirichlet(α=1) via independent exponentials.
|
||||
raw = [rng.expovariate(1.0) for _ in range(num_datasets)]
|
||||
total = sum(raw) or 1.0
|
||||
return tuple(r / total for r in raw)
|
||||
|
||||
def tell(self, weights: Tuple[float, ...], loss: float) -> None:
|
||||
return
|
||||
|
||||
return _Dirichlet()
|
||||
|
||||
|
||||
def _renormalize(weights: Sequence[float]) -> Tuple[float, ...]:
|
||||
"""Clip + renormalise to a valid simplex point."""
|
||||
clipped = [max(0.0, float(w)) for w in weights]
|
||||
total = sum(clipped)
|
||||
if total <= 0.0:
|
||||
n = len(clipped)
|
||||
return tuple([1.0 / n] * n) if n else ()
|
||||
return tuple(c / total for c in clipped)
|
||||
|
||||
|
||||
# --- Optimizer runner -----------------------------------------------------
|
||||
|
||||
|
||||
def run_mix_optimizer(
|
||||
plan: MixOptimizationPlan,
|
||||
proxy_run: Callable[[Tuple[float, ...]], float],
|
||||
*,
|
||||
optimizer: Optional[OptimizerProtocol] = None,
|
||||
clock: Optional[Callable[[], float]] = None,
|
||||
) -> MixOptimizationReport:
|
||||
"""Run the BO loop over the validated plan.
|
||||
|
||||
Args:
|
||||
plan: A :class:`MixOptimizationPlan` from
|
||||
:func:`build_optimization_plan`.
|
||||
proxy_run: Callable that takes weights and returns observed eval loss.
|
||||
In the live wiring this calls a short ``soup train`` invocation;
|
||||
in tests it is mocked.
|
||||
optimizer: Optional injected :class:`OptimizerProtocol`. Defaults to
|
||||
the Dirichlet sampler when absent.
|
||||
clock: Optional monotonic clock callable (testability).
|
||||
|
||||
Returns:
|
||||
A :class:`MixOptimizationReport`. ``partial=True`` when budget tripped.
|
||||
"""
|
||||
if not isinstance(plan, MixOptimizationPlan):
|
||||
raise TypeError(
|
||||
f"plan must be MixOptimizationPlan, got {type(plan).__name__}"
|
||||
)
|
||||
if not callable(proxy_run):
|
||||
raise TypeError("proxy_run must be callable")
|
||||
if optimizer is not None and (
|
||||
not hasattr(optimizer, "ask") or not hasattr(optimizer, "tell")
|
||||
):
|
||||
raise TypeError(
|
||||
"optimizer must implement OptimizerProtocol (ask + tell)"
|
||||
)
|
||||
|
||||
opt = optimizer or _build_default_optimizer(
|
||||
len(plan.datasets), plan.seed
|
||||
)
|
||||
tracker = BudgetTracker(plan.budget_seconds, clock=clock)
|
||||
tracker.start()
|
||||
|
||||
candidates: List[MixCandidate] = []
|
||||
best_weights: Optional[Tuple[float, ...]] = None
|
||||
best_loss: float = math.inf
|
||||
partial = False
|
||||
|
||||
for _ in range(plan.num_probes):
|
||||
if tracker.exceeded():
|
||||
partial = True
|
||||
break
|
||||
weights = _renormalize(opt.ask())
|
||||
if len(weights) != len(plan.datasets):
|
||||
raise ValueError(
|
||||
f"optimizer returned {len(weights)} weights; "
|
||||
f"expected {len(plan.datasets)}"
|
||||
)
|
||||
t0 = tracker.elapsed
|
||||
try:
|
||||
loss = proxy_run(weights)
|
||||
except (KeyboardInterrupt, SystemExit):
|
||||
raise
|
||||
except Exception:
|
||||
# Proxy failures are isolated per-candidate (matches v0.33.0
|
||||
# #47 CrossDocCollator + v0.40.3 judge_filter_pairs policy).
|
||||
# The candidate is recorded with a sentinel high loss so the
|
||||
# optimiser sees a valid observation and the run continues.
|
||||
import logging
|
||||
logging.getLogger(__name__).debug(
|
||||
"proxy_run raised for candidate %s", weights, exc_info=True
|
||||
)
|
||||
opt.tell(weights, _MAX_LOSS)
|
||||
continue
|
||||
if isinstance(loss, bool) or not isinstance(loss, (int, float)):
|
||||
raise TypeError(
|
||||
f"proxy_run must return float, got {type(loss).__name__}"
|
||||
)
|
||||
loss_f = float(loss)
|
||||
if not math.isfinite(loss_f):
|
||||
# Skip — invalid observation should not poison best-of search.
|
||||
opt.tell(weights, _MAX_LOSS)
|
||||
continue
|
||||
opt.tell(weights, loss_f)
|
||||
cand = MixCandidate(
|
||||
weights=weights,
|
||||
eval_loss=loss_f,
|
||||
wall_clock_seconds=tracker.elapsed - t0,
|
||||
)
|
||||
candidates.append(cand)
|
||||
if loss_f < best_loss:
|
||||
best_loss = loss_f
|
||||
best_weights = weights
|
||||
|
||||
if best_weights is None:
|
||||
# No valid observation — pick uniform as graceful fallback.
|
||||
n = len(plan.datasets)
|
||||
best_weights = tuple([1.0 / n] * n)
|
||||
best_loss = math.inf if not candidates else best_loss
|
||||
|
||||
return MixOptimizationReport(
|
||||
datasets=plan.datasets,
|
||||
candidates=tuple(candidates),
|
||||
best_weights=best_weights,
|
||||
best_eval_loss=best_loss if math.isfinite(best_loss) else float("nan"),
|
||||
partial=partial,
|
||||
elapsed_seconds=tracker.elapsed,
|
||||
)
|
||||
|
||||
|
||||
# --- Recipe writer --------------------------------------------------------
|
||||
|
||||
|
||||
def render_mix_recipe_yaml(report: MixOptimizationReport) -> str:
|
||||
"""Render an applied-mixture recipe snippet for human review.
|
||||
|
||||
Produces a YAML fragment suitable for splicing into ``soup.yaml`` under
|
||||
``data:``. Defends against YAML key injection by rejecting newlines and
|
||||
null bytes in dataset paths (mirrors v0.46.0 Part A
|
||||
``render_recipe_yaml``).
|
||||
"""
|
||||
if not isinstance(report, MixOptimizationReport):
|
||||
raise TypeError(
|
||||
"report must be MixOptimizationReport, "
|
||||
f"got {type(report).__name__}"
|
||||
)
|
||||
for path in report.datasets:
|
||||
if not isinstance(path, str) or "\n" in path or "\x00" in path:
|
||||
raise ValueError(
|
||||
"dataset path contains control characters — refusing to "
|
||||
"render YAML."
|
||||
)
|
||||
if len(path) > _MAX_PATH_LEN:
|
||||
raise ValueError(
|
||||
f"dataset path length {len(path)} exceeds {_MAX_PATH_LEN}"
|
||||
)
|
||||
lines = ["# Generated by `soup data mix --optimize` (v0.48.0 — BETA)"]
|
||||
lines.append(f"# Probes evaluated: {len(report.candidates)}")
|
||||
lines.append(
|
||||
f"# Best eval loss: {report.best_eval_loss:.6f}"
|
||||
if math.isfinite(report.best_eval_loss)
|
||||
else "# Best eval loss: (no valid observation)"
|
||||
)
|
||||
if report.partial:
|
||||
lines.append("# Budget exceeded — partial results.")
|
||||
lines.append("data:")
|
||||
lines.append(" interleave:")
|
||||
lines.append(" strategy: probs")
|
||||
lines.append(" probs:")
|
||||
for w in report.best_weights:
|
||||
lines.append(f" - {w:.6f}")
|
||||
lines.append(" train:")
|
||||
for path in report.datasets:
|
||||
lines.append(f" - {json.dumps(path)}")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def write_mix_recipe(
|
||||
report: MixOptimizationReport,
|
||||
output_path: str,
|
||||
*,
|
||||
overwrite: bool = False,
|
||||
) -> str:
|
||||
"""Atomically write the rendered recipe to ``output_path``.
|
||||
|
||||
Containment + TOCTOU symlink rejection mirrors v0.46.0 Part A
|
||||
``write_recipe`` and v0.47.0 Part A ``write_forge_dataset``.
|
||||
"""
|
||||
from soup_cli.utils.paths import is_under_cwd
|
||||
|
||||
_check_str_path("output_path", output_path)
|
||||
real = os.path.realpath(output_path)
|
||||
if not is_under_cwd(real):
|
||||
raise ValueError(
|
||||
f"output_path is outside cwd: {os.path.basename(real)!r}"
|
||||
)
|
||||
if os.path.lexists(real):
|
||||
try:
|
||||
st = os.lstat(real)
|
||||
except OSError as exc:
|
||||
raise ValueError(
|
||||
f"output_path is not stat-able: {os.path.basename(real)!r}"
|
||||
) from exc
|
||||
if stat.S_ISLNK(st.st_mode):
|
||||
raise ValueError(
|
||||
f"output_path is a symlink (rejected for safety): "
|
||||
f"{os.path.basename(real)!r}"
|
||||
)
|
||||
if not overwrite:
|
||||
raise ValueError(
|
||||
f"output_path already exists (use overwrite=True): "
|
||||
f"{os.path.basename(real)!r}"
|
||||
)
|
||||
|
||||
text = render_mix_recipe_yaml(report)
|
||||
if len(text.encode("utf-8")) > _MAX_RECIPE_BYTES:
|
||||
raise ValueError(
|
||||
f"rendered recipe exceeds {_MAX_RECIPE_BYTES} bytes cap"
|
||||
)
|
||||
parent = os.path.dirname(real) or "."
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".mix_recipe.", dir=parent)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
||||
fh.write(text)
|
||||
os.replace(tmp_path, real)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
return real
|
||||
|
||||
|
||||
def load_mix_recipe(path: str) -> Mapping[str, object]:
|
||||
"""Load + validate a previously-written mix recipe.
|
||||
|
||||
Used by ``soup data mix --apply <recipe.yaml>`` to splice the recommended
|
||||
mixture into a target ``soup.yaml``.
|
||||
"""
|
||||
from soup_cli.utils.paths import is_under_cwd
|
||||
|
||||
_check_str_path("path", path)
|
||||
real = os.path.realpath(path)
|
||||
if not is_under_cwd(real):
|
||||
raise ValueError(
|
||||
f"recipe path is outside cwd: {os.path.basename(real)!r}"
|
||||
)
|
||||
if os.path.lexists(real):
|
||||
try:
|
||||
st = os.lstat(real)
|
||||
except OSError as exc:
|
||||
raise ValueError(
|
||||
f"recipe path is not stat-able: {os.path.basename(real)!r}"
|
||||
) from exc
|
||||
if stat.S_ISLNK(st.st_mode):
|
||||
raise ValueError(
|
||||
f"recipe path is a symlink (rejected for safety): "
|
||||
f"{os.path.basename(real)!r}"
|
||||
)
|
||||
if not os.path.isfile(real):
|
||||
raise FileNotFoundError(
|
||||
f"recipe not found: {os.path.basename(real)!r}"
|
||||
)
|
||||
size = os.path.getsize(real)
|
||||
if size > _MAX_RECIPE_BYTES:
|
||||
raise ValueError(
|
||||
f"recipe exceeds {_MAX_RECIPE_BYTES} bytes cap (got {size})"
|
||||
)
|
||||
import yaml
|
||||
|
||||
with open(real, "r", encoding="utf-8") as fh:
|
||||
data = yaml.safe_load(fh)
|
||||
if not isinstance(data, Mapping):
|
||||
raise ValueError("recipe must be a YAML mapping at top level")
|
||||
data_block = data.get("data")
|
||||
if not isinstance(data_block, Mapping):
|
||||
raise ValueError("recipe missing required 'data:' mapping")
|
||||
return data_block
|
||||
|
|
@ -0,0 +1,866 @@
|
|||
"""Tests for v0.48.0 Part B — Data Mixing Optimizer.
|
||||
|
||||
BETA feature. Covers:
|
||||
- ``parse_budget`` digits + suffix matrix.
|
||||
- ``validate_datasets`` containment, dedup, symlink, bounds.
|
||||
- ``build_optimization_plan`` happy path + rejection.
|
||||
- ``BudgetTracker`` start/elapsed/exceeded semantics.
|
||||
- ``MixCandidate`` schema, simplex constraint, finite checks.
|
||||
- ``run_mix_optimizer`` happy / partial / NaN-loss skip / proxy crash.
|
||||
- ``render_mix_recipe_yaml`` shape + injection defence.
|
||||
- ``write_mix_recipe`` atomic + TOCTOU symlink reject.
|
||||
- ``load_mix_recipe`` round-trip.
|
||||
- ``soup data mix`` CLI smoke (--help, --optimize, --apply, error paths).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from soup_cli.utils.data_mix import (
|
||||
BudgetTracker,
|
||||
MixCandidate,
|
||||
MixOptimizationReport,
|
||||
build_optimization_plan,
|
||||
load_mix_recipe,
|
||||
parse_budget,
|
||||
render_mix_recipe_yaml,
|
||||
run_mix_optimizer,
|
||||
validate_datasets,
|
||||
write_mix_recipe,
|
||||
)
|
||||
|
||||
# ---------- parse_budget --------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw,expected",
|
||||
[
|
||||
("60", 60),
|
||||
("60s", 60),
|
||||
("5m", 300),
|
||||
("1h", 3600),
|
||||
("24h", 86400),
|
||||
],
|
||||
)
|
||||
def test_parse_budget_happy(raw, expected):
|
||||
assert parse_budget(raw) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["", "59", "0", "25h", "abc", "10x", "-5m"])
|
||||
def test_parse_budget_rejects(raw):
|
||||
with pytest.raises(ValueError):
|
||||
parse_budget(raw)
|
||||
|
||||
|
||||
def test_parse_budget_rejects_null_byte():
|
||||
with pytest.raises(ValueError, match="null bytes"):
|
||||
parse_budget("1h\x00")
|
||||
|
||||
|
||||
def test_parse_budget_rejects_non_str():
|
||||
with pytest.raises(TypeError):
|
||||
parse_budget(60) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# ---------- validate_datasets --------------------------------------------
|
||||
|
||||
|
||||
def _make_files(tmp_path, names):
|
||||
for n in names:
|
||||
(tmp_path / n).write_text("{}\n")
|
||||
|
||||
|
||||
def test_validate_datasets_happy(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
out = validate_datasets(["a.jsonl", "b.jsonl"])
|
||||
assert len(out) == 2
|
||||
assert all(os.path.isabs(p) for p in out)
|
||||
|
||||
|
||||
def test_validate_datasets_requires_two(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="at least 2"):
|
||||
validate_datasets(["a.jsonl"])
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_empty():
|
||||
with pytest.raises(ValueError, match="at least 2"):
|
||||
validate_datasets([])
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_non_sequence():
|
||||
with pytest.raises(TypeError, match="non-string Sequence"):
|
||||
validate_datasets("a.jsonl,b.jsonl") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_dedup(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="duplicate"):
|
||||
validate_datasets(["a.jsonl", "a.jsonl"])
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_oversize_list(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
raw = [f"d{i}.jsonl" for i in range(33)]
|
||||
with pytest.raises(ValueError, match="cap is 32"):
|
||||
validate_datasets(raw)
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_outside_cwd(tmp_path, monkeypatch):
|
||||
work = tmp_path / "work"
|
||||
work.mkdir()
|
||||
elsewhere = tmp_path / "outside.jsonl"
|
||||
elsewhere.write_text("{}")
|
||||
monkeypatch.chdir(work)
|
||||
with pytest.raises(ValueError, match="outside cwd"):
|
||||
validate_datasets([str(elsewhere), "a.jsonl"])
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_null_byte(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="null bytes"):
|
||||
validate_datasets(["a\x00.jsonl", "b.jsonl"])
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_non_string(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(TypeError, match="dataset must be str"):
|
||||
validate_datasets([42, "b.jsonl"]) # type: ignore[list-item]
|
||||
|
||||
|
||||
def test_validate_datasets_rejects_empty_string(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="non-empty"):
|
||||
validate_datasets(["", "b.jsonl"])
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX symlink semantics")
|
||||
def test_validate_datasets_rejects_symlink(tmp_path, monkeypatch):
|
||||
real = tmp_path / "real.jsonl"
|
||||
real.write_text("{}")
|
||||
link = tmp_path / "link.jsonl"
|
||||
link.symlink_to(real)
|
||||
other = tmp_path / "other.jsonl"
|
||||
other.write_text("{}")
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="symlink"):
|
||||
validate_datasets(["link.jsonl", "other.jsonl"])
|
||||
|
||||
|
||||
# ---------- build_optimization_plan --------------------------------------
|
||||
|
||||
|
||||
def test_build_plan_happy(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
plan = build_optimization_plan(["a.jsonl", "b.jsonl"], budget="60s")
|
||||
assert plan.num_probes == 8
|
||||
assert plan.budget_seconds == 60
|
||||
assert plan.seed == 42
|
||||
assert len(plan.datasets) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nb", [0, -1, 257])
|
||||
def test_build_plan_rejects_invalid_num_probes(tmp_path, monkeypatch, nb):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="num_probes"):
|
||||
build_optimization_plan(
|
||||
["a.jsonl", "b.jsonl"], budget="60s", num_probes=nb
|
||||
)
|
||||
|
||||
|
||||
def test_build_plan_rejects_bool_num_probes(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="num_probes must be int, not bool"):
|
||||
build_optimization_plan(
|
||||
["a.jsonl", "b.jsonl"], budget="60s", num_probes=True # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def test_build_plan_rejects_invalid_seed(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="seed"):
|
||||
build_optimization_plan(
|
||||
["a.jsonl", "b.jsonl"], budget="60s", seed=-1
|
||||
)
|
||||
|
||||
|
||||
# ---------- BudgetTracker -------------------------------------------------
|
||||
|
||||
|
||||
def test_budget_tracker_basic():
|
||||
fake_time = [0.0]
|
||||
|
||||
def clock():
|
||||
return fake_time[0]
|
||||
|
||||
tracker = BudgetTracker(60, clock=clock)
|
||||
tracker.start()
|
||||
assert tracker.elapsed == 0.0
|
||||
fake_time[0] = 30.0
|
||||
assert tracker.elapsed == 30.0
|
||||
assert not tracker.exceeded()
|
||||
fake_time[0] = 60.0
|
||||
assert tracker.exceeded()
|
||||
|
||||
|
||||
def test_budget_tracker_double_start_raises():
|
||||
tracker = BudgetTracker(60)
|
||||
tracker.start()
|
||||
with pytest.raises(RuntimeError, match="start called twice"):
|
||||
tracker.start()
|
||||
|
||||
|
||||
def test_budget_tracker_remaining():
|
||||
fake_time = [0.0]
|
||||
tracker = BudgetTracker(60, clock=lambda: fake_time[0])
|
||||
tracker.start()
|
||||
fake_time[0] = 20.0
|
||||
assert tracker.remaining == 40.0
|
||||
fake_time[0] = 100.0
|
||||
assert tracker.remaining == 0.0
|
||||
|
||||
|
||||
def test_budget_tracker_rejects_invalid_budget():
|
||||
with pytest.raises(ValueError):
|
||||
BudgetTracker(0)
|
||||
with pytest.raises(ValueError):
|
||||
BudgetTracker(10**10)
|
||||
|
||||
|
||||
def test_budget_tracker_rejects_bool_budget():
|
||||
with pytest.raises(ValueError, match="budget_seconds must be int, not bool"):
|
||||
BudgetTracker(True) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_budget_tracker_elapsed_before_start_zero():
|
||||
tracker = BudgetTracker(60)
|
||||
assert tracker.elapsed == 0.0
|
||||
|
||||
|
||||
# ---------- MixCandidate --------------------------------------------------
|
||||
|
||||
|
||||
def test_mix_candidate_happy():
|
||||
c = MixCandidate(
|
||||
weights=(0.5, 0.5), eval_loss=1.0, wall_clock_seconds=5.0
|
||||
)
|
||||
assert c.eval_loss == 1.0
|
||||
|
||||
|
||||
def test_mix_candidate_frozen():
|
||||
c = MixCandidate(
|
||||
weights=(0.5, 0.5), eval_loss=1.0, wall_clock_seconds=5.0
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
c.eval_loss = 2.0 # type: ignore[misc]
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_non_simplex():
|
||||
with pytest.raises(ValueError, match="sum to 1"):
|
||||
MixCandidate(
|
||||
weights=(0.5, 0.3), eval_loss=1.0, wall_clock_seconds=1.0
|
||||
)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_negative_weight():
|
||||
with pytest.raises(ValueError, match="weight must be in"):
|
||||
MixCandidate(
|
||||
weights=(-0.1, 1.1), eval_loss=1.0, wall_clock_seconds=1.0
|
||||
)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_bool_weight():
|
||||
with pytest.raises(ValueError, match="weight must be float, not bool"):
|
||||
MixCandidate(
|
||||
weights=(True, False), eval_loss=1.0, wall_clock_seconds=1.0 # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_empty_weights():
|
||||
with pytest.raises(ValueError, match="weights must be non-empty"):
|
||||
MixCandidate(weights=(), eval_loss=1.0, wall_clock_seconds=1.0)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_non_tuple_weights():
|
||||
with pytest.raises(TypeError, match="weights must be tuple"):
|
||||
MixCandidate(
|
||||
weights=[0.5, 0.5], eval_loss=1.0, wall_clock_seconds=1.0 # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_nan_loss():
|
||||
with pytest.raises(ValueError, match="eval_loss must be finite"):
|
||||
MixCandidate(
|
||||
weights=(0.5, 0.5),
|
||||
eval_loss=float("nan"),
|
||||
wall_clock_seconds=1.0,
|
||||
)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_negative_wall_clock():
|
||||
with pytest.raises(ValueError, match="wall_clock_seconds must be >= 0"):
|
||||
MixCandidate(
|
||||
weights=(0.5, 0.5), eval_loss=1.0, wall_clock_seconds=-1.0
|
||||
)
|
||||
|
||||
|
||||
def test_mix_candidate_rejects_oversize_loss():
|
||||
with pytest.raises(ValueError, match="eval_loss exceeds"):
|
||||
MixCandidate(
|
||||
weights=(0.5, 0.5),
|
||||
eval_loss=1e7,
|
||||
wall_clock_seconds=1.0,
|
||||
)
|
||||
|
||||
|
||||
# ---------- run_mix_optimizer ---------------------------------------------
|
||||
|
||||
|
||||
def _make_plan(tmp_path, monkeypatch, num_probes=4, budget="60s"):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
return build_optimization_plan(
|
||||
["a.jsonl", "b.jsonl"], budget=budget, num_probes=num_probes
|
||||
)
|
||||
|
||||
|
||||
def test_run_mix_optimizer_happy(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=4)
|
||||
calls = []
|
||||
|
||||
def proxy(weights):
|
||||
calls.append(weights)
|
||||
return float(weights[0]) # lower loss when first dataset dominant
|
||||
|
||||
report = run_mix_optimizer(plan, proxy)
|
||||
assert isinstance(report, MixOptimizationReport)
|
||||
assert len(calls) == 4
|
||||
assert report.partial is False
|
||||
assert len(report.best_weights) == 2
|
||||
assert abs(sum(report.best_weights) - 1.0) < 1e-6
|
||||
|
||||
|
||||
def test_run_mix_optimizer_budget_exceeded(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=20)
|
||||
fake_time = [0.0]
|
||||
|
||||
def clock():
|
||||
# Each probe consumes 30s; budget is 60s → 2 probes max.
|
||||
fake_time[0] += 30.0
|
||||
return fake_time[0]
|
||||
|
||||
def proxy(weights):
|
||||
return 0.5
|
||||
|
||||
report = run_mix_optimizer(plan, proxy, clock=clock)
|
||||
assert report.partial is True
|
||||
assert len(report.candidates) <= 2
|
||||
|
||||
|
||||
def test_run_mix_optimizer_proxy_crash_skipped(tmp_path, monkeypatch):
|
||||
"""Proxy exceptions are isolated per-candidate (project pattern)."""
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=2)
|
||||
|
||||
def proxy(weights):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
report = run_mix_optimizer(plan, proxy)
|
||||
# Every candidate raised → no valid observations.
|
||||
assert len(report.candidates) == 0
|
||||
# Uniform fallback used.
|
||||
assert abs(sum(report.best_weights) - 1.0) < 1e-6
|
||||
|
||||
|
||||
def test_run_mix_optimizer_propagates_keyboard_interrupt(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=2)
|
||||
|
||||
def proxy(weights):
|
||||
raise KeyboardInterrupt
|
||||
|
||||
with pytest.raises(KeyboardInterrupt):
|
||||
run_mix_optimizer(plan, proxy)
|
||||
|
||||
|
||||
def test_run_mix_optimizer_skips_nan(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=3)
|
||||
seq = iter([float("nan"), 0.5, 0.3])
|
||||
|
||||
def proxy(weights):
|
||||
return next(seq)
|
||||
|
||||
report = run_mix_optimizer(plan, proxy)
|
||||
# NaN observation skipped → 2 valid candidates recorded.
|
||||
assert len(report.candidates) == 2
|
||||
assert report.best_eval_loss == 0.3
|
||||
|
||||
|
||||
def test_run_mix_optimizer_rejects_non_float_loss(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=2)
|
||||
|
||||
def proxy(weights):
|
||||
return "loss"
|
||||
|
||||
with pytest.raises(TypeError, match="proxy_run must return float"):
|
||||
run_mix_optimizer(plan, proxy)
|
||||
|
||||
|
||||
def test_run_mix_optimizer_rejects_bool_loss(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=2)
|
||||
|
||||
def proxy(weights):
|
||||
return True
|
||||
|
||||
with pytest.raises(TypeError, match="proxy_run must return float"):
|
||||
run_mix_optimizer(plan, proxy)
|
||||
|
||||
|
||||
def test_run_mix_optimizer_rejects_non_plan():
|
||||
with pytest.raises(TypeError, match="plan must be MixOptimizationPlan"):
|
||||
run_mix_optimizer({}, lambda w: 0.0) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_run_mix_optimizer_rejects_non_callable_proxy(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch)
|
||||
with pytest.raises(TypeError, match="proxy_run must be callable"):
|
||||
run_mix_optimizer(plan, "not-callable") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_run_mix_optimizer_custom_optimizer(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=2)
|
||||
|
||||
class StubOpt:
|
||||
def __init__(self):
|
||||
self.told = []
|
||||
|
||||
def ask(self):
|
||||
return (0.7, 0.3)
|
||||
|
||||
def tell(self, weights, loss):
|
||||
self.told.append((weights, loss))
|
||||
|
||||
opt = StubOpt()
|
||||
|
||||
def proxy(w):
|
||||
return 1.0
|
||||
|
||||
report = run_mix_optimizer(plan, proxy, optimizer=opt)
|
||||
assert len(opt.told) == 2
|
||||
assert report.best_weights == (0.7, 0.3)
|
||||
|
||||
|
||||
def test_run_mix_optimizer_rejects_bad_optimizer(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch)
|
||||
with pytest.raises(TypeError, match="OptimizerProtocol"):
|
||||
run_mix_optimizer(plan, lambda w: 0.0, optimizer="not-optimizer") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_run_mix_optimizer_all_nan_uniform_fallback(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch, num_probes=2)
|
||||
|
||||
def proxy(w):
|
||||
return float("inf")
|
||||
|
||||
report = run_mix_optimizer(plan, proxy)
|
||||
assert len(report.best_weights) == 2
|
||||
assert abs(sum(report.best_weights) - 1.0) < 1e-6
|
||||
|
||||
|
||||
# ---------- render_mix_recipe_yaml ---------------------------------------
|
||||
|
||||
|
||||
def _report(tmp_path):
|
||||
return MixOptimizationReport(
|
||||
datasets=(str(tmp_path / "a.jsonl"), str(tmp_path / "b.jsonl")),
|
||||
candidates=(),
|
||||
best_weights=(0.6, 0.4),
|
||||
best_eval_loss=0.123,
|
||||
partial=False,
|
||||
elapsed_seconds=10.0,
|
||||
)
|
||||
|
||||
|
||||
def test_render_recipe_shape(tmp_path):
|
||||
text = render_mix_recipe_yaml(_report(tmp_path))
|
||||
assert "data:" in text
|
||||
assert "interleave:" in text
|
||||
assert "strategy: probs" in text
|
||||
assert "0.600000" in text
|
||||
assert "0.400000" in text
|
||||
assert "a.jsonl" in text
|
||||
|
||||
|
||||
def test_render_recipe_rejects_non_report():
|
||||
with pytest.raises(TypeError, match="MixOptimizationReport"):
|
||||
render_mix_recipe_yaml({}) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_render_recipe_rejects_newline_in_path():
|
||||
bad = MixOptimizationReport(
|
||||
datasets=("a\n.jsonl", "b.jsonl"),
|
||||
candidates=(),
|
||||
best_weights=(0.5, 0.5),
|
||||
best_eval_loss=0.0,
|
||||
partial=False,
|
||||
elapsed_seconds=0.0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="control characters"):
|
||||
render_mix_recipe_yaml(bad)
|
||||
|
||||
|
||||
def test_render_recipe_rejects_null_byte_path():
|
||||
bad = MixOptimizationReport(
|
||||
datasets=("a\x00.jsonl", "b.jsonl"),
|
||||
candidates=(),
|
||||
best_weights=(0.5, 0.5),
|
||||
best_eval_loss=0.0,
|
||||
partial=False,
|
||||
elapsed_seconds=0.0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="control characters"):
|
||||
render_mix_recipe_yaml(bad)
|
||||
|
||||
|
||||
def test_render_recipe_partial_note():
|
||||
rep = MixOptimizationReport(
|
||||
datasets=("a.jsonl", "b.jsonl"),
|
||||
candidates=(),
|
||||
best_weights=(0.5, 0.5),
|
||||
best_eval_loss=0.5,
|
||||
partial=True,
|
||||
elapsed_seconds=10.0,
|
||||
)
|
||||
text = render_mix_recipe_yaml(rep)
|
||||
assert "Budget exceeded" in text
|
||||
|
||||
|
||||
def test_render_recipe_handles_inf_loss():
|
||||
rep = MixOptimizationReport(
|
||||
datasets=("a.jsonl", "b.jsonl"),
|
||||
candidates=(),
|
||||
best_weights=(0.5, 0.5),
|
||||
best_eval_loss=float("inf"),
|
||||
partial=False,
|
||||
elapsed_seconds=10.0,
|
||||
)
|
||||
text = render_mix_recipe_yaml(rep)
|
||||
assert "no valid observation" in text
|
||||
|
||||
|
||||
# ---------- write_mix_recipe + load_mix_recipe ---------------------------
|
||||
|
||||
|
||||
def test_write_recipe_round_trip(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
report = _report(tmp_path)
|
||||
out = write_mix_recipe(report, "recipe.yaml")
|
||||
assert os.path.isfile(out)
|
||||
loaded = load_mix_recipe("recipe.yaml")
|
||||
assert "interleave" in loaded
|
||||
assert "train" in loaded
|
||||
|
||||
|
||||
def test_write_recipe_outside_cwd(tmp_path, monkeypatch):
|
||||
work = tmp_path / "work"
|
||||
work.mkdir()
|
||||
monkeypatch.chdir(work)
|
||||
report = _report(tmp_path)
|
||||
with pytest.raises(ValueError, match="outside cwd"):
|
||||
write_mix_recipe(report, str(tmp_path / "out.yaml"))
|
||||
|
||||
|
||||
def test_write_recipe_refuses_overwrite(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
target = tmp_path / "exists.yaml"
|
||||
target.write_text("# pre")
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
write_mix_recipe(_report(tmp_path), "exists.yaml")
|
||||
|
||||
|
||||
def test_write_recipe_overwrite_ok(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
target = tmp_path / "exists.yaml"
|
||||
target.write_text("# pre")
|
||||
out = write_mix_recipe(_report(tmp_path), "exists.yaml", overwrite=True)
|
||||
text = open(out).read()
|
||||
assert "data:" in text
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX symlinks")
|
||||
def test_write_recipe_symlink_rejected(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
real = tmp_path / "real.yaml"
|
||||
real.write_text("# real")
|
||||
link = tmp_path / "link.yaml"
|
||||
link.symlink_to(real)
|
||||
with pytest.raises(ValueError, match="symlink"):
|
||||
write_mix_recipe(_report(tmp_path), "link.yaml", overwrite=True)
|
||||
|
||||
|
||||
def test_write_recipe_null_byte_path(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="null bytes"):
|
||||
write_mix_recipe(_report(tmp_path), "out\x00.yaml")
|
||||
|
||||
|
||||
def test_write_recipe_non_str(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(TypeError):
|
||||
write_mix_recipe(_report(tmp_path), 42) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_load_recipe_missing(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(FileNotFoundError):
|
||||
load_mix_recipe("missing.yaml")
|
||||
|
||||
|
||||
def test_load_recipe_outside_cwd(tmp_path, monkeypatch):
|
||||
work = tmp_path / "work"
|
||||
work.mkdir()
|
||||
outside = tmp_path / "out.yaml"
|
||||
outside.write_text("data: {interleave: {strategy: probs}}")
|
||||
monkeypatch.chdir(work)
|
||||
with pytest.raises(ValueError, match="outside cwd"):
|
||||
load_mix_recipe(str(outside))
|
||||
|
||||
|
||||
def test_load_recipe_not_yaml_mapping(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
p = tmp_path / "bad.yaml"
|
||||
p.write_text("- just a list\n")
|
||||
with pytest.raises(ValueError, match="YAML mapping"):
|
||||
load_mix_recipe("bad.yaml")
|
||||
|
||||
|
||||
def test_load_recipe_missing_data_block(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
p = tmp_path / "missing.yaml"
|
||||
p.write_text("other: value\n")
|
||||
with pytest.raises(ValueError, match="missing required"):
|
||||
load_mix_recipe("missing.yaml")
|
||||
|
||||
|
||||
# ---------- CLI smoke ----------------------------------------------------
|
||||
|
||||
|
||||
def test_mix_cli_help():
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["data", "mix", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "optimiz" in result.output.lower() or "optimize" in result.output.lower()
|
||||
|
||||
|
||||
def test_mix_cli_requires_mode():
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["data", "mix"])
|
||||
assert result.exit_code == 2
|
||||
|
||||
|
||||
def test_mix_cli_mutual_exclusion(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
(tmp_path / "rec.yaml").write_text("data: {}\n")
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
app,
|
||||
["data", "mix", "--optimize", "--apply", "rec.yaml"],
|
||||
)
|
||||
assert result.exit_code == 2
|
||||
|
||||
|
||||
def test_mix_cli_optimize_requires_datasets(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["data", "mix", "--optimize"])
|
||||
assert result.exit_code == 2
|
||||
|
||||
|
||||
def test_mix_cli_optimize_happy(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"data", "mix",
|
||||
"--optimize",
|
||||
"--datasets", "a.jsonl,b.jsonl",
|
||||
"--budget", "60s",
|
||||
"--num-probes", "2",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, (result.output, repr(result.exception))
|
||||
assert os.path.isfile(tmp_path / "mix_recipe.yaml")
|
||||
|
||||
|
||||
def test_mix_cli_optimize_invalid_datasets(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"data", "mix", "--optimize",
|
||||
"--datasets", "missing-and-only-one.jsonl",
|
||||
"--budget", "60s",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 2
|
||||
assert "validation failed" in result.output.lower()
|
||||
|
||||
|
||||
def test_mix_cli_apply_happy(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
# First generate a recipe.
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"data", "mix",
|
||||
"--optimize",
|
||||
"--datasets", "a.jsonl,b.jsonl",
|
||||
"--budget", "60s",
|
||||
"--num-probes", "2",
|
||||
"--output", "rec.yaml",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
# Now apply it.
|
||||
result = runner.invoke(app, ["data", "mix", "--apply", "rec.yaml"])
|
||||
assert result.exit_code == 0, (result.output, repr(result.exception))
|
||||
assert "interleave" in result.output
|
||||
|
||||
|
||||
def test_mix_cli_apply_outside_cwd(tmp_path, monkeypatch):
|
||||
work = tmp_path / "work"
|
||||
work.mkdir()
|
||||
outside = tmp_path / "rec.yaml"
|
||||
outside.write_text("data: {interleave: {strategy: probs}}\n")
|
||||
monkeypatch.chdir(work)
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
app, ["data", "mix", "--apply", str(outside)]
|
||||
)
|
||||
assert result.exit_code == 2
|
||||
|
||||
|
||||
def test_optimization_plan_frozen(tmp_path, monkeypatch):
|
||||
plan = _make_plan(tmp_path, monkeypatch)
|
||||
with pytest.raises(Exception):
|
||||
plan.num_probes = 99 # type: ignore[misc]
|
||||
|
||||
|
||||
# ---------- Review-fix coverage -------------------------------------------
|
||||
|
||||
|
||||
def test_optimization_report_frozen(tmp_path):
|
||||
rep = MixOptimizationReport(
|
||||
datasets=("a.jsonl", "b.jsonl"),
|
||||
candidates=(),
|
||||
best_weights=(0.5, 0.5),
|
||||
best_eval_loss=0.0,
|
||||
partial=False,
|
||||
elapsed_seconds=0.0,
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
rep.partial = True # type: ignore[misc]
|
||||
|
||||
|
||||
def test_build_plan_rejects_seed_at_upper_boundary(tmp_path, monkeypatch):
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="seed must be in"):
|
||||
build_optimization_plan(
|
||||
["a.jsonl", "b.jsonl"], budget="60s", seed=2**31
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["59s", "0s", "0m", "0h"])
|
||||
def test_parse_budget_rejects_below_min_with_suffix(raw):
|
||||
with pytest.raises(ValueError):
|
||||
parse_budget(raw)
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX symlinks")
|
||||
def test_load_mix_recipe_rejects_symlink(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
real = tmp_path / "real.yaml"
|
||||
real.write_text("data: {interleave: {strategy: probs}}\n")
|
||||
link = tmp_path / "link.yaml"
|
||||
link.symlink_to(real)
|
||||
with pytest.raises(ValueError, match="symlink"):
|
||||
load_mix_recipe("link.yaml")
|
||||
|
||||
|
||||
def test_load_mix_recipe_rejects_oversize(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
big = tmp_path / "big.yaml"
|
||||
big.write_bytes(b"data: {x: y}\n" + b"# pad\n" * 200_000)
|
||||
with pytest.raises(ValueError, match="bytes cap"):
|
||||
load_mix_recipe("big.yaml")
|
||||
|
||||
|
||||
def test_mix_cli_optimize_recipe_content(tmp_path, monkeypatch):
|
||||
"""Recipe file is non-empty and contains the canonical interleave block."""
|
||||
_make_files(tmp_path, ["a.jsonl", "b.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"data", "mix", "--optimize",
|
||||
"--datasets", "a.jsonl,b.jsonl",
|
||||
"--budget", "60s",
|
||||
"--num-probes", "2",
|
||||
"--output", "rec.yaml",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, (result.output, repr(result.exception))
|
||||
text = (tmp_path / "rec.yaml").read_text()
|
||||
assert "interleave:" in text
|
||||
assert "strategy: probs" in text
|
||||
assert "a.jsonl" in text
|
||||
|
||||
|
||||
def test_validate_datasets_single_entry_message(tmp_path, monkeypatch):
|
||||
"""Single-entry input should fail with 'at least 2' immediately."""
|
||||
_make_files(tmp_path, ["a.jsonl"])
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with pytest.raises(ValueError, match="at least 2"):
|
||||
validate_datasets(["a.jsonl"])
|
||||
Loading…
Reference in New Issue