diff --git a/.gitignore b/.gitignore index 4fc2150c1..1a8cf9a2c 100644 --- a/.gitignore +++ b/.gitignore @@ -184,4 +184,12 @@ vector_db* /wandb model_output/ my.py -.DS_store \ No newline at end of file +.DS_store + +# Local advisor validation artifacts + Superpowers process docs (never shipped) +res-adapt-ckeck/ +banking77_*_repredicted.json +claude-issue-investigation.md +review-of-review.md +docs/superpowers/ +.python-version diff --git a/docs/source/advisor.rst b/docs/source/advisor.rst new file mode 100644 index 000000000..6db50249a --- /dev/null +++ b/docs/source/advisor.rst @@ -0,0 +1,89 @@ +Compute feasibility advisor +=========================== + +.. note:: + + **Experimental.** The advisor's estimates are heuristic and calibrated against a limited hardware sample. Treat them as guidance, not guarantees, and read :ref:`advisor-accuracy` before relying on a number. The Python surface may change in a minor release. + +Optimizing a search space can take hours and needs more VRAM than a laptop GPU has. The advisor answers "will this fit, and how long will it take?" *before* anything is downloaded or trained. + +Command line +------------ + +Two subcommands. ``inspect`` prices a specific preset or config; ``recommend`` detects your hardware and picks the heaviest bundled preset that still fits. + +.. code-block:: bash + + # What will transformers-light cost on this machine? + autointent-advisor inspect transformers-light + + # ...against a real dataset rather than placeholder sizes + autointent-advisor inspect transformers-light --dataset banking77 + + # Which preset should I use? + autointent-advisor recommend --dataset banking77 + + # Machine-readable output + autointent-advisor inspect ./my-config.yaml --json + +Without ``--dataset``, the advisor uses placeholder dataset sizes (``--n-samples``, ``--n-classes``, ``--avg-tokens``, ``--task``), so it is useful before you have assembled any data. ``--budget-vram-gb`` overrides hardware detection, and ``recommend`` also accepts ``--budget-time-h``. + +Both subcommands exit non-zero when nothing is feasible, so they work as a CI gate. + +Reading a report +---------------- + +Each finding carries a severity: + +``ample`` + Comfortably within budget. +``tight`` + Fits, but with little headroom โ€” expect swapping or thermal throttling. +``over`` + Exceeds the budget. Any ``over`` finding makes the whole report infeasible. + +The drivers table lists the modules that dominate the cost, so it shows *what* to change. ``low confidence`` on a report means Hub metadata was unavailable or incomplete for at least one model, and conservative large-model defaults were substituted โ€” the numbers are much rougher when you see it. + +Findings are not only about hardware. The advisor also prices your ``DataConfig``: it reports ``over`` when a class has too few samples for the stratified split to succeed, using the same minimum as :py:func:`~autointent.context.data_handler.check_split_readiness`, and when ``LogisticRegressionCV`` would not have ``cv`` samples per class *after* the train/validation split. Per-class counts are measured on the train split you supply, so the advisor discounts them by whatever ``validation_size``, ``n_folds``, and ``separation_ratio`` will take away. + +From Python +----------- + +.. code-block:: python + + from autointent import Dataset + from autointent.advisor import dataset_stats, detect_hardware, estimate, recommend + + report = estimate("transformers-light") + print(report.is_feasible, report.resource.vram_gb) + + result = recommend(stats=dataset_stats(Dataset.from_json(path))) + print(result.chosen) + +``reduce_to_fit`` goes further: it prunes the most expensive scoring module repeatedly until the search space fits, raising ``ReduceToFitError`` if nothing does. + +Inside ``Pipeline.fit`` +----------------------- + +``Pipeline.fit`` accepts a ``preflight`` gate. It defaults to ``"off"``, so the advisor never runs unless you ask โ€” it makes network calls to the Hugging Face Hub for model metadata, which does not belong on every fit by default. + +.. code-block:: python + + pipeline.fit(dataset, preflight="warn") # log findings, always continue + pipeline.fit(dataset, preflight="strict") # raise PreflightError if infeasible + +``"strict"`` raises :class:`autointent.advisor.PreflightError` before allocating any VRAM, which is the useful mode in CI. + +.. _advisor-accuracy: + +How accurate is it? +------------------- + +Validated end to end on one machine class (RTX 3060 Laptop, 6 GB VRAM / 16 GB RAM), where all four fitted presets matched their predicted verdict: both ``over`` predictions did run out of memory, and both feasible predictions did fit. Known limits: + +- **Feasibility verdicts are the reliable part.** That is what the advisor was built and validated for. +- **VRAM is close but not a guaranteed ceiling.** One preset used 1.22ร— its prediction. Leave headroom rather than trusting the figure exactly. +- **Wall-time estimates are indicative only.** Measured error has run in both directions across formula revisions, once by more than an order of magnitude for cross-encoders. ``--budget-time-h`` inherits that uncertainty. +- **CPU parallelism is modelled, not measured.** The CPU coefficients are calibrated single-threaded; core count is then applied as a capped Amdahl speedup (higher for CatBoost, which uses every core by default, than for scikit-learn's L-BFGS, which only threads inside BLAS), divided across concurrent ``hpo_config.n_jobs`` trials. It is a correction for the fact that core count used to change nothing at all, not a validated speedup curve. +- **Preset ranking does not depend on time estimates.** ``recommend`` orders presets by a declared cost ranking, so unstable time figures cannot reorder its choice. +- **Only one hardware class has been validated end to end.** Treat other machines as unverified. diff --git a/docs/source/index.rst b/docs/source/index.rst index 9d0fb5eed..5254e7c6f 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -68,6 +68,9 @@ Reference :doc:`๐ŸŒ Inference servers ` Deploy a trained pipeline behind HTTP (FastAPI) or MCP (FastMCP): installation extras, environment variables, and how to run each server. +:doc:`๐Ÿ” Compute feasibility advisor ` + Estimate VRAM, RAM, disk, and wall-time for a search space before training. Run it from the CLI or gate ``Pipeline.fit`` on it. + :doc:`๐Ÿ”ง API Reference ` Complete technical documentation for all classes, methods, and functions. Essential reference for developers integrating AutoIntent into their applications. @@ -84,4 +87,5 @@ Reference user_guides learn/index server + advisor autoapi/autointent/index \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 06cbda236..280a5e4fd 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" +"autointent-advisor" = "autointent.advisor._cli:main" [build-system] requires = ["uv_build>=0.8.7,<0.9.0"] diff --git a/src/autointent/_pipeline/__init__.py b/src/autointent/_pipeline/__init__.py index 7a8af8259..cc7fe54eb 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, PreflightMode -__all__ = ["Pipeline"] +__all__ = ["Pipeline", "PreflightMode"] diff --git a/src/autointent/_pipeline/_pipeline.py b/src/autointent/_pipeline/_pipeline.py index 7c24c40a6..0c5b9a864 100644 --- a/src/autointent/_pipeline/_pipeline.py +++ b/src/autointent/_pipeline/_pipeline.py @@ -5,7 +5,7 @@ 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 @@ -34,10 +34,14 @@ if TYPE_CHECKING: from autointent import Dataset + from autointent.advisor import Finding, PreflightReport from autointent.custom_types import ListOfGenericLabels, SearchSpacePreset, SearchSpaceValidationMode from autointent.modules.base import BaseDecision, BaseRegex, BaseScorer +PreflightMode = Literal["off", "warn", "strict"] + + class Pipeline: """Pipeline optimizer class. @@ -152,6 +156,45 @@ 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``. + + Imported lazily: the advisor probes the HF Hub, so it must stay off the + ``import autointent`` path. + """ + from autointent.advisor import PreflightError, Severity, dataset_stats, detect_hardware, run_preflight + + config = self._build_advisor_config() + stats = dataset_stats(dataset) + hardware = detect_hardware() + report = run_preflight(config, stats, hardware, refit_after=refit_after) + _log_preflight_report(report, self._logger) + if mode == "strict": + over: list[Finding] = [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,22 +236,42 @@ def fit( dataset: Dataset, refit_after: bool = False, incompatible_search_space: SearchSpaceValidationMode = "filter", + preflight: PreflightMode = "off", ) -> Context: """Optimize the pipeline from dataset. Args: dataset: dataset for optimization. refit_after: whether to refit on whole data after optimization. Valid only for hold-out validaiton. - sampler: sampler type to use. - incompatible_search_space: wow to handle data-incompatible modules occurring in search space. + incompatible_search_space: how to handle data-incompatible modules occurring in search space. + preflight: **experimental** gate that runs + :func:`autointent.advisor.run_preflight` over the pipeline's + effective config + dataset before any heavy work. + ``"off"`` (default) skips it entirely. ``"warn"`` logs findings โ€” + INFO for AMPLE, WARNING for TIGHT, ERROR for OVER โ€” but never + raises; note it probes the HF Hub for model metadata, so it adds + network round-trips. ``"strict"`` additionally raises + :class:`autointent.advisor.PreflightError` when any finding has + severity OVER, so unfeasible runs abort before fit. Raises: RuntimeError: If pipeline is in inference mode. + 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) + # Filter the search space first: ``validate_modules`` drops modules this + # dataset cannot use (e.g. ``mlknn`` on multiclass, ``dnnc`` and its ~6.4 GB + # reranker on multilabel), and preflight must price what will actually run + # rather than what was requested. It takes ``dataset`` only, so it does not + # depend on the ``Context`` built below. + self.validate_modules(dataset, mode=incompatible_search_space) + + if preflight != "off": + self._run_preflight(dataset, refit_after=refit_after, mode=preflight) + context = Context(self._seed) context.set_dataset(dataset, self.data_config) context.configure_logging(self.logging_config) @@ -218,8 +281,6 @@ def fit( context.configure_hpo(self.hpo_config) context.configure_vector_index(self.vector_index_config) - self.validate_modules(dataset, mode=incompatible_search_space) - test_utterances = context.data_handler.test_utterances() if test_utterances is None: self._logger.warning( @@ -472,3 +533,25 @@ 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.""" + # Imported lazily for the same reason as in ``Pipeline._run_preflight``: the + # advisor probes the HF Hub, so it must stay off the ``import autointent`` + # path. Do not hoist to module scope. + from autointent.advisor import Severity + + level_for = { + 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/advisor/__init__.py b/src/autointent/advisor/__init__.py new file mode 100644 index 000000000..42f34b95a --- /dev/null +++ b/src/autointent/advisor/__init__.py @@ -0,0 +1,37 @@ +"""Pre-flight compute feasibility advisor. + +**Experimental.** This subpackage estimates VRAM, RAM, disk, and wall-time for a +search space before any training starts. Estimates are heuristic and calibrated +against a limited hardware sample, so treat them as guidance rather than +guarantees โ€” see the accuracy caveats in the ``advisor`` page of the docs. The +public surface may change in a minor release. + +Two ways in: the ``autointent-advisor`` console script, and the functions below. +``Pipeline.fit(preflight=...)`` wires the same machinery into a fit, opt-in. +""" + +from __future__ import annotations + +from ._errors import PreflightError +from ._hardware import HardwareProfile, detect_hardware +from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity +from ._runner import run_preflight +from ._workflows import ReduceToFitError, dataset_stats, estimate, recommend, reduce_to_fit + +__all__ = [ + "DatasetStats", + "Finding", + "HardwareProfile", + "PreflightError", + "PreflightReport", + "RecommendationResult", + "ReduceToFitError", + "ResourceEstimate", + "Severity", + "dataset_stats", + "detect_hardware", + "estimate", + "recommend", + "reduce_to_fit", + "run_preflight", +] diff --git a/src/autointent/advisor/_cli.py b/src/autointent/advisor/_cli.py new file mode 100644 index 000000000..dd79bd3b8 --- /dev/null +++ b/src/autointent/advisor/_cli.py @@ -0,0 +1,133 @@ +"""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.estimate` 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 ._render import render_json, render_recommendation, render_text +from ._report import DatasetStats +from ._workflows import estimate, recommend, stats_from_dataset + +logger = logging.getLogger("autointent.advisor") + + +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 = estimate( + 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/_errors.py b/src/autointent/advisor/_errors.py new file mode 100644 index 000000000..bd81c48cf --- /dev/null +++ b/src/autointent/advisor/_errors.py @@ -0,0 +1,18 @@ +"""Advisor exceptions raised across the package / Pipeline boundary.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ._report import Finding + + +class PreflightError(RuntimeError): + """Raised when ``Pipeline.fit(preflight="strict")`` finds OVER-budget resources.""" + + def __init__(self, findings: list[Finding]) -> None: + self.findings = findings + lines = "\n".join(f" [{f.phase}] {f.message}" for f in findings) + msg = f"Preflight check failed with {len(findings)} OVER finding(s):\n{lines}" + super().__init__(msg) diff --git a/src/autointent/advisor/_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..368a62e0f --- /dev/null +++ b/src/autointent/advisor/_estimates/_formulas.py @@ -0,0 +1,426 @@ +"""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, since the advisor aims to over- rather than + under-predict cost. +""" + +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: weights + grads + optimizer state. + + Multiplier chosen to over- rather than under-estimate, by mode: 1.3x + inference, 1.3x + 0.5 GB lora adapters, 4.5x full finetune (textbook 4W + + fragmentation/workspaces slack). + """ + weights_gb = meta.weights_gb + if mode == "inference": + 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: + """Activation memory per sample. + + Training: 34 B/token/layer (Korthikanti 2022 upper bound, standard + attention). Inference: 8 B/token (only 1-2 layers' outputs in flight). + """ + hidden = _embedder_dim(meta) + training_bytes_per_token_per_layer = 34 + inference_bytes_per_token = 8 + 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: weights + grads + optimizer state + activations x batch. + + Safety margin: 1.20 for training (backward transients, eval sweep, + allocator fragmentation), 1.10 for inference (no backward). + """ + base = _weights_vram_for_transformer(meta, mode) + if batch_size <= 0: + return base + is_training = mode != "inference" + per_sample = _activations_gb_per_sample(meta, seq_len, is_training=is_training) + safety = 1.20 if is_training else 1.10 + return (base + per_sample * batch_size) * safety + + +def _max_fitting_batch_size( + *, + 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 HF-Trainer MFU (~20% on A100), +# not peak spec sheet. Advisor aims to over- rather than under-predict time, +# so pessimistic (low) values here. +_DEVICE_TFLOPS = { + "high-gpu": 60.0, # A100 / H100 + "mid-gpu": 20.0, # V100 / RTX 3090 / A6000 + "low-gpu": 7.0, # T4 / RTX 3060 / 8 GB consumer card + "apple-silicon": 4.0, # M1/M2/M3 GPU cores + "cpu": 0.05, # single-thread modern x86 with MKL +} +_DEFAULT_TFLOPS = 7.0 # unknown device โ†’ treat as low-GPU + +# HF Trainer overhead: eval sweeps, save syncs, dataloader idle. The raw +# FLOPs formula only counts optimizer steps. +_TRAINER_OVERHEAD_MULT = 1.35 + + +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 wall-time in hours. + + Per-step FLOPs = 6 x params x batch x seq_len (fwd+bwd), รท sustained + device TFLOPS, x total steps x trainer overhead. + """ + steps_per_epoch = max(1, n_samples // max(1, batch_size)) + total_steps = n_trials * epochs * steps_per_epoch + # 6x factor: ~2x for fwd matmul + ~4x for bwd (grad wrt input + grad wrt weight). + 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 * _TRAINER_OVERHEAD_MULT) / 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, *, mode: str = "inference") -> float: + """RAM estimate: weights x mode_mult + tokenized text (n_samples x avg_tokens x 4 B). + + Mode multiplier chosen to over- rather than under-estimate: 1.3 inference, + 1.5 lora, 4.5 full-finetune (Adam mirrors weights on host too). + """ + if mode == "inference": + weights_mult = 1.3 + elif mode == "lora": + weights_mult = 1.5 + else: + weights_mult = 4.5 + tokens_gb = (stats.n_samples * stats.avg_tokens * 4) / _BYTES_PER_GB + return meta.weights_gb * weights_mult + tokens_gb + + +def _embedding_cache_disk_gb(n_samples: int, hidden_size: int) -> float: + """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 on 1-thread CPU (OMP_NUM_THREADS=1), +# seconds per fit-work-unit. Typical L-BFGS iteration count baked in. +_LINEAR_CPU_S_PER_SAMPLE_FEATURE = 1.2e-9 +_CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER = 1e-9 +_CATBOOST_GPU_SPEEDUP = 10.0 +_LOGREG_CS = 10 # LogisticRegressionCV(Cs=10) sklearn default; LinearScorer does not expose it +_LOGREG_DEFAULT_CV = 3 # LinearScorer(cv=3) default +_CATBOOST_DEFAULT_BINS = 254 # CatBoost `border_count` default +_CATBOOST_BYTES_PER_TREE_NODE = 32 + +# CPU parallelism. The coefficients above are calibrated at one thread, which +# left every CPU-bound estimate independent of core count โ€” a 4-core and a +# 64-core box priced identically. Speedup is modelled with Amdahl's law and +# capped: these estimates exist to bound cost from above, and an over-generous +# speedup turns a conservative estimate into an optimistic one. +_CATBOOST_PARALLEL_FRACTION = 0.90 # CatBoost `thread_count` defaults to every core +_LINEAR_PARALLEL_FRACTION = 0.50 # only the BLAS calls inside L-BFGS thread +_MAX_CPU_SPEEDUP = 8.0 # refuse to believe in more than 8x however many cores are reported + + +def _logreg_cv_multiplier(cv: int) -> int: + """Fits per ``LogisticRegressionCV`` run: a ``Cs x cv`` grid plus one final refit.""" + return _LOGREG_CS * max(1, cv) + 1 + + +def _cpu_speedup(cores: int, parallel_fraction: float) -> float: + """Amdahl speedup on ``cores``, capped at :data:`_MAX_CPU_SPEEDUP`.""" + n = max(1, cores) + if n == 1: + return 1.0 + speedup = 1.0 / ((1.0 - parallel_fraction) + parallel_fraction / n) + return min(speedup, _MAX_CPU_SPEEDUP) + + +def _cores_per_trial(cpu_count: int, n_jobs: int) -> int: + """Cores one HPO trial gets when ``n_jobs`` trials run concurrently.""" + return max(1, max(1, cpu_count) // max(1, n_jobs)) + + +def _ram_for_linear(*, stats: DatasetStats, embedder_dim: int) -> float: + """Float64 design matrix dominates; coefficients and L-BFGS history are small.""" + 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 โ€” API stability; typical L-BFGS convergence baked into coeff + cv_multiplier: int, + class_multiplier: int, + cores: int = 1, +) -> float: + """LogisticRegression wall time. + + O(n_samples x features x classes x cv) per fit; typical L-BFGS + convergence absorbed into the calibration constant. ``cores`` divides that + by the modest BLAS-only speedup L-BFGS gets โ€” sklearn's own CV loop runs + single-threaded here, since ``LinearScorer`` leaves ``n_jobs`` unset. + """ + seconds = n_trials * _LINEAR_CPU_S_PER_SAMPLE_FEATURE * n_samples * embedder_dim * cv_multiplier * class_multiplier + return seconds / _cpu_speedup(cores, _LINEAR_PARALLEL_FRACTION) / 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 _ram_for_sklearn( + *, + stats: DatasetStats, + embedder_dim: int, + n_estimators: int, + max_depth: int, + n_jobs: int, +) -> float: + """RandomForest RAM: (feature matrix + trees) x n_jobs. + + joblib workers each hold a full copy. + """ + per_worker_data = stats.n_samples * embedder_dim * 8 # fp64 default + n_leaves = min(2**max_depth, stats.n_samples) if max_depth > 0 else stats.n_samples + per_worker_trees = n_estimators * n_leaves * _CATBOOST_BYTES_PER_TREE_NODE + return float(((per_worker_data + per_worker_trees) * max(1, n_jobs)) / _BYTES_PER_GB) + + +def _embedder_load_ram_gb(meta: ModelMeta | None) -> float: + """Aggregate-level RAM penalty when a classic preset uses an embedder. + + Added on top of the max-driver RAM because embedder + multiple classic + scorers coexist in memory. Uses fp32 weights (transformers up-casts at + load) x 3.5 for weights + activation buffers + framework slack. + """ + if meta is None: + return 0.0 + return ((meta.total_params * 4) / _BYTES_PER_GB) * 3.5 + + +def _time_for_catboost( + *, + n_trials: int, + n_samples: int, + n_features: int, + iterations: int, + depth: int, + class_multiplier: int, + on_gpu: bool, + cores: int = 1, +) -> 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 + + On CPU, ``cores`` divides that: CatBoost's ``thread_count`` defaults to + every core, so core count is the single largest term the one-thread + calibration was missing. Ignored on GPU, where the device is the bottleneck. + """ + coeff = _CATBOOST_CPU_S_PER_SAMPLE_FEATURE_ITER + if on_gpu: + coeff /= _CATBOOST_GPU_SPEEDUP + seconds = n_trials * iterations * coeff * n_samples * n_features * depth * class_multiplier + if not on_gpu: + seconds /= _cpu_speedup(cores, _CATBOOST_PARALLEL_FRACTION) + return seconds / 3600.0 + + +# === CNN / RNN scorers =================================================== +# Small torch models (TextCNN, LSTM) trained from scratch on token ids. +_NN_MAX_VOCAB = 30_000 # upper bound on vocabulary size +_NN_DEFAULT_SEQ_LEN = 50 # VocabConfig.max_seq_length default +_NN_BYTES_PER_PARAM = 4 # fp32 +_NN_TRAIN_ACT_BYTES_PER_UNIT = 16 # ~4x backward + optimizer overhead + + +def _cnn_param_count(*, embed_dim: int, num_filters: int, n_kernels: int, n_classes: int) -> int: + """TextCNN params: embedding + conv (kernel width ~5) + fc.""" + return ( + _NN_MAX_VOCAB * embed_dim + + num_filters * embed_dim * n_kernels * 5 + + num_filters * n_kernels * max(1, n_classes) + ) + + +def _rnn_param_count(*, embed_dim: int, hidden_dim: int, n_classes: int) -> int: + """LSTM classifier params: embedding + 4-gate LSTM cell + fc.""" + return _NN_MAX_VOCAB * embed_dim + 4 * hidden_dim * (embed_dim + hidden_dim + 1) + hidden_dim * max(1, n_classes) + + +def _vram_for_nn(*, params: int, batch_size: int, hidden_dim: int) -> float: + """Weights + 3x optimizer/grads + activations. + + Same fp32, over- rather than under-estimate approach as transformers, + smaller hidden dim (embed_dim / num_filters). + """ + weights_gb = (params * _NN_BYTES_PER_PARAM) / _BYTES_PER_GB + activations_gb = (batch_size * _NN_DEFAULT_SEQ_LEN * hidden_dim * _NN_TRAIN_ACT_BYTES_PER_UNIT) / _BYTES_PER_GB + return 4 * weights_gb + activations_gb + + +def _ram_for_nn(*, params: int, stats: DatasetStats) -> float: + """Weights + tokenized text (int32 ids).""" + return ((params * _NN_BYTES_PER_PARAM) + (stats.n_samples * _NN_DEFAULT_SEQ_LEN * 4)) / _BYTES_PER_GB + + +def _time_for_nn( + *, + n_trials: int, + epochs: int, + batch_size: int, + n_samples: int, + params_millions: float, + device_class: str, +) -> float: + """Reuse transformer FLOPs formula. + + Small models slightly under-predict since they're memory-bandwidth-bound, + but within 2x for cost ranking. + """ + return _time_for_transformer( + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + seq_len=_NN_DEFAULT_SEQ_LEN, + n_samples=n_samples, + params_millions=params_millions, + device_class=device_class, + ) + + +def _floor_to_power_of_two(n: int) -> int: + """Largest power of two <= ``n``; returns 0 when ``n < 1``.""" + if n < 1: + 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..5de6a4407 --- /dev/null +++ b/src/autointent/advisor/_estimates/_resource.py @@ -0,0 +1,915 @@ +"""Resource-phase orchestration. + +Walks the validated search space, asks ``_formulas`` for per-module costs, +aggregates them into a ``ResourceEstimate``, and emits VRAM/RAM/disk/time +findings on the report. + +The public entry is ``_resource_phase`` at the bottom; everything above it is +private machinery. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from autointent.advisor import _hub +from autointent.advisor._report import ResourceEstimate, Severity +from autointent.configs._embedder import ( + EmbedderConfig, + OpenaiEmbeddingConfig, + SentenceTransformerEmbeddingConfig, + VllmEmbeddingConfig, +) + +from ._formulas import ( + _DEFAULT_SEQ_LEN, + _LINEAR_CPU_S_PER_SAMPLE_FEATURE, + _LOGREG_DEFAULT_CV, + _MULTICLASS_THRESHOLD, + _activations_gb_per_sample, + _classify_severity, + _cnn_param_count, + _cores_per_trial, + _embedder_dim, + _embedder_load_ram_gb, + _embedding_cache_disk_gb, + _largest_embedder, + _logreg_cv_multiplier, + _max_fitting_batch_size, + _ram_for_catboost, + _ram_for_linear, + _ram_for_module, + _ram_for_nn, + _ram_for_sklearn, + _rnn_param_count, + _time_for_catboost, + _time_for_linear, + _time_for_nn, + _time_for_transformer, + _vram_for_nn, + _vram_for_transformer, + _weights_vram_for_transformer, +) +from ._search_space import ( + _extract_model_names, + _max_int, + _walk_modules, + _walk_modules_indexed, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from autointent.advisor._hardware import HardwareProfile + from autointent.advisor._hub import ModelMeta + from autointent.advisor._report import DatasetStats, PreflightReport + + +# 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", "sklearn"} 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, mode=mode) + + 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, + hpo_n_jobs: int = 1, +) -> _ModuleEstimate | None: + """Cost row for a linear or catboost scorer (returns ``None`` for any other module).""" + module = entry.get("module_name", "?") + refit = _refit_factor(refit_after=refit_after, n_trials=n_trials) + # Both multinomial (multiclass) and one-vs-rest (multilabel) LR scale linearly in n_classes. + class_multiplier = max(1, stats.n_classes) + # Concurrent HPO trials share the box, so a trial does not get every core. + cores = _cores_per_trial(hardware.cpu_count, hpo_n_jobs) + + if module == "linear": + # cv is per-entry and the grid scales with it; assuming the default 3 + # under-priced every search space that tuned it. + cv = _max_int(entry.get("cv"), _LOGREG_DEFAULT_CV) + cv_multiplier = 1 if stats.multilabel else _logreg_cv_multiplier(cv) + ram = _ram_for_linear(stats=stats, embedder_dim=embedder_dim) + time_h = ( + _time_for_linear( + 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, + cores=cores, + ) + * refit + ) + vram = 0.0 + mode = f"linear-cv{cv}" if cv_multiplier > 1 else "linear" + elif module == "catboost": + on_gpu = entry.get("task_type") == "GPU" and hardware.accelerator == "cuda" + # CatBoost MultiClass loss grows per-class trees only above binary; binary uses + # 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, + cores=cores, + ) + * refit + ) + vram, ram = (ram_total, 0.0) if on_gpu else (0.0, ram_total) + mode = "catboost-gpu" if on_gpu else "catboost" + elif module == "sklearn": + # RandomForest is the common target; joblib replicates data across n_jobs. + n_estimators = _max_int(entry.get("n_estimators"), 100) + max_depth = _max_int(entry.get("max_depth"), 0) + sk_n_jobs = _max_int(entry.get("n_jobs"), 1) + ram = _ram_for_sklearn( + stats=stats, + embedder_dim=embedder_dim, + n_estimators=n_estimators, + max_depth=max_depth, + n_jobs=sk_n_jobs, + ) + # Rough O(n_estimators x n x features / n_jobs); real numbers vary + # wildly by criterion so use the linear coefficient as a proxy. + time_h = ( + n_trials + * _LINEAR_CPU_S_PER_SAMPLE_FEATURE + * stats.n_samples + * embedder_dim + * n_estimators + / max(1, sk_n_jobs) + / 3600.0 + ) * refit + vram = 0.0 + mode = f"sklearn-n_jobs={sk_n_jobs}" + else: + return None + + 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 _estimate_nn_entry( + *, + entry: dict[str, Any], + node_type: str, + stats: DatasetStats, + hardware: HardwareProfile, + n_trials: int, + refit_after: bool, +) -> _ModuleEstimate | None: + """Cost row for cnn / rnn scorers (returns None for anything else). + + Small torch models trained from scratch on token ids. + """ + module = entry.get("module_name", "?") + n_classes = max(1, stats.n_classes) + + if module == "cnn": + embed_dim = _max_int(entry.get("embed_dim"), 128) + num_filters = _max_int(entry.get("num_filters"), 100) + # kernel_sizes is a list of lists (e.g. [[3,4,5]]) โ€” count the largest variant. + kernel_sizes = entry.get("kernel_sizes") + n_kernels = 3 + if isinstance(kernel_sizes, list): + for candidate in kernel_sizes: + if isinstance(candidate, list): + n_kernels = max(n_kernels, len(candidate)) + elif isinstance(candidate, int): + n_kernels = max(n_kernels, 1) + hidden_dim = num_filters + params = _cnn_param_count( + embed_dim=embed_dim, num_filters=num_filters, n_kernels=n_kernels, n_classes=n_classes + ) + elif module == "rnn": + embed_dim = _max_int(entry.get("embed_dim"), 128) + hidden_dim = _max_int(entry.get("hidden_dim"), 512) + params = _rnn_param_count(embed_dim=embed_dim, hidden_dim=hidden_dim, n_classes=n_classes) + else: + return None + + batch_size = _max_int(entry.get("batch_size"), 64) + epochs = _max_int(entry.get("num_train_epochs"), 60) + + vram = _vram_for_nn(params=params, batch_size=batch_size, hidden_dim=hidden_dim) + ram = _ram_for_nn(params=params, stats=stats) + time_h = _time_for_nn( + n_trials=n_trials, + epochs=epochs, + batch_size=batch_size, + n_samples=stats.n_samples, + params_millions=params / 1_000_000, + device_class=hardware.device_class, + ) * _refit_factor(refit_after=refit_after, n_trials=n_trials) + + return _ModuleEstimate( + driver={ + "node_type": node_type, + "module": module, + "model": f"{module}-from-scratch", + "mode": "small-torch-train", + "vram_gb": round(vram, 2), + "ram_gb": round(ram, 2), + "time_hours": round(time_h, 2), + "batch_size": batch_size, + "max_batch_size": None, + "confidence": "heuristic", + }, + vram_gb=vram, + ram_gb=ram, + time_hours=time_h, + ) + + +def _probe_warm_models( + seen_models: dict[str, ModelMeta], + cache_probe: Callable[[str], bool] | None, +) -> set[str]: + """Model names the probe reports as already having embeddings on disk. + + Pre-populating these lets the first-seen module hit the cache-hit branch + instead of paying the forward, and skips them in the disk-cache aggregation + (already on disk). Without a probe nothing is warm โ€” the pessimistic cold + assumption the advisor shipped with. + """ + warm_models: set[str] = set() + if cache_probe is not None: + for name in seen_models: + if cache_probe(name): + warm_models.add(name) + return warm_models + + +def _mark_cache_hit(me: _ModuleEstimate, module: str, *, suffix: str) -> None: + """Zero a transformer entry's forward time because the embedding is cached. + + Only modules whose per-entry estimate bundles the embedder forward into + ``time_hours`` have anything to give back; the rest are left alone. + """ + if module in _EMBEDDER_FORWARD_TRANSFORMER_MODULES: + me.time_hours = 0.0 + me.driver["time_hours"] = 0.0 + me.driver["mode"] = f"{me.driver['mode']}+{suffix}" + + +def _charge_first_forward_if_classic( + me: _ModuleEstimate, + module: str, + model: str, + *, + seen_models: dict[str, ModelMeta], + stats: DatasetStats, + hardware: HardwareProfile, +) -> None: + """Add a synthetic embedder forward to a classic entry that first pays for ``model``. + + Per unique embedder the first cache-honoring entry pays the forward. Only + classic entries (linear / catboost) need it added, because their own cost + model assumes embeddings already exist; a transformer entry's ``time_hours`` + already bundles the forward, so for those this is deliberately a no-op โ€” + hence the guard and the ``_if_classic`` in the name. + """ + if module in {"linear", "catboost"}: + embedder_meta = seen_models.get(model) + forward_h = _time_for_transformer( + n_trials=1, + epochs=1, + batch_size=32, + seq_len=128, + n_samples=stats.n_samples, + params_millions=(embedder_meta.total_params / 1_000_000) if embedder_meta else 100.0, + device_class=hardware.device_class, + ) + me.time_hours += forward_h + me.driver["time_hours"] = round(me.time_hours, 2) + me.driver["mode"] = f"{me.driver['mode']}+embed" + + +def _apply_embedding_cache( + module_estimates: list[_ModuleEstimate], + seen_models: dict[str, ModelMeta], + *, + 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() + warm_models = _probe_warm_models(seen_models, cache_probe) + for me in module_estimates: + module = me.driver["module"] + if module not in _CACHE_HONORING_MODULES: + continue + model = me.driver["model"] + if model not in seen_models: # synthetic / "(no embedder)" rows + continue + if model in warm_models: + _mark_cache_hit(me, module, suffix="warm") + continue + if model in paid: + _mark_cache_hit(me, module, suffix="cached") + else: + paid.add(model) + _charge_first_forward_if_classic(me, module, model, seen_models=seen_models, stats=stats, hardware=hardware) + 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: + cached_meta = seen_models.get(name) + if cached_meta is None: + continue + estimate.disk_embedding_cache_gb += _embedding_cache_disk_gb( + n_samples=stats.n_samples, + hidden_size=_embedder_dim(cached_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", + ) + + +# Process-level memory floors. Every fit reserves ~1.5 GB RSS for torch + +# transformers + datasets + optuna before any weights load. +_PROCESS_BASELINE_RAM_GB = 1.5 +# CUDA driver context + cuDNN/cuBLAS workspace + caching allocator. Bigger +# for training (backward workspaces + optimizer scratch), smaller for +# inference-only. +_CUDA_BASELINE_VRAM_TRAINING_GB = 1.0 +_CUDA_BASELINE_VRAM_INFERENCE_GB = 0.3 + +_UNKNOWN_SCORER_MODULES: frozenset[str] = frozenset() +"""Extension point for scorers with no cost estimator โ€” they get a +not-estimated placeholder row instead of a silent zero.""" + +_NN_SCORER_MODULES = frozenset({"cnn", "rnn"}) + +_EMBEDDER_CONSUMING_MODULES = frozenset( + { + "linear", + "catboost", + "sklearn", + "knn", + "mlknn", + "retrieval", + "description_bi", + "description_cross", + "description_llm", + }, +) + + +def _uses_embedder(search_space: list[dict[str, Any]]) -> bool: + """True when any search-space module consumes the embedder.""" + return any(entry.get("module_name") in _EMBEDDER_CONSUMING_MODULES for _, entry in _walk_modules(search_space)) + + +# 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, + ) + + +@dataclass(frozen=True) +class _ResourceInputs: + """Configuration-shaped inputs to the resource phase. + + Bundled by provenance rather than by use: every field is derived from one + validated ``OptimizationConfig``, plus the caller-injected ``cache_probe``. + Individual passes read only what they need โ€” the classic pass reads + ``refit_after`` and ``n_jobs`` (concurrent trials divide the CPU cores each + one gets), and ``search_space`` / ``n_trials`` / ``dump_modules`` never + leave ``_resource_phase`` โ€” but threading them separately would mean a + dozen keyword arguments down each call. + + ``cross_encoder_model_name`` and ``transformer_model_name`` come from the + pipeline's top-level configs and act as the fallback model for modules that + 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``). + Carrying them is what fixes the disk-download under-count called out in the + follow-up review (missing 6.4 GB reranker in ``zero-shot-encoders``); they + are applied in :func:`_estimate_transformer_entries`, where an entry with no + model of its own falls back to them. + """ + + embedder_config: EmbedderConfig + search_space: list[dict[str, Any]] + n_trials: int + n_jobs: int + dump_modules: bool + refit_after: bool = False + cross_encoder_model_name: str | None = None + transformer_model_name: str | None = None + cache_probe: Callable[[str], bool] | None = None + + +def _estimate_transformer_entries( + transformer_entries: list[tuple[int, str, dict[str, Any]]], + inputs: _ResourceInputs, + stats: DatasetStats, + hardware: HardwareProfile, + seen_models: dict[str, ModelMeta], + effective_trials: Callable[[int, dict[str, Any] | None], int], +) -> tuple[list[_ModuleEstimate], dict[int, float]]: + """First pass: transformer-bearing modules. + + Also populates ``seen_models`` in place, which the classic pass reads to + derive ``embedder_dim`` from the largest model seen โ€” so this must run first. + + Returns ``(module_estimates, node_max_weights)``. + """ + global_embedder = _embedder_model_name(inputs.embedder_config) + cross_encoder_model_name = inputs.cross_encoder_model_name + transformer_model_name = inputs.transformer_model_name + refit_after = inputs.refit_after + + 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 _NN_SCORER_MODULES: + # cnn / rnn โ€” small torch models trained from scratch, no hub + # model to resolve. Route to the small-model heuristic. + nn_estimate = _estimate_nn_entry( + entry=entry, + node_type=node_type, + stats=stats, + hardware=hardware, + n_trials=effective_trials(node_idx, entry), + refit_after=refit_after, + ) + if nn_estimate is not None: + module_estimates.append(nn_estimate) + continue + elif module in _UNKNOWN_SCORER_MODULES: + # Placeholder so the row is visible instead of silently zeroed. + module_estimates.append(_not_estimated_row(node_type=node_type, module=module)) + 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=effective_trials(node_idx, entry), + refit_after=refit_after, + ) + module_estimates.append(me) + # Track heaviest weight per node so dump_modules is bounded by one + # selected variant per node x n_trials, not the sum of all candidates. + node_max_weights[node_idx] = max(node_max_weights.get(node_idx, 0.0), me.model_weights_gb) + return module_estimates, node_max_weights + + +def _estimate_classic_entries( + classic_entries: list[tuple[int, str, dict[str, Any]]], + inputs: _ResourceInputs, + stats: DatasetStats, + hardware: HardwareProfile, + seen_models: dict[str, ModelMeta], + effective_trials: Callable[[int, dict[str, Any] | None], int], +) -> list[_ModuleEstimate]: + """Second pass: linear / catboost / sklearn modules. + + Their cost depends on ``embedder_dim`` rather than on a checkpoint, so this + reads ``seen_models`` as populated by :func:`_estimate_transformer_entries` + and derives the dimension from the largest embedder seen there. + """ + refit_after = inputs.refit_after + + module_estimates: list[_ModuleEstimate] = [] + embedder_meta = _largest_embedder(seen_models) + embedder_dim_val = _embedder_dim(embedder_meta) + for node_idx, node_type, entry in classic_entries: + 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=effective_trials(node_idx, entry), + refit_after=refit_after, + hpo_n_jobs=inputs.n_jobs, + ) + if classic_estimate is not None: + module_estimates.append(classic_estimate) + return module_estimates + + +def _resource_phase( + inputs: _ResourceInputs, + stats: DatasetStats, + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + """Walk the validated search space, fold per-module costs into the report. + + Two passes: transformer-bearing modules first (collects ``seen_models`` so + the largest model can drive ``embedder_dim`` for the classic pass), then + linear / catboost. Disk, VRAM/RAM peak, time sum, and final findings are + folded onto the report. + """ + embedder_config = inputs.embedder_config + search_space = inputs.search_space + n_trials = inputs.n_trials + n_jobs = inputs.n_jobs + dump_modules = inputs.dump_modules + cache_probe = inputs.cache_probe + + seen_models: dict[str, ModelMeta] = {} + global_embedder = _embedder_model_name(embedder_config) + if global_embedder: + seen_models[global_embedder] = _hub.resolve_model(global_embedder) + + transformer_entries, classic_entries = _split_entries(search_space) + + # HPO distributes n_trials evenly across module_name candidates at each + # node, so each variant sees n_trials / n_variants on average. + variants_per_node: dict[int, int] = {} + for node_idx, _node_type, _entry in _walk_modules_indexed(search_space): + variants_per_node[node_idx] = variants_per_node.get(node_idx, 0) + 1 + + def _effective_trials(node_idx: int, entry: dict[str, Any] | None = None) -> int: # noqa: ARG001 + """Trials charged to this module. + + ``entry`` is reserved for future per-module caps (see git history for the + cardinality-cap experiment that broke description-scorer presets). + """ + divisor = max(1, variants_per_node.get(node_idx, 1)) + return max(1, n_trials // divisor) + + # First pass: transformer modules (also populates seen_models for the classic pass). + module_estimates, node_max_weights = _estimate_transformer_entries( + transformer_entries, + inputs, + stats, + hardware, + seen_models, + _effective_trials, + ) + + # Second pass: linear / catboost โ€” cost depends on embedder_dim, not a checkpoint. + embedder_meta = _largest_embedder(seen_models) + module_estimates += _estimate_classic_entries( + classic_entries, + inputs, + stats, + hardware, + seen_models, + _effective_trials, + ) + + # Cache-aware time/disk: must run before the fold below. + cached_embedders = _apply_embedding_cache( + 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) + + # Process baseline: paid once, not per-module. + estimate.ram_gb = max(estimate.ram_gb, 0.0) + _PROCESS_BASELINE_RAM_GB + # Embedder-load penalty: classic presets keep embedder weights + framework + # buffers alongside scorer RAM, which max-of-drivers hides. + if _uses_embedder(search_space) and embedder_meta is not None: + estimate.ram_gb += _embedder_load_ram_gb(embedder_meta) + # CUDA baseline only when we predict some VRAM AND run on CUDA. Mode read + # off drivers: any training row flips to the larger baseline. + if estimate.vram_gb > 0 and hardware.accelerator == "cuda": + is_training = any(d.get("mode") in {"full-finetune", "lora", "small-torch-train"} for d in estimate.drivers) + estimate.vram_gb += _CUDA_BASELINE_VRAM_TRAINING_GB if is_training else _CUDA_BASELINE_VRAM_INFERENCE_GB + + _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 ~2x under-prediction + # of large-model shapes slip past the reviewer. + heuristic_models = [m.name for m in seen_models.values() if m.confidence == "heuristic"] + if heuristic_models: + 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..0f1324600 --- /dev/null +++ b/src/autointent/advisor/_estimates/_search_space.py @@ -0,0 +1,103 @@ +"""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 + + +# Keys that describe the entry rather than a tunable hyperparameter. +_RESERVED_ENTRY_KEYS = frozenset({"module_name", "target_metric"}) +# Above this many combinations the exact count stops mattering for cost. +_CARDINALITY_CAP = 10_000 + + +def _module_cardinality(entry: dict[str, Any]) -> int | None: + """Unique configurations the module entry can produce. + + Returns 1 when every tunable field is a singleton, N for finite list + products (capped at 10_000), None when any field is a continuous + ``{low, high}`` range. + """ + product = 1 + for key, value in entry.items(): + if key in _RESERVED_ENTRY_KEYS: + continue + if isinstance(value, dict): + if "low" in value and "high" in value: + return None # continuous range + continue # non-range dict = fixed + if isinstance(value, list) and value: + product *= max(1, len(value)) + if product >= _CARDINALITY_CAP: + return _CARDINALITY_CAP + return product + + +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..73f247fcd --- /dev/null +++ b/src/autointent/advisor/_hardware.py @@ -0,0 +1,151 @@ +"""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: + """The machine budget every estimate is scored against. + + Produced by :func:`~autointent.advisor.detect_hardware`, or built by hand to + size a search space for a machine you are not currently on. All sizes are in + binary gigabytes (GiB). ``vram_gb`` is 0.0 on CPU-only hosts, and on Apple + silicon it is a fraction of unified memory rather than dedicated VRAM. + ``notes`` carries such caveats and any manual VRAM override that was applied. + """ + + accelerator: Accelerator + device_name: str + vram_gb: float + 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..58366b188 --- /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 ~2x. Because the advisor aims to over- rather than under-predict, the +# offline fallback needs to over-estimate small models rather than under-estimate +# large ones. Callers can still see the fallback happened via ``confidence == +# "heuristic"`` and ``PreflightReport.low_confidence``. +_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 aim to over- 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..1d4954dfe --- /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 guidance, not measurements or guarantees.") + 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..5007780e5 --- /dev/null +++ b/src/autointent/advisor/_report.py @@ -0,0 +1,179 @@ +"""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): + """How much headroom a finding leaves against the detected budget. + + * ``AMPLE`` โ€” comfortably within budget; informational only. + * ``TIGHT`` โ€” expected to fit, but with little margin; the estimate is + heuristic, so treat this as "may not fit". + * ``OVER`` โ€” the budget is expected to be exceeded. + + A single ``OVER`` finding makes the whole report infeasible + (:attr:`PreflightReport.is_feasible` is ``False``) and is what + ``Pipeline.fit(preflight="strict")`` raises on. + """ + + AMPLE = "ample" + TIGHT = "tight" + OVER = "over" + + +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: + """Build stats for a hypothetical dataset, for sizing a search space without data. + + ``p95_tokens`` is derived as ``avg_tokens * 2.5`` and ``class_counts`` is + left empty. Use :func:`~autointent.advisor.dataset_stats` instead when a + real ``Dataset`` is available. + + Args: + n_samples: number of training utterances to assume. + n_classes: number of intent classes to assume. + avg_tokens: average utterance length in whitespace-separated tokens. + multilabel: whether to assume a multilabel task. + + Returns: + Stats with ``source="placeholder"``. + """ + return cls( + n_samples=n_samples, + n_classes=n_classes, + 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: + """Append a :class:`Finding` to this report. + + Args: + phase: which check produced it โ€” ``"resource"``, ``"data"`` or ``"config"``. + severity: headroom level; a single ``OVER`` makes the report infeasible. + message: one-line human-readable explanation, shown as-is in reports. + metric: short budget name the finding is about (``"vram"``, ``"ram"``, + ``"disk"``, ``"time"``), or ``None`` for findings not tied to a budget. + """ + self.findings.append(Finding(phase=phase, severity=severity, message=message, metric=metric)) + + @property + 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: + """Whether the run is expected to fit: True unless some finding is OVER.""" + return self.headroom != Severity.OVER + + def to_dict(self) -> dict[str, Any]: + """Render the report as JSON-serializable data. + + Severities become their string values, and the derived ``headroom`` and + ``is_feasible`` properties are included as keys. This is what the CLI's + ``--json`` output emits. + + Returns: + A plain dict of the report, its findings, resource estimate, + hardware and dataset summaries. + """ + d = asdict(self) + d["findings"] = [{**asdict(f), "severity": f.severity.value} for f in self.findings] + d["headroom"] = self.headroom.value + 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..55c74cbd8 --- /dev/null +++ b/src/autointent/advisor/_runner.py @@ -0,0 +1,284 @@ +"""Public entry point + config validation + data/config phases. + +This file contains the central public function ``run_preflight`` at the top. +Everything below it is supporting machinery for the three phases. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +from pydantic import ValidationError + +from autointent._optimization_config import OptimizationConfig +from autointent.advisor._estimates._resource import _resource_phase, _ResourceInputs +from autointent.advisor._estimates._search_space import _max_int, _module_cardinality, _walk_modules +from autointent.advisor._report import PreflightReport, Severity + +# Imported rather than reimplemented: the advisor must not disagree with the +# splitter about what counts as too few samples per class. `check_split_readiness` +# itself needs a Dataset, which the advisor never has (it works from DatasetStats), +# so the shared piece is the minimum. `test_split_readiness_agreement` pins them together. +from autointent.context.data_handler._readiness_util import _min_samples_per_class_for_config + +if TYPE_CHECKING: + from collections.abc import Callable + + from autointent.advisor._hardware import HardwareProfile + from autointent.advisor._report import DatasetStats + from autointent.configs import DataConfig + + +logger = logging.getLogger(__name__) + + +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( + _ResourceInputs( + embedder_config=cfg.embedder_config, + search_space=cfg.search_space, + n_trials=cfg.hpo_config.n_trials, + n_jobs=cfg.hpo_config.n_jobs, + dump_modules=cfg.logging_config.dump_modules, + refit_after=refit_after, + cross_encoder_model_name=cfg.cross_encoder_config.model_name, + transformer_model_name=cfg.transformer_config.model_name, + cache_probe=embedding_cache_probe, + ), + stats, + hardware, + report, + ) + _data_phase(cfg.search_space, stats, cfg.data_config, report) + _config_phase(cfg.search_space, cfg.hpo_config.n_jobs, cfg.hpo_config.n_trials, 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": []}) + + +# Warn about wasted HPO budget only when trials outnumber unique configs by 4x +# or more; below that the duplicate count is small enough to ignore. +_MIN_DUPLICATE_TRIAL_RATIO = 4 + + +def _config_phase( + search_space: list[dict[str, Any]], + n_jobs: int, + n_trials: int, + hardware: HardwareProfile, + report: PreflightReport, +) -> None: + """Config-phase checks: parallelism vs. hardware mismatches + no-op HPO.""" + 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.", + ) + + # No-op HPO warning: n_trials >> cardinality โ†’ mostly duplicate trials. + # At most one warning per preset to avoid flooding multi-module reports. + for _, entry in _walk_modules(search_space): + module = entry.get("module_name", "?") + if module in {"argmax", "threshold", "jinoos", "tunable", "adaptive"}: + continue # decision modules are cheap and often singleton by design + cardinality = _module_cardinality(entry) + if ( + cardinality is not None + and cardinality < n_trials + and n_trials // max(1, cardinality) >= _MIN_DUPLICATE_TRIAL_RATIO + ): + report.add( + "config", + Severity.TIGHT, + f"'{module}' entry has {cardinality} unique configurations but " + f"hpo_config.n_trials={n_trials} โ€” expect ~{n_trials - cardinality} " + f"duplicate trials unless the sampler dedupes. Reduce n_trials or " + f"widen the search space.", + ) + break + + +def _effective_train_fraction(data_config: DataConfig) -> float: + """Fraction of the train split a scoring module is actually fitted on. + + ``DatasetStats.class_counts`` is measured on the train split as the user + supplies it, but the pipeline carves that up before any module sees it: + hold-out takes ``validation_size`` away for validation, cross-validation + leaves one fold out, and ``separation_ratio`` splits the remaining pool + again into scoring and decision halves. Counting against the raw split is + therefore optimistic, which is the wrong direction for a feasibility gate. + + An approximation of :class:`~autointent.context.data_handler.DataHandler`'s + splitting, not a reimplementation of it โ€” deliberately coarse, and only + used to decide whether a class is at risk. + """ + if data_config.scheme == "cv": + n_folds = max(2, data_config.n_folds) + fraction = (n_folds - 1) / n_folds + else: + fraction = 1.0 - float(data_config.validation_size) + if data_config.separation_ratio is not None: + fraction *= 1.0 - float(data_config.separation_ratio) + return max(0.0, min(1.0, fraction)) + + +def _split_readiness_finding( + stats: DatasetStats, + data_config: DataConfig, + report: PreflightReport, +) -> None: + """Flag classes the stratified splitter will reject, by the splitter's own rule.""" + if not stats.class_counts: + return + min_required = _min_samples_per_class_for_config(config=data_config) + starved = sorted(name for name, count in stats.class_counts.items() if count < min_required) + if not starved: + return + detail = f"scheme={data_config.scheme}" + if data_config.scheme != "ho": + detail += f", n_folds={data_config.n_folds}" + if data_config.separation_ratio is not None: + detail += f", separation_ratio={data_config.separation_ratio}" + report.add( + "data", + Severity.OVER, + f"Stratified splitting will fail before any module is fitted: classes {starved[:5]} " + f"have <{min_required} samples ({detail}). Same minimum as " + f"autointent.context.data_handler.check_split_readiness.", + ) + + +def _data_phase( + search_space: list[dict[str, Any]], + stats: DatasetStats, + data_config: DataConfig, + report: PreflightReport, +) -> None: + """Data-phase checks: token truncation, rare classes, missing intent descriptions.""" + # 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.", + ) + + _split_readiness_finding(stats, data_config, report) + + # sklearn LogisticRegressionCV inner-CV failure: each class needs >= cv samples + # in the split the scorer is fitted on, which is smaller than the train split + # the counts were measured on. cv is configurable per linear entry (default 3); + # use the strictest one across the search space. Multilabel uses + # LogisticRegression (no CV), so skip there. + if not stats.multilabel and stats.class_counts: + linear_cvs = [ + _max_int(e.get("cv"), 3) for _, e in _walk_modules(search_space) if e.get("module_name") == "linear" + ] + if linear_cvs: + cv_max = max(linear_cvs) + fraction = _effective_train_fraction(data_config) + failing = sorted(name for name, count in stats.class_counts.items() if int(count * fraction) < cv_max) + if failing: + note = "" if fraction >= 1.0 else f" after the {fraction:.0%} train/validation split" + report.add( + "data", + Severity.OVER, + f"LogisticRegressionCV (cv={cv_max}) will fail: classes {failing[:5]} " + f"have <{cv_max} samples{note}.", + ) + + # 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..5ed6c53ec --- /dev/null +++ b/src/autointent/advisor/_workflows.py @@ -0,0 +1,485 @@ +"""High-level advisor workflows: ``estimate``, ``recommend``, and ``reduce_to_fit``. + +Each workflow orchestrates the lower-level pieces (``load_config``, +``detect_hardware``, ``stats_from_dataset``, ``run_preflight``) into a single +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 + +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 +# Descending resource cost. Declared explicitly rather than derived from +# get_args(SearchSpacePreset) so that reordering a public type alias cannot +# silently change what `recommend` picks. Kept in sync by +# tests/advisor/test_preset_cost_order.py. +PRESET_COST_ORDER: tuple[SearchSpacePreset, ...] = ( + "transformers-heavy", + "transformers-light", + "nn-heavy", + "zero-shot-llm", + "nn-medium", + "classic-heavy", + "transformers-no-hpo", + "classic-medium", + "zero-shot-encoders", + "classic-light", +) + + +def load_config(target: str) -> tuple[dict[str, Any], str]: + """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 dataset_stats(dataset: Dataset) -> DatasetStats: + """Summarize an in-memory :class:`~autointent.Dataset` for the advisor. + + Reads the train split (``train``, or ``train_0`` once the dataset has been + split) to count samples and measure utterance length โ€” average and 95th + percentile word counts, over at most the first 1000 rows โ€” and takes + ``n_classes``, ``multilabel`` and ``has_descriptions`` from the dataset + itself. Returns a placeholder when no train split is present. + + This is how a caller gets from a ``Dataset`` to the ``DatasetStats`` that + :func:`run_preflight` and :func:`reduce_to_fit` require. Use + :meth:`DatasetStats.placeholder` instead when no dataset exists yet and you + only want to size a search space against hypothetical numbers. + + Args: + dataset: the dataset the pipeline would be fitted on. + + Returns: + Stats describing that dataset, with ``source="dataset:in-memory"``. + """ + from autointent.custom_types import Split + + 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 estimate( + target: str, + *, + stats: DatasetStats | None = None, + budget_vram_gb: float | None = None, +) -> PreflightReport: + """Estimate what a preset (or YAML config path) will cost on the local hardware. + + Args: + target: Bundled preset name (e.g. ``'transformers-light'``) or a YAML + 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 ``PRESET_COST_ORDER``). + budget_vram_gb: Optional VRAM-budget override for the hardware probe. + budget_time_h: Optional wall-time ceiling in hours; presets exceeding it + get an extra ``Severity.OVER`` finding so they drop out of the + 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 list(PRESET_COST_ORDER) + + 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: dict[str, int] = {name: i for i, name in enumerate(PRESET_COST_ORDER)} + feasible = [(name, r) for name, r in results if r.is_feasible] + feasible.sort(key=lambda pair: (cost_rank.get(pair[0], len(PRESET_COST_ORDER)), 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 + + @classmethod + def nothing_droppable(cls, *, pruned_config: dict[str, Any], last_report: PreflightReport) -> ReduceToFitError: + return cls( + "No droppable scoring-node module found; remaining search space cannot be reduced further.", + pruned_config=pruned_config, + last_report=last_report, + ) + + @classmethod + def scoring_exhausted(cls, *, pruned_config: dict[str, Any], last_report: PreflightReport) -> ReduceToFitError: + return cls( + "All scoring modules were pruned to fit the budget; the resulting pipeline would have " + "nothing to run. Raise the budget or add cheaper scoring modules.", + pruned_config=pruned_config, + last_report=last_report, + ) + + @classmethod + def not_converged( + cls, max_iters: int, *, pruned_config: dict[str, Any], last_report: PreflightReport + ) -> ReduceToFitError: + return cls( + f"Search space still infeasible after {max_iters} prune iterations.", + pruned_config=pruned_config, + last_report=last_report, + ) + + +def _drop_module_from_search_space( + search_space: list[dict[str, Any]], + node_type: str, + module_name: str, +) -> list[dict[str, Any]]: + """Return a deep-copied search_space with ``module_name`` removed. + + Only the matching ``node_type`` node is touched. 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 + + +# ``Finding.metric`` uses short names ("vram"); driver rows use suffixed keys +# ("vram_gb"). Mapping the two is what makes the priority walk below work โ€” a +# previous version compared the two namespaces directly, so the lookup never +# matched and every prune silently fell back to VRAM (experiments #40 #3). +_DRIVER_KEY_BY_METRIC = {"vram": "vram_gb", "time": "time_hours", "ram": "ram_gb"} +# Preference order when several budgets are over at once. +_METRIC_PRIORITY = ("vram", "time", "ram") + + +def _pick_module_to_drop(report: PreflightReport) -> tuple[str, str] | None: + """Pick the (node_type, module_name) contributing most to whichever budget is over. + + Drops the driver with the largest cost along the first dimension that has an + OVER finding, preferring VRAM > time > RAM. Disk is deliberately absent from + that walk: driver rows carry no per-module disk figure, so disk pressure + reduces by the VRAM proxy (download size tracks model size). Falls back to + VRAM when nothing is OVER. + + Returns ``None`` when no droppable driver exists โ€” all remaining rows are + decision-node entries or unknown-cost placeholders. + """ + over_metrics = {f.metric for f in report.findings if f.severity == Severity.OVER} + driver_key = next( + (_DRIVER_KEY_BY_METRIC[m] for m in _METRIC_PRIORITY if m in over_metrics), + "vram_gb", + ) + + drivers = report.resource.drivers or [] + # Only drop scoring-node drivers โ€” decision modules are lightweight and + # 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) and removes it + from the search_space, then re-runs preflight. Disk is deliberately not + in that order: driver rows carry no per-module disk figure, so disk + pressure reduces by the VRAM proxy โ€” download size tracks model size. + * Repeats until feasible, ``max_iters`` reached, or no droppable module + remains โ€” in the last two cases raises :class:`ReduceToFitError` + carrying the pruned config and final report. + + 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.nothing_droppable(pruned_config=current, last_report=report) + node_type, module_name = pick + current["search_space"] = _drop_module_from_search_space( + current["search_space"], + node_type, + module_name, + ) + 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.scoring_exhausted(pruned_config=current, last_report=report) + report = run_preflight(current, stats, hardware, refit_after=refit_after) + if report.is_feasible: + return current, report + + raise ReduceToFitError.not_converged(max_iters, pruned_config=current, last_report=report) + + +def _has_scoring_module(config: dict[str, Any]) -> bool: + """Return True when ``config`` has at least one scoring-node entry left. + + Empty scoring is a common outcome of pruning to the bone โ€” ``reduce_to_fit`` + treats it as unfittable rather than "feasible with nothing to do". + """ + for node in config.get("search_space", []): + if node.get("node_type") == "scoring" and node.get("search_space"): + return True + return False diff --git a/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index cbfa82576..f322001d1 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -128,7 +128,15 @@ class Split: "zero-shot-llm", "zero-shot-encoders", ] -"""Some presets that our library supports.""" +"""Bundled search-space presets that our library supports. + +The order here carries no meaning. Resource-cost ranking lives in +``autointent.advisor._workflows.PRESET_COST_ORDER``, which ``recommend`` uses to +pick the heaviest preset that still fits the hardware budget. That is a cost +ranking, **not** a quality ranking: a heavier preset is not strictly better โ€” +``transformers-heavy`` will overfit on tiny datasets where a ``classic-*`` +preset wins on accuracy. +""" class Document(BaseModel): diff --git a/tests/advisor/__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..29f9670f2 --- /dev/null +++ b/tests/advisor/test_estimates_and_cli.py @@ -0,0 +1,274 @@ +"""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 +from pathlib import Path + +import pytest + +from autointent.advisor import DatasetStats, HardwareProfile, run_preflight +from autointent.advisor._cli import build_parser, main +from autointent.advisor._workflows import PRESET_COST_ORDER +from autointent.utils import load_preset + +if sys.version_info >= (3, 11): + import tomllib +else: # pytest depends on tomli below 3.11, so this import is always satisfiable here + import tomli as tomllib + +_PYPROJECT = Path(__file__).resolve().parents[2] / "pyproject.toml" + + +@pytest.fixture(autouse=True) +def _force_offline(monkeypatch: pytest.MonkeyPatch) -> None: + """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", PRESET_COST_ORDER) +def test_every_preset_inspects_without_raising(preset: str) -> None: + cfg = load_preset(preset) # type: ignore[arg-type] + stats = DatasetStats.placeholder(n_samples=500, n_classes=10, avg_tokens=24) + 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 + + +def test_console_script_name_matches_cli_prog() -> None: + """The installed command and the name the CLI prints must be the same string. + + They were not: pyproject registered ``advisor`` while the parser called + itself ``autointent-advisor``, so every usage/error message named a command + that did not exist. Both sides are now asserted against each other. + """ + with _PYPROJECT.open("rb") as f: + scripts = tomllib.load(f)["project"]["scripts"] + + advisor_scripts = {name: target for name, target in scripts.items() if target.startswith("autointent.advisor.")} + assert len(advisor_scripts) == 1, f"expected exactly one advisor console script, got {advisor_scripts}" + + script_name, target = next(iter(advisor_scripts.items())) + assert target == "autointent.advisor._cli:main" + assert script_name == build_parser().prog + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/advisor/test_estimates_internals.py b/tests/advisor/test_estimates_internals.py new file mode 100644 index 000000000..266905781 --- /dev/null +++ b/tests/advisor/test_estimates_internals.py @@ -0,0 +1,1076 @@ +"""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 x hidden_size x 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 + + +class TestCnnRnnHeuristic: + """cnn/rnn get a real small-model estimate, not a not-estimated zero row.""" + + def test_cnn_row_is_nonzero(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "cnn", + "embed_dim": [128], + "num_filters": [128], + "kernel_sizes": [[3, 4, 5]], + "batch_size": [64], + "num_train_epochs": [60], + }, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 10}, + } + report = run_preflight(cfg, DatasetStats.placeholder(n_samples=5000, n_classes=20), _profile()) + cnn_row = next(d for d in report.resource.drivers if d["module"] == "cnn") + # Real numbers (not the not-estimated placeholder) + assert cnn_row["mode"] == "small-torch-train" + assert cnn_row["vram_gb"] > 0 + assert cnn_row["ram_gb"] > 0 + assert cnn_row["time_hours"] > 0 + + def test_rnn_row_uses_hidden_dim(self) -> None: + # Bigger hidden_dim โ†’ bigger VRAM. + base = {"module_name": "rnn", "embed_dim": [128], "batch_size": [64], "num_train_epochs": [30]} + + def _run(hidden: int) -> float: + cfg = { + "search_space": [ + {"node_type": "scoring", "search_space": [{**base, "hidden_dim": [hidden]}]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 5}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + return float(next(d["vram_gb"] for d in report.resource.drivers if d["module"] == "rnn")) + + assert _run(1024) > _run(128), "larger hidden_dim must produce a larger VRAM row" + + +class TestNtrialsSharedAcrossVariants: + """n_trials is a *node* budget shared across module_name candidates.""" + + def test_single_module_gets_full_n_trials(self) -> None: + # Big dataset + embedder so linear time is non-zero and comparable. + embedder_cfg = {"embedder_config": {"model_name": "intfloat/multilingual-e5-large-instruct"}} + cfg = { + **embedder_cfg, + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear"}, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 200}, + } + cfg2 = { + **embedder_cfg, + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "linear"}, + {"module_name": "knn", "k": [5]}, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 200}, + } + stats = DatasetStats.placeholder(n_samples=10000, n_classes=77, avg_tokens=24) + solo = run_preflight(cfg, stats, _profile()) + shared = run_preflight(cfg2, stats, _profile()) + + solo_lin = next(d["time_hours"] for d in solo.resource.drivers if d["module"] == "linear") + shared_lin = next(d["time_hours"] for d in shared.resource.drivers if d["module"] == "linear") + # Same module, same everything, but shared node has 2 variants โ†’ linear + # sees half the trials. + assert solo_lin > 0 + assert shared_lin > 0 + assert solo_lin > shared_lin, ( + f"linear alone should get full n_trials, shared should get half; got solo={solo_lin} shared={shared_lin}" + ) + # Concretely: solo=20 trials, shared=10 trials โ†’ 2x ratio (allow slop for rounding). + ratio = solo_lin / shared_lin + assert 1.5 < ratio < 2.5, f"expected ~2x ratio, got {ratio}" + + +class TestProcessBaselineFloor: + """Every fit reserves ~1.5 GB RAM for torch/transformers/datasets.""" + + def test_ram_estimate_never_below_baseline(self) -> None: + # Minimal preset โ€” no scoring modules that contribute RAM. + cfg = { + "search_space": [ + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 1}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + # The floor is applied as an additive term, so even an empty pipeline + # must report at least the baseline in RAM. + assert report.resource.ram_gb >= 1.0 + + def test_cuda_vram_baseline_only_when_gpu_used(self) -> None: + cfg_cpu_only = { + "search_space": [ + {"node_type": "scoring", "search_space": [{"module_name": "linear"}]}, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 1}, + } + # linear scorer runs on CPU only โ†’ no CUDA baseline should apply. + report = run_preflight(cfg_cpu_only, DatasetStats.placeholder(), _profile(accelerator="cuda")) + assert report.resource.vram_gb == 0, "CPU-only preset must not spend the CUDA VRAM baseline" + + +class TestModuleCardinality: + """1 for all-singleton, N for finite lists, None for continuous ranges.""" + + def test_all_singleton(self) -> None: + from autointent.advisor._estimates._search_space import _module_cardinality + + assert _module_cardinality({"module_name": "bert"}) == 1 + assert _module_cardinality({"module_name": "bert", "batch_size": [64], "epochs": [30]}) == 1 + + def test_multi_list_multiplies(self) -> None: + from autointent.advisor._estimates._search_space import _module_cardinality + + # 2 batch x 3 lr candidates = 6 unique configs + cardinality = _module_cardinality( + {"module_name": "bert", "batch_size": [32, 64], "learning_rate": [1e-5, 5e-5, 1e-4]} + ) + assert cardinality == 6 + + def test_range_dict_is_unbounded(self) -> None: + from autointent.advisor._estimates._search_space import _module_cardinality + + # {low, high} โ†’ continuous โ†’ None (treated as unbounded) + assert _module_cardinality({"module_name": "knn", "k": {"low": 1, "high": 20}}) is None + + def test_reserved_keys_skipped(self) -> None: + from autointent.advisor._estimates._search_space import _module_cardinality + + # module_name / target_metric are not search dimensions + assert _module_cardinality({"module_name": "bert", "target_metric": "scoring_f1", "batch_size": [32, 64]}) == 2 + + +class TestNoOpHpoFinding: + """Config-phase warns when n_trials >> unique configs.""" + + def test_finding_on_singleton_bert_with_high_n_trials(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "num_train_epochs": [30], + "batch_size": [64], + }, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 40}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + no_op = [f for f in report.findings if "unique configurations" in f.message] + assert len(no_op) == 1, f"expected exactly one no-op warning, got {[f.message for f in no_op]}" + assert no_op[0].phase == "config" + assert no_op[0].severity == Severity.TIGHT + assert "bert" in no_op[0].message + + def test_no_finding_when_search_space_has_range(self) -> None: + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "learning_rate": {"low": 1e-5, "high": 1e-4}, + }, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 40}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + no_op = [f for f in report.findings if "unique configurations" in f.message] + assert no_op == [], f"unexpected warning for ranged search space: {[f.message for f in no_op]}" + + def test_no_finding_when_n_trials_matches_cardinality(self) -> None: + # n_trials=4, cardinality=2x2=4 โ†’ not a "no-op" waste + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "batch_size": [32, 64], + "num_train_epochs": [10, 20], + }, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 4}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile()) + no_op = [f for f in report.findings if "unique configurations" in f.message] + assert no_op == [], "n_trials matching cardinality should not warn" + + +class TestModeAwareVramBaseline: + """CUDA baseline + safety margin are mode-aware โ€” training reserves more + cuDNN workspace than inference.""" + + def test_inference_only_preset_gets_smaller_vram_than_training(self) -> None: + # Both use e5-large; only the training config triggers the bigger baseline. + inference_only = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "knn", "k": [5]}, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "embedder_config": {"model_name": "intfloat/multilingual-e5-large-instruct"}, + "hpo_config": {"n_trials": 5}, + } + training = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + { + "module_name": "bert", + "classification_model_config": [{"model_name": "microsoft/deberta-v3-small"}], + "batch_size": [16], + "num_train_epochs": [1], + }, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "hpo_config": {"n_trials": 5}, + } + stats = DatasetStats.placeholder(n_samples=1000, n_classes=10, avg_tokens=24) + + infer_r = run_preflight(inference_only, stats, _profile(vram_gb=16.0)) + train_r = run_preflight(training, stats, _profile(vram_gb=16.0)) + + # The training baseline is 1.0 GB, inference baseline is 0.3 GB โ€” so + # subtracting the driver max should show at least the 0.7 GB gap. + infer_max_driver = max((d.get("vram_gb") or 0 for d in infer_r.resource.drivers), default=0) + train_max_driver = max((d.get("vram_gb") or 0 for d in train_r.resource.drivers), default=0) + infer_baseline = infer_r.resource.vram_gb - infer_max_driver + train_baseline = train_r.resource.vram_gb - train_max_driver + assert infer_baseline < train_baseline, ( + f"inference baseline should be smaller; got infer={infer_baseline:.2f} train={train_baseline:.2f}" + ) + # Should be roughly the 0.3 vs 1.0 gap (small tolerance for rounding). + assert train_baseline - infer_baseline > 0.5 + + def test_inference_only_still_has_a_cuda_baseline(self) -> None: + # Even inference-only should be > 0 on CUDA โ€” a non-zero cuDNN + driver + # context is real. Not zeroing this out would falsely tell users that + # embedder-only presets need no GPU memory. + cfg = { + "search_space": [ + { + "node_type": "scoring", + "search_space": [ + {"module_name": "knn", "k": [5]}, + ], + }, + {"node_type": "decision", "search_space": [{"module_name": "argmax"}]}, + ], + "embedder_config": {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}, + "hpo_config": {"n_trials": 5}, + } + report = run_preflight(cfg, DatasetStats.placeholder(), _profile(vram_gb=16.0)) + assert report.resource.vram_gb > 0 diff --git a/tests/advisor/test_hardware_detection.py b/tests/advisor/test_hardware_detection.py new file mode 100644 index 000000000..24fa51292 --- /dev/null +++ b/tests/advisor/test_hardware_detection.py @@ -0,0 +1,80 @@ +"""Accelerator selection in ``detect_hardware``: CUDA -> MPS -> CPU. + +Each test patches ``_detect_cuda`` / ``_detect_mps`` (and sometimes +``_detect_ram_gb``) to force one branch, then checks the resulting profile โ€” +the CPU fallback when nothing is available, the device_class thresholds, the +MPS unified-memory budget, and the manual VRAM override. + +These do *not* cover a missing ``psutil``: it is a core dependency, imported +unguarded at ``_hardware.py`` module level, and no psutil-absent fallback +exists. The RAM and disk probes are therefore always the real ones. +""" + +from __future__ import annotations + +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..de6714f4f --- /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_pick_module_to_drop.py b/tests/advisor/test_pick_module_to_drop.py new file mode 100644 index 000000000..b27012ae2 --- /dev/null +++ b/tests/advisor/test_pick_module_to_drop.py @@ -0,0 +1,73 @@ +"""Unit tests for ``_pick_module_to_drop``'s constraint selection. + +Built from hand-rolled reports rather than real preflight runs: the rule under +test is "prune the module heaviest along the dimension that is actually over +budget", and that rule should be verifiable without invoking any formula. + +Regression guard for experiments #40 finding 3 โ€” the metric-name mismatch that +made every prune a VRAM prune. +""" + +from __future__ import annotations + +from autointent.advisor._report import PreflightReport, Severity +from autointent.advisor._workflows import _pick_module_to_drop + + +def _report(*over_metrics: str) -> PreflightReport: + """Report whose scoring drivers disagree about which module is heaviest. + + ``bert`` is heaviest on VRAM, ``linear`` on RAM, ``catboost`` on time โ€” so + the module returned identifies which dimension the code actually consulted. + """ + report = PreflightReport() + for metric in ("vram", "ram", "disk", "time"): + severity = Severity.OVER if metric in over_metrics else Severity.AMPLE + report.add("resource", severity, f"{metric} finding", metric=metric) + report.resource.drivers = [ + {"node_type": "scoring", "module": "bert", "vram_gb": 20.0, "ram_gb": 3.0, "time_hours": 2.0}, + {"node_type": "scoring", "module": "linear", "vram_gb": 0.5, "ram_gb": 40.0, "time_hours": 1.0}, + {"node_type": "scoring", "module": "catboost", "vram_gb": 0.2, "ram_gb": 8.0, "time_hours": 90.0}, + # Decision modules are never droppable, however heavy they look. + {"node_type": "decision", "module": "argmax", "vram_gb": 99.0, "ram_gb": 99.0, "time_hours": 99.0}, + ] + return report + + +def test_ram_over_prunes_ram_heaviest() -> None: + assert _pick_module_to_drop(_report("ram")) == ("scoring", "linear") + + +def test_time_over_prunes_time_heaviest() -> None: + assert _pick_module_to_drop(_report("time")) == ("scoring", "catboost") + + +def test_vram_over_prunes_vram_heaviest() -> None: + assert _pick_module_to_drop(_report("vram")) == ("scoring", "bert") + + +def test_vram_wins_when_several_constraints_are_over() -> None: + """Documented preference order is VRAM > time > RAM.""" + assert _pick_module_to_drop(_report("vram", "ram", "time")) == ("scoring", "bert") + + +def test_time_beats_ram_when_both_over() -> None: + assert _pick_module_to_drop(_report("ram", "time")) == ("scoring", "catboost") + + +def test_disk_over_falls_back_to_vram_proxy() -> None: + """Drivers carry no per-module disk figure, so disk reduces by the VRAM proxy.""" + assert _pick_module_to_drop(_report("disk")) == ("scoring", "bert") + + +def test_no_over_findings_falls_back_to_vram() -> None: + assert _pick_module_to_drop(_report()) == ("scoring", "bert") + + +def test_returns_none_when_no_scoring_driver_is_droppable() -> None: + report = _report("ram") + report.resource.drivers = [ + {"node_type": "decision", "module": "argmax", "vram_gb": 1.0, "ram_gb": 1.0, "time_hours": 1.0}, + {"node_type": "scoring", "module": "?", "vram_gb": 1.0, "ram_gb": 1.0, "time_hours": 1.0}, + ] + assert _pick_module_to_drop(report) is None diff --git a/tests/advisor/test_preset_cost_order.py b/tests/advisor/test_preset_cost_order.py new file mode 100644 index 000000000..12906691e --- /dev/null +++ b/tests/advisor/test_preset_cost_order.py @@ -0,0 +1,32 @@ +"""``PRESET_COST_ORDER`` must stay in sync with the preset literal. + +``recommend`` picks the heaviest preset that fits, so a preset missing from the +ranking would silently sort last and effectively never be recommended. This test +turns that into a hard failure at the moment a preset is added. +""" + +from __future__ import annotations + +from typing import get_args + +from autointent.advisor._workflows import PRESET_COST_ORDER +from autointent.custom_types import SearchSpacePreset + + +def test_cost_order_covers_every_preset() -> None: + assert set(PRESET_COST_ORDER) == set(get_args(SearchSpacePreset)) + + +def test_cost_order_has_no_duplicates() -> None: + assert len(PRESET_COST_ORDER) == len(set(PRESET_COST_ORDER)) + + +def test_preset_literal_order_is_not_load_bearing() -> None: + """The literal's declaration order must not be the cost ranking. + + Cost ordering belongs in PRESET_COST_ORDER, where it is commented and + tested. If these two ever coincide exactly, someone has reintroduced the + coupling -- the literal is kept in dev's original (roughly alphabetical) + order precisely so it cannot be mistaken for a ranking. + """ + assert tuple(get_args(SearchSpacePreset)) != PRESET_COST_ORDER diff --git a/tests/advisor/test_public_surface.py b/tests/advisor/test_public_surface.py new file mode 100644 index 000000000..83b31151a --- /dev/null +++ b/tests/advisor/test_public_surface.py @@ -0,0 +1,49 @@ +"""Locks the public surface of ``autointent.advisor``. + +The advisor is marked experimental, but "experimental" is not a licence for the +surface to drift silently. This test is the tripwire: adding or removing a +public name is a deliberate act that updates this list. +""" + +from __future__ import annotations + +from autointent import advisor + +EXPECTED_SURFACE = { + # functions + "dataset_stats", + "detect_hardware", + "estimate", + "recommend", + "reduce_to_fit", + "run_preflight", + # types + "DatasetStats", + "Finding", + "HardwareProfile", + "PreflightError", + "PreflightReport", + "RecommendationResult", + "ReduceToFitError", + "ResourceEstimate", + "Severity", +} + + +def test_all_matches_expected_surface() -> None: + assert set(advisor.__all__) == EXPECTED_SURFACE + + +def test_every_exported_name_resolves() -> None: + missing = [name for name in advisor.__all__ if not hasattr(advisor, name)] + assert missing == [] + + +def test_no_stdlib_shadowing_names() -> None: + """``inspect`` was exported previously and shadows the stdlib module.""" + assert "inspect" not in advisor.__all__ + + +def test_package_documents_experimental_status() -> None: + assert advisor.__doc__ is not None + assert "experimental" in advisor.__doc__.lower() diff --git a/tests/advisor/test_reduce_to_fit.py b/tests/advisor/test_reduce_to_fit.py new file mode 100644 index 000000000..618f52b38 --- /dev/null +++ b/tests/advisor/test_reduce_to_fit.py @@ -0,0 +1,160 @@ +"""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 + +from typing import Any + +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, ram_gb: float = 32.0, free_disk_gb: float = 200.0) -> HardwareProfile: + return HardwareProfile( + accelerator="cuda" if vram_gb > 0 else "cpu", + device_name="test-gpu" if vram_gb > 0 else "test-cpu", + vram_gb=vram_gb, + ram_gb=ram_gb, + free_disk_gb=free_disk_gb, + cpu_count=8, + ) + + +def _cheap_config() -> dict[str, Any]: + 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[str, Any]: + """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[str, Any]: + 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..a6254af08 --- /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 guidance" 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..9885478f3 --- /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/advisor/test_split_and_cpu_awareness.py b/tests/advisor/test_split_and_cpu_awareness.py new file mode 100644 index 000000000..eec9266a7 --- /dev/null +++ b/tests/advisor/test_split_and_cpu_awareness.py @@ -0,0 +1,249 @@ +"""Tests for the review findings on PR #291: split awareness, cv, and CPU cores. + +Three separate complaints, all of which reduced to the advisor ignoring +something it was already being handed: + +* the ``cv`` a linear entry actually declares (the time estimate hardcoded 3), +* the fact that the pipeline splits the train split again before any module + sees it, so raw per-class counts are optimistic, +* ``HardwareProfile.cpu_count``, which was detected and then read by nothing. +""" + +from __future__ import annotations + +from typing import Any, ClassVar + +import pytest + +from autointent.advisor._estimates._formulas import ( + _LINEAR_PARALLEL_FRACTION, + _MAX_CPU_SPEEDUP, + _cores_per_trial, + _cpu_speedup, + _logreg_cv_multiplier, + _time_for_catboost, + _time_for_linear, +) +from autointent.advisor._hardware import HardwareProfile +from autointent.advisor._report import DatasetStats, PreflightReport, Severity +from autointent.advisor._runner import _data_phase, _effective_train_fraction +from autointent.configs import DataConfig +from autointent.context.data_handler._readiness_util import _min_samples_per_class_for_config + + +def _stats(class_counts: dict[str, int], *, multilabel: bool = False) -> DatasetStats: + return DatasetStats( + n_samples=sum(class_counts.values()), + n_classes=len(class_counts), + avg_tokens=16, + p95_tokens=32, + multilabel=multilabel, + class_counts=class_counts, + source="test", + ) + + +def _linear_space(cv: int | None = None) -> list[dict[str, Any]]: + entry: dict[str, Any] = {"module_name": "linear"} + if cv is not None: + entry["cv"] = cv + return [{"node_type": "scoring", "search_space": [entry]}] + + +def _run_data_phase(stats: DatasetStats, data_config: DataConfig, cv: int | None = None) -> PreflightReport: + report = PreflightReport() + _data_phase(_linear_space(cv), stats, data_config, report) + return report + + +def _messages(report: PreflightReport) -> str: + return " | ".join(f.message for f in report.findings) + + +class TestLogregCvMultiplier: + """The time estimate used to hardcode 31 = Cs(10) x cv(3) + 1 refit.""" + + def test_default_cv_reproduces_the_old_constant(self) -> None: + assert _logreg_cv_multiplier(3) == 31 + + @pytest.mark.parametrize(("cv", "expected"), [(2, 21), (5, 51), (10, 101)]) + def test_scales_with_configured_cv(self, cv: int, expected: int) -> None: + assert _logreg_cv_multiplier(cv) == expected + + def test_time_grows_with_cv(self) -> None: + kwargs = { + "n_trials": 5, + "n_samples": 10_000, + "embedder_dim": 768, + "max_iter": 100, + "class_multiplier": 20, + } + cheap = _time_for_linear(cv_multiplier=_logreg_cv_multiplier(3), **kwargs) + dear = _time_for_linear(cv_multiplier=_logreg_cv_multiplier(10), **kwargs) + # cv=10 costs 101/31 as many fits as cv=3; previously both priced the same. + assert dear == pytest.approx(cheap * 101 / 31) + + +class TestCpuSpeedup: + def test_single_core_is_a_no_op(self) -> None: + assert _cpu_speedup(1, 0.9) == 1.0 + + def test_speedup_is_sublinear(self) -> None: + # Amdahl with p=0.9 on 8 cores is ~4.7x, never the naive 8x. + assert 1.0 < _cpu_speedup(8, 0.9) < 8.0 + + def test_capped_however_many_cores(self) -> None: + assert _cpu_speedup(1024, 0.99) == _MAX_CPU_SPEEDUP + + def test_never_optimistic_past_the_cap(self) -> None: + assert _cpu_speedup(10_000, 1.0) <= _MAX_CPU_SPEEDUP + + @pytest.mark.parametrize(("cpu_count", "n_jobs", "expected"), [(16, 1, 16), (16, 4, 4), (16, 32, 1), (0, 1, 1)]) + def test_cores_per_trial_divides_by_concurrent_trials(self, cpu_count: int, n_jobs: int, expected: int) -> None: + assert _cores_per_trial(cpu_count, n_jobs) == expected + + +class TestCpuCountReachesTimeEstimates: + """The complaint was that core count changed nothing. It must now change something.""" + + _CATBOOST: ClassVar[dict[str, int]] = { + "n_trials": 3, + "n_samples": 10_000, + "n_features": 768, + "iterations": 1000, + "depth": 6, + "class_multiplier": 10, + } + + def test_catboost_cpu_time_falls_with_cores(self) -> None: + one = _time_for_catboost(on_gpu=False, cores=1, **self._CATBOOST) + many = _time_for_catboost(on_gpu=False, cores=16, **self._CATBOOST) + assert many < one + + def test_catboost_gpu_time_ignores_cores(self) -> None: + one = _time_for_catboost(on_gpu=True, cores=1, **self._CATBOOST) + many = _time_for_catboost(on_gpu=True, cores=64, **self._CATBOOST) + assert one == many + + def test_linear_time_falls_with_cores_but_less_than_catboost(self) -> None: + kwargs = { + "n_trials": 5, + "n_samples": 10_000, + "embedder_dim": 768, + "max_iter": 100, + "cv_multiplier": 31, + "class_multiplier": 20, + } + one = _time_for_linear(cores=1, **kwargs) + many = _time_for_linear(cores=16, **kwargs) + assert many < one + # L-BFGS only threads inside BLAS, so it must not claim CatBoost's speedup. + assert one / many == pytest.approx(_cpu_speedup(16, _LINEAR_PARALLEL_FRACTION)) + + def test_cpu_count_is_wired_through_run_preflight(self) -> None: + """End to end: two identical configs differing only in cpu_count must differ in time.""" + from autointent.advisor import run_preflight + + config = { + "search_space": [ + {"node_type": "scoring", "search_space": [{"module_name": "catboost", "iterations": 1000}]} + ] + } + stats = DatasetStats(n_samples=10_000, n_classes=20, avg_tokens=16, source="test") + + def profile(cpu_count: int) -> HardwareProfile: + return HardwareProfile( + accelerator="cpu", + device_name="test-cpu", + vram_gb=0.0, + ram_gb=32.0, + free_disk_gb=200.0, + cpu_count=cpu_count, + ) + + small = run_preflight(config, stats, profile(1)) + big = run_preflight(config, stats, profile(32)) + assert big.resource.time_hours < small.resource.time_hours + + +class TestEffectiveTrainFraction: + def test_holdout_removes_the_validation_share(self) -> None: + assert _effective_train_fraction(DataConfig(validation_size=0.2)) == pytest.approx(0.8) + + def test_cross_validation_leaves_one_fold_out(self) -> None: + assert _effective_train_fraction(DataConfig(scheme="cv", n_folds=5)) == pytest.approx(0.8) + + def test_separation_ratio_shrinks_it_further(self) -> None: + cfg = DataConfig(validation_size=0.2, separation_ratio=0.5) + assert _effective_train_fraction(cfg) == pytest.approx(0.4) + + def test_never_leaves_the_unit_interval(self) -> None: + assert 0.0 <= _effective_train_fraction(DataConfig(validation_size=1.0)) <= 1.0 + + +class TestSplitReadinessAgreement: + """The advisor must not green-light a dataset the splitter would reject.""" + + @pytest.mark.parametrize( + "data_config", + [ + DataConfig(), + DataConfig(scheme="cv", n_folds=5), + DataConfig(separation_ratio=0.5), + DataConfig(scheme="cv", n_folds=10, separation_ratio=0.3), + ], + ) + def test_advisor_flags_exactly_what_the_splitter_rejects(self, data_config: DataConfig) -> None: + """Pins the advisor to `check_split_readiness`'s minimum, so the two cannot drift apart.""" + minimum = _min_samples_per_class_for_config(config=data_config) + # One class one sample below the splitter's own threshold. + stats = _stats({"ok": 500, "starved": minimum - 1}) + report = _run_data_phase(stats, data_config) + assert "Stratified splitting will fail" in _messages(report) + assert "starved" in _messages(report) + assert any(f.severity is Severity.OVER for f in report.findings) + + @pytest.mark.parametrize( + "data_config", + [DataConfig(), DataConfig(scheme="cv", n_folds=5), DataConfig(separation_ratio=0.5)], + ) + def test_silent_when_every_class_clears_the_threshold(self, data_config: DataConfig) -> None: + minimum = _min_samples_per_class_for_config(config=data_config) + stats = _stats({"a": minimum * 100, "b": minimum * 100}) + report = _run_data_phase(stats, data_config) + assert "Stratified splitting will fail" not in _messages(report) + + +class TestLogregCheckAccountsForTheSplit: + def test_class_that_only_passes_on_the_raw_split_is_flagged(self) -> None: + """The regression: 4 samples clears cv=3 before splitting, and fails after.""" + stats = _stats({"plenty": 500, "borderline": 4}) + cfg = DataConfig(validation_size=0.2) # 4 * 0.8 = 3.2 -> 3 usable... still >= 3 + assert int(4 * _effective_train_fraction(cfg)) == 3 + + # With separation_ratio the same class drops to 4 * 0.8 * 0.5 = 1 usable sample. + split_cfg = DataConfig(validation_size=0.2, separation_ratio=0.5) + report = _run_data_phase(stats, split_cfg, cv=3) + assert "LogisticRegressionCV (cv=3) will fail" in _messages(report) + assert "borderline" in _messages(report) + + def test_message_names_the_split_when_one_applies(self) -> None: + stats = _stats({"plenty": 500, "thin": 3}) + report = _run_data_phase(stats, DataConfig(validation_size=0.2), cv=3) + assert "after the 80% train/validation split" in _messages(report) + + def test_generous_class_counts_stay_silent(self) -> None: + stats = _stats({"a": 1000, "b": 1000}) + report = _run_data_phase(stats, DataConfig(validation_size=0.2), cv=3) + assert "LogisticRegressionCV" not in _messages(report) + + def test_multilabel_skips_the_cv_check(self) -> None: + """Multilabel uses plain LogisticRegression, which has no inner CV.""" + stats = _stats({"a": 1000, "thin": 1}, multilabel=True) + report = _run_data_phase(stats, DataConfig(validation_size=0.2), cv=3) + assert "LogisticRegressionCV" not in _messages(report) + + def test_declared_cv_is_used_not_the_default(self) -> None: + stats = _stats({"a": 1000, "mid": 40}) + assert "LogisticRegressionCV" not in _messages(_run_data_phase(stats, DataConfig(), cv=3)) + assert "LogisticRegressionCV (cv=50) will fail" in _messages(_run_data_phase(stats, DataConfig(), cv=50)) diff --git a/tests/pipeline/test_preflight.py b/tests/pipeline/test_preflight.py new file mode 100644 index 000000000..c9e689963 --- /dev/null +++ b/tests/pipeline/test_preflight.py @@ -0,0 +1,182 @@ +"""Pipeline.fit preflight integration: default-off, warn, strict. + +``fit()`` is driven with ``Pipeline._fit`` stubbed out, so these tests exercise +the preflight gate (which runs before any heavy work) without training anything. +With ``clear_ram=True, dump_modules=False`` the post-``_fit`` branch returns the +context immediately, so a stubbed ``_fit`` leaves ``fit()`` fully functional. +""" + +from __future__ import annotations + +import logging +import subprocess +import sys +from typing import TYPE_CHECKING + +import pytest + +from autointent import Pipeline +from autointent.advisor import ( + HardwareProfile, + PreflightError, + PreflightReport, + dataset_stats, + detect_hardware, + run_preflight, +) +from autointent.configs import LoggingConfig + +if TYPE_CHECKING: + from typing import Any + + from autointent import Dataset + +_PIPELINE_LOGGER = "autointent._pipeline._pipeline" + + +@pytest.fixture(autouse=True) +def _stub_fit(monkeypatch: pytest.MonkeyPatch) -> None: + """Skip optimization; every test here is about the gate that runs before it.""" + monkeypatch.setattr(Pipeline, "_fit", lambda _self, _context: None) + + +def _tiny_hw() -> HardwareProfile: + """Deterministic, intentionally-infeasible hardware budget.""" + 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 _module_names(config: dict[str, Any]) -> set[str]: + """Every ``module_name`` in an advisor-shaped config's search space.""" + return { + entry["module_name"] + for node in config["search_space"] + for entry in node["search_space"] + if "module_name" in entry + } + + +def test_fit_does_not_run_preflight_by_default(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + """The default is opt-out: no preflight, no Hub round-trips, no log line.""" + p = _classic_light_pipeline() + with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): + p.fit(dataset) + assert not any("Preflight" in r.getMessage() for r in caplog.records) + + +def test_preflight_off_skips_advisor(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + """Asking for ``preflight="off"`` explicitly is honoured, not just the default. + + The previous test pins the default value; this one pins the ``"off"`` branch + of the gate itself, so a change to the default cannot mask a regression here. + """ + p = _classic_light_pipeline() + with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): + p.fit(dataset, preflight="off") + assert not any("Preflight" in r.getMessage() for r in caplog.records) + + +def test_preflight_warn_logs_verdict(dataset: Dataset, caplog: pytest.LogCaptureFixture) -> None: + p = _classic_light_pipeline() + with caplog.at_level(logging.INFO, logger=_PIPELINE_LOGGER): + p.fit(dataset, preflight="warn") + msgs = [r.getMessage() for r in caplog.records] + assert any("Preflight" in m and "verdict=" in m for m in msgs) + + +def test_preflight_strict_raises_on_infeasible(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: + """Patched at the advisor, not the pipeline: the import is lazy now.""" + monkeypatch.setattr("autointent.advisor.detect_hardware", _tiny_hw) + p = _classic_light_pipeline() + with pytest.raises(PreflightError) as exc_info: + p.fit(dataset, preflight="strict") + 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: + monkeypatch.setattr("autointent.advisor.detect_hardware", _tiny_hw) + p = _classic_light_pipeline() + with caplog.at_level(logging.ERROR, logger=_PIPELINE_LOGGER): + p.fit(dataset, preflight="warn") + assert any(r.levelno == logging.ERROR for r in caplog.records) + + +def test_preflight_prices_the_filtered_search_space(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: + """The gate runs after validate_modules, so it never charges for discarded modules. + + ``classic-light`` ships ``mlknn``, which does not support multiclass, so + ``fit()`` drops it from the search space. Pricing the unfiltered space + inflates the estimate โ€” badly so for ``dnnc``'s ~6.4 GB reranker on + multilabel data, where it can flip a strict verdict to OVER. + """ + import autointent.advisor as advisor_pkg + + captured: list[dict[str, Any]] = [] + + def _spy(config: dict[str, Any], *_args: object, **_kwargs: object) -> PreflightReport: + captured.append(config) + return PreflightReport() + + monkeypatch.setattr(advisor_pkg, "run_preflight", _spy) + + p = _classic_light_pipeline() + assert not dataset.multilabel, "fixture must be multiclass for mlknn to be filtered out" + requested = _module_names(p._build_advisor_config()) + assert "mlknn" in requested, f"classic-light should offer mlknn: {requested}" + + p.fit(dataset, preflight="warn") + + assert captured, "preflight did not run" + priced = _module_names(captured[0]) + assert "mlknn" not in priced, f"preflight priced a module fit() discards: {priced}" + assert "linear" in priced, f"preflight lost compatible modules too: {priced}" + + +def test_importing_autointent_does_not_import_the_advisor() -> None: + """The advisor pulls in huggingface_hub probes; it must stay off the import path. + + Checked in a subprocess because pytest has already imported the advisor into + this process. Asserting on ``huggingface_hub`` itself would not work -- + ``datasets`` imports it regardless -- so the subpackage's own absence is the + real invariant. + """ + code = "import autointent, sys; print('autointent.advisor' in sys.modules)" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert result.stdout.strip() == "False", "importing autointent must not import autointent.advisor" + + +def test_pipeline_advisor_config_round_trip(dataset: Dataset) -> None: + """End-to-end: Pipeline -> _build_advisor_config -> run_preflight.""" + p = _classic_light_pipeline() + config = p._build_advisor_config() + stats = dataset_stats(dataset) + hardware = detect_hardware() + + report = run_preflight(config, stats, hardware, preset_name="classic-light") + + assert report.preset_name == "classic-light" + assert report.resource.drivers, "expected at least one driver row for classic-light" + + driver_modules = {d["module"] for d in report.resource.drivers} + assert "linear" in driver_modules, f"missing linear scorer in drivers: {driver_modules}" + + metrics = {f.metric for f in report.findings if f.metric} + assert {"vram", "ram", "disk"} <= metrics, f"missing required metrics: {metrics}" + + assert report.dataset["n_samples"] == stats.n_samples + assert report.dataset["n_classes"] == stats.n_classes