diff --git a/compute-feasibility-advisor-proposal.md b/compute-feasibility-advisor-proposal.md new file mode 100644 index 000000000..7ebf70bd9 --- /dev/null +++ b/compute-feasibility-advisor-proposal.md @@ -0,0 +1,281 @@ +# 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/pyproject.toml b/pyproject.toml index 5fbabaf86..3bc17dc00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ dependencies = [ "aiofiles (>=24.1.0,<25.0.0)", "threadpoolctl (>=3.0.0,<4.0.0)", "packaging (>=23.2)", + "psutil (>=5.9.0,<8.0.0)", ] [project.optional-dependencies] @@ -121,6 +122,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)", @@ -145,6 +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" [build-system] requires = ["uv_build>=0.8.7,<0.9.0"] diff --git a/scripts/calibrate_advisor.py b/scripts/calibrate_advisor.py new file mode 100644 index 000000000..d4dec0f29 --- /dev/null +++ b/scripts/calibrate_advisor.py @@ -0,0 +1,1299 @@ +"""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 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 + 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 + 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" + 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: + 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 + # 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 + + 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), + } + 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 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 new file mode 100755 index 000000000..1c6f0d869 --- /dev/null +++ b/scripts/run_calibration_banking77.sh @@ -0,0 +1,163 @@ +#!/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/src/autointent/_advisor/__init__.py b/src/autointent/_advisor/__init__.py new file mode 100644 index 000000000..4df486f2b --- /dev/null +++ b/src/autointent/_advisor/__init__.py @@ -0,0 +1,42 @@ +"""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. +""" + +from __future__ import annotations + +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, +) + +__all__ = [ + "BUNDLED_PRESETS", + "DatasetStats", + "Finding", + "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/_cli.py b/src/autointent/_advisor/_cli.py new file mode 100644 index 000000000..6c1b2bc2e --- /dev/null +++ b/src/autointent/_advisor/_cli.py @@ -0,0 +1,134 @@ +"""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`` (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. + +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 autointent._advisor import inspect, recommend, stats_from_dataset + +from ._render import render_json, render_recommendation, render_text +from ._report import DatasetStats + +logger = logging.getLogger("autointent.advisor") + + +def _stats_from_args(args: argparse.Namespace) -> DatasetStats: + multilabel = args.task == "multilabel" + if args.dataset: + 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=multilabel, + ) + + +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: + report = inspect( + args.target, + stats=_stats_from_args(args), + budget_vram_gb=args.budget_vram_gb, + ) + if args.json: + sys.stdout.write(render_json(report)) + 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: + result = recommend( + stats=_stats_from_args(args), + budget_vram_gb=args.budget_vram_gb, + budget_time_h=args.budget_time_h, + ) + if args.json: + sys.stdout.write(json.dumps(result.to_dict(), indent=2, default=str)) + sys.stdout.write("\n") + else: + sys.stdout.write(render_recommendation(result.results, result.chosen)) + sys.stdout.write("\n") + if result.chosen: + sys.stdout.write("\n") + sys.stdout.write(render_text(dict(result.results)[result.chosen])) + sys.stdout.write("\n") + return 0 if result.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 int(args.func(args)) + + +if __name__ == "__main__": + main() 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..6fe3985b1 --- /dev/null +++ b/src/autointent/_advisor/_estimates/_formulas.py @@ -0,0 +1,335 @@ +"""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. + + 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_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 + + +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)) + + +# 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( + *, + n_trials: int, + epochs: int, + batch_size: int, + seq_len: int, + n_samples: int, + 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 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_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: + """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 + + +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 + + +# 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 +# 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, # 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 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 + * n_samples + * embedder_dim + * 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..884254faf --- /dev/null +++ b/src/autointent/_advisor/_estimates/_resource.py @@ -0,0 +1,599 @@ +"""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, Callable + +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, + _embedding_cache_disk_gb, + _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", +} + +# 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: + """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, + seq_len=seq_len, + n_samples=stats.n_samples, + 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) + + 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 _apply_embedding_cache( + module_estimates: list[_ModuleEstimate], + seen_models: dict[str, ModelMeta], + *, + stats: DatasetStats, + hardware: HardwareProfile, + cache_probe: Callable[[str], bool] | None = None, +) -> 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. + + ``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: + continue + 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 + me.driver["time_hours"] = 0.0 + me.driver["mode"] = f"{me.driver['mode']}+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" + return paid + + +def _aggregate_disk( + estimate: ResourceEstimate, + seen_models: dict[str, ModelMeta], + node_max_weights: dict[int, float], + *, + dump_modules: bool, + n_trials: int, + cached_embedders: set[str] | None = None, + stats: DatasetStats | None = None, +) -> None: + """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 + 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 + + 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, + 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 + 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") + + if estimate.time_hours > 0: + report.add( + "resource", + Severity.AMPLE, + f"Time ~{estimate.time_hours:.1f} h (worst case, no HPO pruning)", + metric="time", + ) + + +_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, + 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. + + 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. + + ``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) + 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: + 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( + 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) + + # 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, + ) + + 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, + 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). + # 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 + 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 + _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/_hardware.py b/src/autointent/_advisor/_hardware.py new file mode 100644 index 000000000..6aa9741ee --- /dev/null +++ b/src/autointent/_advisor/_hardware.py @@ -0,0 +1,142 @@ +"""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 pathlib import Path +from typing import Literal + +import psutil +import torch + +logger = logging.getLogger(__name__) + +Accelerator = Literal["cuda", "mps", "cpu"] + +# matches macOS PYTORCH_MPS_HIGH_WATERMARK_RATIO default +MPS_DEFAULT_BUDGET_RATIO = 0.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 +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 >= _HIGH_GPU_VRAM_GB: + return "high-gpu" + if self.vram_gb >= _MID_GPU_VRAM_GB: + return "mid-gpu" + return "low-gpu" + + +def _detect_ram_gb() -> float: + return float(psutil.virtual_memory().total) / _BYTES_PER_GB + + +def _detect_free_disk_gb(path: str | None = None) -> float: + 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 / _BYTES_PER_GB + 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: + if not torch.cuda.is_available(): + return None + idx = 0 + try: + _free, total = torch.cuda.mem_get_info(idx) + vram_gb = total / _BYTES_PER_GB + 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: + 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( + *, + 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..3dc48aaab --- /dev/null +++ b/src/autointent/_advisor/_hub.py @@ -0,0 +1,234 @@ +"""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 json +import logging +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +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__) + +# 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 + + +@dataclass +class ModelMeta: + name: str + total_params: int + weight_bytes_per_param: float + total_file_bytes: int + cached_locally: bool + confidence: Confidence + hidden_size: int | None = None + n_layers: int | None = None + + @property + def disk_gb(self) -> float: + return self.total_file_bytes / _BYTES_PER_GB + + @property + def weights_gb(self) -> float: + return (self.total_params * self.weight_bytes_per_param) / _BYTES_PER_GB + + +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"] + for fname in weight_files: + path = try_to_load_from_cache(model_name, fname) + if isinstance(path, str): + 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: + 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 + # 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 when a checkpoint stores tensors in multiple dtypes. + _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: Confidence = "hub" + if total_params == 0: + total_params = _DEFAULT_HEURISTIC_PARAMS + confidence = "heuristic" + + if total_file_bytes == 0: + total_file_bytes = int(total_params * 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 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, + 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), + 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; " + "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( + name=model_name, + 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", + hidden_size=_DEFAULT_HEURISTIC_HIDDEN, + n_layers=_DEFAULT_HEURISTIC_LAYERS, + ) + + +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. + + Always returns a value — never raises — so the advisor can keep going + on offline machines or for unknown checkpoints. + """ + if _looks_like_local_path(model_name): + return ModelMeta( + name=model_name, + total_params=_DEFAULT_HEURISTIC_PARAMS, + weight_bytes_per_param=_DEFAULT_BYTES_PER_PARAM, + total_file_bytes=0, + cached_locally=True, + confidence="heuristic", + ) + + # _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 new file mode 100644 index 000000000..afd541e2b --- /dev/null +++ b/src/autointent/_advisor/_render.py @@ -0,0 +1,154 @@ +"""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, Any + +if TYPE_CHECKING: + from ._report import PreflightReport + +_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 ''.""" + 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[str, Any]]) -> list[str]: + """Format the Drivers of cost section as an aligned table.""" + visible = drivers[:_DRIVERS_LIMIT] + rows: list[tuple[str, ...]] = [ + ( + 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']}]", + ) + for d in visible + ] + + 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" + 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.extend(_render_drivers_table(report.resource.drivers)) + lines.append("") + + if report.notes: + lines.append("Notes:") + lines.extend(f" • {note}" for note in report.notes) + lines.append("") + + summary = f"Verdict: {'feasible' if report.is_feasible else 'INFEASIBLE'} " + summary += f"(headroom: {report.headroom.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} {'Headroom':<10}") + 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.headroom.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..6fc3d8f2f --- /dev/null +++ b/src/autointent/_advisor/_report.py @@ -0,0 +1,132 @@ +"""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): + AMPLE = "ample" + TIGHT = "tight" + OVER = "over" + + +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 + disk_embedding_cache_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 + self.disk_embedding_cache_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 + # 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 + 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 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.AMPLE + return max((f.severity for f in self.findings), key=lambda s: order[s]) + + @property + def is_feasible(self) -> bool: + 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["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/runner.py b/src/autointent/_advisor/runner.py new file mode 100644 index 000000000..4743a3b47 --- /dev/null +++ b/src/autointent/_advisor/runner.py @@ -0,0 +1,187 @@ +"""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, Callable + +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, + embedding_cache_probe: Callable[[str], bool] | None = None, +) -> 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. + 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. + """ + 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, + 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) + + 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.", + ) + + # 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"} + 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.", + ) diff --git a/src/autointent/_advisor/workflows.py b/src/autointent/_advisor/workflows.py new file mode 100644 index 000000000..91c3c24a7 --- /dev/null +++ b/src/autointent/_advisor/workflows.py @@ -0,0 +1,421 @@ +"""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 ._hardware import detect_hardware +from ._report import DatasetStats, RecommendationResult, Severity +from .runner import run_preflight + +if TYPE_CHECKING: + from collections.abc import Iterable + + from autointent import Dataset + + 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, + 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: + 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 _class_counts( + train: Any, # noqa: ANN401 + label_col: str, + multilabel: bool, + n_classes: int, +) -> 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 {} + 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 counts + + +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 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() + 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)) + + 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: (cost_rank.get(pair[0], len(BUNDLED_PRESETS)), pair[0])) + 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/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..491bc57f4 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,20 @@ 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/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/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index cbfa82576..59e6b87a3 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -117,18 +117,25 @@ 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 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/__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..e1ffd77d5 --- /dev/null +++ b/tests/advisor/test_estimates_and_cli.py @@ -0,0 +1,248 @@ +"""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 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 + + _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, + ) + + +@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 + # 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") + 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(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") + assert report.is_feasible + + +def test_n_jobs_doubles_vram_findings() -> None: + 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)) + 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["headroom"] in {"ample", "tight", "over"} + # 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_bi"}, + ], + } + ], + } + 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) + + +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"] == "over" and "exceeds budget" in f["message"] for f in r["report"]["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 + + +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..99bd63243 --- /dev/null +++ b/tests/advisor/test_estimates_internals.py @@ -0,0 +1,771 @@ +"""Targeted tests for `_estimates` helpers + edge cases of `run_preflight`.""" + +from __future__ import annotations + +from typing import Any + +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 + +# 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.resolve_model.cache_clear() + monkeypatch.setattr(_hub, "_is_warm_cached", lambda _name: False) + # Resource phase calls `_hub.resolve_model(...)` via module reference, so + # patching the symbol on `_hub` is enough. + monkeypatch.setattr(_hub, "resolve_model", _fake_resolve) + + +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.AMPLE + + def test_above_yellow_threshold(self) -> None: + 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 + 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.TIGHT + + +class TestVramForTransformer: + @pytest.fixture + def meta(self) -> ModelMeta: + return ModelMeta( + name="x", + total_params=100_000_000, + 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") + 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", 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 + assert train_acts > inf_acts + # 12-layer model: training activations should be at least ~5x inference. + assert train_acts / inf_acts > 5 + + +def test_ram_scales_with_dataset_size() -> None: + meta = ModelMeta( + name="x", + total_params=100_000_000, + 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}, + "logging_config": {"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()) + 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: + 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 + # 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 = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear"}, + ], + } + ] + } + stats = DatasetStats( + n_samples=20, + n_classes=5, + avg_tokens=10, + 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 (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) + + 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.OVER] + 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.TIGHT 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 + # 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 = { + "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 + # 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 = { + "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 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=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=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 + + 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 + assert 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.""" + + 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}, + "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 + # 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}, + "logging_config": {"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) + + +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}, + } + # 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 + 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 + + 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_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..c19018235 --- /dev/null +++ b/tests/advisor/test_hub_heuristics.py @@ -0,0 +1,66 @@ +"""Tests for the offline heuristic fallback in `_hub`. + +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 + +import pytest + +from autointent._advisor import _hub + + +@pytest.fixture(autouse=True) +def _offline(monkeypatch: pytest.MonkeyPatch) -> None: + _hub.resolve_model.cache_clear() + # 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) + + +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.total_params * 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.total_params > 0 + assert meta.disk_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/advisor/test_render.py b/tests/advisor/test_render.py new file mode 100644 index 000000000..7a806c7f2 --- /dev/null +++ b/tests/advisor/test_render.py @@ -0,0 +1,170 @@ +"""Output rendering: text formatting and JSON serialization.""" + +from __future__ import annotations + +import json + +from autointent._advisor._render import _batch_hint, 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.TIGHT, "VRAM ~6 GB vs available 8 GB") + r.add("data", Severity.OVER, "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_headroom(self) -> None: + out = render_text(_populated_report()) + assert "Verdict: INFEASIBLE" in out + assert "headroom: over" 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 {"ample", "tight", "over"} + + def test_headroom_and_feasibility_serialized(self) -> None: + d = json.loads(render_json(_populated_report())) + 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["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.AMPLE, "ok") + b = PreflightReport(preset_name="b", resource=ResourceEstimate(vram_gb=8.0, time_hours=4.0)) + b.add("resource", Severity.OVER, "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 + + +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( + 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..acb2b5bf8 --- /dev/null +++ b/tests/advisor/test_report.py @@ -0,0 +1,87 @@ +"""Unit tests for the report dataclasses.""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from autointent._advisor._report import ( + DatasetStats, + Finding, + PreflightReport, + ResourceEstimate, + Severity, +) + + +class TestSeverityOrdering: + 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.AMPLE, "ok") + r.add("data", Severity.TIGHT, "warn") + assert r.headroom == Severity.TIGHT + r.add("config", Severity.OVER, "fail") + assert r.headroom == Severity.OVER # type: ignore[comparison-overlap] + + def test_is_feasible_flips_on_any_red(self) -> None: + r = PreflightReport() + r.add("resource", Severity.TIGHT, "warn") + assert r.is_feasible is True + r.add("data", Severity.OVER, "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.OVER, "boom") + d = r.to_dict() + assert d["headroom"] == "over" + assert d["is_feasible"] is False + assert d["findings"] == [ + {"phase": "resource", "severity": "over", "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.AMPLE, message="ok") + with pytest.raises(dataclasses.FrozenInstanceError): + f.message = "changed" # type: ignore[misc] 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/pipeline/test_calibration_tracker.py b/tests/pipeline/test_calibration_tracker.py new file mode 100644 index 000000000..7e5640035 --- /dev/null +++ b/tests/pipeline/test_calibration_tracker.py @@ -0,0 +1,150 @@ +"""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'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 diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py new file mode 100644 index 000000000..b9116fbe1 --- /dev/null +++ b/tests/pipeline/test_preflight.py @@ -0,0 +1,120 @@ +"""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() + 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 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())