"""GPU 算力分配管理路由。""" from __future__ import annotations from typing import Any from fastapi import APIRouter, Body, Depends, Request from app.api.v1.endpoints.platform import ok, fail from app.core.auth import get_current_user, is_admin from app.db.platform_store import get_platform_store router = APIRouter(prefix="/compute", tags=["gpu-assignment"]) def _actor_id(request: Request) -> str | None: auth = request.headers.get("Authorization", "") token = auth.replace("Bearer ", "").strip() if token.startswith("platform-token-"): return token[len("platform-token-"):] return None @router.get("/gpu-assignments") def list_assignments(current_user: dict = Depends(get_current_user)) -> dict[str, Any]: """查看全部分配关系(仅 admin)。""" if not is_admin(current_user): raise fail(403, "admin permission required") return ok(get_platform_store().gpu_assignments()) @router.post("/gpu-assignments") def assign_gpus( payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(get_current_user), ) -> dict[str, Any]: """批量分配 GPU(仅 admin)。body: { assignments: [{ node_id, gpu_index, user_id }] }""" if not is_admin(current_user): raise fail(403, "admin permission required") assignments = payload.get("assignments") or [] if not assignments: raise fail(400, "assignments 不能为空") actor = _actor_id(request) if request else None result = get_platform_store().assign_gpus(assignments, assigned_by=actor) get_platform_store().record_audit( action="gpu.assign", actor_id=actor, target_type="gpu", detail=f"count={len(assignments)}", ) return ok(result) @router.delete("/gpu-assignments/{assignment_id}") def unassign_gpu( assignment_id: str, request: Request = None, current_user: dict = Depends(get_current_user), ) -> dict[str, Any]: """撤销 GPU 分配(仅 admin)。""" if not is_admin(current_user): raise fail(403, "admin permission required") get_platform_store().unassign_gpu(assignment_id) actor = _actor_id(request) if request else None get_platform_store().record_audit( action="gpu.unassign", actor_id=actor, target_type="gpu", target_id=assignment_id, ) return ok({"deleted": assignment_id}) @router.get("/my-gpus") def my_gpus(current_user: dict = Depends(get_current_user)) -> dict[str, Any]: """查看当前用户可用的 GPU 列表。""" return ok(get_platform_store().gpu_assignments_for_user(current_user["id"]))