Feat/feasibility check

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
     "aiometer (>=1.0.0,<2.0.0)",
     "aiofiles (>=24.1.0,<25.0.0)",
     "threadpoolctl (>=3.0.0,<4.0.0)",
+    "psutil (>=5.9.0,<8.0.0)",
 ]
 
 [project.optional-dependencies]
@@ -119,6 +120,7 @@ typing = [
     "joblib-stubs (>=1.4.2.5.20240918,<2.0.0)",
     "pandas-stubs (>= 2.2.3.250527, <3.0.0)",
     "types-aiofiles (>=24.1.0.20250606)",
+    "types-psutil>=7.2.2.20260518",
 ]
 docs = [
     "sphinx (>=8.1.3,<9.0.0)",
@@ -143,6 +145,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"]

src/autointent/_advisor/__init__.py

@@ -0,0 +1,29 @@
+"""Pre-flight compute feasibility advisor.
+
+Exposes a small surface used by both ``Pipeline.fit()`` (future integration) and
+the ``autointent-advisor`` CLI script. See ``compute-feasibility-advisor-proposal.md``
+at the repo root for the design document.
+"""
+
+from __future__ import annotations
+
+from ._estimates import run_preflight
+from ._hardware import HardwareProfile, detect_hardware
+from ._report import DatasetStats, Finding, PreflightReport, RecommendationResult, ResourceEstimate, Severity
+from ._workflows import inspect, load_config, recommend, stats_from_dataset
+
+__all__ = [
+    "DatasetStats",
+    "Finding",
+    "HardwareProfile",
+    "PreflightReport",
+    "RecommendationResult",
+    "ResourceEstimate",
+    "Severity",
+    "detect_hardware",
+    "inspect",
+    "load_config",
+    "recommend",
+    "run_preflight",
+    "stats_from_dataset",
+]

src/autointent/_advisor/_cli.py

@@ -0,0 +1,135 @@
+"""Console-script entry point for the pre-flight advisor.
+
+Two subcommands:
+
+* ``inspect`` — show what a given preset / config will cost on this machine.
+* ``recommend`` — pick the best-fitting bundled preset for this machine.
+
+Both subcommands accept either a real ``--dataset`` (Hub id or local
+csv/json/jsonl/parquet path loaded via ``datasets.load_dataset``) or
+``--n-samples / --n-classes / --avg-tokens`` placeholders so the script is
+useful before the user has built a dataset.
+
+The CLI is a thin wrapper around :func:`autointent._advisor.inspect` and
+:func:`autointent._advisor.recommend`; callers that don't need argparse can
+import those helpers directly.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+
+from ._render import render_json, render_recommendation, render_text
+from ._report import DatasetStats
+from ._workflows import BUNDLED_PRESETS, inspect, recommend, stats_from_dataset
+
+__all__ = ["BUNDLED_PRESETS", "build_parser", "cmd_inspect", "cmd_recommend", "main"]
+
+logger = logging.getLogger("autointent.advisor")
+
+
+def _stats_from_args(args: argparse.Namespace) -> DatasetStats:
+    multilabel = args.task == "multilabel"
+    if args.dataset:
+        return stats_from_dataset(args.dataset, multilabel=multilabel)
+    return DatasetStats.placeholder(
+        n_samples=args.n_samples,
+        n_classes=args.n_classes,
+        avg_tokens=args.avg_tokens,
+        multilabel=multilabel,
+    )
+
+
+def _add_common_dataset_args(p: argparse.ArgumentParser) -> None:
+    p.add_argument("--dataset", help="Path or hub id of a dataset; overrides placeholders.")
+    p.add_argument("--n-samples", type=int, default=1_000, help="Placeholder training set size.")
+    p.add_argument("--n-classes", type=int, default=10, help="Placeholder class count.")
+    p.add_argument("--avg-tokens", type=int, default=32, help="Placeholder average token length.")
+    p.add_argument(
+        "--task",
+        choices=("multiclass", "multilabel"),
+        default="multiclass",
+        help="Placeholder task type when --dataset isn't given.",
+    )
+
+
+def cmd_inspect(args: argparse.Namespace) -> int:
+    report = inspect(
+        args.target,
+        stats=_stats_from_args(args),
+        budget_vram_gb=args.budget_vram_gb,
+    )
+    if args.json:
+        sys.stdout.write(render_json(report))
+    else:
+        sys.stdout.write(render_text(report))
+    sys.stdout.write("\n")
+    return 0 if report.is_feasible else 1
+
+
+def cmd_recommend(args: argparse.Namespace) -> int:
+    result = recommend(
+        stats=_stats_from_args(args),
+        budget_vram_gb=args.budget_vram_gb,
+        budget_time_h=args.budget_time_h,
+    )
+    if args.json:
+        sys.stdout.write(json.dumps(result.to_dict(), indent=2, default=str))
+        sys.stdout.write("\n")
+    else:
+        sys.stdout.write(render_recommendation(result.results, result.chosen))
+        sys.stdout.write("\n")
+        if result.chosen:
+            sys.stdout.write("\n")
+            sys.stdout.write(render_text(dict(result.results)[result.chosen]))
+            sys.stdout.write("\n")
+    return 0 if result.chosen else 1
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="autointent-advisor",
+        description="Pre-flight feasibility advisor for AutoIntent search-space optimization.",
+    )
+    parser.add_argument("-v", "--verbose", action="store_true", help="Enable debug logging.")
+
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    p_inspect = sub.add_parser(
+        "inspect",
+        help="Inspect a preset or OptimizationConfig and print a feasibility report.",
+    )
+    p_inspect.add_argument("target", help="Preset name (e.g. transformers-light) or path to a YAML config.")
+    p_inspect.add_argument("--json", action="store_true", help="Emit a structured JSON report.")
+    p_inspect.add_argument("--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget.")
+    _add_common_dataset_args(p_inspect)
+    p_inspect.set_defaults(func=cmd_inspect)
+
+    p_rec = sub.add_parser(
+        "recommend",
+        help="Detect hardware and recommend the best-fitting bundled preset.",
+    )
+    p_rec.add_argument("--json", action="store_true", help="Emit a structured JSON report.")
+    p_rec.add_argument("--budget-vram-gb", type=float, default=None, help="Override detected VRAM budget.")
+    p_rec.add_argument("--budget-time-h", type=float, default=None, help="Optional wall-time ceiling in hours.")
+    _add_common_dataset_args(p_rec)
+    p_rec.set_defaults(func=cmd_recommend)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.WARNING,
+        format="%(levelname)s %(name)s: %(message)s",
+    )
+    return int(args.func(args))
+
+
+if __name__ == "__main__":
+    main()

src/autointent/_advisor/_hardware.py

@@ -0,0 +1,142 @@
+"""Local hardware detection.
+
+Probes CPU / RAM / disk and the highest-priority accelerator available
+(CUDA -> MPS -> CPU). All probes are wrapped to fall back safely on a
+broken install (e.g. CUDA driver mismatch) rather than crash the advisor.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import platform
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+import psutil
+import torch
+
+logger = logging.getLogger(__name__)
+
+Accelerator = Literal["cuda", "mps", "cpu"]
+
+# matches macOS PYTORCH_MPS_HIGH_WATERMARK_RATIO default
+MPS_DEFAULT_BUDGET_RATIO = 0.7
+
+_HIGH_GPU_VRAM_GB = 24
+_MID_GPU_VRAM_GB = 12
+_BYTES_PER_GB = 1024**3  # binary GiB convention; matches all advisor byte->GB conversions
+
+
+@dataclass
+class HardwareProfile:
+    accelerator: Accelerator
+    device_name: str
+    vram_gb: float
+    ram_gb: float
+    free_disk_gb: float
+    cpu_count: int
+    notes: list[str] = field(default_factory=list)
+
+    @property
+    def device_class(self) -> str:
+        if self.accelerator == "cpu":
+            return "cpu"
+        if self.accelerator == "mps":
+            return "apple-silicon"
+        if self.vram_gb >= _HIGH_GPU_VRAM_GB:
+            return "high-gpu"
+        if self.vram_gb >= _MID_GPU_VRAM_GB:
+            return "mid-gpu"
+        return "low-gpu"
+
+
+def _detect_ram_gb() -> float:
+    return float(psutil.virtual_memory().total) / _BYTES_PER_GB
+
+
+def _detect_free_disk_gb(path: str | None = None) -> float:
+    cache = Path(path or os.environ.get("HF_HOME") or Path("~/.cache/huggingface").expanduser())
+    probe_path = cache if cache.exists() else Path("~").expanduser()
+    try:
+        usage = shutil.disk_usage(probe_path)
+        return usage.free / _BYTES_PER_GB
+    except OSError as e:
+        logger.debug("disk usage probe failed at %s: %s", probe_path, e)
+        return 0.0
+
+
+def _detect_cuda() -> tuple[float, str] | None:
+    if not torch.cuda.is_available():
+        return None
+    idx = 0
+    try:
+        _free, total = torch.cuda.mem_get_info(idx)
+        vram_gb = total / _BYTES_PER_GB
+    except (RuntimeError, AttributeError) as e:
+        logger.debug("torch.cuda.mem_get_info failed: %s", e)
+        return None
+    name = torch.cuda.get_device_name(idx)
+    return vram_gb, name
+
+
+def _detect_mps(ram_gb: float, budget_ratio: float = MPS_DEFAULT_BUDGET_RATIO) -> tuple[float, str] | None:
+    if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
+        return None
+    # apple silicon: unified memory; budget is fraction of total RAM
+    return ram_gb * budget_ratio, f"Apple Silicon ({platform.machine()})"
+
+
+def detect_hardware(
+    *,
+    vram_budget_gb: float | None = None,
+    mps_budget_ratio: float = MPS_DEFAULT_BUDGET_RATIO,
+) -> HardwareProfile:
+    """Detect the local hardware, with optional manual overrides.
+
+    Args:
+        vram_budget_gb: when set, overrides the detected VRAM (use for
+            shared-GPU machines where part of the device is taken).
+        mps_budget_ratio: fraction of total RAM treated as the MPS
+            "VRAM" budget on Apple Silicon.
+
+    Returns:
+        HardwareProfile reflecting current machine state.
+    """
+    notes: list[str] = []
+    ram_gb = _detect_ram_gb()
+    free_disk_gb = _detect_free_disk_gb()
+    cpu_count = os.cpu_count() or 1
+
+    cuda = _detect_cuda()
+    if cuda is not None:
+        vram_gb, device_name = cuda
+        accel: Accelerator = "cuda"
+    else:
+        mps = _detect_mps(ram_gb, mps_budget_ratio)
+        if mps is not None:
+            vram_gb, device_name = mps
+            accel = "mps"
+            notes.append(f"MPS unified memory: VRAM budget = {mps_budget_ratio:.0%} of RAM.")
+        else:
+            vram_gb = 0.0
+            device_name = platform.processor() or "cpu"
+            accel = "cpu"
+
+    if vram_budget_gb is not None:
+        if vram_gb and vram_budget_gb > vram_gb:
+            notes.append(f"Manual --budget-vram-gb={vram_budget_gb} exceeds detected {vram_gb:.1f} GB; using override.")
+        notes.append(f"Using manual VRAM budget: {vram_budget_gb} GB.")
+        vram_gb = vram_budget_gb
+
+    return HardwareProfile(
+        accelerator=accel,
+        device_name=device_name,
+        vram_gb=vram_gb,
+        ram_gb=ram_gb,
+        free_disk_gb=free_disk_gb,
+        cpu_count=cpu_count,
+        notes=notes,
+    )