YG_FT/backend/app/modules/gpu/router.py

"""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"]))