39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
|
"""GPU selection routes."""
|
||
|
|
|
||
|
|
from fastapi import APIRouter, Request
|
||
|
|
from utils.logger import get_logger
|
||
|
|
|
||
|
|
logger = get_logger('api')
|
||
|
|
|
||
|
|
router = APIRouter()
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/gpu-status")
|
||
|
|
def get_gpu_status():
|
||
|
|
"""Get current GPU selection"""
|
||
|
|
from config_manager import config_manager
|
||
|
|
return {"gpu": config_manager.get_gpu()}
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/gpu-select")
|
||
|
|
async def select_gpu(request: Request):
|
||
|
|
"""Select which GPU to use for inference"""
|
||
|
|
data = await request.json()
|
||
|
|
gpu = data.get("gpu", "nvidia").lower()
|
||
|
|
|
||
|
|
if gpu not in ["nvidia", "amd"]:
|
||
|
|
return {"status": "error", "message": "Invalid GPU selection. Must be 'nvidia' or 'amd'"}
|
||
|
|
|
||
|
|
try:
|
||
|
|
from config_manager import config_manager
|
||
|
|
success = config_manager.set_gpu(gpu)
|
||
|
|
|
||
|
|
if success:
|
||
|
|
logger.info(f"GPU Selection: Switched to {gpu.upper()} GPU")
|
||
|
|
return {"status": "ok", "message": f"Switched to {gpu.upper()} GPU", "gpu": gpu}
|
||
|
|
else:
|
||
|
|
return {"status": "error", "message": "Failed to save GPU state"}
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"GPU Selection Error: {e}")
|
||
|
|
return {"status": "error", "message": str(e)}
|