assets/sdk/src/controlmt/device.py

"""Device + dtype + quantization auto-detection.

Resolves the user's (device, dtype, quant) preferences into concrete torch
objects. Always returns a usable config — falls back to CPU if GPU is
requested but unavailable, with a warning.
"""

from __future__ import annotations
import warnings
from dataclasses import dataclass


@dataclass
class ResolvedConfig:
    device: str         # "cuda" or "cpu"
    dtype_str: str      # "float32" | "bfloat16" | "float16"
    quant: str          # "none" | "int8-dynamic"
    bf16_cpu: bool      # convenience flag — bf16 works on CPU iff torch supports it

    def describe(self) -> str:
        bits = [self.device, self.dtype_str]
        if self.quant != "none":
            bits.append(self.quant)
        return " · ".join(bits)


def resolve(
    device: str = "auto",       # "auto" | "gpu" | "cuda" | "cpu"
    dtype: str | None = None,   # "float32" | "bfloat16" | "float16" | None
    quant: str = "none",        # "none" | "int8" / "int8-dynamic"
) -> ResolvedConfig:
    """Resolve the user's preferences into a concrete config.

    Logic:
      - device="auto" (default): GPU if available, else CPU with a quiet info note
      - device="gpu" or "cuda": GPU required; fall back to CPU with a WARNING if absent
      - device="cpu": forces CPU
      - dtype=None: pick the best dtype for the resolved device
            GPU: float16 (broadest compatibility — works on Volta+, Pascal too)
            CPU: bfloat16 if torch supports it (~2.8× faster than fp32 in our tests),
                 else float32
      - quant: "int8" / "int8-dynamic" only valid on CPU. GPU + int8 = silently
            falls back to fp16 with a warning (custom-arch incompat with bitsandbytes —
            see DEPLOYMENT.md Section 9).
    """
    import torch

    # ── device ────────────────────────────────────────────────────
    device = device.lower()
    if device in ("gpu", "cuda"):
        if torch.cuda.is_available():
            resolved_device = "cuda"
        else:
            warnings.warn("device='gpu' requested but CUDA unavailable; falling back to CPU.",
                          RuntimeWarning, stacklevel=2)
            resolved_device = "cpu"
    elif device == "cpu":
        resolved_device = "cpu"
    elif device == "auto":
        resolved_device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        raise ValueError(f"unknown device {device!r} — use 'auto', 'gpu', 'cuda', or 'cpu'")

    # ── quant ─────────────────────────────────────────────────────
    quant = (quant or "none").lower().replace("-dynamic", "")
    if quant in ("int8", "int8_dynamic"):
        quant = "int8-dynamic"
    if quant not in ("none", "int8-dynamic"):
        raise ValueError(f"unknown quant {quant!r} — use 'none' or 'int8' (CPU-only)")

    if quant == "int8-dynamic" and resolved_device == "cuda":
        warnings.warn(
            "int8 dynamic quantization is CPU-only on this model "
            "(see DEPLOYMENT.md Section 9). Falling back to fp16 on GPU.",
            RuntimeWarning, stacklevel=2)
        quant = "none"

    # ── dtype ─────────────────────────────────────────────────────
    bf16_cpu_ok = _cpu_supports_bf16(torch)
    if dtype is None:
        if quant == "int8-dynamic":
            resolved_dtype = "float32"       # quantize_dynamic operates on fp32 weights
        elif resolved_device == "cuda":
            resolved_dtype = "float16"       # widest GPU compatibility
        else:
            resolved_dtype = "bfloat16" if bf16_cpu_ok else "float32"
    else:
        dtype = dtype.lower().replace("fp32", "float32").replace("fp16", "float16").replace("bf16", "bfloat16")
        if dtype not in ("float32", "float16", "bfloat16"):
            raise ValueError(f"unknown dtype {dtype!r} — use float32/float16/bfloat16")
        resolved_dtype = dtype

    return ResolvedConfig(
        device=resolved_device,
        dtype_str=resolved_dtype,
        quant=quant,
        bf16_cpu=bf16_cpu_ok,
    )


def _cpu_supports_bf16(torch_mod) -> bool:
    """Quick runtime probe — try a 1-element bf16 add. Some old CPUs and some
    libtorch builds segfault on bf16; this catches that path before model load."""
    try:
        x = torch_mod.zeros(1, dtype=torch_mod.bfloat16)
        _ = x + x
        return True
    except Exception:
        return False


def detect_libraries() -> dict:
    """Probe what's installed. Used for the SDK's status banner + auto-picking
    code paths (e.g. if onnxruntime is present, future ONNX backend can be used)."""
    out = {}
    for name in ("torch", "transformers", "sentencepiece", "safetensors",
                  "huggingface_hub", "accelerate", "bitsandbytes", "onnxruntime"):
        try:
            mod = __import__(name)
            out[name] = getattr(mod, "__version__", "?")
        except ImportError:
            out[name] = None
    try:
        import torch
        out["_cuda"] = torch.cuda.is_available()
        out["_cuda_device"] = torch.cuda.get_device_name(0) if out["_cuda"] else None
    except Exception:
        out["_cuda"] = False
        out["_cuda_device"] = None
    return out