src/chess_tutor/compute.py

"""
Shared compute selection for PyTorch, ONNX Runtime, and LC0 backends.

**Default (no flags):** ``auto`` — use GPU when the stack reports it is available.

**Override from the command line only** (no environment variables required)::

    python -m chess_tutor.export_training_dataset --ort-provider cuda ...
    python -m chess_tutor.heuristic_mlp_torch --device cuda --amp ...
    python -m chess_tutor.analyze_player_games_heuristics --ort-provider cuda --device cuda ...
    python -m chess_tutor.build_player_moves_db --lc0-backend cuda ...

Force CPU::

    python -m chess_tutor.export_training_dataset --ort-provider cpu ...
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from typing import List, Optional, Sequence, Tuple

DEFAULT_SPEC = "auto"

ORT_PROVIDER_CHOICES = ("auto", "cpu", "cuda", "tensorrt")
TORCH_DEVICE_CHOICES = ("auto", "cpu", "cuda", "mps")
LC0_BACKEND_CHOICES = ("auto", "cpu", "eigen", "cuda", "cuda-fp16", "opencl", "blas")


def ort_available_providers() -> Tuple[str, ...]:
    import onnxruntime as ort

    return tuple(ort.get_available_providers())


def cuda_ort_available() -> bool:
    return "CUDAExecutionProvider" in set(ort_available_providers())


def cuda_torch_available() -> bool:
    import torch

    return bool(torch.cuda.is_available())


def _warn(msg: str) -> None:
    print(msg, file=sys.stderr)


def _nvidia_wheel_lib_dir(mod_name: str) -> Optional[str]:
    """``.../site-packages/nvidia/<pkg>/lib`` from a pip NVIDIA wheel (namespace-safe)."""
    import importlib.util

    spec = importlib.util.find_spec(mod_name)
    if spec is None or not spec.submodule_search_locations:
        return None
    lib = Path(spec.submodule_search_locations[0]) / "lib"
    return str(lib) if lib.is_dir() else None


def nvidia_pip_lib_dirs() -> Tuple[str, ...]:
    """
    Directories from NVIDIA pip wheels (``nvidia-cudnn-cu12``, ``nvidia-cu13``, etc.).

    PyTorch installs these but does not require ``ldconfig`` entries; ONNX Runtime
    needs them on ``LD_LIBRARY_PATH`` before the CUDA EP loads.
    """
    dirs: List[str] = []
    for mod_name in (
        "nvidia.cu13",
        "nvidia.cuda_runtime",
        "nvidia.cublas",
        "nvidia.cudnn",
    ):
        lib_dir = _nvidia_wheel_lib_dir(mod_name)
        if lib_dir and lib_dir not in dirs:
            dirs.append(lib_dir)
    return tuple(dirs)


def prepend_ld_library_path(extra_dirs: Sequence[str]) -> None:
    """Prepend directories to ``LD_LIBRARY_PATH`` if not already present."""
    current = os.environ.get("LD_LIBRARY_PATH", "")
    current_parts = [p for p in current.split(":") if p]
    new_parts: List[str] = []
    for d in extra_dirs:
        if d and d not in current_parts and d not in new_parts:
            new_parts.append(d)
    if not new_parts:
        return
    os.environ["LD_LIBRARY_PATH"] = ":".join(new_parts + current_parts)


_ort_gpu_libs_ready = False


def _preload_shared_libs_in_dir(lib_dir: Path) -> None:
    """
    Load pip-shipped ``.so`` files into the global namespace.

    Setting ``LD_LIBRARY_PATH`` from Python is not enough on Linux for ONNX Runtime's
    late-loaded ``providers_cuda.so``; ``ctypes.CDLL(..., RTLD_GLOBAL)`` matches what
    PyTorch effectively does for ``nvidia-*`` wheels.
    """
    import ctypes

    mode = getattr(ctypes, "RTLD_GLOBAL", 0x100)
    if not lib_dir.is_dir():
        return
    for so in sorted(lib_dir.glob("*.so*")):
        if not so.is_file():
            continue
        try:
            ctypes.CDLL(str(so.resolve()), mode=mode)
        except OSError:
            continue


def ensure_ort_gpu_runtime_libs() -> None:
    """Expose pip-shipped CUDA/cuDNN libs to ONNX Runtime (no-op if not installed)."""
    global _ort_gpu_libs_ready
    if _ort_gpu_libs_ready:
        return
    dirs = nvidia_pip_lib_dirs()
    prepend_ld_library_path(dirs)
    for lib_dir in dirs:
        _preload_shared_libs_in_dir(Path(lib_dir))
    _ort_gpu_libs_ready = True


def _normalize_spec(spec: Optional[str]) -> str:
    """CLI value or ``auto`` when the flag was omitted (``None``)."""
    if spec is None or not str(spec).strip():
        return DEFAULT_SPEC
    return str(spec).strip().lower()


def resolve_ort_providers(spec: Optional[str] = None) -> List[str]:
    """
    Return an ONNX Runtime provider list (first entry is preferred).

    ``auto`` uses CUDA when ``CUDAExecutionProvider`` is registered (``onnxruntime-gpu``).
    """
    raw = _normalize_spec(spec)
    available = set(ort_available_providers())

    def _cpu_only() -> List[str]:
        return ["CPUExecutionProvider"]

    if raw == "cpu":
        return _cpu_only()

    if raw == "cuda":
        if "CUDAExecutionProvider" in available:
            return ["CUDAExecutionProvider", "CPUExecutionProvider"]
        _warn(
            "--ort-provider cuda requested but CUDAExecutionProvider is not available; "
            "install onnxruntime-gpu and CUDA. Using CPU."
        )
        return _cpu_only()

    if raw == "tensorrt":
        chain: List[str] = []
        if "TensorrtExecutionProvider" in available:
            chain.append("TensorrtExecutionProvider")
        if "CUDAExecutionProvider" in available:
            chain.append("CUDAExecutionProvider")
        chain.append("CPUExecutionProvider")
        if len(chain) == 1:
            _warn(
                "--ort-provider tensorrt requested but no TensorRT/CUDA EPs found; using CPU."
            )
        return chain

    if raw == "auto":
        if "CUDAExecutionProvider" in available:
            return ["CUDAExecutionProvider", "CPUExecutionProvider"]
        return _cpu_only()

    raise ValueError(
        f"Unknown ORT provider spec {spec!r}; use one of {ORT_PROVIDER_CHOICES}"
    )


def resolve_torch_device(spec: Optional[str] = None) -> str:
    """Return a PyTorch device string (``cpu``, ``cuda``, ``mps``)."""
    import torch

    raw = _normalize_spec(spec)

    if raw == "cpu":
        return "cpu"
    if raw == "cuda":
        if torch.cuda.is_available():
            return "cuda"
        _warn("--device cuda requested but torch.cuda.is_available() is False; using cpu.")
        return "cpu"
    if raw == "mps":
        if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
            return "mps"
        _warn("--device mps requested but MPS is not available; using cpu.")
        return "cpu"
    if raw == "auto":
        if torch.cuda.is_available():
            return "cuda"
        if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
            return "mps"
        return "cpu"

    raise ValueError(
        f"Unknown torch device spec {spec!r}; use one of {TORCH_DEVICE_CHOICES}"
    )


def torch_device_supports_amp(device: str) -> bool:
    """Whether ``torch.amp`` / legacy CUDA AMP is reasonable on this device."""
    return device == "cuda"


def resolve_lc0_bindings_backend(spec: Optional[str] = None) -> str:
    """
    Backend name for ``lczero.backends.Backend(weights, backend=...)``.

    ``auto`` uses ``eigen`` (CPU). Pass ``--lc0-bindings-backend cuda`` when bindings
    were built with CUDA support.
    """
    raw = _normalize_spec(spec)
    if raw == "auto":
        return "eigen"
    if raw in LC0_BACKEND_CHOICES or raw:
        return raw
    raise ValueError(
        f"Unknown LC0 bindings backend {spec!r}; examples: eigen, cuda (see lc0 docs)"
    )


def resolve_lc0_uci_backend(spec: Optional[str] = None) -> Optional[str]:
    """
    LC0 UCI ``--backend=`` CLI value, or ``None`` to omit (lc0 binary default).

    ``auto`` omits the flag (a CUDA-built ``lc0`` defaults to ``cuda-auto``).
    ``eigen`` / ``cpu`` force the CPU Eigen backend (recommended on WSL when CUDA crashes).
    """
    raw = _normalize_spec(spec)
    if raw == "auto":
        return None
    if raw == "cpu":
        return "eigen"
    if raw == "eigen":
        return "eigen"
    return raw


def lc0_uci_extra_args(spec: Optional[str] = None) -> List[str]:
    """Extra CLI args for the ``lc0`` binary (e.g. ``['--backend=cuda']``)."""
    backend = resolve_lc0_uci_backend(spec)
    if backend is None:
        return []
    return [f"--backend={backend}"]


def warn_if_ort_fell_back_to_cpu(
    requested_providers: Sequence[str],
    sess: object,
) -> None:
    """Warn when CUDA/TensorRT was requested but the session is CPU-only."""
    try:
        active = list(sess.get_providers())  # type: ignore[attr-defined]
    except Exception:
        return
    if not active:
        return
    wants_gpu = any(
        p in requested_providers
        for p in ("CUDAExecutionProvider", "TensorrtExecutionProvider")
    )
    if wants_gpu and active[0] == "CPUExecutionProvider":
        pip_dirs = nvidia_pip_lib_dirs()
        hint = (
            f"pip cuDNN dirs: {pip_dirs}" if pip_dirs else "pip install nvidia-cudnn-cu12"
        )
        _warn(
            "ONNX Runtime fell back to CPU despite GPU provider preference. "
            f"Ensure libcudnn.so.9 is on LD_LIBRARY_PATH ({hint}), or use --ort-provider cpu. "
            "See https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements"
        )


def log_ort_session_providers(sess: object, *, label: str = "ONNX Runtime") -> None:
    """Log active providers after creating an ``InferenceSession``."""
    try:
        active = sess.get_providers()  # type: ignore[attr-defined]
    except Exception:
        return
    _warn(f"{label}: session providers = {active}")


def log_compute_plan(
    *,
    torch_spec: Optional[str] = None,
    ort_spec: Optional[str] = None,
    lc0_bindings_spec: Optional[str] = None,
    lc0_uci_spec: Optional[str] = None,
) -> None:
    """One-line stderr summary of resolved compute choices (call after parsing args)."""
    parts: List[str] = []
    if torch_spec is not None:
        parts.append(f"torch={resolve_torch_device(torch_spec)}")
    if ort_spec is not None:
        parts.append(f"ort={resolve_ort_providers(ort_spec)}")
    if lc0_bindings_spec is not None:
        parts.append(f"lc0_bindings={resolve_lc0_bindings_backend(lc0_bindings_spec)}")
    if lc0_uci_spec is not None:
        uci = resolve_lc0_uci_backend(lc0_uci_spec)
        parts.append(f"lc0_uci={uci or '(lc0 default)'}")
    if parts:
        _warn("Compute: " + ", ".join(parts))


def add_compute_arguments(
    parser: argparse.ArgumentParser,
    *,
    include_torch: bool = False,
    include_ort: bool = False,
    include_lc0_bindings: bool = False,
    include_lc0_uci: bool = False,
) -> None:
    """Register ``--device`` / ``--ort-provider`` / LC0 backend flags (default: auto)."""
    if include_torch:
        parser.add_argument(
            "--device",
            default=DEFAULT_SPEC,
            choices=TORCH_DEVICE_CHOICES,
            help="PyTorch device (default: auto — cuda if available, else mps, else cpu).",
        )
    if include_ort:
        parser.add_argument(
            "--ort-provider",
            default=DEFAULT_SPEC,
            choices=ORT_PROVIDER_CHOICES,
            help=(
                "ONNX Runtime EP (default: auto — CUDA if onnxruntime-gpu is installed). "
                "Use cpu to force CPU."
            ),
        )
    if include_lc0_bindings:
        parser.add_argument(
            "--lc0-bindings-backend",
            default=DEFAULT_SPEC,
            metavar="NAME",
            help=(
                "lczero.backends plane encoder (default: auto→eigen). "
                "Use cuda only if bindings were built with CUDA."
            ),
        )
    if include_lc0_uci:
        parser.add_argument(
            "--lc0-backend",
            default=DEFAULT_SPEC,
            choices=LC0_BACKEND_CHOICES,
            help=(
                "lc0 UCI --backend=… (default: auto→lc0 default). "
                "Use cuda or cuda-fp16 with a GPU-built lc0 binary."
            ),
        )