78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
|
"""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"]))
|