From b6be787cc860feed1c1b552228b5eedb0f6645bc Mon Sep 17 00:00:00 2001 From: voorhs Date: Sat, 23 May 2026 19:31:59 +0300 Subject: [PATCH 01/43] add spec --- compute-feasibility-advisor-proposal.md | 201 ++++++++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 compute-feasibility-advisor-proposal.md diff --git a/compute-feasibility-advisor-proposal.md b/compute-feasibility-advisor-proposal.md new file mode 100644 index 000000000..2560d1279 --- /dev/null +++ b/compute-feasibility-advisor-proposal.md @@ -0,0 +1,201 @@ +# Compute Feasibility Advisor for AutoIntent + +- **Date:** 2026-05-23 +- **Status:** Proposal (pre-implementation) +- **Audience:** AutoIntent maintainers / contributor picking up the task + +## Problem + +AutoIntent's main strength is letting a user kick off a full search-space optimization with one call: + +```python +pipeline = Pipeline.from_preset("transformers-heavy") +pipeline.fit(dataset) +``` + +The cost of that convenience is that users — especially those running on a laptop, a single consumer GPU, or a free cloud instance — cannot tell ahead of time whether their hardware can carry the configuration they have just selected. + +Concrete failure cases we see today: + +- `transformers-heavy` fine-tunes `microsoft/deberta-v3-large` for up to 30 epochs across 40 HPO trials. That needs ~12–18 GB VRAM (full fine-tune, fp32) and many hours of wall time on a single GPU. A user with an 8 GB card finds out by OOM, often several minutes into a run. +- Swapping `intfloat/multilingual-e5-large-instruct` (2 GB) for `sentence-transformers/all-MiniLM-L6-v2` (90 MB) changes the resource bill by an order of magnitude — but nothing surfaces this difference up front. +- Disk is a silent failure mode: a search space referencing several large checkpoints can pull >10 GB into the HF cache before any training starts. + +The target audience for this feature is users with limited resources who pick a preset, hit `fit()`, and want to know within a second whether they should change something. + +## Proposed solution: pre-flight resource advisor + +Add a **pre-flight advisor** that, given a parsed search space and a dataset, estimates worst-case disk, RAM, VRAM, and wall-time requirements from public Hugging Face Hub metadata and a small set of formulas, then prints a clear summary with red/yellow/green warnings. By default it is **report-only and never blocks the run**; an opt-in **reduce-to-fit** mode additionally prunes the search space to fit detected hardware. + +### Scope + +The advisor analyses only the **local, model-bearing** modules whose footprint can be derived from HF Hub metadata. Everything else is either trivial or out of band. + + +| Module category | In scope? | Reason | +| -------------------------------------------------------------------------------- | --------- | -------------------------------------------------- | +| `SentenceTransformerEmbeddingConfig` | yes | local transformer, dominant cost on small machines | +| `VllmEmbeddingConfig` | yes | local transformer with extra engine overhead | +| `HFModelConfig`-based scorers (`bert`, `lora`, `ptuning`, `dnnc`, cross-encoder) | yes | the actual heavyweights | +| GCN scorer when configured with a transformer backbone | yes | inherits the backbone cost | +| `OpenaiEmbeddingConfig` | no | no local resources to estimate | +| `HashingVectorizerEmbeddingConfig` | no | trivial cost | +| `knn`, `mlknn`, `linear`, `sklearn`, `catboost`, `description` | no | negligible next to a fine-tune | +| `decision` and `regex` nodes | no | negligible | + + +Rationale: the user's real risk is the heavy transformer-backed modules. A cheap module cannot be the reason a run fails for resource reasons; we don't owe an estimate for it. + +### Inputs + +- The parsed `OptimizationConfig` (search space, HPO config, embedder/transformer configs). +- The training `Dataset` (for `dataset_size` and an approximate token-length distribution). +- Detected local hardware: + - Total / available RAM via `psutil`. + - Free disk on the AutoIntent / HF cache directory via `shutil.disk_usage`. + - Accelerator detection, in priority order: + - **CUDA:** per-GPU VRAM and device name via `torch.cuda`. + - **MPS (Apple Silicon):** detected via `torch.backends.mps.is_available()`. Apple chips use unified memory, so there is no separate VRAM pool — the "VRAM budget" is a fraction of total system RAM. Default budget = 70 % of total RAM (matching the macOS `PYTORCH_MPS_HIGH_WATERMARK_RATIO` default) with the remainder reserved for the OS and other apps. The fraction is exposed as a knob. + - **CPU only:** when neither is available. + +### Output + +A structured estimate plus a human-readable summary printed to the logger. Example: + +``` +Compute feasibility check +───────────────────────── +Available : 8 GB VRAM (NVIDIA RTX 3060), 32 GB RAM, 120 GB free disk +Estimated worst-case requirements for this search space: + Disk : 5.2 GB (3 unique checkpoints) + RAM : ~4 GB + VRAM : ~14 GB ⚠ exceeds available + Time : ~6 h (single-GPU, fp32, rough) + +Drivers of cost: + scoring.bert microsoft/deberta-v3-large full fine-tune × 40 trials × 30 epochs → ~14 GB VRAM, ~5 h + embedder intfloat/multilingual-e5-large-instruct → ~2.2 GB VRAM + +Suggestions: + • Enable mixed precision (fp16/bf16) on the bert scorer + • Reduce batch_size from 64 to 16 or 32 + • Try preset `transformers-light` or `classic-medium` + +These numbers are heuristic upper bounds, not measurements. +``` + +Numbers are reported with honest precision (one significant figure for time, two for memory) and an explicit "estimate, not measurement" disclaimer. + +### Algorithm (proposal, allowed to adjust) + +1. **Collect candidates.** Walk the search space; collect every unique `(module_type, model_name, mode)` triple, where `mode ∈ {inference, lora, full-finetune}`. Also collect HPO knobs that drive cost: `n_trials`, `epochs`, `batch_size`, `max_length`, `dtype` (fp16/bf16/fp32). +2. **Resolve checkpoints.** For each unique `model_name`, query HF Hub for safetensors metadata to read parameter count and weight dtype. Fall back to file-size aggregation if safetensors metadata is missing. Fall back to a "unknown — heuristic only" tag with low-confidence labelling if HF Hub is offline or the repo is private. +3. **Apply formulas.** + - **Disk** = sum over unique checkpoints of total file size, plus a small fixed overhead per checkpoint for tokenizers and config. + - **RAM** = max over modules of `params × dtype_bytes + dataset_tokens × 4 bytes`, treated as a loose upper bound for tokenized buffers. + - **VRAM per module:** + - Inference embedder: `params × dtype_bytes × ~1.3` (small constant for activations). + - Full fine-tune (`bert`, GCN backbone, soft-prompt `ptuning`): `params × dtype_bytes × (1 + 1 + 2)` for weights + grads + Adam state, halved when fp16/bf16 mixed precision is configured. + - LoRA: inference VRAM + a small adapter constant. + - Reranker (cross-encoder, `dnnc`): inference VRAM × small factor for the reranking pass. + - **Time per module** = `n_trials × epochs × (dataset_size / batch_size) × per_step_seconds(params, max_length, device_class)`, where `per_step_seconds` is a small static lookup table keyed on coarse device class (`cpu`, `low-gpu`, `mid-gpu`, `high-gpu`, `apple-silicon`) auto-detected from `torch.cuda.get_device_name` or `platform`/`torch.backends.mps`. Total time = sum across modules. MPS time numbers are coarser than CUDA's (one tier for now); we accept that. +4. **Compare to detected hardware.** Per-dimension status is green / yellow / red against a configurable headroom (defaults: **red** if estimate > 100 % of available, **yellow** if > 70 %). On MPS, "VRAM" and "RAM" estimates draw from the same physical pool; we compare *the larger of the two* against the unified-memory budget rather than each independently. +5. **Render summary.** Log at INFO. If any dimension is red, emit at WARNING so it shows in non-logging contexts. + +### Failure modes + +- **HF Hub offline or private repo:** fall back to "unknown model — name-pattern heuristic only", explicit low-confidence label, never raise. +- **No accelerator (no CUDA and no MPS):** report VRAM as N/A and mark GPU-only modules as "requires GPU" without estimating a (misleading) CPU wall time. +- **MPS configured but a module is incompatible:** vLLM in particular does not run on MPS. Flag the module as "unsupported on MPS" rather than estimating; do not raise. +- **MPS with CPU fallback ops:** some PyTorch ops fall back to CPU on MPS, inflating system-RAM usage and wall time beyond the heuristic. Note this in the disclaimer; we don't try to model it. +- **vLLM configured but not installed:** still estimate (the VRAM accounting is similar), note that the engine itself has additional overhead not captured. +- **Estimate wildly wrong vs. reality:** always-on disclaimer in the printed summary that these are heuristic upper bounds. + +### Reduce-to-fit mode + +The feasibility check has two modes sharing the same estimation pipeline: + +- **Report mode (default).** Print the summary, return the structured estimate, let the run proceed regardless of severity. +- **Reduce-to-fit mode (opt-in).** Additionally prune the search space to fit detected hardware before the run starts. Same estimates, same comparisons — just one extra step that produces a reduced search space. + +Using the same per-module estimates, the pruner applies three least-destructive steps in order: + +1. **Filter discrete-choice hyperparameters.** For lists of cost-driving values (model name, batch size, training epochs), keep only entries whose worst-case estimate fits. +2. **Cap continuous ranges.** For `{low, high}` ranges of cost-driving parameters, lower the upper bound to the largest fitting value. Ranges of non-cost parameters (learning rate, decision thresholds) are not touched. +3. **Drop module variants.** If a module entry has any required hyperparameter with no satisfiable value left, drop that module entry from its node's search space. + +Guard rails: + +- If pruning would leave any node's search space empty, the pruner **raises**. We don't silently produce a non-runnable pipeline, and we don't quietly fall back to report-only — failing loudly is the right contract for a mode whose whole purpose is to make the run feasible. The error message points the user toward a lighter preset. +- Time is not used as a filter — only memory and disk are. Time is still reported. +- Headroom thresholds are intentionally generous to avoid over-pruning and are configurable. + +Alongside the standard estimate, the caller receives a structured description of what was filtered, capped, and dropped, plus the resulting search space and its recomputed (now green) estimate. + +**Drawbacks worth surfacing.** + +- **Silent narrowing of intent.** A search space deliberately written to include heavy/light variants for comparison gets halved. The mode is opt-in for this reason. +- **Over-pruning when our formulas overestimate.** A 30 %-high estimate on a borderline configuration throws away a run that would have succeeded. Generous headroom defaults mitigate; the knob is exposed. +- **Hard failure when nothing fits.** Raising is intentional — silent degradation to report-only would defeat the mode's purpose — but it is a sharper edge than report mode has. +- **Pre-trial only.** The rewrite happens before any HPO trial starts. This is fine because the search space is treated as immutable across a study, but worth calling out so nobody tries to make this dynamic later. + +## Alternatives considered and rejected + +### B. Smoke-test calibration + +Run each unique module for one mini-batch / one step before the real fit, measure peak RAM and VRAM with `psutil`, `tracemalloc`, and `torch.cuda.max_memory_allocated`, time the step, and extrapolate to the full search space. + +Rejected because: + +- It **downloads weights just to estimate** — the disk-headroom check we wanted to provide is defeated by the act of performing it. +- It can **OOM while predicting OOM**, exactly on the constrained hardware that is the target audience. +- It adds **seconds to minutes** of wall time before `fit()` does anything, surprising users. +- It needs per-module "tiny run" hooks; not every scorer has a clean "stop after one step" path. +- For OpenAI- or vLLM-served embedders, a smoke test costs real money or starts the engine. +- Still not accurate due to CUDA and CPU cache, memory heating and so on. + +### C. Curated benchmark table + +Ship a JSON in the package with measured VRAM and per-step time for the bundled-preset checkpoints, broken out by hardware class (cpu / mid-gpu / high-gpu) and mode (inference / lora / full-finetune). Fall back to heuristics for unknown checkpoints. + +Rejected because: + +- **Maintenance burden:** every new model added to a preset would need entries across the hardware × precision × mode matrix. +- Numbers **go stale** when `transformers` updates change defaults (attention impl, dtype, gradient checkpointing). +- It still needs the chosen-solution heuristics as a long-tail fallback — so it adds work on top of Option A without replacing it. +- **Confident-but-wrong is worse than honest-but-fuzzy.** A table that says "4 GB on 4090" when the user OOMs at 4.5 GB damages trust more than a clearly-labelled range would. + +### D. Layered (A by default, opt-in B, embedded table from C, local actuals cache) + +Combine all three: ship A as the fast path, allow `calibrate=True` to trigger B for heavy modules only, embed a small table from C for the bundled-preset checkpoints, and write actuals from every real run to a local cache that feeds back into future estimates. + +Rejected because: + +- **Implementation surface multiplies:** two estimation code paths to keep consistent, a cache schema with versioning and eviction, two failure modes to document. +- **Discoverability:** users may not learn about `calibrate=True` and the realized value compresses back to roughly Option A anyway. +- The team's bandwidth doesn't justify the marginal accuracy gain over A for the target audience. + +## Comparison + + +| Dimension | A (chosen) | B (smoke-test) | C (benchmark table) | D (layered) | +| -------------------------------- | ------------------------------ | ---------------------- | ---------------------------------- | ------------------------------------- | +| Wall time at pre-flight | < 1 s | seconds–minutes | < 1 s | < 1 s default, s–min when calibrating | +| Accuracy on common checkpoints | medium | high | high | high | +| Accuracy on custom checkpoints | medium | high | medium (fallback) | medium–high | +| Time-estimate quality | low–medium | high | high | high | +| Disk pre-download required | no | yes | no | only when calibrating | +| Risk of OOM during the check | none | real | none | only when calibrating | +| Network usage | 1 cached call per unique model | none beyond normal fit | none | combination | +| Implementation effort | small | large | medium + ongoing benchmark refresh | large + cache infra | +| Ongoing maintenance | low (formulas only) | low | high | high | +| Friendly to offline / air-gapped | with fallback | yes | yes | partial | + + +The chosen solution accepts a real accuracy gap on time and a moderate accuracy gap on VRAM in exchange for the only profile that fits the target audience's constraints: zero added wall time, zero added downloads, zero added failure modes, and a small one-time implementation cost. + +## Out of scope (possible follow-ups) + +- Live resource observability during `fit()` (peak RAM / VRAM per trial, abort on overrun). +- A learned calibration cache from real runs to refine estimates over time. + From dceb9854e0dd51171dd42291086266cfb3c269f4 Mon Sep 17 00:00:00 2001 From: voorhs Date: Fri, 5 Jun 2026 11:05:27 +0300 Subject: [PATCH 02/43] upd tech spec --- compute-feasibility-advisor-proposal.md | 37 +++++++++++++++++++------ 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/compute-feasibility-advisor-proposal.md b/compute-feasibility-advisor-proposal.md index 2560d1279..9e7833e3d 100644 --- a/compute-feasibility-advisor-proposal.md +++ b/compute-feasibility-advisor-proposal.md @@ -3,6 +3,7 @@ - **Date:** 2026-05-23 - **Status:** Proposal (pre-implementation) - **Audience:** AutoIntent maintainers / contributor picking up the task +- **Scope of this document:** technical specification — *what* the advisor estimates and the formulas it uses. Architectural and system-design choices (where the advisor lives in the codebase, how it integrates with the optimizer, the public API surface, file/module layout) are deliberately left to the implementer. ## Problem @@ -38,13 +39,15 @@ The advisor analyses only the **local, model-bearing** modules whose footprint c | `VllmEmbeddingConfig` | yes | local transformer with extra engine overhead | | `HFModelConfig`-based scorers (`bert`, `lora`, `ptuning`, `dnnc`, cross-encoder) | yes | the actual heavyweights | | GCN scorer when configured with a transformer backbone | yes | inherits the backbone cost | +| `LinearScorer` (sklearn `LogisticRegression` / `LogisticRegressionCV`) | yes | dominant cost on presets with no transformer fine-tune; the CV path multiplies a single fit by ~30 | +| `CatBoostScorer` | yes | dominant cost on presets with no transformer fine-tune; high default `iterations` | | `OpenaiEmbeddingConfig` | no | no local resources to estimate | | `HashingVectorizerEmbeddingConfig` | no | trivial cost | -| `knn`, `mlknn`, `linear`, `sklearn`, `catboost`, `description` | no | negligible next to a fine-tune | +| `knn`, `mlknn`, generic `sklearn` classifiers via `SklearnScorer`, `description` | no | bounded so far below any in-scope module that they cannot plausibly be the bottleneck | | `decision` and `regex` nodes | no | negligible | -Rationale: the user's real risk is the heavy transformer-backed modules. A cheap module cannot be the reason a run fails for resource reasons; we don't owe an estimate for it. +Rationale: the user's real risk is whichever module is the actual bottleneck. On heavy presets that is a transformer fine-tune; on light presets it shifts to `linear` (CV-multiplied) or `catboost` (1000 default iterations × dataset shape). Modules left out of scope are ones whose cost is bounded so far below any in-scope module that they cannot plausibly be the reason a run fails. ### Inputs @@ -88,17 +91,33 @@ Numbers are reported with honest precision (one significant figure for time, two ### Algorithm (proposal, allowed to adjust) -1. **Collect candidates.** Walk the search space; collect every unique `(module_type, model_name, mode)` triple, where `mode ∈ {inference, lora, full-finetune}`. Also collect HPO knobs that drive cost: `n_trials`, `epochs`, `batch_size`, `max_length`, `dtype` (fp16/bf16/fp32). -2. **Resolve checkpoints.** For each unique `model_name`, query HF Hub for safetensors metadata to read parameter count and weight dtype. Fall back to file-size aggregation if safetensors metadata is missing. Fall back to a "unknown — heuristic only" tag with low-confidence labelling if HF Hub is offline or the repo is private. -3. **Apply formulas.** - - **Disk** = sum over unique checkpoints of total file size, plus a small fixed overhead per checkpoint for tokenizers and config. - - **RAM** = max over modules of `params × dtype_bytes + dataset_tokens × 4 bytes`, treated as a loose upper bound for tokenized buffers. +1. **Collect candidates.** Walk the search space; collect every unique in-scope module. For transformer-bearing modules the identity is `(module_type, model_name, mode)` with `mode ∈ {inference, lora, full-finetune}`. For `linear` and `catboost` the identity is `(module_type, embedder_name, task_kind)` with `task_kind ∈ {multiclass, multilabel}` — the routing through `LogisticRegressionCV` vs `MultiOutputClassifier`, and CatBoost's per-class trees, both depend on it. Also collect the HPO knobs that drive cost: `n_trials` plus per-module knobs — transformer (`epochs`, `batch_size`, `max_length`, `dtype` ∈ {fp16, bf16, fp32}), `linear` (`cv`, `max_iter`), `catboost` (`iterations`, `depth`, `task_type`, `features_type`). +2. **Resolve checkpoints.** For each unique `model_name`, query HF Hub for safetensors metadata to read parameter count and weight dtype. Fall back to file-size aggregation if safetensors metadata is missing. Fall back to a "unknown — heuristic only" tag with low-confidence labelling if HF Hub is offline or the repo is private. `LinearScorer` and `CatBoostScorer` have no checkpoint of their own; they reuse the embedder resolved by this step in their formulas (their cost is parameterised by `embedder_dim`, not parameter count). +3. **Apply formulas.** All values are honest upper bounds; convergence and early stopping often terminate well below them. + - **Disk** = sum over unique downloadable checkpoints of total file size, plus a small fixed overhead per checkpoint for tokenizers and config. `LinearScorer` and `CatBoostScorer` contribute zero (they consume embedder output that is already accounted for upstream). + - **RAM per module:** + - Transformer modules (any mode): `params × dtype_bytes + dataset_tokens × 4 bytes`, treated as a loose upper bound for tokenized buffers. + - `LinearScorer`: `8 × n_samples × embedder_dim` (float64 data matrix — the dominant term) `+ 8 × n_classes × embedder_dim` (coefficients) `+ ~10 × 8 × embedder_dim` (L-BFGS history). + - `CatBoostScorer`: `4 × n_samples × n_features` (data, float32 internally) `+ 4 × n_features × n_bins` (histograms; default `n_bins = 254`) `+ iterations × 2^depth × ~32 bytes` (tree storage). For `features_type ∈ {embedding, both}`, `n_features = embedder_dim`. For `features_type = text`, `n_features` is the BoW vocab discovered at fit; bound with a coarse default (e.g. 50 000) and tag the estimate low-confidence. + - For `linear` and `catboost`, `embedder_dim` is taken from the largest embedder in the same node group — same worst-case stance as the rest of the estimate. - **VRAM per module:** - Inference embedder: `params × dtype_bytes × ~1.3` (small constant for activations). - Full fine-tune (`bert`, GCN backbone, soft-prompt `ptuning`): `params × dtype_bytes × (1 + 1 + 2)` for weights + grads + Adam state, halved when fp16/bf16 mixed precision is configured. - LoRA: inference VRAM + a small adapter constant. - Reranker (cross-encoder, `dnnc`): inference VRAM × small factor for the reranking pass. - - **Time per module** = `n_trials × epochs × (dataset_size / batch_size) × per_step_seconds(params, max_length, device_class)`, where `per_step_seconds` is a small static lookup table keyed on coarse device class (`cpu`, `low-gpu`, `mid-gpu`, `high-gpu`, `apple-silicon`) auto-detected from `torch.cuda.get_device_name` or `platform`/`torch.backends.mps`. Total time = sum across modules. MPS time numbers are coarser than CUDA's (one tier for now); we accept that. + - `LinearScorer`: N/A (sklearn is CPU-only). + - `CatBoostScorer`: 0 by default; if `task_type="GPU"` is configured, the RAM formula above lives on device instead. + - **Time per module:** + - Transformer modules: `n_trials × epochs × (dataset_size / batch_size) × per_step_seconds(params, max_length, device_class)`, where `per_step_seconds` is a small static lookup keyed on coarse device class (`cpu`, `low-gpu`, `mid-gpu`, `high-gpu`, `apple-silicon`) auto-detected from `torch.cuda.get_device_name` or `platform`/`torch.backends.mps`. + - `LinearScorer`: `n_trials × C_cpu × n_samples × embedder_dim × max_iter × cv_multiplier × class_multiplier`, where: + - `C_cpu ≈ 1e-8 s` per `(sample × feature × iteration)` on a single modern CPU core. + - `cv_multiplier = Cs × cv + 1 ≈ 31` for the multiclass path (`LogisticRegressionCV` with default `Cs = 10`, repo default `cv = 3`, plus one final refit). `cv_multiplier = 1` for the multilabel path (no inner CV). + - `class_multiplier = n_classes` for the multilabel path (`MultiOutputClassifier` fits one binary LogReg per class); `class_multiplier = 1` otherwise. + - `CatBoostScorer`: `n_trials × iterations × C_device × n_samples × n_features × depth × class_multiplier`, where: + - `C_device ≈ 1e-9 s` on CPU, ~5–20× faster on GPU. Resolve `C_device` via the same `device_class` lookup as the transformer time formula. + - `class_multiplier = n_classes` for both the multiclass `MultiClass` loss (per-class trees per iteration) and the multilabel routing (one CatBoost per class). + - Early stopping is not modelled; `iterations` is treated as the upper bound. + - Total time = sum across modules. MPS time numbers are coarser than CUDA's (one tier for now); we accept that. 4. **Compare to detected hardware.** Per-dimension status is green / yellow / red against a configurable headroom (defaults: **red** if estimate > 100 % of available, **yellow** if > 70 %). On MPS, "VRAM" and "RAM" estimates draw from the same physical pool; we compare *the larger of the two* against the unified-memory budget rather than each independently. 5. **Render summary.** Log at INFO. If any dimension is red, emit at WARNING so it shows in non-logging contexts. @@ -120,7 +139,7 @@ The feasibility check has two modes sharing the same estimation pipeline: Using the same per-module estimates, the pruner applies three least-destructive steps in order: -1. **Filter discrete-choice hyperparameters.** For lists of cost-driving values (model name, batch size, training epochs), keep only entries whose worst-case estimate fits. +1. **Filter discrete-choice hyperparameters.** For lists of cost-driving values (model name, batch size, training epochs, CatBoost `iterations` / `depth`, sklearn `cv`), keep only entries whose worst-case estimate fits. 2. **Cap continuous ranges.** For `{low, high}` ranges of cost-driving parameters, lower the upper bound to the largest fitting value. Ranges of non-cost parameters (learning rate, decision thresholds) are not touched. 3. **Drop module variants.** If a module entry has any required hyperparameter with no satisfiable value left, drop that module entry from its node's search space. From 94b4e121a7cad8527ceab66984d2e23086799775 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Wed, 10 Jun 2026 02:16:01 +0300 Subject: [PATCH 03/43] add feasibility advisor: CLI script, package, tests; expand proposal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - proposal: introduce 3-phase framing (resource/data/config), add resource-phase refinements (warm cache, n_jobs × VRAM, refit_after, Hub reachability, CatBoost GPU sanity), data-quality phase (token truncation, split readiness, partial descriptions, embedder dim), config sanity phase, updated example output, CLI surface, out-of- scope deferrals - _advisor package: hardware detection (CUDA/MPS/CPU with broken-CUDA fallback), HF Hub metadata + warm-cache probe + offline heuristics, three-phase run_preflight returning structured PreflightReport, text + JSON renderers - autointent-advisor CLI: inspect and recommend subcommands; placeholder dataset stats when no --dataset given - 88 offline tests covering hardware fallbacks, every bundled preset, severity routing, report serialization, name-pattern heuristics, AMP invariant, dump_modules / refit_after, CLI flows Co-Authored-By: Claude Opus 4.7 --- compute-feasibility-advisor-proposal.md | 71 +++- pyproject.toml | 1 + src/autointent/_advisor/__init__.py | 23 ++ src/autointent/_advisor/_cli.py | 243 ++++++++++++++ src/autointent/_advisor/_estimates.py | 382 ++++++++++++++++++++++ src/autointent/_advisor/_hardware.py | 160 +++++++++ src/autointent/_advisor/_hub.py | 183 +++++++++++ src/autointent/_advisor/_render.py | 104 ++++++ src/autointent/_advisor/_report.py | 113 +++++++ tests/advisor/__init__.py | 0 tests/advisor/test_estimates_and_cli.py | 198 +++++++++++ tests/advisor/test_estimates_internals.py | 319 ++++++++++++++++++ tests/advisor/test_hardware_detection.py | 72 ++++ tests/advisor/test_hub_heuristics.py | 81 +++++ tests/advisor/test_render.py | 151 +++++++++ tests/advisor/test_report.py | 85 +++++ 16 files changed, 2181 insertions(+), 5 deletions(-) create mode 100644 src/autointent/_advisor/__init__.py create mode 100644 src/autointent/_advisor/_cli.py create mode 100644 src/autointent/_advisor/_estimates.py create mode 100644 src/autointent/_advisor/_hardware.py create mode 100644 src/autointent/_advisor/_hub.py create mode 100644 src/autointent/_advisor/_render.py create mode 100644 src/autointent/_advisor/_report.py create mode 100644 tests/advisor/__init__.py create mode 100644 tests/advisor/test_estimates_and_cli.py create mode 100644 tests/advisor/test_estimates_internals.py create mode 100644 tests/advisor/test_hardware_detection.py create mode 100644 tests/advisor/test_hub_heuristics.py create mode 100644 tests/advisor/test_render.py create mode 100644 tests/advisor/test_report.py diff --git a/compute-feasibility-advisor-proposal.md b/compute-feasibility-advisor-proposal.md index 9e7833e3d..7ebf70bd9 100644 --- a/compute-feasibility-advisor-proposal.md +++ b/compute-feasibility-advisor-proposal.md @@ -49,6 +49,16 @@ The advisor analyses only the **local, model-bearing** modules whose footprint c Rationale: the user's real risk is whichever module is the actual bottleneck. On heavy presets that is a transformer fine-tune; on light presets it shifts to `linear` (CV-multiplied) or `catboost` (1000 default iterations × dataset shape). Modules left out of scope are ones whose cost is bounded so far below any in-scope module that they cannot plausibly be the reason a run fails. +### Phases + +The advisor is one entry point, but internally splits work into three phases that share a single `PreflightReport` object. The split is internal organization — all three run at the same hook point (after `validate_modules`, before `_fit(context)`) and the user sees one summary. Separating them keeps each phase's inputs, formulas, and failure modes scoped: + +- **Resource phase.** Disk / RAM / VRAM / wall-time estimates and comparisons against detected hardware. Most of the formulas in this document live here. This is the only phase consumed by the reduce-to-fit pruner. +- **Data quality phase.** Findings derived from the dataset jointly with the active search space — token-length truncation, split readiness (auto-invokes the existing `check_split_readiness` utility rather than re-implementing it), partial intent descriptions paired with the `description` scorer, embedder/scorer dimension consistency. Reports red/yellow lines but never prunes the search space; the user fixes the dataset or the config. +- **Configuration sanity phase.** Joint checks across dataset + search-space + hardware that don't slot cleanly into the other two — e.g., `hpo_config.n_jobs > 1` × per-trial VRAM contention, CatBoost `task_type="GPU"` with no CUDA. Pydantic schema validation already runs upstream on `OptimizationConfig`; this phase only adds checks that need joint inspection. + +The advisor consumes `validate_modules`'s *post-filter* view of `self.nodes` — it does not duplicate that mutating filter. + ### Inputs - The parsed `OptimizationConfig` (search space, HPO config, embedder/transformer configs). @@ -68,12 +78,19 @@ A structured estimate plus a human-readable summary printed to the logger. Examp ``` Compute feasibility check ───────────────────────── -Available : 8 GB VRAM (NVIDIA RTX 3060), 32 GB RAM, 120 GB free disk -Estimated worst-case requirements for this search space: - Disk : 5.2 GB (3 unique checkpoints) +Resource: + Available : 8 GB VRAM (NVIDIA RTX 3060), 32 GB RAM, 120 GB free disk + Disk : 5.2 GB to download, 1.1 GB already cached (3 unique checkpoints) RAM : ~4 GB - VRAM : ~14 GB ⚠ exceeds available - Time : ~6 h (single-GPU, fp32, rough) + VRAM : ~14 GB × 2 parallel trials (n_jobs=2) ⚠ exceeds available + Time : ~6 h (+~12 min for refit_after) (single-GPU, fp32, rough) + +Data: + Train tokens p95 : 612 (exceeds bert.max_length=512) ⚠ ~7% truncated + Split readiness : 2 classes have <3 samples — LogisticRegressionCV cv=3 will fail ✗ + +Config: + CatBoost task_type=GPU but no CUDA detected — will fall back to CPU ⚠ Drivers of cost: scoring.bert microsoft/deberta-v3-large full fine-tune × 40 trials × 30 epochs → ~14 GB VRAM, ~5 h @@ -82,6 +99,7 @@ Drivers of cost: Suggestions: • Enable mixed precision (fp16/bf16) on the bert scorer • Reduce batch_size from 64 to 16 or 32 + • Set hpo_config.n_jobs=1 — parallel trials are doubling VRAM demand • Try preset `transformers-light` or `classic-medium` These numbers are heuristic upper bounds, not measurements. @@ -121,6 +139,34 @@ Numbers are reported with honest precision (one significant figure for time, two 4. **Compare to detected hardware.** Per-dimension status is green / yellow / red against a configurable headroom (defaults: **red** if estimate > 100 % of available, **yellow** if > 70 %). On MPS, "VRAM" and "RAM" estimates draw from the same physical pool; we compare *the larger of the two* against the unified-memory budget rather than each independently. 5. **Render summary.** Log at INFO. If any dimension is red, emit at WARNING so it shows in non-logging contexts. +#### Resource-phase refinements + +These adjust the formulas above for situations that look fine in single-trial isolation but blow up in practice: + +- **Cold-vs-warm HF cache (Tier 1).** Before reporting disk, probe each unique `model_name` against the local HF cache via `huggingface_hub.try_to_load_from_cache` / `scan_cache_dir`, keyed off `HF_HOME`. Split the disk line into `to_download` vs `already_cached`. Treat a repo as cached only if the weight shard (`model.safetensors` or equivalent) is present — not just config/tokenizer files. Without this, a repeated run on the same machine alarms the user about gigabytes they already have. +- **Concurrent-trial × per-trial VRAM (Tier 1).** Multiply the per-trial VRAM estimate by `hpo_config.n_jobs` when `n_jobs > 1` and the active accelerator is GPU. Same for the `dump_modules=True` path on disk: each trial writes module weights to the dump dir, so multiply per-module dump-disk by `n_trials`. vLLM is process-isolated and its contention model differs; note this in the disclaimer. +- **`refit_after=True` time delta (Tier 2).** When `Pipeline.fit(refit_after=True)`, add one full-data training pass per node to the time estimate. Small term but easy to forget; users running close to their time budget care about it. +- **HF Hub reachability probe (Tier 2).** One up-front `HfApi().whoami()` (or unauthenticated `HEAD` to `huggingface.co`) at the start of the phase. On failure, consistently downgrade *all* model entries to the "unknown — heuristic only" path instead of timing out per-model 10× on a 10-model search space. +- **CatBoost `task_type="GPU"` sanity (Tier 2).** When CatBoost is in the search space with `task_type="GPU"` but `torch.cuda.is_available()` is false, tag yellow — CatBoost silently falls back to CPU and the user otherwise sees CPU speeds with no warning. + +### Data quality phase + +The resource phase predicts whether the run *fits*. The data quality phase predicts whether the run *produces a meaningful result*. Both are caught at the same hook point because both have the same failure mode from the user's perspective: hours of compute followed by a cryptic error or a silently degraded model. + +- **Token-length truncation (Tier 1).** Sample ~1000 utterances from the train split, tokenize against each unique transformer's tokenizer, compute `p95_tokens` and `% truncated` against the module's `max_length`. Yellow when >1% truncated; red when >10%. Reuse the tokenizer the resource phase already loaded for parameter-count resolution — don't double-fetch. The existing pipeline silently truncates (sentence-transformers and the HF Trainer both default to `truncation=True`); there is no warning anywhere today. +- **Auto-invoke `check_split_readiness` (Tier 1).** Call the existing utility at `context/data_handler/_readiness_util.py:44–109` with the active `data_config` and surface its `SplitReadinessResult` — it already returns `underpopulated_classes`, `ready`, and a `reason` string, but is not called anywhere from `Pipeline.fit()` today. When `LinearScorer` with CV is in the search space and any class has `n < cv`, name the module by name in the red line ("`LogisticRegressionCV` cv=3 will fail: classes [X, Y] have <3 samples") rather than emitting a generic split-readiness message. +- **Partial intent descriptions × `description` scorer (Tier 1).** The dataset constructor already warns once at import when *some* but not all intents have descriptions (`_dataset/_dataset.py:199–207`). The advisor escalates this to red when the `description` scorer is also present in the active search space — otherwise the run will produce NaN embeddings for the missing intents. Action message: "fill in N missing descriptions", not "drop the scorer". +- **Embedder ↔ scorer dimension consistency (Tier 2).** For `LinearScorer` / `CatBoostScorer` with `features_type="both"`, verify the embedder reachable from the same node group exposes a stable, expected dimension. Cross-node walk; surface as yellow when the resolved dimension cannot be confirmed pre-flight. + +### Configuration sanity phase + +Pydantic schema validation on `OptimizationConfig` runs upstream at config-load time; this phase only adds checks that require *joint* inspection of dataset + search-space + hardware. With Tier 1 + Tier 2 in scope today, this phase holds two items: + +- The `n_jobs × VRAM` callout, surfaced jointly with the resource phase (single line in the rendered output). +- The CatBoost `task_type="GPU"` without CUDA check, same. + +Both could live entirely in the resource phase; they get their own phase because future additions — joint scorer↔decision shape checks, OOS-support mismatches detected up front rather than at module instantiation, embedder-dimension mismatches — slot here naturally. Keep the phase scaffold even if it is currently thin. + ### Failure modes - **HF Hub offline or private repo:** fall back to "unknown model — name-pattern heuristic only", explicit low-confidence label, never raise. @@ -137,6 +183,8 @@ The feasibility check has two modes sharing the same estimation pipeline: - **Report mode (default).** Print the summary, return the structured estimate, let the run proceed regardless of severity. - **Reduce-to-fit mode (opt-in).** Additionally prune the search space to fit detected hardware before the run starts. Same estimates, same comparisons — just one extra step that produces a reduced search space. +Reduce-to-fit consumes only the **resource phase** output. Data-quality and config-sanity findings are reported but never trigger pruning — they require user action (fix the dataset, change a config flag), not search-space narrowing. + Using the same per-module estimates, the pruner applies three least-destructive steps in order: 1. **Filter discrete-choice hyperparameters.** For lists of cost-driving values (model name, batch size, training epochs, CatBoost `iterations` / `depth`, sklearn `cv`), keep only entries whose worst-case estimate fits. @@ -158,6 +206,15 @@ Alongside the standard estimate, the caller receives a structured description of - **Hard failure when nothing fits.** Raising is intentional — silent degradation to report-only would defeat the mode's purpose — but it is a sharper edge than report mode has. - **Pre-trial only.** The rewrite happens before any HPO trial starts. This is fine because the search space is treated as immutable across a study, but worth calling out so nobody tries to make this dynamic later. +### CLI surface + +The advisor is also exposed as a console script (`autointent-advisor`) so users can answer "what will this cost?" and "what should I run?" without writing Python. Two subcommands: + +- **`autointent-advisor inspect `.** Resolves the preset (or a user-supplied `OptimizationConfig`), detects local hardware, runs the same three-phase advisor that `Pipeline.fit()` runs, and prints the same report. Accepts `--dataset` for a real dataset, or `--n-samples / --n-classes / --avg-tokens` placeholders when the dataset is not yet built — so the script is useful before any training data exists. `--json` emits the structured `PreflightReport` for scripting. +- **`autointent-advisor recommend [--n-samples ... | --dataset ...] [--budget-time 12h] [--budget-vram-gb 8]`.** Detects local hardware (with manual overrides applied), iterates over the bundled presets in `_presets/`, and tags each as `feasible` / `feasible-with-reduce` / `infeasible`. Ranks feasible presets by quality tier (`heavy > medium > light`) then estimated wall-time; picks the top one as the recommendation. For the heaviest infeasible preset, surfaces the single most-impactful knob change that would make it fit (e.g., "`transformers-heavy` would fit if `batch_size` ≤ 16 and `dtype=fp16`"), reusing the reduce-to-fit pruner's per-knob delta info. + +**Constraints (both subcommands).** No model downloads — only HF Hub metadata endpoints (`HfApi().model_info`); never `from_pretrained`. Offline-safe — on Hub unreachability, fall back to the same "heuristic only" path and mark the report low-confidence; do not raise. Hardware-detection failures (broken CUDA install where `torch.cuda.mem_get_info()` raises) fall back to CPU detection and tag the report rather than crashing. + ## Alternatives considered and rejected ### B. Smoke-test calibration @@ -217,4 +274,8 @@ The chosen solution accepts a real accuracy gap on time and a moderate accuracy - Live resource observability during `fit()` (peak RAM / VRAM per trial, abort on overrun). - A learned calibration cache from real runs to refine estimates over time. +- **Determinism / `cudnn.deterministic` check.** Belongs in seed-setting code (`set_seed` utility, `Pipeline.__init__`), not in a feasibility advisor — reproducibility is not a hardware-budget question. +- **OpenAI / Generator token-cost ($) estimation.** Real value, but pricing tables age badly, the `StructuredOutputCache` hit rate is unknowable upfront, and the API-paying audience overlaps poorly with this advisor's stated audience (resource-constrained local users). Push to a separate `cost_estimator` tool. +- **Predictive CO₂ / emissions.** `_callbacks/emissions_tracker.py` already does this retrospectively, accurately. A predictive version multiplies our (loose) time estimate by a regional kWh/CO₂ factor — two sources of imprecision compounded. The retrospective number is the trustworthy one. +- **vLLM startup compile time.** Minutes of overhead before any work, but vLLM is unsupported on MPS, isn't the dominant cost on CUDA once running, and modelling it needs a startup-time lookup table. Note once in the disclaimer; do not model. diff --git a/pyproject.toml b/pyproject.toml index 202a8cb3f..b47993faf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,6 +141,7 @@ Documentation = "https://deeppavlov.github.io/AutoIntent/" [project.scripts] "basic-aug" = "autointent.generation.utterances.basic.cli:main" "evolution-aug" = "autointent.generation.utterances.evolution.cli:main" +"autointent-advisor" = "autointent._advisor._cli:main" [build-system] requires = ["uv_build>=0.8.7,<0.9.0"] diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py new file mode 100644 index 000000000..5f29b028e --- /dev/null +++ b/src/autointent/_advisor/__init__.py @@ -0,0 +1,23 @@ +"""Pre-flight compute feasibility advisor. + +Exposes a small surface used by both ``Pipeline.fit()`` (future integration) and +the ``autointent-advisor`` CLI script. See ``compute-feasibility-advisor-proposal.md`` +at the repo root for the design document. +""" + +from __future__ import annotations + +from ._hardware import HardwareProfile, detect_hardware +from ._report import DatasetStats, Finding, PreflightReport, ResourceEstimate, Severity +from ._estimates import run_preflight + +__all__ = [ + "DatasetStats", + "Finding", + "HardwareProfile", + "PreflightReport", + "ResourceEstimate", + "Severity", + "detect_hardware", + "run_preflight", +] diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py new file mode 100644 index 000000000..4e7eae000 --- /dev/null +++ b/src/autointent/_advisor/_cli.py @@ -0,0 +1,243 @@ +"""Console-script entry point for the pre-flight advisor. + +Two subcommands: + +* ``inspect`` — show what a given preset / config will cost on this machine. +* ``recommend`` — pick the best-fitting bundled preset for this machine. + +Both subcommands accept either a real ``--dataset`` (path to load with +``Dataset.from_*`` constructors) or ``--n-samples / --n-classes / --avg-tokens`` +placeholders so the script is useful before the user has built a dataset. +""" + +from __future__ import annotations + +import argparse +import logging +import sys +from pathlib import Path +from typing import Any + +import yaml + +from ._estimates import run_preflight +from ._hardware import detect_hardware +from ._render import render_json, render_recommendation, render_text +from ._report import DatasetStats, PreflightReport + +logger = logging.getLogger("autointent.advisor") + +BUNDLED_PRESETS = [ + "transformers-heavy", + "transformers-light", + "transformers-no-hpo", + "nn-heavy", + "nn-medium", + "classic-heavy", + "classic-medium", + "classic-light", + "zero-shot-encoders", + "zero-shot-llm", +] + +# rough quality tiering used by `recommend` +_QUALITY_TIER = { + "transformers-heavy": 5, + "nn-heavy": 4, + "transformers-light": 4, + "nn-medium": 3, + "classic-heavy": 3, + "transformers-no-hpo": 3, + "classic-medium": 2, + "classic-light": 1, + "zero-shot-encoders": 2, + "zero-shot-llm": 4, +} + + +def _load_config(target: str) -> tuple[dict[str, Any], str]: + """Return (config_dict, friendly_name) for either a preset or a path.""" + path = Path(target) + if path.is_file(): + with path.open(encoding="utf-8") as f: + return yaml.safe_load(f), path.stem + # treat as a bundled preset name + from autointent.utils import load_preset + + return load_preset(target), target # type: ignore[arg-type] + + +def _stats_from_args(args: argparse.Namespace) -> DatasetStats: + if args.dataset: + return _stats_from_dataset(args.dataset, multilabel=args.task == "multilabel") + return DatasetStats.placeholder( + n_samples=args.n_samples, + n_classes=args.n_classes, + avg_tokens=args.avg_tokens, + multilabel=args.task == "multilabel", + ) + + +def _stats_from_dataset(path: str, *, multilabel: bool) -> DatasetStats: + """Best-effort: load a dataset from disk via the existing Dataset constructor.""" + try: + from autointent import Dataset + except ImportError: + logger.warning("autointent.Dataset unavailable; falling back to placeholders.") + return DatasetStats.placeholder(multilabel=multilabel) + + try: + ds = Dataset.from_json(path) if path.endswith(".json") else Dataset.from_hub(path) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to load dataset %s: %s", path, e) + return DatasetStats.placeholder(multilabel=multilabel) + + train = ds.get("train") or next(iter(ds.values()), None) + if train is None: + return DatasetStats.placeholder(multilabel=multilabel) + + utt_col = getattr(ds, "utterance_feature", "utterance") + sample = train[:1000] if len(train) > 1000 else train[:] + lengths = [len(str(s).split()) for s in sample.get(utt_col, [])] + avg_tokens = int(sum(lengths) / max(1, len(lengths))) if lengths else 32 + p95 = sorted(lengths)[int(len(lengths) * 0.95)] if lengths else avg_tokens * 2 + + return DatasetStats( + n_samples=len(train), + n_classes=getattr(ds, "n_classes", 0) or 0, + avg_tokens=avg_tokens, + p95_tokens=p95, + multilabel=getattr(ds, "multilabel", multilabel), + has_descriptions=getattr(ds, "has_descriptions", None), + source=f"dataset:{path}", + ) + + +def _add_common_dataset_args(p: argparse.ArgumentParser) -> None: + p.add_argument("--dataset", help="Path or hub id of a dataset; overrides placeholders.") + p.add_argument("--n-samples", type=int, default=1_000, help="Placeholder training set size.") + p.add_argument("--n-classes", type=int, default=10, help="Placeholder class count.") + p.add_argument("--avg-tokens", type=int, default=32, help="Placeholder average token length.") + p.add_argument( + "--task", + choices=("multiclass", "multilabel"), + default="multiclass", + help="Placeholder task type when --dataset isn't given.", + ) + + +def cmd_inspect(args: argparse.Namespace) -> int: + config, name = _load_config(args.target) + hardware = detect_hardware( + vram_budget_gb=args.budget_vram_gb, + ) + stats = _stats_from_args(args) + report = run_preflight(config, stats, hardware, preset_name=name) + if args.json: + sys.stdout.write(render_json(report)) + sys.stdout.write("\n") + else: + sys.stdout.write(render_text(report)) + sys.stdout.write("\n") + return 0 if report.is_feasible else 1 + + +def cmd_recommend(args: argparse.Namespace) -> int: + hardware = detect_hardware(vram_budget_gb=args.budget_vram_gb) + stats = _stats_from_args(args) + + results: list[tuple[str, PreflightReport]] = [] + from autointent.utils import load_preset + + for preset in BUNDLED_PRESETS: + try: + cfg = load_preset(preset) # type: ignore[arg-type] + except Exception as e: # noqa: BLE001 + logger.debug("Skipping preset %s: %s", preset, e) + continue + report = run_preflight(cfg, stats, hardware, preset_name=preset) + if args.budget_time_h is not None and report.resource.time_hours > args.budget_time_h: + report.add( + "resource", + report.worst_severity if report.worst_severity.value == "red" else report.worst_severity, # noqa: PLW0125 - explicit + f"Estimated time {report.resource.time_hours:.1f} h exceeds budget {args.budget_time_h} h.", + ) + results.append((preset, report)) + + feasible = [(name, r) for name, r in results if r.is_feasible] + feasible.sort( + key=lambda pair: (-_QUALITY_TIER.get(pair[0], 0), pair[1].resource.time_hours, pair[0]) + ) + chosen = feasible[0][0] if feasible else None + + if args.json: + import json + + out = { + "chosen": chosen, + "results": [ + {"preset": name, "report": r.to_dict()} for name, r in results + ], + } + sys.stdout.write(json.dumps(out, indent=2, default=str)) + sys.stdout.write("\n") + else: + sys.stdout.write(render_recommendation(results, chosen)) + sys.stdout.write("\n") + if chosen: + sys.stdout.write("\n") + sys.stdout.write(render_text(dict(results)[chosen])) + sys.stdout.write("\n") + return 0 if chosen else 1 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="autointent-advisor", + description="Pre-flight feasibility advisor for AutoIntent search-space optimization.", + ) + parser.add_argument("-v", "--verbose", action="store_true", help="Enable debug logging.") + + sub = parser.add_subparsers(dest="cmd", required=True) + + p_inspect = sub.add_parser( + "inspect", + help="Inspect a preset or OptimizationConfig and print a feasibility report.", + ) + p_inspect.add_argument("target", help="Preset name (e.g. transformers-light) or path to a YAML config.") + p_inspect.add_argument("--json", action="store_true", help="Emit a structured JSON report.") + p_inspect.add_argument( + "--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget." + ) + _add_common_dataset_args(p_inspect) + p_inspect.set_defaults(func=cmd_inspect) + + p_rec = sub.add_parser( + "recommend", + help="Detect hardware and recommend the best-fitting bundled preset.", + ) + p_rec.add_argument("--json", action="store_true", help="Emit a structured JSON report.") + p_rec.add_argument( + "--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget." + ) + p_rec.add_argument( + "--budget-time-h", type=float, default=None, help="Optional wall-time ceiling in hours." + ) + _add_common_dataset_args(p_rec) + p_rec.set_defaults(func=cmd_recommend) + + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.WARNING, + format="%(levelname)s %(name)s: %(message)s", + ) + return args.func(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py new file mode 100644 index 000000000..f60f940a6 --- /dev/null +++ b/src/autointent/_advisor/_estimates.py @@ -0,0 +1,382 @@ +"""Resource-phase estimation: walk the search space and aggregate cost. + +Implements an honest worst-case for the modules the proposal lists as +in-scope. Formulas are intentionally coarse — the advisor's contract is +"heuristic upper bound, not measurement". Time and VRAM are the noisiest; +treat them as ballparks, not budgets. +""" + +from __future__ import annotations + +import logging +from typing import Any, Iterable + +from ._hardware import HardwareProfile +from ._hub import ModelMeta, hub_reachable, resolve_model +from ._report import DatasetStats, PreflightReport, ResourceEstimate, Severity + +logger = logging.getLogger(__name__) + +# yellow / red thresholds as fraction of available budget +_YELLOW = 0.7 +_RED = 1.0 + +# rough per-step seconds, keyed on device class. Scaled by params_millions / 100. +_PER_STEP_BASELINE_S = { + "cpu": 0.5, + "low-gpu": 0.04, + "mid-gpu": 0.02, + "high-gpu": 0.01, + "apple-silicon": 0.08, +} + +TRANSFORMER_SCORER_MODULES = {"bert", "lora", "ptuning", "dnnc"} + + +def _extract_model_names(module_entry: dict[str, Any]) -> list[str]: + """Pull model name(s) from a search-space module entry.""" + candidates: list[str] = [] + cfg = module_entry.get("classification_model_config") + if isinstance(cfg, list): + for c in cfg: + if isinstance(c, dict) and c.get("model_name"): + candidates.append(c["model_name"]) + elif isinstance(cfg, dict) and cfg.get("model_name"): + candidates.append(cfg["model_name"]) + embedder_cfg = module_entry.get("embedder_config") + if isinstance(embedder_cfg, list): + for c in embedder_cfg: + if isinstance(c, dict) and c.get("model_name"): + candidates.append(c["model_name"]) + elif isinstance(embedder_cfg, dict) and embedder_cfg.get("model_name"): + candidates.append(embedder_cfg["model_name"]) + return candidates + + +def _max_int(value: Any, default: int) -> int: + if value is None: + return default + if isinstance(value, list) and value: + return max(int(x) for x in value) + if isinstance(value, dict): + return int(value.get("high", default)) + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dict[str, Any]]]: + """Yield (node_type, module_entry) pairs.""" + for node in search_space or []: + node_type = node.get("node_type", "?") + for entry in node.get("search_space", []) or []: + yield node_type, entry + + +def _vram_for_transformer(meta: ModelMeta, mode: str, mixed_precision: bool) -> float: + """VRAM in GB for one trial of a transformer-based module. + + Conservative AMP accounting (the proposal flags the prior naive halving + as too generous; keep optimizer state at fp32 even in AMP). + """ + weights_gb = meta.weights_gb + if mode == "inference": + return weights_gb * 1.3 + if mode == "lora": + return weights_gb * 1.3 + 0.5 + if mode == "reranker": + return weights_gb * 1.5 + # full fine-tune (bert, ptuning, gcn-with-backbone) + if mixed_precision: + # fp16 weights+grads + fp32 master+adam moments + return (weights_gb * 0.5) * 2 + weights_gb * 1 + weights_gb * 2 + return weights_gb * (1 + 1 + 2) + + +def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: + """RAM in GB. Loose upper bound.""" + return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / (1024**3) + + +def _time_for_transformer( + *, + meta: ModelMeta, + n_trials: int, + epochs: int, + batch_size: int, + n_samples: int, + device_class: str, +) -> float: + per_step = _PER_STEP_BASELINE_S[device_class] * (meta.params_millions / 100.0) + steps = max(1, (n_samples // max(1, batch_size))) * epochs + return (n_trials * steps * per_step) / 3600.0 + + +def _classify_severity(estimate: float, budget: float) -> Severity: + if budget <= 0: + return Severity.YELLOW + ratio = estimate / budget + if ratio >= _RED: + return Severity.RED + if ratio >= _YELLOW: + return Severity.YELLOW + return Severity.GREEN + + +def _resource_phase( # noqa: PLR0912 - kept linear for clarity + config: dict[str, Any], + stats: DatasetStats, + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + hpo = config.get("hpo_config") or {} + n_trials = int(hpo.get("n_trials", 1)) + n_jobs = int(hpo.get("n_jobs", 1)) + refit_after = bool(config.get("refit_after", False)) + dump_modules = bool(config.get("dump_modules", False)) + + if not hub_reachable(): + report.low_confidence = True + report.notes.append("HF Hub unreachable — all model sizes are name-pattern heuristics.") + + seen_models: dict[str, ModelMeta] = {} + estimate = ResourceEstimate(parallel_factor=max(1, n_jobs)) + + embedder_cfg = config.get("embedder_config") or {} + global_embedder = embedder_cfg.get("model_name") if isinstance(embedder_cfg, dict) else None + if global_embedder: + seen_models[global_embedder] = resolve_model(global_embedder) + + for node_type, entry in _walk_modules(config.get("search_space") or []): + module = entry.get("module_name", "?") + model_names = _extract_model_names(entry) + if not model_names and global_embedder and module in {"linear", "catboost", "knn", "mlknn"}: + model_names = [global_embedder] + + for name in model_names: + meta = seen_models.setdefault(name, resolve_model(name)) + + mixed_precision = entry.get("dtype") in {"fp16", "bf16"} + if module == "bert": + mode = "full-finetune" + elif module == "lora": + mode = "lora" + elif module == "dnnc": + mode = "reranker" + elif module == "ptuning": + mode = "full-finetune" + else: + mode = "inference" + + batch_size = _max_int(entry.get("batch_size"), 32) + epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) + + vram = _vram_for_transformer(meta, mode, mixed_precision) + ram = _ram_for_module(meta, stats) + + time_h = 0.0 + if mode != "inference": + time_h = _time_for_transformer( + meta=meta, + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + device_class=hardware.device_class, + ) + if refit_after and mode != "inference": + time_h *= 1 + 1.0 / max(1, n_trials) + + estimate.vram_gb = max(estimate.vram_gb, vram) + estimate.ram_gb = max(estimate.ram_gb, ram) + estimate.time_hours += time_h + estimate.drivers.append( + { + "node_type": node_type, + "module": module, + "model": name, + "mode": mode, + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "confidence": meta.confidence, + } + ) + + for meta in seen_models.values(): + if meta.cached_locally: + estimate.disk_cached_gb += meta.disk_gb + else: + estimate.disk_download_gb += meta.disk_gb + + if dump_modules: + weights_total = sum(m.weights_gb for m in seen_models.values()) + estimate.disk_dump_gb = weights_total * n_trials + + if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: + effective_vram = estimate.vram_gb * n_jobs + else: + effective_vram = estimate.vram_gb + + report.resource = estimate + + # render findings + vram_sev = _classify_severity(effective_vram, hardware.vram_gb) + if hardware.accelerator == "cpu" and effective_vram > 0: + report.add( + "resource", + Severity.YELLOW, + f"No GPU detected; transformer modules will be very slow (worst case ~{estimate.time_hours:.1f} h).", + metric="vram", + ) + else: + msg = f"VRAM ~{effective_vram:.1f} GB" + if n_jobs > 1: + msg += f" (= per-trial {estimate.vram_gb:.1f} GB × {n_jobs} parallel trials)" + msg += f" vs available {hardware.vram_gb:.1f} GB" + report.add("resource", vram_sev, msg, metric="vram") + + ram_sev = _classify_severity(estimate.ram_gb, hardware.ram_gb) + report.add( + "resource", + ram_sev, + f"RAM ~{estimate.ram_gb:.1f} GB vs available {hardware.ram_gb:.1f} GB", + metric="ram", + ) + + disk_total = estimate.disk_download_gb + estimate.disk_dump_gb + disk_sev = _classify_severity(disk_total, hardware.free_disk_gb) + disk_msg = f"Disk ~{estimate.disk_download_gb:.1f} GB to download" + if estimate.disk_cached_gb > 0: + disk_msg += f", {estimate.disk_cached_gb:.1f} GB already cached" + if estimate.disk_dump_gb > 0: + disk_msg += f", +{estimate.disk_dump_gb:.1f} GB during training (dump_modules=True)" + disk_msg += f" vs {hardware.free_disk_gb:.0f} GB free" + report.add("resource", disk_sev, disk_msg, metric="disk") + + if estimate.time_hours > 0: + time_msg = f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)" + report.add("resource", Severity.GREEN, time_msg, metric="time") + + +def _config_phase( + config: dict[str, Any], + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + hpo = config.get("hpo_config") or {} + n_jobs = int(hpo.get("n_jobs", 1)) + + if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: + report.add( + "config", + Severity.YELLOW, + f"hpo_config.n_jobs={n_jobs} on a single GPU multiplies VRAM demand by {n_jobs}×.", + ) + + uses_catboost_gpu = False + for _, entry in _walk_modules(config.get("search_space") or []): + if entry.get("module_name") == "catboost" and entry.get("task_type") == "GPU": + uses_catboost_gpu = True + break + if uses_catboost_gpu and hardware.accelerator != "cuda": + report.add( + "config", + Severity.YELLOW, + "CatBoost task_type=GPU configured but no CUDA detected — will fall back to CPU.", + ) + + +def _data_phase( + config: dict[str, Any], + stats: DatasetStats, + report: PreflightReport, +) -> None: + # token-length truncation (heuristic — we use stats.p95_tokens vs configured max_length) + p95 = stats.p95_tokens or int(stats.avg_tokens * 2.5) + for _, entry in _walk_modules(config.get("search_space") or []): + max_len_value = entry.get("max_length") + if max_len_value is None: + continue + max_len = _max_int(max_len_value, 512) + if p95 > max_len: + severity = Severity.RED if p95 > max_len * 1.5 else Severity.YELLOW + report.add( + "data", + severity, + f"Train tokens p95~{p95} exceeds {entry.get('module_name', '?')}.max_length={max_len}; expect silent truncation.", + ) + + # rare class × linear-CV + has_linear = any( + e.get("module_name") == "linear" for _, e in _walk_modules(config.get("search_space") or []) + ) + if has_linear and stats.rare_classes: + report.add( + "data", + Severity.RED, + ( + "LogisticRegressionCV (cv=3) will fail: classes " + f"{stats.rare_classes[:5]} have <3 samples." + ), + ) + + # partial descriptions × description scorer + has_description = any( + e.get("module_name") == "description" + for _, e in _walk_modules(config.get("search_space") or []) + ) + if has_description and stats.has_descriptions is False: + report.add( + "data", + Severity.RED, + "description scorer present but intent descriptions are missing — fill them in or drop the scorer.", + ) + + +def run_preflight( + config: dict[str, Any], + stats: DatasetStats, + hardware: HardwareProfile, + *, + preset_name: str | None = None, +) -> PreflightReport: + """Run all three phases and return one report. + + Args: + config: parsed preset / OptimizationConfig dict (top-level keys: + ``search_space``, ``hpo_config``, optional ``embedder_config``). + stats: dataset statistics (real or placeholder). + hardware: detected hardware profile. + preset_name: optional friendly name for the report header. + + Returns: + PreflightReport with findings across resource/data/config phases. + """ + report = PreflightReport( + preset_name=preset_name, + hardware={ + "accelerator": hardware.accelerator, + "device_name": hardware.device_name, + "vram_gb": round(hardware.vram_gb, 2), + "ram_gb": round(hardware.ram_gb, 2), + "free_disk_gb": round(hardware.free_disk_gb, 2), + "device_class": hardware.device_class, + }, + dataset={ + "n_samples": stats.n_samples, + "n_classes": stats.n_classes, + "avg_tokens": stats.avg_tokens, + "p95_tokens": stats.p95_tokens, + "multilabel": stats.multilabel, + "source": stats.source, + }, + ) + report.notes.extend(hardware.notes) + + _resource_phase(config, stats, hardware, report) + _data_phase(config, stats, report) + _config_phase(config, hardware, report) + + return report diff --git a/src/autointent/_advisor/_hardware.py b/src/autointent/_advisor/_hardware.py new file mode 100644 index 000000000..2bda6120f --- /dev/null +++ b/src/autointent/_advisor/_hardware.py @@ -0,0 +1,160 @@ +"""Local hardware detection. + +Probes CPU / RAM / disk and the highest-priority accelerator available +(CUDA → MPS → CPU). All probes are wrapped to fall back safely on a +broken install (e.g. CUDA driver mismatch) rather than crash the advisor. +""" + +from __future__ import annotations + +import logging +import os +import platform +import shutil +from dataclasses import dataclass, field +from typing import Literal + +logger = logging.getLogger(__name__) + +Accelerator = Literal["cuda", "mps", "cpu"] + +# matches macOS PYTORCH_MPS_HIGH_WATERMARK_RATIO default +MPS_DEFAULT_BUDGET_RATIO = 0.7 + + +@dataclass +class HardwareProfile: + accelerator: Accelerator + device_name: str + vram_gb: float + ram_gb: float + free_disk_gb: float + cpu_count: int + notes: list[str] = field(default_factory=list) + + @property + def device_class(self) -> str: + if self.accelerator == "cpu": + return "cpu" + if self.accelerator == "mps": + return "apple-silicon" + if self.vram_gb >= 24: + return "high-gpu" + if self.vram_gb >= 12: + return "mid-gpu" + return "low-gpu" + + +def _detect_ram_gb() -> float: + try: + import psutil + + return psutil.virtual_memory().total / (1024**3) + except ImportError: + logger.debug("psutil unavailable; RAM unknown") + return 0.0 + + +def _detect_free_disk_gb(path: str | None = None) -> float: + cache = path or os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface") + probe_path = cache if os.path.exists(cache) else os.path.expanduser("~") + try: + usage = shutil.disk_usage(probe_path) + return usage.free / (1024**3) + except OSError as e: + logger.debug("disk usage probe failed at %s: %s", probe_path, e) + return 0.0 + + +def _detect_cuda() -> tuple[float, str] | None: + try: + import torch + + if not torch.cuda.is_available(): + return None + idx = 0 + try: + free, total = torch.cuda.mem_get_info(idx) + vram_gb = total / (1024**3) + except (RuntimeError, AttributeError) as e: + logger.debug("torch.cuda.mem_get_info failed: %s", e) + return None + name = torch.cuda.get_device_name(idx) + return vram_gb, name + except ImportError: + return None + except Exception as e: # noqa: BLE001 - protect the advisor from torch quirks + logger.debug("CUDA detection raised: %s", e) + return None + + +def _detect_mps(ram_gb: float, budget_ratio: float = MPS_DEFAULT_BUDGET_RATIO) -> tuple[float, str] | None: + try: + import torch + + if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): + return None + # apple silicon: unified memory; budget is fraction of total RAM + return ram_gb * budget_ratio, f"Apple Silicon ({platform.machine()})" + except ImportError: + return None + except Exception as e: # noqa: BLE001 + logger.debug("MPS detection raised: %s", e) + return None + + +def detect_hardware( + *, + vram_budget_gb: float | None = None, + mps_budget_ratio: float = MPS_DEFAULT_BUDGET_RATIO, +) -> HardwareProfile: + """Detect the local hardware, with optional manual overrides. + + Args: + vram_budget_gb: when set, overrides the detected VRAM (use for + shared-GPU machines where part of the device is taken). + mps_budget_ratio: fraction of total RAM treated as the MPS + "VRAM" budget on Apple Silicon. + + Returns: + HardwareProfile reflecting current machine state. + """ + notes: list[str] = [] + ram_gb = _detect_ram_gb() + free_disk_gb = _detect_free_disk_gb() + cpu_count = os.cpu_count() or 1 + + cuda = _detect_cuda() + if cuda is not None: + vram_gb, device_name = cuda + accel: Accelerator = "cuda" + else: + mps = _detect_mps(ram_gb, mps_budget_ratio) + if mps is not None: + vram_gb, device_name = mps + accel = "mps" + notes.append( + f"MPS unified memory: VRAM budget = {mps_budget_ratio:.0%} of RAM." + ) + else: + vram_gb = 0.0 + device_name = platform.processor() or "cpu" + accel = "cpu" + + if vram_budget_gb is not None: + if vram_gb and vram_budget_gb > vram_gb: + notes.append( + f"Manual --budget-vram-gb={vram_budget_gb} exceeds detected {vram_gb:.1f} GB; using override." + ) + notes.append(f"Using manual VRAM budget: {vram_budget_gb} GB.") + vram_gb = vram_budget_gb + + return HardwareProfile( + accelerator=accel, + device_name=device_name, + vram_gb=vram_gb, + ram_gb=ram_gb, + free_disk_gb=free_disk_gb, + cpu_count=cpu_count, + notes=notes, + ) diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py new file mode 100644 index 000000000..80ccb7133 --- /dev/null +++ b/src/autointent/_advisor/_hub.py @@ -0,0 +1,183 @@ +"""HF Hub metadata lookups + warm-cache probe. + +Memoized per-process. Offline-safe: every probe falls back to a +heuristic value rather than raising. The advisor flips the report's +``low_confidence`` flag when a fallback is taken. +""" + +from __future__ import annotations + +import logging +import os +import re +from dataclasses import dataclass +from functools import lru_cache +from typing import Any + +logger = logging.getLogger(__name__) + +# Coarse heuristic estimates keyed on name fragments. Used only when HF Hub +# is unreachable and we can't get safetensors metadata. Values in millions. +_NAME_HEURISTICS = [ + (re.compile(r"(?i)(deberta|roberta|bert).*(xxlarge|huge)"), 1_500), + (re.compile(r"(?i)(deberta|roberta|bert).*xlarge"), 750), + (re.compile(r"(?i)(deberta|roberta|bert).*large"), 350), + (re.compile(r"(?i)e5.*large"), 560), + (re.compile(r"(?i)e5.*small"), 33), + (re.compile(r"(?i)mpnet"), 110), + (re.compile(r"(?i)minilm"), 33), + (re.compile(r"(?i)distil"), 66), + (re.compile(r"(?i)small"), 60), + (re.compile(r"(?i)base"), 110), + (re.compile(r"(?i)large"), 350), +] + + +@dataclass +class ModelMeta: + name: str + params_millions: float + weight_bytes_per_param: int + total_file_bytes: int + cached_locally: bool + confidence: str # "hub" | "heuristic" + + @property + def disk_gb(self) -> float: + return self.total_file_bytes / (1024**3) + + @property + def weights_gb(self) -> float: + return (self.params_millions * 1_000_000 * self.weight_bytes_per_param) / (1024**3) + + +@lru_cache(maxsize=1) +def hub_reachable(timeout_s: float = 2.0) -> bool: + """Single up-front probe. Memoized per process.""" + try: + from huggingface_hub import HfApi + + HfApi().list_models(limit=1) + except ImportError: + logger.debug("huggingface_hub not installed; assuming offline") + return False + except Exception as e: # noqa: BLE001 + logger.debug("HF Hub probe failed: %s", e) + return False + else: + return True + + +def _heuristic_params_millions(model_name: str) -> float: + for pattern, m in _NAME_HEURISTICS: + if pattern.search(model_name): + return float(m) + return 110.0 # generic BERT-base default + + +def _is_warm_cached(model_name: str) -> bool: + """True when the weight shard is present in the local HF cache.""" + try: + from huggingface_hub import scan_cache_dir, try_to_load_from_cache + except ImportError: + return False + + weight_files = ["model.safetensors", "pytorch_model.bin", "model.safetensors.index.json"] + for fname in weight_files: + path = try_to_load_from_cache(model_name, fname) + if path is not None and path is not False: + return True + + # sharded models won't match the single-file probe; fall back to a scan + try: + cache = scan_cache_dir() + except Exception as e: # noqa: BLE001 + logger.debug("scan_cache_dir failed: %s", e) + return False + return any(repo.repo_id == model_name for repo in cache.repos) + + +def _hub_metadata(model_name: str) -> ModelMeta | None: + try: + from huggingface_hub import HfApi + except ImportError: + return None + + try: + info = HfApi().model_info(model_name, files_metadata=True) + except Exception as e: # noqa: BLE001 + logger.debug("model_info(%s) failed: %s", model_name, e) + return None + + params_millions = 0.0 + weight_bytes_per_param = 4 + safetensors = getattr(info, "safetensors", None) + if safetensors is not None: + params_total = getattr(safetensors, "total", None) or sum( + getattr(safetensors, "parameters", {}).values() or [0] + ) + if params_total: + params_millions = params_total / 1_000_000 + params_map: dict[str, Any] = getattr(safetensors, "parameters", {}) or {} + if any("F16" in k or "BF16" in k for k in params_map): + weight_bytes_per_param = 2 + + total_file_bytes = 0 + for sibling in getattr(info, "siblings", []) or []: + size = getattr(sibling, "size", None) + if size: + total_file_bytes += int(size) + + if params_millions == 0: + params_millions = _heuristic_params_millions(model_name) + + if total_file_bytes == 0: + total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) + + return ModelMeta( + name=model_name, + params_millions=params_millions, + weight_bytes_per_param=weight_bytes_per_param, + total_file_bytes=total_file_bytes, + cached_locally=_is_warm_cached(model_name), + confidence="hub", + ) + + +def _heuristic_metadata(model_name: str) -> ModelMeta: + params_millions = _heuristic_params_millions(model_name) + weight_bytes_per_param = 4 + total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) + return ModelMeta( + name=model_name, + params_millions=params_millions, + weight_bytes_per_param=weight_bytes_per_param, + total_file_bytes=total_file_bytes, + cached_locally=_is_warm_cached(model_name), + confidence="heuristic", + ) + + +@lru_cache(maxsize=64) +def resolve_model(model_name: str) -> ModelMeta: + """Resolve metadata for a single model name. Memoized per process. + + Always returns a value — never raises — so the advisor can keep going + on offline machines or for unknown checkpoints. + """ + if model_name.startswith("local:") or os.path.isabs(model_name): + return ModelMeta( + name=model_name, + params_millions=_heuristic_params_millions(model_name), + weight_bytes_per_param=4, + total_file_bytes=0, + cached_locally=True, + confidence="heuristic", + ) + + if hub_reachable(): + meta = _hub_metadata(model_name) + if meta is not None: + return meta + + return _heuristic_metadata(model_name) diff --git a/src/autointent/_advisor/_render.py b/src/autointent/_advisor/_render.py new file mode 100644 index 000000000..52168aa75 --- /dev/null +++ b/src/autointent/_advisor/_render.py @@ -0,0 +1,104 @@ +"""Rendering for the pre-flight report. + +Text output is grouped by phase (Resource / Data / Config) plus a Drivers +section and the always-on disclaimer. JSON output dumps the structured +report straight through. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ._report import PreflightReport + +_SEVERITY_TAG = {"green": "✓", "yellow": "⚠", "red": "✗"} +_PHASE_ORDER = ("resource", "data", "config") +_PHASE_LABEL = {"resource": "Resource", "data": "Data", "config": "Config"} + + +def render_text(report: "PreflightReport") -> str: + lines: list[str] = [] + title = "Compute feasibility check" + if report.preset_name: + title += f" — {report.preset_name}" + lines.append(title) + lines.append("─" * len(title)) + + hw = report.hardware + lines.append( + f"Hardware: {hw.get('accelerator', '?')} ({hw.get('device_name', '?')})," + f" {hw.get('vram_gb', 0):.1f} GB VRAM, {hw.get('ram_gb', 0):.0f} GB RAM," + f" {hw.get('free_disk_gb', 0):.0f} GB free disk" + ) + ds = report.dataset + lines.append( + f"Dataset: n_samples={ds.get('n_samples')}, n_classes={ds.get('n_classes')}," + f" avg_tokens={ds.get('avg_tokens')} ({ds.get('source')})" + ) + lines.append("") + + for phase in _PHASE_ORDER: + bucket = [f for f in report.findings if f.phase == phase] + if not bucket: + continue + lines.append(f"{_PHASE_LABEL[phase]}:") + for f in bucket: + tag = _SEVERITY_TAG.get(f.severity.value, "·") + lines.append(f" {tag} {f.message}") + lines.append("") + + if report.resource.drivers: + lines.append("Drivers of cost:") + for d in report.resource.drivers[:8]: + lines.append( + f" {d['node_type']}.{d['module']:<10} {d['model']:<48}" + f" {d['mode']:<14} VRAM ~{d['vram_gb']} GB, time ~{d['time_hours']} h" + f" [{d['confidence']}]" + ) + if len(report.resource.drivers) > 8: + lines.append(f" … and {len(report.resource.drivers) - 8} more") + lines.append("") + + if report.notes: + lines.append("Notes:") + for note in report.notes: + lines.append(f" • {note}") + lines.append("") + + summary = f"Verdict: {'feasible' if report.is_feasible else 'INFEASIBLE'} " + summary += f"(worst severity: {report.worst_severity.value})" + if report.low_confidence: + summary += " — low-confidence (heuristic fallback in use)" + lines.append(summary) + lines.append("Note: estimates are heuristic upper bounds, not measurements.") + return "\n".join(lines) + + +def render_json(report: "PreflightReport") -> str: + return json.dumps(report.to_dict(), indent=2, default=str) + + +def render_recommendation( + results: list[tuple[str, "PreflightReport"]], + chosen: str | None, +) -> str: + """Compact table for the ``recommend`` subcommand.""" + lines = ["", "Recommendation:"] + if chosen: + lines.append(f" → {chosen}") + else: + lines.append(" → none of the bundled presets fit your hardware as-is.") + lines.append("") + lines.append(f"{'Preset':<24} {'Status':<14} {'VRAM':<10} {'Time':<10} {'Worst':<8}") + lines.append("-" * 68) + for name, report in results: + verdict = "feasible" if report.is_feasible else "infeasible" + lines.append( + f"{name:<24} {verdict:<14} " + f"{report.resource.vram_gb:>4.1f} GB " + f"{report.resource.time_hours:>4.1f} h " + f"{report.worst_severity.value:<8}" + ) + return "\n".join(lines) diff --git a/src/autointent/_advisor/_report.py b/src/autointent/_advisor/_report.py new file mode 100644 index 000000000..0250482a5 --- /dev/null +++ b/src/autointent/_advisor/_report.py @@ -0,0 +1,113 @@ +"""Dataclasses for the pre-flight advisor's structured report.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from enum import Enum +from typing import Any, Literal + + +class Severity(str, Enum): + GREEN = "green" + YELLOW = "yellow" + RED = "red" + + +Phase = Literal["resource", "data", "config"] + + +@dataclass(frozen=True) +class Finding: + """A single advisor finding rendered as one line in the summary.""" + + phase: Phase + severity: Severity + message: str + metric: str | None = None + + +@dataclass +class ResourceEstimate: + """Aggregated resource numbers across the search space.""" + + disk_download_gb: float = 0.0 + disk_cached_gb: float = 0.0 + disk_dump_gb: float = 0.0 + ram_gb: float = 0.0 + vram_gb: float = 0.0 + time_hours: float = 0.0 + parallel_factor: int = 1 + drivers: list[dict[str, Any]] = field(default_factory=list) + + @property + def total_disk_gb(self) -> float: + return self.disk_download_gb + self.disk_dump_gb + + +@dataclass +class DatasetStats: + """Minimal stats the advisor needs about the user's dataset. + + Built either from a real ``Dataset`` or from CLI placeholder flags. + """ + + n_samples: int + n_classes: int + avg_tokens: int + p95_tokens: int | None = None + multilabel: bool = False + has_descriptions: bool | None = None + rare_classes: list[str] = field(default_factory=list) + source: str = "placeholder" + + @classmethod + def placeholder( + cls, + n_samples: int = 1_000, + n_classes: int = 10, + avg_tokens: int = 32, + multilabel: bool = False, + ) -> "DatasetStats": + return cls( + n_samples=n_samples, + n_classes=n_classes, + avg_tokens=avg_tokens, + p95_tokens=int(avg_tokens * 2.5), + multilabel=multilabel, + ) + + +@dataclass +class PreflightReport: + """One report covering all three phases.""" + + findings: list[Finding] = field(default_factory=list) + resource: ResourceEstimate = field(default_factory=ResourceEstimate) + hardware: dict[str, Any] = field(default_factory=dict) + dataset: dict[str, Any] = field(default_factory=dict) + preset_name: str | None = None + low_confidence: bool = False + notes: list[str] = field(default_factory=list) + + def add(self, phase: Phase, severity: Severity, message: str, metric: str | None = None) -> None: + self.findings.append(Finding(phase=phase, severity=severity, message=message, metric=metric)) + + @property + def worst_severity(self) -> Severity: + order = {Severity.GREEN: 0, Severity.YELLOW: 1, Severity.RED: 2} + if not self.findings: + return Severity.GREEN + return max((f.severity for f in self.findings), key=lambda s: order[s]) + + @property + def is_feasible(self) -> bool: + return self.worst_severity != Severity.RED + + def to_dict(self) -> dict[str, Any]: + d = asdict(self) + d["findings"] = [ + {**asdict(f), "severity": f.severity.value} for f in self.findings + ] + d["worst_severity"] = self.worst_severity.value + d["is_feasible"] = self.is_feasible + return d diff --git a/tests/advisor/__init__.py b/tests/advisor/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py new file mode 100644 index 000000000..18c2615a6 --- /dev/null +++ b/tests/advisor/test_estimates_and_cli.py @@ -0,0 +1,198 @@ +"""End-to-end smoke tests for the advisor. + +These run offline — HF Hub probes are monkeypatched to fail so the +advisor falls back to its name-pattern heuristics. Verifies that: + +* every bundled preset can be inspected without raising; +* the recommend subcommand picks something on a generous budget and + nothing on a hostile one; +* ``--json`` emits parseable JSON. +""" + +from __future__ import annotations + +import json +import sys + +import pytest + +from autointent._advisor import DatasetStats, HardwareProfile, run_preflight +from autointent._advisor._cli import BUNDLED_PRESETS, main +from autointent.utils import load_preset + + +@pytest.fixture(autouse=True) +def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: + """Pin the HF Hub probe to "offline" so tests don't hit the network.""" + from autointent._advisor import _estimates, _hub + + _hub.hub_reachable.cache_clear() + _hub.resolve_model.cache_clear() + offline = lambda *_a, **_kw: False # noqa: E731 + monkeypatch.setattr(_hub, "hub_reachable", offline) + monkeypatch.setattr(_estimates, "hub_reachable", offline) + + +def _profile(vram_gb: float = 16.0) -> HardwareProfile: + return HardwareProfile( + accelerator="cuda" if vram_gb > 0 else "cpu", + device_name="test-gpu" if vram_gb > 0 else "test-cpu", + vram_gb=vram_gb, + ram_gb=32.0, + free_disk_gb=200.0, + cpu_count=8, + ) + + +@pytest.mark.parametrize("preset", BUNDLED_PRESETS) +def test_every_preset_inspects_without_raising(preset: str) -> None: + cfg = load_preset(preset) # type: ignore[arg-type] + stats = DatasetStats.placeholder(n_samples=500, n_classes=10, avg_tokens=24) + report = run_preflight(cfg, stats, _profile(vram_gb=16.0), preset_name=preset) + assert report.preset_name == preset + assert report.low_confidence is True # we forced offline + # always at least one resource-phase finding + assert any(f.phase == "resource" for f in report.findings) + + +def test_heavy_preset_is_infeasible_on_2gb_budget() -> None: + cfg = load_preset("transformers-heavy") # type: ignore[arg-type] + stats = DatasetStats.placeholder(n_samples=5000, n_classes=20, avg_tokens=40) + report = run_preflight(cfg, stats, _profile(vram_gb=2.0), preset_name="transformers-heavy") + assert not report.is_feasible, "deberta-v3-large should not fit in 2 GB" + + +def test_light_preset_is_feasible_on_8gb_budget() -> None: + cfg = load_preset("transformers-light") # type: ignore[arg-type] + stats = DatasetStats.placeholder(n_samples=1000, n_classes=10, avg_tokens=24) + report = run_preflight(cfg, stats, _profile(vram_gb=8.0), preset_name="transformers-light") + assert report.is_feasible + + +def test_n_jobs_doubles_vram_findings() -> None: + cfg = load_preset("transformers-light") # type: ignore[arg-type] + cfg = {**cfg, "hpo_config": {**(cfg.get("hpo_config") or {}), "n_jobs": 4}} + stats = DatasetStats.placeholder() + report = run_preflight(cfg, stats, _profile(vram_gb=4.0)) + assert any("parallel trials" in f.message for f in report.findings) + assert any(f.phase == "config" and "n_jobs" in f.message for f in report.findings) + + +def test_cli_inspect_json_is_parseable(capsys: pytest.CaptureFixture[str]) -> None: + rc = main( + [ + "inspect", + "transformers-light", + "--n-samples", + "500", + "--n-classes", + "5", + "--avg-tokens", + "20", + "--json", + "--budget-vram-gb", + "16", + ] + ) + captured = capsys.readouterr() + payload = json.loads(captured.out) + assert payload["preset_name"] == "transformers-light" + assert "findings" in payload + assert payload["worst_severity"] in {"green", "yellow", "red"} + # rc is 0 on feasible, 1 otherwise + assert rc in (0, 1) + + +def test_cli_inspect_text_runs(capsys: pytest.CaptureFixture[str]) -> None: + main( + [ + "inspect", + "transformers-light", + "--n-samples", + "200", + "--n-classes", + "5", + "--avg-tokens", + "15", + "--budget-vram-gb", + "16", + ] + ) + out = capsys.readouterr().out + assert "Compute feasibility check" in out + assert "Verdict:" in out + + +def test_cli_recommend_picks_a_preset_on_generous_hardware( + capsys: pytest.CaptureFixture[str], +) -> None: + rc = main( + [ + "recommend", + "--n-samples", + "1000", + "--n-classes", + "10", + "--avg-tokens", + "20", + "--budget-vram-gb", + "24", + ] + ) + out = capsys.readouterr().out + assert "Recommendation:" in out + assert rc == 0 + + +def test_partial_descriptions_with_description_scorer_flags_red() -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "description"}, + ], + } + ], + } + stats = DatasetStats( + n_samples=500, + n_classes=10, + avg_tokens=24, + has_descriptions=False, + ) + report = run_preflight(cfg, stats, _profile(vram_gb=16.0)) + assert any( + f.phase == "data" and "description" in f.message.lower() for f in report.findings + ) + + +def test_long_dataset_triggers_truncation_warning() -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [ + {"model_name": "microsoft/deberta-v3-small"} + ], + "max_length": [128], + } + ], + } + ], + } + stats = DatasetStats( + n_samples=500, + n_classes=10, + avg_tokens=80, + p95_tokens=512, # well over 128 + ) + report = run_preflight(cfg, stats, _profile(vram_gb=16.0)) + assert any("truncation" in f.message.lower() for f in report.findings) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py new file mode 100644 index 000000000..0317ff27b --- /dev/null +++ b/tests/advisor/test_estimates_internals.py @@ -0,0 +1,319 @@ +"""Targeted tests for `_estimates` helpers + edge cases of `run_preflight`.""" + +from __future__ import annotations + +import pytest + +from autointent._advisor import _estimates, _hub +from autointent._advisor._estimates import ( + _classify_severity, + _extract_model_names, + _max_int, + _ram_for_module, + _vram_for_transformer, + run_preflight, +) +from autointent._advisor._hardware import HardwareProfile +from autointent._advisor._hub import ModelMeta +from autointent._advisor._report import DatasetStats, Severity + + +@pytest.fixture(autouse=True) +def _offline(monkeypatch: pytest.MonkeyPatch) -> None: + _hub.hub_reachable.cache_clear() + _hub.resolve_model.cache_clear() + offline = lambda *_a, **_kw: False # noqa: E731 + monkeypatch.setattr(_hub, "hub_reachable", offline) + monkeypatch.setattr(_estimates, "hub_reachable", offline) + monkeypatch.setattr(_hub, "_is_warm_cached", lambda _name: False) + + +def _profile(vram_gb: float = 16.0, accelerator: str = "cuda") -> HardwareProfile: + return HardwareProfile( + accelerator=accelerator, # type: ignore[arg-type] + device_name=f"test-{accelerator}", + vram_gb=vram_gb, + ram_gb=32.0, + free_disk_gb=200.0, + cpu_count=8, + ) + + +class TestMaxInt: + def test_none_returns_default(self) -> None: + assert _max_int(None, 7) == 7 + + def test_list_picks_max(self) -> None: + assert _max_int([1, 5, 3], 0) == 5 + + def test_range_dict_uses_high(self) -> None: + assert _max_int({"low": 1, "high": 9}, 0) == 9 + + def test_scalar_int_passes_through(self) -> None: + assert _max_int(42, 0) == 42 + + def test_garbage_returns_default(self) -> None: + assert _max_int("not-a-number", 11) == 11 + + +class TestExtractModelNames: + def test_classification_model_config_as_list(self) -> None: + entry = {"classification_model_config": [{"model_name": "foo/bar"}]} + assert _extract_model_names(entry) == ["foo/bar"] + + def test_classification_model_config_as_dict(self) -> None: + entry = {"classification_model_config": {"model_name": "foo/bar"}} + assert _extract_model_names(entry) == ["foo/bar"] + + def test_embedder_config_picked_up(self) -> None: + entry = {"embedder_config": [{"model_name": "e/b"}]} + assert _extract_model_names(entry) == ["e/b"] + + def test_multiple_choices_all_returned(self) -> None: + entry = { + "classification_model_config": [ + {"model_name": "a/x"}, + {"model_name": "b/y"}, + ] + } + assert _extract_model_names(entry) == ["a/x", "b/y"] + + def test_empty_entry(self) -> None: + assert _extract_model_names({}) == [] + + +class TestClassifySeverity: + def test_below_yellow_is_green(self) -> None: + assert _classify_severity(estimate=1.0, budget=10.0) == Severity.GREEN + + def test_above_yellow_threshold(self) -> None: + assert _classify_severity(estimate=8.0, budget=10.0) == Severity.YELLOW + + def test_at_or_above_red_threshold(self) -> None: + assert _classify_severity(estimate=10.0, budget=10.0) == Severity.RED + assert _classify_severity(estimate=12.0, budget=10.0) == Severity.RED + + def test_zero_budget_returns_yellow(self) -> None: + assert _classify_severity(estimate=1.0, budget=0.0) == Severity.YELLOW + + +class TestVramForTransformer: + @pytest.fixture + def meta(self) -> ModelMeta: + return ModelMeta( + name="x", + params_millions=100.0, + weight_bytes_per_param=4, + total_file_bytes=0, + cached_locally=False, + confidence="hub", + ) + + def test_full_finetune_is_larger_than_lora_is_larger_than_inference( + self, meta: ModelMeta + ) -> None: + inference = _vram_for_transformer(meta, "inference", mixed_precision=False) + lora = _vram_for_transformer(meta, "lora", mixed_precision=False) + full = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) + assert inference < lora < full + + def test_amp_does_not_naively_halve(self, meta: ModelMeta) -> None: + """The proposal calls out that AMP doesn't halve total VRAM — fp32 master + weights and Adam moments don't shrink. Weight-side accounting comes out + equal to fp32; the only savings (activations) aren't modeled by us.""" + full_fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) + full_amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True) + assert full_amp / full_fp32 == pytest.approx(1.0) + assert full_amp / full_fp32 > 0.5 # explicit check vs the naive-halving formula + + def test_reranker_uses_inference_class(self, meta: ModelMeta) -> None: + inference = _vram_for_transformer(meta, "inference", mixed_precision=False) + reranker = _vram_for_transformer(meta, "reranker", mixed_precision=False) + assert reranker > inference + + +def test_ram_scales_with_dataset_size() -> None: + meta = ModelMeta( + name="x", + params_millions=100.0, + weight_bytes_per_param=4, + total_file_bytes=0, + cached_locally=False, + confidence="hub", + ) + small = _ram_for_module(meta, DatasetStats.placeholder(n_samples=100)) + big = _ram_for_module(meta, DatasetStats.placeholder(n_samples=10_000_000, avg_tokens=128)) + assert big > small + + +class TestRunPreflightFeatures: + def test_dump_modules_adds_disk_during_training(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [ + {"model_name": "microsoft/deberta-v3-small"} + ], + "num_train_epochs": [3], + "batch_size": [16], + } + ], + } + ], + "hpo_config": {"n_trials": 5}, + "dump_modules": True, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + assert report.resource.disk_dump_gb > 0 + assert any("during training" in f.message for f in report.findings) + + def test_refit_after_increases_time(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [ + {"model_name": "microsoft/deberta-v3-small"} + ], + "num_train_epochs": [3], + "batch_size": [16], + } + ], + } + ], + "hpo_config": {"n_trials": 10}, + } + baseline = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + cfg_refit = {**cfg, "refit_after": True} + bumped = run_preflight(cfg_refit, DatasetStats.placeholder(), _profile()) + assert bumped.resource.time_hours > baseline.resource.time_hours + + def test_catboost_gpu_without_cuda_flags_config(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "catboost", "task_type": "GPU"}, + ], + } + ], + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile(accelerator="cpu")) + assert any( + f.phase == "config" and "CatBoost" in f.message for f in report.findings + ) + + def test_catboost_gpu_with_cuda_is_silent(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "catboost", "task_type": "GPU"}, + ], + } + ], + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile(accelerator="cuda")) + assert not any( + f.phase == "config" and "CatBoost" in f.message for f in report.findings + ) + + def test_offline_flips_low_confidence(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "any/model"}], + } + ], + } + ] + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + assert report.low_confidence is True + assert any("HF Hub unreachable" in n for n in report.notes) + + def test_rare_classes_with_linear_scorer_flag_red(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear"}, + ], + } + ] + } + stats = DatasetStats( + n_samples=20, + n_classes=5, + avg_tokens=10, + rare_classes=["intent_a", "intent_b"], + ) + report = run_preflight(cfg, stats, _profile()) + assert any( + f.phase == "data" and "LogisticRegressionCV" in f.message and f.severity == Severity.RED + for f in report.findings + ) + + def test_truncation_red_when_p95_dominates_max_length(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "max_length": [128], + "classification_model_config": [ + {"model_name": "some/model"} + ], + } + ], + } + ] + } + stats = DatasetStats(n_samples=500, n_classes=5, avg_tokens=50, p95_tokens=400) + report = run_preflight(cfg, stats, _profile()) + red = [f for f in report.findings if f.phase == "data" and f.severity == Severity.RED] + assert red, "p95=400 > 1.5 * max_length=128 should be red" + + def test_truncation_yellow_when_p95_only_slightly_exceeds(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "max_length": [128], + "classification_model_config": [ + {"model_name": "some/model"} + ], + } + ], + } + ] + } + stats = DatasetStats(n_samples=500, n_classes=5, avg_tokens=50, p95_tokens=140) + report = run_preflight(cfg, stats, _profile()) + yellows = [ + f + for f in report.findings + if f.phase == "data" + and f.severity == Severity.YELLOW + and "truncation" in f.message.lower() + ] + assert yellows diff --git a/tests/advisor/test_hardware_detection.py b/tests/advisor/test_hardware_detection.py new file mode 100644 index 000000000..d8131fb19 --- /dev/null +++ b/tests/advisor/test_hardware_detection.py @@ -0,0 +1,72 @@ +"""Hardware detection has to be safe on every machine — broken CUDA, no GPU, +no psutil. Verify the fallbacks work without raising. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from autointent._advisor._hardware import detect_hardware + + +def test_cpu_fallback_when_no_accelerator() -> None: + with ( + patch("autointent._advisor._hardware._detect_cuda", return_value=None), + patch("autointent._advisor._hardware._detect_mps", return_value=None), + ): + hw = detect_hardware() + assert hw.accelerator == "cpu" + assert hw.vram_gb == 0.0 + assert hw.device_class == "cpu" + + +def test_cuda_branch_classifies_low_gpu() -> None: + with ( + patch( + "autointent._advisor._hardware._detect_cuda", + return_value=(8.0, "NVIDIA RTX 3060"), + ), + ): + hw = detect_hardware() + assert hw.accelerator == "cuda" + assert hw.vram_gb == pytest.approx(8.0) + assert hw.device_class == "low-gpu" + + +def test_mps_budget_uses_ram_fraction() -> None: + with ( + patch("autointent._advisor._hardware._detect_cuda", return_value=None), + patch("autointent._advisor._hardware._detect_ram_gb", return_value=32.0), + patch( + "autointent._advisor._hardware._detect_mps", + side_effect=lambda ram, ratio: (ram * ratio, "Apple Silicon (arm64)"), + ), + ): + hw = detect_hardware() + assert hw.accelerator == "mps" + assert hw.vram_gb == pytest.approx(32.0 * 0.7) + assert any("MPS unified memory" in n for n in hw.notes) + + +def test_vram_budget_override_applies() -> None: + with ( + patch( + "autointent._advisor._hardware._detect_cuda", + return_value=(24.0, "NVIDIA RTX 4090"), + ), + ): + hw = detect_hardware(vram_budget_gb=8.0) + assert hw.vram_gb == pytest.approx(8.0) + assert any("manual VRAM budget" in n for n in hw.notes) + + +def test_broken_cuda_returns_none_does_not_crash() -> None: + # _detect_cuda swallows torch quirks already; verify the wrapper holds. + with ( + patch("autointent._advisor._hardware._detect_cuda", return_value=None), + patch("autointent._advisor._hardware._detect_mps", return_value=None), + ): + hw = detect_hardware() + assert hw.accelerator == "cpu" diff --git a/tests/advisor/test_hub_heuristics.py b/tests/advisor/test_hub_heuristics.py new file mode 100644 index 000000000..54a03431d --- /dev/null +++ b/tests/advisor/test_hub_heuristics.py @@ -0,0 +1,81 @@ +"""Tests for the offline name-pattern heuristics in `_hub`. + +The advisor must produce a sensible estimate even when HF Hub is +unreachable, so these tests pin the public `hub_reachable` to False and +exercise the heuristic path directly. +""" + +from __future__ import annotations + +import pytest + +from autointent._advisor import _hub + + +@pytest.fixture(autouse=True) +def _offline(monkeypatch: pytest.MonkeyPatch) -> None: + _hub.hub_reachable.cache_clear() + _hub.resolve_model.cache_clear() + monkeypatch.setattr(_hub, "hub_reachable", lambda *_a, **_kw: False) + monkeypatch.setattr(_hub, "_is_warm_cached", lambda _name: False) + + +@pytest.mark.parametrize( + ("name", "expected_min_m", "expected_max_m"), + [ + ("microsoft/deberta-v3-large", 200, 500), + ("microsoft/deberta-v3-small", 30, 200), + ("sentence-transformers/all-MiniLM-L6-v2", 20, 80), + ("intfloat/multilingual-e5-large-instruct", 300, 700), + ("intfloat/e5-small", 20, 80), + ("distilbert-base-uncased", 40, 150), + ("bert-base-uncased", 70, 200), + ], +) +def test_name_heuristic_picks_reasonable_bucket( + name: str, expected_min_m: int, expected_max_m: int +) -> None: + meta = _hub.resolve_model(name) + assert meta.confidence == "heuristic" + assert expected_min_m <= meta.params_millions <= expected_max_m, ( + f"{name} got {meta.params_millions}M; expected [{expected_min_m}, {expected_max_m}]" + ) + + +def test_unknown_name_falls_back_to_bert_base() -> None: + meta = _hub.resolve_model("totally-made-up/no-such-model") + assert meta.confidence == "heuristic" + assert meta.params_millions == pytest.approx(110.0) + + +def test_weights_gb_matches_params_times_bytes() -> None: + meta = _hub.resolve_model("microsoft/deberta-v3-large") + expected_gb = meta.params_millions * 1_000_000 * meta.weight_bytes_per_param / (1024**3) + assert meta.weights_gb == pytest.approx(expected_gb) + + +def test_local_path_returns_zero_disk() -> None: + meta = _hub.resolve_model("/tmp/local/path/to/model") + assert meta.total_file_bytes == 0 + assert meta.cached_locally is True + + +def test_disk_gb_falls_back_to_param_size_when_siblings_unknown() -> None: + meta = _hub.resolve_model("intfloat/multilingual-e5-large-instruct") + assert meta.disk_gb > 0 + assert meta.disk_gb == pytest.approx(meta.weights_gb, rel=0.01) + + +def test_resolve_is_memoized() -> None: + a = _hub.resolve_model("microsoft/deberta-v3-large") + b = _hub.resolve_model("microsoft/deberta-v3-large") + assert a is b + + +def test_metadata_fallback_uses_heuristic_when_hub_unreachable() -> None: + """End-to-end: resolve_model must return a usable ModelMeta even when + the live Hub is unreachable (autouse fixture forces offline).""" + meta = _hub.resolve_model("microsoft/deberta-v3-large") + assert meta.confidence == "heuristic" + assert meta.params_millions > 0 + assert meta.disk_gb > 0 diff --git a/tests/advisor/test_render.py b/tests/advisor/test_render.py new file mode 100644 index 000000000..e82d7573b --- /dev/null +++ b/tests/advisor/test_render.py @@ -0,0 +1,151 @@ +"""Output rendering: text formatting and JSON serialization.""" + +from __future__ import annotations + +import json + +import pytest + +from autointent._advisor._render import render_json, render_recommendation, render_text +from autointent._advisor._report import ( + DatasetStats, + PreflightReport, + ResourceEstimate, + Severity, +) + + +def _populated_report() -> PreflightReport: + r = PreflightReport( + preset_name="example", + hardware={ + "accelerator": "cuda", + "device_name": "RTX 3060", + "vram_gb": 8.0, + "ram_gb": 32.0, + "free_disk_gb": 100.0, + "device_class": "low-gpu", + }, + dataset={"n_samples": 500, "n_classes": 10, "avg_tokens": 30, "source": "placeholder"}, + resource=ResourceEstimate( + disk_download_gb=2.5, + disk_cached_gb=0.5, + ram_gb=1.0, + vram_gb=4.0, + time_hours=1.2, + drivers=[ + { + "node_type": "scoring", + "module": "bert", + "model": "x/y", + "mode": "full-finetune", + "vram_gb": 4.0, + "ram_gb": 1.0, + "time_hours": 1.2, + "confidence": "hub", + } + ], + ), + notes=["MPS unified memory note"], + ) + r.add("resource", Severity.YELLOW, "VRAM ~6 GB vs available 8 GB") + r.add("data", Severity.RED, "rare classes blocked") + return r + + +class TestRenderText: + def test_contains_phase_blocks(self) -> None: + out = render_text(_populated_report()) + assert "Resource:" in out + assert "Data:" in out + # Config phase has no findings → block omitted + assert "Config:" not in out + + def test_includes_drivers_block(self) -> None: + out = render_text(_populated_report()) + assert "Drivers of cost:" in out + assert "x/y" in out + + def test_verdict_reflects_worst_severity(self) -> None: + out = render_text(_populated_report()) + assert "Verdict: INFEASIBLE" in out + assert "worst severity: red" in out + + def test_disclaimer_always_present(self) -> None: + out = render_text(_populated_report()) + assert "heuristic upper bounds" in out + + def test_low_confidence_tag_when_offline(self) -> None: + r = _populated_report() + r.low_confidence = True + out = render_text(r) + assert "low-confidence" in out + + def test_preset_name_in_title(self) -> None: + out = render_text(_populated_report()) + assert "Compute feasibility check — example" in out + + def test_empty_report_still_renders(self) -> None: + out = render_text(PreflightReport()) + assert "Compute feasibility check" in out + assert "Verdict: feasible" in out + + +class TestRenderJson: + def test_is_valid_json(self) -> None: + json.loads(render_json(_populated_report())) + + def test_findings_have_string_severity(self) -> None: + d = json.loads(render_json(_populated_report())) + for f in d["findings"]: + assert f["severity"] in {"green", "yellow", "red"} + + def test_worst_severity_and_feasibility_serialized(self) -> None: + d = json.loads(render_json(_populated_report())) + assert d["worst_severity"] == "red" + assert d["is_feasible"] is False + + def test_empty_report_serializes(self) -> None: + d = json.loads(render_json(PreflightReport())) + assert d["worst_severity"] == "green" + assert d["is_feasible"] is True + + +class TestRenderRecommendation: + def _two_reports(self) -> list[tuple[str, PreflightReport]]: + a = PreflightReport(preset_name="a", resource=ResourceEstimate(vram_gb=2.0, time_hours=0.5)) + a.add("resource", Severity.GREEN, "ok") + b = PreflightReport(preset_name="b", resource=ResourceEstimate(vram_gb=8.0, time_hours=4.0)) + b.add("resource", Severity.RED, "too big") + return [("a", a), ("b", b)] + + def test_lists_chosen_preset_when_present(self) -> None: + out = render_recommendation(self._two_reports(), chosen="a") + assert "→ a" in out + + def test_handles_no_chosen(self) -> None: + out = render_recommendation(self._two_reports(), chosen=None) + assert "none of the bundled presets" in out + + def test_includes_all_presets_in_table(self) -> None: + out = render_recommendation(self._two_reports(), chosen="a") + assert "a " in out # preset name + assert "b " in out + + def test_shows_status_per_preset(self) -> None: + out = render_recommendation(self._two_reports(), chosen="a") + assert "feasible" in out + assert "infeasible" in out + + +def test_dataset_stats_in_text_block() -> None: + stats = DatasetStats.placeholder(n_samples=777, n_classes=4) + r = PreflightReport(dataset={ + "n_samples": stats.n_samples, + "n_classes": stats.n_classes, + "avg_tokens": stats.avg_tokens, + "source": stats.source, + }) + out = render_text(r) + assert "777" in out + assert "n_classes=4" in out diff --git a/tests/advisor/test_report.py b/tests/advisor/test_report.py new file mode 100644 index 000000000..52f2e675e --- /dev/null +++ b/tests/advisor/test_report.py @@ -0,0 +1,85 @@ +"""Unit tests for the report dataclasses.""" + +from __future__ import annotations + +import pytest + +from autointent._advisor._report import ( + DatasetStats, + Finding, + PreflightReport, + ResourceEstimate, + Severity, +) + + +class TestSeverityOrdering: + def test_worst_severity_on_empty_report_is_green(self) -> None: + assert PreflightReport().worst_severity == Severity.GREEN + + def test_red_beats_yellow_beats_green(self) -> None: + r = PreflightReport() + r.add("resource", Severity.GREEN, "ok") + r.add("data", Severity.YELLOW, "warn") + assert r.worst_severity == Severity.YELLOW + r.add("config", Severity.RED, "fail") + assert r.worst_severity == Severity.RED + + def test_is_feasible_flips_on_any_red(self) -> None: + r = PreflightReport() + r.add("resource", Severity.YELLOW, "warn") + assert r.is_feasible is True + r.add("data", Severity.RED, "fail") + assert r.is_feasible is False + + +class TestDatasetStatsPlaceholder: + def test_defaults_populate_p95_above_avg(self) -> None: + stats = DatasetStats.placeholder() + assert stats.n_samples == 1_000 + assert stats.p95_tokens is not None + assert stats.p95_tokens > stats.avg_tokens + assert stats.source == "placeholder" + + def test_overrides_propagate(self) -> None: + stats = DatasetStats.placeholder(n_samples=42, n_classes=3, avg_tokens=80, multilabel=True) + assert stats.n_samples == 42 + assert stats.n_classes == 3 + assert stats.avg_tokens == 80 + assert stats.multilabel is True + + +class TestResourceEstimate: + def test_total_disk_sums_download_and_dump(self) -> None: + e = ResourceEstimate(disk_download_gb=2.5, disk_dump_gb=4.0) + assert e.total_disk_gb == pytest.approx(6.5) + + def test_total_disk_ignores_cached(self) -> None: + e = ResourceEstimate(disk_download_gb=1.0, disk_cached_gb=100.0, disk_dump_gb=0.5) + assert e.total_disk_gb == pytest.approx(1.5) + + +class TestToDictSerialization: + def test_findings_round_trip_severity_as_string(self) -> None: + r = PreflightReport() + r.add("resource", Severity.RED, "boom") + d = r.to_dict() + assert d["worst_severity"] == "red" + assert d["is_feasible"] is False + assert d["findings"] == [ + {"phase": "resource", "severity": "red", "message": "boom", "metric": None}, + ] + + def test_hardware_and_dataset_pass_through(self) -> None: + r = PreflightReport( + hardware={"accelerator": "cuda", "vram_gb": 8.0}, + dataset={"n_samples": 100, "n_classes": 5}, + ) + d = r.to_dict() + assert d["hardware"]["accelerator"] == "cuda" + assert d["dataset"]["n_samples"] == 100 + + def test_finding_is_frozen(self) -> None: + f = Finding(phase="resource", severity=Severity.GREEN, message="ok") + with pytest.raises(Exception): # noqa: PT011 - dataclass.FrozenInstanceError varies + f.message = "changed" # type: ignore[misc] From c8675b9a69cf5c1492380a1bdac2c8a3885281b2 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Mon, 15 Jun 2026 23:58:59 +0300 Subject: [PATCH 04/43] fix --- src/autointent/_advisor/__init__.py | 2 +- src/autointent/_advisor/_cli.py | 40 +-- src/autointent/_advisor/_estimates.py | 217 ++++++++++++++-- src/autointent/_advisor/_hardware.py | 59 ++--- src/autointent/_advisor/_hub.py | 20 +- src/autointent/_advisor/_render.py | 6 +- src/autointent/_advisor/_report.py | 6 +- tests/advisor/test_estimates_and_cli.py | 43 +++- tests/advisor/test_estimates_internals.py | 242 +++++++++++++++--- tests/advisor/test_hub_heuristics.py | 4 +- tests/advisor/test_render.py | 16 +- .../advanced/02_embedder_configuration.py | 4 +- 12 files changed, 495 insertions(+), 164 deletions(-) diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index 5f29b028e..3ff898816 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -7,9 +7,9 @@ from __future__ import annotations +from ._estimates import run_preflight from ._hardware import HardwareProfile, detect_hardware from ._report import DatasetStats, Finding, PreflightReport, ResourceEstimate, Severity -from ._estimates import run_preflight __all__ = [ "DatasetStats", diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py index 4e7eae000..b3f43aab5 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/_advisor/_cli.py @@ -20,10 +20,13 @@ import yaml +from autointent import Dataset +from autointent.utils import load_preset + from ._estimates import run_preflight from ._hardware import detect_hardware from ._render import render_json, render_recommendation, render_text -from ._report import DatasetStats, PreflightReport +from ._report import DatasetStats, PreflightReport, Severity logger = logging.getLogger("autointent.advisor") @@ -62,8 +65,6 @@ def _load_config(target: str) -> tuple[dict[str, Any], str]: with path.open(encoding="utf-8") as f: return yaml.safe_load(f), path.stem # treat as a bundled preset name - from autointent.utils import load_preset - return load_preset(target), target # type: ignore[arg-type] @@ -80,15 +81,9 @@ def _stats_from_args(args: argparse.Namespace) -> DatasetStats: def _stats_from_dataset(path: str, *, multilabel: bool) -> DatasetStats: """Best-effort: load a dataset from disk via the existing Dataset constructor.""" - try: - from autointent import Dataset - except ImportError: - logger.warning("autointent.Dataset unavailable; falling back to placeholders.") - return DatasetStats.placeholder(multilabel=multilabel) - try: ds = Dataset.from_json(path) if path.endswith(".json") else Dataset.from_hub(path) - except Exception as e: # noqa: BLE001 + except (OSError, ValueError) as e: logger.warning("Failed to load dataset %s: %s", path, e) return DatasetStats.placeholder(multilabel=multilabel) @@ -147,27 +142,24 @@ def cmd_recommend(args: argparse.Namespace) -> int: stats = _stats_from_args(args) results: list[tuple[str, PreflightReport]] = [] - from autointent.utils import load_preset for preset in BUNDLED_PRESETS: try: cfg = load_preset(preset) # type: ignore[arg-type] - except Exception as e: # noqa: BLE001 + except (OSError, ValueError, KeyError) as e: logger.debug("Skipping preset %s: %s", preset, e) continue report = run_preflight(cfg, stats, hardware, preset_name=preset) if args.budget_time_h is not None and report.resource.time_hours > args.budget_time_h: report.add( "resource", - report.worst_severity if report.worst_severity.value == "red" else report.worst_severity, # noqa: PLW0125 - explicit + Severity.RED, f"Estimated time {report.resource.time_hours:.1f} h exceeds budget {args.budget_time_h} h.", ) results.append((preset, report)) feasible = [(name, r) for name, r in results if r.is_feasible] - feasible.sort( - key=lambda pair: (-_QUALITY_TIER.get(pair[0], 0), pair[1].resource.time_hours, pair[0]) - ) + feasible.sort(key=lambda pair: (-_QUALITY_TIER.get(pair[0], 0), pair[1].resource.time_hours, pair[0])) chosen = feasible[0][0] if feasible else None if args.json: @@ -175,9 +167,7 @@ def cmd_recommend(args: argparse.Namespace) -> int: out = { "chosen": chosen, - "results": [ - {"preset": name, "report": r.to_dict()} for name, r in results - ], + "results": [{"preset": name, "report": r.to_dict()} for name, r in results], } sys.stdout.write(json.dumps(out, indent=2, default=str)) sys.stdout.write("\n") @@ -206,9 +196,7 @@ def build_parser() -> argparse.ArgumentParser: ) p_inspect.add_argument("target", help="Preset name (e.g. transformers-light) or path to a YAML config.") p_inspect.add_argument("--json", action="store_true", help="Emit a structured JSON report.") - p_inspect.add_argument( - "--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget." - ) + p_inspect.add_argument("--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget.") _add_common_dataset_args(p_inspect) p_inspect.set_defaults(func=cmd_inspect) @@ -217,12 +205,8 @@ def build_parser() -> argparse.ArgumentParser: help="Detect hardware and recommend the best-fitting bundled preset.", ) p_rec.add_argument("--json", action="store_true", help="Emit a structured JSON report.") - p_rec.add_argument( - "--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget." - ) - p_rec.add_argument( - "--budget-time-h", type=float, default=None, help="Optional wall-time ceiling in hours." - ) + p_rec.add_argument("--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget.") + p_rec.add_argument("--budget-time-h", type=float, default=None, help="Optional wall-time ceiling in hours.") _add_common_dataset_args(p_rec) p_rec.set_defaults(func=cmd_recommend) diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py index f60f940a6..e8f619303 100644 --- a/src/autointent/_advisor/_estimates.py +++ b/src/autointent/_advisor/_estimates.py @@ -9,7 +9,8 @@ from __future__ import annotations import logging -from typing import Any, Iterable +from collections.abc import Iterable +from typing import Any from ._hardware import HardwareProfile from ._hub import ModelMeta, hub_reachable, resolve_model @@ -32,6 +33,16 @@ TRANSFORMER_SCORER_MODULES = {"bert", "lora", "ptuning", "dnnc"} +# Coefficients for the linear / catboost time formulas (proposal §"Algorithm"). +_LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-8 +_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 +_CATBOOST_GPU_SPEEDUP = 10.0 +# LogisticRegressionCV defaults: Cs=10, cv=3 → 31 inner fits + 1 final refit. +_LOGREG_CV_MULTIPLIER = 31 +_CATBOOST_DEFAULT_BINS = 254 +# Bytes per histogram bucket / tree node — order-of-magnitude constants. +_CATBOOST_BYTES_PER_TREE_NODE = 32 + def _extract_model_names(module_entry: dict[str, Any]) -> list[str]: """Pull model name(s) from a search-space module entry.""" @@ -74,11 +85,22 @@ def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dic yield node_type, entry +def _walk_modules_indexed( + search_space: list[dict[str, Any]], +) -> Iterable[tuple[int, str, dict[str, Any]]]: + """Yield (node_index, node_type, module_entry) — index lets us bound per-node max cost.""" + for node_idx, node in enumerate(search_space or []): + node_type = node.get("node_type", "?") + for entry in node.get("search_space", []) or []: + yield node_idx, node_type, entry + + def _vram_for_transformer(meta: ModelMeta, mode: str, mixed_precision: bool) -> float: """VRAM in GB for one trial of a transformer-based module. - Conservative AMP accounting (the proposal flags the prior naive halving - as too generous; keep optimizer state at fp32 even in AMP). + Full fine-tune fp32: weights + grads + Adam (m, v) = 4W. + Full fine-tune AMP: fp16 weights + fp16 grads + fp32 master copy + fp32 Adam = 3W. + (Activations are not modeled separately.) """ weights_gb = meta.weights_gb if mode == "inference": @@ -87,11 +109,9 @@ def _vram_for_transformer(meta: ModelMeta, mode: str, mixed_precision: bool) -> return weights_gb * 1.3 + 0.5 if mode == "reranker": return weights_gb * 1.5 - # full fine-tune (bert, ptuning, gcn-with-backbone) if mixed_precision: - # fp16 weights+grads + fp32 master+adam moments - return (weights_gb * 0.5) * 2 + weights_gb * 1 + weights_gb * 2 - return weights_gb * (1 + 1 + 2) + return weights_gb * 3.0 + return weights_gb * 4.0 def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: @@ -99,6 +119,82 @@ def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / (1024**3) +def _embedder_dim(meta: ModelMeta | None) -> int: + """Coarse hidden-size guess from parameter count. + + Concrete points: MiniLM (33M) ~384, BERT-base (110M) ~768, BERT-large (350M) ~1024. + """ + if meta is None: + return 768 + params = meta.params_millions + if params >= 300: + return 1024 + if params >= 100: + return 768 + if params >= 50: + return 512 + return 384 + + +def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: + if not seen_models: + return None + return max(seen_models.values(), key=lambda m: m.params_millions) + + +def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: + """Float64 design matrix dominates; coefficients and L-BFGS history are small.""" + data_bytes = 8.0 * stats.n_samples * embedder_dim + coef_bytes = 8.0 * max(1, stats.n_classes) * embedder_dim + lbfgs_bytes = 10.0 * 8.0 * embedder_dim + return (data_bytes + coef_bytes + lbfgs_bytes) / (1024**3) + + +def _time_for_linear( + *, + n_trials: int, + n_samples: int, + embedder_dim: int, + max_iter: int, + cv_multiplier: int, + class_multiplier: int, +) -> float: + seconds = ( + n_trials + * _LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER + * n_samples + * embedder_dim + * max_iter + * cv_multiplier + * class_multiplier + ) + return seconds / 3600.0 + + +def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, depth: int) -> float: + data_bytes = 4.0 * stats.n_samples * n_features + histograms_bytes = 4.0 * n_features * _CATBOOST_DEFAULT_BINS + trees_bytes = iterations * (2**depth) * _CATBOOST_BYTES_PER_TREE_NODE + return (data_bytes + histograms_bytes + trees_bytes) / (1024**3) + + +def _time_for_catboost( + *, + n_trials: int, + n_samples: int, + n_features: int, + iterations: int, + depth: int, + class_multiplier: int, + on_gpu: bool, +) -> float: + coeff = _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER + if on_gpu: + coeff /= _CATBOOST_GPU_SPEEDUP + seconds = n_trials * iterations * coeff * n_samples * n_features * depth * class_multiplier + return seconds / 3600.0 + + def _time_for_transformer( *, meta: ModelMeta, @@ -148,10 +244,24 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity if global_embedder: seen_models[global_embedder] = resolve_model(global_embedder) - for node_type, entry in _walk_modules(config.get("search_space") or []): + # First pass: walk transformer-bearing modules (collects seen_models for embedder_dim lookup). + transformer_entries: list[tuple[int, str, dict[str, Any]]] = [] + classic_entries: list[tuple[int, str, dict[str, Any]]] = [] + for node_idx, node_type, entry in _walk_modules_indexed(config.get("search_space") or []): + module = entry.get("module_name", "?") + if module in {"linear", "catboost"}: + classic_entries.append((node_idx, node_type, entry)) + else: + transformer_entries.append((node_idx, node_type, entry)) + + # Track the heaviest module per node so dump_modules accounting is bounded by + # "one selected variant per node × n_trials", not "sum of every candidate". + node_max_weights: dict[int, float] = {} + + for node_idx, node_type, entry in transformer_entries: module = entry.get("module_name", "?") model_names = _extract_model_names(entry) - if not model_names and global_embedder and module in {"linear", "catboost", "knn", "mlknn"}: + if not model_names and global_embedder and module in {"knn", "mlknn"}: model_names = [global_embedder] for name in model_names: @@ -191,6 +301,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity estimate.vram_gb = max(estimate.vram_gb, vram) estimate.ram_gb = max(estimate.ram_gb, ram) estimate.time_hours += time_h + node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), meta.weights_gb) estimate.drivers.append( { "node_type": node_type, @@ -204,6 +315,76 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity } ) + # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. + embedder_meta = _largest_embedder(seen_models) + embedder_dim = _embedder_dim(embedder_meta) + class_multiplier_classic = max(1, stats.n_classes) if stats.multilabel else 1 + for _node_idx, node_type, entry in classic_entries: + module = entry.get("module_name", "?") + if module == "linear": + max_iter = _max_int(entry.get("max_iter"), 100) + cv_multiplier = 1 if stats.multilabel else _LOGREG_CV_MULTIPLIER + ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) + time_h = _time_for_linear( + n_trials=n_trials, + n_samples=stats.n_samples, + embedder_dim=embedder_dim, + max_iter=max_iter, + cv_multiplier=cv_multiplier, + class_multiplier=class_multiplier_classic, + ) + if refit_after: + time_h *= 1 + 1.0 / max(1, n_trials) + vram = 0.0 + mode = "linear-cv" if cv_multiplier > 1 else "linear" + confidence = embedder_meta.confidence if embedder_meta else "heuristic" + elif module == "catboost": + iterations = _max_int(entry.get("iterations"), 1000) + depth = _max_int(entry.get("depth"), 6) + on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" + # CatBoost's multiclass MultiClass loss already grows per-class trees. + cb_class_mult = max(1, stats.n_classes) + ram = _ram_for_catboost( + stats=stats, + n_features=embedder_dim, + iterations=iterations, + depth=depth, + ) + time_h = _time_for_catboost( + n_trials=n_trials, + n_samples=stats.n_samples, + n_features=embedder_dim, + iterations=iterations, + depth=depth, + class_multiplier=cb_class_mult, + on_gpu=on_gpu, + ) + if refit_after: + time_h *= 1 + 1.0 / max(1, n_trials) + vram = ram if on_gpu else 0.0 + if on_gpu: + ram = 0.0 + mode = "catboost-gpu" if on_gpu else "catboost" + confidence = embedder_meta.confidence if embedder_meta else "heuristic" + else: + continue + + estimate.vram_gb = max(estimate.vram_gb, vram) + estimate.ram_gb = max(estimate.ram_gb, ram) + estimate.time_hours += time_h + estimate.drivers.append( + { + "node_type": node_type, + "module": module, + "model": embedder_meta.name if embedder_meta else "(no embedder)", + "mode": mode, + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "confidence": confidence, + } + ) + for meta in seen_models.values(): if meta.cached_locally: estimate.disk_cached_gb += meta.disk_gb @@ -211,8 +392,10 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity estimate.disk_download_gb += meta.disk_gb if dump_modules: - weights_total = sum(m.weights_gb for m in seen_models.values()) - estimate.disk_dump_gb = weights_total * n_trials + # Each trial selects one variant per node, so per-trial dumped weights + # are bounded by the heaviest module in each node, summed across nodes. + per_trial_dump_gb = sum(node_max_weights.values()) + estimate.disk_dump_gb = per_trial_dump_gb * n_trials if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: effective_vram = estimate.vram_gb * n_jobs @@ -309,23 +492,17 @@ def _data_phase( ) # rare class × linear-CV - has_linear = any( - e.get("module_name") == "linear" for _, e in _walk_modules(config.get("search_space") or []) - ) + has_linear = any(e.get("module_name") == "linear" for _, e in _walk_modules(config.get("search_space") or [])) if has_linear and stats.rare_classes: report.add( "data", Severity.RED, - ( - "LogisticRegressionCV (cv=3) will fail: classes " - f"{stats.rare_classes[:5]} have <3 samples." - ), + (f"LogisticRegressionCV (cv=3) will fail: classes {stats.rare_classes[:5]} have <3 samples."), ) # partial descriptions × description scorer has_description = any( - e.get("module_name") == "description" - for _, e in _walk_modules(config.get("search_space") or []) + e.get("module_name") == "description" for _, e in _walk_modules(config.get("search_space") or []) ) if has_description and stats.has_descriptions is False: report.add( diff --git a/src/autointent/_advisor/_hardware.py b/src/autointent/_advisor/_hardware.py index 2bda6120f..9c0cae049 100644 --- a/src/autointent/_advisor/_hardware.py +++ b/src/autointent/_advisor/_hardware.py @@ -14,6 +14,9 @@ from dataclasses import dataclass, field from typing import Literal +import psutil +import torch + logger = logging.getLogger(__name__) Accelerator = Literal["cuda", "mps", "cpu"] @@ -46,13 +49,7 @@ def device_class(self) -> str: def _detect_ram_gb() -> float: - try: - import psutil - - return psutil.virtual_memory().total / (1024**3) - except ImportError: - logger.debug("psutil unavailable; RAM unknown") - return 0.0 + return psutil.virtual_memory().total / (1024**3) def _detect_free_disk_gb(path: str | None = None) -> float: @@ -67,40 +64,24 @@ def _detect_free_disk_gb(path: str | None = None) -> float: def _detect_cuda() -> tuple[float, str] | None: - try: - import torch - - if not torch.cuda.is_available(): - return None - idx = 0 - try: - free, total = torch.cuda.mem_get_info(idx) - vram_gb = total / (1024**3) - except (RuntimeError, AttributeError) as e: - logger.debug("torch.cuda.mem_get_info failed: %s", e) - return None - name = torch.cuda.get_device_name(idx) - return vram_gb, name - except ImportError: + if not torch.cuda.is_available(): return None - except Exception as e: # noqa: BLE001 - protect the advisor from torch quirks - logger.debug("CUDA detection raised: %s", e) + idx = 0 + try: + _free, total = torch.cuda.mem_get_info(idx) + vram_gb = total / (1024**3) + except (RuntimeError, AttributeError) as e: + logger.debug("torch.cuda.mem_get_info failed: %s", e) return None + name = torch.cuda.get_device_name(idx) + return vram_gb, name def _detect_mps(ram_gb: float, budget_ratio: float = MPS_DEFAULT_BUDGET_RATIO) -> tuple[float, str] | None: - try: - import torch - - if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): - return None - # apple silicon: unified memory; budget is fraction of total RAM - return ram_gb * budget_ratio, f"Apple Silicon ({platform.machine()})" - except ImportError: - return None - except Exception as e: # noqa: BLE001 - logger.debug("MPS detection raised: %s", e) + if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()): return None + # apple silicon: unified memory; budget is fraction of total RAM + return ram_gb * budget_ratio, f"Apple Silicon ({platform.machine()})" def detect_hardware( @@ -133,9 +114,7 @@ def detect_hardware( if mps is not None: vram_gb, device_name = mps accel = "mps" - notes.append( - f"MPS unified memory: VRAM budget = {mps_budget_ratio:.0%} of RAM." - ) + notes.append(f"MPS unified memory: VRAM budget = {mps_budget_ratio:.0%} of RAM.") else: vram_gb = 0.0 device_name = platform.processor() or "cpu" @@ -143,9 +122,7 @@ def detect_hardware( if vram_budget_gb is not None: if vram_gb and vram_budget_gb > vram_gb: - notes.append( - f"Manual --budget-vram-gb={vram_budget_gb} exceeds detected {vram_gb:.1f} GB; using override." - ) + notes.append(f"Manual --budget-vram-gb={vram_budget_gb} exceeds detected {vram_gb:.1f} GB; using override.") notes.append(f"Using manual VRAM budget: {vram_budget_gb} GB.") vram_gb = vram_budget_gb diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index 80ccb7133..613ab6b40 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -14,6 +14,8 @@ from functools import lru_cache from typing import Any +from huggingface_hub import HfApi, scan_cache_dir, try_to_load_from_cache + logger = logging.getLogger(__name__) # Coarse heuristic estimates keyed on name fragments. Used only when HF Hub @@ -55,17 +57,11 @@ def weights_gb(self) -> float: def hub_reachable(timeout_s: float = 2.0) -> bool: """Single up-front probe. Memoized per process.""" try: - from huggingface_hub import HfApi - HfApi().list_models(limit=1) - except ImportError: - logger.debug("huggingface_hub not installed; assuming offline") - return False except Exception as e: # noqa: BLE001 logger.debug("HF Hub probe failed: %s", e) return False - else: - return True + return True def _heuristic_params_millions(model_name: str) -> float: @@ -77,11 +73,6 @@ def _heuristic_params_millions(model_name: str) -> float: def _is_warm_cached(model_name: str) -> bool: """True when the weight shard is present in the local HF cache.""" - try: - from huggingface_hub import scan_cache_dir, try_to_load_from_cache - except ImportError: - return False - weight_files = ["model.safetensors", "pytorch_model.bin", "model.safetensors.index.json"] for fname in weight_files: path = try_to_load_from_cache(model_name, fname) @@ -98,11 +89,6 @@ def _is_warm_cached(model_name: str) -> bool: def _hub_metadata(model_name: str) -> ModelMeta | None: - try: - from huggingface_hub import HfApi - except ImportError: - return None - try: info = HfApi().model_info(model_name, files_metadata=True) except Exception as e: # noqa: BLE001 diff --git a/src/autointent/_advisor/_render.py b/src/autointent/_advisor/_render.py index 52168aa75..fe0f32dd7 100644 --- a/src/autointent/_advisor/_render.py +++ b/src/autointent/_advisor/_render.py @@ -18,7 +18,7 @@ _PHASE_LABEL = {"resource": "Resource", "data": "Data", "config": "Config"} -def render_text(report: "PreflightReport") -> str: +def render_text(report: PreflightReport) -> str: lines: list[str] = [] title = "Compute feasibility check" if report.preset_name: @@ -76,12 +76,12 @@ def render_text(report: "PreflightReport") -> str: return "\n".join(lines) -def render_json(report: "PreflightReport") -> str: +def render_json(report: PreflightReport) -> str: return json.dumps(report.to_dict(), indent=2, default=str) def render_recommendation( - results: list[tuple[str, "PreflightReport"]], + results: list[tuple[str, PreflightReport]], chosen: str | None, ) -> str: """Compact table for the ``recommend`` subcommand.""" diff --git a/src/autointent/_advisor/_report.py b/src/autointent/_advisor/_report.py index 0250482a5..6b930db95 100644 --- a/src/autointent/_advisor/_report.py +++ b/src/autointent/_advisor/_report.py @@ -67,7 +67,7 @@ def placeholder( n_classes: int = 10, avg_tokens: int = 32, multilabel: bool = False, - ) -> "DatasetStats": + ) -> DatasetStats: return cls( n_samples=n_samples, n_classes=n_classes, @@ -105,9 +105,7 @@ def is_feasible(self) -> bool: def to_dict(self) -> dict[str, Any]: d = asdict(self) - d["findings"] = [ - {**asdict(f), "severity": f.severity.value} for f in self.findings - ] + d["findings"] = [{**asdict(f), "severity": f.severity.value} for f in self.findings] d["worst_severity"] = self.worst_severity.value d["is_feasible"] = self.is_feasible return d diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 18c2615a6..00537226d 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -162,9 +162,7 @@ def test_partial_descriptions_with_description_scorer_flags_red() -> None: has_descriptions=False, ) report = run_preflight(cfg, stats, _profile(vram_gb=16.0)) - assert any( - f.phase == "data" and "description" in f.message.lower() for f in report.findings - ) + assert any(f.phase == "data" and "description" in f.message.lower() for f in report.findings) def test_long_dataset_triggers_truncation_warning() -> None: @@ -175,9 +173,7 @@ def test_long_dataset_triggers_truncation_warning() -> None: "search_space": [ { "module_name": "bert", - "classification_model_config": [ - {"model_name": "microsoft/deberta-v3-small"} - ], + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], "max_length": [128], } ], @@ -194,5 +190,40 @@ def test_long_dataset_triggers_truncation_warning() -> None: assert any("truncation" in f.message.lower() for f in report.findings) +def test_cli_recommend_budget_time_flags_red_for_overbudget_presets( + capsys: pytest.CaptureFixture[str], +) -> None: + """Tight time budget must flag every preset that exceeds it with RED severity. + + Previously the budget path used a tautological severity expression and the + breach never escalated the finding — covers the regression.""" + main( + [ + "recommend", + "--n-samples", + "1000", + "--n-classes", + "10", + "--avg-tokens", + "20", + "--budget-vram-gb", + "48", + "--budget-time-h", + "0.0001", + "--json", + ] + ) + payload = json.loads(capsys.readouterr().out) + flagged = [ + r + for r in payload["results"] + if any(f["severity"] == "red" and "exceeds budget" in f["message"] for f in r["report"]["findings"]) + ] + assert flagged, "budget-time-h breach should produce RED severity findings" + # Any preset above the budget must be marked infeasible. + for r in flagged: + assert r["report"]["is_feasible"] is False + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 0317ff27b..713db27f7 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Any + import pytest from autointent._advisor import _estimates, _hub @@ -109,22 +111,19 @@ def meta(self) -> ModelMeta: confidence="hub", ) - def test_full_finetune_is_larger_than_lora_is_larger_than_inference( - self, meta: ModelMeta - ) -> None: + def test_full_finetune_is_larger_than_lora_is_larger_than_inference(self, meta: ModelMeta) -> None: inference = _vram_for_transformer(meta, "inference", mixed_precision=False) lora = _vram_for_transformer(meta, "lora", mixed_precision=False) full = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) assert inference < lora < full - def test_amp_does_not_naively_halve(self, meta: ModelMeta) -> None: - """The proposal calls out that AMP doesn't halve total VRAM — fp32 master - weights and Adam moments don't shrink. Weight-side accounting comes out - equal to fp32; the only savings (activations) aren't modeled by us.""" + def test_amp_partially_reduces_full_finetune_vram(self, meta: ModelMeta) -> None: + """AMP saves on fp16 weights+grads (W down from 2W); Adam state stays + fp32 (2W). Total 3W vs fp32's 4W — real but not a full halving.""" full_fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) full_amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True) - assert full_amp / full_fp32 == pytest.approx(1.0) - assert full_amp / full_fp32 > 0.5 # explicit check vs the naive-halving formula + assert full_amp < full_fp32 + assert full_amp / full_fp32 == pytest.approx(0.75) def test_reranker_uses_inference_class(self, meta: ModelMeta) -> None: inference = _vram_for_transformer(meta, "inference", mixed_precision=False) @@ -155,9 +154,7 @@ def test_dump_modules_adds_disk_during_training(self) -> None: "search_space": [ { "module_name": "bert", - "classification_model_config": [ - {"model_name": "microsoft/deberta-v3-small"} - ], + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], "num_train_epochs": [3], "batch_size": [16], } @@ -179,9 +176,7 @@ def test_refit_after_increases_time(self) -> None: "search_space": [ { "module_name": "bert", - "classification_model_config": [ - {"model_name": "microsoft/deberta-v3-small"} - ], + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], "num_train_epochs": [3], "batch_size": [16], } @@ -207,9 +202,7 @@ def test_catboost_gpu_without_cuda_flags_config(self) -> None: ], } report = run_preflight(cfg, DatasetStats.placeholder(), _profile(accelerator="cpu")) - assert any( - f.phase == "config" and "CatBoost" in f.message for f in report.findings - ) + assert any(f.phase == "config" and "CatBoost" in f.message for f in report.findings) def test_catboost_gpu_with_cuda_is_silent(self) -> None: cfg = { @@ -223,9 +216,7 @@ def test_catboost_gpu_with_cuda_is_silent(self) -> None: ], } report = run_preflight(cfg, DatasetStats.placeholder(), _profile(accelerator="cuda")) - assert not any( - f.phase == "config" and "CatBoost" in f.message for f in report.findings - ) + assert not any(f.phase == "config" and "CatBoost" in f.message for f in report.findings) def test_offline_flips_low_confidence(self) -> None: cfg = { @@ -277,9 +268,7 @@ def test_truncation_red_when_p95_dominates_max_length(self) -> None: { "module_name": "bert", "max_length": [128], - "classification_model_config": [ - {"model_name": "some/model"} - ], + "classification_model_config": [{"model_name": "some/model"}], } ], } @@ -299,9 +288,7 @@ def test_truncation_yellow_when_p95_only_slightly_exceeds(self) -> None: { "module_name": "bert", "max_length": [128], - "classification_model_config": [ - {"model_name": "some/model"} - ], + "classification_model_config": [{"model_name": "some/model"}], } ], } @@ -312,8 +299,203 @@ def test_truncation_yellow_when_p95_only_slightly_exceeds(self) -> None: yellows = [ f for f in report.findings - if f.phase == "data" - and f.severity == Severity.YELLOW - and "truncation" in f.message.lower() + if f.phase == "data" and f.severity == Severity.YELLOW and "truncation" in f.message.lower() ] assert yellows + + +class TestLinearCatboostFormulas: + """Cost surfaces for the classic (sklearn / catboost) scorers.""" + + def _embedder_node(self) -> dict[str, Any]: + return { + "node_type": "embedder", + "search_space": [ + { + "module_name": "sentence_transformer", + "embedder_config": [{"model_name": "sentence-transformers/all-MiniLM-L6-v2"}], + } + ], + } + + def test_linear_contributes_ram_and_time(self) -> None: + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [{"module_name": "linear", "max_iter": [200]}], + }, + ], + "hpo_config": {"n_trials": 5}, + } + stats = DatasetStats.placeholder(n_samples=100_000, n_classes=10, avg_tokens=24) + report = run_preflight(cfg, stats, _profile()) + linear_drivers = [d for d in report.resource.drivers if d["module"] == "linear"] + assert len(linear_drivers) == 1 + assert report.resource.ram_gb > 0 + assert report.resource.time_hours > 0 + assert linear_drivers[0]["vram_gb"] == 0 # sklearn is CPU-only + + def test_logreg_cv_multiplier_dominates_multiclass_time(self) -> None: + """Multiclass linear uses LogisticRegressionCV (Cs*cv+1 ≈ 31 inner fits); + multilabel uses one LogReg per class (cv_multiplier=1). At equal n_classes, + multiclass must be much slower than the per-class multilabel path.""" + base = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [{"module_name": "linear", "max_iter": [1000]}], + }, + ], + "hpo_config": {"n_trials": 1}, + } + multiclass = run_preflight( + base, + DatasetStats.placeholder(n_samples=100_000, n_classes=10, multilabel=False), + _profile(), + ) + multilabel = run_preflight( + base, + DatasetStats.placeholder(n_samples=100_000, n_classes=10, multilabel=True), + _profile(), + ) + # multiclass: 31 inner fits x 1 model; multilabel: 1 fit x n_classes=10 models. + # 31 > 10 => multiclass is the slower path. + assert multiclass.resource.time_hours > multilabel.resource.time_hours + + def test_catboost_contributes_ram_and_time_on_cpu(self) -> None: + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "catboost", + "iterations": [1000], + "depth": [6], + } + ], + }, + ], + "hpo_config": {"n_trials": 3}, + } + stats = DatasetStats.placeholder(n_samples=100_000, n_classes=8, avg_tokens=24) + report = run_preflight(cfg, stats, _profile(accelerator="cpu")) + cb = next(d for d in report.resource.drivers if d["module"] == "catboost") + assert report.resource.ram_gb > 0 + assert report.resource.time_hours > 0 + assert cb["vram_gb"] == 0 + assert cb["mode"] == "catboost" + + def test_catboost_gpu_moves_cost_to_vram(self) -> None: + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "catboost", + "iterations": [1000], + "depth": [6], + "task_type": "GPU", + } + ], + }, + ], + "hpo_config": {"n_trials": 2}, + } + stats = DatasetStats.placeholder(n_samples=100_000, n_classes=8, avg_tokens=24) + report = run_preflight(cfg, stats, _profile(accelerator="cuda")) + cb = next(d for d in report.resource.drivers if d["module"] == "catboost") + assert report.resource.vram_gb > 0 + assert cb["ram_gb"] == 0 + assert cb["mode"] == "catboost-gpu" + + def test_linear_scales_with_n_samples(self) -> None: + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [{"module_name": "linear"}], + }, + ], + } + small = run_preflight(cfg, DatasetStats.placeholder(n_samples=500), _profile()) + big = run_preflight(cfg, DatasetStats.placeholder(n_samples=500_000), _profile()) + assert big.resource.time_hours > small.resource.time_hours + assert big.resource.ram_gb > small.resource.ram_gb + + +class TestDumpModulesBounding: + """`dump_modules=True` writes one selected variant per node per trial — not + every candidate. The estimate must be bounded by sum-of-max-per-node x n_trials.""" + + def test_dump_disk_is_bounded_by_per_node_max_not_sum_of_all_variants(self) -> None: + # Two BERT candidates in the same node: only one is selected per trial. + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [ + {"model_name": "microsoft/deberta-v3-small"}, + {"model_name": "microsoft/deberta-v3-large"}, + ], + "num_train_epochs": [3], + "batch_size": [16], + } + ], + } + ], + "hpo_config": {"n_trials": 4}, + "dump_modules": True, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + # Per-node max ~ deberta-v3-large weights (~350M x 4 ~ 1.3 GB). Two-candidate + # sum would be roughly doubled. Verify we used the per-node-max bound. + small_meta = _hub.resolve_model("microsoft/deberta-v3-small") + large_meta = _hub.resolve_model("microsoft/deberta-v3-large") + expected = large_meta.weights_gb * 4 + naive_sum = (small_meta.weights_gb + large_meta.weights_gb) * 4 + assert report.resource.disk_dump_gb == pytest.approx(expected, rel=0.01) + assert report.resource.disk_dump_gb < naive_sum + + def test_dump_disk_sums_across_nodes(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "embedder", + "search_space": [ + { + "module_name": "sentence_transformer", + "embedder_config": [{"model_name": "sentence-transformers/all-MiniLM-L6-v2"}], + } + ], + }, + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "num_train_epochs": [3], + "batch_size": [16], + } + ], + }, + ], + "hpo_config": {"n_trials": 2}, + "dump_modules": True, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + embedder = _hub.resolve_model("sentence-transformers/all-MiniLM-L6-v2") + bert = _hub.resolve_model("microsoft/deberta-v3-small") + expected = (embedder.weights_gb + bert.weights_gb) * 2 + assert report.resource.disk_dump_gb == pytest.approx(expected, rel=0.01) diff --git a/tests/advisor/test_hub_heuristics.py b/tests/advisor/test_hub_heuristics.py index 54a03431d..b43b95522 100644 --- a/tests/advisor/test_hub_heuristics.py +++ b/tests/advisor/test_hub_heuristics.py @@ -32,9 +32,7 @@ def _offline(monkeypatch: pytest.MonkeyPatch) -> None: ("bert-base-uncased", 70, 200), ], ) -def test_name_heuristic_picks_reasonable_bucket( - name: str, expected_min_m: int, expected_max_m: int -) -> None: +def test_name_heuristic_picks_reasonable_bucket(name: str, expected_min_m: int, expected_max_m: int) -> None: meta = _hub.resolve_model(name) assert meta.confidence == "heuristic" assert expected_min_m <= meta.params_millions <= expected_max_m, ( diff --git a/tests/advisor/test_render.py b/tests/advisor/test_render.py index e82d7573b..55a2b0ce3 100644 --- a/tests/advisor/test_render.py +++ b/tests/advisor/test_render.py @@ -4,8 +4,6 @@ import json -import pytest - from autointent._advisor._render import render_json, render_recommendation, render_text from autointent._advisor._report import ( DatasetStats, @@ -140,12 +138,14 @@ def test_shows_status_per_preset(self) -> None: def test_dataset_stats_in_text_block() -> None: stats = DatasetStats.placeholder(n_samples=777, n_classes=4) - r = PreflightReport(dataset={ - "n_samples": stats.n_samples, - "n_classes": stats.n_classes, - "avg_tokens": stats.avg_tokens, - "source": stats.source, - }) + r = PreflightReport( + dataset={ + "n_samples": stats.n_samples, + "n_classes": stats.n_classes, + "avg_tokens": stats.avg_tokens, + "source": stats.source, + } + ) out = render_text(r) assert "777" in out assert "n_classes=4" in out diff --git a/user_guides/advanced/02_embedder_configuration.py b/user_guides/advanced/02_embedder_configuration.py index 32118fed2..43ce27278 100644 --- a/user_guides/advanced/02_embedder_configuration.py +++ b/user_guides/advanced/02_embedder_configuration.py @@ -261,9 +261,7 @@ ) # Example (does not run training here): construct an embedder and call train when you have data. -_embedder_for_ft = Embedder( - SentenceTransformerEmbeddingConfig(model_name="sentence-transformers/all-MiniLM-L6-v2") -) +_embedder_for_ft = Embedder(SentenceTransformerEmbeddingConfig(model_name="sentence-transformers/all-MiniLM-L6-v2")) # _embedder_for_ft.train(utterances=[...], labels=[...], config=ft_cfg) # %% From f927729e2c0d736023d7e7c88bbb1d84bcaf5a57 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 01:45:25 +0300 Subject: [PATCH 05/43] add more handling --- src/autointent/_advisor/_cli.py | 2 +- src/autointent/_advisor/_estimates.py | 171 +++++++++++++++--- src/autointent/_advisor/_render.py | 76 ++++++-- src/autointent/_advisor/_report.py | 17 +- .../_presets/transformers-heavy.yaml | 7 + tests/advisor/test_estimates_and_cli.py | 6 +- tests/advisor/test_estimates_internals.py | 131 ++++++++++++-- tests/advisor/test_render.py | 41 +++-- tests/advisor/test_report.py | 26 +-- 9 files changed, 385 insertions(+), 92 deletions(-) diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py index b3f43aab5..d300ad6fa 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/_advisor/_cli.py @@ -153,7 +153,7 @@ def cmd_recommend(args: argparse.Namespace) -> int: if args.budget_time_h is not None and report.resource.time_hours > args.budget_time_h: report.add( "resource", - Severity.RED, + Severity.OVER, f"Estimated time {report.resource.time_hours:.1f} h exceeds budget {args.budget_time_h} h.", ) results.append((preset, report)) diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py index e8f619303..06d0e4fe1 100644 --- a/src/autointent/_advisor/_estimates.py +++ b/src/autointent/_advisor/_estimates.py @@ -33,6 +33,10 @@ TRANSFORMER_SCORER_MODULES = {"bert", "lora", "ptuning", "dnnc"} +# Fallback max_length when the search-space entry doesn't pin it. Used both as +# the default in _vram_for_transformer and in the entry-walk seq_len resolution. +_DEFAULT_SEQ_LEN = 128 + # Coefficients for the linear / catboost time formulas (proposal §"Algorithm"). _LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-8 _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 @@ -95,12 +99,12 @@ def _walk_modules_indexed( yield node_idx, node_type, entry -def _vram_for_transformer(meta: ModelMeta, mode: str, mixed_precision: bool) -> float: - """VRAM in GB for one trial of a transformer-based module. +def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: + """Weight-side VRAM in GB — weights + grads + Adam optimizer state. Excludes activations. - Full fine-tune fp32: weights + grads + Adam (m, v) = 4W. - Full fine-tune AMP: fp16 weights + fp16 grads + fp32 master copy + fp32 Adam = 3W. - (Activations are not modeled separately.) + Full fine-tune fp32: W + W + 2W (Adam m, v) = 4W. + Full fine-tune AMP: 0.5W (fp16 weights) + 0.5W (fp16 grads) + W (fp32 master) + 2W (fp32 Adam) = 4W. + AMP's savings live in activations, not the optimizer — the weight side is identical. """ weights_gb = meta.weights_gb if mode == "inference": @@ -109,16 +113,111 @@ def _vram_for_transformer(meta: ModelMeta, mode: str, mixed_precision: bool) -> return weights_gb * 1.3 + 0.5 if mode == "reranker": return weights_gb * 1.5 - if mixed_precision: - return weights_gb * 3.0 return weights_gb * 4.0 +def _vram_for_transformer( + meta: ModelMeta, + mode: str, + mixed_precision: bool, + *, + batch_size: int = 0, + seq_len: int = _DEFAULT_SEQ_LEN, +) -> float: + """Total VRAM in GB: weights + grads + optimizer state + activations × batch. + + Activation accounting differs by mode — training keeps per-layer outputs for + backward; inference only needs one or two layers in flight. + """ + base = _weights_vram_for_transformer(meta, mode) + if batch_size <= 0: + return base + per_sample = _activations_gb_per_sample( + meta, seq_len, mixed_precision=mixed_precision, is_training=mode != "inference" + ) + return base + per_sample * batch_size + + def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: """RAM in GB. Loose upper bound.""" return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / (1024**3) +def _floor_to_power_of_two(n: int) -> int: + """Largest power of two ≤ n; returns 0 when n < 1.""" + if n < 1: + return 0 + power = 1 + while power * 2 <= n: + power *= 2 + return power + + +def _n_layers(meta: ModelMeta | None) -> int: + """Coarse layer-count guess from parameter count. + + MiniLM (33M) ~6, BERT-base (110M) ~12, BERT-large (350M) ~24. + """ + if meta is None: + return 12 + params = meta.params_millions + if params >= 300: + return 24 + if params >= 100: + return 12 + if params >= 50: + return 8 + return 6 + + +def _activations_gb_per_sample( + meta: ModelMeta | None, + seq_len: int, + *, + mixed_precision: bool, + is_training: bool, +) -> float: + """Heuristic activation memory per sample. + + Training: ``seq_len × hidden × layers × const`` — per-layer outputs are kept + for backward. + Inference: ``seq_len × hidden × const`` — only one or two layers' outputs in + flight at once. + Mixed precision halves activation bytes. + """ + hidden = _embedder_dim(meta) + if is_training: + # Training keeps every layer's outputs for backward → scales × n_layers. + # The 16-byte/token/layer coefficient bundles fp32 activation + ~4× backward overhead. + bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 + else: + # Inference only holds ~1-2 layers' outputs in flight at once. + bytes_per_sample = seq_len * hidden * 8 + if mixed_precision: + bytes_per_sample //= 2 + return bytes_per_sample / (1024**3) + + +def _max_fitting_batch_size( + *, + weight_vram_gb: float, + vram_budget_gb: float, + per_sample_gb: float, +) -> int: + """Largest batch that keeps total VRAM under the AMPLE/TIGHT threshold. + + Returns 0 when even the weights blow the budget. Result is rounded down to + the nearest power of two. + """ + if per_sample_gb <= 0: + return 0 + target_vram = vram_budget_gb * _YELLOW + available_for_activations = target_vram - weight_vram_gb + if available_for_activations <= 0: + return 0 + return _floor_to_power_of_two(int(available_for_activations / per_sample_gb)) + + def _embedder_dim(meta: ModelMeta | None) -> int: """Coarse hidden-size guess from parameter count. @@ -211,13 +310,13 @@ def _time_for_transformer( def _classify_severity(estimate: float, budget: float) -> Severity: if budget <= 0: - return Severity.YELLOW + return Severity.TIGHT ratio = estimate / budget if ratio >= _RED: - return Severity.RED + return Severity.OVER if ratio >= _YELLOW: - return Severity.YELLOW - return Severity.GREEN + return Severity.TIGHT + return Severity.AMPLE def _resource_phase( # noqa: PLR0912 - kept linear for clarity @@ -281,20 +380,31 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity batch_size = _max_int(entry.get("batch_size"), 32) epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) + seq_len = _max_int(entry.get("max_length"), _DEFAULT_SEQ_LEN) - vram = _vram_for_transformer(meta, mode, mixed_precision) + vram = _vram_for_transformer(meta, mode, mixed_precision, batch_size=batch_size, seq_len=seq_len) ram = _ram_for_module(meta, stats) - time_h = 0.0 - if mode != "inference": - time_h = _time_for_transformer( - meta=meta, - n_trials=n_trials, - epochs=epochs, - batch_size=batch_size, - n_samples=stats.n_samples, - device_class=hardware.device_class, + driver_max_batch: int | None = None + if hardware.vram_gb > 0: + weights_vram = _weights_vram_for_transformer(meta, mode) + per_sample_gb = _activations_gb_per_sample( + meta, seq_len, mixed_precision=mixed_precision, is_training=mode != "inference" ) + driver_max_batch = _max_fitting_batch_size( + weight_vram_gb=weights_vram, + vram_budget_gb=hardware.vram_gb, + per_sample_gb=per_sample_gb, + ) + + time_h = _time_for_transformer( + meta=meta, + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + device_class=hardware.device_class, + ) if refit_after and mode != "inference": time_h *= 1 + 1.0 / max(1, n_trials) @@ -311,6 +421,8 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity "vram_gb": round(vram, 2), "ram_gb": round(ram, 2), "time_hours": round(time_h, 2), + "batch_size": batch_size, + "max_batch_size": driver_max_batch, "confidence": meta.confidence, } ) @@ -381,6 +493,8 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity "vram_gb": round(vram, 2), "ram_gb": round(ram, 2), "time_hours": round(time_h, 2), + "batch_size": None, + "max_batch_size": None, "confidence": confidence, } ) @@ -409,7 +523,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity if hardware.accelerator == "cpu" and effective_vram > 0: report.add( "resource", - Severity.YELLOW, + Severity.TIGHT, f"No GPU detected; transformer modules will be very slow (worst case ~{estimate.time_hours:.1f} h).", metric="vram", ) @@ -420,6 +534,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity msg += f" vs available {hardware.vram_gb:.1f} GB" report.add("resource", vram_sev, msg, metric="vram") + ram_sev = _classify_severity(estimate.ram_gb, hardware.ram_gb) report.add( "resource", @@ -440,7 +555,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity if estimate.time_hours > 0: time_msg = f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)" - report.add("resource", Severity.GREEN, time_msg, metric="time") + report.add("resource", Severity.AMPLE, time_msg, metric="time") def _config_phase( @@ -454,7 +569,7 @@ def _config_phase( if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: report.add( "config", - Severity.YELLOW, + Severity.TIGHT, f"hpo_config.n_jobs={n_jobs} on a single GPU multiplies VRAM demand by {n_jobs}×.", ) @@ -466,7 +581,7 @@ def _config_phase( if uses_catboost_gpu and hardware.accelerator != "cuda": report.add( "config", - Severity.YELLOW, + Severity.TIGHT, "CatBoost task_type=GPU configured but no CUDA detected — will fall back to CPU.", ) @@ -484,7 +599,7 @@ def _data_phase( continue max_len = _max_int(max_len_value, 512) if p95 > max_len: - severity = Severity.RED if p95 > max_len * 1.5 else Severity.YELLOW + severity = Severity.OVER if p95 > max_len * 1.5 else Severity.TIGHT report.add( "data", severity, @@ -496,7 +611,7 @@ def _data_phase( if has_linear and stats.rare_classes: report.add( "data", - Severity.RED, + Severity.OVER, (f"LogisticRegressionCV (cv=3) will fail: classes {stats.rare_classes[:5]} have <3 samples."), ) @@ -507,7 +622,7 @@ def _data_phase( if has_description and stats.has_descriptions is False: report.add( "data", - Severity.RED, + Severity.OVER, "description scorer present but intent descriptions are missing — fill them in or drop the scorer.", ) diff --git a/src/autointent/_advisor/_render.py b/src/autointent/_advisor/_render.py index fe0f32dd7..a3778d307 100644 --- a/src/autointent/_advisor/_render.py +++ b/src/autointent/_advisor/_render.py @@ -13,11 +13,69 @@ if TYPE_CHECKING: from ._report import PreflightReport -_SEVERITY_TAG = {"green": "✓", "yellow": "⚠", "red": "✗"} +_SEVERITY_TAG = {"ample": "✓", "tight": "⚠", "over": "✗"} _PHASE_ORDER = ("resource", "data", "config") _PHASE_LABEL = {"resource": "Resource", "data": "Data", "config": "Config"} +def _batch_hint(driver: dict) -> str: + """Per-driver batch annotation: '64 → 32', '64', '64 (no fit)', or ''.""" + bs = driver.get("batch_size") + if bs is None: + return "" + mx = driver.get("max_batch_size") + if mx is None: + return str(bs) + if mx == 0: + return f"{bs} (no fit)" + if mx == bs: + return str(bs) + return f"{bs} → {mx}" + + +_DRIVERS_LIMIT = 8 +_DRIVERS_HEADERS = ("Node", "Model", "Mode", "VRAM", "Time", "Batch", "Source") + + +def _render_drivers_table(drivers: list[dict]) -> list[str]: + """Format the Drivers of cost section as an aligned table.""" + visible = drivers[:_DRIVERS_LIMIT] + rows: list[tuple[str, ...]] = [] + for d in visible: + rows.append(( + f"{d['node_type']}.{d['module']}", + str(d["model"]), + str(d["mode"]), + f"{d['vram_gb']:.2f} GB", + f"{d['time_hours']:.2f} h", + _batch_hint(d), + f"[{d['confidence']}]", + )) + + widths = [len(h) for h in _DRIVERS_HEADERS] + for row in rows: + for i, cell in enumerate(row): + widths[i] = max(widths[i], len(cell)) + + # Right-align numeric columns (VRAM @ idx 3, Time @ idx 4); left-align the rest. + right_align = {3, 4} + + def fmt(row: tuple[str, ...]) -> str: + cells = [] + for i, cell in enumerate(row): + if i in right_align: + cells.append(cell.rjust(widths[i])) + else: + cells.append(cell.ljust(widths[i])) + return " " + " ".join(cells).rstrip() + + lines = ["Drivers of cost:", fmt(_DRIVERS_HEADERS), " " + " ".join("─" * w for w in widths)] + lines.extend(fmt(r) for r in rows) + if len(drivers) > _DRIVERS_LIMIT: + lines.append(f" … and {len(drivers) - _DRIVERS_LIMIT} more") + return lines + + def render_text(report: PreflightReport) -> str: lines: list[str] = [] title = "Compute feasibility check" @@ -50,15 +108,7 @@ def render_text(report: PreflightReport) -> str: lines.append("") if report.resource.drivers: - lines.append("Drivers of cost:") - for d in report.resource.drivers[:8]: - lines.append( - f" {d['node_type']}.{d['module']:<10} {d['model']:<48}" - f" {d['mode']:<14} VRAM ~{d['vram_gb']} GB, time ~{d['time_hours']} h" - f" [{d['confidence']}]" - ) - if len(report.resource.drivers) > 8: - lines.append(f" … and {len(report.resource.drivers) - 8} more") + lines.extend(_render_drivers_table(report.resource.drivers)) lines.append("") if report.notes: @@ -68,7 +118,7 @@ def render_text(report: PreflightReport) -> str: lines.append("") summary = f"Verdict: {'feasible' if report.is_feasible else 'INFEASIBLE'} " - summary += f"(worst severity: {report.worst_severity.value})" + summary += f"(headroom: {report.headroom.value})" if report.low_confidence: summary += " — low-confidence (heuristic fallback in use)" lines.append(summary) @@ -91,7 +141,7 @@ def render_recommendation( else: lines.append(" → none of the bundled presets fit your hardware as-is.") lines.append("") - lines.append(f"{'Preset':<24} {'Status':<14} {'VRAM':<10} {'Time':<10} {'Worst':<8}") + lines.append(f"{'Preset':<24} {'Status':<14} {'VRAM':<10} {'Time':<10} {'Headroom':<10}") lines.append("-" * 68) for name, report in results: verdict = "feasible" if report.is_feasible else "infeasible" @@ -99,6 +149,6 @@ def render_recommendation( f"{name:<24} {verdict:<14} " f"{report.resource.vram_gb:>4.1f} GB " f"{report.resource.time_hours:>4.1f} h " - f"{report.worst_severity.value:<8}" + f"{report.headroom.value:<8}" ) return "\n".join(lines) diff --git a/src/autointent/_advisor/_report.py b/src/autointent/_advisor/_report.py index 6b930db95..9b4a319c8 100644 --- a/src/autointent/_advisor/_report.py +++ b/src/autointent/_advisor/_report.py @@ -8,9 +8,9 @@ class Severity(str, Enum): - GREEN = "green" - YELLOW = "yellow" - RED = "red" + AMPLE = "ample" + TIGHT = "tight" + OVER = "over" Phase = Literal["resource", "data", "config"] @@ -93,19 +93,20 @@ def add(self, phase: Phase, severity: Severity, message: str, metric: str | None self.findings.append(Finding(phase=phase, severity=severity, message=message, metric=metric)) @property - def worst_severity(self) -> Severity: - order = {Severity.GREEN: 0, Severity.YELLOW: 1, Severity.RED: 2} + def headroom(self) -> Severity: + """Worst headroom level across all findings — the column shown in CLI reports.""" + order = {Severity.AMPLE: 0, Severity.TIGHT: 1, Severity.OVER: 2} if not self.findings: - return Severity.GREEN + return Severity.AMPLE return max((f.severity for f in self.findings), key=lambda s: order[s]) @property def is_feasible(self) -> bool: - return self.worst_severity != Severity.RED + return self.headroom != Severity.OVER def to_dict(self) -> dict[str, Any]: d = asdict(self) d["findings"] = [{**asdict(f), "severity": f.severity.value} for f in self.findings] - d["worst_severity"] = self.worst_severity.value + d["headroom"] = self.headroom.value d["is_feasible"] = self.is_feasible return d diff --git a/src/autointent/_presets/transformers-heavy.yaml b/src/autointent/_presets/transformers-heavy.yaml index 2576fbc82..cd15d791e 100644 --- a/src/autointent/_presets/transformers-heavy.yaml +++ b/src/autointent/_presets/transformers-heavy.yaml @@ -5,12 +5,19 @@ search_space: - module_name: bert classification_model_config: - model_name: microsoft/deberta-v3-large + - model_name: intfloat/multilingual-e5-large-instruct + - model_name: microsoft/harrier-oss-v1-27b num_train_epochs: [30] batch_size: [32, 64] learning_rate: low: 1.0e-5 high: 1.0e-4 log: True + - module_name: description_bi + embedder_config: + - model_name: microsoft/deberta-v3-large + - model_name: intfloat/multilingual-e5-large-instruct + - model_name: microsoft/harrier-oss-v1-27b - node_type: decision target_metric: decision_accuracy search_space: diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 00537226d..15a087e07 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -98,7 +98,7 @@ def test_cli_inspect_json_is_parseable(capsys: pytest.CaptureFixture[str]) -> No payload = json.loads(captured.out) assert payload["preset_name"] == "transformers-light" assert "findings" in payload - assert payload["worst_severity"] in {"green", "yellow", "red"} + assert payload["headroom"] in {"ample", "tight", "over"} # rc is 0 on feasible, 1 otherwise assert rc in (0, 1) @@ -217,9 +217,9 @@ def test_cli_recommend_budget_time_flags_red_for_overbudget_presets( flagged = [ r for r in payload["results"] - if any(f["severity"] == "red" and "exceeds budget" in f["message"] for f in r["report"]["findings"]) + if any(f["severity"] == "over" and "exceeds budget" in f["message"] for f in r["report"]["findings"]) ] - assert flagged, "budget-time-h breach should produce RED severity findings" + assert flagged, "budget-time-h breach should produce OVER severity findings" # Any preset above the budget must be marked infeasible. for r in flagged: assert r["report"]["is_feasible"] is False diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 713db27f7..5ac66af2f 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -86,17 +86,17 @@ def test_empty_entry(self) -> None: class TestClassifySeverity: def test_below_yellow_is_green(self) -> None: - assert _classify_severity(estimate=1.0, budget=10.0) == Severity.GREEN + assert _classify_severity(estimate=1.0, budget=10.0) == Severity.AMPLE def test_above_yellow_threshold(self) -> None: - assert _classify_severity(estimate=8.0, budget=10.0) == Severity.YELLOW + assert _classify_severity(estimate=8.0, budget=10.0) == Severity.TIGHT def test_at_or_above_red_threshold(self) -> None: - assert _classify_severity(estimate=10.0, budget=10.0) == Severity.RED - assert _classify_severity(estimate=12.0, budget=10.0) == Severity.RED + assert _classify_severity(estimate=10.0, budget=10.0) == Severity.OVER + assert _classify_severity(estimate=12.0, budget=10.0) == Severity.OVER def test_zero_budget_returns_yellow(self) -> None: - assert _classify_severity(estimate=1.0, budget=0.0) == Severity.YELLOW + assert _classify_severity(estimate=1.0, budget=0.0) == Severity.TIGHT class TestVramForTransformer: @@ -117,13 +117,34 @@ def test_full_finetune_is_larger_than_lora_is_larger_than_inference(self, meta: full = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) assert inference < lora < full - def test_amp_partially_reduces_full_finetune_vram(self, meta: ModelMeta) -> None: - """AMP saves on fp16 weights+grads (W down from 2W); Adam state stays - fp32 (2W). Total 3W vs fp32's 4W — real but not a full halving.""" - full_fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) - full_amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True) - assert full_amp < full_fp32 - assert full_amp / full_fp32 == pytest.approx(0.75) + def test_inference_activations_are_smaller_than_training(self, meta: ModelMeta) -> None: + """Inference doesn't store per-layer outputs for backward — activation memory + should be many times smaller than training at the same batch_size.""" + train_total = _vram_for_transformer(meta, "full-finetune", False, batch_size=64, seq_len=128) + train_weights = _vram_for_transformer(meta, "full-finetune", False, batch_size=0) + inf_total = _vram_for_transformer(meta, "inference", False, batch_size=64, seq_len=128) + inf_weights = _vram_for_transformer(meta, "inference", False, batch_size=0) + train_acts = train_total - train_weights + inf_acts = inf_total - inf_weights + assert inf_acts > 0 + assert train_acts > inf_acts + # 12-layer model: training activations should be at least ~5× inference. + assert train_acts / inf_acts > 5 + + def test_amp_does_not_reduce_weight_side_vram(self, meta: ModelMeta) -> None: + """Weight-side AMP accounting: fp16 weights+grads (W) + fp32 master copy (W) + + fp32 Adam moments (2W) = 4W, identical to pure fp32. AMP's savings live + in activations, not the optimizer.""" + full_fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False, batch_size=0) + full_amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True, batch_size=0) + assert full_amp == pytest.approx(full_fp32) + + def test_amp_does_reduce_activation_side_vram(self, meta: ModelMeta) -> None: + """When a batch is configured, AMP halves activation bytes — total VRAM + with batch should be strictly smaller under AMP than fp32.""" + fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False, batch_size=64, seq_len=128) + amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True, batch_size=64, seq_len=128) + assert amp < fp32 def test_reranker_uses_inference_class(self, meta: ModelMeta) -> None: inference = _vram_for_transformer(meta, "inference", mixed_precision=False) @@ -255,7 +276,7 @@ def test_rare_classes_with_linear_scorer_flag_red(self) -> None: ) report = run_preflight(cfg, stats, _profile()) assert any( - f.phase == "data" and "LogisticRegressionCV" in f.message and f.severity == Severity.RED + f.phase == "data" and "LogisticRegressionCV" in f.message and f.severity == Severity.OVER for f in report.findings ) @@ -276,7 +297,7 @@ def test_truncation_red_when_p95_dominates_max_length(self) -> None: } stats = DatasetStats(n_samples=500, n_classes=5, avg_tokens=50, p95_tokens=400) report = run_preflight(cfg, stats, _profile()) - red = [f for f in report.findings if f.phase == "data" and f.severity == Severity.RED] + red = [f for f in report.findings if f.phase == "data" and f.severity == Severity.OVER] assert red, "p95=400 > 1.5 * max_length=128 should be red" def test_truncation_yellow_when_p95_only_slightly_exceeds(self) -> None: @@ -299,7 +320,7 @@ def test_truncation_yellow_when_p95_only_slightly_exceeds(self) -> None: yellows = [ f for f in report.findings - if f.phase == "data" and f.severity == Severity.YELLOW and "truncation" in f.message.lower() + if f.phase == "data" and f.severity == Severity.TIGHT and "truncation" in f.message.lower() ] assert yellows @@ -431,6 +452,86 @@ def test_linear_scales_with_n_samples(self) -> None: assert big.resource.ram_gb > small.resource.ram_gb +class TestPerDriverBatchHint: + """Each transformer driver carries its own (batch_size, max_batch_size) for rendering.""" + + def _bert_cfg(self, model_name: str, batch_size: int) -> dict[str, Any]: + return { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": model_name}], + "num_train_epochs": [3], + "batch_size": [batch_size], + } + ], + } + ], + "hpo_config": {"n_trials": 1}, + } + + def test_driver_records_current_and_max_batch(self) -> None: + report = run_preflight( + self._bert_cfg("microsoft/deberta-v3-large", batch_size=64), + DatasetStats.placeholder(), + _profile(vram_gb=10.0), + ) + drivers = [d for d in report.resource.drivers if d["module"] == "bert"] + assert drivers + d = drivers[0] + assert d["batch_size"] == 64 + # vram_gb=10 + 5 GB weights → some room for activations, max < 64. + assert d["max_batch_size"] is not None + assert 0 < d["max_batch_size"] < 64 + + def test_max_batch_zero_when_weights_alone_overflow(self) -> None: + report = run_preflight( + self._bert_cfg("microsoft/deberta-v3-large", batch_size=64), + DatasetStats.placeholder(), + _profile(vram_gb=2.0), + ) + d = next(d for d in report.resource.drivers if d["module"] == "bert") + assert d["max_batch_size"] == 0 + + def test_max_batch_can_be_larger_than_current(self) -> None: + report = run_preflight( + self._bert_cfg("microsoft/deberta-v3-large", batch_size=32), + DatasetStats.placeholder(), + _profile(vram_gb=64.0), + ) + d = next(d for d in report.resource.drivers if d["module"] == "bert") + assert d["max_batch_size"] is not None and d["max_batch_size"] > 32 + + def test_multiple_drivers_carry_independent_max_batch(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [ + {"model_name": "microsoft/deberta-v3-small"}, + {"model_name": "microsoft/deberta-v3-large"}, + ], + "num_train_epochs": [3], + "batch_size": [64], + } + ], + } + ], + "hpo_config": {"n_trials": 1}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile(vram_gb=10.0)) + small = next(d for d in report.resource.drivers if "small" in d["model"]) + large = next(d for d in report.resource.drivers if "large" in d["model"]) + # The smaller model has more headroom → larger max batch (or equal-cap when both saturate). + assert small["max_batch_size"] >= large["max_batch_size"] + + class TestDumpModulesBounding: """`dump_modules=True` writes one selected variant per node per trial — not every candidate. The estimate must be bounded by sum-of-max-per-node x n_trials.""" diff --git a/tests/advisor/test_render.py b/tests/advisor/test_render.py index 55a2b0ce3..2c0604a11 100644 --- a/tests/advisor/test_render.py +++ b/tests/advisor/test_render.py @@ -4,7 +4,7 @@ import json -from autointent._advisor._render import render_json, render_recommendation, render_text +from autointent._advisor._render import _batch_hint, render_json, render_recommendation, render_text from autointent._advisor._report import ( DatasetStats, PreflightReport, @@ -46,8 +46,8 @@ def _populated_report() -> PreflightReport: ), notes=["MPS unified memory note"], ) - r.add("resource", Severity.YELLOW, "VRAM ~6 GB vs available 8 GB") - r.add("data", Severity.RED, "rare classes blocked") + r.add("resource", Severity.TIGHT, "VRAM ~6 GB vs available 8 GB") + r.add("data", Severity.OVER, "rare classes blocked") return r @@ -64,10 +64,10 @@ def test_includes_drivers_block(self) -> None: assert "Drivers of cost:" in out assert "x/y" in out - def test_verdict_reflects_worst_severity(self) -> None: + def test_verdict_reflects_headroom(self) -> None: out = render_text(_populated_report()) assert "Verdict: INFEASIBLE" in out - assert "worst severity: red" in out + assert "headroom: over" in out def test_disclaimer_always_present(self) -> None: out = render_text(_populated_report()) @@ -96,25 +96,25 @@ def test_is_valid_json(self) -> None: def test_findings_have_string_severity(self) -> None: d = json.loads(render_json(_populated_report())) for f in d["findings"]: - assert f["severity"] in {"green", "yellow", "red"} + assert f["severity"] in {"ample", "tight", "over"} - def test_worst_severity_and_feasibility_serialized(self) -> None: + def test_headroom_and_feasibility_serialized(self) -> None: d = json.loads(render_json(_populated_report())) - assert d["worst_severity"] == "red" + assert d["headroom"] == "over" assert d["is_feasible"] is False def test_empty_report_serializes(self) -> None: d = json.loads(render_json(PreflightReport())) - assert d["worst_severity"] == "green" + assert d["headroom"] == "ample" assert d["is_feasible"] is True class TestRenderRecommendation: def _two_reports(self) -> list[tuple[str, PreflightReport]]: a = PreflightReport(preset_name="a", resource=ResourceEstimate(vram_gb=2.0, time_hours=0.5)) - a.add("resource", Severity.GREEN, "ok") + a.add("resource", Severity.AMPLE, "ok") b = PreflightReport(preset_name="b", resource=ResourceEstimate(vram_gb=8.0, time_hours=4.0)) - b.add("resource", Severity.RED, "too big") + b.add("resource", Severity.OVER, "too big") return [("a", a), ("b", b)] def test_lists_chosen_preset_when_present(self) -> None: @@ -136,6 +136,25 @@ def test_shows_status_per_preset(self) -> None: assert "infeasible" in out +class TestBatchHint: + """Per-driver batch cell rendered in the Drivers-of-cost table.""" + + def test_arrow_when_max_differs(self) -> None: + assert _batch_hint({"batch_size": 64, "max_batch_size": 32}) == "64 → 32" + + def test_plain_when_max_equals_current(self) -> None: + assert _batch_hint({"batch_size": 64, "max_batch_size": 64}) == "64" + + def test_no_fit_label_when_max_zero(self) -> None: + assert _batch_hint({"batch_size": 64, "max_batch_size": 0}) == "64 (no fit)" + + def test_empty_when_no_batch(self) -> None: + assert _batch_hint({"batch_size": None, "max_batch_size": None}) == "" + + def test_increase_arrow(self) -> None: + assert _batch_hint({"batch_size": 32, "max_batch_size": 128}) == "32 → 128" + + def test_dataset_stats_in_text_block() -> None: stats = DatasetStats.placeholder(n_samples=777, n_classes=4) r = PreflightReport( diff --git a/tests/advisor/test_report.py b/tests/advisor/test_report.py index 52f2e675e..28adbfc34 100644 --- a/tests/advisor/test_report.py +++ b/tests/advisor/test_report.py @@ -14,22 +14,22 @@ class TestSeverityOrdering: - def test_worst_severity_on_empty_report_is_green(self) -> None: - assert PreflightReport().worst_severity == Severity.GREEN + def test_headroom_on_empty_report_is_green(self) -> None: + assert PreflightReport().headroom == Severity.AMPLE def test_red_beats_yellow_beats_green(self) -> None: r = PreflightReport() - r.add("resource", Severity.GREEN, "ok") - r.add("data", Severity.YELLOW, "warn") - assert r.worst_severity == Severity.YELLOW - r.add("config", Severity.RED, "fail") - assert r.worst_severity == Severity.RED + r.add("resource", Severity.AMPLE, "ok") + r.add("data", Severity.TIGHT, "warn") + assert r.headroom == Severity.TIGHT + r.add("config", Severity.OVER, "fail") + assert r.headroom == Severity.OVER def test_is_feasible_flips_on_any_red(self) -> None: r = PreflightReport() - r.add("resource", Severity.YELLOW, "warn") + r.add("resource", Severity.TIGHT, "warn") assert r.is_feasible is True - r.add("data", Severity.RED, "fail") + r.add("data", Severity.OVER, "fail") assert r.is_feasible is False @@ -62,12 +62,12 @@ def test_total_disk_ignores_cached(self) -> None: class TestToDictSerialization: def test_findings_round_trip_severity_as_string(self) -> None: r = PreflightReport() - r.add("resource", Severity.RED, "boom") + r.add("resource", Severity.OVER, "boom") d = r.to_dict() - assert d["worst_severity"] == "red" + assert d["headroom"] == "over" assert d["is_feasible"] is False assert d["findings"] == [ - {"phase": "resource", "severity": "red", "message": "boom", "metric": None}, + {"phase": "resource", "severity": "over", "message": "boom", "metric": None}, ] def test_hardware_and_dataset_pass_through(self) -> None: @@ -80,6 +80,6 @@ def test_hardware_and_dataset_pass_through(self) -> None: assert d["dataset"]["n_samples"] == 100 def test_finding_is_frozen(self) -> None: - f = Finding(phase="resource", severity=Severity.GREEN, message="ok") + f = Finding(phase="resource", severity=Severity.AMPLE, message="ok") with pytest.raises(Exception): # noqa: PT011 - dataclass.FrozenInstanceError varies f.message = "changed" # type: ignore[misc] From 82a78287bdbb1bec9e4bce6dda9804f7a801fe19 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 02:26:01 +0300 Subject: [PATCH 06/43] add more handling --- pyproject.toml | 1 + src/autointent/_advisor/_cli.py | 101 +++++++++++++++--- src/autointent/_advisor/_estimates.py | 135 +++++++++++++++--------- src/autointent/_advisor/_hub.py | 10 +- tests/advisor/test_estimates_and_cli.py | 2 +- 5 files changed, 183 insertions(+), 66 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8675e7ea3..4ee34b9c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "aiometer (>=1.0.0,<2.0.0)", "aiofiles (>=24.1.0,<25.0.0)", "threadpoolctl (>=3.0.0,<4.0.0)", + "psutil (>=5.9.0,<8.0.0)", ] [project.optional-dependencies] diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py index d300ad6fa..c9485cc81 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/_advisor/_cli.py @@ -5,9 +5,10 @@ * ``inspect`` — show what a given preset / config will cost on this machine. * ``recommend`` — pick the best-fitting bundled preset for this machine. -Both subcommands accept either a real ``--dataset`` (path to load with -``Dataset.from_*`` constructors) or ``--n-samples / --n-classes / --avg-tokens`` -placeholders so the script is useful before the user has built a dataset. +Both subcommands accept either a real ``--dataset`` (Hub id or local +csv/json/jsonl/parquet path loaded via ``datasets.load_dataset``) or +``--n-samples / --n-classes / --avg-tokens`` placeholders so the script is +useful before the user has built a dataset. """ from __future__ import annotations @@ -19,8 +20,8 @@ from typing import Any import yaml +from datasets import ClassLabel, Sequence, load_dataset -from autointent import Dataset from autointent.utils import load_preset from ._estimates import run_preflight @@ -79,35 +80,107 @@ def _stats_from_args(args: argparse.Namespace) -> DatasetStats: ) +_UTTERANCE_COLS = ("utterance", "text", "sentence", "query", "input") +_LABEL_COLS = ("label", "labels", "intent", "target") +# Map file extension → datasets builder name. Anything else is treated as a Hub +# repo id or a directory and passed to load_dataset directly. +_FILE_BUILDERS = {".csv": "csv", ".tsv": "csv", ".json": "json", ".jsonl": "json", ".parquet": "parquet"} + + def _stats_from_dataset(path: str, *, multilabel: bool) -> DatasetStats: - """Best-effort: load a dataset from disk via the existing Dataset constructor.""" + """Best-effort: load via HF ``datasets.load_dataset``. + + Accepts a Hub repo id ('DeepPavlov/clinc150') or a local file path + (.csv / .json / .jsonl / .parquet) / dataset directory. Falls back to a + placeholder on any loader error so the advisor stays best-effort. + """ + builder = _FILE_BUILDERS.get(Path(path).suffix.lower()) try: - ds = Dataset.from_json(path) if path.endswith(".json") else Dataset.from_hub(path) - except (OSError, ValueError) as e: + ds = load_dataset(builder, data_files=path) if builder else load_dataset(path) + except (OSError, ValueError, FileNotFoundError) as e: logger.warning("Failed to load dataset %s: %s", path, e) return DatasetStats.placeholder(multilabel=multilabel) - train = ds.get("train") or next(iter(ds.values()), None) + train = ds["train"] if "train" in ds else next(iter(ds.values()), None) if train is None: return DatasetStats.placeholder(multilabel=multilabel) - utt_col = getattr(ds, "utterance_feature", "utterance") + cols = train.column_names + utt_col = next((c for c in _UTTERANCE_COLS if c in cols), cols[0] if cols else None) + label_col = next((c for c in _LABEL_COLS if c in cols), None) + + detected_multilabel, n_classes = _label_shape(train, label_col, fallback_multilabel=multilabel) + sample = train[:1000] if len(train) > 1000 else train[:] - lengths = [len(str(s).split()) for s in sample.get(utt_col, [])] + lengths = [len(str(s).split()) for s in (sample.get(utt_col, []) if utt_col else [])] avg_tokens = int(sum(lengths) / max(1, len(lengths))) if lengths else 32 - p95 = sorted(lengths)[int(len(lengths) * 0.95)] if lengths else avg_tokens * 2 + if lengths: + sorted_lengths = sorted(lengths) + idx = max(0, min(len(sorted_lengths) - 1, int(round((len(sorted_lengths) - 1) * 0.95)))) + p95 = sorted_lengths[idx] + else: + p95 = avg_tokens * 2 return DatasetStats( n_samples=len(train), - n_classes=getattr(ds, "n_classes", 0) or 0, + n_classes=n_classes, avg_tokens=avg_tokens, p95_tokens=p95, - multilabel=getattr(ds, "multilabel", multilabel), - has_descriptions=getattr(ds, "has_descriptions", None), + multilabel=detected_multilabel, + has_descriptions=None, + rare_classes=_rare_classes(train, label_col, detected_multilabel, n_classes) if label_col else [], source=f"dataset:{path}", ) +def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool) -> tuple[bool, int]: + """Derive (multilabel, n_classes) from the HF feature schema, with a value-based fallback.""" + if label_col is None: + return fallback_multilabel, 0 + feature = train.features.get(label_col) + if isinstance(feature, Sequence): + inner = feature.feature + if isinstance(inner, ClassLabel): + return True, inner.num_classes + # Sequence of plain ints — n_classes = max label index + 1. + max_idx = max((max(row) for row in train[label_col] if row), default=-1) + return True, max_idx + 1 + if isinstance(feature, ClassLabel): + return False, feature.num_classes + # Plain int/string column. Detect multilabel from the first non-empty row, then count uniques. + is_multi = len(train) > 0 and isinstance(train[0][label_col], (list, tuple)) + if is_multi: + max_idx = max((max(row) for row in train[label_col] if row), default=-1) + return True, max_idx + 1 + return False, len({label for label in train[label_col] if label is not None}) + + +def _rare_classes(train: Any, label_col: str, multilabel: bool, n_classes: int, min_count: int = 3) -> list[str]: + """Return labels with fewer than ``min_count`` samples in the train split. + + Used to surface the LogisticRegressionCV(cv=3) failure case before fit. + Returns an empty list on any error so the advisor stays best-effort. + """ + try: + labels = train[label_col] + except (KeyError, AttributeError, TypeError): + return [] + counts: dict[str, int] = {} + if multilabel: + for row in labels: + if not row: + continue + for i, v in enumerate(row): + if v: + counts[str(i)] = counts.get(str(i), 0) + 1 + for i in range(n_classes): + counts.setdefault(str(i), 0) + else: + for label in labels: + counts[str(label)] = counts.get(str(label), 0) + 1 + return sorted(name for name, c in counts.items() if c < min_count) + + def _add_common_dataset_args(p: argparse.ArgumentParser) -> None: p.add_argument("--dataset", help="Path or hub id of a dataset; overrides placeholders.") p.add_argument("--n-samples", type=int, default=1_000, help="Placeholder training set size.") diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py index 06d0e4fe1..e88fffd0a 100644 --- a/src/autointent/_advisor/_estimates.py +++ b/src/autointent/_advisor/_estimates.py @@ -12,15 +12,49 @@ from collections.abc import Iterable from typing import Any +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from autointent.configs._optimization import HPOConfig + from ._hardware import HardwareProfile from ._hub import ModelMeta, hub_reachable, resolve_model from ._report import DatasetStats, PreflightReport, ResourceEstimate, Severity logger = logging.getLogger(__name__) -# yellow / red thresholds as fraction of available budget -_YELLOW = 0.7 -_RED = 1.0 + +class _AdvisorConfig(BaseModel): + """Validated view of the advisor's input config. + + Wraps the four top-level keys the phase helpers read. Unknown top-level + keys are ignored (preset YAMLs carry extra metadata the advisor doesn't model). + """ + + model_config = ConfigDict(extra="ignore") + + hpo_config: HPOConfig = Field(default_factory=HPOConfig) + search_space: list[dict[str, Any]] = Field(default_factory=list) + refit_after: bool = False + dump_modules: bool = False + embedder_config: dict[str, Any] | None = None + + +def _validated_config(config: dict[str, Any]) -> _AdvisorConfig: + """Validate ``config`` against ``_AdvisorConfig``; fall back to defaults on any error. + + The advisor is best-effort: a malformed user config should still produce a + report (with placeholder costs) rather than crashing. + """ + try: + return _AdvisorConfig.model_validate(config) + except ValidationError as e: + logger.warning("Advisor config failed validation; falling back to defaults: %s", e) + return _AdvisorConfig() + +# Severity thresholds as a fraction of available budget: at or above _TIGHT +# downgrades to Severity.TIGHT; at or above _OVER downgrades to Severity.OVER. +_TIGHT_RATIO = 0.7 +_OVER_RATIO = 1.0 # rough per-step seconds, keyed on device class. Scaled by params_millions / 100. _PER_STEP_BASELINE_S = { @@ -31,7 +65,14 @@ "apple-silicon": 0.08, } -TRANSFORMER_SCORER_MODULES = {"bert", "lora", "ptuning", "dnnc"} +# Maps each fine-tunable transformer module to its training-mode label. +# Modules not listed are treated as inference-only. +_TRANSFORMER_TRAINING_MODE = { + "bert": "full-finetune", + "ptuning": "lora", + "lora": "lora", + "dnnc": "reranker", +} # Fallback max_length when the search-space entry doesn't pin it. Used both as # the default in _vram_for_transformer and in the entry-walk seq_len resolution. @@ -81,14 +122,6 @@ def _max_int(value: Any, default: int) -> int: return default -def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dict[str, Any]]]: - """Yield (node_type, module_entry) pairs.""" - for node in search_space or []: - node_type = node.get("node_type", "?") - for entry in node.get("search_space", []) or []: - yield node_type, entry - - def _walk_modules_indexed( search_space: list[dict[str, Any]], ) -> Iterable[tuple[int, str, dict[str, Any]]]: @@ -99,6 +132,12 @@ def _walk_modules_indexed( yield node_idx, node_type, entry +def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dict[str, Any]]]: + """Yield (node_type, module_entry) pairs — index-agnostic view over `_walk_modules_indexed`.""" + for _, node_type, entry in _walk_modules_indexed(search_space): + yield node_type, entry + + def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: """Weight-side VRAM in GB — weights + grads + Adam optimizer state. Excludes activations. @@ -211,7 +250,7 @@ def _max_fitting_batch_size( """ if per_sample_gb <= 0: return 0 - target_vram = vram_budget_gb * _YELLOW + target_vram = vram_budget_gb * _TIGHT_RATIO available_for_activations = target_vram - weight_vram_gb if available_for_activations <= 0: return 0 @@ -309,12 +348,14 @@ def _time_for_transformer( def _classify_severity(estimate: float, budget: float) -> Severity: + if estimate <= 0: + return Severity.AMPLE if budget <= 0: return Severity.TIGHT ratio = estimate / budget - if ratio >= _RED: + if ratio >= _OVER_RATIO: return Severity.OVER - if ratio >= _YELLOW: + if ratio >= _TIGHT_RATIO: return Severity.TIGHT return Severity.AMPLE @@ -325,28 +366,27 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity hardware: HardwareProfile, report: PreflightReport, ) -> None: - hpo = config.get("hpo_config") or {} - n_trials = int(hpo.get("n_trials", 1)) - n_jobs = int(hpo.get("n_jobs", 1)) - refit_after = bool(config.get("refit_after", False)) - dump_modules = bool(config.get("dump_modules", False)) + cfg = _validated_config(config) + n_trials = max(1, cfg.hpo_config.n_trials) + n_jobs = max(1, cfg.hpo_config.n_jobs) + refit_after = cfg.refit_after + dump_modules = cfg.dump_modules if not hub_reachable(): report.low_confidence = True report.notes.append("HF Hub unreachable — all model sizes are name-pattern heuristics.") seen_models: dict[str, ModelMeta] = {} - estimate = ResourceEstimate(parallel_factor=max(1, n_jobs)) + estimate = ResourceEstimate(parallel_factor=n_jobs) - embedder_cfg = config.get("embedder_config") or {} - global_embedder = embedder_cfg.get("model_name") if isinstance(embedder_cfg, dict) else None + global_embedder = (cfg.embedder_config or {}).get("model_name") if global_embedder: seen_models[global_embedder] = resolve_model(global_embedder) # First pass: walk transformer-bearing modules (collects seen_models for embedder_dim lookup). transformer_entries: list[tuple[int, str, dict[str, Any]]] = [] classic_entries: list[tuple[int, str, dict[str, Any]]] = [] - for node_idx, node_type, entry in _walk_modules_indexed(config.get("search_space") or []): + for node_idx, node_type, entry in _walk_modules_indexed(cfg.search_space): module = entry.get("module_name", "?") if module in {"linear", "catboost"}: classic_entries.append((node_idx, node_type, entry)) @@ -367,16 +407,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity meta = seen_models.setdefault(name, resolve_model(name)) mixed_precision = entry.get("dtype") in {"fp16", "bf16"} - if module == "bert": - mode = "full-finetune" - elif module == "lora": - mode = "lora" - elif module == "dnnc": - mode = "reranker" - elif module == "ptuning": - mode = "full-finetune" - else: - mode = "inference" + mode = _TRANSFORMER_TRAINING_MODE.get(module, "inference") batch_size = _max_int(entry.get("batch_size"), 32) epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) @@ -430,7 +461,11 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. embedder_meta = _largest_embedder(seen_models) embedder_dim = _embedder_dim(embedder_meta) - class_multiplier_classic = max(1, stats.n_classes) if stats.multilabel else 1 + # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes; + # the multiclass path additionally pays the LogisticRegressionCV inner-fit multiplier. + class_multiplier_classic = max(1, stats.n_classes) + confidence = embedder_meta.confidence if embedder_meta else "heuristic" + embedder_label = embedder_meta.name if embedder_meta else "(no embedder)" for _node_idx, node_type, entry in classic_entries: module = entry.get("module_name", "?") if module == "linear": @@ -449,14 +484,14 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity time_h *= 1 + 1.0 / max(1, n_trials) vram = 0.0 mode = "linear-cv" if cv_multiplier > 1 else "linear" - confidence = embedder_meta.confidence if embedder_meta else "heuristic" elif module == "catboost": iterations = _max_int(entry.get("iterations"), 1000) depth = _max_int(entry.get("depth"), 6) on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" - # CatBoost's multiclass MultiClass loss already grows per-class trees. - cb_class_mult = max(1, stats.n_classes) - ram = _ram_for_catboost( + # CatBoost's MultiClass loss grows per-class trees only above binary; + # binary uses Logloss with one tree per iteration. + cb_class_mult = max(1, stats.n_classes) if stats.n_classes > 2 or stats.multilabel else 1 + ram_total = _ram_for_catboost( stats=stats, n_features=embedder_dim, iterations=iterations, @@ -473,11 +508,8 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity ) if refit_after: time_h *= 1 + 1.0 / max(1, n_trials) - vram = ram if on_gpu else 0.0 - if on_gpu: - ram = 0.0 + vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) mode = "catboost-gpu" if on_gpu else "catboost" - confidence = embedder_meta.confidence if embedder_meta else "heuristic" else: continue @@ -488,7 +520,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity { "node_type": node_type, "module": module, - "model": embedder_meta.name if embedder_meta else "(no embedder)", + "model": embedder_label, "mode": mode, "vram_gb": round(vram, 2), "ram_gb": round(ram, 2), @@ -515,6 +547,9 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity effective_vram = estimate.vram_gb * n_jobs else: effective_vram = estimate.vram_gb + # MPS shares one unified pool: parallel workers each allocate weights+activations + # in RAM, so peak RAM also scales with n_jobs on Apple Silicon. + effective_ram = estimate.ram_gb * n_jobs if n_jobs > 1 and hardware.accelerator == "mps" else estimate.ram_gb report.resource = estimate @@ -535,11 +570,11 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity report.add("resource", vram_sev, msg, metric="vram") - ram_sev = _classify_severity(estimate.ram_gb, hardware.ram_gb) + ram_sev = _classify_severity(effective_ram, hardware.ram_gb) report.add( "resource", ram_sev, - f"RAM ~{estimate.ram_gb:.1f} GB vs available {hardware.ram_gb:.1f} GB", + f"RAM ~{effective_ram:.1f} GB vs available {hardware.ram_gb:.1f} GB", metric="ram", ) @@ -606,9 +641,10 @@ def _data_phase( f"Train tokens p95~{p95} exceeds {entry.get('module_name', '?')}.max_length={max_len}; expect silent truncation.", ) - # rare class × linear-CV + # rare class × linear-CV (LogisticRegressionCV cv=3 needs ≥3 samples/class; + # multilabel path uses one-vs-rest without CV so the failure can't occur there) has_linear = any(e.get("module_name") == "linear" for _, e in _walk_modules(config.get("search_space") or [])) - if has_linear and stats.rare_classes: + if has_linear and stats.rare_classes and not stats.multilabel: report.add( "data", Severity.OVER, @@ -616,8 +652,9 @@ def _data_phase( ) # partial descriptions × description scorer + description_modules = {"description_bi", "description_cross", "description_llm"} has_description = any( - e.get("module_name") == "description" for _, e in _walk_modules(config.get("search_space") or []) + e.get("module_name") in description_modules for _, e in _walk_modules(config.get("search_space") or []) ) if has_description and stats.has_descriptions is False: report.add( diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index 613ab6b40..1c559ee2f 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -76,7 +76,7 @@ def _is_warm_cached(model_name: str) -> bool: weight_files = ["model.safetensors", "pytorch_model.bin", "model.safetensors.index.json"] for fname in weight_files: path = try_to_load_from_cache(model_name, fname) - if path is not None and path is not False: + if isinstance(path, str): return True # sharded models won't match the single-file probe; fall back to a scan @@ -114,11 +114,17 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: if size: total_file_bytes += int(size) + # Track whether either size came from the Hub or from the name-pattern fallback; + # if any field was filled by heuristic, downgrade confidence so the report flips + # low_confidence rather than misreporting hub-grade accuracy. + confidence = "hub" if params_millions == 0: params_millions = _heuristic_params_millions(model_name) + confidence = "heuristic" if total_file_bytes == 0: total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) + confidence = "heuristic" return ModelMeta( name=model_name, @@ -126,7 +132,7 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: weight_bytes_per_param=weight_bytes_per_param, total_file_bytes=total_file_bytes, cached_locally=_is_warm_cached(model_name), - confidence="hub", + confidence=confidence, ) diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 15a087e07..3092dce9e 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -150,7 +150,7 @@ def test_partial_descriptions_with_description_scorer_flags_red() -> None: { "node_type": "scoring", "search_space": [ - {"module_name": "description"}, + {"module_name": "description_bi"}, ], } ], From bbb039e576467b479e8148a87bbfd4b34299976e Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 17:41:08 +0300 Subject: [PATCH 07/43] fix typing & lint --- pyproject.toml | 6 ++ src/autointent/_advisor/_cli.py | 26 ++++++-- src/autointent/_advisor/_estimates.py | 78 ++++++++++++----------- src/autointent/_advisor/_hardware.py | 14 ++-- src/autointent/_advisor/_hub.py | 6 +- src/autointent/_advisor/_render.py | 18 +++--- tests/advisor/test_estimates_internals.py | 5 +- tests/advisor/test_report.py | 4 +- 8 files changed, 94 insertions(+), 63 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4ee34b9c5..971de2655 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -296,6 +296,12 @@ module = [ "dspy.evaluate.auto_evaluation", "codecarbon", "catboost", + "openai", + "openai.*", + "tiktoken", + "peft", + "sentence_transformers", + "psutil", ] ignore_missing_imports = true diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py index c9485cc81..8c8b7b9d2 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/_advisor/_cli.py @@ -17,7 +17,7 @@ import logging import sys from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import yaml from datasets import ClassLabel, Sequence, load_dataset @@ -27,10 +27,16 @@ from ._estimates import run_preflight from ._hardware import detect_hardware from ._render import render_json, render_recommendation, render_text -from ._report import DatasetStats, PreflightReport, Severity +from ._report import DatasetStats, Severity + +if TYPE_CHECKING: + from ._report import PreflightReport logger = logging.getLogger("autointent.advisor") +_SAMPLE_LIMIT = 1000 +_P95_PERCENTILE = 0.95 + BUNDLED_PRESETS = [ "transformers-heavy", "transformers-light", @@ -111,12 +117,12 @@ def _stats_from_dataset(path: str, *, multilabel: bool) -> DatasetStats: detected_multilabel, n_classes = _label_shape(train, label_col, fallback_multilabel=multilabel) - sample = train[:1000] if len(train) > 1000 else train[:] + sample = train[:_SAMPLE_LIMIT] if len(train) > _SAMPLE_LIMIT else train[:] lengths = [len(str(s).split()) for s in (sample.get(utt_col, []) if utt_col else [])] avg_tokens = int(sum(lengths) / max(1, len(lengths))) if lengths else 32 if lengths: sorted_lengths = sorted(lengths) - idx = max(0, min(len(sorted_lengths) - 1, int(round((len(sorted_lengths) - 1) * 0.95)))) + idx = max(0, min(len(sorted_lengths) - 1, round((len(sorted_lengths) - 1) * _P95_PERCENTILE))) p95 = sorted_lengths[idx] else: p95 = avg_tokens * 2 @@ -133,7 +139,7 @@ def _stats_from_dataset(path: str, *, multilabel: bool) -> DatasetStats: ) -def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool) -> tuple[bool, int]: +def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool) -> tuple[bool, int]: # noqa: ANN401 """Derive (multilabel, n_classes) from the HF feature schema, with a value-based fallback.""" if label_col is None: return fallback_multilabel, 0 @@ -155,7 +161,13 @@ def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool return False, len({label for label in train[label_col] if label is not None}) -def _rare_classes(train: Any, label_col: str, multilabel: bool, n_classes: int, min_count: int = 3) -> list[str]: +def _rare_classes( + train: Any, # noqa: ANN401 + label_col: str, + multilabel: bool, + n_classes: int, + min_count: int = 3, +) -> list[str]: """Return labels with fewer than ``min_count`` samples in the train split. Used to surface the LogisticRegressionCV(cv=3) failure case before fit. @@ -293,7 +305,7 @@ def main(argv: list[str] | None = None) -> int: level=logging.DEBUG if args.verbose else logging.WARNING, format="%(levelname)s %(name)s: %(message)s", ) - return args.func(args) + return int(args.func(args)) if __name__ == "__main__": diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py index e88fffd0a..93dfcedaf 100644 --- a/src/autointent/_advisor/_estimates.py +++ b/src/autointent/_advisor/_estimates.py @@ -9,16 +9,26 @@ from __future__ import annotations import logging -from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import BaseModel, ConfigDict, Field, ValidationError from autointent.configs._optimization import HPOConfig -from ._hardware import HardwareProfile -from ._hub import ModelMeta, hub_reachable, resolve_model -from ._report import DatasetStats, PreflightReport, ResourceEstimate, Severity +from ._hub import hub_reachable, resolve_model +from ._report import PreflightReport, ResourceEstimate, Severity + +if TYPE_CHECKING: + from collections.abc import Iterable + + from ._hardware import HardwareProfile + from ._hub import ModelMeta + from ._report import DatasetStats + +_MULTICLASS_THRESHOLD = 2 +_PARAMS_LARGE = 300 +_PARAMS_BASE = 100 +_PARAMS_SMALL = 50 logger = logging.getLogger(__name__) @@ -51,6 +61,7 @@ def _validated_config(config: dict[str, Any]) -> _AdvisorConfig: logger.warning("Advisor config failed validation; falling back to defaults: %s", e) return _AdvisorConfig() + # Severity thresholds as a fraction of available budget: at or above _TIGHT # downgrades to Severity.TIGHT; at or above _OVER downgrades to Severity.OVER. _TIGHT_RATIO = 0.7 @@ -94,22 +105,18 @@ def _extract_model_names(module_entry: dict[str, Any]) -> list[str]: candidates: list[str] = [] cfg = module_entry.get("classification_model_config") if isinstance(cfg, list): - for c in cfg: - if isinstance(c, dict) and c.get("model_name"): - candidates.append(c["model_name"]) + candidates.extend(c["model_name"] for c in cfg if isinstance(c, dict) and c.get("model_name")) elif isinstance(cfg, dict) and cfg.get("model_name"): candidates.append(cfg["model_name"]) embedder_cfg = module_entry.get("embedder_config") if isinstance(embedder_cfg, list): - for c in embedder_cfg: - if isinstance(c, dict) and c.get("model_name"): - candidates.append(c["model_name"]) + candidates.extend(c["model_name"] for c in embedder_cfg if isinstance(c, dict) and c.get("model_name")) elif isinstance(embedder_cfg, dict) and embedder_cfg.get("model_name"): candidates.append(embedder_cfg["model_name"]) return candidates -def _max_int(value: Any, default: int) -> int: +def _max_int(value: Any, default: int) -> int: # noqa: ANN401 if value is None: return default if isinstance(value, list) and value: @@ -163,7 +170,7 @@ def _vram_for_transformer( batch_size: int = 0, seq_len: int = _DEFAULT_SEQ_LEN, ) -> float: - """Total VRAM in GB: weights + grads + optimizer state + activations × batch. + """Total VRAM in GB: weights + grads + optimizer state + activations x batch. Activation accounting differs by mode — training keeps per-layer outputs for backward; inference only needs one or two layers in flight. @@ -200,11 +207,11 @@ def _n_layers(meta: ModelMeta | None) -> int: if meta is None: return 12 params = meta.params_millions - if params >= 300: + if params >= _PARAMS_LARGE: return 24 - if params >= 100: + if params >= _PARAMS_BASE: return 12 - if params >= 50: + if params >= _PARAMS_SMALL: return 8 return 6 @@ -218,20 +225,17 @@ def _activations_gb_per_sample( ) -> float: """Heuristic activation memory per sample. - Training: ``seq_len × hidden × layers × const`` — per-layer outputs are kept + Training: ``seq_len x hidden x layers x const`` — per-layer outputs are kept for backward. - Inference: ``seq_len × hidden × const`` — only one or two layers' outputs in + Inference: ``seq_len x hidden x const`` — only one or two layers' outputs in flight at once. Mixed precision halves activation bytes. """ hidden = _embedder_dim(meta) - if is_training: - # Training keeps every layer's outputs for backward → scales × n_layers. - # The 16-byte/token/layer coefficient bundles fp32 activation + ~4× backward overhead. - bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 - else: - # Inference only holds ~1-2 layers' outputs in flight at once. - bytes_per_sample = seq_len * hidden * 8 + # Training keeps every layer's outputs for backward -> scales x n_layers. + # The 16-byte/token/layer coefficient bundles fp32 activation + ~4x backward overhead. + # Inference only holds ~1-2 layers' outputs in flight at once. + bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 if is_training else seq_len * hidden * 8 if mixed_precision: bytes_per_sample //= 2 return bytes_per_sample / (1024**3) @@ -265,11 +269,11 @@ def _embedder_dim(meta: ModelMeta | None) -> int: if meta is None: return 768 params = meta.params_millions - if params >= 300: + if params >= _PARAMS_LARGE: return 1024 - if params >= 100: + if params >= _PARAMS_BASE: return 768 - if params >= 50: + if params >= _PARAMS_SMALL: return 512 return 384 @@ -313,7 +317,7 @@ def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, data_bytes = 4.0 * stats.n_samples * n_features histograms_bytes = 4.0 * n_features * _CATBOOST_DEFAULT_BINS trees_bytes = iterations * (2**depth) * _CATBOOST_BYTES_PER_TREE_NODE - return (data_bytes + histograms_bytes + trees_bytes) / (1024**3) + return float((data_bytes + histograms_bytes + trees_bytes) / (1024**3)) def _time_for_catboost( @@ -360,7 +364,7 @@ def _classify_severity(estimate: float, budget: float) -> Severity: return Severity.AMPLE -def _resource_phase( # noqa: PLR0912 - kept linear for clarity +def _resource_phase( # noqa: PLR0912, C901, PLR0915 - kept linear for clarity config: dict[str, Any], stats: DatasetStats, hardware: HardwareProfile, @@ -394,7 +398,7 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity transformer_entries.append((node_idx, node_type, entry)) # Track the heaviest module per node so dump_modules accounting is bounded by - # "one selected variant per node × n_trials", not "sum of every candidate". + # "one selected variant per node x n_trials", not "sum of every candidate". node_max_weights: dict[int, float] = {} for node_idx, node_type, entry in transformer_entries: @@ -490,7 +494,9 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" # CatBoost's MultiClass loss grows per-class trees only above binary; # binary uses Logloss with one tree per iteration. - cb_class_mult = max(1, stats.n_classes) if stats.n_classes > 2 or stats.multilabel else 1 + cb_class_mult = ( + max(1, stats.n_classes) if stats.n_classes > _MULTICLASS_THRESHOLD or stats.multilabel else 1 + ) ram_total = _ram_for_catboost( stats=stats, n_features=embedder_dim, @@ -569,7 +575,6 @@ def _resource_phase( # noqa: PLR0912 - kept linear for clarity msg += f" vs available {hardware.vram_gb:.1f} GB" report.add("resource", vram_sev, msg, metric="vram") - ram_sev = _classify_severity(effective_ram, hardware.ram_gb) report.add( "resource", @@ -635,13 +640,14 @@ def _data_phase( max_len = _max_int(max_len_value, 512) if p95 > max_len: severity = Severity.OVER if p95 > max_len * 1.5 else Severity.TIGHT + module_name = entry.get("module_name", "?") report.add( "data", severity, - f"Train tokens p95~{p95} exceeds {entry.get('module_name', '?')}.max_length={max_len}; expect silent truncation.", + f"Train tokens p95~{p95} exceeds {module_name}.max_length={max_len}; expect silent truncation.", ) - # rare class × linear-CV (LogisticRegressionCV cv=3 needs ≥3 samples/class; + # rare class x linear-CV (LogisticRegressionCV cv=3 needs >=3 samples/class; # multilabel path uses one-vs-rest without CV so the failure can't occur there) has_linear = any(e.get("module_name") == "linear" for _, e in _walk_modules(config.get("search_space") or [])) if has_linear and stats.rare_classes and not stats.multilabel: @@ -651,7 +657,7 @@ def _data_phase( (f"LogisticRegressionCV (cv=3) will fail: classes {stats.rare_classes[:5]} have <3 samples."), ) - # partial descriptions × description scorer + # partial descriptions x description scorer description_modules = {"description_bi", "description_cross", "description_llm"} has_description = any( e.get("module_name") in description_modules for _, e in _walk_modules(config.get("search_space") or []) diff --git a/src/autointent/_advisor/_hardware.py b/src/autointent/_advisor/_hardware.py index 9c0cae049..e959b6ebf 100644 --- a/src/autointent/_advisor/_hardware.py +++ b/src/autointent/_advisor/_hardware.py @@ -12,6 +12,7 @@ import platform import shutil from dataclasses import dataclass, field +from pathlib import Path from typing import Literal import psutil @@ -24,6 +25,9 @@ # matches macOS PYTORCH_MPS_HIGH_WATERMARK_RATIO default MPS_DEFAULT_BUDGET_RATIO = 0.7 +_HIGH_GPU_VRAM_GB = 24 +_MID_GPU_VRAM_GB = 12 + @dataclass class HardwareProfile: @@ -41,20 +45,20 @@ def device_class(self) -> str: return "cpu" if self.accelerator == "mps": return "apple-silicon" - if self.vram_gb >= 24: + if self.vram_gb >= _HIGH_GPU_VRAM_GB: return "high-gpu" - if self.vram_gb >= 12: + if self.vram_gb >= _MID_GPU_VRAM_GB: return "mid-gpu" return "low-gpu" def _detect_ram_gb() -> float: - return psutil.virtual_memory().total / (1024**3) + return float(psutil.virtual_memory().total) / (1024**3) def _detect_free_disk_gb(path: str | None = None) -> float: - cache = path or os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface") - probe_path = cache if os.path.exists(cache) else os.path.expanduser("~") + cache = Path(path or os.environ.get("HF_HOME") or Path("~/.cache/huggingface").expanduser()) + probe_path = cache if cache.exists() else Path("~").expanduser() try: usage = shutil.disk_usage(probe_path) return usage.free / (1024**3) diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index 1c559ee2f..9b351952a 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -8,10 +8,10 @@ from __future__ import annotations import logging -import os import re from dataclasses import dataclass from functools import lru_cache +from pathlib import Path from typing import Any from huggingface_hub import HfApi, scan_cache_dir, try_to_load_from_cache @@ -54,7 +54,7 @@ def weights_gb(self) -> float: @lru_cache(maxsize=1) -def hub_reachable(timeout_s: float = 2.0) -> bool: +def hub_reachable() -> bool: """Single up-front probe. Memoized per process.""" try: HfApi().list_models(limit=1) @@ -157,7 +157,7 @@ def resolve_model(model_name: str) -> ModelMeta: Always returns a value — never raises — so the advisor can keep going on offline machines or for unknown checkpoints. """ - if model_name.startswith("local:") or os.path.isabs(model_name): + if model_name.startswith("local:") or Path(model_name).is_absolute(): return ModelMeta( name=model_name, params_millions=_heuristic_params_millions(model_name), diff --git a/src/autointent/_advisor/_render.py b/src/autointent/_advisor/_render.py index a3778d307..82771ef9f 100644 --- a/src/autointent/_advisor/_render.py +++ b/src/autointent/_advisor/_render.py @@ -8,7 +8,7 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from ._report import PreflightReport @@ -18,7 +18,7 @@ _PHASE_LABEL = {"resource": "Resource", "data": "Data", "config": "Config"} -def _batch_hint(driver: dict) -> str: +def _batch_hint(driver: dict[str, Any]) -> str: """Per-driver batch annotation: '64 → 32', '64', '64 (no fit)', or ''.""" bs = driver.get("batch_size") if bs is None: @@ -37,12 +37,11 @@ def _batch_hint(driver: dict) -> str: _DRIVERS_HEADERS = ("Node", "Model", "Mode", "VRAM", "Time", "Batch", "Source") -def _render_drivers_table(drivers: list[dict]) -> list[str]: +def _render_drivers_table(drivers: list[dict[str, Any]]) -> list[str]: """Format the Drivers of cost section as an aligned table.""" visible = drivers[:_DRIVERS_LIMIT] - rows: list[tuple[str, ...]] = [] - for d in visible: - rows.append(( + rows: list[tuple[str, ...]] = [ + ( f"{d['node_type']}.{d['module']}", str(d["model"]), str(d["mode"]), @@ -50,7 +49,9 @@ def _render_drivers_table(drivers: list[dict]) -> list[str]: f"{d['time_hours']:.2f} h", _batch_hint(d), f"[{d['confidence']}]", - )) + ) + for d in visible + ] widths = [len(h) for h in _DRIVERS_HEADERS] for row in rows: @@ -113,8 +114,7 @@ def render_text(report: PreflightReport) -> str: if report.notes: lines.append("Notes:") - for note in report.notes: - lines.append(f" • {note}") + lines.extend(f" • {note}" for note in report.notes) lines.append("") summary = f"Verdict: {'feasible' if report.is_feasible else 'INFEASIBLE'} " diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 5ac66af2f..9b4881611 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -128,7 +128,7 @@ def test_inference_activations_are_smaller_than_training(self, meta: ModelMeta) inf_acts = inf_total - inf_weights assert inf_acts > 0 assert train_acts > inf_acts - # 12-layer model: training activations should be at least ~5× inference. + # 12-layer model: training activations should be at least ~5x inference. assert train_acts / inf_acts > 5 def test_amp_does_not_reduce_weight_side_vram(self, meta: ModelMeta) -> None: @@ -503,7 +503,8 @@ def test_max_batch_can_be_larger_than_current(self) -> None: _profile(vram_gb=64.0), ) d = next(d for d in report.resource.drivers if d["module"] == "bert") - assert d["max_batch_size"] is not None and d["max_batch_size"] > 32 + assert d["max_batch_size"] is not None + assert d["max_batch_size"] > 32 def test_multiple_drivers_carry_independent_max_batch(self) -> None: cfg = { diff --git a/tests/advisor/test_report.py b/tests/advisor/test_report.py index 28adbfc34..dbfc7adf6 100644 --- a/tests/advisor/test_report.py +++ b/tests/advisor/test_report.py @@ -2,6 +2,8 @@ from __future__ import annotations +import dataclasses + import pytest from autointent._advisor._report import ( @@ -81,5 +83,5 @@ def test_hardware_and_dataset_pass_through(self) -> None: def test_finding_is_frozen(self) -> None: f = Finding(phase="resource", severity=Severity.AMPLE, message="ok") - with pytest.raises(Exception): # noqa: PT011 - dataclass.FrozenInstanceError varies + with pytest.raises(dataclasses.FrozenInstanceError): f.message = "changed" # type: ignore[misc] From 334783c77f02b084a47d2e504782b55b296565ac Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 18:01:40 +0300 Subject: [PATCH 08/43] try to fix typing --- tests/advisor/test_estimates_and_cli.py | 6 +++--- tests/advisor/test_report.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 3092dce9e..2f16555ae 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -56,21 +56,21 @@ def test_every_preset_inspects_without_raising(preset: str) -> None: def test_heavy_preset_is_infeasible_on_2gb_budget() -> None: - cfg = load_preset("transformers-heavy") # type: ignore[arg-type] + cfg = load_preset("transformers-heavy") stats = DatasetStats.placeholder(n_samples=5000, n_classes=20, avg_tokens=40) report = run_preflight(cfg, stats, _profile(vram_gb=2.0), preset_name="transformers-heavy") assert not report.is_feasible, "deberta-v3-large should not fit in 2 GB" def test_light_preset_is_feasible_on_8gb_budget() -> None: - cfg = load_preset("transformers-light") # type: ignore[arg-type] + cfg = load_preset("transformers-light") stats = DatasetStats.placeholder(n_samples=1000, n_classes=10, avg_tokens=24) report = run_preflight(cfg, stats, _profile(vram_gb=8.0), preset_name="transformers-light") assert report.is_feasible def test_n_jobs_doubles_vram_findings() -> None: - cfg = load_preset("transformers-light") # type: ignore[arg-type] + cfg = load_preset("transformers-light") cfg = {**cfg, "hpo_config": {**(cfg.get("hpo_config") or {}), "n_jobs": 4}} stats = DatasetStats.placeholder() report = run_preflight(cfg, stats, _profile(vram_gb=4.0)) diff --git a/tests/advisor/test_report.py b/tests/advisor/test_report.py index dbfc7adf6..acb2b5bf8 100644 --- a/tests/advisor/test_report.py +++ b/tests/advisor/test_report.py @@ -25,7 +25,7 @@ def test_red_beats_yellow_beats_green(self) -> None: r.add("data", Severity.TIGHT, "warn") assert r.headroom == Severity.TIGHT r.add("config", Severity.OVER, "fail") - assert r.headroom == Severity.OVER + assert r.headroom == Severity.OVER # type: ignore[comparison-overlap] def test_is_feasible_flips_on_any_red(self) -> None: r = PreflightReport() From 4e4da91edc5712017ad134b22db8b8260a1b3ca9 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 18:10:21 +0300 Subject: [PATCH 09/43] roll back config changes --- src/autointent/_presets/transformers-heavy.yaml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/autointent/_presets/transformers-heavy.yaml b/src/autointent/_presets/transformers-heavy.yaml index cd15d791e..2576fbc82 100644 --- a/src/autointent/_presets/transformers-heavy.yaml +++ b/src/autointent/_presets/transformers-heavy.yaml @@ -5,19 +5,12 @@ search_space: - module_name: bert classification_model_config: - model_name: microsoft/deberta-v3-large - - model_name: intfloat/multilingual-e5-large-instruct - - model_name: microsoft/harrier-oss-v1-27b num_train_epochs: [30] batch_size: [32, 64] learning_rate: low: 1.0e-5 high: 1.0e-4 log: True - - module_name: description_bi - embedder_config: - - model_name: microsoft/deberta-v3-large - - model_name: intfloat/multilingual-e5-large-instruct - - model_name: microsoft/harrier-oss-v1-27b - node_type: decision target_metric: decision_accuracy search_space: From 8bd0b018ede3c89175cca41e6d71f0676d3f58d9 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 18:26:19 +0300 Subject: [PATCH 10/43] move cli logic --- src/autointent/_advisor/__init__.py | 9 +- src/autointent/_advisor/_cli.py | 231 +++----------------------- src/autointent/_advisor/_report.py | 18 ++ src/autointent/_advisor/_workflows.py | 231 ++++++++++++++++++++++++++ src/autointent/custom_types/_types.py | 17 +- 5 files changed, 294 insertions(+), 212 deletions(-) create mode 100644 src/autointent/_advisor/_workflows.py diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index 3ff898816..28422c78d 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -9,15 +9,22 @@ from ._estimates import run_preflight from ._hardware import HardwareProfile, detect_hardware -from ._report import DatasetStats, Finding, PreflightReport, ResourceEstimate, Severity +from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity +from ._workflows import BUNDLED_PRESETS, inspect, load_config, recommend, stats_from_dataset __all__ = [ + "BUNDLED_PRESETS", "DatasetStats", "Finding", "HardwareProfile", "PreflightReport", + "RecommendationResult", "ResourceEstimate", "Severity", "detect_hardware", + "inspect", + "load_config", + "recommend", "run_preflight", + "stats_from_dataset", ] diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py index 8c8b7b9d2..315d5c443 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/_advisor/_cli.py @@ -9,190 +9,40 @@ csv/json/jsonl/parquet path loaded via ``datasets.load_dataset``) or ``--n-samples / --n-classes / --avg-tokens`` placeholders so the script is useful before the user has built a dataset. + +The CLI is a thin wrapper around :func:`autointent._advisor.inspect` and +:func:`autointent._advisor.recommend`; callers that don't need argparse can +import those helpers directly. """ from __future__ import annotations import argparse +import json import logging import sys -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import yaml -from datasets import ClassLabel, Sequence, load_dataset - -from autointent.utils import load_preset -from ._estimates import run_preflight -from ._hardware import detect_hardware from ._render import render_json, render_recommendation, render_text -from ._report import DatasetStats, Severity +from ._report import DatasetStats +from ._workflows import BUNDLED_PRESETS, inspect, recommend, stats_from_dataset -if TYPE_CHECKING: - from ._report import PreflightReport +__all__ = ["BUNDLED_PRESETS", "build_parser", "cmd_inspect", "cmd_recommend", "main"] logger = logging.getLogger("autointent.advisor") -_SAMPLE_LIMIT = 1000 -_P95_PERCENTILE = 0.95 - -BUNDLED_PRESETS = [ - "transformers-heavy", - "transformers-light", - "transformers-no-hpo", - "nn-heavy", - "nn-medium", - "classic-heavy", - "classic-medium", - "classic-light", - "zero-shot-encoders", - "zero-shot-llm", -] - -# rough quality tiering used by `recommend` -_QUALITY_TIER = { - "transformers-heavy": 5, - "nn-heavy": 4, - "transformers-light": 4, - "nn-medium": 3, - "classic-heavy": 3, - "transformers-no-hpo": 3, - "classic-medium": 2, - "classic-light": 1, - "zero-shot-encoders": 2, - "zero-shot-llm": 4, -} - - -def _load_config(target: str) -> tuple[dict[str, Any], str]: - """Return (config_dict, friendly_name) for either a preset or a path.""" - path = Path(target) - if path.is_file(): - with path.open(encoding="utf-8") as f: - return yaml.safe_load(f), path.stem - # treat as a bundled preset name - return load_preset(target), target # type: ignore[arg-type] - def _stats_from_args(args: argparse.Namespace) -> DatasetStats: + multilabel = args.task == "multilabel" if args.dataset: - return _stats_from_dataset(args.dataset, multilabel=args.task == "multilabel") + return stats_from_dataset(args.dataset, multilabel=multilabel) return DatasetStats.placeholder( n_samples=args.n_samples, n_classes=args.n_classes, avg_tokens=args.avg_tokens, - multilabel=args.task == "multilabel", - ) - - -_UTTERANCE_COLS = ("utterance", "text", "sentence", "query", "input") -_LABEL_COLS = ("label", "labels", "intent", "target") -# Map file extension → datasets builder name. Anything else is treated as a Hub -# repo id or a directory and passed to load_dataset directly. -_FILE_BUILDERS = {".csv": "csv", ".tsv": "csv", ".json": "json", ".jsonl": "json", ".parquet": "parquet"} - - -def _stats_from_dataset(path: str, *, multilabel: bool) -> DatasetStats: - """Best-effort: load via HF ``datasets.load_dataset``. - - Accepts a Hub repo id ('DeepPavlov/clinc150') or a local file path - (.csv / .json / .jsonl / .parquet) / dataset directory. Falls back to a - placeholder on any loader error so the advisor stays best-effort. - """ - builder = _FILE_BUILDERS.get(Path(path).suffix.lower()) - try: - ds = load_dataset(builder, data_files=path) if builder else load_dataset(path) - except (OSError, ValueError, FileNotFoundError) as e: - logger.warning("Failed to load dataset %s: %s", path, e) - return DatasetStats.placeholder(multilabel=multilabel) - - train = ds["train"] if "train" in ds else next(iter(ds.values()), None) - if train is None: - return DatasetStats.placeholder(multilabel=multilabel) - - cols = train.column_names - utt_col = next((c for c in _UTTERANCE_COLS if c in cols), cols[0] if cols else None) - label_col = next((c for c in _LABEL_COLS if c in cols), None) - - detected_multilabel, n_classes = _label_shape(train, label_col, fallback_multilabel=multilabel) - - sample = train[:_SAMPLE_LIMIT] if len(train) > _SAMPLE_LIMIT else train[:] - lengths = [len(str(s).split()) for s in (sample.get(utt_col, []) if utt_col else [])] - avg_tokens = int(sum(lengths) / max(1, len(lengths))) if lengths else 32 - if lengths: - sorted_lengths = sorted(lengths) - idx = max(0, min(len(sorted_lengths) - 1, round((len(sorted_lengths) - 1) * _P95_PERCENTILE))) - p95 = sorted_lengths[idx] - else: - p95 = avg_tokens * 2 - - return DatasetStats( - n_samples=len(train), - n_classes=n_classes, - avg_tokens=avg_tokens, - p95_tokens=p95, - multilabel=detected_multilabel, - has_descriptions=None, - rare_classes=_rare_classes(train, label_col, detected_multilabel, n_classes) if label_col else [], - source=f"dataset:{path}", + multilabel=multilabel, ) -def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool) -> tuple[bool, int]: # noqa: ANN401 - """Derive (multilabel, n_classes) from the HF feature schema, with a value-based fallback.""" - if label_col is None: - return fallback_multilabel, 0 - feature = train.features.get(label_col) - if isinstance(feature, Sequence): - inner = feature.feature - if isinstance(inner, ClassLabel): - return True, inner.num_classes - # Sequence of plain ints — n_classes = max label index + 1. - max_idx = max((max(row) for row in train[label_col] if row), default=-1) - return True, max_idx + 1 - if isinstance(feature, ClassLabel): - return False, feature.num_classes - # Plain int/string column. Detect multilabel from the first non-empty row, then count uniques. - is_multi = len(train) > 0 and isinstance(train[0][label_col], (list, tuple)) - if is_multi: - max_idx = max((max(row) for row in train[label_col] if row), default=-1) - return True, max_idx + 1 - return False, len({label for label in train[label_col] if label is not None}) - - -def _rare_classes( - train: Any, # noqa: ANN401 - label_col: str, - multilabel: bool, - n_classes: int, - min_count: int = 3, -) -> list[str]: - """Return labels with fewer than ``min_count`` samples in the train split. - - Used to surface the LogisticRegressionCV(cv=3) failure case before fit. - Returns an empty list on any error so the advisor stays best-effort. - """ - try: - labels = train[label_col] - except (KeyError, AttributeError, TypeError): - return [] - counts: dict[str, int] = {} - if multilabel: - for row in labels: - if not row: - continue - for i, v in enumerate(row): - if v: - counts[str(i)] = counts.get(str(i), 0) + 1 - for i in range(n_classes): - counts.setdefault(str(i), 0) - else: - for label in labels: - counts[str(label)] = counts.get(str(label), 0) + 1 - return sorted(name for name, c in counts.items() if c < min_count) - - def _add_common_dataset_args(p: argparse.ArgumentParser) -> None: p.add_argument("--dataset", help="Path or hub id of a dataset; overrides placeholders.") p.add_argument("--n-samples", type=int, default=1_000, help="Placeholder training set size.") @@ -207,63 +57,36 @@ def _add_common_dataset_args(p: argparse.ArgumentParser) -> None: def cmd_inspect(args: argparse.Namespace) -> int: - config, name = _load_config(args.target) - hardware = detect_hardware( - vram_budget_gb=args.budget_vram_gb, + report = inspect( + args.target, + stats=_stats_from_args(args), + budget_vram_gb=args.budget_vram_gb, ) - stats = _stats_from_args(args) - report = run_preflight(config, stats, hardware, preset_name=name) if args.json: sys.stdout.write(render_json(report)) - sys.stdout.write("\n") else: sys.stdout.write(render_text(report)) - sys.stdout.write("\n") + sys.stdout.write("\n") return 0 if report.is_feasible else 1 def cmd_recommend(args: argparse.Namespace) -> int: - hardware = detect_hardware(vram_budget_gb=args.budget_vram_gb) - stats = _stats_from_args(args) - - results: list[tuple[str, PreflightReport]] = [] - - for preset in BUNDLED_PRESETS: - try: - cfg = load_preset(preset) # type: ignore[arg-type] - except (OSError, ValueError, KeyError) as e: - logger.debug("Skipping preset %s: %s", preset, e) - continue - report = run_preflight(cfg, stats, hardware, preset_name=preset) - if args.budget_time_h is not None and report.resource.time_hours > args.budget_time_h: - report.add( - "resource", - Severity.OVER, - f"Estimated time {report.resource.time_hours:.1f} h exceeds budget {args.budget_time_h} h.", - ) - results.append((preset, report)) - - feasible = [(name, r) for name, r in results if r.is_feasible] - feasible.sort(key=lambda pair: (-_QUALITY_TIER.get(pair[0], 0), pair[1].resource.time_hours, pair[0])) - chosen = feasible[0][0] if feasible else None - + result = recommend( + stats=_stats_from_args(args), + budget_vram_gb=args.budget_vram_gb, + budget_time_h=args.budget_time_h, + ) if args.json: - import json - - out = { - "chosen": chosen, - "results": [{"preset": name, "report": r.to_dict()} for name, r in results], - } - sys.stdout.write(json.dumps(out, indent=2, default=str)) + sys.stdout.write(json.dumps(result.to_dict(), indent=2, default=str)) sys.stdout.write("\n") else: - sys.stdout.write(render_recommendation(results, chosen)) + sys.stdout.write(render_recommendation(result.results, result.chosen)) sys.stdout.write("\n") - if chosen: + if result.chosen: sys.stdout.write("\n") - sys.stdout.write(render_text(dict(results)[chosen])) + sys.stdout.write(render_text(dict(result.results)[result.chosen])) sys.stdout.write("\n") - return 0 if chosen else 1 + return 0 if result.chosen else 1 def build_parser() -> argparse.ArgumentParser: @@ -309,4 +132,4 @@ def main(argv: list[str] | None = None) -> int: if __name__ == "__main__": - raise SystemExit(main()) + main() diff --git a/src/autointent/_advisor/_report.py b/src/autointent/_advisor/_report.py index 9b4a319c8..c9fd920f4 100644 --- a/src/autointent/_advisor/_report.py +++ b/src/autointent/_advisor/_report.py @@ -110,3 +110,21 @@ def to_dict(self) -> dict[str, Any]: d["headroom"] = self.headroom.value d["is_feasible"] = self.is_feasible return d + + +@dataclass +class RecommendationResult: + """Output of the recommend workflow: ranked per-preset reports plus the pick. + + ``chosen`` is the best feasible preset name, or ``None`` if none fit. + ``results`` is the full per-preset report list in evaluation order. + """ + + chosen: str | None + results: list[tuple[str, PreflightReport]] + + def to_dict(self) -> dict[str, Any]: + return { + "chosen": self.chosen, + "results": [{"preset": name, "report": r.to_dict()} for name, r in self.results], + } diff --git a/src/autointent/_advisor/_workflows.py b/src/autointent/_advisor/_workflows.py new file mode 100644 index 000000000..0bd7ee5bf --- /dev/null +++ b/src/autointent/_advisor/_workflows.py @@ -0,0 +1,231 @@ +"""High-level advisor workflows: ``inspect`` and ``recommend``. + +Each workflow orchestrates the lower-level pieces (``load_config``, +``detect_hardware``, ``stats_from_dataset``, ``run_preflight``) into a single +typed call. They expose the same logic the CLI uses but accept Python +arguments instead of an ``argparse.Namespace`` — useful from notebooks, +integration tests, or any caller that wants a ``PreflightReport`` / +``RecommendationResult`` directly. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import TYPE_CHECKING, Any, get_args + +import yaml +from datasets import ClassLabel, Sequence, load_dataset + +from autointent.custom_types import SearchSpacePreset +from autointent.utils import load_preset + +from ._estimates import run_preflight +from ._hardware import detect_hardware +from ._report import DatasetStats, RecommendationResult, Severity + +if TYPE_CHECKING: + from collections.abc import Iterable + + from ._report import PreflightReport + + +logger = logging.getLogger("autointent.advisor") + +_SAMPLE_LIMIT = 1000 +_P95_PERCENTILE = 0.95 +BUNDLED_PRESETS: tuple[str, ...] = get_args(SearchSpacePreset) + + +def load_config(target: str) -> tuple[dict[str, Any], str]: + """Return ``(config_dict, friendly_name)`` for either a preset name or a YAML path.""" + path = Path(target) + if path.is_file(): + with path.open(encoding="utf-8") as f: + return yaml.safe_load(f), path.stem + return load_preset(target), target # type: ignore[arg-type] + + +def stats_from_dataset(path: str, *, multilabel: bool = False) -> DatasetStats: + """Best-effort: load a dataset via HF ``datasets.load_dataset`` and derive advisor stats. + + Accepts a Hub repo id (``DeepPavlov/clinc150``) or a local file path + (``.csv`` / ``.json`` / ``.jsonl`` / ``.parquet``) / dataset directory. Falls + back to a placeholder on any loader error so callers stay best-effort. + """ + # Anything not in this map (no suffix, unknown suffix) is treated as a Hub + # repo id or a dataset directory and passed to load_dataset directly. + file_builders = {".csv": "csv", ".tsv": "csv", ".json": "json", ".jsonl": "json", ".parquet": "parquet"} + builder = file_builders.get(Path(path).suffix.lower()) + try: + ds = load_dataset(builder, data_files=path) if builder else load_dataset(path) + except (OSError, ValueError, FileNotFoundError) as e: + logger.warning("Failed to load dataset %s: %s", path, e) + return DatasetStats.placeholder(multilabel=multilabel) + + train = ds["train"] if "train" in ds else next(iter(ds.values()), None) + if train is None: + return DatasetStats.placeholder(multilabel=multilabel) + + cols = train.column_names + utt_col = next( + (c for c in ("utterance", "text", "sentence", "query", "input") if c in cols), cols[0] if cols else None + ) + label_col = next((c for c in ("label", "labels", "intent", "target") if c in cols), None) + + detected_multilabel, n_classes = _label_shape(train, label_col, fallback_multilabel=multilabel) + + sample = train[:_SAMPLE_LIMIT] if len(train) > _SAMPLE_LIMIT else train[:] + lengths = [len(str(s).split()) for s in (sample.get(utt_col, []) if utt_col else [])] + avg_tokens = int(sum(lengths) / max(1, len(lengths))) if lengths else 32 + if lengths: + sorted_lengths = sorted(lengths) + idx = max(0, min(len(sorted_lengths) - 1, round((len(sorted_lengths) - 1) * _P95_PERCENTILE))) + p95 = sorted_lengths[idx] + else: + p95 = avg_tokens * 2 + + return DatasetStats( + n_samples=len(train), + n_classes=n_classes, + avg_tokens=avg_tokens, + p95_tokens=p95, + multilabel=detected_multilabel, + has_descriptions=None, + rare_classes=_rare_classes(train, label_col, detected_multilabel, n_classes) if label_col else [], + source=f"dataset:{path}", + ) + + +def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool) -> tuple[bool, int]: # noqa: ANN401 + """Derive ``(multilabel, n_classes)`` from the HF feature schema with a value-based fallback.""" + if label_col is None: + return fallback_multilabel, 0 + feature = train.features.get(label_col) + if isinstance(feature, Sequence): + inner = feature.feature + if isinstance(inner, ClassLabel): + return True, inner.num_classes + # Sequence of plain ints — n_classes = max label index + 1. + max_idx = max((max(row) for row in train[label_col] if row), default=-1) + return True, max_idx + 1 + if isinstance(feature, ClassLabel): + return False, feature.num_classes + # Plain int/string column. Detect multilabel from the first non-empty row, then count uniques. + is_multi = len(train) > 0 and isinstance(train[0][label_col], (list, tuple)) + if is_multi: + max_idx = max((max(row) for row in train[label_col] if row), default=-1) + return True, max_idx + 1 + return False, len({label for label in train[label_col] if label is not None}) + + +def _rare_classes( + train: Any, # noqa: ANN401 + label_col: str, + multilabel: bool, + n_classes: int, + min_count: int = 3, +) -> list[str]: + """Return labels with fewer than ``min_count`` samples in the train split. + + Used to surface the LogisticRegressionCV(cv=3) failure case before fit. + Returns an empty list on any error so the advisor stays best-effort. + """ + try: + labels = train[label_col] + except (KeyError, AttributeError, TypeError): + return [] + counts: dict[str, int] = {} + if multilabel: + for row in labels: + if not row: + continue + for i, v in enumerate(row): + if v: + counts[str(i)] = counts.get(str(i), 0) + 1 + for i in range(n_classes): + counts.setdefault(str(i), 0) + else: + for label in labels: + counts[str(label)] = counts.get(str(label), 0) + 1 + return sorted(name for name, c in counts.items() if c < min_count) + + +def inspect( + target: str, + *, + stats: DatasetStats | None = None, + budget_vram_gb: float | None = None, +) -> PreflightReport: + """Inspect a preset (or YAML config path) against the local hardware. + + Args: + target: Bundled preset name (e.g. ``'transformers-light'``) or a YAML + config path. The friendly name surfaced in the report is the file + stem for paths and the preset name otherwise. + stats: Dataset stats to score against. Defaults to a placeholder if + ``None``. + budget_vram_gb: Optional VRAM-budget override for the hardware probe. + + Returns: + ``PreflightReport`` covering resource / data / config phases. + """ + config, name = load_config(target) + hardware = detect_hardware(vram_budget_gb=budget_vram_gb) + return run_preflight(config, stats or DatasetStats.placeholder(), hardware, preset_name=name) + + +def recommend( + *, + stats: DatasetStats | None = None, + presets: Iterable[str] | None = None, + budget_vram_gb: float | None = None, + budget_time_h: float | None = None, +) -> RecommendationResult: + """Walk bundled presets and return the best feasible fit plus all per-preset reports. + + Args: + stats: Dataset stats to score against. Defaults to a placeholder if ``None``. + presets: Override of the preset list (defaults to ``BUNDLED_PRESETS``). + budget_vram_gb: Optional VRAM-budget override for the hardware probe. + budget_time_h: Optional wall-time ceiling in hours; presets exceeding it + get an extra ``Severity.OVER`` finding so they drop out of the + feasible ranking. + + Returns: + ``RecommendationResult`` with the chosen preset name and full results list. + + Note: + Among feasible presets we pick the one with the largest estimated + ``time_hours`` (ties broken alphabetically). Higher-quality presets cost + more wall-time, so the slowest feasible preset is also the heaviest + preset that still fits the hardware — i.e. "use what you have". + """ + hardware = detect_hardware(vram_budget_gb=budget_vram_gb) + stats = stats or DatasetStats.placeholder() + preset_iter = list(presets) if presets is not None else BUNDLED_PRESETS + + results: list[tuple[str, PreflightReport]] = [] + for preset in preset_iter: + try: + cfg = load_preset(preset) # type: ignore[arg-type] + except (OSError, ValueError, KeyError) as e: + logger.debug("Skipping preset %s: %s", preset, e) + continue + report = run_preflight(cfg, stats, hardware, preset_name=preset) + if budget_time_h is not None and report.resource.time_hours > budget_time_h: + report.add( + "resource", + Severity.OVER, + f"Estimated time {report.resource.time_hours:.1f} h exceeds budget {budget_time_h} h.", + ) + results.append((preset, report)) + + # Rank by Literal position (lower index = higher quality); presets the user + # passed via the ``presets`` override but not in BUNDLED_PRESETS sort last. + quality_rank = {name: i for i, name in enumerate(BUNDLED_PRESETS)} + feasible = [(name, r) for name, r in results if r.is_feasible] + feasible.sort(key=lambda pair: (quality_rank.get(pair[0], len(BUNDLED_PRESETS)), pair[0])) + chosen = feasible[0][0] if feasible else None + + return RecommendationResult(chosen=chosen, results=results) diff --git a/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index cbfa82576..a54da368d 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -117,18 +117,21 @@ class Split: """ SearchSpacePreset = Literal[ - "classic-heavy", - "classic-light", - "classic-medium", - "nn-heavy", - "nn-medium", "transformers-heavy", "transformers-light", - "transformers-no-hpo", + "nn-heavy", "zero-shot-llm", + "nn-medium", + "classic-heavy", + "transformers-no-hpo", + "classic-medium", "zero-shot-encoders", + "classic-light", ] -"""Some presets that our library supports.""" +"""Bundled search-space presets, listed in descending quality order. + +The order is consumed by ``autointent._advisor.recommend`` to pick the +highest-quality feasible preset (lower index = higher quality).""" class Document(BaseModel): From b77d57586028af16a01b210161af702dcf34e789 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 18:57:07 +0300 Subject: [PATCH 11/43] simplify logic --- src/autointent/_advisor/_estimates.py | 477 ++++++++++++---------- src/autointent/_advisor/_hub.py | 47 ++- tests/advisor/test_estimates_internals.py | 10 +- 3 files changed, 320 insertions(+), 214 deletions(-) diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py index 93dfcedaf..274407c26 100644 --- a/src/autointent/_advisor/_estimates.py +++ b/src/autointent/_advisor/_estimates.py @@ -9,6 +9,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -26,9 +27,11 @@ from ._report import DatasetStats _MULTICLASS_THRESHOLD = 2 -_PARAMS_LARGE = 300 -_PARAMS_BASE = 100 -_PARAMS_SMALL = 50 + +# Fallback architecture shape (BERT-base) used only when the model's actual +# config.json couldn't be fetched from HF Hub — see _hub._shape_from_config. +_DEFAULT_HIDDEN = 768 +_DEFAULT_LAYERS = 12 logger = logging.getLogger(__name__) @@ -77,12 +80,14 @@ def _validated_config(config: dict[str, Any]) -> _AdvisorConfig: } # Maps each fine-tunable transformer module to its training-mode label. -# Modules not listed are treated as inference-only. +# Modules not listed (or listed as "inference") run the encoder forward-only. +# Note: dnnc keeps the cross-encoder frozen and trains an sklearn LogisticRegressionCV +# head on top of its features (see autointent._wrappers.ranker.Ranker._fit), so the +# encoder's VRAM profile matches inference rather than fine-tuning. _TRANSFORMER_TRAINING_MODE = { "bert": "full-finetune", "ptuning": "lora", "lora": "lora", - "dnnc": "reranker", } # Fallback max_length when the search-space entry doesn't pin it. Used both as @@ -148,17 +153,16 @@ def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dic def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: """Weight-side VRAM in GB — weights + grads + Adam optimizer state. Excludes activations. - Full fine-tune fp32: W + W + 2W (Adam m, v) = 4W. - Full fine-tune AMP: 0.5W (fp16 weights) + 0.5W (fp16 grads) + W (fp32 master) + 2W (fp32 Adam) = 4W. - AMP's savings live in activations, not the optimizer — the weight side is identical. + Modes: + * ``inference``: forward only — weights + ~30% intermediate-tensor overhead. + * ``lora``: frozen base + small trainable adapters + their grads/optimizer (~0.5 GB). + * ``full-finetune`` (default): weights + grads + Adam (m, v) = 4x weights. """ weights_gb = meta.weights_gb if mode == "inference": return weights_gb * 1.3 if mode == "lora": return weights_gb * 1.3 + 0.5 - if mode == "reranker": - return weights_gb * 1.5 return weights_gb * 4.0 @@ -200,20 +204,10 @@ def _floor_to_power_of_two(n: int) -> int: def _n_layers(meta: ModelMeta | None) -> int: - """Coarse layer-count guess from parameter count. - - MiniLM (33M) ~6, BERT-base (110M) ~12, BERT-large (350M) ~24. - """ - if meta is None: - return 12 - params = meta.params_millions - if params >= _PARAMS_LARGE: - return 24 - if params >= _PARAMS_BASE: - return 12 - if params >= _PARAMS_SMALL: - return 8 - return 6 + """Layer count from the model's ``config.json``; falls back to BERT-base when absent.""" + if meta is not None and meta.n_layers is not None: + return meta.n_layers + return _DEFAULT_LAYERS def _activations_gb_per_sample( @@ -262,20 +256,10 @@ def _max_fitting_batch_size( def _embedder_dim(meta: ModelMeta | None) -> int: - """Coarse hidden-size guess from parameter count. - - Concrete points: MiniLM (33M) ~384, BERT-base (110M) ~768, BERT-large (350M) ~1024. - """ - if meta is None: - return 768 - params = meta.params_millions - if params >= _PARAMS_LARGE: - return 1024 - if params >= _PARAMS_BASE: - return 768 - if params >= _PARAMS_SMALL: - return 512 - return 384 + """Hidden size from the model's ``config.json``; falls back to BERT-base when absent.""" + if meta is not None and meta.hidden_size is not None: + return meta.hidden_size + return _DEFAULT_HIDDEN def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: @@ -364,146 +348,139 @@ def _classify_severity(estimate: float, budget: float) -> Severity: return Severity.AMPLE -def _resource_phase( # noqa: PLR0912, C901, PLR0915 - kept linear for clarity - config: dict[str, Any], - stats: DatasetStats, - hardware: HardwareProfile, - report: PreflightReport, -) -> None: - cfg = _validated_config(config) - n_trials = max(1, cfg.hpo_config.n_trials) - n_jobs = max(1, cfg.hpo_config.n_jobs) - refit_after = cfg.refit_after - dump_modules = cfg.dump_modules +@dataclass +class _ModuleEstimate: + """Per-module cost contribution + the dict that gets rendered in the report.""" - if not hub_reachable(): - report.low_confidence = True - report.notes.append("HF Hub unreachable — all model sizes are name-pattern heuristics.") + driver: dict[str, Any] + vram_gb: float + ram_gb: float + time_hours: float + model_weights_gb: float = 0.0 - seen_models: dict[str, ModelMeta] = {} - estimate = ResourceEstimate(parallel_factor=n_jobs) - global_embedder = (cfg.embedder_config or {}).get("model_name") - if global_embedder: - seen_models[global_embedder] = resolve_model(global_embedder) +def _refit_factor(*, refit_after: bool, n_trials: int) -> float: + """Wall-time multiplier for ``refit_after=True`` (amortized 1/n_trials extra).""" + return 1 + 1.0 / max(1, n_trials) if refit_after else 1.0 - # First pass: walk transformer-bearing modules (collects seen_models for embedder_dim lookup). - transformer_entries: list[tuple[int, str, dict[str, Any]]] = [] - classic_entries: list[tuple[int, str, dict[str, Any]]] = [] - for node_idx, node_type, entry in _walk_modules_indexed(cfg.search_space): - module = entry.get("module_name", "?") - if module in {"linear", "catboost"}: - classic_entries.append((node_idx, node_type, entry)) - else: - transformer_entries.append((node_idx, node_type, entry)) - # Track the heaviest module per node so dump_modules accounting is bounded by - # "one selected variant per node x n_trials", not "sum of every candidate". - node_max_weights: dict[int, float] = {} +def _split_entries( + search_space: list[dict[str, Any]], +) -> tuple[list[tuple[int, str, dict[str, Any]]], list[tuple[int, str, dict[str, Any]]]]: + """Partition search-space entries into (transformer-bearing, classic).""" + transformer, classic = [], [] + for node_idx, node_type, entry in _walk_modules_indexed(search_space): + bucket = classic if entry.get("module_name") in {"linear", "catboost"} else transformer + bucket.append((node_idx, node_type, entry)) + return transformer, classic - for node_idx, node_type, entry in transformer_entries: - module = entry.get("module_name", "?") - model_names = _extract_model_names(entry) - if not model_names and global_embedder and module in {"knn", "mlknn"}: - model_names = [global_embedder] - for name in model_names: - meta = seen_models.setdefault(name, resolve_model(name)) +def _estimate_transformer_model( + *, + meta: ModelMeta, + entry: dict[str, Any], + node_type: str, + module: str, + name: str, + stats: DatasetStats, + hardware: HardwareProfile, + n_trials: int, + refit_after: bool, +) -> _ModuleEstimate: + """One row of cost for a transformer module + a specific model checkpoint.""" + mixed_precision = entry.get("dtype") in {"fp16", "bf16"} + mode = _TRANSFORMER_TRAINING_MODE.get(module, "inference") + batch_size = _max_int(entry.get("batch_size"), 32) + epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) + seq_len = _max_int(entry.get("max_length"), _DEFAULT_SEQ_LEN) + + vram = _vram_for_transformer(meta, mode, mixed_precision, batch_size=batch_size, seq_len=seq_len) + ram = _ram_for_module(meta, stats) + + driver_max_batch: int | None = None + if hardware.vram_gb > 0: + driver_max_batch = _max_fitting_batch_size( + weight_vram_gb=_weights_vram_for_transformer(meta, mode), + vram_budget_gb=hardware.vram_gb, + per_sample_gb=_activations_gb_per_sample( + meta, seq_len, mixed_precision=mixed_precision, is_training=mode != "inference" + ), + ) - mixed_precision = entry.get("dtype") in {"fp16", "bf16"} - mode = _TRANSFORMER_TRAINING_MODE.get(module, "inference") - - batch_size = _max_int(entry.get("batch_size"), 32) - epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) - seq_len = _max_int(entry.get("max_length"), _DEFAULT_SEQ_LEN) - - vram = _vram_for_transformer(meta, mode, mixed_precision, batch_size=batch_size, seq_len=seq_len) - ram = _ram_for_module(meta, stats) - - driver_max_batch: int | None = None - if hardware.vram_gb > 0: - weights_vram = _weights_vram_for_transformer(meta, mode) - per_sample_gb = _activations_gb_per_sample( - meta, seq_len, mixed_precision=mixed_precision, is_training=mode != "inference" - ) - driver_max_batch = _max_fitting_batch_size( - weight_vram_gb=weights_vram, - vram_budget_gb=hardware.vram_gb, - per_sample_gb=per_sample_gb, - ) - - time_h = _time_for_transformer( - meta=meta, - n_trials=n_trials, - epochs=epochs, - batch_size=batch_size, - n_samples=stats.n_samples, - device_class=hardware.device_class, - ) - if refit_after and mode != "inference": - time_h *= 1 + 1.0 / max(1, n_trials) - - estimate.vram_gb = max(estimate.vram_gb, vram) - estimate.ram_gb = max(estimate.ram_gb, ram) - estimate.time_hours += time_h - node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), meta.weights_gb) - estimate.drivers.append( - { - "node_type": node_type, - "module": module, - "model": name, - "mode": mode, - "vram_gb": round(vram, 2), - "ram_gb": round(ram, 2), - "time_hours": round(time_h, 2), - "batch_size": batch_size, - "max_batch_size": driver_max_batch, - "confidence": meta.confidence, - } - ) + time_h = _time_for_transformer( + meta=meta, + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + device_class=hardware.device_class, + ) + if mode != "inference": + time_h *= _refit_factor(refit_after=refit_after, n_trials=n_trials) + + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": name, + "mode": mode, + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "batch_size": batch_size, + "max_batch_size": driver_max_batch, + "confidence": meta.confidence, + }, + vram_gb=vram, + ram_gb=ram, + time_hours=time_h, + model_weights_gb=meta.weights_gb, + ) - # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. - embedder_meta = _largest_embedder(seen_models) - embedder_dim = _embedder_dim(embedder_meta) - # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes; - # the multiclass path additionally pays the LogisticRegressionCV inner-fit multiplier. - class_multiplier_classic = max(1, stats.n_classes) - confidence = embedder_meta.confidence if embedder_meta else "heuristic" - embedder_label = embedder_meta.name if embedder_meta else "(no embedder)" - for _node_idx, node_type, entry in classic_entries: - module = entry.get("module_name", "?") - if module == "linear": - max_iter = _max_int(entry.get("max_iter"), 100) - cv_multiplier = 1 if stats.multilabel else _LOGREG_CV_MULTIPLIER - ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) - time_h = _time_for_linear( + +def _estimate_classic_entry( + *, + entry: dict[str, Any], + node_type: str, + embedder_meta: ModelMeta | None, + embedder_dim: int, + stats: DatasetStats, + hardware: HardwareProfile, + n_trials: int, + refit_after: bool, +) -> _ModuleEstimate | None: + """Cost row for a linear or catboost scorer (returns ``None`` for any other module).""" + module = entry.get("module_name", "?") + refit = _refit_factor(refit_after=refit_after, n_trials=n_trials) + # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes. + class_multiplier = max(1, stats.n_classes) + + if module == "linear": + cv_multiplier = 1 if stats.multilabel else _LOGREG_CV_MULTIPLIER + ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) + time_h = ( + _time_for_linear( n_trials=n_trials, n_samples=stats.n_samples, embedder_dim=embedder_dim, - max_iter=max_iter, + max_iter=_max_int(entry.get("max_iter"), 100), cv_multiplier=cv_multiplier, - class_multiplier=class_multiplier_classic, - ) - if refit_after: - time_h *= 1 + 1.0 / max(1, n_trials) - vram = 0.0 - mode = "linear-cv" if cv_multiplier > 1 else "linear" - elif module == "catboost": - iterations = _max_int(entry.get("iterations"), 1000) - depth = _max_int(entry.get("depth"), 6) - on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" - # CatBoost's MultiClass loss grows per-class trees only above binary; - # binary uses Logloss with one tree per iteration. - cb_class_mult = ( - max(1, stats.n_classes) if stats.n_classes > _MULTICLASS_THRESHOLD or stats.multilabel else 1 - ) - ram_total = _ram_for_catboost( - stats=stats, - n_features=embedder_dim, - iterations=iterations, - depth=depth, + class_multiplier=class_multiplier, ) - time_h = _time_for_catboost( + * refit + ) + vram = 0.0 + mode = "linear-cv" if cv_multiplier > 1 else "linear" + elif module == "catboost": + on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" + # CatBoost MultiClass loss grows per-class trees only above binary; binary uses + # Logloss with one tree per iteration. + cb_class_mult = class_multiplier if stats.n_classes > _MULTICLASS_THRESHOLD or stats.multilabel else 1 + iterations = _max_int(entry.get("iterations"), 1000) + depth = _max_int(entry.get("depth"), 6) + ram_total = _ram_for_catboost(stats=stats, n_features=embedder_dim, iterations=iterations, depth=depth) + time_h = ( + _time_for_catboost( n_trials=n_trials, n_samples=stats.n_samples, n_features=embedder_dim, @@ -512,55 +489,66 @@ def _resource_phase( # noqa: PLR0912, C901, PLR0915 - kept linear for clarity class_multiplier=cb_class_mult, on_gpu=on_gpu, ) - if refit_after: - time_h *= 1 + 1.0 / max(1, n_trials) - vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) - mode = "catboost-gpu" if on_gpu else "catboost" - else: - continue - - estimate.vram_gb = max(estimate.vram_gb, vram) - estimate.ram_gb = max(estimate.ram_gb, ram) - estimate.time_hours += time_h - estimate.drivers.append( - { - "node_type": node_type, - "module": module, - "model": embedder_label, - "mode": mode, - "vram_gb": round(vram, 2), - "ram_gb": round(ram, 2), - "time_hours": round(time_h, 2), - "batch_size": None, - "max_batch_size": None, - "confidence": confidence, - } + * refit ) + vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) + mode = "catboost-gpu" if on_gpu else "catboost" + else: + return None + + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": embedder_meta.name if embedder_meta else "(no embedder)", + "mode": mode, + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "batch_size": None, + "max_batch_size": None, + "confidence": embedder_meta.confidence if embedder_meta else "heuristic", + }, + vram_gb=vram, + ram_gb=ram, + time_hours=time_h, + ) + +def _aggregate_disk( + estimate: ResourceEstimate, + seen_models: dict[str, ModelMeta], + node_max_weights: dict[int, float], + *, + dump_modules: bool, + n_trials: int, +) -> None: + """Fold per-model download/cached sizes into ``estimate`` and apply dump-modules accounting.""" for meta in seen_models.values(): if meta.cached_locally: estimate.disk_cached_gb += meta.disk_gb else: estimate.disk_download_gb += meta.disk_gb - if dump_modules: # Each trial selects one variant per node, so per-trial dumped weights # are bounded by the heaviest module in each node, summed across nodes. - per_trial_dump_gb = sum(node_max_weights.values()) - estimate.disk_dump_gb = per_trial_dump_gb * n_trials + estimate.disk_dump_gb = sum(node_max_weights.values()) * n_trials - if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: - effective_vram = estimate.vram_gb * n_jobs - else: - effective_vram = estimate.vram_gb + +def _emit_resource_findings( + report: PreflightReport, + estimate: ResourceEstimate, + hardware: HardwareProfile, + *, + n_jobs: int, +) -> None: + """Translate aggregated estimates into VRAM/RAM/disk/time findings on the report.""" + parallel_gpu = n_jobs > 1 and hardware.accelerator in {"cuda", "mps"} + effective_vram = estimate.vram_gb * n_jobs if parallel_gpu else estimate.vram_gb # MPS shares one unified pool: parallel workers each allocate weights+activations # in RAM, so peak RAM also scales with n_jobs on Apple Silicon. effective_ram = estimate.ram_gb * n_jobs if n_jobs > 1 and hardware.accelerator == "mps" else estimate.ram_gb - report.resource = estimate - - # render findings - vram_sev = _classify_severity(effective_vram, hardware.vram_gb) if hardware.accelerator == "cpu" and effective_vram > 0: report.add( "resource", @@ -573,29 +561,108 @@ def _resource_phase( # noqa: PLR0912, C901, PLR0915 - kept linear for clarity if n_jobs > 1: msg += f" (= per-trial {estimate.vram_gb:.1f} GB × {n_jobs} parallel trials)" msg += f" vs available {hardware.vram_gb:.1f} GB" - report.add("resource", vram_sev, msg, metric="vram") + report.add("resource", _classify_severity(effective_vram, hardware.vram_gb), msg, metric="vram") - ram_sev = _classify_severity(effective_ram, hardware.ram_gb) report.add( "resource", - ram_sev, + _classify_severity(effective_ram, hardware.ram_gb), f"RAM ~{effective_ram:.1f} GB vs available {hardware.ram_gb:.1f} GB", metric="ram", ) disk_total = estimate.disk_download_gb + estimate.disk_dump_gb - disk_sev = _classify_severity(disk_total, hardware.free_disk_gb) disk_msg = f"Disk ~{estimate.disk_download_gb:.1f} GB to download" if estimate.disk_cached_gb > 0: disk_msg += f", {estimate.disk_cached_gb:.1f} GB already cached" if estimate.disk_dump_gb > 0: disk_msg += f", +{estimate.disk_dump_gb:.1f} GB during training (dump_modules=True)" disk_msg += f" vs {hardware.free_disk_gb:.0f} GB free" - report.add("resource", disk_sev, disk_msg, metric="disk") + report.add("resource", _classify_severity(disk_total, hardware.free_disk_gb), disk_msg, metric="disk") if estimate.time_hours > 0: - time_msg = f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)" - report.add("resource", Severity.AMPLE, time_msg, metric="time") + report.add( + "resource", + Severity.AMPLE, + f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)", + metric="time", + ) + + +def _resource_phase( + config: dict[str, Any], + stats: DatasetStats, + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + cfg = _validated_config(config) + n_trials = max(1, cfg.hpo_config.n_trials) + n_jobs = max(1, cfg.hpo_config.n_jobs) + + if not hub_reachable(): + report.low_confidence = True + report.notes.append("HF Hub unreachable — all model sizes are name-pattern heuristics.") + + seen_models: dict[str, ModelMeta] = {} + global_embedder = (cfg.embedder_config or {}).get("model_name") + if global_embedder: + seen_models[global_embedder] = resolve_model(global_embedder) + + transformer_entries, classic_entries = _split_entries(cfg.search_space) + + # First pass: transformer modules (also populates seen_models for the classic pass). + module_estimates: list[_ModuleEstimate] = [] + node_max_weights: dict[int, float] = {} + for node_idx, node_type, entry in transformer_entries: + module = entry.get("module_name", "?") + model_names = _extract_model_names(entry) + if not model_names and global_embedder and module in {"knn", "mlknn"}: + model_names = [global_embedder] + for name in model_names: + meta = seen_models.setdefault(name, resolve_model(name)) + me = _estimate_transformer_model( + meta=meta, + entry=entry, + node_type=node_type, + module=module, + name=name, + stats=stats, + hardware=hardware, + n_trials=n_trials, + refit_after=cfg.refit_after, + ) + module_estimates.append(me) + # Track heaviest weight per node so dump_modules is bounded by one + # selected variant per node x n_trials, not the sum of all candidates. + node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), me.model_weights_gb) + + # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. + embedder_meta = _largest_embedder(seen_models) + embedder_dim = _embedder_dim(embedder_meta) + for _, node_type, entry in classic_entries: + me = _estimate_classic_entry( + entry=entry, + node_type=node_type, + embedder_meta=embedder_meta, + embedder_dim=embedder_dim, + stats=stats, + hardware=hardware, + n_trials=n_trials, + refit_after=cfg.refit_after, + ) + if me is not None: + module_estimates.append(me) + + estimate = ResourceEstimate(parallel_factor=n_jobs) + for me in module_estimates: + estimate.vram_gb = max(estimate.vram_gb, me.vram_gb) + estimate.ram_gb = max(estimate.ram_gb, me.ram_gb) + estimate.time_hours += me.time_hours + estimate.drivers.append(me.driver) + + _aggregate_disk(estimate, seen_models, node_max_weights, dump_modules=cfg.dump_modules, n_trials=n_trials) + + report.resource = estimate + _emit_resource_findings(report, estimate, hardware, n_jobs=n_jobs) def _config_phase( diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index 9b351952a..b459adf9c 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -7,6 +7,7 @@ from __future__ import annotations +import json import logging import re from dataclasses import dataclass @@ -14,7 +15,7 @@ from pathlib import Path from typing import Any -from huggingface_hub import HfApi, scan_cache_dir, try_to_load_from_cache +from huggingface_hub import HfApi, hf_hub_download, scan_cache_dir, try_to_load_from_cache logger = logging.getLogger(__name__) @@ -43,6 +44,11 @@ class ModelMeta: total_file_bytes: int cached_locally: bool confidence: str # "hub" | "heuristic" + # Architecture shape read straight from the model's config.json when reachable; + # None when the file couldn't be fetched/parsed. Estimates fall back to a + # BERT-base default in that case. + hidden_size: int | None = None + n_layers: int | None = None @property def disk_gb(self) -> float: @@ -71,6 +77,30 @@ def _heuristic_params_millions(model_name: str) -> float: return 110.0 # generic BERT-base default +def _shape_from_config(model_name: str) -> tuple[int | None, int | None]: + """Return ``(hidden_size, num_hidden_layers)`` straight from the model's config.json. + + ``hf_hub_download`` caches the file after the first call, so repeated lookups + in the same process (or across CLI invocations) hit local disk. Returns + ``(None, None)`` on any failure — the advisor stays best-effort. + """ + try: + path = hf_hub_download(model_name, "config.json") + except Exception as e: # noqa: BLE001 + logger.debug("config.json download(%s) failed: %s", model_name, e) + return None, None + try: + cfg = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as e: + logger.debug("config.json parse(%s) failed: %s", model_name, e) + return None, None + # Cover the common HF naming variants: BERT/Llama/Gemma use hidden_size + + # num_hidden_layers; T5/MT5 use d_model + num_layers; GPT-2/Neo use n_embd + n_layer. + hidden = cfg.get("hidden_size") or cfg.get("d_model") or cfg.get("n_embd") + layers = cfg.get("num_hidden_layers") or cfg.get("num_layers") or cfg.get("n_layer") + return (int(hidden) if hidden else None, int(layers) if layers else None) + + def _is_warm_cached(model_name: str) -> bool: """True when the weight shard is present in the local HF cache.""" weight_files = ["model.safetensors", "pytorch_model.bin", "model.safetensors.index.json"] @@ -126,6 +156,14 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) confidence = "heuristic" + hidden_size, n_layers = _shape_from_config(model_name) + if hidden_size is None or n_layers is None: + logger.warning( + "Could not read hidden_size / num_hidden_layers from config.json for %s; " + "activation-memory estimates will fall back to BERT-base defaults (768 / 12).", + model_name, + ) + return ModelMeta( name=model_name, params_millions=params_millions, @@ -133,10 +171,17 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: total_file_bytes=total_file_bytes, cached_locally=_is_warm_cached(model_name), confidence=confidence, + hidden_size=hidden_size, + n_layers=n_layers, ) def _heuristic_metadata(model_name: str) -> ModelMeta: + logger.warning( + "Falling back to name-pattern heuristic for %s; " + "activation-memory estimates will use BERT-base defaults (hidden=768, layers=12).", + model_name, + ) params_millions = _heuristic_params_millions(model_name) weight_bytes_per_param = 4 total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 9b4881611..c63acde9d 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -146,12 +146,6 @@ def test_amp_does_reduce_activation_side_vram(self, meta: ModelMeta) -> None: amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True, batch_size=64, seq_len=128) assert amp < fp32 - def test_reranker_uses_inference_class(self, meta: ModelMeta) -> None: - inference = _vram_for_transformer(meta, "inference", mixed_precision=False) - reranker = _vram_for_transformer(meta, "reranker", mixed_precision=False) - assert reranker > inference - - def test_ram_scales_with_dataset_size() -> None: meta = ModelMeta( name="x", @@ -477,13 +471,13 @@ def test_driver_records_current_and_max_batch(self) -> None: report = run_preflight( self._bert_cfg("microsoft/deberta-v3-large", batch_size=64), DatasetStats.placeholder(), - _profile(vram_gb=10.0), + _profile(vram_gb=8.0), ) drivers = [d for d in report.resource.drivers if d["module"] == "bert"] assert drivers d = drivers[0] assert d["batch_size"] == 64 - # vram_gb=10 + 5 GB weights → some room for activations, max < 64. + # vram_gb=8 with ~5 GB weights leaves little room for activations → max < 64. assert d["max_batch_size"] is not None assert 0 < d["max_batch_size"] < 64 From 1f1778a509f0f8c4fff8277d43e307ffccbbed41 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 21:24:25 +0300 Subject: [PATCH 12/43] simplify logic --- src/autointent/_advisor/_estimates.py | 144 +++++++++++--------- src/autointent/_advisor/_hardware.py | 9 +- src/autointent/_advisor/_hub.py | 153 +++++++++++----------- src/autointent/_advisor/_render.py | 10 +- tests/advisor/test_estimates_and_cli.py | 10 +- tests/advisor/test_estimates_internals.py | 58 +++++--- tests/advisor/test_hub_heuristics.py | 51 +++----- tests/advisor/test_render.py | 8 +- 8 files changed, 230 insertions(+), 213 deletions(-) diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py index 274407c26..faeb4b1ee 100644 --- a/src/autointent/_advisor/_estimates.py +++ b/src/autointent/_advisor/_estimates.py @@ -12,11 +12,17 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import ValidationError -from autointent.configs._optimization import HPOConfig +from autointent._optimization_config import OptimizationConfig +from autointent.configs._embedder import ( + EmbedderConfig, + OpenaiEmbeddingConfig, + SentenceTransformerEmbeddingConfig, + VllmEmbeddingConfig, +) -from ._hub import hub_reachable, resolve_model +from ._hub import resolve_model from ._report import PreflightReport, ResourceEstimate, Severity if TYPE_CHECKING: @@ -27,6 +33,7 @@ from ._report import DatasetStats _MULTICLASS_THRESHOLD = 2 +_BYTES_PER_GB = 1024**3 # binary GiB convention; matches all advisor byte->GB conversions # Fallback architecture shape (BERT-base) used only when the model's actual # config.json couldn't be fetched from HF Hub — see _hub._shape_from_config. @@ -36,48 +43,39 @@ logger = logging.getLogger(__name__) -class _AdvisorConfig(BaseModel): - """Validated view of the advisor's input config. - - Wraps the four top-level keys the phase helpers read. Unknown top-level - keys are ignored (preset YAMLs carry extra metadata the advisor doesn't model). - """ - - model_config = ConfigDict(extra="ignore") - - hpo_config: HPOConfig = Field(default_factory=HPOConfig) - search_space: list[dict[str, Any]] = Field(default_factory=list) - refit_after: bool = False - dump_modules: bool = False - embedder_config: dict[str, Any] | None = None - - -def _validated_config(config: dict[str, Any]) -> _AdvisorConfig: - """Validate ``config`` against ``_AdvisorConfig``; fall back to defaults on any error. +def _validated_config(config: dict[str, Any]) -> OptimizationConfig: + """Validate ``config`` against the project's canonical ``OptimizationConfig``. The advisor is best-effort: a malformed user config should still produce a - report (with placeholder costs) rather than crashing. + report (with placeholder costs) rather than crashing, so any validation + error falls back to the model defaults. """ try: - return _AdvisorConfig.model_validate(config) + return OptimizationConfig.model_validate(config) except ValidationError as e: logger.warning("Advisor config failed validation; falling back to defaults: %s", e) - return _AdvisorConfig() + # OptimizationConfig requires `search_space`; build a minimal valid default. + return OptimizationConfig.model_validate({"search_space": []}) -# Severity thresholds as a fraction of available budget: at or above _TIGHT -# downgrades to Severity.TIGHT; at or above _OVER downgrades to Severity.OVER. -_TIGHT_RATIO = 0.7 -_OVER_RATIO = 1.0 +_TIGHT_RATIO = 0.9 + +# Union variants of EmbedderConfig that carry a model_name attribute. +# HashingVectorizerEmbeddingConfig and the bare BaseEmbedderConfig don't have +# one (sklearn vectorizer / abstract base), so we filter them out below. +_MODEL_BACKED_EMBEDDERS = ( + SentenceTransformerEmbeddingConfig, + OpenaiEmbeddingConfig, + VllmEmbeddingConfig, +) + + +def _embedder_model_name(embedder: EmbedderConfig) -> str | None: + """Return the embedder's model_name when the config variant carries one.""" + if isinstance(embedder, _MODEL_BACKED_EMBEDDERS): + return embedder.model_name + return None -# rough per-step seconds, keyed on device class. Scaled by params_millions / 100. -_PER_STEP_BASELINE_S = { - "cpu": 0.5, - "low-gpu": 0.04, - "mid-gpu": 0.02, - "high-gpu": 0.01, - "apple-silicon": 0.08, -} # Maps each fine-tunable transformer module to its training-mode label. # Modules not listed (or listed as "inference") run the encoder forward-only. @@ -98,7 +96,7 @@ def _validated_config(config: dict[str, Any]) -> _AdvisorConfig: _LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-8 _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 _CATBOOST_GPU_SPEEDUP = 10.0 -# LogisticRegressionCV defaults: Cs=10, cv=3 → 31 inner fits + 1 final refit. +# LogisticRegressionCV defaults: Cs=10, cv=3 -> 31 inner fits + 1 final refit. _LOGREG_CV_MULTIPLIER = 31 _CATBOOST_DEFAULT_BINS = 254 # Bytes per histogram bucket / tree node — order-of-magnitude constants. @@ -190,7 +188,7 @@ def _vram_for_transformer( def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: """RAM in GB. Loose upper bound.""" - return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / (1024**3) + return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB def _floor_to_power_of_two(n: int) -> int: @@ -232,7 +230,7 @@ def _activations_gb_per_sample( bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 if is_training else seq_len * hidden * 8 if mixed_precision: bytes_per_sample //= 2 - return bytes_per_sample / (1024**3) + return bytes_per_sample / _BYTES_PER_GB def _max_fitting_batch_size( @@ -265,7 +263,7 @@ def _embedder_dim(meta: ModelMeta | None) -> int: def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: if not seen_models: return None - return max(seen_models.values(), key=lambda m: m.params_millions) + return max(seen_models.values(), key=lambda m: m.total_params) def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: @@ -273,7 +271,7 @@ def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: data_bytes = 8.0 * stats.n_samples * embedder_dim coef_bytes = 8.0 * max(1, stats.n_classes) * embedder_dim lbfgs_bytes = 10.0 * 8.0 * embedder_dim - return (data_bytes + coef_bytes + lbfgs_bytes) / (1024**3) + return (data_bytes + coef_bytes + lbfgs_bytes) / _BYTES_PER_GB def _time_for_linear( @@ -301,7 +299,7 @@ def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, data_bytes = 4.0 * stats.n_samples * n_features histograms_bytes = 4.0 * n_features * _CATBOOST_DEFAULT_BINS trees_bytes = iterations * (2**depth) * _CATBOOST_BYTES_PER_TREE_NODE - return float((data_bytes + histograms_bytes + trees_bytes) / (1024**3)) + return float((data_bytes + histograms_bytes + trees_bytes) / _BYTES_PER_GB) def _time_for_catboost( @@ -323,16 +321,20 @@ def _time_for_catboost( def _time_for_transformer( *, - meta: ModelMeta, n_trials: int, epochs: int, batch_size: int, n_samples: int, - device_class: str, ) -> float: - per_step = _PER_STEP_BASELINE_S[device_class] * (meta.params_millions / 100.0) + """Transformer training time in hours, assuming a flat 1 second per step. + + The advisor has no real wall-time calibration across hardware tiers / model + sizes, so the report uses ``time_hours`` as a step-count proxy rather than + pretending to estimate seconds. Users should treat the number as ordering / + ballpark information, not a budget. + """ steps = max(1, (n_samples // max(1, batch_size))) * epochs - return (n_trials * steps * per_step) / 3600.0 + return (n_trials * steps) / 3600.0 def _classify_severity(estimate: float, budget: float) -> Severity: @@ -341,7 +343,7 @@ def _classify_severity(estimate: float, budget: float) -> Severity: if budget <= 0: return Severity.TIGHT ratio = estimate / budget - if ratio >= _OVER_RATIO: + if ratio >= 1: return Severity.OVER if ratio >= _TIGHT_RATIO: return Severity.TIGHT @@ -368,7 +370,8 @@ def _split_entries( search_space: list[dict[str, Any]], ) -> tuple[list[tuple[int, str, dict[str, Any]]], list[tuple[int, str, dict[str, Any]]]]: """Partition search-space entries into (transformer-bearing, classic).""" - transformer, classic = [], [] + transformer: list[tuple[int, str, dict[str, Any]]] = [] + classic: list[tuple[int, str, dict[str, Any]]] = [] for node_idx, node_type, entry in _walk_modules_indexed(search_space): bucket = classic if entry.get("module_name") in {"linear", "catboost"} else transformer bucket.append((node_idx, node_type, entry)) @@ -408,12 +411,10 @@ def _estimate_transformer_model( ) time_h = _time_for_transformer( - meta=meta, n_trials=n_trials, epochs=epochs, batch_size=batch_size, n_samples=stats.n_samples, - device_class=hardware.device_class, ) if mode != "inference": time_h *= _refit_factor(refit_after=refit_after, n_trials=n_trials) @@ -593,17 +594,16 @@ def _resource_phase( stats: DatasetStats, hardware: HardwareProfile, report: PreflightReport, + *, + refit_after: bool = False, ) -> None: cfg = _validated_config(config) - n_trials = max(1, cfg.hpo_config.n_trials) - n_jobs = max(1, cfg.hpo_config.n_jobs) - - if not hub_reachable(): - report.low_confidence = True - report.notes.append("HF Hub unreachable — all model sizes are name-pattern heuristics.") + n_trials = cfg.hpo_config.n_trials + n_jobs = cfg.hpo_config.n_jobs + dump_modules = cfg.logging_config.dump_modules seen_models: dict[str, ModelMeta] = {} - global_embedder = (cfg.embedder_config or {}).get("model_name") + global_embedder = _embedder_model_name(cfg.embedder_config) if global_embedder: seen_models[global_embedder] = resolve_model(global_embedder) @@ -628,7 +628,7 @@ def _resource_phase( stats=stats, hardware=hardware, n_trials=n_trials, - refit_after=cfg.refit_after, + refit_after=refit_after, ) module_estimates.append(me) # Track heaviest weight per node so dump_modules is bounded by one @@ -639,7 +639,7 @@ def _resource_phase( embedder_meta = _largest_embedder(seen_models) embedder_dim = _embedder_dim(embedder_meta) for _, node_type, entry in classic_entries: - me = _estimate_classic_entry( + classic_estimate = _estimate_classic_entry( entry=entry, node_type=node_type, embedder_meta=embedder_meta, @@ -647,10 +647,10 @@ def _resource_phase( stats=stats, hardware=hardware, n_trials=n_trials, - refit_after=cfg.refit_after, + refit_after=refit_after, ) - if me is not None: - module_estimates.append(me) + if classic_estimate is not None: + module_estimates.append(classic_estimate) estimate = ResourceEstimate(parallel_factor=n_jobs) for me in module_estimates: @@ -659,7 +659,17 @@ def _resource_phase( estimate.time_hours += me.time_hours estimate.drivers.append(me.driver) - _aggregate_disk(estimate, seen_models, node_max_weights, dump_modules=cfg.dump_modules, n_trials=n_trials) + _aggregate_disk(estimate, seen_models, node_max_weights, dump_modules=dump_modules, n_trials=n_trials) + + # Flip low_confidence if any model fell back to the heuristic path (Hub + # unreachable, repo missing safetensors metadata, local-path checkpoint). + heuristic_models = [m.name for m in seen_models.values() if m.confidence == "heuristic"] + if heuristic_models: + report.low_confidence = True + report.notes.append( + f"Heuristic fallback used for {len(heuristic_models)} model(s) — sizes are BERT-base " + f"defaults: {', '.join(heuristic_models[:3])}{'...' if len(heuristic_models) > 3 else ''}", # noqa: PLR2004 + ) report.resource = estimate _emit_resource_findings(report, estimate, hardware, n_jobs=n_jobs) @@ -743,15 +753,19 @@ def run_preflight( hardware: HardwareProfile, *, preset_name: str | None = None, + refit_after: bool = False, ) -> PreflightReport: """Run all three phases and return one report. Args: config: parsed preset / OptimizationConfig dict (top-level keys: - ``search_space``, ``hpo_config``, optional ``embedder_config``). + ``search_space``, ``hpo_config``, optional ``embedder_config``, + optional ``logging_config.dump_modules``). stats: dataset statistics (real or placeholder). hardware: detected hardware profile. preset_name: optional friendly name for the report header. + refit_after: matches the ``Pipeline.fit(refit_after=...)`` argument. + When True, time estimates include the extra refit-on-full-data pass. Returns: PreflightReport with findings across resource/data/config phases. @@ -777,7 +791,7 @@ def run_preflight( ) report.notes.extend(hardware.notes) - _resource_phase(config, stats, hardware, report) + _resource_phase(config, stats, hardware, report, refit_after=refit_after) _data_phase(config, stats, report) _config_phase(config, hardware, report) diff --git a/src/autointent/_advisor/_hardware.py b/src/autointent/_advisor/_hardware.py index e959b6ebf..6aa9741ee 100644 --- a/src/autointent/_advisor/_hardware.py +++ b/src/autointent/_advisor/_hardware.py @@ -1,7 +1,7 @@ """Local hardware detection. Probes CPU / RAM / disk and the highest-priority accelerator available -(CUDA → MPS → CPU). All probes are wrapped to fall back safely on a +(CUDA -> MPS -> CPU). All probes are wrapped to fall back safely on a broken install (e.g. CUDA driver mismatch) rather than crash the advisor. """ @@ -27,6 +27,7 @@ _HIGH_GPU_VRAM_GB = 24 _MID_GPU_VRAM_GB = 12 +_BYTES_PER_GB = 1024**3 # binary GiB convention; matches all advisor byte->GB conversions @dataclass @@ -53,7 +54,7 @@ def device_class(self) -> str: def _detect_ram_gb() -> float: - return float(psutil.virtual_memory().total) / (1024**3) + return float(psutil.virtual_memory().total) / _BYTES_PER_GB def _detect_free_disk_gb(path: str | None = None) -> float: @@ -61,7 +62,7 @@ def _detect_free_disk_gb(path: str | None = None) -> float: probe_path = cache if cache.exists() else Path("~").expanduser() try: usage = shutil.disk_usage(probe_path) - return usage.free / (1024**3) + return usage.free / _BYTES_PER_GB except OSError as e: logger.debug("disk usage probe failed at %s: %s", probe_path, e) return 0.0 @@ -73,7 +74,7 @@ def _detect_cuda() -> tuple[float, str] | None: idx = 0 try: _free, total = torch.cuda.mem_get_info(idx) - vram_gb = total / (1024**3) + vram_gb = total / _BYTES_PER_GB except (RuntimeError, AttributeError) as e: logger.debug("torch.cuda.mem_get_info failed: %s", e) return None diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index b459adf9c..a5b6126a5 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -9,72 +9,40 @@ import json import logging -import re from dataclasses import dataclass from functools import lru_cache from pathlib import Path -from typing import Any +from typing import Literal from huggingface_hub import HfApi, hf_hub_download, scan_cache_dir, try_to_load_from_cache +Confidence = Literal["hub", "heuristic"] + logger = logging.getLogger(__name__) -# Coarse heuristic estimates keyed on name fragments. Used only when HF Hub -# is unreachable and we can't get safetensors metadata. Values in millions. -_NAME_HEURISTICS = [ - (re.compile(r"(?i)(deberta|roberta|bert).*(xxlarge|huge)"), 1_500), - (re.compile(r"(?i)(deberta|roberta|bert).*xlarge"), 750), - (re.compile(r"(?i)(deberta|roberta|bert).*large"), 350), - (re.compile(r"(?i)e5.*large"), 560), - (re.compile(r"(?i)e5.*small"), 33), - (re.compile(r"(?i)mpnet"), 110), - (re.compile(r"(?i)minilm"), 33), - (re.compile(r"(?i)distil"), 66), - (re.compile(r"(?i)small"), 60), - (re.compile(r"(?i)base"), 110), - (re.compile(r"(?i)large"), 350), -] +_DEFAULT_HEURISTIC_PARAMS = 110_000_000 +_DEFAULT_BYTES_PER_PARAM = 4 +_BYTES_PER_GB = 1024**3 # using the binary GiB convention everywhere in the advisor @dataclass class ModelMeta: name: str - params_millions: float - weight_bytes_per_param: int + total_params: int + weight_bytes_per_param: float total_file_bytes: int cached_locally: bool - confidence: str # "hub" | "heuristic" - # Architecture shape read straight from the model's config.json when reachable; - # None when the file couldn't be fetched/parsed. Estimates fall back to a - # BERT-base default in that case. + confidence: Confidence hidden_size: int | None = None n_layers: int | None = None @property def disk_gb(self) -> float: - return self.total_file_bytes / (1024**3) + return self.total_file_bytes / _BYTES_PER_GB @property def weights_gb(self) -> float: - return (self.params_millions * 1_000_000 * self.weight_bytes_per_param) / (1024**3) - - -@lru_cache(maxsize=1) -def hub_reachable() -> bool: - """Single up-front probe. Memoized per process.""" - try: - HfApi().list_models(limit=1) - except Exception as e: # noqa: BLE001 - logger.debug("HF Hub probe failed: %s", e) - return False - return True - - -def _heuristic_params_millions(model_name: str) -> float: - for pattern, m in _NAME_HEURISTICS: - if pattern.search(model_name): - return float(m) - return 110.0 # generic BERT-base default + return (self.total_params * self.weight_bytes_per_param) / _BYTES_PER_GB def _shape_from_config(model_name: str) -> tuple[int | None, int | None]: @@ -98,7 +66,7 @@ def _shape_from_config(model_name: str) -> tuple[int | None, int | None]: # num_hidden_layers; T5/MT5 use d_model + num_layers; GPT-2/Neo use n_embd + n_layer. hidden = cfg.get("hidden_size") or cfg.get("d_model") or cfg.get("n_embd") layers = cfg.get("num_hidden_layers") or cfg.get("num_layers") or cfg.get("n_layer") - return (int(hidden) if hidden else None, int(layers) if layers else None) + return int(hidden) if hidden else None, int(layers) if layers else None def _is_warm_cached(model_name: str) -> bool: @@ -124,36 +92,46 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: except Exception as e: # noqa: BLE001 logger.debug("model_info(%s) failed: %s", model_name, e) return None - - params_millions = 0.0 - weight_bytes_per_param = 4 - safetensors = getattr(info, "safetensors", None) - if safetensors is not None: - params_total = getattr(safetensors, "total", None) or sum( - getattr(safetensors, "parameters", {}).values() or [0] - ) - if params_total: - params_millions = params_total / 1_000_000 - params_map: dict[str, Any] = getattr(safetensors, "parameters", {}) or {} - if any("F16" in k or "BF16" in k for k in params_map): - weight_bytes_per_param = 2 - - total_file_bytes = 0 - for sibling in getattr(info, "siblings", []) or []: - size = getattr(sibling, "size", None) - if size: - total_file_bytes += int(size) + # Bytes-per-element for safetensors dtype strings. Used to convert the per-dtype + # parameter counts (info.safetensors.parameters) into a weighted average + # bytes-per-param for mixed-precision repos. + _dtype_bytes: dict[str, int] = { + "F64": 8, + "F32": 4, + "F16": 2, + "BF16": 2, + "I64": 8, + "I32": 4, + "I16": 2, + "I8": 1, + "U8": 1, + "BOOL": 1, + } + + total_params = 0 + weight_bytes_per_param: float = _DEFAULT_BYTES_PER_PARAM + if info.safetensors is not None: + params_by_dtype = info.safetensors.parameters or {} + total_params = info.safetensors.total or sum(params_by_dtype.values()) + if total_params: + total_weight_bytes = sum( + _dtype_bytes.get(dtype, _DEFAULT_BYTES_PER_PARAM) * count for dtype, count in params_by_dtype.items() + ) + if total_weight_bytes: + weight_bytes_per_param = total_weight_bytes / total_params + + total_file_bytes = sum(s.size for s in (info.siblings or []) if s.size) # Track whether either size came from the Hub or from the name-pattern fallback; # if any field was filled by heuristic, downgrade confidence so the report flips # low_confidence rather than misreporting hub-grade accuracy. - confidence = "hub" - if params_millions == 0: - params_millions = _heuristic_params_millions(model_name) + confidence: Confidence = "hub" + if total_params == 0: + total_params = _DEFAULT_HEURISTIC_PARAMS confidence = "heuristic" if total_file_bytes == 0: - total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) + total_file_bytes = int(total_params * weight_bytes_per_param) confidence = "heuristic" hidden_size, n_layers = _shape_from_config(model_name) @@ -166,7 +144,7 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: return ModelMeta( name=model_name, - params_millions=params_millions, + total_params=total_params, weight_bytes_per_param=weight_bytes_per_param, total_file_bytes=total_file_bytes, cached_locally=_is_warm_cached(model_name), @@ -182,19 +160,33 @@ def _heuristic_metadata(model_name: str) -> ModelMeta: "activation-memory estimates will use BERT-base defaults (hidden=768, layers=12).", model_name, ) - params_millions = _heuristic_params_millions(model_name) - weight_bytes_per_param = 4 - total_file_bytes = int(params_millions * 1_000_000 * weight_bytes_per_param) + total_file_bytes = _DEFAULT_HEURISTIC_PARAMS * _DEFAULT_BYTES_PER_PARAM return ModelMeta( name=model_name, - params_millions=params_millions, - weight_bytes_per_param=weight_bytes_per_param, + total_params=_DEFAULT_HEURISTIC_PARAMS, + weight_bytes_per_param=_DEFAULT_BYTES_PER_PARAM, total_file_bytes=total_file_bytes, cached_locally=_is_warm_cached(model_name), confidence="heuristic", ) +def _looks_like_local_path(model_name: str) -> bool: + """True when ``model_name`` is a filesystem path rather than an HF Hub repo id. + + Hub repo ids match ``org/repo``; anything that starts with a path separator, + ``~``, a relative-path prefix, or a Windows drive letter, or contains a + backslash, is treated as a local path. We can't rely on ``Path.is_absolute()`` + alone because POSIX-style absolute paths (``/tmp/...``) are *not* absolute + on Windows. + """ + if model_name.startswith(("local:", "/", "~", "./", "../", "\\\\")): + return True + if "\\" in model_name: + return True + return len(model_name) >= 2 and model_name[1] == ":" and model_name[0].isalpha() # noqa: PLR2004 + + @lru_cache(maxsize=64) def resolve_model(model_name: str) -> ModelMeta: """Resolve metadata for a single model name. Memoized per process. @@ -202,19 +194,20 @@ def resolve_model(model_name: str) -> ModelMeta: Always returns a value — never raises — so the advisor can keep going on offline machines or for unknown checkpoints. """ - if model_name.startswith("local:") or Path(model_name).is_absolute(): + if _looks_like_local_path(model_name): return ModelMeta( name=model_name, - params_millions=_heuristic_params_millions(model_name), - weight_bytes_per_param=4, + total_params=_DEFAULT_HEURISTIC_PARAMS, + weight_bytes_per_param=_DEFAULT_BYTES_PER_PARAM, total_file_bytes=0, cached_locally=True, confidence="heuristic", ) - if hub_reachable(): - meta = _hub_metadata(model_name) - if meta is not None: - return meta + # _hub_metadata returns None on any failure (network outage, missing repo, + # SDK exception) so we don't need a separate up-front probe. + meta = _hub_metadata(model_name) + if meta is not None: + return meta return _heuristic_metadata(model_name) diff --git a/src/autointent/_advisor/_render.py b/src/autointent/_advisor/_render.py index 82771ef9f..afd541e2b 100644 --- a/src/autointent/_advisor/_render.py +++ b/src/autointent/_advisor/_render.py @@ -13,13 +13,13 @@ if TYPE_CHECKING: from ._report import PreflightReport -_SEVERITY_TAG = {"ample": "✓", "tight": "⚠", "over": "✗"} +_SEVERITY_TAG = {"ample": "✓", "tight": "⚠", "over": "x"} _PHASE_ORDER = ("resource", "data", "config") _PHASE_LABEL = {"resource": "Resource", "data": "Data", "config": "Config"} def _batch_hint(driver: dict[str, Any]) -> str: - """Per-driver batch annotation: '64 → 32', '64', '64 (no fit)', or ''.""" + """Per-driver batch annotation: '64 -> 32', '64', '64 (no fit)', or ''.""" bs = driver.get("batch_size") if bs is None: return "" @@ -30,7 +30,7 @@ def _batch_hint(driver: dict[str, Any]) -> str: return f"{bs} (no fit)" if mx == bs: return str(bs) - return f"{bs} → {mx}" + return f"{bs} -> {mx}" _DRIVERS_LIMIT = 8 @@ -137,9 +137,9 @@ def render_recommendation( """Compact table for the ``recommend`` subcommand.""" lines = ["", "Recommendation:"] if chosen: - lines.append(f" → {chosen}") + lines.append(f" -> {chosen}") else: - lines.append(" → none of the bundled presets fit your hardware as-is.") + lines.append(" -> none of the bundled presets fit your hardware as-is.") lines.append("") lines.append(f"{'Preset':<24} {'Status':<14} {'VRAM':<10} {'Time':<10} {'Headroom':<10}") lines.append("-" * 68) diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 2f16555ae..d87f90740 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -23,14 +23,11 @@ @pytest.fixture(autouse=True) def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: - """Pin the HF Hub probe to "offline" so tests don't hit the network.""" - from autointent._advisor import _estimates, _hub + """Force HF Hub lookups to fail so tests don't hit the network.""" + from autointent._advisor import _hub - _hub.hub_reachable.cache_clear() _hub.resolve_model.cache_clear() - offline = lambda *_a, **_kw: False # noqa: E731 - monkeypatch.setattr(_hub, "hub_reachable", offline) - monkeypatch.setattr(_estimates, "hub_reachable", offline) + monkeypatch.setattr(_hub, "_hub_metadata", lambda _name: None) def _profile(vram_gb: float = 16.0) -> HardwareProfile: @@ -50,7 +47,6 @@ def test_every_preset_inspects_without_raising(preset: str) -> None: stats = DatasetStats.placeholder(n_samples=500, n_classes=10, avg_tokens=24) report = run_preflight(cfg, stats, _profile(vram_gb=16.0), preset_name=preset) assert report.preset_name == preset - assert report.low_confidence is True # we forced offline # always at least one resource-phase finding assert any(f.phase == "resource" for f in report.findings) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index c63acde9d..3fa293b7e 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -19,15 +19,41 @@ from autointent._advisor._hub import ModelMeta from autointent._advisor._report import DatasetStats, Severity +# Per-name ModelMeta fixtures used by the offline tests. Production resolution +# (HF Hub config.json + safetensors metadata) is mocked away so the batch-fit +# math doesn't depend on whatever fallback the heuristic path returns. +_FAKE_SHAPES: dict[str, tuple[int, int, int]] = { + # (total_params, hidden_size, n_layers) + "microsoft/deberta-v3-large": (350_000_000, 1024, 24), + "microsoft/deberta-v3-small": (140_000_000, 768, 6), + "sentence-transformers/all-MiniLM-L6-v2": (33_000_000, 384, 6), + "intfloat/multilingual-e5-large-instruct": (560_000_000, 1024, 24), +} + + +def _fake_resolve(model_name: str) -> ModelMeta: + known = _FAKE_SHAPES.get(model_name) + params, hidden, layers = known or (110_000_000, 768, 12) + return ModelMeta( + name=model_name, + total_params=params, + weight_bytes_per_param=4, + total_file_bytes=params * 4, + cached_locally=False, + confidence="hub" if known else "heuristic", + hidden_size=hidden, + n_layers=layers, + ) + @pytest.fixture(autouse=True) def _offline(monkeypatch: pytest.MonkeyPatch) -> None: - _hub.hub_reachable.cache_clear() _hub.resolve_model.cache_clear() - offline = lambda *_a, **_kw: False # noqa: E731 - monkeypatch.setattr(_hub, "hub_reachable", offline) - monkeypatch.setattr(_estimates, "hub_reachable", offline) monkeypatch.setattr(_hub, "_is_warm_cached", lambda _name: False) + # Inject deterministic ModelMeta per name; both the _hub re-export and the + # _estimates rebinding need to be replaced for run_preflight to pick it up. + monkeypatch.setattr(_hub, "resolve_model", _fake_resolve) + monkeypatch.setattr(_estimates, "resolve_model", _fake_resolve) def _profile(vram_gb: float = 16.0, accelerator: str = "cuda") -> HardwareProfile: @@ -89,7 +115,7 @@ def test_below_yellow_is_green(self) -> None: assert _classify_severity(estimate=1.0, budget=10.0) == Severity.AMPLE def test_above_yellow_threshold(self) -> None: - assert _classify_severity(estimate=8.0, budget=10.0) == Severity.TIGHT + assert _classify_severity(estimate=9.5, budget=10.0) == Severity.TIGHT def test_at_or_above_red_threshold(self) -> None: assert _classify_severity(estimate=10.0, budget=10.0) == Severity.OVER @@ -104,7 +130,7 @@ class TestVramForTransformer: def meta(self) -> ModelMeta: return ModelMeta( name="x", - params_millions=100.0, + total_params=100_000_000, weight_bytes_per_param=4, total_file_bytes=0, cached_locally=False, @@ -146,10 +172,11 @@ def test_amp_does_reduce_activation_side_vram(self, meta: ModelMeta) -> None: amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True, batch_size=64, seq_len=128) assert amp < fp32 + def test_ram_scales_with_dataset_size() -> None: meta = ModelMeta( name="x", - params_millions=100.0, + total_params=100_000_000, weight_bytes_per_param=4, total_file_bytes=0, cached_locally=False, @@ -177,7 +204,7 @@ def test_dump_modules_adds_disk_during_training(self) -> None: } ], "hpo_config": {"n_trials": 5}, - "dump_modules": True, + "logging_config": {"dump_modules": True}, } report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) assert report.resource.disk_dump_gb > 0 @@ -201,8 +228,7 @@ def test_refit_after_increases_time(self) -> None: "hpo_config": {"n_trials": 10}, } baseline = run_preflight(cfg, DatasetStats.placeholder(), _profile()) - cfg_refit = {**cfg, "refit_after": True} - bumped = run_preflight(cfg_refit, DatasetStats.placeholder(), _profile()) + bumped = run_preflight(cfg, DatasetStats.placeholder(), _profile(), refit_after=True) assert bumped.resource.time_hours > baseline.resource.time_hours def test_catboost_gpu_without_cuda_flags_config(self) -> None: @@ -249,7 +275,7 @@ def test_offline_flips_low_confidence(self) -> None: } report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) assert report.low_confidence is True - assert any("HF Hub unreachable" in n for n in report.notes) + assert any("Heuristic fallback" in n for n in report.notes) def test_rare_classes_with_linear_scorer_flag_red(self) -> None: cfg = { @@ -471,13 +497,13 @@ def test_driver_records_current_and_max_batch(self) -> None: report = run_preflight( self._bert_cfg("microsoft/deberta-v3-large", batch_size=64), DatasetStats.placeholder(), - _profile(vram_gb=8.0), + _profile(vram_gb=6.5), ) drivers = [d for d in report.resource.drivers if d["module"] == "bert"] assert drivers d = drivers[0] assert d["batch_size"] == 64 - # vram_gb=8 with ~5 GB weights leaves little room for activations → max < 64. + # vram_gb=6.5 against ~5 GB weights x 0.9 tight ratio -> little activation room, max < 64. assert d["max_batch_size"] is not None assert 0 < d["max_batch_size"] < 64 @@ -523,7 +549,7 @@ def test_multiple_drivers_carry_independent_max_batch(self) -> None: report = run_preflight(cfg, DatasetStats.placeholder(), _profile(vram_gb=10.0)) small = next(d for d in report.resource.drivers if "small" in d["model"]) large = next(d for d in report.resource.drivers if "large" in d["model"]) - # The smaller model has more headroom → larger max batch (or equal-cap when both saturate). + # The smaller model has more headroom -> larger max batch (or equal-cap when both saturate). assert small["max_batch_size"] >= large["max_batch_size"] @@ -551,7 +577,7 @@ def test_dump_disk_is_bounded_by_per_node_max_not_sum_of_all_variants(self) -> N } ], "hpo_config": {"n_trials": 4}, - "dump_modules": True, + "logging_config": {"dump_modules": True}, } report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) # Per-node max ~ deberta-v3-large weights (~350M x 4 ~ 1.3 GB). Two-candidate @@ -588,7 +614,7 @@ def test_dump_disk_sums_across_nodes(self) -> None: }, ], "hpo_config": {"n_trials": 2}, - "dump_modules": True, + "logging_config": {"dump_modules": True}, } report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) embedder = _hub.resolve_model("sentence-transformers/all-MiniLM-L6-v2") diff --git a/tests/advisor/test_hub_heuristics.py b/tests/advisor/test_hub_heuristics.py index b43b95522..c19018235 100644 --- a/tests/advisor/test_hub_heuristics.py +++ b/tests/advisor/test_hub_heuristics.py @@ -1,8 +1,8 @@ -"""Tests for the offline name-pattern heuristics in `_hub`. +"""Tests for the offline heuristic fallback in `_hub`. -The advisor must produce a sensible estimate even when HF Hub is -unreachable, so these tests pin the public `hub_reachable` to False and -exercise the heuristic path directly. +The advisor must produce a sensible estimate even when HF Hub is unreachable. +Without a per-name heuristic, every offline lookup collapses to a single +BERT-base-sized default — these tests pin that contract. """ from __future__ import annotations @@ -14,41 +14,28 @@ @pytest.fixture(autouse=True) def _offline(monkeypatch: pytest.MonkeyPatch) -> None: - _hub.hub_reachable.cache_clear() _hub.resolve_model.cache_clear() - monkeypatch.setattr(_hub, "hub_reachable", lambda *_a, **_kw: False) + # Force `_hub_metadata` to behave as if the live Hub were unreachable so + # resolve_model falls through to `_heuristic_metadata`. + monkeypatch.setattr(_hub, "_hub_metadata", lambda _name: None) monkeypatch.setattr(_hub, "_is_warm_cached", lambda _name: False) -@pytest.mark.parametrize( - ("name", "expected_min_m", "expected_max_m"), - [ - ("microsoft/deberta-v3-large", 200, 500), - ("microsoft/deberta-v3-small", 30, 200), - ("sentence-transformers/all-MiniLM-L6-v2", 20, 80), - ("intfloat/multilingual-e5-large-instruct", 300, 700), - ("intfloat/e5-small", 20, 80), - ("distilbert-base-uncased", 40, 150), - ("bert-base-uncased", 70, 200), - ], -) -def test_name_heuristic_picks_reasonable_bucket(name: str, expected_min_m: int, expected_max_m: int) -> None: - meta = _hub.resolve_model(name) - assert meta.confidence == "heuristic" - assert expected_min_m <= meta.params_millions <= expected_max_m, ( - f"{name} got {meta.params_millions}M; expected [{expected_min_m}, {expected_max_m}]" - ) - - -def test_unknown_name_falls_back_to_bert_base() -> None: - meta = _hub.resolve_model("totally-made-up/no-such-model") - assert meta.confidence == "heuristic" - assert meta.params_millions == pytest.approx(110.0) +def test_offline_lookup_uses_bert_base_default() -> None: + """Every offline lookup returns the same BERT-base-sized fallback.""" + for name in ( + "microsoft/deberta-v3-large", + "sentence-transformers/all-MiniLM-L6-v2", + "totally-made-up/no-such-model", + ): + meta = _hub.resolve_model(name) + assert meta.confidence == "heuristic" + assert meta.total_params == _hub._DEFAULT_HEURISTIC_PARAMS def test_weights_gb_matches_params_times_bytes() -> None: meta = _hub.resolve_model("microsoft/deberta-v3-large") - expected_gb = meta.params_millions * 1_000_000 * meta.weight_bytes_per_param / (1024**3) + expected_gb = meta.total_params * meta.weight_bytes_per_param / (1024**3) assert meta.weights_gb == pytest.approx(expected_gb) @@ -75,5 +62,5 @@ def test_metadata_fallback_uses_heuristic_when_hub_unreachable() -> None: the live Hub is unreachable (autouse fixture forces offline).""" meta = _hub.resolve_model("microsoft/deberta-v3-large") assert meta.confidence == "heuristic" - assert meta.params_millions > 0 + assert meta.total_params > 0 assert meta.disk_gb > 0 diff --git a/tests/advisor/test_render.py b/tests/advisor/test_render.py index 2c0604a11..7a806c7f2 100644 --- a/tests/advisor/test_render.py +++ b/tests/advisor/test_render.py @@ -56,7 +56,7 @@ def test_contains_phase_blocks(self) -> None: out = render_text(_populated_report()) assert "Resource:" in out assert "Data:" in out - # Config phase has no findings → block omitted + # Config phase has no findings -> block omitted assert "Config:" not in out def test_includes_drivers_block(self) -> None: @@ -119,7 +119,7 @@ def _two_reports(self) -> list[tuple[str, PreflightReport]]: def test_lists_chosen_preset_when_present(self) -> None: out = render_recommendation(self._two_reports(), chosen="a") - assert "→ a" in out + assert "-> a" in out def test_handles_no_chosen(self) -> None: out = render_recommendation(self._two_reports(), chosen=None) @@ -140,7 +140,7 @@ class TestBatchHint: """Per-driver batch cell rendered in the Drivers-of-cost table.""" def test_arrow_when_max_differs(self) -> None: - assert _batch_hint({"batch_size": 64, "max_batch_size": 32}) == "64 → 32" + assert _batch_hint({"batch_size": 64, "max_batch_size": 32}) == "64 -> 32" def test_plain_when_max_equals_current(self) -> None: assert _batch_hint({"batch_size": 64, "max_batch_size": 64}) == "64" @@ -152,7 +152,7 @@ def test_empty_when_no_batch(self) -> None: assert _batch_hint({"batch_size": None, "max_batch_size": None}) == "" def test_increase_arrow(self) -> None: - assert _batch_hint({"batch_size": 32, "max_batch_size": 128}) == "32 → 128" + assert _batch_hint({"batch_size": 32, "max_batch_size": 128}) == "32 -> 128" def test_dataset_stats_in_text_block() -> None: From e0f14866a898714df279f73ff1dccd9f288725ed Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 21:26:39 +0300 Subject: [PATCH 13/43] remove from init --- src/autointent/_advisor/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index 28422c78d..d8eb6f5ba 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -10,10 +10,9 @@ from ._estimates import run_preflight from ._hardware import HardwareProfile, detect_hardware from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity -from ._workflows import BUNDLED_PRESETS, inspect, load_config, recommend, stats_from_dataset +from ._workflows import inspect, load_config, recommend, stats_from_dataset __all__ = [ - "BUNDLED_PRESETS", "DatasetStats", "Finding", "HardwareProfile", From 6496b4e1873d93ef5e5cc53cdb43b2612611f4d5 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 21:29:13 +0300 Subject: [PATCH 14/43] revert pyproject.toml --- pyproject.toml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ace1a3c77..276361669 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -298,12 +298,6 @@ module = [ "dspy.evaluate.auto_evaluation", "codecarbon", "catboost", - "openai", - "openai.*", - "tiktoken", - "peft", - "sentence_transformers", - "psutil", ] ignore_missing_imports = true From bc3df74af98ae1e63aac3b3f69391b223d03585d Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Jun 2026 21:37:55 +0300 Subject: [PATCH 15/43] update typing --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 276361669..7677ac2aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,6 +120,7 @@ typing = [ "joblib-stubs (>=1.4.2.5.20240918,<2.0.0)", "pandas-stubs (>= 2.2.3.250527, <3.0.0)", "types-aiofiles (>=24.1.0.20250606)", + "types-psutil>=7.2.2.20260518", ] docs = [ "sphinx (>=8.1.3,<9.0.0)", From 7cb0f53e9929465637290dbfce27c7661ee9bd61 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Mon, 22 Jun 2026 13:10:14 +0300 Subject: [PATCH 16/43] refactor --- src/autointent/_advisor/_estimates.py | 798 ------------------ src/autointent/_advisor/_hub.py | 2 +- .../_advisor/{_workflows.py => workflows.py} | 17 +- src/autointent/custom_types/_types.py | 12 +- tests/advisor/test_estimates_internals.py | 33 +- 5 files changed, 26 insertions(+), 836 deletions(-) delete mode 100644 src/autointent/_advisor/_estimates.py rename src/autointent/_advisor/{_workflows.py => workflows.py} (92%) diff --git a/src/autointent/_advisor/_estimates.py b/src/autointent/_advisor/_estimates.py deleted file mode 100644 index faeb4b1ee..000000000 --- a/src/autointent/_advisor/_estimates.py +++ /dev/null @@ -1,798 +0,0 @@ -"""Resource-phase estimation: walk the search space and aggregate cost. - -Implements an honest worst-case for the modules the proposal lists as -in-scope. Formulas are intentionally coarse — the advisor's contract is -"heuristic upper bound, not measurement". Time and VRAM are the noisiest; -treat them as ballparks, not budgets. -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -from pydantic import ValidationError - -from autointent._optimization_config import OptimizationConfig -from autointent.configs._embedder import ( - EmbedderConfig, - OpenaiEmbeddingConfig, - SentenceTransformerEmbeddingConfig, - VllmEmbeddingConfig, -) - -from ._hub import resolve_model -from ._report import PreflightReport, ResourceEstimate, Severity - -if TYPE_CHECKING: - from collections.abc import Iterable - - from ._hardware import HardwareProfile - from ._hub import ModelMeta - from ._report import DatasetStats - -_MULTICLASS_THRESHOLD = 2 -_BYTES_PER_GB = 1024**3 # binary GiB convention; matches all advisor byte->GB conversions - -# Fallback architecture shape (BERT-base) used only when the model's actual -# config.json couldn't be fetched from HF Hub — see _hub._shape_from_config. -_DEFAULT_HIDDEN = 768 -_DEFAULT_LAYERS = 12 - -logger = logging.getLogger(__name__) - - -def _validated_config(config: dict[str, Any]) -> OptimizationConfig: - """Validate ``config`` against the project's canonical ``OptimizationConfig``. - - The advisor is best-effort: a malformed user config should still produce a - report (with placeholder costs) rather than crashing, so any validation - error falls back to the model defaults. - """ - try: - return OptimizationConfig.model_validate(config) - except ValidationError as e: - logger.warning("Advisor config failed validation; falling back to defaults: %s", e) - # OptimizationConfig requires `search_space`; build a minimal valid default. - return OptimizationConfig.model_validate({"search_space": []}) - - -_TIGHT_RATIO = 0.9 - -# Union variants of EmbedderConfig that carry a model_name attribute. -# HashingVectorizerEmbeddingConfig and the bare BaseEmbedderConfig don't have -# one (sklearn vectorizer / abstract base), so we filter them out below. -_MODEL_BACKED_EMBEDDERS = ( - SentenceTransformerEmbeddingConfig, - OpenaiEmbeddingConfig, - VllmEmbeddingConfig, -) - - -def _embedder_model_name(embedder: EmbedderConfig) -> str | None: - """Return the embedder's model_name when the config variant carries one.""" - if isinstance(embedder, _MODEL_BACKED_EMBEDDERS): - return embedder.model_name - return None - - -# Maps each fine-tunable transformer module to its training-mode label. -# Modules not listed (or listed as "inference") run the encoder forward-only. -# Note: dnnc keeps the cross-encoder frozen and trains an sklearn LogisticRegressionCV -# head on top of its features (see autointent._wrappers.ranker.Ranker._fit), so the -# encoder's VRAM profile matches inference rather than fine-tuning. -_TRANSFORMER_TRAINING_MODE = { - "bert": "full-finetune", - "ptuning": "lora", - "lora": "lora", -} - -# Fallback max_length when the search-space entry doesn't pin it. Used both as -# the default in _vram_for_transformer and in the entry-walk seq_len resolution. -_DEFAULT_SEQ_LEN = 128 - -# Coefficients for the linear / catboost time formulas (proposal §"Algorithm"). -_LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-8 -_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 -_CATBOOST_GPU_SPEEDUP = 10.0 -# LogisticRegressionCV defaults: Cs=10, cv=3 -> 31 inner fits + 1 final refit. -_LOGREG_CV_MULTIPLIER = 31 -_CATBOOST_DEFAULT_BINS = 254 -# Bytes per histogram bucket / tree node — order-of-magnitude constants. -_CATBOOST_BYTES_PER_TREE_NODE = 32 - - -def _extract_model_names(module_entry: dict[str, Any]) -> list[str]: - """Pull model name(s) from a search-space module entry.""" - candidates: list[str] = [] - cfg = module_entry.get("classification_model_config") - if isinstance(cfg, list): - candidates.extend(c["model_name"] for c in cfg if isinstance(c, dict) and c.get("model_name")) - elif isinstance(cfg, dict) and cfg.get("model_name"): - candidates.append(cfg["model_name"]) - embedder_cfg = module_entry.get("embedder_config") - if isinstance(embedder_cfg, list): - candidates.extend(c["model_name"] for c in embedder_cfg if isinstance(c, dict) and c.get("model_name")) - elif isinstance(embedder_cfg, dict) and embedder_cfg.get("model_name"): - candidates.append(embedder_cfg["model_name"]) - return candidates - - -def _max_int(value: Any, default: int) -> int: # noqa: ANN401 - if value is None: - return default - if isinstance(value, list) and value: - return max(int(x) for x in value) - if isinstance(value, dict): - return int(value.get("high", default)) - try: - return int(value) - except (TypeError, ValueError): - return default - - -def _walk_modules_indexed( - search_space: list[dict[str, Any]], -) -> Iterable[tuple[int, str, dict[str, Any]]]: - """Yield (node_index, node_type, module_entry) — index lets us bound per-node max cost.""" - for node_idx, node in enumerate(search_space or []): - node_type = node.get("node_type", "?") - for entry in node.get("search_space", []) or []: - yield node_idx, node_type, entry - - -def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dict[str, Any]]]: - """Yield (node_type, module_entry) pairs — index-agnostic view over `_walk_modules_indexed`.""" - for _, node_type, entry in _walk_modules_indexed(search_space): - yield node_type, entry - - -def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: - """Weight-side VRAM in GB — weights + grads + Adam optimizer state. Excludes activations. - - Modes: - * ``inference``: forward only — weights + ~30% intermediate-tensor overhead. - * ``lora``: frozen base + small trainable adapters + their grads/optimizer (~0.5 GB). - * ``full-finetune`` (default): weights + grads + Adam (m, v) = 4x weights. - """ - weights_gb = meta.weights_gb - if mode == "inference": - return weights_gb * 1.3 - if mode == "lora": - return weights_gb * 1.3 + 0.5 - return weights_gb * 4.0 - - -def _vram_for_transformer( - meta: ModelMeta, - mode: str, - mixed_precision: bool, - *, - batch_size: int = 0, - seq_len: int = _DEFAULT_SEQ_LEN, -) -> float: - """Total VRAM in GB: weights + grads + optimizer state + activations x batch. - - Activation accounting differs by mode — training keeps per-layer outputs for - backward; inference only needs one or two layers in flight. - """ - base = _weights_vram_for_transformer(meta, mode) - if batch_size <= 0: - return base - per_sample = _activations_gb_per_sample( - meta, seq_len, mixed_precision=mixed_precision, is_training=mode != "inference" - ) - return base + per_sample * batch_size - - -def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: - """RAM in GB. Loose upper bound.""" - return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB - - -def _floor_to_power_of_two(n: int) -> int: - """Largest power of two ≤ n; returns 0 when n < 1.""" - if n < 1: - return 0 - power = 1 - while power * 2 <= n: - power *= 2 - return power - - -def _n_layers(meta: ModelMeta | None) -> int: - """Layer count from the model's ``config.json``; falls back to BERT-base when absent.""" - if meta is not None and meta.n_layers is not None: - return meta.n_layers - return _DEFAULT_LAYERS - - -def _activations_gb_per_sample( - meta: ModelMeta | None, - seq_len: int, - *, - mixed_precision: bool, - is_training: bool, -) -> float: - """Heuristic activation memory per sample. - - Training: ``seq_len x hidden x layers x const`` — per-layer outputs are kept - for backward. - Inference: ``seq_len x hidden x const`` — only one or two layers' outputs in - flight at once. - Mixed precision halves activation bytes. - """ - hidden = _embedder_dim(meta) - # Training keeps every layer's outputs for backward -> scales x n_layers. - # The 16-byte/token/layer coefficient bundles fp32 activation + ~4x backward overhead. - # Inference only holds ~1-2 layers' outputs in flight at once. - bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 if is_training else seq_len * hidden * 8 - if mixed_precision: - bytes_per_sample //= 2 - return bytes_per_sample / _BYTES_PER_GB - - -def _max_fitting_batch_size( - *, - weight_vram_gb: float, - vram_budget_gb: float, - per_sample_gb: float, -) -> int: - """Largest batch that keeps total VRAM under the AMPLE/TIGHT threshold. - - Returns 0 when even the weights blow the budget. Result is rounded down to - the nearest power of two. - """ - if per_sample_gb <= 0: - return 0 - target_vram = vram_budget_gb * _TIGHT_RATIO - available_for_activations = target_vram - weight_vram_gb - if available_for_activations <= 0: - return 0 - return _floor_to_power_of_two(int(available_for_activations / per_sample_gb)) - - -def _embedder_dim(meta: ModelMeta | None) -> int: - """Hidden size from the model's ``config.json``; falls back to BERT-base when absent.""" - if meta is not None and meta.hidden_size is not None: - return meta.hidden_size - return _DEFAULT_HIDDEN - - -def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: - if not seen_models: - return None - return max(seen_models.values(), key=lambda m: m.total_params) - - -def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: - """Float64 design matrix dominates; coefficients and L-BFGS history are small.""" - data_bytes = 8.0 * stats.n_samples * embedder_dim - coef_bytes = 8.0 * max(1, stats.n_classes) * embedder_dim - lbfgs_bytes = 10.0 * 8.0 * embedder_dim - return (data_bytes + coef_bytes + lbfgs_bytes) / _BYTES_PER_GB - - -def _time_for_linear( - *, - n_trials: int, - n_samples: int, - embedder_dim: int, - max_iter: int, - cv_multiplier: int, - class_multiplier: int, -) -> float: - seconds = ( - n_trials - * _LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER - * n_samples - * embedder_dim - * max_iter - * cv_multiplier - * class_multiplier - ) - return seconds / 3600.0 - - -def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, depth: int) -> float: - data_bytes = 4.0 * stats.n_samples * n_features - histograms_bytes = 4.0 * n_features * _CATBOOST_DEFAULT_BINS - trees_bytes = iterations * (2**depth) * _CATBOOST_BYTES_PER_TREE_NODE - return float((data_bytes + histograms_bytes + trees_bytes) / _BYTES_PER_GB) - - -def _time_for_catboost( - *, - n_trials: int, - n_samples: int, - n_features: int, - iterations: int, - depth: int, - class_multiplier: int, - on_gpu: bool, -) -> float: - coeff = _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER - if on_gpu: - coeff /= _CATBOOST_GPU_SPEEDUP - seconds = n_trials * iterations * coeff * n_samples * n_features * depth * class_multiplier - return seconds / 3600.0 - - -def _time_for_transformer( - *, - n_trials: int, - epochs: int, - batch_size: int, - n_samples: int, -) -> float: - """Transformer training time in hours, assuming a flat 1 second per step. - - The advisor has no real wall-time calibration across hardware tiers / model - sizes, so the report uses ``time_hours`` as a step-count proxy rather than - pretending to estimate seconds. Users should treat the number as ordering / - ballpark information, not a budget. - """ - steps = max(1, (n_samples // max(1, batch_size))) * epochs - return (n_trials * steps) / 3600.0 - - -def _classify_severity(estimate: float, budget: float) -> Severity: - if estimate <= 0: - return Severity.AMPLE - if budget <= 0: - return Severity.TIGHT - ratio = estimate / budget - if ratio >= 1: - return Severity.OVER - if ratio >= _TIGHT_RATIO: - return Severity.TIGHT - return Severity.AMPLE - - -@dataclass -class _ModuleEstimate: - """Per-module cost contribution + the dict that gets rendered in the report.""" - - driver: dict[str, Any] - vram_gb: float - ram_gb: float - time_hours: float - model_weights_gb: float = 0.0 - - -def _refit_factor(*, refit_after: bool, n_trials: int) -> float: - """Wall-time multiplier for ``refit_after=True`` (amortized 1/n_trials extra).""" - return 1 + 1.0 / max(1, n_trials) if refit_after else 1.0 - - -def _split_entries( - search_space: list[dict[str, Any]], -) -> tuple[list[tuple[int, str, dict[str, Any]]], list[tuple[int, str, dict[str, Any]]]]: - """Partition search-space entries into (transformer-bearing, classic).""" - transformer: list[tuple[int, str, dict[str, Any]]] = [] - classic: list[tuple[int, str, dict[str, Any]]] = [] - for node_idx, node_type, entry in _walk_modules_indexed(search_space): - bucket = classic if entry.get("module_name") in {"linear", "catboost"} else transformer - bucket.append((node_idx, node_type, entry)) - return transformer, classic - - -def _estimate_transformer_model( - *, - meta: ModelMeta, - entry: dict[str, Any], - node_type: str, - module: str, - name: str, - stats: DatasetStats, - hardware: HardwareProfile, - n_trials: int, - refit_after: bool, -) -> _ModuleEstimate: - """One row of cost for a transformer module + a specific model checkpoint.""" - mixed_precision = entry.get("dtype") in {"fp16", "bf16"} - mode = _TRANSFORMER_TRAINING_MODE.get(module, "inference") - batch_size = _max_int(entry.get("batch_size"), 32) - epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) - seq_len = _max_int(entry.get("max_length"), _DEFAULT_SEQ_LEN) - - vram = _vram_for_transformer(meta, mode, mixed_precision, batch_size=batch_size, seq_len=seq_len) - ram = _ram_for_module(meta, stats) - - driver_max_batch: int | None = None - if hardware.vram_gb > 0: - driver_max_batch = _max_fitting_batch_size( - weight_vram_gb=_weights_vram_for_transformer(meta, mode), - vram_budget_gb=hardware.vram_gb, - per_sample_gb=_activations_gb_per_sample( - meta, seq_len, mixed_precision=mixed_precision, is_training=mode != "inference" - ), - ) - - time_h = _time_for_transformer( - n_trials=n_trials, - epochs=epochs, - batch_size=batch_size, - n_samples=stats.n_samples, - ) - if mode != "inference": - time_h *= _refit_factor(refit_after=refit_after, n_trials=n_trials) - - return _ModuleEstimate( - driver={ - "node_type": node_type, - "module": module, - "model": name, - "mode": mode, - "vram_gb": round(vram, 2), - "ram_gb": round(ram, 2), - "time_hours": round(time_h, 2), - "batch_size": batch_size, - "max_batch_size": driver_max_batch, - "confidence": meta.confidence, - }, - vram_gb=vram, - ram_gb=ram, - time_hours=time_h, - model_weights_gb=meta.weights_gb, - ) - - -def _estimate_classic_entry( - *, - entry: dict[str, Any], - node_type: str, - embedder_meta: ModelMeta | None, - embedder_dim: int, - stats: DatasetStats, - hardware: HardwareProfile, - n_trials: int, - refit_after: bool, -) -> _ModuleEstimate | None: - """Cost row for a linear or catboost scorer (returns ``None`` for any other module).""" - module = entry.get("module_name", "?") - refit = _refit_factor(refit_after=refit_after, n_trials=n_trials) - # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes. - class_multiplier = max(1, stats.n_classes) - - if module == "linear": - cv_multiplier = 1 if stats.multilabel else _LOGREG_CV_MULTIPLIER - ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) - time_h = ( - _time_for_linear( - n_trials=n_trials, - n_samples=stats.n_samples, - embedder_dim=embedder_dim, - max_iter=_max_int(entry.get("max_iter"), 100), - cv_multiplier=cv_multiplier, - class_multiplier=class_multiplier, - ) - * refit - ) - vram = 0.0 - mode = "linear-cv" if cv_multiplier > 1 else "linear" - elif module == "catboost": - on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" - # CatBoost MultiClass loss grows per-class trees only above binary; binary uses - # Logloss with one tree per iteration. - cb_class_mult = class_multiplier if stats.n_classes > _MULTICLASS_THRESHOLD or stats.multilabel else 1 - iterations = _max_int(entry.get("iterations"), 1000) - depth = _max_int(entry.get("depth"), 6) - ram_total = _ram_for_catboost(stats=stats, n_features=embedder_dim, iterations=iterations, depth=depth) - time_h = ( - _time_for_catboost( - n_trials=n_trials, - n_samples=stats.n_samples, - n_features=embedder_dim, - iterations=iterations, - depth=depth, - class_multiplier=cb_class_mult, - on_gpu=on_gpu, - ) - * refit - ) - vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) - mode = "catboost-gpu" if on_gpu else "catboost" - else: - return None - - return _ModuleEstimate( - driver={ - "node_type": node_type, - "module": module, - "model": embedder_meta.name if embedder_meta else "(no embedder)", - "mode": mode, - "vram_gb": round(vram, 2), - "ram_gb": round(ram, 2), - "time_hours": round(time_h, 2), - "batch_size": None, - "max_batch_size": None, - "confidence": embedder_meta.confidence if embedder_meta else "heuristic", - }, - vram_gb=vram, - ram_gb=ram, - time_hours=time_h, - ) - - -def _aggregate_disk( - estimate: ResourceEstimate, - seen_models: dict[str, ModelMeta], - node_max_weights: dict[int, float], - *, - dump_modules: bool, - n_trials: int, -) -> None: - """Fold per-model download/cached sizes into ``estimate`` and apply dump-modules accounting.""" - for meta in seen_models.values(): - if meta.cached_locally: - estimate.disk_cached_gb += meta.disk_gb - else: - estimate.disk_download_gb += meta.disk_gb - if dump_modules: - # Each trial selects one variant per node, so per-trial dumped weights - # are bounded by the heaviest module in each node, summed across nodes. - estimate.disk_dump_gb = sum(node_max_weights.values()) * n_trials - - -def _emit_resource_findings( - report: PreflightReport, - estimate: ResourceEstimate, - hardware: HardwareProfile, - *, - n_jobs: int, -) -> None: - """Translate aggregated estimates into VRAM/RAM/disk/time findings on the report.""" - parallel_gpu = n_jobs > 1 and hardware.accelerator in {"cuda", "mps"} - effective_vram = estimate.vram_gb * n_jobs if parallel_gpu else estimate.vram_gb - # MPS shares one unified pool: parallel workers each allocate weights+activations - # in RAM, so peak RAM also scales with n_jobs on Apple Silicon. - effective_ram = estimate.ram_gb * n_jobs if n_jobs > 1 and hardware.accelerator == "mps" else estimate.ram_gb - - if hardware.accelerator == "cpu" and effective_vram > 0: - report.add( - "resource", - Severity.TIGHT, - f"No GPU detected; transformer modules will be very slow (worst case ~{estimate.time_hours:.1f} h).", - metric="vram", - ) - else: - msg = f"VRAM ~{effective_vram:.1f} GB" - if n_jobs > 1: - msg += f" (= per-trial {estimate.vram_gb:.1f} GB × {n_jobs} parallel trials)" - msg += f" vs available {hardware.vram_gb:.1f} GB" - report.add("resource", _classify_severity(effective_vram, hardware.vram_gb), msg, metric="vram") - - report.add( - "resource", - _classify_severity(effective_ram, hardware.ram_gb), - f"RAM ~{effective_ram:.1f} GB vs available {hardware.ram_gb:.1f} GB", - metric="ram", - ) - - disk_total = estimate.disk_download_gb + estimate.disk_dump_gb - disk_msg = f"Disk ~{estimate.disk_download_gb:.1f} GB to download" - if estimate.disk_cached_gb > 0: - disk_msg += f", {estimate.disk_cached_gb:.1f} GB already cached" - if estimate.disk_dump_gb > 0: - disk_msg += f", +{estimate.disk_dump_gb:.1f} GB during training (dump_modules=True)" - disk_msg += f" vs {hardware.free_disk_gb:.0f} GB free" - report.add("resource", _classify_severity(disk_total, hardware.free_disk_gb), disk_msg, metric="disk") - - if estimate.time_hours > 0: - report.add( - "resource", - Severity.AMPLE, - f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)", - metric="time", - ) - - -def _resource_phase( - config: dict[str, Any], - stats: DatasetStats, - hardware: HardwareProfile, - report: PreflightReport, - *, - refit_after: bool = False, -) -> None: - cfg = _validated_config(config) - n_trials = cfg.hpo_config.n_trials - n_jobs = cfg.hpo_config.n_jobs - dump_modules = cfg.logging_config.dump_modules - - seen_models: dict[str, ModelMeta] = {} - global_embedder = _embedder_model_name(cfg.embedder_config) - if global_embedder: - seen_models[global_embedder] = resolve_model(global_embedder) - - transformer_entries, classic_entries = _split_entries(cfg.search_space) - - # First pass: transformer modules (also populates seen_models for the classic pass). - module_estimates: list[_ModuleEstimate] = [] - node_max_weights: dict[int, float] = {} - for node_idx, node_type, entry in transformer_entries: - module = entry.get("module_name", "?") - model_names = _extract_model_names(entry) - if not model_names and global_embedder and module in {"knn", "mlknn"}: - model_names = [global_embedder] - for name in model_names: - meta = seen_models.setdefault(name, resolve_model(name)) - me = _estimate_transformer_model( - meta=meta, - entry=entry, - node_type=node_type, - module=module, - name=name, - stats=stats, - hardware=hardware, - n_trials=n_trials, - refit_after=refit_after, - ) - module_estimates.append(me) - # Track heaviest weight per node so dump_modules is bounded by one - # selected variant per node x n_trials, not the sum of all candidates. - node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), me.model_weights_gb) - - # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. - embedder_meta = _largest_embedder(seen_models) - embedder_dim = _embedder_dim(embedder_meta) - for _, node_type, entry in classic_entries: - classic_estimate = _estimate_classic_entry( - entry=entry, - node_type=node_type, - embedder_meta=embedder_meta, - embedder_dim=embedder_dim, - stats=stats, - hardware=hardware, - n_trials=n_trials, - refit_after=refit_after, - ) - if classic_estimate is not None: - module_estimates.append(classic_estimate) - - estimate = ResourceEstimate(parallel_factor=n_jobs) - for me in module_estimates: - estimate.vram_gb = max(estimate.vram_gb, me.vram_gb) - estimate.ram_gb = max(estimate.ram_gb, me.ram_gb) - estimate.time_hours += me.time_hours - estimate.drivers.append(me.driver) - - _aggregate_disk(estimate, seen_models, node_max_weights, dump_modules=dump_modules, n_trials=n_trials) - - # Flip low_confidence if any model fell back to the heuristic path (Hub - # unreachable, repo missing safetensors metadata, local-path checkpoint). - heuristic_models = [m.name for m in seen_models.values() if m.confidence == "heuristic"] - if heuristic_models: - report.low_confidence = True - report.notes.append( - f"Heuristic fallback used for {len(heuristic_models)} model(s) — sizes are BERT-base " - f"defaults: {', '.join(heuristic_models[:3])}{'...' if len(heuristic_models) > 3 else ''}", # noqa: PLR2004 - ) - - report.resource = estimate - _emit_resource_findings(report, estimate, hardware, n_jobs=n_jobs) - - -def _config_phase( - config: dict[str, Any], - hardware: HardwareProfile, - report: PreflightReport, -) -> None: - hpo = config.get("hpo_config") or {} - n_jobs = int(hpo.get("n_jobs", 1)) - - if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: - report.add( - "config", - Severity.TIGHT, - f"hpo_config.n_jobs={n_jobs} on a single GPU multiplies VRAM demand by {n_jobs}×.", - ) - - uses_catboost_gpu = False - for _, entry in _walk_modules(config.get("search_space") or []): - if entry.get("module_name") == "catboost" and entry.get("task_type") == "GPU": - uses_catboost_gpu = True - break - if uses_catboost_gpu and hardware.accelerator != "cuda": - report.add( - "config", - Severity.TIGHT, - "CatBoost task_type=GPU configured but no CUDA detected — will fall back to CPU.", - ) - - -def _data_phase( - config: dict[str, Any], - stats: DatasetStats, - report: PreflightReport, -) -> None: - # token-length truncation (heuristic — we use stats.p95_tokens vs configured max_length) - p95 = stats.p95_tokens or int(stats.avg_tokens * 2.5) - for _, entry in _walk_modules(config.get("search_space") or []): - max_len_value = entry.get("max_length") - if max_len_value is None: - continue - max_len = _max_int(max_len_value, 512) - if p95 > max_len: - severity = Severity.OVER if p95 > max_len * 1.5 else Severity.TIGHT - module_name = entry.get("module_name", "?") - report.add( - "data", - severity, - f"Train tokens p95~{p95} exceeds {module_name}.max_length={max_len}; expect silent truncation.", - ) - - # rare class x linear-CV (LogisticRegressionCV cv=3 needs >=3 samples/class; - # multilabel path uses one-vs-rest without CV so the failure can't occur there) - has_linear = any(e.get("module_name") == "linear" for _, e in _walk_modules(config.get("search_space") or [])) - if has_linear and stats.rare_classes and not stats.multilabel: - report.add( - "data", - Severity.OVER, - (f"LogisticRegressionCV (cv=3) will fail: classes {stats.rare_classes[:5]} have <3 samples."), - ) - - # partial descriptions x description scorer - description_modules = {"description_bi", "description_cross", "description_llm"} - has_description = any( - e.get("module_name") in description_modules for _, e in _walk_modules(config.get("search_space") or []) - ) - if has_description and stats.has_descriptions is False: - report.add( - "data", - Severity.OVER, - "description scorer present but intent descriptions are missing — fill them in or drop the scorer.", - ) - - -def run_preflight( - config: dict[str, Any], - stats: DatasetStats, - hardware: HardwareProfile, - *, - preset_name: str | None = None, - refit_after: bool = False, -) -> PreflightReport: - """Run all three phases and return one report. - - Args: - config: parsed preset / OptimizationConfig dict (top-level keys: - ``search_space``, ``hpo_config``, optional ``embedder_config``, - optional ``logging_config.dump_modules``). - stats: dataset statistics (real or placeholder). - hardware: detected hardware profile. - preset_name: optional friendly name for the report header. - refit_after: matches the ``Pipeline.fit(refit_after=...)`` argument. - When True, time estimates include the extra refit-on-full-data pass. - - Returns: - PreflightReport with findings across resource/data/config phases. - """ - report = PreflightReport( - preset_name=preset_name, - hardware={ - "accelerator": hardware.accelerator, - "device_name": hardware.device_name, - "vram_gb": round(hardware.vram_gb, 2), - "ram_gb": round(hardware.ram_gb, 2), - "free_disk_gb": round(hardware.free_disk_gb, 2), - "device_class": hardware.device_class, - }, - dataset={ - "n_samples": stats.n_samples, - "n_classes": stats.n_classes, - "avg_tokens": stats.avg_tokens, - "p95_tokens": stats.p95_tokens, - "multilabel": stats.multilabel, - "source": stats.source, - }, - ) - report.notes.extend(hardware.notes) - - _resource_phase(config, stats, hardware, report, refit_after=refit_after) - _data_phase(config, stats, report) - _config_phase(config, hardware, report) - - return report diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index a5b6126a5..677e6045c 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -94,7 +94,7 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: return None # Bytes-per-element for safetensors dtype strings. Used to convert the per-dtype # parameter counts (info.safetensors.parameters) into a weighted average - # bytes-per-param for mixed-precision repos. + # bytes-per-param when a checkpoint stores tensors in multiple dtypes. _dtype_bytes: dict[str, int] = { "F64": 8, "F32": 4, diff --git a/src/autointent/_advisor/_workflows.py b/src/autointent/_advisor/workflows.py similarity index 92% rename from src/autointent/_advisor/_workflows.py rename to src/autointent/_advisor/workflows.py index 0bd7ee5bf..2331e225a 100644 --- a/src/autointent/_advisor/_workflows.py +++ b/src/autointent/_advisor/workflows.py @@ -20,9 +20,9 @@ from autointent.custom_types import SearchSpacePreset from autointent.utils import load_preset -from ._estimates import run_preflight from ._hardware import detect_hardware from ._report import DatasetStats, RecommendationResult, Severity +from .runner import run_preflight if TYPE_CHECKING: from collections.abc import Iterable @@ -196,10 +196,11 @@ def recommend( ``RecommendationResult`` with the chosen preset name and full results list. Note: - Among feasible presets we pick the one with the largest estimated - ``time_hours`` (ties broken alphabetically). Higher-quality presets cost - more wall-time, so the slowest feasible preset is also the heaviest - preset that still fits the hardware — i.e. "use what you have". + Among feasible presets we pick the heaviest one that still fits the + hardware budget — "use what you have" semantics. This is a *cost* + ranking, not a quality ranking: a heavier preset is not strictly better + and may overfit on small datasets where a classic-* preset would win on + accuracy. Override ``presets=`` if you want a different ranking. """ hardware = detect_hardware(vram_budget_gb=budget_vram_gb) stats = stats or DatasetStats.placeholder() @@ -221,11 +222,9 @@ def recommend( ) results.append((preset, report)) - # Rank by Literal position (lower index = higher quality); presets the user - # passed via the ``presets`` override but not in BUNDLED_PRESETS sort last. - quality_rank = {name: i for i, name in enumerate(BUNDLED_PRESETS)} + cost_rank = {name: i for i, name in enumerate(BUNDLED_PRESETS)} feasible = [(name, r) for name, r in results if r.is_feasible] - feasible.sort(key=lambda pair: (quality_rank.get(pair[0], len(BUNDLED_PRESETS)), pair[0])) + feasible.sort(key=lambda pair: (cost_rank.get(pair[0], len(BUNDLED_PRESETS)), pair[0])) chosen = feasible[0][0] if feasible else None return RecommendationResult(chosen=chosen, results=results) diff --git a/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index a54da368d..59e6b87a3 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -128,10 +128,14 @@ class Split: "zero-shot-encoders", "classic-light", ] -"""Bundled search-space presets, listed in descending quality order. - -The order is consumed by ``autointent._advisor.recommend`` to pick the -highest-quality feasible preset (lower index = higher quality).""" +"""Bundled search-space presets, listed in descending resource-cost order. + +Heavier presets explore more / larger models and take longer to run. The order +is a cost ranking, **not** a quality ranking: a heavier preset is not strictly +better — e.g. ``transformers-heavy`` will overfit on tiny datasets where a +classic-* preset wins on accuracy. ``autointent._advisor.recommend`` uses this +ordering to pick the heaviest preset that still fits the hardware budget, +which is a reasonable default but not always the right choice for the data.""" class Document(BaseModel): diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 3fa293b7e..3d4e4e526 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -138,18 +138,18 @@ def meta(self) -> ModelMeta: ) def test_full_finetune_is_larger_than_lora_is_larger_than_inference(self, meta: ModelMeta) -> None: - inference = _vram_for_transformer(meta, "inference", mixed_precision=False) - lora = _vram_for_transformer(meta, "lora", mixed_precision=False) - full = _vram_for_transformer(meta, "full-finetune", mixed_precision=False) + inference = _vram_for_transformer(meta, "inference") + lora = _vram_for_transformer(meta, "lora") + full = _vram_for_transformer(meta, "full-finetune") assert inference < lora < full def test_inference_activations_are_smaller_than_training(self, meta: ModelMeta) -> None: """Inference doesn't store per-layer outputs for backward — activation memory should be many times smaller than training at the same batch_size.""" - train_total = _vram_for_transformer(meta, "full-finetune", False, batch_size=64, seq_len=128) - train_weights = _vram_for_transformer(meta, "full-finetune", False, batch_size=0) - inf_total = _vram_for_transformer(meta, "inference", False, batch_size=64, seq_len=128) - inf_weights = _vram_for_transformer(meta, "inference", False, batch_size=0) + train_total = _vram_for_transformer(meta, "full-finetune", batch_size=64, seq_len=128) + train_weights = _vram_for_transformer(meta, "full-finetune", batch_size=0) + inf_total = _vram_for_transformer(meta, "inference", batch_size=64, seq_len=128) + inf_weights = _vram_for_transformer(meta, "inference", batch_size=0) train_acts = train_total - train_weights inf_acts = inf_total - inf_weights assert inf_acts > 0 @@ -157,21 +157,6 @@ def test_inference_activations_are_smaller_than_training(self, meta: ModelMeta) # 12-layer model: training activations should be at least ~5x inference. assert train_acts / inf_acts > 5 - def test_amp_does_not_reduce_weight_side_vram(self, meta: ModelMeta) -> None: - """Weight-side AMP accounting: fp16 weights+grads (W) + fp32 master copy (W) - + fp32 Adam moments (2W) = 4W, identical to pure fp32. AMP's savings live - in activations, not the optimizer.""" - full_fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False, batch_size=0) - full_amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True, batch_size=0) - assert full_amp == pytest.approx(full_fp32) - - def test_amp_does_reduce_activation_side_vram(self, meta: ModelMeta) -> None: - """When a batch is configured, AMP halves activation bytes — total VRAM - with batch should be strictly smaller under AMP than fp32.""" - fp32 = _vram_for_transformer(meta, "full-finetune", mixed_precision=False, batch_size=64, seq_len=128) - amp = _vram_for_transformer(meta, "full-finetune", mixed_precision=True, batch_size=64, seq_len=128) - assert amp < fp32 - def test_ram_scales_with_dataset_size() -> None: meta = ModelMeta( @@ -497,13 +482,13 @@ def test_driver_records_current_and_max_batch(self) -> None: report = run_preflight( self._bert_cfg("microsoft/deberta-v3-large", batch_size=64), DatasetStats.placeholder(), - _profile(vram_gb=6.5), + _profile(vram_gb=7.5), ) drivers = [d for d in report.resource.drivers if d["module"] == "bert"] assert drivers d = drivers[0] assert d["batch_size"] == 64 - # vram_gb=6.5 against ~5 GB weights x 0.9 tight ratio -> little activation room, max < 64. + # vram_gb=7.5 against ~5.9 GB weights x 0.9 tight ratio -> little activation room, max < 64. assert d["max_batch_size"] is not None assert 0 < d["max_batch_size"] < 64 From 1841d0896a442d2ac22fea12975f1f4d60f1df1c Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Wed, 24 Jun 2026 15:38:37 +0300 Subject: [PATCH 17/43] commit missing files --- .../_advisor/_estimates/__init__.py | 0 .../_advisor/_estimates/_formulas.py | 293 +++++++++++++ .../_advisor/_estimates/_resource.py | 406 ++++++++++++++++++ .../_advisor/_estimates/_search_space.py | 75 ++++ src/autointent/_advisor/runner.py | 170 ++++++++ 5 files changed, 944 insertions(+) create mode 100644 src/autointent/_advisor/_estimates/__init__.py create mode 100644 src/autointent/_advisor/_estimates/_formulas.py create mode 100644 src/autointent/_advisor/_estimates/_resource.py create mode 100644 src/autointent/_advisor/_estimates/_search_space.py create mode 100644 src/autointent/_advisor/runner.py diff --git a/src/autointent/_advisor/_estimates/__init__.py b/src/autointent/_advisor/_estimates/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/autointent/_advisor/_estimates/_formulas.py b/src/autointent/_advisor/_estimates/_formulas.py new file mode 100644 index 000000000..8bbbdcebf --- /dev/null +++ b/src/autointent/_advisor/_estimates/_formulas.py @@ -0,0 +1,293 @@ +"""Pure cost-estimate formulas — VRAM, RAM, time, severity, model shape. + +No I/O, no logging, no orchestration. Each formula docstring links to the +reference it was calibrated against so a reviewer can follow each coefficient +back to its source. + +Conventions: + * All ``*_gb`` results use the binary GiB convention (1024**3 bytes per GB) — + matches the rest of the advisor's byte->GB conversions. + * All ``*_hours`` results assume the GPU baseline of ~1 second per step; + CPU runs pay a flat slowdown factor (see ``_time_for_transformer``). + * "fp32 worst case" — we deliberately ignore lower-precision / FlashAttention / + quantization optimizations, per the advisor's "pessimistic upper bound" contract. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from autointent._advisor._report import Severity + +if TYPE_CHECKING: + from autointent._advisor._hub import ModelMeta + from autointent._advisor._report import DatasetStats + + +_BYTES_PER_GB = 1024**3 +_DEFAULT_SEQ_LEN = 128 + +# Fallback architecture shape (BERT-base) used only when the model's actual +# config.json couldn't be fetched from HF Hub — see _hub._shape_from_config. +_DEFAULT_HIDDEN = 768 +_DEFAULT_LAYERS = 12 + +_TIGHT_RATIO = 0.9 +_MULTICLASS_THRESHOLD = 2 + + +def _classify_severity(estimate: float, budget: float) -> Severity: + """Map a ``(estimate, budget)`` pair onto a Severity bucket. + + * AMPLE: ``estimate <= 0`` OR ``ratio < _TIGHT_RATIO`` + * TIGHT: ``budget <= 0`` OR ``_TIGHT_RATIO <= ratio < 1`` + * OVER: ``ratio >= 1`` + """ + if estimate <= 0: + return Severity.AMPLE + if budget <= 0: + return Severity.TIGHT + ratio = estimate / budget + if ratio >= 1: + return Severity.OVER + if ratio >= _TIGHT_RATIO: + return Severity.TIGHT + return Severity.AMPLE + + +def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: + """Weight-side VRAM in GB — weights + grads + optimizer state. Excludes activations. + + Returns a deliberately pessimistic upper bound, matching the advisor's + "heuristic upper bound, not measurement" contract. + + Modes: + * ``inference``: forward only — weights + ~30% intermediate-tensor overhead. + * ``lora``: frozen base + small trainable adapters + their grads/optimizer (~0.5 GB). + * ``full-finetune`` (default): the textbook 4W (weights + grads + Adam m + Adam v). + We use 4.5W to leave headroom for loss-scale buffers, allocator fragmentation, + cuDNN workspaces, and gradient-accumulation buffers — none of which the textbook + 4W accounting captures. + """ + weights_gb = meta.weights_gb + if mode == "inference": + return weights_gb * 1.3 + if mode == "lora": + return weights_gb * 1.3 + 0.5 + return weights_gb * 4.5 + + +def _activations_gb_per_sample( + meta: ModelMeta | None, + seq_len: int, + *, + is_training: bool, +) -> float: + """Heuristic activation memory per sample, assuming a fp32 worst case. + + Training: ``seq_len x hidden x layers x const`` — per-layer outputs are kept + for backward. + Inference: ``seq_len x hidden x const`` — only one or two layers' outputs in + flight at once. + """ + hidden = _embedder_dim(meta) + # Training keeps every layer's outputs for backward -> scales x n_layers. + # 16 bytes/token/layer ~ fp32 activation (4B) x ~4x backward overhead (Korthikanti et al.). + # Inference only holds ~1-2 layers' outputs in flight at once. + bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 if is_training else seq_len * hidden * 8 + return bytes_per_sample / _BYTES_PER_GB + + +def _vram_for_transformer( + meta: ModelMeta, + mode: str, + *, + batch_size: int = 0, + seq_len: int = _DEFAULT_SEQ_LEN, +) -> float: + """Total VRAM in GB: weights + grads + optimizer state + activations x batch. + + Activation accounting differs by mode — training keeps per-layer outputs for + backward; inference only needs one or two layers in flight. + """ + base = _weights_vram_for_transformer(meta, mode) + if batch_size <= 0: + return base + per_sample = _activations_gb_per_sample(meta, seq_len, is_training=mode != "inference") + return base + per_sample * batch_size + + +def _max_fitting_batch_size( + *, + weight_vram_gb: float, + vram_budget_gb: float, + per_sample_gb: float, +) -> int: + """Largest batch that keeps total VRAM under the AMPLE/TIGHT threshold. + + Returns 0 when even the weights blow the budget. Result is rounded down to + the nearest power of two + """ + if per_sample_gb <= 0: + return 0 + target_vram = vram_budget_gb * _TIGHT_RATIO + available_for_activations = target_vram - weight_vram_gb + if available_for_activations <= 0: + return 0 + return _floor_to_power_of_two(int(available_for_activations / per_sample_gb)) + + +_CPU_SLOWDOWN_FACTOR = 50.0 +"""Rough multiplier for transformer training on CPU vs. a modern GPU. + +Real benchmarks vary widely (30x for small BERTs on AVX-512 boxes, 100x+ for +billion-scale models on a stock laptop). A single 50x constant is a pessimistic +upper bound that's good enough to make the CPU/GPU distinction visible without +re-introducing the per-device tier table.""" + + +def _time_for_transformer( + *, + n_trials: int, + epochs: int, + batch_size: int, + n_samples: int, + accelerator: str, +) -> float: + """Transformer training time in hours. + + Baseline is "1 second per step" on a GPU (CUDA / MPS) — a step-count proxy, + not a real wall-time calibration. CPU training pays a flat ``_CPU_SLOWDOWN_FACTOR`` + so the report doesn't hide the fact that the same workload is dramatically + slower without a GPU. Users should treat absolute numbers as ordering / + ballpark information, not a budget. + """ + steps = max(1, (n_samples // max(1, batch_size))) * epochs + h = (n_trials * steps) / 3600.0 + if accelerator == "cpu": + h *= _CPU_SLOWDOWN_FACTOR + return h + + +def _n_layers(meta: ModelMeta | None) -> int: + """Layer count from the model's ``config.json``; falls back to BERT-base when absent.""" + if meta is not None and meta.n_layers is not None: + return meta.n_layers + return _DEFAULT_LAYERS + + +def _embedder_dim(meta: ModelMeta | None) -> int: + """Hidden size from the model's ``config.json``; falls back to BERT-base when absent.""" + if meta is not None and meta.hidden_size is not None: + return meta.hidden_size + return _DEFAULT_HIDDEN + + +def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: + """Return the largest model in ``seen_models`` by parameter count, or None if empty.""" + if not seen_models: + return None + return max(seen_models.values(), key=lambda m: m.total_params) + + +def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: + """RAM in GB. Loose upper bound: weights + tokenized text in memory. + + Tokenized text is approximated as ``n_samples x avg_tokens x 4 bytes`` + (BPE/WordPiece token ids fit in int32). The 4 bytes/token bound is tight + enough for the report's purposes and intentionally ignores any preprocessing + artefacts (attention masks, position ids, etc.) since they're bounded by the + same factor. + """ + return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB + + +# Coefficients are dimensional (per-sample-per-feature-per-iteration seconds) +# rather than empirically tuned constants — they give relative-cost ordering +# across configurations and absolute ballpark wall-times. +_LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-8 +_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 +_CATBOOST_GPU_SPEEDUP = 10.0 +# LogisticRegressionCV defaults: Cs=10, cv=3 -> 10x3 inner fits + 1 final refit = 31. +_LOGREG_CV_MULTIPLIER = 31 +# Default value of `border_count` in CatBoost (number of histogram buckets per feature). +_CATBOOST_DEFAULT_BINS = 254 +# Bytes per histogram bucket / tree node — order-of-magnitude constant. +_CATBOOST_BYTES_PER_TREE_NODE = 32 + + +def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: + """Float64 design matrix dominates; coefficients and L-BFGS history are small.""" + data_bytes = 8.0 * stats.n_samples * embedder_dim + coef_bytes = 8.0 * max(1, stats.n_classes) * embedder_dim + lbfgs_bytes = 10.0 * 8.0 * embedder_dim + return (data_bytes + coef_bytes + lbfgs_bytes) / _BYTES_PER_GB + + +def _time_for_linear( + *, + n_trials: int, + n_samples: int, + embedder_dim: int, + max_iter: int, + cv_multiplier: int, + class_multiplier: int, +) -> float: + """LogisticRegression wall time, in hours. + + Cost is ``O(n_samples x n_features x max_iter x n_classes)`` per fit + (sklearn's L-BFGS solver), multiplied by the CV inner-fit count (31 for the + default LogisticRegressionCV). + """ + seconds = ( + n_trials + * _LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER + * n_samples + * embedder_dim + * max_iter + * cv_multiplier + * class_multiplier + ) + return seconds / 3600.0 + + +def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, depth: int) -> float: + """CatBoost RAM = quantized data matrix + histograms + tree storage.""" + data_bytes = 4.0 * stats.n_samples * n_features + histograms_bytes = 4.0 * n_features * _CATBOOST_DEFAULT_BINS + trees_bytes = iterations * (2**depth) * _CATBOOST_BYTES_PER_TREE_NODE + return float((data_bytes + histograms_bytes + trees_bytes) / _BYTES_PER_GB) + + +def _time_for_catboost( + *, + n_trials: int, + n_samples: int, + n_features: int, + iterations: int, + depth: int, + class_multiplier: int, + on_gpu: bool, +) -> float: + """CatBoost wall time, in hours. + + Cost is ``O(iterations x n_samples x n_features x depth x n_classes)`` per + fit. GPU training is ~10x faster than CPU for typical workloads per + CatBoost's published benchmarks. + https://catboost.ai/en/docs/concepts/speed-up-training + """ + coeff = _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER + if on_gpu: + coeff /= _CATBOOST_GPU_SPEEDUP + seconds = n_trials * iterations * coeff * n_samples * n_features * depth * class_multiplier + return seconds / 3600.0 + + +def _floor_to_power_of_two(n: int) -> int: + """Largest power of two <= ``n``; returns 0 when ``n < 1``.""" + if n < 1: + return 0 + power = 1 + while power * 2 <= n: + power *= 2 + return power diff --git a/src/autointent/_advisor/_estimates/_resource.py b/src/autointent/_advisor/_estimates/_resource.py new file mode 100644 index 000000000..6ed93578c --- /dev/null +++ b/src/autointent/_advisor/_estimates/_resource.py @@ -0,0 +1,406 @@ +"""Resource-phase orchestration. + +Walks the validated search space, asks ``_formulas`` for per-module costs, +aggregates them into a ``ResourceEstimate``, and emits VRAM/RAM/disk/time +findings on the report. + +The public entry is ``_resource_phase`` at the bottom; everything above it is +private machinery. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from autointent._advisor import _hub +from autointent._advisor._report import ResourceEstimate, Severity +from autointent.configs._embedder import ( + EmbedderConfig, + OpenaiEmbeddingConfig, + SentenceTransformerEmbeddingConfig, + VllmEmbeddingConfig, +) + +from ._formulas import ( + _DEFAULT_SEQ_LEN, + _LOGREG_CV_MULTIPLIER, + _MULTICLASS_THRESHOLD, + _activations_gb_per_sample, + _classify_severity, + _embedder_dim, + _largest_embedder, + _max_fitting_batch_size, + _ram_for_catboost, + _ram_for_linear, + _ram_for_module, + _time_for_catboost, + _time_for_linear, + _time_for_transformer, + _vram_for_transformer, + _weights_vram_for_transformer, +) +from ._search_space import _extract_model_names, _max_int, _walk_modules_indexed + +if TYPE_CHECKING: + from autointent._advisor._hardware import HardwareProfile + from autointent._advisor._hub import ModelMeta + from autointent._advisor._report import DatasetStats, PreflightReport + + +# Union variants of EmbedderConfig that carry a model_name attribute. +# HashingVectorizerEmbeddingConfig and the bare BaseEmbedderConfig don't have +# one (sklearn vectorizer / abstract base), so we filter them out below. +_MODEL_BACKED_EMBEDDERS = ( + SentenceTransformerEmbeddingConfig, + OpenaiEmbeddingConfig, + VllmEmbeddingConfig, +) + + +def _embedder_model_name(embedder: EmbedderConfig) -> str | None: + """Return the embedder's model_name when the config variant carries one.""" + if isinstance(embedder, _MODEL_BACKED_EMBEDDERS): + return embedder.model_name + return None + + +# Maps each fine-tunable transformer module to its training-mode label. +# Modules not listed (or listed as "inference") run the encoder forward-only. +# Note: dnnc keeps the cross-encoder frozen and trains an sklearn LogisticRegressionCV +# head on top of its features (see autointent._wrappers.ranker.Ranker._fit), so the +# encoder's VRAM profile matches inference rather than fine-tuning. +_TRANSFORMER_TRAINING_MODE = { + "bert": "full-finetune", + "ptuning": "lora", + "lora": "lora", +} + + +@dataclass +class _ModuleEstimate: + """Per-module cost contribution + the dict that gets rendered in the report.""" + + driver: dict[str, Any] + vram_gb: float + ram_gb: float + time_hours: float + model_weights_gb: float = 0.0 + + +def _refit_factor(*, refit_after: bool, n_trials: int) -> float: + """Wall-time multiplier for ``refit_after=True`` (amortized 1/n_trials extra).""" + return 1 + 1.0 / max(1, n_trials) if refit_after else 1.0 + + +def _split_entries( + search_space: list[dict[str, Any]], +) -> tuple[list[tuple[int, str, dict[str, Any]]], list[tuple[int, str, dict[str, Any]]]]: + """Partition search-space entries into (transformer-bearing, classic).""" + transformer: list[tuple[int, str, dict[str, Any]]] = [] + classic: list[tuple[int, str, dict[str, Any]]] = [] + for node_idx, node_type, entry in _walk_modules_indexed(search_space): + bucket = classic if entry.get("module_name") in {"linear", "catboost"} else transformer + bucket.append((node_idx, node_type, entry)) + return transformer, classic + + +def _estimate_transformer_model( + *, + meta: ModelMeta, + entry: dict[str, Any], + node_type: str, + module: str, + name: str, + stats: DatasetStats, + hardware: HardwareProfile, + n_trials: int, + refit_after: bool, +) -> _ModuleEstimate: + """One row of cost for a transformer module + a specific model checkpoint.""" + mode = _TRANSFORMER_TRAINING_MODE.get(module, "inference") + batch_size = _max_int(entry.get("batch_size"), 32) + epochs = _max_int(entry.get("num_train_epochs"), 1 if mode == "inference" else 10) + seq_len = _max_int(entry.get("max_length"), _DEFAULT_SEQ_LEN) + + vram = _vram_for_transformer(meta, mode, batch_size=batch_size, seq_len=seq_len) + ram = _ram_for_module(meta, stats) + + driver_max_batch: int | None = None + if hardware.vram_gb > 0: + driver_max_batch = _max_fitting_batch_size( + weight_vram_gb=_weights_vram_for_transformer(meta, mode), + vram_budget_gb=hardware.vram_gb, + per_sample_gb=_activations_gb_per_sample(meta, seq_len, is_training=mode != "inference"), + ) + + time_h = _time_for_transformer( + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + accelerator=hardware.accelerator, + ) + if mode != "inference": + time_h *= _refit_factor(refit_after=refit_after, n_trials=n_trials) + + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": name, + "mode": mode, + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "batch_size": batch_size, + "max_batch_size": driver_max_batch, + "confidence": meta.confidence, + }, + vram_gb=vram, + ram_gb=ram, + time_hours=time_h, + model_weights_gb=meta.weights_gb, + ) + + +def _estimate_classic_entry( + *, + entry: dict[str, Any], + node_type: str, + embedder_meta: ModelMeta | None, + embedder_dim: int, + stats: DatasetStats, + hardware: HardwareProfile, + n_trials: int, + refit_after: bool, +) -> _ModuleEstimate | None: + """Cost row for a linear or catboost scorer (returns ``None`` for any other module).""" + module = entry.get("module_name", "?") + refit = _refit_factor(refit_after=refit_after, n_trials=n_trials) + # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes. + class_multiplier = max(1, stats.n_classes) + + if module == "linear": + cv_multiplier = 1 if stats.multilabel else _LOGREG_CV_MULTIPLIER + ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) + time_h = ( + _time_for_linear( + n_trials=n_trials, + n_samples=stats.n_samples, + embedder_dim=embedder_dim, + max_iter=_max_int(entry.get("max_iter"), 100), + cv_multiplier=cv_multiplier, + class_multiplier=class_multiplier, + ) + * refit + ) + vram = 0.0 + mode = "linear-cv" if cv_multiplier > 1 else "linear" + elif module == "catboost": + on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" + # CatBoost MultiClass loss grows per-class trees only above binary; binary uses + # Logloss with one tree per iteration. + cb_class_mult = class_multiplier if stats.n_classes > _MULTICLASS_THRESHOLD or stats.multilabel else 1 + iterations = _max_int(entry.get("iterations"), 1000) + depth = _max_int(entry.get("depth"), 6) + ram_total = _ram_for_catboost(stats=stats, n_features=embedder_dim, iterations=iterations, depth=depth) + time_h = ( + _time_for_catboost( + n_trials=n_trials, + n_samples=stats.n_samples, + n_features=embedder_dim, + iterations=iterations, + depth=depth, + class_multiplier=cb_class_mult, + on_gpu=on_gpu, + ) + * refit + ) + vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) + mode = "catboost-gpu" if on_gpu else "catboost" + else: + return None + + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": embedder_meta.name if embedder_meta else "(no embedder)", + "mode": mode, + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "batch_size": None, + "max_batch_size": None, + "confidence": embedder_meta.confidence if embedder_meta else "heuristic", + }, + vram_gb=vram, + ram_gb=ram, + time_hours=time_h, + ) + + +def _aggregate_disk( + estimate: ResourceEstimate, + seen_models: dict[str, ModelMeta], + node_max_weights: dict[int, float], + *, + dump_modules: bool, + n_trials: int, +) -> None: + """Fold per-model download/cached sizes into ``estimate`` and apply dump-modules accounting.""" + for meta in seen_models.values(): + if meta.cached_locally: + estimate.disk_cached_gb += meta.disk_gb + else: + estimate.disk_download_gb += meta.disk_gb + if dump_modules: + # Each trial selects one variant per node, so per-trial dumped weights + # are bounded by the heaviest module in each node, summed across nodes. + estimate.disk_dump_gb = sum(node_max_weights.values()) * n_trials + + +def _emit_resource_findings( + report: PreflightReport, + estimate: ResourceEstimate, + hardware: HardwareProfile, + *, + n_jobs: int, +) -> None: + """Translate aggregated estimates into VRAM/RAM/disk/time findings on the report.""" + parallel_gpu = n_jobs > 1 and hardware.accelerator in {"cuda", "mps"} + effective_vram = estimate.vram_gb * n_jobs if parallel_gpu else estimate.vram_gb + # MPS shares one unified pool: parallel workers each allocate weights+activations + # in RAM, so peak RAM also scales with n_jobs on Apple Silicon. + effective_ram = estimate.ram_gb * n_jobs if n_jobs > 1 and hardware.accelerator == "mps" else estimate.ram_gb + + if hardware.accelerator == "cpu" and effective_vram > 0: + report.add( + "resource", + Severity.TIGHT, + f"No GPU detected; transformer modules will be very slow (worst case ~{estimate.time_hours:.1f} h).", + metric="vram", + ) + else: + msg = f"VRAM ~{effective_vram:.1f} GB" + if n_jobs > 1: + msg += f" (= per-trial {estimate.vram_gb:.1f} GB x {n_jobs} parallel trials)" + msg += f" vs available {hardware.vram_gb:.1f} GB" + report.add("resource", _classify_severity(effective_vram, hardware.vram_gb), msg, metric="vram") + + report.add( + "resource", + _classify_severity(effective_ram, hardware.ram_gb), + f"RAM ~{effective_ram:.1f} GB vs available {hardware.ram_gb:.1f} GB", + metric="ram", + ) + + disk_total = estimate.disk_download_gb + estimate.disk_dump_gb + disk_msg = f"Disk ~{estimate.disk_download_gb:.1f} GB to download" + if estimate.disk_cached_gb > 0: + disk_msg += f", {estimate.disk_cached_gb:.1f} GB already cached" + if estimate.disk_dump_gb > 0: + disk_msg += f", +{estimate.disk_dump_gb:.1f} GB during training (dump_modules=True)" + disk_msg += f" vs {hardware.free_disk_gb:.0f} GB free" + report.add("resource", _classify_severity(disk_total, hardware.free_disk_gb), disk_msg, metric="disk") + + if estimate.time_hours > 0: + report.add( + "resource", + Severity.AMPLE, + f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)", + metric="time", + ) + + +def _resource_phase( + *, + embedder_config: EmbedderConfig, + search_space: list[dict[str, Any]], + n_trials: int, + n_jobs: int, + dump_modules: bool, + stats: DatasetStats, + hardware: HardwareProfile, + report: PreflightReport, + refit_after: bool = False, +) -> None: + """Walk the validated search space, fold per-module costs into the report. + + Two passes: transformer-bearing modules first (collects ``seen_models`` so + the largest model can drive ``embedder_dim`` for the classic pass), then + linear / catboost. Disk, VRAM/RAM peak, time sum, and final findings are + folded onto the report. + """ + seen_models: dict[str, ModelMeta] = {} + global_embedder = _embedder_model_name(embedder_config) + if global_embedder: + seen_models[global_embedder] = _hub.resolve_model(global_embedder) + + transformer_entries, classic_entries = _split_entries(search_space) + + # First pass: transformer modules (also populates seen_models for the classic pass). + module_estimates: list[_ModuleEstimate] = [] + node_max_weights: dict[int, float] = {} + for node_idx, node_type, entry in transformer_entries: + module = entry.get("module_name", "?") + model_names = _extract_model_names(entry) + if not model_names and global_embedder and module in {"knn", "mlknn"}: + model_names = [global_embedder] + for name in model_names: + meta = seen_models.setdefault(name, _hub.resolve_model(name)) + me = _estimate_transformer_model( + meta=meta, + entry=entry, + node_type=node_type, + module=module, + name=name, + stats=stats, + hardware=hardware, + n_trials=n_trials, + refit_after=refit_after, + ) + module_estimates.append(me) + # Track heaviest weight per node so dump_modules is bounded by one + # selected variant per node x n_trials, not the sum of all candidates. + node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), me.model_weights_gb) + + # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. + embedder_meta = _largest_embedder(seen_models) + embedder_dim_val = _embedder_dim(embedder_meta) + for _, node_type, entry in classic_entries: + classic_estimate = _estimate_classic_entry( + entry=entry, + node_type=node_type, + embedder_meta=embedder_meta, + embedder_dim=embedder_dim_val, + stats=stats, + hardware=hardware, + n_trials=n_trials, + refit_after=refit_after, + ) + if classic_estimate is not None: + module_estimates.append(classic_estimate) + + estimate = ResourceEstimate(parallel_factor=n_jobs) + for me in module_estimates: + estimate.vram_gb = max(estimate.vram_gb, me.vram_gb) + estimate.ram_gb = max(estimate.ram_gb, me.ram_gb) + estimate.time_hours += me.time_hours + estimate.drivers.append(me.driver) + + _aggregate_disk(estimate, seen_models, node_max_weights, dump_modules=dump_modules, n_trials=n_trials) + + # Flip low_confidence if any model fell back to the heuristic path (Hub + # unreachable, repo missing safetensors metadata, local-path checkpoint). + heuristic_models = [m.name for m in seen_models.values() if m.confidence == "heuristic"] + if heuristic_models: + report.low_confidence = True + report.notes.append( + f"Heuristic fallback used for {len(heuristic_models)} model(s) - sizes are BERT-base " + f"defaults: {', '.join(heuristic_models[:3])}{'...' if len(heuristic_models) > 3 else ''}", # noqa: PLR2004 + ) + + report.resource = estimate + _emit_resource_findings(report, estimate, hardware, n_jobs=n_jobs) diff --git a/src/autointent/_advisor/_estimates/_search_space.py b/src/autointent/_advisor/_estimates/_search_space.py new file mode 100644 index 000000000..577f500e6 --- /dev/null +++ b/src/autointent/_advisor/_estimates/_search_space.py @@ -0,0 +1,75 @@ +"""Walk preset / OptimizationConfig search-space dicts and extract module info. + +This module is the only place that knows the nested shape of the preset YAML: +``search_space -> list of nodes -> each node has its own search_space -> list of +module entries``. All other modules in the package consume the flattened +``(node_idx, node_type, entry)`` triples this file yields. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterable + + +def _extract_model_names(module_entry: dict[str, Any]) -> list[str]: + """Pull model name(s) from a search-space module entry. + + Each module entry can declare zero or more model candidates under + ``classification_model_config`` and/or ``embedder_config``; both keys may be + a single dict or a list of dicts, and only entries with ``model_name`` are + kept. + """ + candidates: list[str] = [] + cfg = module_entry.get("classification_model_config") + if isinstance(cfg, list): + candidates.extend(c["model_name"] for c in cfg if isinstance(c, dict) and c.get("model_name")) + elif isinstance(cfg, dict) and cfg.get("model_name"): + candidates.append(cfg["model_name"]) + embedder_cfg = module_entry.get("embedder_config") + if isinstance(embedder_cfg, list): + candidates.extend(c["model_name"] for c in embedder_cfg if isinstance(c, dict) and c.get("model_name")) + elif isinstance(embedder_cfg, dict) and embedder_cfg.get("model_name"): + candidates.append(embedder_cfg["model_name"]) + return candidates + + +def _max_int(value: Any, default: int) -> int: # noqa: ANN401 + """Coerce a search-space distribution descriptor into an int upper bound. + + Accepts a plain int, a list of candidate values (returns the max), or an + Optuna-style ``{"low": ..., "high": ...}`` range dict (returns the high end). + Anything unparseable falls back to ``default``. + """ + if value is None: + return default + if isinstance(value, list) and value: + return max(int(x) for x in value) + if isinstance(value, dict): + return int(value.get("high", default)) + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _walk_modules_indexed( + search_space: list[dict[str, Any]], +) -> Iterable[tuple[int, str, dict[str, Any]]]: + """Yield ``(node_index, node_type, module_entry)`` triples. + + The index lets the resource phase bound per-node max cost — see + ``dump_modules`` accounting in ``_resource.py``. + """ + for node_idx, node in enumerate(search_space or []): + node_type = node.get("node_type", "?") + for entry in node.get("search_space", []) or []: + yield node_idx, node_type, entry + + +def _walk_modules(search_space: list[dict[str, Any]]) -> Iterable[tuple[str, dict[str, Any]]]: + """Yield ``(node_type, module_entry)`` pairs — index-agnostic view.""" + for _, node_type, entry in _walk_modules_indexed(search_space): + yield node_type, entry diff --git a/src/autointent/_advisor/runner.py b/src/autointent/_advisor/runner.py new file mode 100644 index 000000000..36fb4fa94 --- /dev/null +++ b/src/autointent/_advisor/runner.py @@ -0,0 +1,170 @@ +"""Public entry point + config validation + data/config phases. + +This file contains the central public function ``run_preflight`` at the top. +Everything below it is supporting machinery for the three phases. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +from pydantic import ValidationError + +from autointent._advisor._estimates._resource import _resource_phase +from autointent._advisor._estimates._search_space import _max_int, _walk_modules +from autointent._advisor._report import PreflightReport, Severity +from autointent._optimization_config import OptimizationConfig + +if TYPE_CHECKING: + from autointent._advisor._hardware import HardwareProfile + from autointent._advisor._report import DatasetStats + + +logger = logging.getLogger(__name__) + + +def run_preflight( + config: dict[str, Any], + stats: DatasetStats, + hardware: HardwareProfile, + *, + preset_name: str | None = None, + refit_after: bool = False, +) -> PreflightReport: + """Run all three preflight phases and return one report. + + Args: + config: parsed preset / ``OptimizationConfig`` dict (top-level keys: + ``search_space``, ``hpo_config``, optional ``embedder_config``, + optional ``logging_config.dump_modules``). + stats: dataset statistics (real or placeholder). + hardware: detected hardware profile. + preset_name: optional friendly name for the report header. + refit_after: matches the ``Pipeline.fit(refit_after=...)`` argument. + When True, time estimates include the extra refit-on-full-data pass. + + Returns: + ``PreflightReport`` with findings across resource / data / config phases. + """ + cfg = _validated_config(config) + report = PreflightReport( + preset_name=preset_name, + hardware={ + "accelerator": hardware.accelerator, + "device_name": hardware.device_name, + "vram_gb": round(hardware.vram_gb, 2), + "ram_gb": round(hardware.ram_gb, 2), + "free_disk_gb": round(hardware.free_disk_gb, 2), + "device_class": hardware.device_class, + }, + dataset={ + "n_samples": stats.n_samples, + "n_classes": stats.n_classes, + "avg_tokens": stats.avg_tokens, + "p95_tokens": stats.p95_tokens, + "multilabel": stats.multilabel, + "source": stats.source, + }, + ) + report.notes.extend(hardware.notes) + + _resource_phase( + embedder_config=cfg.embedder_config, + search_space=cfg.search_space, + n_trials=cfg.hpo_config.n_trials, + n_jobs=cfg.hpo_config.n_jobs, + dump_modules=cfg.logging_config.dump_modules, + stats=stats, + hardware=hardware, + report=report, + refit_after=refit_after, + ) + _data_phase(cfg.search_space, stats, report) + _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, hardware, report) + + return report + + +def _validated_config(config: dict[str, Any]) -> OptimizationConfig: + """Validate ``config`` against the project's canonical ``OptimizationConfig``. + + The advisor is best-effort: a malformed user config should still produce a + report (with placeholder costs) rather than crashing, so any validation + error falls back to the model defaults. + """ + try: + return OptimizationConfig.model_validate(config) + except ValidationError as e: + logger.warning("Advisor config failed validation; falling back to defaults: %s", e) + # OptimizationConfig requires `search_space`; build a minimal valid default. + return OptimizationConfig.model_validate({"search_space": []}) + + +def _config_phase( + search_space: list[dict[str, Any]], + n_jobs: int, + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + """Config-phase checks: parallelism vs. hardware mismatches.""" + if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: + report.add( + "config", + Severity.TIGHT, + f"hpo_config.n_jobs={n_jobs} on a single GPU multiplies VRAM demand by {n_jobs}x.", + ) + + uses_catboost_gpu = any( + entry.get("module_name") == "catboost" and entry.get("task_type") == "GPU" + for _, entry in _walk_modules(search_space) + ) + if uses_catboost_gpu and hardware.accelerator != "cuda": + report.add( + "config", + Severity.TIGHT, + "CatBoost task_type=GPU configured but no CUDA detected - will fall back to CPU.", + ) + + +def _data_phase( + search_space: list[dict[str, Any]], + stats: DatasetStats, + report: PreflightReport, +) -> None: + """Data-phase checks: token truncation, rare classes, missing intent descriptions.""" + # token-length truncation (heuristic — we use stats.p95_tokens vs configured max_length) + p95 = stats.p95_tokens or int(stats.avg_tokens * 2.5) + for _, entry in _walk_modules(search_space): + max_len_value = entry.get("max_length") + if max_len_value is None: + continue + max_len = _max_int(max_len_value, 512) + if p95 > max_len: + severity = Severity.OVER if p95 > max_len * 1.5 else Severity.TIGHT + module_name = entry.get("module_name", "?") + report.add( + "data", + severity, + f"Train tokens p95~{p95} exceeds {module_name}.max_length={max_len}; expect silent truncation.", + ) + + # rare class x linear-CV (LogisticRegressionCV cv=3 needs >=3 samples/class; + # multilabel path uses one-vs-rest without CV so the failure can't occur there) + has_linear = any(e.get("module_name") == "linear" for _, e in _walk_modules(search_space)) + if has_linear and stats.rare_classes and not stats.multilabel: + report.add( + "data", + Severity.OVER, + f"LogisticRegressionCV (cv=3) will fail: classes {stats.rare_classes[:5]} have <3 samples.", + ) + + # partial descriptions x description scorer + description_modules = {"description_bi", "description_cross", "description_llm"} + has_description = any(e.get("module_name") in description_modules for _, e in _walk_modules(search_space)) + if has_description and stats.has_descriptions is False: + report.add( + "data", + Severity.OVER, + "description scorer present but intent descriptions are missing - fill them in or drop the scorer.", + ) From bfd5e0c74d4667d22f59a634b59d5683b123ef3f Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Wed, 24 Jun 2026 16:08:36 +0300 Subject: [PATCH 18/43] fix typing --- src/autointent/_advisor/__init__.py | 4 ++-- src/autointent/_advisor/_cli.py | 5 ++--- src/autointent/_utils.py | 2 -- tests/advisor/test_estimates_internals.py | 12 +++-------- tests/ci/test_compute_matrix.py | 8 ++------ tests/test_deps.py | 25 ++++++++++++++--------- 6 files changed, 24 insertions(+), 32 deletions(-) diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index d8eb6f5ba..04c71007f 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -7,10 +7,10 @@ from __future__ import annotations -from ._estimates import run_preflight from ._hardware import HardwareProfile, detect_hardware from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity -from ._workflows import inspect, load_config, recommend, stats_from_dataset +from .runner import run_preflight +from .workflows import inspect, load_config, recommend, stats_from_dataset __all__ = [ "DatasetStats", diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/_advisor/_cli.py index 315d5c443..6c1b2bc2e 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/_advisor/_cli.py @@ -22,11 +22,10 @@ import logging import sys +from autointent._advisor import inspect, recommend, stats_from_dataset + from ._render import render_json, render_recommendation, render_text from ._report import DatasetStats -from ._workflows import BUNDLED_PRESETS, inspect, recommend, stats_from_dataset - -__all__ = ["BUNDLED_PRESETS", "build_parser", "cmd_inspect", "cmd_recommend", "main"] logger = logging.getLogger("autointent.advisor") diff --git a/src/autointent/_utils.py b/src/autointent/_utils.py index c8a8614b7..92c81431b 100644 --- a/src/autointent/_utils.py +++ b/src/autointent/_utils.py @@ -25,5 +25,3 @@ def detect_device() -> str: if torch.mps.is_available(): return "mps" return "cpu" - - diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 3d4e4e526..67f6d3f80 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -6,15 +6,9 @@ import pytest -from autointent._advisor import _estimates, _hub -from autointent._advisor._estimates import ( - _classify_severity, - _extract_model_names, - _max_int, - _ram_for_module, - _vram_for_transformer, - run_preflight, -) +from autointent._advisor import _estimates, _hub, run_preflight +from autointent._advisor._estimates._formulas import _classify_severity, _ram_for_module, _vram_for_transformer +from autointent._advisor._estimates._search_space import _extract_model_names, _max_int from autointent._advisor._hardware import HardwareProfile from autointent._advisor._hub import ModelMeta from autointent._advisor._report import DatasetStats, Severity diff --git a/tests/ci/test_compute_matrix.py b/tests/ci/test_compute_matrix.py index c03049815..f8f6899b8 100644 --- a/tests/ci/test_compute_matrix.py +++ b/tests/ci/test_compute_matrix.py @@ -88,9 +88,7 @@ def test_push_writes_full_matrix(self, monkeypatch: pytest.MonkeyPatch, tmp_path assert json.loads(outputs["matrix"]) == cm.FULL_MATRIX assert json.loads(outputs["warm_os"]) == ["ubuntu-latest", "windows-latest"] - def test_pr_without_label_writes_minimal_matrix( - self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path - ) -> None: + def test_pr_without_label_writes_minimal_matrix(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: out = tmp_path / "out.txt" monkeypatch.setenv("EVENT_NAME", "pull_request") monkeypatch.setenv("LABELS_JSON", '["bug"]') @@ -103,9 +101,7 @@ def test_pr_without_label_writes_minimal_matrix( assert json.loads(outputs["matrix"]) == cm.MINIMAL_MATRIX assert json.loads(outputs["warm_os"]) == ["ubuntu-latest"] - def test_pr_with_full_ci_label_writes_full_matrix( - self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path - ) -> None: + def test_pr_with_full_ci_label_writes_full_matrix(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: out = tmp_path / "out.txt" monkeypatch.setenv("EVENT_NAME", "pull_request") monkeypatch.setenv("LABELS_JSON", '["full-ci"]') diff --git a/tests/test_deps.py b/tests/test_deps.py index a5c74a716..1be235e62 100644 --- a/tests/test_deps.py +++ b/tests/test_deps.py @@ -35,6 +35,7 @@ def _patch_metadata( requires_map: {dist_name: [PEP 508 requirement string, ...]} versions: {dist_name: installed_version_string} (absent key => not installed) """ + def fake_requires(dist: str) -> list[str]: # Mirror the real importlib.metadata.requires: a dist with no metadata # (i.e. not installed) raises PackageNotFoundError rather than returning []. @@ -81,12 +82,14 @@ def test_check_reports_outdated(monkeypatch: pytest.MonkeyPatch) -> None: def test_iter_extra_reqs_selects_only_extra_members(monkeypatch: pytest.MonkeyPatch) -> None: _patch_metadata( monkeypatch, - {"autointent": [ - "numpy>=1.0 ; python_version >= '3.0'", # base dep w/ env marker -> excluded - "torch>=2.0", # base dep, no marker -> excluded - "catboost>=1.2.8,<2.0.0 ; extra == 'catboost'", # extra member -> included - "peft>=0.10.0 ; extra == 'peft'", # different extra -> excluded - ]}, + { + "autointent": [ + "numpy>=1.0 ; python_version >= '3.0'", # base dep w/ env marker -> excluded + "torch>=2.0", # base dep, no marker -> excluded + "catboost>=1.2.8,<2.0.0 ; extra == 'catboost'", # extra member -> included + "peft>=0.10.0 ; extra == 'peft'", # different extra -> excluded + ] + }, {}, ) reqs = deps._iter_extra_reqs("autointent", "catboost") @@ -119,10 +122,12 @@ def test_resolve_recurses_into_nested_extra(monkeypatch: pytest.MonkeyPatch) -> def test_resolve_terminates_on_cycle(monkeypatch: pytest.MonkeyPatch) -> None: _patch_metadata( monkeypatch, - {"pkg": [ - "pkg[b]>=1.0 ; extra == 'a'", - "pkg[a]>=1.0 ; extra == 'b'", - ]}, + { + "pkg": [ + "pkg[b]>=1.0 ; extra == 'a'", + "pkg[a]>=1.0 ; extra == 'b'", + ] + }, {}, ) reqs = deps._resolve("pkg", "a", set()) From 319d88edc9cc59319a26d1583ec710459aba1221 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Wed, 24 Jun 2026 16:11:16 +0300 Subject: [PATCH 19/43] fix typing --- tests/advisor/test_estimates_and_cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index d87f90740..cbaee2f81 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -17,7 +17,8 @@ import pytest from autointent._advisor import DatasetStats, HardwareProfile, run_preflight -from autointent._advisor._cli import BUNDLED_PRESETS, main +from autointent._advisor._cli import main +from autointent._advisor.workflows import BUNDLED_PRESETS from autointent.utils import load_preset From 908b28deebd17bbd335b726022be3077d20a616b Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Wed, 24 Jun 2026 16:31:48 +0300 Subject: [PATCH 20/43] fix test --- tests/advisor/test_estimates_internals.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 67f6d3f80..02c254f87 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -6,7 +6,7 @@ import pytest -from autointent._advisor import _estimates, _hub, run_preflight +from autointent._advisor import _hub, run_preflight from autointent._advisor._estimates._formulas import _classify_severity, _ram_for_module, _vram_for_transformer from autointent._advisor._estimates._search_space import _extract_model_names, _max_int from autointent._advisor._hardware import HardwareProfile @@ -44,10 +44,9 @@ def _fake_resolve(model_name: str) -> ModelMeta: def _offline(monkeypatch: pytest.MonkeyPatch) -> None: _hub.resolve_model.cache_clear() monkeypatch.setattr(_hub, "_is_warm_cached", lambda _name: False) - # Inject deterministic ModelMeta per name; both the _hub re-export and the - # _estimates rebinding need to be replaced for run_preflight to pick it up. + # Resource phase calls `_hub.resolve_model(...)` via module reference, so + # patching the symbol on `_hub` is enough. monkeypatch.setattr(_hub, "resolve_model", _fake_resolve) - monkeypatch.setattr(_estimates, "resolve_model", _fake_resolve) def _profile(vram_gb: float = 16.0, accelerator: str = "cuda") -> HardwareProfile: From 3c676ef65e4670102c45c6ef7e30e67fc3aceb35 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Mon, 6 Jul 2026 20:03:15 +0300 Subject: [PATCH 21/43] upd --- src/autointent/_advisor/__init__.py | 17 ++- .../_advisor/_estimates/_formulas.py | 5 + .../_advisor/_estimates/_resource.py | 98 ++++++++++++- src/autointent/_advisor/_report.py | 6 +- src/autointent/_advisor/runner.py | 27 ++-- src/autointent/_advisor/workflows.py | 57 ++++++-- src/autointent/_pipeline/__init__.py | 4 +- src/autointent/_pipeline/_pipeline.py | 81 +++++++++- tests/advisor/test_estimates_internals.py | 138 +++++++++++++++++- 9 files changed, 397 insertions(+), 36 deletions(-) diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index 04c71007f..681363afd 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -1,8 +1,8 @@ """Pre-flight compute feasibility advisor. -Exposes a small surface used by both ``Pipeline.fit()`` (future integration) and -the ``autointent-advisor`` CLI script. See ``compute-feasibility-advisor-proposal.md`` -at the repo root for the design document. +Exposes a small surface used by both ``Pipeline.fit()`` (see the ``preflight=`` +kwarg) and the ``autointent-advisor`` CLI script. See +``compute-feasibility-advisor-proposal.md`` at the repo root for the design. """ from __future__ import annotations @@ -10,9 +10,17 @@ from ._hardware import HardwareProfile, detect_hardware from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity from .runner import run_preflight -from .workflows import inspect, load_config, recommend, stats_from_dataset +from .workflows import ( + BUNDLED_PRESETS, + inspect, + load_config, + recommend, + stats_from_dataset, + stats_from_dataset_obj, +) __all__ = [ + "BUNDLED_PRESETS", "DatasetStats", "Finding", "HardwareProfile", @@ -26,4 +34,5 @@ "recommend", "run_preflight", "stats_from_dataset", + "stats_from_dataset_obj", ] diff --git a/src/autointent/_advisor/_estimates/_formulas.py b/src/autointent/_advisor/_estimates/_formulas.py index 8bbbdcebf..2b99b7c57 100644 --- a/src/autointent/_advisor/_estimates/_formulas.py +++ b/src/autointent/_advisor/_estimates/_formulas.py @@ -202,6 +202,11 @@ def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB +def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: + """Disk footprint of one fp32 cached embedding file: ``n_samples x hidden_size x 4``.""" + return (n_samples * hidden_size * 4) / _BYTES_PER_GB + + # Coefficients are dimensional (per-sample-per-feature-per-iteration seconds) # rather than empirically tuned constants — they give relative-cost ordering # across configurations and absolute ballpark wall-times. diff --git a/src/autointent/_advisor/_estimates/_resource.py b/src/autointent/_advisor/_estimates/_resource.py index 6ed93578c..d14dbd186 100644 --- a/src/autointent/_advisor/_estimates/_resource.py +++ b/src/autointent/_advisor/_estimates/_resource.py @@ -29,6 +29,7 @@ _activations_gb_per_sample, _classify_severity, _embedder_dim, + _embedding_cache_disk_gb, _largest_embedder, _max_fitting_batch_size, _ram_for_catboost, @@ -76,6 +77,27 @@ def _embedder_model_name(embedder: EmbedderConfig) -> str | None: "lora": "lora", } +# Scorers that consume embeddings (cache key = model + utterances + prompt) but +# don't train the encoder — embedder forward is shared via the persistent cache. +_CACHE_HONORING_MODULES = frozenset( + { + "linear", + "catboost", + "knn", + "mlknn", + "retrieval", + "description_bi", + "description_cross", + "description_llm", + }, +) + +# Cache-honoring modules whose per-entry estimate already bundles the embedder +# forward into `time_hours` (vs. classic linear/catboost which don't). +_EMBEDDER_FORWARD_TRANSFORMER_MODULES = frozenset( + {"knn", "mlknn", "retrieval", "description_bi", "description_cross", "description_llm"}, +) + @dataclass class _ModuleEstimate: @@ -241,6 +263,51 @@ def _estimate_classic_entry( ) +def _apply_embedding_cache( + module_estimates: list[_ModuleEstimate], + seen_models: dict[str, ModelMeta], + *, + stats: DatasetStats, + hardware: HardwareProfile, +) -> set[str]: + """Adjust ``module_estimates`` in-place for autointent's persistent embedding cache. + + Per unique embedder, the first cache-honoring entry pays the forward; later + transformer entries get ``time_hours`` zeroed (cache hit), and classic + entries (linear/catboost) get a synthetic forward added since their + per-entry estimate doesn't include one. + + Returns the set of unique embedder model names whose forward was charged. + """ + paid: set[str] = set() + for me in module_estimates: + module = me.driver["module"] + if module not in _CACHE_HONORING_MODULES: + continue + model = me.driver["model"] + if model not in seen_models: # synthetic / "(no embedder)" rows + continue + if model in paid: + if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: + me.time_hours = 0.0 + me.driver["time_hours"] = 0.0 + me.driver["mode"] = f"{me.driver['mode']}+cached" + else: + paid.add(model) + if module in {"linear", "catboost"}: + forward_h = _time_for_transformer( + n_trials=1, + epochs=1, + batch_size=32, + n_samples=stats.n_samples, + accelerator=hardware.accelerator, + ) + me.time_hours += forward_h + me.driver["time_hours"] = round(me.time_hours, 2) + me.driver["mode"] = f"{me.driver['mode']}+embed" + return paid + + def _aggregate_disk( estimate: ResourceEstimate, seen_models: dict[str, ModelMeta], @@ -248,8 +315,10 @@ def _aggregate_disk( *, dump_modules: bool, n_trials: int, + cached_embedders: set[str] | None = None, + stats: DatasetStats | None = None, ) -> None: - """Fold per-model download/cached sizes into ``estimate`` and apply dump-modules accounting.""" + """Fold per-model download/cached/embedding-cache sizes into ``estimate``.""" for meta in seen_models.values(): if meta.cached_locally: estimate.disk_cached_gb += meta.disk_gb @@ -260,6 +329,16 @@ def _aggregate_disk( # are bounded by the heaviest module in each node, summed across nodes. estimate.disk_dump_gb = sum(node_max_weights.values()) * n_trials + if cached_embedders and stats is not None: + for name in cached_embedders: + meta = seen_models.get(name) + if meta is None: + continue + estimate.disk_embedding_cache_gb += _embedding_cache_disk_gb( + n_samples=stats.n_samples, + hidden_size=_embedder_dim(meta), + ) + def _emit_resource_findings( report: PreflightReport, @@ -296,12 +375,14 @@ def _emit_resource_findings( metric="ram", ) - disk_total = estimate.disk_download_gb + estimate.disk_dump_gb + disk_total = estimate.disk_download_gb + estimate.disk_dump_gb + estimate.disk_embedding_cache_gb disk_msg = f"Disk ~{estimate.disk_download_gb:.1f} GB to download" if estimate.disk_cached_gb > 0: disk_msg += f", {estimate.disk_cached_gb:.1f} GB already cached" if estimate.disk_dump_gb > 0: disk_msg += f", +{estimate.disk_dump_gb:.1f} GB during training (dump_modules=True)" + if estimate.disk_embedding_cache_gb > 0: + disk_msg += f", +{estimate.disk_embedding_cache_gb:.2f} GB embedding cache" disk_msg += f" vs {hardware.free_disk_gb:.0f} GB free" report.add("resource", _classify_severity(disk_total, hardware.free_disk_gb), disk_msg, metric="disk") @@ -383,6 +464,9 @@ def _resource_phase( if classic_estimate is not None: module_estimates.append(classic_estimate) + # Cache-aware time/disk: must run before the fold below. + cached_embedders = _apply_embedding_cache(module_estimates, seen_models, stats=stats, hardware=hardware) + estimate = ResourceEstimate(parallel_factor=n_jobs) for me in module_estimates: estimate.vram_gb = max(estimate.vram_gb, me.vram_gb) @@ -390,7 +474,15 @@ def _resource_phase( estimate.time_hours += me.time_hours estimate.drivers.append(me.driver) - _aggregate_disk(estimate, seen_models, node_max_weights, dump_modules=dump_modules, n_trials=n_trials) + _aggregate_disk( + estimate, + seen_models, + node_max_weights, + dump_modules=dump_modules, + n_trials=n_trials, + cached_embedders=cached_embedders, + stats=stats, + ) # Flip low_confidence if any model fell back to the heuristic path (Hub # unreachable, repo missing safetensors metadata, local-path checkpoint). diff --git a/src/autointent/_advisor/_report.py b/src/autointent/_advisor/_report.py index c9fd920f4..6fc3d8f2f 100644 --- a/src/autointent/_advisor/_report.py +++ b/src/autointent/_advisor/_report.py @@ -33,6 +33,7 @@ class ResourceEstimate: disk_download_gb: float = 0.0 disk_cached_gb: float = 0.0 disk_dump_gb: float = 0.0 + disk_embedding_cache_gb: float = 0.0 ram_gb: float = 0.0 vram_gb: float = 0.0 time_hours: float = 0.0 @@ -41,7 +42,7 @@ class ResourceEstimate: @property def total_disk_gb(self) -> float: - return self.disk_download_gb + self.disk_dump_gb + return self.disk_download_gb + self.disk_dump_gb + self.disk_embedding_cache_gb @dataclass @@ -57,7 +58,8 @@ class DatasetStats: p95_tokens: int | None = None multilabel: bool = False has_descriptions: bool | None = None - rare_classes: list[str] = field(default_factory=list) + # Per-class train-split sample counts; empty when no real dataset was provided. + class_counts: dict[str, int] = field(default_factory=dict) source: str = "placeholder" @classmethod diff --git a/src/autointent/_advisor/runner.py b/src/autointent/_advisor/runner.py index 36fb4fa94..168157550 100644 --- a/src/autointent/_advisor/runner.py +++ b/src/autointent/_advisor/runner.py @@ -149,15 +149,24 @@ def _data_phase( f"Train tokens p95~{p95} exceeds {module_name}.max_length={max_len}; expect silent truncation.", ) - # rare class x linear-CV (LogisticRegressionCV cv=3 needs >=3 samples/class; - # multilabel path uses one-vs-rest without CV so the failure can't occur there) - has_linear = any(e.get("module_name") == "linear" for _, e in _walk_modules(search_space)) - if has_linear and stats.rare_classes and not stats.multilabel: - report.add( - "data", - Severity.OVER, - f"LogisticRegressionCV (cv=3) will fail: classes {stats.rare_classes[:5]} have <3 samples.", - ) + # sklearn LogisticRegressionCV inner-CV failure: each class needs >= cv samples. + # cv is configurable per linear entry (default 3); use the strictest one across + # the search space. Multilabel uses LogisticRegression (no CV), so skip there. + if not stats.multilabel and stats.class_counts: + linear_cvs = [ + _max_int(e.get("cv"), 3) + for _, e in _walk_modules(search_space) + if e.get("module_name") == "linear" + ] + if linear_cvs: + cv_max = max(linear_cvs) + failing = sorted(name for name, count in stats.class_counts.items() if count < cv_max) + if failing: + report.add( + "data", + Severity.OVER, + f"LogisticRegressionCV (cv={cv_max}) will fail: classes {failing[:5]} have <{cv_max} samples.", + ) # partial descriptions x description scorer description_modules = {"description_bi", "description_cross", "description_llm"} diff --git a/src/autointent/_advisor/workflows.py b/src/autointent/_advisor/workflows.py index 2331e225a..fad0562d6 100644 --- a/src/autointent/_advisor/workflows.py +++ b/src/autointent/_advisor/workflows.py @@ -27,6 +27,8 @@ if TYPE_CHECKING: from collections.abc import Iterable + from autointent import Dataset + from ._report import PreflightReport @@ -92,11 +94,49 @@ def stats_from_dataset(path: str, *, multilabel: bool = False) -> DatasetStats: p95_tokens=p95, multilabel=detected_multilabel, has_descriptions=None, - rare_classes=_rare_classes(train, label_col, detected_multilabel, n_classes) if label_col else [], + class_counts=_class_counts(train, label_col, detected_multilabel, n_classes) if label_col else {}, source=f"dataset:{path}", ) +def stats_from_dataset_obj(dataset: Dataset) -> DatasetStats: + """Build :class:`DatasetStats` straight from an in-memory ``Dataset``. + + Counterpart of :func:`stats_from_dataset` that skips HF ``load_dataset`` + and reads the train split + autointent-specific attributes (``n_classes``, + ``multilabel``, ``has_descriptions``) directly. + """ + from autointent.custom_types import Split + + train_key = Split.TRAIN if Split.TRAIN in dataset else f"{Split.TRAIN}_0" + if train_key not in dataset: + return DatasetStats.placeholder() + train = dataset[train_key] + utt_col = dataset.utterance_feature + label_col = dataset.label_feature + + sample = train[:_SAMPLE_LIMIT] if len(train) > _SAMPLE_LIMIT else train[:] + lengths = [len(str(s).split()) for s in sample.get(utt_col, [])] + avg_tokens = int(sum(lengths) / max(1, len(lengths))) if lengths else 32 + if lengths: + sorted_lengths = sorted(lengths) + idx = max(0, min(len(sorted_lengths) - 1, round((len(sorted_lengths) - 1) * _P95_PERCENTILE))) + p95 = sorted_lengths[idx] + else: + p95 = avg_tokens * 2 + + return DatasetStats( + n_samples=len(train), + n_classes=dataset.n_classes, + avg_tokens=avg_tokens, + p95_tokens=p95, + multilabel=dataset.multilabel, + has_descriptions=dataset.has_descriptions, + class_counts=_class_counts(train, label_col, dataset.multilabel, dataset.n_classes), + source="dataset:in-memory", + ) + + def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool) -> tuple[bool, int]: # noqa: ANN401 """Derive ``(multilabel, n_classes)`` from the HF feature schema with a value-based fallback.""" if label_col is None: @@ -119,22 +159,17 @@ def _label_shape(train: Any, label_col: str | None, *, fallback_multilabel: bool return False, len({label for label in train[label_col] if label is not None}) -def _rare_classes( +def _class_counts( train: Any, # noqa: ANN401 label_col: str, multilabel: bool, n_classes: int, - min_count: int = 3, -) -> list[str]: - """Return labels with fewer than ``min_count`` samples in the train split. - - Used to surface the LogisticRegressionCV(cv=3) failure case before fit. - Returns an empty list on any error so the advisor stays best-effort. - """ +) -> dict[str, int]: + """Per-class sample counts in the train split; empty on any error.""" try: labels = train[label_col] except (KeyError, AttributeError, TypeError): - return [] + return {} counts: dict[str, int] = {} if multilabel: for row in labels: @@ -148,7 +183,7 @@ def _rare_classes( else: for label in labels: counts[str(label)] = counts.get(str(label), 0) + 1 - return sorted(name for name, c in counts.items() if c < min_count) + return counts def inspect( diff --git a/src/autointent/_pipeline/__init__.py b/src/autointent/_pipeline/__init__.py index 7a8af8259..50d58f7d1 100644 --- a/src/autointent/_pipeline/__init__.py +++ b/src/autointent/_pipeline/__init__.py @@ -1,3 +1,3 @@ -from ._pipeline import Pipeline +from ._pipeline import Pipeline, PreflightError, PreflightMode -__all__ = ["Pipeline"] +__all__ = ["Pipeline", "PreflightError", "PreflightMode"] diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index a764f1d78..2fe352ebd 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -5,13 +5,19 @@ import json import logging from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import numpy as np import yaml from typing_extensions import assert_never from autointent import Context, OptimizationConfig +from autointent._advisor import ( + Severity, + detect_hardware, + run_preflight, + stats_from_dataset_obj, +) from autointent.configs import ( CrossEncoderConfig, DataConfig, @@ -34,10 +40,23 @@ if TYPE_CHECKING: from autointent import Dataset + from autointent._advisor import PreflightReport from autointent.custom_types import ListOfGenericLabels, SearchSpacePreset, SearchSpaceValidationMode from autointent.modules.base import BaseDecision, BaseRegex, BaseScorer +PreflightMode = Literal["off", "warn", "strict"] + + +class PreflightError(RuntimeError): + """Raised when ``Pipeline.fit(preflight="strict")`` finds OVER-budget resources.""" + + def __init__(self, findings: list[Any]) -> None: + self.findings = findings + lines = "\n".join(f" [{f.phase}] {f.message}" for f in findings) + super().__init__(f"Preflight check failed with {len(findings)} OVER finding(s):\n{lines}") + + class Pipeline: """Pipeline optimizer class. @@ -152,6 +171,40 @@ def from_optimization_config(cls, config: dict[str, Any] | Path | str | Optimiza pipeline.set_config(optimization_config.hpo_config) return pipeline + def _build_advisor_config(self) -> dict[str, Any]: + """Reconstruct an ``OptimizationConfig``-shaped dict for the advisor.""" + search_space = [ + {"node_type": opt.node_type, "search_space": opt.modules_search_spaces} + for opt in self.nodes.values() + if isinstance(opt, NodeOptimizer) + ] + return { + "search_space": search_space, + "data_config": self.data_config.model_dump(), + "logging_config": self.logging_config.model_dump(), + "embedder_config": self.embedder_config.model_dump(), + "cross_encoder_config": self.cross_encoder_config.model_dump(), + "transformer_config": self.transformer_config.model_dump(), + "hpo_config": self.hpo_config.model_dump(), + } + + def _run_preflight(self, dataset: Dataset, *, refit_after: bool, mode: PreflightMode) -> PreflightReport: + """Run the advisor against this pipeline's effective config + dataset. + + Logs each finding at INFO/WARNING/ERROR (by severity). When ``mode`` is + ``"strict"`` and any OVER finding is produced, raises ``PreflightError``. + """ + config = self._build_advisor_config() + stats = stats_from_dataset_obj(dataset) + hardware = detect_hardware() + report = run_preflight(config, stats, hardware, refit_after=refit_after) + _log_preflight_report(report, self._logger) + if mode == "strict": + over = [f for f in report.findings if f.severity == Severity.OVER] + if over: + raise PreflightError(over) + return report + def _fit(self, context: Context) -> None: """Optimize the pipeline. @@ -193,6 +246,7 @@ def fit( dataset: Dataset, refit_after: bool = False, incompatible_search_space: SearchSpaceValidationMode = "filter", + preflight: PreflightMode = "warn", ) -> Context: """Optimize the pipeline from dataset. @@ -201,14 +255,24 @@ def fit( refit_after: whether to refit on whole data after optimization. Valid only for hold-out validaiton. sampler: sampler type to use. incompatible_search_space: wow to handle data-incompatible modules occurring in search space. + preflight: gate that runs :func:`autointent._advisor.run_preflight` over the + pipeline's effective config + dataset before any heavy work. + ``"off"`` skips it. ``"warn"`` (default) logs findings — INFO for + AMPLE, WARNING for TIGHT, ERROR for OVER — but never raises. + ``"strict"`` additionally raises :class:`PreflightError` when any + finding has severity OVER, so unfeasible runs abort before fit. Raises: RuntimeError: If pipeline is in inference mode. + PreflightError: If ``preflight="strict"`` and any OVER finding is produced. """ if self._is_inference(): msg = "Pipeline in inference mode cannot be fitted" raise RuntimeError(msg) + if preflight != "off": + self._run_preflight(dataset, refit_after=refit_after, mode=preflight) + context = Context(self._seed) context.set_dataset(dataset, self.data_config) context.configure_logging(self.logging_config) @@ -472,3 +536,18 @@ def make_report(logs: dict[str, Any], nodes: list[NodeType]) -> str: messages = [json.dumps(c, indent=4) for c in configs] msg = "\n".join(messages) return "resulting pipeline configuration is the following:\n" + msg + + +def _log_preflight_report(report: PreflightReport, logger: logging.Logger) -> None: + """Log each preflight finding at the appropriate level.""" + level_for = { + Severity.AMPLE: logging.INFO, + Severity.TIGHT: logging.WARNING, + Severity.OVER: logging.ERROR, + } + header = f"Preflight ({report.preset_name or 'pipeline'}): verdict={'feasible' if report.is_feasible else 'INFEASIBLE'}" + logger.info(header) + for finding in report.findings: + logger.log(level_for[finding.severity], "[%s] %s", finding.phase, finding.message) + if report.low_confidence: + logger.info("Preflight: low-confidence (heuristic fallback in use)") diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 02c254f87..a62bce45b 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -270,11 +270,37 @@ def test_rare_classes_with_linear_scorer_flag_red(self) -> None: n_samples=20, n_classes=5, avg_tokens=10, - rare_classes=["intent_a", "intent_b"], + class_counts={"intent_a": 1, "intent_b": 2, "intent_c": 6, "intent_d": 6, "intent_e": 5}, ) report = run_preflight(cfg, stats, _profile()) assert any( - f.phase == "data" and "LogisticRegressionCV" in f.message and f.severity == Severity.OVER + f.phase == "data" and "LogisticRegressionCV (cv=3)" in f.message and f.severity == Severity.OVER + for f in report.findings + ) + + def test_rare_classes_threshold_follows_entry_cv(self) -> None: + """When a linear entry sets cv=5, classes with 4 samples should still fail.""" + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear", "cv": 5}, + ], + } + ] + } + # All classes have >=3 samples, so a cv=3 check would pass — but cv=5 + # needs >=5, so intent_a (4 samples) must be flagged. + stats = DatasetStats( + n_samples=20, + n_classes=3, + avg_tokens=10, + class_counts={"intent_a": 4, "intent_b": 8, "intent_c": 8}, + ) + report = run_preflight(cfg, stats, _profile()) + assert any( + f.phase == "data" and "cv=5" in f.message and "intent_a" in f.message for f in report.findings ) @@ -407,7 +433,9 @@ def test_catboost_contributes_ram_and_time_on_cpu(self) -> None: assert report.resource.ram_gb > 0 assert report.resource.time_hours > 0 assert cb["vram_gb"] == 0 - assert cb["mode"] == "catboost" + # The "+embed" suffix is added when the embedder forward is folded into + # this classic entry via the embedding-cache adjustment. + assert cb["mode"].startswith("catboost") def test_catboost_gpu_moves_cost_to_vram(self) -> None: cfg = { @@ -432,7 +460,9 @@ def test_catboost_gpu_moves_cost_to_vram(self) -> None: cb = next(d for d in report.resource.drivers if d["module"] == "catboost") assert report.resource.vram_gb > 0 assert cb["ram_gb"] == 0 - assert cb["mode"] == "catboost-gpu" + # The "+embed" suffix is added when the embedder forward is folded into + # this classic entry via the embedding-cache adjustment. + assert cb["mode"].startswith("catboost-gpu") def test_linear_scales_with_n_samples(self) -> None: cfg = { @@ -599,3 +629,103 @@ def test_dump_disk_sums_across_nodes(self) -> None: bert = _hub.resolve_model("microsoft/deberta-v3-small") expected = (embedder.weights_gb + bert.weights_gb) * 2 assert report.resource.disk_dump_gb == pytest.approx(expected, rel=0.01) + + +class TestEmbeddingCache: + """Cache-aware time + disk accounting for embedder-honoring scorers. + + autointent's ``SentenceTransformerEmbedding`` (``use_cache=True`` by default) + persists per-(model, utterances, prompt) embeddings to disk, so subsequent + trials/modules that reuse the same embedder hit the cache instead of + re-running the forward pass. + """ + + def _embedder_node(self) -> dict[str, Any]: + return { + "node_type": "embedder", + "search_space": [ + { + "module_name": "sentence_transformer", + "embedder_config": [{"model_name": "sentence-transformers/all-MiniLM-L6-v2"}], + } + ], + } + + def test_duplicate_knn_entries_zero_time_after_first(self) -> None: + """Two knn entries sharing an embedder: the second one's forward is free.""" + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "knn", + "embedder_config": [{"model_name": "sentence-transformers/all-MiniLM-L6-v2"}], + "batch_size": [32], + "max_length": [128], + }, + { + "module_name": "knn", + "embedder_config": [{"model_name": "sentence-transformers/all-MiniLM-L6-v2"}], + "batch_size": [32], + "max_length": [128], + }, + ], + }, + ], + "hpo_config": {"n_trials": 5}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + knn_drivers = [d for d in report.resource.drivers if d["module"] == "knn"] + assert len(knn_drivers) == 2 + first, second = knn_drivers + assert first["time_hours"] > 0 + assert second["time_hours"] == 0 + assert "cached" in second["mode"] + + def test_classic_entry_gets_synthetic_embedder_forward(self) -> None: + """A linear scorer alone with an embedder: the embedder forward is added once.""" + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [{"module_name": "linear"}], + }, + ], + "hpo_config": {"n_trials": 3}, + } + # Re-run with the embedder node removed to compare cleanly. + cfg_no_embed = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [{"module_name": "linear"}], + }, + ], + "hpo_config": {"n_trials": 3}, + } + with_embed = run_preflight(cfg, DatasetStats.placeholder(n_samples=10_000), _profile()) + no_embed = run_preflight(cfg_no_embed, DatasetStats.placeholder(n_samples=10_000), _profile()) + # The linear row gets a "+embed" suffix when an embedder is present. + linear_with = next(d for d in with_embed.resource.drivers if d["module"] == "linear") + linear_no = next(d for d in no_embed.resource.drivers if d["module"] == "linear") + assert "embed" in linear_with["mode"] + assert linear_with["time_hours"] >= linear_no["time_hours"] + + def test_disk_embedding_cache_scales_with_n_samples(self) -> None: + """``disk_embedding_cache_gb`` ~ n_samples × hidden_size × 4 bytes per embedder.""" + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [{"module_name": "linear"}], + }, + ], + } + small = run_preflight(cfg, DatasetStats.placeholder(n_samples=1_000), _profile()) + big = run_preflight(cfg, DatasetStats.placeholder(n_samples=1_000_000), _profile()) + assert small.resource.disk_embedding_cache_gb > 0 + assert big.resource.disk_embedding_cache_gb > small.resource.disk_embedding_cache_gb * 100 From 570a135048ac83cdf3e9c1fec29d249b4c7370b4 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:43:50 +0300 Subject: [PATCH 22/43] add scripts --- scripts/calibrate_advisor.py | 541 +++++++++++++++++++++++++++ scripts/run_calibration_banking77.sh | 84 +++++ tests/pipeline/test_preflight.py | 122 ++++++ 3 files changed, 747 insertions(+) create mode 100644 scripts/calibrate_advisor.py create mode 100755 scripts/run_calibration_banking77.sh create mode 100644 tests/pipeline/test_preflight.py diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py new file mode 100644 index 000000000..13dd76836 --- /dev/null +++ b/scripts/calibrate_advisor.py @@ -0,0 +1,541 @@ +"""Calibrate advisor preflight estimates against real Pipeline.fit measurements. + +Runs each requested preset twice: first through ``run_preflight`` to capture the +heuristic estimate, then through ``Pipeline.from_preset(...).fit(...)`` while +measuring wall-time, peak RAM (RSS), peak VRAM (CUDA only — MPS has no exact +peak API), and the disk delta in the HF Hub cache. + +The output is a JSON file with per-preset predicted vs. actual values plus +ratios, and a side-by-side table on stdout for quick eyeballing. + +Usage: + python scripts/calibrate_advisor.py \\ + --dataset tests/assets/data/clinc_subset.json \\ + --presets classic-light classic-medium \\ + --output calibration.json \\ + --max-trials 3 + +The ``--skip-fit`` flag runs only the predicted side, useful for sanity-checking +the preflight numbers across presets without paying for fits. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +import threading +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import psutil + +from autointent import Dataset, Pipeline +from autointent._advisor import ( + BUNDLED_PRESETS, + PreflightReport, + detect_hardware, + run_preflight, + stats_from_dataset_obj, +) +from autointent._callbacks.base import OptimizerCallback +from autointent.configs import HPOConfig, LoggingConfig + +logger = logging.getLogger("calibrate_advisor") + +_BYTES_PER_GB = 1024**3 + + +@dataclass +class CalibrationRow: + """One preset's predicted vs. actual numbers.""" + + preset: str + predicted: dict[str, float] = field(default_factory=dict) + actual: dict[str, float | None] = field(default_factory=dict) + ratios: dict[str, float | None] = field(default_factory=dict) + findings: int = 0 + findings_over: int = 0 + # Per-module records from _ModuleTracker: [{module, num, config, duration_s, peak_vram_gb?}, ...] + modules: list[dict[str, Any]] = field(default_factory=list) + error: str | None = None + notes: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="calibrate_advisor", + description="Compare advisor preflight estimates to real Pipeline.fit measurements.", + ) + p.add_argument( + "--dataset", + required=True, + type=str, + help=( + "Either a local JSON path (loaded via ``Dataset.from_json``) or an HF Hub repo id " + "such as ``DeepPavlov/banking77`` (loaded via ``Dataset.from_hub``)." + ), + ) + p.add_argument( + "--presets", + nargs="+", + default=None, + help="Preset names to run (default: every preset in BUNDLED_PRESETS).", + ) + p.add_argument("--output", type=Path, default=Path("calibration.json"), help="Where to write the JSON report.") + p.add_argument("--max-trials", type=int, default=None, help="Override hpo_config.n_trials for faster runs.") + p.add_argument( + "--skip-fit", + action="store_true", + help="Only run preflight (no fit) — useful for sanity-checking estimates.", + ) + p.add_argument( + "--poll-interval-ms", + type=int, + default=100, + help="RSS polling interval during fit (ms). Lower is more accurate but more overhead.", + ) + p.add_argument( + "--wandb", + action="store_true", + help=( + "Attach the W&B reporter so per-step GPU/system metrics land in wandb.ai. " + "Requires ``wandb`` installed + ``WANDB_API_KEY`` in the environment." + ), + ) + p.add_argument("-v", "--verbose", action="store_true") + return p + + +# === measurement helpers ================================================= + + +def _hf_cache_dir() -> Path: + """Return the active HF Hub cache directory ($HF_HOME / ~/.cache/huggingface).""" + return Path(os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")) + + +def _dir_size_gb(path: Path) -> float: + """Disk usage of ``path`` in GB; 0 when the directory is missing.""" + if not path.exists(): + return 0.0 + total = 0 + for entry in path.rglob("*"): + try: + if entry.is_file(): + total += entry.stat().st_size + except OSError: + continue + return total / _BYTES_PER_GB + + +class _PeakSampler: + """Background thread tracking peak RSS and (on MPS) peak GPU allocation. + + CUDA has an accurate native peak-memory API and doesn't need polling; we + still read it after the fit. MPS lacks a peak API, so the sampler polls + ``torch.mps.current_allocated_memory()`` alongside RSS and keeps the max. + """ + + def __init__(self, interval_s: float = 0.1, *, sample_mps: bool = False) -> None: + self._interval_s = interval_s + self._proc = psutil.Process() + self.peak_ram_gb = self._proc.memory_info().rss / _BYTES_PER_GB + self.peak_mps_gb: float | None = 0.0 if sample_mps else None + self._sample_mps = sample_mps + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + def __enter__(self) -> _PeakSampler: + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + return self + + def __exit__(self, *_exc: object) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=1.0) + + def _run(self) -> None: + try: + import torch # noqa: PLC0415 + except ImportError: + torch = None # type: ignore[assignment] + while not self._stop.is_set(): + try: + rss = self._proc.memory_info().rss / _BYTES_PER_GB + if rss > self.peak_ram_gb: + self.peak_ram_gb = rss + if self._sample_mps and torch is not None: + mps = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB + if self.peak_mps_gb is None or mps > self.peak_mps_gb: + self.peak_mps_gb = mps + except (psutil.NoSuchProcess, psutil.AccessDenied): + break + self._stop.wait(self._interval_s) + + +def _reset_vram_peak() -> None: + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + except ImportError: + pass + + +def _read_vram_peak_gb(accelerator: str) -> float | None: + """Peak VRAM/GPU in GB. CUDA uses the native peak API; MPS uses the polled sampler value (caller-side).""" + try: + import torch + except ImportError: + return None + if accelerator == "cuda" and torch.cuda.is_available(): + return float(torch.cuda.max_memory_allocated()) / _BYTES_PER_GB + return None + + +# === per-module tracking ================================================= + + +class _ModuleTracker(OptimizerCallback): + """Records per-module wall time and peak VRAM. + + Hooks ``start_module`` / ``end_module`` on the CallbackHandler so we get + one record per (module_name, trial_num). CUDA peak VRAM is reset per module + via ``torch.cuda.reset_peak_memory_stats``; MPS is sampled at ``end_module`` + (no per-module peak API, so it's the moment-in-time allocation). + """ + + name = "calibration_tracker" + + def __init__(self) -> None: # noqa: D401 + self.records: list[dict[str, Any]] = [] + self._current: dict[str, Any] | None = None + + def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: # noqa: ARG002 + pass + + def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None: + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + except ImportError: + pass + # Only capture JSON-safe scalars in the config snapshot. + safe_config = { + k: v for k, v in module_kwargs.items() if isinstance(v, (str, int, float, bool)) or v is None + } + self._current = { + "module": module_name, + "num": num, + "config": safe_config, + "_start": time.perf_counter(), + } + + def log_value(self, **kwargs: Any) -> None: # noqa: ANN401, ARG002 + pass + + def log_metrics(self, metrics: dict[str, Any]) -> None: # noqa: ARG002 + pass + + def end_module(self) -> None: + if self._current is None: + return + rec = self._current + rec["duration_s"] = time.perf_counter() - rec.pop("_start") + try: + import torch + + if torch.cuda.is_available(): + rec["peak_vram_gb"] = float(torch.cuda.max_memory_allocated()) / _BYTES_PER_GB + elif torch.backends.mps.is_available(): + # MPS has no per-module peak API — snapshot the current allocation. + rec["peak_vram_gb"] = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB + except (ImportError, AttributeError): + pass + self.records.append(rec) + self._current = None + + def end_run(self) -> None: + pass + + def log_final_metrics(self, metrics: dict[str, Any]) -> None: # noqa: ARG002 + pass + + +def _attach_tracker(pipeline: Pipeline, tracker: _ModuleTracker) -> None: + """Instance-patch ``pipeline._fit`` so ``tracker`` is appended to the callback chain.""" + original_fit = pipeline._fit # noqa: SLF001 + + def patched(context: Any) -> Any: # noqa: ANN401 + context.callback_handler.callbacks.append(tracker) + return original_fit(context) + + pipeline._fit = patched # type: ignore[method-assign] # noqa: SLF001 + + +# === per-preset run ====================================================== + + +def _override_trials(pipeline: Pipeline, max_trials: int | None, *, enable_wandb: bool) -> None: + """Cap n_trials, disable dumping, optionally enable W&B for post-run analysis.""" + updates: dict[str, Any] = {} + if max_trials is not None: + updates["n_trials"] = max_trials + if enable_wandb: + # Trigger built-in per-run system-metrics collection in W&B. + updates["report_to"] = ["wandb"] + if updates: + pipeline.set_config(pipeline.hpo_config.model_copy(update=updates)) + # We don't want the calibration run to leave dumped module weights on disk. + pipeline.set_config(LoggingConfig(dump_modules=False, clear_ram=True)) + + +def _calibrate_one( + *, + preset: str, + dataset: Dataset, + stats: Any, # noqa: ANN401 + hardware: Any, # noqa: ANN401 + max_trials: int | None, + skip_fit: bool, + poll_interval_ms: int, + enable_wandb: bool, +) -> CalibrationRow: + row = CalibrationRow(preset=preset) + + # === predicted ====================================================== + try: + pipeline = Pipeline.from_preset(preset) + except Exception as e: # noqa: BLE001 + row.error = f"from_preset failed: {e}" + return row + + _override_trials(pipeline, max_trials, enable_wandb=enable_wandb) + + try: + report: PreflightReport = run_preflight( + pipeline._build_advisor_config(), # noqa: SLF001 + stats, + hardware, + ) + except Exception as e: # noqa: BLE001 + row.error = f"preflight failed: {e}" + return row + + row.predicted = { + "time_h": report.resource.time_hours, + "ram_gb": report.resource.ram_gb, + "vram_gb": report.resource.vram_gb, + "disk_download_gb": report.resource.disk_download_gb, + "disk_cached_gb": report.resource.disk_cached_gb, + "disk_embedding_cache_gb": report.resource.disk_embedding_cache_gb, + } + row.findings = len(report.findings) + row.findings_over = sum(1 for f in report.findings if f.severity.value == "over") + if report.low_confidence: + row.notes.append("low-confidence (heuristic fallback in use)") + + if skip_fit: + return row + + # === actual ========================================================= + hf_cache = _hf_cache_dir() + cache_before = _dir_size_gb(hf_cache) + _reset_vram_peak() + + tracker = _ModuleTracker() + _attach_tracker(pipeline, tracker) + + is_mps = hardware.accelerator == "mps" + start = time.perf_counter() + try: + with _PeakSampler(interval_s=poll_interval_ms / 1000.0, sample_mps=is_mps) as sampler: + pipeline.fit(dataset, preflight="off") + except Exception as e: # noqa: BLE001 + row.error = f"fit failed: {e}" + row.modules = tracker.records # keep whatever we collected + return row + elapsed_s = time.perf_counter() - start + + cache_after = _dir_size_gb(hf_cache) + actual_time_h = elapsed_s / 3600.0 + actual_ram_gb = sampler.peak_ram_gb + actual_vram_gb = _read_vram_peak_gb(hardware.accelerator) + if actual_vram_gb is None and is_mps: + actual_vram_gb = sampler.peak_mps_gb + actual_disk_download_gb = max(0.0, cache_after - cache_before) + + row.actual = { + "time_h": actual_time_h, + "ram_gb": actual_ram_gb, + "vram_gb": actual_vram_gb, + "disk_download_gb": actual_disk_download_gb, + } + row.modules = tracker.records + if enable_wandb: + row.notes.append("W&B reporter enabled — inspect wandb.ai run group for per-step GPU/system metrics") + + def _ratio(actual: float | None, predicted: float) -> float | None: + if actual is None or predicted <= 0: + return None + return actual / predicted + + row.ratios = { + "time": _ratio(actual_time_h, row.predicted["time_h"]), + "ram": _ratio(actual_ram_gb, row.predicted["ram_gb"]), + "vram": _ratio(actual_vram_gb, row.predicted["vram_gb"]), + "disk_download": _ratio(actual_disk_download_gb, row.predicted["disk_download_gb"]), + } + return row + + +# === rendering =========================================================== + + +_COLS = [ + ("preset", "Preset", 22), + ("pred_time", "pred_time_h", 12), + ("act_time", "act_time_h", 12), + ("r_time", "ratio_t", 8), + ("pred_ram", "pred_ram_gb", 12), + ("act_ram", "act_ram_gb", 12), + ("r_ram", "ratio_r", 8), + ("pred_vram", "pred_vram_gb", 13), + ("act_vram", "act_vram_gb", 13), + ("r_vram", "ratio_v", 8), +] + + +def _fmt_cell(value: Any) -> str: # noqa: ANN401 + if value is None: + return "-" + if isinstance(value, float): + if value == 0: + return "0.00" + return f"{value:.2f}" if abs(value) >= 0.01 else f"{value:.4f}" + return str(value) + + +def _print_summary(rows: list[CalibrationRow]) -> None: + """Pretty side-by-side table for stdout.""" + header = " ".join(label.ljust(width) for _, label, width in _COLS) + print(header) + print("-" * len(header)) + for row in rows: + cells = { + "preset": row.preset, + "pred_time": row.predicted.get("time_h"), + "act_time": row.actual.get("time_h"), + "r_time": row.ratios.get("time"), + "pred_ram": row.predicted.get("ram_gb"), + "act_ram": row.actual.get("ram_gb"), + "r_ram": row.ratios.get("ram"), + "pred_vram": row.predicted.get("vram_gb"), + "act_vram": row.actual.get("vram_gb"), + "r_vram": row.ratios.get("vram"), + } + print(" ".join(_fmt_cell(cells[key]).ljust(width) for key, _, width in _COLS)) + if row.error: + print(f" ! {row.error}") + for note in row.notes: + print(f" * {note}") + for mod in row.modules: + duration = mod.get("duration_s") + vram = mod.get("peak_vram_gb") + duration_s = f"{duration:.2f}s" if duration is not None else "-" + vram_s = f"{vram:.2f} GB" if vram is not None else "-" + print(f" · {mod.get('module', '?')}#{mod.get('num', '?')} {duration_s} vram={vram_s}") + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s %(name)s: %(message)s", + ) + + presets = args.presets or list(BUNDLED_PRESETS) + unknown = [p for p in presets if p not in BUNDLED_PRESETS] + if unknown: + parser.error(f"Unknown preset(s): {', '.join(unknown)}. Known: {', '.join(BUNDLED_PRESETS)}") + + dataset_path = Path(args.dataset) + if dataset_path.is_file(): + logger.info("Loading dataset from local file %s", dataset_path) + dataset = Dataset.from_json(dataset_path) + dataset_source = str(dataset_path) + else: + logger.info("Loading dataset from HF Hub: %s", args.dataset) + try: + dataset = Dataset.from_hub(args.dataset) + except Exception as e: # noqa: BLE001 + parser.error(f"Could not load '{args.dataset}' as a local JSON file or as a Hub repo id: {e}") + dataset_source = f"hub:{args.dataset}" + stats = stats_from_dataset_obj(dataset) + hardware = detect_hardware() + logger.info( + "Hardware: %s (%s) — %.1f GB VRAM, %.0f GB RAM, %.0f GB free disk", + hardware.accelerator, + hardware.device_name, + hardware.vram_gb, + hardware.ram_gb, + hardware.free_disk_gb, + ) + + rows: list[CalibrationRow] = [] + for preset in presets: + logger.info("=== %s ===", preset) + row = _calibrate_one( + preset=preset, + dataset=dataset, + stats=stats, + hardware=hardware, + max_trials=args.max_trials, + skip_fit=args.skip_fit, + poll_interval_ms=args.poll_interval_ms, + enable_wandb=args.wandb, + ) + rows.append(row) + + payload = { + "hardware": { + "accelerator": hardware.accelerator, + "device_name": hardware.device_name, + "vram_gb": hardware.vram_gb, + "ram_gb": hardware.ram_gb, + "free_disk_gb": hardware.free_disk_gb, + }, + "dataset": { + "path": dataset_source, + "n_samples": stats.n_samples, + "n_classes": stats.n_classes, + "avg_tokens": stats.avg_tokens, + "multilabel": stats.multilabel, + }, + "max_trials_override": args.max_trials, + "skip_fit": args.skip_fit, + "rows": [r.to_dict() for r in rows], + } + args.output.write_text(json.dumps(payload, indent=2, default=str)) + logger.info("Wrote report to %s", args.output) + + print() + _print_summary(rows) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/run_calibration_banking77.sh b/scripts/run_calibration_banking77.sh new file mode 100755 index 000000000..5450c6716 --- /dev/null +++ b/scripts/run_calibration_banking77.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# Run the advisor calibration across every bundled preset on DeepPavlov/banking77. +# +# WARNING: transformers-heavy on banking77 (10k train samples, 77 classes) can +# take *many* hours on a single GPU. Set MAX_TRIALS to a small number for a +# fast sanity check, or leave it unset to let each preset use its bundled +# ``hpo_config.n_trials``. +# +# Environment overrides: +# DATASET HF Hub repo id (default: DeepPavlov/banking77) +# PRESETS Space-separated preset names (default: every bundled preset) +# MAX_TRIALS Cap for hpo_config.n_trials (default: unset -> preset default) +# WANDB If non-empty, pass --wandb so system metrics land in wandb.ai +# OUTPUT_DIR Where JSON reports + logs land (default: ./calibration_runs) +# SKIP_FIT If non-empty, only run preflight (no real fit) for a fast sanity check +# +# Examples: +# scripts/run_calibration_banking77.sh # full sweep +# MAX_TRIALS=3 scripts/run_calibration_banking77.sh # quick sweep +# PRESETS="classic-light nn-medium" scripts/run_calibration_banking77.sh +# WANDB=1 MAX_TRIALS=5 scripts/run_calibration_banking77.sh + +set -euo pipefail + +# Resolve repo root even when the script is called from anywhere. +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$REPO_ROOT" + +DATASET="${DATASET:-DeepPavlov/banking77}" +OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/calibration_runs}" +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" +OUTPUT_JSON="$OUTPUT_DIR/banking77_$TIMESTAMP.json" +LOG_FILE="$OUTPUT_DIR/banking77_$TIMESTAMP.log" + +mkdir -p "$OUTPUT_DIR" + +# Assemble optional flags. +EXTRA_FLAGS=() +if [[ -n "${MAX_TRIALS:-}" ]]; then + EXTRA_FLAGS+=("--max-trials" "$MAX_TRIALS") +fi +if [[ -n "${WANDB:-}" ]]; then + EXTRA_FLAGS+=("--wandb") +fi +if [[ -n "${SKIP_FIT:-}" ]]; then + EXTRA_FLAGS+=("--skip-fit") +fi + +# Preset list: pull it from the advisor package at runtime unless overridden, +# so the script auto-discovers presets that are added later. +if [[ -n "${PRESETS:-}" ]]; then + # shellcheck disable=SC2206 # intentional word-split from env + PRESET_ARR=($PRESETS) +else + PRESET_ARR=() + while IFS= read -r preset; do + PRESET_ARR+=("$preset") + done < <( +python - <<'PY' +from autointent._advisor import BUNDLED_PRESETS +for name in BUNDLED_PRESETS: + print(name) +PY + ) +fi + +echo "Repo: $REPO_ROOT" +echo "Dataset: $DATASET" +echo "Presets: ${PRESET_ARR[*]}" +echo "Output: $OUTPUT_JSON" +echo "Log: $LOG_FILE" +echo "Flags: ${EXTRA_FLAGS[*]:-}" +echo + +uv run --no-sync python scripts/calibrate_advisor.py \ + --dataset "$DATASET" \ + --presets "${PRESET_ARR[@]}" \ + --output "$OUTPUT_JSON" \ + "${EXTRA_FLAGS[@]}" \ + 2>&1 | tee "$LOG_FILE" + +echo +echo "Done. JSON: $OUTPUT_JSON" +echo " Log: $LOG_FILE" diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py new file mode 100644 index 000000000..1bd1f26db --- /dev/null +++ b/tests/pipeline/test_preflight.py @@ -0,0 +1,122 @@ +"""Pipeline.fit preflight integration: off / warn / strict modes.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import pytest + +from autointent import Pipeline +from autointent._advisor import HardwareProfile, detect_hardware, run_preflight, stats_from_dataset_obj +from autointent._pipeline import PreflightError +from autointent.configs import LoggingConfig + +if TYPE_CHECKING: + from autointent import Dataset + + +def _tiny_hw() -> HardwareProfile: + """Deterministic, intentionally-infeasible hardware budget.""" + return HardwareProfile( + accelerator="cuda", + device_name="test-tiny", + vram_gb=0.1, + ram_gb=0.5, + free_disk_gb=1.0, + cpu_count=2, + ) + + +def _classic_light_pipeline() -> Pipeline: + p = Pipeline.from_preset("classic-light") + p.set_config(LoggingConfig(dump_modules=False, clear_ram=True)) + return p + + +def test_preflight_off_skips_advisor(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + """preflight='off' must not run the advisor (no Preflight log line).""" + p = _classic_light_pipeline() + with caplog.at_level(logging.INFO, logger="autointent._pipeline._pipeline"): + try: + p.fit(dataset, preflight="off") + except Exception: # noqa: BLE001 — fit may fail in test env; we only care about preflight side effect + pass + assert not any("Preflight" in r.getMessage() for r in caplog.records) + + +def test_preflight_warn_logs_findings(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + """preflight='warn' logs a Preflight verdict line.""" + p = _classic_light_pipeline() + with caplog.at_level(logging.INFO, logger="autointent._pipeline._pipeline"): + try: + p.fit(dataset, preflight="warn") + except Exception: # noqa: BLE001 + pass + msgs = [r.getMessage() for r in caplog.records] + assert any("Preflight" in m and "verdict=" in m for m in msgs) + + +def test_preflight_strict_raises_on_infeasible( + dataset: Dataset, monkeypatch: pytest.MonkeyPatch +) -> None: + """preflight='strict' raises PreflightError when findings include OVER. + + Forces a tiny hardware budget so even cheap presets blow it. + """ + monkeypatch.setattr("autointent._pipeline._pipeline.detect_hardware", _tiny_hw) + p = _classic_light_pipeline() + with pytest.raises(PreflightError) as exc_info: + p.fit(dataset, preflight="strict") + assert exc_info.value.findings + assert all(f.severity.value == "over" for f in exc_info.value.findings) + + +def test_preflight_warn_does_not_raise_on_infeasible( + dataset: Dataset, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """Tiny hardware + warn mode logs an ERROR but doesn't raise.""" + monkeypatch.setattr("autointent._pipeline._pipeline.detect_hardware", _tiny_hw) + p = _classic_light_pipeline() + with caplog.at_level(logging.ERROR, logger="autointent._pipeline._pipeline"): + try: + p.fit(dataset, preflight="warn") + except PreflightError: + pytest.fail("warn mode must not raise PreflightError") + except Exception: # noqa: BLE001 — downstream fit errors are out of scope + pass + assert any(r.levelno == logging.ERROR for r in caplog.records) + + +def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: + """End-to-end integration: Pipeline -> _build_advisor_config -> run_preflight. + + Asserts the round-trip is wired correctly: the dict ``Pipeline`` exposes to + the advisor validates against ``OptimizationConfig``, the advisor produces a + well-formed report, and the driver list reflects the actual modules from the + preset's search space (not silently empty). + """ + p = _classic_light_pipeline() + config = p._build_advisor_config() # noqa: SLF001 + stats = stats_from_dataset_obj(dataset) + hardware = detect_hardware() + + report = run_preflight(config, stats, hardware, preset_name="classic-light") + + # The advisor accepted the pipeline-built config and produced findings. + assert report.preset_name == "classic-light" + assert report.resource.drivers, "expected at least one driver row for classic-light" + + # classic-light's scoring node has knn / linear / mlknn — at least linear + # should always end up in drivers (knn variants don't always carry an + # explicit model_name, so they're allowed to be absent). + driver_modules = {d["module"] for d in report.resource.drivers} + assert "linear" in driver_modules, f"missing linear scorer in drivers: {driver_modules}" + + # The advisor must always emit the three resource findings. + metrics = {f.metric for f in report.findings if f.metric} + assert {"vram", "ram", "disk"} <= metrics, f"missing required metrics: {metrics}" + + # Dataset stats round-trip into the report. + assert report.dataset["n_samples"] == stats.n_samples + assert report.dataset["n_classes"] == stats.n_classes From 7eb73f83cb185b3c11be498a41a31be04a0b1588 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:53:33 +0300 Subject: [PATCH 23/43] upd w&b --- scripts/calibrate_advisor.py | 34 +++++++++++++++------------- scripts/run_calibration_banking77.sh | 2 ++ 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py index 13dd76836..9e2746a57 100644 --- a/scripts/calibrate_advisor.py +++ b/scripts/calibrate_advisor.py @@ -274,12 +274,12 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None: # noqa: ARG002 pass -def _attach_tracker(pipeline: Pipeline, tracker: _ModuleTracker) -> None: - """Instance-patch ``pipeline._fit`` so ``tracker`` is appended to the callback chain.""" +def _attach_callbacks(pipeline: Pipeline, callbacks: list[OptimizerCallback]) -> None: + """Instance-patch ``pipeline._fit`` to append ``callbacks`` to the callback chain.""" original_fit = pipeline._fit # noqa: SLF001 def patched(context: Any) -> Any: # noqa: ANN401 - context.callback_handler.callbacks.append(tracker) + context.callback_handler.callbacks.extend(callbacks) return original_fit(context) pipeline._fit = patched # type: ignore[method-assign] # noqa: SLF001 @@ -288,16 +288,10 @@ def patched(context: Any) -> Any: # noqa: ANN401 # === per-preset run ====================================================== -def _override_trials(pipeline: Pipeline, max_trials: int | None, *, enable_wandb: bool) -> None: - """Cap n_trials, disable dumping, optionally enable W&B for post-run analysis.""" - updates: dict[str, Any] = {} +def _override_trials(pipeline: Pipeline, max_trials: int | None) -> None: + """Cap n_trials and disable dumping for the calibration run.""" if max_trials is not None: - updates["n_trials"] = max_trials - if enable_wandb: - # Trigger built-in per-run system-metrics collection in W&B. - updates["report_to"] = ["wandb"] - if updates: - pipeline.set_config(pipeline.hpo_config.model_copy(update=updates)) + pipeline.set_config(pipeline.hpo_config.model_copy(update={"n_trials": max_trials})) # We don't want the calibration run to leave dumped module weights on disk. pipeline.set_config(LoggingConfig(dump_modules=False, clear_ram=True)) @@ -322,7 +316,7 @@ def _calibrate_one( row.error = f"from_preset failed: {e}" return row - _override_trials(pipeline, max_trials, enable_wandb=enable_wandb) + _override_trials(pipeline, max_trials) try: report: PreflightReport = run_preflight( @@ -356,7 +350,15 @@ def _calibrate_one( _reset_vram_peak() tracker = _ModuleTracker() - _attach_tracker(pipeline, tracker) + callbacks: list[OptimizerCallback] = [tracker] + if enable_wandb: + try: + from autointent._callbacks.wandb import WandbCallback + + callbacks.append(WandbCallback()) + except ImportError as e: + row.notes.append(f"W&B requested but not available: {e}") + _attach_callbacks(pipeline, callbacks) is_mps = hardware.accelerator == "mps" start = time.perf_counter() @@ -384,8 +386,8 @@ def _calibrate_one( "disk_download_gb": actual_disk_download_gb, } row.modules = tracker.records - if enable_wandb: - row.notes.append("W&B reporter enabled — inspect wandb.ai run group for per-step GPU/system metrics") + if enable_wandb and not any("W&B requested but not available" in n for n in row.notes): + row.notes.append("W&B reporter attached — inspect wandb.ai run group for per-step GPU/system metrics") def _ratio(actual: float | None, predicted: float) -> float | None: if actual is None or predicted <= 0: diff --git a/scripts/run_calibration_banking77.sh b/scripts/run_calibration_banking77.sh index 5450c6716..460d4432c 100755 --- a/scripts/run_calibration_banking77.sh +++ b/scripts/run_calibration_banking77.sh @@ -32,6 +32,8 @@ TIMESTAMP="$(date +%Y%m%d_%H%M%S)" OUTPUT_JSON="$OUTPUT_DIR/banking77_$TIMESTAMP.json" LOG_FILE="$OUTPUT_DIR/banking77_$TIMESTAMP.log" +export WANDB_PROJECT="autointent_feasibility" + mkdir -p "$OUTPUT_DIR" # Assemble optional flags. From 953262b0e1175690db9df5e373fedaa4df0f7280 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:58:00 +0300 Subject: [PATCH 24/43] improve w&b run name --- scripts/calibrate_advisor.py | 23 +++++++++++++++++++---- scripts/run_calibration_banking77.sh | 7 +++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py index 9e2746a57..a908bfbff 100644 --- a/scripts/calibrate_advisor.py +++ b/scripts/calibrate_advisor.py @@ -110,6 +110,16 @@ def _build_parser() -> argparse.ArgumentParser: "Requires ``wandb`` installed + ``WANDB_API_KEY`` in the environment." ), ) + p.add_argument( + "--run-name", + type=str, + default=None, + help=( + "Suffix appended to each preset's LoggingConfig.run_name — the resulting " + "value is ``{preset}_{run_name}`` and is used by the LoggingHandler as the " + "W&B run group / on-disk dump directory name." + ), + ) p.add_argument("-v", "--verbose", action="store_true") return p @@ -288,12 +298,14 @@ def patched(context: Any) -> Any: # noqa: ANN401 # === per-preset run ====================================================== -def _override_trials(pipeline: Pipeline, max_trials: int | None) -> None: - """Cap n_trials and disable dumping for the calibration run.""" +def _override_trials(pipeline: Pipeline, max_trials: int | None, *, run_name: str | None = None) -> None: + """Cap n_trials and disable dumping. When ``run_name`` is set, tag the + ``LoggingConfig.run_name`` (used by W&B run groups / dump dir names).""" if max_trials is not None: pipeline.set_config(pipeline.hpo_config.model_copy(update={"n_trials": max_trials})) # We don't want the calibration run to leave dumped module weights on disk. - pipeline.set_config(LoggingConfig(dump_modules=False, clear_ram=True)) + logging_config = LoggingConfig(dump_modules=False, clear_ram=True, run_name=run_name) + pipeline.set_config(logging_config) def _calibrate_one( @@ -306,6 +318,7 @@ def _calibrate_one( skip_fit: bool, poll_interval_ms: int, enable_wandb: bool, + run_name: str | None, ) -> CalibrationRow: row = CalibrationRow(preset=preset) @@ -316,7 +329,8 @@ def _calibrate_one( row.error = f"from_preset failed: {e}" return row - _override_trials(pipeline, max_trials) + tagged_run_name = f"{preset}_{run_name}" if run_name else None + _override_trials(pipeline, max_trials, run_name=tagged_run_name) try: report: PreflightReport = run_preflight( @@ -509,6 +523,7 @@ def main(argv: list[str] | None = None) -> int: skip_fit=args.skip_fit, poll_interval_ms=args.poll_interval_ms, enable_wandb=args.wandb, + run_name=args.run_name, ) rows.append(row) diff --git a/scripts/run_calibration_banking77.sh b/scripts/run_calibration_banking77.sh index 460d4432c..1bcf2add2 100755 --- a/scripts/run_calibration_banking77.sh +++ b/scripts/run_calibration_banking77.sh @@ -11,6 +11,9 @@ # PRESETS Space-separated preset names (default: every bundled preset) # MAX_TRIALS Cap for hpo_config.n_trials (default: unset -> preset default) # WANDB If non-empty, pass --wandb so system metrics land in wandb.ai +# RUN_NAME Suffix appended to each preset's LoggingConfig.run_name — the +# resulting name is ``{preset}_{RUN_NAME}`` (default: unset -> +# autointent generates a random name) # OUTPUT_DIR Where JSON reports + logs land (default: ./calibration_runs) # SKIP_FIT If non-empty, only run preflight (no real fit) for a fast sanity check # @@ -19,6 +22,7 @@ # MAX_TRIALS=3 scripts/run_calibration_banking77.sh # quick sweep # PRESETS="classic-light nn-medium" scripts/run_calibration_banking77.sh # WANDB=1 MAX_TRIALS=5 scripts/run_calibration_banking77.sh +# RUN_NAME=calib_2026_07 WANDB=1 scripts/run_calibration_banking77.sh set -euo pipefail @@ -47,6 +51,9 @@ fi if [[ -n "${SKIP_FIT:-}" ]]; then EXTRA_FLAGS+=("--skip-fit") fi +if [[ -n "${RUN_NAME:-}" ]]; then + EXTRA_FLAGS+=("--run-name" "$RUN_NAME") +fi # Preset list: pull it from the advisor package at runtime unless overridden, # so the script auto-discovers presets that are added later. From ac8cc67dde9d4371af2a73049462da564f8f9867 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 7 Jul 2026 02:40:11 +0300 Subject: [PATCH 25/43] cap cpu --- scripts/calibrate_advisor.py | 64 +++++++++++++++++------ scripts/run_calibration_banking77.sh | 54 +++++++++++++------ src/autointent/_advisor/runner.py | 4 +- src/autointent/_pipeline/_pipeline.py | 4 +- tests/advisor/test_estimates_internals.py | 5 +- tests/pipeline/test_preflight.py | 6 +-- 6 files changed, 93 insertions(+), 44 deletions(-) diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py index a908bfbff..6057fc0dd 100644 --- a/scripts/calibrate_advisor.py +++ b/scripts/calibrate_advisor.py @@ -43,7 +43,11 @@ stats_from_dataset_obj, ) from autointent._callbacks.base import OptimizerCallback -from autointent.configs import HPOConfig, LoggingConfig +from autointent.configs import LoggingConfig +from autointent import setup_logging + +setup_logging("INFO", log_filename="logs.log") +logging.basicConfig(level=logging.INFO) logger = logging.getLogger("calibrate_advisor") @@ -175,14 +179,13 @@ def __exit__(self, *_exc: object) -> None: def _run(self) -> None: try: - import torch # noqa: PLC0415 + import torch except ImportError: torch = None # type: ignore[assignment] while not self._stop.is_set(): try: rss = self._proc.memory_info().rss / _BYTES_PER_GB - if rss > self.peak_ram_gb: - self.peak_ram_gb = rss + self.peak_ram_gb = max(self.peak_ram_gb, rss) if self._sample_mps and torch is not None: mps = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB if self.peak_mps_gb is None or mps > self.peak_mps_gb: @@ -227,11 +230,11 @@ class _ModuleTracker(OptimizerCallback): name = "calibration_tracker" - def __init__(self) -> None: # noqa: D401 + def __init__(self) -> None: self.records: list[dict[str, Any]] = [] self._current: dict[str, Any] | None = None - def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: # noqa: ARG002 + def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: pass def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None: @@ -243,9 +246,7 @@ def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any] except ImportError: pass # Only capture JSON-safe scalars in the config snapshot. - safe_config = { - k: v for k, v in module_kwargs.items() if isinstance(v, (str, int, float, bool)) or v is None - } + safe_config = {k: v for k, v in module_kwargs.items() if isinstance(v, (str, int, float, bool)) or v is None} self._current = { "module": module_name, "num": num, @@ -253,10 +254,10 @@ def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any] "_start": time.perf_counter(), } - def log_value(self, **kwargs: Any) -> None: # noqa: ANN401, ARG002 + def log_value(self, **kwargs: Any) -> None: # noqa: ANN401 pass - def log_metrics(self, metrics: dict[str, Any]) -> None: # noqa: ARG002 + def log_metrics(self, metrics: dict[str, Any]) -> None: pass def end_module(self) -> None: @@ -280,7 +281,7 @@ def end_module(self) -> None: def end_run(self) -> None: pass - def log_final_metrics(self, metrics: dict[str, Any]) -> None: # noqa: ARG002 + def log_final_metrics(self, metrics: dict[str, Any]) -> None: pass @@ -299,11 +300,15 @@ def patched(context: Any) -> Any: # noqa: ANN401 def _override_trials(pipeline: Pipeline, max_trials: int | None, *, run_name: str | None = None) -> None: - """Cap n_trials and disable dumping. When ``run_name`` is set, tag the - ``LoggingConfig.run_name`` (used by W&B run groups / dump dir names).""" + """Cap n_trials, force ``n_jobs=1`` (serial HPO to keep wall-time measurements clean + and to prevent CPU oversubscription with sklearn's own ``n_jobs``), disable dumping. + When ``run_name`` is set, tag ``LoggingConfig.run_name`` (used as the W&B group / + dump-dir name). + """ + updates: dict[str, Any] = {"n_jobs": 1} if max_trials is not None: - pipeline.set_config(pipeline.hpo_config.model_copy(update={"n_trials": max_trials})) - # We don't want the calibration run to leave dumped module weights on disk. + updates["n_trials"] = max_trials + pipeline.set_config(pipeline.hpo_config.model_copy(update=updates)) logging_config = LoggingConfig(dump_modules=False, clear_ram=True, run_name=run_name) pipeline.set_config(logging_config) @@ -475,6 +480,25 @@ def _print_summary(rows: list[CalibrationRow]) -> None: print(f" · {mod.get('module', '?')}#{mod.get('num', '?')} {duration_s} vram={vram_s}") +def _apply_thread_cap() -> None: + """Cap torch intra-op threads to the same value as OMP_NUM_THREADS. + + Env vars (OMP/MKL/OpenBLAS) MUST be set before Python starts to be effective — + that's the bash wrapper's job. This function is belt-and-braces: torch reads + OMP_NUM_THREADS on init, but ``set_num_threads`` also caps its C++ intra-op + pool if a caller forgets the env var. + """ + n = int(os.environ.get("OMP_NUM_THREADS", "0") or 0) + if n <= 0: + return + try: + import torch + + torch.set_num_threads(n) + except ImportError: + pass + + def main(argv: list[str] | None = None) -> int: parser = _build_parser() args = parser.parse_args(argv) @@ -482,6 +506,7 @@ def main(argv: list[str] | None = None) -> int: level=logging.DEBUG if args.verbose else logging.INFO, format="%(levelname)s %(name)s: %(message)s", ) + _apply_thread_cap() presets = args.presets or list(BUNDLED_PRESETS) unknown = [p for p in presets if p not in BUNDLED_PRESETS] @@ -510,6 +535,13 @@ def main(argv: list[str] | None = None) -> int: hardware.ram_gb, hardware.free_disk_gb, ) + logger.info( + "Thread caps: OMP=%s MKL=%s OPENBLAS=%s TOKENIZERS_PARALLELISM=%s", + os.environ.get("OMP_NUM_THREADS", ""), + os.environ.get("MKL_NUM_THREADS", ""), + os.environ.get("OPENBLAS_NUM_THREADS", ""), + os.environ.get("TOKENIZERS_PARALLELISM", ""), + ) rows: list[CalibrationRow] = [] for preset in presets: diff --git a/scripts/run_calibration_banking77.sh b/scripts/run_calibration_banking77.sh index 1bcf2add2..7db4d9120 100755 --- a/scripts/run_calibration_banking77.sh +++ b/scripts/run_calibration_banking77.sh @@ -7,22 +7,27 @@ # ``hpo_config.n_trials``. # # Environment overrides: -# DATASET HF Hub repo id (default: DeepPavlov/banking77) -# PRESETS Space-separated preset names (default: every bundled preset) -# MAX_TRIALS Cap for hpo_config.n_trials (default: unset -> preset default) -# WANDB If non-empty, pass --wandb so system metrics land in wandb.ai -# RUN_NAME Suffix appended to each preset's LoggingConfig.run_name — the -# resulting name is ``{preset}_{RUN_NAME}`` (default: unset -> -# autointent generates a random name) -# OUTPUT_DIR Where JSON reports + logs land (default: ./calibration_runs) -# SKIP_FIT If non-empty, only run preflight (no real fit) for a fast sanity check +# DATASET HF Hub repo id (default: DeepPavlov/banking77) +# PRESETS Space-separated preset names (default: every bundled preset) +# MAX_TRIALS Cap for hpo_config.n_trials (default: unset -> preset default) +# WANDB If non-empty, pass --wandb so system metrics land in wandb.ai +# RUN_NAME Suffix appended to each preset's LoggingConfig.run_name — the +# resulting name is ``{preset}_{RUN_NAME}`` (default: unset -> +# autointent generates a random name) +# OUTPUT_DIR Where JSON reports + logs land (default: ./calibration_runs) +# SKIP_FIT If non-empty, only run preflight (no real fit) — fast sanity check +# THREADS_PER_JOB Cap for BLAS/OpenMP/torch intra-op threads per HPO trial +# (default: 1). Increase carefully — sklearn's own ``n_jobs`` and +# HPO parallelism multiply on top, so oversubscription is easy +# on many-core boxes. # # Examples: -# scripts/run_calibration_banking77.sh # full sweep +# scripts/run_calibration_banking77.sh # full sweep, serial # MAX_TRIALS=3 scripts/run_calibration_banking77.sh # quick sweep # PRESETS="classic-light nn-medium" scripts/run_calibration_banking77.sh # WANDB=1 MAX_TRIALS=5 scripts/run_calibration_banking77.sh # RUN_NAME=calib_2026_07 WANDB=1 scripts/run_calibration_banking77.sh +# THREADS_PER_JOB=4 scripts/run_calibration_banking77.sh # 4-thread BLAS set -euo pipefail @@ -38,6 +43,22 @@ LOG_FILE="$OUTPUT_DIR/banking77_$TIMESTAMP.log" export WANDB_PROJECT="autointent_feasibility" +# --------------------------------------------------------------------------- +# CPU thread caps — set BEFORE python starts, because numpy/torch/sklearn read +# them at import time. Without these, on a 16+ core box each BLAS-backed +# operation defaults to N-thread pools which multiply with sklearn's own +# ``n_jobs`` and HPO parallelism → the machine oversubscribes and stalls. +# --------------------------------------------------------------------------- +THREADS_PER_JOB="${THREADS_PER_JOB:-1}" +export OMP_NUM_THREADS="$THREADS_PER_JOB" +export MKL_NUM_THREADS="$THREADS_PER_JOB" +export OPENBLAS_NUM_THREADS="$THREADS_PER_JOB" +export NUMEXPR_NUM_THREADS="$THREADS_PER_JOB" +# HF tokenizers deadlock on fork if left in parallel mode. +export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}" +# torch reads OMP_NUM_THREADS for intra-op, but set explicitly too — belt-and-braces. +export PYTORCH_NUM_THREADS="$THREADS_PER_JOB" + mkdir -p "$OUTPUT_DIR" # Assemble optional flags. @@ -73,12 +94,13 @@ PY ) fi -echo "Repo: $REPO_ROOT" -echo "Dataset: $DATASET" -echo "Presets: ${PRESET_ARR[*]}" -echo "Output: $OUTPUT_JSON" -echo "Log: $LOG_FILE" -echo "Flags: ${EXTRA_FLAGS[*]:-}" +echo "Repo: $REPO_ROOT" +echo "Dataset: $DATASET" +echo "Presets: ${PRESET_ARR[*]}" +echo "Output: $OUTPUT_JSON" +echo "Log: $LOG_FILE" +echo "Flags: ${EXTRA_FLAGS[*]:-}" +echo "Threads per job: $THREADS_PER_JOB (OMP/MKL/OpenBLAS/torch)" echo uv run --no-sync python scripts/calibrate_advisor.py \ diff --git a/src/autointent/_advisor/runner.py b/src/autointent/_advisor/runner.py index 168157550..0a9a56f02 100644 --- a/src/autointent/_advisor/runner.py +++ b/src/autointent/_advisor/runner.py @@ -154,9 +154,7 @@ def _data_phase( # the search space. Multilabel uses LogisticRegression (no CV), so skip there. if not stats.multilabel and stats.class_counts: linear_cvs = [ - _max_int(e.get("cv"), 3) - for _, e in _walk_modules(search_space) - if e.get("module_name") == "linear" + _max_int(e.get("cv"), 3) for _, e in _walk_modules(search_space) if e.get("module_name") == "linear" ] if linear_cvs: cv_max = max(linear_cvs) diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index 2fe352ebd..491bc57f4 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -545,7 +545,9 @@ def _log_preflight_report(report: PreflightReport, logger: logging.Logger) -> No Severity.TIGHT: logging.WARNING, Severity.OVER: logging.ERROR, } - header = f"Preflight ({report.preset_name or 'pipeline'}): verdict={'feasible' if report.is_feasible else 'INFEASIBLE'}" + header = ( + f"Preflight ({report.preset_name or 'pipeline'}): verdict={'feasible' if report.is_feasible else 'INFEASIBLE'}" + ) logger.info(header) for finding in report.findings: logger.log(level_for[finding.severity], "[%s] %s", finding.phase, finding.message) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index a62bce45b..fdcdf45f7 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -299,10 +299,7 @@ def test_rare_classes_threshold_follows_entry_cv(self) -> None: class_counts={"intent_a": 4, "intent_b": 8, "intent_c": 8}, ) report = run_preflight(cfg, stats, _profile()) - assert any( - f.phase == "data" and "cv=5" in f.message and "intent_a" in f.message - for f in report.findings - ) + assert any(f.phase == "data" and "cv=5" in f.message and "intent_a" in f.message for f in report.findings) def test_truncation_red_when_p95_dominates_max_length(self) -> None: cfg = { diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py index 1bd1f26db..b9116fbe1 100644 --- a/tests/pipeline/test_preflight.py +++ b/tests/pipeline/test_preflight.py @@ -57,9 +57,7 @@ def test_preflight_warn_logs_findings(dataset: Dataset, caplog: pytest.LogCaptur assert any("Preflight" in m and "verdict=" in m for m in msgs) -def test_preflight_strict_raises_on_infeasible( - dataset: Dataset, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_preflight_strict_raises_on_infeasible(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: """preflight='strict' raises PreflightError when findings include OVER. Forces a tiny hardware budget so even cheap presets blow it. @@ -97,7 +95,7 @@ def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: preset's search space (not silently empty). """ p = _classic_light_pipeline() - config = p._build_advisor_config() # noqa: SLF001 + config = p._build_advisor_config() stats = stats_from_dataset_obj(dataset) hardware = detect_hardware() From 970a7f387ff15fc84c42b5e6feae036819d44285 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 4 Aug 2026 12:16:11 +0300 Subject: [PATCH 26/43] address follow-up review P0/P1/P2 items across advisor + calibrator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Advisor (src/autointent/_advisor/): - fix linear/classic time formula and transformer VRAM under-prediction (34 B/token/layer upper bound; batch-scaled) - device-class per-step transformer time lookup - optional embedding-cache warmth probe: predict 0 forward + 0 disk_embedding_cache when caller certifies warmth - stop silent-zero estimates for cnn/rnn/sklearn (emit not-estimated row) - count cross-encoder / reranker downloads via cross_encoder_config + transformer_config fallback - conservative + loud low-confidence fallback (large-model defaults; TIGHT finding instead of buried note) - new reduce_to_fit + ReduceToFitError workflow with empty-scoring guard Calibrator (scripts/calibrate_advisor.py + run_calibration_banking77.sh): - fix CUDA VRAM measurement (per-module reset was clobbering peak); ratios computed at serialization time - --clear-embedding-cache + cache-policy tagging - --budget-vram-gb, --require-cuda, --subsample-per-class, --repeats, --dataset nargs="+" for constrained-hardware and shape sweeps - incremental atomic per-preset JSON checkpoint - role classification (embedder/scorer/decision) + time_by_role_s - optional-extras skip (peft/catboost/openai) — clean skip row instead of fit-failed - --presets accepts .yaml paths; coverage_preset.yaml packs lora, ptuning, dnnc, gcn, description_cross for module coverage - in-process CLI smoke (autointent-advisor inspect --json) compared to direct-API report every preset - per-step timing captured via monkey-patched HF Trainer callback → step_timings on each module record Tests: 111 passing (test_reduce_to_fit + test_calibration_tracker new). Co-Authored-By: Claude Opus 4.7 --- scripts/calibrate_advisor.py | 828 ++++++++++++++++-- scripts/coverage_preset.yaml | 42 + scripts/run_calibration_banking77.sh | 54 +- src/autointent/_advisor/__init__.py | 4 + .../_advisor/_estimates/_formulas.py | 117 ++- .../_advisor/_estimates/_resource.py | 121 ++- src/autointent/_advisor/_hub.py | 27 +- src/autointent/_advisor/runner.py | 12 +- src/autointent/_advisor/workflows.py | 156 ++++ tests/advisor/test_estimates_and_cli.py | 24 +- tests/advisor/test_estimates_internals.py | 47 +- tests/advisor/test_reduce_to_fit.py | 164 ++++ tests/pipeline/test_calibration_tracker.py | 116 +++ 13 files changed, 1585 insertions(+), 127 deletions(-) create mode 100644 scripts/coverage_preset.yaml create mode 100644 tests/advisor/test_reduce_to_fit.py create mode 100644 tests/pipeline/test_calibration_tracker.py diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py index 6057fc0dd..3f18b1cec 100644 --- a/scripts/calibrate_advisor.py +++ b/scripts/calibrate_advisor.py @@ -61,16 +61,55 @@ class CalibrationRow: preset: str predicted: dict[str, float] = field(default_factory=dict) actual: dict[str, float | None] = field(default_factory=dict) - ratios: dict[str, float | None] = field(default_factory=dict) findings: int = 0 findings_over: int = 0 # Per-module records from _ModuleTracker: [{module, num, config, duration_s, peak_vram_gb?}, ...] modules: list[dict[str, Any]] = field(default_factory=list) + cache_policy: str = "unknown" # "cold" (embeddings cache cleared) | "warm" (kept as-is) + low_confidence: bool = False # advisor fell back to heuristic HF-metadata for one+ models + repeat_idx: int = 0 # 0-based index within a (preset, dataset) repeat group + # ``skipped`` is set when the preset needs an optional extra that isn't + # installed (peft / catboost / openai / ...). We still populate ``error`` + # for the summary, but callers analysing the JSON should treat + # ``skipped=True`` rows separately from ``error != None && skipped=False`` + # rows (real crashes) — the former are expected and shouldn't count as + # advisor failures. + skipped: bool = False + # Snapshot of ``autointent-advisor inspect --json`` run in-process + # under the same stats + budget as the direct-API preflight. Lets us catch + # a CLI-wrapper regression (JSON schema drift, feasibility verdict flip) + # without a separate subprocess round-trip. None means we didn't run it + # (e.g. skipped row, or CLI itself crashed — see notes for the reason). + cli_smoke: dict[str, Any] | None = None error: str | None = None notes: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: - return asdict(self) + """Serialize with ratios + role-decomposed timings computed at read time. + + Storing ratios in the row is a bug magnet: any late edit to + ``predicted``/``actual`` gets missed. We compute them from the current + row values at serialization time so consumers can trust ``row["ratios"]``. + ``time_by_role_s`` splits the measured wall-time across embedder / + scorer / decision so classic-preset time can be interpreted (embedder + forward vs sklearn fit) without re-walking ``modules``. + """ + payload = asdict(self) + payload["ratios"] = self._ratios() + payload["time_by_role_s"] = _sum_time_by_role(self.modules) + return payload + + def _ratios(self) -> dict[str, float | None]: + keys = ("time_h", "ram_gb", "vram_gb", "disk_download_gb", "disk_embedding_cache_gb") + out: dict[str, float | None] = {} + for key in keys: + actual = self.actual.get(key) + predicted = self.predicted.get(key) + if actual is None or predicted is None or predicted <= 0: + out[key] = None + else: + out[key] = actual / predicted + return out def _build_parser() -> argparse.ArgumentParser: @@ -81,17 +120,48 @@ def _build_parser() -> argparse.ArgumentParser: p.add_argument( "--dataset", required=True, + nargs="+", type=str, help=( - "Either a local JSON path (loaded via ``Dataset.from_json``) or an HF Hub repo id " - "such as ``DeepPavlov/banking77`` (loaded via ``Dataset.from_hub``)." + "One or more datasets — each is either a local JSON path (loaded via " + "``Dataset.from_json``) or an HF Hub repo id such as ``DeepPavlov/banking77`` " + "(loaded via ``Dataset.from_hub``). Every preset runs against every dataset, " + "so pairing a multilabel + long-token + small + large dataset exercises the " + "n_samples / n_classes / avg_tokens surfaces of the advisor's formulas." + ), + ) + p.add_argument( + "--subsample-per-class", + type=int, + default=None, + help=( + "Cap each class to at most N training samples (deterministic first-N slice) " + "before running. Lets one big dataset stand in as a 'small' shape — enough to " + "exercise ``LogisticRegressionCV cv=3`` split-readiness and rare-class findings." + ), + ) + p.add_argument( + "--repeats", + type=int, + default=1, + help=( + "Run each (preset, dataset) N times so ratio gaps have variance bars. " + "The summary prints mean ± stdev for the actual measurements across " + "repeats; individual repeat rows are still written to the JSON with " + "``repeat_idx`` so consumers can compute their own aggregates. Default: 1." ), ) p.add_argument( "--presets", nargs="+", default=None, - help="Preset names to run (default: every preset in BUNDLED_PRESETS).", + help=( + "Preset names to run (default: every preset in BUNDLED_PRESETS). " + "Items ending in .yaml/.yml are treated as paths to a preset file — " + "used to run e.g. ``scripts/coverage_preset.yaml`` which packs " + "lora/ptuning/dnnc/gcn/cross-encoder into one small run for module " + "coverage without touching the shipped presets." + ), ) p.add_argument("--output", type=Path, default=Path("calibration.json"), help="Where to write the JSON report.") p.add_argument("--max-trials", type=int, default=None, help="Override hpo_config.n_trials for faster runs.") @@ -124,6 +194,33 @@ def _build_parser() -> argparse.ArgumentParser: "W&B run group / on-disk dump directory name." ), ) + p.add_argument( + "--clear-embedding-cache", + action="store_true", + help=( + "Wipe ``/autointent/embeddings/`` before each preset so every " + "measurement reflects a COLD run (embedder forward not skipped). Without this " + "flag, the cross-run cache silently makes later runs look artificially cheap." + ), + ) + p.add_argument( + "--budget-vram-gb", + type=float, + default=None, + help=( + "Override the detected VRAM budget passed to run_preflight, e.g. ``--budget-vram-gb 8`` " + "to exercise the constrained-hardware / severity paths on a big box without needing " + "a small GPU. Does NOT affect the real fit — only the predicted-side estimate." + ), + ) + p.add_argument( + "--require-cuda", + action="store_true", + help=( + "Fail fast if PyTorch can't initialize CUDA (guards against the silent " + "'2 GPUs detected, but torch runs on CPU' driver-mismatch trap)." + ), + ) p.add_argument("-v", "--verbose", action="store_true") return p @@ -136,6 +233,33 @@ def _hf_cache_dir() -> Path: return Path(os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")) +def _embeddings_cache_dir() -> Path: + """Return autointent's embeddings-cache dir (``/autointent/embeddings/``). + + Uses the same ``appdirs.user_cache_dir("autointent")`` path as + :func:`autointent._wrappers.embedder.utils.get_embeddings_path` so the + harness reads/clears the same directory the runtime writes to. + """ + from autointent._wrappers.embedder.utils import get_embeddings_path + + return get_embeddings_path("_probe").parent + + +def _clear_embeddings_cache() -> int: + """Delete every ``*.npy`` file in the embeddings cache. Returns count removed.""" + cache = _embeddings_cache_dir() + if not cache.exists(): + return 0 + removed = 0 + for path in cache.glob("*.npy"): + try: + path.unlink() + removed += 1 + except OSError: + continue + return removed + + def _dir_size_gb(path: Path) -> float: """Disk usage of ``path`` in GB; 0 when the directory is missing.""" if not path.exists(): @@ -219,6 +343,114 @@ def _read_vram_peak_gb(accelerator: str) -> float | None: # === per-module tracking ================================================= +# Static module_name → role classification. Used to tag tracker records so +# downstream analysis can decompose classic-preset wall-time into +# embedder-forward vs scorer-fit vs decision-search — the follow-up review's +# R4-P1 #30 asked for this because today classic wall-time conflates the two. +_EMBEDDER_MODULE_NAMES = frozenset( + {"sentence_transformer", "openai_embedder", "vllm_embedder", "hashing_vectorizer"}, +) +_DECISION_MODULE_NAMES = frozenset({"threshold", "argmax", "jinoos", "tunable", "adaptive"}) + + +def _classify_module_role(module_name: str) -> str: + """Bucket module_name into ``embedder`` / ``decision`` / ``scorer``. + + Everything not in the known embedder or decision sets is treated as a + scorer — so newly added scorer modules land in the right bucket by default + and only new decision/embedder modules would need to update the sets. + """ + if module_name in _EMBEDDER_MODULE_NAMES: + return "embedder" + if module_name in _DECISION_MODULE_NAMES: + return "decision" + return "scorer" + + +class _StepTimingCallback: + """HF ``TrainerCallback`` that appends the wall-time of each optimizer step + to a caller-owned list. + + Injected into every ``transformers.Trainer`` for the duration of a fit via + :func:`_patch_trainer_for_step_timing`. The sink is the current module's + step buffer on :class:`_ModuleTracker`, so the transformer's per-step + latency lands in that module's record automatically — no plumbing across + module boundaries. + + Duck-typed (not a ``TrainerCallback`` subclass) so importing transformers + stays lazy — the harness must work on classic-only runs without the + transformers extra. + """ + + def __init__(self, sink: list[float]) -> None: + self._sink = sink + self._t0: float | None = None + + # HF's CallbackHandler calls these positionally with (args, state, control, **kwargs) + def on_step_begin(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None: # noqa: ANN401, ARG002 + self._t0 = time.perf_counter() + + def on_step_end(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None: # noqa: ANN401, ARG002 + if self._t0 is not None: + self._sink.append(time.perf_counter() - self._t0) + self._t0 = None + + # HF's CallbackHandler probes each callback with hasattr; leave the rest unset. + + +def _summarize_step_times(step_times: list[float]) -> dict[str, float]: + """Fold a list of per-step wall-times into summary stats for the row. + + ``seconds_per_step`` is what the advisor's transformer-time baseline + encodes (currently a flat ~1 s constant); logging measured ``mean`` and + ``p95`` lets the baseline be recalibrated directly from row data instead + of eyeballed off a wandb dashboard. + """ + import statistics as _stats + + if not step_times: + return {} + if len(step_times) == 1: + return {"n_steps": 1, "mean_step_s": step_times[0], "p95_step_s": step_times[0]} + sorted_st = sorted(step_times) + p95_idx = min(len(sorted_st) - 1, int(round(0.95 * (len(sorted_st) - 1)))) + return { + "n_steps": len(step_times), + "mean_step_s": _stats.fmean(step_times), + "p95_step_s": sorted_st[p95_idx], + "total_step_s": sum(step_times), + } + + +def _patch_trainer_for_step_timing(tracker: _ModuleTracker) -> Any: # noqa: ANN401 + """Monkey-patch ``transformers.Trainer.__init__`` to inject a step-timing + callback bound to ``tracker._current_step_buffer`` — the list on the + module record currently being tracked. + + Returns a callable that undoes the patch. No-ops (returns a no-op undoer) + when transformers isn't importable, so classic-only presets aren't blocked. + """ + try: + from transformers import Trainer # type: ignore[import-not-found] + except ImportError: + return lambda: None + + original_init = Trainer.__init__ + + def patched(self: Any, *args: Any, **kwargs: Any) -> None: # noqa: ANN401 + original_init(self, *args, **kwargs) + buffer = tracker.current_step_buffer() + if buffer is not None: + self.add_callback(_StepTimingCallback(buffer)) + + Trainer.__init__ = patched # type: ignore[method-assign] + + def _undo() -> None: + Trainer.__init__ = original_init # type: ignore[method-assign] + + return _undo + + class _ModuleTracker(OptimizerCallback): """Records per-module wall time and peak VRAM. @@ -226,6 +458,12 @@ class _ModuleTracker(OptimizerCallback): one record per (module_name, trial_num). CUDA peak VRAM is reset per module via ``torch.cuda.reset_peak_memory_stats``; MPS is sampled at ``end_module`` (no per-module peak API, so it's the moment-in-time allocation). + + Because per-module CUDA resets clobber the global ``max_memory_allocated`` + counter, the tracker also keeps ``self.peak_vram_gb_overall`` — the max + across every recorded module. The calibration script reads this instead of + the post-fit ``torch.cuda.max_memory_allocated()`` value, which by then + reflects only the last (usually CPU-only decision) module. """ name = "calibration_tracker" @@ -233,6 +471,14 @@ class _ModuleTracker(OptimizerCallback): def __init__(self) -> None: self.records: list[dict[str, Any]] = [] self._current: dict[str, Any] | None = None + self._current_step_buffer: list[float] | None = None + self.peak_vram_gb_overall: float = 0.0 + + def current_step_buffer(self) -> list[float] | None: + """Return the per-step wall-time list the ``_StepTimingCallback`` + should append to. ``None`` when no module is currently being tracked + (e.g. between modules) — the callback then skips.""" + return self._current_step_buffer def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: pass @@ -247,8 +493,10 @@ def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any] pass # Only capture JSON-safe scalars in the config snapshot. safe_config = {k: v for k, v in module_kwargs.items() if isinstance(v, (str, int, float, bool)) or v is None} + self._current_step_buffer = [] self._current = { "module": module_name, + "role": _classify_module_role(module_name), "num": num, "config": safe_config, "_start": time.perf_counter(), @@ -275,6 +523,18 @@ def end_module(self) -> None: rec["peak_vram_gb"] = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB except (ImportError, AttributeError): pass + peak = rec.get("peak_vram_gb") + if peak is not None and peak > self.peak_vram_gb_overall: + self.peak_vram_gb_overall = peak + # Fold per-step timings into the module record so a transformer's + # trial exposes ``mean_step_s`` / ``p95_step_s`` next to its total + # duration — the advisor's flat 1 s/step baseline can then be + # recalibrated per device_class directly from row data. + step_times = self._current_step_buffer or [] + step_summary = _summarize_step_times(step_times) + if step_summary: + rec["step_timings"] = step_summary + self._current_step_buffer = None self.records.append(rec) self._current = None @@ -296,6 +556,145 @@ def patched(context: Any) -> Any: # noqa: ANN401 pipeline._fit = patched # type: ignore[method-assign] # noqa: SLF001 +# === preset resolution & optional-extras skip ============================ + + +# module_name → the ``autointent[extra]`` that must be installed for the +# module's __init__ to succeed. Sourced from ``require(...)`` calls in +# ``src/autointent/modules/scoring/`` — keep in sync. +_MODULE_TO_EXTRA: dict[str, str] = { + "bert": "transformers", + "catboost": "catboost", + "lora": "peft", + "ptuning": "peft", + "description_llm": "openai", +} + + +def _missing_extras_for_config(cfg: dict[str, Any]) -> list[str]: + """Return every optional extra that ``cfg``'s search_space needs but isn't installed. + + Uses the same ``_deps.require`` validator the modules use at runtime, so + what the harness pre-checks matches what would fail inside ``fit``. + A missing extra returns an ``ImportError``; anything else (e.g. unknown + extra) propagates. + """ + from autointent._deps import require # type: ignore[import-not-found] + + needed: set[str] = set() + for node in cfg.get("search_space") or []: + for entry in node.get("search_space") or []: + name = entry.get("module_name") if isinstance(entry, dict) else None + extra = _MODULE_TO_EXTRA.get(name) if isinstance(name, str) else None + if extra: + needed.add(extra) + + missing: list[str] = [] + for extra in sorted(needed): + try: + require(extra) # type: ignore[arg-type] + except ImportError: + missing.append(extra) + return missing + + +def _load_pipeline_from_preset_ref(ref: str) -> tuple[Pipeline, str]: + """Resolve ``ref`` as either a bundled preset name or a YAML file path. + + Returns ``(pipeline, display_name)`` where ``display_name`` is the bundled + name for name refs or the file stem for path refs. Path refs let the + harness exercise modules not in any bundled preset (LoRA / ptuning / + dnnc / gcn / cross-encoder scorer) without polluting the shipped + ``SearchSpacePreset`` literal. + """ + if ref.endswith((".yaml", ".yml")): + path = Path(ref).expanduser() + if not path.exists(): + raise FileNotFoundError(f"Preset file not found: {path}") + pipeline = Pipeline.from_optimization_config(path) + return pipeline, path.stem + return Pipeline.from_preset(ref), ref # type: ignore[arg-type] + + +def _run_cli_smoke( + preset_ref: str, + stats: Any, # noqa: ANN401 + budget_vram_gb: float | None, +) -> dict[str, Any]: + """Invoke ``autointent-advisor inspect --json`` in-process. + + Fed the same stats (as placeholder args) and budget the direct-API path + saw, so a divergence in ``is_feasible`` / predicted numbers points at the + CLI wrapper or the JSON renderer, not at differing inputs. + + Returns a dict with: + * ``payload`` — the parsed JSON body from the CLI (or ``None`` on crash) + * ``rc`` — the CLI return code + * ``error`` — traceback string when the CLI or JSON parse failed + * ``divergence`` — dict of |cli - direct| deltas populated by the caller + + Runs in-process (no subprocess) so we don't pay the interpreter-startup + cost on every preset — the review only asked for a wrapper smoke, not a + full subprocess isolation test. + """ + import contextlib + import io as _io + import traceback + + from autointent._advisor._cli import main as cli_main + + argv = [ + "inspect", + preset_ref, + "--n-samples", + str(int(stats.n_samples)), + "--n-classes", + str(int(stats.n_classes)), + "--avg-tokens", + str(int(stats.avg_tokens)), + "--task", + "multilabel" if getattr(stats, "multilabel", False) else "multiclass", + "--json", + ] + if budget_vram_gb is not None: + argv += ["--budget-vram-gb", str(budget_vram_gb)] + + buf = _io.StringIO() + result: dict[str, Any] = {"payload": None, "rc": None, "error": None} + try: + with contextlib.redirect_stdout(buf): + result["rc"] = cli_main(argv) + except Exception: # noqa: BLE001 + result["error"] = traceback.format_exc(limit=3) + return result + + raw = buf.getvalue().strip() + if not raw: + result["error"] = "CLI produced empty stdout" + return result + try: + result["payload"] = json.loads(raw) + except json.JSONDecodeError as e: + result["error"] = f"CLI --json output not parseable: {e}" + return result + + +def _load_config_from_preset_ref(ref: str) -> dict[str, Any]: + """Load the raw preset config dict without instantiating a Pipeline. + + Used for pre-fit extras checks so a preset whose modules would need a + missing extra (e.g. ``lora`` → ``peft``) never reaches ``from_preset``. + """ + if ref.endswith((".yaml", ".yml")): + import yaml + + with Path(ref).expanduser().open(encoding="utf-8") as f: + return yaml.safe_load(f) + from autointent.utils import load_preset # local import to avoid top-level cost + + return load_preset(ref) # type: ignore[arg-type] + + # === per-preset run ====================================================== @@ -324,24 +723,61 @@ def _calibrate_one( poll_interval_ms: int, enable_wandb: bool, run_name: str | None, + budget_vram_gb: float | None, + clear_embedding_cache: bool, ) -> CalibrationRow: - row = CalibrationRow(preset=preset) + # ``preset`` may be a bundled name OR a path to a YAML file (the coverage + # preset). Resolve early so we can pre-check extras against the raw config + # before touching the Pipeline machinery. + try: + raw_cfg = _load_config_from_preset_ref(preset) + except Exception as e: # noqa: BLE001 + display_name = Path(preset).stem if preset.endswith((".yaml", ".yml")) else preset + row = CalibrationRow(preset=display_name) + row.error = f"load-preset failed: {e}" + return row + display_name = Path(preset).stem if preset.endswith((".yaml", ".yml")) else preset + + row = CalibrationRow(preset=display_name) + row.cache_policy = "cold" if clear_embedding_cache else "warm" + + # Detect missing optional extras BEFORE the fit — otherwise the trial + # would raise ImportError deep inside HPO, producing a fit-failed row + # indistinguishable from a real bug. + missing = _missing_extras_for_config(raw_cfg) + if missing: + row.skipped = True + row.error = f"skipped: missing extras {sorted(missing)}" + row.notes.append( + "install with: uv pip install " + " ".join(f"'autointent[{e}]'" for e in sorted(missing)) + ) + return row + + if clear_embedding_cache: + removed = _clear_embeddings_cache() + logger.info("Cleared %d embedding cache files for cold-cache measurement", removed) # === predicted ====================================================== try: - pipeline = Pipeline.from_preset(preset) + pipeline, _ = _load_pipeline_from_preset_ref(preset) except Exception as e: # noqa: BLE001 row.error = f"from_preset failed: {e}" return row - tagged_run_name = f"{preset}_{run_name}" if run_name else None + tagged_run_name = f"{display_name}_{run_name}" if run_name else None _override_trials(pipeline, max_trials, run_name=tagged_run_name) try: + # Optionally override hardware.vram_gb to exercise severity paths without a small GPU. + preflight_hw = hardware + if budget_vram_gb is not None: + from dataclasses import replace + + preflight_hw = replace(hardware, vram_gb=budget_vram_gb) report: PreflightReport = run_preflight( pipeline._build_advisor_config(), # noqa: SLF001 stats, - hardware, + preflight_hw, ) except Exception as e: # noqa: BLE001 row.error = f"preflight failed: {e}" @@ -357,15 +793,50 @@ def _calibrate_one( } row.findings = len(report.findings) row.findings_over = sum(1 for f in report.findings if f.severity.value == "over") + row.low_confidence = report.low_confidence if report.low_confidence: - row.notes.append("low-confidence (heuristic fallback in use)") + row.notes.append("low-confidence (heuristic HF metadata fallback in use)") + + # CLI wrapper smoke — same preset, same stats, same budget. Any divergence + # in ``is_feasible`` or the top-line predicted numbers means the CLI / + # JSON renderer drifted from the direct API. Runs unconditionally so a + # regression shows up on every calibration run. + smoke = _run_cli_smoke(preset, stats, budget_vram_gb) + if smoke["error"]: + row.notes.append(f"cli-smoke FAILED: {smoke['error'].splitlines()[-1] if smoke['error'] else '?'}") + elif smoke["payload"]: + cli_pred = smoke["payload"].get("resource") or {} + divergence: dict[str, float] = {} + for cli_key, direct_val in ( + ("time_hours", report.resource.time_hours), + ("ram_gb", report.resource.ram_gb), + ("vram_gb", report.resource.vram_gb), + ("disk_download_gb", report.resource.disk_download_gb), + ): + cli_val = cli_pred.get(cli_key) + if cli_val is None or direct_val is None: + continue + delta = abs(float(cli_val) - float(direct_val)) + if delta > 1e-6: + divergence[cli_key] = delta + smoke["divergence"] = divergence + cli_feasible = smoke["payload"].get("is_feasible") + if cli_feasible is not None and cli_feasible != report.is_feasible: + row.notes.append( + f"cli-smoke VERDICT MISMATCH: cli.is_feasible={cli_feasible} vs direct={report.is_feasible}" + ) + elif divergence: + row.notes.append(f"cli-smoke numeric drift on {sorted(divergence)} (see cli_smoke.divergence)") + row.cli_smoke = smoke if skip_fit: return row # === actual ========================================================= hf_cache = _hf_cache_dir() - cache_before = _dir_size_gb(hf_cache) + embed_cache = _embeddings_cache_dir() + hf_before = _dir_size_gb(hf_cache) + embed_before = _dir_size_gb(embed_cache) _reset_vram_peak() tracker = _ModuleTracker() @@ -380,6 +851,7 @@ def _calibrate_one( _attach_callbacks(pipeline, callbacks) is_mps = hardware.accelerator == "mps" + undo_step_patch = _patch_trainer_for_step_timing(tracker) start = time.perf_counter() try: with _PeakSampler(interval_s=poll_interval_ms / 1000.0, sample_mps=is_mps) as sampler: @@ -388,37 +860,35 @@ def _calibrate_one( row.error = f"fit failed: {e}" row.modules = tracker.records # keep whatever we collected return row + finally: + undo_step_patch() elapsed_s = time.perf_counter() - start - cache_after = _dir_size_gb(hf_cache) + hf_after = _dir_size_gb(hf_cache) + embed_after = _dir_size_gb(embed_cache) actual_time_h = elapsed_s / 3600.0 actual_ram_gb = sampler.peak_ram_gb - actual_vram_gb = _read_vram_peak_gb(hardware.accelerator) + # Prefer the tracker's per-module max: the fit-level torch.cuda peak is + # clobbered by the per-module reset_peak_memory_stats calls, so the final + # reading only reflects VRAM used since the last (usually CPU-only) module. + actual_vram_gb: float | None + if tracker.peak_vram_gb_overall > 0: + actual_vram_gb = tracker.peak_vram_gb_overall + else: + actual_vram_gb = _read_vram_peak_gb(hardware.accelerator) if actual_vram_gb is None and is_mps: actual_vram_gb = sampler.peak_mps_gb - actual_disk_download_gb = max(0.0, cache_after - cache_before) row.actual = { "time_h": actual_time_h, "ram_gb": actual_ram_gb, "vram_gb": actual_vram_gb, - "disk_download_gb": actual_disk_download_gb, + "disk_download_gb": max(0.0, hf_after - hf_before), + "disk_embedding_cache_gb": max(0.0, embed_after - embed_before), } row.modules = tracker.records if enable_wandb and not any("W&B requested but not available" in n for n in row.notes): row.notes.append("W&B reporter attached — inspect wandb.ai run group for per-step GPU/system metrics") - - def _ratio(actual: float | None, predicted: float) -> float | None: - if actual is None or predicted <= 0: - return None - return actual / predicted - - row.ratios = { - "time": _ratio(actual_time_h, row.predicted["time_h"]), - "ram": _ratio(actual_ram_gb, row.predicted["ram_gb"]), - "vram": _ratio(actual_vram_gb, row.predicted["vram_gb"]), - "disk_download": _ratio(actual_disk_download_gb, row.predicted["disk_download_gb"]), - } return row @@ -454,30 +924,108 @@ def _print_summary(rows: list[CalibrationRow]) -> None: header = " ".join(label.ljust(width) for _, label, width in _COLS) print(header) print("-" * len(header)) + _print_repeat_aggregates(rows) for row in rows: + # Ratios are always computed at read time (see CalibrationRow.to_dict). + ratios = row._ratios() # noqa: SLF001 cells = { "preset": row.preset, "pred_time": row.predicted.get("time_h"), "act_time": row.actual.get("time_h"), - "r_time": row.ratios.get("time"), + "r_time": ratios.get("time_h"), "pred_ram": row.predicted.get("ram_gb"), "act_ram": row.actual.get("ram_gb"), - "r_ram": row.ratios.get("ram"), + "r_ram": ratios.get("ram_gb"), "pred_vram": row.predicted.get("vram_gb"), "act_vram": row.actual.get("vram_gb"), - "r_vram": row.ratios.get("vram"), + "r_vram": ratios.get("vram_gb"), } print(" ".join(_fmt_cell(cells[key]).ljust(width) for key, _, width in _COLS)) if row.error: - print(f" ! {row.error}") + marker = "~" if row.skipped else "!" + print(f" {marker} {row.error}") + if row.low_confidence: + print(f" ! LOW-CONFIDENCE — advisor used heuristic HF metadata (exclude from prediction-accuracy stats)") + print(f" · cache-policy={row.cache_policy}") + role_totals = _sum_time_by_role(row.modules) + if role_totals: + breakdown = " ".join(f"{role}={total:.2f}s" for role, total in role_totals.items()) + print(f" · time-by-role: {breakdown}") for note in row.notes: print(f" * {note}") for mod in row.modules: duration = mod.get("duration_s") vram = mod.get("peak_vram_gb") + role = mod.get("role", "?") duration_s = f"{duration:.2f}s" if duration is not None else "-" vram_s = f"{vram:.2f} GB" if vram is not None else "-" - print(f" · {mod.get('module', '?')}#{mod.get('num', '?')} {duration_s} vram={vram_s}") + line = ( + f" · [{role}] {mod.get('module', '?')}#{mod.get('num', '?')} {duration_s} vram={vram_s}" + ) + step = mod.get("step_timings") + if step: + line += ( + f" n_steps={step['n_steps']} mean_step_s={step['mean_step_s']:.3f} " + f"p95_step_s={step['p95_step_s']:.3f}" + ) + print(line) + + +def _sum_time_by_role(modules: list[dict[str, Any]]) -> dict[str, float]: + """Fold per-module durations into ``{role: total_seconds}`` — used both for + the printed breakdown and for the top-level ``time_by_role`` row field.""" + totals: dict[str, float] = {} + for mod in modules: + role = mod.get("role", "?") + duration = mod.get("duration_s") + if duration is None: + continue + totals[role] = totals.get(role, 0.0) + float(duration) + return totals + + +def _print_repeat_aggregates(rows: list[CalibrationRow]) -> None: + """When any (preset, dataset) has more than one repeat, print a mean±stdev + block up-front so small ratio gaps are judgeable at a glance. + + Groups by ``(preset, first-note)`` — the dataset marker is inserted as the + first note in main() so this key is stable across repeats. + """ + import statistics + + groups: dict[tuple[str, str], list[CalibrationRow]] = {} + for row in rows: + dataset_note = row.notes[0] if row.notes else "dataset=?" + groups.setdefault((row.preset, dataset_note), []).append(row) + + multi_groups = [(k, v) for k, v in groups.items() if len(v) > 1] + if not multi_groups: + return + print(">>> repeats aggregation (mean ± stdev, successful runs only):") + for (preset, dataset_note), group in multi_groups: + # Skipped rows are expected — separate them from real failures so the + # aggregate isn't polluted by "all repeats failed" when the actual + # cause is a missing optional extra. + skipped = [r for r in group if r.skipped] + successful = [r for r in group if r.error is None] + real_failures = len(group) - len(successful) - len(skipped) + n = len(successful) + if n == 0: + reason = f"{real_failures} failed" + if skipped: + reason += f", {len(skipped)} skipped" + print(f" {preset} [{dataset_note}] {reason} (no successful repeats)") + continue + parts = [f" {preset} [{dataset_note}] n={n}"] + for metric in ("time_h", "ram_gb", "vram_gb"): + values = [r.actual.get(metric) for r in successful if r.actual.get(metric) is not None] + if not values: + continue + mean = statistics.fmean(values) + stdev = statistics.stdev(values) if len(values) > 1 else 0.0 + parts.append(f"{metric}={mean:.2f}±{stdev:.2f}") + print(" " + " ".join(parts)) + print() def _apply_thread_cap() -> None: @@ -499,6 +1047,79 @@ def _apply_thread_cap() -> None: pass +def _guard_cuda_init(*, required: bool) -> None: + """When ``required`` is True, fail fast if PyTorch can't initialize CUDA. + + Guards against the silent 'nvidia-smi shows 2 GPUs but torch runs on CPU' + trap that happens when the CUDA driver is older than what the installed + torch wheel was built against. + """ + if not required: + return + try: + import torch + except ImportError: + msg = "--require-cuda passed but torch isn't installed" + raise SystemExit(msg) from None + if not torch.cuda.is_available(): + # Get the underlying reason if we can — usually a warning on import time. + msg = ( + "--require-cuda passed but torch.cuda.is_available() is False. " + "Check `nvidia-smi` vs `python -c 'import torch; print(torch.version.cuda)'` — " + "you likely need a torch wheel built against a matching CUDA runtime." + ) + raise SystemExit(msg) + + +def _load_dataset(dataset_arg: str, parser: argparse.ArgumentParser) -> tuple[Dataset, str]: + """Load one dataset from a local JSON path or an HF Hub repo id, returning + ``(dataset, source_label)`` — the label mirrors what the calibrator writes + to the report so different sources are distinguishable in aggregate output. + """ + dataset_path = Path(dataset_arg) + if dataset_path.is_file(): + logger.info("Loading dataset from local file %s", dataset_path) + return Dataset.from_json(dataset_path), str(dataset_path) + logger.info("Loading dataset from HF Hub: %s", dataset_arg) + try: + dataset = Dataset.from_hub(dataset_arg) + except Exception as e: # noqa: BLE001 + parser.error(f"Could not load '{dataset_arg}' as a local JSON file or as a Hub repo id: {e}") + return dataset, f"hub:{dataset_arg}" + + +def _subsample_per_class(dataset: Dataset, cap: int) -> Dataset: + """Cap each class in the train split to at most ``cap`` samples (first-N slice). + + Uses a deterministic first-N slice per class — reproducible across runs + without seeding, and keeps class-ordering intuitive when inspecting the + subset. Only rewrites the train split; validation/test are left as-is so + the metric baselines remain comparable. + """ + from autointent.custom_types import Split + + train_key = Split.TRAIN if Split.TRAIN in dataset else next( + (k for k in dataset if str(k).startswith(str(Split.TRAIN))), None, + ) + if train_key is None: + return dataset + train = dataset[train_key] + label_feature = dataset.label_feature + seen: dict[Any, int] = {} + keep: list[int] = [] + for idx, row in enumerate(train): + label = row[label_feature] + # For multilabel, key on the tuple so a sample with a rare-class tag + # still contributes toward that class's cap. + key = tuple(label) if isinstance(label, list) else label + count = seen.get(key, 0) + if count < cap: + keep.append(idx) + seen[key] = count + 1 + dataset[train_key] = train.select(keep) + return dataset + + def main(argv: list[str] | None = None) -> int: parser = _build_parser() args = parser.parse_args(argv) @@ -507,25 +1128,23 @@ def main(argv: list[str] | None = None) -> int: format="%(levelname)s %(name)s: %(message)s", ) _apply_thread_cap() + _guard_cuda_init(required=args.require_cuda) presets = args.presets or list(BUNDLED_PRESETS) - unknown = [p for p in presets if p not in BUNDLED_PRESETS] + unknown = [ + p + for p in presets + if not p.endswith((".yaml", ".yml")) and p not in BUNDLED_PRESETS + ] if unknown: - parser.error(f"Unknown preset(s): {', '.join(unknown)}. Known: {', '.join(BUNDLED_PRESETS)}") + parser.error( + f"Unknown preset(s): {', '.join(unknown)}. Known: {', '.join(BUNDLED_PRESETS)}, " + "or pass a path to a .yaml file (e.g. scripts/coverage_preset.yaml)." + ) + for p in presets: + if p.endswith((".yaml", ".yml")) and not Path(p).expanduser().exists(): + parser.error(f"Preset file not found: {p}") - dataset_path = Path(args.dataset) - if dataset_path.is_file(): - logger.info("Loading dataset from local file %s", dataset_path) - dataset = Dataset.from_json(dataset_path) - dataset_source = str(dataset_path) - else: - logger.info("Loading dataset from HF Hub: %s", args.dataset) - try: - dataset = Dataset.from_hub(args.dataset) - except Exception as e: # noqa: BLE001 - parser.error(f"Could not load '{args.dataset}' as a local JSON file or as a Hub repo id: {e}") - dataset_source = f"hub:{args.dataset}" - stats = stats_from_dataset_obj(dataset) hardware = detect_hardware() logger.info( "Hardware: %s (%s) — %.1f GB VRAM, %.0f GB RAM, %.0f GB free disk", @@ -544,20 +1163,91 @@ def main(argv: list[str] | None = None) -> int: ) rows: list[CalibrationRow] = [] - for preset in presets: - logger.info("=== %s ===", preset) - row = _calibrate_one( - preset=preset, - dataset=dataset, - stats=stats, - hardware=hardware, - max_trials=args.max_trials, - skip_fit=args.skip_fit, - poll_interval_ms=args.poll_interval_ms, - enable_wandb=args.wandb, - run_name=args.run_name, + datasets_meta: list[dict[str, Any]] = [] + + def _write_payload() -> None: + """Serialize the current in-memory rows to ``args.output``. Called + after each preset finishes so a mid-sweep crash / Broken pipe leaves + a valid partial report behind rather than losing everything. + """ + payload_now = { + "hardware": { + "accelerator": hardware.accelerator, + "device_name": hardware.device_name, + "vram_gb": hardware.vram_gb, + "ram_gb": hardware.ram_gb, + "free_disk_gb": hardware.free_disk_gb, + }, + "datasets": datasets_meta, + "max_trials_override": args.max_trials, + "skip_fit": args.skip_fit, + "cache_policy": "cold" if args.clear_embedding_cache else "warm", + "budget_vram_gb_override": args.budget_vram_gb, + "subsample_per_class": args.subsample_per_class, + "thread_caps": { + "OMP_NUM_THREADS": os.environ.get("OMP_NUM_THREADS"), + "MKL_NUM_THREADS": os.environ.get("MKL_NUM_THREADS"), + "OPENBLAS_NUM_THREADS": os.environ.get("OPENBLAS_NUM_THREADS"), + "TOKENIZERS_PARALLELISM": os.environ.get("TOKENIZERS_PARALLELISM"), + }, + "in_progress": True, + "rows": [r.to_dict() for r in rows], + } + # Atomic write: dump to a sibling file, then rename. Prevents readers + # from seeing a half-written JSON if the run is killed mid-serialize. + tmp = args.output.with_suffix(args.output.suffix + ".partial") + tmp.write_text(json.dumps(payload_now, indent=2, default=str)) + tmp.replace(args.output) + + for dataset_arg in args.dataset: + dataset, dataset_source = _load_dataset(dataset_arg, parser) + if args.subsample_per_class is not None: + dataset = _subsample_per_class(dataset, args.subsample_per_class) + dataset_source += f"|subsample-per-class={args.subsample_per_class}" + stats = stats_from_dataset_obj(dataset) + datasets_meta.append( + { + "path": dataset_source, + "n_samples": stats.n_samples, + "n_classes": stats.n_classes, + "avg_tokens": stats.avg_tokens, + "multilabel": stats.multilabel, + }, + ) + logger.info( + "Dataset %s: n_samples=%d n_classes=%d avg_tokens=%.1f multilabel=%s", + dataset_source, + stats.n_samples, + stats.n_classes, + stats.avg_tokens, + stats.multilabel, ) - rows.append(row) + for preset in presets: + for repeat_idx in range(max(1, args.repeats)): + header = f"=== {preset} @ {dataset_source}" + if args.repeats > 1: + header += f" (repeat {repeat_idx + 1}/{args.repeats})" + header += " ===" + logger.info(header) + row = _calibrate_one( + preset=preset, + dataset=dataset, + stats=stats, + hardware=hardware, + max_trials=args.max_trials, + skip_fit=args.skip_fit, + poll_interval_ms=args.poll_interval_ms, + enable_wandb=args.wandb, + run_name=( + f"{args.run_name}_r{repeat_idx}" if args.run_name and args.repeats > 1 else args.run_name + ), + budget_vram_gb=args.budget_vram_gb, + clear_embedding_cache=args.clear_embedding_cache, + ) + row.repeat_idx = repeat_idx + row.notes.insert(0, f"dataset={dataset_source}") + rows.append(row) + _write_payload() payload = { "hardware": { @@ -567,15 +1257,19 @@ def main(argv: list[str] | None = None) -> int: "ram_gb": hardware.ram_gb, "free_disk_gb": hardware.free_disk_gb, }, - "dataset": { - "path": dataset_source, - "n_samples": stats.n_samples, - "n_classes": stats.n_classes, - "avg_tokens": stats.avg_tokens, - "multilabel": stats.multilabel, - }, + "datasets": datasets_meta, "max_trials_override": args.max_trials, "skip_fit": args.skip_fit, + "cache_policy": "cold" if args.clear_embedding_cache else "warm", + "budget_vram_gb_override": args.budget_vram_gb, + "subsample_per_class": args.subsample_per_class, + "thread_caps": { + "OMP_NUM_THREADS": os.environ.get("OMP_NUM_THREADS"), + "MKL_NUM_THREADS": os.environ.get("MKL_NUM_THREADS"), + "OPENBLAS_NUM_THREADS": os.environ.get("OPENBLAS_NUM_THREADS"), + "TOKENIZERS_PARALLELISM": os.environ.get("TOKENIZERS_PARALLELISM"), + }, + "in_progress": False, "rows": [r.to_dict() for r in rows], } args.output.write_text(json.dumps(payload, indent=2, default=str)) diff --git a/scripts/coverage_preset.yaml b/scripts/coverage_preset.yaml new file mode 100644 index 000000000..02844aa50 --- /dev/null +++ b/scripts/coverage_preset.yaml @@ -0,0 +1,42 @@ +## Coverage-only preset for the calibration harness. +# +# Bundled presets don't touch lora / ptuning / dnnc / gcn / a plain +# cross-encoder scorer, so the advisor's estimates for those modules go +# unvalidated. This preset packs one of each into a single small run so +# `scripts/calibrate_advisor.py` can exercise them against a real fit. +# +# Intentionally cheap (n_trials: 1 per scorer, single decision module) — +# the point is coverage, not tuning quality. Skip cleanly when the peft +# extra is missing (lora/ptuning will be marked skipped by the harness). +search_space: + - node_type: scoring + target_metric: scoring_f1 + search_space: + - module_name: lora + classification_model_config: + - model_name: microsoft/deberta-v3-small + num_train_epochs: [1] + batch_size: [16] + learning_rate: [5.0e-5] + - module_name: ptuning + classification_model_config: + - model_name: microsoft/deberta-v3-small + num_train_epochs: [1] + batch_size: [16] + learning_rate: [5.0e-5] + num_virtual_tokens: [8] + - module_name: dnnc + k: [3] + - module_name: gcn + num_train_epochs: [1] + batch_size: [16] + learning_rate: [1.0e-3] + - module_name: description_cross + - node_type: decision + target_metric: decision_accuracy + search_space: + - module_name: argmax +hpo_config: + sampler: tpe + n_trials: 5 + n_startup_trials: 2 diff --git a/scripts/run_calibration_banking77.sh b/scripts/run_calibration_banking77.sh index 7db4d9120..1c6f0d869 100755 --- a/scripts/run_calibration_banking77.sh +++ b/scripts/run_calibration_banking77.sh @@ -8,7 +8,20 @@ # # Environment overrides: # DATASET HF Hub repo id (default: DeepPavlov/banking77) -# PRESETS Space-separated preset names (default: every bundled preset) +# DATASETS Space-separated list of dataset ids (default: unset -> use +# DATASET). Every preset runs against every dataset — useful +# to sweep a small + large + multilabel + long-token shape in +# a single invocation. +# SUBSAMPLE_PER_CLASS Cap each class to N training samples (deterministic +# first-N slice) before running — turns banking77 into a +# small dataset without needing a separate corpus. +# REPEATS Run each (preset, dataset) N times. The summary prints +# mean ± stdev of actual measurements across repeats so +# small ratio gaps become judgeable. Default: 1. +# PRESETS Space-separated preset names OR paths (default: every bundled preset). +# Items ending in .yaml/.yml are treated as file paths — use +# scripts/coverage_preset.yaml to exercise lora/ptuning/dnnc/gcn/ +# cross-encoder in one small run without touching bundled presets. # MAX_TRIALS Cap for hpo_config.n_trials (default: unset -> preset default) # WANDB If non-empty, pass --wandb so system metrics land in wandb.ai # RUN_NAME Suffix appended to each preset's LoggingConfig.run_name — the @@ -20,6 +33,14 @@ # (default: 1). Increase carefully — sklearn's own ``n_jobs`` and # HPO parallelism multiply on top, so oversubscription is easy # on many-core boxes. +# COLD If non-empty, pass --clear-embedding-cache so each preset +# starts with an empty embeddings cache (measure COLD cost). +# BUDGET_VRAM_GB Force run_preflight to see a specific VRAM budget instead +# of what the box exposes — lets you exercise severity paths +# (red/yellow/green + findings_over) on a big box. +# REQUIRE_CUDA If non-empty, pass --require-cuda so the run fails fast when +# torch.cuda.is_available() is False (guards against silent +# CPU-fallback caused by CUDA-driver / torch-wheel mismatch). # # Examples: # scripts/run_calibration_banking77.sh # full sweep, serial @@ -28,6 +49,11 @@ # WANDB=1 MAX_TRIALS=5 scripts/run_calibration_banking77.sh # RUN_NAME=calib_2026_07 WANDB=1 scripts/run_calibration_banking77.sh # THREADS_PER_JOB=4 scripts/run_calibration_banking77.sh # 4-thread BLAS +# SUBSAMPLE_PER_CLASS=5 scripts/run_calibration_banking77.sh # small-dataset shape +# DATASETS="DeepPavlov/banking77 DeepPavlov/clinc150" scripts/run_calibration_banking77.sh +# REPEATS=3 MAX_TRIALS=3 scripts/run_calibration_banking77.sh # variance bars +# PRESETS="scripts/coverage_preset.yaml" MAX_TRIALS=1 scripts/run_calibration_banking77.sh +# # exercise lora/ptuning/dnnc/gcn/description_cross — modules no bundled preset touches set -euo pipefail @@ -36,6 +62,13 @@ REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$REPO_ROOT" DATASET="${DATASET:-DeepPavlov/banking77}" +# DATASETS overrides DATASET when set; enables multi-dataset sweeps. +if [[ -n "${DATASETS:-}" ]]; then + # shellcheck disable=SC2206 # intentional word-split from env + DATASET_ARR=($DATASETS) +else + DATASET_ARR=("$DATASET") +fi OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/calibration_runs}" TIMESTAMP="$(date +%Y%m%d_%H%M%S)" OUTPUT_JSON="$OUTPUT_DIR/banking77_$TIMESTAMP.json" @@ -75,6 +108,21 @@ fi if [[ -n "${RUN_NAME:-}" ]]; then EXTRA_FLAGS+=("--run-name" "$RUN_NAME") fi +if [[ -n "${COLD:-}" ]]; then + EXTRA_FLAGS+=("--clear-embedding-cache") +fi +if [[ -n "${BUDGET_VRAM_GB:-}" ]]; then + EXTRA_FLAGS+=("--budget-vram-gb" "$BUDGET_VRAM_GB") +fi +if [[ -n "${REQUIRE_CUDA:-}" ]]; then + EXTRA_FLAGS+=("--require-cuda") +fi +if [[ -n "${SUBSAMPLE_PER_CLASS:-}" ]]; then + EXTRA_FLAGS+=("--subsample-per-class" "$SUBSAMPLE_PER_CLASS") +fi +if [[ -n "${REPEATS:-}" ]]; then + EXTRA_FLAGS+=("--repeats" "$REPEATS") +fi # Preset list: pull it from the advisor package at runtime unless overridden, # so the script auto-discovers presets that are added later. @@ -95,7 +143,7 @@ PY fi echo "Repo: $REPO_ROOT" -echo "Dataset: $DATASET" +echo "Datasets: ${DATASET_ARR[*]}" echo "Presets: ${PRESET_ARR[*]}" echo "Output: $OUTPUT_JSON" echo "Log: $LOG_FILE" @@ -104,7 +152,7 @@ echo "Threads per job: $THREADS_PER_JOB (OMP/MKL/OpenBLAS/torch)" echo uv run --no-sync python scripts/calibrate_advisor.py \ - --dataset "$DATASET" \ + --dataset "${DATASET_ARR[@]}" \ --presets "${PRESET_ARR[@]}" \ --output "$OUTPUT_JSON" \ "${EXTRA_FLAGS[@]}" \ diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index 681363afd..4df486f2b 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -12,9 +12,11 @@ from .runner import run_preflight from .workflows import ( BUNDLED_PRESETS, + ReduceToFitError, inspect, load_config, recommend, + reduce_to_fit, stats_from_dataset, stats_from_dataset_obj, ) @@ -26,12 +28,14 @@ "HardwareProfile", "PreflightReport", "RecommendationResult", + "ReduceToFitError", "ResourceEstimate", "Severity", "detect_hardware", "inspect", "load_config", "recommend", + "reduce_to_fit", "run_preflight", "stats_from_dataset", "stats_from_dataset_obj", diff --git a/src/autointent/_advisor/_estimates/_formulas.py b/src/autointent/_advisor/_estimates/_formulas.py index 2b99b7c57..6fe3985b1 100644 --- a/src/autointent/_advisor/_estimates/_formulas.py +++ b/src/autointent/_advisor/_estimates/_formulas.py @@ -83,18 +83,31 @@ def _activations_gb_per_sample( *, is_training: bool, ) -> float: - """Heuristic activation memory per sample, assuming a fp32 worst case. - - Training: ``seq_len x hidden x layers x const`` — per-layer outputs are kept - for backward. - Inference: ``seq_len x hidden x const`` — only one or two layers' outputs in - flight at once. + """Heuristic activation memory per sample. + + Training uses **34 bytes/token/layer** as a pessimistic upper bound — + Korthikanti et al. (2022, "Reducing Activation Recomputation ...") derive + this for standard attention: the linear-layer activations account for ~11B + and the attention matrix + intermediate tensors add ~23B. FlashAttention + kernels drop the attention-matrix term (~12 B/token/layer total), but we + can't detect at preflight time whether the user's stack will use them, so + the upper bound is the safe choice. + + Inference: only 1-2 layers' outputs are kept in flight at once. 8 B/token + covers fp32 hidden (4B) plus a bit of intermediate slack. + + An earlier revision used 16 B/token/layer for training; that under-predicted + real deberta-v3-large VRAM by ~2x at bs=128 (measured 13.1 GB, predicted + ~11.5 GB), which is unsafe for an OOM-avoidance tool. See ``interpretation.md`` + (2026-07-19) for the calibration data. """ hidden = _embedder_dim(meta) - # Training keeps every layer's outputs for backward -> scales x n_layers. - # 16 bytes/token/layer ~ fp32 activation (4B) x ~4x backward overhead (Korthikanti et al.). - # Inference only holds ~1-2 layers' outputs in flight at once. - bytes_per_sample = seq_len * hidden * _n_layers(meta) * 16 if is_training else seq_len * hidden * 8 + training_bytes_per_token_per_layer = 34 + inference_bytes_per_token = 8 + if is_training: + bytes_per_sample = seq_len * hidden * _n_layers(meta) * training_bytes_per_token_per_layer + else: + bytes_per_sample = seq_len * hidden * inference_bytes_per_token return bytes_per_sample / _BYTES_PER_GB @@ -137,13 +150,18 @@ def _max_fitting_batch_size( return _floor_to_power_of_two(int(available_for_activations / per_sample_gb)) -_CPU_SLOWDOWN_FACTOR = 50.0 -"""Rough multiplier for transformer training on CPU vs. a modern GPU. - -Real benchmarks vary widely (30x for small BERTs on AVX-512 boxes, 100x+ for -billion-scale models on a stock laptop). A single 50x constant is a pessimistic -upper bound that's good enough to make the CPU/GPU distinction visible without -re-introducing the per-device tier table.""" +# Sustained TFLOPS per device class — real MFU (model-FLOPs utilization) at +# training batch sizes, NOT peak spec sheet numbers. Numbers reflect ~30-50% +# MFU which is typical for BERT-scale training with FA2 / cuDNN kernels. +# Source: MLPerf training results + community benchmarks (2024-2025). +_DEVICE_TFLOPS = { + "high-gpu": 150.0, # A100 / H100 fp16 + "mid-gpu": 45.0, # V100 / RTX 3090 / A6000 + "low-gpu": 15.0, # T4 / RTX 3060 / 8 GB consumer card + "apple-silicon": 8.0, # M1/M2/M3 GPU cores + "cpu": 0.1, # single-thread modern x86 with MKL +} +_DEFAULT_TFLOPS = 15.0 # unknown device → treat as low-GPU def _time_for_transformer( @@ -151,22 +169,28 @@ def _time_for_transformer( n_trials: int, epochs: int, batch_size: int, + seq_len: int, n_samples: int, - accelerator: str, + params_millions: float, + device_class: str, ) -> float: - """Transformer training time in hours. + """Transformer training time in hours, from per-step FLOPs / device TFLOPS. - Baseline is "1 second per step" on a GPU (CUDA / MPS) — a step-count proxy, - not a real wall-time calibration. CPU training pays a flat ``_CPU_SLOWDOWN_FACTOR`` - so the report doesn't hide the fact that the same workload is dramatically - slower without a GPU. Users should treat absolute numbers as ordering / - ballpark information, not a budget. + Per-step FLOPs ≈ 6 x params x batch_size x seq_len (2 for forward mul-add, + 3-4x for backward). Divided by sustained device TFLOPS to get wall-time per + step, then multiplied by (steps x epochs x n_trials). + + Replaces an earlier "1 second per step" heuristic, which was ~10x too high + on A100 and identical for MPS vs CUDA (predicted times were the same on + both while real times differed ~7x — see interpretation.md 2026-07-19). """ - steps = max(1, (n_samples // max(1, batch_size))) * epochs - h = (n_trials * steps) / 3600.0 - if accelerator == "cpu": - h *= _CPU_SLOWDOWN_FACTOR - return h + steps_per_epoch = max(1, n_samples // max(1, batch_size)) + total_steps = n_trials * epochs * steps_per_epoch + # 6x factor: ~2x for fwd matmul + ~4x for bwd (grad wrt input + grad wrt weight). + step_flops = 6.0 * params_millions * 1e6 * batch_size * seq_len + tflops = _DEVICE_TFLOPS.get(device_class, _DEFAULT_TFLOPS) + step_seconds = step_flops / (tflops * 1e12) + return (total_steps * step_seconds) / 3600.0 def _n_layers(meta: ModelMeta | None) -> int: @@ -207,11 +231,20 @@ def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: return (n_samples * hidden_size * 4) / _BYTES_PER_GB -# Coefficients are dimensional (per-sample-per-feature-per-iteration seconds) -# rather than empirically tuned constants — they give relative-cost ordering -# across configurations and absolute ballpark wall-times. -_LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-8 -_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 +# Wall-time coefficients calibrated against measured fits on 1-thread CPU +# (OMP_NUM_THREADS=1). Values represent seconds per per-fit-work-unit and +# already absorb the number of L-BFGS iterations the optimizer typically +# takes to converge (~50) — so the ``max_iter`` upper bound does NOT enter +# the formula directly. Historical formula also scaled by ``max_iter`` which +# multi-cent-ordered-over-predicted (137 h vs measured ~30 s = ~15000x on +# banking77 × 1024-dim e5-large × 77 classes × cv=3). +# +# Calibration point (reviewer's res-adapt-ckeck/a100 run, warm cache): +# classic-light linear on banking77 (n=10003, dim=1024, cls=77, cv_mult=31) +# measured ~30 s per fit x n_trials=20 = ~10 min total = ~0.17 h. +# Formula: 20 x 1.2e-9 x 10003 x 1024 x 31 x 77 = ~588 s = ~0.16 h ✓ +_LINEAR_CPU_S_PER_SAMPLE_FEATURE = 1.2e-9 +_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 # catboost is measured per iteration _CATBOOST_GPU_SPEEDUP = 10.0 # LogisticRegressionCV defaults: Cs=10, cv=3 -> 10x3 inner fits + 1 final refit = 31. _LOGREG_CV_MULTIPLIER = 31 @@ -234,22 +267,26 @@ def _time_for_linear( n_trials: int, n_samples: int, embedder_dim: int, - max_iter: int, + max_iter: int, # noqa: ARG001 — kept in signature for API stability; typical L-BFGS convergence is absorbed into the coefficient cv_multiplier: int, class_multiplier: int, ) -> float: """LogisticRegression wall time, in hours. - Cost is ``O(n_samples x n_features x max_iter x n_classes)`` per fit - (sklearn's L-BFGS solver), multiplied by the CV inner-fit count (31 for the - default LogisticRegressionCV). + Cost is ``O(n_samples x n_features x n_classes)`` per fit (sklearn's L-BFGS + solver, iterations absorbed into the calibration constant), multiplied by the + CV inner-fit count (31 for the default LogisticRegressionCV). + + ``max_iter`` is a per-fit upper bound, not the typical work — L-BFGS on a + well-conditioned classifier converges long before it. Older versions of + this formula scaled by ``max_iter`` and predicted ~1000x higher than + reality; the constant now bakes in a typical convergence-iteration count. """ seconds = ( n_trials - * _LINEAR_CPU_S_PER_SAMPLE_FEATURE_ITER + * _LINEAR_CPU_S_PER_SAMPLE_FEATURE * n_samples * embedder_dim - * max_iter * cv_multiplier * class_multiplier ) diff --git a/src/autointent/_advisor/_estimates/_resource.py b/src/autointent/_advisor/_estimates/_resource.py index d14dbd186..884254faf 100644 --- a/src/autointent/_advisor/_estimates/_resource.py +++ b/src/autointent/_advisor/_estimates/_resource.py @@ -11,7 +11,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable from autointent._advisor import _hub from autointent._advisor._report import ResourceEstimate, Severity @@ -160,8 +160,10 @@ def _estimate_transformer_model( n_trials=n_trials, epochs=epochs, batch_size=batch_size, + seq_len=seq_len, n_samples=stats.n_samples, - accelerator=hardware.accelerator, + params_millions=meta.total_params / 1_000_000, + device_class=hardware.device_class, ) if mode != "inference": time_h *= _refit_factor(refit_after=refit_after, n_trials=n_trials) @@ -269,6 +271,7 @@ def _apply_embedding_cache( *, stats: DatasetStats, hardware: HardwareProfile, + cache_probe: Callable[[str], bool] | None = None, ) -> set[str]: """Adjust ``module_estimates`` in-place for autointent's persistent embedding cache. @@ -277,9 +280,24 @@ def _apply_embedding_cache( entries (linear/catboost) get a synthetic forward added since their per-entry estimate doesn't include one. - Returns the set of unique embedder model names whose forward was charged. + ``cache_probe`` (optional): callable that takes an embedder model_name and + returns True if the embedding is already cached on disk. When it returns + True, the model is treated as pre-paid — forward is zero and disk cache + delta is zero. Default (None) preserves the pessimistic cold assumption + the advisor shipped with — every embedder pays once. + + Returns the set of unique embedder model names whose forward was charged + (i.e. contributed to ``disk_embedding_cache_gb`` in the disk aggregation). """ paid: set[str] = set() + # Models the probe reports as already-warm — pre-populate ``paid`` so the + # first-seen module also hits the cache-hit branch instead of paying the + # forward, and skip them in the disk-cache aggregation (already on disk). + warm_models: set[str] = set() + if cache_probe is not None: + for name in seen_models: + if cache_probe(name): + warm_models.add(name) for me in module_estimates: module = me.driver["module"] if module not in _CACHE_HONORING_MODULES: @@ -287,6 +305,12 @@ def _apply_embedding_cache( model = me.driver["model"] if model not in seen_models: # synthetic / "(no embedder)" rows continue + if model in warm_models: + if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: + me.time_hours = 0.0 + me.driver["time_hours"] = 0.0 + me.driver["mode"] = f"{me.driver['mode']}+warm" + continue if model in paid: if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: me.time_hours = 0.0 @@ -295,12 +319,15 @@ def _apply_embedding_cache( else: paid.add(model) if module in {"linear", "catboost"}: + embedder_meta = seen_models.get(model) forward_h = _time_for_transformer( n_trials=1, epochs=1, batch_size=32, + seq_len=128, n_samples=stats.n_samples, - accelerator=hardware.accelerator, + params_millions=(embedder_meta.total_params / 1_000_000) if embedder_meta else 100.0, + device_class=hardware.device_class, ) me.time_hours += forward_h me.driver["time_hours"] = round(me.time_hours, 2) @@ -395,6 +422,51 @@ def _emit_resource_findings( ) +_UNKNOWN_SCORER_MODULES = frozenset({"cnn", "rnn", "sklearn"}) +"""Scorer modules the advisor has no cost estimator for. + +These get a placeholder ``not-estimated`` driver row so they never appear as +"free/safe" in the report — silent-zero was the ``nn-heavy`` predicted 0h/0GB +bug that hid a real 0.52 h + 1.45 GB RAM cost. +""" + +# Modules that consume the top-level ``cross_encoder_config.model_name`` as +# their scoring model (see zero-shot-encoders preset: description_cross pulls +# BAAI/bge-reranker-v2-m3 from that config, not from its per-entry dict). +_CROSS_ENCODER_CONSUMERS = frozenset({"description_cross", "dnnc", "retrieval"}) + +# Modules that fall back to the top-level ``transformer_config.model_name`` +# when no per-entry ``classification_model_config`` is given. +_TRANSFORMER_CONFIG_CONSUMERS = frozenset({"bert"}) + + +def _not_estimated_row(*, node_type: str, module: str) -> _ModuleEstimate: + """Placeholder row for a module the advisor has no cost formula for. + + Renders as ``mode="not-estimated"`` in the report so the module isn't + silently absent (would read as "free/safe") — a call to action for whoever + reads the JSON that the actual cost is unknown, not zero. + """ + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": "(not estimated)", + "mode": "not-estimated", + "vram_gb": 0.0, + "ram_gb": 0.0, + "time_hours": 0.0, + "batch_size": None, + "max_batch_size": None, + "confidence": "unknown", + "note": "advisor has no cost estimator for this module; treat as unknown, not zero", + }, + vram_gb=0.0, + ram_gb=0.0, + time_hours=0.0, + ) + + def _resource_phase( *, embedder_config: EmbedderConfig, @@ -406,6 +478,9 @@ def _resource_phase( hardware: HardwareProfile, report: PreflightReport, refit_after: bool = False, + cross_encoder_model_name: str | None = None, + transformer_model_name: str | None = None, + cache_probe: Callable[[str], bool] | None = None, ) -> None: """Walk the validated search space, fold per-module costs into the report. @@ -413,6 +488,15 @@ def _resource_phase( the largest model can drive ``embedder_dim`` for the classic pass), then linear / catboost. Disk, VRAM/RAM peak, time sum, and final findings are folded onto the report. + + ``cross_encoder_model_name`` and ``transformer_model_name`` come from the + pipeline's top-level configs. They're used as the fallback model for + modules that don't declare a per-entry ``classification_model_config`` but + still consume one at runtime (``description_cross`` / ``dnnc`` / + ``retrieval`` pull from ``cross_encoder_config``; ``bert`` falls back to + ``transformer_config``). Seeding them here fixes the disk-download + under-count called out in the follow-up review (missing 6.4 GB reranker in + ``zero-shot-encoders``). """ seen_models: dict[str, ModelMeta] = {} global_embedder = _embedder_model_name(embedder_config) @@ -427,8 +511,17 @@ def _resource_phase( for node_idx, node_type, entry in transformer_entries: module = entry.get("module_name", "?") model_names = _extract_model_names(entry) - if not model_names and global_embedder and module in {"knn", "mlknn"}: - model_names = [global_embedder] + if not model_names: + if module in {"knn", "mlknn"} and global_embedder: + model_names = [global_embedder] + elif module in _CROSS_ENCODER_CONSUMERS and cross_encoder_model_name: + model_names = [cross_encoder_model_name] + elif module in _TRANSFORMER_CONFIG_CONSUMERS and transformer_model_name: + model_names = [transformer_model_name] + elif module in _UNKNOWN_SCORER_MODULES: + # Placeholder so the row is visible instead of silently zeroed. + module_estimates.append(_not_estimated_row(node_type=node_type, module=module)) + continue for name in model_names: meta = seen_models.setdefault(name, _hub.resolve_model(name)) me = _estimate_transformer_model( @@ -465,7 +558,9 @@ def _resource_phase( module_estimates.append(classic_estimate) # Cache-aware time/disk: must run before the fold below. - cached_embedders = _apply_embedding_cache(module_estimates, seen_models, stats=stats, hardware=hardware) + cached_embedders = _apply_embedding_cache( + module_estimates, seen_models, stats=stats, hardware=hardware, cache_probe=cache_probe, + ) estimate = ResourceEstimate(parallel_factor=n_jobs) for me in module_estimates: @@ -486,12 +581,18 @@ def _resource_phase( # Flip low_confidence if any model fell back to the heuristic path (Hub # unreachable, repo missing safetensors metadata, local-path checkpoint). + # Emit as a TIGHT finding (not just a note) so it shows up in the main + # rendered findings block — buried notes previously let ~2× under-prediction + # of large-model shapes slip past the reviewer. heuristic_models = [m.name for m in seen_models.values() if m.confidence == "heuristic"] if heuristic_models: report.low_confidence = True - report.notes.append( - f"Heuristic fallback used for {len(heuristic_models)} model(s) - sizes are BERT-base " - f"defaults: {', '.join(heuristic_models[:3])}{'...' if len(heuristic_models) > 3 else ''}", # noqa: PLR2004 + sample = ", ".join(heuristic_models[:3]) + ("..." if len(heuristic_models) > 3 else "") # noqa: PLR2004 + report.add( + "resource", + Severity.TIGHT, + f"LOW CONFIDENCE - Hub metadata unavailable for {len(heuristic_models)} model(s); " + f"cost estimates use conservative large-model defaults (may over-predict small models): {sample}", ) report.resource = estimate diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/_advisor/_hub.py index 677e6045c..3dc48aaab 100644 --- a/src/autointent/_advisor/_hub.py +++ b/src/autointent/_advisor/_hub.py @@ -20,8 +20,17 @@ logger = logging.getLogger(__name__) -_DEFAULT_HEURISTIC_PARAMS = 110_000_000 +# Conservative "large-model" shape used when Hub metadata is unavailable — +# roughly deberta-v3-large / bert-large sized. Previously we defaulted to a +# BERT-base shape (110M / 768 / 12), which *under*-predicted a real deberta-large +# fit by ~2×. Because the advisor's contract is a pessimistic upper bound, the +# offline fallback needs to over-estimate small models rather than under-estimate +# large ones. Callers can still see the fallback happened via ``confidence == +# "heuristic"`` and ``PreflightReport.low_confidence``. +_DEFAULT_HEURISTIC_PARAMS = 350_000_000 _DEFAULT_BYTES_PER_PARAM = 4 +_DEFAULT_HEURISTIC_HIDDEN = 1024 +_DEFAULT_HEURISTIC_LAYERS = 24 _BYTES_PER_GB = 1024**3 # using the binary GiB convention everywhere in the advisor @@ -138,9 +147,15 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: if hidden_size is None or n_layers is None: logger.warning( "Could not read hidden_size / num_hidden_layers from config.json for %s; " - "activation-memory estimates will fall back to BERT-base defaults (768 / 12).", + "activation-memory estimates will fall back to CONSERVATIVE large-model " + "defaults (hidden=%d, layers=%d) to avoid under-predicting.", model_name, + _DEFAULT_HEURISTIC_HIDDEN, + _DEFAULT_HEURISTIC_LAYERS, ) + hidden_size = hidden_size or _DEFAULT_HEURISTIC_HIDDEN + n_layers = n_layers or _DEFAULT_HEURISTIC_LAYERS + confidence = "heuristic" return ModelMeta( name=model_name, @@ -157,8 +172,12 @@ def _hub_metadata(model_name: str) -> ModelMeta | None: def _heuristic_metadata(model_name: str) -> ModelMeta: logger.warning( "Falling back to name-pattern heuristic for %s; " - "activation-memory estimates will use BERT-base defaults (hidden=768, layers=12).", + "using CONSERVATIVE large-model defaults (params=%dM, hidden=%d, layers=%d) " + "so cost estimates upper-bound rather than under-predict.", model_name, + _DEFAULT_HEURISTIC_PARAMS // 1_000_000, + _DEFAULT_HEURISTIC_HIDDEN, + _DEFAULT_HEURISTIC_LAYERS, ) total_file_bytes = _DEFAULT_HEURISTIC_PARAMS * _DEFAULT_BYTES_PER_PARAM return ModelMeta( @@ -168,6 +187,8 @@ def _heuristic_metadata(model_name: str) -> ModelMeta: total_file_bytes=total_file_bytes, cached_locally=_is_warm_cached(model_name), confidence="heuristic", + hidden_size=_DEFAULT_HEURISTIC_HIDDEN, + n_layers=_DEFAULT_HEURISTIC_LAYERS, ) diff --git a/src/autointent/_advisor/runner.py b/src/autointent/_advisor/runner.py index 0a9a56f02..4743a3b47 100644 --- a/src/autointent/_advisor/runner.py +++ b/src/autointent/_advisor/runner.py @@ -7,7 +7,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable from pydantic import ValidationError @@ -31,6 +31,7 @@ def run_preflight( *, preset_name: str | None = None, refit_after: bool = False, + embedding_cache_probe: Callable[[str], bool] | None = None, ) -> PreflightReport: """Run all three preflight phases and return one report. @@ -43,6 +44,12 @@ def run_preflight( preset_name: optional friendly name for the report header. refit_after: matches the ``Pipeline.fit(refit_after=...)`` argument. When True, time estimates include the extra refit-on-full-data pass. + embedding_cache_probe: optional callable ``(embedder_model_name) -> bool``. + Return True when the embedding cache already holds this model's + embeddings for the current dataset — the advisor then predicts 0 + forward time and 0 ``disk_embedding_cache_gb`` for that embedder + (mirrors the ``cached_locally`` treatment for HF weights). Default + is the pessimistic cold assumption every embedder pays once. Returns: ``PreflightReport`` with findings across resource / data / config phases. @@ -79,6 +86,9 @@ def run_preflight( hardware=hardware, report=report, refit_after=refit_after, + cross_encoder_model_name=cfg.cross_encoder_config.model_name, + transformer_model_name=cfg.transformer_config.model_name, + cache_probe=embedding_cache_probe, ) _data_phase(cfg.search_space, stats, report) _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, hardware, report) diff --git a/src/autointent/_advisor/workflows.py b/src/autointent/_advisor/workflows.py index fad0562d6..91c3c24a7 100644 --- a/src/autointent/_advisor/workflows.py +++ b/src/autointent/_advisor/workflows.py @@ -263,3 +263,159 @@ def recommend( chosen = feasible[0][0] if feasible else None return RecommendationResult(chosen=chosen, results=results) + + +class ReduceToFitError(RuntimeError): + """Raised by :func:`reduce_to_fit` when no subset of the search space fits. + + The exception carries the final pruned config and the last report so callers + can still inspect what was tried — the review's contract was "raise, don't + silently degrade," which is exactly what this signals: even after removing + every module the advisor knows how to drop, at least one scoring node has + an OVER finding that no further pruning can resolve. + """ + + def __init__(self, message: str, *, pruned_config: dict[str, Any], last_report: PreflightReport) -> None: + super().__init__(message) + self.pruned_config = pruned_config + self.last_report = last_report + + +def _drop_module_from_search_space( + search_space: list[dict[str, Any]], node_type: str, module_name: str, +) -> list[dict[str, Any]]: + """Return a deep-copied search_space with ``module_name`` removed from the + matching ``node_type`` node. Nodes whose ``search_space`` becomes empty are + dropped entirely so the pipeline stays valid. + """ + import copy + + out: list[dict[str, Any]] = [] + for node in search_space: + node_copy = copy.deepcopy(node) + if node_copy.get("node_type") == node_type: + entries = [e for e in node_copy.get("search_space") or [] if e.get("module_name") != module_name] + node_copy["search_space"] = entries + if not entries: + # Node has nothing left to try — drop it. A missing decision or + # scoring node will surface as an OVER finding on the next + # preflight, terminating the loop cleanly. + continue + out.append(node_copy) + return out + + +def _pick_module_to_drop(report: PreflightReport) -> tuple[str, str] | None: + """Pick the (node_type, module_name) that contributes the most to whichever + budget is over. Returns ``None`` when no droppable driver exists (all + remaining rows are decision-node entries or unknown-cost placeholders). + + Preference order: VRAM > time > RAM > disk. We drop the driver with the + largest cost along the *first* dimension that has at least one OVER + finding — otherwise (edge case: is_feasible False without an OVER, which + shouldn't happen) fall back to VRAM. + """ + findings_by_metric = {f.metric for f in report.findings if f.severity == Severity.OVER} + priority = ["vram_gb", "time_hours", "ram_gb", "disk_download_gb"] + driver_key = next((k for k in priority if k in findings_by_metric), None) or "vram_gb" + + drivers = report.resource.drivers or [] + # Only drop scoring-node drivers — decision modules are lightweight and + # dropping the last one would leave the pipeline unable to make decisions. + candidates = [d for d in drivers if d.get("node_type") == "scoring" and d.get("module") not in {None, "?"}] + if not candidates: + return None + + def _cost(driver: dict[str, Any]) -> float: + raw = driver.get(driver_key) + return float(raw) if raw is not None else 0.0 + + heaviest = max(candidates, key=_cost) + module = heaviest.get("module") + if not isinstance(module, str): + return None + return "scoring", module + + +def reduce_to_fit( + config: dict[str, Any], + stats: DatasetStats, + hardware: Any, # noqa: ANN401 + *, + max_iters: int = 20, + refit_after: bool = False, +) -> tuple[dict[str, Any], PreflightReport]: + """Iteratively drop the most expensive infeasible module until the search space fits. + + Behavior: + * If ``config`` is already feasible, returns ``(config, report)`` unchanged. + * Otherwise picks the OVER-driving scoring-node module with the largest + cost along whichever budget breached (VRAM > time > RAM > disk) and + removes it from the search_space, then re-runs preflight. + * Repeats until feasible, ``max_iters`` reached, or no droppable module + remains — in the last two cases raises :class:`ReduceToFitError` + carrying the pruned config and final report. + + Args: + config: an OptimizationConfig-shaped dict (same input as :func:`run_preflight`). + stats: dataset stats to score against. + hardware: detected hardware profile. + max_iters: safety cap; a valid pipeline has ≤ ~10 scoring modules so + hitting the default cap means the picker is stuck (raises). + refit_after: forwarded to :func:`run_preflight`. + + Returns: + ``(pruned_config, report)`` where ``report.is_feasible`` is True. + + Raises: + ReduceToFitError: nothing fits after pruning. + """ + import copy + + current = copy.deepcopy(config) + report = run_preflight(current, stats, hardware, refit_after=refit_after) + if report.is_feasible and _has_scoring_module(current): + return current, report + + for _ in range(max_iters): + pick = _pick_module_to_drop(report) + if pick is None: + raise ReduceToFitError( + "No droppable scoring-node module found; remaining search space cannot be reduced further.", + pruned_config=current, + last_report=report, + ) + node_type, module_name = pick + current["search_space"] = _drop_module_from_search_space( + current["search_space"], node_type, module_name, + ) + logger.info("reduce_to_fit: dropped %s/%s to fit budget", node_type, module_name) + # An empty scoring node — after dropping the last scoring module — + # would look "feasible" to run_preflight (no drivers, no findings), so + # explicitly rule it out: an empty pipeline can't score anything. + if not _has_scoring_module(current): + raise ReduceToFitError( + "All scoring modules were pruned to fit the budget; the resulting pipeline " + "would have nothing to run. Raise the budget or add cheaper scoring modules.", + pruned_config=current, + last_report=report, + ) + report = run_preflight(current, stats, hardware, refit_after=refit_after) + if report.is_feasible: + return current, report + + raise ReduceToFitError( + f"Search space still infeasible after {max_iters} prune iterations.", + pruned_config=current, + last_report=report, + ) + + +def _has_scoring_module(config: dict[str, Any]) -> bool: + """True when ``config`` has at least one scoring-node entry left. Empty + scoring is a common outcome of pruning to the bone — reduce_to_fit treats + it as unfittable rather than "feasible with nothing to do.""" + for node in config.get("search_space", []): + if node.get("node_type") == "scoring" and node.get("search_space"): + return True + return False diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index cbaee2f81..e1ffd77d5 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -59,7 +59,29 @@ def test_heavy_preset_is_infeasible_on_2gb_budget() -> None: assert not report.is_feasible, "deberta-v3-large should not fit in 2 GB" -def test_light_preset_is_feasible_on_8gb_budget() -> None: +def test_light_preset_is_feasible_on_8gb_budget(monkeypatch: pytest.MonkeyPatch) -> None: + # This test runs under the offline fixture, which now returns + # ``_heuristic_metadata`` (conservative large-model shape) — that's + # deliberately pessimistic, so "light" would look infeasible on 8 GB. + # Restore small-model resolution just for this test so we're verifying + # the "light on 8 GB" contract, not the fallback pessimism. + from autointent._advisor import _hub + + def _small_model(name: str) -> _hub.ModelMeta: + return _hub.ModelMeta( + name=name, + total_params=140_000_000, + weight_bytes_per_param=4, + total_file_bytes=140_000_000 * 4, + cached_locally=False, + confidence="hub", + hidden_size=768, + n_layers=6, + ) + + _hub.resolve_model.cache_clear() + monkeypatch.setattr(_hub, "resolve_model", _small_model) + cfg = load_preset("transformers-light") stats = DatasetStats.placeholder(n_samples=1000, n_classes=10, avg_tokens=24) report = run_preflight(cfg, stats, _profile(vram_gb=8.0), preset_name="transformers-light") diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index fdcdf45f7..99bd63243 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -253,7 +253,9 @@ def test_offline_flips_low_confidence(self) -> None: } report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) assert report.low_confidence is True - assert any("Heuristic fallback" in n for n in report.notes) + # Low-confidence used to be a note; it's now a prominent finding so + # reviewers of the report see it in the main findings block. + assert any("LOW CONFIDENCE" in f.message for f in report.findings) def test_rare_classes_with_linear_scorer_flag_red(self) -> None: cfg = { @@ -673,7 +675,10 @@ def test_duplicate_knn_entries_zero_time_after_first(self) -> None: ], "hpo_config": {"n_trials": 5}, } - report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + # Use a large placeholder so per-step FLOPs are enough to register as + # non-zero rounded time even for tiny MiniLM. Behavior we're testing is + # "first entry pays, second is cached" — needs first > 0 to be visible. + report = run_preflight(cfg, DatasetStats.placeholder(n_samples=1_000_000), _profile()) knn_drivers = [d for d in report.resource.drivers if d["module"] == "knn"] assert len(knn_drivers) == 2 first, second = knn_drivers @@ -726,3 +731,41 @@ def test_disk_embedding_cache_scales_with_n_samples(self) -> None: big = run_preflight(cfg, DatasetStats.placeholder(n_samples=1_000_000), _profile()) assert small.resource.disk_embedding_cache_gb > 0 assert big.resource.disk_embedding_cache_gb > small.resource.disk_embedding_cache_gb * 100 + + def test_warm_cache_probe_zeroes_forward_and_disk(self) -> None: + """When ``embedding_cache_probe`` reports the embedder is warm, the + advisor must predict 0 forward time AND 0 ``disk_embedding_cache_gb`` + for that model — mirrors HF-weights ``cached_locally`` behavior.""" + cfg = { + "search_space": [ + self._embedder_node(), + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "knn", + "embedder_config": [ + {"model_name": "sentence-transformers/all-MiniLM-L6-v2"} + ], + "batch_size": [32], + "max_length": [128], + } + ], + }, + ], + "hpo_config": {"n_trials": 1}, + } + stats = DatasetStats.placeholder(n_samples=1_000_000) + cold = run_preflight(cfg, stats, _profile()) + warm = run_preflight(cfg, stats, _profile(), embedding_cache_probe=lambda _name: True) + + cold_knn = next(d for d in cold.resource.drivers if d["module"] == "knn") + warm_knn = next(d for d in warm.resource.drivers if d["module"] == "knn") + + assert cold_knn["time_hours"] > 0 + assert warm_knn["time_hours"] == 0 + assert "warm" in warm_knn["mode"] + assert cold.resource.disk_embedding_cache_gb > 0 + # Warm: forward wasn't charged → model isn't in ``cached_embedders`` → + # no disk_embedding_cache contribution. + assert warm.resource.disk_embedding_cache_gb == 0 diff --git a/tests/advisor/test_reduce_to_fit.py b/tests/advisor/test_reduce_to_fit.py new file mode 100644 index 000000000..08a7a9cb3 --- /dev/null +++ b/tests/advisor/test_reduce_to_fit.py @@ -0,0 +1,164 @@ +"""Tests for ``autointent._advisor.reduce_to_fit``. + +Covers the three review-mandated contracts: + +* a feasible config passes through unchanged; +* an infeasible config gets pruned to a config the advisor calls feasible; +* when nothing fits, we raise :class:`ReduceToFitError` — no silent degradation. + +Runs fully offline: the same ``_force_offline`` fixture pattern as the sibling +smoke tests, so HF Hub probes fall back to the heuristic large-model shape. +""" + +from __future__ import annotations + +import pytest + +from autointent._advisor import ( + DatasetStats, + HardwareProfile, + ReduceToFitError, + reduce_to_fit, + run_preflight, +) + + +@pytest.fixture(autouse=True) +def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: + from autointent._advisor import _hub + + _hub.resolve_model.cache_clear() + monkeypatch.setattr(_hub, "_hub_metadata", lambda _name: None) + + +def _profile(vram_gb: float = 16.0) -> HardwareProfile: + return HardwareProfile( + accelerator="cuda" if vram_gb > 0 else "cpu", + device_name="test-gpu" if vram_gb > 0 else "test-cpu", + vram_gb=vram_gb, + ram_gb=32.0, + free_disk_gb=200.0, + cpu_count=8, + ) + + +def _cheap_config() -> dict: + return { + "search_space": [ + { + "node_type": "scoring", + "target_metric": "scoring_f1", + "search_space": [{"module_name": "linear"}], + }, + { + "node_type": "decision", + "target_metric": "decision_accuracy", + "search_space": [{"module_name": "argmax"}], + }, + ], + } + + +def _big_and_cheap_config() -> dict: + """One expensive transformer + one cheap classic scorer. + + On a tiny (1 GB) VRAM budget, the transformer trips OVER; ``reduce_to_fit`` + should drop it and leave the classic one behind. + """ + return { + "search_space": [ + { + "node_type": "scoring", + "target_metric": "scoring_f1", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-large"}], + "batch_size": [128], + "max_length": [256], + }, + {"module_name": "linear"}, + ], + }, + { + "node_type": "decision", + "target_metric": "decision_accuracy", + "search_space": [{"module_name": "argmax"}], + }, + ], + } + + +def _unfittable_config() -> dict: + return { + "search_space": [ + { + "node_type": "scoring", + "target_metric": "scoring_f1", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-large"}], + "batch_size": [128], + "max_length": [512], + }, + ], + }, + { + "node_type": "decision", + "target_metric": "decision_accuracy", + "search_space": [{"module_name": "argmax"}], + }, + ], + } + + +def test_feasible_config_returns_unchanged() -> None: + stats = DatasetStats.placeholder(n_samples=500, n_classes=10, avg_tokens=24) + config = _cheap_config() + pruned, report = reduce_to_fit(config, stats, _profile(vram_gb=16.0)) + assert report.is_feasible + # Passthrough: same module still present. + modules = [ + e["module_name"] for node in pruned["search_space"] for e in node["search_space"] + ] + assert "linear" in modules + assert "argmax" in modules + + +def test_prunes_infeasible_transformer_to_classic() -> None: + stats = DatasetStats.placeholder(n_samples=2000, n_classes=20, avg_tokens=48) + config = _big_and_cheap_config() + + # Sanity check: base config must be infeasible on a tiny budget, otherwise + # this test isn't exercising the prune path. + base = run_preflight(config, stats, _profile(vram_gb=1.0)) + assert not base.is_feasible + + pruned, report = reduce_to_fit(config, stats, _profile(vram_gb=1.0)) + assert report.is_feasible + modules = [ + e["module_name"] for node in pruned["search_space"] for e in node["search_space"] + ] + assert "bert" not in modules, "expensive transformer should have been dropped" + assert "linear" in modules, "cheap classic scorer should be preserved" + + +def test_raises_when_nothing_fits() -> None: + stats = DatasetStats.placeholder(n_samples=2000, n_classes=20, avg_tokens=48) + config = _unfittable_config() + + with pytest.raises(ReduceToFitError) as exc_info: + reduce_to_fit(config, stats, _profile(vram_gb=0.5)) + + # The exception carries the final pruned config + last report so callers + # can inspect what was tried — contract from the review's follow-up. + err = exc_info.value + assert err.pruned_config is not None + assert err.last_report is not None + # After pruning the only scoring module, the config's scoring node should + # be gone entirely (or empty), leaving an unfittable pipeline. + scoring_nodes = [n for n in err.pruned_config["search_space"] if n.get("node_type") == "scoring"] + assert scoring_nodes == [] or all( + not n.get("search_space") for n in scoring_nodes + ) diff --git a/tests/pipeline/test_calibration_tracker.py b/tests/pipeline/test_calibration_tracker.py new file mode 100644 index 000000000..24196ce49 --- /dev/null +++ b/tests/pipeline/test_calibration_tracker.py @@ -0,0 +1,116 @@ +"""Tests for the calibration script's _ModuleTracker callback.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +# Add scripts/ to sys.path so the test can import calibrate_advisor. +_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "scripts" +sys.path.insert(0, str(_SCRIPTS_DIR)) + +from calibrate_advisor import _ModuleTracker, _classify_module_role, _sum_time_by_role # noqa: E402 + + +def test_tracker_records_wall_time_per_module() -> None: + """One (module, num) → one record with a positive duration.""" + tracker = _ModuleTracker() + + tracker.start_module("linear", 0, {"cv": 3}) + tracker.end_module() + + tracker.start_module("catboost", 1, {"iterations": 100, "depth": 6}) + tracker.end_module() + + assert len(tracker.records) == 2 + assert tracker.records[0]["module"] == "linear" + assert tracker.records[0]["num"] == 0 + assert tracker.records[0]["config"] == {"cv": 3} + assert tracker.records[0]["duration_s"] >= 0 + assert tracker.records[1]["module"] == "catboost" + assert tracker.records[1]["config"] == {"iterations": 100, "depth": 6} + + +def test_tracker_filters_non_scalar_config_values() -> None: + """Complex objects in module_kwargs must not appear in the recorded config.""" + tracker = _ModuleTracker() + tracker.start_module( + "bert", + 0, + {"cv": 3, "classification_model_config": {"model_name": "microsoft/deberta"}, "flag": True}, + ) + tracker.end_module() + assert tracker.records[0]["config"] == {"cv": 3, "flag": True} + + +def test_end_module_without_start_is_noop() -> None: + """Defensive: no crash when end_module is called without a matching start.""" + tracker = _ModuleTracker() + tracker.end_module() # must not raise + assert tracker.records == [] + + +def test_records_are_json_serialisable() -> None: + """Records must survive round-tripping through json.dumps for the CalibrationRow output.""" + import json + + tracker = _ModuleTracker() + tracker.start_module("linear", 0, {"cv": 3, "unused_none": None}) + tracker.end_module() + payload = json.dumps(tracker.records) + assert "linear" in payload + assert "duration_s" in payload + + +def test_running_peak_survives_low_last_module() -> None: + """peak_vram_gb_overall must retain the max across all modules. + + Pins the fix for a bug where torch.cuda's per-module reset_peak_memory_stats + clobbered the top-level VRAM reading with the last (usually CPU-only) module. + """ + tracker = _ModuleTracker() + + # Simulate a big embedder module. Populate _current directly to sidestep + # the real torch.cuda call and inject a synthetic peak. + tracker.start_module("linear", 0, {"cv": 3}) + tracker.end_module() + tracker.records[-1]["peak_vram_gb"] = 2.5 + tracker.peak_vram_gb_overall = max(tracker.peak_vram_gb_overall, 2.5) + + # Then a decision module that touches no VRAM. + tracker.start_module("threshold", 1, {"thresh": 0.5}) + tracker.end_module() + tracker.records[-1]["peak_vram_gb"] = 0.01 + tracker.peak_vram_gb_overall = max(tracker.peak_vram_gb_overall, 0.01) + + assert tracker.peak_vram_gb_overall == 2.5, "running peak clobbered by later small module" + + +def test_role_classification_and_time_decomposition() -> None: + """Each record carries a role, and _sum_time_by_role folds durations correctly. + + Pins R4-P1 #30: classic-preset wall-time must be decomposable into + embedder-forward vs scorer-fit vs decision-search so a consumer can + validate the advisor's per-role predictions without post-hoc classification. + """ + assert _classify_module_role("sentence_transformer") == "embedder" + assert _classify_module_role("hashing_vectorizer") == "embedder" + assert _classify_module_role("linear") == "scorer" + assert _classify_module_role("bert") == "scorer" + assert _classify_module_role("threshold") == "decision" + assert _classify_module_role("argmax") == "decision" + + tracker = _ModuleTracker() + tracker.start_module("sentence_transformer", 0, {}) + tracker.end_module() + tracker.records[-1]["duration_s"] = 8.0 + tracker.start_module("linear", 1, {}) + tracker.end_module() + tracker.records[-1]["duration_s"] = 2.0 + tracker.start_module("threshold", 2, {}) + tracker.end_module() + tracker.records[-1]["duration_s"] = 0.5 + + assert [r["role"] for r in tracker.records] == ["embedder", "scorer", "decision"] + totals = _sum_time_by_role(tracker.records) + assert totals == {"embedder": 8.0, "scorer": 2.0, "decision": 0.5} From 0dfcd6dd0de59d6328af000b14b28cfa1a2d3fbb Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:42:13 +0300 Subject: [PATCH 27/43] fix _StepTimingCallback missing on_train_begin (and every other HF hook) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HF's CallbackHandler.call_event dispatches with a bare ``getattr(callback, event)(...)`` — no hasattr probe — so a plain class that only implements on_step_begin/on_step_end crashes with ``AttributeError: '_StepTimingCallback' object has no attribute 'on_train_begin'`` the moment a real bert trial runs through the calibrator's step-timing patch. Fix by subclassing ``transformers.TrainerCallback`` directly: every ``on_*`` hook is inherited as a proper no-op, so we only override the two we time. Lazy try/except on the import keeps the module loadable in classic-only environments — the fallback base is only used for the class definition (the callback is never instantiated there because ``_patch_trainer_for_step_timing`` bails out in the same ImportError branch). Regression test in test_calibration_tracker.py pins the isinstance contract. Co-Authored-By: Claude Opus 4.7 --- scripts/calibrate_advisor.py | 31 ++++++++++++++----- tests/pipeline/test_calibration_tracker.py | 36 +++++++++++++++++++++- 2 files changed, 58 insertions(+), 9 deletions(-) diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py index 3f18b1cec..d4dec0f29 100644 --- a/scripts/calibrate_advisor.py +++ b/scripts/calibrate_advisor.py @@ -367,7 +367,26 @@ def _classify_module_role(module_name: str) -> str: return "scorer" -class _StepTimingCallback: +# Base class for _StepTimingCallback. We inherit from HF's real +# ``TrainerCallback`` when transformers is installed — that gives us the +# correct no-op default for every ``on_*`` hook (on_train_begin/on_log/ +# on_save/...) automatically, so we only override the two we time. HF's +# ``CallbackHandler.call_event`` dispatches with a bare ``getattr`` (no +# hasattr probe), so a plain class missing hooks would ``AttributeError`` +# the moment a real trial calls e.g. ``on_train_begin``. +# +# When transformers isn't installed we fall back to ``object`` so the +# harness still imports on classic-only runs. In that case the callback is +# never actually instantiated (``_patch_trainer_for_step_timing`` bails out +# in the same ``ImportError`` branch), so the fallback base is only needed +# to make the class definition itself succeed. +try: + from transformers import TrainerCallback as _StepTimingBase # type: ignore[import-not-found] +except ImportError: + _StepTimingBase = object # type: ignore[assignment,misc] + + +class _StepTimingCallback(_StepTimingBase): # type: ignore[misc,valid-type] """HF ``TrainerCallback`` that appends the wall-time of each optimizer step to a caller-owned list. @@ -376,17 +395,15 @@ class _StepTimingCallback: step buffer on :class:`_ModuleTracker`, so the transformer's per-step latency lands in that module's record automatically — no plumbing across module boundaries. - - Duck-typed (not a ``TrainerCallback`` subclass) so importing transformers - stays lazy — the harness must work on classic-only runs without the - transformers extra. """ def __init__(self, sink: list[float]) -> None: + # TrainerCallback.__init__ takes (*args, **kwargs); calling super is + # safe both when the base is the real HF class and when it's ``object``. + super().__init__() self._sink = sink self._t0: float | None = None - # HF's CallbackHandler calls these positionally with (args, state, control, **kwargs) def on_step_begin(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None: # noqa: ANN401, ARG002 self._t0 = time.perf_counter() @@ -395,8 +412,6 @@ def on_step_end(self, args: Any, state: Any, control: Any, **kwargs: Any) -> Non self._sink.append(time.perf_counter() - self._t0) self._t0 = None - # HF's CallbackHandler probes each callback with hasattr; leave the rest unset. - def _summarize_step_times(step_times: list[float]) -> dict[str, float]: """Fold a list of per-step wall-times into summary stats for the row. diff --git a/tests/pipeline/test_calibration_tracker.py b/tests/pipeline/test_calibration_tracker.py index 24196ce49..7e5640035 100644 --- a/tests/pipeline/test_calibration_tracker.py +++ b/tests/pipeline/test_calibration_tracker.py @@ -9,7 +9,12 @@ _SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "scripts" sys.path.insert(0, str(_SCRIPTS_DIR)) -from calibrate_advisor import _ModuleTracker, _classify_module_role, _sum_time_by_role # noqa: E402 +from calibrate_advisor import ( # noqa: E402 + _ModuleTracker, + _StepTimingCallback, + _classify_module_role, + _sum_time_by_role, +) def test_tracker_records_wall_time_per_module() -> None: @@ -114,3 +119,32 @@ def test_role_classification_and_time_decomposition() -> None: assert [r["role"] for r in tracker.records] == ["embedder", "scorer", "decision"] totals = _sum_time_by_role(tracker.records) assert totals == {"embedder": 8.0, "scorer": 2.0, "decision": 0.5} + + +def test_step_timing_callback_answers_every_hf_hook() -> None: + """HF's CallbackHandler dispatches with a bare ``getattr(cb, event)`` — no + hasattr probe — so the callback MUST answer every ``on_*`` hook it may + ever ask for, even ones we don't time. + + Pins the regression that surfaced as + ``AttributeError: '_StepTimingCallback' object has no attribute 'on_train_begin'`` + when a bert scorer trial fired the harness's step-timing patch on real HF + Trainer machinery. The fix is to subclass ``transformers.TrainerCallback`` + directly so every default hook is inherited as a no-op — no ``__getattr__`` + trickery, no per-hook boilerplate. + """ + from transformers import TrainerCallback + + cb = _StepTimingCallback(sink=[]) + # Must actually be an HF TrainerCallback — this is the guarantee that + # every hook HF may dispatch resolves to an inherited pass-through. + assert isinstance(cb, TrainerCallback) + + # Sanity-check a representative slice of hooks (every documented HF hook + # subclass has one, and the isinstance above already proves the rest are + # inherited). Calling them with (args, state, control, **kwargs) must + # succeed and return None — HF's call_event keeps the incoming control + # unchanged when the result is None. + for name in ["on_init_end", "on_train_begin", "on_epoch_begin", "on_log", "on_save", "on_train_end"]: + hook = getattr(cb, name) + assert hook(None, None, "control", model=None) is None, name From b8f1b240856413b40f24cb26aea9102428a3b6a6 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:59:24 +0300 Subject: [PATCH 28/43] improve advisor --- .../_advisor/_estimates/_formulas.py | 232 ++++++++++++++-- .../_advisor/_estimates/_resource.py | 250 +++++++++++++++++- .../_advisor/_estimates/_search_space.py | 44 +++ src/autointent/_advisor/runner.py | 33 ++- 4 files changed, 520 insertions(+), 39 deletions(-) diff --git a/src/autointent/_advisor/_estimates/_formulas.py b/src/autointent/_advisor/_estimates/_formulas.py index 6fe3985b1..795390f17 100644 --- a/src/autointent/_advisor/_estimates/_formulas.py +++ b/src/autointent/_advisor/_estimates/_formulas.py @@ -122,12 +122,21 @@ def _vram_for_transformer( Activation accounting differs by mode — training keeps per-layer outputs for backward; inference only needs one or two layers in flight. + + Final ``* 1.20`` is a safety margin covering allocator fragmentation, + peak transient tensors during backward, and HF-Trainer's eval-loop + double-forward that the textbook accounting above misses. Advisor should + upper-bound: earlier 15% margin left transformers-light on banking77 at + 8.12 GB predicted vs 8.75 GB measured (1.08x UNDER — unsafe for an OOM + tool). Bumped to 20% + a bigger fixed CUDA baseline (see + ``_CUDA_BASELINE_VRAM_GB``) to close the gap on large-batch training runs + without over-inflating small ones. """ base = _weights_vram_for_transformer(meta, mode) if batch_size <= 0: return base per_sample = _activations_gb_per_sample(meta, seq_len, is_training=mode != "inference") - return base + per_sample * batch_size + return (base + per_sample * batch_size) * 1.20 def _max_fitting_batch_size( @@ -151,17 +160,31 @@ def _max_fitting_batch_size( # Sustained TFLOPS per device class — real MFU (model-FLOPs utilization) at -# training batch sizes, NOT peak spec sheet numbers. Numbers reflect ~30-50% -# MFU which is typical for BERT-scale training with FA2 / cuDNN kernels. -# Source: MLPerf training results + community benchmarks (2024-2025). +# training batch sizes, NOT peak spec sheet numbers. Numbers reflect ~20-30% +# MFU which is what HF Trainer actually achieves on BERT-scale workloads once +# you factor in dataloader idle, tokenization warmup, per-epoch eval, and +# checkpoint saves — all of which the raw FLOPs formula ignores. Earlier +# values (150 / 45 / 15 for high/mid/low) were closer to peak spec numbers +# and under-predicted transformers-heavy on banking77 by 1.7x (measured +# 3.17 h vs predicted 1.88 h, calibration_runs2 2026-08-07); an advisor +# should upper-bound, so pick sustained numbers that err on the side of +# over-predicting. Source: MLPerf training results + measured banking77 +# runs where mean_step_s / p95_step_s → 154ms / 246ms for bert-base bs=64. _DEVICE_TFLOPS = { - "high-gpu": 150.0, # A100 / H100 fp16 - "mid-gpu": 45.0, # V100 / RTX 3090 / A6000 - "low-gpu": 15.0, # T4 / RTX 3060 / 8 GB consumer card - "apple-silicon": 8.0, # M1/M2/M3 GPU cores - "cpu": 0.1, # single-thread modern x86 with MKL + "high-gpu": 60.0, # A100 / H100 — sustained ~19% MFU under HF Trainer + "mid-gpu": 20.0, # V100 / RTX 3090 / A6000 + "low-gpu": 7.0, # T4 / RTX 3060 / 8 GB consumer card + "apple-silicon": 4.0, # M1/M2/M3 GPU cores + "cpu": 0.05, # single-thread modern x86 with MKL } -_DEFAULT_TFLOPS = 15.0 # unknown device → treat as low-GPU +_DEFAULT_TFLOPS = 7.0 # unknown device → treat as low-GPU + +# HF Trainer overhead — the FLOPs formula only counts optimizer steps; real +# wall-time also includes per-epoch eval sweeps, save-checkpoint syncs, +# tokenizer warmup, dataloader queue idle, and gradient-accumulation gaps. +# Factor calibrated so transformers-heavy predicts ~1.2-1.5x the measured +# 3.17 h (upper-bound stance). +_TRAINER_OVERHEAD_MULT = 1.35 def _time_for_transformer( @@ -177,12 +200,14 @@ def _time_for_transformer( """Transformer training time in hours, from per-step FLOPs / device TFLOPS. Per-step FLOPs ≈ 6 x params x batch_size x seq_len (2 for forward mul-add, - 3-4x for backward). Divided by sustained device TFLOPS to get wall-time per - step, then multiplied by (steps x epochs x n_trials). - - Replaces an earlier "1 second per step" heuristic, which was ~10x too high - on A100 and identical for MPS vs CUDA (predicted times were the same on - both while real times differed ~7x — see interpretation.md 2026-07-19). + 3-4x for backward). Divided by *sustained* device TFLOPS (not peak spec) + to get wall-time per step, then multiplied by (steps x epochs x n_trials) + x ``_TRAINER_OVERHEAD_MULT`` for HF-Trainer wall-clock overhead. + + Advisor contract: err on the side of over-prediction. Under-predicting + time makes users blow through wall-clock budgets; over-predicting only + biases them toward smaller / cheaper presets. See ``_DEVICE_TFLOPS`` + docstring for the sustained-MFU calibration. """ steps_per_epoch = max(1, n_samples // max(1, batch_size)) total_steps = n_trials * epochs * steps_per_epoch @@ -190,7 +215,7 @@ def _time_for_transformer( step_flops = 6.0 * params_millions * 1e6 * batch_size * seq_len tflops = _DEVICE_TFLOPS.get(device_class, _DEFAULT_TFLOPS) step_seconds = step_flops / (tflops * 1e12) - return (total_steps * step_seconds) / 3600.0 + return (total_steps * step_seconds * _TRAINER_OVERHEAD_MULT) / 3600.0 def _n_layers(meta: ModelMeta | None) -> int: @@ -214,16 +239,29 @@ def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: return max(seen_models.values(), key=lambda m: m.total_params) -def _ram_for_module(meta: ModelMeta, stats: DatasetStats) -> float: - """RAM in GB. Loose upper bound: weights + tokenized text in memory. +def _ram_for_module(meta: ModelMeta, stats: DatasetStats, *, mode: str = "inference") -> float: + """RAM in GB. Loose upper bound: weights + optimizer/grads + tokenized text. Tokenized text is approximated as ``n_samples x avg_tokens x 4 bytes`` - (BPE/WordPiece token ids fit in int32). The 4 bytes/token bound is tight - enough for the report's purposes and intentionally ignores any preprocessing - artefacts (attention masks, position ids, etc.) since they're bounded by the - same factor. + (BPE/WordPiece token ids fit in int32). + + ``mode``-dependent multiplier on the weights term: + * ``inference``: 1.3x (weights + intermediate-tensor slack) + * ``lora``: 1.5x (frozen base + trainable adapters) + * ``full-finetune`` / anything else: 4.5x (weights + grads + Adam m + v + + framework slack) — matches the VRAM-side ``4.5W`` accounting so the + host-pinned optimizer state (Adam mirrors weights) shows up in the + RAM estimate too. Under-predicting RAM lets a training preset OOM the + host well before it OOMs the GPU; advisor should upper-bound. """ - return meta.weights_gb + (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB + if mode == "inference": + weights_mult = 1.3 + elif mode == "lora": + weights_mult = 1.5 + else: + weights_mult = 4.5 + tokens_gb = (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB + return meta.weights_gb * weights_mult + tokens_gb def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: @@ -301,6 +339,65 @@ def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, return float((data_bytes + histograms_bytes + trees_bytes) / _BYTES_PER_GB) +def _ram_for_sklearn( + *, + stats: DatasetStats, + embedder_dim: int, + n_estimators: int, + max_depth: int, + n_jobs: int, +) -> float: + """RandomForest/similar RAM upper bound, aware of ``n_jobs`` replication. + + sklearn spawns ``n_jobs`` worker processes with joblib's loky backend by + default; each worker holds its own copy of the training feature matrix and + the trees it grew, so a preset with ``n_jobs=8`` on a 10 k × 1024 embedder + dataset multiplies the base RAM 8x. Previously sklearn was in + ``_UNKNOWN_SCORER_MODULES`` and emitted a zero row — classic-heavy's + real ~2 GB sklearn contribution slipped through invisibly. + """ + # Per-worker feature matrix (fp64 in sklearn by default). + per_worker_data = stats.n_samples * embedder_dim * 8 + # Per-worker tree storage: n_estimators × n_leaves × ~32 B/node. Cap + # n_leaves at n_samples (a tree with max_depth 150 on 10k samples can't + # actually have 2**150 leaves). + n_leaves = min(2**max_depth, stats.n_samples) if max_depth > 0 else stats.n_samples + per_worker_trees = n_estimators * n_leaves * _CATBOOST_BYTES_PER_TREE_NODE + per_worker = per_worker_data + per_worker_trees + return float((per_worker * max(1, n_jobs)) / _BYTES_PER_GB) + + +def _embedder_load_ram_gb(meta: ModelMeta | None) -> float: + """Extra RAM the process holds when an embedder is loaded — separately from + any per-driver row that already accounts for it. + + Rationale: classic presets pre-compute embeddings via the embedder, then + train sklearn/catboost/linear scorers on top. During and after that + forward pass the process holds: the embedder weights on the compute + device, a copy in CPU RAM (fp32 from the safetensors load), the tokenizer + state, HF Trainer buffers, and the cached embeddings themselves. The + per-driver ``_ram_for_module`` already captures weights x 1.3 for the + knn/mlknn rows, but the aggregate ``max`` across drivers hides + contributions from the other classic scorers that are simultaneously in + memory. This term is added *on top* of the max-driver RAM so classic + presets like classic-heavy stop under-predicting by ~4x. + + Uses ``total_params × 4`` (fp32) as the weight footprint even when the + hub reports fp16 storage (``weight_bytes_per_param=2``) — transformers + up-casts to fp32 at load time by default, so the fp16 disk size + under-counts real RAM usage by 2x. + """ + if meta is None: + return 0.0 + fp32_weights_gb = (meta.total_params * 4) / _BYTES_PER_GB + # 3.5x factor: raw weights + activation buffers + HF/tokenizer/loader + # slack. Empirical: real classic-heavy on banking77 (e5-large embedder, + # sklearn RF n_jobs=8) measured 10.22 GB RAM. At 3.0x we predicted 9.56 + # (1.07x under — still unsafe); at 3.5x we predict 10.85 (0.94x — safely + # over). Advisor's OOM-avoidance contract requires "err over". + return fp32_weights_gb * 3.5 + + def _time_for_catboost( *, n_trials: int, @@ -325,6 +422,93 @@ def _time_for_catboost( return seconds / 3600.0 +# === CNN / RNN scorers =================================================== +# +# Small torch models (Kim's TextCNN, LSTM classifier) trained from scratch +# on token ids. The advisor previously emitted a ``not-estimated`` placeholder +# for these — which read as "free/safe" on nn-heavy / nn-medium (predicted 0h +# / 0GB, real 0.3 h + 2.3 GB RAM + 0.7 GB VRAM on banking77). +# +# Cost model: embedding table + head weights (small) + activations that scale +# with batch × seq_len × hidden. We assume a bounded vocabulary (~30k) — one +# HPO trial for TextCNN embeds every token that appears in the training set; +# banking77 tops out around a few thousand unique tokens so 30k is a safe +# upper bound. +_NN_MAX_VOCAB = 30_000 +_NN_DEFAULT_SEQ_LEN = 50 # VocabConfig.max_seq_length default +_NN_BYTES_PER_PARAM = 4 # fp32 weights +# fp32 activation storage per (batch, token, hidden) unit, factor absorbs +# ~4x backward overhead + optimizer + gradient state for these tiny models. +_NN_TRAIN_ACT_BYTES_PER_UNIT = 16 + + +def _cnn_param_count(*, embed_dim: int, num_filters: int, n_kernels: int, n_classes: int) -> int: + """Approximate CNN parameter count: embedding + conv + fc layers.""" + vocab_params = _NN_MAX_VOCAB * embed_dim + conv_params = num_filters * embed_dim * n_kernels * 5 # avg kernel width ~5 + fc_params = num_filters * n_kernels * max(1, n_classes) + return vocab_params + conv_params + fc_params + + +def _rnn_param_count(*, embed_dim: int, hidden_dim: int, n_classes: int) -> int: + """Approximate LSTM classifier parameter count: embedding + LSTM + fc.""" + vocab_params = _NN_MAX_VOCAB * embed_dim + # LSTM cell has 4 gates, each with (embed+hidden+1) × hidden params. + lstm_params = 4 * hidden_dim * (embed_dim + hidden_dim + 1) + fc_params = hidden_dim * max(1, n_classes) + return vocab_params + lstm_params + fc_params + + +def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: + """Weights + activations for a small torch scorer in training mode. + + Activation term uses ``batch × seq_len × hidden × const`` per the same + fp32 upper bound as transformers, but with a much smaller effective + hidden dim (embed_dim/num_filters, not model dim). + """ + weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB + # Optimizer state (Adam has 2x weights) + gradients (1x weights) = 4x weights total. + optimizer_gb = 3 * weights_gb + activations_gb = ( + batch_size * _NN_DEFAULT_SEQ_LEN * hidden_dim * _NN_TRAIN_ACT_BYTES_PER_UNIT + ) / _BYTES_PER_GB + return weights_gb + optimizer_gb + activations_gb + + +def _ram_for_nn(*, params: int, stats: DatasetStats) -> float: + """CPU-side memory: weights + tokenized text (int32 ids).""" + weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB + tokens_gb = (stats.n_samples * _NN_DEFAULT_SEQ_LEN * 4) / _BYTES_PER_GB + return weights_gb + tokens_gb + + +def _time_for_nn( + *, + n_trials: int, + epochs: int, + batch_size: int, + n_samples: int, + params_millions: float, + device_class: str, +) -> float: + """Reuse the transformer FLOPs formula for a small torch model. + + Small models are memory-bandwidth-bound, not compute-bound, so the FLOPs + formula slightly *under*-predicts wall-time. Empirically for banking77 + nn-heavy we measured 0.32 h across 55 trials — the formula lands within + 2x of that, which is enough for a cost-ranking estimate. + """ + return _time_for_transformer( + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + seq_len=_NN_DEFAULT_SEQ_LEN, + n_samples=n_samples, + params_millions=params_millions, + device_class=device_class, + ) + + def _floor_to_power_of_two(n: int) -> int: """Largest power of two <= ``n``; returns 0 when ``n < 1``.""" if n < 1: diff --git a/src/autointent/_advisor/_estimates/_resource.py b/src/autointent/_advisor/_estimates/_resource.py index 884254faf..acf9fbb3d 100644 --- a/src/autointent/_advisor/_estimates/_resource.py +++ b/src/autointent/_advisor/_estimates/_resource.py @@ -24,24 +24,38 @@ from ._formulas import ( _DEFAULT_SEQ_LEN, + _LINEAR_CPU_S_PER_SAMPLE_FEATURE, _LOGREG_CV_MULTIPLIER, _MULTICLASS_THRESHOLD, _activations_gb_per_sample, _classify_severity, + _cnn_param_count, _embedder_dim, + _embedder_load_ram_gb, _embedding_cache_disk_gb, _largest_embedder, _max_fitting_batch_size, _ram_for_catboost, _ram_for_linear, _ram_for_module, + _ram_for_nn, + _ram_for_sklearn, + _rnn_param_count, _time_for_catboost, _time_for_linear, + _time_for_nn, _time_for_transformer, + _vram_for_nn, _vram_for_transformer, _weights_vram_for_transformer, ) -from ._search_space import _extract_model_names, _max_int, _walk_modules_indexed +from ._search_space import ( + _extract_model_names, + _max_int, + _module_cardinality, + _walk_modules, + _walk_modules_indexed, +) if TYPE_CHECKING: from autointent._advisor._hardware import HardwareProfile @@ -122,7 +136,7 @@ def _split_entries( transformer: list[tuple[int, str, dict[str, Any]]] = [] classic: list[tuple[int, str, dict[str, Any]]] = [] for node_idx, node_type, entry in _walk_modules_indexed(search_space): - bucket = classic if entry.get("module_name") in {"linear", "catboost"} else transformer + bucket = classic if entry.get("module_name") in {"linear", "catboost", "sklearn"} else transformer bucket.append((node_idx, node_type, entry)) return transformer, classic @@ -146,7 +160,7 @@ def _estimate_transformer_model( seq_len = _max_int(entry.get("max_length"), _DEFAULT_SEQ_LEN) vram = _vram_for_transformer(meta, mode, batch_size=batch_size, seq_len=seq_len) - ram = _ram_for_module(meta, stats) + ram = _ram_for_module(meta, stats, mode=mode) driver_max_batch: int | None = None if hardware.vram_gb > 0: @@ -243,6 +257,35 @@ def _estimate_classic_entry( ) vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) mode = "catboost-gpu" if on_gpu else "catboost" + elif module == "sklearn": + # RandomForestClassifier is the most common target; joblib spawns + # ``n_jobs`` worker processes each replicating the feature matrix + + # trees. Predicting 0 here (previous "not-estimated" behaviour) hid + # classic-heavy's real 1-2 GB sklearn contribution. + n_estimators = _max_int(entry.get("n_estimators"), 100) + max_depth = _max_int(entry.get("max_depth"), 0) + sk_n_jobs = _max_int(entry.get("n_jobs"), 1) + ram = _ram_for_sklearn( + stats=stats, + embedder_dim=embedder_dim, + n_estimators=n_estimators, + max_depth=max_depth, + n_jobs=sk_n_jobs, + ) + # Time: rough O(n_estimators × n_samples × sqrt(features) × log2 n) + # per fit, divided by n_jobs. Absorbed into the linear coefficient + # since real numbers vary wildly by criterion / max_features. + time_h = ( + n_trials + * _LINEAR_CPU_S_PER_SAMPLE_FEATURE + * stats.n_samples + * embedder_dim + * n_estimators + / max(1, sk_n_jobs) + / 3600.0 + ) * refit + vram = 0.0 + mode = f"sklearn-n_jobs={sk_n_jobs}" else: return None @@ -265,6 +308,86 @@ def _estimate_classic_entry( ) +def _estimate_nn_entry( + *, + entry: dict[str, Any], + node_type: str, + stats: DatasetStats, + hardware: HardwareProfile, + n_trials: int, + refit_after: bool, +) -> _ModuleEstimate | None: + """Cost row for a cnn / rnn scorer (returns ``None`` for anything else). + + These are small torch models trained from scratch on token ids. Previously + the advisor emitted a ``not-estimated`` placeholder for them, which read + as "free/safe" — nn-heavy on banking77 predicted 0h/0GB but actually used + 0.32 h + 2.3 GB RAM + 0.7 GB VRAM. This restores a real estimate using + small-model parameter counts + the transformer FLOPs formula for time. + """ + module = entry.get("module_name", "?") + n_classes = max(1, stats.n_classes) + + if module == "cnn": + embed_dim = _max_int(entry.get("embed_dim"), 128) + num_filters = _max_int(entry.get("num_filters"), 100) + kernel_sizes = entry.get("kernel_sizes") + # Kernel sizes are a list of lists in the search space + # (e.g. [[3, 4, 5]]); count entries in the largest variant. + n_kernels = 3 + if isinstance(kernel_sizes, list): + for candidate in kernel_sizes: + if isinstance(candidate, list): + n_kernels = max(n_kernels, len(candidate)) + elif isinstance(candidate, int): + n_kernels = max(n_kernels, 1) + hidden_dim = num_filters + params = _cnn_param_count( + embed_dim=embed_dim, num_filters=num_filters, n_kernels=n_kernels, n_classes=n_classes + ) + elif module == "rnn": + embed_dim = _max_int(entry.get("embed_dim"), 128) + hidden_dim = _max_int(entry.get("hidden_dim"), 512) + params = _rnn_param_count(embed_dim=embed_dim, hidden_dim=hidden_dim, n_classes=n_classes) + else: + return None + + batch_size = _max_int(entry.get("batch_size"), 64) + epochs = _max_int(entry.get("num_train_epochs"), 60) + + vram = _vram_for_nn(params=params, batch_size=batch_size, hidden_dim=hidden_dim) + ram = _ram_for_nn(params=params, stats=stats) + time_h = ( + _time_for_nn( + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + params_millions=params / 1_000_000, + device_class=hardware.device_class, + ) + * _refit_factor(refit_after=refit_after, n_trials=n_trials) + ) + + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": f"{module}-from-scratch", + "mode": "small-torch-train", + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "batch_size": batch_size, + "max_batch_size": None, + "confidence": "heuristic", + }, + vram_gb=vram, + ram_gb=ram, + time_hours=time_h, + ) + + def _apply_embedding_cache( module_estimates: list[_ModuleEstimate], seen_models: dict[str, ModelMeta], @@ -422,13 +545,48 @@ def _emit_resource_findings( ) -_UNKNOWN_SCORER_MODULES = frozenset({"cnn", "rnn", "sklearn"}) -"""Scorer modules the advisor has no cost estimator for. +# Process-level memory floors. Every autointent fit imports torch + +# transformers + datasets + optuna, which reserve resident memory the moment +# they load. Measured against calibration_runs2 (2026-08-07 banking77 sweep): +# every preset used 2-10 GB RAM but the per-module estimates alone predicted +# 0.4-2 GB → advisor was systematically 4-13x LOW on RAM. A ~1.5 GB baseline +# lifts the estimate into range without over-inflating heavy presets. +_PROCESS_BASELINE_RAM_GB = 1.5 +# CUDA driver context + cuDNN/cuBLAS workspace pools + caching allocator +# fragmentation. A real training process on A100 reserves ~1 GB the moment +# torch initializes CUDA + the first tensor lands, regardless of model size. +# Earlier 0.5 GB baseline left transformers-light on banking77 at 8.12 GB +# predicted vs 8.75 GB measured (unsafe under-prediction for OOM avoidance); +# 1.0 GB closes the gap with room to spare. Only added when we already +# predict some VRAM usage so CPU-only presets aren't spuriously flagged as +# GPU users. +_CUDA_BASELINE_VRAM_GB = 1.0 + +_UNKNOWN_SCORER_MODULES: frozenset[str] = frozenset() +"""Scorer modules the advisor has no cost estimator for — kept as an empty +extension point. cnn / rnn moved to :func:`_estimate_nn_entry`; sklearn moved +to the classic branch of :func:`_estimate_classic_entry`. New unknown-cost +scorers should still register here so they emit a not-estimated placeholder +row instead of silently reporting zero.""" + +_NN_SCORER_MODULES = frozenset({"cnn", "rnn"}) +"""Small torch scorers routed through :func:`_estimate_nn_entry`.""" + + +_EMBEDDER_CONSUMING_MODULES = frozenset( + {"linear", "catboost", "sklearn", "knn", "mlknn", "retrieval", + "description_bi", "description_cross", "description_llm"}, +) -These get a placeholder ``not-estimated`` driver row so they never appear as -"free/safe" in the report — silent-zero was the ``nn-heavy`` predicted 0h/0GB -bug that hid a real 0.52 h + 1.45 GB RAM cost. -""" + +def _uses_embedder(search_space: list[dict[str, Any]]) -> bool: + """True when any module in the search space consumes an embedder — signals + that the aggregate RAM should include the embedder-load penalty on top + of the per-driver max.""" + for _, entry in _walk_modules(search_space): + if entry.get("module_name") in _EMBEDDER_CONSUMING_MODULES: + return True + return False # Modules that consume the top-level ``cross_encoder_config.model_name`` as # their scoring model (see zero-shot-encoders preset: description_cross pulls @@ -505,6 +663,34 @@ def _resource_phase( transformer_entries, classic_entries = _split_entries(search_space) + # Per-node module-variant count. HPO distributes ``n_trials`` across the + # module_name candidates at each node roughly evenly (TPE's sampler bias + # aside), so a module_name that shares its node with 4 others sees on + # average ``n_trials / 5`` trials — not ``n_trials``. Previously every + # per-module estimate used the full ``n_trials``, which inflated + # classic-heavy's catboost row to 120 h vs measured 12 h across the whole + # node. Divide once here and pass the effective share down to every + # per-module estimator (transformer + classic + nn). + variants_per_node: dict[int, int] = {} + for node_idx, _node_type, _entry in _walk_modules_indexed(search_space): + variants_per_node[node_idx] = variants_per_node.get(node_idx, 0) + 1 + + def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int: # noqa: ARG001 + """Trials this specific module should be charged for. + + Divides ``n_trials`` evenly across the module_name candidates at the + node. ``entry`` is accepted for API stability and future extensions + (e.g. a per-module cardinality cap) — we tried capping by + :func:`_module_cardinality` earlier but it produced wrong estimates + for description-scorer presets where Optuna's TPE runs every + declared trial regardless of parameter-space size (no automatic + dedup). The advisor charges for the declared work; if a real run + crashes or dedupes, that's an artefact of the runtime, not something + the advisor should try to predict. + """ + divisor = max(1, variants_per_node.get(node_idx, 1)) + return max(1, n_trials // divisor) + # First pass: transformer modules (also populates seen_models for the classic pass). module_estimates: list[_ModuleEstimate] = [] node_max_weights: dict[int, float] = {} @@ -518,6 +704,20 @@ def _resource_phase( model_names = [cross_encoder_model_name] elif module in _TRANSFORMER_CONFIG_CONSUMERS and transformer_model_name: model_names = [transformer_model_name] + elif module in _NN_SCORER_MODULES: + # cnn / rnn — small torch models trained from scratch, no hub + # model to resolve. Route to the small-model heuristic. + nn_estimate = _estimate_nn_entry( + entry=entry, + node_type=node_type, + stats=stats, + hardware=hardware, + n_trials=_effective_trials(node_idx, entry), + refit_after=refit_after, + ) + if nn_estimate is not None: + module_estimates.append(nn_estimate) + continue elif module in _UNKNOWN_SCORER_MODULES: # Placeholder so the row is visible instead of silently zeroed. module_estimates.append(_not_estimated_row(node_type=node_type, module=module)) @@ -532,7 +732,7 @@ def _resource_phase( name=name, stats=stats, hardware=hardware, - n_trials=n_trials, + n_trials=_effective_trials(node_idx, entry), refit_after=refit_after, ) module_estimates.append(me) @@ -543,7 +743,7 @@ def _resource_phase( # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. embedder_meta = _largest_embedder(seen_models) embedder_dim_val = _embedder_dim(embedder_meta) - for _, node_type, entry in classic_entries: + for node_idx, node_type, entry in classic_entries: classic_estimate = _estimate_classic_entry( entry=entry, node_type=node_type, @@ -551,7 +751,7 @@ def _resource_phase( embedder_dim=embedder_dim_val, stats=stats, hardware=hardware, - n_trials=n_trials, + n_trials=_effective_trials(node_idx, entry), refit_after=refit_after, ) if classic_estimate is not None: @@ -569,6 +769,32 @@ def _resource_phase( estimate.time_hours += me.time_hours estimate.drivers.append(me.driver) + # Process-level baselines. Every fit — even a trivial one — imports + # torch / transformers / datasets, which on their own take ~1.5 GB of RSS + # before any model weights load. Real runs on banking77 measure 2-10 GB + # RAM across every preset while the per-module estimates alone predicted + # 0.4-2 GB (systematically 4-13x low, see calibration_runs2 2026-08-07). + # Adding a floor here (rather than per-module) means the estimate stays + # accurate when multiple modules coexist — the floor is paid once, not N times. + estimate.ram_gb = max(estimate.ram_gb, 0.0) + _PROCESS_BASELINE_RAM_GB + # Embedder-load penalty for classic presets: when at least one classic + # scorer (linear / catboost / sklearn / knn / mlknn) sits on top of an + # embedder, the process holds the embedder weights + tokenizer + HF + # buffers *in addition to* whatever the per-driver ``max`` reported. + # classic-heavy on banking77 predicted 2.86 GB RAM against a measured + # 10.22 GB (3.6x under) because the max-of-drivers hides the fact that + # multiple scorers coexist in RAM. Only added when the search space + # actually consumes an embedder. + if _uses_embedder(search_space) and embedder_meta is not None: + estimate.ram_gb += _embedder_load_ram_gb(embedder_meta) + # Same story on the CUDA side: torch's caching allocator, cuBLAS/cuDNN + # workspaces, and driver context together reserve ~0.5 GB the moment the + # first tensor lands on the device — regardless of model size. Only apply + # when we actually predict some GPU usage AND running on CUDA hardware, + # so CPU-only presets stay honest. + if estimate.vram_gb > 0 and hardware.accelerator == "cuda": + estimate.vram_gb += _CUDA_BASELINE_VRAM_GB + _aggregate_disk( estimate, seen_models, diff --git a/src/autointent/_advisor/_estimates/_search_space.py b/src/autointent/_advisor/_estimates/_search_space.py index 577f500e6..a2b708650 100644 --- a/src/autointent/_advisor/_estimates/_search_space.py +++ b/src/autointent/_advisor/_estimates/_search_space.py @@ -55,6 +55,50 @@ def _max_int(value: Any, default: int) -> int: # noqa: ANN401 return default +def _module_cardinality(entry: dict[str, Any]) -> int | None: + """Approximate number of unique configurations the module entry can produce. + + Returns: + * ``1`` when every tunable field is a singleton (list of length 1, or a + plain scalar). Signal that HPO would rediscover the same config every + trial — real optuna dedupes, so effective n_trials = 1. + * ``N`` when the cardinality is finite and bounded (product of list + lengths across categorical fields, capped at :data:`_CARDINALITY_CAP` + to avoid overflow on large multi-list grids). + * ``None`` when any field declares a continuous range (``{low, high}`` + dict) — treated as "unbounded" so the caller falls back to the full + node-level n_trials. + + Ignores ``module_name`` (fixed, not a search dim), reserved keys like + ``target_metric``, and any non-tunable scalar keys that are already + single values. + """ + _RESERVED = {"module_name", "target_metric"} + _CARDINALITY_CAP = 10_000 + product = 1 + for key, value in entry.items(): + if key in _RESERVED: + continue + if isinstance(value, dict): + # Optuna-style range descriptor with low/high => continuous; treat + # as unbounded (many possible samples). + if "low" in value and "high" in value: + return None + # Non-range dict (e.g. nested config) counts as 1 — the dict + # itself is fixed unless it wraps a list. + continue + if isinstance(value, list): + if not value: + continue + # A list of dicts (like classification_model_config: [{...}, {...}]) + # still counts as N candidates. Length 1 = singleton. + product *= max(1, len(value)) + if product >= _CARDINALITY_CAP: + return _CARDINALITY_CAP + # Plain scalar (str/int/float/bool/None) is a singleton — contributes 1. + return product + + def _walk_modules_indexed( search_space: list[dict[str, Any]], ) -> Iterable[tuple[int, str, dict[str, Any]]]: diff --git a/src/autointent/_advisor/runner.py b/src/autointent/_advisor/runner.py index 4743a3b47..f83a649ba 100644 --- a/src/autointent/_advisor/runner.py +++ b/src/autointent/_advisor/runner.py @@ -12,7 +12,7 @@ from pydantic import ValidationError from autointent._advisor._estimates._resource import _resource_phase -from autointent._advisor._estimates._search_space import _max_int, _walk_modules +from autointent._advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules from autointent._advisor._report import PreflightReport, Severity from autointent._optimization_config import OptimizationConfig @@ -91,7 +91,7 @@ def run_preflight( cache_probe=embedding_cache_probe, ) _data_phase(cfg.search_space, stats, report) - _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, hardware, report) + _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, cfg.hpo_config.n_trials, hardware, report) return report @@ -114,10 +114,11 @@ def _validated_config(config: dict[str, Any]) -> OptimizationConfig: def _config_phase( search_space: list[dict[str, Any]], n_jobs: int, + n_trials: int, hardware: HardwareProfile, report: PreflightReport, ) -> None: - """Config-phase checks: parallelism vs. hardware mismatches.""" + """Config-phase checks: parallelism vs. hardware mismatches + no-op HPO.""" if n_jobs > 1 and hardware.accelerator in {"cuda", "mps"}: report.add( "config", @@ -136,6 +137,32 @@ def _config_phase( "CatBoost task_type=GPU configured but no CUDA detected - will fall back to CPU.", ) + # No-op HPO detector: n_trials >> search-space cardinality means most + # trials will be exact duplicates of previous ones. Optuna's TPE sampler + # doesn't auto-dedupe, so real runtime = n_trials × per-trial cost — the + # advisor charges honestly for that. But the user probably didn't intend + # this, so surface it as a finding: transformers-no-hpo on banking77 + # declares n_trials=40 with a single-value grid → 40x the useful work. + for _, entry in _walk_modules(search_space): + module = entry.get("module_name", "?") + # Skip decision-node entries; they're cheap and often intentionally + # singleton-configured. + if module in {"argmax", "threshold", "jinoos", "tunable", "adaptive"}: + continue + cardinality = _module_cardinality(entry) + if cardinality is not None and cardinality < n_trials and n_trials // max(1, cardinality) >= 4: + report.add( + "config", + Severity.TIGHT, + f"'{module}' entry has {cardinality} unique configurations but " + f"hpo_config.n_trials={n_trials} — expect ~{n_trials - cardinality} " + f"duplicate trials unless the sampler dedupes. Reduce n_trials or " + f"widen the search space.", + ) + # Emit at most one warning per preset — otherwise multi-module + # presets with several singleton entries flood the report. + break + def _data_phase( search_space: list[dict[str, Any]], From 85848f27fd0cc8e473815fdabb9a7231c8c86f4c Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:50:27 +0300 Subject: [PATCH 29/43] improve calibrate --- scripts/calibrate_advisor.py | 47 ++- .../_advisor/_estimates/_formulas.py | 274 +++++------------- .../_advisor/_estimates/_resource.py | 124 +++----- .../_advisor/_estimates/_search_space.py | 36 +-- src/autointent/_advisor/runner.py | 14 +- tests/advisor/test_estimates_internals.py | 268 +++++++++++++++++ tests/pipeline/test_calibration_tracker.py | 36 ++- 7 files changed, 430 insertions(+), 369 deletions(-) diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py index d4dec0f29..8e7d54cf5 100644 --- a/scripts/calibrate_advisor.py +++ b/scripts/calibrate_advisor.py @@ -275,19 +275,23 @@ def _dir_size_gb(path: Path) -> float: class _PeakSampler: - """Background thread tracking peak RSS and (on MPS) peak GPU allocation. - - CUDA has an accurate native peak-memory API and doesn't need polling; we - still read it after the fit. MPS lacks a peak API, so the sampler polls - ``torch.mps.current_allocated_memory()`` alongside RSS and keeps the max. - """ - - def __init__(self, interval_s: float = 0.1, *, sample_mps: bool = False) -> None: + """Background thread polling peak RSS + (optionally) MPS / CUDA current + allocation. CUDA polling catches allocations that fall outside any + module bracket — the per-module tracker's peak counter gets reset at each + start_module, losing anything allocated before it (e.g. the embedder + forward during pipeline setup). Best-effort: sub-poll-interval spikes + can be missed.""" + + def __init__( + self, interval_s: float = 0.1, *, sample_mps: bool = False, sample_cuda: bool = False, + ) -> None: self._interval_s = interval_s self._proc = psutil.Process() self.peak_ram_gb = self._proc.memory_info().rss / _BYTES_PER_GB self.peak_mps_gb: float | None = 0.0 if sample_mps else None + self.peak_cuda_gb: float | None = 0.0 if sample_cuda else None self._sample_mps = sample_mps + self._sample_cuda = sample_cuda self._stop = threading.Event() self._thread: threading.Thread | None = None @@ -314,6 +318,10 @@ def _run(self) -> None: mps = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB if self.peak_mps_gb is None or mps > self.peak_mps_gb: self.peak_mps_gb = mps + if self._sample_cuda and torch is not None and torch.cuda.is_available(): + cuda = float(torch.cuda.memory_allocated()) / _BYTES_PER_GB + if self.peak_cuda_gb is None or cuda > self.peak_cuda_gb: + self.peak_cuda_gb = cuda except (psutil.NoSuchProcess, psutil.AccessDenied): break self._stop.wait(self._interval_s) @@ -866,10 +874,13 @@ def _calibrate_one( _attach_callbacks(pipeline, callbacks) is_mps = hardware.accelerator == "mps" + is_cuda = hardware.accelerator == "cuda" undo_step_patch = _patch_trainer_for_step_timing(tracker) start = time.perf_counter() try: - with _PeakSampler(interval_s=poll_interval_ms / 1000.0, sample_mps=is_mps) as sampler: + with _PeakSampler( + interval_s=poll_interval_ms / 1000.0, sample_mps=is_mps, sample_cuda=is_cuda, + ) as sampler: pipeline.fit(dataset, preflight="off") except Exception as e: # noqa: BLE001 row.error = f"fit failed: {e}" @@ -883,12 +894,15 @@ def _calibrate_one( embed_after = _dir_size_gb(embed_cache) actual_time_h = elapsed_s / 3600.0 actual_ram_gb = sampler.peak_ram_gb - # Prefer the tracker's per-module max: the fit-level torch.cuda peak is - # clobbered by the per-module reset_peak_memory_stats calls, so the final - # reading only reflects VRAM used since the last (usually CPU-only) module. - actual_vram_gb: float | None - if tracker.peak_vram_gb_overall > 0: - actual_vram_gb = tracker.peak_vram_gb_overall + # VRAM: take max of per-module tracker (inside brackets) and background + # sampler (outside brackets, e.g. classic-preset embedder forward). + # Fallback to a raw peak read only if both are zero. + actual_vram_gb: float | None = None + tracker_peak = tracker.peak_vram_gb_overall if tracker.peak_vram_gb_overall > 0 else None + sampler_peak = sampler.peak_cuda_gb if sampler.peak_cuda_gb and sampler.peak_cuda_gb > 0 else None + candidates = [x for x in (tracker_peak, sampler_peak) if x is not None] + if candidates: + actual_vram_gb = max(candidates) else: actual_vram_gb = _read_vram_peak_gb(hardware.accelerator) if actual_vram_gb is None and is_mps: @@ -900,6 +914,9 @@ def _calibrate_one( "vram_gb": actual_vram_gb, "disk_download_gb": max(0.0, hf_after - hf_before), "disk_embedding_cache_gb": max(0.0, embed_after - embed_before), + # Per-signal breakdown; classic presets expect sampler > tracker. + "vram_gb_tracker": tracker_peak, + "vram_gb_sampler": sampler_peak, } row.modules = tracker.records if enable_wandb and not any("W&B requested but not available" in n for n in row.notes): diff --git a/src/autointent/_advisor/_estimates/_formulas.py b/src/autointent/_advisor/_estimates/_formulas.py index 795390f17..05b68d076 100644 --- a/src/autointent/_advisor/_estimates/_formulas.py +++ b/src/autointent/_advisor/_estimates/_formulas.py @@ -56,19 +56,9 @@ def _classify_severity(estimate: float, budget: float) -> Severity: def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: - """Weight-side VRAM in GB — weights + grads + optimizer state. Excludes activations. - - Returns a deliberately pessimistic upper bound, matching the advisor's - "heuristic upper bound, not measurement" contract. - - Modes: - * ``inference``: forward only — weights + ~30% intermediate-tensor overhead. - * ``lora``: frozen base + small trainable adapters + their grads/optimizer (~0.5 GB). - * ``full-finetune`` (default): the textbook 4W (weights + grads + Adam m + Adam v). - We use 4.5W to leave headroom for loss-scale buffers, allocator fragmentation, - cuDNN workspaces, and gradient-accumulation buffers — none of which the textbook - 4W accounting captures. - """ + """Weight-side VRAM: weights + grads + optimizer state. Pessimistic upper + bound by mode: 1.3× inference, 1.3× + 0.5 GB lora adapters, 4.5× full + finetune (textbook 4W + fragmentation/workspaces slack).""" weights_gb = meta.weights_gb if mode == "inference": return weights_gb * 1.3 @@ -83,24 +73,9 @@ def _activations_gb_per_sample( *, is_training: bool, ) -> float: - """Heuristic activation memory per sample. - - Training uses **34 bytes/token/layer** as a pessimistic upper bound — - Korthikanti et al. (2022, "Reducing Activation Recomputation ...") derive - this for standard attention: the linear-layer activations account for ~11B - and the attention matrix + intermediate tensors add ~23B. FlashAttention - kernels drop the attention-matrix term (~12 B/token/layer total), but we - can't detect at preflight time whether the user's stack will use them, so - the upper bound is the safe choice. - - Inference: only 1-2 layers' outputs are kept in flight at once. 8 B/token - covers fp32 hidden (4B) plus a bit of intermediate slack. - - An earlier revision used 16 B/token/layer for training; that under-predicted - real deberta-v3-large VRAM by ~2x at bs=128 (measured 13.1 GB, predicted - ~11.5 GB), which is unsafe for an OOM-avoidance tool. See ``interpretation.md`` - (2026-07-19) for the calibration data. - """ + """Activation memory per sample. Training: 34 B/token/layer (Korthikanti + 2022 upper bound, standard attention). Inference: 8 B/token (only 1-2 + layers' outputs in flight).""" hidden = _embedder_dim(meta) training_bytes_per_token_per_layer = 34 inference_bytes_per_token = 8 @@ -118,25 +93,17 @@ def _vram_for_transformer( batch_size: int = 0, seq_len: int = _DEFAULT_SEQ_LEN, ) -> float: - """Total VRAM in GB: weights + grads + optimizer state + activations x batch. - - Activation accounting differs by mode — training keeps per-layer outputs for - backward; inference only needs one or two layers in flight. - - Final ``* 1.20`` is a safety margin covering allocator fragmentation, - peak transient tensors during backward, and HF-Trainer's eval-loop - double-forward that the textbook accounting above misses. Advisor should - upper-bound: earlier 15% margin left transformers-light on banking77 at - 8.12 GB predicted vs 8.75 GB measured (1.08x UNDER — unsafe for an OOM - tool). Bumped to 20% + a bigger fixed CUDA baseline (see - ``_CUDA_BASELINE_VRAM_GB``) to close the gap on large-batch training runs - without over-inflating small ones. - """ + """Total VRAM: weights + grads + optimizer state + activations × batch. + + Safety margin: 1.20 for training (backward transients, eval sweep, + allocator fragmentation), 1.10 for inference (no backward).""" base = _weights_vram_for_transformer(meta, mode) if batch_size <= 0: return base - per_sample = _activations_gb_per_sample(meta, seq_len, is_training=mode != "inference") - return (base + per_sample * batch_size) * 1.20 + is_training = mode != "inference" + per_sample = _activations_gb_per_sample(meta, seq_len, is_training=is_training) + safety = 1.20 if is_training else 1.10 + return (base + per_sample * batch_size) * safety def _max_fitting_batch_size( @@ -159,19 +126,10 @@ def _max_fitting_batch_size( return _floor_to_power_of_two(int(available_for_activations / per_sample_gb)) -# Sustained TFLOPS per device class — real MFU (model-FLOPs utilization) at -# training batch sizes, NOT peak spec sheet numbers. Numbers reflect ~20-30% -# MFU which is what HF Trainer actually achieves on BERT-scale workloads once -# you factor in dataloader idle, tokenization warmup, per-epoch eval, and -# checkpoint saves — all of which the raw FLOPs formula ignores. Earlier -# values (150 / 45 / 15 for high/mid/low) were closer to peak spec numbers -# and under-predicted transformers-heavy on banking77 by 1.7x (measured -# 3.17 h vs predicted 1.88 h, calibration_runs2 2026-08-07); an advisor -# should upper-bound, so pick sustained numbers that err on the side of -# over-predicting. Source: MLPerf training results + measured banking77 -# runs where mean_step_s / p95_step_s → 154ms / 246ms for bert-base bs=64. +# Sustained TFLOPS per device class — real HF-Trainer MFU (~20% on A100), +# not peak spec sheet. Advisor upper-bounds, so pessimistic values here. _DEVICE_TFLOPS = { - "high-gpu": 60.0, # A100 / H100 — sustained ~19% MFU under HF Trainer + "high-gpu": 60.0, # A100 / H100 "mid-gpu": 20.0, # V100 / RTX 3090 / A6000 "low-gpu": 7.0, # T4 / RTX 3060 / 8 GB consumer card "apple-silicon": 4.0, # M1/M2/M3 GPU cores @@ -179,11 +137,8 @@ def _max_fitting_batch_size( } _DEFAULT_TFLOPS = 7.0 # unknown device → treat as low-GPU -# HF Trainer overhead — the FLOPs formula only counts optimizer steps; real -# wall-time also includes per-epoch eval sweeps, save-checkpoint syncs, -# tokenizer warmup, dataloader queue idle, and gradient-accumulation gaps. -# Factor calibrated so transformers-heavy predicts ~1.2-1.5x the measured -# 3.17 h (upper-bound stance). +# HF Trainer overhead: eval sweeps, save syncs, dataloader idle. The raw +# FLOPs formula only counts optimizer steps. _TRAINER_OVERHEAD_MULT = 1.35 @@ -197,18 +152,9 @@ def _time_for_transformer( params_millions: float, device_class: str, ) -> float: - """Transformer training time in hours, from per-step FLOPs / device TFLOPS. - - Per-step FLOPs ≈ 6 x params x batch_size x seq_len (2 for forward mul-add, - 3-4x for backward). Divided by *sustained* device TFLOPS (not peak spec) - to get wall-time per step, then multiplied by (steps x epochs x n_trials) - x ``_TRAINER_OVERHEAD_MULT`` for HF-Trainer wall-clock overhead. - - Advisor contract: err on the side of over-prediction. Under-predicting - time makes users blow through wall-clock budgets; over-predicting only - biases them toward smaller / cheaper presets. See ``_DEVICE_TFLOPS`` - docstring for the sustained-MFU calibration. - """ + """Transformer training wall-time in hours. Per-step FLOPs = 6 × params × + batch × seq_len (fwd+bwd), ÷ sustained device TFLOPS, × total steps × + trainer overhead.""" steps_per_epoch = max(1, n_samples // max(1, batch_size)) total_steps = n_trials * epochs * steps_per_epoch # 6x factor: ~2x for fwd matmul + ~4x for bwd (grad wrt input + grad wrt weight). @@ -240,20 +186,9 @@ def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: def _ram_for_module(meta: ModelMeta, stats: DatasetStats, *, mode: str = "inference") -> float: - """RAM in GB. Loose upper bound: weights + optimizer/grads + tokenized text. - - Tokenized text is approximated as ``n_samples x avg_tokens x 4 bytes`` - (BPE/WordPiece token ids fit in int32). - - ``mode``-dependent multiplier on the weights term: - * ``inference``: 1.3x (weights + intermediate-tensor slack) - * ``lora``: 1.5x (frozen base + trainable adapters) - * ``full-finetune`` / anything else: 4.5x (weights + grads + Adam m + v - + framework slack) — matches the VRAM-side ``4.5W`` accounting so the - host-pinned optimizer state (Adam mirrors weights) shows up in the - RAM estimate too. Under-predicting RAM lets a training preset OOM the - host well before it OOMs the GPU; advisor should upper-bound. - """ + """RAM upper bound: weights × mode_mult + tokenized text (n_samples × + avg_tokens × 4 B). Mode multiplier: 1.3 inference, 1.5 lora, 4.5 + full-finetune (Adam mirrors weights on host too).""" if mode == "inference": weights_mult = 1.3 elif mode == "lora": @@ -269,26 +204,13 @@ def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: return (n_samples * hidden_size * 4) / _BYTES_PER_GB -# Wall-time coefficients calibrated against measured fits on 1-thread CPU -# (OMP_NUM_THREADS=1). Values represent seconds per per-fit-work-unit and -# already absorb the number of L-BFGS iterations the optimizer typically -# takes to converge (~50) — so the ``max_iter`` upper bound does NOT enter -# the formula directly. Historical formula also scaled by ``max_iter`` which -# multi-cent-ordered-over-predicted (137 h vs measured ~30 s = ~15000x on -# banking77 × 1024-dim e5-large × 77 classes × cv=3). -# -# Calibration point (reviewer's res-adapt-ckeck/a100 run, warm cache): -# classic-light linear on banking77 (n=10003, dim=1024, cls=77, cv_mult=31) -# measured ~30 s per fit x n_trials=20 = ~10 min total = ~0.17 h. -# Formula: 20 x 1.2e-9 x 10003 x 1024 x 31 x 77 = ~588 s = ~0.16 h ✓ +# Wall-time coefficients calibrated on 1-thread CPU (OMP_NUM_THREADS=1), +# seconds per fit-work-unit. Typical L-BFGS iteration count baked in. _LINEAR_CPU_S_PER_SAMPLE_FEATURE = 1.2e-9 -_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 # catboost is measured per iteration +_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 _CATBOOST_GPU_SPEEDUP = 10.0 -# LogisticRegressionCV defaults: Cs=10, cv=3 -> 10x3 inner fits + 1 final refit = 31. -_LOGREG_CV_MULTIPLIER = 31 -# Default value of `border_count` in CatBoost (number of histogram buckets per feature). -_CATBOOST_DEFAULT_BINS = 254 -# Bytes per histogram bucket / tree node — order-of-magnitude constant. +_LOGREG_CV_MULTIPLIER = 31 # sklearn default: Cs=10 × cv=3 + 1 final refit +_CATBOOST_DEFAULT_BINS = 254 # CatBoost `border_count` default _CATBOOST_BYTES_PER_TREE_NODE = 32 @@ -305,21 +227,12 @@ def _time_for_linear( n_trials: int, n_samples: int, embedder_dim: int, - max_iter: int, # noqa: ARG001 — kept in signature for API stability; typical L-BFGS convergence is absorbed into the coefficient + max_iter: int, # noqa: ARG001 — API stability; typical L-BFGS convergence baked into coeff cv_multiplier: int, class_multiplier: int, ) -> float: - """LogisticRegression wall time, in hours. - - Cost is ``O(n_samples x n_features x n_classes)`` per fit (sklearn's L-BFGS - solver, iterations absorbed into the calibration constant), multiplied by the - CV inner-fit count (31 for the default LogisticRegressionCV). - - ``max_iter`` is a per-fit upper bound, not the typical work — L-BFGS on a - well-conditioned classifier converges long before it. Older versions of - this formula scaled by ``max_iter`` and predicted ~1000x higher than - reality; the constant now bakes in a typical convergence-iteration count. - """ + """LogisticRegression wall time. O(n_samples × features × classes × cv) + per fit; typical L-BFGS convergence absorbed into the calibration constant.""" seconds = ( n_trials * _LINEAR_CPU_S_PER_SAMPLE_FEATURE @@ -347,55 +260,22 @@ def _ram_for_sklearn( max_depth: int, n_jobs: int, ) -> float: - """RandomForest/similar RAM upper bound, aware of ``n_jobs`` replication. - - sklearn spawns ``n_jobs`` worker processes with joblib's loky backend by - default; each worker holds its own copy of the training feature matrix and - the trees it grew, so a preset with ``n_jobs=8`` on a 10 k × 1024 embedder - dataset multiplies the base RAM 8x. Previously sklearn was in - ``_UNKNOWN_SCORER_MODULES`` and emitted a zero row — classic-heavy's - real ~2 GB sklearn contribution slipped through invisibly. - """ - # Per-worker feature matrix (fp64 in sklearn by default). - per_worker_data = stats.n_samples * embedder_dim * 8 - # Per-worker tree storage: n_estimators × n_leaves × ~32 B/node. Cap - # n_leaves at n_samples (a tree with max_depth 150 on 10k samples can't - # actually have 2**150 leaves). + """RandomForest RAM: (feature matrix + trees) × n_jobs. joblib workers + each hold a full copy.""" + per_worker_data = stats.n_samples * embedder_dim * 8 # fp64 default n_leaves = min(2**max_depth, stats.n_samples) if max_depth > 0 else stats.n_samples per_worker_trees = n_estimators * n_leaves * _CATBOOST_BYTES_PER_TREE_NODE - per_worker = per_worker_data + per_worker_trees - return float((per_worker * max(1, n_jobs)) / _BYTES_PER_GB) + return float(((per_worker_data + per_worker_trees) * max(1, n_jobs)) / _BYTES_PER_GB) def _embedder_load_ram_gb(meta: ModelMeta | None) -> float: - """Extra RAM the process holds when an embedder is loaded — separately from - any per-driver row that already accounts for it. - - Rationale: classic presets pre-compute embeddings via the embedder, then - train sklearn/catboost/linear scorers on top. During and after that - forward pass the process holds: the embedder weights on the compute - device, a copy in CPU RAM (fp32 from the safetensors load), the tokenizer - state, HF Trainer buffers, and the cached embeddings themselves. The - per-driver ``_ram_for_module`` already captures weights x 1.3 for the - knn/mlknn rows, but the aggregate ``max`` across drivers hides - contributions from the other classic scorers that are simultaneously in - memory. This term is added *on top* of the max-driver RAM so classic - presets like classic-heavy stop under-predicting by ~4x. - - Uses ``total_params × 4`` (fp32) as the weight footprint even when the - hub reports fp16 storage (``weight_bytes_per_param=2``) — transformers - up-casts to fp32 at load time by default, so the fp16 disk size - under-counts real RAM usage by 2x. - """ + """Aggregate-level RAM penalty when a classic preset uses an embedder. + Added on top of the max-driver RAM because embedder + multiple classic + scorers coexist in memory. Uses fp32 weights (transformers up-casts at + load) × 3.5 for weights + activation buffers + framework slack.""" if meta is None: return 0.0 - fp32_weights_gb = (meta.total_params * 4) / _BYTES_PER_GB - # 3.5x factor: raw weights + activation buffers + HF/tokenizer/loader - # slack. Empirical: real classic-heavy on banking77 (e5-large embedder, - # sklearn RF n_jobs=8) measured 10.22 GB RAM. At 3.0x we predicted 9.56 - # (1.07x under — still unsafe); at 3.5x we predict 10.85 (0.94x — safely - # over). Advisor's OOM-avoidance contract requires "err over". - return fp32_weights_gb * 3.5 + return ((meta.total_params * 4) / _BYTES_PER_GB) * 3.5 def _time_for_catboost( @@ -423,63 +303,44 @@ def _time_for_catboost( # === CNN / RNN scorers =================================================== -# -# Small torch models (Kim's TextCNN, LSTM classifier) trained from scratch -# on token ids. The advisor previously emitted a ``not-estimated`` placeholder -# for these — which read as "free/safe" on nn-heavy / nn-medium (predicted 0h -# / 0GB, real 0.3 h + 2.3 GB RAM + 0.7 GB VRAM on banking77). -# -# Cost model: embedding table + head weights (small) + activations that scale -# with batch × seq_len × hidden. We assume a bounded vocabulary (~30k) — one -# HPO trial for TextCNN embeds every token that appears in the training set; -# banking77 tops out around a few thousand unique tokens so 30k is a safe -# upper bound. -_NN_MAX_VOCAB = 30_000 +# Small torch models (TextCNN, LSTM) trained from scratch on token ids. +_NN_MAX_VOCAB = 30_000 # upper bound on vocabulary size _NN_DEFAULT_SEQ_LEN = 50 # VocabConfig.max_seq_length default -_NN_BYTES_PER_PARAM = 4 # fp32 weights -# fp32 activation storage per (batch, token, hidden) unit, factor absorbs -# ~4x backward overhead + optimizer + gradient state for these tiny models. -_NN_TRAIN_ACT_BYTES_PER_UNIT = 16 +_NN_BYTES_PER_PARAM = 4 # fp32 +_NN_TRAIN_ACT_BYTES_PER_UNIT = 16 # ~4x backward + optimizer overhead def _cnn_param_count(*, embed_dim: int, num_filters: int, n_kernels: int, n_classes: int) -> int: - """Approximate CNN parameter count: embedding + conv + fc layers.""" - vocab_params = _NN_MAX_VOCAB * embed_dim - conv_params = num_filters * embed_dim * n_kernels * 5 # avg kernel width ~5 - fc_params = num_filters * n_kernels * max(1, n_classes) - return vocab_params + conv_params + fc_params + """TextCNN params: embedding + conv (kernel width ~5) + fc.""" + return ( + _NN_MAX_VOCAB * embed_dim + + num_filters * embed_dim * n_kernels * 5 + + num_filters * n_kernels * max(1, n_classes) + ) def _rnn_param_count(*, embed_dim: int, hidden_dim: int, n_classes: int) -> int: - """Approximate LSTM classifier parameter count: embedding + LSTM + fc.""" - vocab_params = _NN_MAX_VOCAB * embed_dim - # LSTM cell has 4 gates, each with (embed+hidden+1) × hidden params. - lstm_params = 4 * hidden_dim * (embed_dim + hidden_dim + 1) - fc_params = hidden_dim * max(1, n_classes) - return vocab_params + lstm_params + fc_params + """LSTM classifier params: embedding + 4-gate LSTM cell + fc.""" + return ( + _NN_MAX_VOCAB * embed_dim + + 4 * hidden_dim * (embed_dim + hidden_dim + 1) + + hidden_dim * max(1, n_classes) + ) def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: - """Weights + activations for a small torch scorer in training mode. - - Activation term uses ``batch × seq_len × hidden × const`` per the same - fp32 upper bound as transformers, but with a much smaller effective - hidden dim (embed_dim/num_filters, not model dim). - """ + """Weights + 3× optimizer/grads + activations. Same fp32 upper bound as + transformers, smaller hidden dim (embed_dim / num_filters).""" weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB - # Optimizer state (Adam has 2x weights) + gradients (1x weights) = 4x weights total. - optimizer_gb = 3 * weights_gb activations_gb = ( batch_size * _NN_DEFAULT_SEQ_LEN * hidden_dim * _NN_TRAIN_ACT_BYTES_PER_UNIT ) / _BYTES_PER_GB - return weights_gb + optimizer_gb + activations_gb + return 4 * weights_gb + activations_gb def _ram_for_nn(*, params: int, stats: DatasetStats) -> float: - """CPU-side memory: weights + tokenized text (int32 ids).""" - weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB - tokens_gb = (stats.n_samples * _NN_DEFAULT_SEQ_LEN * 4) / _BYTES_PER_GB - return weights_gb + tokens_gb + """Weights + tokenized text (int32 ids).""" + return ((params * _NN_BYTES_PER_PARAM) + (stats.n_samples * _NN_DEFAULT_SEQ_LEN * 4)) / _BYTES_PER_GB def _time_for_nn( @@ -491,13 +352,8 @@ def _time_for_nn( params_millions: float, device_class: str, ) -> float: - """Reuse the transformer FLOPs formula for a small torch model. - - Small models are memory-bandwidth-bound, not compute-bound, so the FLOPs - formula slightly *under*-predicts wall-time. Empirically for banking77 - nn-heavy we measured 0.32 h across 55 trials — the formula lands within - 2x of that, which is enough for a cost-ranking estimate. - """ + """Reuse transformer FLOPs formula; small models slightly under-predict + since they're memory-bandwidth-bound, but within 2x for cost ranking.""" return _time_for_transformer( n_trials=n_trials, epochs=epochs, diff --git a/src/autointent/_advisor/_estimates/_resource.py b/src/autointent/_advisor/_estimates/_resource.py index acf9fbb3d..05895f1d4 100644 --- a/src/autointent/_advisor/_estimates/_resource.py +++ b/src/autointent/_advisor/_estimates/_resource.py @@ -258,10 +258,7 @@ def _estimate_classic_entry( vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) mode = "catboost-gpu" if on_gpu else "catboost" elif module == "sklearn": - # RandomForestClassifier is the most common target; joblib spawns - # ``n_jobs`` worker processes each replicating the feature matrix + - # trees. Predicting 0 here (previous "not-estimated" behaviour) hid - # classic-heavy's real 1-2 GB sklearn contribution. + # RandomForest is the common target; joblib replicates data across n_jobs. n_estimators = _max_int(entry.get("n_estimators"), 100) max_depth = _max_int(entry.get("max_depth"), 0) sk_n_jobs = _max_int(entry.get("n_jobs"), 1) @@ -272,9 +269,8 @@ def _estimate_classic_entry( max_depth=max_depth, n_jobs=sk_n_jobs, ) - # Time: rough O(n_estimators × n_samples × sqrt(features) × log2 n) - # per fit, divided by n_jobs. Absorbed into the linear coefficient - # since real numbers vary wildly by criterion / max_features. + # Rough O(n_estimators × n × features / n_jobs); real numbers vary + # wildly by criterion so use the linear coefficient as a proxy. time_h = ( n_trials * _LINEAR_CPU_S_PER_SAMPLE_FEATURE @@ -317,23 +313,16 @@ def _estimate_nn_entry( n_trials: int, refit_after: bool, ) -> _ModuleEstimate | None: - """Cost row for a cnn / rnn scorer (returns ``None`` for anything else). - - These are small torch models trained from scratch on token ids. Previously - the advisor emitted a ``not-estimated`` placeholder for them, which read - as "free/safe" — nn-heavy on banking77 predicted 0h/0GB but actually used - 0.32 h + 2.3 GB RAM + 0.7 GB VRAM. This restores a real estimate using - small-model parameter counts + the transformer FLOPs formula for time. - """ + """Cost row for cnn / rnn scorers (returns None for anything else). + Small torch models trained from scratch on token ids.""" module = entry.get("module_name", "?") n_classes = max(1, stats.n_classes) if module == "cnn": embed_dim = _max_int(entry.get("embed_dim"), 128) num_filters = _max_int(entry.get("num_filters"), 100) + # kernel_sizes is a list of lists (e.g. [[3,4,5]]) — count the largest variant. kernel_sizes = entry.get("kernel_sizes") - # Kernel sizes are a list of lists in the search space - # (e.g. [[3, 4, 5]]); count entries in the largest variant. n_kernels = 3 if isinstance(kernel_sizes, list): for candidate in kernel_sizes: @@ -545,33 +534,20 @@ def _emit_resource_findings( ) -# Process-level memory floors. Every autointent fit imports torch + -# transformers + datasets + optuna, which reserve resident memory the moment -# they load. Measured against calibration_runs2 (2026-08-07 banking77 sweep): -# every preset used 2-10 GB RAM but the per-module estimates alone predicted -# 0.4-2 GB → advisor was systematically 4-13x LOW on RAM. A ~1.5 GB baseline -# lifts the estimate into range without over-inflating heavy presets. +# Process-level memory floors. Every fit reserves ~1.5 GB RSS for torch + +# transformers + datasets + optuna before any weights load. _PROCESS_BASELINE_RAM_GB = 1.5 -# CUDA driver context + cuDNN/cuBLAS workspace pools + caching allocator -# fragmentation. A real training process on A100 reserves ~1 GB the moment -# torch initializes CUDA + the first tensor lands, regardless of model size. -# Earlier 0.5 GB baseline left transformers-light on banking77 at 8.12 GB -# predicted vs 8.75 GB measured (unsafe under-prediction for OOM avoidance); -# 1.0 GB closes the gap with room to spare. Only added when we already -# predict some VRAM usage so CPU-only presets aren't spuriously flagged as -# GPU users. -_CUDA_BASELINE_VRAM_GB = 1.0 +# CUDA driver context + cuDNN/cuBLAS workspace + caching allocator. Bigger +# for training (backward workspaces + optimizer scratch), smaller for +# inference-only. +_CUDA_BASELINE_VRAM_TRAINING_GB = 1.0 +_CUDA_BASELINE_VRAM_INFERENCE_GB = 0.3 _UNKNOWN_SCORER_MODULES: frozenset[str] = frozenset() -"""Scorer modules the advisor has no cost estimator for — kept as an empty -extension point. cnn / rnn moved to :func:`_estimate_nn_entry`; sklearn moved -to the classic branch of :func:`_estimate_classic_entry`. New unknown-cost -scorers should still register here so they emit a not-estimated placeholder -row instead of silently reporting zero.""" +"""Extension point for scorers with no cost estimator — they get a +not-estimated placeholder row instead of a silent zero.""" _NN_SCORER_MODULES = frozenset({"cnn", "rnn"}) -"""Small torch scorers routed through :func:`_estimate_nn_entry`.""" - _EMBEDDER_CONSUMING_MODULES = frozenset( {"linear", "catboost", "sklearn", "knn", "mlknn", "retrieval", @@ -580,13 +556,11 @@ def _emit_resource_findings( def _uses_embedder(search_space: list[dict[str, Any]]) -> bool: - """True when any module in the search space consumes an embedder — signals - that the aggregate RAM should include the embedder-load penalty on top - of the per-driver max.""" - for _, entry in _walk_modules(search_space): - if entry.get("module_name") in _EMBEDDER_CONSUMING_MODULES: - return True - return False + """True when any search-space module consumes the embedder.""" + return any( + entry.get("module_name") in _EMBEDDER_CONSUMING_MODULES + for _, entry in _walk_modules(search_space) + ) # Modules that consume the top-level ``cross_encoder_config.model_name`` as # their scoring model (see zero-shot-encoders preset: description_cross pulls @@ -663,31 +637,16 @@ def _resource_phase( transformer_entries, classic_entries = _split_entries(search_space) - # Per-node module-variant count. HPO distributes ``n_trials`` across the - # module_name candidates at each node roughly evenly (TPE's sampler bias - # aside), so a module_name that shares its node with 4 others sees on - # average ``n_trials / 5`` trials — not ``n_trials``. Previously every - # per-module estimate used the full ``n_trials``, which inflated - # classic-heavy's catboost row to 120 h vs measured 12 h across the whole - # node. Divide once here and pass the effective share down to every - # per-module estimator (transformer + classic + nn). + # HPO distributes n_trials evenly across module_name candidates at each + # node, so each variant sees n_trials / n_variants on average. variants_per_node: dict[int, int] = {} for node_idx, _node_type, _entry in _walk_modules_indexed(search_space): variants_per_node[node_idx] = variants_per_node.get(node_idx, 0) + 1 def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int: # noqa: ARG001 - """Trials this specific module should be charged for. - - Divides ``n_trials`` evenly across the module_name candidates at the - node. ``entry`` is accepted for API stability and future extensions - (e.g. a per-module cardinality cap) — we tried capping by - :func:`_module_cardinality` earlier but it produced wrong estimates - for description-scorer presets where Optuna's TPE runs every - declared trial regardless of parameter-space size (no automatic - dedup). The advisor charges for the declared work; if a real run - crashes or dedupes, that's an artefact of the runtime, not something - the advisor should try to predict. - """ + """Trials charged to this module. ``entry`` reserved for future + per-module caps (see git history for the cardinality-cap experiment + that broke description-scorer presets).""" divisor = max(1, variants_per_node.get(node_idx, 1)) return max(1, n_trials // divisor) @@ -769,31 +728,22 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int estimate.time_hours += me.time_hours estimate.drivers.append(me.driver) - # Process-level baselines. Every fit — even a trivial one — imports - # torch / transformers / datasets, which on their own take ~1.5 GB of RSS - # before any model weights load. Real runs on banking77 measure 2-10 GB - # RAM across every preset while the per-module estimates alone predicted - # 0.4-2 GB (systematically 4-13x low, see calibration_runs2 2026-08-07). - # Adding a floor here (rather than per-module) means the estimate stays - # accurate when multiple modules coexist — the floor is paid once, not N times. + # Process baseline: paid once, not per-module. estimate.ram_gb = max(estimate.ram_gb, 0.0) + _PROCESS_BASELINE_RAM_GB - # Embedder-load penalty for classic presets: when at least one classic - # scorer (linear / catboost / sklearn / knn / mlknn) sits on top of an - # embedder, the process holds the embedder weights + tokenizer + HF - # buffers *in addition to* whatever the per-driver ``max`` reported. - # classic-heavy on banking77 predicted 2.86 GB RAM against a measured - # 10.22 GB (3.6x under) because the max-of-drivers hides the fact that - # multiple scorers coexist in RAM. Only added when the search space - # actually consumes an embedder. + # Embedder-load penalty: classic presets keep embedder weights + framework + # buffers alongside scorer RAM, which max-of-drivers hides. if _uses_embedder(search_space) and embedder_meta is not None: estimate.ram_gb += _embedder_load_ram_gb(embedder_meta) - # Same story on the CUDA side: torch's caching allocator, cuBLAS/cuDNN - # workspaces, and driver context together reserve ~0.5 GB the moment the - # first tensor lands on the device — regardless of model size. Only apply - # when we actually predict some GPU usage AND running on CUDA hardware, - # so CPU-only presets stay honest. + # CUDA baseline only when we predict some VRAM AND run on CUDA. Mode read + # off drivers: any training row flips to the larger baseline. if estimate.vram_gb > 0 and hardware.accelerator == "cuda": - estimate.vram_gb += _CUDA_BASELINE_VRAM_GB + is_training = any( + d.get("mode") in {"full-finetune", "lora", "small-torch-train"} + for d in estimate.drivers + ) + estimate.vram_gb += ( + _CUDA_BASELINE_VRAM_TRAINING_GB if is_training else _CUDA_BASELINE_VRAM_INFERENCE_GB + ) _aggregate_disk( estimate, diff --git a/src/autointent/_advisor/_estimates/_search_space.py b/src/autointent/_advisor/_estimates/_search_space.py index a2b708650..3792d2f23 100644 --- a/src/autointent/_advisor/_estimates/_search_space.py +++ b/src/autointent/_advisor/_estimates/_search_space.py @@ -56,22 +56,11 @@ def _max_int(value: Any, default: int) -> int: # noqa: ANN401 def _module_cardinality(entry: dict[str, Any]) -> int | None: - """Approximate number of unique configurations the module entry can produce. - - Returns: - * ``1`` when every tunable field is a singleton (list of length 1, or a - plain scalar). Signal that HPO would rediscover the same config every - trial — real optuna dedupes, so effective n_trials = 1. - * ``N`` when the cardinality is finite and bounded (product of list - lengths across categorical fields, capped at :data:`_CARDINALITY_CAP` - to avoid overflow on large multi-list grids). - * ``None`` when any field declares a continuous range (``{low, high}`` - dict) — treated as "unbounded" so the caller falls back to the full - node-level n_trials. - - Ignores ``module_name`` (fixed, not a search dim), reserved keys like - ``target_metric``, and any non-tunable scalar keys that are already - single values. + """Unique configurations the module entry can produce. + + Returns 1 when every tunable field is a singleton, N for finite list + products (capped at 10_000), None when any field is a continuous + ``{low, high}`` range. """ _RESERVED = {"module_name", "target_metric"} _CARDINALITY_CAP = 10_000 @@ -80,22 +69,13 @@ def _module_cardinality(entry: dict[str, Any]) -> int | None: if key in _RESERVED: continue if isinstance(value, dict): - # Optuna-style range descriptor with low/high => continuous; treat - # as unbounded (many possible samples). if "low" in value and "high" in value: - return None - # Non-range dict (e.g. nested config) counts as 1 — the dict - # itself is fixed unless it wraps a list. - continue - if isinstance(value, list): - if not value: - continue - # A list of dicts (like classification_model_config: [{...}, {...}]) - # still counts as N candidates. Length 1 = singleton. + return None # continuous range + continue # non-range dict = fixed + if isinstance(value, list) and value: product *= max(1, len(value)) if product >= _CARDINALITY_CAP: return _CARDINALITY_CAP - # Plain scalar (str/int/float/bool/None) is a singleton — contributes 1. return product diff --git a/src/autointent/_advisor/runner.py b/src/autointent/_advisor/runner.py index f83a649ba..584f3b75c 100644 --- a/src/autointent/_advisor/runner.py +++ b/src/autointent/_advisor/runner.py @@ -137,18 +137,12 @@ def _config_phase( "CatBoost task_type=GPU configured but no CUDA detected - will fall back to CPU.", ) - # No-op HPO detector: n_trials >> search-space cardinality means most - # trials will be exact duplicates of previous ones. Optuna's TPE sampler - # doesn't auto-dedupe, so real runtime = n_trials × per-trial cost — the - # advisor charges honestly for that. But the user probably didn't intend - # this, so surface it as a finding: transformers-no-hpo on banking77 - # declares n_trials=40 with a single-value grid → 40x the useful work. + # No-op HPO warning: n_trials >> cardinality → mostly duplicate trials. + # At most one warning per preset to avoid flooding multi-module reports. for _, entry in _walk_modules(search_space): module = entry.get("module_name", "?") - # Skip decision-node entries; they're cheap and often intentionally - # singleton-configured. if module in {"argmax", "threshold", "jinoos", "tunable", "adaptive"}: - continue + continue # decision modules are cheap and often singleton by design cardinality = _module_cardinality(entry) if cardinality is not None and cardinality < n_trials and n_trials // max(1, cardinality) >= 4: report.add( @@ -159,8 +153,6 @@ def _config_phase( f"duplicate trials unless the sampler dedupes. Reduce n_trials or " f"widen the search space.", ) - # Emit at most one warning per preset — otherwise multi-module - # presets with several singleton entries flood the report. break diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 99bd63243..b90997cbb 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -769,3 +769,271 @@ def test_warm_cache_probe_zeroes_forward_and_disk(self) -> None: # Warm: forward wasn't charged → model isn't in ``cached_embedders`` → # no disk_embedding_cache contribution. assert warm.resource.disk_embedding_cache_gb == 0 + + +class TestCnnRnnHeuristic: + """cnn/rnn get a real small-model estimate, not a not-estimated zero row.""" + + def test_cnn_row_is_nonzero(self) -> None: + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "cnn", "embed_dim": [128], "num_filters": [128], + "kernel_sizes": [[3, 4, 5]], "batch_size": [64], "num_train_epochs": [60]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 10}, + } + report = run_preflight(cfg, DatasetStats.placeholder(n_samples=5000, n_classes=20), _profile()) + cnn_row = next(d for d in report.resource.drivers if d["module"] == "cnn") + # Real numbers (not the not-estimated placeholder) + assert cnn_row["mode"] == "small-torch-train" + assert cnn_row["vram_gb"] > 0 + assert cnn_row["ram_gb"] > 0 + assert cnn_row["time_hours"] > 0 + + def test_rnn_row_uses_hidden_dim(self) -> None: + # Bigger hidden_dim → bigger VRAM. + base = {"module_name": "rnn", "embed_dim": [128], "batch_size": [64], "num_train_epochs": [30]} + + def _run(hidden: int) -> float: + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [{**base, "hidden_dim": [hidden]}]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 5}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + return next(d["vram_gb"] for d in report.resource.drivers if d["module"] == "rnn") + + assert _run(1024) > _run(128), "larger hidden_dim must produce a larger VRAM row" + + +class TestNtrialsSharedAcrossVariants: + """n_trials is a *node* budget shared across module_name candidates.""" + + def test_single_module_gets_full_n_trials(self) -> None: + # Big dataset + embedder so linear time is non-zero and comparable. + embedder_cfg = {"embedder_config": {"model_name": "intfloat/multilingual-e5-large-instruct"}} + cfg = { + **embedder_cfg, + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "linear"}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 200}, + } + cfg2 = { + **embedder_cfg, + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "linear"}, + {"module_name": "knn", "k": [5]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 200}, + } + stats = DatasetStats.placeholder(n_samples=10000, n_classes=77, avg_tokens=24) + solo = run_preflight(cfg, stats, _profile()) + shared = run_preflight(cfg2, stats, _profile()) + + solo_lin = next(d["time_hours"] for d in solo.resource.drivers if d["module"] == "linear") + shared_lin = next(d["time_hours"] for d in shared.resource.drivers if d["module"] == "linear") + # Same module, same everything, but shared node has 2 variants → linear + # sees half the trials. + assert solo_lin > 0 + assert shared_lin > 0 + assert solo_lin > shared_lin, ( + f"linear alone should get full n_trials, shared should get half; got solo={solo_lin} shared={shared_lin}" + ) + # Concretely: solo=20 trials, shared=10 trials → 2x ratio (allow slop for rounding). + ratio = solo_lin / shared_lin + assert 1.5 < ratio < 2.5, f"expected ~2x ratio, got {ratio}" + + +class TestProcessBaselineFloor: + """Every fit reserves ~1.5 GB RAM for torch/transformers/datasets.""" + + def test_ram_estimate_never_below_baseline(self) -> None: + # Minimal preset — no scoring modules that contribute RAM. + cfg = { + "search_space": [ + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 1}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + # The floor is applied as an additive term, so even an empty pipeline + # must report at least the baseline in RAM. + assert report.resource.ram_gb >= 1.0 + + def test_cuda_vram_baseline_only_when_gpu_used(self) -> None: + cfg_cpu_only = { + "search_space": [ + {"node_type": "scoring", "search_space": [{"module_name": "linear"}]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 1}, + } + # linear scorer runs on CPU only → no CUDA baseline should apply. + report = run_preflight(cfg_cpu_only, DatasetStats.placeholder(), _profile(accelerator="cuda")) + assert report.resource.vram_gb == 0, "CPU-only preset must not spend the CUDA VRAM baseline" + + +class TestModuleCardinality: + """1 for all-singleton, N for finite lists, None for continuous ranges.""" + + def test_all_singleton(self) -> None: + from autointent._advisor._estimates._search_space import _module_cardinality + + assert _module_cardinality({"module_name": "bert"}) == 1 + assert _module_cardinality({"module_name": "bert", "batch_size": [64], "epochs": [30]}) == 1 + + def test_multi_list_multiplies(self) -> None: + from autointent._advisor._estimates._search_space import _module_cardinality + + # 2 batch × 3 lr candidates = 6 unique configs + cardinality = _module_cardinality( + {"module_name": "bert", "batch_size": [32, 64], "learning_rate": [1e-5, 5e-5, 1e-4]} + ) + assert cardinality == 6 + + def test_range_dict_is_unbounded(self) -> None: + from autointent._advisor._estimates._search_space import _module_cardinality + + # {low, high} → continuous → None (treated as unbounded) + assert _module_cardinality({"module_name": "knn", "k": {"low": 1, "high": 20}}) is None + + def test_reserved_keys_skipped(self) -> None: + from autointent._advisor._estimates._search_space import _module_cardinality + + # module_name / target_metric are not search dimensions + assert ( + _module_cardinality( + {"module_name": "bert", "target_metric": "scoring_f1", "batch_size": [32, 64]} + ) + == 2 + ) + + +class TestNoOpHpoFinding: + """Config-phase warns when n_trials >> unique configs.""" + + def test_finding_on_singleton_bert_with_high_n_trials(self) -> None: + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "num_train_epochs": [30], "batch_size": [64]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 40}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + no_op = [f for f in report.findings if "unique configurations" in f.message] + assert len(no_op) == 1, f"expected exactly one no-op warning, got {[f.message for f in no_op]}" + assert no_op[0].phase == "config" + assert no_op[0].severity == Severity.TIGHT + assert "bert" in no_op[0].message + + def test_no_finding_when_search_space_has_range(self) -> None: + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "learning_rate": {"low": 1e-5, "high": 1e-4}}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 40}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + no_op = [f for f in report.findings if "unique configurations" in f.message] + assert no_op == [], f"unexpected warning for ranged search space: {[f.message for f in no_op]}" + + def test_no_finding_when_n_trials_matches_cardinality(self) -> None: + # n_trials=4, cardinality=2×2=4 → not a "no-op" waste + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "batch_size": [32, 64], "num_train_epochs": [10, 20]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 4}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + no_op = [f for f in report.findings if "unique configurations" in f.message] + assert no_op == [], "n_trials matching cardinality should not warn" + + +class TestModeAwareVramBaseline: + """CUDA baseline + safety margin are mode-aware — training reserves more + cuDNN workspace than inference.""" + + def test_inference_only_preset_gets_smaller_vram_than_training(self) -> None: + # Both use e5-large; only the training config triggers the bigger baseline. + inference_only = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "knn", "k": [5]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "embedder_config": {"model_name": "intfloat/multilingual-e5-large-instruct"}, + "hpo_config": {"n_trials": 5}, + } + training = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "batch_size": [16], "num_train_epochs": [1]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 5}, + } + stats = DatasetStats.placeholder(n_samples=1000, n_classes=10, avg_tokens=24) + + infer_r = run_preflight(inference_only, stats, _profile(vram_gb=16.0)) + train_r = run_preflight(training, stats, _profile(vram_gb=16.0)) + + # The training baseline is 1.0 GB, inference baseline is 0.3 GB — so + # subtracting the driver max should show at least the 0.7 GB gap. + infer_max_driver = max((d.get("vram_gb") or 0 for d in infer_r.resource.drivers), default=0) + train_max_driver = max((d.get("vram_gb") or 0 for d in train_r.resource.drivers), default=0) + infer_baseline = infer_r.resource.vram_gb - infer_max_driver + train_baseline = train_r.resource.vram_gb - train_max_driver + assert infer_baseline < train_baseline, ( + f"inference baseline should be smaller; got infer={infer_baseline:.2f} train={train_baseline:.2f}" + ) + # Should be roughly the 0.3 vs 1.0 gap (small tolerance for rounding). + assert train_baseline - infer_baseline > 0.5 + + def test_inference_only_still_has_a_cuda_baseline(self) -> None: + # Even inference-only should be > 0 on CUDA — a non-zero cuDNN + driver + # context is real. Not zeroing this out would falsely tell users that + # embedder-only presets need no GPU memory. + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [ + {"module_name": "knn", "k": [5]}, + ]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "embedder_config": {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}, + "hpo_config": {"n_trials": 5}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile(vram_gb=16.0)) + assert report.resource.vram_gb > 0 diff --git a/tests/pipeline/test_calibration_tracker.py b/tests/pipeline/test_calibration_tracker.py index 7e5640035..d538f6f6f 100644 --- a/tests/pipeline/test_calibration_tracker.py +++ b/tests/pipeline/test_calibration_tracker.py @@ -122,29 +122,27 @@ def test_role_classification_and_time_decomposition() -> None: def test_step_timing_callback_answers_every_hf_hook() -> None: - """HF's CallbackHandler dispatches with a bare ``getattr(cb, event)`` — no - hasattr probe — so the callback MUST answer every ``on_*`` hook it may - ever ask for, even ones we don't time. - - Pins the regression that surfaced as - ``AttributeError: '_StepTimingCallback' object has no attribute 'on_train_begin'`` - when a bert scorer trial fired the harness's step-timing patch on real HF - Trainer machinery. The fix is to subclass ``transformers.TrainerCallback`` - directly so every default hook is inherited as a no-op — no ``__getattr__`` - trickery, no per-hook boilerplate. - """ + """HF dispatches every hook via bare getattr — callback must be a real + TrainerCallback subclass so all defaults are inherited as no-ops. + Regression: AttributeError on on_train_begin during a bert trial.""" from transformers import TrainerCallback cb = _StepTimingCallback(sink=[]) - # Must actually be an HF TrainerCallback — this is the guarantee that - # every hook HF may dispatch resolves to an inherited pass-through. assert isinstance(cb, TrainerCallback) - - # Sanity-check a representative slice of hooks (every documented HF hook - # subclass has one, and the isinstance above already proves the rest are - # inherited). Calling them with (args, state, control, **kwargs) must - # succeed and return None — HF's call_event keeps the incoming control - # unchanged when the result is None. for name in ["on_init_end", "on_train_begin", "on_epoch_begin", "on_log", "on_save", "on_train_end"]: hook = getattr(cb, name) assert hook(None, None, "control", model=None) is None, name + + +def test_peak_sampler_records_cuda_current_allocation() -> None: + """CUDA polling API: sample_cuda=True → 0.0 float, False → None. Thread + starts and stops cleanly with no CUDA hardware present.""" + from calibrate_advisor import _PeakSampler + + s = _PeakSampler(sample_cuda=True) + assert s.peak_cuda_gb == 0.0 + assert _PeakSampler(sample_cuda=False).peak_cuda_gb is None + import time + with s: + time.sleep(0.15) + assert s.peak_cuda_gb == 0.0 From d84b4be428f6a7aa4005eb22e695ff1cabfb143a Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 17:53:18 +0300 Subject: [PATCH 30/43] chore: move advisor calibration harness to experiments repo The harness is a validation instrument, not library code. It now lives in Darinochka/AutoIntent-experiments PR #40 beside the results it produced. - drops the tests/ -> scripts/ import that broke mypy (scripts/ is not a package) - reverts ruff-format-only churn in tests/test_deps.py, tests/ci/test_compute_matrix.py - gitignores local validation artifacts and Superpowers process docs --- .gitignore | 10 +- compute-feasibility-advisor-proposal.md | 281 ----- scripts/calibrate_advisor.py | 1316 -------------------- scripts/coverage_preset.yaml | 42 - scripts/run_calibration_banking77.sh | 163 --- tests/pipeline/test_calibration_tracker.py | 148 --- 6 files changed, 9 insertions(+), 1951 deletions(-) delete mode 100644 compute-feasibility-advisor-proposal.md delete mode 100644 scripts/calibrate_advisor.py delete mode 100644 scripts/coverage_preset.yaml delete mode 100755 scripts/run_calibration_banking77.sh delete mode 100644 tests/pipeline/test_calibration_tracker.py diff --git a/.gitignore b/.gitignore index 4fc2150c1..1a8cf9a2c 100644 --- a/.gitignore +++ b/.gitignore @@ -184,4 +184,12 @@ vector_db* /wandb model_output/ my.py -.DS_store \ No newline at end of file +.DS_store + +# Local advisor validation artifacts + Superpowers process docs (never shipped) +res-adapt-ckeck/ +banking77_*_repredicted.json +claude-issue-investigation.md +review-of-review.md +docs/superpowers/ +.python-version diff --git a/compute-feasibility-advisor-proposal.md b/compute-feasibility-advisor-proposal.md deleted file mode 100644 index 7ebf70bd9..000000000 --- a/compute-feasibility-advisor-proposal.md +++ /dev/null @@ -1,281 +0,0 @@ -# Compute Feasibility Advisor for AutoIntent - -- **Date:** 2026-05-23 -- **Status:** Proposal (pre-implementation) -- **Audience:** AutoIntent maintainers / contributor picking up the task -- **Scope of this document:** technical specification — *what* the advisor estimates and the formulas it uses. Architectural and system-design choices (where the advisor lives in the codebase, how it integrates with the optimizer, the public API surface, file/module layout) are deliberately left to the implementer. - -## Problem - -AutoIntent's main strength is letting a user kick off a full search-space optimization with one call: - -```python -pipeline = Pipeline.from_preset("transformers-heavy") -pipeline.fit(dataset) -``` - -The cost of that convenience is that users — especially those running on a laptop, a single consumer GPU, or a free cloud instance — cannot tell ahead of time whether their hardware can carry the configuration they have just selected. - -Concrete failure cases we see today: - -- `transformers-heavy` fine-tunes `microsoft/deberta-v3-large` for up to 30 epochs across 40 HPO trials. That needs ~12–18 GB VRAM (full fine-tune, fp32) and many hours of wall time on a single GPU. A user with an 8 GB card finds out by OOM, often several minutes into a run. -- Swapping `intfloat/multilingual-e5-large-instruct` (2 GB) for `sentence-transformers/all-MiniLM-L6-v2` (90 MB) changes the resource bill by an order of magnitude — but nothing surfaces this difference up front. -- Disk is a silent failure mode: a search space referencing several large checkpoints can pull >10 GB into the HF cache before any training starts. - -The target audience for this feature is users with limited resources who pick a preset, hit `fit()`, and want to know within a second whether they should change something. - -## Proposed solution: pre-flight resource advisor - -Add a **pre-flight advisor** that, given a parsed search space and a dataset, estimates worst-case disk, RAM, VRAM, and wall-time requirements from public Hugging Face Hub metadata and a small set of formulas, then prints a clear summary with red/yellow/green warnings. By default it is **report-only and never blocks the run**; an opt-in **reduce-to-fit** mode additionally prunes the search space to fit detected hardware. - -### Scope - -The advisor analyses only the **local, model-bearing** modules whose footprint can be derived from HF Hub metadata. Everything else is either trivial or out of band. - - -| Module category | In scope? | Reason | -| -------------------------------------------------------------------------------- | --------- | -------------------------------------------------- | -| `SentenceTransformerEmbeddingConfig` | yes | local transformer, dominant cost on small machines | -| `VllmEmbeddingConfig` | yes | local transformer with extra engine overhead | -| `HFModelConfig`-based scorers (`bert`, `lora`, `ptuning`, `dnnc`, cross-encoder) | yes | the actual heavyweights | -| GCN scorer when configured with a transformer backbone | yes | inherits the backbone cost | -| `LinearScorer` (sklearn `LogisticRegression` / `LogisticRegressionCV`) | yes | dominant cost on presets with no transformer fine-tune; the CV path multiplies a single fit by ~30 | -| `CatBoostScorer` | yes | dominant cost on presets with no transformer fine-tune; high default `iterations` | -| `OpenaiEmbeddingConfig` | no | no local resources to estimate | -| `HashingVectorizerEmbeddingConfig` | no | trivial cost | -| `knn`, `mlknn`, generic `sklearn` classifiers via `SklearnScorer`, `description` | no | bounded so far below any in-scope module that they cannot plausibly be the bottleneck | -| `decision` and `regex` nodes | no | negligible | - - -Rationale: the user's real risk is whichever module is the actual bottleneck. On heavy presets that is a transformer fine-tune; on light presets it shifts to `linear` (CV-multiplied) or `catboost` (1000 default iterations × dataset shape). Modules left out of scope are ones whose cost is bounded so far below any in-scope module that they cannot plausibly be the reason a run fails. - -### Phases - -The advisor is one entry point, but internally splits work into three phases that share a single `PreflightReport` object. The split is internal organization — all three run at the same hook point (after `validate_modules`, before `_fit(context)`) and the user sees one summary. Separating them keeps each phase's inputs, formulas, and failure modes scoped: - -- **Resource phase.** Disk / RAM / VRAM / wall-time estimates and comparisons against detected hardware. Most of the formulas in this document live here. This is the only phase consumed by the reduce-to-fit pruner. -- **Data quality phase.** Findings derived from the dataset jointly with the active search space — token-length truncation, split readiness (auto-invokes the existing `check_split_readiness` utility rather than re-implementing it), partial intent descriptions paired with the `description` scorer, embedder/scorer dimension consistency. Reports red/yellow lines but never prunes the search space; the user fixes the dataset or the config. -- **Configuration sanity phase.** Joint checks across dataset + search-space + hardware that don't slot cleanly into the other two — e.g., `hpo_config.n_jobs > 1` × per-trial VRAM contention, CatBoost `task_type="GPU"` with no CUDA. Pydantic schema validation already runs upstream on `OptimizationConfig`; this phase only adds checks that need joint inspection. - -The advisor consumes `validate_modules`'s *post-filter* view of `self.nodes` — it does not duplicate that mutating filter. - -### Inputs - -- The parsed `OptimizationConfig` (search space, HPO config, embedder/transformer configs). -- The training `Dataset` (for `dataset_size` and an approximate token-length distribution). -- Detected local hardware: - - Total / available RAM via `psutil`. - - Free disk on the AutoIntent / HF cache directory via `shutil.disk_usage`. - - Accelerator detection, in priority order: - - **CUDA:** per-GPU VRAM and device name via `torch.cuda`. - - **MPS (Apple Silicon):** detected via `torch.backends.mps.is_available()`. Apple chips use unified memory, so there is no separate VRAM pool — the "VRAM budget" is a fraction of total system RAM. Default budget = 70 % of total RAM (matching the macOS `PYTORCH_MPS_HIGH_WATERMARK_RATIO` default) with the remainder reserved for the OS and other apps. The fraction is exposed as a knob. - - **CPU only:** when neither is available. - -### Output - -A structured estimate plus a human-readable summary printed to the logger. Example: - -``` -Compute feasibility check -───────────────────────── -Resource: - Available : 8 GB VRAM (NVIDIA RTX 3060), 32 GB RAM, 120 GB free disk - Disk : 5.2 GB to download, 1.1 GB already cached (3 unique checkpoints) - RAM : ~4 GB - VRAM : ~14 GB × 2 parallel trials (n_jobs=2) ⚠ exceeds available - Time : ~6 h (+~12 min for refit_after) (single-GPU, fp32, rough) - -Data: - Train tokens p95 : 612 (exceeds bert.max_length=512) ⚠ ~7% truncated - Split readiness : 2 classes have <3 samples — LogisticRegressionCV cv=3 will fail ✗ - -Config: - CatBoost task_type=GPU but no CUDA detected — will fall back to CPU ⚠ - -Drivers of cost: - scoring.bert microsoft/deberta-v3-large full fine-tune × 40 trials × 30 epochs → ~14 GB VRAM, ~5 h - embedder intfloat/multilingual-e5-large-instruct → ~2.2 GB VRAM - -Suggestions: - • Enable mixed precision (fp16/bf16) on the bert scorer - • Reduce batch_size from 64 to 16 or 32 - • Set hpo_config.n_jobs=1 — parallel trials are doubling VRAM demand - • Try preset `transformers-light` or `classic-medium` - -These numbers are heuristic upper bounds, not measurements. -``` - -Numbers are reported with honest precision (one significant figure for time, two for memory) and an explicit "estimate, not measurement" disclaimer. - -### Algorithm (proposal, allowed to adjust) - -1. **Collect candidates.** Walk the search space; collect every unique in-scope module. For transformer-bearing modules the identity is `(module_type, model_name, mode)` with `mode ∈ {inference, lora, full-finetune}`. For `linear` and `catboost` the identity is `(module_type, embedder_name, task_kind)` with `task_kind ∈ {multiclass, multilabel}` — the routing through `LogisticRegressionCV` vs `MultiOutputClassifier`, and CatBoost's per-class trees, both depend on it. Also collect the HPO knobs that drive cost: `n_trials` plus per-module knobs — transformer (`epochs`, `batch_size`, `max_length`, `dtype` ∈ {fp16, bf16, fp32}), `linear` (`cv`, `max_iter`), `catboost` (`iterations`, `depth`, `task_type`, `features_type`). -2. **Resolve checkpoints.** For each unique `model_name`, query HF Hub for safetensors metadata to read parameter count and weight dtype. Fall back to file-size aggregation if safetensors metadata is missing. Fall back to a "unknown — heuristic only" tag with low-confidence labelling if HF Hub is offline or the repo is private. `LinearScorer` and `CatBoostScorer` have no checkpoint of their own; they reuse the embedder resolved by this step in their formulas (their cost is parameterised by `embedder_dim`, not parameter count). -3. **Apply formulas.** All values are honest upper bounds; convergence and early stopping often terminate well below them. - - **Disk** = sum over unique downloadable checkpoints of total file size, plus a small fixed overhead per checkpoint for tokenizers and config. `LinearScorer` and `CatBoostScorer` contribute zero (they consume embedder output that is already accounted for upstream). - - **RAM per module:** - - Transformer modules (any mode): `params × dtype_bytes + dataset_tokens × 4 bytes`, treated as a loose upper bound for tokenized buffers. - - `LinearScorer`: `8 × n_samples × embedder_dim` (float64 data matrix — the dominant term) `+ 8 × n_classes × embedder_dim` (coefficients) `+ ~10 × 8 × embedder_dim` (L-BFGS history). - - `CatBoostScorer`: `4 × n_samples × n_features` (data, float32 internally) `+ 4 × n_features × n_bins` (histograms; default `n_bins = 254`) `+ iterations × 2^depth × ~32 bytes` (tree storage). For `features_type ∈ {embedding, both}`, `n_features = embedder_dim`. For `features_type = text`, `n_features` is the BoW vocab discovered at fit; bound with a coarse default (e.g. 50 000) and tag the estimate low-confidence. - - For `linear` and `catboost`, `embedder_dim` is taken from the largest embedder in the same node group — same worst-case stance as the rest of the estimate. - - **VRAM per module:** - - Inference embedder: `params × dtype_bytes × ~1.3` (small constant for activations). - - Full fine-tune (`bert`, GCN backbone, soft-prompt `ptuning`): `params × dtype_bytes × (1 + 1 + 2)` for weights + grads + Adam state, halved when fp16/bf16 mixed precision is configured. - - LoRA: inference VRAM + a small adapter constant. - - Reranker (cross-encoder, `dnnc`): inference VRAM × small factor for the reranking pass. - - `LinearScorer`: N/A (sklearn is CPU-only). - - `CatBoostScorer`: 0 by default; if `task_type="GPU"` is configured, the RAM formula above lives on device instead. - - **Time per module:** - - Transformer modules: `n_trials × epochs × (dataset_size / batch_size) × per_step_seconds(params, max_length, device_class)`, where `per_step_seconds` is a small static lookup keyed on coarse device class (`cpu`, `low-gpu`, `mid-gpu`, `high-gpu`, `apple-silicon`) auto-detected from `torch.cuda.get_device_name` or `platform`/`torch.backends.mps`. - - `LinearScorer`: `n_trials × C_cpu × n_samples × embedder_dim × max_iter × cv_multiplier × class_multiplier`, where: - - `C_cpu ≈ 1e-8 s` per `(sample × feature × iteration)` on a single modern CPU core. - - `cv_multiplier = Cs × cv + 1 ≈ 31` for the multiclass path (`LogisticRegressionCV` with default `Cs = 10`, repo default `cv = 3`, plus one final refit). `cv_multiplier = 1` for the multilabel path (no inner CV). - - `class_multiplier = n_classes` for the multilabel path (`MultiOutputClassifier` fits one binary LogReg per class); `class_multiplier = 1` otherwise. - - `CatBoostScorer`: `n_trials × iterations × C_device × n_samples × n_features × depth × class_multiplier`, where: - - `C_device ≈ 1e-9 s` on CPU, ~5–20× faster on GPU. Resolve `C_device` via the same `device_class` lookup as the transformer time formula. - - `class_multiplier = n_classes` for both the multiclass `MultiClass` loss (per-class trees per iteration) and the multilabel routing (one CatBoost per class). - - Early stopping is not modelled; `iterations` is treated as the upper bound. - - Total time = sum across modules. MPS time numbers are coarser than CUDA's (one tier for now); we accept that. -4. **Compare to detected hardware.** Per-dimension status is green / yellow / red against a configurable headroom (defaults: **red** if estimate > 100 % of available, **yellow** if > 70 %). On MPS, "VRAM" and "RAM" estimates draw from the same physical pool; we compare *the larger of the two* against the unified-memory budget rather than each independently. -5. **Render summary.** Log at INFO. If any dimension is red, emit at WARNING so it shows in non-logging contexts. - -#### Resource-phase refinements - -These adjust the formulas above for situations that look fine in single-trial isolation but blow up in practice: - -- **Cold-vs-warm HF cache (Tier 1).** Before reporting disk, probe each unique `model_name` against the local HF cache via `huggingface_hub.try_to_load_from_cache` / `scan_cache_dir`, keyed off `HF_HOME`. Split the disk line into `to_download` vs `already_cached`. Treat a repo as cached only if the weight shard (`model.safetensors` or equivalent) is present — not just config/tokenizer files. Without this, a repeated run on the same machine alarms the user about gigabytes they already have. -- **Concurrent-trial × per-trial VRAM (Tier 1).** Multiply the per-trial VRAM estimate by `hpo_config.n_jobs` when `n_jobs > 1` and the active accelerator is GPU. Same for the `dump_modules=True` path on disk: each trial writes module weights to the dump dir, so multiply per-module dump-disk by `n_trials`. vLLM is process-isolated and its contention model differs; note this in the disclaimer. -- **`refit_after=True` time delta (Tier 2).** When `Pipeline.fit(refit_after=True)`, add one full-data training pass per node to the time estimate. Small term but easy to forget; users running close to their time budget care about it. -- **HF Hub reachability probe (Tier 2).** One up-front `HfApi().whoami()` (or unauthenticated `HEAD` to `huggingface.co`) at the start of the phase. On failure, consistently downgrade *all* model entries to the "unknown — heuristic only" path instead of timing out per-model 10× on a 10-model search space. -- **CatBoost `task_type="GPU"` sanity (Tier 2).** When CatBoost is in the search space with `task_type="GPU"` but `torch.cuda.is_available()` is false, tag yellow — CatBoost silently falls back to CPU and the user otherwise sees CPU speeds with no warning. - -### Data quality phase - -The resource phase predicts whether the run *fits*. The data quality phase predicts whether the run *produces a meaningful result*. Both are caught at the same hook point because both have the same failure mode from the user's perspective: hours of compute followed by a cryptic error or a silently degraded model. - -- **Token-length truncation (Tier 1).** Sample ~1000 utterances from the train split, tokenize against each unique transformer's tokenizer, compute `p95_tokens` and `% truncated` against the module's `max_length`. Yellow when >1% truncated; red when >10%. Reuse the tokenizer the resource phase already loaded for parameter-count resolution — don't double-fetch. The existing pipeline silently truncates (sentence-transformers and the HF Trainer both default to `truncation=True`); there is no warning anywhere today. -- **Auto-invoke `check_split_readiness` (Tier 1).** Call the existing utility at `context/data_handler/_readiness_util.py:44–109` with the active `data_config` and surface its `SplitReadinessResult` — it already returns `underpopulated_classes`, `ready`, and a `reason` string, but is not called anywhere from `Pipeline.fit()` today. When `LinearScorer` with CV is in the search space and any class has `n < cv`, name the module by name in the red line ("`LogisticRegressionCV` cv=3 will fail: classes [X, Y] have <3 samples") rather than emitting a generic split-readiness message. -- **Partial intent descriptions × `description` scorer (Tier 1).** The dataset constructor already warns once at import when *some* but not all intents have descriptions (`_dataset/_dataset.py:199–207`). The advisor escalates this to red when the `description` scorer is also present in the active search space — otherwise the run will produce NaN embeddings for the missing intents. Action message: "fill in N missing descriptions", not "drop the scorer". -- **Embedder ↔ scorer dimension consistency (Tier 2).** For `LinearScorer` / `CatBoostScorer` with `features_type="both"`, verify the embedder reachable from the same node group exposes a stable, expected dimension. Cross-node walk; surface as yellow when the resolved dimension cannot be confirmed pre-flight. - -### Configuration sanity phase - -Pydantic schema validation on `OptimizationConfig` runs upstream at config-load time; this phase only adds checks that require *joint* inspection of dataset + search-space + hardware. With Tier 1 + Tier 2 in scope today, this phase holds two items: - -- The `n_jobs × VRAM` callout, surfaced jointly with the resource phase (single line in the rendered output). -- The CatBoost `task_type="GPU"` without CUDA check, same. - -Both could live entirely in the resource phase; they get their own phase because future additions — joint scorer↔decision shape checks, OOS-support mismatches detected up front rather than at module instantiation, embedder-dimension mismatches — slot here naturally. Keep the phase scaffold even if it is currently thin. - -### Failure modes - -- **HF Hub offline or private repo:** fall back to "unknown model — name-pattern heuristic only", explicit low-confidence label, never raise. -- **No accelerator (no CUDA and no MPS):** report VRAM as N/A and mark GPU-only modules as "requires GPU" without estimating a (misleading) CPU wall time. -- **MPS configured but a module is incompatible:** vLLM in particular does not run on MPS. Flag the module as "unsupported on MPS" rather than estimating; do not raise. -- **MPS with CPU fallback ops:** some PyTorch ops fall back to CPU on MPS, inflating system-RAM usage and wall time beyond the heuristic. Note this in the disclaimer; we don't try to model it. -- **vLLM configured but not installed:** still estimate (the VRAM accounting is similar), note that the engine itself has additional overhead not captured. -- **Estimate wildly wrong vs. reality:** always-on disclaimer in the printed summary that these are heuristic upper bounds. - -### Reduce-to-fit mode - -The feasibility check has two modes sharing the same estimation pipeline: - -- **Report mode (default).** Print the summary, return the structured estimate, let the run proceed regardless of severity. -- **Reduce-to-fit mode (opt-in).** Additionally prune the search space to fit detected hardware before the run starts. Same estimates, same comparisons — just one extra step that produces a reduced search space. - -Reduce-to-fit consumes only the **resource phase** output. Data-quality and config-sanity findings are reported but never trigger pruning — they require user action (fix the dataset, change a config flag), not search-space narrowing. - -Using the same per-module estimates, the pruner applies three least-destructive steps in order: - -1. **Filter discrete-choice hyperparameters.** For lists of cost-driving values (model name, batch size, training epochs, CatBoost `iterations` / `depth`, sklearn `cv`), keep only entries whose worst-case estimate fits. -2. **Cap continuous ranges.** For `{low, high}` ranges of cost-driving parameters, lower the upper bound to the largest fitting value. Ranges of non-cost parameters (learning rate, decision thresholds) are not touched. -3. **Drop module variants.** If a module entry has any required hyperparameter with no satisfiable value left, drop that module entry from its node's search space. - -Guard rails: - -- If pruning would leave any node's search space empty, the pruner **raises**. We don't silently produce a non-runnable pipeline, and we don't quietly fall back to report-only — failing loudly is the right contract for a mode whose whole purpose is to make the run feasible. The error message points the user toward a lighter preset. -- Time is not used as a filter — only memory and disk are. Time is still reported. -- Headroom thresholds are intentionally generous to avoid over-pruning and are configurable. - -Alongside the standard estimate, the caller receives a structured description of what was filtered, capped, and dropped, plus the resulting search space and its recomputed (now green) estimate. - -**Drawbacks worth surfacing.** - -- **Silent narrowing of intent.** A search space deliberately written to include heavy/light variants for comparison gets halved. The mode is opt-in for this reason. -- **Over-pruning when our formulas overestimate.** A 30 %-high estimate on a borderline configuration throws away a run that would have succeeded. Generous headroom defaults mitigate; the knob is exposed. -- **Hard failure when nothing fits.** Raising is intentional — silent degradation to report-only would defeat the mode's purpose — but it is a sharper edge than report mode has. -- **Pre-trial only.** The rewrite happens before any HPO trial starts. This is fine because the search space is treated as immutable across a study, but worth calling out so nobody tries to make this dynamic later. - -### CLI surface - -The advisor is also exposed as a console script (`autointent-advisor`) so users can answer "what will this cost?" and "what should I run?" without writing Python. Two subcommands: - -- **`autointent-advisor inspect `.** Resolves the preset (or a user-supplied `OptimizationConfig`), detects local hardware, runs the same three-phase advisor that `Pipeline.fit()` runs, and prints the same report. Accepts `--dataset` for a real dataset, or `--n-samples / --n-classes / --avg-tokens` placeholders when the dataset is not yet built — so the script is useful before any training data exists. `--json` emits the structured `PreflightReport` for scripting. -- **`autointent-advisor recommend [--n-samples ... | --dataset ...] [--budget-time 12h] [--budget-vram-gb 8]`.** Detects local hardware (with manual overrides applied), iterates over the bundled presets in `_presets/`, and tags each as `feasible` / `feasible-with-reduce` / `infeasible`. Ranks feasible presets by quality tier (`heavy > medium > light`) then estimated wall-time; picks the top one as the recommendation. For the heaviest infeasible preset, surfaces the single most-impactful knob change that would make it fit (e.g., "`transformers-heavy` would fit if `batch_size` ≤ 16 and `dtype=fp16`"), reusing the reduce-to-fit pruner's per-knob delta info. - -**Constraints (both subcommands).** No model downloads — only HF Hub metadata endpoints (`HfApi().model_info`); never `from_pretrained`. Offline-safe — on Hub unreachability, fall back to the same "heuristic only" path and mark the report low-confidence; do not raise. Hardware-detection failures (broken CUDA install where `torch.cuda.mem_get_info()` raises) fall back to CPU detection and tag the report rather than crashing. - -## Alternatives considered and rejected - -### B. Smoke-test calibration - -Run each unique module for one mini-batch / one step before the real fit, measure peak RAM and VRAM with `psutil`, `tracemalloc`, and `torch.cuda.max_memory_allocated`, time the step, and extrapolate to the full search space. - -Rejected because: - -- It **downloads weights just to estimate** — the disk-headroom check we wanted to provide is defeated by the act of performing it. -- It can **OOM while predicting OOM**, exactly on the constrained hardware that is the target audience. -- It adds **seconds to minutes** of wall time before `fit()` does anything, surprising users. -- It needs per-module "tiny run" hooks; not every scorer has a clean "stop after one step" path. -- For OpenAI- or vLLM-served embedders, a smoke test costs real money or starts the engine. -- Still not accurate due to CUDA and CPU cache, memory heating and so on. - -### C. Curated benchmark table - -Ship a JSON in the package with measured VRAM and per-step time for the bundled-preset checkpoints, broken out by hardware class (cpu / mid-gpu / high-gpu) and mode (inference / lora / full-finetune). Fall back to heuristics for unknown checkpoints. - -Rejected because: - -- **Maintenance burden:** every new model added to a preset would need entries across the hardware × precision × mode matrix. -- Numbers **go stale** when `transformers` updates change defaults (attention impl, dtype, gradient checkpointing). -- It still needs the chosen-solution heuristics as a long-tail fallback — so it adds work on top of Option A without replacing it. -- **Confident-but-wrong is worse than honest-but-fuzzy.** A table that says "4 GB on 4090" when the user OOMs at 4.5 GB damages trust more than a clearly-labelled range would. - -### D. Layered (A by default, opt-in B, embedded table from C, local actuals cache) - -Combine all three: ship A as the fast path, allow `calibrate=True` to trigger B for heavy modules only, embed a small table from C for the bundled-preset checkpoints, and write actuals from every real run to a local cache that feeds back into future estimates. - -Rejected because: - -- **Implementation surface multiplies:** two estimation code paths to keep consistent, a cache schema with versioning and eviction, two failure modes to document. -- **Discoverability:** users may not learn about `calibrate=True` and the realized value compresses back to roughly Option A anyway. -- The team's bandwidth doesn't justify the marginal accuracy gain over A for the target audience. - -## Comparison - - -| Dimension | A (chosen) | B (smoke-test) | C (benchmark table) | D (layered) | -| -------------------------------- | ------------------------------ | ---------------------- | ---------------------------------- | ------------------------------------- | -| Wall time at pre-flight | < 1 s | seconds–minutes | < 1 s | < 1 s default, s–min when calibrating | -| Accuracy on common checkpoints | medium | high | high | high | -| Accuracy on custom checkpoints | medium | high | medium (fallback) | medium–high | -| Time-estimate quality | low–medium | high | high | high | -| Disk pre-download required | no | yes | no | only when calibrating | -| Risk of OOM during the check | none | real | none | only when calibrating | -| Network usage | 1 cached call per unique model | none beyond normal fit | none | combination | -| Implementation effort | small | large | medium + ongoing benchmark refresh | large + cache infra | -| Ongoing maintenance | low (formulas only) | low | high | high | -| Friendly to offline / air-gapped | with fallback | yes | yes | partial | - - -The chosen solution accepts a real accuracy gap on time and a moderate accuracy gap on VRAM in exchange for the only profile that fits the target audience's constraints: zero added wall time, zero added downloads, zero added failure modes, and a small one-time implementation cost. - -## Out of scope (possible follow-ups) - -- Live resource observability during `fit()` (peak RAM / VRAM per trial, abort on overrun). -- A learned calibration cache from real runs to refine estimates over time. -- **Determinism / `cudnn.deterministic` check.** Belongs in seed-setting code (`set_seed` utility, `Pipeline.__init__`), not in a feasibility advisor — reproducibility is not a hardware-budget question. -- **OpenAI / Generator token-cost ($) estimation.** Real value, but pricing tables age badly, the `StructuredOutputCache` hit rate is unknowable upfront, and the API-paying audience overlaps poorly with this advisor's stated audience (resource-constrained local users). Push to a separate `cost_estimator` tool. -- **Predictive CO₂ / emissions.** `_callbacks/emissions_tracker.py` already does this retrospectively, accurately. A predictive version multiplies our (loose) time estimate by a regional kWh/CO₂ factor — two sources of imprecision compounded. The retrospective number is the trustworthy one. -- **vLLM startup compile time.** Minutes of overhead before any work, but vLLM is unsupported on MPS, isn't the dominant cost on CUDA once running, and modelling it needs a startup-time lookup table. Note once in the disclaimer; do not model. - diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py deleted file mode 100644 index 8e7d54cf5..000000000 --- a/scripts/calibrate_advisor.py +++ /dev/null @@ -1,1316 +0,0 @@ -"""Calibrate advisor preflight estimates against real Pipeline.fit measurements. - -Runs each requested preset twice: first through ``run_preflight`` to capture the -heuristic estimate, then through ``Pipeline.from_preset(...).fit(...)`` while -measuring wall-time, peak RAM (RSS), peak VRAM (CUDA only — MPS has no exact -peak API), and the disk delta in the HF Hub cache. - -The output is a JSON file with per-preset predicted vs. actual values plus -ratios, and a side-by-side table on stdout for quick eyeballing. - -Usage: - python scripts/calibrate_advisor.py \\ - --dataset tests/assets/data/clinc_subset.json \\ - --presets classic-light classic-medium \\ - --output calibration.json \\ - --max-trials 3 - -The ``--skip-fit`` flag runs only the predicted side, useful for sanity-checking -the preflight numbers across presets without paying for fits. -""" - -from __future__ import annotations - -import argparse -import json -import logging -import os -import sys -import threading -import time -from dataclasses import asdict, dataclass, field -from pathlib import Path -from typing import Any - -import psutil - -from autointent import Dataset, Pipeline -from autointent._advisor import ( - BUNDLED_PRESETS, - PreflightReport, - detect_hardware, - run_preflight, - stats_from_dataset_obj, -) -from autointent._callbacks.base import OptimizerCallback -from autointent.configs import LoggingConfig -from autointent import setup_logging - -setup_logging("INFO", log_filename="logs.log") -logging.basicConfig(level=logging.INFO) - -logger = logging.getLogger("calibrate_advisor") - -_BYTES_PER_GB = 1024**3 - - -@dataclass -class CalibrationRow: - """One preset's predicted vs. actual numbers.""" - - preset: str - predicted: dict[str, float] = field(default_factory=dict) - actual: dict[str, float | None] = field(default_factory=dict) - findings: int = 0 - findings_over: int = 0 - # Per-module records from _ModuleTracker: [{module, num, config, duration_s, peak_vram_gb?}, ...] - modules: list[dict[str, Any]] = field(default_factory=list) - cache_policy: str = "unknown" # "cold" (embeddings cache cleared) | "warm" (kept as-is) - low_confidence: bool = False # advisor fell back to heuristic HF-metadata for one+ models - repeat_idx: int = 0 # 0-based index within a (preset, dataset) repeat group - # ``skipped`` is set when the preset needs an optional extra that isn't - # installed (peft / catboost / openai / ...). We still populate ``error`` - # for the summary, but callers analysing the JSON should treat - # ``skipped=True`` rows separately from ``error != None && skipped=False`` - # rows (real crashes) — the former are expected and shouldn't count as - # advisor failures. - skipped: bool = False - # Snapshot of ``autointent-advisor inspect --json`` run in-process - # under the same stats + budget as the direct-API preflight. Lets us catch - # a CLI-wrapper regression (JSON schema drift, feasibility verdict flip) - # without a separate subprocess round-trip. None means we didn't run it - # (e.g. skipped row, or CLI itself crashed — see notes for the reason). - cli_smoke: dict[str, Any] | None = None - error: str | None = None - notes: list[str] = field(default_factory=list) - - def to_dict(self) -> dict[str, Any]: - """Serialize with ratios + role-decomposed timings computed at read time. - - Storing ratios in the row is a bug magnet: any late edit to - ``predicted``/``actual`` gets missed. We compute them from the current - row values at serialization time so consumers can trust ``row["ratios"]``. - ``time_by_role_s`` splits the measured wall-time across embedder / - scorer / decision so classic-preset time can be interpreted (embedder - forward vs sklearn fit) without re-walking ``modules``. - """ - payload = asdict(self) - payload["ratios"] = self._ratios() - payload["time_by_role_s"] = _sum_time_by_role(self.modules) - return payload - - def _ratios(self) -> dict[str, float | None]: - keys = ("time_h", "ram_gb", "vram_gb", "disk_download_gb", "disk_embedding_cache_gb") - out: dict[str, float | None] = {} - for key in keys: - actual = self.actual.get(key) - predicted = self.predicted.get(key) - if actual is None or predicted is None or predicted <= 0: - out[key] = None - else: - out[key] = actual / predicted - return out - - -def _build_parser() -> argparse.ArgumentParser: - p = argparse.ArgumentParser( - prog="calibrate_advisor", - description="Compare advisor preflight estimates to real Pipeline.fit measurements.", - ) - p.add_argument( - "--dataset", - required=True, - nargs="+", - type=str, - help=( - "One or more datasets — each is either a local JSON path (loaded via " - "``Dataset.from_json``) or an HF Hub repo id such as ``DeepPavlov/banking77`` " - "(loaded via ``Dataset.from_hub``). Every preset runs against every dataset, " - "so pairing a multilabel + long-token + small + large dataset exercises the " - "n_samples / n_classes / avg_tokens surfaces of the advisor's formulas." - ), - ) - p.add_argument( - "--subsample-per-class", - type=int, - default=None, - help=( - "Cap each class to at most N training samples (deterministic first-N slice) " - "before running. Lets one big dataset stand in as a 'small' shape — enough to " - "exercise ``LogisticRegressionCV cv=3`` split-readiness and rare-class findings." - ), - ) - p.add_argument( - "--repeats", - type=int, - default=1, - help=( - "Run each (preset, dataset) N times so ratio gaps have variance bars. " - "The summary prints mean ± stdev for the actual measurements across " - "repeats; individual repeat rows are still written to the JSON with " - "``repeat_idx`` so consumers can compute their own aggregates. Default: 1." - ), - ) - p.add_argument( - "--presets", - nargs="+", - default=None, - help=( - "Preset names to run (default: every preset in BUNDLED_PRESETS). " - "Items ending in .yaml/.yml are treated as paths to a preset file — " - "used to run e.g. ``scripts/coverage_preset.yaml`` which packs " - "lora/ptuning/dnnc/gcn/cross-encoder into one small run for module " - "coverage without touching the shipped presets." - ), - ) - p.add_argument("--output", type=Path, default=Path("calibration.json"), help="Where to write the JSON report.") - p.add_argument("--max-trials", type=int, default=None, help="Override hpo_config.n_trials for faster runs.") - p.add_argument( - "--skip-fit", - action="store_true", - help="Only run preflight (no fit) — useful for sanity-checking estimates.", - ) - p.add_argument( - "--poll-interval-ms", - type=int, - default=100, - help="RSS polling interval during fit (ms). Lower is more accurate but more overhead.", - ) - p.add_argument( - "--wandb", - action="store_true", - help=( - "Attach the W&B reporter so per-step GPU/system metrics land in wandb.ai. " - "Requires ``wandb`` installed + ``WANDB_API_KEY`` in the environment." - ), - ) - p.add_argument( - "--run-name", - type=str, - default=None, - help=( - "Suffix appended to each preset's LoggingConfig.run_name — the resulting " - "value is ``{preset}_{run_name}`` and is used by the LoggingHandler as the " - "W&B run group / on-disk dump directory name." - ), - ) - p.add_argument( - "--clear-embedding-cache", - action="store_true", - help=( - "Wipe ``/autointent/embeddings/`` before each preset so every " - "measurement reflects a COLD run (embedder forward not skipped). Without this " - "flag, the cross-run cache silently makes later runs look artificially cheap." - ), - ) - p.add_argument( - "--budget-vram-gb", - type=float, - default=None, - help=( - "Override the detected VRAM budget passed to run_preflight, e.g. ``--budget-vram-gb 8`` " - "to exercise the constrained-hardware / severity paths on a big box without needing " - "a small GPU. Does NOT affect the real fit — only the predicted-side estimate." - ), - ) - p.add_argument( - "--require-cuda", - action="store_true", - help=( - "Fail fast if PyTorch can't initialize CUDA (guards against the silent " - "'2 GPUs detected, but torch runs on CPU' driver-mismatch trap)." - ), - ) - p.add_argument("-v", "--verbose", action="store_true") - return p - - -# === measurement helpers ================================================= - - -def _hf_cache_dir() -> Path: - """Return the active HF Hub cache directory ($HF_HOME / ~/.cache/huggingface).""" - return Path(os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")) - - -def _embeddings_cache_dir() -> Path: - """Return autointent's embeddings-cache dir (``/autointent/embeddings/``). - - Uses the same ``appdirs.user_cache_dir("autointent")`` path as - :func:`autointent._wrappers.embedder.utils.get_embeddings_path` so the - harness reads/clears the same directory the runtime writes to. - """ - from autointent._wrappers.embedder.utils import get_embeddings_path - - return get_embeddings_path("_probe").parent - - -def _clear_embeddings_cache() -> int: - """Delete every ``*.npy`` file in the embeddings cache. Returns count removed.""" - cache = _embeddings_cache_dir() - if not cache.exists(): - return 0 - removed = 0 - for path in cache.glob("*.npy"): - try: - path.unlink() - removed += 1 - except OSError: - continue - return removed - - -def _dir_size_gb(path: Path) -> float: - """Disk usage of ``path`` in GB; 0 when the directory is missing.""" - if not path.exists(): - return 0.0 - total = 0 - for entry in path.rglob("*"): - try: - if entry.is_file(): - total += entry.stat().st_size - except OSError: - continue - return total / _BYTES_PER_GB - - -class _PeakSampler: - """Background thread polling peak RSS + (optionally) MPS / CUDA current - allocation. CUDA polling catches allocations that fall outside any - module bracket — the per-module tracker's peak counter gets reset at each - start_module, losing anything allocated before it (e.g. the embedder - forward during pipeline setup). Best-effort: sub-poll-interval spikes - can be missed.""" - - def __init__( - self, interval_s: float = 0.1, *, sample_mps: bool = False, sample_cuda: bool = False, - ) -> None: - self._interval_s = interval_s - self._proc = psutil.Process() - self.peak_ram_gb = self._proc.memory_info().rss / _BYTES_PER_GB - self.peak_mps_gb: float | None = 0.0 if sample_mps else None - self.peak_cuda_gb: float | None = 0.0 if sample_cuda else None - self._sample_mps = sample_mps - self._sample_cuda = sample_cuda - self._stop = threading.Event() - self._thread: threading.Thread | None = None - - def __enter__(self) -> _PeakSampler: - self._thread = threading.Thread(target=self._run, daemon=True) - self._thread.start() - return self - - def __exit__(self, *_exc: object) -> None: - self._stop.set() - if self._thread is not None: - self._thread.join(timeout=1.0) - - def _run(self) -> None: - try: - import torch - except ImportError: - torch = None # type: ignore[assignment] - while not self._stop.is_set(): - try: - rss = self._proc.memory_info().rss / _BYTES_PER_GB - self.peak_ram_gb = max(self.peak_ram_gb, rss) - if self._sample_mps and torch is not None: - mps = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB - if self.peak_mps_gb is None or mps > self.peak_mps_gb: - self.peak_mps_gb = mps - if self._sample_cuda and torch is not None and torch.cuda.is_available(): - cuda = float(torch.cuda.memory_allocated()) / _BYTES_PER_GB - if self.peak_cuda_gb is None or cuda > self.peak_cuda_gb: - self.peak_cuda_gb = cuda - except (psutil.NoSuchProcess, psutil.AccessDenied): - break - self._stop.wait(self._interval_s) - - -def _reset_vram_peak() -> None: - try: - import torch - - if torch.cuda.is_available(): - torch.cuda.reset_peak_memory_stats() - except ImportError: - pass - - -def _read_vram_peak_gb(accelerator: str) -> float | None: - """Peak VRAM/GPU in GB. CUDA uses the native peak API; MPS uses the polled sampler value (caller-side).""" - try: - import torch - except ImportError: - return None - if accelerator == "cuda" and torch.cuda.is_available(): - return float(torch.cuda.max_memory_allocated()) / _BYTES_PER_GB - return None - - -# === per-module tracking ================================================= - - -# Static module_name → role classification. Used to tag tracker records so -# downstream analysis can decompose classic-preset wall-time into -# embedder-forward vs scorer-fit vs decision-search — the follow-up review's -# R4-P1 #30 asked for this because today classic wall-time conflates the two. -_EMBEDDER_MODULE_NAMES = frozenset( - {"sentence_transformer", "openai_embedder", "vllm_embedder", "hashing_vectorizer"}, -) -_DECISION_MODULE_NAMES = frozenset({"threshold", "argmax", "jinoos", "tunable", "adaptive"}) - - -def _classify_module_role(module_name: str) -> str: - """Bucket module_name into ``embedder`` / ``decision`` / ``scorer``. - - Everything not in the known embedder or decision sets is treated as a - scorer — so newly added scorer modules land in the right bucket by default - and only new decision/embedder modules would need to update the sets. - """ - if module_name in _EMBEDDER_MODULE_NAMES: - return "embedder" - if module_name in _DECISION_MODULE_NAMES: - return "decision" - return "scorer" - - -# Base class for _StepTimingCallback. We inherit from HF's real -# ``TrainerCallback`` when transformers is installed — that gives us the -# correct no-op default for every ``on_*`` hook (on_train_begin/on_log/ -# on_save/...) automatically, so we only override the two we time. HF's -# ``CallbackHandler.call_event`` dispatches with a bare ``getattr`` (no -# hasattr probe), so a plain class missing hooks would ``AttributeError`` -# the moment a real trial calls e.g. ``on_train_begin``. -# -# When transformers isn't installed we fall back to ``object`` so the -# harness still imports on classic-only runs. In that case the callback is -# never actually instantiated (``_patch_trainer_for_step_timing`` bails out -# in the same ``ImportError`` branch), so the fallback base is only needed -# to make the class definition itself succeed. -try: - from transformers import TrainerCallback as _StepTimingBase # type: ignore[import-not-found] -except ImportError: - _StepTimingBase = object # type: ignore[assignment,misc] - - -class _StepTimingCallback(_StepTimingBase): # type: ignore[misc,valid-type] - """HF ``TrainerCallback`` that appends the wall-time of each optimizer step - to a caller-owned list. - - Injected into every ``transformers.Trainer`` for the duration of a fit via - :func:`_patch_trainer_for_step_timing`. The sink is the current module's - step buffer on :class:`_ModuleTracker`, so the transformer's per-step - latency lands in that module's record automatically — no plumbing across - module boundaries. - """ - - def __init__(self, sink: list[float]) -> None: - # TrainerCallback.__init__ takes (*args, **kwargs); calling super is - # safe both when the base is the real HF class and when it's ``object``. - super().__init__() - self._sink = sink - self._t0: float | None = None - - def on_step_begin(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None: # noqa: ANN401, ARG002 - self._t0 = time.perf_counter() - - def on_step_end(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None: # noqa: ANN401, ARG002 - if self._t0 is not None: - self._sink.append(time.perf_counter() - self._t0) - self._t0 = None - - -def _summarize_step_times(step_times: list[float]) -> dict[str, float]: - """Fold a list of per-step wall-times into summary stats for the row. - - ``seconds_per_step`` is what the advisor's transformer-time baseline - encodes (currently a flat ~1 s constant); logging measured ``mean`` and - ``p95`` lets the baseline be recalibrated directly from row data instead - of eyeballed off a wandb dashboard. - """ - import statistics as _stats - - if not step_times: - return {} - if len(step_times) == 1: - return {"n_steps": 1, "mean_step_s": step_times[0], "p95_step_s": step_times[0]} - sorted_st = sorted(step_times) - p95_idx = min(len(sorted_st) - 1, int(round(0.95 * (len(sorted_st) - 1)))) - return { - "n_steps": len(step_times), - "mean_step_s": _stats.fmean(step_times), - "p95_step_s": sorted_st[p95_idx], - "total_step_s": sum(step_times), - } - - -def _patch_trainer_for_step_timing(tracker: _ModuleTracker) -> Any: # noqa: ANN401 - """Monkey-patch ``transformers.Trainer.__init__`` to inject a step-timing - callback bound to ``tracker._current_step_buffer`` — the list on the - module record currently being tracked. - - Returns a callable that undoes the patch. No-ops (returns a no-op undoer) - when transformers isn't importable, so classic-only presets aren't blocked. - """ - try: - from transformers import Trainer # type: ignore[import-not-found] - except ImportError: - return lambda: None - - original_init = Trainer.__init__ - - def patched(self: Any, *args: Any, **kwargs: Any) -> None: # noqa: ANN401 - original_init(self, *args, **kwargs) - buffer = tracker.current_step_buffer() - if buffer is not None: - self.add_callback(_StepTimingCallback(buffer)) - - Trainer.__init__ = patched # type: ignore[method-assign] - - def _undo() -> None: - Trainer.__init__ = original_init # type: ignore[method-assign] - - return _undo - - -class _ModuleTracker(OptimizerCallback): - """Records per-module wall time and peak VRAM. - - Hooks ``start_module`` / ``end_module`` on the CallbackHandler so we get - one record per (module_name, trial_num). CUDA peak VRAM is reset per module - via ``torch.cuda.reset_peak_memory_stats``; MPS is sampled at ``end_module`` - (no per-module peak API, so it's the moment-in-time allocation). - - Because per-module CUDA resets clobber the global ``max_memory_allocated`` - counter, the tracker also keeps ``self.peak_vram_gb_overall`` — the max - across every recorded module. The calibration script reads this instead of - the post-fit ``torch.cuda.max_memory_allocated()`` value, which by then - reflects only the last (usually CPU-only decision) module. - """ - - name = "calibration_tracker" - - def __init__(self) -> None: - self.records: list[dict[str, Any]] = [] - self._current: dict[str, Any] | None = None - self._current_step_buffer: list[float] | None = None - self.peak_vram_gb_overall: float = 0.0 - - def current_step_buffer(self) -> list[float] | None: - """Return the per-step wall-time list the ``_StepTimingCallback`` - should append to. ``None`` when no module is currently being tracked - (e.g. between modules) — the callback then skips.""" - return self._current_step_buffer - - def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None: - pass - - def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None: - try: - import torch - - if torch.cuda.is_available(): - torch.cuda.reset_peak_memory_stats() - except ImportError: - pass - # Only capture JSON-safe scalars in the config snapshot. - safe_config = {k: v for k, v in module_kwargs.items() if isinstance(v, (str, int, float, bool)) or v is None} - self._current_step_buffer = [] - self._current = { - "module": module_name, - "role": _classify_module_role(module_name), - "num": num, - "config": safe_config, - "_start": time.perf_counter(), - } - - def log_value(self, **kwargs: Any) -> None: # noqa: ANN401 - pass - - def log_metrics(self, metrics: dict[str, Any]) -> None: - pass - - def end_module(self) -> None: - if self._current is None: - return - rec = self._current - rec["duration_s"] = time.perf_counter() - rec.pop("_start") - try: - import torch - - if torch.cuda.is_available(): - rec["peak_vram_gb"] = float(torch.cuda.max_memory_allocated()) / _BYTES_PER_GB - elif torch.backends.mps.is_available(): - # MPS has no per-module peak API — snapshot the current allocation. - rec["peak_vram_gb"] = float(torch.mps.current_allocated_memory()) / _BYTES_PER_GB - except (ImportError, AttributeError): - pass - peak = rec.get("peak_vram_gb") - if peak is not None and peak > self.peak_vram_gb_overall: - self.peak_vram_gb_overall = peak - # Fold per-step timings into the module record so a transformer's - # trial exposes ``mean_step_s`` / ``p95_step_s`` next to its total - # duration — the advisor's flat 1 s/step baseline can then be - # recalibrated per device_class directly from row data. - step_times = self._current_step_buffer or [] - step_summary = _summarize_step_times(step_times) - if step_summary: - rec["step_timings"] = step_summary - self._current_step_buffer = None - self.records.append(rec) - self._current = None - - def end_run(self) -> None: - pass - - def log_final_metrics(self, metrics: dict[str, Any]) -> None: - pass - - -def _attach_callbacks(pipeline: Pipeline, callbacks: list[OptimizerCallback]) -> None: - """Instance-patch ``pipeline._fit`` to append ``callbacks`` to the callback chain.""" - original_fit = pipeline._fit # noqa: SLF001 - - def patched(context: Any) -> Any: # noqa: ANN401 - context.callback_handler.callbacks.extend(callbacks) - return original_fit(context) - - pipeline._fit = patched # type: ignore[method-assign] # noqa: SLF001 - - -# === preset resolution & optional-extras skip ============================ - - -# module_name → the ``autointent[extra]`` that must be installed for the -# module's __init__ to succeed. Sourced from ``require(...)`` calls in -# ``src/autointent/modules/scoring/`` — keep in sync. -_MODULE_TO_EXTRA: dict[str, str] = { - "bert": "transformers", - "catboost": "catboost", - "lora": "peft", - "ptuning": "peft", - "description_llm": "openai", -} - - -def _missing_extras_for_config(cfg: dict[str, Any]) -> list[str]: - """Return every optional extra that ``cfg``'s search_space needs but isn't installed. - - Uses the same ``_deps.require`` validator the modules use at runtime, so - what the harness pre-checks matches what would fail inside ``fit``. - A missing extra returns an ``ImportError``; anything else (e.g. unknown - extra) propagates. - """ - from autointent._deps import require # type: ignore[import-not-found] - - needed: set[str] = set() - for node in cfg.get("search_space") or []: - for entry in node.get("search_space") or []: - name = entry.get("module_name") if isinstance(entry, dict) else None - extra = _MODULE_TO_EXTRA.get(name) if isinstance(name, str) else None - if extra: - needed.add(extra) - - missing: list[str] = [] - for extra in sorted(needed): - try: - require(extra) # type: ignore[arg-type] - except ImportError: - missing.append(extra) - return missing - - -def _load_pipeline_from_preset_ref(ref: str) -> tuple[Pipeline, str]: - """Resolve ``ref`` as either a bundled preset name or a YAML file path. - - Returns ``(pipeline, display_name)`` where ``display_name`` is the bundled - name for name refs or the file stem for path refs. Path refs let the - harness exercise modules not in any bundled preset (LoRA / ptuning / - dnnc / gcn / cross-encoder scorer) without polluting the shipped - ``SearchSpacePreset`` literal. - """ - if ref.endswith((".yaml", ".yml")): - path = Path(ref).expanduser() - if not path.exists(): - raise FileNotFoundError(f"Preset file not found: {path}") - pipeline = Pipeline.from_optimization_config(path) - return pipeline, path.stem - return Pipeline.from_preset(ref), ref # type: ignore[arg-type] - - -def _run_cli_smoke( - preset_ref: str, - stats: Any, # noqa: ANN401 - budget_vram_gb: float | None, -) -> dict[str, Any]: - """Invoke ``autointent-advisor inspect --json`` in-process. - - Fed the same stats (as placeholder args) and budget the direct-API path - saw, so a divergence in ``is_feasible`` / predicted numbers points at the - CLI wrapper or the JSON renderer, not at differing inputs. - - Returns a dict with: - * ``payload`` — the parsed JSON body from the CLI (or ``None`` on crash) - * ``rc`` — the CLI return code - * ``error`` — traceback string when the CLI or JSON parse failed - * ``divergence`` — dict of |cli - direct| deltas populated by the caller - - Runs in-process (no subprocess) so we don't pay the interpreter-startup - cost on every preset — the review only asked for a wrapper smoke, not a - full subprocess isolation test. - """ - import contextlib - import io as _io - import traceback - - from autointent._advisor._cli import main as cli_main - - argv = [ - "inspect", - preset_ref, - "--n-samples", - str(int(stats.n_samples)), - "--n-classes", - str(int(stats.n_classes)), - "--avg-tokens", - str(int(stats.avg_tokens)), - "--task", - "multilabel" if getattr(stats, "multilabel", False) else "multiclass", - "--json", - ] - if budget_vram_gb is not None: - argv += ["--budget-vram-gb", str(budget_vram_gb)] - - buf = _io.StringIO() - result: dict[str, Any] = {"payload": None, "rc": None, "error": None} - try: - with contextlib.redirect_stdout(buf): - result["rc"] = cli_main(argv) - except Exception: # noqa: BLE001 - result["error"] = traceback.format_exc(limit=3) - return result - - raw = buf.getvalue().strip() - if not raw: - result["error"] = "CLI produced empty stdout" - return result - try: - result["payload"] = json.loads(raw) - except json.JSONDecodeError as e: - result["error"] = f"CLI --json output not parseable: {e}" - return result - - -def _load_config_from_preset_ref(ref: str) -> dict[str, Any]: - """Load the raw preset config dict without instantiating a Pipeline. - - Used for pre-fit extras checks so a preset whose modules would need a - missing extra (e.g. ``lora`` → ``peft``) never reaches ``from_preset``. - """ - if ref.endswith((".yaml", ".yml")): - import yaml - - with Path(ref).expanduser().open(encoding="utf-8") as f: - return yaml.safe_load(f) - from autointent.utils import load_preset # local import to avoid top-level cost - - return load_preset(ref) # type: ignore[arg-type] - - -# === per-preset run ====================================================== - - -def _override_trials(pipeline: Pipeline, max_trials: int | None, *, run_name: str | None = None) -> None: - """Cap n_trials, force ``n_jobs=1`` (serial HPO to keep wall-time measurements clean - and to prevent CPU oversubscription with sklearn's own ``n_jobs``), disable dumping. - When ``run_name`` is set, tag ``LoggingConfig.run_name`` (used as the W&B group / - dump-dir name). - """ - updates: dict[str, Any] = {"n_jobs": 1} - if max_trials is not None: - updates["n_trials"] = max_trials - pipeline.set_config(pipeline.hpo_config.model_copy(update=updates)) - logging_config = LoggingConfig(dump_modules=False, clear_ram=True, run_name=run_name) - pipeline.set_config(logging_config) - - -def _calibrate_one( - *, - preset: str, - dataset: Dataset, - stats: Any, # noqa: ANN401 - hardware: Any, # noqa: ANN401 - max_trials: int | None, - skip_fit: bool, - poll_interval_ms: int, - enable_wandb: bool, - run_name: str | None, - budget_vram_gb: float | None, - clear_embedding_cache: bool, -) -> CalibrationRow: - # ``preset`` may be a bundled name OR a path to a YAML file (the coverage - # preset). Resolve early so we can pre-check extras against the raw config - # before touching the Pipeline machinery. - try: - raw_cfg = _load_config_from_preset_ref(preset) - except Exception as e: # noqa: BLE001 - display_name = Path(preset).stem if preset.endswith((".yaml", ".yml")) else preset - row = CalibrationRow(preset=display_name) - row.error = f"load-preset failed: {e}" - return row - display_name = Path(preset).stem if preset.endswith((".yaml", ".yml")) else preset - - row = CalibrationRow(preset=display_name) - row.cache_policy = "cold" if clear_embedding_cache else "warm" - - # Detect missing optional extras BEFORE the fit — otherwise the trial - # would raise ImportError deep inside HPO, producing a fit-failed row - # indistinguishable from a real bug. - missing = _missing_extras_for_config(raw_cfg) - if missing: - row.skipped = True - row.error = f"skipped: missing extras {sorted(missing)}" - row.notes.append( - "install with: uv pip install " + " ".join(f"'autointent[{e}]'" for e in sorted(missing)) - ) - return row - - if clear_embedding_cache: - removed = _clear_embeddings_cache() - logger.info("Cleared %d embedding cache files for cold-cache measurement", removed) - - # === predicted ====================================================== - try: - pipeline, _ = _load_pipeline_from_preset_ref(preset) - except Exception as e: # noqa: BLE001 - row.error = f"from_preset failed: {e}" - return row - - tagged_run_name = f"{display_name}_{run_name}" if run_name else None - _override_trials(pipeline, max_trials, run_name=tagged_run_name) - - try: - # Optionally override hardware.vram_gb to exercise severity paths without a small GPU. - preflight_hw = hardware - if budget_vram_gb is not None: - from dataclasses import replace - - preflight_hw = replace(hardware, vram_gb=budget_vram_gb) - report: PreflightReport = run_preflight( - pipeline._build_advisor_config(), # noqa: SLF001 - stats, - preflight_hw, - ) - except Exception as e: # noqa: BLE001 - row.error = f"preflight failed: {e}" - return row - - row.predicted = { - "time_h": report.resource.time_hours, - "ram_gb": report.resource.ram_gb, - "vram_gb": report.resource.vram_gb, - "disk_download_gb": report.resource.disk_download_gb, - "disk_cached_gb": report.resource.disk_cached_gb, - "disk_embedding_cache_gb": report.resource.disk_embedding_cache_gb, - } - row.findings = len(report.findings) - row.findings_over = sum(1 for f in report.findings if f.severity.value == "over") - row.low_confidence = report.low_confidence - if report.low_confidence: - row.notes.append("low-confidence (heuristic HF metadata fallback in use)") - - # CLI wrapper smoke — same preset, same stats, same budget. Any divergence - # in ``is_feasible`` or the top-line predicted numbers means the CLI / - # JSON renderer drifted from the direct API. Runs unconditionally so a - # regression shows up on every calibration run. - smoke = _run_cli_smoke(preset, stats, budget_vram_gb) - if smoke["error"]: - row.notes.append(f"cli-smoke FAILED: {smoke['error'].splitlines()[-1] if smoke['error'] else '?'}") - elif smoke["payload"]: - cli_pred = smoke["payload"].get("resource") or {} - divergence: dict[str, float] = {} - for cli_key, direct_val in ( - ("time_hours", report.resource.time_hours), - ("ram_gb", report.resource.ram_gb), - ("vram_gb", report.resource.vram_gb), - ("disk_download_gb", report.resource.disk_download_gb), - ): - cli_val = cli_pred.get(cli_key) - if cli_val is None or direct_val is None: - continue - delta = abs(float(cli_val) - float(direct_val)) - if delta > 1e-6: - divergence[cli_key] = delta - smoke["divergence"] = divergence - cli_feasible = smoke["payload"].get("is_feasible") - if cli_feasible is not None and cli_feasible != report.is_feasible: - row.notes.append( - f"cli-smoke VERDICT MISMATCH: cli.is_feasible={cli_feasible} vs direct={report.is_feasible}" - ) - elif divergence: - row.notes.append(f"cli-smoke numeric drift on {sorted(divergence)} (see cli_smoke.divergence)") - row.cli_smoke = smoke - - if skip_fit: - return row - - # === actual ========================================================= - hf_cache = _hf_cache_dir() - embed_cache = _embeddings_cache_dir() - hf_before = _dir_size_gb(hf_cache) - embed_before = _dir_size_gb(embed_cache) - _reset_vram_peak() - - tracker = _ModuleTracker() - callbacks: list[OptimizerCallback] = [tracker] - if enable_wandb: - try: - from autointent._callbacks.wandb import WandbCallback - - callbacks.append(WandbCallback()) - except ImportError as e: - row.notes.append(f"W&B requested but not available: {e}") - _attach_callbacks(pipeline, callbacks) - - is_mps = hardware.accelerator == "mps" - is_cuda = hardware.accelerator == "cuda" - undo_step_patch = _patch_trainer_for_step_timing(tracker) - start = time.perf_counter() - try: - with _PeakSampler( - interval_s=poll_interval_ms / 1000.0, sample_mps=is_mps, sample_cuda=is_cuda, - ) as sampler: - pipeline.fit(dataset, preflight="off") - except Exception as e: # noqa: BLE001 - row.error = f"fit failed: {e}" - row.modules = tracker.records # keep whatever we collected - return row - finally: - undo_step_patch() - elapsed_s = time.perf_counter() - start - - hf_after = _dir_size_gb(hf_cache) - embed_after = _dir_size_gb(embed_cache) - actual_time_h = elapsed_s / 3600.0 - actual_ram_gb = sampler.peak_ram_gb - # VRAM: take max of per-module tracker (inside brackets) and background - # sampler (outside brackets, e.g. classic-preset embedder forward). - # Fallback to a raw peak read only if both are zero. - actual_vram_gb: float | None = None - tracker_peak = tracker.peak_vram_gb_overall if tracker.peak_vram_gb_overall > 0 else None - sampler_peak = sampler.peak_cuda_gb if sampler.peak_cuda_gb and sampler.peak_cuda_gb > 0 else None - candidates = [x for x in (tracker_peak, sampler_peak) if x is not None] - if candidates: - actual_vram_gb = max(candidates) - else: - actual_vram_gb = _read_vram_peak_gb(hardware.accelerator) - if actual_vram_gb is None and is_mps: - actual_vram_gb = sampler.peak_mps_gb - - row.actual = { - "time_h": actual_time_h, - "ram_gb": actual_ram_gb, - "vram_gb": actual_vram_gb, - "disk_download_gb": max(0.0, hf_after - hf_before), - "disk_embedding_cache_gb": max(0.0, embed_after - embed_before), - # Per-signal breakdown; classic presets expect sampler > tracker. - "vram_gb_tracker": tracker_peak, - "vram_gb_sampler": sampler_peak, - } - row.modules = tracker.records - if enable_wandb and not any("W&B requested but not available" in n for n in row.notes): - row.notes.append("W&B reporter attached — inspect wandb.ai run group for per-step GPU/system metrics") - return row - - -# === rendering =========================================================== - - -_COLS = [ - ("preset", "Preset", 22), - ("pred_time", "pred_time_h", 12), - ("act_time", "act_time_h", 12), - ("r_time", "ratio_t", 8), - ("pred_ram", "pred_ram_gb", 12), - ("act_ram", "act_ram_gb", 12), - ("r_ram", "ratio_r", 8), - ("pred_vram", "pred_vram_gb", 13), - ("act_vram", "act_vram_gb", 13), - ("r_vram", "ratio_v", 8), -] - - -def _fmt_cell(value: Any) -> str: # noqa: ANN401 - if value is None: - return "-" - if isinstance(value, float): - if value == 0: - return "0.00" - return f"{value:.2f}" if abs(value) >= 0.01 else f"{value:.4f}" - return str(value) - - -def _print_summary(rows: list[CalibrationRow]) -> None: - """Pretty side-by-side table for stdout.""" - header = " ".join(label.ljust(width) for _, label, width in _COLS) - print(header) - print("-" * len(header)) - _print_repeat_aggregates(rows) - for row in rows: - # Ratios are always computed at read time (see CalibrationRow.to_dict). - ratios = row._ratios() # noqa: SLF001 - cells = { - "preset": row.preset, - "pred_time": row.predicted.get("time_h"), - "act_time": row.actual.get("time_h"), - "r_time": ratios.get("time_h"), - "pred_ram": row.predicted.get("ram_gb"), - "act_ram": row.actual.get("ram_gb"), - "r_ram": ratios.get("ram_gb"), - "pred_vram": row.predicted.get("vram_gb"), - "act_vram": row.actual.get("vram_gb"), - "r_vram": ratios.get("vram_gb"), - } - print(" ".join(_fmt_cell(cells[key]).ljust(width) for key, _, width in _COLS)) - if row.error: - marker = "~" if row.skipped else "!" - print(f" {marker} {row.error}") - if row.low_confidence: - print(f" ! LOW-CONFIDENCE — advisor used heuristic HF metadata (exclude from prediction-accuracy stats)") - print(f" · cache-policy={row.cache_policy}") - role_totals = _sum_time_by_role(row.modules) - if role_totals: - breakdown = " ".join(f"{role}={total:.2f}s" for role, total in role_totals.items()) - print(f" · time-by-role: {breakdown}") - for note in row.notes: - print(f" * {note}") - for mod in row.modules: - duration = mod.get("duration_s") - vram = mod.get("peak_vram_gb") - role = mod.get("role", "?") - duration_s = f"{duration:.2f}s" if duration is not None else "-" - vram_s = f"{vram:.2f} GB" if vram is not None else "-" - line = ( - f" · [{role}] {mod.get('module', '?')}#{mod.get('num', '?')} {duration_s} vram={vram_s}" - ) - step = mod.get("step_timings") - if step: - line += ( - f" n_steps={step['n_steps']} mean_step_s={step['mean_step_s']:.3f} " - f"p95_step_s={step['p95_step_s']:.3f}" - ) - print(line) - - -def _sum_time_by_role(modules: list[dict[str, Any]]) -> dict[str, float]: - """Fold per-module durations into ``{role: total_seconds}`` — used both for - the printed breakdown and for the top-level ``time_by_role`` row field.""" - totals: dict[str, float] = {} - for mod in modules: - role = mod.get("role", "?") - duration = mod.get("duration_s") - if duration is None: - continue - totals[role] = totals.get(role, 0.0) + float(duration) - return totals - - -def _print_repeat_aggregates(rows: list[CalibrationRow]) -> None: - """When any (preset, dataset) has more than one repeat, print a mean±stdev - block up-front so small ratio gaps are judgeable at a glance. - - Groups by ``(preset, first-note)`` — the dataset marker is inserted as the - first note in main() so this key is stable across repeats. - """ - import statistics - - groups: dict[tuple[str, str], list[CalibrationRow]] = {} - for row in rows: - dataset_note = row.notes[0] if row.notes else "dataset=?" - groups.setdefault((row.preset, dataset_note), []).append(row) - - multi_groups = [(k, v) for k, v in groups.items() if len(v) > 1] - if not multi_groups: - return - print(">>> repeats aggregation (mean ± stdev, successful runs only):") - for (preset, dataset_note), group in multi_groups: - # Skipped rows are expected — separate them from real failures so the - # aggregate isn't polluted by "all repeats failed" when the actual - # cause is a missing optional extra. - skipped = [r for r in group if r.skipped] - successful = [r for r in group if r.error is None] - real_failures = len(group) - len(successful) - len(skipped) - n = len(successful) - if n == 0: - reason = f"{real_failures} failed" - if skipped: - reason += f", {len(skipped)} skipped" - print(f" {preset} [{dataset_note}] {reason} (no successful repeats)") - continue - parts = [f" {preset} [{dataset_note}] n={n}"] - for metric in ("time_h", "ram_gb", "vram_gb"): - values = [r.actual.get(metric) for r in successful if r.actual.get(metric) is not None] - if not values: - continue - mean = statistics.fmean(values) - stdev = statistics.stdev(values) if len(values) > 1 else 0.0 - parts.append(f"{metric}={mean:.2f}±{stdev:.2f}") - print(" " + " ".join(parts)) - print() - - -def _apply_thread_cap() -> None: - """Cap torch intra-op threads to the same value as OMP_NUM_THREADS. - - Env vars (OMP/MKL/OpenBLAS) MUST be set before Python starts to be effective — - that's the bash wrapper's job. This function is belt-and-braces: torch reads - OMP_NUM_THREADS on init, but ``set_num_threads`` also caps its C++ intra-op - pool if a caller forgets the env var. - """ - n = int(os.environ.get("OMP_NUM_THREADS", "0") or 0) - if n <= 0: - return - try: - import torch - - torch.set_num_threads(n) - except ImportError: - pass - - -def _guard_cuda_init(*, required: bool) -> None: - """When ``required`` is True, fail fast if PyTorch can't initialize CUDA. - - Guards against the silent 'nvidia-smi shows 2 GPUs but torch runs on CPU' - trap that happens when the CUDA driver is older than what the installed - torch wheel was built against. - """ - if not required: - return - try: - import torch - except ImportError: - msg = "--require-cuda passed but torch isn't installed" - raise SystemExit(msg) from None - if not torch.cuda.is_available(): - # Get the underlying reason if we can — usually a warning on import time. - msg = ( - "--require-cuda passed but torch.cuda.is_available() is False. " - "Check `nvidia-smi` vs `python -c 'import torch; print(torch.version.cuda)'` — " - "you likely need a torch wheel built against a matching CUDA runtime." - ) - raise SystemExit(msg) - - -def _load_dataset(dataset_arg: str, parser: argparse.ArgumentParser) -> tuple[Dataset, str]: - """Load one dataset from a local JSON path or an HF Hub repo id, returning - ``(dataset, source_label)`` — the label mirrors what the calibrator writes - to the report so different sources are distinguishable in aggregate output. - """ - dataset_path = Path(dataset_arg) - if dataset_path.is_file(): - logger.info("Loading dataset from local file %s", dataset_path) - return Dataset.from_json(dataset_path), str(dataset_path) - logger.info("Loading dataset from HF Hub: %s", dataset_arg) - try: - dataset = Dataset.from_hub(dataset_arg) - except Exception as e: # noqa: BLE001 - parser.error(f"Could not load '{dataset_arg}' as a local JSON file or as a Hub repo id: {e}") - return dataset, f"hub:{dataset_arg}" - - -def _subsample_per_class(dataset: Dataset, cap: int) -> Dataset: - """Cap each class in the train split to at most ``cap`` samples (first-N slice). - - Uses a deterministic first-N slice per class — reproducible across runs - without seeding, and keeps class-ordering intuitive when inspecting the - subset. Only rewrites the train split; validation/test are left as-is so - the metric baselines remain comparable. - """ - from autointent.custom_types import Split - - train_key = Split.TRAIN if Split.TRAIN in dataset else next( - (k for k in dataset if str(k).startswith(str(Split.TRAIN))), None, - ) - if train_key is None: - return dataset - train = dataset[train_key] - label_feature = dataset.label_feature - seen: dict[Any, int] = {} - keep: list[int] = [] - for idx, row in enumerate(train): - label = row[label_feature] - # For multilabel, key on the tuple so a sample with a rare-class tag - # still contributes toward that class's cap. - key = tuple(label) if isinstance(label, list) else label - count = seen.get(key, 0) - if count < cap: - keep.append(idx) - seen[key] = count + 1 - dataset[train_key] = train.select(keep) - return dataset - - -def main(argv: list[str] | None = None) -> int: - parser = _build_parser() - args = parser.parse_args(argv) - logging.basicConfig( - level=logging.DEBUG if args.verbose else logging.INFO, - format="%(levelname)s %(name)s: %(message)s", - ) - _apply_thread_cap() - _guard_cuda_init(required=args.require_cuda) - - presets = args.presets or list(BUNDLED_PRESETS) - unknown = [ - p - for p in presets - if not p.endswith((".yaml", ".yml")) and p not in BUNDLED_PRESETS - ] - if unknown: - parser.error( - f"Unknown preset(s): {', '.join(unknown)}. Known: {', '.join(BUNDLED_PRESETS)}, " - "or pass a path to a .yaml file (e.g. scripts/coverage_preset.yaml)." - ) - for p in presets: - if p.endswith((".yaml", ".yml")) and not Path(p).expanduser().exists(): - parser.error(f"Preset file not found: {p}") - - hardware = detect_hardware() - logger.info( - "Hardware: %s (%s) — %.1f GB VRAM, %.0f GB RAM, %.0f GB free disk", - hardware.accelerator, - hardware.device_name, - hardware.vram_gb, - hardware.ram_gb, - hardware.free_disk_gb, - ) - logger.info( - "Thread caps: OMP=%s MKL=%s OPENBLAS=%s TOKENIZERS_PARALLELISM=%s", - os.environ.get("OMP_NUM_THREADS", ""), - os.environ.get("MKL_NUM_THREADS", ""), - os.environ.get("OPENBLAS_NUM_THREADS", ""), - os.environ.get("TOKENIZERS_PARALLELISM", ""), - ) - - rows: list[CalibrationRow] = [] - datasets_meta: list[dict[str, Any]] = [] - - def _write_payload() -> None: - """Serialize the current in-memory rows to ``args.output``. Called - after each preset finishes so a mid-sweep crash / Broken pipe leaves - a valid partial report behind rather than losing everything. - """ - payload_now = { - "hardware": { - "accelerator": hardware.accelerator, - "device_name": hardware.device_name, - "vram_gb": hardware.vram_gb, - "ram_gb": hardware.ram_gb, - "free_disk_gb": hardware.free_disk_gb, - }, - "datasets": datasets_meta, - "max_trials_override": args.max_trials, - "skip_fit": args.skip_fit, - "cache_policy": "cold" if args.clear_embedding_cache else "warm", - "budget_vram_gb_override": args.budget_vram_gb, - "subsample_per_class": args.subsample_per_class, - "thread_caps": { - "OMP_NUM_THREADS": os.environ.get("OMP_NUM_THREADS"), - "MKL_NUM_THREADS": os.environ.get("MKL_NUM_THREADS"), - "OPENBLAS_NUM_THREADS": os.environ.get("OPENBLAS_NUM_THREADS"), - "TOKENIZERS_PARALLELISM": os.environ.get("TOKENIZERS_PARALLELISM"), - }, - "in_progress": True, - "rows": [r.to_dict() for r in rows], - } - # Atomic write: dump to a sibling file, then rename. Prevents readers - # from seeing a half-written JSON if the run is killed mid-serialize. - tmp = args.output.with_suffix(args.output.suffix + ".partial") - tmp.write_text(json.dumps(payload_now, indent=2, default=str)) - tmp.replace(args.output) - - for dataset_arg in args.dataset: - dataset, dataset_source = _load_dataset(dataset_arg, parser) - if args.subsample_per_class is not None: - dataset = _subsample_per_class(dataset, args.subsample_per_class) - dataset_source += f"|subsample-per-class={args.subsample_per_class}" - stats = stats_from_dataset_obj(dataset) - datasets_meta.append( - { - "path": dataset_source, - "n_samples": stats.n_samples, - "n_classes": stats.n_classes, - "avg_tokens": stats.avg_tokens, - "multilabel": stats.multilabel, - }, - ) - logger.info( - "Dataset %s: n_samples=%d n_classes=%d avg_tokens=%.1f multilabel=%s", - dataset_source, - stats.n_samples, - stats.n_classes, - stats.avg_tokens, - stats.multilabel, - ) - for preset in presets: - for repeat_idx in range(max(1, args.repeats)): - header = f"=== {preset} @ {dataset_source}" - if args.repeats > 1: - header += f" (repeat {repeat_idx + 1}/{args.repeats})" - header += " ===" - logger.info(header) - row = _calibrate_one( - preset=preset, - dataset=dataset, - stats=stats, - hardware=hardware, - max_trials=args.max_trials, - skip_fit=args.skip_fit, - poll_interval_ms=args.poll_interval_ms, - enable_wandb=args.wandb, - run_name=( - f"{args.run_name}_r{repeat_idx}" if args.run_name and args.repeats > 1 else args.run_name - ), - budget_vram_gb=args.budget_vram_gb, - clear_embedding_cache=args.clear_embedding_cache, - ) - row.repeat_idx = repeat_idx - row.notes.insert(0, f"dataset={dataset_source}") - rows.append(row) - _write_payload() - - payload = { - "hardware": { - "accelerator": hardware.accelerator, - "device_name": hardware.device_name, - "vram_gb": hardware.vram_gb, - "ram_gb": hardware.ram_gb, - "free_disk_gb": hardware.free_disk_gb, - }, - "datasets": datasets_meta, - "max_trials_override": args.max_trials, - "skip_fit": args.skip_fit, - "cache_policy": "cold" if args.clear_embedding_cache else "warm", - "budget_vram_gb_override": args.budget_vram_gb, - "subsample_per_class": args.subsample_per_class, - "thread_caps": { - "OMP_NUM_THREADS": os.environ.get("OMP_NUM_THREADS"), - "MKL_NUM_THREADS": os.environ.get("MKL_NUM_THREADS"), - "OPENBLAS_NUM_THREADS": os.environ.get("OPENBLAS_NUM_THREADS"), - "TOKENIZERS_PARALLELISM": os.environ.get("TOKENIZERS_PARALLELISM"), - }, - "in_progress": False, - "rows": [r.to_dict() for r in rows], - } - args.output.write_text(json.dumps(payload, indent=2, default=str)) - logger.info("Wrote report to %s", args.output) - - print() - _print_summary(rows) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/scripts/coverage_preset.yaml b/scripts/coverage_preset.yaml deleted file mode 100644 index 02844aa50..000000000 --- a/scripts/coverage_preset.yaml +++ /dev/null @@ -1,42 +0,0 @@ -## Coverage-only preset for the calibration harness. -# -# Bundled presets don't touch lora / ptuning / dnnc / gcn / a plain -# cross-encoder scorer, so the advisor's estimates for those modules go -# unvalidated. This preset packs one of each into a single small run so -# `scripts/calibrate_advisor.py` can exercise them against a real fit. -# -# Intentionally cheap (n_trials: 1 per scorer, single decision module) — -# the point is coverage, not tuning quality. Skip cleanly when the peft -# extra is missing (lora/ptuning will be marked skipped by the harness). -search_space: - - node_type: scoring - target_metric: scoring_f1 - search_space: - - module_name: lora - classification_model_config: - - model_name: microsoft/deberta-v3-small - num_train_epochs: [1] - batch_size: [16] - learning_rate: [5.0e-5] - - module_name: ptuning - classification_model_config: - - model_name: microsoft/deberta-v3-small - num_train_epochs: [1] - batch_size: [16] - learning_rate: [5.0e-5] - num_virtual_tokens: [8] - - module_name: dnnc - k: [3] - - module_name: gcn - num_train_epochs: [1] - batch_size: [16] - learning_rate: [1.0e-3] - - module_name: description_cross - - node_type: decision - target_metric: decision_accuracy - search_space: - - module_name: argmax -hpo_config: - sampler: tpe - n_trials: 5 - n_startup_trials: 2 diff --git a/scripts/run_calibration_banking77.sh b/scripts/run_calibration_banking77.sh deleted file mode 100755 index 1c6f0d869..000000000 --- a/scripts/run_calibration_banking77.sh +++ /dev/null @@ -1,163 +0,0 @@ -#!/usr/bin/env bash -# Run the advisor calibration across every bundled preset on DeepPavlov/banking77. -# -# WARNING: transformers-heavy on banking77 (10k train samples, 77 classes) can -# take *many* hours on a single GPU. Set MAX_TRIALS to a small number for a -# fast sanity check, or leave it unset to let each preset use its bundled -# ``hpo_config.n_trials``. -# -# Environment overrides: -# DATASET HF Hub repo id (default: DeepPavlov/banking77) -# DATASETS Space-separated list of dataset ids (default: unset -> use -# DATASET). Every preset runs against every dataset — useful -# to sweep a small + large + multilabel + long-token shape in -# a single invocation. -# SUBSAMPLE_PER_CLASS Cap each class to N training samples (deterministic -# first-N slice) before running — turns banking77 into a -# small dataset without needing a separate corpus. -# REPEATS Run each (preset, dataset) N times. The summary prints -# mean ± stdev of actual measurements across repeats so -# small ratio gaps become judgeable. Default: 1. -# PRESETS Space-separated preset names OR paths (default: every bundled preset). -# Items ending in .yaml/.yml are treated as file paths — use -# scripts/coverage_preset.yaml to exercise lora/ptuning/dnnc/gcn/ -# cross-encoder in one small run without touching bundled presets. -# MAX_TRIALS Cap for hpo_config.n_trials (default: unset -> preset default) -# WANDB If non-empty, pass --wandb so system metrics land in wandb.ai -# RUN_NAME Suffix appended to each preset's LoggingConfig.run_name — the -# resulting name is ``{preset}_{RUN_NAME}`` (default: unset -> -# autointent generates a random name) -# OUTPUT_DIR Where JSON reports + logs land (default: ./calibration_runs) -# SKIP_FIT If non-empty, only run preflight (no real fit) — fast sanity check -# THREADS_PER_JOB Cap for BLAS/OpenMP/torch intra-op threads per HPO trial -# (default: 1). Increase carefully — sklearn's own ``n_jobs`` and -# HPO parallelism multiply on top, so oversubscription is easy -# on many-core boxes. -# COLD If non-empty, pass --clear-embedding-cache so each preset -# starts with an empty embeddings cache (measure COLD cost). -# BUDGET_VRAM_GB Force run_preflight to see a specific VRAM budget instead -# of what the box exposes — lets you exercise severity paths -# (red/yellow/green + findings_over) on a big box. -# REQUIRE_CUDA If non-empty, pass --require-cuda so the run fails fast when -# torch.cuda.is_available() is False (guards against silent -# CPU-fallback caused by CUDA-driver / torch-wheel mismatch). -# -# Examples: -# scripts/run_calibration_banking77.sh # full sweep, serial -# MAX_TRIALS=3 scripts/run_calibration_banking77.sh # quick sweep -# PRESETS="classic-light nn-medium" scripts/run_calibration_banking77.sh -# WANDB=1 MAX_TRIALS=5 scripts/run_calibration_banking77.sh -# RUN_NAME=calib_2026_07 WANDB=1 scripts/run_calibration_banking77.sh -# THREADS_PER_JOB=4 scripts/run_calibration_banking77.sh # 4-thread BLAS -# SUBSAMPLE_PER_CLASS=5 scripts/run_calibration_banking77.sh # small-dataset shape -# DATASETS="DeepPavlov/banking77 DeepPavlov/clinc150" scripts/run_calibration_banking77.sh -# REPEATS=3 MAX_TRIALS=3 scripts/run_calibration_banking77.sh # variance bars -# PRESETS="scripts/coverage_preset.yaml" MAX_TRIALS=1 scripts/run_calibration_banking77.sh -# # exercise lora/ptuning/dnnc/gcn/description_cross — modules no bundled preset touches - -set -euo pipefail - -# Resolve repo root even when the script is called from anywhere. -REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -cd "$REPO_ROOT" - -DATASET="${DATASET:-DeepPavlov/banking77}" -# DATASETS overrides DATASET when set; enables multi-dataset sweeps. -if [[ -n "${DATASETS:-}" ]]; then - # shellcheck disable=SC2206 # intentional word-split from env - DATASET_ARR=($DATASETS) -else - DATASET_ARR=("$DATASET") -fi -OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/calibration_runs}" -TIMESTAMP="$(date +%Y%m%d_%H%M%S)" -OUTPUT_JSON="$OUTPUT_DIR/banking77_$TIMESTAMP.json" -LOG_FILE="$OUTPUT_DIR/banking77_$TIMESTAMP.log" - -export WANDB_PROJECT="autointent_feasibility" - -# --------------------------------------------------------------------------- -# CPU thread caps — set BEFORE python starts, because numpy/torch/sklearn read -# them at import time. Without these, on a 16+ core box each BLAS-backed -# operation defaults to N-thread pools which multiply with sklearn's own -# ``n_jobs`` and HPO parallelism → the machine oversubscribes and stalls. -# --------------------------------------------------------------------------- -THREADS_PER_JOB="${THREADS_PER_JOB:-1}" -export OMP_NUM_THREADS="$THREADS_PER_JOB" -export MKL_NUM_THREADS="$THREADS_PER_JOB" -export OPENBLAS_NUM_THREADS="$THREADS_PER_JOB" -export NUMEXPR_NUM_THREADS="$THREADS_PER_JOB" -# HF tokenizers deadlock on fork if left in parallel mode. -export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}" -# torch reads OMP_NUM_THREADS for intra-op, but set explicitly too — belt-and-braces. -export PYTORCH_NUM_THREADS="$THREADS_PER_JOB" - -mkdir -p "$OUTPUT_DIR" - -# Assemble optional flags. -EXTRA_FLAGS=() -if [[ -n "${MAX_TRIALS:-}" ]]; then - EXTRA_FLAGS+=("--max-trials" "$MAX_TRIALS") -fi -if [[ -n "${WANDB:-}" ]]; then - EXTRA_FLAGS+=("--wandb") -fi -if [[ -n "${SKIP_FIT:-}" ]]; then - EXTRA_FLAGS+=("--skip-fit") -fi -if [[ -n "${RUN_NAME:-}" ]]; then - EXTRA_FLAGS+=("--run-name" "$RUN_NAME") -fi -if [[ -n "${COLD:-}" ]]; then - EXTRA_FLAGS+=("--clear-embedding-cache") -fi -if [[ -n "${BUDGET_VRAM_GB:-}" ]]; then - EXTRA_FLAGS+=("--budget-vram-gb" "$BUDGET_VRAM_GB") -fi -if [[ -n "${REQUIRE_CUDA:-}" ]]; then - EXTRA_FLAGS+=("--require-cuda") -fi -if [[ -n "${SUBSAMPLE_PER_CLASS:-}" ]]; then - EXTRA_FLAGS+=("--subsample-per-class" "$SUBSAMPLE_PER_CLASS") -fi -if [[ -n "${REPEATS:-}" ]]; then - EXTRA_FLAGS+=("--repeats" "$REPEATS") -fi - -# Preset list: pull it from the advisor package at runtime unless overridden, -# so the script auto-discovers presets that are added later. -if [[ -n "${PRESETS:-}" ]]; then - # shellcheck disable=SC2206 # intentional word-split from env - PRESET_ARR=($PRESETS) -else - PRESET_ARR=() - while IFS= read -r preset; do - PRESET_ARR+=("$preset") - done < <( -python - <<'PY' -from autointent._advisor import BUNDLED_PRESETS -for name in BUNDLED_PRESETS: - print(name) -PY - ) -fi - -echo "Repo: $REPO_ROOT" -echo "Datasets: ${DATASET_ARR[*]}" -echo "Presets: ${PRESET_ARR[*]}" -echo "Output: $OUTPUT_JSON" -echo "Log: $LOG_FILE" -echo "Flags: ${EXTRA_FLAGS[*]:-}" -echo "Threads per job: $THREADS_PER_JOB (OMP/MKL/OpenBLAS/torch)" -echo - -uv run --no-sync python scripts/calibrate_advisor.py \ - --dataset "${DATASET_ARR[@]}" \ - --presets "${PRESET_ARR[@]}" \ - --output "$OUTPUT_JSON" \ - "${EXTRA_FLAGS[@]}" \ - 2>&1 | tee "$LOG_FILE" - -echo -echo "Done. JSON: $OUTPUT_JSON" -echo " Log: $LOG_FILE" diff --git a/tests/pipeline/test_calibration_tracker.py b/tests/pipeline/test_calibration_tracker.py deleted file mode 100644 index d538f6f6f..000000000 --- a/tests/pipeline/test_calibration_tracker.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Tests for the calibration script's _ModuleTracker callback.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -# Add scripts/ to sys.path so the test can import calibrate_advisor. -_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "scripts" -sys.path.insert(0, str(_SCRIPTS_DIR)) - -from calibrate_advisor import ( # noqa: E402 - _ModuleTracker, - _StepTimingCallback, - _classify_module_role, - _sum_time_by_role, -) - - -def test_tracker_records_wall_time_per_module() -> None: - """One (module, num) → one record with a positive duration.""" - tracker = _ModuleTracker() - - tracker.start_module("linear", 0, {"cv": 3}) - tracker.end_module() - - tracker.start_module("catboost", 1, {"iterations": 100, "depth": 6}) - tracker.end_module() - - assert len(tracker.records) == 2 - assert tracker.records[0]["module"] == "linear" - assert tracker.records[0]["num"] == 0 - assert tracker.records[0]["config"] == {"cv": 3} - assert tracker.records[0]["duration_s"] >= 0 - assert tracker.records[1]["module"] == "catboost" - assert tracker.records[1]["config"] == {"iterations": 100, "depth": 6} - - -def test_tracker_filters_non_scalar_config_values() -> None: - """Complex objects in module_kwargs must not appear in the recorded config.""" - tracker = _ModuleTracker() - tracker.start_module( - "bert", - 0, - {"cv": 3, "classification_model_config": {"model_name": "microsoft/deberta"}, "flag": True}, - ) - tracker.end_module() - assert tracker.records[0]["config"] == {"cv": 3, "flag": True} - - -def test_end_module_without_start_is_noop() -> None: - """Defensive: no crash when end_module is called without a matching start.""" - tracker = _ModuleTracker() - tracker.end_module() # must not raise - assert tracker.records == [] - - -def test_records_are_json_serialisable() -> None: - """Records must survive round-tripping through json.dumps for the CalibrationRow output.""" - import json - - tracker = _ModuleTracker() - tracker.start_module("linear", 0, {"cv": 3, "unused_none": None}) - tracker.end_module() - payload = json.dumps(tracker.records) - assert "linear" in payload - assert "duration_s" in payload - - -def test_running_peak_survives_low_last_module() -> None: - """peak_vram_gb_overall must retain the max across all modules. - - Pins the fix for a bug where torch.cuda's per-module reset_peak_memory_stats - clobbered the top-level VRAM reading with the last (usually CPU-only) module. - """ - tracker = _ModuleTracker() - - # Simulate a big embedder module. Populate _current directly to sidestep - # the real torch.cuda call and inject a synthetic peak. - tracker.start_module("linear", 0, {"cv": 3}) - tracker.end_module() - tracker.records[-1]["peak_vram_gb"] = 2.5 - tracker.peak_vram_gb_overall = max(tracker.peak_vram_gb_overall, 2.5) - - # Then a decision module that touches no VRAM. - tracker.start_module("threshold", 1, {"thresh": 0.5}) - tracker.end_module() - tracker.records[-1]["peak_vram_gb"] = 0.01 - tracker.peak_vram_gb_overall = max(tracker.peak_vram_gb_overall, 0.01) - - assert tracker.peak_vram_gb_overall == 2.5, "running peak clobbered by later small module" - - -def test_role_classification_and_time_decomposition() -> None: - """Each record carries a role, and _sum_time_by_role folds durations correctly. - - Pins R4-P1 #30: classic-preset wall-time must be decomposable into - embedder-forward vs scorer-fit vs decision-search so a consumer can - validate the advisor's per-role predictions without post-hoc classification. - """ - assert _classify_module_role("sentence_transformer") == "embedder" - assert _classify_module_role("hashing_vectorizer") == "embedder" - assert _classify_module_role("linear") == "scorer" - assert _classify_module_role("bert") == "scorer" - assert _classify_module_role("threshold") == "decision" - assert _classify_module_role("argmax") == "decision" - - tracker = _ModuleTracker() - tracker.start_module("sentence_transformer", 0, {}) - tracker.end_module() - tracker.records[-1]["duration_s"] = 8.0 - tracker.start_module("linear", 1, {}) - tracker.end_module() - tracker.records[-1]["duration_s"] = 2.0 - tracker.start_module("threshold", 2, {}) - tracker.end_module() - tracker.records[-1]["duration_s"] = 0.5 - - assert [r["role"] for r in tracker.records] == ["embedder", "scorer", "decision"] - totals = _sum_time_by_role(tracker.records) - assert totals == {"embedder": 8.0, "scorer": 2.0, "decision": 0.5} - - -def test_step_timing_callback_answers_every_hf_hook() -> None: - """HF dispatches every hook via bare getattr — callback must be a real - TrainerCallback subclass so all defaults are inherited as no-ops. - Regression: AttributeError on on_train_begin during a bert trial.""" - from transformers import TrainerCallback - - cb = _StepTimingCallback(sink=[]) - assert isinstance(cb, TrainerCallback) - for name in ["on_init_end", "on_train_begin", "on_epoch_begin", "on_log", "on_save", "on_train_end"]: - hook = getattr(cb, name) - assert hook(None, None, "control", model=None) is None, name - - -def test_peak_sampler_records_cuda_current_allocation() -> None: - """CUDA polling API: sample_cuda=True → 0.0 float, False → None. Thread - starts and stops cleanly with no CUDA hardware present.""" - from calibrate_advisor import _PeakSampler - - s = _PeakSampler(sample_cuda=True) - assert s.peak_cuda_gb == 0.0 - assert _PeakSampler(sample_cuda=False).peak_cuda_gb is None - import time - with s: - time.sleep(0.15) - assert s.peak_cuda_gb == 0.0 From 3cb6ca50bcb4f61cd283931817ec8ac93cb706e2 Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 17:59:52 +0300 Subject: [PATCH 31/43] fix: drop dangling reference to deleted proposal doc in advisor docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compute-feasibility-advisor-proposal.md was removed in the harness-relocation commit but the advisor package's module docstring still pointed to it. Strip the sentence rather than repoint it — a later task rewrites this docstring wholesale. --- src/autointent/_advisor/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py index 4df486f2b..ecd6d9b5e 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/_advisor/__init__.py @@ -1,8 +1,7 @@ """Pre-flight compute feasibility advisor. Exposes a small surface used by both ``Pipeline.fit()`` (see the ``preflight=`` -kwarg) and the ``autointent-advisor`` CLI script. See -``compute-feasibility-advisor-proposal.md`` at the repo root for the design. +kwarg) and the ``autointent-advisor`` CLI script. """ from __future__ import annotations From 5444a3870313ae9ec94870b0eb45cebf4926d18e Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 18:14:53 +0300 Subject: [PATCH 32/43] refactor: promote _advisor to public autointent.advisor package Pipeline.fit's docstring already pointed users at autointent._advisor.run_preflight, so the only usable entry point was behind a leading underscore. Promote the package, keep internals private (runner.py -> _runner.py, workflows.py -> _workflows.py), and rename the console script to match the prog= the CLI already reports. --- pyproject.toml | 2 +- src/autointent/_pipeline/_pipeline.py | 6 +++--- .../{_advisor => advisor}/__init__.py | 4 ++-- src/autointent/{_advisor => advisor}/_cli.py | 6 +++--- .../_estimates/__init__.py | 0 .../_estimates/_formulas.py | 6 +++--- .../_estimates/_resource.py | 10 +++++----- .../_estimates/_search_space.py | 0 .../{_advisor => advisor}/_hardware.py | 0 src/autointent/{_advisor => advisor}/_hub.py | 0 .../{_advisor => advisor}/_render.py | 0 .../{_advisor => advisor}/_report.py | 0 .../runner.py => advisor/_runner.py} | 10 +++++----- .../workflows.py => advisor/_workflows.py} | 2 +- src/autointent/custom_types/_types.py | 2 +- tests/advisor/test_estimates_and_cli.py | 10 +++++----- tests/advisor/test_estimates_internals.py | 20 +++++++++---------- tests/advisor/test_hardware_detection.py | 20 +++++++++---------- tests/advisor/test_hub_heuristics.py | 2 +- tests/advisor/test_reduce_to_fit.py | 6 +++--- tests/advisor/test_render.py | 4 ++-- tests/advisor/test_report.py | 2 +- tests/pipeline/test_preflight.py | 2 +- 23 files changed, 57 insertions(+), 57 deletions(-) rename src/autointent/{_advisor => advisor}/__init__.py (93%) rename src/autointent/{_advisor => advisor}/_cli.py (95%) rename src/autointent/{_advisor => advisor}/_estimates/__init__.py (100%) rename src/autointent/{_advisor => advisor}/_estimates/_formulas.py (98%) rename src/autointent/{_advisor => advisor}/_estimates/_resource.py (99%) rename src/autointent/{_advisor => advisor}/_estimates/_search_space.py (100%) rename src/autointent/{_advisor => advisor}/_hardware.py (100%) rename src/autointent/{_advisor => advisor}/_hub.py (100%) rename src/autointent/{_advisor => advisor}/_render.py (100%) rename src/autointent/{_advisor => advisor}/_report.py (100%) rename src/autointent/{_advisor/runner.py => advisor/_runner.py} (95%) rename src/autointent/{_advisor/workflows.py => advisor/_workflows.py} (99%) diff --git a/pyproject.toml b/pyproject.toml index 3bc17dc00..0466cf4a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,7 +147,7 @@ Documentation = "https://deeppavlov.github.io/AutoIntent/" [project.scripts] "basic-aug" = "autointent.generation.utterances._basic.cli:main" "evolution-aug" = "autointent.generation.utterances._evolution.cli:main" -"advisor" = "autointent._advisor._cli:main" +"autointent-advisor" = "autointent.advisor._cli:main" [build-system] requires = ["uv_build>=0.8.7,<0.9.0"] diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index 491bc57f4..0249649dd 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -12,7 +12,7 @@ from typing_extensions import assert_never from autointent import Context, OptimizationConfig -from autointent._advisor import ( +from autointent.advisor import ( Severity, detect_hardware, run_preflight, @@ -40,7 +40,7 @@ if TYPE_CHECKING: from autointent import Dataset - from autointent._advisor import PreflightReport + from autointent.advisor import PreflightReport from autointent.custom_types import ListOfGenericLabels, SearchSpacePreset, SearchSpaceValidationMode from autointent.modules.base import BaseDecision, BaseRegex, BaseScorer @@ -255,7 +255,7 @@ def fit( refit_after: whether to refit on whole data after optimization. Valid only for hold-out validaiton. sampler: sampler type to use. incompatible_search_space: wow to handle data-incompatible modules occurring in search space. - preflight: gate that runs :func:`autointent._advisor.run_preflight` over the + preflight: gate that runs :func:`autointent.advisor.run_preflight` over the pipeline's effective config + dataset before any heavy work. ``"off"`` skips it. ``"warn"`` (default) logs findings — INFO for AMPLE, WARNING for TIGHT, ERROR for OVER — but never raises. diff --git a/src/autointent/_advisor/__init__.py b/src/autointent/advisor/__init__.py similarity index 93% rename from src/autointent/_advisor/__init__.py rename to src/autointent/advisor/__init__.py index ecd6d9b5e..a7d9d1b87 100644 --- a/src/autointent/_advisor/__init__.py +++ b/src/autointent/advisor/__init__.py @@ -8,8 +8,8 @@ from ._hardware import HardwareProfile, detect_hardware from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity -from .runner import run_preflight -from .workflows import ( +from ._runner import run_preflight +from ._workflows import ( BUNDLED_PRESETS, ReduceToFitError, inspect, diff --git a/src/autointent/_advisor/_cli.py b/src/autointent/advisor/_cli.py similarity index 95% rename from src/autointent/_advisor/_cli.py rename to src/autointent/advisor/_cli.py index 6c1b2bc2e..2c37525f2 100644 --- a/src/autointent/_advisor/_cli.py +++ b/src/autointent/advisor/_cli.py @@ -10,8 +10,8 @@ ``--n-samples / --n-classes / --avg-tokens`` placeholders so the script is useful before the user has built a dataset. -The CLI is a thin wrapper around :func:`autointent._advisor.inspect` and -:func:`autointent._advisor.recommend`; callers that don't need argparse can +The CLI is a thin wrapper around :func:`autointent.advisor.inspect` and +:func:`autointent.advisor.recommend`; callers that don't need argparse can import those helpers directly. """ @@ -22,7 +22,7 @@ import logging import sys -from autointent._advisor import inspect, recommend, stats_from_dataset +from autointent.advisor import inspect, recommend, stats_from_dataset from ._render import render_json, render_recommendation, render_text from ._report import DatasetStats diff --git a/src/autointent/_advisor/_estimates/__init__.py b/src/autointent/advisor/_estimates/__init__.py similarity index 100% rename from src/autointent/_advisor/_estimates/__init__.py rename to src/autointent/advisor/_estimates/__init__.py diff --git a/src/autointent/_advisor/_estimates/_formulas.py b/src/autointent/advisor/_estimates/_formulas.py similarity index 98% rename from src/autointent/_advisor/_estimates/_formulas.py rename to src/autointent/advisor/_estimates/_formulas.py index 05b68d076..f9300be80 100644 --- a/src/autointent/_advisor/_estimates/_formulas.py +++ b/src/autointent/advisor/_estimates/_formulas.py @@ -17,11 +17,11 @@ from typing import TYPE_CHECKING -from autointent._advisor._report import Severity +from autointent.advisor._report import Severity if TYPE_CHECKING: - from autointent._advisor._hub import ModelMeta - from autointent._advisor._report import DatasetStats + from autointent.advisor._hub import ModelMeta + from autointent.advisor._report import DatasetStats _BYTES_PER_GB = 1024**3 diff --git a/src/autointent/_advisor/_estimates/_resource.py b/src/autointent/advisor/_estimates/_resource.py similarity index 99% rename from src/autointent/_advisor/_estimates/_resource.py rename to src/autointent/advisor/_estimates/_resource.py index 05895f1d4..8d7ee6854 100644 --- a/src/autointent/_advisor/_estimates/_resource.py +++ b/src/autointent/advisor/_estimates/_resource.py @@ -13,8 +13,8 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable -from autointent._advisor import _hub -from autointent._advisor._report import ResourceEstimate, Severity +from autointent.advisor import _hub +from autointent.advisor._report import ResourceEstimate, Severity from autointent.configs._embedder import ( EmbedderConfig, OpenaiEmbeddingConfig, @@ -58,9 +58,9 @@ ) if TYPE_CHECKING: - from autointent._advisor._hardware import HardwareProfile - from autointent._advisor._hub import ModelMeta - from autointent._advisor._report import DatasetStats, PreflightReport + from autointent.advisor._hardware import HardwareProfile + from autointent.advisor._hub import ModelMeta + from autointent.advisor._report import DatasetStats, PreflightReport # Union variants of EmbedderConfig that carry a model_name attribute. diff --git a/src/autointent/_advisor/_estimates/_search_space.py b/src/autointent/advisor/_estimates/_search_space.py similarity index 100% rename from src/autointent/_advisor/_estimates/_search_space.py rename to src/autointent/advisor/_estimates/_search_space.py diff --git a/src/autointent/_advisor/_hardware.py b/src/autointent/advisor/_hardware.py similarity index 100% rename from src/autointent/_advisor/_hardware.py rename to src/autointent/advisor/_hardware.py diff --git a/src/autointent/_advisor/_hub.py b/src/autointent/advisor/_hub.py similarity index 100% rename from src/autointent/_advisor/_hub.py rename to src/autointent/advisor/_hub.py diff --git a/src/autointent/_advisor/_render.py b/src/autointent/advisor/_render.py similarity index 100% rename from src/autointent/_advisor/_render.py rename to src/autointent/advisor/_render.py diff --git a/src/autointent/_advisor/_report.py b/src/autointent/advisor/_report.py similarity index 100% rename from src/autointent/_advisor/_report.py rename to src/autointent/advisor/_report.py diff --git a/src/autointent/_advisor/runner.py b/src/autointent/advisor/_runner.py similarity index 95% rename from src/autointent/_advisor/runner.py rename to src/autointent/advisor/_runner.py index 584f3b75c..a7bb8d6fb 100644 --- a/src/autointent/_advisor/runner.py +++ b/src/autointent/advisor/_runner.py @@ -11,14 +11,14 @@ from pydantic import ValidationError -from autointent._advisor._estimates._resource import _resource_phase -from autointent._advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules -from autointent._advisor._report import PreflightReport, Severity +from autointent.advisor._estimates._resource import _resource_phase +from autointent.advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules +from autointent.advisor._report import PreflightReport, Severity from autointent._optimization_config import OptimizationConfig if TYPE_CHECKING: - from autointent._advisor._hardware import HardwareProfile - from autointent._advisor._report import DatasetStats + from autointent.advisor._hardware import HardwareProfile + from autointent.advisor._report import DatasetStats logger = logging.getLogger(__name__) diff --git a/src/autointent/_advisor/workflows.py b/src/autointent/advisor/_workflows.py similarity index 99% rename from src/autointent/_advisor/workflows.py rename to src/autointent/advisor/_workflows.py index 91c3c24a7..fed8b268c 100644 --- a/src/autointent/_advisor/workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -22,7 +22,7 @@ from ._hardware import detect_hardware from ._report import DatasetStats, RecommendationResult, Severity -from .runner import run_preflight +from ._runner import run_preflight if TYPE_CHECKING: from collections.abc import Iterable diff --git a/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index 59e6b87a3..d476531d5 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -133,7 +133,7 @@ class Split: Heavier presets explore more / larger models and take longer to run. The order is a cost ranking, **not** a quality ranking: a heavier preset is not strictly better — e.g. ``transformers-heavy`` will overfit on tiny datasets where a -classic-* preset wins on accuracy. ``autointent._advisor.recommend`` uses this +classic-* preset wins on accuracy. ``autointent.advisor.recommend`` uses this ordering to pick the heaviest preset that still fits the hardware budget, which is a reasonable default but not always the right choice for the data.""" diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index e1ffd77d5..9aad99717 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -16,16 +16,16 @@ import pytest -from autointent._advisor import DatasetStats, HardwareProfile, run_preflight -from autointent._advisor._cli import main -from autointent._advisor.workflows import BUNDLED_PRESETS +from autointent.advisor import DatasetStats, HardwareProfile, run_preflight +from autointent.advisor._cli import main +from autointent.advisor._workflows import BUNDLED_PRESETS from autointent.utils import load_preset @pytest.fixture(autouse=True) def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: """Force HF Hub lookups to fail so tests don't hit the network.""" - from autointent._advisor import _hub + from autointent.advisor import _hub _hub.resolve_model.cache_clear() monkeypatch.setattr(_hub, "_hub_metadata", lambda _name: None) @@ -65,7 +65,7 @@ def test_light_preset_is_feasible_on_8gb_budget(monkeypatch: pytest.MonkeyPatch) # deliberately pessimistic, so "light" would look infeasible on 8 GB. # Restore small-model resolution just for this test so we're verifying # the "light on 8 GB" contract, not the fallback pessimism. - from autointent._advisor import _hub + from autointent.advisor import _hub def _small_model(name: str) -> _hub.ModelMeta: return _hub.ModelMeta( diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index b90997cbb..318530dde 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -6,12 +6,12 @@ import pytest -from autointent._advisor import _hub, run_preflight -from autointent._advisor._estimates._formulas import _classify_severity, _ram_for_module, _vram_for_transformer -from autointent._advisor._estimates._search_space import _extract_model_names, _max_int -from autointent._advisor._hardware import HardwareProfile -from autointent._advisor._hub import ModelMeta -from autointent._advisor._report import DatasetStats, Severity +from autointent.advisor import _hub, run_preflight +from autointent.advisor._estimates._formulas import _classify_severity, _ram_for_module, _vram_for_transformer +from autointent.advisor._estimates._search_space import _extract_model_names, _max_int +from autointent.advisor._hardware import HardwareProfile +from autointent.advisor._hub import ModelMeta +from autointent.advisor._report import DatasetStats, Severity # Per-name ModelMeta fixtures used by the offline tests. Production resolution # (HF Hub config.json + safetensors metadata) is mocked away so the batch-fit @@ -889,13 +889,13 @@ class TestModuleCardinality: """1 for all-singleton, N for finite lists, None for continuous ranges.""" def test_all_singleton(self) -> None: - from autointent._advisor._estimates._search_space import _module_cardinality + from autointent.advisor._estimates._search_space import _module_cardinality assert _module_cardinality({"module_name": "bert"}) == 1 assert _module_cardinality({"module_name": "bert", "batch_size": [64], "epochs": [30]}) == 1 def test_multi_list_multiplies(self) -> None: - from autointent._advisor._estimates._search_space import _module_cardinality + from autointent.advisor._estimates._search_space import _module_cardinality # 2 batch × 3 lr candidates = 6 unique configs cardinality = _module_cardinality( @@ -904,13 +904,13 @@ def test_multi_list_multiplies(self) -> None: assert cardinality == 6 def test_range_dict_is_unbounded(self) -> None: - from autointent._advisor._estimates._search_space import _module_cardinality + from autointent.advisor._estimates._search_space import _module_cardinality # {low, high} → continuous → None (treated as unbounded) assert _module_cardinality({"module_name": "knn", "k": {"low": 1, "high": 20}}) is None def test_reserved_keys_skipped(self) -> None: - from autointent._advisor._estimates._search_space import _module_cardinality + from autointent.advisor._estimates._search_space import _module_cardinality # module_name / target_metric are not search dimensions assert ( diff --git a/tests/advisor/test_hardware_detection.py b/tests/advisor/test_hardware_detection.py index d8131fb19..3926afd8d 100644 --- a/tests/advisor/test_hardware_detection.py +++ b/tests/advisor/test_hardware_detection.py @@ -8,13 +8,13 @@ import pytest -from autointent._advisor._hardware import detect_hardware +from autointent.advisor._hardware import detect_hardware def test_cpu_fallback_when_no_accelerator() -> None: with ( - patch("autointent._advisor._hardware._detect_cuda", return_value=None), - patch("autointent._advisor._hardware._detect_mps", return_value=None), + patch("autointent.advisor._hardware._detect_cuda", return_value=None), + patch("autointent.advisor._hardware._detect_mps", return_value=None), ): hw = detect_hardware() assert hw.accelerator == "cpu" @@ -25,7 +25,7 @@ def test_cpu_fallback_when_no_accelerator() -> None: def test_cuda_branch_classifies_low_gpu() -> None: with ( patch( - "autointent._advisor._hardware._detect_cuda", + "autointent.advisor._hardware._detect_cuda", return_value=(8.0, "NVIDIA RTX 3060"), ), ): @@ -37,10 +37,10 @@ def test_cuda_branch_classifies_low_gpu() -> None: def test_mps_budget_uses_ram_fraction() -> None: with ( - patch("autointent._advisor._hardware._detect_cuda", return_value=None), - patch("autointent._advisor._hardware._detect_ram_gb", return_value=32.0), + patch("autointent.advisor._hardware._detect_cuda", return_value=None), + patch("autointent.advisor._hardware._detect_ram_gb", return_value=32.0), patch( - "autointent._advisor._hardware._detect_mps", + "autointent.advisor._hardware._detect_mps", side_effect=lambda ram, ratio: (ram * ratio, "Apple Silicon (arm64)"), ), ): @@ -53,7 +53,7 @@ def test_mps_budget_uses_ram_fraction() -> None: def test_vram_budget_override_applies() -> None: with ( patch( - "autointent._advisor._hardware._detect_cuda", + "autointent.advisor._hardware._detect_cuda", return_value=(24.0, "NVIDIA RTX 4090"), ), ): @@ -65,8 +65,8 @@ def test_vram_budget_override_applies() -> None: def test_broken_cuda_returns_none_does_not_crash() -> None: # _detect_cuda swallows torch quirks already; verify the wrapper holds. with ( - patch("autointent._advisor._hardware._detect_cuda", return_value=None), - patch("autointent._advisor._hardware._detect_mps", return_value=None), + patch("autointent.advisor._hardware._detect_cuda", return_value=None), + patch("autointent.advisor._hardware._detect_mps", return_value=None), ): hw = detect_hardware() assert hw.accelerator == "cpu" diff --git a/tests/advisor/test_hub_heuristics.py b/tests/advisor/test_hub_heuristics.py index c19018235..de6714f4f 100644 --- a/tests/advisor/test_hub_heuristics.py +++ b/tests/advisor/test_hub_heuristics.py @@ -9,7 +9,7 @@ import pytest -from autointent._advisor import _hub +from autointent.advisor import _hub @pytest.fixture(autouse=True) diff --git a/tests/advisor/test_reduce_to_fit.py b/tests/advisor/test_reduce_to_fit.py index 08a7a9cb3..17cfb9abd 100644 --- a/tests/advisor/test_reduce_to_fit.py +++ b/tests/advisor/test_reduce_to_fit.py @@ -1,4 +1,4 @@ -"""Tests for ``autointent._advisor.reduce_to_fit``. +"""Tests for ``autointent.advisor.reduce_to_fit``. Covers the three review-mandated contracts: @@ -14,7 +14,7 @@ import pytest -from autointent._advisor import ( +from autointent.advisor import ( DatasetStats, HardwareProfile, ReduceToFitError, @@ -25,7 +25,7 @@ @pytest.fixture(autouse=True) def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: - from autointent._advisor import _hub + from autointent.advisor import _hub _hub.resolve_model.cache_clear() monkeypatch.setattr(_hub, "_hub_metadata", lambda _name: None) diff --git a/tests/advisor/test_render.py b/tests/advisor/test_render.py index 7a806c7f2..8aa47ad88 100644 --- a/tests/advisor/test_render.py +++ b/tests/advisor/test_render.py @@ -4,8 +4,8 @@ import json -from autointent._advisor._render import _batch_hint, render_json, render_recommendation, render_text -from autointent._advisor._report import ( +from autointent.advisor._render import _batch_hint, render_json, render_recommendation, render_text +from autointent.advisor._report import ( DatasetStats, PreflightReport, ResourceEstimate, diff --git a/tests/advisor/test_report.py b/tests/advisor/test_report.py index acb2b5bf8..9885478f3 100644 --- a/tests/advisor/test_report.py +++ b/tests/advisor/test_report.py @@ -6,7 +6,7 @@ import pytest -from autointent._advisor._report import ( +from autointent.advisor._report import ( DatasetStats, Finding, PreflightReport, diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py index b9116fbe1..990adf418 100644 --- a/tests/pipeline/test_preflight.py +++ b/tests/pipeline/test_preflight.py @@ -8,7 +8,7 @@ import pytest from autointent import Pipeline -from autointent._advisor import HardwareProfile, detect_hardware, run_preflight, stats_from_dataset_obj +from autointent.advisor import HardwareProfile, detect_hardware, run_preflight, stats_from_dataset_obj from autointent._pipeline import PreflightError from autointent.configs import LoggingConfig From b0619a265c23845b720a73cda66cb521981c6e10 Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 18:35:38 +0300 Subject: [PATCH 33/43] refactor: narrow advisor public surface to 15 names - inspect -> estimate (the old name shadowed the stdlib inspect module) - stats_from_dataset_obj -> dataset_stats - drop BUNDLED_PRESETS, load_config, stats_from_dataset from __all__ (CLI plumbing) - move PreflightError into advisor/_errors.py so there is one import path for it - mark the package experimental in its docstring - lock the surface with tests/advisor/test_public_surface.py --- src/autointent/_pipeline/_pipeline.py | 14 ++------ src/autointent/advisor/__init__.py | 30 +++++++--------- src/autointent/advisor/_cli.py | 7 ++-- src/autointent/advisor/_errors.py | 18 ++++++++++ src/autointent/advisor/_workflows.py | 8 ++--- tests/advisor/test_public_surface.py | 49 +++++++++++++++++++++++++++ tests/pipeline/test_preflight.py | 4 +-- 7 files changed, 92 insertions(+), 38 deletions(-) create mode 100644 src/autointent/advisor/_errors.py create mode 100644 tests/advisor/test_public_surface.py diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index 0249649dd..3c1bbfd68 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -13,10 +13,11 @@ from autointent import Context, OptimizationConfig from autointent.advisor import ( + PreflightError, Severity, + dataset_stats, detect_hardware, run_preflight, - stats_from_dataset_obj, ) from autointent.configs import ( CrossEncoderConfig, @@ -48,15 +49,6 @@ PreflightMode = Literal["off", "warn", "strict"] -class PreflightError(RuntimeError): - """Raised when ``Pipeline.fit(preflight="strict")`` finds OVER-budget resources.""" - - def __init__(self, findings: list[Any]) -> None: - self.findings = findings - lines = "\n".join(f" [{f.phase}] {f.message}" for f in findings) - super().__init__(f"Preflight check failed with {len(findings)} OVER finding(s):\n{lines}") - - class Pipeline: """Pipeline optimizer class. @@ -195,7 +187,7 @@ def _run_preflight(self, dataset: Dataset, *, refit_after: bool, mode: Preflight ``"strict"`` and any OVER finding is produced, raises ``PreflightError``. """ config = self._build_advisor_config() - stats = stats_from_dataset_obj(dataset) + stats = dataset_stats(dataset) hardware = detect_hardware() report = run_preflight(config, stats, hardware, refit_after=refit_after) _log_preflight_report(report, self._logger) diff --git a/src/autointent/advisor/__init__.py b/src/autointent/advisor/__init__.py index a7d9d1b87..42f34b95a 100644 --- a/src/autointent/advisor/__init__.py +++ b/src/autointent/advisor/__init__.py @@ -1,41 +1,37 @@ """Pre-flight compute feasibility advisor. -Exposes a small surface used by both ``Pipeline.fit()`` (see the ``preflight=`` -kwarg) and the ``autointent-advisor`` CLI script. +**Experimental.** This subpackage estimates VRAM, RAM, disk, and wall-time for a +search space before any training starts. Estimates are heuristic and calibrated +against a limited hardware sample, so treat them as guidance rather than +guarantees — see the accuracy caveats in the ``advisor`` page of the docs. The +public surface may change in a minor release. + +Two ways in: the ``autointent-advisor`` console script, and the functions below. +``Pipeline.fit(preflight=...)`` wires the same machinery into a fit, opt-in. """ from __future__ import annotations +from ._errors import PreflightError from ._hardware import HardwareProfile, detect_hardware from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity from ._runner import run_preflight -from ._workflows import ( - BUNDLED_PRESETS, - ReduceToFitError, - inspect, - load_config, - recommend, - reduce_to_fit, - stats_from_dataset, - stats_from_dataset_obj, -) +from ._workflows import ReduceToFitError, dataset_stats, estimate, recommend, reduce_to_fit __all__ = [ - "BUNDLED_PRESETS", "DatasetStats", "Finding", "HardwareProfile", + "PreflightError", "PreflightReport", "RecommendationResult", "ReduceToFitError", "ResourceEstimate", "Severity", + "dataset_stats", "detect_hardware", - "inspect", - "load_config", + "estimate", "recommend", "reduce_to_fit", "run_preflight", - "stats_from_dataset", - "stats_from_dataset_obj", ] diff --git a/src/autointent/advisor/_cli.py b/src/autointent/advisor/_cli.py index 2c37525f2..dd79bd3b8 100644 --- a/src/autointent/advisor/_cli.py +++ b/src/autointent/advisor/_cli.py @@ -10,7 +10,7 @@ ``--n-samples / --n-classes / --avg-tokens`` placeholders so the script is useful before the user has built a dataset. -The CLI is a thin wrapper around :func:`autointent.advisor.inspect` and +The CLI is a thin wrapper around :func:`autointent.advisor.estimate` and :func:`autointent.advisor.recommend`; callers that don't need argparse can import those helpers directly. """ @@ -22,10 +22,9 @@ import logging import sys -from autointent.advisor import inspect, recommend, stats_from_dataset - from ._render import render_json, render_recommendation, render_text from ._report import DatasetStats +from ._workflows import estimate, recommend, stats_from_dataset logger = logging.getLogger("autointent.advisor") @@ -56,7 +55,7 @@ def _add_common_dataset_args(p: argparse.ArgumentParser) -> None: def cmd_inspect(args: argparse.Namespace) -> int: - report = inspect( + report = estimate( args.target, stats=_stats_from_args(args), budget_vram_gb=args.budget_vram_gb, diff --git a/src/autointent/advisor/_errors.py b/src/autointent/advisor/_errors.py new file mode 100644 index 000000000..bd81c48cf --- /dev/null +++ b/src/autointent/advisor/_errors.py @@ -0,0 +1,18 @@ +"""Advisor exceptions raised across the package / Pipeline boundary.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ._report import Finding + + +class PreflightError(RuntimeError): + """Raised when ``Pipeline.fit(preflight="strict")`` finds OVER-budget resources.""" + + def __init__(self, findings: list[Finding]) -> None: + self.findings = findings + lines = "\n".join(f" [{f.phase}] {f.message}" for f in findings) + msg = f"Preflight check failed with {len(findings)} OVER finding(s):\n{lines}" + super().__init__(msg) diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index fed8b268c..6c34ef93c 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -1,4 +1,4 @@ -"""High-level advisor workflows: ``inspect`` and ``recommend``. +"""High-level advisor workflows: ``estimate``, ``recommend``, and ``reduce_to_fit``. Each workflow orchestrates the lower-level pieces (``load_config``, ``detect_hardware``, ``stats_from_dataset``, ``run_preflight``) into a single @@ -99,7 +99,7 @@ def stats_from_dataset(path: str, *, multilabel: bool = False) -> DatasetStats: ) -def stats_from_dataset_obj(dataset: Dataset) -> DatasetStats: +def dataset_stats(dataset: Dataset) -> DatasetStats: """Build :class:`DatasetStats` straight from an in-memory ``Dataset``. Counterpart of :func:`stats_from_dataset` that skips HF ``load_dataset`` @@ -186,13 +186,13 @@ def _class_counts( return counts -def inspect( +def estimate( target: str, *, stats: DatasetStats | None = None, budget_vram_gb: float | None = None, ) -> PreflightReport: - """Inspect a preset (or YAML config path) against the local hardware. + """Estimate what a preset (or YAML config path) will cost on the local hardware. Args: target: Bundled preset name (e.g. ``'transformers-light'``) or a YAML diff --git a/tests/advisor/test_public_surface.py b/tests/advisor/test_public_surface.py new file mode 100644 index 000000000..59b848921 --- /dev/null +++ b/tests/advisor/test_public_surface.py @@ -0,0 +1,49 @@ +"""Locks the public surface of ``autointent.advisor``. + +The advisor is marked experimental, but "experimental" is not a licence for the +surface to drift silently. This test is the tripwire: adding or removing a +public name is a deliberate act that updates this list. +""" + +from __future__ import annotations + +import autointent.advisor as advisor + +EXPECTED_SURFACE = { + # functions + "dataset_stats", + "detect_hardware", + "estimate", + "recommend", + "reduce_to_fit", + "run_preflight", + # types + "DatasetStats", + "Finding", + "HardwareProfile", + "PreflightError", + "PreflightReport", + "RecommendationResult", + "ReduceToFitError", + "ResourceEstimate", + "Severity", +} + + +def test_all_matches_expected_surface() -> None: + assert set(advisor.__all__) == EXPECTED_SURFACE + + +def test_every_exported_name_resolves() -> None: + missing = [name for name in advisor.__all__ if not hasattr(advisor, name)] + assert missing == [] + + +def test_no_stdlib_shadowing_names() -> None: + """``inspect`` was exported previously and shadows the stdlib module.""" + assert "inspect" not in advisor.__all__ + + +def test_package_documents_experimental_status() -> None: + assert advisor.__doc__ is not None + assert "experimental" in advisor.__doc__.lower() diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py index 990adf418..9c366cf1f 100644 --- a/tests/pipeline/test_preflight.py +++ b/tests/pipeline/test_preflight.py @@ -8,7 +8,7 @@ import pytest from autointent import Pipeline -from autointent.advisor import HardwareProfile, detect_hardware, run_preflight, stats_from_dataset_obj +from autointent.advisor import HardwareProfile, dataset_stats, detect_hardware, run_preflight from autointent._pipeline import PreflightError from autointent.configs import LoggingConfig @@ -96,7 +96,7 @@ def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: """ p = _classic_light_pipeline() config = p._build_advisor_config() - stats = stats_from_dataset_obj(dataset) + stats = dataset_stats(dataset) hardware = detect_hardware() report = run_preflight(config, stats, hardware, preset_name="classic-light") From 751f0e1b60f7ae207e9a82b5a1896f7bcd7c7a84 Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 18:46:23 +0300 Subject: [PATCH 34/43] fix: reduce_to_fit pruned by VRAM regardless of the binding constraint Finding.metric holds short names ('vram', 'ram', 'disk', 'time') but _pick_module_to_drop tested membership against ['vram_gb', 'time_hours', 'ram_gb', 'disk_download_gb']. The sets are disjoint, so the lookup never matched and driver_key always fell through to 'vram_gb' -- correct by accident on VRAM-bound machines, wrong everywhere else. Map the two namespaces explicitly and drop disk from the priority walk, since driver rows carry no per-module disk figure (documented as a VRAM proxy). Every existing test used a _profile() with hardcoded ram_gb/free_disk_gb, which is why a green suite missed this; the helper now parameterizes both. Found during validation: Darinochka/AutoIntent-experiments#40, finding 3. --- src/autointent/advisor/_workflows.py | 35 +++++++---- tests/advisor/test_pick_module_to_drop.py | 73 +++++++++++++++++++++++ tests/advisor/test_reduce_to_fit.py | 6 +- 3 files changed, 100 insertions(+), 14 deletions(-) create mode 100644 tests/advisor/test_pick_module_to_drop.py diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index 6c34ef93c..a9fbca7eb 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -305,19 +305,32 @@ def _drop_module_from_search_space( return out +# ``Finding.metric`` uses short names ("vram"); driver rows use suffixed keys +# ("vram_gb"). Mapping the two is what makes the priority walk below work — a +# previous version compared the two namespaces directly, so the lookup never +# matched and every prune silently fell back to VRAM (experiments #40 #3). +_DRIVER_KEY_BY_METRIC = {"vram": "vram_gb", "time": "time_hours", "ram": "ram_gb"} +# Preference order when several budgets are over at once. +_METRIC_PRIORITY = ("vram", "time", "ram") + + def _pick_module_to_drop(report: PreflightReport) -> tuple[str, str] | None: - """Pick the (node_type, module_name) that contributes the most to whichever - budget is over. Returns ``None`` when no droppable driver exists (all - remaining rows are decision-node entries or unknown-cost placeholders). - - Preference order: VRAM > time > RAM > disk. We drop the driver with the - largest cost along the *first* dimension that has at least one OVER - finding — otherwise (edge case: is_feasible False without an OVER, which - shouldn't happen) fall back to VRAM. + """Pick the (node_type, module_name) contributing most to whichever budget is over. + + Drops the driver with the largest cost along the first dimension that has an + OVER finding, preferring VRAM > time > RAM. Disk is deliberately absent from + that walk: driver rows carry no per-module disk figure, so disk pressure + reduces by the VRAM proxy (download size tracks model size). Falls back to + VRAM when nothing is OVER. + + Returns ``None`` when no droppable driver exists — all remaining rows are + decision-node entries or unknown-cost placeholders. """ - findings_by_metric = {f.metric for f in report.findings if f.severity == Severity.OVER} - priority = ["vram_gb", "time_hours", "ram_gb", "disk_download_gb"] - driver_key = next((k for k in priority if k in findings_by_metric), None) or "vram_gb" + over_metrics = {f.metric for f in report.findings if f.severity == Severity.OVER} + driver_key = next( + (_DRIVER_KEY_BY_METRIC[m] for m in _METRIC_PRIORITY if m in over_metrics), + "vram_gb", + ) drivers = report.resource.drivers or [] # Only drop scoring-node drivers — decision modules are lightweight and diff --git a/tests/advisor/test_pick_module_to_drop.py b/tests/advisor/test_pick_module_to_drop.py new file mode 100644 index 000000000..b27012ae2 --- /dev/null +++ b/tests/advisor/test_pick_module_to_drop.py @@ -0,0 +1,73 @@ +"""Unit tests for ``_pick_module_to_drop``'s constraint selection. + +Built from hand-rolled reports rather than real preflight runs: the rule under +test is "prune the module heaviest along the dimension that is actually over +budget", and that rule should be verifiable without invoking any formula. + +Regression guard for experiments #40 finding 3 — the metric-name mismatch that +made every prune a VRAM prune. +""" + +from __future__ import annotations + +from autointent.advisor._report import PreflightReport, Severity +from autointent.advisor._workflows import _pick_module_to_drop + + +def _report(*over_metrics: str) -> PreflightReport: + """Report whose scoring drivers disagree about which module is heaviest. + + ``bert`` is heaviest on VRAM, ``linear`` on RAM, ``catboost`` on time — so + the module returned identifies which dimension the code actually consulted. + """ + report = PreflightReport() + for metric in ("vram", "ram", "disk", "time"): + severity = Severity.OVER if metric in over_metrics else Severity.AMPLE + report.add("resource", severity, f"{metric} finding", metric=metric) + report.resource.drivers = [ + {"node_type": "scoring", "module": "bert", "vram_gb": 20.0, "ram_gb": 3.0, "time_hours": 2.0}, + {"node_type": "scoring", "module": "linear", "vram_gb": 0.5, "ram_gb": 40.0, "time_hours": 1.0}, + {"node_type": "scoring", "module": "catboost", "vram_gb": 0.2, "ram_gb": 8.0, "time_hours": 90.0}, + # Decision modules are never droppable, however heavy they look. + {"node_type": "decision", "module": "argmax", "vram_gb": 99.0, "ram_gb": 99.0, "time_hours": 99.0}, + ] + return report + + +def test_ram_over_prunes_ram_heaviest() -> None: + assert _pick_module_to_drop(_report("ram")) == ("scoring", "linear") + + +def test_time_over_prunes_time_heaviest() -> None: + assert _pick_module_to_drop(_report("time")) == ("scoring", "catboost") + + +def test_vram_over_prunes_vram_heaviest() -> None: + assert _pick_module_to_drop(_report("vram")) == ("scoring", "bert") + + +def test_vram_wins_when_several_constraints_are_over() -> None: + """Documented preference order is VRAM > time > RAM.""" + assert _pick_module_to_drop(_report("vram", "ram", "time")) == ("scoring", "bert") + + +def test_time_beats_ram_when_both_over() -> None: + assert _pick_module_to_drop(_report("ram", "time")) == ("scoring", "catboost") + + +def test_disk_over_falls_back_to_vram_proxy() -> None: + """Drivers carry no per-module disk figure, so disk reduces by the VRAM proxy.""" + assert _pick_module_to_drop(_report("disk")) == ("scoring", "bert") + + +def test_no_over_findings_falls_back_to_vram() -> None: + assert _pick_module_to_drop(_report()) == ("scoring", "bert") + + +def test_returns_none_when_no_scoring_driver_is_droppable() -> None: + report = _report("ram") + report.resource.drivers = [ + {"node_type": "decision", "module": "argmax", "vram_gb": 1.0, "ram_gb": 1.0, "time_hours": 1.0}, + {"node_type": "scoring", "module": "?", "vram_gb": 1.0, "ram_gb": 1.0, "time_hours": 1.0}, + ] + assert _pick_module_to_drop(report) is None diff --git a/tests/advisor/test_reduce_to_fit.py b/tests/advisor/test_reduce_to_fit.py index 17cfb9abd..4dd05ffe1 100644 --- a/tests/advisor/test_reduce_to_fit.py +++ b/tests/advisor/test_reduce_to_fit.py @@ -31,13 +31,13 @@ def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(_hub, "_hub_metadata", lambda _name: None) -def _profile(vram_gb: float = 16.0) -> HardwareProfile: +def _profile(vram_gb: float = 16.0, ram_gb: float = 32.0, free_disk_gb: float = 200.0) -> HardwareProfile: return HardwareProfile( accelerator="cuda" if vram_gb > 0 else "cpu", device_name="test-gpu" if vram_gb > 0 else "test-cpu", vram_gb=vram_gb, - ram_gb=32.0, - free_disk_gb=200.0, + ram_gb=ram_gb, + free_disk_gb=free_disk_gb, cpu_count=8, ) From 9e12b0cc0233271f2d100011cf1c0e3e9b3893ca Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 19:00:17 +0300 Subject: [PATCH 35/43] refactor: declare preset cost ranking explicitly The branch reordered the public SearchSpacePreset literal into a cost ranking and derived cost_rank from get_args() declaration order, making a public type alias's element order load-bearing. Restore dev's ordering and move the ranking into PRESET_COST_ORDER, covered by a test so adding a preset fails loudly instead of silently sorting it cheapest-last. --- src/autointent/advisor/_workflows.py | 27 ++++++++++++++++----- src/autointent/custom_types/_types.py | 29 +++++++++++----------- tests/advisor/test_estimates_and_cli.py | 4 ++-- tests/advisor/test_preset_cost_order.py | 32 +++++++++++++++++++++++++ 4 files changed, 70 insertions(+), 22 deletions(-) create mode 100644 tests/advisor/test_preset_cost_order.py diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index a9fbca7eb..c65f0bd17 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -12,7 +12,7 @@ import logging from pathlib import Path -from typing import TYPE_CHECKING, Any, get_args +from typing import TYPE_CHECKING, Any import yaml from datasets import ClassLabel, Sequence, load_dataset @@ -36,7 +36,22 @@ _SAMPLE_LIMIT = 1000 _P95_PERCENTILE = 0.95 -BUNDLED_PRESETS: tuple[str, ...] = get_args(SearchSpacePreset) +# Descending resource cost. Declared explicitly rather than derived from +# get_args(SearchSpacePreset) so that reordering a public type alias cannot +# silently change what `recommend` picks. Kept in sync by +# tests/advisor/test_preset_cost_order.py. +PRESET_COST_ORDER: tuple[SearchSpacePreset, ...] = ( + "transformers-heavy", + "transformers-light", + "nn-heavy", + "zero-shot-llm", + "nn-medium", + "classic-heavy", + "transformers-no-hpo", + "classic-medium", + "zero-shot-encoders", + "classic-light", +) def load_config(target: str) -> tuple[dict[str, Any], str]: @@ -221,7 +236,7 @@ def recommend( Args: stats: Dataset stats to score against. Defaults to a placeholder if ``None``. - presets: Override of the preset list (defaults to ``BUNDLED_PRESETS``). + presets: Override of the preset list (defaults to ``PRESET_COST_ORDER``). budget_vram_gb: Optional VRAM-budget override for the hardware probe. budget_time_h: Optional wall-time ceiling in hours; presets exceeding it get an extra ``Severity.OVER`` finding so they drop out of the @@ -239,7 +254,7 @@ def recommend( """ hardware = detect_hardware(vram_budget_gb=budget_vram_gb) stats = stats or DatasetStats.placeholder() - preset_iter = list(presets) if presets is not None else BUNDLED_PRESETS + preset_iter = list(presets) if presets is not None else list(PRESET_COST_ORDER) results: list[tuple[str, PreflightReport]] = [] for preset in preset_iter: @@ -257,9 +272,9 @@ def recommend( ) results.append((preset, report)) - cost_rank = {name: i for i, name in enumerate(BUNDLED_PRESETS)} + cost_rank: dict[str, int] = {name: i for i, name in enumerate(PRESET_COST_ORDER)} feasible = [(name, r) for name, r in results if r.is_feasible] - feasible.sort(key=lambda pair: (cost_rank.get(pair[0], len(BUNDLED_PRESETS)), pair[0])) + feasible.sort(key=lambda pair: (cost_rank.get(pair[0], len(PRESET_COST_ORDER)), pair[0])) chosen = feasible[0][0] if feasible else None return RecommendationResult(chosen=chosen, results=results) diff --git a/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index d476531d5..f322001d1 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -117,25 +117,26 @@ class Split: """ SearchSpacePreset = Literal[ - "transformers-heavy", - "transformers-light", + "classic-heavy", + "classic-light", + "classic-medium", "nn-heavy", - "zero-shot-llm", "nn-medium", - "classic-heavy", + "transformers-heavy", + "transformers-light", "transformers-no-hpo", - "classic-medium", + "zero-shot-llm", "zero-shot-encoders", - "classic-light", ] -"""Bundled search-space presets, listed in descending resource-cost order. - -Heavier presets explore more / larger models and take longer to run. The order -is a cost ranking, **not** a quality ranking: a heavier preset is not strictly -better — e.g. ``transformers-heavy`` will overfit on tiny datasets where a -classic-* preset wins on accuracy. ``autointent.advisor.recommend`` uses this -ordering to pick the heaviest preset that still fits the hardware budget, -which is a reasonable default but not always the right choice for the data.""" +"""Bundled search-space presets that our library supports. + +The order here carries no meaning. Resource-cost ranking lives in +``autointent.advisor._workflows.PRESET_COST_ORDER``, which ``recommend`` uses to +pick the heaviest preset that still fits the hardware budget. That is a cost +ranking, **not** a quality ranking: a heavier preset is not strictly better — +``transformers-heavy`` will overfit on tiny datasets where a ``classic-*`` +preset wins on accuracy. +""" class Document(BaseModel): diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 9aad99717..608c7420f 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -18,7 +18,7 @@ from autointent.advisor import DatasetStats, HardwareProfile, run_preflight from autointent.advisor._cli import main -from autointent.advisor._workflows import BUNDLED_PRESETS +from autointent.advisor._workflows import PRESET_COST_ORDER from autointent.utils import load_preset @@ -42,7 +42,7 @@ def _profile(vram_gb: float = 16.0) -> HardwareProfile: ) -@pytest.mark.parametrize("preset", BUNDLED_PRESETS) +@pytest.mark.parametrize("preset", PRESET_COST_ORDER) def test_every_preset_inspects_without_raising(preset: str) -> None: cfg = load_preset(preset) # type: ignore[arg-type] stats = DatasetStats.placeholder(n_samples=500, n_classes=10, avg_tokens=24) diff --git a/tests/advisor/test_preset_cost_order.py b/tests/advisor/test_preset_cost_order.py new file mode 100644 index 000000000..12906691e --- /dev/null +++ b/tests/advisor/test_preset_cost_order.py @@ -0,0 +1,32 @@ +"""``PRESET_COST_ORDER`` must stay in sync with the preset literal. + +``recommend`` picks the heaviest preset that fits, so a preset missing from the +ranking would silently sort last and effectively never be recommended. This test +turns that into a hard failure at the moment a preset is added. +""" + +from __future__ import annotations + +from typing import get_args + +from autointent.advisor._workflows import PRESET_COST_ORDER +from autointent.custom_types import SearchSpacePreset + + +def test_cost_order_covers_every_preset() -> None: + assert set(PRESET_COST_ORDER) == set(get_args(SearchSpacePreset)) + + +def test_cost_order_has_no_duplicates() -> None: + assert len(PRESET_COST_ORDER) == len(set(PRESET_COST_ORDER)) + + +def test_preset_literal_order_is_not_load_bearing() -> None: + """The literal's declaration order must not be the cost ranking. + + Cost ordering belongs in PRESET_COST_ORDER, where it is commented and + tested. If these two ever coincide exactly, someone has reintroduced the + coupling -- the literal is kept in dev's original (roughly alphabetical) + order precisely so it cannot be mistaken for a ranking. + """ + assert tuple(get_args(SearchSpacePreset)) != PRESET_COST_ORDER From 60d7b231acfbeca9839ce59530804f767db101c1 Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 19:20:15 +0300 Subject: [PATCH 36/43] style: clear mechanical ruff and mypy findings in the advisor 28 ruff findings (UP035, F401, RUF002/003, D205/D209, N806, PLR2004, EM101/EM102/TRY003) and 6 mypy errors. No behaviour change. The ModelMeta fix is a real latent bug: _fold_disk_costs reused the name 'meta' for both a ModelMeta loop variable and a ModelMeta | None lookup, which is why mypy also reported the None guard as unreachable. Co-Authored-By: Claude Opus 5 (1M context) --- .../advisor/_estimates/_formulas.py | 71 ++++++++++++------- .../advisor/_estimates/_resource.py | 27 ++++--- .../advisor/_estimates/_search_space.py | 10 ++- src/autointent/advisor/_hub.py | 2 +- src/autointent/advisor/_runner.py | 17 ++++- src/autointent/advisor/_workflows.py | 61 ++++++++++------ tests/advisor/test_estimates_internals.py | 8 +-- tests/advisor/test_public_surface.py | 2 +- tests/advisor/test_reduce_to_fit.py | 8 ++- 9 files changed, 134 insertions(+), 72 deletions(-) diff --git a/src/autointent/advisor/_estimates/_formulas.py b/src/autointent/advisor/_estimates/_formulas.py index f9300be80..9fabf2c11 100644 --- a/src/autointent/advisor/_estimates/_formulas.py +++ b/src/autointent/advisor/_estimates/_formulas.py @@ -56,9 +56,12 @@ def _classify_severity(estimate: float, budget: float) -> Severity: def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: - """Weight-side VRAM: weights + grads + optimizer state. Pessimistic upper - bound by mode: 1.3× inference, 1.3× + 0.5 GB lora adapters, 4.5× full - finetune (textbook 4W + fragmentation/workspaces slack).""" + """Weight-side VRAM: weights + grads + optimizer state. + + Pessimistic upper bound by mode: 1.3x inference, 1.3x + 0.5 GB lora + adapters, 4.5x full finetune (textbook 4W + fragmentation/workspaces + slack). + """ weights_gb = meta.weights_gb if mode == "inference": return weights_gb * 1.3 @@ -73,9 +76,11 @@ def _activations_gb_per_sample( *, is_training: bool, ) -> float: - """Activation memory per sample. Training: 34 B/token/layer (Korthikanti - 2022 upper bound, standard attention). Inference: 8 B/token (only 1-2 - layers' outputs in flight).""" + """Activation memory per sample. + + Training: 34 B/token/layer (Korthikanti 2022 upper bound, standard + attention). Inference: 8 B/token (only 1-2 layers' outputs in flight). + """ hidden = _embedder_dim(meta) training_bytes_per_token_per_layer = 34 inference_bytes_per_token = 8 @@ -93,10 +98,11 @@ def _vram_for_transformer( batch_size: int = 0, seq_len: int = _DEFAULT_SEQ_LEN, ) -> float: - """Total VRAM: weights + grads + optimizer state + activations × batch. + """Total VRAM: weights + grads + optimizer state + activations x batch. Safety margin: 1.20 for training (backward transients, eval sweep, - allocator fragmentation), 1.10 for inference (no backward).""" + allocator fragmentation), 1.10 for inference (no backward). + """ base = _weights_vram_for_transformer(meta, mode) if batch_size <= 0: return base @@ -152,9 +158,11 @@ def _time_for_transformer( params_millions: float, device_class: str, ) -> float: - """Transformer training wall-time in hours. Per-step FLOPs = 6 × params × - batch × seq_len (fwd+bwd), ÷ sustained device TFLOPS, × total steps × - trainer overhead.""" + """Transformer training wall-time in hours. + + Per-step FLOPs = 6 x params x batch x seq_len (fwd+bwd), ÷ sustained + device TFLOPS, x total steps x trainer overhead. + """ steps_per_epoch = max(1, n_samples // max(1, batch_size)) total_steps = n_trials * epochs * steps_per_epoch # 6x factor: ~2x for fwd matmul + ~4x for bwd (grad wrt input + grad wrt weight). @@ -186,9 +194,11 @@ def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: def _ram_for_module(meta: ModelMeta, stats: DatasetStats, *, mode: str = "inference") -> float: - """RAM upper bound: weights × mode_mult + tokenized text (n_samples × - avg_tokens × 4 B). Mode multiplier: 1.3 inference, 1.5 lora, 4.5 - full-finetune (Adam mirrors weights on host too).""" + """RAM upper bound: weights x mode_mult + tokenized text (n_samples x avg_tokens x 4 B). + + Mode multiplier: 1.3 inference, 1.5 lora, 4.5 full-finetune (Adam + mirrors weights on host too). + """ if mode == "inference": weights_mult = 1.3 elif mode == "lora": @@ -209,7 +219,7 @@ def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: _LINEAR_CPU_S_PER_SAMPLE_FEATURE = 1.2e-9 _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 _CATBOOST_GPU_SPEEDUP = 10.0 -_LOGREG_CV_MULTIPLIER = 31 # sklearn default: Cs=10 × cv=3 + 1 final refit +_LOGREG_CV_MULTIPLIER = 31 # sklearn default: Cs=10 x cv=3 + 1 final refit _CATBOOST_DEFAULT_BINS = 254 # CatBoost `border_count` default _CATBOOST_BYTES_PER_TREE_NODE = 32 @@ -231,8 +241,11 @@ def _time_for_linear( cv_multiplier: int, class_multiplier: int, ) -> float: - """LogisticRegression wall time. O(n_samples × features × classes × cv) - per fit; typical L-BFGS convergence absorbed into the calibration constant.""" + """LogisticRegression wall time. + + O(n_samples x features x classes x cv) per fit; typical L-BFGS + convergence absorbed into the calibration constant. + """ seconds = ( n_trials * _LINEAR_CPU_S_PER_SAMPLE_FEATURE @@ -260,8 +273,10 @@ def _ram_for_sklearn( max_depth: int, n_jobs: int, ) -> float: - """RandomForest RAM: (feature matrix + trees) × n_jobs. joblib workers - each hold a full copy.""" + """RandomForest RAM: (feature matrix + trees) x n_jobs. + + joblib workers each hold a full copy. + """ per_worker_data = stats.n_samples * embedder_dim * 8 # fp64 default n_leaves = min(2**max_depth, stats.n_samples) if max_depth > 0 else stats.n_samples per_worker_trees = n_estimators * n_leaves * _CATBOOST_BYTES_PER_TREE_NODE @@ -270,9 +285,11 @@ def _ram_for_sklearn( def _embedder_load_ram_gb(meta: ModelMeta | None) -> float: """Aggregate-level RAM penalty when a classic preset uses an embedder. + Added on top of the max-driver RAM because embedder + multiple classic scorers coexist in memory. Uses fp32 weights (transformers up-casts at - load) × 3.5 for weights + activation buffers + framework slack.""" + load) x 3.5 for weights + activation buffers + framework slack. + """ if meta is None: return 0.0 return ((meta.total_params * 4) / _BYTES_PER_GB) * 3.5 @@ -329,8 +346,11 @@ def _rnn_param_count(*, embed_dim: int, hidden_dim: int, n_classes: int) -> int: def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: - """Weights + 3× optimizer/grads + activations. Same fp32 upper bound as - transformers, smaller hidden dim (embed_dim / num_filters).""" + """Weights + 3x optimizer/grads + activations. + + Same fp32 upper bound as transformers, smaller hidden dim (embed_dim / + num_filters). + """ weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB activations_gb = ( batch_size * _NN_DEFAULT_SEQ_LEN * hidden_dim * _NN_TRAIN_ACT_BYTES_PER_UNIT @@ -352,8 +372,11 @@ def _time_for_nn( params_millions: float, device_class: str, ) -> float: - """Reuse transformer FLOPs formula; small models slightly under-predict - since they're memory-bandwidth-bound, but within 2x for cost ranking.""" + """Reuse transformer FLOPs formula. + + Small models slightly under-predict since they're memory-bandwidth-bound, + but within 2x for cost ranking. + """ return _time_for_transformer( n_trials=n_trials, epochs=epochs, diff --git a/src/autointent/advisor/_estimates/_resource.py b/src/autointent/advisor/_estimates/_resource.py index 8d7ee6854..d994d1f67 100644 --- a/src/autointent/advisor/_estimates/_resource.py +++ b/src/autointent/advisor/_estimates/_resource.py @@ -11,7 +11,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable +from typing import TYPE_CHECKING, Any from autointent.advisor import _hub from autointent.advisor._report import ResourceEstimate, Severity @@ -52,12 +52,13 @@ from ._search_space import ( _extract_model_names, _max_int, - _module_cardinality, _walk_modules, _walk_modules_indexed, ) if TYPE_CHECKING: + from collections.abc import Callable + from autointent.advisor._hardware import HardwareProfile from autointent.advisor._hub import ModelMeta from autointent.advisor._report import DatasetStats, PreflightReport @@ -269,7 +270,7 @@ def _estimate_classic_entry( max_depth=max_depth, n_jobs=sk_n_jobs, ) - # Rough O(n_estimators × n × features / n_jobs); real numbers vary + # Rough O(n_estimators x n x features / n_jobs); real numbers vary # wildly by criterion so use the linear coefficient as a proxy. time_h = ( n_trials @@ -314,7 +315,9 @@ def _estimate_nn_entry( refit_after: bool, ) -> _ModuleEstimate | None: """Cost row for cnn / rnn scorers (returns None for anything else). - Small torch models trained from scratch on token ids.""" + + Small torch models trained from scratch on token ids. + """ module = entry.get("module_name", "?") n_classes = max(1, stats.n_classes) @@ -470,12 +473,12 @@ def _aggregate_disk( if cached_embedders and stats is not None: for name in cached_embedders: - meta = seen_models.get(name) - if meta is None: + cached_meta = seen_models.get(name) + if cached_meta is None: continue estimate.disk_embedding_cache_gb += _embedding_cache_disk_gb( n_samples=stats.n_samples, - hidden_size=_embedder_dim(meta), + hidden_size=_embedder_dim(cached_meta), ) @@ -644,9 +647,11 @@ def _resource_phase( variants_per_node[node_idx] = variants_per_node.get(node_idx, 0) + 1 def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int: # noqa: ARG001 - """Trials charged to this module. ``entry`` reserved for future - per-module caps (see git history for the cardinality-cap experiment - that broke description-scorer presets).""" + """Trials charged to this module. + + ``entry`` is reserved for future per-module caps (see git history for the + cardinality-cap experiment that broke description-scorer presets). + """ divisor = max(1, variants_per_node.get(node_idx, 1)) return max(1, n_trials // divisor) @@ -758,7 +763,7 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int # Flip low_confidence if any model fell back to the heuristic path (Hub # unreachable, repo missing safetensors metadata, local-path checkpoint). # Emit as a TIGHT finding (not just a note) so it shows up in the main - # rendered findings block — buried notes previously let ~2× under-prediction + # rendered findings block — buried notes previously let ~2x under-prediction # of large-model shapes slip past the reviewer. heuristic_models = [m.name for m in seen_models.values() if m.confidence == "heuristic"] if heuristic_models: diff --git a/src/autointent/advisor/_estimates/_search_space.py b/src/autointent/advisor/_estimates/_search_space.py index 3792d2f23..0f1324600 100644 --- a/src/autointent/advisor/_estimates/_search_space.py +++ b/src/autointent/advisor/_estimates/_search_space.py @@ -55,6 +55,12 @@ def _max_int(value: Any, default: int) -> int: # noqa: ANN401 return default +# Keys that describe the entry rather than a tunable hyperparameter. +_RESERVED_ENTRY_KEYS = frozenset({"module_name", "target_metric"}) +# Above this many combinations the exact count stops mattering for cost. +_CARDINALITY_CAP = 10_000 + + def _module_cardinality(entry: dict[str, Any]) -> int | None: """Unique configurations the module entry can produce. @@ -62,11 +68,9 @@ def _module_cardinality(entry: dict[str, Any]) -> int | None: products (capped at 10_000), None when any field is a continuous ``{low, high}`` range. """ - _RESERVED = {"module_name", "target_metric"} - _CARDINALITY_CAP = 10_000 product = 1 for key, value in entry.items(): - if key in _RESERVED: + if key in _RESERVED_ENTRY_KEYS: continue if isinstance(value, dict): if "low" in value and "high" in value: diff --git a/src/autointent/advisor/_hub.py b/src/autointent/advisor/_hub.py index 3dc48aaab..f277a0819 100644 --- a/src/autointent/advisor/_hub.py +++ b/src/autointent/advisor/_hub.py @@ -23,7 +23,7 @@ # Conservative "large-model" shape used when Hub metadata is unavailable — # roughly deberta-v3-large / bert-large sized. Previously we defaulted to a # BERT-base shape (110M / 768 / 12), which *under*-predicted a real deberta-large -# fit by ~2×. Because the advisor's contract is a pessimistic upper bound, the +# fit by ~2x. Because the advisor's contract is a pessimistic upper bound, the # offline fallback needs to over-estimate small models rather than under-estimate # large ones. Callers can still see the fallback happened via ``confidence == # "heuristic"`` and ``PreflightReport.low_confidence``. diff --git a/src/autointent/advisor/_runner.py b/src/autointent/advisor/_runner.py index a7bb8d6fb..464f471a1 100644 --- a/src/autointent/advisor/_runner.py +++ b/src/autointent/advisor/_runner.py @@ -7,16 +7,18 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Callable +from typing import TYPE_CHECKING, Any from pydantic import ValidationError +from autointent._optimization_config import OptimizationConfig from autointent.advisor._estimates._resource import _resource_phase from autointent.advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules from autointent.advisor._report import PreflightReport, Severity -from autointent._optimization_config import OptimizationConfig if TYPE_CHECKING: + from collections.abc import Callable + from autointent.advisor._hardware import HardwareProfile from autointent.advisor._report import DatasetStats @@ -111,6 +113,11 @@ def _validated_config(config: dict[str, Any]) -> OptimizationConfig: return OptimizationConfig.model_validate({"search_space": []}) +# Warn about wasted HPO budget only when trials outnumber unique configs by 4x +# or more; below that the duplicate count is small enough to ignore. +_MIN_DUPLICATE_TRIAL_RATIO = 4 + + def _config_phase( search_space: list[dict[str, Any]], n_jobs: int, @@ -144,7 +151,11 @@ def _config_phase( if module in {"argmax", "threshold", "jinoos", "tunable", "adaptive"}: continue # decision modules are cheap and often singleton by design cardinality = _module_cardinality(entry) - if cardinality is not None and cardinality < n_trials and n_trials // max(1, cardinality) >= 4: + if ( + cardinality is not None + and cardinality < n_trials + and n_trials // max(1, cardinality) >= _MIN_DUPLICATE_TRIAL_RATIO + ): report.add( "config", Severity.TIGHT, diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index c65f0bd17..27c7d705e 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -295,13 +295,41 @@ def __init__(self, message: str, *, pruned_config: dict[str, Any], last_report: self.pruned_config = pruned_config self.last_report = last_report + @classmethod + def nothing_droppable(cls, *, pruned_config: dict[str, Any], last_report: PreflightReport) -> ReduceToFitError: + return cls( + "No droppable scoring-node module found; remaining search space cannot be reduced further.", + pruned_config=pruned_config, + last_report=last_report, + ) + + @classmethod + def scoring_exhausted(cls, *, pruned_config: dict[str, Any], last_report: PreflightReport) -> ReduceToFitError: + return cls( + "All scoring modules were pruned to fit the budget; the resulting pipeline would have " + "nothing to run. Raise the budget or add cheaper scoring modules.", + pruned_config=pruned_config, + last_report=last_report, + ) + + @classmethod + def not_converged( + cls, max_iters: int, *, pruned_config: dict[str, Any], last_report: PreflightReport + ) -> ReduceToFitError: + return cls( + f"Search space still infeasible after {max_iters} prune iterations.", + pruned_config=pruned_config, + last_report=last_report, + ) + def _drop_module_from_search_space( search_space: list[dict[str, Any]], node_type: str, module_name: str, ) -> list[dict[str, Any]]: - """Return a deep-copied search_space with ``module_name`` removed from the - matching ``node_type`` node. Nodes whose ``search_space`` becomes empty are - dropped entirely so the pipeline stays valid. + """Return a deep-copied search_space with ``module_name`` removed. + + Only the matching ``node_type`` node is touched. Nodes whose ``search_space`` + becomes empty are dropped entirely so the pipeline stays valid. """ import copy @@ -408,11 +436,7 @@ def reduce_to_fit( for _ in range(max_iters): pick = _pick_module_to_drop(report) if pick is None: - raise ReduceToFitError( - "No droppable scoring-node module found; remaining search space cannot be reduced further.", - pruned_config=current, - last_report=report, - ) + raise ReduceToFitError.nothing_droppable(pruned_config=current, last_report=report) node_type, module_name = pick current["search_space"] = _drop_module_from_search_space( current["search_space"], node_type, module_name, @@ -422,27 +446,20 @@ def reduce_to_fit( # would look "feasible" to run_preflight (no drivers, no findings), so # explicitly rule it out: an empty pipeline can't score anything. if not _has_scoring_module(current): - raise ReduceToFitError( - "All scoring modules were pruned to fit the budget; the resulting pipeline " - "would have nothing to run. Raise the budget or add cheaper scoring modules.", - pruned_config=current, - last_report=report, - ) + raise ReduceToFitError.scoring_exhausted(pruned_config=current, last_report=report) report = run_preflight(current, stats, hardware, refit_after=refit_after) if report.is_feasible: return current, report - raise ReduceToFitError( - f"Search space still infeasible after {max_iters} prune iterations.", - pruned_config=current, - last_report=report, - ) + raise ReduceToFitError.not_converged(max_iters, pruned_config=current, last_report=report) def _has_scoring_module(config: dict[str, Any]) -> bool: - """True when ``config`` has at least one scoring-node entry left. Empty - scoring is a common outcome of pruning to the bone — reduce_to_fit treats - it as unfittable rather than "feasible with nothing to do.""" + """Return True when ``config`` has at least one scoring-node entry left. + + Empty scoring is a common outcome of pruning to the bone — ``reduce_to_fit`` + treats it as unfittable rather than "feasible with nothing to do". + """ for node in config.get("search_space", []): if node.get("node_type") == "scoring" and node.get("search_space"): return True diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index 318530dde..d552759ac 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -717,7 +717,7 @@ def test_classic_entry_gets_synthetic_embedder_forward(self) -> None: assert linear_with["time_hours"] >= linear_no["time_hours"] def test_disk_embedding_cache_scales_with_n_samples(self) -> None: - """``disk_embedding_cache_gb`` ~ n_samples × hidden_size × 4 bytes per embedder.""" + """``disk_embedding_cache_gb`` ~ n_samples x hidden_size x 4 bytes per embedder.""" cfg = { "search_space": [ self._embedder_node(), @@ -806,7 +806,7 @@ def _run(hidden: int) -> float: "hpo_config": {"n_trials": 5}, } report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) - return next(d["vram_gb"] for d in report.resource.drivers if d["module"] == "rnn") + return float(next(d["vram_gb"] for d in report.resource.drivers if d["module"] == "rnn")) assert _run(1024) > _run(128), "larger hidden_dim must produce a larger VRAM row" @@ -897,7 +897,7 @@ def test_all_singleton(self) -> None: def test_multi_list_multiplies(self) -> None: from autointent.advisor._estimates._search_space import _module_cardinality - # 2 batch × 3 lr candidates = 6 unique configs + # 2 batch x 3 lr candidates = 6 unique configs cardinality = _module_cardinality( {"module_name": "bert", "batch_size": [32, 64], "learning_rate": [1e-5, 5e-5, 1e-4]} ) @@ -960,7 +960,7 @@ def test_no_finding_when_search_space_has_range(self) -> None: assert no_op == [], f"unexpected warning for ranged search space: {[f.message for f in no_op]}" def test_no_finding_when_n_trials_matches_cardinality(self) -> None: - # n_trials=4, cardinality=2×2=4 → not a "no-op" waste + # n_trials=4, cardinality=2x2=4 → not a "no-op" waste cfg = { "search_space": [ {"node_type": "scoring", "search_space": [ diff --git a/tests/advisor/test_public_surface.py b/tests/advisor/test_public_surface.py index 59b848921..83b31151a 100644 --- a/tests/advisor/test_public_surface.py +++ b/tests/advisor/test_public_surface.py @@ -7,7 +7,7 @@ from __future__ import annotations -import autointent.advisor as advisor +from autointent import advisor EXPECTED_SURFACE = { # functions diff --git a/tests/advisor/test_reduce_to_fit.py b/tests/advisor/test_reduce_to_fit.py index 4dd05ffe1..1a4d5e161 100644 --- a/tests/advisor/test_reduce_to_fit.py +++ b/tests/advisor/test_reduce_to_fit.py @@ -12,6 +12,8 @@ from __future__ import annotations +from typing import Any + import pytest from autointent.advisor import ( @@ -42,7 +44,7 @@ def _profile(vram_gb: float = 16.0, ram_gb: float = 32.0, free_disk_gb: float = ) -def _cheap_config() -> dict: +def _cheap_config() -> dict[str, Any]: return { "search_space": [ { @@ -59,7 +61,7 @@ def _cheap_config() -> dict: } -def _big_and_cheap_config() -> dict: +def _big_and_cheap_config() -> dict[str, Any]: """One expensive transformer + one cheap classic scorer. On a tiny (1 GB) VRAM budget, the transformer trips OVER; ``reduce_to_fit`` @@ -89,7 +91,7 @@ def _big_and_cheap_config() -> dict: } -def _unfittable_config() -> dict: +def _unfittable_config() -> dict[str, Any]: return { "search_space": [ { From 14ef1cc6c34e8c4dafeab49f3aeac007631a331d Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 19:49:27 +0300 Subject: [PATCH 37/43] refactor: split _resource_phase and _apply_embedding_cache _resource_phase took 12 keyword arguments and ran to 59 statements / 17 branches / complexity 19; _apply_embedding_cache hit complexity 12. Bundle the config-shaped inputs into a frozen _ResourceInputs and extract the two estimation passes plus the embedding-cache first-pay bookkeeping. Pure restructuring -- verified byte-identical estimates across all 10 bundled presets before and after. No noqa suppressions. Co-Authored-By: Claude Opus 5 (1M context) --- .../advisor/_estimates/_resource.py | 260 ++++++++++++------ src/autointent/advisor/_runner.py | 28 +- 2 files changed, 196 insertions(+), 92 deletions(-) diff --git a/src/autointent/advisor/_estimates/_resource.py b/src/autointent/advisor/_estimates/_resource.py index d994d1f67..89205a256 100644 --- a/src/autointent/advisor/_estimates/_resource.py +++ b/src/autointent/advisor/_estimates/_resource.py @@ -380,6 +380,68 @@ def _estimate_nn_entry( ) +def _probe_warm_models( + seen_models: dict[str, ModelMeta], + cache_probe: Callable[[str], bool] | None, +) -> set[str]: + """Model names the probe reports as already having embeddings on disk. + + Pre-populating these lets the first-seen module hit the cache-hit branch + instead of paying the forward, and skips them in the disk-cache aggregation + (already on disk). Without a probe nothing is warm — the pessimistic cold + assumption the advisor shipped with. + """ + warm_models: set[str] = set() + if cache_probe is not None: + for name in seen_models: + if cache_probe(name): + warm_models.add(name) + return warm_models + + +def _mark_cache_hit(me: _ModuleEstimate, module: str, *, suffix: str) -> None: + """Zero a transformer entry's forward time because the embedding is cached. + + Only modules whose per-entry estimate bundles the embedder forward into + ``time_hours`` have anything to give back; the rest are left alone. + """ + if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: + me.time_hours = 0.0 + me.driver["time_hours"] = 0.0 + me.driver["mode"] = f"{me.driver['mode']}+{suffix}" + + +def _charge_first_forward( + me: _ModuleEstimate, + module: str, + model: str, + *, + seen_models: dict[str, ModelMeta], + stats: DatasetStats, + hardware: HardwareProfile, +) -> None: + """Add a synthetic embedder forward to the entry that first pays for ``model``. + + Per unique embedder the first cache-honoring entry pays the forward; later + transformer entries hit the cache, and classic entries need a forward added + because their own cost model assumes embeddings already exist. + """ + if module in {"linear", "catboost"}: + embedder_meta = seen_models.get(model) + forward_h = _time_for_transformer( + n_trials=1, + epochs=1, + batch_size=32, + seq_len=128, + n_samples=stats.n_samples, + params_millions=(embedder_meta.total_params / 1_000_000) if embedder_meta else 100.0, + device_class=hardware.device_class, + ) + me.time_hours += forward_h + me.driver["time_hours"] = round(me.time_hours, 2) + me.driver["mode"] = f"{me.driver['mode']}+embed" + + def _apply_embedding_cache( module_estimates: list[_ModuleEstimate], seen_models: dict[str, ModelMeta], @@ -405,14 +467,7 @@ def _apply_embedding_cache( (i.e. contributed to ``disk_embedding_cache_gb`` in the disk aggregation). """ paid: set[str] = set() - # Models the probe reports as already-warm — pre-populate ``paid`` so the - # first-seen module also hits the cache-hit branch instead of paying the - # forward, and skip them in the disk-cache aggregation (already on disk). - warm_models: set[str] = set() - if cache_probe is not None: - for name in seen_models: - if cache_probe(name): - warm_models.add(name) + warm_models = _probe_warm_models(seen_models, cache_probe) for me in module_estimates: module = me.driver["module"] if module not in _CACHE_HONORING_MODULES: @@ -421,32 +476,13 @@ def _apply_embedding_cache( if model not in seen_models: # synthetic / "(no embedder)" rows continue if model in warm_models: - if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: - me.time_hours = 0.0 - me.driver["time_hours"] = 0.0 - me.driver["mode"] = f"{me.driver['mode']}+warm" + _mark_cache_hit(me, module, suffix="warm") continue if model in paid: - if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: - me.time_hours = 0.0 - me.driver["time_hours"] = 0.0 - me.driver["mode"] = f"{me.driver['mode']}+cached" + _mark_cache_hit(me, module, suffix="cached") else: paid.add(model) - if module in {"linear", "catboost"}: - embedder_meta = seen_models.get(model) - forward_h = _time_for_transformer( - n_trials=1, - epochs=1, - batch_size=32, - seq_len=128, - n_samples=stats.n_samples, - params_millions=(embedder_meta.total_params / 1_000_000) if embedder_meta else 100.0, - device_class=hardware.device_class, - ) - me.time_hours += forward_h - me.driver["time_hours"] = round(me.time_hours, 2) - me.driver["mode"] = f"{me.driver['mode']}+embed" + _charge_first_forward(me, module, model, seen_models=seen_models, stats=stats, hardware=hardware) return paid @@ -602,60 +638,53 @@ def _not_estimated_row(*, node_type: str, module: str) -> _ModuleEstimate: ) -def _resource_phase( - *, - embedder_config: EmbedderConfig, - search_space: list[dict[str, Any]], - n_trials: int, - n_jobs: int, - dump_modules: bool, - stats: DatasetStats, - hardware: HardwareProfile, - report: PreflightReport, - refit_after: bool = False, - cross_encoder_model_name: str | None = None, - transformer_model_name: str | None = None, - cache_probe: Callable[[str], bool] | None = None, -) -> None: - """Walk the validated search space, fold per-module costs into the report. +@dataclass(frozen=True) +class _ResourceInputs: + """Configuration-shaped inputs to the resource phase. - Two passes: transformer-bearing modules first (collects ``seen_models`` so - the largest model can drive ``embedder_dim`` for the classic pass), then - linear / catboost. Disk, VRAM/RAM peak, time sum, and final findings are - folded onto the report. + Bundled because these travel together through both estimation passes and + passing nine keyword arguments down each one is unreadable. ``cross_encoder_model_name`` and ``transformer_model_name`` come from the - pipeline's top-level configs. They're used as the fallback model for - modules that don't declare a per-entry ``classification_model_config`` but - still consume one at runtime (``description_cross`` / ``dnnc`` / - ``retrieval`` pull from ``cross_encoder_config``; ``bert`` falls back to - ``transformer_config``). Seeding them here fixes the disk-download - under-count called out in the follow-up review (missing 6.4 GB reranker in - ``zero-shot-encoders``). + pipeline's top-level configs and act as the fallback model for modules that + don't declare a per-entry ``classification_model_config`` but still consume + one at runtime (``description_cross`` / ``dnnc`` / ``retrieval`` pull from + ``cross_encoder_config``; ``bert`` falls back to ``transformer_config``). + Seeding them here fixes the disk-download under-count called out in the + follow-up review (missing 6.4 GB reranker in ``zero-shot-encoders``). """ - seen_models: dict[str, ModelMeta] = {} - global_embedder = _embedder_model_name(embedder_config) - if global_embedder: - seen_models[global_embedder] = _hub.resolve_model(global_embedder) - transformer_entries, classic_entries = _split_entries(search_space) + embedder_config: EmbedderConfig + search_space: list[dict[str, Any]] + n_trials: int + n_jobs: int + dump_modules: bool + refit_after: bool = False + cross_encoder_model_name: str | None = None + transformer_model_name: str | None = None + cache_probe: Callable[[str], bool] | None = None - # HPO distributes n_trials evenly across module_name candidates at each - # node, so each variant sees n_trials / n_variants on average. - variants_per_node: dict[int, int] = {} - for node_idx, _node_type, _entry in _walk_modules_indexed(search_space): - variants_per_node[node_idx] = variants_per_node.get(node_idx, 0) + 1 - def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int: # noqa: ARG001 - """Trials charged to this module. +def _estimate_transformer_entries( + transformer_entries: list[tuple[int, str, dict[str, Any]]], + inputs: _ResourceInputs, + stats: DatasetStats, + hardware: HardwareProfile, + seen_models: dict[str, ModelMeta], + effective_trials: Callable[[int, dict[str, Any] | None], int], +) -> tuple[list[_ModuleEstimate], dict[int, float]]: + """First pass: transformer-bearing modules. - ``entry`` is reserved for future per-module caps (see git history for the - cardinality-cap experiment that broke description-scorer presets). - """ - divisor = max(1, variants_per_node.get(node_idx, 1)) - return max(1, n_trials // divisor) + Also populates ``seen_models`` in place, which the classic pass reads to + derive ``embedder_dim`` from the largest model seen — so this must run first. + + Returns ``(module_estimates, node_max_weights)``. + """ + global_embedder = _embedder_model_name(inputs.embedder_config) + cross_encoder_model_name = inputs.cross_encoder_model_name + transformer_model_name = inputs.transformer_model_name + refit_after = inputs.refit_after - # First pass: transformer modules (also populates seen_models for the classic pass). module_estimates: list[_ModuleEstimate] = [] node_max_weights: dict[int, float] = {} for node_idx, node_type, entry in transformer_entries: @@ -676,7 +705,7 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int node_type=node_type, stats=stats, hardware=hardware, - n_trials=_effective_trials(node_idx, entry), + n_trials=effective_trials(node_idx, entry), refit_after=refit_after, ) if nn_estimate is not None: @@ -696,15 +725,33 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int name=name, stats=stats, hardware=hardware, - n_trials=_effective_trials(node_idx, entry), + n_trials=effective_trials(node_idx, entry), refit_after=refit_after, ) module_estimates.append(me) # Track heaviest weight per node so dump_modules is bounded by one # selected variant per node x n_trials, not the sum of all candidates. node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), me.model_weights_gb) + return module_estimates, node_max_weights - # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. + +def _estimate_classic_entries( + classic_entries: list[tuple[int, str, dict[str, Any]]], + inputs: _ResourceInputs, + stats: DatasetStats, + hardware: HardwareProfile, + seen_models: dict[str, ModelMeta], + effective_trials: Callable[[int, dict[str, Any] | None], int], +) -> list[_ModuleEstimate]: + """Second pass: linear / catboost / sklearn modules. + + Their cost depends on ``embedder_dim`` rather than on a checkpoint, so this + reads ``seen_models`` as populated by :func:`_estimate_transformer_entries` + and derives the dimension from the largest embedder seen there. + """ + refit_after = inputs.refit_after + + module_estimates: list[_ModuleEstimate] = [] embedder_meta = _largest_embedder(seen_models) embedder_dim_val = _embedder_dim(embedder_meta) for node_idx, node_type, entry in classic_entries: @@ -715,11 +762,66 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int embedder_dim=embedder_dim_val, stats=stats, hardware=hardware, - n_trials=_effective_trials(node_idx, entry), + n_trials=effective_trials(node_idx, entry), refit_after=refit_after, ) if classic_estimate is not None: module_estimates.append(classic_estimate) + return module_estimates + + +def _resource_phase( + inputs: _ResourceInputs, + stats: DatasetStats, + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + """Walk the validated search space, fold per-module costs into the report. + + Two passes: transformer-bearing modules first (collects ``seen_models`` so + the largest model can drive ``embedder_dim`` for the classic pass), then + linear / catboost. Disk, VRAM/RAM peak, time sum, and final findings are + folded onto the report. + """ + embedder_config = inputs.embedder_config + search_space = inputs.search_space + n_trials = inputs.n_trials + n_jobs = inputs.n_jobs + dump_modules = inputs.dump_modules + cache_probe = inputs.cache_probe + + seen_models: dict[str, ModelMeta] = {} + global_embedder = _embedder_model_name(embedder_config) + if global_embedder: + seen_models[global_embedder] = _hub.resolve_model(global_embedder) + + transformer_entries, classic_entries = _split_entries(search_space) + + # HPO distributes n_trials evenly across module_name candidates at each + # node, so each variant sees n_trials / n_variants on average. + variants_per_node: dict[int, int] = {} + for node_idx, _node_type, _entry in _walk_modules_indexed(search_space): + variants_per_node[node_idx] = variants_per_node.get(node_idx, 0) + 1 + + def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int: # noqa: ARG001 + """Trials charged to this module. + + ``entry`` is reserved for future per-module caps (see git history for the + cardinality-cap experiment that broke description-scorer presets). + """ + divisor = max(1, variants_per_node.get(node_idx, 1)) + return max(1, n_trials // divisor) + + # First pass: transformer modules (also populates seen_models for the classic pass). + module_estimates, node_max_weights = _estimate_transformer_entries( + transformer_entries, inputs, stats, hardware, seen_models, _effective_trials, + ) + + # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. + embedder_meta = _largest_embedder(seen_models) + module_estimates += _estimate_classic_entries( + classic_entries, inputs, stats, hardware, seen_models, _effective_trials, + ) # Cache-aware time/disk: must run before the fold below. cached_embedders = _apply_embedding_cache( diff --git a/src/autointent/advisor/_runner.py b/src/autointent/advisor/_runner.py index 464f471a1..416cca80a 100644 --- a/src/autointent/advisor/_runner.py +++ b/src/autointent/advisor/_runner.py @@ -12,7 +12,7 @@ from pydantic import ValidationError from autointent._optimization_config import OptimizationConfig -from autointent.advisor._estimates._resource import _resource_phase +from autointent.advisor._estimates._resource import _resource_phase, _ResourceInputs from autointent.advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules from autointent.advisor._report import PreflightReport, Severity @@ -79,18 +79,20 @@ def run_preflight( report.notes.extend(hardware.notes) _resource_phase( - embedder_config=cfg.embedder_config, - search_space=cfg.search_space, - n_trials=cfg.hpo_config.n_trials, - n_jobs=cfg.hpo_config.n_jobs, - dump_modules=cfg.logging_config.dump_modules, - stats=stats, - hardware=hardware, - report=report, - refit_after=refit_after, - cross_encoder_model_name=cfg.cross_encoder_config.model_name, - transformer_model_name=cfg.transformer_config.model_name, - cache_probe=embedding_cache_probe, + _ResourceInputs( + embedder_config=cfg.embedder_config, + search_space=cfg.search_space, + n_trials=cfg.hpo_config.n_trials, + n_jobs=cfg.hpo_config.n_jobs, + dump_modules=cfg.logging_config.dump_modules, + refit_after=refit_after, + cross_encoder_model_name=cfg.cross_encoder_config.model_name, + transformer_model_name=cfg.transformer_config.model_name, + cache_probe=embedding_cache_probe, + ), + stats, + hardware, + report, ) _data_phase(cfg.search_space, stats, report) _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, cfg.hpo_config.n_trials, hardware, report) From 285160b4dbcc6c639356c854d614870f40b4edad Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 20:10:34 +0300 Subject: [PATCH 38/43] feat: make the Pipeline.fit preflight gate opt-in preflight defaulted to 'warn', so every fit() ran the advisor -- and resolve_model() calls HfApi().model_info() unconditionally per distinct model name, with no cache-first short-circuit. That put N Hub round-trips and, when offline, one WARNING per model on the library's hottest path, for estimates that are explicitly heuristic. Default to 'off'; all three modes still work. Also make the advisor import lazy so 'import autointent' never pulls in huggingface_hub probes, and rewrite the preflight tests: they previously ran a real classic-light fit inside 'except Exception: pass', so they passed even when the fit failed for unrelated reasons. --- src/autointent/_pipeline/__init__.py | 4 +- src/autointent/_pipeline/_pipeline.py | 38 +++++------ tests/pipeline/test_preflight.py | 91 ++++++++++++++------------- 3 files changed, 71 insertions(+), 62 deletions(-) diff --git a/src/autointent/_pipeline/__init__.py b/src/autointent/_pipeline/__init__.py index 50d58f7d1..cc7fe54eb 100644 --- a/src/autointent/_pipeline/__init__.py +++ b/src/autointent/_pipeline/__init__.py @@ -1,3 +1,3 @@ -from ._pipeline import Pipeline, PreflightError, PreflightMode +from ._pipeline import Pipeline, PreflightMode -__all__ = ["Pipeline", "PreflightError", "PreflightMode"] +__all__ = ["Pipeline", "PreflightMode"] diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index 3c1bbfd68..8704aab2c 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -12,13 +12,6 @@ from typing_extensions import assert_never from autointent import Context, OptimizationConfig -from autointent.advisor import ( - PreflightError, - Severity, - dataset_stats, - detect_hardware, - run_preflight, -) from autointent.configs import ( CrossEncoderConfig, DataConfig, @@ -41,7 +34,7 @@ if TYPE_CHECKING: from autointent import Dataset - from autointent.advisor import PreflightReport + from autointent.advisor import Finding, PreflightReport from autointent.custom_types import ListOfGenericLabels, SearchSpacePreset, SearchSpaceValidationMode from autointent.modules.base import BaseDecision, BaseRegex, BaseScorer @@ -185,14 +178,19 @@ def _run_preflight(self, dataset: Dataset, *, refit_after: bool, mode: Preflight Logs each finding at INFO/WARNING/ERROR (by severity). When ``mode`` is ``"strict"`` and any OVER finding is produced, raises ``PreflightError``. + + Imported lazily: the advisor probes the HF Hub, so it must stay off the + ``import autointent`` path. """ + from autointent.advisor import PreflightError, Severity, dataset_stats, detect_hardware, run_preflight + config = self._build_advisor_config() stats = dataset_stats(dataset) hardware = detect_hardware() report = run_preflight(config, stats, hardware, refit_after=refit_after) _log_preflight_report(report, self._logger) if mode == "strict": - over = [f for f in report.findings if f.severity == Severity.OVER] + over: list[Finding] = [f for f in report.findings if f.severity == Severity.OVER] if over: raise PreflightError(over) return report @@ -238,21 +236,23 @@ def fit( dataset: Dataset, refit_after: bool = False, incompatible_search_space: SearchSpaceValidationMode = "filter", - preflight: PreflightMode = "warn", + preflight: PreflightMode = "off", ) -> Context: """Optimize the pipeline from dataset. Args: dataset: dataset for optimization. refit_after: whether to refit on whole data after optimization. Valid only for hold-out validaiton. - sampler: sampler type to use. - incompatible_search_space: wow to handle data-incompatible modules occurring in search space. - preflight: gate that runs :func:`autointent.advisor.run_preflight` over the - pipeline's effective config + dataset before any heavy work. - ``"off"`` skips it. ``"warn"`` (default) logs findings — INFO for - AMPLE, WARNING for TIGHT, ERROR for OVER — but never raises. - ``"strict"`` additionally raises :class:`PreflightError` when any - finding has severity OVER, so unfeasible runs abort before fit. + incompatible_search_space: how to handle data-incompatible modules occurring in search space. + preflight: **experimental** gate that runs + :func:`autointent.advisor.run_preflight` over the pipeline's + effective config + dataset before any heavy work. + ``"off"`` (default) skips it entirely. ``"warn"`` logs findings — + INFO for AMPLE, WARNING for TIGHT, ERROR for OVER — but never + raises; note it probes the HF Hub for model metadata, so it adds + network round-trips. ``"strict"`` additionally raises + :class:`autointent.advisor.PreflightError` when any finding has + severity OVER, so unfeasible runs abort before fit. Raises: RuntimeError: If pipeline is in inference mode. @@ -532,6 +532,8 @@ def make_report(logs: dict[str, Any], nodes: list[NodeType]) -> str: def _log_preflight_report(report: PreflightReport, logger: logging.Logger) -> None: """Log each preflight finding at the appropriate level.""" + from autointent.advisor import Severity + level_for = { Severity.AMPLE: logging.INFO, Severity.TIGHT: logging.WARNING, diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py index 9c366cf1f..747241d87 100644 --- a/tests/pipeline/test_preflight.py +++ b/tests/pipeline/test_preflight.py @@ -1,20 +1,35 @@ -"""Pipeline.fit preflight integration: off / warn / strict modes.""" +"""Pipeline.fit preflight integration: default-off, warn, strict. + +``fit()`` is driven with ``Pipeline._fit`` stubbed out, so these tests exercise +the preflight gate (which runs before any heavy work) without training anything. +With ``clear_ram=True, dump_modules=False`` the post-``_fit`` branch returns the +context immediately, so a stubbed ``_fit`` leaves ``fit()`` fully functional. +""" from __future__ import annotations import logging +import subprocess +import sys from typing import TYPE_CHECKING import pytest from autointent import Pipeline -from autointent.advisor import HardwareProfile, dataset_stats, detect_hardware, run_preflight -from autointent._pipeline import PreflightError +from autointent.advisor import HardwareProfile, PreflightError, dataset_stats, detect_hardware, run_preflight from autointent.configs import LoggingConfig if TYPE_CHECKING: from autointent import Dataset +_PIPELINE_LOGGER = "autointent._pipeline._pipeline" + + +@pytest.fixture(autouse=True) +def _stub_fit(monkeypatch: pytest.MonkeyPatch) -> None: + """Skip optimization; every test here is about the gate that runs before it.""" + monkeypatch.setattr(Pipeline, "_fit", lambda _self, _context: None) + def _tiny_hw() -> HardwareProfile: """Deterministic, intentionally-infeasible hardware budget.""" @@ -34,35 +49,32 @@ def _classic_light_pipeline() -> Pipeline: return p +def test_fit_does_not_run_preflight_by_default(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + """The default is opt-out: no preflight, no Hub round-trips, no log line.""" + p = _classic_light_pipeline() + with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): + p.fit(dataset) + assert not any("Preflight" in r.getMessage() for r in caplog.records) + + def test_preflight_off_skips_advisor(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: - """preflight='off' must not run the advisor (no Preflight log line).""" p = _classic_light_pipeline() - with caplog.at_level(logging.INFO, logger="autointent._pipeline._pipeline"): - try: - p.fit(dataset, preflight="off") - except Exception: # noqa: BLE001 — fit may fail in test env; we only care about preflight side effect - pass + with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): + p.fit(dataset, preflight="off") assert not any("Preflight" in r.getMessage() for r in caplog.records) -def test_preflight_warn_logs_findings(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: - """preflight='warn' logs a Preflight verdict line.""" +def test_preflight_warn_logs_verdict(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: p = _classic_light_pipeline() - with caplog.at_level(logging.INFO, logger="autointent._pipeline._pipeline"): - try: - p.fit(dataset, preflight="warn") - except Exception: # noqa: BLE001 - pass + with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): + p.fit(dataset, preflight="warn") msgs = [r.getMessage() for r in caplog.records] assert any("Preflight" in m and "verdict=" in m for m in msgs) def test_preflight_strict_raises_on_infeasible(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: - """preflight='strict' raises PreflightError when findings include OVER. - - Forces a tiny hardware budget so even cheap presets blow it. - """ - monkeypatch.setattr("autointent._pipeline._pipeline.detect_hardware", _tiny_hw) + """Patched at the advisor, not the pipeline: the import is lazy now.""" + monkeypatch.setattr("autointent.advisor.detect_hardware", _tiny_hw) p = _classic_light_pipeline() with pytest.raises(PreflightError) as exc_info: p.fit(dataset, preflight="strict") @@ -73,27 +85,28 @@ def test_preflight_strict_raises_on_infeasible(dataset: Dataset, monkeypatch: py def test_preflight_warn_does_not_raise_on_infeasible( dataset: Dataset, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: - """Tiny hardware + warn mode logs an ERROR but doesn't raise.""" - monkeypatch.setattr("autointent._pipeline._pipeline.detect_hardware", _tiny_hw) + monkeypatch.setattr("autointent.advisor.detect_hardware", _tiny_hw) p = _classic_light_pipeline() - with caplog.at_level(logging.ERROR, logger="autointent._pipeline._pipeline"): - try: - p.fit(dataset, preflight="warn") - except PreflightError: - pytest.fail("warn mode must not raise PreflightError") - except Exception: # noqa: BLE001 — downstream fit errors are out of scope - pass + with caplog.at_level(logging.ERROR, logger=_PIPELINE_LOGGER): + p.fit(dataset, preflight="warn") assert any(r.levelno == logging.ERROR for r in caplog.records) -def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: - """End-to-end integration: Pipeline -> _build_advisor_config -> run_preflight. +def test_importing_autointent_does_not_import_the_advisor() -> None: + """The advisor pulls in huggingface_hub probes; it must stay off the import path. - Asserts the round-trip is wired correctly: the dict ``Pipeline`` exposes to - the advisor validates against ``OptimizationConfig``, the advisor produces a - well-formed report, and the driver list reflects the actual modules from the - preset's search space (not silently empty). + Checked in a subprocess because pytest has already imported the advisor into + this process. Asserting on ``huggingface_hub`` itself would not work -- + ``datasets`` imports it regardless -- so the subpackage's own absence is the + real invariant. """ + code = "import autointent, sys; print('autointent.advisor' in sys.modules)" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert result.stdout.strip() == "False", "importing autointent must not import autointent.advisor" + + +def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: + """End-to-end: Pipeline -> _build_advisor_config -> run_preflight.""" p = _classic_light_pipeline() config = p._build_advisor_config() stats = dataset_stats(dataset) @@ -101,20 +114,14 @@ def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: report = run_preflight(config, stats, hardware, preset_name="classic-light") - # The advisor accepted the pipeline-built config and produced findings. assert report.preset_name == "classic-light" assert report.resource.drivers, "expected at least one driver row for classic-light" - # classic-light's scoring node has knn / linear / mlknn — at least linear - # should always end up in drivers (knn variants don't always carry an - # explicit model_name, so they're allowed to be absent). driver_modules = {d["module"] for d in report.resource.drivers} assert "linear" in driver_modules, f"missing linear scorer in drivers: {driver_modules}" - # The advisor must always emit the three resource findings. metrics = {f.metric for f in report.findings if f.metric} assert {"vram", "ram", "disk"} <= metrics, f"missing required metrics: {metrics}" - # Dataset stats round-trip into the report. assert report.dataset["n_samples"] == stats.n_samples assert report.dataset["n_classes"] == stats.n_classes From ff76750631fde04ffec4dcb6236d34e0746a9120 Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 20:30:37 +0300 Subject: [PATCH 39/43] docs: add the compute feasibility advisor page Renaming _advisor to advisor makes autoapi publish the package, so the public docstrings are now user-facing. Add a prose page covering the CLI, reading a report, the Python API, and the Pipeline.fit gate. Documents the accuracy limits measured in Darinochka/AutoIntent-experiments#40 and softens the 'pessimistic upper bound' claims -- one preset measured 1.22x its predicted VRAM, so that guarantee doesn't hold. Written as docs/source/advisor.rst rather than a user_guides page: those are jupytext-executed during the docs build and would have to run real fits. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/advisor.rst | 114 ++++++++++++++++++ docs/source/index.rst | 4 + .../advisor/_estimates/_formulas.py | 22 ++-- src/autointent/advisor/_hub.py | 4 +- src/autointent/advisor/_render.py | 2 +- tests/advisor/test_render.py | 2 +- 6 files changed, 134 insertions(+), 14 deletions(-) create mode 100644 docs/source/advisor.rst diff --git a/docs/source/advisor.rst b/docs/source/advisor.rst new file mode 100644 index 000000000..b5517e585 --- /dev/null +++ b/docs/source/advisor.rst @@ -0,0 +1,114 @@ +Compute feasibility advisor +=========================== + +.. note:: + + **Experimental.** The advisor's estimates are heuristic and calibrated + against a limited hardware sample. Treat them as guidance, not guarantees, + and read :ref:`advisor-accuracy` before relying on a number. The Python + surface may change in a minor release. + +Optimizing a search space can take hours and needs more VRAM than a laptop GPU +has. The advisor answers "will this fit, and how long will it take?" *before* +anything is downloaded or trained. + +Command line +------------ + +Two subcommands. ``inspect`` prices a specific preset or config; ``recommend`` +detects your hardware and picks the heaviest bundled preset that still fits. + +.. code-block:: bash + + # What will transformers-light cost on this machine? + autointent-advisor inspect transformers-light + + # ...against a real dataset rather than placeholder sizes + autointent-advisor inspect transformers-light --dataset banking77 + + # Which preset should I use? + autointent-advisor recommend --dataset banking77 + + # Machine-readable output + autointent-advisor inspect ./my-config.yaml --json + +Without ``--dataset``, the advisor uses placeholder dataset sizes +(``--n-samples``, ``--n-classes``, ``--avg-tokens``, ``--task``), so it is +useful before you have assembled any data. ``--budget-vram-gb`` overrides +hardware detection, and ``recommend`` also accepts ``--budget-time-h``. + +Both subcommands exit non-zero when nothing is feasible, so they work as a CI +gate. + +Reading a report +---------------- + +Each finding carries a severity: + +``ample`` + Comfortably within budget. +``tight`` + Fits, but with little headroom — expect swapping or thermal throttling. +``over`` + Exceeds the budget. Any ``over`` finding makes the whole report infeasible. + +The drivers table lists the modules that dominate the cost, so it shows *what* +to change. ``low confidence`` on a report means Hub metadata was unavailable or +incomplete for at least one model, and conservative large-model defaults were +substituted — the numbers are much rougher when you see it. + +From Python +----------- + +.. code-block:: python + + from autointent import Dataset + from autointent.advisor import dataset_stats, detect_hardware, estimate, recommend + + report = estimate("transformers-light") + print(report.is_feasible, report.resource.vram_gb) + + result = recommend(stats=dataset_stats(Dataset.from_json(path))) + print(result.chosen) + +``reduce_to_fit`` goes further: it prunes the most expensive scoring module +repeatedly until the search space fits, raising ``ReduceToFitError`` if nothing +does. + +Inside ``Pipeline.fit`` +----------------------- + +``Pipeline.fit`` accepts a ``preflight`` gate. It defaults to ``"off"``, so the +advisor never runs unless you ask — it makes network calls to the Hugging Face +Hub for model metadata, which does not belong on every fit by default. + +.. code-block:: python + + pipeline.fit(dataset, preflight="warn") # log findings, always continue + pipeline.fit(dataset, preflight="strict") # raise PreflightError if infeasible + +``"strict"`` raises :class:`autointent.advisor.PreflightError` before allocating +any VRAM, which is the useful mode in CI. + +.. _advisor-accuracy: + +How accurate is it? +------------------- + +Validated end to end on one machine class (RTX 3060 Laptop, 6 GB VRAM / 16 GB +RAM), where all four fitted presets matched their predicted verdict: both +``over`` predictions did run out of memory, and both feasible predictions did +fit. Known limits: + +- **Feasibility verdicts are the reliable part.** That is what the advisor was + built and validated for. +- **VRAM is close but not a guaranteed ceiling.** One preset used 1.22× its + prediction. Leave headroom rather than trusting the figure exactly. +- **Wall-time estimates are indicative only.** Measured error has run in both + directions across formula revisions, once by more than an order of magnitude + for cross-encoders. ``--budget-time-h`` inherits that uncertainty. +- **Preset ranking does not depend on time estimates.** ``recommend`` orders + presets by a declared cost ranking, so unstable time figures cannot reorder + its choice. +- **Only one hardware class has been validated end to end.** Treat other + machines as unverified. diff --git a/docs/source/index.rst b/docs/source/index.rst index 9d0fb5eed..5254e7c6f 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -68,6 +68,9 @@ Reference :doc:`🌐 Inference servers ` Deploy a trained pipeline behind HTTP (FastAPI) or MCP (FastMCP): installation extras, environment variables, and how to run each server. +:doc:`🔍 Compute feasibility advisor ` + Estimate VRAM, RAM, disk, and wall-time for a search space before training. Run it from the CLI or gate ``Pipeline.fit`` on it. + :doc:`🔧 API Reference ` Complete technical documentation for all classes, methods, and functions. Essential reference for developers integrating AutoIntent into their applications. @@ -84,4 +87,5 @@ Reference user_guides learn/index server + advisor autoapi/autointent/index \ No newline at end of file diff --git a/src/autointent/advisor/_estimates/_formulas.py b/src/autointent/advisor/_estimates/_formulas.py index 9fabf2c11..25daf6fb8 100644 --- a/src/autointent/advisor/_estimates/_formulas.py +++ b/src/autointent/advisor/_estimates/_formulas.py @@ -10,7 +10,8 @@ * All ``*_hours`` results assume the GPU baseline of ~1 second per step; CPU runs pay a flat slowdown factor (see ``_time_for_transformer``). * "fp32 worst case" — we deliberately ignore lower-precision / FlashAttention / - quantization optimizations, per the advisor's "pessimistic upper bound" contract. + quantization optimizations, since the advisor aims to over- rather than + under-predict cost. """ from __future__ import annotations @@ -58,9 +59,9 @@ def _classify_severity(estimate: float, budget: float) -> Severity: def _weights_vram_for_transformer(meta: ModelMeta, mode: str) -> float: """Weight-side VRAM: weights + grads + optimizer state. - Pessimistic upper bound by mode: 1.3x inference, 1.3x + 0.5 GB lora - adapters, 4.5x full finetune (textbook 4W + fragmentation/workspaces - slack). + Multiplier chosen to over- rather than under-estimate, by mode: 1.3x + inference, 1.3x + 0.5 GB lora adapters, 4.5x full finetune (textbook 4W + + fragmentation/workspaces slack). """ weights_gb = meta.weights_gb if mode == "inference": @@ -133,7 +134,8 @@ def _max_fitting_batch_size( # Sustained TFLOPS per device class — real HF-Trainer MFU (~20% on A100), -# not peak spec sheet. Advisor upper-bounds, so pessimistic values here. +# not peak spec sheet. Advisor aims to over- rather than under-predict time, +# so pessimistic (low) values here. _DEVICE_TFLOPS = { "high-gpu": 60.0, # A100 / H100 "mid-gpu": 20.0, # V100 / RTX 3090 / A6000 @@ -194,10 +196,10 @@ def _largest_embedder(seen_models: dict[str, ModelMeta]) -> ModelMeta | None: def _ram_for_module(meta: ModelMeta, stats: DatasetStats, *, mode: str = "inference") -> float: - """RAM upper bound: weights x mode_mult + tokenized text (n_samples x avg_tokens x 4 B). + """RAM estimate: weights x mode_mult + tokenized text (n_samples x avg_tokens x 4 B). - Mode multiplier: 1.3 inference, 1.5 lora, 4.5 full-finetune (Adam - mirrors weights on host too). + Mode multiplier chosen to over- rather than under-estimate: 1.3 inference, + 1.5 lora, 4.5 full-finetune (Adam mirrors weights on host too). """ if mode == "inference": weights_mult = 1.3 @@ -348,8 +350,8 @@ def _rnn_param_count(*, embed_dim: int, hidden_dim: int, n_classes: int) -> int: def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: """Weights + 3x optimizer/grads + activations. - Same fp32 upper bound as transformers, smaller hidden dim (embed_dim / - num_filters). + Same fp32, over- rather than under-estimate approach as transformers, + smaller hidden dim (embed_dim / num_filters). """ weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB activations_gb = ( diff --git a/src/autointent/advisor/_hub.py b/src/autointent/advisor/_hub.py index f277a0819..58366b188 100644 --- a/src/autointent/advisor/_hub.py +++ b/src/autointent/advisor/_hub.py @@ -23,7 +23,7 @@ # Conservative "large-model" shape used when Hub metadata is unavailable — # roughly deberta-v3-large / bert-large sized. Previously we defaulted to a # BERT-base shape (110M / 768 / 12), which *under*-predicted a real deberta-large -# fit by ~2x. Because the advisor's contract is a pessimistic upper bound, the +# fit by ~2x. Because the advisor aims to over- rather than under-predict, the # offline fallback needs to over-estimate small models rather than under-estimate # large ones. Callers can still see the fallback happened via ``confidence == # "heuristic"`` and ``PreflightReport.low_confidence``. @@ -173,7 +173,7 @@ def _heuristic_metadata(model_name: str) -> ModelMeta: logger.warning( "Falling back to name-pattern heuristic for %s; " "using CONSERVATIVE large-model defaults (params=%dM, hidden=%d, layers=%d) " - "so cost estimates upper-bound rather than under-predict.", + "so cost estimates aim to over- rather than under-predict.", model_name, _DEFAULT_HEURISTIC_PARAMS // 1_000_000, _DEFAULT_HEURISTIC_HIDDEN, diff --git a/src/autointent/advisor/_render.py b/src/autointent/advisor/_render.py index afd541e2b..1d4954dfe 100644 --- a/src/autointent/advisor/_render.py +++ b/src/autointent/advisor/_render.py @@ -122,7 +122,7 @@ def render_text(report: PreflightReport) -> str: if report.low_confidence: summary += " — low-confidence (heuristic fallback in use)" lines.append(summary) - lines.append("Note: estimates are heuristic upper bounds, not measurements.") + lines.append("Note: estimates are heuristic guidance, not measurements or guarantees.") return "\n".join(lines) diff --git a/tests/advisor/test_render.py b/tests/advisor/test_render.py index 8aa47ad88..a6254af08 100644 --- a/tests/advisor/test_render.py +++ b/tests/advisor/test_render.py @@ -71,7 +71,7 @@ def test_verdict_reflects_headroom(self) -> None: def test_disclaimer_always_present(self) -> None: out = render_text(_populated_report()) - assert "heuristic upper bounds" in out + assert "heuristic guidance" in out def test_low_confidence_tag_when_offline(self) -> None: r = _populated_report() From aa7e9fd1563390fb8076f6631180f9cd30524c7a Mon Sep 17 00:00:00 2001 From: voorhs Date: Mon, 17 Aug 2026 21:12:22 +0300 Subject: [PATCH 40/43] fix: address final review findings - price the filtered search space: run validate_modules before the preflight gate, so preflight no longer charges for modules fit() will discard (mlknn on multiclass, dnnc on multilabel and its ~6.4 GB reranker) - correct reduce_to_fit's public docstring, which still advertised the pruning order this branch fixed - document Severity, HardwareProfile and four published members; Severity's API page was rendering str.__doc__ - rewrite dataset_stats' docstring, which referenced a non-public name - add the console-script name regression test and correct test_hardware_detection.py's "no psutil" claim (both spec section F) - comment the second lazy-import site so it is not tidied back to module scope - rename _charge_first_forward to _charge_first_forward_if_classic and fix three stale sentences in _resource.py; no arithmetic touched Co-Authored-By: Claude Opus 5 (1M context) --- src/autointent/_pipeline/_pipeline.py | 12 +++- .../advisor/_estimates/_resource.py | 30 ++++++---- src/autointent/advisor/_hardware.py | 9 +++ src/autointent/advisor/_report.py | 47 +++++++++++++++ src/autointent/advisor/_workflows.py | 27 +++++++-- tests/advisor/test_estimates_and_cli.py | 28 ++++++++- tests/advisor/test_hardware_detection.py | 12 +++- tests/pipeline/test_preflight.py | 57 ++++++++++++++++++- 8 files changed, 200 insertions(+), 22 deletions(-) diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index d5d25686d..0c5b9a864 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -262,6 +262,13 @@ def fit( msg = "Pipeline in inference mode cannot be fitted" raise RuntimeError(msg) + # Filter the search space first: ``validate_modules`` drops modules this + # dataset cannot use (e.g. ``mlknn`` on multiclass, ``dnnc`` and its ~6.4 GB + # reranker on multilabel), and preflight must price what will actually run + # rather than what was requested. It takes ``dataset`` only, so it does not + # depend on the ``Context`` built below. + self.validate_modules(dataset, mode=incompatible_search_space) + if preflight != "off": self._run_preflight(dataset, refit_after=refit_after, mode=preflight) @@ -274,8 +281,6 @@ def fit( context.configure_hpo(self.hpo_config) context.configure_vector_index(self.vector_index_config) - self.validate_modules(dataset, mode=incompatible_search_space) - test_utterances = context.data_handler.test_utterances() if test_utterances is None: self._logger.warning( @@ -532,6 +537,9 @@ def make_report(logs: dict[str, Any], nodes: list[NodeType]) -> str: def _log_preflight_report(report: PreflightReport, logger: logging.Logger) -> None: """Log each preflight finding at the appropriate level.""" + # Imported lazily for the same reason as in ``Pipeline._run_preflight``: the + # advisor probes the HF Hub, so it must stay off the ``import autointent`` + # path. Do not hoist to module scope. from autointent.advisor import Severity level_for = { diff --git a/src/autointent/advisor/_estimates/_resource.py b/src/autointent/advisor/_estimates/_resource.py index 89205a256..ec6a4949c 100644 --- a/src/autointent/advisor/_estimates/_resource.py +++ b/src/autointent/advisor/_estimates/_resource.py @@ -411,7 +411,7 @@ def _mark_cache_hit(me: _ModuleEstimate, module: str, *, suffix: str) -> None: me.driver["mode"] = f"{me.driver['mode']}+{suffix}" -def _charge_first_forward( +def _charge_first_forward_if_classic( me: _ModuleEstimate, module: str, model: str, @@ -420,11 +420,13 @@ def _charge_first_forward( stats: DatasetStats, hardware: HardwareProfile, ) -> None: - """Add a synthetic embedder forward to the entry that first pays for ``model``. + """Add a synthetic embedder forward to a classic entry that first pays for ``model``. - Per unique embedder the first cache-honoring entry pays the forward; later - transformer entries hit the cache, and classic entries need a forward added - because their own cost model assumes embeddings already exist. + Per unique embedder the first cache-honoring entry pays the forward. Only + classic entries (linear / catboost) need it added, because their own cost + model assumes embeddings already exist; a transformer entry's ``time_hours`` + already bundles the forward, so for those this is deliberately a no-op — + hence the guard and the ``_if_classic`` in the name. """ if module in {"linear", "catboost"}: embedder_meta = seen_models.get(model) @@ -482,7 +484,9 @@ def _apply_embedding_cache( _mark_cache_hit(me, module, suffix="cached") else: paid.add(model) - _charge_first_forward(me, module, model, seen_models=seen_models, stats=stats, hardware=hardware) + _charge_first_forward_if_classic( + me, module, model, seen_models=seen_models, stats=stats, hardware=hardware + ) return paid @@ -642,16 +646,22 @@ def _not_estimated_row(*, node_type: str, module: str) -> _ModuleEstimate: class _ResourceInputs: """Configuration-shaped inputs to the resource phase. - Bundled because these travel together through both estimation passes and - passing nine keyword arguments down each one is unreadable. + Bundled by provenance rather than by use: every field is derived from one + validated ``OptimizationConfig``, plus the caller-injected ``cache_probe``. + Individual passes read only what they need — the classic pass reads + ``refit_after`` alone, and ``search_space`` / ``n_trials`` / ``n_jobs`` / + ``dump_modules`` never leave ``_resource_phase`` — but threading them + separately would mean a dozen keyword arguments down each call. ``cross_encoder_model_name`` and ``transformer_model_name`` come from the pipeline's top-level configs and act as the fallback model for modules that don't declare a per-entry ``classification_model_config`` but still consume one at runtime (``description_cross`` / ``dnnc`` / ``retrieval`` pull from ``cross_encoder_config``; ``bert`` falls back to ``transformer_config``). - Seeding them here fixes the disk-download under-count called out in the - follow-up review (missing 6.4 GB reranker in ``zero-shot-encoders``). + Carrying them is what fixes the disk-download under-count called out in the + follow-up review (missing 6.4 GB reranker in ``zero-shot-encoders``); they + are applied in :func:`_estimate_transformer_entries`, where an entry with no + model of its own falls back to them. """ embedder_config: EmbedderConfig diff --git a/src/autointent/advisor/_hardware.py b/src/autointent/advisor/_hardware.py index 6aa9741ee..73f247fcd 100644 --- a/src/autointent/advisor/_hardware.py +++ b/src/autointent/advisor/_hardware.py @@ -32,6 +32,15 @@ @dataclass class HardwareProfile: + """The machine budget every estimate is scored against. + + Produced by :func:`~autointent.advisor.detect_hardware`, or built by hand to + size a search space for a machine you are not currently on. All sizes are in + binary gigabytes (GiB). ``vram_gb`` is 0.0 on CPU-only hosts, and on Apple + silicon it is a fraction of unified memory rather than dedicated VRAM. + ``notes`` carries such caveats and any manual VRAM override that was applied. + """ + accelerator: Accelerator device_name: str vram_gb: float diff --git a/src/autointent/advisor/_report.py b/src/autointent/advisor/_report.py index 6fc3d8f2f..5007780e5 100644 --- a/src/autointent/advisor/_report.py +++ b/src/autointent/advisor/_report.py @@ -8,6 +8,18 @@ class Severity(str, Enum): + """How much headroom a finding leaves against the detected budget. + + * ``AMPLE`` — comfortably within budget; informational only. + * ``TIGHT`` — expected to fit, but with little margin; the estimate is + heuristic, so treat this as "may not fit". + * ``OVER`` — the budget is expected to be exceeded. + + A single ``OVER`` finding makes the whole report infeasible + (:attr:`PreflightReport.is_feasible` is ``False``) and is what + ``Pipeline.fit(preflight="strict")`` raises on. + """ + AMPLE = "ample" TIGHT = "tight" OVER = "over" @@ -70,6 +82,21 @@ def placeholder( avg_tokens: int = 32, multilabel: bool = False, ) -> DatasetStats: + """Build stats for a hypothetical dataset, for sizing a search space without data. + + ``p95_tokens`` is derived as ``avg_tokens * 2.5`` and ``class_counts`` is + left empty. Use :func:`~autointent.advisor.dataset_stats` instead when a + real ``Dataset`` is available. + + Args: + n_samples: number of training utterances to assume. + n_classes: number of intent classes to assume. + avg_tokens: average utterance length in whitespace-separated tokens. + multilabel: whether to assume a multilabel task. + + Returns: + Stats with ``source="placeholder"``. + """ return cls( n_samples=n_samples, n_classes=n_classes, @@ -92,6 +119,15 @@ class PreflightReport: notes: list[str] = field(default_factory=list) def add(self, phase: Phase, severity: Severity, message: str, metric: str | None = None) -> None: + """Append a :class:`Finding` to this report. + + Args: + phase: which check produced it — ``"resource"``, ``"data"`` or ``"config"``. + severity: headroom level; a single ``OVER`` makes the report infeasible. + message: one-line human-readable explanation, shown as-is in reports. + metric: short budget name the finding is about (``"vram"``, ``"ram"``, + ``"disk"``, ``"time"``), or ``None`` for findings not tied to a budget. + """ self.findings.append(Finding(phase=phase, severity=severity, message=message, metric=metric)) @property @@ -104,9 +140,20 @@ def headroom(self) -> Severity: @property def is_feasible(self) -> bool: + """Whether the run is expected to fit: True unless some finding is OVER.""" return self.headroom != Severity.OVER def to_dict(self) -> dict[str, Any]: + """Render the report as JSON-serializable data. + + Severities become their string values, and the derived ``headroom`` and + ``is_feasible`` properties are included as keys. This is what the CLI's + ``--json`` output emits. + + Returns: + A plain dict of the report, its findings, resource estimate, + hardware and dataset summaries. + """ d = asdict(self) d["findings"] = [{**asdict(f), "severity": f.severity.value} for f in self.findings] d["headroom"] = self.headroom.value diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index 27c7d705e..7a638899a 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -115,11 +115,24 @@ def stats_from_dataset(path: str, *, multilabel: bool = False) -> DatasetStats: def dataset_stats(dataset: Dataset) -> DatasetStats: - """Build :class:`DatasetStats` straight from an in-memory ``Dataset``. + """Summarize an in-memory :class:`~autointent.Dataset` for the advisor. - Counterpart of :func:`stats_from_dataset` that skips HF ``load_dataset`` - and reads the train split + autointent-specific attributes (``n_classes``, - ``multilabel``, ``has_descriptions``) directly. + Reads the train split (``train``, or ``train_0`` once the dataset has been + split) to count samples and measure utterance length — average and 95th + percentile word counts, over at most the first 1000 rows — and takes + ``n_classes``, ``multilabel`` and ``has_descriptions`` from the dataset + itself. Returns a placeholder when no train split is present. + + This is how a caller gets from a ``Dataset`` to the ``DatasetStats`` that + :func:`run_preflight` and :func:`reduce_to_fit` require. Use + :meth:`DatasetStats.placeholder` instead when no dataset exists yet and you + only want to size a search space against hypothetical numbers. + + Args: + dataset: the dataset the pipeline would be fitted on. + + Returns: + Stats describing that dataset, with ``source="dataset:in-memory"``. """ from autointent.custom_types import Split @@ -406,8 +419,10 @@ def reduce_to_fit( Behavior: * If ``config`` is already feasible, returns ``(config, report)`` unchanged. * Otherwise picks the OVER-driving scoring-node module with the largest - cost along whichever budget breached (VRAM > time > RAM > disk) and - removes it from the search_space, then re-runs preflight. + cost along whichever budget breached (VRAM > time > RAM) and removes it + from the search_space, then re-runs preflight. Disk is deliberately not + in that order: driver rows carry no per-module disk figure, so disk + pressure reduces by the VRAM proxy — download size tracks model size. * Repeats until feasible, ``max_iters`` reached, or no droppable module remains — in the last two cases raises :class:`ReduceToFitError` carrying the pruned config and final report. diff --git a/tests/advisor/test_estimates_and_cli.py b/tests/advisor/test_estimates_and_cli.py index 608c7420f..29f9670f2 100644 --- a/tests/advisor/test_estimates_and_cli.py +++ b/tests/advisor/test_estimates_and_cli.py @@ -13,14 +13,22 @@ import json import sys +from pathlib import Path import pytest from autointent.advisor import DatasetStats, HardwareProfile, run_preflight -from autointent.advisor._cli import main +from autointent.advisor._cli import build_parser, main from autointent.advisor._workflows import PRESET_COST_ORDER from autointent.utils import load_preset +if sys.version_info >= (3, 11): + import tomllib +else: # pytest depends on tomli below 3.11, so this import is always satisfiable here + import tomli as tomllib + +_PYPROJECT = Path(__file__).resolve().parents[2] / "pyproject.toml" + @pytest.fixture(autouse=True) def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: @@ -244,5 +252,23 @@ def test_cli_recommend_budget_time_flags_red_for_overbudget_presets( assert r["report"]["is_feasible"] is False +def test_console_script_name_matches_cli_prog() -> None: + """The installed command and the name the CLI prints must be the same string. + + They were not: pyproject registered ``advisor`` while the parser called + itself ``autointent-advisor``, so every usage/error message named a command + that did not exist. Both sides are now asserted against each other. + """ + with _PYPROJECT.open("rb") as f: + scripts = tomllib.load(f)["project"]["scripts"] + + advisor_scripts = {name: target for name, target in scripts.items() if target.startswith("autointent.advisor.")} + assert len(advisor_scripts) == 1, f"expected exactly one advisor console script, got {advisor_scripts}" + + script_name, target = next(iter(advisor_scripts.items())) + assert target == "autointent.advisor._cli:main" + assert script_name == build_parser().prog + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/advisor/test_hardware_detection.py b/tests/advisor/test_hardware_detection.py index 3926afd8d..24fa51292 100644 --- a/tests/advisor/test_hardware_detection.py +++ b/tests/advisor/test_hardware_detection.py @@ -1,5 +1,13 @@ -"""Hardware detection has to be safe on every machine — broken CUDA, no GPU, -no psutil. Verify the fallbacks work without raising. +"""Accelerator selection in ``detect_hardware``: CUDA -> MPS -> CPU. + +Each test patches ``_detect_cuda`` / ``_detect_mps`` (and sometimes +``_detect_ram_gb``) to force one branch, then checks the resulting profile — +the CPU fallback when nothing is available, the device_class thresholds, the +MPS unified-memory budget, and the manual VRAM override. + +These do *not* cover a missing ``psutil``: it is a core dependency, imported +unguarded at ``_hardware.py`` module level, and no psutil-absent fallback +exists. The RAM and disk probes are therefore always the real ones. """ from __future__ import annotations diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py index 747241d87..c9e689963 100644 --- a/tests/pipeline/test_preflight.py +++ b/tests/pipeline/test_preflight.py @@ -16,10 +16,19 @@ import pytest from autointent import Pipeline -from autointent.advisor import HardwareProfile, PreflightError, dataset_stats, detect_hardware, run_preflight +from autointent.advisor import ( + HardwareProfile, + PreflightError, + PreflightReport, + dataset_stats, + detect_hardware, + run_preflight, +) from autointent.configs import LoggingConfig if TYPE_CHECKING: + from typing import Any + from autointent import Dataset _PIPELINE_LOGGER = "autointent._pipeline._pipeline" @@ -49,6 +58,16 @@ def _classic_light_pipeline() -> Pipeline: return p +def _module_names(config: dict[str, Any]) -> set[str]: + """Every ``module_name`` in an advisor-shaped config's search space.""" + return { + entry["module_name"] + for node in config["search_space"] + for entry in node["search_space"] + if "module_name" in entry + } + + def test_fit_does_not_run_preflight_by_default(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: """The default is opt-out: no preflight, no Hub round-trips, no log line.""" p = _classic_light_pipeline() @@ -58,6 +77,11 @@ def test_fit_does_not_run_preflight_by_default(dataset: Dataset, caplog: pytest. def test_preflight_off_skips_advisor(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + """Asking for ``preflight="off"`` explicitly is honoured, not just the default. + + The previous test pins the default value; this one pins the ``"off"`` branch + of the gate itself, so a change to the default cannot mask a regression here. + """ p = _classic_light_pipeline() with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): p.fit(dataset, preflight="off") @@ -92,6 +116,37 @@ def test_preflight_warn_does_not_raise_on_infeasible( assert any(r.levelno == logging.ERROR for r in caplog.records) +def test_preflight_prices_the_filtered_search_space(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: + """The gate runs after validate_modules, so it never charges for discarded modules. + + ``classic-light`` ships ``mlknn``, which does not support multiclass, so + ``fit()`` drops it from the search space. Pricing the unfiltered space + inflates the estimate — badly so for ``dnnc``'s ~6.4 GB reranker on + multilabel data, where it can flip a strict verdict to OVER. + """ + import autointent.advisor as advisor_pkg + + captured: list[dict[str, Any]] = [] + + def _spy(config: dict[str, Any], *_args: object, **_kwargs: object) -> PreflightReport: + captured.append(config) + return PreflightReport() + + monkeypatch.setattr(advisor_pkg, "run_preflight", _spy) + + p = _classic_light_pipeline() + assert not dataset.multilabel, "fixture must be multiclass for mlknn to be filtered out" + requested = _module_names(p._build_advisor_config()) + assert "mlknn" in requested, f"classic-light should offer mlknn: {requested}" + + p.fit(dataset, preflight="warn") + + assert captured, "preflight did not run" + priced = _module_names(captured[0]) + assert "mlknn" not in priced, f"preflight priced a module fit() discards: {priced}" + assert "linear" in priced, f"preflight lost compatible modules too: {priced}" + + def test_importing_autointent_does_not_import_the_advisor() -> None: """The advisor pulls in huggingface_hub probes; it must stay off the import path. From 7a075302e8860a3fbd2b74fcf2f88054ae3997cb Mon Sep 17 00:00:00 2001 From: voorhs Date: Tue, 18 Aug 2026 16:34:06 +0300 Subject: [PATCH 41/43] docs: unwrap hard-wrapped lines in advisor.rst The prose was manually wrapped at ~79 columns, which makes every edit churn unrelated lines. Sphinx reflows paragraphs itself, so the breaks buy nothing. One line per paragraph, list item, and definition body. Code blocks, the note directive, and section underlines are untouched. Verified the docutils doctree is identical to the previous version once the shifted line numbers in system messages are normalised. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/advisor.rst | 60 +++++++++++------------------------------ 1 file changed, 16 insertions(+), 44 deletions(-) diff --git a/docs/source/advisor.rst b/docs/source/advisor.rst index b5517e585..b2d9c1fe8 100644 --- a/docs/source/advisor.rst +++ b/docs/source/advisor.rst @@ -3,20 +3,14 @@ Compute feasibility advisor .. note:: - **Experimental.** The advisor's estimates are heuristic and calibrated - against a limited hardware sample. Treat them as guidance, not guarantees, - and read :ref:`advisor-accuracy` before relying on a number. The Python - surface may change in a minor release. + **Experimental.** The advisor's estimates are heuristic and calibrated against a limited hardware sample. Treat them as guidance, not guarantees, and read :ref:`advisor-accuracy` before relying on a number. The Python surface may change in a minor release. -Optimizing a search space can take hours and needs more VRAM than a laptop GPU -has. The advisor answers "will this fit, and how long will it take?" *before* -anything is downloaded or trained. +Optimizing a search space can take hours and needs more VRAM than a laptop GPU has. The advisor answers "will this fit, and how long will it take?" *before* anything is downloaded or trained. Command line ------------ -Two subcommands. ``inspect`` prices a specific preset or config; ``recommend`` -detects your hardware and picks the heaviest bundled preset that still fits. +Two subcommands. ``inspect`` prices a specific preset or config; ``recommend`` detects your hardware and picks the heaviest bundled preset that still fits. .. code-block:: bash @@ -32,13 +26,9 @@ detects your hardware and picks the heaviest bundled preset that still fits. # Machine-readable output autointent-advisor inspect ./my-config.yaml --json -Without ``--dataset``, the advisor uses placeholder dataset sizes -(``--n-samples``, ``--n-classes``, ``--avg-tokens``, ``--task``), so it is -useful before you have assembled any data. ``--budget-vram-gb`` overrides -hardware detection, and ``recommend`` also accepts ``--budget-time-h``. +Without ``--dataset``, the advisor uses placeholder dataset sizes (``--n-samples``, ``--n-classes``, ``--avg-tokens``, ``--task``), so it is useful before you have assembled any data. ``--budget-vram-gb`` overrides hardware detection, and ``recommend`` also accepts ``--budget-time-h``. -Both subcommands exit non-zero when nothing is feasible, so they work as a CI -gate. +Both subcommands exit non-zero when nothing is feasible, so they work as a CI gate. Reading a report ---------------- @@ -52,10 +42,7 @@ Each finding carries a severity: ``over`` Exceeds the budget. Any ``over`` finding makes the whole report infeasible. -The drivers table lists the modules that dominate the cost, so it shows *what* -to change. ``low confidence`` on a report means Hub metadata was unavailable or -incomplete for at least one model, and conservative large-model defaults were -substituted — the numbers are much rougher when you see it. +The drivers table lists the modules that dominate the cost, so it shows *what* to change. ``low confidence`` on a report means Hub metadata was unavailable or incomplete for at least one model, and conservative large-model defaults were substituted — the numbers are much rougher when you see it. From Python ----------- @@ -71,44 +58,29 @@ From Python result = recommend(stats=dataset_stats(Dataset.from_json(path))) print(result.chosen) -``reduce_to_fit`` goes further: it prunes the most expensive scoring module -repeatedly until the search space fits, raising ``ReduceToFitError`` if nothing -does. +``reduce_to_fit`` goes further: it prunes the most expensive scoring module repeatedly until the search space fits, raising ``ReduceToFitError`` if nothing does. Inside ``Pipeline.fit`` ----------------------- -``Pipeline.fit`` accepts a ``preflight`` gate. It defaults to ``"off"``, so the -advisor never runs unless you ask — it makes network calls to the Hugging Face -Hub for model metadata, which does not belong on every fit by default. +``Pipeline.fit`` accepts a ``preflight`` gate. It defaults to ``"off"``, so the advisor never runs unless you ask — it makes network calls to the Hugging Face Hub for model metadata, which does not belong on every fit by default. .. code-block:: python pipeline.fit(dataset, preflight="warn") # log findings, always continue pipeline.fit(dataset, preflight="strict") # raise PreflightError if infeasible -``"strict"`` raises :class:`autointent.advisor.PreflightError` before allocating -any VRAM, which is the useful mode in CI. +``"strict"`` raises :class:`autointent.advisor.PreflightError` before allocating any VRAM, which is the useful mode in CI. .. _advisor-accuracy: How accurate is it? ------------------- -Validated end to end on one machine class (RTX 3060 Laptop, 6 GB VRAM / 16 GB -RAM), where all four fitted presets matched their predicted verdict: both -``over`` predictions did run out of memory, and both feasible predictions did -fit. Known limits: - -- **Feasibility verdicts are the reliable part.** That is what the advisor was - built and validated for. -- **VRAM is close but not a guaranteed ceiling.** One preset used 1.22× its - prediction. Leave headroom rather than trusting the figure exactly. -- **Wall-time estimates are indicative only.** Measured error has run in both - directions across formula revisions, once by more than an order of magnitude - for cross-encoders. ``--budget-time-h`` inherits that uncertainty. -- **Preset ranking does not depend on time estimates.** ``recommend`` orders - presets by a declared cost ranking, so unstable time figures cannot reorder - its choice. -- **Only one hardware class has been validated end to end.** Treat other - machines as unverified. +Validated end to end on one machine class (RTX 3060 Laptop, 6 GB VRAM / 16 GB RAM), where all four fitted presets matched their predicted verdict: both ``over`` predictions did run out of memory, and both feasible predictions did fit. Known limits: + +- **Feasibility verdicts are the reliable part.** That is what the advisor was built and validated for. +- **VRAM is close but not a guaranteed ceiling.** One preset used 1.22× its prediction. Leave headroom rather than trusting the figure exactly. +- **Wall-time estimates are indicative only.** Measured error has run in both directions across formula revisions, once by more than an order of magnitude for cross-encoders. ``--budget-time-h`` inherits that uncertainty. +- **Preset ranking does not depend on time estimates.** ``recommend`` orders presets by a declared cost ranking, so unstable time figures cannot reorder its choice. +- **Only one hardware class has been validated end to end.** Treat other machines as unverified. From f52d19c288932b99cad887037670f25a0d99b7ab Mon Sep 17 00:00:00 2001 From: voorhs Date: Tue, 18 Aug 2026 16:51:59 +0300 Subject: [PATCH 42/43] style: apply ruff format to three files that had drifted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ruff format --check` flagged these on the branch before any of the review fixes touched them — magic-trailing-comma expansions an earlier commit left behind. No behaviour change; separated so the review-fix commit is readable. Co-Authored-By: Claude Opus 5 (1M context) --- src/autointent/advisor/_workflows.py | 8 +- tests/advisor/test_estimates_internals.py | 129 ++++++++++++++-------- tests/advisor/test_reduce_to_fit.py | 12 +- 3 files changed, 92 insertions(+), 57 deletions(-) diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index 7a638899a..5ed6c53ec 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -337,7 +337,9 @@ def not_converged( def _drop_module_from_search_space( - search_space: list[dict[str, Any]], node_type: str, module_name: str, + search_space: list[dict[str, Any]], + node_type: str, + module_name: str, ) -> list[dict[str, Any]]: """Return a deep-copied search_space with ``module_name`` removed. @@ -454,7 +456,9 @@ def reduce_to_fit( raise ReduceToFitError.nothing_droppable(pruned_config=current, last_report=report) node_type, module_name = pick current["search_space"] = _drop_module_from_search_space( - current["search_space"], node_type, module_name, + current["search_space"], + node_type, + module_name, ) logger.info("reduce_to_fit: dropped %s/%s to fit budget", node_type, module_name) # An empty scoring node — after dropping the last scoring module — diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py index d552759ac..266905781 100644 --- a/tests/advisor/test_estimates_internals.py +++ b/tests/advisor/test_estimates_internals.py @@ -744,9 +744,7 @@ def test_warm_cache_probe_zeroes_forward_and_disk(self) -> None: "search_space": [ { "module_name": "knn", - "embedder_config": [ - {"model_name": "sentence-transformers/all-MiniLM-L6-v2"} - ], + "embedder_config": [{"model_name": "sentence-transformers/all-MiniLM-L6-v2"}], "batch_size": [32], "max_length": [128], } @@ -777,10 +775,19 @@ class TestCnnRnnHeuristic: def test_cnn_row_is_nonzero(self) -> None: cfg = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "cnn", "embed_dim": [128], "num_filters": [128], - "kernel_sizes": [[3, 4, 5]], "batch_size": [64], "num_train_epochs": [60]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "cnn", + "embed_dim": [128], + "num_filters": [128], + "kernel_sizes": [[3, 4, 5]], + "batch_size": [64], + "num_train_epochs": [60], + }, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 10}, @@ -820,9 +827,12 @@ def test_single_module_gets_full_n_trials(self) -> None: cfg = { **embedder_cfg, "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "linear"}, - ]}, + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear"}, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 200}, @@ -830,10 +840,13 @@ def test_single_module_gets_full_n_trials(self) -> None: cfg2 = { **embedder_cfg, "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "linear"}, - {"module_name": "knn", "k": [5]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear"}, + {"module_name": "knn", "k": [5]}, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 200}, @@ -913,12 +926,7 @@ def test_reserved_keys_skipped(self) -> None: from autointent.advisor._estimates._search_space import _module_cardinality # module_name / target_metric are not search dimensions - assert ( - _module_cardinality( - {"module_name": "bert", "target_metric": "scoring_f1", "batch_size": [32, 64]} - ) - == 2 - ) + assert _module_cardinality({"module_name": "bert", "target_metric": "scoring_f1", "batch_size": [32, 64]}) == 2 class TestNoOpHpoFinding: @@ -927,11 +935,17 @@ class TestNoOpHpoFinding: def test_finding_on_singleton_bert_with_high_n_trials(self) -> None: cfg = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "bert", - "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], - "num_train_epochs": [30], "batch_size": [64]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "num_train_epochs": [30], + "batch_size": [64], + }, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 40}, @@ -946,11 +960,16 @@ def test_finding_on_singleton_bert_with_high_n_trials(self) -> None: def test_no_finding_when_search_space_has_range(self) -> None: cfg = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "bert", - "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], - "learning_rate": {"low": 1e-5, "high": 1e-4}}, - ]}, + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "learning_rate": {"low": 1e-5, "high": 1e-4}, + }, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 40}, @@ -963,11 +982,17 @@ def test_no_finding_when_n_trials_matches_cardinality(self) -> None: # n_trials=4, cardinality=2x2=4 → not a "no-op" waste cfg = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "bert", - "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], - "batch_size": [32, 64], "num_train_epochs": [10, 20]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "batch_size": [32, 64], + "num_train_epochs": [10, 20], + }, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 4}, @@ -985,9 +1010,12 @@ def test_inference_only_preset_gets_smaller_vram_than_training(self) -> None: # Both use e5-large; only the training config triggers the bigger baseline. inference_only = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "knn", "k": [5]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + {"module_name": "knn", "k": [5]}, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "embedder_config": {"model_name": "intfloat/multilingual-e5-large-instruct"}, @@ -995,11 +1023,17 @@ def test_inference_only_preset_gets_smaller_vram_than_training(self) -> None: } training = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "bert", - "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], - "batch_size": [16], "num_train_epochs": [1]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "batch_size": [16], + "num_train_epochs": [1], + }, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "hpo_config": {"n_trials": 5}, @@ -1027,9 +1061,12 @@ def test_inference_only_still_has_a_cuda_baseline(self) -> None: # embedder-only presets need no GPU memory. cfg = { "search_space": [ - {"node_type": "scoring", "search_space": [ - {"module_name": "knn", "k": [5]}, - ]}, + { + "node_type": "scoring", + "search_space": [ + {"module_name": "knn", "k": [5]}, + ], + }, {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, ], "embedder_config": {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}, diff --git a/tests/advisor/test_reduce_to_fit.py b/tests/advisor/test_reduce_to_fit.py index 1a4d5e161..618f52b38 100644 --- a/tests/advisor/test_reduce_to_fit.py +++ b/tests/advisor/test_reduce_to_fit.py @@ -121,9 +121,7 @@ def test_feasible_config_returns_unchanged() -> None: pruned, report = reduce_to_fit(config, stats, _profile(vram_gb=16.0)) assert report.is_feasible # Passthrough: same module still present. - modules = [ - e["module_name"] for node in pruned["search_space"] for e in node["search_space"] - ] + modules = [e["module_name"] for node in pruned["search_space"] for e in node["search_space"]] assert "linear" in modules assert "argmax" in modules @@ -139,9 +137,7 @@ def test_prunes_infeasible_transformer_to_classic() -> None: pruned, report = reduce_to_fit(config, stats, _profile(vram_gb=1.0)) assert report.is_feasible - modules = [ - e["module_name"] for node in pruned["search_space"] for e in node["search_space"] - ] + modules = [e["module_name"] for node in pruned["search_space"] for e in node["search_space"]] assert "bert" not in modules, "expensive transformer should have been dropped" assert "linear" in modules, "cheap classic scorer should be preserved" @@ -161,6 +157,4 @@ def test_raises_when_nothing_fits() -> None: # After pruning the only scoring module, the config's scoring node should # be gone entirely (or empty), leaving an unfittable pipeline. scoring_nodes = [n for n in err.pruned_config["search_space"] if n.get("node_type") == "scoring"] - assert scoring_nodes == [] or all( - not n.get("search_space") for n in scoring_nodes - ) + assert scoring_nodes == [] or all(not n.get("search_space") for n in scoring_nodes) From 48109b8623ef30516ce118936fd0fca7d8e4bccb Mon Sep 17 00:00:00 2001 From: voorhs Date: Tue, 18 Aug 2026 16:52:29 +0300 Subject: [PATCH 43/43] fix: make the advisor read the split config, cv, and CPU count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three review findings on #291, all the same shape: the advisor was already being handed the input and threw it away. Splits (voorhs, runner.py). DatasetStats.class_counts is measured on the train split as supplied, but DataHandler carves that up before any module sees it — validation_size off the top, one fold out under cv, and separation_ratio splitting the remainder into scoring and decision. The LogisticRegressionCV check counted against the raw split, so it could pass where the real fit fails. Pipeline._build_advisor_config already passed "data_config" into run_preflight and nothing in the advisor read it; _data_phase now takes it and discounts the counts. Same commit adds the check that was missing entirely: classes below the stratified splitter's own minimum. It imports _min_samples_per_class_for_config rather than restating the rule, so the advisor cannot drift away from check_split_readiness, and a parametrised test pins the two together. cv (voorhs, runner.py). The feasibility gate already read the declared cv, but the time estimate used _LOGREG_CV_MULTIPLIER = 31, hardcoding Cs=10 x cv=3 + 1. Anyone tuning cv got a cost for cv=3. Now derived: _logreg_cv_multiplier(cv) = Cs * cv + 1, still 31 at the default. CPU (voorhs, _estimates.py). HardwareProfile.cpu_count was detected and read by nothing, so a 4-core and a 64-core box priced identically. Time now divides by a capped Amdahl speedup — 0.90 parallel for CatBoost, which defaults to every core, 0.50 for L-BFGS, which only threads inside BLAS, ignored on GPU — over cores/n_jobs, since concurrent HPO trials share the machine. Capped at 8x deliberately: these estimates bound cost from above and an over-generous speedup would flip that. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/advisor.rst | 3 + .../advisor/_estimates/_formulas.py | 72 +++-- .../advisor/_estimates/_resource.py | 95 ++++--- src/autointent/advisor/_runner.py | 77 +++++- tests/advisor/test_split_and_cpu_awareness.py | 249 ++++++++++++++++++ 5 files changed, 431 insertions(+), 65 deletions(-) create mode 100644 tests/advisor/test_split_and_cpu_awareness.py diff --git a/docs/source/advisor.rst b/docs/source/advisor.rst index b2d9c1fe8..6db50249a 100644 --- a/docs/source/advisor.rst +++ b/docs/source/advisor.rst @@ -44,6 +44,8 @@ Each finding carries a severity: The drivers table lists the modules that dominate the cost, so it shows *what* to change. ``low confidence`` on a report means Hub metadata was unavailable or incomplete for at least one model, and conservative large-model defaults were substituted — the numbers are much rougher when you see it. +Findings are not only about hardware. The advisor also prices your ``DataConfig``: it reports ``over`` when a class has too few samples for the stratified split to succeed, using the same minimum as :py:func:`~autointent.context.data_handler.check_split_readiness`, and when ``LogisticRegressionCV`` would not have ``cv`` samples per class *after* the train/validation split. Per-class counts are measured on the train split you supply, so the advisor discounts them by whatever ``validation_size``, ``n_folds``, and ``separation_ratio`` will take away. + From Python ----------- @@ -82,5 +84,6 @@ Validated end to end on one machine class (RTX 3060 Laptop, 6 GB VRAM / 16 GB RA - **Feasibility verdicts are the reliable part.** That is what the advisor was built and validated for. - **VRAM is close but not a guaranteed ceiling.** One preset used 1.22× its prediction. Leave headroom rather than trusting the figure exactly. - **Wall-time estimates are indicative only.** Measured error has run in both directions across formula revisions, once by more than an order of magnitude for cross-encoders. ``--budget-time-h`` inherits that uncertainty. +- **CPU parallelism is modelled, not measured.** The CPU coefficients are calibrated single-threaded; core count is then applied as a capped Amdahl speedup (higher for CatBoost, which uses every core by default, than for scikit-learn's L-BFGS, which only threads inside BLAS), divided across concurrent ``hpo_config.n_jobs`` trials. It is a correction for the fact that core count used to change nothing at all, not a validated speedup curve. - **Preset ranking does not depend on time estimates.** ``recommend`` orders presets by a declared cost ranking, so unstable time figures cannot reorder its choice. - **Only one hardware class has been validated end to end.** Treat other machines as unverified. diff --git a/src/autointent/advisor/_estimates/_formulas.py b/src/autointent/advisor/_estimates/_formulas.py index 25daf6fb8..368a62e0f 100644 --- a/src/autointent/advisor/_estimates/_formulas.py +++ b/src/autointent/advisor/_estimates/_formulas.py @@ -137,11 +137,11 @@ def _max_fitting_batch_size( # not peak spec sheet. Advisor aims to over- rather than under-predict time, # so pessimistic (low) values here. _DEVICE_TFLOPS = { - "high-gpu": 60.0, # A100 / H100 - "mid-gpu": 20.0, # V100 / RTX 3090 / A6000 - "low-gpu": 7.0, # T4 / RTX 3060 / 8 GB consumer card + "high-gpu": 60.0, # A100 / H100 + "mid-gpu": 20.0, # V100 / RTX 3090 / A6000 + "low-gpu": 7.0, # T4 / RTX 3060 / 8 GB consumer card "apple-silicon": 4.0, # M1/M2/M3 GPU cores - "cpu": 0.05, # single-thread modern x86 with MKL + "cpu": 0.05, # single-thread modern x86 with MKL } _DEFAULT_TFLOPS = 7.0 # unknown device → treat as low-GPU @@ -221,10 +221,39 @@ def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: _LINEAR_CPU_S_PER_SAMPLE_FEATURE = 1.2e-9 _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 _CATBOOST_GPU_SPEEDUP = 10.0 -_LOGREG_CV_MULTIPLIER = 31 # sklearn default: Cs=10 x cv=3 + 1 final refit +_LOGREG_CS = 10 # LogisticRegressionCV(Cs=10) sklearn default; LinearScorer does not expose it +_LOGREG_DEFAULT_CV = 3 # LinearScorer(cv=3) default _CATBOOST_DEFAULT_BINS = 254 # CatBoost `border_count` default _CATBOOST_BYTES_PER_TREE_NODE = 32 +# CPU parallelism. The coefficients above are calibrated at one thread, which +# left every CPU-bound estimate independent of core count — a 4-core and a +# 64-core box priced identically. Speedup is modelled with Amdahl's law and +# capped: these estimates exist to bound cost from above, and an over-generous +# speedup turns a conservative estimate into an optimistic one. +_CATBOOST_PARALLEL_FRACTION = 0.90 # CatBoost `thread_count` defaults to every core +_LINEAR_PARALLEL_FRACTION = 0.50 # only the BLAS calls inside L-BFGS thread +_MAX_CPU_SPEEDUP = 8.0 # refuse to believe in more than 8x however many cores are reported + + +def _logreg_cv_multiplier(cv: int) -> int: + """Fits per ``LogisticRegressionCV`` run: a ``Cs x cv`` grid plus one final refit.""" + return _LOGREG_CS * max(1, cv) + 1 + + +def _cpu_speedup(cores: int, parallel_fraction: float) -> float: + """Amdahl speedup on ``cores``, capped at :data:`_MAX_CPU_SPEEDUP`.""" + n = max(1, cores) + if n == 1: + return 1.0 + speedup = 1.0 / ((1.0 - parallel_fraction) + parallel_fraction / n) + return min(speedup, _MAX_CPU_SPEEDUP) + + +def _cores_per_trial(cpu_count: int, n_jobs: int) -> int: + """Cores one HPO trial gets when ``n_jobs`` trials run concurrently.""" + return max(1, max(1, cpu_count) // max(1, n_jobs)) + def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: """Float64 design matrix dominates; coefficients and L-BFGS history are small.""" @@ -242,21 +271,17 @@ def _time_for_linear( max_iter: int, # noqa: ARG001 — API stability; typical L-BFGS convergence baked into coeff cv_multiplier: int, class_multiplier: int, + cores: int = 1, ) -> float: """LogisticRegression wall time. O(n_samples x features x classes x cv) per fit; typical L-BFGS - convergence absorbed into the calibration constant. + convergence absorbed into the calibration constant. ``cores`` divides that + by the modest BLAS-only speedup L-BFGS gets — sklearn's own CV loop runs + single-threaded here, since ``LinearScorer`` leaves ``n_jobs`` unset. """ - seconds = ( - n_trials - * _LINEAR_CPU_S_PER_SAMPLE_FEATURE - * n_samples - * embedder_dim - * cv_multiplier - * class_multiplier - ) - return seconds / 3600.0 + seconds = n_trials * _LINEAR_CPU_S_PER_SAMPLE_FEATURE * n_samples * embedder_dim * cv_multiplier * class_multiplier + return seconds / _cpu_speedup(cores, _LINEAR_PARALLEL_FRACTION) / 3600.0 def _ram_for_catboost(*, stats: DatasetStats, n_features: int, iterations: int, depth: int) -> float: @@ -306,6 +331,7 @@ def _time_for_catboost( depth: int, class_multiplier: int, on_gpu: bool, + cores: int = 1, ) -> float: """CatBoost wall time, in hours. @@ -313,11 +339,17 @@ def _time_for_catboost( fit. GPU training is ~10x faster than CPU for typical workloads per CatBoost's published benchmarks. https://catboost.ai/en/docs/concepts/speed-up-training + + On CPU, ``cores`` divides that: CatBoost's ``thread_count`` defaults to + every core, so core count is the single largest term the one-thread + calibration was missing. Ignored on GPU, where the device is the bottleneck. """ coeff = _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER if on_gpu: coeff /= _CATBOOST_GPU_SPEEDUP seconds = n_trials * iterations * coeff * n_samples * n_features * depth * class_multiplier + if not on_gpu: + seconds /= _cpu_speedup(cores, _CATBOOST_PARALLEL_FRACTION) return seconds / 3600.0 @@ -340,11 +372,7 @@ def _cnn_param_count(*, embed_dim: int, num_filters: int, n_kernels: int, n_clas def _rnn_param_count(*, embed_dim: int, hidden_dim: int, n_classes: int) -> int: """LSTM classifier params: embedding + 4-gate LSTM cell + fc.""" - return ( - _NN_MAX_VOCAB * embed_dim - + 4 * hidden_dim * (embed_dim + hidden_dim + 1) - + hidden_dim * max(1, n_classes) - ) + return _NN_MAX_VOCAB * embed_dim + 4 * hidden_dim * (embed_dim + hidden_dim + 1) + hidden_dim * max(1, n_classes) def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: @@ -354,9 +382,7 @@ def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: smaller hidden dim (embed_dim / num_filters). """ weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB - activations_gb = ( - batch_size * _NN_DEFAULT_SEQ_LEN * hidden_dim * _NN_TRAIN_ACT_BYTES_PER_UNIT - ) / _BYTES_PER_GB + activations_gb = (batch_size * _NN_DEFAULT_SEQ_LEN * hidden_dim * _NN_TRAIN_ACT_BYTES_PER_UNIT) / _BYTES_PER_GB return 4 * weights_gb + activations_gb diff --git a/src/autointent/advisor/_estimates/_resource.py b/src/autointent/advisor/_estimates/_resource.py index ec6a4949c..5de6a4407 100644 --- a/src/autointent/advisor/_estimates/_resource.py +++ b/src/autointent/advisor/_estimates/_resource.py @@ -25,15 +25,17 @@ from ._formulas import ( _DEFAULT_SEQ_LEN, _LINEAR_CPU_S_PER_SAMPLE_FEATURE, - _LOGREG_CV_MULTIPLIER, + _LOGREG_DEFAULT_CV, _MULTICLASS_THRESHOLD, _activations_gb_per_sample, _classify_severity, _cnn_param_count, + _cores_per_trial, _embedder_dim, _embedder_load_ram_gb, _embedding_cache_disk_gb, _largest_embedder, + _logreg_cv_multiplier, _max_fitting_batch_size, _ram_for_catboost, _ram_for_linear, @@ -213,15 +215,21 @@ def _estimate_classic_entry( hardware: HardwareProfile, n_trials: int, refit_after: bool, + hpo_n_jobs: int = 1, ) -> _ModuleEstimate | None: """Cost row for a linear or catboost scorer (returns ``None`` for any other module).""" module = entry.get("module_name", "?") refit = _refit_factor(refit_after=refit_after, n_trials=n_trials) # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes. class_multiplier = max(1, stats.n_classes) + # Concurrent HPO trials share the box, so a trial does not get every core. + cores = _cores_per_trial(hardware.cpu_count, hpo_n_jobs) if module == "linear": - cv_multiplier = 1 if stats.multilabel else _LOGREG_CV_MULTIPLIER + # cv is per-entry and the grid scales with it; assuming the default 3 + # under-priced every search space that tuned it. + cv = _max_int(entry.get("cv"), _LOGREG_DEFAULT_CV) + cv_multiplier = 1 if stats.multilabel else _logreg_cv_multiplier(cv) ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) time_h = ( _time_for_linear( @@ -231,11 +239,12 @@ def _estimate_classic_entry( max_iter=_max_int(entry.get("max_iter"), 100), cv_multiplier=cv_multiplier, class_multiplier=class_multiplier, + cores=cores, ) * refit ) vram = 0.0 - mode = "linear-cv" if cv_multiplier > 1 else "linear" + mode = f"linear-cv{cv}" if cv_multiplier > 1 else "linear" elif module == "catboost": on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" # CatBoost MultiClass loss grows per-class trees only above binary; binary uses @@ -253,6 +262,7 @@ def _estimate_classic_entry( depth=depth, class_multiplier=cb_class_mult, on_gpu=on_gpu, + cores=cores, ) * refit ) @@ -349,17 +359,14 @@ def _estimate_nn_entry( vram = _vram_for_nn(params=params, batch_size=batch_size, hidden_dim=hidden_dim) ram = _ram_for_nn(params=params, stats=stats) - time_h = ( - _time_for_nn( - n_trials=n_trials, - epochs=epochs, - batch_size=batch_size, - n_samples=stats.n_samples, - params_millions=params / 1_000_000, - device_class=hardware.device_class, - ) - * _refit_factor(refit_after=refit_after, n_trials=n_trials) - ) + time_h = _time_for_nn( + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + params_millions=params / 1_000_000, + device_class=hardware.device_class, + ) * _refit_factor(refit_after=refit_after, n_trials=n_trials) return _ModuleEstimate( driver={ @@ -484,9 +491,7 @@ def _apply_embedding_cache( _mark_cache_hit(me, module, suffix="cached") else: paid.add(model) - _charge_first_forward_if_classic( - me, module, model, seen_models=seen_models, stats=stats, hardware=hardware - ) + _charge_first_forward_if_classic(me, module, model, seen_models=seen_models, stats=stats, hardware=hardware) return paid @@ -593,17 +598,24 @@ def _emit_resource_findings( _NN_SCORER_MODULES = frozenset({"cnn", "rnn"}) _EMBEDDER_CONSUMING_MODULES = frozenset( - {"linear", "catboost", "sklearn", "knn", "mlknn", "retrieval", - "description_bi", "description_cross", "description_llm"}, + { + "linear", + "catboost", + "sklearn", + "knn", + "mlknn", + "retrieval", + "description_bi", + "description_cross", + "description_llm", + }, ) def _uses_embedder(search_space: list[dict[str, Any]]) -> bool: """True when any search-space module consumes the embedder.""" - return any( - entry.get("module_name") in _EMBEDDER_CONSUMING_MODULES - for _, entry in _walk_modules(search_space) - ) + return any(entry.get("module_name") in _EMBEDDER_CONSUMING_MODULES for _, entry in _walk_modules(search_space)) + # Modules that consume the top-level ``cross_encoder_config.model_name`` as # their scoring model (see zero-shot-encoders preset: description_cross pulls @@ -649,9 +661,10 @@ class _ResourceInputs: Bundled by provenance rather than by use: every field is derived from one validated ``OptimizationConfig``, plus the caller-injected ``cache_probe``. Individual passes read only what they need — the classic pass reads - ``refit_after`` alone, and ``search_space`` / ``n_trials`` / ``n_jobs`` / - ``dump_modules`` never leave ``_resource_phase`` — but threading them - separately would mean a dozen keyword arguments down each call. + ``refit_after`` and ``n_jobs`` (concurrent trials divide the CPU cores each + one gets), and ``search_space`` / ``n_trials`` / ``dump_modules`` never + leave ``_resource_phase`` — but threading them separately would mean a + dozen keyword arguments down each call. ``cross_encoder_model_name`` and ``transformer_model_name`` come from the pipeline's top-level configs and act as the fallback model for modules that @@ -774,6 +787,7 @@ def _estimate_classic_entries( hardware=hardware, n_trials=effective_trials(node_idx, entry), refit_after=refit_after, + hpo_n_jobs=inputs.n_jobs, ) if classic_estimate is not None: module_estimates.append(classic_estimate) @@ -824,18 +838,32 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int # First pass: transformer modules (also populates seen_models for the classic pass). module_estimates, node_max_weights = _estimate_transformer_entries( - transformer_entries, inputs, stats, hardware, seen_models, _effective_trials, + transformer_entries, + inputs, + stats, + hardware, + seen_models, + _effective_trials, ) # Second pass: linear / catboost — cost depends on embedder_dim, not a checkpoint. embedder_meta = _largest_embedder(seen_models) module_estimates += _estimate_classic_entries( - classic_entries, inputs, stats, hardware, seen_models, _effective_trials, + classic_entries, + inputs, + stats, + hardware, + seen_models, + _effective_trials, ) # Cache-aware time/disk: must run before the fold below. cached_embedders = _apply_embedding_cache( - module_estimates, seen_models, stats=stats, hardware=hardware, cache_probe=cache_probe, + module_estimates, + seen_models, + stats=stats, + hardware=hardware, + cache_probe=cache_probe, ) estimate = ResourceEstimate(parallel_factor=n_jobs) @@ -854,13 +882,8 @@ def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int # CUDA baseline only when we predict some VRAM AND run on CUDA. Mode read # off drivers: any training row flips to the larger baseline. if estimate.vram_gb > 0 and hardware.accelerator == "cuda": - is_training = any( - d.get("mode") in {"full-finetune", "lora", "small-torch-train"} - for d in estimate.drivers - ) - estimate.vram_gb += ( - _CUDA_BASELINE_VRAM_TRAINING_GB if is_training else _CUDA_BASELINE_VRAM_INFERENCE_GB - ) + is_training = any(d.get("mode") in {"full-finetune", "lora", "small-torch-train"} for d in estimate.drivers) + estimate.vram_gb += _CUDA_BASELINE_VRAM_TRAINING_GB if is_training else _CUDA_BASELINE_VRAM_INFERENCE_GB _aggregate_disk( estimate, diff --git a/src/autointent/advisor/_runner.py b/src/autointent/advisor/_runner.py index 416cca80a..55c74cbd8 100644 --- a/src/autointent/advisor/_runner.py +++ b/src/autointent/advisor/_runner.py @@ -16,11 +16,18 @@ from autointent.advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules from autointent.advisor._report import PreflightReport, Severity +# Imported rather than reimplemented: the advisor must not disagree with the +# splitter about what counts as too few samples per class. `check_split_readiness` +# itself needs a Dataset, which the advisor never has (it works from DatasetStats), +# so the shared piece is the minimum. `test_split_readiness_agreement` pins them together. +from autointent.context.data_handler._readiness_util import _min_samples_per_class_for_config + if TYPE_CHECKING: from collections.abc import Callable from autointent.advisor._hardware import HardwareProfile from autointent.advisor._report import DatasetStats + from autointent.configs import DataConfig logger = logging.getLogger(__name__) @@ -94,7 +101,7 @@ def run_preflight( hardware, report, ) - _data_phase(cfg.search_space, stats, report) + _data_phase(cfg.search_space, stats, cfg.data_config, report) _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, cfg.hpo_config.n_trials, hardware, report) return report @@ -169,9 +176,60 @@ def _config_phase( break +def _effective_train_fraction(data_config: DataConfig) -> float: + """Fraction of the train split a scoring module is actually fitted on. + + ``DatasetStats.class_counts`` is measured on the train split as the user + supplies it, but the pipeline carves that up before any module sees it: + hold-out takes ``validation_size`` away for validation, cross-validation + leaves one fold out, and ``separation_ratio`` splits the remaining pool + again into scoring and decision halves. Counting against the raw split is + therefore optimistic, which is the wrong direction for a feasibility gate. + + An approximation of :class:`~autointent.context.data_handler.DataHandler`'s + splitting, not a reimplementation of it — deliberately coarse, and only + used to decide whether a class is at risk. + """ + if data_config.scheme == "cv": + n_folds = max(2, data_config.n_folds) + fraction = (n_folds - 1) / n_folds + else: + fraction = 1.0 - float(data_config.validation_size) + if data_config.separation_ratio is not None: + fraction *= 1.0 - float(data_config.separation_ratio) + return max(0.0, min(1.0, fraction)) + + +def _split_readiness_finding( + stats: DatasetStats, + data_config: DataConfig, + report: PreflightReport, +) -> None: + """Flag classes the stratified splitter will reject, by the splitter's own rule.""" + if not stats.class_counts: + return + min_required = _min_samples_per_class_for_config(config=data_config) + starved = sorted(name for name, count in stats.class_counts.items() if count < min_required) + if not starved: + return + detail = f"scheme={data_config.scheme}" + if data_config.scheme != "ho": + detail += f", n_folds={data_config.n_folds}" + if data_config.separation_ratio is not None: + detail += f", separation_ratio={data_config.separation_ratio}" + report.add( + "data", + Severity.OVER, + f"Stratified splitting will fail before any module is fitted: classes {starved[:5]} " + f"have <{min_required} samples ({detail}). Same minimum as " + f"autointent.context.data_handler.check_split_readiness.", + ) + + def _data_phase( search_space: list[dict[str, Any]], stats: DatasetStats, + data_config: DataConfig, report: PreflightReport, ) -> None: """Data-phase checks: token truncation, rare classes, missing intent descriptions.""" @@ -191,21 +249,28 @@ def _data_phase( f"Train tokens p95~{p95} exceeds {module_name}.max_length={max_len}; expect silent truncation.", ) - # sklearn LogisticRegressionCV inner-CV failure: each class needs >= cv samples. - # cv is configurable per linear entry (default 3); use the strictest one across - # the search space. Multilabel uses LogisticRegression (no CV), so skip there. + _split_readiness_finding(stats, data_config, report) + + # sklearn LogisticRegressionCV inner-CV failure: each class needs >= cv samples + # in the split the scorer is fitted on, which is smaller than the train split + # the counts were measured on. cv is configurable per linear entry (default 3); + # use the strictest one across the search space. Multilabel uses + # LogisticRegression (no CV), so skip there. if not stats.multilabel and stats.class_counts: linear_cvs = [ _max_int(e.get("cv"), 3) for _, e in _walk_modules(search_space) if e.get("module_name") == "linear" ] if linear_cvs: cv_max = max(linear_cvs) - failing = sorted(name for name, count in stats.class_counts.items() if count < cv_max) + fraction = _effective_train_fraction(data_config) + failing = sorted(name for name, count in stats.class_counts.items() if int(count * fraction) < cv_max) if failing: + note = "" if fraction >= 1.0 else f" after the {fraction:.0%} train/validation split" report.add( "data", Severity.OVER, - f"LogisticRegressionCV (cv={cv_max}) will fail: classes {failing[:5]} have <{cv_max} samples.", + f"LogisticRegressionCV (cv={cv_max}) will fail: classes {failing[:5]} " + f"have <{cv_max} samples{note}.", ) # partial descriptions x description scorer diff --git a/tests/advisor/test_split_and_cpu_awareness.py b/tests/advisor/test_split_and_cpu_awareness.py new file mode 100644 index 000000000..eec9266a7 --- /dev/null +++ b/tests/advisor/test_split_and_cpu_awareness.py @@ -0,0 +1,249 @@ +"""Tests for the review findings on PR #291: split awareness, cv, and CPU cores. + +Three separate complaints, all of which reduced to the advisor ignoring +something it was already being handed: + +* the ``cv`` a linear entry actually declares (the time estimate hardcoded 3), +* the fact that the pipeline splits the train split again before any module + sees it, so raw per-class counts are optimistic, +* ``HardwareProfile.cpu_count``, which was detected and then read by nothing. +""" + +from __future__ import annotations + +from typing import Any, ClassVar + +import pytest + +from autointent.advisor._estimates._formulas import ( + _LINEAR_PARALLEL_FRACTION, + _MAX_CPU_SPEEDUP, + _cores_per_trial, + _cpu_speedup, + _logreg_cv_multiplier, + _time_for_catboost, + _time_for_linear, +) +from autointent.advisor._hardware import HardwareProfile +from autointent.advisor._report import DatasetStats, PreflightReport, Severity +from autointent.advisor._runner import _data_phase, _effective_train_fraction +from autointent.configs import DataConfig +from autointent.context.data_handler._readiness_util import _min_samples_per_class_for_config + + +def _stats(class_counts: dict[str, int], *, multilabel: bool = False) -> DatasetStats: + return DatasetStats( + n_samples=sum(class_counts.values()), + n_classes=len(class_counts), + avg_tokens=16, + p95_tokens=32, + multilabel=multilabel, + class_counts=class_counts, + source="test", + ) + + +def _linear_space(cv: int | None = None) -> list[dict[str, Any]]: + entry: dict[str, Any] = {"module_name": "linear"} + if cv is not None: + entry["cv"] = cv + return [{"node_type": "scoring", "search_space": [entry]}] + + +def _run_data_phase(stats: DatasetStats, data_config: DataConfig, cv: int | None = None) -> PreflightReport: + report = PreflightReport() + _data_phase(_linear_space(cv), stats, data_config, report) + return report + + +def _messages(report: PreflightReport) -> str: + return " | ".join(f.message for f in report.findings) + + +class TestLogregCvMultiplier: + """The time estimate used to hardcode 31 = Cs(10) x cv(3) + 1 refit.""" + + def test_default_cv_reproduces_the_old_constant(self) -> None: + assert _logreg_cv_multiplier(3) == 31 + + @pytest.mark.parametrize(("cv", "expected"), [(2, 21), (5, 51), (10, 101)]) + def test_scales_with_configured_cv(self, cv: int, expected: int) -> None: + assert _logreg_cv_multiplier(cv) == expected + + def test_time_grows_with_cv(self) -> None: + kwargs = { + "n_trials": 5, + "n_samples": 10_000, + "embedder_dim": 768, + "max_iter": 100, + "class_multiplier": 20, + } + cheap = _time_for_linear(cv_multiplier=_logreg_cv_multiplier(3), **kwargs) + dear = _time_for_linear(cv_multiplier=_logreg_cv_multiplier(10), **kwargs) + # cv=10 costs 101/31 as many fits as cv=3; previously both priced the same. + assert dear == pytest.approx(cheap * 101 / 31) + + +class TestCpuSpeedup: + def test_single_core_is_a_no_op(self) -> None: + assert _cpu_speedup(1, 0.9) == 1.0 + + def test_speedup_is_sublinear(self) -> None: + # Amdahl with p=0.9 on 8 cores is ~4.7x, never the naive 8x. + assert 1.0 < _cpu_speedup(8, 0.9) < 8.0 + + def test_capped_however_many_cores(self) -> None: + assert _cpu_speedup(1024, 0.99) == _MAX_CPU_SPEEDUP + + def test_never_optimistic_past_the_cap(self) -> None: + assert _cpu_speedup(10_000, 1.0) <= _MAX_CPU_SPEEDUP + + @pytest.mark.parametrize(("cpu_count", "n_jobs", "expected"), [(16, 1, 16), (16, 4, 4), (16, 32, 1), (0, 1, 1)]) + def test_cores_per_trial_divides_by_concurrent_trials(self, cpu_count: int, n_jobs: int, expected: int) -> None: + assert _cores_per_trial(cpu_count, n_jobs) == expected + + +class TestCpuCountReachesTimeEstimates: + """The complaint was that core count changed nothing. It must now change something.""" + + _CATBOOST: ClassVar[dict[str, int]] = { + "n_trials": 3, + "n_samples": 10_000, + "n_features": 768, + "iterations": 1000, + "depth": 6, + "class_multiplier": 10, + } + + def test_catboost_cpu_time_falls_with_cores(self) -> None: + one = _time_for_catboost(on_gpu=False, cores=1, **self._CATBOOST) + many = _time_for_catboost(on_gpu=False, cores=16, **self._CATBOOST) + assert many < one + + def test_catboost_gpu_time_ignores_cores(self) -> None: + one = _time_for_catboost(on_gpu=True, cores=1, **self._CATBOOST) + many = _time_for_catboost(on_gpu=True, cores=64, **self._CATBOOST) + assert one == many + + def test_linear_time_falls_with_cores_but_less_than_catboost(self) -> None: + kwargs = { + "n_trials": 5, + "n_samples": 10_000, + "embedder_dim": 768, + "max_iter": 100, + "cv_multiplier": 31, + "class_multiplier": 20, + } + one = _time_for_linear(cores=1, **kwargs) + many = _time_for_linear(cores=16, **kwargs) + assert many < one + # L-BFGS only threads inside BLAS, so it must not claim CatBoost's speedup. + assert one / many == pytest.approx(_cpu_speedup(16, _LINEAR_PARALLEL_FRACTION)) + + def test_cpu_count_is_wired_through_run_preflight(self) -> None: + """End to end: two identical configs differing only in cpu_count must differ in time.""" + from autointent.advisor import run_preflight + + config = { + "search_space": [ + {"node_type": "scoring", "search_space": [{"module_name": "catboost", "iterations": 1000}]} + ] + } + stats = DatasetStats(n_samples=10_000, n_classes=20, avg_tokens=16, source="test") + + def profile(cpu_count: int) -> HardwareProfile: + return HardwareProfile( + accelerator="cpu", + device_name="test-cpu", + vram_gb=0.0, + ram_gb=32.0, + free_disk_gb=200.0, + cpu_count=cpu_count, + ) + + small = run_preflight(config, stats, profile(1)) + big = run_preflight(config, stats, profile(32)) + assert big.resource.time_hours < small.resource.time_hours + + +class TestEffectiveTrainFraction: + def test_holdout_removes_the_validation_share(self) -> None: + assert _effective_train_fraction(DataConfig(validation_size=0.2)) == pytest.approx(0.8) + + def test_cross_validation_leaves_one_fold_out(self) -> None: + assert _effective_train_fraction(DataConfig(scheme="cv", n_folds=5)) == pytest.approx(0.8) + + def test_separation_ratio_shrinks_it_further(self) -> None: + cfg = DataConfig(validation_size=0.2, separation_ratio=0.5) + assert _effective_train_fraction(cfg) == pytest.approx(0.4) + + def test_never_leaves_the_unit_interval(self) -> None: + assert 0.0 <= _effective_train_fraction(DataConfig(validation_size=1.0)) <= 1.0 + + +class TestSplitReadinessAgreement: + """The advisor must not green-light a dataset the splitter would reject.""" + + @pytest.mark.parametrize( + "data_config", + [ + DataConfig(), + DataConfig(scheme="cv", n_folds=5), + DataConfig(separation_ratio=0.5), + DataConfig(scheme="cv", n_folds=10, separation_ratio=0.3), + ], + ) + def test_advisor_flags_exactly_what_the_splitter_rejects(self, data_config: DataConfig) -> None: + """Pins the advisor to `check_split_readiness`'s minimum, so the two cannot drift apart.""" + minimum = _min_samples_per_class_for_config(config=data_config) + # One class one sample below the splitter's own threshold. + stats = _stats({"ok": 500, "starved": minimum - 1}) + report = _run_data_phase(stats, data_config) + assert "Stratified splitting will fail" in _messages(report) + assert "starved" in _messages(report) + assert any(f.severity is Severity.OVER for f in report.findings) + + @pytest.mark.parametrize( + "data_config", + [DataConfig(), DataConfig(scheme="cv", n_folds=5), DataConfig(separation_ratio=0.5)], + ) + def test_silent_when_every_class_clears_the_threshold(self, data_config: DataConfig) -> None: + minimum = _min_samples_per_class_for_config(config=data_config) + stats = _stats({"a": minimum * 100, "b": minimum * 100}) + report = _run_data_phase(stats, data_config) + assert "Stratified splitting will fail" not in _messages(report) + + +class TestLogregCheckAccountsForTheSplit: + def test_class_that_only_passes_on_the_raw_split_is_flagged(self) -> None: + """The regression: 4 samples clears cv=3 before splitting, and fails after.""" + stats = _stats({"plenty": 500, "borderline": 4}) + cfg = DataConfig(validation_size=0.2) # 4 * 0.8 = 3.2 -> 3 usable... still >= 3 + assert int(4 * _effective_train_fraction(cfg)) == 3 + + # With separation_ratio the same class drops to 4 * 0.8 * 0.5 = 1 usable sample. + split_cfg = DataConfig(validation_size=0.2, separation_ratio=0.5) + report = _run_data_phase(stats, split_cfg, cv=3) + assert "LogisticRegressionCV (cv=3) will fail" in _messages(report) + assert "borderline" in _messages(report) + + def test_message_names_the_split_when_one_applies(self) -> None: + stats = _stats({"plenty": 500, "thin": 3}) + report = _run_data_phase(stats, DataConfig(validation_size=0.2), cv=3) + assert "after the 80% train/validation split" in _messages(report) + + def test_generous_class_counts_stay_silent(self) -> None: + stats = _stats({"a": 1000, "b": 1000}) + report = _run_data_phase(stats, DataConfig(validation_size=0.2), cv=3) + assert "LogisticRegressionCV" not in _messages(report) + + def test_multilabel_skips_the_cv_check(self) -> None: + """Multilabel uses plain LogisticRegression, which has no inner CV.""" + stats = _stats({"a": 1000, "thin": 1}, multilabel=True) + report = _run_data_phase(stats, DataConfig(validation_size=0.2), cv=3) + assert "LogisticRegressionCV" not in _messages(report) + + def test_declared_cv_is_used_not_the_default(self) -> None: + stats = _stats({"a": 1000, "mid": 40}) + assert "LogisticRegressionCV" not in _messages(_run_data_phase(stats, DataConfig(), cv=3)) + assert "LogisticRegressionCV (cv=50) will fail" in _messages(_run_data_phase(stats, DataConfig(), cv=50))