src/gpu_setup.py

"""Cài torch bản CUDA phù hợp để bật GPU cho CTranslate2.

Engine mặc định (CT2) cần thư viện cuBLAS/cuDNN để chạy GPU; cách đơn giản nhất
là cài bản torch CUDA (đã đóng gói sẵn các DLL đó). Module này:
  - chọn channel cu1xx cao nhất mà driver hỗ trợ,
  - chạy `pip install torch --index-url ...` vào CHÍNH python đang chạy (sys.executable),
  - stream log để UI hiển thị tiến trình.

Sau khi cài xong PHẢI khởi động lại app: torch phải có mặt TRƯỚC khi import
ctranslate2 (xem hardware._guard_ct2_cuda_before_import) thì GPU mới được bật.
"""

from __future__ import annotations

import subprocess
import os
import sys
from typing import Callable, Iterator

# Các channel CUDA mà PyTorch stable phát hành (cao → thấp). Xác minh tại
# https://pytorch.org/get-started/locally/ (hiện: cu118, cu126, cu128).
# (major, minor, "cuXXX")
_TORCH_CUDA_CHANNELS = [
    (12, 8, "cu128"),
    (12, 6, "cu126"),
    (11, 8, "cu118"),
]


def choose_cuda_channel(driver_cuda: str | None) -> str | None:
    """Chọn channel torch cao nhất mà driver còn hỗ trợ (driver_cuda dạng '13.2').

    NVIDIA backward-compatible: driver hỗ trợ tới CUDA X chạy được mọi wheel <= X.
    Trả None nếu driver quá cũ hơn cả bản thấp nhất (cu118).
    """
    if not driver_cuda:
        # Không biết driver → chọn bản phổ biến tương thích rộng nhất.
        return "cu118"
    try:
        major, minor = (int(part) for part in driver_cuda.split(".")[:2])
    except (ValueError, TypeError):
        return "cu118"
    for ch_major, ch_minor, channel in _TORCH_CUDA_CHANNELS:
        if (major, minor) >= (ch_major, ch_minor):
            return channel
    return None


def torch_install_command(channel: str) -> list[str]:
    # --upgrade --force-reinstall: BẮT BUỘC. Nếu user đã có torch-CPU (vd từ setup
    # cũ / requirements-pytorch), `pip install torch` thường báo "already satisfied"
    # và KHÔNG ghi đè → cài xong vẫn là CPU. Force-reinstall đảm bảo lấy bản CUDA.
    return [
        sys.executable,
        "-m",
        "pip",
        "install",
        "--upgrade",
        "--force-reinstall",
        "torch",
        "--index-url",
        f"https://download.pytorch.org/whl/{channel}",
    ]


def verify_torch_cuda() -> tuple[bool, str]:
    """Sau khi cài, kiểm tra torch có thật sự thấy CUDA không (subprocess sạch).

    Chạy trong tiến trình con KHÔNG bị mask CUDA_VISIBLE_DEVICES=-1 (guard của app
    có thể đã set ở tiến trình hiện tại). Bắt trường hợp 'cài xong nhưng vẫn CPU'.
    """
    code = (
        "import torch,sys;"
        "print('TORCH_VERSION='+torch.__version__);"
        "print('CUDA_OK='+str(torch.cuda.is_available()))"
    )
    env = dict(os.environ)
    env.pop("CUDA_VISIBLE_DEVICES", None)  # bỏ mask để torch nhìn thấy GPU thật
    try:
        result = subprocess.run(
            [sys.executable, "-c", code],
            capture_output=True,
            text=True,
            timeout=120,
            env=env,
        )
    except Exception as exc:
        return False, f"Không kiểm tra được torch sau cài: {exc}"
    out = result.stdout
    version = ""
    for line in out.splitlines():
        if line.startswith("TORCH_VERSION="):
            version = line.split("=", 1)[1]
    cuda_ok = "CUDA_OK=True" in out
    if cuda_ok:
        return True, f"torch {version} đã nhận GPU."
    return False, (
        f"Đã cài torch {version or '(?)'} nhưng torch.cuda vẫn = False — "
        "có thể driver chưa phù hợp hoặc bản torch không khớp. Xem README."
    )


def verify_ct2_cuda() -> tuple[bool, str]:
    """Kiểm tra tiến trình mới có thể import torch CUDA rồi import CT2 CUDA.

    Đây là smoke test đúng đường app sau restart: torch phải được import trước để
    đăng ký thư mục DLL CUDA, rồi CTranslate2 mới dò và nạp CUDA/cuBLAS.
    """
    code = (
        "import os,sys;"
        "os.environ.pop('CUDA_VISIBLE_DEVICES', None);"
        "import torch;"
        "print('TORCH_VERSION='+torch.__version__);"
        "print('TORCH_CUDA_OK='+str(torch.cuda.is_available()));"
        "import ctranslate2;"
        "count=ctranslate2.get_cuda_device_count();"
        "print('CT2_VERSION='+ctranslate2.__version__);"
        "print('CT2_CUDA_COUNT='+str(count));"
        "types=ctranslate2.get_supported_compute_types('cuda') if count else set();"
        "print('CT2_CUDA_TYPES='+','.join(sorted(types)));"
        "sys.exit(0 if count > 0 else 2)"
    )
    env = dict(os.environ)
    env.pop("CUDA_VISIBLE_DEVICES", None)
    try:
        result = subprocess.run(
            [sys.executable, "-c", code],
            capture_output=True,
            text=True,
            timeout=120,
            env=env,
        )
    except Exception as exc:
        return False, f"Không kiểm tra được CTranslate2 CUDA: {exc}"

    stdout = result.stdout.strip()
    stderr = result.stderr.strip()
    count = ""
    version = ""
    types = ""
    for line in stdout.splitlines():
        if line.startswith("CT2_CUDA_COUNT="):
            count = line.split("=", 1)[1]
        elif line.startswith("CT2_VERSION="):
            version = line.split("=", 1)[1]
        elif line.startswith("CT2_CUDA_TYPES="):
            types = line.split("=", 1)[1]

    try:
        cuda_count = int(count)
    except ValueError:
        cuda_count = 0

    if result.returncode == 0 and cuda_count > 0:
        type_text = f", compute={types}" if types else ""
        return True, f"CTranslate2 {version or '(?)'} đã thấy {count} GPU CUDA{type_text}."

    detail = stdout or stderr or f"exit={result.returncode}"
    return False, (
        "torch đã nhận GPU nhưng CTranslate2 chưa qua smoke test CUDA. "
        f"Chi tiết: {detail}"
    )


def _stream_pip(cmd: list[str]) -> Iterator[str]:
    """Chạy pip, yield từng dòng output (cả stdout/stderr gộp)."""
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )
    assert proc.stdout is not None
    for line in proc.stdout:
        yield line.rstrip()
    proc.wait()
    yield f"__EXIT__:{proc.returncode}"


def install_torch_cuda(
    driver_cuda: str | None,
    on_log: Callable[[str], None] | None = None,
) -> tuple[bool, str]:
    """Cài torch CUDA vào env hiện tại. Trả (thành công, thông điệp cuối).

    on_log nhận từng dòng log (để UI cập nhật). Đây là hàm blocking — gọi trong
    thread/generator của Gradio, đừng gọi thẳng trên event loop chính.
    """
    channel = choose_cuda_channel(driver_cuda)
    if channel is None:
        return False, (
            "Driver NVIDIA quá cũ so với các bản torch CUDA hiện có. "
            "Hãy cập nhật driver rồi thử lại, hoặc cài torch thủ công."
        )

    cmd = torch_install_command(channel)
    if on_log:
        on_log(f"Cài torch CUDA ({channel}) — tải ~2–3 GB, cần ~5 GB ổ trống, vui lòng đợi…")
        on_log(" ".join(cmd))

    exit_code: int | None = None
    for line in _stream_pip(cmd):
        if line.startswith("__EXIT__:"):
            exit_code = int(line.split(":", 1)[1])
            continue
        if on_log and line:
            on_log(line)

    if exit_code != 0:
        return False, (
            f"Cài torch thất bại (mã lỗi {exit_code}). "
            "Kiểm tra mạng/dung lượng đĩa, hoặc cài thủ công theo README."
        )

    # Cài xong chưa đủ — xác minh torch và CT2 THẬT SỰ thấy CUDA (bắt 'already
    # satisfied', bản không khớp driver, hoặc CT2 không nạp được CUDA/cuBLAS).
    ok, verify_msg = verify_torch_cuda()
    if not ok:
        return False, verify_msg
    ct2_ok, ct2_msg = verify_ct2_cuda()
    if not ct2_ok:
        return False, f"{verify_msg} {ct2_msg}"
    return True, (
        f"Đã cài torch CUDA ({channel}) — {verify_msg} {ct2_msg} "
        "Hãy TẮT và MỞ LẠI app (stop rồi start) để bật GPU."
    )