src/device_utils.py - master - Chaizee/ZenithCode_Incident-LLM-analytics

src/device_utils.py
124 строки3 KB
Chaizee
fix: fix speed
12 июн 2026, 16:28
12 июн 2026, 16:28b2745cc
from __future__ import annotations

import os

from config import (
    LOCAL_LLM_N_GPU_LAYERS,
    ML_CLASSIFY_CHUNK,
    ML_CLASSIFY_CHUNK_GPU,
    ML_EMBED_BATCH_SIZE,
    ML_EMBED_BATCH_SIZE_GPU,
    ML_TORCH_THREADS,
    USE_GPU,
)


def cuda_available() -> bool:
    if not USE_GPU:
        return False
    try:
        import torch

        return bool(torch.cuda.is_available())
    except ImportError:
        return False


_cpu_threads_configured = False


def cpu_thread_count() -> int:
    return max(1, int(ML_TORCH_THREADS or os.cpu_count() or 4))


def configure_cpu_threads() -> int:
    global _cpu_threads_configured
    count = cpu_thread_count()
    if _cpu_threads_configured:
        return count
    try:
        import torch

        torch.set_num_threads(count)
        if hasattr(torch, "set_num_interop_threads"):
            try:
                torch.set_num_interop_threads(max(1, min(4, count // 2 or 1)))
            except RuntimeError:
                pass
    except (ImportError, RuntimeError):
        pass
    _cpu_threads_configured = True
    return count


def gpu_device_name() -> str:
    if not cuda_available():
        return ""
    try:
        import torch

        return str(torch.cuda.get_device_name(0))
    except Exception:
        return "CUDA"


def ml_device() -> str:
    if cuda_available():
        return "cuda"
    return "cpu"


def ml_embed_batch_size() -> int:
    if ml_device() == "cuda":
        return ML_EMBED_BATCH_SIZE_GPU
    return ML_EMBED_BATCH_SIZE


def ml_classify_chunk() -> int:
    if ml_device() == "cuda":
        return ML_CLASSIFY_CHUNK_GPU
    return ML_CLASSIFY_CHUNK


def llm_gpu_layers() -> int:
    if cuda_available():
        return LOCAL_LLM_N_GPU_LAYERS
    return 0


def llm_on_gpu() -> bool:
    return llm_gpu_layers() != 0


def gpu_status_label() -> str:
    if ml_device() == "cuda":
        name = gpu_device_name()
        return f"GPU: {name}" if name else "GPU"
    if USE_GPU:
        return "CPU (CUDA недоступна, fallback)"
    return "CPU"


def gpu_setup_hint() -> str:
    if not USE_GPU or cuda_available():
        return ""
    return (
        "В config включён USE_GPU=True, но CUDA не найдена — работаем на CPU.\n\n"
        "Для GPU установите CUDA-сборки:\n"
        "    scripts\\install_gpu.ps1  (Windows)\n"
        "    scripts/install_gpu.sh   (Linux)"
    )


def gpu_runtime_summary() -> dict[str, str | int | bool]:
    return {
        "use_gpu_config": USE_GPU,
        "cuda_available": cuda_available(),
        "device_ml": ml_device(),
        "gpu_label": gpu_status_label(),
        "llm_gpu_layers": llm_gpu_layers(),
        "llm_on_gpu": llm_on_gpu(),
        "embed_batch_size": ml_embed_batch_size(),
        "classify_chunk": ml_classify_chunk(),
        "cpu_threads": cpu_thread_count(),
    }