"""GPU 算力分配管理路由。""" from __future__ import annotations import json 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, user_tenant_ids from app.db.platform_store import get_platform_store router = APIRouter(prefix="/compute", tags=["gpu-assignment"]) @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 = current_user.get("id") try: result = get_platform_store().assign_gpus(assignments, assigned_by=actor) except ValueError as exc: raise fail(409, str(exc)) get_platform_store().record_audit( action="gpu.assign", actor_id=actor, target_type="gpu", detail=f"count={len(assignments)}", ) return ok(result) @router.post("/gpu-assignments/request") def request_gpu_assignment( payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user), ) -> dict[str, Any]: assignments = payload.get("assignments") or [] if not assignments: raise fail(400, "assignments 不能为空") if is_admin(current_user): try: return ok(get_platform_store().assign_gpus(assignments, assigned_by=current_user.get("id"))) except ValueError as exc: raise fail(409, str(exc)) user_id = str(current_user.get("id") or "") normalized = [] for item in assignments: if not isinstance(item, dict) or not item.get("node_id") or item.get("gpu_index") is None: raise fail(400, "每项必须包含 node_id 和 gpu_index") normalized.append({**item, "user_id": user_id}) instance = get_platform_store().create_approval_instance({ "resource_type": "gpu", "resource_id": f"batch:{user_id}", "applicant_id": user_id, "action": "gpu.assign", "tenant_id": current_user.get("tenant_id") or "default", "reason": json.dumps({"assignments": normalized}, ensure_ascii=False), }) get_platform_store().record_audit( action="gpu.assign.request", actor_id=user_id, target_type="gpu", target_id=instance["id"], tenant_id=current_user.get("tenant_id") or "default", detail=f"count={len(normalized)}", ) return ok({"approval_required": True, "approval_id": instance["id"], "approval": instance}) @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 = current_user.get("id") 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"]))