Allocate 回调与容器 GPU 注入机制

func (plugin *NvidiaDevicePlugin) Allocate(ctx context.Context,
    reqs *kubeletdevicepluginv1beta1.AllocateRequest) (
    *kubeletdevicepluginv1beta1.AllocateResponse, error) {

    nodename := os.Getenv(util.NodeNameEnvName)
    // 1. 获取正在等待分配的 Pod
    current, err := util.GetPendingPod(ctx, nodename)

    for idx, req := range reqs.ContainerRequests {
        // 2. 判断是 MIG 设备还是普通 GPU
        if strings.Contains(req.DevicesIDs[0], "MIG") {
            // MIG 设备分配路径
            response, _ := plugin.getAllocateResponse(req.DevicesIDs)
        } else {
            // 3. 从 Pod Annotation 读取调度决策
            currentCtr, devreq, _ := GetNextDeviceRequest(
                nvidia.NvidiaGPUDevice, *current)

            // 4. 构建基础 AllocateResponse
            response, _ := plugin.getAllocateResponse(
                plugin.GetContainerDeviceStrArray(devreq))

            // 5. 注入 HAMi 特有的环境变量
            for i, dev := range devreq {
                limitKey := fmt.Sprintf("CUDA_DEVICE_MEMORY_LIMIT_%v", i)
                response.Envs[limitKey] = fmt.Sprintf("%vm", dev.Usedmem)
            }
            response.Envs["CUDA_DEVICE_SM_LIMIT"] = fmt.Sprint(devreq[0].Usedcores)

            // 6. 注入 Volume Mounts
            response.Mounts = append(response.Mounts, ...)

            // 7. 清除已处理的 Annotation
            EraseNextDeviceTypeFromAnnotation(nvidia.NvidiaGPUDevice, *current)
        }
    }

    // 8. 标记分配成功，释放节点锁
    PodAllocationTrySuccess(nodename, nvidia.NvidiaGPUDevice,
        NodeLockNvidia, current)
    return &responses, nil
}

// pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go
func GetNextDeviceRequest(dtype string, p corev1.Pod) (
    corev1.Container, device.ContainerDevices, error) {

    // 从 Pod Annotation 解码设备请求
    pdevices, _ := device.DecodePodDevices(
        device.InRequestDevices, p.Annotations)

    // 找到第一个尚未处理的容器设备请求
    pd := pdevices[dtype]
    for ctridx, ctrDevice := range pd {
        if len(ctrDevice) > 0 {
            return p.Spec.Containers[ctridx], ctrDevice, nil
        }
    }
}

// pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go - Allocate()
response.Mounts = append(response.Mounts,
    // 1. libvgpu.so - HAMi-core 拦截库
    &kubeletdevicepluginv1beta1.Mount{
        ContainerPath: fmt.Sprintf("%s/vgpu/libvgpu.so", hostHookPath),
        HostPath:      GetLibPath(),
        ReadOnly:      true,
    },
    // 2. 容器级缓存目录
    &kubeletdevicepluginv1beta1.Mount{
        ContainerPath: fmt.Sprintf("%s/vgpu", hostHookPath),
        HostPath:      cacheFileHostDirectory,
        ReadOnly:      false,
    },
    // 3. vGPU 进程间锁
    &kubeletdevicepluginv1beta1.Mount{
        ContainerPath: "/tmp/vgpulock",
        HostPath:      "/tmp/vgpulock",
        ReadOnly:      false,
    },
)

// 4. ld.so.preload（条件挂载 - 检查 CUDA_DISABLE_CONTROL）
if !found {
    response.Mounts = append(response.Mounts,
        &kubeletdevicepluginv1beta1.Mount{
            ContainerPath: "/etc/ld.so.preload",
            HostPath:      hostHookPath + "/vgpu/ld.so.preload",
            ReadOnly:      true,
        },
    )
}

// pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go
func GetLibPath() string {
    // 优先使用带版本号的库文件
    libPath := hostHookPath + "/vgpu/libvgpu.so." + info.GetVersion()
    if _, err := os.Stat(libPath); os.IsNotExist(err) {
        // 回退到不带版本号的文件
        libPath = hostHookPath + "/vgpu/libvgpu.so"
    }
    return libPath
}

for _, val := range currentCtr.Env {
    if strings.Compare(val.Name, "CUDA_DISABLE_CONTROL") == 0 {
        t, _ := strconv.ParseBool(val.Value)
        if t {
            found = true  // 跳过 ld.so.preload 挂载
            break
        }
    }
}

// HAMi-core/src/libvgpu.c
fp_dlsym real_dlsym = NULL;

void init_dlsym() {
    pthread_mutex_init(&dlsym_lock, NULL);
    dlmap_count = 0;
}

// HAMi-core/src/cuda/hook.c
cuda_entry_t cuda_library_entry[] = {
    {.name = "cuInit"},
    {.name = "cuDeviceGetCount"},
    {.name = "cuMemAlloc_v2"},
    {.name = "cuMemFree_v2"},
    {.name = "cuMemGetInfo_v2"},
    {.name = "cuLaunchKernel"},
    // ... 数十个 CUDA API
};

func PodAllocationTrySuccess(nodeName, devName, lockName string,
    pod *corev1.Pod) {
    // 检查是否所有设备类型都已分配完成
    refreshed, _ := client.GetClient().CoreV1().Pods(
        pod.Namespace).Get(ctx, pod.Name, metav1.GetOptions{})
    annos := refreshed.Annotations[device.InRequestDevices[devName]]
    for _, val := range device.DevicesToHandle {
        if strings.Contains(annos, val) {
            return  // 还有设备未分配，等待
        }
    }
    // 所有设备分配完成
    PodAllocationSuccess(nodeName, pod, lockName)
}

func PodAllocationSuccess(nodeName string, pod *corev1.Pod,
    lockName string) {
    // 1. 更新 Pod Annotation: hami.io/bind-phase = "success"
    // 2. 释放节点互斥锁
    updatePodAnnotationsAndReleaseLock(nodeName, pod, lockName,
        util.DeviceBindSuccess)
}

func PodAllocationFailed(nodeName string, pod *corev1.Pod,
    lockName string) {
    // 1. 更新 Pod Annotation: hami.io/bind-phase = "failed"
    // 2. 释放节点互斥锁
    updatePodAnnotationsAndReleaseLock(nodeName, pod, lockName,
        util.DeviceBindFailed)
}