pyai/route/gpu.py

import logging
from fastapi import APIRouter, HTTPException
import service.pygpu as pygpu

logger = logging.getLogger(__name__)

# 创建路由对象，可以统一设置前缀 (prefix) 和 标签 (tags)
router = APIRouter(
    prefix="/api1/gpu",
    tags=["gpu"]
)

@router.get("/info")
async def gpu_info():
    try:
        gpu_result = pygpu.get_gpu_memory_info(0)
        torch_result = pygpu.get_torch_memory_usage()
        result = pygpu.get_ollama_resource()
        return {
            "gpu": gpu_result,
            "torch": torch_result,
            "ollama": result,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))