From 4d86341dac3febdffec08041a919bd1972af2bbc Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Wed, 5 Aug 2026 15:09:31 +0800 Subject: [PATCH 01/11] feat(nvidia): support topology-aware dynamic MIG allocation test(nvidia): add dynamic MIG workload coverage docs(nvidia): describe dynamic MIG lifecycle Signed-off-by: Jifei Wang --- .../device-plugin/daemonsetnvidia.yaml | 7 + .../templates/scheduler/device-configmap.yaml | 218 +----- docker/Dockerfile.no-core-test | 14 + docs/develop/mig-dynamic-deallocate_cn.md | 254 +++++++ go.mod | 2 +- go.sum | 47 +- hack/hami-mig-e2e.sh | 273 ++++++++ .../nvinternal/plugin/mig_startup.go | 330 ++++++++++ .../nvinternal/plugin/mig_startup_test.go | 138 ++++ .../nvidiadevice/nvinternal/plugin/migmgr.go | 623 ++++++++++++++++++ .../plugin/migmgr_placement_test.go | 59 ++ .../nvinternal/plugin/register.go | 65 +- .../nvidiadevice/nvinternal/plugin/server.go | 265 ++++++-- .../nvinternal/plugin/server_test.go | 15 +- .../nvidiadevice/nvinternal/plugin/util.go | 285 ++------ .../nvinternal/plugin/util_test.go | 157 ----- .../nvinternal/podresources/watcher.go | 264 ++++++++ .../nvinternal/podresources/watcher_test.go | 161 +++++ pkg/device/common/common.go | 1 + pkg/device/devices.go | 207 +++--- pkg/device/nvidia/device.go | 183 ++--- pkg/device/nvidia/device_test.go | 456 +------------ pkg/device/nvidia/mig_allocations.go | 81 +++ pkg/device/nvidia/mig_capability_test.go | 51 ++ pkg/device/nvidia/mig_profiles.go | 25 + pkg/device/nvidia/mig_profiles_test.go | 50 ++ pkg/device/nvidia/mig_topology.go | 116 ++++ pkg/device/nvidia/mig_topology_test.go | 52 ++ pkg/device/pod_test.go | 26 +- pkg/device/pods.go | 3 + pkg/scheduler/mig_allocation_test.go | 60 ++ pkg/scheduler/scheduler.go | 66 +- pkg/scheduler/scheduler_test.go | 235 ------- pkg/scheduler/score_test.go | 16 +- 34 files changed, 3151 insertions(+), 1654 deletions(-) create mode 100644 docker/Dockerfile.no-core-test create mode 100644 docs/develop/mig-dynamic-deallocate_cn.md create mode 100644 hack/hami-mig-e2e.sh create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go create mode 100644 pkg/device/nvidia/mig_allocations.go create mode 100644 pkg/device/nvidia/mig_capability_test.go create mode 100644 pkg/device/nvidia/mig_profiles.go create mode 100644 pkg/device/nvidia/mig_profiles_test.go create mode 100644 pkg/device/nvidia/mig_topology.go create mode 100644 pkg/device/nvidia/mig_topology_test.go create mode 100644 pkg/scheduler/mig_allocation_test.go diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index 0e785c363b..d185a0f8ad 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -134,6 +134,9 @@ spec: volumeMounts: - name: device-plugin mountPath: /var/lib/kubelet/device-plugins + - name: pod-resources + mountPath: /var/lib/kubelet/pod-resources + readOnly: true - name: lib mountPath: {{ printf "%s%s" .Values.global.gpuHookPath "/vgpu" }} - name: usrbin @@ -220,6 +223,10 @@ spec: - name: device-plugin hostPath: path: {{ .Values.devicePlugin.pluginPath }} + - name: pod-resources + hostPath: + path: /var/lib/kubelet/pod-resources + type: Directory - name: lib hostPath: path: {{ .Values.devicePlugin.libPath }} diff --git a/charts/hami/templates/scheduler/device-configmap.yaml b/charts/hami/templates/scheduler/device-configmap.yaml index d748089e8f..a6de8075e8 100644 --- a/charts/hami/templates/scheduler/device-configmap.yaml +++ b/charts/hami/templates/scheduler/device-configmap.yaml @@ -32,223 +32,25 @@ data: gpuCorePolicy: {{ .Values.devices.nvidia.gpuCorePolicy }} libCudaLogLevel: {{ .Values.devices.nvidia.libCudaLogLevel }} runtimeClassName: "{{ .Values.devicePlugin.runtimeClassName }}" - knownMigGeometries: + migProfileAllowlist: - models: [ "A30" ] - allowedGeometries: - - - - name: 1g.6gb - core: 25 - memory: 6144 - count: 4 - - - - name: 2g.12gb - core: 50 - memory: 12288 - count: 2 - - - - name: 4g.24gb - core: 100 - memory: 24576 - count: 1 + profiles: [ "1g.6gb", "2g.12gb", "4g.24gb" ] - models: [ "A100-SXM4-40GB", "A100-40GB-PCIe", "A100-PCIE-40GB"] - allowedGeometries: - - - - name: 1g.5gb - core: 14 - memory: 5120 - count: 7 - - - - name: 1g.5gb - core: 14 - memory: 5120 - count: 1 - - name: 2g.10gb - core: 28 - memory: 10240 - count: 3 - - - - name: 3g.20gb - core: 42 - memory: 20480 - count: 2 - - - - name: 7g.40gb - core: 100 - memory: 40960 - count: 1 + profiles: [ "1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb" ] - models: [ "A100-SXM4-80GB", "A100-80GB-PCIe", "A100-PCIE-80GB"] - allowedGeometries: - - - - name: 1g.10gb - core: 14 - memory: 9728 - count: 7 - - - - name: 1g.10gb - core: 14 - memory: 9728 - count: 1 - - name: 2g.20gb - core: 28 - memory: 19968 - count: 3 - - - - name: 3g.40gb - core: 42 - memory: 40192 - count: 2 - - - - name: 7g.80gb - core: 98 - memory: 80896 - count: 1 + profiles: [ "1g.10gb", "2g.20gb", "3g.40gb", "7g.79gb" ] - models: [ "H100-PCIE-80GB", "H100-SXM5-80GB"] - allowedGeometries: - - - - name: 1g.10gb - core: 14 - memory: 10240 - count: 7 - - - - name: 1g.10gb - core: 14 - memory: 10240 - count: 1 - - name: 2g.20gb - core: 28 - memory: 20480 - count: 3 - - - - name: 3g.40gb - core: 42 - memory: 40960 - count: 2 - - - - name: 7g.80gb - core: 100 - memory: 81920 - count: 1 + profiles: [ "1g.10gb", "2g.20gb", "3g.40gb", "7g.80gb" ] - models: [ "H100-PCIE-94GB", "H100-SXM5-94GB"] - allowedGeometries: - - - - name: 1g.12gb - core: 14 - memory: 12288 - count: 7 - - - - name: 1g.12gb - core: 14 - memory: 12288 - count: 1 - - name: 2g.24gb - core: 28 - memory: 24576 - count: 3 - - - - name: 3g.47gb - core: 42 - memory: 48128 - count: 2 - - - - name: 7g.94gb - core: 100 - memory: 96256 - count: 1 + profiles: [ "1g.12gb", "2g.24gb", "3g.47gb", "7g.94gb" ] - models: [ "H20", "H100 on GH200"] - allowedGeometries: - - - - name: 1g.12gb - core: 14 - memory: 12288 - count: 7 - - - - name: 1g.12gb - core: 14 - memory: 12288 - count: 1 - - name: 2g.24gb - core: 28 - memory: 24576 - count: 3 - - - - name: 3g.48gb - core: 42 - memory: 49152 - count: 2 - - - - name: 7g.96gb - core: 100 - memory: 98304 - count: 1 + profiles: [ "1g.12gb", "2g.24gb", "3g.48gb", "7g.96gb" ] - models: [ "H200 NVL", "H200-SXM5"] - allowedGeometries: - - - - name: 1g.18gb - core: 14 - memory: 18432 - count: 7 - - - - name: 1g.18gb - core: 14 - memory: 18432 - count: 1 - - name: 2g.35gb - core: 28 - memory: 35840 - count: 3 - - - - name: 3g.71gb - core: 42 - memory: 72704 - count: 2 - - - - name: 7g.141gb - core: 100 - memory: 144384 - count: 1 + profiles: [ "1g.18gb", "2g.35gb", "3g.71gb", "7g.141gb" ] - models: [ "B200" ] - allowedGeometries: - - - - name: 1g.23gb - core: 14 - memory: 23552 - count: 7 - - - - name: 1g.23gb - core: 14 - memory: 23552 - count: 1 - - name: 2g.45gb - core: 28 - memory: 46080 - count: 3 - - - - name: 3g.90gb - core: 42 - memory: 92160 - count: 2 - - - - name: 7g.180gb - core: 100 - memory: 184320 - count: 1 + profiles: [ "1g.23gb", "2g.45gb", "3g.90gb", "7g.180gb" ] - models: [ "RTX PRO 6000 Blackwell Server Edition" ] - allowedGeometries: - - - - name: 1g.24gb - core: 25 - memory: 24576 - count: 4 - - - - name: 2g.48gb - core: 50 - memory: 49152 - count: 2 - - - - name: 4g.96gb - core: 100 - memory: 98304 - count: 1 + profiles: [ "1g.24gb", "2g.48gb", "4g.96gb" ] cambricon: resourceCountName: {{ .Values.mluResourceName }} resourceMemoryName: {{ .Values.mluResourceMem }} diff --git a/docker/Dockerfile.no-core-test b/docker/Dockerfile.no-core-test new file mode 100644 index 0000000000..4839ff6adb --- /dev/null +++ b/docker/Dockerfile.no-core-test @@ -0,0 +1,14 @@ +ARG GOLANG_IMAGE=golang:1.26.5-bookworm +ARG BASE_IMAGE + +FROM --platform=$TARGETPLATFORM $GOLANG_IMAGE AS build +ARG TARGETOS +ARG TARGETARCH +ARG VERSION +ARG GOPROXY=https://goproxy.cn,direct +WORKDIR /src +COPY . . +RUN GOOS=$TARGETOS GOARCH=$TARGETARCH CGO_ENABLED=1 GOFLAGS=-mod=mod GOPROXY=$GOPROXY \ + make build VERSION=$VERSION +FROM $BASE_IMAGE +COPY --from=build /src/bin /k8s-vgpu/bin diff --git a/docs/develop/mig-dynamic-deallocate_cn.md b/docs/develop/mig-dynamic-deallocate_cn.md new file mode 100644 index 0000000000..c28b0b33df --- /dev/null +++ b/docs/develop/mig-dynamic-deallocate_cn.md @@ -0,0 +1,254 @@ +# 动态 MIG 实例释放与按需重建设计 + +## 背景 + +HAMi 的动态 MIG 模式通过调度器选择 MIG 模板，并由 NVIDIA device plugin 在容器启动前创建对应的 MIG 实例。旧实现主要依赖 `nvidia-mig-parted` 对整张 GPU 应用模板：当新任务需要不同模板时，需要等待卡上所有已有 MIG 任务结束后才能重新切分。 + +本次改动的目标是把 MIG 实例生命周期从“整卡模板重切”细化到“按 slot 创建和释放”： + +- 调度器仍然按照稳定的 MIG slot 视图做资源分配。 +- device plugin 在 `Allocate` 阶段只创建被当前容器实际使用的 MIG GI/CI。 +- 容器结束后，device plugin 通过 kubelet pod-resources API 感知 MIG UUID 释放，并销毁对应 GI/CI。 +- 空闲 slot 保留其模板、位置和 profile 信息，下一个任务命中同一 slot 时可原位重建。 +- 配置层只使用 `profiles` 定义常见 MIG 模板，由代码内置 profile catalog 统一展开。 + +## 设计目标 + +1. **减少整卡重切次数** + + 对同一 MIG 模板下的任务，只按需创建或销毁单个 slot 对应的 MIG 实例，避免每次分配都重新 apply 整卡模板。 + +2. **保护运行中任务** + + 当某张 GPU 上存在其他模板下的活跃 MIG 实例时，device plugin 拒绝切换模板，避免销毁仍被容器使用的 GI/CI。 + +3. **保持调度视图稳定** + + 调度器继续使用 `GPU-UUID[templateIdx-slotIdx]` 形式表达 MIG slot。即使底层 MIG UUID 因销毁和重建发生变化，调度器与 Pod 注解中的 slot 标识仍保持稳定。 + +4. **支持 plugin 重启恢复** + + device plugin 启动时保留忙碌 GPU 的现有 MIG 布局；空闲 GPU 被重置为 “MIG enabled, no partitions”，为后续按需创建留出干净状态。 + +5. **降低默认配置维护成本** + + 默认 Helm 配置从冗长的旧模板字段切换为 `profiles` 列表，由代码内置 profile catalog 自动展开为调度器需要的 geometry。 + +## 配置模型 + +`AllowedMigGeometries` 新增 `profiles` 字段： + +```yaml +knownMigGeometries: +- models: [ "A100-SXM4-40GB", "A100-40GB-PCIe", "A100-PCIE-40GB" ] + profiles: [ "1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb" ] +``` + +归一化逻辑位于 `pkg/device/nvidia/mig_profiles.go`： + +- `profiles` 会被展开为 `[]Geometry`，每个 profile 形成一个单 profile 模板。 +- profile 的 `core`、`memory`、`count` 从内置 catalog 读取。 +- 未设置 `profiles` 的配置会被拒绝，不再兼容旧模板字段。 +- 未知 profile 会返回配置错误，scheduler/device plugin 初始化失败，避免运行时产生不一致调度视图。 + +Helm 默认配置同步简化了 `charts/hami/templates/scheduler/device-configmap.yaml` 中的 NVIDIA `knownMigGeometries`，覆盖 A30、A100、H100、H20、H200、B200 等型号。 + +## 核心数据结构 + +### MIG slot 注解 + +调度器分配的 MIG slot 会写入 Pod 注解 `hami.io/vgpu-mig-slots`： + +```json +[ + { + "deviceUUID": "GPU-xxx[0-1]", + "gpuUUID": "GPU-xxx", + "templateIdx": 0, + "slotIdx": 1 + } +] +``` + +实现位于 `pkg/device/nvidia/mig_slots.go`： + +- `EncodeMigSlotAllocations` 从调度结果中提取 `templateIdx` 和 `slotIdx`。 +- `DecodeMigSlotAllocations` 供 scheduler 和 device plugin 恢复 slot 语义。 + +这个注解解决了一个关键问题：底层 `MIG-*` UUID 是动态创建出来的，不适合作为调度器长期状态；`templateIdx + slotIdx` 才是 HAMi 资源模型里的稳定身份。 + +### MigInstanceManager + +`pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go` 新增 `MigInstanceManager`，作为节点内 MIG GI/CI 状态的单一管理者。 + +主要索引： + +- `bySlot`: `slotKey -> migInstance` +- `byMigUUID`: `MIG UUID -> slotKey` +- `gpuLocks`: 每张 GPU 一个互斥锁，串行化同一物理卡上的 NVML 操作 + +`slotKey` 由三元组组成： + +- `GPUIndex` +- `TemplateIdx` +- `PositionIdx` + +`migInstance` 保存： + +- profile 切片类型，例如 `1g`、`2g` +- NVML placement +- GI ID / CI ID +- 当前绑定的 MIG UUID +- `Present` 状态 + +当 `Present=false` 时，表示 slot 当前没有实际 GI/CI，但 manager 仍记住 profile 和 placement，后续可以原位创建。 + +## 分配流程 + +1. scheduler 在 `Fit` 阶段选择某张 GPU 的 MIG template 和 slot。 +2. scheduler 通过 Pod 注解写入原有设备分配结果，同时额外写入 `hami.io/vgpu-mig-slots`。 +3. kubelet 调用 device plugin `Allocate`。 +4. `GetContainerDeviceStrArray` 解析 `GPU-UUID[templateIdx-slotIdx]`。 +5. 如果当前 plugin 运行在 `mig` 模式，优先走 `resolveMigUUIDOnDemand`： + - 解析 template 和 slot。 + - 根据 GPU 型号和 template index 查找 geometry。 + - 调用 `PrepareGPU` 准备物理 GPU。 + - 调用 `EnsureSlot` 创建或复用该 slot 的 GI/CI。 + - 返回真实 `MIG-*` UUID 给 kubelet。 +6. 如果按需路径不可用或失败，Allocate 直接返回错误，不再回退到整卡 `nvidia-mig-parted` apply 路径。 + +简化流程如下： + +```mermaid +sequenceDiagram + participant S as Scheduler + participant P as Pod Annotation + participant D as Device Plugin + participant M as MigInstanceManager + participant N as NVML + + S->>P: 写入 GPU-UUID[templateIdx-slotIdx] + S->>P: 写入 hami.io/vgpu-mig-slots + D->>P: Allocate 读取调度结果 + D->>M: PrepareGPU(gpu, template, geometry) + M->>N: 启用 MIG / 采纳现有实例 / 检查模板冲突 + D->>M: EnsureSlot(gpu, template, slot) + M->>N: 创建 GI + CI 或返回已有 MIG UUID + M-->>D: MIG-* UUID + D-->>Kubelet: 返回容器可见设备 +``` + +## 释放流程 + +device plugin 在 `mig` 模式启动 `podresources.Watcher`，周期性调用 kubelet pod-resources API： + +- socket 挂载路径：`/var/lib/kubelet/pod-resources` +- 默认轮询周期：10 秒 +- 只关注当前 resource name，例如 `nvidia.com/gpu` + +Watcher 保存上一轮 kubelet 视图，并与当前视图做 diff： + +1. 某个 `MIG-*` device ID 从 kubelet pod-resources 中消失。 +2. 回调 `MigInstanceManager.Release(MIG UUID)`。 +3. manager 通过 `byMigUUID` 找到 slot。 +4. 使用 NVML 销毁对应 CI 和 GI。 +5. 将 slot 标记为 `Present=false`，清理 MIG UUID 反向索引。 + +释放回调还会通过当前节点 Pod 注解重建活跃 slot 集合，并调用 `ReconcileActiveSlots` 清理 manager 中仍 `Present` 但已不属于活跃 Pod 的 stale 实例。 + +## 启动恢复策略 + +device plugin 启动时，`mig_startup.go` 会执行一次 best-effort 检测，找出仍在使用的 GPU： + +- 通过 kubelet pod-resources List 查询正在使用的 `MIG-*` UUID。 +- 通过 NVML 查询父 GPU 或 MIG device 上的运行中 compute/graphics process。 + +随后 `resetIdleMigGPUs` 会修改启动 MIG spec： + +- 忙碌 GPU 保留当前布局，不销毁运行中任务使用的 GI/CI。 +- 空闲 GPU 设置为 `MigEnabled=true` 且 `MigDevices={}`，即开启 MIG 但不预创建分区。 + +这样可以在 plugin 重启后同时满足两点： + +- 不影响已有容器。 +- 空闲卡回到适合按需创建的初始状态。 + +## 调度器状态同步 + +调度器在统计节点用量时新增对 `hami.io/vgpu-mig-slots` 的解析： + +- 对带 slot 注解的 Pod，直接根据 `DeviceUUID -> slot` 标记 `MigUsage.UsageList[slotIdx].InUse=true`。 +- 如果设备还没有对应 `MigUsage`，使用 `device.PlatternMIG` 按 `templateIdx` 初始化模板视图。 +- 如果 MIG 任务错误地落在 `hami-core` 模式 GPU 上，标记设备不健康并跳过。 + +这让 scheduler 不依赖底层 MIG UUID 的稳定性，而是依赖 HAMi 分配时写入的 slot 元数据恢复资源占用。 + +## 失败处理与兼容性 + +### 模板冲突 + +`PrepareGPU` 会检查同一 GPU 是否存在其他 template 下的活跃 slot： + +- 如果存在 `Present=true` 的其他模板实例，则拒绝切换模板。 +- 如果其他模板只剩 absent slot 记录，则清理旧记录并销毁旧布局残留后重建新模板 slot map。 + +调度侧将 MIG 模式下的 `CustomFilterRule` 失败原因细化为 `CardMigTopologyInfeasible`，用于区分： + +- 设备资源不足。 +- 自定义过滤失败。 +- MIG 拓扑在不销毁活跃实例的前提下不可行。 + +### kubelet pod-resources 不可用 + +启动探测和 watcher 都对 pod-resources API 失败做降级处理： + +- 启动探测失败时，继续使用 NVML 进程检测结果。 +- watcher tick 失败时保留上一轮快照，避免误触发 release。 +- kubelet socket 重建时，watcher 会关闭旧连接并在下一轮重连。 + +## 部署变更 + +`charts/hami/templates/device-plugin/daemonsetnvidia.yaml` 新增 hostPath 挂载： + +```yaml +- name: pod-resources + hostPath: + path: /var/lib/kubelet/pod-resources + type: Directory +``` + +容器内挂载为只读： + +```yaml +- name: pod-resources + mountPath: /var/lib/kubelet/pod-resources + readOnly: true +``` + +该挂载是 watcher 访问 kubelet pod-resources unix socket 的前提。 + +## 测试覆盖 + +本分支新增和调整的测试覆盖以下场景： + +- `NormalizeMigGeometries`： + - `profiles` 展开为 geometry。 + - 缺少 `profiles` 的配置会被拒绝。 +- `EncodeMigSlotAllocations` / `DecodeMigSlotAllocations`： + - 从 `GPU-UUID[templateIdx-slotIdx]` 正确生成 slot 注解。 +- `resetIdleMigGPUs`： + - 空闲 GPU 被重置为空 MIG 布局。 + - 忙碌 GPU 保留原布局。 + - 无 devices entry 的 spec 不被修改。 +- `podresources.Watcher`： + - 能从 fake kubelet pod-resources server 读取设备快照。 + - 当前快照缺少上一轮设备时触发 release 回调。 +- scheduler score 测试： + - MIG 拓扑不可行时返回 `CardMigTopologyInfeasible`。 + +## 后续演进 + +- 为 `MigInstanceManager` 增加可注入 NVML backend，降低对真实 GPU 环境的单测依赖。 +- 将 watcher 的轮询周期、超时时间暴露为配置项，适配大规模节点和 kubelet 压力场景。 +- 在 metrics 中区分 slot absent、present、stale、release failed 等状态，便于定位资源回收问题。 +- 为 `CardMigTopologyInfeasible` 增加更明确的事件或调度失败提示，帮助用户理解是 MIG 拓扑限制而不是普通资源不足。 diff --git a/go.mod b/go.mod index 6ee832f562..80ee5a975b 100644 --- a/go.mod +++ b/go.mod @@ -26,7 +26,6 @@ require ( golang.org/x/tools v0.48.0 google.golang.org/grpc v1.83.0 gopkg.in/yaml.v2 v2.4.0 - gopkg.in/yaml.v3 v3.0.1 gotest.tools/v3 v3.5.2 k8s.io/api v0.36.3 k8s.io/apimachinery v0.36.3 @@ -90,6 +89,7 @@ require ( google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a // indirect k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 // indirect sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect diff --git a/go.sum b/go.sum index bff5cf3b07..43bc6f930b 100644 --- a/go.sum +++ b/go.sum @@ -88,10 +88,8 @@ github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnr github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/julienschmidt/httprouter v1.3.0 h1:U0609e9tgbseu3rBINet9P48AI/D3oJs4dN7jwJOQ1U= github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8IZAc4RVcycCCAKdM= -github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= -github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= -github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= +github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -133,22 +131,12 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= -github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg= -github.com/prometheus/client_golang v1.24.0 h1:5XStIklKuAtJSNpdD3s8XJj/Yv78IQmE1kbNk87JrAI= -github.com/prometheus/client_golang v1.24.0/go.mod h1:QcsNdotprC2nS4BTM2ucbcqxd2CeXTEa9jW7zHO9iDE= github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU= github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE= github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= -github.com/prometheus/common v0.67.5 h1:pIgK94WWlQt1WLwAC5j2ynLaBRDiinoAb86HZHTUGI4= -github.com/prometheus/common v0.67.5/go.mod h1:SjE/0MzDEEAyrdr5Gqc6G+sXI67maCxzaT3A2+HqjUw= -github.com/prometheus/common v0.70.0 h1:bcpru3tWPVnxGnETLgOV5jbp/JRXgYEyv65CuBLAMMI= -github.com/prometheus/common v0.70.0/go.mod h1:S/SFasQmgGiYH6C81LKCtYa8QACgthGg5zxL2udV7SY= github.com/prometheus/common v0.70.1 h1:1HvjP4D5oL3t8RsPlwxA9onvvStjtIHYE5XuuwOi/PY= github.com/prometheus/common v0.70.1/go.mod h1:VdFUQDMZK3VLkurFUVhia6uys/0suUp86TJz5qbJRhc= -github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= -github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI= github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= @@ -190,29 +178,22 @@ github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 h1:gEOO8jv9F4OT7lGC github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= -go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0= go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= -go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM= -go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY= +go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= -go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= -go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= +go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58= -go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= -go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= +go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0= go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI= -go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= -go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= +go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA= go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= +go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/zap v1.27.1 h1:08RqriUEv8+ArZRYSTXy1LeBScaMpVSTBhCeaZYfMYc= go.uber.org/zap v1.27.1/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= -go.yaml.in/yaml/v2 v2.4.3 h1:6gvOSjQoTB3vt1l+CU+tSyi/HOjfOjRLJ4YwYZGwRO0= -go.yaml.in/yaml/v2 v2.4.3/go.mod h1:zSxWcmIDjOzPXpjlTTbAsKokqkDNAVtZO0WOMiT90s8= go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= @@ -241,12 +222,8 @@ gomodules.xyz/jsonpatch/v2 v2.4.0 h1:Ci3iUJyx9UeRx7CeFN8ARgGbkESwJK+KB9lLcWxY/Zw gomodules.xyz/jsonpatch/v2 v2.4.0/go.mod h1:AH3dM2RI6uoBZxn3LVrfvJ3E0/9dG4cSrbuBJT4moAY= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk= google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.82.1 h1:NnAxzGRA0677vCa4BUkOAnO5+FfQqVl9iUXeD0IqcGE= -google.golang.org/grpc v1.82.1/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= google.golang.org/grpc v1.83.0 h1:JeNZEKJFbQxArAMl+hiytHauacDNqJUllNfmIMmpqnQ= google.golang.org/grpc v1.83.0/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= @@ -264,30 +241,20 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q= gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA= -k8s.io/api v0.36.2 h1:TF6YDLIzKfccK7cq9YpTcGX8TJmEkHVRv78DM51fRYY= -k8s.io/api v0.36.2/go.mod h1:F4LbMO4brjZYh7yFkXWhynSvtB7YauxV4c+HHkNRGNg= k8s.io/api v0.36.3 h1:NxB+05W2UGqXWFXcLO0RB5cnqnUPP5v5sVlaOH0Iz4w= k8s.io/api v0.36.3/go.mod h1:JzLQKqRHC5+I8RVj/lS3lCg0mg6nWI9Fo/Sk3ElxHzg= k8s.io/apiextensions-apiserver v0.36.0 h1:Wt7E8J+VBCbj4FjiBfDTK/neXDDjyJVJc7xfuOHImZ0= k8s.io/apiextensions-apiserver v0.36.0/go.mod h1:kGDjH0msuiIB3tgsYRV0kS9GqpMYMUsQ3GHv7TApyug= -k8s.io/apimachinery v0.36.2 h1:0PE/W/WNy1UX61NLbXY5TMbJ6UwLL6E6lAPkYrKFxbQ= -k8s.io/apimachinery v0.36.2/go.mod h1:fvf/HOLXq9RId0rnDIbN1OEBvHXdQbLMM8nu0LcBUf4= k8s.io/apimachinery v0.36.3 h1:PkzMRBRG8joFD8EhCuQAtNPvJlxb82FwplP26HIzvAM= k8s.io/apimachinery v0.36.3/go.mod h1:cTSjBWgPe/6CQyBKzY/hDIRWCQQQeK0mfLbml0UYFHE= -k8s.io/client-go v0.36.2 h1:bfgxmFKc9CgqsgX4xKLAAdmTQlWee7Ob/HlDOrJ5TBI= -k8s.io/client-go v0.36.2/go.mod h1:1vgO4OAlfPnoLcb+Rze2GF5rAr14w8qjrYMoyXJzQj0= k8s.io/client-go v0.36.3 h1:M4JdVzXxYcZk4fGpfDdYnxSwhLKWCFoQsHW6t+z8Hfg= k8s.io/client-go v0.36.3/go.mod h1:gcPwr0c87vjjG6HB6pWEqOeuYVoXSsREjzux2j6GF30= k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= k8s.io/klog/v2 v2.140.0/go.mod h1:o+/RWfJ6PwpnFn7OyAG3QnO47BFsymfEfrz6XyYSSp0= k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a h1:xCeOEAOoGYl2jnJoHkC3hkbPJgdATINPMAxaynU2Ovg= k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a/go.mod h1:uGBT7iTA6c6MvqUvSXIaYZo9ukscABYi2btjhvgKGZ0= -k8s.io/kube-scheduler v0.36.2 h1:Lf3j9NGNGA/MBFT11bLAWKZzFDS6IVRxOGP/Hw8YVLU= -k8s.io/kube-scheduler v0.36.2/go.mod h1:3Fb6kMLmawSKg8NQ6q2E6O1cQzVUIqTpn6h9AHz3hUU= k8s.io/kube-scheduler v0.36.3 h1:sc18quI2CgvH23oU1eJIQ21ivJENzjH4cxATqLSSubM= k8s.io/kube-scheduler v0.36.3/go.mod h1:M7zaLPp1Q3S6ddqZYYLiQiGw21t5bLD0FaPRXyomP7c= -k8s.io/kubelet v0.36.2 h1:9x+Tf8TEFYCcHdClzYL+IgDpfqbi+qqSdIIcXVKvr7k= -k8s.io/kubelet v0.36.2/go.mod h1:APRnAz9lmKmKsQunzUrZgQOm0k0f+NG9YxIrFCYYxcU= k8s.io/kubelet v0.36.3 h1:dRzEnhHk35Opy6wjWR4YBcN5RI9lB2npUY37TghFuPU= k8s.io/kubelet v0.36.3/go.mod h1:4USFGr21Ioka+b964Beq0NvV5b5aca3RWJ1/kfq+RLw= k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 h1:AZYQSJemyQB5eRxqcPky+/7EdBj0xi3g0ZcxxJ7vbWU= @@ -298,8 +265,6 @@ sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 h1:IpInykpT6ceI+QxKBbEflcR5E sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730/go.mod h1:mdzfpAEoE6DHQEN0uh9ZbOCuHbLK5wOm7dK4ctXE9Tg= sigs.k8s.io/randfill v1.0.0 h1:JfjMILfT8A6RbawdsK2JXGBR5AQVfd+9TbzrlneTyrU= sigs.k8s.io/randfill v1.0.0/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY= -sigs.k8s.io/structured-merge-diff/v6 v6.3.2 h1:kwVWMx5yS1CrnFWA/2QHyRVJ8jM6dBA80uLmm0wJkk8= -sigs.k8s.io/structured-merge-diff/v6 v6.3.2/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= sigs.k8s.io/structured-merge-diff/v6 v6.3.3 h1:u08YRbVUi59ri4YD6cg0UqNM4Dimn0sIl+wldcx5PYw= sigs.k8s.io/structured-merge-diff/v6 v6.3.3/go.mod h1:M3W8sfWvn2HhQDIbGWj3S099YozAsymCo/wrT5ohRUE= sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= diff --git a/hack/hami-mig-e2e.sh b/hack/hami-mig-e2e.sh new file mode 100644 index 0000000000..2b42dbc47c --- /dev/null +++ b/hack/hami-mig-e2e.sh @@ -0,0 +1,273 @@ +#!/usr/bin/env bash +set -euo pipefail + +NS=hami-mig-final-e2e +HAMI_NS=hami-system +IMAGE=${GPU_TEST_IMAGE:-nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda12.5.0-ubuntu22.04} +GPU_PROGRESS_WAIT=${GPU_PROGRESS_WAIT:-5} + +log() { printf '\n[%s] %s\n' "$(date -u +%H:%M:%S)" "$*"; } +fail() { echo "FAIL: $*" >&2; exit 1; } + +cleanup() { + kubectl delete namespace "$NS" --wait=false >/dev/null 2>&1 || true + for _ in $(seq 1 90); do + kubectl get namespace "$NS" >/dev/null 2>&1 || return 0 + sleep 2 + done + fail "namespace cleanup timed out" +} + +create_pod() { + local name=$1 memory=$2 + printf '%s\n' "{\"apiVersion\":\"v1\",\"kind\":\"Pod\",\"metadata\":{\"name\":\"${name}\",\"namespace\":\"${NS}\",\"annotations\":{\"nvidia.com/vgpu-mode\":\"mig\"}},\"spec\":{\"schedulerName\":\"hami-scheduler\",\"restartPolicy\":\"Never\",\"containers\":[{\"name\":\"cuda\",\"image\":\"${IMAGE}\",\"imagePullPolicy\":\"IfNotPresent\",\"command\":[\"bash\",\"-lc\",\"set -euo pipefail; nvidia-smi -L; n=0; echo 0 > /tmp/gpu-progress; while true; do if ! /cuda-samples/vectorAdd > /tmp/vectoradd.last 2>&1; then cat /tmp/vectoradd.last >&2; exit 1; fi; n=\$((n + 1)); echo \\\"\$n\\\" > /tmp/gpu-progress.next; mv /tmp/gpu-progress.next /tmp/gpu-progress; done\"],\"resources\":{\"limits\":{\"nvidia.com/gpu\":1,\"nvidia.com/gpumem\":${memory}}}}]}}" | kubectl apply -f - +} + +wait_ready() { kubectl wait -n "$NS" --for=condition=Ready "pod/$1" --timeout=180s; } +mig_count() { nvidia-smi -L | grep -c 'MIG ' || true; } +profile_count() { nvidia-smi -L | grep -c "MIG $1" || true; } +pod_uuid() { kubectl exec -n "$NS" "$1" -- nvidia-smi -L | sed -n 's/.*UUID: \(MIG-[^)]*\)).*/\1/p' | head -1; } +gpu_progress() { kubectl exec -n "$NS" "$1" -- cat /tmp/gpu-progress; } + +assert_gpu_progress() { + local pod=$1 before after phase restarts + phase=$(kubectl get pod -n "$NS" "$pod" -o jsonpath='{.status.phase}') + restarts=$(kubectl get pod -n "$NS" "$pod" -o jsonpath='{.status.containerStatuses[0].restartCount}') + [[ "$phase" == Running && "$restarts" == 0 ]] || fail "$pod workload unhealthy: phase=${phase} restarts=${restarts}" + before=$(gpu_progress "$pod") + sleep "$GPU_PROGRESS_WAIT" + after=$(gpu_progress "$pod") + [[ "$before" =~ ^[0-9]+$ && "$after" =~ ^[0-9]+$ && "$after" -gt "$before" ]] || { + kubectl exec -n "$NS" "$pod" -- cat /tmp/vectoradd.last >&2 || true + fail "$pod CUDA progress stalled: before=${before} after=${after}" + } + echo "GPU_PROGRESS pod=${pod} before=${before} after=${after}" +} + +snapshot_gpu_progress() { + local pod + for pod in "$@"; do + printf '%s=%s\n' "$pod" "$(gpu_progress "$pod")" + done +} + +assert_gpu_progress_since() { + local snapshot=$1 pod before after phase restarts + shift + for pod in "$@"; do + before=$(sed -n "s/^${pod}=//p" <<<"$snapshot") + after=$(gpu_progress "$pod") + phase=$(kubectl get pod -n "$NS" "$pod" -o jsonpath='{.status.phase}') + restarts=$(kubectl get pod -n "$NS" "$pod" -o jsonpath='{.status.containerStatuses[0].restartCount}') + [[ "$phase" == Running && "$restarts" == 0 && "$before" =~ ^[0-9]+$ && "$after" =~ ^[0-9]+$ && "$after" -gt "$before" ]] || { + kubectl exec -n "$NS" "$pod" -- cat /tmp/vectoradd.last >&2 || true + fail "$pod CUDA workload did not survive operation: before=${before} after=${after} phase=${phase} restarts=${restarts}" + } + echo "GPU_PROGRESS_ACROSS_OPERATION pod=${pod} before=${before} after=${after}" + done +} + +wait_count() { + local want=$1 timeout=${2:-120} start count + start=$(date +%s) + while true; do + count=$(mig_count) + [[ "$count" == "$want" ]] && { echo "MIG_COUNT=${want}"; return; } + (( $(date +%s) - start < timeout )) || { nvidia-smi -L; fail "MIG count=${count}, want=${want}"; } + sleep 2 + done +} + +assert_uuid() { + local pod=$1 want=$2 got + got=$(pod_uuid "$pod") + [[ -n "$got" && "$got" == "$want" ]] || fail "$pod UUID changed: got=${got} want=${want}" +} + +assert_runtime_annotation() { + local pod=$1 raw + raw=$(kubectl get pod -n "$NS" "$pod" -o json | jq -r '.metadata.annotations["hami.io/vgpu-mig-allocations"]') + jq -e 'length > 0 and all(.[]; (.migUUID | startswith("MIG-")) and (.profile | length > 0) and (.placement.size > 0))' <<<"$raw" >/dev/null || fail "$pod lacks concrete MIG runtime placement annotation" +} + +assert_pending_unbound() { + local pod=$1 phase node + sleep 20 + phase=$(kubectl get pod -n "$NS" "$pod" -o jsonpath='{.status.phase}') + node=$(kubectl get pod -n "$NS" "$pod" -o jsonpath='{.spec.nodeName}') + [[ "$phase" == Pending && -z "$node" ]] || fail "$pod expected Pending/unbound, got phase=${phase} node=${node}" +} + +log "baseline cleanup" +cleanup +kubectl create namespace "$NS" +wait_count 0 + +log "CASE 1: two concurrent 1g plus a mixed 2g" +create_pod one-a 4500 +create_pod one-b 4500 +wait_ready one-a & p1=$! +wait_ready one-b & p2=$! +wait "$p1"; wait "$p2" +wait_count 2 +[[ "$(profile_count 1g.5gb)" == 2 ]] || fail "expected two 1g.5gb" +uuid_one_a=$(pod_uuid one-a) +uuid_one_b=$(pod_uuid one-b) +create_pod two-a 9500 +wait_ready two-a +wait_count 3 +[[ "$(profile_count 1g.5gb)" == 2 && "$(profile_count 2g.10gb)" == 1 ]] || fail "1g/2g mixed layout missing" +uuid_two_a=$(pod_uuid two-a) +assert_runtime_annotation one-a +assert_runtime_annotation one-b +assert_runtime_annotation two-a +assert_gpu_progress one-a +assert_gpu_progress one-b +assert_gpu_progress two-a +echo "PASS CASE 1 one-a=${uuid_one_a} one-b=${uuid_one_b} two-a=${uuid_two_a}" + +log "CASE 2: reach exact 1x1g + 3x2g capacity and reject overflow" +kubectl delete pod -n "$NS" one-b --wait=true +wait_count 2 +assert_uuid one-a "$uuid_one_a" +assert_uuid two-a "$uuid_two_a" +create_pod two-b 9500 +create_pod two-c 9500 +wait_ready two-b & p1=$! +wait_ready two-c & p2=$! +wait "$p1"; wait "$p2" +wait_count 4 +[[ "$(profile_count 1g.5gb)" == 1 && "$(profile_count 2g.10gb)" == 3 ]] || fail "expected 1x1g + 3x2g" +uuid_two_b=$(pod_uuid two-b) +uuid_two_c=$(pod_uuid two-c) +assert_gpu_progress one-a +assert_gpu_progress two-a +assert_gpu_progress two-b +assert_gpu_progress two-c +create_pod overflow-one 4500 +assert_pending_unbound overflow-one +wait_count 4 +echo "PASS CASE 2 capacity enforced" + +log "CASE 3: busy device-plugin restart preserves every MIG UUID" +progress_before_restart=$(snapshot_gpu_progress one-a two-a two-b two-c) +kubectl rollout restart daemonset/hami-device-plugin -n "$HAMI_NS" +kubectl rollout status daemonset/hami-device-plugin -n "$HAMI_NS" --timeout=180s +sleep 15 +assert_gpu_progress_since "$progress_before_restart" one-a two-a two-b two-c +assert_uuid one-a "$uuid_one_a" +assert_uuid two-a "$uuid_two_a" +assert_uuid two-b "$uuid_two_b" +assert_uuid two-c "$uuid_two_c" +wait_count 4 +[[ "$(profile_count 1g.5gb)" == 1 && "$(profile_count 2g.10gb)" == 3 ]] || fail "layout changed during restart" +kubectl logs -n "$HAMI_NS" daemonset/hami-device-plugin -c device-plugin --since=3m | grep 'inUseGPUs=\[0\]' >/dev/null || fail "startup did not detect busy GPU" +echo "PASS CASE 3 all UUIDs preserved" + +log "CASE 4: immediate delete/replacement after restart" +progress_before_reclaim=$(snapshot_gpu_progress one-a two-a two-c) +kubectl delete pod -n "$NS" two-b --wait=true +wait_ready overflow-one +assert_gpu_progress_since "$progress_before_reclaim" one-a two-a two-c +assert_gpu_progress overflow-one +wait_count 4 +[[ "$(profile_count 1g.5gb)" == 2 && "$(profile_count 2g.10gb)" == 2 ]] || fail "replacement layout incorrect" +assert_uuid one-a "$uuid_one_a" +assert_uuid two-a "$uuid_two_a" +assert_uuid two-c "$uuid_two_c" +uuid_overflow=$(pod_uuid overflow-one) +[[ -n "$uuid_overflow" ]] || fail "replacement 1g has no MIG UUID" +echo "PASS CASE 4 replacement=${uuid_overflow}" + +log "reset before 3g topology" +kubectl delete pod -n "$NS" one-a two-a two-c overflow-one --wait=false +for pod in one-a two-a two-c overflow-one; do kubectl wait -n "$NS" --for=delete "pod/$pod" --timeout=120s || true; done +wait_count 0 + +log "CASE 5: scheduler rejects topology-infeasible 1g beside 2x3g" +create_pod three-a 19000 +create_pod three-b 19000 +wait_ready three-a & p1=$! +wait_ready three-b & p2=$! +wait "$p1"; wait "$p2" +wait_count 2 +[[ "$(profile_count 3g.20gb)" == 2 ]] || fail "expected two 3g instances" +assert_gpu_progress three-a +assert_gpu_progress three-b +create_pod blocked-one 4500 +assert_pending_unbound blocked-one +[[ "$(kubectl get pod -n "$NS" blocked-one -o jsonpath='{.status.reason}')" != UnexpectedAdmissionError ]] || fail "topology-infeasible Pod reached admission" +echo "PASS CASE 5 topology rejected before Bind" + +log "reset before balanced topology" +kubectl delete pod -n "$NS" three-a three-b blocked-one --wait=false +for pod in three-a three-b blocked-one; do kubectl wait -n "$NS" --for=delete "pod/$pod" --timeout=120s || true; done +wait_count 0 + +log "CASE 6: 1x3g + 1x2g + 2x1g exact capacity, overflow, and immediate replacement" +create_pod three-a 19000 +create_pod two-d 9500 +create_pod one-c 4500 +create_pod one-d 4500 +wait_ready three-a & p1=$! +wait_ready two-d & p2=$! +wait_ready one-c & p3=$! +wait_ready one-d & p4=$! +wait "$p1"; wait "$p2"; wait "$p3"; wait "$p4" +wait_count 4 +[[ "$(profile_count 3g.20gb)" == 1 && "$(profile_count 2g.10gb)" == 1 && "$(profile_count 1g.5gb)" == 2 ]] || fail "expected 1x3g + 1x2g + 2x1g" +uuid_three_a=$(pod_uuid three-a) +uuid_two_d=$(pod_uuid two-d) +uuid_one_d=$(pod_uuid one-d) +assert_gpu_progress three-a +assert_gpu_progress two-d +assert_gpu_progress one-c +assert_gpu_progress one-d +create_pod overflow-two 4500 +assert_pending_unbound overflow-two +progress_before_mixed_reclaim=$(snapshot_gpu_progress three-a two-d one-d) +kubectl delete pod -n "$NS" one-c --wait=true +wait_ready overflow-two +assert_gpu_progress_since "$progress_before_mixed_reclaim" three-a two-d one-d +assert_gpu_progress overflow-two +wait_count 4 +[[ "$(profile_count 3g.20gb)" == 1 && "$(profile_count 2g.10gb)" == 1 && "$(profile_count 1g.5gb)" == 2 ]] || fail "mixed replacement layout incorrect" +assert_uuid three-a "$uuid_three_a" +assert_uuid two-d "$uuid_two_d" +assert_uuid one-d "$uuid_one_d" +echo "PASS CASE 6" + +log "reset before seven-way burst" +kubectl delete pod -n "$NS" three-a two-d one-d overflow-two --wait=false +for pod in three-a two-d one-d overflow-two; do kubectl wait -n "$NS" --for=delete "pod/$pod" --timeout=120s || true; done +wait_count 0 + +log "CASE 7: seven concurrent 1g instances, partial reclaim and refill" +for i in $(seq 1 7); do create_pod "burst-${i}" 4500; done +for i in $(seq 1 7); do wait_ready "burst-${i}" & done +wait +wait_count 7 180 +[[ "$(profile_count 1g.5gb)" == 7 ]] || fail "expected seven 1g instances" +for i in $(seq 1 7); do assert_gpu_progress "burst-${i}"; done +uuid_b2=$(pod_uuid burst-2); uuid_b4=$(pod_uuid burst-4); uuid_b6=$(pod_uuid burst-6); uuid_b7=$(pod_uuid burst-7) +progress_before_burst_reclaim=$(snapshot_gpu_progress burst-2 burst-4 burst-6 burst-7) +kubectl delete pod -n "$NS" burst-1 burst-3 burst-5 --wait=false +for i in 1 3 5; do kubectl wait -n "$NS" --for=delete "pod/burst-${i}" --timeout=120s || true; done +wait_count 4 +assert_gpu_progress_since "$progress_before_burst_reclaim" burst-2 burst-4 burst-6 burst-7 +assert_uuid burst-2 "$uuid_b2"; assert_uuid burst-4 "$uuid_b4"; assert_uuid burst-6 "$uuid_b6"; assert_uuid burst-7 "$uuid_b7" +for i in 8 9 10; do create_pod "burst-${i}" 4500; done +for i in 8 9 10; do wait_ready "burst-${i}" & done +wait +wait_count 7 180 +[[ "$(profile_count 1g.5gb)" == 7 ]] || fail "refill did not restore seven 1g instances" +for i in 2 4 6 7 8 9 10; do assert_gpu_progress "burst-${i}"; done +echo "PASS CASE 7" + +log "final cleanup and health" +cleanup +wait_count 0 +kubectl get nodes +kubectl get pods -n "$HAMI_NS" +nvidia-smi -q | grep -A3 'MIG Mode' +echo "ALL_FIXED_MIG_E2E_TESTS_PASSED" diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go new file mode 100644 index 0000000000..fea0eadd3e --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2026, HAMi. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +package plugin + +import ( + "context" + "fmt" + "net" + "sort" + "strings" + "time" + + "github.com/NVIDIA/go-nvml/pkg/nvml" + "google.golang.org/grpc" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/credentials/insecure" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/klog/v2" + podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" + + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" + "github.com/Project-HAMi/HAMi/pkg/util/client" +) + +// Tunables for the one-shot startup poll of kubelet's pod-resources API. +// Separate from the long-lived watcher's dialWait/listTimeout because the +// startup path runs before kubelet is known to have seeded pod state, so +// falling back to NVML-only detection is acceptable. +const ( + migStartupPodresourcesSocket = "/var/lib/kubelet/pod-resources/kubelet.sock" + migStartupDialWait = 5 * time.Second + migStartupListTimeout = 10 * time.Second + migStartupMaxMsgSize = 16 * 1024 * 1024 +) + +func normalizeUnixDialAddr(addr string) string { + trimmed := strings.TrimPrefix(addr, "unix://") + if trimmed != addr { + return trimmed + } + return strings.TrimPrefix(addr, "unix:") +} + +// sortedIntSetKeys returns the keys of a set-style map sorted ascending. +// Small helper kept here so logging at startup emits stable key order. +func sortedIntSetKeys(s map[int]struct{}) []int { + out := make([]int, 0, len(s)) + for k := range s { + out = append(out, k) + } + sort.Ints(out) + return out +} + +// resetIdleMigGPUs edits the per-device MIG spec in place: GPUs that show no +// sign of in-use compute are returned to "MIG-on with no partitions" so the +// on-demand migMgr path can reshape them per request without destroying live +// GIs on busy cards. Returns the set of GPU indexes that were reset. +// +// cfg is expected to already be in per-device form (one MigConfigSpec per +// Devices=[i]); this is how processMigConfigs arranges it. +func resetIdleMigGPUs(cfg nvidia.MigConfigSpecSlice, inUse map[int]struct{}) []int { + reset := []int{} + for i := range cfg { + devs := cfg[i].Devices + if len(devs) == 0 { + continue + } + gpu := int(devs[0]) + if _, busy := inUse[gpu]; busy { + continue + } + cfg[i].MigEnabled = true + cfg[i].MigDevices = map[string]int32{} + reset = append(reset, gpu) + } + sort.Ints(reset) + return reset +} + +// collectInUseGPUs returns the set of GPU indexes that have at least one +// in-use MIG instance, unioned from two best-effort sources: +// - kubelet's pod-resources List (authoritative for k8s-managed usage). +// - NVML running processes on each MIG instance or the parent card (catches +// usage that bypasses kubelet, e.g. bare processes on the node). +// +// Failures in either source are logged and downgrade to the other source. +// When both fail, an empty set is returned and the caller treats every GPU +// as idle; the very first apply after a failed detection window will +// reshape cards conservatively because idle means "MIG on, no partitions". +func collectInUseGPUs(ctx context.Context, resourceName, nodeName string) (map[int]struct{}, error) { + out := make(map[int]struct{}) + + if uuids, err := listPodResourcesMigUUIDs(ctx, resourceName); err != nil { + klog.InfoS("mig init: pod-resources List skipped", "err", err) + } else { + for uuid := range uuids { + if gpu, ok := migUUIDToGPUIndex(uuid); ok { + out[gpu] = struct{}{} + } + } + } + + annotated, err := kubernetesAllocatedMigGPUs(ctx, nodeName) + if err != nil { + return out, fmt.Errorf("list Kubernetes MIG allocations: %w", err) + } + for g := range annotated { + out[g] = struct{}{} + } + + if busy, err := nvmlBusyGPUs(); err != nil { + klog.InfoS("mig init: NVML busy-GPU detection skipped", "err", err) + } else { + for g := range busy { + out[g] = struct{}{} + } + } + + return out, nil +} + +// activeMigGPUUUIDs returns physical GPU UUIDs referenced by live HAMi MIG +// allocations. HAMi exposes virtual resource IDs to kubelet and passes the +// dynamically created MIG UUID through NVIDIA_VISIBLE_DEVICES, so kubelet's +// pod-resources API cannot by itself identify these allocations after a +// device-plugin restart. +func activeMigGPUUUIDs(pods []corev1.Pod) map[string]struct{} { + out := make(map[string]struct{}) + for i := range pods { + pod := &pods[i] + if pod.DeletionTimestamp != nil || pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + continue + } + allocations, err := nvidia.DecodeMigAllocations(pod.Annotations[nvidia.MigAllocationsAnnotation]) + if err != nil { + continue + } + for _, allocation := range allocations { + if strings.HasPrefix(allocation.GPUUUID, "GPU-") { + out[allocation.GPUUUID] = struct{}{} + } + } + } + return out +} + +func kubernetesAllocatedMigGPUs(ctx context.Context, nodeName string) (map[int]struct{}, error) { + kubeClient := client.GetClient() + if kubeClient == nil { + return nil, fmt.Errorf("Kubernetes client is not initialized") + } + pods, err := kubeClient.CoreV1().Pods("").List(ctx, metav1.ListOptions{ + FieldSelector: "spec.nodeName=" + nodeName, + }) + if err != nil { + return nil, err + } + + out := make(map[int]struct{}) + for gpuUUID := range activeMigGPUUUIDs(pods.Items) { + idx, ok := gpuUUIDToIndex(gpuUUID) + if !ok { + return nil, fmt.Errorf("resolve GPU UUID %s", gpuUUID) + } + out[idx] = struct{}{} + } + return out, nil +} + +func gpuUUIDToIndex(gpuUUID string) (int, bool) { + if nvret := nvml.Init(); nvret != nvml.SUCCESS { + return 0, false + } + dev, ret := nvml.DeviceGetHandleByUUID(gpuUUID) + if ret != nvml.SUCCESS { + return 0, false + } + idx, ret := dev.GetIndex() + return idx, ret == nvml.SUCCESS +} + +// listPodResourcesMigUUIDs issues a single List on the kubelet pod-resources +// API and returns the set of MIG device IDs currently attached to containers +// under the given resource name. The call uses bounded dial and RPC timeouts +// so a kubelet that isn't accepting yet doesn't block plugin startup. +func listPodResourcesMigUUIDs(ctx context.Context, resourceName string) (map[string]struct{}, error) { + conn, err := grpc.NewClient( + "unix://"+migStartupPodresourcesSocket, + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(migStartupMaxMsgSize)), + grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, "unix", normalizeUnixDialAddr(addr)) + }), + ) + if err != nil { + return nil, err + } + defer conn.Close() + + conn.Connect() + dialCtx, cancelDial := context.WithTimeout(ctx, migStartupDialWait) + defer cancelDial() + for { + s := conn.GetState() + if s == connectivity.Ready { + break + } + if !conn.WaitForStateChange(dialCtx, s) { + return nil, fmt.Errorf("pod-resources dial: %w", dialCtx.Err()) + } + } + + listCtx, cancelList := context.WithTimeout(ctx, migStartupListTimeout) + defer cancelList() + cl := podresourcesv1.NewPodResourcesListerClient(conn) + resp, err := cl.List(listCtx, &podresourcesv1.ListPodResourcesRequest{}) + if err != nil { + return nil, err + } + + out := make(map[string]struct{}) + for _, pod := range resp.GetPodResources() { + for _, c := range pod.GetContainers() { + for _, d := range c.GetDevices() { + if !strings.EqualFold(d.GetResourceName(), resourceName) { + continue + } + for _, id := range d.GetDeviceIds() { + if strings.HasPrefix(id, "MIG-") { + out[id] = struct{}{} + } + } + } + } + } + return out, nil +} + +// migUUIDToGPUIndex resolves a MIG device UUID to its parent GPU's NVML +// index. Missing MIG UUIDs (e.g. stale kubelet state) return false so the +// caller skips them rather than mis-attributing to GPU 0. +func migUUIDToGPUIndex(migUUID string) (int, bool) { + if nvret := nvml.Init(); nvret != nvml.SUCCESS { + klog.InfoS("mig init: nvml.Init failed", "err", nvml.ErrorString(nvret)) + return 0, false + } + migDev, ret := nvml.DeviceGetHandleByUUID(migUUID) + if ret != nvml.SUCCESS { + return 0, false + } + parent, ret := nvml.DeviceGetDeviceHandleFromMigDeviceHandle(migDev) + if ret != nvml.SUCCESS { + return 0, false + } + idx, ret := parent.GetIndex() + if ret != nvml.SUCCESS { + return 0, false + } + return idx, true +} + +// nvmlBusyGPUs returns the set of GPU indexes with at least one running +// compute or graphics process. For MIG-enabled cards every live MIG instance +// is inspected; for non-MIG cards the parent device is inspected directly. +func nvmlBusyGPUs() (map[int]struct{}, error) { + if nvret := nvml.Init(); nvret != nvml.SUCCESS { + return nil, fmt.Errorf("nvml Init: %s", nvml.ErrorString(nvret)) + } + count, ret := nvml.DeviceGetCount() + if ret != nvml.SUCCESS { + return nil, fmt.Errorf("DeviceGetCount: %s", nvml.ErrorString(ret)) + } + + out := make(map[int]struct{}) + for i := 0; i < count; i++ { + dev, ret := nvml.DeviceGetHandleByIndex(i) + if ret != nvml.SUCCESS { + continue + } + + curMode, _, ret := dev.GetMigMode() + if ret == nvml.ERROR_NOT_SUPPORTED || ret != nvml.SUCCESS || curMode != nvml.DEVICE_MIG_ENABLE { + if deviceHasProcesses(dev) { + out[i] = struct{}{} + } + continue + } + + maxCount, ret := dev.GetMaxMigDeviceCount() + if ret != nvml.SUCCESS { + continue + } + for j := 0; j < maxCount; j++ { + migDev, ret := dev.GetMigDeviceHandleByIndex(j) + if ret != nvml.SUCCESS { + continue + } + if deviceHasProcesses(migDev) { + out[i] = struct{}{} + break + } + } + } + return out, nil +} + +func deviceHasProcesses(dev nvml.Device) bool { + if procs, ret := dev.GetComputeRunningProcesses(); ret == nvml.SUCCESS { + if len(procs) > 0 { + return true + } + } else { + // NVML query failed; assume busy so startup does not reset in-use GPUs. + return true + } + if gprocs, ret := dev.GetGraphicsRunningProcesses(); ret == nvml.SUCCESS { + return len(gprocs) > 0 + } + return true +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go new file mode 100644 index 0000000000..d22ce48db4 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go @@ -0,0 +1,138 @@ +/* + * Copyright (c) 2026, HAMi. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +package plugin + +import ( + "reflect" + "testing" + "time" + + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +func TestResetIdleMigGPUs(t *testing.T) { + tests := []struct { + name string + input nvidia.MigConfigSpecSlice + inUse map[int]struct{} + wantReset []int + wantLayout nvidia.MigConfigSpecSlice + }{ + { + name: "all idle gets reset to MIG-on empty partitions", + input: nvidia.MigConfigSpecSlice{ + {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{"1g.5gb": 4}}, + {MigEnabled: false, Devices: []int32{1}, MigDevices: map[string]int32{}}, + }, + inUse: map[int]struct{}{}, + wantReset: []int{0, 1}, + wantLayout: nvidia.MigConfigSpecSlice{ + {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{}}, + {MigEnabled: true, Devices: []int32{1}, MigDevices: map[string]int32{}}, + }, + }, + { + name: "busy gpu is preserved, idle is reset", + input: nvidia.MigConfigSpecSlice{ + {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{"1g.5gb": 4}}, + {MigEnabled: true, Devices: []int32{1}, MigDevices: map[string]int32{"3g.20gb": 2}}, + }, + inUse: map[int]struct{}{0: {}}, + wantReset: []int{1}, + wantLayout: nvidia.MigConfigSpecSlice{ + {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{"1g.5gb": 4}}, + {MigEnabled: true, Devices: []int32{1}, MigDevices: map[string]int32{}}, + }, + }, + { + name: "spec entry without devices is left alone", + input: nvidia.MigConfigSpecSlice{ + {MigEnabled: true, Devices: nil, MigDevices: map[string]int32{"1g.5gb": 4}}, + }, + inUse: map[int]struct{}{}, + wantReset: []int{}, + wantLayout: nvidia.MigConfigSpecSlice{ + {MigEnabled: true, Devices: nil, MigDevices: map[string]int32{"1g.5gb": 4}}, + }, + }, + { + name: "empty input", + input: nvidia.MigConfigSpecSlice{}, + inUse: map[int]struct{}{}, + wantReset: []int{}, + wantLayout: nvidia.MigConfigSpecSlice{}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := resetIdleMigGPUs(tc.input, tc.inUse) + if !reflect.DeepEqual(got, tc.wantReset) { + t.Errorf("resetIdleMigGPUs reset set = %v, want %v", got, tc.wantReset) + } + if !reflect.DeepEqual(tc.input, tc.wantLayout) { + t.Errorf("resetIdleMigGPUs mutated layout = %+v, want %+v", tc.input, tc.wantLayout) + } + }) + } +} + +func TestSortedIntSetKeys(t *testing.T) { + got := sortedIntSetKeys(map[int]struct{}{3: {}, 0: {}, 1: {}}) + want := []int{0, 1, 3} + if !reflect.DeepEqual(got, want) { + t.Errorf("sortedIntSetKeys = %v, want %v", got, want) + } + if got := sortedIntSetKeys(nil); len(got) != 0 { + t.Errorf("sortedIntSetKeys(nil) = %v, want empty", got) + } +} + +func TestActiveMigGPUUUIDs(t *testing.T) { + now := metav1.NewTime(time.Now()) + pods := []corev1.Pod{ + { + ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{ + nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-live","profile":"1g.5gb","placement":{"start":6,"size":1}}]`, + }}, + Status: corev1.PodStatus{Phase: corev1.PodRunning}, + }, + { + ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{ + nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-pending","profile":"2g.10gb","placement":{"start":0,"size":2}}]`, + }}, + Status: corev1.PodStatus{Phase: corev1.PodPending}, + }, + { + ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{ + nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-complete","profile":"1g.5gb","placement":{"start":0,"size":1}}]`, + }}, + Status: corev1.PodStatus{Phase: corev1.PodSucceeded}, + }, + { + ObjectMeta: metav1.ObjectMeta{ + DeletionTimestamp: &now, + Annotations: map[string]string{ + nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-deleting","profile":"1g.5gb","placement":{"start":0,"size":1}}]`, + }, + }, + Status: corev1.PodStatus{Phase: corev1.PodRunning}, + }, + } + + got := activeMigGPUUUIDs(pods) + want := map[string]struct{}{"GPU-live": {}, "GPU-pending": {}} + if !reflect.DeepEqual(got, want) { + t.Fatalf("activeMigGPUUUIDs() = %v, want %v", got, want) + } +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go new file mode 100644 index 0000000000..a31312e335 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go @@ -0,0 +1,623 @@ +/* + * Copyright (c) 2026, HAMi. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +package plugin + +import ( + "fmt" + "sort" + "strings" + "sync" + + "github.com/NVIDIA/go-nvml/pkg/nvml" + "k8s.io/klog/v2" +) + +var profileNameToGIProfileID = map[string]int{ + "1g": nvml.GPU_INSTANCE_PROFILE_1_SLICE, + "2g": nvml.GPU_INSTANCE_PROFILE_2_SLICE, + "3g": nvml.GPU_INSTANCE_PROFILE_3_SLICE, + "4g": nvml.GPU_INSTANCE_PROFILE_4_SLICE, + "6g": nvml.GPU_INSTANCE_PROFILE_6_SLICE, + "7g": nvml.GPU_INSTANCE_PROFILE_7_SLICE, + "8g": nvml.GPU_INSTANCE_PROFILE_8_SLICE, +} + +var profileNameToCIProfileID = map[string]int{ + "1g": nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE, + "2g": nvml.COMPUTE_INSTANCE_PROFILE_2_SLICE, + "3g": nvml.COMPUTE_INSTANCE_PROFILE_3_SLICE, + "4g": nvml.COMPUTE_INSTANCE_PROFILE_4_SLICE, + "6g": nvml.COMPUTE_INSTANCE_PROFILE_6_SLICE, + "7g": nvml.COMPUTE_INSTANCE_PROFILE_7_SLICE, + "8g": nvml.COMPUTE_INSTANCE_PROFILE_8_SLICE, +} + +type migAllocationKey struct { + GPUIndex int + Profile string + Start uint32 + Size uint32 +} + +// migInstance tracks the nvml-level identity of a MIG GI+CI pair bound to a +// slot. Absent means the slot's GI+CI have been destroyed (e.g. on task end) +// but we remember the profile and placement so we can recreate the instance +// at the same physical slice when the next task claims this slot. +type migInstance struct { + Profile string // slice group, e.g. "1g" + Placement nvml.GpuInstancePlacement + Present bool + GIID uint32 + CIID uint32 + MigUUID string +} + +// MigInstanceManager is the single authority over live MIG GI+CI state on a +// node. Keys are the scheduler-reserved profile and physical placement. +type MigInstanceManager struct { + mu sync.Mutex + gpuLocks map[int]*sync.Mutex + byAllocation map[migAllocationKey]*migInstance + byAllocationMigUUID map[string]migAllocationKey +} + +func NewMigInstanceManager() *MigInstanceManager { + return &MigInstanceManager{ + gpuLocks: make(map[int]*sync.Mutex), + byAllocation: make(map[migAllocationKey]*migInstance), + byAllocationMigUUID: make(map[string]migAllocationKey), + } +} + +func (m *MigInstanceManager) gpuLock(gpuIndex int) *sync.Mutex { + m.mu.Lock() + defer m.mu.Unlock() + lk, ok := m.gpuLocks[gpuIndex] + if !ok { + lk = &sync.Mutex{} + m.gpuLocks[gpuIndex] = lk + } + return lk +} + +func profileSliceKey(profile string) string { + if idx := strings.Index(profile, "."); idx > 0 { + return profile[:idx] + } + return profile +} + +// ResetIdleGPUs prepares idle MIG-capable GPUs for on-demand slot creation +// without calling nvidia-mig-parted. Busy GPUs are left untouched; idle GPUs +// have MIG mode enabled and all existing GI/CI instances destroyed. +func (m *MigInstanceManager) ResetIdleGPUs(deviceCount int, inUse map[int]struct{}) ([]int, error) { + reset := []int{} + for gpuIndex := 0; gpuIndex < deviceCount; gpuIndex++ { + if _, busy := inUse[gpuIndex]; busy { + continue + } + + lk := m.gpuLock(gpuIndex) + lk.Lock() + if err := ensureMigModeEnabled(gpuIndex); err != nil { + lk.Unlock() + return reset, err + } + dev, err := deviceHandleByIndex(gpuIndex) + if err != nil { + lk.Unlock() + return reset, err + } + if err := destroyAllMigInstances(dev); err != nil { + lk.Unlock() + return reset, err + } + + m.mu.Lock() + for key := range m.byAllocation { + if key.GPUIndex == gpuIndex { + delete(m.byAllocation, key) + } + } + for uuid, key := range m.byAllocationMigUUID { + if key.GPUIndex == gpuIndex { + delete(m.byAllocationMigUUID, uuid) + } + } + m.mu.Unlock() + lk.Unlock() + reset = append(reset, gpuIndex) + } + sort.Ints(reset) + return reset, nil +} + +func deviceHandleByIndex(gpuIndex int) (nvml.Device, error) { + dev, ret := nvml.DeviceGetHandleByIndex(gpuIndex) + if ret != nvml.SUCCESS { + return nil, fmt.Errorf("nvml get handle by index %d: %s", gpuIndex, nvml.ErrorString(ret)) + } + return dev, nil +} + +// ensureMigModeEnabled turns on MIG mode via NVML when the card is currently +// in non-MIG mode. No-op when MIG mode is unsupported (non-MIG cards) so the +// caller can invoke it uniformly. +// +// SetMigMode may reset/unbind the device; callers must re-fetch the device +// handle after this returns successfully before further NVML operations. +func ensureMigModeEnabled(gpuIndex int) error { + dev, err := deviceHandleByIndex(gpuIndex) + if err != nil { + return err + } + curMode, pendingMode, ret := dev.GetMigMode() + if ret == nvml.ERROR_NOT_SUPPORTED { + return nil + } + if ret != nvml.SUCCESS { + return fmt.Errorf("gpu %d get mig mode: %s", gpuIndex, nvml.ErrorString(ret)) + } + if curMode == nvml.DEVICE_MIG_ENABLE { + if pendingMode == nvml.DEVICE_MIG_ENABLE { + return nil + } + return fmt.Errorf("gpu %d mig mode disable is pending (current=enable pending=%d)", gpuIndex, pendingMode) + } + if pendingMode == nvml.DEVICE_MIG_ENABLE { + return fmt.Errorf("gpu %d mig mode enable is pending; GPU reset required", gpuIndex) + } + + activation, ret := dev.SetMigMode(nvml.DEVICE_MIG_ENABLE) + if ret != nvml.SUCCESS { + return fmt.Errorf("gpu %d set mig mode: %s", gpuIndex, nvml.ErrorString(ret)) + } + if activation != nvml.SUCCESS { + return fmt.Errorf("gpu %d activate mig mode: %s", gpuIndex, nvml.ErrorString(activation)) + } + + dev, err = deviceHandleByIndex(gpuIndex) + if err != nil { + return err + } + curMode, pendingMode, ret = dev.GetMigMode() + if ret != nvml.SUCCESS { + return fmt.Errorf("gpu %d verify mig mode after set: %s", gpuIndex, nvml.ErrorString(ret)) + } + if curMode == nvml.DEVICE_MIG_ENABLE { + return nil + } + if pendingMode == nvml.DEVICE_MIG_ENABLE { + return fmt.Errorf("gpu %d mig mode enable is pending after set; GPU reset required", gpuIndex) + } + return fmt.Errorf("gpu %d mig mode is not enabled after set (current=%d pending=%d)", gpuIndex, curMode, pendingMode) +} + +// destroyPresentMigInstance destroys the tracked GI+CI on hardware. Returns +// nil when the instance is already gone or was destroyed successfully. +func destroyPresentMigInstance(gpuIndex int, inst *migInstance) error { + if inst == nil || !inst.Present { + return nil + } + dev, err := deviceHandleByIndex(gpuIndex) + if err != nil { + return err + } + gi, ret := dev.GetGpuInstanceById(int(inst.GIID)) + if ret == nvml.ERROR_NOT_FOUND { + return nil + } + if ret != nvml.SUCCESS { + return fmt.Errorf("get GI %d on gpu %d: %s", inst.GIID, gpuIndex, nvml.ErrorString(ret)) + } + if ci, r := gi.GetComputeInstanceById(int(inst.CIID)); r == nvml.SUCCESS { + if d := ci.Destroy(); d != nvml.SUCCESS { + return fmt.Errorf("destroy CI %d on gpu %d: %s", inst.CIID, gpuIndex, nvml.ErrorString(d)) + } + } else if r != nvml.ERROR_NOT_FOUND { + return fmt.Errorf("get CI %d on gpu %d: %s", inst.CIID, gpuIndex, nvml.ErrorString(r)) + } + if d := gi.Destroy(); d != nvml.SUCCESS { + return fmt.Errorf("destroy GI %d on gpu %d: %s", inst.GIID, gpuIndex, nvml.ErrorString(d)) + } + return nil +} + +// destroyAllMigInstances enumerates and destroys every GI+CI on the device. +// Used on template switches when no scheduler-allocated slot is in use. +func destroyAllMigInstances(dev nvml.Device) error { + for _, giProfileID := range []int{ + nvml.GPU_INSTANCE_PROFILE_1_SLICE, + nvml.GPU_INSTANCE_PROFILE_2_SLICE, + nvml.GPU_INSTANCE_PROFILE_3_SLICE, + nvml.GPU_INSTANCE_PROFILE_4_SLICE, + nvml.GPU_INSTANCE_PROFILE_6_SLICE, + nvml.GPU_INSTANCE_PROFILE_7_SLICE, + nvml.GPU_INSTANCE_PROFILE_8_SLICE, + } { + info, ret := dev.GetGpuInstanceProfileInfo(giProfileID) + if ret != nvml.SUCCESS { + continue + } + gis, ret := dev.GetGpuInstances(&info) + if ret != nvml.SUCCESS { + continue + } + for _, gi := range gis { + ciInfoRet := profileIDToCIProfileID(giProfileID) + if ciInfo, r := gi.GetComputeInstanceProfileInfo(ciInfoRet, nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED); r == nvml.SUCCESS { + if cis, r2 := gi.GetComputeInstances(&ciInfo); r2 == nvml.SUCCESS { + for _, ci := range cis { + if d := ci.Destroy(); d != nvml.SUCCESS { + return fmt.Errorf("destroy compute instance: %s", nvml.ErrorString(d)) + } + } + } + } + if d := gi.Destroy(); d != nvml.SUCCESS { + return fmt.Errorf("destroy gpu instance: %s", nvml.ErrorString(d)) + } + } + } + return nil +} + +// Release destroys the GI+CI bound to the given MIG UUID and marks the slot +// absent (preserving its profile and placement). Invoked by the podresources +// watcher when kubelet reports the device is no longer in use. +func (m *MigInstanceManager) Release(migUUID string) error { + m.mu.Lock() + key, ok := m.byAllocationMigUUID[migUUID] + m.mu.Unlock() + if !ok { + klog.V(5).InfoS("release: unknown MIG UUID, skipping", "uuid", migUUID) + return nil + } + lk := m.gpuLock(key.GPUIndex) + lk.Lock() + defer lk.Unlock() + m.mu.Lock() + inst := m.byAllocation[key] + m.mu.Unlock() + if inst == nil || !inst.Present { + return nil + } + if err := destroyPresentMigInstance(key.GPUIndex, inst); err != nil { + return err + } + m.mu.Lock() + delete(m.byAllocation, key) + delete(m.byAllocationMigUUID, migUUID) + m.mu.Unlock() + klog.InfoS("released MIG allocation", "uuid", migUUID, "gpu", key.GPUIndex, "profile", key.Profile, "start", key.Start) + return nil +} +func allocationKey(gpuIndex int, profile string, placement nvml.GpuInstancePlacement) migAllocationKey { + return migAllocationKey{GPUIndex: gpuIndex, Profile: profile, Start: placement.Start, Size: placement.Size} +} + +// EnsureAllocation realizes exactly the scheduler-reserved profile and +// placement. It never retries another placement. +func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, placement nvml.GpuInstancePlacement) (string, error) { + key := allocationKey(gpuIndex, profile, placement) + lk := m.gpuLock(gpuIndex) + lk.Lock() + defer lk.Unlock() + + m.mu.Lock() + if inst := m.byAllocation[key]; inst != nil && inst.Present { + uuid := inst.MigUUID + m.mu.Unlock() + return uuid, nil + } + m.mu.Unlock() + + if err := ensureMigModeEnabled(gpuIndex); err != nil { + return "", err + } + profileKey := profileSliceKey(profile) + giProfileID, ok := profileNameToGIProfileID[profileKey] + if !ok { + return "", fmt.Errorf("unsupported MIG profile %q", profile) + } + ciProfileID, ok := profileNameToCIProfileID[profileKey] + if !ok { + return "", fmt.Errorf("unsupported MIG compute profile %q", profile) + } + dev, err := deviceHandleByIndex(gpuIndex) + if err != nil { + return "", err + } + giInfo, ret := dev.GetGpuInstanceProfileInfo(giProfileID) + if ret != nvml.SUCCESS { + return "", fmt.Errorf("get GI profile %s: %s", profile, nvml.ErrorString(ret)) + } + possible, ret := dev.GetGpuInstancePossiblePlacements(&giInfo) + if ret != nvml.SUCCESS { + return "", fmt.Errorf("get placements for %s: %s", profile, nvml.ErrorString(ret)) + } + valid := false + for _, candidate := range possible { + if candidate == placement { + valid = true + break + } + } + if !valid { + return "", fmt.Errorf("scheduler selected invalid placement %+v for profile %s", placement, profile) + } + gi, ret := dev.CreateGpuInstanceWithPlacement(&giInfo, &placement) + if ret != nvml.SUCCESS { + return "", fmt.Errorf("create GI profile=%s placement=%+v: %s", profile, placement, nvml.ErrorString(ret)) + } + giData, ret := gi.GetInfo() + if ret != nvml.SUCCESS { + gi.Destroy() + return "", fmt.Errorf("get GI info: %s", nvml.ErrorString(ret)) + } + ciInfo, ret := gi.GetComputeInstanceProfileInfo(ciProfileID, nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED) + if ret != nvml.SUCCESS { + gi.Destroy() + return "", fmt.Errorf("get CI profile info: %s", nvml.ErrorString(ret)) + } + ci, ret := gi.CreateComputeInstance(&ciInfo) + if ret != nvml.SUCCESS { + gi.Destroy() + return "", fmt.Errorf("create CI: %s", nvml.ErrorString(ret)) + } + ciData, ret := ci.GetInfo() + if ret != nvml.SUCCESS { + ci.Destroy() + gi.Destroy() + return "", fmt.Errorf("get CI info: %s", nvml.ErrorString(ret)) + } + migUUID, err := findMigUUIDForGI(dev, giData.Id) + if err != nil { + ci.Destroy() + gi.Destroy() + return "", err + } + inst := &migInstance{Profile: profile, Placement: placement, Present: true, GIID: giData.Id, CIID: ciData.Id, MigUUID: migUUID} + m.mu.Lock() + m.byAllocation[key] = inst + m.byAllocationMigUUID[migUUID] = key + m.mu.Unlock() + klog.InfoS("created scheduler-reserved MIG allocation", "uuid", migUUID, "gpu", gpuIndex, "profile", profile, "start", placement.Start, "size", placement.Size) + return migUUID, nil +} + +func (m *MigInstanceManager) AllocationRuntimeInfo(gpuIndex int, profile string, placement nvml.GpuInstancePlacement) (migAllocationRuntimeInfo, bool) { + key := allocationKey(gpuIndex, profile, placement) + m.mu.Lock() + defer m.mu.Unlock() + inst := m.byAllocation[key] + if inst == nil || !inst.Present { + return migAllocationRuntimeInfo{}, false + } + return migAllocationRuntimeInfo{MigUUID: inst.MigUUID, Profile: inst.Profile, Placement: inst.Placement}, true +} + +func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID string, placement nvml.GpuInstancePlacement) error { + lk := m.gpuLock(gpuIndex) + lk.Lock() + defer lk.Unlock() + dev, err := deviceHandleByIndex(gpuIndex) + if err != nil { + return err + } + giProfileID, ok := profileNameToGIProfileID[profileSliceKey(profile)] + if !ok { + return fmt.Errorf("unsupported MIG profile %q", profile) + } + profileInfo, ret := dev.GetGpuInstanceProfileInfo(giProfileID) + if ret != nvml.SUCCESS { + return fmt.Errorf("get GI profile %s: %s", profile, nvml.ErrorString(ret)) + } + instances, ret := dev.GetGpuInstances(&profileInfo) + if ret != nvml.SUCCESS { + return fmt.Errorf("list GI profile %s: %s", profile, nvml.ErrorString(ret)) + } + ciProfileID := profileIDToCIProfileID(giProfileID) + for _, gi := range instances { + giInfo, r := gi.GetInfo() + if r != nvml.SUCCESS || giInfo.Placement != placement { + continue + } + actualUUID, findErr := findMigUUIDForGI(dev, giInfo.Id) + if findErr != nil || actualUUID != migUUID { + continue + } + ciInfo, r := gi.GetComputeInstanceProfileInfo(ciProfileID, nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED) + if r != nvml.SUCCESS { + continue + } + cis, r := gi.GetComputeInstances(&ciInfo) + if r != nvml.SUCCESS || len(cis) == 0 { + continue + } + ciData, r := cis[0].GetInfo() + if r != nvml.SUCCESS { + continue + } + key := allocationKey(gpuIndex, profile, placement) + m.mu.Lock() + m.byAllocation[key] = &migInstance{Profile: profile, Placement: placement, Present: true, GIID: giInfo.Id, CIID: ciData.Id, MigUUID: migUUID} + m.byAllocationMigUUID[migUUID] = key + m.mu.Unlock() + return nil + } + return fmt.Errorf("annotated MIG allocation %s profile=%s placement=%+v is not live", migUUID, profile, placement) +} + +func (m *MigInstanceManager) ReconcileActiveAllocations(active map[migAllocationKey]struct{}) error { + m.mu.Lock() + keys := make([]migAllocationKey, 0, len(m.byAllocation)) + for key, inst := range m.byAllocation { + if inst.Present { + keys = append(keys, key) + } + } + m.mu.Unlock() + for _, key := range keys { + if _, ok := active[key]; ok { + continue + } + lk := m.gpuLock(key.GPUIndex) + lk.Lock() + m.mu.Lock() + inst := m.byAllocation[key] + m.mu.Unlock() + if inst != nil && inst.Present { + oldUUID := inst.MigUUID + if err := destroyPresentMigInstance(key.GPUIndex, inst); err != nil { + lk.Unlock() + return err + } + m.mu.Lock() + delete(m.byAllocation, key) + delete(m.byAllocationMigUUID, oldUUID) + m.mu.Unlock() + } + lk.Unlock() + } + return nil +} + +type migAllocationRuntimeInfo struct { + MigUUID string + Profile string + Placement nvml.GpuInstancePlacement +} + +func findMigUUIDForGI(dev nvml.Device, giID uint32) (string, error) { + maxCount, ret := dev.GetMaxMigDeviceCount() + if ret != nvml.SUCCESS { + return "", fmt.Errorf("get max MIG device count: %s", nvml.ErrorString(ret)) + } + for i := 0; i < maxCount; i++ { + migDev, ret := dev.GetMigDeviceHandleByIndex(i) + if ret != nvml.SUCCESS { + continue + } + gotGI, ret := migDev.GetGpuInstanceId() + if ret != nvml.SUCCESS { + continue + } + if uint32(gotGI) == giID { + uuid, ret := migDev.GetUUID() + if ret != nvml.SUCCESS { + return "", fmt.Errorf("get MIG UUID: %s", nvml.ErrorString(ret)) + } + return uuid, nil + } + } + return "", fmt.Errorf("no MIG device found for GI %d", giID) +} + +// pickFreePlacement returns a placement for the given GI profile that does +// not overlap with any of the placements already in use on this GPU. +func pickFreePlacement(dev nvml.Device, info *nvml.GpuInstanceProfileInfo, inUse map[uint32]uint32) (nvml.GpuInstancePlacement, error) { + possible, ret := dev.GetGpuInstancePossiblePlacements(info) + if ret != nvml.SUCCESS { + return nvml.GpuInstancePlacement{}, fmt.Errorf("get possible placements: %s", nvml.ErrorString(ret)) + } + return chooseFreePlacement(possible, inUse, preferHighPlacement(info.SliceCount)) +} + +func preferHighPlacement(sliceCount uint32) bool { + return sliceCount == 1 || sliceCount == 3 +} + +func sortPlacements(possible []nvml.GpuInstancePlacement, preferHigh bool) { + sort.SliceStable(possible, func(i, j int) bool { + if preferHigh { + return possible[i].Start > possible[j].Start + } + return possible[i].Start < possible[j].Start + }) +} + +func placementCandidates(previous nvml.GpuInstancePlacement, possible []nvml.GpuInstancePlacement) []nvml.GpuInstancePlacement { + candidates := make([]nvml.GpuInstancePlacement, 0, len(possible)+1) + if previous.Size != 0 { + candidates = append(candidates, previous) + } + for _, candidate := range possible { + if previous.Size != 0 && candidate == previous { + continue + } + candidates = append(candidates, candidate) + } + return candidates +} + +// chooseFreePlacement packs 1g and 3g instances from high addresses while 2g +// instances use low addresses. This matches NVIDIA's A100 balanced placement +// (2g at 0:2, 1g at 2:1 and 3:1, 3g at 4:4) while retaining the canonical +// 1x1g + 3x2g layout. +func chooseFreePlacement(possible []nvml.GpuInstancePlacement, inUse map[uint32]uint32, preferHigh bool) (nvml.GpuInstancePlacement, error) { + sortPlacements(possible, preferHigh) + for _, p := range possible { + if !placementOverlaps(p, inUse) { + return p, nil + } + } + return nvml.GpuInstancePlacement{}, fmt.Errorf("no free placement for profile") +} + +func placementOverlaps(p nvml.GpuInstancePlacement, inUse map[uint32]uint32) bool { + for start, size := range inUse { + if p.Start < start+size && start < p.Start+p.Size { + return true + } + } + return false +} + +func profileIDToCIProfileID(giProfileID int) int { + switch giProfileID { + case nvml.GPU_INSTANCE_PROFILE_1_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE + case nvml.GPU_INSTANCE_PROFILE_2_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_2_SLICE + case nvml.GPU_INSTANCE_PROFILE_3_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_3_SLICE + case nvml.GPU_INSTANCE_PROFILE_4_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_4_SLICE + case nvml.GPU_INSTANCE_PROFILE_6_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_6_SLICE + case nvml.GPU_INSTANCE_PROFILE_7_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_7_SLICE + case nvml.GPU_INSTANCE_PROFILE_8_SLICE: + return nvml.COMPUTE_INSTANCE_PROFILE_8_SLICE + } + return nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE +} + +func giProfileIDToSliceKey(giProfileID int) string { + switch giProfileID { + case nvml.GPU_INSTANCE_PROFILE_1_SLICE: + return "1g" + case nvml.GPU_INSTANCE_PROFILE_2_SLICE: + return "2g" + case nvml.GPU_INSTANCE_PROFILE_3_SLICE: + return "3g" + case nvml.GPU_INSTANCE_PROFILE_4_SLICE: + return "4g" + case nvml.GPU_INSTANCE_PROFILE_6_SLICE: + return "6g" + case nvml.GPU_INSTANCE_PROFILE_7_SLICE: + return "7g" + case nvml.GPU_INSTANCE_PROFILE_8_SLICE: + return "8g" + } + return "" +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go new file mode 100644 index 0000000000..7e697bff33 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go @@ -0,0 +1,59 @@ +package plugin + +import ( + "testing" + + "github.com/NVIDIA/go-nvml/pkg/nvml" +) + +func TestChooseFreePlacementPackingDirection(t *testing.T) { + possible := []nvml.GpuInstancePlacement{ + {Start: 0, Size: 1}, + {Start: 6, Size: 1}, + {Start: 3, Size: 1}, + } + high, err := chooseFreePlacement(append([]nvml.GpuInstancePlacement(nil), possible...), nil, true) + if err != nil || high.Start != 6 { + t.Fatalf("high packing = %+v, %v; want start 6", high, err) + } + low, err := chooseFreePlacement(append([]nvml.GpuInstancePlacement(nil), possible...), nil, false) + if err != nil || low.Start != 0 { + t.Fatalf("low packing = %+v, %v; want start 0", low, err) + } + free, err := chooseFreePlacement(append([]nvml.GpuInstancePlacement(nil), possible...), map[uint32]uint32{6: 1}, true) + if err != nil || free.Start != 3 { + t.Fatalf("overlap-aware high packing = %+v, %v; want start 3", free, err) + } +} + +func TestPreferHighPlacementMatchesA100BalancedLayout(t *testing.T) { + if !preferHighPlacement(1) || !preferHighPlacement(3) { + t.Fatal("1g and 3g profiles must pack from high placements") + } + if preferHighPlacement(2) { + t.Fatal("2g profiles must pack from low placements") + } +} + +func TestPlacementCandidatesPreferPreviousThenFallback(t *testing.T) { + previous := nvml.GpuInstancePlacement{Start: 6, Size: 1} + possible := []nvml.GpuInstancePlacement{ + {Start: 5, Size: 1}, + {Start: 3, Size: 1}, + {Start: 6, Size: 1}, + } + want := []nvml.GpuInstancePlacement{ + {Start: 6, Size: 1}, + {Start: 5, Size: 1}, + {Start: 3, Size: 1}, + } + got := placementCandidates(previous, possible) + if len(got) != len(want) { + t.Fatalf("candidate count = %d, want %d: %+v", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("candidate %d = %+v, want %+v", i, got[i], want[i]) + } + } +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go index 4149bc54aa..781e126b06 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go @@ -176,7 +176,7 @@ func (plugin *NvidiaDevicePlugin) getAPIDevices() *[]*device.DeviceInfo { if !isMigMode { devcore = int32(*plugin.schedulerConfig.DeviceCoreScaling * 100) } - res = append(res, &device.DeviceInfo{ + info := &device.DeviceInfo{ ID: UUID, Index: uint(idx), Count: int32(*plugin.schedulerConfig.DeviceSplitCount), @@ -186,12 +186,73 @@ func (plugin *NvidiaDevicePlugin) getAPIDevices() *[]*device.DeviceInfo { Numa: numa, Mode: plugin.operatingMode, Health: health, - }) + } + if isMigMode { + info.MIGProfiles = plugin.discoverMigProfiles(ndev, Model) + info.Count = 0 + for _, profile := range info.MIGProfiles { + if int32(profile.InstanceCount) > info.Count { + info.Count = int32(profile.InstanceCount) + } + } + } + res = append(res, info) klog.V(3).Infof("Registered device id=%v, memory=%vMB, type=%v, numa=%v, health=%v", idx, registeredmem, Model, numa, health) } return &res } +func (plugin *NvidiaDevicePlugin) discoverMigProfiles(dev nvml.Device, model string) []device.MigProfile { + out := make([]device.MigProfile, 0) + var fullGPUMultiprocessors uint32 + for _, profileID := range profileNameToGIProfileID { + profileInfo, ret := dev.GetGpuInstanceProfileInfo(profileID) + if ret == nvml.SUCCESS && profileInfo.MultiprocessorCount > fullGPUMultiprocessors { + fullGPUMultiprocessors = profileInfo.MultiprocessorCount + } + } + for _, allowed := range plugin.schedulerConfig.MigProfileAllowlist { + if !containsModel(model, allowed.Models) { + continue + } + klog.InfoS("discovering MIG profile capabilities", "model", model, "profiles", allowed.Profiles) + for _, profileName := range allowed.Profiles { + profileID, ok := profileNameToGIProfileID[profileSliceKey(profileName)] + if !ok { + continue + } + profileInfo, ret := dev.GetGpuInstanceProfileInfo(profileID) + if ret != nvml.SUCCESS { + klog.InfoS("skip MIG profile placement discovery", "profile", profileName, "step", "profile-info", "err", nvml.ErrorString(ret)) + continue + } + placements, ret := dev.GetGpuInstancePossiblePlacements(&profileInfo) + if ret != nvml.SUCCESS { + klog.InfoS("skip MIG profile placement discovery", "profile", profileName, "step", "possible-placements", "err", nvml.ErrorString(ret)) + continue + } + profile := device.MigProfile{ + Name: profileName, GIProfileID: profileInfo.Id, MemoryMB: int32(profileInfo.MemorySizeMB), + SliceCount: profileInfo.SliceCount, InstanceCount: profileInfo.InstanceCount, + MultiprocessorCount: profileInfo.MultiprocessorCount, + P2PSupported: profileInfo.IsP2pSupported != 0, + CopyEngineCount: profileInfo.CopyEngineCount, DecoderCount: profileInfo.DecoderCount, + EncoderCount: profileInfo.EncoderCount, JPEGCount: profileInfo.JpegCount, OFACount: profileInfo.OfaCount, + } + if fullGPUMultiprocessors > 0 { + profile.Core = int32((profileInfo.MultiprocessorCount*100 + fullGPUMultiprocessors - 1) / fullGPUMultiprocessors) + } + for _, placement := range placements { + profile.Placements = append(profile.Placements, device.MigPlacement{Start: placement.Start, Size: placement.Size}) + } + out = append(out, profile) + } + break + } + klog.InfoS("discovered MIG profile capabilities", "model", model, "profiles", out) + return out +} + // RegisterInAnnotation scans devices and patches node annotations. // Returns (changed, error) where changed indicates whether the annotation was actually updated. func (plugin *NvidiaDevicePlugin) RegisterInAnnotation() (bool, error) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index e320f79c9f..01dfd5e2ea 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -33,14 +33,12 @@ package plugin import ( - "bytes" "context" "encoding/json" "errors" "fmt" "net" "os" - "os/exec" "path" "path/filepath" "strconv" @@ -48,12 +46,14 @@ import ( "sync" "time" + "github.com/NVIDIA/go-nvml/pkg/nvml" spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" "github.com/google/uuid" "github.com/imdario/mergo" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" - "k8s.io/apimachinery/pkg/util/yaml" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/klog/v2" kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" cdiapi "tags.cncf.io/container-device-interface/pkg/cdi" @@ -61,10 +61,12 @@ import ( "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/cdi" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/imex" + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/podresources" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/rm" "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/scheduler/config" "github.com/Project-HAMi/HAMi/pkg/util" + "github.com/Project-HAMi/HAMi/pkg/util/client" ) // Constants for use by the 'volume-mounts' device list strategy @@ -112,6 +114,11 @@ type NvidiaDevicePlugin struct { migCurrent nvidia.MigPartedSpec deviceCache string + // migMgr tracks live MIG GI+CI instances so we can destroy and recreate + // them per-task rather than resharding the whole card. Only set when + // operatingMode == "mig". + migMgr *MigInstanceManager + imexChannels imex.Channels server *grpc.Server @@ -181,6 +188,9 @@ func (o *options) devicePluginForResource(ctx context.Context, nvconfig *nvidia. if err := config.InitDevicesWithConfig(sConfig); err != nil { klog.Fatalf("failed to initialize devices: %v", err) } + if err := nvidia.ValidateMigProfileAllowlist(sConfig.NvidiaConfig.MigProfileAllowlist); err != nil { + return nil, fmt.Errorf("validate MIG profile allowlist: %w", err) + } return &NvidiaDevicePlugin{ ctx: ctx, rm: resourceManager, @@ -275,80 +285,204 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { } klog.Infof("Registered device plugin for '%s' with Kubelet", plugin.rm.Resource()) - migApplied := false - if plugin.operatingMode == "mig" { - deviceSupportMig := true - for _, name := range deviceNames { - supported := false - for _, migTemplate := range plugin.schedulerConfig.MigGeometriesList { - if containsModel(name, migTemplate.Models) { - supported = true - break - } - } - if !supported { - deviceSupportMig = false + // Prepare the lock directory before any dynamic MIG operation. A stale + // lock can be left behind when the previous plugin process exits midway. + if err = CreateMigApplyLockDir(); err != nil { + klog.Fatalf("CreateMIGLockSubDir failed: %v", err) + } + if err = RemoveMigApplyLock(); err != nil { + klog.Fatalf("RemoveMigApplyLock failed: %v", err) + } + + deviceSupportMig := true + for _, name := range deviceNames { + supported := false + for _, allowlist := range plugin.schedulerConfig.MigProfileAllowlist { + if containsModel(name, allowlist.Models) { + supported = true break } } + if !supported { + deviceSupportMig = false + break + } + } + go func() { + err := plugin.rm.CheckHealth(plugin.stop, plugin.health, plugin.disableHealthChecks, plugin.ackDisableHealthChecks) + if err != nil { + klog.Infof("Failed to start health check: %v; continuing with health checks disabled", err) + } + }() + go func() { + plugin.WatchAndRegister(plugin.disableWatchAndRegister, plugin.ackDisableWatchAndRegister) + }() + if plugin.operatingMode == "mig" { + plugin.migMgr = NewMigInstanceManager() if deviceSupportMig { - err = CreateMigApplyLockDir() - if err != nil { - klog.Fatalf("CreateMIGLockSubDir failed: %v", err) + inUse, detectErr := collectInUseGPUs(plugin.ctx, string(plugin.rm.Resource()), os.Getenv(util.NodeNameEnvName)) + if detectErr != nil { + // Startup reset is destructive. If Kubernetes allocation state + // cannot be read reliably, preserve every GPU and retry cleanup + // later through the normal reclaim watcher. + klog.InfoS("mig init: allocation detection failed; preserving all GPUs", "err", detectErr) + for i := 0; i < deviceNumbers; i++ { + inUse[i] = struct{}{} + } } - err = RemoveMigApplyLock() + reset, err := plugin.migMgr.ResetIdleGPUs(deviceNumbers, inUse) if err != nil { - klog.Fatalf("RemoveMigApplyLock failed: %v", err) + klog.InfoS("mig init: failed to reset idle GPUs", "err", err) + } + klog.InfoS("mig init: resolved startup layout", + "inUseGPUs", sortedIntSetKeys(inUse), + "resetGPUs", reset) + if err := plugin.primeMigManagerFromAnnotations(deviceNames); err != nil { + klog.InfoS("mig init: failed to adopt active MIG allocations", "err", err) + } + } + // New allocations are realized lazily from the scheduler's exact + // profile and placement reservation. The watcher's release callback + // tolerates an unknown MIG UUID during early kubelet startup. + watcher := podresources.NewWatcher("", 0, []string{string(plugin.rm.Resource())}, func(_ string, deviceID string) { + if strings.HasPrefix(deviceID, "MIG-") { + if err := plugin.migMgr.Release(deviceID); err != nil { + klog.InfoS("failed to release MIG instance on reclaim", "deviceID", deviceID, "err", err) + } + } + if err := plugin.reconcileActiveMigAllocations(); err != nil { + klog.InfoS("failed to reconcile MIG allocations on reclaim", "deviceID", deviceID, "err", err) } + }) + go watcher.Run(plugin.ctx) + go plugin.runMigAnnotationReconciler(5 * time.Second) + } - cmd := exec.Command("nvidia-mig-parted", "export") - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - err := cmd.Run() - if err != nil { - klog.Errorf("nvidia-mig-parted failed: %v (stderr: %s)", err, stderr.String()) - klog.Warning("Falling back to non‑MIG configuration") - } else { - outStr := stdout.Bytes() - yaml.Unmarshal(outStr, &plugin.migCurrent) - writeMigConfig(outStr) + return nil +} - HamiInitMigConfig, err := plugin.processMigConfigs(plugin.migCurrent.MigConfigs, deviceNumbers) - if err != nil { - klog.Infof("no device in node: %v", err) - } else { - plugin.migCurrent.MigConfigs["current"] = HamiInitMigConfig - migApplied = true - } +func activeMigAllocationKeys(pods []corev1.Pod) (map[migAllocationKey]struct{}, error) { + active := make(map[migAllocationKey]struct{}) + for i := range pods { + pod := &pods[i] + if pod.DeletionTimestamp != nil || pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + continue + } + allocations, err := nvidia.DecodeMigAllocations(pod.Annotations[nvidia.MigAllocationsAnnotation]) + if err != nil { + return nil, fmt.Errorf("decode MIG allocations for pod %s/%s: %w", pod.Namespace, pod.Name, err) + } + for _, allocation := range allocations { + gpuIndex, ok := gpuUUIDToIndex(allocation.GPUUUID) + if !ok { + return nil, fmt.Errorf("resolve active MIG parent GPU %q", allocation.GPUUUID) } + active[allocationKey(gpuIndex, allocation.Profile, nvml.GpuInstancePlacement{Start: allocation.Placement.Start, Size: allocation.Placement.Size})] = struct{}{} } } + return active, nil +} - if !migApplied { - plugin.buildFallbackMigConfig(deviceNumbers) - klog.Infoln("Using non‑MIG configuration") +func (plugin *NvidiaDevicePlugin) annotateMigRuntimeInfo(pod *corev1.Pod) error { + allocations, err := nvidia.DecodeMigAllocations(pod.Annotations[nvidia.MigAllocationsAnnotation]) + if err != nil { + return err } - - go func() { - err := plugin.rm.CheckHealth(plugin.stop, plugin.health, plugin.disableHealthChecks, plugin.ackDisableHealthChecks) - if err != nil { - klog.Infof("Failed to start health check: %v; continuing with health checks disabled", err) + updated := false + for i := range allocations { + gpuIndex, ok := gpuUUIDToIndex(allocations[i].GPUUUID) + if !ok { + return fmt.Errorf("resolve MIG parent GPU %q", allocations[i].GPUUUID) } - }() + info, ok := plugin.migMgr.AllocationRuntimeInfo(gpuIndex, allocations[i].Profile, nvml.GpuInstancePlacement{Start: allocations[i].Placement.Start, Size: allocations[i].Placement.Size}) + if !ok { + continue // A later Allocate call may own another container's allocation. + } + allocations[i].MigUUID = info.MigUUID + updated = true + } + if !updated { + return errors.New("no realised MIG allocation found for pod") + } + raw, err := json.Marshal(allocations) + if err != nil { + return err + } + if err := util.PatchPodAnnotations(pod, map[string]string{nvidia.MigAllocationsAnnotation: string(raw)}); err != nil { + return err + } + pod.Annotations[nvidia.MigAllocationsAnnotation] = string(raw) + return nil +} - go func() { - plugin.WatchAndRegister(plugin.disableWatchAndRegister, plugin.ackDisableWatchAndRegister) - }() +func (plugin *NvidiaDevicePlugin) reconcileActiveMigAllocations() error { + active, err := plugin.listActiveMigAllocationKeys() + if err != nil { + return err + } + return plugin.migMgr.ReconcileActiveAllocations(active) +} - if migApplied { - plugin.ApplyMigTemplate() +func (plugin *NvidiaDevicePlugin) listActiveMigAllocationKeys() (map[migAllocationKey]struct{}, error) { + pods, err := client.GetClient().CoreV1().Pods("").List(plugin.ctx, metav1.ListOptions{ + FieldSelector: "spec.nodeName=" + os.Getenv(util.NodeNameEnvName), + }) + if err != nil { + return nil, err } + return activeMigAllocationKeys(pods.Items) +} +func (plugin *NvidiaDevicePlugin) primeMigManagerFromAnnotations(_ []string) error { + pods, err := client.GetClient().CoreV1().Pods("").List(plugin.ctx, metav1.ListOptions{ + FieldSelector: "spec.nodeName=" + os.Getenv(util.NodeNameEnvName), + }) + if err != nil { + return err + } + for i := range pods.Items { + pod := &pods.Items[i] + if pod.DeletionTimestamp != nil || pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + continue + } + allocations, err := nvidia.DecodeMigAllocations(pod.Annotations[nvidia.MigAllocationsAnnotation]) + if err != nil { + return err + } + for _, allocation := range allocations { + if allocation.MigUUID == "" { + return fmt.Errorf("active pod %s/%s MIG allocation lacks runtime UUID", pod.Namespace, pod.Name) + } + gpuIndex, ok := gpuUUIDToIndex(allocation.GPUUUID) + if !ok { + return fmt.Errorf("resolve active MIG parent GPU %q", allocation.GPUUUID) + } + if err := plugin.migMgr.AdoptAllocation(gpuIndex, allocation.Profile, allocation.MigUUID, nvml.GpuInstancePlacement{Start: allocation.Placement.Start, Size: allocation.Placement.Size}); err != nil { + return fmt.Errorf("adopt pod %s/%s MIG allocation: %w", pod.Namespace, pod.Name, err) + } + } + } return nil } +func (plugin *NvidiaDevicePlugin) runMigAnnotationReconciler(interval time.Duration) { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-plugin.ctx.Done(): + return + case <-ticker.C: + if err := plugin.reconcileActiveMigAllocations(); err != nil { + // Reconciliation is destructive, so API or annotation errors are + // fail-closed and leave current MIG instances untouched. + klog.InfoS("periodic MIG reconciliation skipped", "err", err) + } + } + } +} + // Stop stops the gRPC server. func (plugin *NvidiaDevicePlugin) Stop() error { if plugin == nil || plugin.server == nil { @@ -565,10 +699,6 @@ func physicalDeviceID(id string) string { if strings.Contains(id, "::") { return rm.AnnotatedID(id).GetID() } - // Handle MIG format: GPU-UUID[tidx-idx] -> GPU-UUID - if bracketIdx := strings.Index(id, "["); bracketIdx != -1 { - return id[:bracketIdx] - } // Handle virtual device format: GPU-UUID-N -> GPU-UUID // NVIDIA GPU UUID has exactly 5 dashes (GPU-xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx) // Virtual devices append "-N" suffix, resulting in exactly 6 dashes @@ -598,6 +728,12 @@ func (plugin *NvidiaDevicePlugin) alignContainerDevicesWithAllocatedIDs(devreq d // Allocate which return list of devices. func (plugin *NvidiaDevicePlugin) Allocate(ctx context.Context, reqs *kubeletdevicepluginv1beta1.AllocateRequest) (*kubeletdevicepluginv1beta1.AllocateResponse, error) { + // Kubelet may issue Allocate calls concurrently. The pending-pod + // annotation protocol and dynamic MIG preparation are node-global, so keep + // pod selection, GI/CI creation, and annotation consumption atomic. + plugin.applyMutex.Lock() + defer plugin.applyMutex.Unlock() + klog.InfoS("Allocate", "request", reqs) responses := kubeletdevicepluginv1beta1.AllocateResponse{} nodename := os.Getenv(util.NodeNameEnvName) @@ -659,7 +795,18 @@ func (plugin *NvidiaDevicePlugin) Allocate(ctx context.Context, reqs *kubeletdev } devreq = alignedDevreq } - response, err := plugin.getAllocateResponse(plugin.GetContainerDeviceStrArray(devreq)) + requestIDs, err := plugin.GetContainerDeviceStrArray(devreq, current, currentCtr.Name) + if err != nil { + PodAllocationFailed(nodename, current, NodeLockNvidia) + return nil, fmt.Errorf("resolve allocated NVIDIA devices: %w", err) + } + if plugin.operatingMode == "mig" { + if err := plugin.annotateMigRuntimeInfo(current); err != nil { + PodAllocationFailed(nodename, current, NodeLockNvidia) + return nil, fmt.Errorf("record MIG runtime placement: %w", err) + } + } + response, err := plugin.getAllocateResponse(requestIDs) if err != nil { PodAllocationFailed(nodename, current, NodeLockNvidia) return nil, fmt.Errorf("failed to get allocate response: %v", err) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go index 482d0efad7..969d3d8eb2 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go @@ -697,9 +697,7 @@ func TestGetPreferredAllocationSkipsEmptyAnnotations(t *testing.T) { require.ElementsMatch(t, []string{"GPU-03f69c50-207a-2038-9b45-23cac89cb67a-0", "GPU-03f69c50-207a-2038-9b45-23cac89cb67b-0"}, response.ContainerResponses[0].DeviceIDs) } -func TestPhysicalDeviceIDHandlesMIGFormat(t *testing.T) { - // Use real NVIDIA GPU UUID format: GPU-xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx (5 dashes) - // Virtual devices have 6 dashes: GPU-xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx-N +func TestPhysicalDeviceIDHandlesVirtualFormats(t *testing.T) { tests := []struct { input string expected string @@ -707,10 +705,6 @@ func TestPhysicalDeviceIDHandlesMIGFormat(t *testing.T) { // Virtual device format (6 dashes) {"GPU-03f69c50-207a-2038-9b45-23cac89cb67a-0", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, {"GPU-03f69c50-207a-2038-9b45-23cac89cb67a-10", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, - // MIG format with template index - {"GPU-03f69c50-207a-2038-9b45-23cac89cb67a[0-1]", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, - {"GPU-03f69c50-207a-2038-9b45-23cac89cb67a[1-2]", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, - // Replica format {"GPU-03f69c50-207a-2038-9b45-23cac89cb67a::replica-1", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, // Plain UUID (5 dashes, should not be modified) {"GPU-03f69c50-207a-2038-9b45-23cac89cb67a", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, @@ -726,18 +720,17 @@ func TestPhysicalDeviceIDHandlesMIGFormat(t *testing.T) { } } -func TestSelectPreferredDeviceIDsWithMIGUUIDs(t *testing.T) { +func TestSelectPreferredDeviceIDsWithPhysicalMIGReservations(t *testing.T) { plugin := &NvidiaDevicePlugin{} - // Use real NVIDIA GPU UUID format available := []string{ "GPU-03f69c50-207a-2038-9b45-23cac89cb67a-0", "GPU-03f69c50-207a-2038-9b45-23cac89cb67a-1", "GPU-03f69c50-207a-2038-9b45-23cac89cb67b-0", "GPU-03f69c50-207a-2038-9b45-23cac89cb67c-0", } desired := device.ContainerDevices{ - {UUID: "GPU-03f69c50-207a-2038-9b45-23cac89cb67a[0-1]"}, // MIG format + {UUID: "GPU-03f69c50-207a-2038-9b45-23cac89cb67a"}, {UUID: "GPU-03f69c50-207a-2038-9b45-23cac89cb67b"}, - {UUID: "GPU-03f69c50-207a-2038-9b45-23cac89cb67c[1-2]"}, // MIG format with different index + {UUID: "GPU-03f69c50-207a-2038-9b45-23cac89cb67c"}, } got, err := plugin.selectPreferredDeviceIDsFromAnnotatedDevices(available, nil, desired, 3) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go index 2d64448a7d..a4b2a2e988 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go @@ -17,7 +17,6 @@ package plugin import ( - "bytes" "context" "encoding/json" "errors" @@ -25,13 +24,10 @@ import ( "os" "os/exec" "path/filepath" - "strconv" + "sort" "strings" - "time" "github.com/NVIDIA/go-nvml/pkg/nvml" - spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" - "gopkg.in/yaml.v3" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/klog/v2" @@ -172,73 +168,6 @@ func GetIndexAndTypeFromUUID(uuid string) (string, int) { return Model, index } -func GetMigUUIDFromSmiOutput(output string, uuid string, idx int) string { - migmode := false - for val := range strings.SplitSeq(output, "\n") { - if !strings.Contains(val, "MIG") && strings.Contains(val, uuid) { - migmode = true - continue - } - if !strings.Contains(val, "MIG") && !strings.Contains(val, uuid) { - migmode = false - continue - } - if !migmode { - continue - } - klog.Infoln("inspecting", val) - num := strings.Split(val, "Device")[1] - num = strings.Split(num, ":")[0] - num = strings.TrimSpace(num) - index, err := strconv.Atoi(num) - if err != nil { - klog.Fatal("atoi failed num=", num) - } - if index == idx { - outputStr := strings.Split(val, ":")[2] - outputStr = strings.TrimSpace(outputStr) - outputStr = strings.TrimRight(outputStr, ")") - return outputStr - } - } - return "" -} - -func GetMigUUIDFromIndex(uuid string, idx int) string { - defer nvml.Shutdown() - if nvret := nvml.Init(); nvret != nvml.SUCCESS { - klog.Errorln("nvml Init err: ", nvret) - panic(0) - } - originuuid := strings.Split(uuid, "[")[0] - ndev, ret := nvml.DeviceGetHandleByUUID(originuuid) - if ret != nvml.SUCCESS { - klog.Error(`nvml get device uuid error ret=`, ret) - panic(0) - } - migdev, ret := nvml.DeviceGetMigDeviceHandleByIndex(ndev, idx) - if ret != nvml.SUCCESS { - klog.Error("nvml get mig dev error ret=", ret, ",idx=", idx, "using nvidia-smi -L for query") - cmd := exec.Command("nvidia-smi", "-L") - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - err := cmd.Run() - if err != nil { - klog.Fatalf("nvidia-smi -L failed with %s\n", err) - } - outStr := stdout.String() - uuid := GetMigUUIDFromSmiOutput(outStr, originuuid, idx) - return uuid - } - res, ret := migdev.GetUUID() - if ret != nvml.SUCCESS { - klog.Error(`nvml get mig uuid error ret=`, ret) - panic(0) - } - return res -} - func GetMigGpuInstanceIdFromIndex(uuid string, idx int) (int, error) { if nvret := nvml.Init(); nvret != nvml.SUCCESS { klog.Errorln("nvml Init err: ", nvret) @@ -349,104 +278,6 @@ func (nv *NvidiaDevicePlugin) EnableOtherNVMLOperation() { nv.disableWatchAndRegister <- false } -// migConfigPath is where nvidia-mig-parted reads the rendered MIG config. -// Variable rather than const so tests can redirect the write. -var migConfigPath = "/tmp/migconfig.yaml" - -// writeMigConfig persists the rendered MIG config with owner-only permissions. -// On write failure any previous file is removed, so a later -// nvidia-mig-parted apply fails on a missing file instead of silently -// applying a stale config. -func writeMigConfig(data []byte) { - if err := os.WriteFile(migConfigPath, data, 0o600); err != nil { - klog.Errorf("failed to write %s: %v", migConfigPath, err) - _ = os.Remove(migConfigPath) - return - } - klog.V(4).InfoS("wrote MIG config", "path", migConfigPath, "bytes", len(data)) -} - -func (nv *NvidiaDevicePlugin) ApplyMigTemplate() { - nv.applyMutex.Lock() - nv.DisableOtherNVMLOperation() - defer func() { - nv.EnableOtherNVMLOperation() - nv.applyMutex.Unlock() - }() - data, err := yaml.Marshal(nv.migCurrent) - if err != nil { - klog.Error("marshal failed", err.Error()) - } - klog.Infoln("Applying data=", string(data)) - writeMigConfig(data) - cmd := exec.Command("nvidia-mig-parted", "apply", "-f", migConfigPath) - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - err = cmd.Run() - if err != nil { - klog.Fatalf("nvidia-mig-parted failed with %s,reason:%s\n", err, stderr.String()) - } - outStr := stdout.String() - klog.Infoln("Mig apply", outStr) -} - -func (nv *NvidiaDevicePlugin) GenerateMigTemplate(devtype string, devindex int, val device.ContainerDevice) (int, bool) { - needsreset := false - position := -1 // Initialize to an invalid position - - for _, migTemplate := range nv.schedulerConfig.MigGeometriesList { - if containsModel(devtype, migTemplate.Models) { - klog.InfoS("type found", "Type", devtype, "Models", strings.Join(migTemplate.Models, ", ")) - - templateIdx, pos, err := device.ExtractMigTemplatesFromUUID(val.UUID) - if err != nil { - klog.ErrorS(err, "failed to extract template index from UUID", "UUID", val.UUID) - return -1, false - } - position = pos - - if templateIdx < 0 || templateIdx >= len(migTemplate.Geometries) { - klog.ErrorS(nil, "invalid template index extracted from UUID", "UUID", val.UUID, "Index", templateIdx) - return -1, false - } - - v := migTemplate.Geometries[templateIdx] - - for migidx, migpartedDev := range nv.migCurrent.MigConfigs["current"] { - if containsDevice(devindex, migpartedDev.Devices) { - for _, migTemplateEntry := range v { - currentCount, ok := migpartedDev.MigDevices[migTemplateEntry.Name] - expectedCount := migTemplateEntry.Count - - if !ok || currentCount != expectedCount { - needsreset = true - klog.InfoS("updated mig device count", "Template", v) - } else { - klog.InfoS("incremented mig device count", "TemplateName", migTemplateEntry.Name, "Count", currentCount+1) - } - } - - if needsreset { - for k := range nv.migCurrent.MigConfigs["current"][migidx].MigDevices { - delete(nv.migCurrent.MigConfigs["current"][migidx].MigDevices, k) - } - - for _, migTemplateEntry := range v { - nv.migCurrent.MigConfigs["current"][migidx].MigDevices[migTemplateEntry.Name] = migTemplateEntry.Count - nv.migCurrent.MigConfigs["current"][migidx].MigEnabled = true - } - } - break - } - } - break - } - } - - return position, needsreset -} - // Helper function to check if a model is in the list of models. func containsModel(target string, models []string) bool { for _, model := range models { @@ -457,16 +288,6 @@ func containsModel(target string, models []string) bool { return false } -// Helper function to check if a device index is in the list of devices. -func containsDevice(target int, devices []int32) bool { - for _, device := range devices { - if int(device) == target { - return true - } - } - return false -} - // Helper function to deepcopy new mig spec func deepCopyMigConfig(src nvidia.MigConfigSpec) nvidia.MigConfigSpec { dst := src @@ -483,54 +304,68 @@ func deepCopyMigConfig(src nvidia.MigConfigSpec) nvidia.MigConfigSpec { return dst } -func (nv *NvidiaDevicePlugin) GetContainerDeviceStrArray(c device.ContainerDevices) []string { - tmp := []string{} - needsreset := false - position := 0 - for _, val := range c { - if !strings.Contains(val.UUID, "[") { - tmp = append(tmp, val.UUID) - } else { - devtype, devindex := GetIndexAndTypeFromUUID(val.UUID) - position, needsreset = nv.GenerateMigTemplate(devtype, devindex, val) - if needsreset { - nv.ApplyMigTemplate() - if nv.deviceListStrategies.Includes(spec.DeviceListStrategyVolumeMounts) || - nv.deviceListStrategies.Includes(spec.DeviceListStrategyCDIAnnotations) || - nv.deviceListStrategies.Includes(spec.DeviceListStrategyCDICRI) { - klog.V(3).Infoln("generate CDI spec file") - const ( - maxTryTimes = 5 - waitTimeInterval = 5 * time.Second - specFilePath = "/var/run/cdi/k8s.device-plugin.nvidia.com-gpu.json" - kind = "k8s.device-plugin.nvidia.com/gpu" - ) - for i := 0; i < maxTryTimes; i++ { - if err := createSpecFile(specFilePath); err != nil { - klog.Warningf("failed to create CDI spec file: %v", err) - } else { - klog.Infof("createSpecFile ok. file path %s", specFilePath) - } - if err := checkCDISpecFile(specFilePath, kind); err != nil { - klog.Warningf("check CDI spec file failed. %v", err) - if i == maxTryTimes-1 { - klog.Fatalf("exceed the max trytime %d", maxTryTimes) - } else { - time.Sleep(waitTimeInterval) - klog.Warningf("try to create CDI spec file again. try times: %d", i) - continue - } - } - klog.Infof("check CDI spec file ok") - break - } - } +func (nv *NvidiaDevicePlugin) GetContainerDeviceStrArray(c device.ContainerDevices, pod *corev1.Pod, containerName string) ([]string, error) { + if nv.operatingMode != "mig" { + out := make([]string, 0, len(c)) + for _, value := range c { + out = append(out, value.UUID) + } + return out, nil + } + if nv.migMgr == nil || pod == nil { + return nil, fmt.Errorf("MIG allocation requires an initialized manager and Pod") + } + containerIndex := -1 + for i := range pod.Spec.InitContainers { + if pod.Spec.InitContainers[i].Name == containerName { + containerIndex = i + break + } + } + if containerIndex < 0 { + for i := range pod.Spec.Containers { + if pod.Spec.Containers[i].Name == containerName { + containerIndex = len(pod.Spec.InitContainers) + i + break } - tmp = append(tmp, GetMigUUIDFromIndex(val.UUID, position)) } } - klog.V(3).Infoln("mig current=", nv.migCurrent, ":", needsreset, "position=", position, "uuid lists", tmp) - return tmp + if containerIndex < 0 { + return nil, fmt.Errorf("container %q not found in Pod", containerName) + } + allocations, err := nvidia.DecodeMigAllocations(pod.Annotations[nvidia.MigAllocationsAnnotation]) + if err != nil { + return nil, fmt.Errorf("decode scheduler MIG allocations: %w", err) + } + containerAllocations := make([]nvidia.MigAllocation, 0, len(c)) + for _, allocation := range allocations { + if allocation.ContainerIndex == containerIndex { + containerAllocations = append(containerAllocations, allocation) + } + } + sort.Slice(containerAllocations, func(i, j int) bool { return containerAllocations[i].DeviceIndex < containerAllocations[j].DeviceIndex }) + if len(containerAllocations) != len(c) { + return nil, fmt.Errorf("container %s has %d MIG reservations, requested %d devices", containerName, len(containerAllocations), len(c)) + } + if err := nv.reconcileActiveMigAllocations(); err != nil { + return nil, fmt.Errorf("reconcile MIG allocations before allocation: %w", err) + } + out := make([]string, 0, len(c)) + for i, reservation := range containerAllocations { + if reservation.GPUUUID != c[i].UUID { + return nil, fmt.Errorf("MIG reservation GPU %s does not match allocated GPU %s", reservation.GPUUUID, c[i].UUID) + } + gpuIndex, ok := gpuUUIDToIndex(reservation.GPUUUID) + if !ok { + return nil, fmt.Errorf("resolve parent GPU %s", reservation.GPUUUID) + } + migUUID, err := nv.migMgr.EnsureAllocation(gpuIndex, reservation.Profile, nvml.GpuInstancePlacement{Start: reservation.Placement.Start, Size: reservation.Placement.Size}) + if err != nil { + return nil, err + } + out = append(out, migUUID) + } + return out, nil } var podAllocationTrySuccess = func(nodeName string, devName string, lockName string, pod *corev1.Pod) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go index e2055c13cd..ea0eaac0ff 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go @@ -29,168 +29,11 @@ import ( "tags.cncf.io/container-device-interface/specs-go" "github.com/Project-HAMi/HAMi/pkg/device" - "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/util" "github.com/Project-HAMi/HAMi/pkg/util/client" "github.com/Project-HAMi/HAMi/pkg/util/nodelock" ) -func TestGenerateMigTemplate(t *testing.T) { - sconfig := nvidia.NvidiaConfig{ - MigGeometriesList: []device.AllowedMigGeometries{ - { - Models: []string{"A30"}, - Geometries: []device.Geometry{ - {device.MigTemplate{Name: "1g.6gb", Core: 25, Memory: 6144, Count: 4}}, - {device.MigTemplate{Name: "2g.12gb", Core: 50, Memory: 12288, Count: 2}}, - {device.MigTemplate{Name: "4g.24gb", Core: 100, Memory: 24576, Count: 1}}, - }, - }, - { - Models: []string{"A100-SXM4-40GB", "A100-40GB-PCIe", "A100-PCIE-40GB", "A100-SXM4-40GB"}, - Geometries: []device.Geometry{ - {device.MigTemplate{Name: "1g.5gb", Core: 14, Memory: 5120, Count: 7}}, - {device.MigTemplate{Name: "1g.5gb", Core: 14, Memory: 5120, Count: 1}, device.MigTemplate{Name: "2g.10gb", Core: 28, Memory: 10240, Count: 3}}, - {device.MigTemplate{Name: "3g.20gb", Core: 42, Memory: 20480, Count: 2}}, - {device.MigTemplate{Name: "7g.40gb", Core: 100, Memory: 40960, Count: 1}}, - }, - }, - { - Models: []string{"A100-SXM4-80GB", "A100-80GB-PCIe", "A100-PCIE-80GB"}, - Geometries: []device.Geometry{ - {device.MigTemplate{Name: "1g.10gb", Core: 14, Memory: 10240, Count: 7}}, - {device.MigTemplate{Name: "1g.10gb", Core: 14, Memory: 10240, Count: 1}, device.MigTemplate{Name: "2g.20gb", Core: 28, Memory: 20480, Count: 3}}, - {device.MigTemplate{Name: "3g.40gb", Core: 42, Memory: 40960, Count: 2}}, - {device.MigTemplate{Name: "7g.80gb", Core: 100, Memory: 81920, Count: 1}}, - }, - }, - { - Models: []string{"RTX PRO 6000 Blackwell Server Edition"}, - Geometries: []device.Geometry{ - {device.MigTemplate{Name: "1g.24gb", Core: 25, Memory: 24576, Count: 4}}, - {device.MigTemplate{Name: "2g.48gb", Core: 50, Memory: 49152, Count: 2}}, - {device.MigTemplate{Name: "4g.96gb", Core: 100, Memory: 98304, Count: 1}}, - }, - }, - }, - } - - plugin := NvidiaDevicePlugin{ - operatingMode: "mig", - schedulerConfig: sconfig, - } - plugin.migCurrent = nvidia.MigPartedSpec{ - Version: "v1", - MigConfigs: make(map[string]nvidia.MigConfigSpecSlice), - } - plugin.migCurrent.MigConfigs["current"] = nvidia.MigConfigSpecSlice{ - nvidia.MigConfigSpec{ - Devices: []int32{0, 1}, - MigEnabled: true, - MigDevices: make(map[string]int32), // Ensure this map is initialized - }, - } - - testCases := []struct { - name string - model string - deviceIdx int - containerDev device.ContainerDevice - expectedPos int - expectedReset bool - expectedMig map[string]int32 - }{ - { - name: "2g.10gb template", - model: "A100-SXM4-40GB", - deviceIdx: 0, - containerDev: device.ContainerDevice{ - Idx: 0, - UUID: "aaaaabbbb[1-1]", - Usedmem: 8000, - }, - expectedPos: 1, - expectedReset: true, - expectedMig: map[string]int32{ - "1g.5gb": 1, - "2g.10gb": 3, - }, - }, - { - name: "1g.5gb template", - model: "A100-SXM4-40GB", - deviceIdx: 0, - containerDev: device.ContainerDevice{ - Idx: 0, - UUID: "aaaaabbbb[0-1]", - Usedmem: 3000, - }, - expectedPos: 1, - expectedReset: true, - expectedMig: map[string]int32{ - "1g.5gb": 7, - }, - }, - { - name: "no reset needed", - model: "A100-SXM4-40GB", - deviceIdx: 0, - containerDev: device.ContainerDevice{ - Idx: 0, - UUID: "aaaaabbbb[0-2]", - Usedmem: 3000, - }, - expectedPos: 2, - expectedReset: false, - expectedMig: map[string]int32{ - "1g.5gb": 7, - }, - }, - { - // The full NVML model string must match the shorter configured - // model via substring matching (RTX PRO 6000 Blackwell Server Edition). - name: "rtx pro 6000 blackwell 1g.24gb template", - model: "NVIDIA RTX PRO 6000 Blackwell Server Edition", - deviceIdx: 0, - containerDev: device.ContainerDevice{ - Idx: 0, - UUID: "ccccdddd[0-3]", - Usedmem: 20000, - }, - expectedPos: 3, - expectedReset: true, - expectedMig: map[string]int32{ - "1g.24gb": 4, - }, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - pos, needsreset := plugin.GenerateMigTemplate(tc.model, tc.deviceIdx, tc.containerDev) - - // Check if the position matches the expected value - if pos != tc.expectedPos { - t.Errorf("expected position %d, got %d", tc.expectedPos, pos) - } - - // Check if the reset flag matches the expected value - if needsreset != tc.expectedReset { - t.Errorf("expected reset %v, got %v", tc.expectedReset, needsreset) - } - - // Check if the mig devices match the expected values - migDevices := plugin.migCurrent.MigConfigs["current"][0].MigDevices - for k, v := range tc.expectedMig { - actual, ok := migDevices[k] - if !ok || actual != v { - t.Errorf("expected %s count %d, got %d", k, v, actual) - } - } - }) - } -} - func TestGetNextDeviceRequest_DeviceInRegularContainer(t *testing.T) { // Save and restore InRequestDevices oldInRequestDevices := device.InRequestDevices diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go b/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go new file mode 100644 index 0000000000..b15d1bbd51 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go @@ -0,0 +1,264 @@ +/* + * Copyright (c) 2026, HAMi. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + */ + +// Package podresources polls the kubelet pod-resources API and reports when +// devices that were previously allocated to a container disappear from the +// kubelet's view, which happens when the container ends. +// +// This is the signal HAMi's MIG reclaim path uses to destroy individual GPU +// instances on task completion instead of waiting for every task on a GPU to +// finish before the next nvidia-mig-parted apply can re-shape the card. +package podresources + +import ( + "context" + "net" + "strings" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/credentials/insecure" + "k8s.io/klog/v2" + podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" +) + +const ( + defaultSocketPath = "/var/lib/kubelet/pod-resources/kubelet.sock" + defaultPollEvery = 10 * time.Second + // dialWait bounds how long we wait for the kubelet-side socket to reach + // Ready before the first List after (re)connect. The connection itself + // is created lazily by grpc.NewClient; this timeout only gates the + // explicit wait performed before the first RPC on a fresh conn. + dialWait = 5 * time.Second + // listTimeout bounds a single List RPC. This is intentionally smaller + // than the poll interval so a slow kubelet can't cause back-to-back + // ticks to pile up. On very busy nodes with many pods, raising this + // (and the poll interval) is the right lever. + listTimeout = 5 * time.Second + // maxMsgSize matches the kubelet default for this API; on busy nodes the + // List response can exceed the grpc default 4 MiB. + maxMsgSize = 16 * 1024 * 1024 +) + +// ReleaseHandler is invoked once per deviceID that was present in the previous +// snapshot and is missing in the current one. resourceName is the resource the +// device was allocated under (e.g. "nvidia.com/gpu") so a handler serving +// several resources can filter quickly. +type ReleaseHandler func(resourceName, deviceID string) + +// Watcher polls kubelet's pod-resources API and fires a ReleaseHandler when a +// previously-allocated device is no longer in use by any container. +type Watcher struct { + socketPath string + interval time.Duration + + // resourceNames restricts the release callback to devices under these + // resource names. An empty slice means "all resources". + resourceNames []string + + onRelease ReleaseHandler + + // Previous snapshot: set of active deviceIDs indexed by resource name. + prev map[string]map[string]struct{} + + // Long-lived gRPC client; rebuilt on failure. kubelet's pod-resources + // socket is stable across the plugin's lifetime, so reusing the + // connection avoids eating our per-tick budget on DNS/dial/handshake. + conn *grpc.ClientConn + client podresourcesv1.PodResourcesListerClient +} + +// NewWatcher constructs a Watcher. Pass an empty socketPath to use the default +// kubelet location; pass 0 interval to use the default poll cadence. +func NewWatcher(socketPath string, interval time.Duration, resourceNames []string, onRelease ReleaseHandler) *Watcher { + if socketPath == "" { + socketPath = defaultSocketPath + } + if interval <= 0 { + interval = defaultPollEvery + } + return &Watcher{ + socketPath: socketPath, + interval: interval, + resourceNames: resourceNames, + onRelease: onRelease, + prev: make(map[string]map[string]struct{}), + } +} + +// Run polls the kubelet pod-resources API in a loop until ctx is cancelled. +// It never returns an error; transient gRPC failures are logged and the +// previous snapshot is preserved so a missed tick doesn't produce spurious +// release events. +func (w *Watcher) Run(ctx context.Context) { + klog.InfoS("starting podresources watcher", "socket", w.socketPath, "interval", w.interval, "resources", w.resourceNames) + defer w.closeConn() + + // Prime the snapshot before starting to diff. If this first call fails + // we start with an empty map and the first successful tick will just + // record — no spurious release events for the state at plugin start. + if err := w.tick(ctx, true); err != nil { + klog.InfoS("podresources initial List failed; will retry on interval", "err", err) + } + + t := time.NewTicker(w.interval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if err := w.tick(ctx, false); err != nil { + klog.InfoS("podresources tick failed; keeping previous snapshot", "err", err) + } + } + } +} + +func (w *Watcher) tick(ctx context.Context, prime bool) error { + if err := w.ensureConn(ctx); err != nil { + return err + } + + callCtx, cancel := context.WithTimeout(ctx, listTimeout) + defer cancel() + + resp, err := w.client.List(callCtx, &podresourcesv1.ListPodResourcesRequest{}) + if err != nil { + // Drop the connection so the next tick reconnects; the kubelet + // socket can be recreated (e.g. kubelet restart) without the + // plugin restarting, and sticking to a dead conn wastes ticks. + w.closeConn() + return err + } + + current := w.collect(resp) + if !prime { + w.diff(current) + } + w.prev = current + return nil +} + +// ensureConn makes sure w.conn/w.client are usable, creating them lazily. +// The first call after (re)connect waits up to dialWait for the channel to +// reach Ready so the subsequent List has the full listTimeout budget. +func (w *Watcher) ensureConn(ctx context.Context) error { + if w.conn != nil { + switch w.conn.GetState() { + case connectivity.Shutdown: + w.closeConn() + default: + return nil + } + } + + conn, err := grpc.NewClient( + "unix://"+w.socketPath, + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxMsgSize)), + grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, "unix", normalizeUnixDialAddr(addr)) + }), + ) + if err != nil { + return err + } + + // Kick the channel and wait briefly for Ready. If the kubelet socket + // isn't accepting yet, let the caller surface a clean error instead + // of eating the full List budget on handshake. + conn.Connect() + waitCtx, cancel := context.WithTimeout(ctx, dialWait) + defer cancel() + for { + s := conn.GetState() + if s == connectivity.Ready { + break + } + if !conn.WaitForStateChange(waitCtx, s) { + _ = conn.Close() + return waitCtx.Err() + } + } + + w.conn = conn + w.client = podresourcesv1.NewPodResourcesListerClient(conn) + return nil +} + +func normalizeUnixDialAddr(addr string) string { + trimmed := strings.TrimPrefix(addr, "unix://") + if trimmed != addr { + return trimmed + } + return strings.TrimPrefix(addr, "unix:") +} + +func (w *Watcher) closeConn() { + if w.conn != nil { + _ = w.conn.Close() + w.conn = nil + w.client = nil + } +} + +// collect flattens the List response into "resourceName -> set(deviceID)". +func (w *Watcher) collect(resp *podresourcesv1.ListPodResourcesResponse) map[string]map[string]struct{} { + out := make(map[string]map[string]struct{}) + for _, pod := range resp.GetPodResources() { + for _, c := range pod.GetContainers() { + for _, d := range c.GetDevices() { + rn := d.GetResourceName() + if !w.resourceMatch(rn) { + continue + } + set, ok := out[rn] + if !ok { + set = make(map[string]struct{}) + out[rn] = set + } + for _, id := range d.GetDeviceIds() { + set[id] = struct{}{} + } + } + } + } + return out +} + +// diff fires onRelease for every deviceID that was present in the previous +// snapshot under a given resource but is absent in the current snapshot. +func (w *Watcher) diff(current map[string]map[string]struct{}) { + for rn, prevSet := range w.prev { + currSet := current[rn] + for id := range prevSet { + if _, stillUsed := currSet[id]; stillUsed { + continue + } + if w.onRelease != nil { + w.onRelease(rn, id) + } + } + } +} + +func (w *Watcher) resourceMatch(rn string) bool { + if len(w.resourceNames) == 0 { + return true + } + for _, want := range w.resourceNames { + if strings.EqualFold(rn, want) { + return true + } + } + return false +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go new file mode 100644 index 0000000000..d35e55f171 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go @@ -0,0 +1,161 @@ +package podresources + +import ( + "context" + "net" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "google.golang.org/grpc" + podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" +) + +type fakePodResourcesServer struct { + podresourcesv1.UnimplementedPodResourcesListerServer + mu sync.RWMutex + resp *podresourcesv1.ListPodResourcesResponse +} +func (s *fakePodResourcesServer) setResponse(resp *podresourcesv1.ListPodResourcesResponse) { + s.mu.Lock() + defer s.mu.Unlock() + s.resp = resp +} + +func (s *fakePodResourcesServer) List(context.Context, *podresourcesv1.ListPodResourcesRequest) (*podresourcesv1.ListPodResourcesResponse, error) { + s.mu.RLock() + defer s.mu.RUnlock() + return s.resp, nil +} + +func startFakePodResourcesServer(t *testing.T, resp *podresourcesv1.ListPodResourcesResponse) (socketPath string, setResp func(*podresourcesv1.ListPodResourcesResponse), stop func()) { + t.Helper() + dir := t.TempDir() + socketPath = filepath.Join(dir, "kubelet.sock") + l, err := net.Listen("unix", socketPath) + if err != nil { + t.Fatalf("listen unix socket: %v", err) + } + + srv := grpc.NewServer() + fake := &fakePodResourcesServer{resp: resp} + podresourcesv1.RegisterPodResourcesListerServer(srv, fake) + go func() { + _ = srv.Serve(l) + }() + + setResp = fake.setResponse + stop = func() { + srv.Stop() + _ = l.Close() + _ = os.Remove(socketPath) + } + return socketPath, setResp, stop +} + +func TestTickCanListPodResources(t *testing.T) { + initial := &podresourcesv1.ListPodResourcesResponse{ + PodResources: []*podresourcesv1.PodResources{ + { + Name: "p1", + Namespace: "default", + Containers: []*podresourcesv1.ContainerResources{ + { + Name: "c1", + Devices: []*podresourcesv1.ContainerDevices{ + {ResourceName: "nvidia.com/gpu", DeviceIds: []string{"MIG-a", "MIG-b"}}, + }, + }, + }, + }, + }, + } + socketPath, _, stop := startFakePodResourcesServer(t, initial) + defer stop() + + w := NewWatcher(socketPath, time.Second, []string{"nvidia.com/gpu"}, nil) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if err := w.tick(ctx, true); err != nil { + t.Fatalf("tick should succeed, got error: %v", err) + } + + got := w.prev["nvidia.com/gpu"] + if got == nil { + t.Fatalf("expected snapshot for nvidia.com/gpu") + } + if _, ok := got["MIG-a"]; !ok { + t.Fatalf("expected MIG-a in snapshot") + } + if _, ok := got["MIG-b"]; !ok { + t.Fatalf("expected MIG-b in snapshot") + } +} + +func TestTickDiffTriggersRelease(t *testing.T) { + resp1 := &podresourcesv1.ListPodResourcesResponse{ + PodResources: []*podresourcesv1.PodResources{ + { + Name: "p1", + Namespace: "default", + Containers: []*podresourcesv1.ContainerResources{ + { + Name: "c1", + Devices: []*podresourcesv1.ContainerDevices{ + {ResourceName: "nvidia.com/gpu", DeviceIds: []string{"MIG-a", "MIG-b"}}, + }, + }, + }, + }, + }, + } + resp2 := &podresourcesv1.ListPodResourcesResponse{ + PodResources: []*podresourcesv1.PodResources{ + { + Name: "p1", + Namespace: "default", + Containers: []*podresourcesv1.ContainerResources{ + { + Name: "c1", + Devices: []*podresourcesv1.ContainerDevices{ + {ResourceName: "nvidia.com/gpu", DeviceIds: []string{"MIG-a"}}, + }, + }, + }, + }, + }, + } + + socketPath, setResp, stop := startFakePodResourcesServer(t, resp1) + defer stop() + + var ( + mu sync.Mutex + released []string + ) + w := NewWatcher(socketPath, time.Second, []string{"nvidia.com/gpu"}, func(_ string, deviceID string) { + mu.Lock() + defer mu.Unlock() + released = append(released, deviceID) + }) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if err := w.tick(ctx, true); err != nil { + t.Fatalf("prime tick failed: %v", err) + } + setResp(resp2) + if err := w.tick(ctx, false); err != nil { + t.Fatalf("diff tick failed: %v", err) + } + + mu.Lock() + defer mu.Unlock() + if len(released) != 1 || released[0] != "MIG-b" { + t.Fatalf("expected one released device MIG-b, got %v", released) + } +} diff --git a/pkg/device/common/common.go b/pkg/device/common/common.go index 174be2bf1c..c13ede643c 100644 --- a/pkg/device/common/common.go +++ b/pkg/device/common/common.go @@ -33,6 +33,7 @@ const ( NumaNotFit = "NumaNotFit" ExclusiveDeviceAllocateConflict = "ExclusiveDeviceAllocateConflict" CardNotFoundCustomFilterRule = "CardNotFoundCustomFilterRule" + CardMigTopologyInfeasible = "CardMigTopologyInfeasible" NodeInsufficientDevice = "NodeInsufficientDevice" AllocatedCardsInsufficientRequest = "AllocatedCardsInsufficientRequest" NodeUnfitPod = "NodeUnfitPod" diff --git a/pkg/device/devices.go b/pkg/device/devices.go index 8435f0e42f..072abfa2ae 100644 --- a/pkg/device/devices.go +++ b/pkg/device/devices.go @@ -25,8 +25,6 @@ import ( "strings" "time" - "github.com/ccoveille/go-safecast/v2" - corev1 "k8s.io/api/core/v1" "k8s.io/klog/v2" @@ -49,6 +47,14 @@ type Devices interface { Fit(devices []*DeviceUsage, request ContainerDeviceRequest, pod *corev1.Pod, nodeInfo *NodeInfo, allocated *PodDevices) (bool, map[string]ContainerDevices, string) } +type MigPlacement struct { + Start uint32 `json:"start"` + Size uint32 `json:"size"` +} + +// Deprecated scheduler-template types are retained for source compatibility +// with non-NVIDIA device implementations and tests. Dynamic NVIDIA MIG does +// not populate or consume them. type MigTemplate struct { Name string `yaml:"name"` Core int32 `yaml:"core"` @@ -57,43 +63,65 @@ type MigTemplate struct { } type MigTemplateUsage struct { - Name string `json:"name,omitempty"` - Core int32 `json:"core,omitempty"` - Memory int32 `json:"memory,omitempty"` - InUse bool `json:"inuse,omitempty"` + Name string + Core int32 + Memory int32 + InUse bool } type Geometry []MigTemplate - type MIGS []MigTemplateUsage - type MigInUse struct { Index int32 UsageList MIGS } -type AllowedMigGeometries struct { - Models []string `yaml:"models"` - Geometries []Geometry `yaml:"allowedGeometries"` +type MigProfile struct { + Name string `json:"name"` + GIProfileID uint32 `json:"giProfileID"` + MemoryMB int32 `json:"memoryMB"` + Core int32 `json:"core"` + SliceCount uint32 `json:"sliceCount"` + InstanceCount uint32 `json:"instanceCount"` + MultiprocessorCount uint32 `json:"multiprocessorCount"` + P2PSupported bool `json:"p2pSupported"` + CopyEngineCount uint32 `json:"copyEngineCount"` + DecoderCount uint32 `json:"decoderCount"` + EncoderCount uint32 `json:"encoderCount"` + JPEGCount uint32 `json:"jpegCount"` + OFACount uint32 `json:"ofaCount"` + Placements []MigPlacement `json:"placements"` +} + +type MigAllocation struct { + Profile string `json:"profile"` + Placement MigPlacement `json:"placement"` +} + +type AllowedMigProfiles struct { + Models []string `yaml:"models"` + Profiles []string `yaml:"profiles"` } type DeviceUsage struct { - ID string - Index uint - Used int32 - Count int32 - Usedmem int32 - Totalmem int32 - Totalcore int32 - Usedcores int32 - Mode string - MigTemplate []Geometry - MigUsage MigInUse - Numa int - Type string - Health bool - PodInfos []*PodInfo - CustomInfo map[string]any + ID string + Index uint + Used int32 + Count int32 + Usedmem int32 + Totalmem int32 + Totalcore int32 + Usedcores int32 + Mode string + MigProfiles []MigProfile + MigAllocationsInUse []MigAllocation + MigTemplate []Geometry // Deprecated: unused by dynamic NVIDIA MIG. + MigUsage MigInUse // Deprecated: unused by dynamic NVIDIA MIG. + Numa int + Type string + Health bool + PodInfos []*PodInfo + CustomInfo map[string]any } type DeviceInfo struct { @@ -105,7 +133,8 @@ type DeviceInfo struct { Type string `json:"type,omitempty"` Numa int `json:"numa,omitempty"` Mode string `json:"mode,omitempty"` - MIGTemplate []Geometry `json:"migtemplate,omitempty"` + MIGProfiles []MigProfile `json:"migProfiles,omitempty"` + MIGTemplate []Geometry `json:"migtemplate,omitempty"` // Deprecated. Health bool `json:"health,omitempty"` DeviceVendor string `json:"devicevendor,omitempty"` CustomInfo map[string]any `json:"custominfo,omitempty"` @@ -121,10 +150,17 @@ type DevicePairScore struct { func (d DeviceInfo) DeepCopy() DeviceInfo { dup := d + if d.MIGProfiles != nil { + dup.MIGProfiles = make([]MigProfile, len(d.MIGProfiles)) + copy(dup.MIGProfiles, d.MIGProfiles) + for i := range dup.MIGProfiles { + dup.MIGProfiles[i].Placements = slices.Clone(d.MIGProfiles[i].Placements) + } + } if d.MIGTemplate != nil { dup.MIGTemplate = make([]Geometry, len(d.MIGTemplate)) - for i, g := range d.MIGTemplate { - dup.MIGTemplate[i] = slices.Clone(g) + for i, geometry := range d.MIGTemplate { + dup.MIGTemplate[i] = slices.Clone(geometry) } } @@ -234,14 +270,20 @@ func (d *DeviceUsage) DeepCopy() *DeviceUsage { Health: d.Health, } + if d.MigProfiles != nil { + dup.MigProfiles = make([]MigProfile, len(d.MigProfiles)) + copy(dup.MigProfiles, d.MigProfiles) + for i := range dup.MigProfiles { + dup.MigProfiles[i].Placements = slices.Clone(d.MigProfiles[i].Placements) + } + } + dup.MigAllocationsInUse = slices.Clone(d.MigAllocationsInUse) if d.MigTemplate != nil { dup.MigTemplate = make([]Geometry, len(d.MigTemplate)) - for i, g := range d.MigTemplate { - dup.MigTemplate[i] = make(Geometry, len(g)) - copy(dup.MigTemplate[i], g) + for i, geometry := range d.MigTemplate { + dup.MigTemplate[i] = slices.Clone(geometry) } } - dup.MigUsage = d.MigUsage.DeepCopy() if d.PodInfos != nil { @@ -260,15 +302,7 @@ func (d *DeviceUsage) DeepCopy() *DeviceUsage { } func (m MigInUse) DeepCopy() MigInUse { - var usageList MIGS - if m.UsageList != nil { - usageList = make(MIGS, len(m.UsageList)) - copy(usageList, m.UsageList) - } - return MigInUse{ - Index: m.Index, - UsageList: usageList, - } + return MigInUse{Index: m.Index, UsageList: slices.Clone(m.UsageList)} } func GetDevices() map[string]Devices { @@ -377,20 +411,23 @@ func EncodeNodeDevices(dlist []*DeviceInfo) string { return tmp } -// MarshalNodeDevices will only marshal general information, customInfo is neglected. +// MarshalNodeDevices marshals scheduler-visible device information. Arbitrary +// customInfo is deliberately excluded; MIG placement capabilities are part of +// the scheduling contract and must be preserved. func MarshalNodeDevices(dlist []*DeviceInfo) string { devAnnos := []*DeviceInfo{} for _, val := range dlist { devAnnos = append(devAnnos, &DeviceInfo{ - ID: val.ID, - Count: val.Count, - Devmem: val.Devmem, - Devcore: val.Devcore, - Type: val.Type, - Numa: val.Numa, - Health: val.Health, - Index: val.Index, - Mode: val.Mode, + ID: val.ID, + Count: val.Count, + Devmem: val.Devmem, + Devcore: val.Devcore, + Type: val.Type, + Numa: val.Numa, + Health: val.Health, + Index: val.Index, + Mode: val.Mode, + MIGProfiles: val.MIGProfiles, }) } data, err := json.Marshal(devAnnos) @@ -513,26 +550,6 @@ func DecodePodDevices(checklist map[string]string, annos map[string]string) (Pod return pd, nil } -func PlatternMIG(n *MigInUse, templates []Geometry, templateIdx int) { - var err error - for _, val := range templates[templateIdx] { - count := 0 - for count < int(val.Count) { - n.Index, err = safecast.Convert[int32](templateIdx) - if err != nil { - continue - } - n.UsageList = append(n.UsageList, MigTemplateUsage{ - Name: val.Name, - Memory: val.Memory, - Core: val.Core, - InUse: false, - }) - count++ - } - } -} - func GetDevicesUUIDList(infos []*DeviceInfo) []string { uuids := make([]string, 0) for _, info := range infos { @@ -541,6 +558,16 @@ func GetDevicesUUIDList(infos []*DeviceInfo) []string { return uuids } +// Deprecated: dynamic NVIDIA MIG constructs candidates from MigProfiles. +func PlatternMIG(n *MigInUse, templates []Geometry, templateIdx int) { + for _, value := range templates[templateIdx] { + for count := int32(0); count < value.Count; count++ { + n.Index = int32(templateIdx) + n.UsageList = append(n.UsageList, MigTemplateUsage{Name: value.Name, Memory: value.Memory, Core: value.Core}) + } + } +} + func CheckHealth(devType string, resourceCountName string, node *corev1.Node) (bool, bool) { handshake := node.Annotations[util.HandshakeAnnos[devType]] if strings.Contains(handshake, "Requesting") { @@ -582,36 +609,26 @@ func CheckHealth(devType string, resourceCountName string, node *corev1.Node) (b } } -// Enhanced ExtractMigTemplatesFromUUID with error handling. +// Deprecated: new MIG allocations never encode template/slot indexes in UUIDs. func ExtractMigTemplatesFromUUID(uuid string) (int, int, error) { - parts := strings.Split(uuid, "[") - if len(parts) < 2 { - return -1, -1, fmt.Errorf("invalid UUID format: missing '[' delimiter") - } - - tmp := parts[1] - parts = strings.Split(tmp, "]") - if len(parts) < 2 { - return -1, -1, fmt.Errorf("invalid UUID format: missing ']' delimiter") + left := strings.Index(uuid, "[") + right := strings.Index(uuid, "]") + if left < 0 || right <= left { + return -1, -1, fmt.Errorf("invalid legacy MIG UUID %q", uuid) } - - tmp = parts[0] - parts = strings.Split(tmp, "-") - if len(parts) < 2 { - return -1, -1, fmt.Errorf("invalid UUID format: missing '-' delimiter") + parts := strings.Split(uuid[left+1:right], "-") + if len(parts) != 2 { + return -1, -1, fmt.Errorf("invalid legacy MIG UUID %q", uuid) } - templateIdx, err := strconv.Atoi(parts[0]) if err != nil { - return -1, -1, fmt.Errorf("invalid template index: %v", err) + return -1, -1, err } - - pos, err := strconv.Atoi(parts[1]) + slotIdx, err := strconv.Atoi(parts[1]) if err != nil { - return -1, -1, fmt.Errorf("invalid position: %v", err) + return -1, -1, err } - - return templateIdx, pos, nil + return templateIdx, slotIdx, nil } func Resourcereqs(pod *corev1.Pod) (counts PodDeviceRequests) { diff --git a/pkg/device/nvidia/device.go b/pkg/device/nvidia/device.go index f1b7ae2e56..311113dafe 100644 --- a/pkg/device/nvidia/device.go +++ b/pkg/device/nvidia/device.go @@ -20,6 +20,7 @@ import ( "errors" "flag" "fmt" + "sort" "strconv" "strings" "sync" @@ -113,8 +114,8 @@ type NvidiaConfig struct { DefaultGPUNum int32 `yaml:"defaultGPUNum"` MemoryFactor int32 `yaml:"memoryFactor"` // TODO Whether these should be removed - DisableCoreLimit bool `yaml:"disableCoreLimit"` - MigGeometriesList []device.AllowedMigGeometries `yaml:"knownMigGeometries"` + DisableCoreLimit bool `yaml:"disableCoreLimit"` + MigProfileAllowlist []device.AllowedMigProfiles `yaml:"migProfileAllowlist"` // GPUCorePolicy through webhook automatic injected to container env GPUCorePolicy GPUCoreUtilizationPolicy `yaml:"gpuCorePolicy"` // RuntimeClassName is the name of the runtime class to be added to pod.spec.runtimeClassName @@ -176,6 +177,9 @@ func InitNvidiaDevice(nvconfig NvidiaConfig) *NvidiaGPUDevices { device.SupportDevices[NvidiaGPUDevice] = "hami.io/vgpu-devices-allocated" util.HandshakeAnnos[NvidiaGPUDevice] = HandshakeAnnos } + if err := ValidateMigProfileAllowlist(nvconfig.MigProfileAllowlist); err != nil { + klog.Fatalf("invalid MIG profile allowlist: %v", err) + } MemoryFactor = nvconfig.MemoryFactor return &NvidiaGPUDevices{ config: nvconfig, @@ -309,25 +313,6 @@ func (dev *NvidiaGPUDevices) GetNodeDevices(n corev1.Node) ([]*device.DeviceInfo for idx := range nodedevices { nodedevices[idx].DeviceVendor = dev.CommonWord() } - for _, val := range nodedevices { - if val.Mode == MigMode { - val.MIGTemplate = make([]device.Geometry, 0) - for _, migTemplates := range dev.config.MigGeometriesList { - found := false - for _, migDevices := range migTemplates.Models { - if strings.Contains(val.Type, migDevices) { - found = true - break - } - } - if found { - val.MIGTemplate = append(val.MIGTemplate, migTemplates.Geometries...) - break - } - } - } - } - pairScores, ok := n.Annotations[RegisterGPUPairScore] if !ok { klog.V(5).InfoS("no topology score found", "node", n.Name) @@ -520,6 +505,9 @@ func (dev *NvidiaGPUDevices) PatchAnnotations(pod *corev1.Pod, annoinput *map[st deviceStr := device.EncodePodSingleDevice(devlist) (*annoinput)[device.InRequestDevices[NvidiaGPUDevice]] = deviceStr (*annoinput)[device.SupportDevices[NvidiaGPUDevice]] = deviceStr + if allocationData, hasAllocations := EncodeMigAllocations(devlist); hasAllocations { + (*annoinput)[MigAllocationsAnnotation] = allocationData + } klog.V(5).Infof("pod add notation key [%s], values is [%s]", device.InRequestDevices[NvidiaGPUDevice], deviceStr) klog.V(5).Infof("pod add notation key [%s], values is [%s]", device.SupportDevices[NvidiaGPUDevice], deviceStr) } @@ -605,127 +593,59 @@ func (dev *NvidiaGPUDevices) GenerateResourceRequests(ctr *corev1.Container) dev } func (dev *NvidiaGPUDevices) CustomFilterRule(allocated *device.PodDevices, request device.ContainerDeviceRequest, toAllocate device.ContainerDevices, devusage *device.DeviceUsage) bool { - //memreq := request.Memreq - deviceUsageSnapshot := devusage.MigUsage - deviceUsageCurrent := device.MigInUse{ - UsageList: make(device.MIGS, 0), - } - deviceUsageCurrent.UsageList = append(deviceUsageCurrent.UsageList, deviceUsageSnapshot.UsageList...) if devusage.Mode == MigMode { - // The same logic as in AddResourceUsage - if len(deviceUsageCurrent.UsageList) == 0 { - tmpfound := false - for tidx, templates := range devusage.MigTemplate { - for _, template := range templates { - if template.Memory < request.Memreq { - continue - } else { - device.PlatternMIG(&deviceUsageCurrent, devusage.MigTemplate, tidx) - tmpfound = true - break - } - } - if tmpfound { - break - } - } - if !tmpfound { - klog.Infoln("MIG entry no template fit", deviceUsageCurrent.UsageList, "request=", request) - } - } - for _, val := range toAllocate { - found := false - for idx := range deviceUsageCurrent.UsageList { - if !deviceUsageCurrent.UsageList[idx].InUse && deviceUsageCurrent.UsageList[idx].Memory >= val.Usedmem { - deviceUsageCurrent.UsageList[idx].InUse = true - found = true - break - } - } - if !found { - klog.Infoln("MIG entry not found", deviceUsageCurrent.UsageList) + occupied := occupiedMigPlacements(devusage.MigAllocationsInUse) + for _, existing := range toAllocate { + _, placement, ok := selectMigCandidate(devusage.MigProfiles, occupied, existing.Usedmem) + if !ok { return false } + occupied = append(occupied, placement) } - for idx := range deviceUsageCurrent.UsageList { - if !deviceUsageCurrent.UsageList[idx].InUse && deviceUsageCurrent.UsageList[idx].Memory >= request.Memreq { - deviceUsageCurrent.UsageList[idx].InUse = true - klog.Infoln("MIG entry device usage true=", deviceUsageCurrent.UsageList, "request", request, "toAllocate", toAllocate) - return true - } - } - klog.Infoln("MIG entry device usage false=", deviceUsageCurrent.UsageList) - return false + _, _, ok := selectMigCandidate(devusage.MigProfiles, occupied, request.Memreq) + return ok } return true } -func (dev *NvidiaGPUDevices) ScoreNode(node *corev1.Node, podDevices device.PodSingleDevice, previous []*device.DeviceUsage, policy string) float32 { - return 0 -} - -func (dev *NvidiaGPUDevices) migNeedsReset(n *device.DeviceUsage) bool { - if len(n.MigUsage.UsageList) == 0 { - return true - } - for _, val := range n.MigUsage.UsageList { - if val.InUse { - return false +func selectMigCandidate(profiles []device.MigProfile, occupied []device.MigPlacement, memory int32) (device.MigProfile, device.MigPlacement, bool) { + candidates := append([]device.MigProfile(nil), profiles...) + sort.SliceStable(candidates, func(i, j int) bool { + if candidates[i].MemoryMB != candidates[j].MemoryMB { + return candidates[i].MemoryMB < candidates[j].MemoryMB + } + return candidates[i].SliceCount < candidates[j].SliceCount + }) + for _, profile := range candidates { + if profile.MemoryMB < memory { + continue + } + placement, ok := selectMigPlacement(profiles, occupied, profile.Name) + if ok { + return profile, placement, true } } - n.MigUsage.UsageList = make(device.MIGS, 0) - return true + return device.MigProfile{}, device.MigPlacement{}, false +} + +func (dev *NvidiaGPUDevices) ScoreNode(node *corev1.Node, podDevices device.PodSingleDevice, previous []*device.DeviceUsage, policy string) float32 { + return 0 } func (dev *NvidiaGPUDevices) AddResourceUsage(pod *corev1.Pod, n *device.DeviceUsage, ctr *device.ContainerDevice) error { if n.Mode == MigMode { - if dev.migNeedsReset(n) { - found := false - OuterLoop: - for tidx, templates := range n.MigTemplate { - for idx, template := range templates { - if template.Memory < ctr.Usedmem { - continue - } else { - device.PlatternMIG(&n.MigUsage, n.MigTemplate, tidx) - // Calculate the correct UsageList index by summing Count of all templates before idx - usageListIdx := 0 - for i := range idx { - usageListIdx += int(templates[i].Count) - } - ctr.Usedmem = n.MigUsage.UsageList[usageListIdx].Memory - ctr.Usedcores = n.MigUsage.UsageList[usageListIdx].Core - if !strings.Contains(ctr.UUID, "[") { - ctr.UUID = ctr.UUID + "[" + fmt.Sprint(tidx) + "-" + fmt.Sprint(idx) + "]" - } - n.MigUsage.Index = int32(tidx) - n.MigUsage.UsageList[usageListIdx].InUse = true - found = true - break OuterLoop - } - } - } - if !found { - return errors.New("mig template allocate resource fail") - } - } else { - found := false - for idx, val := range n.MigUsage.UsageList { - if !val.InUse && val.Memory >= ctr.Usedmem { - n.MigUsage.UsageList[idx].InUse = true - ctr.Usedmem = n.MigUsage.UsageList[idx].Memory - ctr.Usedcores = n.MigUsage.UsageList[idx].Core - if !strings.Contains(ctr.UUID, "[") { - ctr.UUID = ctr.UUID + "[" + fmt.Sprint(n.MigUsage.Index) + "-" + fmt.Sprint(idx) + "]" - } - found = true - break - } - } - if !found { - return errors.New("mig template allocate resource fail") - } + profile, placement, ok := selectMigCandidate(n.MigProfiles, occupiedMigPlacements(n.MigAllocationsInUse), ctr.Usedmem) + if !ok { + return errors.New("MIG profile and placement allocation failed") } + ctr.Usedmem = profile.MemoryMB + ctr.Usedcores = profile.Core + if ctr.CustomInfo == nil { + ctr.CustomInfo = make(map[string]any) + } + ctr.CustomInfo[MigProfileCustomInfo] = profile.Name + ctr.CustomInfo[MigPlacementCustomInfo] = placement + n.MigAllocationsInUse = append(n.MigAllocationsInUse, device.MigAllocation{Profile: profile.Name, Placement: placement}) } n.Used++ n.Usedcores += ctr.Usedcores @@ -854,6 +774,15 @@ func (nv *NvidiaGPUDevices) Fit(devices []*device.DeviceUsage, request device.Co resolvedReq := request resolvedReq.Memreq = memreq if !nv.CustomFilterRule(allocated, resolvedReq, tmpDevs[k.Type], dev) { + // In MIG mode, CustomFilterRule rejects when the requested memory + // does not fit an allowed profile with a free placement on this + // device. Surface this as a distinct reason so users can tell + // placement infeasibility apart from generic filter failure. + if dev.Mode == MigMode { + reason[common.CardMigTopologyInfeasible]++ + klog.V(5).InfoS(common.CardMigTopologyInfeasible, "pod", klog.KObj(pod), "device", dev.ID, "device index", i, "allocations", dev.MigAllocationsInUse) + continue + } reason[common.CardNotFoundCustomFilterRule]++ klog.V(5).InfoS(common.CardNotFoundCustomFilterRule, "pod", klog.KObj(pod), "device", dev.ID, "device index", i) continue diff --git a/pkg/device/nvidia/device_test.go b/pkg/device/nvidia/device_test.go index 933ff0c06c..bc820f607d 100644 --- a/pkg/device/nvidia/device_test.go +++ b/pkg/device/nvidia/device_test.go @@ -19,8 +19,6 @@ package nvidia import ( "context" "errors" - "strconv" - "strings" "testing" "time" @@ -1225,234 +1223,16 @@ func TestDevices_Fit(t *testing.T) { } func TestDevices_AddResourceUsage(t *testing.T) { - tests := []struct { - name string - deviceUsage *device.DeviceUsage - ctr *device.ContainerDevice - wantErr bool - wantUsage *device.DeviceUsage - checkMig bool - wantMigIdx int32 - wantUUID string - }{ - { - name: "test add resource usage", - deviceUsage: &device.DeviceUsage{ - ID: "dev-0", - Used: 0, - Usedcores: 15, - Usedmem: 2000, - }, - ctr: &device.ContainerDevice{ - UUID: "dev-0", - Usedcores: 50, - Usedmem: 1024, - }, - wantUsage: &device.DeviceUsage{ - ID: "dev-0", - Used: 1, - Usedcores: 65, - Usedmem: 3024, - }, - wantErr: false, - checkMig: false, - }, - { - name: "test MIG mode with migNeedsReset true - first template matches", - deviceUsage: &device.DeviceUsage{ - ID: "dev-0", - Used: 0, - Usedcores: 0, - Usedmem: 0, - Mode: MigMode, - MigTemplate: []device.Geometry{ - { - {Name: "1g.5gb", Core: 25, Memory: 5120, Count: 1}, - }, - { - {Name: "1g.5gb", Core: 25, Memory: 5120, Count: 1}, - {Name: "2g.10gb", Core: 50, Memory: 10240, Count: 1}, - }, - }, - MigUsage: device.MigInUse{ - UsageList: make(device.MIGS, 0), - }, - }, - ctr: &device.ContainerDevice{ - UUID: "dev-0", - Usedmem: 6000, - }, - wantUsage: &device.DeviceUsage{ - Used: 1, - Usedcores: 50, - Usedmem: 10240, - }, - wantErr: false, - checkMig: true, - wantMigIdx: 1, - wantUUID: "dev-0[1-1]", - }, - { - name: "test MIG mode with migNeedsReset true - second template matches with correct idx", - deviceUsage: &device.DeviceUsage{ - ID: "dev-1", - Used: 0, - Usedcores: 0, - Usedmem: 0, - Mode: MigMode, - MigTemplate: []device.Geometry{ - { - {Name: "1g.3gb", Core: 25, Memory: 3072, Count: 1}, - }, - { - {Name: "1g.5gb", Core: 25, Memory: 5120, Count: 1}, - {Name: "2g.10gb", Core: 50, Memory: 10240, Count: 1}, - }, - }, - MigUsage: device.MigInUse{ - UsageList: make(device.MIGS, 0), - }, - }, - ctr: &device.ContainerDevice{ - UUID: "dev-1", - Usedmem: 8000, - }, - wantUsage: &device.DeviceUsage{ - Used: 1, - Usedcores: 50, - Usedmem: 10240, - }, - wantErr: false, - checkMig: true, - wantMigIdx: 1, - wantUUID: "dev-1[1-1]", - }, - { - name: "test MIG mode with migNeedsReset true - verify outer loop break", - deviceUsage: &device.DeviceUsage{ - ID: "dev-2", - Used: 0, - Usedcores: 0, - Usedmem: 0, - Mode: MigMode, - MigTemplate: []device.Geometry{ - { - {Name: "1g.5gb", Core: 25, Memory: 5120, Count: 1}, - {Name: "2g.10gb", Core: 50, Memory: 10240, Count: 1}, - }, - { - {Name: "3g.20gb", Core: 100, Memory: 20480, Count: 1}, - }, - }, - MigUsage: device.MigInUse{ - UsageList: make(device.MIGS, 0), - }, - }, - ctr: &device.ContainerDevice{ - UUID: "dev-2", - Usedmem: 6000, - }, - wantUsage: &device.DeviceUsage{ - Used: 1, - Usedcores: 50, - Usedmem: 10240, - }, - wantErr: false, - checkMig: true, - wantMigIdx: 0, - wantUUID: "dev-2[0-1]", - }, - { - name: "test MIG mode with migNeedsReset true - template with Count > 1", - deviceUsage: &device.DeviceUsage{ - ID: "dev-3", - Used: 0, - Usedcores: 0, - Usedmem: 0, - Mode: MigMode, - MigTemplate: []device.Geometry{ - { - // Template index 0: first template has Count=2, second template has Count=1 - {Name: "1g.5gb", Core: 50, Memory: 5120, Count: 2}, - {Name: "2g.10gb", Core: 100, Memory: 10240, Count: 1}, - }, - }, - MigUsage: device.MigInUse{ - UsageList: make(device.MIGS, 0), - }, - }, - ctr: &device.ContainerDevice{ - UUID: "dev-3", - Usedmem: 8000, // Requires 8GB, matches second template (idx=1) which should be at UsageList[2] - }, - wantUsage: &device.DeviceUsage{ - Used: 1, - Usedcores: 100, - Usedmem: 10240, // Should be set to the matched template's memory - }, - wantErr: false, - checkMig: true, - wantMigIdx: 0, - wantUUID: "dev-3[0-1]", - }, + dev := &NvidiaGPUDevices{} + usage := &device.DeviceUsage{ID: "dev-0", Usedcores: 15, Usedmem: 2000} + ctr := &device.ContainerDevice{UUID: "dev-0", Usedcores: 50, Usedmem: 1024} + if err := dev.AddResourceUsage(&corev1.Pod{}, usage, ctr); err != nil { + t.Fatal(err) } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - dev := &NvidiaGPUDevices{} - if err := dev.AddResourceUsage(&corev1.Pod{}, tt.deviceUsage, tt.ctr); (err != nil) != tt.wantErr { - t.Errorf("AddResourceUsage() error=%v, wantErr %v", err, tt.wantErr) - } - if !tt.wantErr { - if tt.deviceUsage.Usedcores != tt.wantUsage.Usedcores { - t.Errorf("expected used cores: %d, got used cores %d", tt.wantUsage.Usedcores, tt.deviceUsage.Usedcores) - } - if tt.deviceUsage.Usedmem != tt.wantUsage.Usedmem { - t.Errorf("expected used mem: %d, got used mem %d", tt.wantUsage.Usedmem, tt.deviceUsage.Usedmem) - } - if tt.deviceUsage.Used != tt.wantUsage.Used { - t.Errorf("expected used: %d, got used %d", tt.wantUsage.Used, tt.deviceUsage.Used) - } - if tt.checkMig { - // Verify MIG-related fields - if tt.deviceUsage.MigUsage.Index != tt.wantMigIdx { - t.Errorf("expected MigUsage.Index: %d, got %d", tt.wantMigIdx, tt.deviceUsage.MigUsage.Index) - } - if tt.ctr.UUID != tt.wantUUID { - t.Errorf("expected UUID: %s, got %s", tt.wantUUID, tt.ctr.UUID) - } - // Verify that the entry at the corresponding index in UsageList is marked as InUse - // According to the modified code, should calculate usageListIdx by summing Count of all templates before idx - expectedUsageListIdx := -1 - if strings.Contains(tt.wantUUID, "[") { - parts := strings.Split(strings.TrimSuffix(strings.Split(tt.wantUUID, "[")[1], "]"), "-") - if len(parts) == 2 { - if tidx, err1 := strconv.Atoi(parts[0]); err1 == nil { - if idx, err2 := strconv.Atoi(parts[1]); err2 == nil { - // Calculate usageListIdx by summing Count of all templates before idx - if tidx >= 0 && tidx < len(tt.deviceUsage.MigTemplate) { - expectedUsageListIdx = 0 - for i := 0; i < idx && i < len(tt.deviceUsage.MigTemplate[tidx]); i++ { - expectedUsageListIdx += int(tt.deviceUsage.MigTemplate[tidx][i].Count) - } - } - } - } - } - } - if expectedUsageListIdx >= 0 && expectedUsageListIdx < len(tt.deviceUsage.MigUsage.UsageList) { - if !tt.deviceUsage.MigUsage.UsageList[expectedUsageListIdx].InUse { - t.Errorf("expected UsageList[%d].InUse to be true, got false", expectedUsageListIdx) - } - if tt.deviceUsage.MigUsage.UsageList[expectedUsageListIdx].Memory != tt.ctr.Usedmem { - t.Errorf("expected UsageList[%d].Memory: %d, got %d", expectedUsageListIdx, tt.ctr.Usedmem, tt.deviceUsage.MigUsage.UsageList[expectedUsageListIdx].Memory) - } - } - } - } - }) + if usage.Used != 1 || usage.Usedcores != 65 || usage.Usedmem != 3024 { + t.Fatalf("unexpected usage: %+v", usage) } } - func TestFitQuota(t *testing.T) { NvidiaGPUDevice := "NVIDIA" config := NvidiaConfig{ @@ -2209,58 +1989,6 @@ func TestZeroMemoryPercentageIsAccountedAsWholeCard(t *testing.T) { assert.Assert(t, strings.Contains(reason, "CardInsufficientMemory"), "reason: %s", reason) } -func TestMigNeedsReset(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - - tests := []struct { - name string - usage *device.DeviceUsage - want bool - wantEmpty bool - }{ - { - name: "empty UsageList — needs reset", - usage: &device.DeviceUsage{MigUsage: device.MigInUse{UsageList: device.MIGS{}}}, - want: true, - }, - { - name: "all entries not InUse — resets and clears list", - usage: &device.DeviceUsage{ - MigUsage: device.MigInUse{ - UsageList: device.MIGS{ - {Memory: 10, InUse: false}, - {Memory: 20, InUse: false}, - }, - }, - }, - want: true, - wantEmpty: true, - }, - { - name: "one entry InUse=true — no reset", - usage: &device.DeviceUsage{ - MigUsage: device.MigInUse{ - UsageList: device.MIGS{ - {Memory: 10, InUse: true}, - {Memory: 20, InUse: false}, - }, - }, - }, - want: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := dev.migNeedsReset(tt.usage) - assert.Equal(t, result, tt.want) - if tt.wantEmpty { - assert.Equal(t, len(tt.usage.MigUsage.UsageList), 0) - } - }) - } -} - func TestGenerateCombinations(t *testing.T) { devs := device.ContainerDevices{ {UUID: "gpu0"}, @@ -2370,59 +2098,6 @@ func TestCustomFilterRule_NonMig(t *testing.T) { assert.Equal(t, result, true) } -func TestCustomFilterRule_Mig(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - - tests := []struct { - name string - usageList device.MIGS - toAllocate device.ContainerDevices - memreq int32 - want bool - }{ - { - name: "slot available after allocating toAllocate", - usageList: device.MIGS{ - {Memory: 1000, InUse: false}, - {Memory: 1000, InUse: false}, - }, - toAllocate: device.ContainerDevices{{Usedmem: 500}}, - memreq: 500, - want: true, - }, - { - name: "no remaining slot after allocating toAllocate", - usageList: device.MIGS{ - {Memory: 1000, InUse: false}, - }, - toAllocate: device.ContainerDevices{{Usedmem: 500}}, - memreq: 500, - want: false, - }, - { - name: "toAllocate does not fit any slot", - usageList: device.MIGS{ - {Memory: 100, InUse: false}, - }, - toAllocate: device.ContainerDevices{{Usedmem: 500}}, - memreq: 500, - want: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - devusage := &device.DeviceUsage{ - Mode: MigMode, - MigUsage: device.MigInUse{UsageList: tt.usageList}, - } - req := device.ContainerDeviceRequest{Memreq: tt.memreq} - result := dev.CustomFilterRule(nil, req, tt.toAllocate, devusage) - assert.Equal(t, result, tt.want) - }) - } -} - func TestNodeCleanUp(t *testing.T) { client.KubeClient = fake.NewClientset() node := &corev1.Node{ @@ -2685,34 +2360,22 @@ func TestGetNodeDevices_InvalidJSON(t *testing.T) { assert.Assert(t, err != nil) } -func TestGetNodeDevices_MigTemplate(t *testing.T) { - dev := &NvidiaGPUDevices{ - config: NvidiaConfig{ - MigGeometriesList: []device.AllowedMigGeometries{ - { - Models: []string{"A100"}, - Geometries: []device.Geometry{ - { - {Name: "1g.5gb", Memory: 5120, Core: 14, Count: 7}, - }, - }, - }, - }, - }, - } +func TestGetNodeDevices_MigProfilesFromNode(t *testing.T) { + dev := &NvidiaGPUDevices{} node := corev1.Node{ ObjectMeta: metav1.ObjectMeta{ Name: "node-mig", Annotations: map[string]string{ - RegisterAnnos: `[{"id":"GPU-0","count":7,"devmem":40960,"devcore":100,"type":"NVIDIA-A100-SXM4-40GB","numa":0,"health":true,"mode":"mig"}]`, + RegisterAnnos: `[{"id":"GPU-0","count":7,"devmem":40960,"devcore":100,"type":"NVIDIA-A100-SXM4-40GB","numa":0,"health":true,"mode":"mig","migProfiles":[{"name":"1g.5gb","memoryMB":5120,"core":14,"sliceCount":1,"instanceCount":7,"multiprocessorCount":14,"placements":[{"start":6,"size":1}]}]}]`, }, }, } result, err := dev.GetNodeDevices(node) assert.NilError(t, err) assert.Equal(t, len(result), 1) - assert.Equal(t, len(result[0].MIGTemplate), 1) - assert.Equal(t, result[0].MIGTemplate[0][0].Name, "1g.5gb") + assert.Equal(t, len(result[0].MIGProfiles), 1) + assert.Equal(t, result[0].MIGProfiles[0].Name, "1g.5gb") + assert.DeepEqual(t, result[0].MIGProfiles[0].Placements, []device.MigPlacement{{Start: 6, Size: 1}}) } func TestGetNodeDevices_PairScores(t *testing.T) { @@ -2747,99 +2410,6 @@ func TestGetNodeDevices_InvalidPairScores(t *testing.T) { assert.Assert(t, err != nil) } -func TestAddResourceUsage_MigNonReset(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - usage := &device.DeviceUsage{ - Mode: MigMode, - MigUsage: device.MigInUse{ - Index: 0, - UsageList: device.MIGS{ - {Name: "1g.5gb", Memory: 5120, Core: 14, InUse: true}, - {Name: "1g.5gb", Memory: 5120, Core: 14, InUse: false}, - }, - }, - } - ctr := &device.ContainerDevice{UUID: "GPU-0", Usedmem: 4096, Usedcores: 10} - err := dev.AddResourceUsage(&corev1.Pod{}, usage, ctr) - assert.NilError(t, err) - assert.Assert(t, usage.MigUsage.UsageList[1].InUse) - assert.Equal(t, ctr.Usedmem, int32(5120)) - assert.Assert(t, strings.Contains(ctr.UUID, "[")) -} - -func TestAddResourceUsage_MigNonResetNoSlot(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - usage := &device.DeviceUsage{ - Mode: MigMode, - MigUsage: device.MigInUse{ - UsageList: device.MIGS{ - {Name: "1g.5gb", Memory: 5120, Core: 14, InUse: true}, - }, - }, - } - ctr := &device.ContainerDevice{UUID: "GPU-0", Usedmem: 4096} - err := dev.AddResourceUsage(&corev1.Pod{}, usage, ctr) - assert.Assert(t, err != nil) - assert.Assert(t, strings.Contains(err.Error(), "mig template allocate resource fail")) - assert.Equal(t, usage.Used, int32(0)) -} - -func TestAddResourceUsage_MigResetNoFit(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - usage := &device.DeviceUsage{ - Mode: MigMode, - MigTemplate: []device.Geometry{ - { - {Name: "1g.5gb", Memory: 1024, Core: 14, Count: 1}, - }, - }, - } - ctr := &device.ContainerDevice{UUID: "GPU-0", Usedmem: 4096} - err := dev.AddResourceUsage(&corev1.Pod{}, usage, ctr) - assert.Assert(t, err != nil) - assert.Assert(t, strings.Contains(err.Error(), "mig template allocate resource fail")) - // No template fit: usage counters must not reflect a phantom allocation. - assert.Equal(t, usage.Usedmem, int32(0)) - assert.Equal(t, usage.Used, int32(0)) - assert.Assert(t, !strings.Contains(ctr.UUID, "[")) -} - -func TestCustomFilterRule_MigEmptyUsageWithTemplate(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - devusage := &device.DeviceUsage{ - Mode: MigMode, - MigUsage: device.MigInUse{ - UsageList: device.MIGS{}, - }, - MigTemplate: []device.Geometry{ - { - {Name: "1g.5gb", Memory: 5120, Core: 14, Count: 2}, - }, - }, - } - req := device.ContainerDeviceRequest{Memreq: 4096} - result := dev.CustomFilterRule(nil, req, nil, devusage) - assert.Equal(t, result, true) -} - -func TestCustomFilterRule_MigEmptyUsageNoFitTemplate(t *testing.T) { - dev := InitNvidiaDevice(NvidiaConfig{}) - devusage := &device.DeviceUsage{ - Mode: MigMode, - MigUsage: device.MigInUse{ - UsageList: device.MIGS{}, - }, - MigTemplate: []device.Geometry{ - { - {Name: "1g.5gb", Memory: 100, Core: 14, Count: 1}, - }, - }, - } - req := device.ContainerDeviceRequest{Memreq: 4096} - result := dev.CustomFilterRule(nil, req, nil, devusage) - assert.Equal(t, result, false) -} - func TestReleaseNodeLock(t *testing.T) { config := NvidiaConfig{ ResourceCountName: "nvidia.com/gpu", diff --git a/pkg/device/nvidia/mig_allocations.go b/pkg/device/nvidia/mig_allocations.go new file mode 100644 index 0000000000..fd31b4daa1 --- /dev/null +++ b/pkg/device/nvidia/mig_allocations.go @@ -0,0 +1,81 @@ +package nvidia + +import ( + "encoding/json" + "fmt" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +const MigAllocationsAnnotation = "hami.io/vgpu-mig-allocations" + +const ( + MigProfileCustomInfo = "migProfile" + MigPlacementCustomInfo = "migPlacement" +) + +// MigAllocation is the complete scheduler reservation for one MIG device. +// The device plugin only realizes this exact hardware placement and fills in +// MigUUID. ContainerIndex and DeviceIndex make repeated allocations on the +// same physical GPU unambiguous without synthesizing logical device IDs. +type MigAllocation struct { + ContainerIndex int `json:"containerIndex"` + DeviceIndex int `json:"deviceIndex"` + GPUUUID string `json:"gpuUUID"` + Profile string `json:"profile"` + Placement device.MigPlacement `json:"placement"` + MigUUID string `json:"migUUID,omitempty"` +} + +func EncodeMigAllocations(pd device.PodSingleDevice) (string, bool) { + out := make([]MigAllocation, 0) + for containerIndex, ctr := range pd { + for deviceIndex, dev := range ctr { + if dev.CustomInfo == nil { + continue + } + profile, profileOK := dev.CustomInfo[MigProfileCustomInfo].(string) + placement, placementOK := dev.CustomInfo[MigPlacementCustomInfo].(device.MigPlacement) + if !profileOK || profile == "" || !placementOK || placement.Size == 0 { + continue + } + out = append(out, MigAllocation{ + ContainerIndex: containerIndex, + DeviceIndex: deviceIndex, + GPUUUID: dev.UUID, + Profile: profile, + Placement: placement, + }) + } + } + if len(out) == 0 { + return "", false + } + raw, err := json.Marshal(out) + if err != nil { + return "", false + } + return string(raw), true +} + +func DecodeMigAllocations(raw string) ([]MigAllocation, error) { + if raw == "" { + return nil, nil + } + var out []MigAllocation + if err := json.Unmarshal([]byte(raw), &out); err != nil { + return nil, err + } + seen := make(map[[2]int]struct{}, len(out)) + for i, allocation := range out { + if allocation.ContainerIndex < 0 || allocation.DeviceIndex < 0 || allocation.GPUUUID == "" || allocation.Profile == "" || allocation.Placement.Size == 0 { + return nil, fmt.Errorf("MIG allocation %d is incomplete", i) + } + key := [2]int{allocation.ContainerIndex, allocation.DeviceIndex} + if _, ok := seen[key]; ok { + return nil, fmt.Errorf("duplicate MIG allocation for container %d device %d", allocation.ContainerIndex, allocation.DeviceIndex) + } + seen[key] = struct{}{} + } + return out, nil +} diff --git a/pkg/device/nvidia/mig_capability_test.go b/pkg/device/nvidia/mig_capability_test.go new file mode 100644 index 0000000000..a8c3c4e68b --- /dev/null +++ b/pkg/device/nvidia/mig_capability_test.go @@ -0,0 +1,51 @@ +package nvidia + +import ( + "testing" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +func a100MigProfiles() []device.MigProfile { + return []device.MigProfile{ + {Name: "1g.5gb", MemoryMB: 5120, Core: 14, SliceCount: 1, InstanceCount: 7, Placements: []device.MigPlacement{{Start: 0, Size: 1}, {Start: 1, Size: 1}, {Start: 2, Size: 1}, {Start: 3, Size: 1}, {Start: 4, Size: 1}, {Start: 5, Size: 1}, {Start: 6, Size: 1}}}, + {Name: "2g.10gb", MemoryMB: 10240, Core: 28, SliceCount: 2, InstanceCount: 3, Placements: []device.MigPlacement{{Start: 0, Size: 2}, {Start: 2, Size: 2}, {Start: 4, Size: 2}}}, + {Name: "3g.20gb", MemoryMB: 20480, Core: 42, SliceCount: 3, InstanceCount: 2, Placements: []device.MigPlacement{{Start: 0, Size: 4}, {Start: 4, Size: 4}}}, + } +} + +func TestAddResourceUsageUsesReportedProfileAndPlacement(t *testing.T) { + dev := &NvidiaGPUDevices{} + usage := &device.DeviceUsage{ID: "GPU-a", Mode: MigMode, MigProfiles: a100MigProfiles()} + ctr := &device.ContainerDevice{UUID: "GPU-a", Usedmem: 10000} + if err := dev.AddResourceUsage(nil, usage, ctr); err != nil { + t.Fatalf("allocate MIG capability: %v", err) + } + if ctr.UUID != "GPU-a" { + t.Fatalf("physical UUID was rewritten: %s", ctr.UUID) + } + if ctr.Usedmem != 10240 || ctr.Usedcores != 28 { + t.Fatalf("allocated resources=(%d,%d), want reported 2g metadata", ctr.Usedmem, ctr.Usedcores) + } + if ctr.CustomInfo[MigProfileCustomInfo] != "2g.10gb" || ctr.CustomInfo[MigPlacementCustomInfo] != (device.MigPlacement{Start: 0, Size: 2}) { + t.Fatalf("unexpected scheduler reservation: %+v", ctr.CustomInfo) + } +} + +func TestCustomFilterUsesReportedPlacementCapacity(t *testing.T) { + dev := &NvidiaGPUDevices{} + usage := &device.DeviceUsage{ + Mode: MigMode, MigProfiles: a100MigProfiles(), + MigAllocationsInUse: []device.MigAllocation{ + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 0, Size: 2}}, + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 2, Size: 2}}, + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 4, Size: 2}}, + }, + } + if dev.CustomFilterRule(nil, device.ContainerDeviceRequest{Memreq: 10000}, nil, usage) { + t.Fatal("fourth 2g request should not fit reported placements") + } + if !dev.CustomFilterRule(nil, device.ContainerDeviceRequest{Memreq: 5000}, nil, usage) { + t.Fatal("1g request should fit the remaining placement") + } +} diff --git a/pkg/device/nvidia/mig_profiles.go b/pkg/device/nvidia/mig_profiles.go new file mode 100644 index 0000000000..cec0ab5eec --- /dev/null +++ b/pkg/device/nvidia/mig_profiles.go @@ -0,0 +1,25 @@ +package nvidia + +import ( + "fmt" + "strings" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +// ValidateMigProfileAllowlist validates policy only. Profile capacity, +// memory, compute metadata and placements always come from NVML on the node. +func ValidateMigProfileAllowlist(in []device.AllowedMigProfiles) error { + for _, cfg := range in { + if len(cfg.Models) == 0 || len(cfg.Profiles) == 0 { + return fmt.Errorf("MIG profile allowlist must define models and profiles") + } + for _, profile := range cfg.Profiles { + parts := strings.SplitN(profile, ".", 2) + if len(parts) != 2 || len(parts[0]) < 2 || !strings.HasSuffix(parts[0], "g") { + return fmt.Errorf("invalid MIG profile %q", profile) + } + } + } + return nil +} diff --git a/pkg/device/nvidia/mig_profiles_test.go b/pkg/device/nvidia/mig_profiles_test.go new file mode 100644 index 0000000000..a7b51d7576 --- /dev/null +++ b/pkg/device/nvidia/mig_profiles_test.go @@ -0,0 +1,50 @@ +package nvidia + +import ( + "testing" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +func TestValidateMigProfileAllowlist(t *testing.T) { + valid := []device.AllowedMigProfiles{{Models: []string{"A100"}, Profiles: []string{"1g.5gb", "2g.10gb"}}} + if err := ValidateMigProfileAllowlist(valid); err != nil { + t.Fatalf("valid allowlist rejected: %v", err) + } + invalid := []device.AllowedMigProfiles{{Models: []string{"A100"}}} + if err := ValidateMigProfileAllowlist(invalid); err == nil { + t.Fatal("allowlist without profiles should fail") + } +} + +func TestEncodeDecodeMigAllocations(t *testing.T) { + pd := device.PodSingleDevice{{{ + UUID: "GPU-a", Type: NvidiaGPUDevice, Usedmem: 5120, Usedcores: 14, + CustomInfo: map[string]any{ + MigProfileCustomInfo: "1g.5gb", + MigPlacementCustomInfo: device.MigPlacement{Start: 6, Size: 1}, + }, + }}} + raw, ok := EncodeMigAllocations(pd) + if !ok { + t.Fatal("expected MIG allocation annotation") + } + allocations, err := DecodeMigAllocations(raw) + if err != nil { + t.Fatalf("decode allocation: %v", err) + } + if len(allocations) != 1 { + t.Fatalf("allocations=%d, want 1", len(allocations)) + } + allocation := allocations[0] + if allocation.ContainerIndex != 0 || allocation.DeviceIndex != 0 || allocation.GPUUUID != "GPU-a" || allocation.Profile != "1g.5gb" || allocation.Placement != (device.MigPlacement{Start: 6, Size: 1}) { + t.Fatalf("unexpected allocation: %+v", allocation) + } +} + +func TestDecodeMigAllocationsRejectsDuplicateDeviceIndex(t *testing.T) { + raw := `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":6,"size":1}},{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":5,"size":1}}]` + if _, err := DecodeMigAllocations(raw); err == nil { + t.Fatal("duplicate container/device allocation index should fail") + } +} diff --git a/pkg/device/nvidia/mig_topology.go b/pkg/device/nvidia/mig_topology.go new file mode 100644 index 0000000000..4dbc017357 --- /dev/null +++ b/pkg/device/nvidia/mig_topology.go @@ -0,0 +1,116 @@ +package nvidia + +import ( + "sort" + "strings" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +func preferHighMigPlacement(profile string) bool { + return strings.HasPrefix(profile, "1g.") || strings.HasPrefix(profile, "3g.") +} + +func orderedMigPlacements(profile string, placements []device.MigPlacement) []device.MigPlacement { + ordered := append([]device.MigPlacement(nil), placements...) + sort.SliceStable(ordered, func(i, j int) bool { + if preferHighMigPlacement(profile) { + return ordered[i].Start > ordered[j].Start + } + return ordered[i].Start < ordered[j].Start + }) + return ordered +} + +func findMigProfile(profiles []device.MigProfile, name string) (device.MigProfile, bool) { + for _, profile := range profiles { + if profile.Name == name { + return profile, true + } + } + return device.MigProfile{}, false +} + +func occupiedMigPlacements(allocations []device.MigAllocation) []device.MigPlacement { + out := make([]device.MigPlacement, 0, len(allocations)) + for _, allocation := range allocations { + out = append(out, allocation.Placement) + } + return out +} + +func selectMigPlacement(profiles []device.MigProfile, occupied []device.MigPlacement, profileName string) (device.MigPlacement, bool) { + profile, ok := findMigProfile(profiles, profileName) + if !ok { + return device.MigPlacement{}, false + } + for _, candidate := range orderedMigPlacements(profileName, profile.Placements) { + conflict := false + for _, existing := range occupied { + if migPlacementsOverlap(candidate, existing) { + conflict = true + break + } + } + if !conflict { + return candidate, true + } + } + return device.MigPlacement{}, false +} + +func migPlacementsOverlap(a, b device.MigPlacement) bool { + return a.Start < b.Start+b.Size && b.Start < a.Start+a.Size +} + +// canPlaceMigProfiles determines whether every requested profile can be +// assigned one of the profile-specific placements reported by NVML without +// overlapping an already-running or newly-selected instance. +func canPlaceMigProfiles(capabilities []device.MigProfile, occupied []device.MigPlacement, requested []string) bool { + if len(requested) == 0 { + return true + } + placements := make(map[string][]device.MigPlacement, len(capabilities)) + for _, profile := range capabilities { + placements[profile.Name] = profile.Placements + } + ordered := append([]string(nil), requested...) + sort.SliceStable(ordered, func(i, j int) bool { + left, right := placements[ordered[i]], placements[ordered[j]] + if len(left) != len(right) { + return len(left) < len(right) + } + if len(left) == 0 { + return false + } + return left[0].Size > right[0].Size + }) + + used := append([]device.MigPlacement(nil), occupied...) + var place func(int) bool + place = func(idx int) bool { + if idx == len(ordered) { + return true + } + candidates := orderedMigPlacements(ordered[idx], placements[ordered[idx]]) + for _, candidate := range candidates { + conflict := false + for _, existing := range used { + if migPlacementsOverlap(candidate, existing) { + conflict = true + break + } + } + if conflict { + continue + } + used = append(used, candidate) + if place(idx + 1) { + return true + } + used = used[:len(used)-1] + } + return false + } + return place(0) +} diff --git a/pkg/device/nvidia/mig_topology_test.go b/pkg/device/nvidia/mig_topology_test.go new file mode 100644 index 0000000000..8363b70ba3 --- /dev/null +++ b/pkg/device/nvidia/mig_topology_test.go @@ -0,0 +1,52 @@ +package nvidia + +import ( + "testing" + + "github.com/Project-HAMi/HAMi/pkg/device" +) + +func TestCanPlaceMigProfilesA100(t *testing.T) { + allowed := []device.MigProfile{ + {Name: "1g.5gb", Placements: []device.MigPlacement{{Start: 0, Size: 1}, {Start: 1, Size: 1}, {Start: 2, Size: 1}, {Start: 3, Size: 1}, {Start: 4, Size: 1}, {Start: 5, Size: 1}, {Start: 6, Size: 1}}}, + {Name: "2g.10gb", Placements: []device.MigPlacement{{Start: 0, Size: 2}, {Start: 2, Size: 2}, {Start: 4, Size: 2}}}, + {Name: "3g.20gb", Placements: []device.MigPlacement{{Start: 0, Size: 4}, {Start: 4, Size: 4}}}, + } + + if !canPlaceMigProfiles(allowed, nil, []string{"3g.20gb", "2g.10gb", "1g.5gb", "1g.5gb"}) { + t.Fatal("A100 all-balanced geometry should be feasible") + } + if !canPlaceMigProfiles(allowed, []device.MigPlacement{{Start: 4, Size: 4}}, []string{"2g.10gb", "1g.5gb", "1g.5gb"}) { + t.Fatal("remaining all-balanced placements should be feasible") + } + if canPlaceMigProfiles(allowed, []device.MigPlacement{{Start: 0, Size: 4}, {Start: 4, Size: 4}}, []string{"1g.5gb"}) { + t.Fatal("no profile should overlap two active 3g placements") + } + if canPlaceMigProfiles(allowed, []device.MigPlacement{{Start: 0, Size: 2}, {Start: 2, Size: 2}, {Start: 4, Size: 2}}, []string{"2g.10gb"}) { + t.Fatal("fourth 2g placement should be rejected") + } +} + +func TestSelectMigPlacementUsesBalancedPacking(t *testing.T) { + allowed := []device.MigProfile{ + {Name: "1g.5gb", Placements: []device.MigPlacement{{Start: 0, Size: 1}, {Start: 1, Size: 1}, {Start: 2, Size: 1}, {Start: 3, Size: 1}, {Start: 4, Size: 1}, {Start: 5, Size: 1}, {Start: 6, Size: 1}}}, + {Name: "2g.10gb", Placements: []device.MigPlacement{{Start: 0, Size: 2}, {Start: 2, Size: 2}, {Start: 4, Size: 2}}}, + {Name: "3g.20gb", Placements: []device.MigPlacement{{Start: 0, Size: 4}, {Start: 4, Size: 4}}}, + } + occupied := []device.MigPlacement{} + for _, tc := range []struct { + profile string + start uint32 + }{ + {profile: "3g.20gb", start: 4}, + {profile: "2g.10gb", start: 0}, + {profile: "1g.5gb", start: 3}, + {profile: "1g.5gb", start: 2}, + } { + placement, ok := selectMigPlacement(allowed, occupied, tc.profile) + if !ok || placement.Start != tc.start { + t.Fatalf("profile %s placement = %+v, %v; want start %d", tc.profile, placement, ok, tc.start) + } + occupied = append(occupied, placement) + } +} diff --git a/pkg/device/pod_test.go b/pkg/device/pod_test.go index 1d0e21962a..74eef79ebf 100644 --- a/pkg/device/pod_test.go +++ b/pkg/device/pod_test.go @@ -318,7 +318,7 @@ func TestAddPod_UpdateExistingDevices(t *testing.T) { pi, ok := podManager.GetPod(pod) assert.True(t, ok) - assert.Equal(t, "node1", pi.NodeID, "NodeID should be unchanged on update") + assert.Equal(t, "node2", pi.NodeID, "NodeID should be refreshed with the latest pod state") assert.Equal(t, PodDevices{"device1": {{{UUID: "GPU-1"}}}}, pi.Devices, "Devices should be replaced with the new value") } @@ -393,6 +393,30 @@ func TestListPodsUID(t *testing.T) { assert.True(t, gotUIDs[pod2.UID], "expected pod2's UID to be present") } +func TestAddPodRefreshesExistingPodObject(t *testing.T) { + podManager := NewPodManager() + uid := k8stypes.UID("uid1") + original := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Namespace: "default", + Name: "pod1", + UID: uid, + Annotations: map[string]string{ + "hami.io/vgpu-mig-allocations": `[{"profile":"1g.5gb"}]`, + }, + }} + assert.True(t, podManager.AddPod(original, "node1", PodDevices{"device1": {{}}})) + + updated := original.DeepCopy() + updated.Annotations["hami.io/vgpu-mig-allocations"] = `[{"profile":"1g.5gb","placement":{"start":6,"size":1}}]` + devices := PodDevices{"device1": {{{UUID: "GPU-1"}}}} + assert.False(t, podManager.AddPod(updated, "node1", devices), "refreshing an existing pod must not add quota usage again") + + cached, ok := podManager.GetPod(updated) + assert.True(t, ok) + assert.Equal(t, updated.Annotations, cached.Annotations) + assert.Equal(t, devices, cached.Devices) +} + func TestUpdatePod(t *testing.T) { podManager := NewPodManager() diff --git a/pkg/device/pods.go b/pkg/device/pods.go index fc2d2b85f7..c4e1d05924 100644 --- a/pkg/device/pods.go +++ b/pkg/device/pods.go @@ -75,6 +75,7 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) } else { pi := m.pods[pod.UID] pi.Pod = pod + pi.NodeID = nodeID if pi.InitContainerResourceReleased { // Usage was already shrunk after init containers finished; a re-add // (e.g. an informer resync decoding the full annotation) must not @@ -86,6 +87,8 @@ func (m *PodManager) AddPod(pod *corev1.Pod, nodeID string, devices PodDevices) pi.Devices = devices klog.V(5).InfoS("Pod already exists; devices updated", "pod", klog.KRef(pod.Namespace, pod.Name), + "nodeID", nodeID, + "devices", devices, ) } } diff --git a/pkg/scheduler/mig_allocation_test.go b/pkg/scheduler/mig_allocation_test.go new file mode 100644 index 0000000000..e680b3b7a5 --- /dev/null +++ b/pkg/scheduler/mig_allocation_test.go @@ -0,0 +1,60 @@ +package scheduler + +import ( + "testing" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" +) + +func TestGetNodesUsageRestoresMigAllocationByProfileAndPlacement(t *testing.T) { + nodes := newNodeManager() + nodes.addNode("node1", &device.NodeInfo{ + ID: "node1", Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: {{ + ID: "GPU-a", Count: 7, Devmem: 40960, Devcore: 100, Mode: nvidia.MigMode, Health: true, + MIGProfiles: []device.MigProfile{{Name: "1g.5gb", MemoryMB: 5120, Core: 14, Placements: []device.MigPlacement{{Start: 6, Size: 1}}}}, + }}}, + }) + pods := device.NewPodManager() + pods.AddPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + UID: "pod-1", Name: "pod-1", Namespace: "default", + Annotations: map[string]string{nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":6,"size":1},"migUUID":"MIG-a"}]`}, + }}, "node1", device.PodDevices{nvidia.NvidiaGPUDevice: {{{UUID: "GPU-a", Usedmem: 5120, Usedcores: 14}}}}) + s := Scheduler{nodeManager: nodes, podManager: pods} + nodeNames := []string{"node1"} + usage, _, _, err := s.getNodesUsage(&nodeNames, nil) + if err != nil { + t.Fatal(err) + } + allocations := (*usage)["node1"].Devices.DeviceLists[0].Device.MigAllocationsInUse + if len(allocations) != 1 || allocations[0].Profile != "1g.5gb" || allocations[0].Placement != (device.MigPlacement{Start: 6, Size: 1}) { + t.Fatalf("restored allocations: %+v", allocations) + } +} + +func TestGetNodesUsageFailsClosedWithoutMigAllocation(t *testing.T) { + nodes := newNodeManager() + nodes.addNode("node1", &device.NodeInfo{ + ID: "node1", Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: {{ + ID: "GPU-a", Count: 7, Devmem: 40960, Devcore: 100, Mode: nvidia.MigMode, Health: true, + MIGProfiles: []device.MigProfile{{Name: "1g.5gb", MemoryMB: 5120, Core: 14, Placements: []device.MigPlacement{{Start: 6, Size: 1}}}}, + }}}, + }) + pods := device.NewPodManager() + pods.AddPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-1", Name: "pod-1", Namespace: "default"}}, "node1", + device.PodDevices{nvidia.NvidiaGPUDevice: {{{UUID: "GPU-a", Usedmem: 5120, Usedcores: 14}}}}) + s := Scheduler{nodeManager: nodes, podManager: pods} + nodeNames := []string{"node1"} + usage, _, _, err := s.getNodesUsage(&nodeNames, nil) + if err != nil { + t.Fatal(err) + } + if (*usage)["node1"].Devices.DeviceLists[0].Device.Health { + t.Fatal("MIG device must be fail-closed when an allocated Pod lacks profile/placement") + } +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 457a92e3a6..905ca5ba2c 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -643,19 +643,15 @@ func buildNodeUsage(node *device.NodeInfo, task *corev1.Pod) *NodeUsage { nodeUsage.Devices.DeviceLists = append(nodeUsage.Devices.DeviceLists, &policy.DeviceListsScore{ Score: 0, Device: &device.DeviceUsage{ - ID: d.ID, - Index: d.Index, - Used: 0, - Count: d.Count, - Usedmem: 0, - Totalmem: d.Devmem, - Totalcore: d.Devcore, - Usedcores: 0, - MigUsage: device.MigInUse{ - Index: 0, - UsageList: make(device.MIGS, 0), - }, - MigTemplate: d.MIGTemplate, + ID: d.ID, + Index: d.Index, + Used: 0, + Count: d.Count, + Usedmem: 0, + Totalmem: d.Devmem, + Totalcore: d.Devcore, + Usedcores: 0, + MigProfiles: d.MIGProfiles, Mode: d.Mode, Type: d.Type, Numa: d.Numa, @@ -721,6 +717,14 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin podsInfo := s.podManager.ListPodsInfo() for _, p := range podsInfo { + allocationsByGPU := map[string][]nvidia.MigAllocation{} + if slotRaw, ok := p.Annotations[nvidia.MigAllocationsAnnotation]; ok { + if allocations, err := nvidia.DecodeMigAllocations(slotRaw); err == nil { + for _, allocation := range allocations { + allocationsByGPU[allocation.GPUUUID] = append(allocationsByGPU[allocation.GPUUUID], allocation) + } + } + } node, ok := overallnodeMap[p.NodeID] if !ok { klog.V(5).InfoS("pod allocated unknown node resources", @@ -732,42 +736,28 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin for _, udevice := range ctrdevs { for _, d := range node.Devices.DeviceLists { deviceID := udevice.UUID - if strings.Contains(deviceID, "[") { - deviceID = strings.Split(deviceID, "[")[0] - } if d.Device.ID == deviceID { d.Device.Used++ d.Device.Usedmem += udevice.Usedmem d.Device.Usedcores += udevice.Usedcores d.Device.PodInfos = append(d.Device.PodInfos, p) - if strings.Contains(udevice.UUID, "[") { + if allocations := allocationsByGPU[udevice.UUID]; len(allocations) > 0 { if strings.Compare(d.Device.Mode, "hami-core") == 0 { klog.Errorf("found a mig task running on a hami-core GPU\n") d.Device.Health = false continue } - tmpIdx, instanceIdx, err := device.ExtractMigTemplatesFromUUID(udevice.UUID) - if err != nil { - klog.Errorf("failed to extract mig templates from uuid %s: %v", udevice.UUID, err) - continue - } - if tmpIdx < 0 || tmpIdx >= len(d.Device.MigTemplate) { - klog.Errorf("invalid mig template index %d in uuid %s (templates length: %d)", tmpIdx, udevice.UUID, len(d.Device.MigTemplate)) - continue - } - if len(d.Device.MigUsage.UsageList) == 0 { - device.PlatternMIG(&d.Device.MigUsage, d.Device.MigTemplate, tmpIdx) - } else if tmpIdx != int(d.Device.MigUsage.Index) { - klog.Errorf("mig template index mismatch in uuid %s: expected %d, got %d", udevice.UUID, d.Device.MigUsage.Index, tmpIdx) - continue - } - if instanceIdx < 0 || instanceIdx >= len(d.Device.MigUsage.UsageList) { - klog.Errorf("invalid mig instance in uuid %s", udevice.UUID) - continue - } - d.Device.MigUsage.UsageList[instanceIdx].InUse = true - klog.V(5).Infoln("add mig usage", d.Device.MigUsage, "template=", d.Device.MigTemplate, "uuid=", d.Device.ID) + allocation := allocations[0] + allocationsByGPU[udevice.UUID] = allocations[1:] + d.Device.MigAllocationsInUse = append(d.Device.MigAllocationsInUse, device.MigAllocation{ + Profile: allocation.Profile, Placement: allocation.Placement, + }) + continue + } + if d.Device.Mode == nvidia.MigMode { + klog.ErrorS(nil, "MIG Pod lacks a matching profile/placement reservation", "pod", klog.KRef(p.Namespace, p.Name), "gpuUUID", udevice.UUID) + d.Device.Health = false } } } diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go index 6646b4cd0f..23383a9a08 100644 --- a/pkg/scheduler/scheduler_test.go +++ b/pkg/scheduler/scheduler_test.go @@ -129,241 +129,6 @@ func Test_getNodesUsage(t *testing.T) { assert.Equal(t, v.Devices.DeviceLists[0].Device.Usedcores, int32(20)) } -func Test_getNodesUsage_StaleMigIndexDoesNotPanic(t *testing.T) { - nodeMage := newNodeManager() - nodeMage.addNode("node1", &device.NodeInfo{ - ID: "node1", - Node: &corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node1", - }, - }, - Devices: map[string][]device.DeviceInfo{ - nvidia.NvidiaGPUDevice: {{ - ID: "GPU0", - Index: 0, - Count: 10, - Devmem: 1024, - Devcore: 100, - Numa: 1, - Mode: "mig", - Health: true, - MIGTemplate: []device.Geometry{ - {{Name: "1g.5gb", Memory: 5, Count: 1}}, - }, - }}, - }, - }) - // tmpIdx=99 is far past len(MIGTemplate)==1: a stale/corrupt annotation - // must not panic PlatternMIG or the UsageList index write. - podDevces := device.PodDevices{ - "NVIDIA": device.PodSingleDevice{ - []device.ContainerDevice{ - { - Idx: 0, - UUID: "GPU0[99-0]", - Usedmem: 100, - Usedcores: 10, - }, - }, - }, - } - podMap := device.NewPodManager() - podMap.AddPod(&corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - UID: "1111", - Name: "test1", - Namespace: "default", - }, - }, "node1", podDevces) - s := Scheduler{ - nodeManager: nodeMage, - podManager: podMap, - } - nodes := []string{"node1"} - cachenodeMap, _, _, err := s.getNodesUsage(&nodes, nil) - if err != nil { - t.Fatal(err) - } - v, ok := (*cachenodeMap)["node1"] - assert.Assert(t, ok) - assert.Equal(t, int32(1), v.Devices.DeviceLists[0].Device.Used) - assert.Equal(t, 0, len(v.Devices.DeviceLists[0].Device.MigUsage.UsageList)) -} - -func Test_getNodesUsage_UnparsableMigUUIDDoesNotPanic(t *testing.T) { - nodeMage := newNodeManager() - nodeMage.addNode("node1", &device.NodeInfo{ - ID: "node1", - Node: &corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node1", - }, - }, - Devices: map[string][]device.DeviceInfo{ - nvidia.NvidiaGPUDevice: {{ - ID: "GPU0", - Index: 0, - Count: 10, - Devmem: 1024, - Devcore: 100, - Numa: 1, - Mode: "mig", - Health: true, - MIGTemplate: []device.Geometry{ - {{Name: "1g.5gb", Memory: 5, Count: 1}}, - }, - }}, - }, - }) - // "abc" fails strconv.Atoi inside ExtractMigTemplatesFromUUID: exercises - // the parse-error branch, distinct from the out-of-range branch above. - podMap := device.NewPodManager() - podMap.AddPod(&corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{UID: "1111", Name: "test1", Namespace: "default"}, - }, "node1", device.PodDevices{ - nvidia.NvidiaGPUDevice: device.PodSingleDevice{ - []device.ContainerDevice{{Idx: 0, UUID: "GPU0[abc-0]", Usedmem: 100, Usedcores: 10}}, - }, - }) - s := Scheduler{ - nodeManager: nodeMage, - podManager: podMap, - } - nodes := []string{"node1"} - cachenodeMap, _, _, err := s.getNodesUsage(&nodes, nil) - if err != nil { - t.Fatal(err) - } - v, ok := (*cachenodeMap)["node1"] - assert.Assert(t, ok) - assert.Equal(t, int32(1), v.Devices.DeviceLists[0].Device.Used) - assert.Equal(t, 0, len(v.Devices.DeviceLists[0].Device.MigUsage.UsageList)) -} - -func Test_getNodesUsage_OutOfRangeMigInstanceSkipped(t *testing.T) { - nodeMage := newNodeManager() - nodeMage.addNode("node1", &device.NodeInfo{ - ID: "node1", - Node: &corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node1", - }, - }, - Devices: map[string][]device.DeviceInfo{ - nvidia.NvidiaGPUDevice: {{ - ID: "GPU0", - Index: 0, - Count: 10, - Devmem: 1024, - Devcore: 100, - Numa: 1, - Mode: "mig", - Health: true, - MIGTemplate: []device.Geometry{ - {{Name: "1g.5gb", Memory: 5, Count: 1}}, - }, - }}, - }, - }) - // Template index 0 is valid and populates a 1-entry UsageList, but the - // instance position "5" is past its end: exercises the instance bounds - // check separately from the template bounds check above. - podMap := device.NewPodManager() - podMap.AddPod(&corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{UID: "1111", Name: "test1", Namespace: "default"}, - }, "node1", device.PodDevices{ - nvidia.NvidiaGPUDevice: device.PodSingleDevice{ - []device.ContainerDevice{{Idx: 0, UUID: "GPU0[0-5]", Usedmem: 100, Usedcores: 10}}, - }, - }) - s := Scheduler{ - nodeManager: nodeMage, - podManager: podMap, - } - nodes := []string{"node1"} - cachenodeMap, _, _, err := s.getNodesUsage(&nodes, nil) - if err != nil { - t.Fatal(err) - } - v, ok := (*cachenodeMap)["node1"] - assert.Assert(t, ok) - assert.Equal(t, int32(1), v.Devices.DeviceLists[0].Device.Used) - assert.Equal(t, 1, len(v.Devices.DeviceLists[0].Device.MigUsage.UsageList)) - assert.Assert(t, !v.Devices.DeviceLists[0].Device.MigUsage.UsageList[0].InUse) -} - -func Test_getNodesUsage_MismatchedMigIndexSkipped(t *testing.T) { - nodeMage := newNodeManager() - nodeMage.addNode("node1", &device.NodeInfo{ - ID: "node1", - Node: &corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node1", - }, - }, - Devices: map[string][]device.DeviceInfo{ - nvidia.NvidiaGPUDevice: {{ - ID: "GPU0", - Index: 0, - Count: 10, - Devmem: 1024, - Devcore: 100, - Numa: 1, - Mode: "mig", - Health: true, - MIGTemplate: []device.Geometry{ - {{Name: "1g.5gb", Memory: 5, Count: 1}}, - {{Name: "2g.10gb", Memory: 10, Count: 1}}, - }, - }}, - }, - }) - // Two pods on the same device disagree on which geometry (template index) - // is active. Whichever is processed first wins; the other must be - // skipped rather than writing its Instance into the wrong UsageList. - podMap := device.NewPodManager() - podMap.AddPod(&corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{UID: "1111", Name: "test1", Namespace: "default"}, - }, "node1", device.PodDevices{ - nvidia.NvidiaGPUDevice: device.PodSingleDevice{ - []device.ContainerDevice{{Idx: 0, UUID: "GPU0[0-0]", Usedmem: 100, Usedcores: 10}}, - }, - }) - podMap.AddPod(&corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{UID: "2222", Name: "test2", Namespace: "default"}, - }, "node1", device.PodDevices{ - nvidia.NvidiaGPUDevice: device.PodSingleDevice{ - []device.ContainerDevice{{Idx: 0, UUID: "GPU0[1-0]", Usedmem: 100, Usedcores: 10}}, - }, - }) - s := Scheduler{ - nodeManager: nodeMage, - podManager: podMap, - } - nodes := []string{"node1"} - cachenodeMap, _, _, err := s.getNodesUsage(&nodes, nil) - if err != nil { - t.Fatal(err) - } - v, ok := (*cachenodeMap)["node1"] - assert.Assert(t, ok) - dev := v.Devices.DeviceLists[0].Device - assert.Equal(t, int32(2), dev.Used) - assert.Equal(t, 1, len(dev.MigUsage.UsageList)) - assert.Assert(t, dev.MigUsage.UsageList[0].InUse) -} - -// test case matrix -/** -| pod name | node name| pod status | annotations | result | -|--------------|----------|-------------|---------------------------- |-------------------------------------| -| test-pod-1 | node11 | Succeeded | hami.io/bind-phase:success | node11:{TotalPod:1,UseDevicePod:1} | -| test-pod-2 | node12 | Running | none | node12:{TotalPod:0;UseDevicePod:0} | -| test-pod-3 | node13 | Succeeded | none | node13:{TotalPod:1;UseDevicePod:0} | -test case matrix. -*/ - func Test_getPodUsage(t *testing.T) { s := NewScheduler() t.Cleanup(func() { close(s.stopCh) }) diff --git a/pkg/scheduler/score_test.go b/pkg/scheduler/score_test.go index e35191e6cd..d9c6a2ea45 100644 --- a/pkg/scheduler/score_test.go +++ b/pkg/scheduler/score_test.go @@ -3634,16 +3634,10 @@ func Test_fitInCertainDevice(t *testing.T) { Usedcores: int32(1), Totalcore: int32(4), Mode: "mig", - MigUsage: device.MigInUse{ - Index: int32(1), - UsageList: device.MIGS{ - { - Name: "test6", - Memory: int32(2048), - InUse: false, - }, - }, - }, + MigProfiles: []device.MigProfile{{ + Name: "1g.test", MemoryMB: 2048, Core: 1, + Placements: []device.MigPlacement{{Start: 0, Size: 1}}, + }}, Health: true, }, }, @@ -3671,7 +3665,7 @@ func Test_fitInCertainDevice(t *testing.T) { }, }, }, - want3: map[string]int{common.CardNotFoundCustomFilterRule: 1, common.AllocatedCardsInsufficientRequest: 1}, + want3: map[string]int{common.CardMigTopologyInfeasible: 1, common.AllocatedCardsInsufficientRequest: 1}, }, { name: "card uuid don't match", From bf7b35b034e748d90f3914fcb6a2913c4cca55e0 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Wed, 5 Aug 2026 15:30:01 +0800 Subject: [PATCH 02/11] clean pod-resource & fix test Signed-off-by: Jifei Wang --- .../device-plugin/daemonsetnvidia.yaml | 7 - docs/develop/mig-dynamic-deallocate_cn.md | 301 ++++++------------ .../nvinternal/plugin/mig_startup.go | 134 +------- .../nvidiadevice/nvinternal/plugin/migmgr.go | 4 +- .../nvidiadevice/nvinternal/plugin/server.go | 24 +- .../nvinternal/plugin/server_test.go | 7 +- .../nvinternal/plugin/util_test.go | 52 --- .../nvinternal/podresources/watcher.go | 264 --------------- .../nvinternal/podresources/watcher_test.go | 161 ---------- 9 files changed, 124 insertions(+), 830 deletions(-) delete mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go delete mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index d185a0f8ad..0e785c363b 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -134,9 +134,6 @@ spec: volumeMounts: - name: device-plugin mountPath: /var/lib/kubelet/device-plugins - - name: pod-resources - mountPath: /var/lib/kubelet/pod-resources - readOnly: true - name: lib mountPath: {{ printf "%s%s" .Values.global.gpuHookPath "/vgpu" }} - name: usrbin @@ -223,10 +220,6 @@ spec: - name: device-plugin hostPath: path: {{ .Values.devicePlugin.pluginPath }} - - name: pod-resources - hostPath: - path: /var/lib/kubelet/pod-resources - type: Directory - name: lib hostPath: path: {{ .Values.devicePlugin.libPath }} diff --git a/docs/develop/mig-dynamic-deallocate_cn.md b/docs/develop/mig-dynamic-deallocate_cn.md index c28b0b33df..9fb7c84492 100644 --- a/docs/develop/mig-dynamic-deallocate_cn.md +++ b/docs/develop/mig-dynamic-deallocate_cn.md @@ -1,254 +1,161 @@ -# 动态 MIG 实例释放与按需重建设计 +# 动态 MIG 按需分配与回收 -## 背景 +## 目标 -HAMi 的动态 MIG 模式通过调度器选择 MIG 模板，并由 NVIDIA device plugin 在容器启动前创建对应的 MIG 实例。旧实现主要依赖 `nvidia-mig-parted` 对整张 GPU 应用模板：当新任务需要不同模板时，需要等待卡上所有已有 MIG 任务结束后才能重新切分。 +动态 MIG 模式把硬件事实与调度策略分离： -本次改动的目标是把 MIG 实例生命周期从“整卡模板重切”细化到“按 slot 创建和释放”： +- device plugin 通过 NVML 上报 GPU 支持的完整 profile 元数据和 placement。 +- scheduler 根据显存请求、profile allowlist 和当前 placement 占用选择候选实例。 +- device plugin 信任 scheduler 的 `profile + placement` 结果，在 `Allocate` 阶段精确创建 GI/CI。 +- Pod 删除或结束后，根据节点上仍存活 Pod 的 allocation annotation 回收不再使用的实例。 +- device plugin 重启时，通过 annotation 中的 `profile + placement + migUUID` 恢复管理状态。 -- 调度器仍然按照稳定的 MIG slot 视图做资源分配。 -- device plugin 在 `Allocate` 阶段只创建被当前容器实际使用的 MIG GI/CI。 -- 容器结束后，device plugin 通过 kubelet pod-resources API 感知 MIG UUID 释放，并销毁对应 GI/CI。 -- 空闲 slot 保留其模板、位置和 profile 信息，下一个任务命中同一 slot 时可原位重建。 -- 配置层只使用 `profiles` 定义常见 MIG 模板，由代码内置 profile catalog 统一展开。 +该模型不再依赖静态 geometry/count 模板，也不再使用 `templateIdx/slotIdx`。 -## 设计目标 +## Node 能力上报 -1. **减少整卡重切次数** +MIG 模式下，`hami.io/node-nvidia-register` 中的每张 GPU 包含 `migProfiles`。每个 profile 上报： - 对同一 MIG 模板下的任务，只按需创建或销毁单个 slot 对应的 MIG 实例，避免每次分配都重新 apply 整卡模板。 +- profile 名称和 NVML GI profile ID +- 实际可用显存、SM 数量和计算比例 +- slice 数量和最大实例数 +- copy engine、decoder、encoder、JPEG、OFA 等能力 +- NVML 返回的全部合法 placement -2. **保护运行中任务** +例如 A100 40GB 的 `2g.10gb` 通常包含： - 当某张 GPU 上存在其他模板下的活跃 MIG 实例时，device plugin 拒绝切换模板，避免销毁仍被容器使用的 GI/CI。 - -3. **保持调度视图稳定** - - 调度器继续使用 `GPU-UUID[templateIdx-slotIdx]` 形式表达 MIG slot。即使底层 MIG UUID 因销毁和重建发生变化，调度器与 Pod 注解中的 slot 标识仍保持稳定。 - -4. **支持 plugin 重启恢复** - - device plugin 启动时保留忙碌 GPU 的现有 MIG 布局；空闲 GPU 被重置为 “MIG enabled, no partitions”，为后续按需创建留出干净状态。 - -5. **降低默认配置维护成本** +```json +{ + "name": "2g.10gb", + "memoryMB": 9984, + "sliceCount": 2, + "instanceCount": 3, + "placements": [ + {"start": 0, "size": 2}, + {"start": 2, "size": 2}, + {"start": 4, "size": 2} + ] +} +``` - 默认 Helm 配置从冗长的旧模板字段切换为 `profiles` 列表，由代码内置 profile catalog 自动展开为调度器需要的 geometry。 +scheduler 只从该上报构造候选实例，不根据配置中的 count 展开虚拟 slot。 -## 配置模型 +## 配置 -`AllowedMigGeometries` 新增 `profiles` 字段： +配置仅保留 profile allowlist： ```yaml -knownMigGeometries: -- models: [ "A100-SXM4-40GB", "A100-40GB-PCIe", "A100-PCIE-40GB" ] - profiles: [ "1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb" ] +nvidia: + migProfileAllowlist: + - models: ["A100-SXM4-40GB", "A100-40GB-PCIe", "A100-PCIE-40GB"] + profiles: ["1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb"] ``` -归一化逻辑位于 `pkg/device/nvidia/mig_profiles.go`： - -- `profiles` 会被展开为 `[]Geometry`，每个 profile 形成一个单 profile 模板。 -- profile 的 `core`、`memory`、`count` 从内置 catalog 读取。 -- 未设置 `profiles` 的配置会被拒绝，不再兼容旧模板字段。 -- 未知 profile 会返回配置错误，scheduler/device plugin 初始化失败，避免运行时产生不一致调度视图。 - -Helm 默认配置同步简化了 `charts/hami/templates/scheduler/device-configmap.yaml` 中的 NVIDIA `knownMigGeometries`，覆盖 A30、A100、H100、H20、H200、B200 等型号。 +allowlist 决定哪些硬件 profile 可以参与调度，不定义 placement 或实例数量。 -## 核心数据结构 +## Pod allocation annotation -### MIG slot 注解 - -调度器分配的 MIG slot 会写入 Pod 注解 `hami.io/vgpu-mig-slots`： +scheduler 将分配结果写入 `hami.io/vgpu-mig-allocations`： ```json [ { - "deviceUUID": "GPU-xxx[0-1]", + "containerIndex": 0, + "deviceIndex": 0, "gpuUUID": "GPU-xxx", - "templateIdx": 0, - "slotIdx": 1 + "profile": "2g.10gb", + "placement": {"start": 2, "size": 2} } ] ``` -实现位于 `pkg/device/nvidia/mig_slots.go`： - -- `EncodeMigSlotAllocations` 从调度结果中提取 `templateIdx` 和 `slotIdx`。 -- `DecodeMigSlotAllocations` 供 scheduler 和 device plugin 恢复 slot 语义。 - -这个注解解决了一个关键问题：底层 `MIG-*` UUID 是动态创建出来的，不适合作为调度器长期状态；`templateIdx + slotIdx` 才是 HAMi 资源模型里的稳定身份。 - -### MigInstanceManager - -`pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go` 新增 `MigInstanceManager`，作为节点内 MIG GI/CI 状态的单一管理者。 - -主要索引： - -- `bySlot`: `slotKey -> migInstance` -- `byMigUUID`: `MIG UUID -> slotKey` -- `gpuLocks`: 每张 GPU 一个互斥锁，串行化同一物理卡上的 NVML 操作 - -`slotKey` 由三元组组成： - -- `GPUIndex` -- `TemplateIdx` -- `PositionIdx` - -`migInstance` 保存： - -- profile 切片类型，例如 `1g`、`2g` -- NVML placement -- GI ID / CI ID -- 当前绑定的 MIG UUID -- `Present` 状态 - -当 `Present=false` 时，表示 slot 当前没有实际 GI/CI，但 manager 仍记住 profile 和 placement，后续可以原位创建。 - -## 分配流程 +device plugin 创建实例后写回实际 `migUUID`： -1. scheduler 在 `Fit` 阶段选择某张 GPU 的 MIG template 和 slot。 -2. scheduler 通过 Pod 注解写入原有设备分配结果，同时额外写入 `hami.io/vgpu-mig-slots`。 -3. kubelet 调用 device plugin `Allocate`。 -4. `GetContainerDeviceStrArray` 解析 `GPU-UUID[templateIdx-slotIdx]`。 -5. 如果当前 plugin 运行在 `mig` 模式，优先走 `resolveMigUUIDOnDemand`： - - 解析 template 和 slot。 - - 根据 GPU 型号和 template index 查找 geometry。 - - 调用 `PrepareGPU` 准备物理 GPU。 - - 调用 `EnsureSlot` 创建或复用该 slot 的 GI/CI。 - - 返回真实 `MIG-*` UUID 给 kubelet。 -6. 如果按需路径不可用或失败，Allocate 直接返回错误，不再回退到整卡 `nvidia-mig-parted` apply 路径。 - -简化流程如下： - -```mermaid -sequenceDiagram - participant S as Scheduler - participant P as Pod Annotation - participant D as Device Plugin - participant M as MigInstanceManager - participant N as NVML - - S->>P: 写入 GPU-UUID[templateIdx-slotIdx] - S->>P: 写入 hami.io/vgpu-mig-slots - D->>P: Allocate 读取调度结果 - D->>M: PrepareGPU(gpu, template, geometry) - M->>N: 启用 MIG / 采纳现有实例 / 检查模板冲突 - D->>M: EnsureSlot(gpu, template, slot) - M->>N: 创建 GI + CI 或返回已有 MIG UUID - M-->>D: MIG-* UUID - D-->>Kubelet: 返回容器可见设备 +```json +[ + { + "containerIndex": 0, + "deviceIndex": 0, + "gpuUUID": "GPU-xxx", + "profile": "2g.10gb", + "placement": {"start": 2, "size": 2}, + "migUUID": "MIG-xxx" + } +] ``` -## 释放流程 - -device plugin 在 `mig` 模式启动 `podresources.Watcher`，周期性调用 kubelet pod-resources API： - -- socket 挂载路径：`/var/lib/kubelet/pod-resources` -- 默认轮询周期：10 秒 -- 只关注当前 resource name，例如 `nvidia.com/gpu` +当前实现不兼容旧的 `templateIdx/slotIdx` Pod annotation。 -Watcher 保存上一轮 kubelet 视图，并与当前视图做 diff： +## 调度流程 -1. 某个 `MIG-*` device ID 从 kubelet pod-resources 中消失。 -2. 回调 `MigInstanceManager.Release(MIG UUID)`。 -3. manager 通过 `byMigUUID` 找到 slot。 -4. 使用 NVML 销毁对应 CI 和 GI。 -5. 将 slot 标记为 `Present=false`，清理 MIG UUID 反向索引。 +1. scheduler 读取 Node 上报的 profile 和 placement。 +2. 根据 Pod 请求选择能够满足实际显存的最小允许 profile。 +3. 排除与已分配实例重叠的 placement。 +4. 将选中的 `profile + placement` 写入 Pod annotation。 +5. 如果不存在连续且合法的 placement，Pod 保持 Pending，不进入 Bind。 -释放回调还会通过当前节点 Pod 注解重建活跃 slot 集合，并调用 `ReconcileActiveSlots` 清理 manager 中仍 `Present` 但已不属于活跃 Pod 的 stale 实例。 +placement 是否可创建以 NVML 的 `GetGpuInstancePossiblePlacements` 返回结果为准。scheduler 对多个待分配实例进行组合检查，避免局部选择造成后续实例无法放置。 -## 启动恢复策略 +## 创建流程 -device plugin 启动时，`mig_startup.go` 会执行一次 best-effort 检测，找出仍在使用的 GPU： +1. kubelet 调用 device plugin `Allocate`。 +2. device plugin 按 container/device index 读取 scheduler reservation。 +3. 校验 profile 和 placement 是否属于当前 GPU 的 NVML 能力集合。 +4. `MigInstanceManager.EnsureAllocation` 在指定 placement 创建 GI 和 CI，不尝试其他位置。 +5. 将生成的 MIG UUID 返回给容器运行时，并写回 Pod annotation。 -- 通过 kubelet pod-resources List 查询正在使用的 `MIG-*` UUID。 -- 通过 NVML 查询父 GPU 或 MIG device 上的运行中 compute/graphics process。 +同一 GPU 的 NVML 修改由卡级锁串行化，不同 GPU 可以并行操作。 -随后 `resetIdleMigGPUs` 会修改启动 MIG spec： +## 回收流程 -- 忙碌 GPU 保留当前布局，不销毁运行中任务使用的 GI/CI。 -- 空闲 GPU 设置为 `MigEnabled=true` 且 `MigDevices={}`，即开启 MIG 但不预创建分区。 +Pod annotation 是 HAMi allocation 的事实源。device plugin 每 5 秒执行一次 fail-closed reconciliation： -这样可以在 plugin 重启后同时满足两点： +1. 列出调度到当前节点的未结束 Pod。 +2. 从 `hami.io/vgpu-mig-allocations` 构造活跃的 `profile + placement` 集合。 +3. 与 `MigInstanceManager` 中已创建的实例比较。 +4. 销毁不再属于任何活跃 Pod 的 CI 和 GI。 -- 不影响已有容器。 -- 空闲卡回到适合按需创建的初始状态。 +Kubernetes API 查询失败或 annotation 无法解析时跳过本轮回收，保留现有实例，避免误删运行中任务。 -## 调度器状态同步 +动态 MIG UUID 不作为 kubelet 注册的资源 ID，因此 kubelet pod-resources API 无法可靠表达该 UUID。当前实现不再挂载或轮询 pod-resources socket。 -调度器在统计节点用量时新增对 `hami.io/vgpu-mig-slots` 的解析： +## 重启恢复 -- 对带 slot 注解的 Pod，直接根据 `DeviceUUID -> slot` 标记 `MigUsage.UsageList[slotIdx].InUse=true`。 -- 如果设备还没有对应 `MigUsage`，使用 `device.PlatternMIG` 按 `templateIdx` 初始化模板视图。 -- 如果 MIG 任务错误地落在 `hami-core` 模式 GPU 上，标记设备不健康并跳过。 +device plugin 启动时使用两个来源保护已有任务： -这让 scheduler 不依赖底层 MIG UUID 的稳定性，而是依赖 HAMi 分配时写入的 slot 元数据恢复资源占用。 +- 活跃 Pod annotation：确定 Kubernetes 管理的物理 GPU、profile、placement 和 migUUID。 +- NVML 进程查询：保护绕过 Kubernetes 或 annotation 尚未完成写回的活动实例。 -## 失败处理与兼容性 +启动过程： -### 模板冲突 +1. 对没有活跃 allocation、也没有 NVML 进程的 GPU 清理残留 MIG 实例。 +2. 对活跃 Pod annotation 中的实例校验实际 profile、placement 和 migUUID。 +3. 将验证通过的实例采纳到 `MigInstanceManager`。 +4. 后续周期 reconciler 继续清理 stale allocation。 -`PrepareGPU` 会检查同一 GPU 是否存在其他 template 下的活跃 slot： +如果 Kubernetes allocation 状态无法可靠读取，启动清理会 fail closed，保留 GPU 当前布局。 -- 如果存在 `Present=true` 的其他模板实例，则拒绝切换模板。 -- 如果其他模板只剩 absent slot 记录，则清理旧记录并销毁旧布局残留后重建新模板 slot map。 +## 容器运行时 -调度侧将 MIG 模式下的 `CustomFilterRule` 失败原因细化为 `CardMigTopologyInfeasible`，用于区分： - -- 设备资源不足。 -- 自定义过滤失败。 -- MIG 拓扑在不销毁活跃实例的前提下不可行。 - -### kubelet pod-resources 不可用 - -启动探测和 watcher 都对 pod-resources API 失败做降级处理： - -- 启动探测失败时，继续使用 NVML 进程检测结果。 -- watcher tick 失败时保留上一轮快照，避免误触发 release。 -- kubelet socket 重建时，watcher 会关闭旧连接并在下一轮重连。 - -## 部署变更 - -`charts/hami/templates/device-plugin/daemonsetnvidia.yaml` 新增 hostPath 挂载： +动态创建 MIG 实例后，GPU Operator 的 CDI spec 不一定在容器创建前完成刷新。当前验证环境使用 GPU Operator 提供的 `nvidia-legacy` RuntimeClass： ```yaml -- name: pod-resources - hostPath: - path: /var/lib/kubelet/pod-resources - type: Directory +devicePlugin: + runtimeClassName: nvidia-legacy ``` -容器内挂载为只读： +若使用 CDI runtime，需要额外保证创建 MIG 实例后同步刷新 CDI spec，并解决容器创建时序问题。 -```yaml -- name: pod-resources - mountPath: /var/lib/kubelet/pod-resources - readOnly: true -``` +## 测试覆盖 -该挂载是 watcher 访问 kubelet pod-resources unix socket 的前提。 +`hack/hami-mig-e2e.sh` 使用持续执行 CUDA `vectorAdd` 的 Pod 验证： -## 测试覆盖 +- 1g 和 2g 实例同时运行。 +- 混合 profile 达到精确 placement 容量后拒绝溢出请求。 +- device plugin 重启时 MIG UUID 不变，已有 CUDA 负载持续前进。 +- Pod 删除后立即回收并补位，其他任务不受影响。 +- 两个 3g 占满拓扑后，1g 请求在 scheduler 阶段保持 Pending。 +- 七个 1g 并发、部分回收和重新填满。 +- 最终所有 MIG 实例回收为零。 -本分支新增和调整的测试覆盖以下场景： - -- `NormalizeMigGeometries`： - - `profiles` 展开为 geometry。 - - 缺少 `profiles` 的配置会被拒绝。 -- `EncodeMigSlotAllocations` / `DecodeMigSlotAllocations`： - - 从 `GPU-UUID[templateIdx-slotIdx]` 正确生成 slot 注解。 -- `resetIdleMigGPUs`： - - 空闲 GPU 被重置为空 MIG 布局。 - - 忙碌 GPU 保留原布局。 - - 无 devices entry 的 spec 不被修改。 -- `podresources.Watcher`： - - 能从 fake kubelet pod-resources server 读取设备快照。 - - 当前快照缺少上一轮设备时触发 release 回调。 -- scheduler score 测试： - - MIG 拓扑不可行时返回 `CardMigTopologyInfeasible`。 - -## 后续演进 - -- 为 `MigInstanceManager` 增加可注入 NVML backend，降低对真实 GPU 环境的单测依赖。 -- 将 watcher 的轮询周期、超时时间暴露为配置项，适配大规模节点和 kubelet 压力场景。 -- 在 metrics 中区分 slot absent、present、stale、release failed 等状态，便于定位资源回收问题。 -- 为 `CardMigTopologyInfeasible` 增加更明确的事件或调度失败提示，帮助用户理解是 MIG 拓扑限制而不是普通资源不足。 +显存请求应使用 NVML 上报的实际容量。例如 A100 40GB 上 `1g.5gb` 实际为 4864 MiB，测试请求使用 4500 MiB；`2g.10gb` 实际为 9984 MiB，测试请求使用 9500 MiB。 diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go index fea0eadd3e..519f96a66b 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go @@ -13,43 +13,18 @@ package plugin import ( "context" "fmt" - "net" "sort" "strings" - "time" "github.com/NVIDIA/go-nvml/pkg/nvml" - "google.golang.org/grpc" - "google.golang.org/grpc/connectivity" - "google.golang.org/grpc/credentials/insecure" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/klog/v2" - podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/util/client" ) -// Tunables for the one-shot startup poll of kubelet's pod-resources API. -// Separate from the long-lived watcher's dialWait/listTimeout because the -// startup path runs before kubelet is known to have seeded pod state, so -// falling back to NVML-only detection is acceptable. -const ( - migStartupPodresourcesSocket = "/var/lib/kubelet/pod-resources/kubelet.sock" - migStartupDialWait = 5 * time.Second - migStartupListTimeout = 10 * time.Second - migStartupMaxMsgSize = 16 * 1024 * 1024 -) - -func normalizeUnixDialAddr(addr string) string { - trimmed := strings.TrimPrefix(addr, "unix://") - if trimmed != addr { - return trimmed - } - return strings.TrimPrefix(addr, "unix:") -} - // sortedIntSetKeys returns the keys of a set-style map sorted ascending. // Small helper kept here so logging at startup emits stable key order. func sortedIntSetKeys(s map[int]struct{}) []int { @@ -88,28 +63,17 @@ func resetIdleMigGPUs(cfg nvidia.MigConfigSpecSlice, inUse map[int]struct{}) []i } // collectInUseGPUs returns the set of GPU indexes that have at least one -// in-use MIG instance, unioned from two best-effort sources: -// - kubelet's pod-resources List (authoritative for k8s-managed usage). +// in-use MIG instance, unioned from two sources: +// - live Pod allocation annotations (authoritative for HAMi allocations). // - NVML running processes on each MIG instance or the parent card (catches // usage that bypasses kubelet, e.g. bare processes on the node). // -// Failures in either source are logged and downgrade to the other source. -// When both fail, an empty set is returned and the caller treats every GPU -// as idle; the very first apply after a failed detection window will -// reshape cards conservatively because idle means "MIG on, no partitions". -func collectInUseGPUs(ctx context.Context, resourceName, nodeName string) (map[int]struct{}, error) { +// Failure to read Pod annotations is returned to the caller because startup +// reset must not proceed without the authoritative allocation state. NVML +// process detection is an additional safeguard and remains best effort. +func collectInUseGPUs(ctx context.Context, nodeName string) (map[int]struct{}, error) { out := make(map[int]struct{}) - if uuids, err := listPodResourcesMigUUIDs(ctx, resourceName); err != nil { - klog.InfoS("mig init: pod-resources List skipped", "err", err) - } else { - for uuid := range uuids { - if gpu, ok := migUUIDToGPUIndex(uuid); ok { - out[gpu] = struct{}{} - } - } - } - annotated, err := kubernetesAllocatedMigGPUs(ctx, nodeName) if err != nil { return out, fmt.Errorf("list Kubernetes MIG allocations: %w", err) @@ -130,10 +94,8 @@ func collectInUseGPUs(ctx context.Context, resourceName, nodeName string) (map[i } // activeMigGPUUUIDs returns physical GPU UUIDs referenced by live HAMi MIG -// allocations. HAMi exposes virtual resource IDs to kubelet and passes the -// dynamically created MIG UUID through NVIDIA_VISIBLE_DEVICES, so kubelet's -// pod-resources API cannot by itself identify these allocations after a -// device-plugin restart. +// allocations. The annotation preserves the physical GPU identity across a +// device-plugin restart even though the MIG UUID is created at Allocate time. func activeMigGPUUUIDs(pods []corev1.Pod) map[string]struct{} { out := make(map[string]struct{}) for i := range pods { @@ -189,86 +151,6 @@ func gpuUUIDToIndex(gpuUUID string) (int, bool) { return idx, ret == nvml.SUCCESS } -// listPodResourcesMigUUIDs issues a single List on the kubelet pod-resources -// API and returns the set of MIG device IDs currently attached to containers -// under the given resource name. The call uses bounded dial and RPC timeouts -// so a kubelet that isn't accepting yet doesn't block plugin startup. -func listPodResourcesMigUUIDs(ctx context.Context, resourceName string) (map[string]struct{}, error) { - conn, err := grpc.NewClient( - "unix://"+migStartupPodresourcesSocket, - grpc.WithTransportCredentials(insecure.NewCredentials()), - grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(migStartupMaxMsgSize)), - grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) { - return (&net.Dialer{}).DialContext(ctx, "unix", normalizeUnixDialAddr(addr)) - }), - ) - if err != nil { - return nil, err - } - defer conn.Close() - - conn.Connect() - dialCtx, cancelDial := context.WithTimeout(ctx, migStartupDialWait) - defer cancelDial() - for { - s := conn.GetState() - if s == connectivity.Ready { - break - } - if !conn.WaitForStateChange(dialCtx, s) { - return nil, fmt.Errorf("pod-resources dial: %w", dialCtx.Err()) - } - } - - listCtx, cancelList := context.WithTimeout(ctx, migStartupListTimeout) - defer cancelList() - cl := podresourcesv1.NewPodResourcesListerClient(conn) - resp, err := cl.List(listCtx, &podresourcesv1.ListPodResourcesRequest{}) - if err != nil { - return nil, err - } - - out := make(map[string]struct{}) - for _, pod := range resp.GetPodResources() { - for _, c := range pod.GetContainers() { - for _, d := range c.GetDevices() { - if !strings.EqualFold(d.GetResourceName(), resourceName) { - continue - } - for _, id := range d.GetDeviceIds() { - if strings.HasPrefix(id, "MIG-") { - out[id] = struct{}{} - } - } - } - } - } - return out, nil -} - -// migUUIDToGPUIndex resolves a MIG device UUID to its parent GPU's NVML -// index. Missing MIG UUIDs (e.g. stale kubelet state) return false so the -// caller skips them rather than mis-attributing to GPU 0. -func migUUIDToGPUIndex(migUUID string) (int, bool) { - if nvret := nvml.Init(); nvret != nvml.SUCCESS { - klog.InfoS("mig init: nvml.Init failed", "err", nvml.ErrorString(nvret)) - return 0, false - } - migDev, ret := nvml.DeviceGetHandleByUUID(migUUID) - if ret != nvml.SUCCESS { - return 0, false - } - parent, ret := nvml.DeviceGetDeviceHandleFromMigDeviceHandle(migDev) - if ret != nvml.SUCCESS { - return 0, false - } - idx, ret := parent.GetIndex() - if ret != nvml.SUCCESS { - return 0, false - } - return idx, true -} - // nvmlBusyGPUs returns the set of GPU indexes with at least one running // compute or graphics process. For MIG-enabled cards every live MIG instance // is inspected; for non-MIG cards the parent device is inspected directly. diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go index a31312e335..bfbb95a100 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go @@ -270,9 +270,7 @@ func destroyAllMigInstances(dev nvml.Device) error { return nil } -// Release destroys the GI+CI bound to the given MIG UUID and marks the slot -// absent (preserving its profile and placement). Invoked by the podresources -// watcher when kubelet reports the device is no longer in use. +// Release destroys the GI+CI bound to the given MIG UUID. func (m *MigInstanceManager) Release(migUUID string) error { m.mu.Lock() key, ok := m.byAllocationMigUUID[migUUID] diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index 01dfd5e2ea..a35304450a 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -61,7 +61,6 @@ import ( "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/cdi" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/imex" - "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/podresources" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/rm" "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/scheduler/config" @@ -321,11 +320,11 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { if plugin.operatingMode == "mig" { plugin.migMgr = NewMigInstanceManager() if deviceSupportMig { - inUse, detectErr := collectInUseGPUs(plugin.ctx, string(plugin.rm.Resource()), os.Getenv(util.NodeNameEnvName)) + inUse, detectErr := collectInUseGPUs(plugin.ctx, os.Getenv(util.NodeNameEnvName)) if detectErr != nil { // Startup reset is destructive. If Kubernetes allocation state - // cannot be read reliably, preserve every GPU and retry cleanup - // later through the normal reclaim watcher. + // cannot be read reliably, preserve every GPU rather than risk + // removing an allocation that is still active. klog.InfoS("mig init: allocation detection failed; preserving all GPUs", "err", detectErr) for i := 0; i < deviceNumbers; i++ { inUse[i] = struct{}{} @@ -342,20 +341,9 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { klog.InfoS("mig init: failed to adopt active MIG allocations", "err", err) } } - // New allocations are realized lazily from the scheduler's exact - // profile and placement reservation. The watcher's release callback - // tolerates an unknown MIG UUID during early kubelet startup. - watcher := podresources.NewWatcher("", 0, []string{string(plugin.rm.Resource())}, func(_ string, deviceID string) { - if strings.HasPrefix(deviceID, "MIG-") { - if err := plugin.migMgr.Release(deviceID); err != nil { - klog.InfoS("failed to release MIG instance on reclaim", "deviceID", deviceID, "err", err) - } - } - if err := plugin.reconcileActiveMigAllocations(); err != nil { - klog.InfoS("failed to reconcile MIG allocations on reclaim", "deviceID", deviceID, "err", err) - } - }) - go watcher.Run(plugin.ctx) + // Pod annotations are the allocation source of truth. Periodically + // reconcile the manager with live Pods so completed or deleted Pods + // release their exact profile+placement allocation. go plugin.runMigAnnotationReconciler(5 * time.Second) } diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go index 969d3d8eb2..0adcadd006 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go @@ -51,8 +51,8 @@ import ( "github.com/stretchr/testify/require" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" "k8s.io/client-go/kubernetes/fake" + kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" ) func TestCDIAllocateResponse(t *testing.T) { @@ -494,7 +494,6 @@ func TestGetPreferredAllocationAlignsWithAnnotatedDevices(t *testing.T) { Containers: []corev1.Container{{Name: "main"}}, }, } - plugin := &NvidiaDevicePlugin{} t.Setenv(util.NodeNameEnvName, "node-a") previousGetPendingPod := getPendingPod @@ -1012,6 +1011,10 @@ func TestAllocatePreservesContainerOrderWhenOneContainerFallsBack(t *testing.T) }, Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "c0"}, {Name: "c1"}}}, } + fakeClient := fake.NewSimpleClientset(pod.DeepCopy()) + previousKubeClient := client.KubeClient + client.KubeClient = fakeClient + defer func() { client.KubeClient = previousKubeClient }() previousGetPendingPod := getPendingPod getPendingPod = func(context.Context, string) (*corev1.Pod, error) { return pod, nil } diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go index ea0eaac0ff..35d000b2e6 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go @@ -17,8 +17,6 @@ package plugin import ( - "os" - "path/filepath" "strings" "testing" @@ -733,53 +731,3 @@ func TestCheckCDISpec(t *testing.T) { }) } } - -func TestWriteMigConfig(t *testing.T) { - orig := migConfigPath - defer func() { migConfigPath = orig }() - - migConfigPath = filepath.Join(t.TempDir(), "migconfig.yaml") - writeMigConfig([]byte("version: v1")) - - info, err := os.Stat(migConfigPath) - if err != nil { - t.Fatalf("expected config file to exist: %v", err) - } - if perm := info.Mode().Perm(); perm != 0o600 { - t.Errorf("expected permissions 0600, got %o", perm) - } - data, err := os.ReadFile(migConfigPath) - if err != nil { - t.Fatalf("read failed: %v", err) - } - if string(data) != "version: v1" { - t.Errorf("unexpected content: %q", string(data)) - } - - // Write into a missing directory must not panic; the error is only logged. - migConfigPath = filepath.Join(t.TempDir(), "missing", "migconfig.yaml") - writeMigConfig([]byte("x")) - if _, err := os.Stat(migConfigPath); err == nil { - t.Errorf("expected write to fail for missing directory") - } -} - -func TestWriteMigConfig_RemovesStaleFileOnFailure(t *testing.T) { - if os.Geteuid() == 0 { - t.Skip("file permissions do not block root") - } - orig := migConfigPath - defer func() { migConfigPath = orig }() - - migConfigPath = filepath.Join(t.TempDir(), "migconfig.yaml") - if err := os.WriteFile(migConfigPath, []byte("stale"), 0o400); err != nil { - t.Fatalf("setup failed: %v", err) - } - - // Read-only file makes the write fail; the stale file must be removed so - // nvidia-mig-parted cannot pick it up afterwards. - writeMigConfig([]byte("fresh")) - if _, err := os.Stat(migConfigPath); !os.IsNotExist(err) { - t.Errorf("expected stale config to be removed, stat err: %v", err) - } -} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go b/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go deleted file mode 100644 index b15d1bbd51..0000000000 --- a/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher.go +++ /dev/null @@ -1,264 +0,0 @@ -/* - * Copyright (c) 2026, HAMi. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - */ - -// Package podresources polls the kubelet pod-resources API and reports when -// devices that were previously allocated to a container disappear from the -// kubelet's view, which happens when the container ends. -// -// This is the signal HAMi's MIG reclaim path uses to destroy individual GPU -// instances on task completion instead of waiting for every task on a GPU to -// finish before the next nvidia-mig-parted apply can re-shape the card. -package podresources - -import ( - "context" - "net" - "strings" - "time" - - "google.golang.org/grpc" - "google.golang.org/grpc/connectivity" - "google.golang.org/grpc/credentials/insecure" - "k8s.io/klog/v2" - podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" -) - -const ( - defaultSocketPath = "/var/lib/kubelet/pod-resources/kubelet.sock" - defaultPollEvery = 10 * time.Second - // dialWait bounds how long we wait for the kubelet-side socket to reach - // Ready before the first List after (re)connect. The connection itself - // is created lazily by grpc.NewClient; this timeout only gates the - // explicit wait performed before the first RPC on a fresh conn. - dialWait = 5 * time.Second - // listTimeout bounds a single List RPC. This is intentionally smaller - // than the poll interval so a slow kubelet can't cause back-to-back - // ticks to pile up. On very busy nodes with many pods, raising this - // (and the poll interval) is the right lever. - listTimeout = 5 * time.Second - // maxMsgSize matches the kubelet default for this API; on busy nodes the - // List response can exceed the grpc default 4 MiB. - maxMsgSize = 16 * 1024 * 1024 -) - -// ReleaseHandler is invoked once per deviceID that was present in the previous -// snapshot and is missing in the current one. resourceName is the resource the -// device was allocated under (e.g. "nvidia.com/gpu") so a handler serving -// several resources can filter quickly. -type ReleaseHandler func(resourceName, deviceID string) - -// Watcher polls kubelet's pod-resources API and fires a ReleaseHandler when a -// previously-allocated device is no longer in use by any container. -type Watcher struct { - socketPath string - interval time.Duration - - // resourceNames restricts the release callback to devices under these - // resource names. An empty slice means "all resources". - resourceNames []string - - onRelease ReleaseHandler - - // Previous snapshot: set of active deviceIDs indexed by resource name. - prev map[string]map[string]struct{} - - // Long-lived gRPC client; rebuilt on failure. kubelet's pod-resources - // socket is stable across the plugin's lifetime, so reusing the - // connection avoids eating our per-tick budget on DNS/dial/handshake. - conn *grpc.ClientConn - client podresourcesv1.PodResourcesListerClient -} - -// NewWatcher constructs a Watcher. Pass an empty socketPath to use the default -// kubelet location; pass 0 interval to use the default poll cadence. -func NewWatcher(socketPath string, interval time.Duration, resourceNames []string, onRelease ReleaseHandler) *Watcher { - if socketPath == "" { - socketPath = defaultSocketPath - } - if interval <= 0 { - interval = defaultPollEvery - } - return &Watcher{ - socketPath: socketPath, - interval: interval, - resourceNames: resourceNames, - onRelease: onRelease, - prev: make(map[string]map[string]struct{}), - } -} - -// Run polls the kubelet pod-resources API in a loop until ctx is cancelled. -// It never returns an error; transient gRPC failures are logged and the -// previous snapshot is preserved so a missed tick doesn't produce spurious -// release events. -func (w *Watcher) Run(ctx context.Context) { - klog.InfoS("starting podresources watcher", "socket", w.socketPath, "interval", w.interval, "resources", w.resourceNames) - defer w.closeConn() - - // Prime the snapshot before starting to diff. If this first call fails - // we start with an empty map and the first successful tick will just - // record — no spurious release events for the state at plugin start. - if err := w.tick(ctx, true); err != nil { - klog.InfoS("podresources initial List failed; will retry on interval", "err", err) - } - - t := time.NewTicker(w.interval) - defer t.Stop() - for { - select { - case <-ctx.Done(): - return - case <-t.C: - if err := w.tick(ctx, false); err != nil { - klog.InfoS("podresources tick failed; keeping previous snapshot", "err", err) - } - } - } -} - -func (w *Watcher) tick(ctx context.Context, prime bool) error { - if err := w.ensureConn(ctx); err != nil { - return err - } - - callCtx, cancel := context.WithTimeout(ctx, listTimeout) - defer cancel() - - resp, err := w.client.List(callCtx, &podresourcesv1.ListPodResourcesRequest{}) - if err != nil { - // Drop the connection so the next tick reconnects; the kubelet - // socket can be recreated (e.g. kubelet restart) without the - // plugin restarting, and sticking to a dead conn wastes ticks. - w.closeConn() - return err - } - - current := w.collect(resp) - if !prime { - w.diff(current) - } - w.prev = current - return nil -} - -// ensureConn makes sure w.conn/w.client are usable, creating them lazily. -// The first call after (re)connect waits up to dialWait for the channel to -// reach Ready so the subsequent List has the full listTimeout budget. -func (w *Watcher) ensureConn(ctx context.Context) error { - if w.conn != nil { - switch w.conn.GetState() { - case connectivity.Shutdown: - w.closeConn() - default: - return nil - } - } - - conn, err := grpc.NewClient( - "unix://"+w.socketPath, - grpc.WithTransportCredentials(insecure.NewCredentials()), - grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxMsgSize)), - grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) { - return (&net.Dialer{}).DialContext(ctx, "unix", normalizeUnixDialAddr(addr)) - }), - ) - if err != nil { - return err - } - - // Kick the channel and wait briefly for Ready. If the kubelet socket - // isn't accepting yet, let the caller surface a clean error instead - // of eating the full List budget on handshake. - conn.Connect() - waitCtx, cancel := context.WithTimeout(ctx, dialWait) - defer cancel() - for { - s := conn.GetState() - if s == connectivity.Ready { - break - } - if !conn.WaitForStateChange(waitCtx, s) { - _ = conn.Close() - return waitCtx.Err() - } - } - - w.conn = conn - w.client = podresourcesv1.NewPodResourcesListerClient(conn) - return nil -} - -func normalizeUnixDialAddr(addr string) string { - trimmed := strings.TrimPrefix(addr, "unix://") - if trimmed != addr { - return trimmed - } - return strings.TrimPrefix(addr, "unix:") -} - -func (w *Watcher) closeConn() { - if w.conn != nil { - _ = w.conn.Close() - w.conn = nil - w.client = nil - } -} - -// collect flattens the List response into "resourceName -> set(deviceID)". -func (w *Watcher) collect(resp *podresourcesv1.ListPodResourcesResponse) map[string]map[string]struct{} { - out := make(map[string]map[string]struct{}) - for _, pod := range resp.GetPodResources() { - for _, c := range pod.GetContainers() { - for _, d := range c.GetDevices() { - rn := d.GetResourceName() - if !w.resourceMatch(rn) { - continue - } - set, ok := out[rn] - if !ok { - set = make(map[string]struct{}) - out[rn] = set - } - for _, id := range d.GetDeviceIds() { - set[id] = struct{}{} - } - } - } - } - return out -} - -// diff fires onRelease for every deviceID that was present in the previous -// snapshot under a given resource but is absent in the current snapshot. -func (w *Watcher) diff(current map[string]map[string]struct{}) { - for rn, prevSet := range w.prev { - currSet := current[rn] - for id := range prevSet { - if _, stillUsed := currSet[id]; stillUsed { - continue - } - if w.onRelease != nil { - w.onRelease(rn, id) - } - } - } -} - -func (w *Watcher) resourceMatch(rn string) bool { - if len(w.resourceNames) == 0 { - return true - } - for _, want := range w.resourceNames { - if strings.EqualFold(rn, want) { - return true - } - } - return false -} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go deleted file mode 100644 index d35e55f171..0000000000 --- a/pkg/device-plugin/nvidiadevice/nvinternal/podresources/watcher_test.go +++ /dev/null @@ -1,161 +0,0 @@ -package podresources - -import ( - "context" - "net" - "os" - "path/filepath" - "sync" - "testing" - "time" - - "google.golang.org/grpc" - podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" -) - -type fakePodResourcesServer struct { - podresourcesv1.UnimplementedPodResourcesListerServer - mu sync.RWMutex - resp *podresourcesv1.ListPodResourcesResponse -} -func (s *fakePodResourcesServer) setResponse(resp *podresourcesv1.ListPodResourcesResponse) { - s.mu.Lock() - defer s.mu.Unlock() - s.resp = resp -} - -func (s *fakePodResourcesServer) List(context.Context, *podresourcesv1.ListPodResourcesRequest) (*podresourcesv1.ListPodResourcesResponse, error) { - s.mu.RLock() - defer s.mu.RUnlock() - return s.resp, nil -} - -func startFakePodResourcesServer(t *testing.T, resp *podresourcesv1.ListPodResourcesResponse) (socketPath string, setResp func(*podresourcesv1.ListPodResourcesResponse), stop func()) { - t.Helper() - dir := t.TempDir() - socketPath = filepath.Join(dir, "kubelet.sock") - l, err := net.Listen("unix", socketPath) - if err != nil { - t.Fatalf("listen unix socket: %v", err) - } - - srv := grpc.NewServer() - fake := &fakePodResourcesServer{resp: resp} - podresourcesv1.RegisterPodResourcesListerServer(srv, fake) - go func() { - _ = srv.Serve(l) - }() - - setResp = fake.setResponse - stop = func() { - srv.Stop() - _ = l.Close() - _ = os.Remove(socketPath) - } - return socketPath, setResp, stop -} - -func TestTickCanListPodResources(t *testing.T) { - initial := &podresourcesv1.ListPodResourcesResponse{ - PodResources: []*podresourcesv1.PodResources{ - { - Name: "p1", - Namespace: "default", - Containers: []*podresourcesv1.ContainerResources{ - { - Name: "c1", - Devices: []*podresourcesv1.ContainerDevices{ - {ResourceName: "nvidia.com/gpu", DeviceIds: []string{"MIG-a", "MIG-b"}}, - }, - }, - }, - }, - }, - } - socketPath, _, stop := startFakePodResourcesServer(t, initial) - defer stop() - - w := NewWatcher(socketPath, time.Second, []string{"nvidia.com/gpu"}, nil) - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - if err := w.tick(ctx, true); err != nil { - t.Fatalf("tick should succeed, got error: %v", err) - } - - got := w.prev["nvidia.com/gpu"] - if got == nil { - t.Fatalf("expected snapshot for nvidia.com/gpu") - } - if _, ok := got["MIG-a"]; !ok { - t.Fatalf("expected MIG-a in snapshot") - } - if _, ok := got["MIG-b"]; !ok { - t.Fatalf("expected MIG-b in snapshot") - } -} - -func TestTickDiffTriggersRelease(t *testing.T) { - resp1 := &podresourcesv1.ListPodResourcesResponse{ - PodResources: []*podresourcesv1.PodResources{ - { - Name: "p1", - Namespace: "default", - Containers: []*podresourcesv1.ContainerResources{ - { - Name: "c1", - Devices: []*podresourcesv1.ContainerDevices{ - {ResourceName: "nvidia.com/gpu", DeviceIds: []string{"MIG-a", "MIG-b"}}, - }, - }, - }, - }, - }, - } - resp2 := &podresourcesv1.ListPodResourcesResponse{ - PodResources: []*podresourcesv1.PodResources{ - { - Name: "p1", - Namespace: "default", - Containers: []*podresourcesv1.ContainerResources{ - { - Name: "c1", - Devices: []*podresourcesv1.ContainerDevices{ - {ResourceName: "nvidia.com/gpu", DeviceIds: []string{"MIG-a"}}, - }, - }, - }, - }, - }, - } - - socketPath, setResp, stop := startFakePodResourcesServer(t, resp1) - defer stop() - - var ( - mu sync.Mutex - released []string - ) - w := NewWatcher(socketPath, time.Second, []string{"nvidia.com/gpu"}, func(_ string, deviceID string) { - mu.Lock() - defer mu.Unlock() - released = append(released, deviceID) - }) - - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - if err := w.tick(ctx, true); err != nil { - t.Fatalf("prime tick failed: %v", err) - } - setResp(resp2) - if err := w.tick(ctx, false); err != nil { - t.Fatalf("diff tick failed: %v", err) - } - - mu.Lock() - defer mu.Unlock() - if len(released) != 1 || released[0] != "MIG-b" { - t.Fatalf("expected one released device MIG-b, got %v", released) - } -} From 175a5d3b15df47dd2c646908448212ad0cf828dc Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Wed, 5 Aug 2026 17:00:23 +0800 Subject: [PATCH 03/11] update design docs and simplify node report annotation Signed-off-by: Jifei Wang --- docs/develop/mig-dynamic-deallocate.md | 197 ++++++++++++++++++ docs/develop/mig-dynamic-deallocate_cn.md | 161 -------------- .../nvinternal/plugin/register.go | 6 +- pkg/device/devices.go | 54 +++-- pkg/device/devices_test.go | 15 ++ 5 files changed, 239 insertions(+), 194 deletions(-) create mode 100644 docs/develop/mig-dynamic-deallocate.md delete mode 100644 docs/develop/mig-dynamic-deallocate_cn.md diff --git a/docs/develop/mig-dynamic-deallocate.md b/docs/develop/mig-dynamic-deallocate.md new file mode 100644 index 0000000000..35bfe374f3 --- /dev/null +++ b/docs/develop/mig-dynamic-deallocate.md @@ -0,0 +1,197 @@ +# Dynamic MIG Architecture + +## Document Status + +| Item | Value | +| --- | --- | +| Status | Active design | +| Scope | NVIDIA Dynamic MIG in HAMi | +| Audience | Maintainers, reviewers, operators, and contributors | +| Source branch | `feat/mig-dynamic-deallocate` | + +## Context + +NVIDIA MIG divides a physical GPU into hardware-isolated compute instances. GPU models expose profile capacity and placement rules through NVML, while Kubernetes schedules workloads through declarative resources. + +Dynamic MIG connects these layers through a reservation-first architecture. The node publishes hardware capability, the scheduler reserves an exact profile and placement, and the device plugin realizes that reservation during container allocation. Pod metadata carries the allocation identity across the workload lifecycle. + +## Design Goals + +The architecture pursues five outcomes: + +1. Hardware capability originates from the node that owns the GPU. +2. Scheduling decisions include physical MIG placement. +3. Runtime realization follows the scheduler reservation exactly. +4. Workload metadata supports reconciliation and restart recovery. +5. Cross-component contracts remain compact and stable as GPU density grows. + +## Principles + +### Hardware authority + +NVML defines profile capacity and legal placement. Device plugin discovery translates this information into a scheduler-facing capability contract. + +### Explicit reservation + +The scheduler selects the physical GPU, MIG profile, and placement before binding. Each accepted allocation immediately contributes to scheduler occupancy. + +### Clear ownership + +The scheduler owns placement policy. The device plugin owns hardware mutation. Pod annotations provide the durable handoff between both responsibilities. + +### Convergent lifecycle + +A stable reservation key supports idempotent realization. Reconciliation aligns managed hardware with active workload reservations. + +## Architecture + +```text + Kubernetes control plane + + +----------------+ Node capability +----------------+ + | Device Plugin | --------------------------> | HAMi Scheduler | + | | | | + | NVML discovery | Pod reservation | Placement | + | GI/CI manager | <-------------------------- | policy | + | Reconciler | | Capacity model | + +-------+--------+ +--------+-------+ + | | + | exact GI/CI realization | bind + v v + +----------------+ +----------------+ + | NVIDIA GPU | | Workload Pod | + | MIG topology | | Allocation | + | and instances | | annotation | + +----------------+ +----------------+ +``` + +### Device plugin + +The device plugin is the node hardware authority. It discovers allowlisted profiles, publishes compact capability, prepares MIG-ready GPUs, realizes reservations, records runtime MIG UUIDs, adopts active instances after restart, and reconciles managed instances with active Pods. + +### HAMi scheduler + +The scheduler is the policy and reservation authority. It reads node capability, reconstructs topology occupancy, matches workload demand to a profile, selects a legal placement, and persists the reservation before binding. + +### Pod allocation record + +The Pod allocation annotation is the persistent system record connecting scheduling, runtime realization, reconciliation, and restart recovery. + +## Capability Contract + +The device plugin publishes MIG profiles inside the NVIDIA node registration annotation. The contract carries scheduler-relevant fields: + +```json +{ + "name": "2g.10gb", + "memoryMB": 9984, + "core": 29, + "sliceCount": 2, + "placements": [ + {"start": 0, "size": 2}, + {"start": 2, "size": 2}, + {"start": 4, "size": 2} + ] +} +``` + +| Field | Purpose | +| --- | --- | +| `name` | Profile identity across scheduler and device plugin | +| `memoryMB` | Workload capacity matching | +| `core` | Scheduler resource accounting | +| `sliceCount` | Deterministic profile ordering | +| `placements` | Legal topology choices reported by NVML | + +Device-local discovery data stays inside the device plugin process. This compact boundary controls node annotation growth and keeps the shared API aligned with scheduling needs. + +The profile allowlist expresses cluster policy, while NVML supplies capacity and topology. Together they define the capability visible to the scheduler. + +## Reservation Contract + +The scheduler stores reservations in `hami.io/vgpu-mig-allocations`: + +```json +{ + "containerIndex": 0, + "deviceIndex": 0, + "gpuUUID": "GPU-xxxxxxxx", + "profile": "2g.10gb", + "placement": {"start": 2, "size": 2}, + "migUUID": "MIG-xxxxxxxx" +} +``` + +The scheduler writes container identity, parent GPU, profile, and placement. The device plugin adds the MIG UUID after realization. These fields form one allocation identity across control-plane and node lifecycle operations. + +## Core Workflows + +### Capability publication + +The device plugin discovers the allowlisted profile set through NVML and publishes the compact capability contract in the node registration annotation. Scheduler state refresh turns this contract into per-GPU topology capacity. + +### Scheduling and reservation + +The scheduler rebuilds occupancy from active Pod reservations. Each placement occupies the interval `[start, start + size)`. Profile selection follows workload capacity, and placement selection follows deterministic packing. Capacity pressure keeps the workload in the Kubernetes Pending phase. + +### Runtime realization + +During kubelet allocation, the device plugin resolves the reservation and verifies it against current NVML capability. The MIG instance manager serializes mutation per physical GPU, creates the GI and CI at the selected placement, resolves the MIG UUID, and enriches the Pod allocation record. + +The reservation key combines physical GPU, profile, and placement. Repeated allocation requests converge on the same managed instance. + +### Lifecycle reconciliation + +The reconciler derives desired allocations from active Pods assigned to the node. It compares this desired set with managed instances and releases completed workload allocations. Reconciliation runs periodically and during new allocation activity, supporting steady-state cleanup and prompt capacity reuse. + +A complete Kubernetes state snapshot authorizes each cleanup cycle and preserves conservative hardware lifecycle behavior. + +### Restart recovery + +Startup combines active Pod reservations with NVML process activity to identify GPUs carrying active work. Idle GPUs enter a clean MIG-ready state. Active records containing profile, placement, and MIG UUID are verified through NVML and adopted into the new manager process. + +This flow preserves workload identity and re-establishes lifecycle ownership after a device plugin restart. + +## Consistency Model + +Dynamic MIG follows a reservation-first sequence: + +1. Node capability becomes scheduler input. +2. The scheduler persists an exact reservation. +3. The device plugin realizes that reservation. +4. The device plugin enriches the record with runtime identity. +5. Reconciliation converges managed hardware toward active reservations. + +Each phase has one authority and one durable handoff. This separation keeps placement policy independent from NVML mutation while preserving a shared allocation identity. + +## Operational Model + +Dynamic MIG mode is selected per node through device plugin configuration. Workloads select MIG allocation through the `nvidia.com/vgpu-mode: mig` annotation. GPU Operator integration supplies the NVIDIA driver and container runtime path. + +Operational visibility centers on: + +- discovered profiles and placements per GPU; +- scheduler reservation decisions and capacity pressure; +- GI/CI realization and runtime UUID assignment; +- reconciliation activity and capacity reuse; +- startup adoption results; +- workload GPU progress across lifecycle events. + +End-to-end validation uses sustained CUDA workloads across mixed profiles, capacity saturation, workload replacement, device plugin restart, and burst allocation. UUID stability and continuous GPU progress demonstrate lifecycle continuity. + +## Evolution Direction + +The architecture provides stable extension points for: + +- pluggable placement policies; +- node-level fragmentation scoring; +- placement-aware scheduler metrics; +- event-driven lifecycle acceleration; +- CDI synchronization for dynamic MIG identities; +- richer multi-GPU and heterogeneous-node validation. + +The capability contract, reservation contract, and hardware ownership boundary remain the foundation for these extensions. + +## Decision Summary + +Dynamic MIG treats hardware topology as node-owned capability, placement as a scheduler reservation, and GI/CI mutation as a device plugin responsibility. Pod metadata carries allocation identity across the control plane and runtime. This model gives HAMi a deterministic, topology-aware foundation for dynamic MIG scheduling and lifecycle management. diff --git a/docs/develop/mig-dynamic-deallocate_cn.md b/docs/develop/mig-dynamic-deallocate_cn.md deleted file mode 100644 index 9fb7c84492..0000000000 --- a/docs/develop/mig-dynamic-deallocate_cn.md +++ /dev/null @@ -1,161 +0,0 @@ -# 动态 MIG 按需分配与回收 - -## 目标 - -动态 MIG 模式把硬件事实与调度策略分离： - -- device plugin 通过 NVML 上报 GPU 支持的完整 profile 元数据和 placement。 -- scheduler 根据显存请求、profile allowlist 和当前 placement 占用选择候选实例。 -- device plugin 信任 scheduler 的 `profile + placement` 结果，在 `Allocate` 阶段精确创建 GI/CI。 -- Pod 删除或结束后，根据节点上仍存活 Pod 的 allocation annotation 回收不再使用的实例。 -- device plugin 重启时，通过 annotation 中的 `profile + placement + migUUID` 恢复管理状态。 - -该模型不再依赖静态 geometry/count 模板，也不再使用 `templateIdx/slotIdx`。 - -## Node 能力上报 - -MIG 模式下，`hami.io/node-nvidia-register` 中的每张 GPU 包含 `migProfiles`。每个 profile 上报： - -- profile 名称和 NVML GI profile ID -- 实际可用显存、SM 数量和计算比例 -- slice 数量和最大实例数 -- copy engine、decoder、encoder、JPEG、OFA 等能力 -- NVML 返回的全部合法 placement - -例如 A100 40GB 的 `2g.10gb` 通常包含： - -```json -{ - "name": "2g.10gb", - "memoryMB": 9984, - "sliceCount": 2, - "instanceCount": 3, - "placements": [ - {"start": 0, "size": 2}, - {"start": 2, "size": 2}, - {"start": 4, "size": 2} - ] -} -``` - -scheduler 只从该上报构造候选实例，不根据配置中的 count 展开虚拟 slot。 - -## 配置 - -配置仅保留 profile allowlist： - -```yaml -nvidia: - migProfileAllowlist: - - models: ["A100-SXM4-40GB", "A100-40GB-PCIe", "A100-PCIE-40GB"] - profiles: ["1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb"] -``` - -allowlist 决定哪些硬件 profile 可以参与调度，不定义 placement 或实例数量。 - -## Pod allocation annotation - -scheduler 将分配结果写入 `hami.io/vgpu-mig-allocations`： - -```json -[ - { - "containerIndex": 0, - "deviceIndex": 0, - "gpuUUID": "GPU-xxx", - "profile": "2g.10gb", - "placement": {"start": 2, "size": 2} - } -] -``` - -device plugin 创建实例后写回实际 `migUUID`： - -```json -[ - { - "containerIndex": 0, - "deviceIndex": 0, - "gpuUUID": "GPU-xxx", - "profile": "2g.10gb", - "placement": {"start": 2, "size": 2}, - "migUUID": "MIG-xxx" - } -] -``` - -当前实现不兼容旧的 `templateIdx/slotIdx` Pod annotation。 - -## 调度流程 - -1. scheduler 读取 Node 上报的 profile 和 placement。 -2. 根据 Pod 请求选择能够满足实际显存的最小允许 profile。 -3. 排除与已分配实例重叠的 placement。 -4. 将选中的 `profile + placement` 写入 Pod annotation。 -5. 如果不存在连续且合法的 placement，Pod 保持 Pending，不进入 Bind。 - -placement 是否可创建以 NVML 的 `GetGpuInstancePossiblePlacements` 返回结果为准。scheduler 对多个待分配实例进行组合检查，避免局部选择造成后续实例无法放置。 - -## 创建流程 - -1. kubelet 调用 device plugin `Allocate`。 -2. device plugin 按 container/device index 读取 scheduler reservation。 -3. 校验 profile 和 placement 是否属于当前 GPU 的 NVML 能力集合。 -4. `MigInstanceManager.EnsureAllocation` 在指定 placement 创建 GI 和 CI，不尝试其他位置。 -5. 将生成的 MIG UUID 返回给容器运行时，并写回 Pod annotation。 - -同一 GPU 的 NVML 修改由卡级锁串行化，不同 GPU 可以并行操作。 - -## 回收流程 - -Pod annotation 是 HAMi allocation 的事实源。device plugin 每 5 秒执行一次 fail-closed reconciliation： - -1. 列出调度到当前节点的未结束 Pod。 -2. 从 `hami.io/vgpu-mig-allocations` 构造活跃的 `profile + placement` 集合。 -3. 与 `MigInstanceManager` 中已创建的实例比较。 -4. 销毁不再属于任何活跃 Pod 的 CI 和 GI。 - -Kubernetes API 查询失败或 annotation 无法解析时跳过本轮回收，保留现有实例，避免误删运行中任务。 - -动态 MIG UUID 不作为 kubelet 注册的资源 ID，因此 kubelet pod-resources API 无法可靠表达该 UUID。当前实现不再挂载或轮询 pod-resources socket。 - -## 重启恢复 - -device plugin 启动时使用两个来源保护已有任务： - -- 活跃 Pod annotation：确定 Kubernetes 管理的物理 GPU、profile、placement 和 migUUID。 -- NVML 进程查询：保护绕过 Kubernetes 或 annotation 尚未完成写回的活动实例。 - -启动过程： - -1. 对没有活跃 allocation、也没有 NVML 进程的 GPU 清理残留 MIG 实例。 -2. 对活跃 Pod annotation 中的实例校验实际 profile、placement 和 migUUID。 -3. 将验证通过的实例采纳到 `MigInstanceManager`。 -4. 后续周期 reconciler 继续清理 stale allocation。 - -如果 Kubernetes allocation 状态无法可靠读取，启动清理会 fail closed，保留 GPU 当前布局。 - -## 容器运行时 - -动态创建 MIG 实例后，GPU Operator 的 CDI spec 不一定在容器创建前完成刷新。当前验证环境使用 GPU Operator 提供的 `nvidia-legacy` RuntimeClass： - -```yaml -devicePlugin: - runtimeClassName: nvidia-legacy -``` - -若使用 CDI runtime，需要额外保证创建 MIG 实例后同步刷新 CDI spec，并解决容器创建时序问题。 - -## 测试覆盖 - -`hack/hami-mig-e2e.sh` 使用持续执行 CUDA `vectorAdd` 的 Pod 验证： - -- 1g 和 2g 实例同时运行。 -- 混合 profile 达到精确 placement 容量后拒绝溢出请求。 -- device plugin 重启时 MIG UUID 不变，已有 CUDA 负载持续前进。 -- Pod 删除后立即回收并补位，其他任务不受影响。 -- 两个 3g 占满拓扑后，1g 请求在 scheduler 阶段保持 Pending。 -- 七个 1g 并发、部分回收和重新填满。 -- 最终所有 MIG 实例回收为零。 - -显存请求应使用 NVML 上报的实际容量。例如 A100 40GB 上 `1g.5gb` 实际为 4864 MiB，测试请求使用 4500 MiB；`2g.10gb` 实际为 9984 MiB，测试请求使用 9500 MiB。 diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go index 781e126b06..83104e1566 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go @@ -232,12 +232,8 @@ func (plugin *NvidiaDevicePlugin) discoverMigProfiles(dev nvml.Device, model str continue } profile := device.MigProfile{ - Name: profileName, GIProfileID: profileInfo.Id, MemoryMB: int32(profileInfo.MemorySizeMB), + Name: profileName, MemoryMB: int32(profileInfo.MemorySizeMB), SliceCount: profileInfo.SliceCount, InstanceCount: profileInfo.InstanceCount, - MultiprocessorCount: profileInfo.MultiprocessorCount, - P2PSupported: profileInfo.IsP2pSupported != 0, - CopyEngineCount: profileInfo.CopyEngineCount, DecoderCount: profileInfo.DecoderCount, - EncoderCount: profileInfo.EncoderCount, JPEGCount: profileInfo.JpegCount, OFACount: profileInfo.OfaCount, } if fullGPUMultiprocessors > 0 { profile.Core = int32((profileInfo.MultiprocessorCount*100 + fullGPUMultiprocessors - 1) / fullGPUMultiprocessors) diff --git a/pkg/device/devices.go b/pkg/device/devices.go index 072abfa2ae..58683f09bb 100644 --- a/pkg/device/devices.go +++ b/pkg/device/devices.go @@ -77,20 +77,16 @@ type MigInUse struct { } type MigProfile struct { - Name string `json:"name"` - GIProfileID uint32 `json:"giProfileID"` - MemoryMB int32 `json:"memoryMB"` - Core int32 `json:"core"` - SliceCount uint32 `json:"sliceCount"` - InstanceCount uint32 `json:"instanceCount"` - MultiprocessorCount uint32 `json:"multiprocessorCount"` - P2PSupported bool `json:"p2pSupported"` - CopyEngineCount uint32 `json:"copyEngineCount"` - DecoderCount uint32 `json:"decoderCount"` - EncoderCount uint32 `json:"encoderCount"` - JPEGCount uint32 `json:"jpegCount"` - OFACount uint32 `json:"ofaCount"` - Placements []MigPlacement `json:"placements"` + Name string `json:"name"` + MemoryMB int32 `json:"memoryMB"` + Core int32 `json:"core"` + SliceCount uint32 `json:"sliceCount"` + Placements []MigPlacement `json:"placements"` + + // InstanceCount is only used by the device plugin to derive the physical + // GPU's advertised replica count. It is not part of the scheduler wire + // format because placements already describe schedulable MIG capacity. + InstanceCount uint32 `json:"-"` } type MigAllocation struct { @@ -125,20 +121,22 @@ type DeviceUsage struct { } type DeviceInfo struct { - ID string `json:"id,omitempty"` - Index uint `json:"index,omitempty"` - Count int32 `json:"count,omitempty"` - Devmem int32 `json:"devmem,omitempty"` - Devcore int32 `json:"devcore,omitempty"` - Type string `json:"type,omitempty"` - Numa int `json:"numa,omitempty"` - Mode string `json:"mode,omitempty"` - MIGProfiles []MigProfile `json:"migProfiles,omitempty"` - MIGTemplate []Geometry `json:"migtemplate,omitempty"` // Deprecated. - Health bool `json:"health,omitempty"` - DeviceVendor string `json:"devicevendor,omitempty"` - CustomInfo map[string]any `json:"custominfo,omitempty"` - DevicePairScore DevicePairScore `json:"devicepairscore,omitempty"` + ID string `json:"id,omitempty"` + Index uint `json:"index,omitempty"` + Count int32 `json:"count,omitempty"` + Devmem int32 `json:"devmem,omitempty"` + Devcore int32 `json:"devcore,omitempty"` + Type string `json:"type,omitempty"` + Numa int `json:"numa,omitempty"` + Mode string `json:"mode,omitempty"` + MIGProfiles []MigProfile `json:"migProfiles,omitempty"` + MIGTemplate []Geometry `json:"migtemplate,omitempty"` // Deprecated. + Health bool `json:"health,omitempty"` + DeviceVendor string `json:"devicevendor,omitempty"` + CustomInfo map[string]any `json:"custominfo,omitempty"` + // Device pair scores use a dedicated node annotation and must not inflate + // the per-device registration payload. + DevicePairScore DevicePairScore `json:"-"` } type DevicePairScores []DevicePairScore diff --git a/pkg/device/devices_test.go b/pkg/device/devices_test.go index 727b1897d3..305e265fcf 100644 --- a/pkg/device/devices_test.go +++ b/pkg/device/devices_test.go @@ -380,6 +380,21 @@ func TestMarshalNodeDevices(t *testing.T) { } } +func TestMigProfileUsesCompactWireFormat(t *testing.T) { + raw, err := json.Marshal(MigProfile{ + Name: "1g.5gb", + MemoryMB: 4864, + Core: 15, + SliceCount: 1, + InstanceCount: 7, + Placements: []MigPlacement{{Start: 6, Size: 1}}, + }) + assert.NilError(t, err) + + want := `{"name":"1g.5gb","memoryMB":4864,"core":15,"sliceCount":1,"placements":[{"start":6,"size":1}]}` + assert.Equal(t, string(raw), want) +} + func TestUnMarshalNodeDevices(t *testing.T) { type args struct { str string From 76e046ace3d7ab0388f6ba025088e9538a376e5a Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Wed, 5 Aug 2026 18:31:50 +0800 Subject: [PATCH 04/11] update metrics to match dcgm Signed-off-by: Jifei Wang --- cmd/scheduler/metrics.go | 41 +++++++--- cmd/scheduler/metrics_test.go | 31 ++++++++ cmd/vGPUmonitor/metrics.go | 75 +++++++++++-------- cmd/vGPUmonitor/metrics_mig_test.go | 47 ++++++++++++ docs/develop/mig-dynamic-deallocate.md | 10 ++- hack/hami-mig-e2e.sh | 2 +- .../nvidiadevice/nvinternal/plugin/migmgr.go | 18 +++-- .../nvidiadevice/nvinternal/plugin/server.go | 19 ++++- pkg/device/devices.go | 8 +- pkg/device/nvidia/mig_allocations.go | 33 +++++--- pkg/device/nvidia/mig_profiles_test.go | 21 ++++++ pkg/scheduler/mig_allocation_test.go | 4 +- pkg/scheduler/scheduler.go | 13 +++- 13 files changed, 252 insertions(+), 70 deletions(-) create mode 100644 cmd/vGPUmonitor/metrics_mig_test.go diff --git a/cmd/scheduler/metrics.go b/cmd/scheduler/metrics.go index 5412256225..00e0764bed 100644 --- a/cmd/scheduler/metrics.go +++ b/cmd/scheduler/metrics.go @@ -141,8 +141,8 @@ func (cc ClusterManagerCollector) collectNodeMetrics(ch chan<- prometheus.Metric ) nodeGPUMigInstance := prometheus.NewDesc( "hami_node_gpu_mig_instance_info", - "GPU Sharing mode. 0 for hami-core, 1 for mig, 2 for mps", - []string{"node", "device_uuid", "device_index", "mig_name"}, nil, + "Realized MIG instance identity and scheduler placement", + []string{"node", "device_uuid", "device_index", "mig_uuid", "profile", "gpu_instance_id", "compute_instance_id", "placement_start", "placement_size"}, nil, ) // Legacy metric descriptors (only created when legacy mode is enabled) @@ -203,17 +203,40 @@ func (cc ClusterManagerCollector) collectNodeMetrics(ch chan<- prometheus.Metric for _, devs := range val.Devices.DeviceLists { coreLimit, coreAllocated := normalizeAMDCoreMetrics(devs.Device.Type, devs.Device.Totalcore, devs.Device.Usedcores) if devs.Device.Mode == "mig" { - for idx, migs := range devs.Device.MigUsage.UsageList { - klog.V(3).Infoln("mig instances=", devs.Device.MigUsage) - inuse := 0 - if migs.InUse { - inuse = 1 + for _, allocation := range devs.Device.MigAllocationsInUse { + if !allocation.RuntimeReady { + continue } - if err := sendMetric(ch, nodeGPUMigInstance, prometheus.GaugeValue, float64(inuse), nodeID, devs.Device.ID, fmt.Sprint(devs.Device.Index), migs.Name+"-"+fmt.Sprint(idx)); err != nil { + klog.V(3).InfoS("MIG instance allocation", + "profile", allocation.Profile, + "gpuInstanceID", allocation.GPUInstanceID, + "computeInstanceID", allocation.ComputeInstanceID, + "migUUID", allocation.MigUUID) + if err := sendMetric( + ch, + nodeGPUMigInstance, + prometheus.GaugeValue, + 1, + nodeID, + devs.Device.ID, + fmt.Sprint(devs.Device.Index), + allocation.MigUUID, + allocation.Profile, + fmt.Sprint(allocation.GPUInstanceID), + fmt.Sprint(allocation.ComputeInstanceID), + fmt.Sprint(allocation.Placement.Start), + fmt.Sprint(allocation.Placement.Size), + ); err != nil { klog.V(4).Infof("Failed to send nodeGPUMigInstance metric: %v", err) } if legacy { - sendLegacyMetric(ch, legacyMigInstance, prometheus.GaugeValue, float64(inuse), nodeID, devs.Device.ID, fmt.Sprint(devs.Device.Index), migs.Name+"-"+fmt.Sprint(idx)) + sendLegacyMetric( + ch, + legacyMigInstance, + prometheus.GaugeValue, + 1, + nodeID, devs.Device.ID, fmt.Sprint(devs.Device.Index), allocation.Profile+"-"+fmt.Sprint(allocation.GPUInstanceID), + ) } } } diff --git a/cmd/scheduler/metrics_test.go b/cmd/scheduler/metrics_test.go index 3fcff7813c..0c7a1f9adf 100644 --- a/cmd/scheduler/metrics_test.go +++ b/cmd/scheduler/metrics_test.go @@ -377,3 +377,34 @@ nodeGPUMemoryPercentage{deviceidx="2",deviceuuid="normal-memory",nodeid="node-1" t.Fatalf("unexpected collecting result:\n%s", err) } } + +func TestClusterManagerCollectorExposesDynamicMigRuntimeIdentity(t *testing.T) { + nodeUsage := map[string]*schedulerpkg.NodeUsage{ + "node-1": { + Devices: policy.DeviceUsageList{DeviceLists: []*policy.DeviceListsScore{{ + Device: &device.DeviceUsage{ + ID: "GPU-parent", Index: 0, Mode: "mig", Type: "NVIDIA", + MigAllocationsInUse: []device.MigAllocation{{ + Profile: "2g.10gb", Placement: device.MigPlacement{Start: 0, Size: 2}, + MigUUID: "MIG-runtime", GPUInstanceID: 4, ComputeInstanceID: 0, RuntimeReady: true, + }}, + }, + }}}, + }, + } + collector := ClusterManagerCollector{ + ClusterManager: &ClusterManager{}, + metricsProvider: &fakeMetricsProvider{ + nodeUsage: nodeUsage, quotaManager: device.NewQuotaManager(), podManager: device.NewPodManager(), + }, + } + + want := ` +# HELP hami_node_gpu_mig_instance_info Realized MIG instance identity and scheduler placement +# TYPE hami_node_gpu_mig_instance_info gauge +hami_node_gpu_mig_instance_info{compute_instance_id="0",device_index="0",device_uuid="GPU-parent",gpu_instance_id="4",mig_uuid="MIG-runtime",node="node-1",placement_size="2",placement_start="0",profile="2g.10gb"} 1 +` + if err := promtestutil.CollectAndCompare(collector, strings.NewReader(want), "hami_node_gpu_mig_instance_info"); err != nil { + t.Fatalf("unexpected collecting result:\n%s", err) + } +} diff --git a/cmd/vGPUmonitor/metrics.go b/cmd/vGPUmonitor/metrics.go index f943a6349e..f93e3c803b 100644 --- a/cmd/vGPUmonitor/metrics.go +++ b/cmd/vGPUmonitor/metrics.go @@ -19,12 +19,9 @@ package main import ( "fmt" "os" - "strings" "time" "unicode/utf8" - "github.com/Project-HAMi/HAMi/pkg/device" - dp "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/plugin" nv "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia" "github.com/Project-HAMi/HAMi/pkg/util" @@ -98,8 +95,8 @@ var ( ) ctrDeviceMigInfo = prometheus.NewDesc( "hami_mig_device_info", - "MIG device information for container", - []string{"namespace", "pod", "container", "vdevice_index", "device_uuid", "instance_id"}, nil, + "MIG runtime identity for a container allocation", + []string{"namespace", "pod", "container", "vdevice_index", "device_uuid", "mig_uuid", "profile", "gpu_instance_id", "compute_instance_id"}, nil, ) ctrDeviceMemoryContextDesc = prometheus.NewDesc( "hami_vgpu_memory_context_bytes", @@ -515,43 +512,55 @@ func (cc ClusterManagerCollector) collectPodAndContainerMigInfo(ch chan<- promet return fmt.Errorf("failed to list pods: %w", err) } for _, pod := range pods { - pdevices, err := device.DecodePodDevices(device.SupportDevices, pod.Annotations) + allocations, err := nv.DecodeMigAllocations(pod.Annotations[nv.MigAllocationsAnnotation]) if err != nil { - klog.Errorf("failed to decode pod devices for pod %s/%s: %v", pod.Namespace, pod.Name, err) + klog.Errorf("failed to decode MIG allocations for pod %s/%s: %v", pod.Namespace, pod.Name, err) continue } - for ctrIdx, container := range pod.Spec.Containers { - for ctrDevIdx, ctrDevices := range pdevices[nv.NvidiaGPUDevice] { - if len(ctrDevices) == 0 || ctrIdx != ctrDevIdx { - continue - } - for _, ctrDev := range ctrDevices { - if strings.Contains(ctrDev.UUID, "[") { - uuid := strings.Split(ctrDev.UUID, "[")[0] - _, idx, err := device.ExtractMigTemplatesFromUUID(ctrDev.UUID) - if err != nil { - klog.Errorf("Failed to get mig template for device %s in Pod %s/%s, container %s: %v", ctrDev.UUID, pod.Namespace, pod.Name, container.Name, err) - continue - } - gpuInstanceId, err := dp.GetMigGpuInstanceIdFromIndex(ctrDev.UUID, idx) - if err != nil { - klog.Errorf("Failed to get mig InstanceId for device %s in Pod %s/%s, container %s: %v", ctrDev.UUID, pod.Namespace, pod.Name, container.Name, err) - continue - } - labels := []string{pod.Namespace, pod.Name, container.Name, fmt.Sprint(idx), uuid, fmt.Sprint(gpuInstanceId)} - if err := sendMetric(ch, ctrDeviceMigInfo, prometheus.GaugeValue, 1, labels...); err != nil { - klog.Errorf("Failed to send mig info metric for device %s in Pod %s/%s, container %s: %v", ctrDev.UUID, pod.Namespace, pod.Name, container.Name, err) - return err - } - sendLegacyMetric(ch, legacyCtrDeviceMigInfo, prometheus.GaugeValue, 1, labels...) - } - } + for _, allocation := range allocations { + if allocation.MigUUID == "" || allocation.GPUInstanceID == nil || allocation.ComputeInstanceID == nil { + continue + } + containerName, ok := migAllocationContainerName(pod, allocation.ContainerIndex) + if !ok { + klog.Errorf("MIG allocation container index %d out of range for Pod %s/%s", allocation.ContainerIndex, pod.Namespace, pod.Name) + continue + } + metricLabels := []string{ + pod.Namespace, + pod.Name, + containerName, + fmt.Sprint(allocation.DeviceIndex), + allocation.GPUUUID, + allocation.MigUUID, + allocation.Profile, + fmt.Sprint(*allocation.GPUInstanceID), + fmt.Sprint(*allocation.ComputeInstanceID), + } + if err := sendMetric(ch, ctrDeviceMigInfo, prometheus.GaugeValue, 1, metricLabels...); err != nil { + return fmt.Errorf("send MIG info metric for pod %s/%s: %w", pod.Namespace, pod.Name, err) } + sendLegacyMetric(ch, legacyCtrDeviceMigInfo, prometheus.GaugeValue, 1, + pod.Namespace, pod.Name, containerName, fmt.Sprint(allocation.DeviceIndex), allocation.GPUUUID, fmt.Sprint(*allocation.GPUInstanceID)) } } return nil } +func migAllocationContainerName(pod *corev1.Pod, containerIndex int) (string, bool) { + if containerIndex < 0 { + return "", false + } + if containerIndex < len(pod.Spec.InitContainers) { + return pod.Spec.InitContainers[containerIndex].Name, true + } + containerIndex -= len(pod.Spec.InitContainers) + if containerIndex >= len(pod.Spec.Containers) { + return "", false + } + return pod.Spec.Containers[containerIndex].Name, true +} + func sendMetric(ch chan<- prometheus.Metric, desc *prometheus.Desc, valueType prometheus.ValueType, value float64, labels ...string) error { metric, err := prometheus.NewConstMetric(desc, valueType, value, labels...) if err != nil { diff --git a/cmd/vGPUmonitor/metrics_mig_test.go b/cmd/vGPUmonitor/metrics_mig_test.go new file mode 100644 index 0000000000..4ce5a16a33 --- /dev/null +++ b/cmd/vGPUmonitor/metrics_mig_test.go @@ -0,0 +1,47 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "testing" + + corev1 "k8s.io/api/core/v1" +) + +func TestMigAllocationContainerName(t *testing.T) { + pod := &corev1.Pod{Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{{Name: "init"}}, + Containers: []corev1.Container{{Name: "main-a"}, {Name: "main-b"}}, + }} + tests := []struct { + index int + name string + ok bool + }{ + {index: 0, name: "init", ok: true}, + {index: 1, name: "main-a", ok: true}, + {index: 2, name: "main-b", ok: true}, + {index: 3, ok: false}, + {index: -1, ok: false}, + } + for _, test := range tests { + name, ok := migAllocationContainerName(pod, test.index) + if name != test.name || ok != test.ok { + t.Fatalf("index %d: got (%q, %t), want (%q, %t)", test.index, name, ok, test.name, test.ok) + } + } +} diff --git a/docs/develop/mig-dynamic-deallocate.md b/docs/develop/mig-dynamic-deallocate.md index 35bfe374f3..d654759f9f 100644 --- a/docs/develop/mig-dynamic-deallocate.md +++ b/docs/develop/mig-dynamic-deallocate.md @@ -118,11 +118,13 @@ The scheduler stores reservations in `hami.io/vgpu-mig-allocations`: "gpuUUID": "GPU-xxxxxxxx", "profile": "2g.10gb", "placement": {"start": 2, "size": 2}, - "migUUID": "MIG-xxxxxxxx" + "migUUID": "MIG-xxxxxxxx", + "gpuInstanceID": 4, + "computeInstanceID": 0 } ``` -The scheduler writes container identity, parent GPU, profile, and placement. The device plugin adds the MIG UUID after realization. These fields form one allocation identity across control-plane and node lifecycle operations. +The scheduler writes container identity, parent GPU, profile, and placement. The device plugin adds the MIG UUID, GPU Instance ID, and Compute Instance ID after realization. These fields form one allocation identity across control-plane and node lifecycle operations. The parent GPU UUID and GPU Instance ID provide a direct correlation key for DCGM metrics carrying `UUID` and `GPU_I_ID` labels. ## Core Workflows @@ -136,7 +138,7 @@ The scheduler rebuilds occupancy from active Pod reservations. Each placement oc ### Runtime realization -During kubelet allocation, the device plugin resolves the reservation and verifies it against current NVML capability. The MIG instance manager serializes mutation per physical GPU, creates the GI and CI at the selected placement, resolves the MIG UUID, and enriches the Pod allocation record. +During kubelet allocation, the device plugin resolves the reservation and verifies it against current NVML capability. The MIG instance manager serializes mutation per physical GPU, creates the GI and CI at the selected placement, resolves their runtime IDs, and enriches the Pod allocation record. The reservation key combines physical GPU, profile, and placement. Repeated allocation requests converge on the same managed instance. @@ -172,7 +174,7 @@ Operational visibility centers on: - discovered profiles and placements per GPU; - scheduler reservation decisions and capacity pressure; -- GI/CI realization and runtime UUID assignment; +- GI/CI realization and runtime identity assignment; - reconciliation activity and capacity reuse; - startup adoption results; - workload GPU progress across lifecycle events. diff --git a/hack/hami-mig-e2e.sh b/hack/hami-mig-e2e.sh index 2b42dbc47c..7062cfbe13 100644 --- a/hack/hami-mig-e2e.sh +++ b/hack/hami-mig-e2e.sh @@ -87,7 +87,7 @@ assert_uuid() { assert_runtime_annotation() { local pod=$1 raw raw=$(kubectl get pod -n "$NS" "$pod" -o json | jq -r '.metadata.annotations["hami.io/vgpu-mig-allocations"]') - jq -e 'length > 0 and all(.[]; (.migUUID | startswith("MIG-")) and (.profile | length > 0) and (.placement.size > 0))' <<<"$raw" >/dev/null || fail "$pod lacks concrete MIG runtime placement annotation" + jq -e 'length > 0 and all(.[]; (.migUUID | startswith("MIG-")) and (.profile | length > 0) and (.placement.size > 0) and (.gpuInstanceID | type == "number") and (.computeInstanceID | type == "number"))' <<<"$raw" >/dev/null || fail "$pod lacks concrete MIG runtime identity annotation" } assert_pending_unbound() { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go index bfbb95a100..0ed811c455 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go @@ -388,7 +388,7 @@ func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, plac m.byAllocation[key] = inst m.byAllocationMigUUID[migUUID] = key m.mu.Unlock() - klog.InfoS("created scheduler-reserved MIG allocation", "uuid", migUUID, "gpu", gpuIndex, "profile", profile, "start", placement.Start, "size", placement.Size) + klog.InfoS("created scheduler-reserved MIG allocation", "uuid", migUUID, "gpu", gpuIndex, "profile", profile, "start", placement.Start, "size", placement.Size, "gpuInstanceID", giData.Id, "computeInstanceID", ciData.Id) return migUUID, nil } @@ -400,10 +400,16 @@ func (m *MigInstanceManager) AllocationRuntimeInfo(gpuIndex int, profile string, if inst == nil || !inst.Present { return migAllocationRuntimeInfo{}, false } - return migAllocationRuntimeInfo{MigUUID: inst.MigUUID, Profile: inst.Profile, Placement: inst.Placement}, true + return migAllocationRuntimeInfo{ + MigUUID: inst.MigUUID, + Profile: inst.Profile, + Placement: inst.Placement, + GIID: inst.GIID, + CIID: inst.CIID, + }, true } -func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID string, placement nvml.GpuInstancePlacement) error { +func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID string, placement nvml.GpuInstancePlacement, gpuInstanceID, computeInstanceID uint32) error { lk := m.gpuLock(gpuIndex) lk.Lock() defer lk.Unlock() @@ -426,7 +432,7 @@ func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID stri ciProfileID := profileIDToCIProfileID(giProfileID) for _, gi := range instances { giInfo, r := gi.GetInfo() - if r != nvml.SUCCESS || giInfo.Placement != placement { + if r != nvml.SUCCESS || giInfo.Placement != placement || giInfo.Id != gpuInstanceID { continue } actualUUID, findErr := findMigUUIDForGI(dev, giInfo.Id) @@ -442,7 +448,7 @@ func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID stri continue } ciData, r := cis[0].GetInfo() - if r != nvml.SUCCESS { + if r != nvml.SUCCESS || ciData.Id != computeInstanceID { continue } key := allocationKey(gpuIndex, profile, placement) @@ -493,6 +499,8 @@ type migAllocationRuntimeInfo struct { MigUUID string Profile string Placement nvml.GpuInstancePlacement + GIID uint32 + CIID uint32 } func findMigUUIDForGI(dev nvml.Device, giID uint32) (string, error) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index a35304450a..b6874b98ec 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -388,6 +388,8 @@ func (plugin *NvidiaDevicePlugin) annotateMigRuntimeInfo(pod *corev1.Pod) error continue // A later Allocate call may own another container's allocation. } allocations[i].MigUUID = info.MigUUID + allocations[i].GPUInstanceID = uint32Ptr(info.GIID) + allocations[i].ComputeInstanceID = uint32Ptr(info.CIID) updated = true } if !updated { @@ -404,6 +406,10 @@ func (plugin *NvidiaDevicePlugin) annotateMigRuntimeInfo(pod *corev1.Pod) error return nil } +func uint32Ptr(value uint32) *uint32 { + return &value +} + func (plugin *NvidiaDevicePlugin) reconcileActiveMigAllocations() error { active, err := plugin.listActiveMigAllocationKeys() if err != nil { @@ -439,14 +445,21 @@ func (plugin *NvidiaDevicePlugin) primeMigManagerFromAnnotations(_ []string) err return err } for _, allocation := range allocations { - if allocation.MigUUID == "" { - return fmt.Errorf("active pod %s/%s MIG allocation lacks runtime UUID", pod.Namespace, pod.Name) + if allocation.MigUUID == "" || allocation.GPUInstanceID == nil || allocation.ComputeInstanceID == nil { + return fmt.Errorf("active pod %s/%s MIG allocation lacks runtime identity", pod.Namespace, pod.Name) } gpuIndex, ok := gpuUUIDToIndex(allocation.GPUUUID) if !ok { return fmt.Errorf("resolve active MIG parent GPU %q", allocation.GPUUUID) } - if err := plugin.migMgr.AdoptAllocation(gpuIndex, allocation.Profile, allocation.MigUUID, nvml.GpuInstancePlacement{Start: allocation.Placement.Start, Size: allocation.Placement.Size}); err != nil { + if err := plugin.migMgr.AdoptAllocation( + gpuIndex, + allocation.Profile, + allocation.MigUUID, + nvml.GpuInstancePlacement{Start: allocation.Placement.Start, Size: allocation.Placement.Size}, + *allocation.GPUInstanceID, + *allocation.ComputeInstanceID, + ); err != nil { return fmt.Errorf("adopt pod %s/%s MIG allocation: %w", pod.Namespace, pod.Name, err) } } diff --git a/pkg/device/devices.go b/pkg/device/devices.go index 58683f09bb..8a1e87d8db 100644 --- a/pkg/device/devices.go +++ b/pkg/device/devices.go @@ -90,8 +90,12 @@ type MigProfile struct { } type MigAllocation struct { - Profile string `json:"profile"` - Placement MigPlacement `json:"placement"` + Profile string + Placement MigPlacement + MigUUID string + GPUInstanceID uint32 + ComputeInstanceID uint32 + RuntimeReady bool } type AllowedMigProfiles struct { diff --git a/pkg/device/nvidia/mig_allocations.go b/pkg/device/nvidia/mig_allocations.go index fd31b4daa1..85689fc1c0 100644 --- a/pkg/device/nvidia/mig_allocations.go +++ b/pkg/device/nvidia/mig_allocations.go @@ -15,16 +15,18 @@ const ( ) // MigAllocation is the complete scheduler reservation for one MIG device. -// The device plugin only realizes this exact hardware placement and fills in -// MigUUID. ContainerIndex and DeviceIndex make repeated allocations on the -// same physical GPU unambiguous without synthesizing logical device IDs. +// The device plugin realizes this exact hardware placement and fills in its +// MIG UUID plus GI/CI runtime identity. ContainerIndex and DeviceIndex make +// repeated allocations on the same physical GPU unambiguous. type MigAllocation struct { - ContainerIndex int `json:"containerIndex"` - DeviceIndex int `json:"deviceIndex"` - GPUUUID string `json:"gpuUUID"` - Profile string `json:"profile"` - Placement device.MigPlacement `json:"placement"` - MigUUID string `json:"migUUID,omitempty"` + ContainerIndex int `json:"containerIndex"` + DeviceIndex int `json:"deviceIndex"` + GPUUUID string `json:"gpuUUID"` + Profile string `json:"profile"` + Placement device.MigPlacement `json:"placement"` + MigUUID string `json:"migUUID,omitempty"` + GPUInstanceID *uint32 `json:"gpuInstanceID,omitempty"` + ComputeInstanceID *uint32 `json:"computeInstanceID,omitempty"` } func EncodeMigAllocations(pd device.PodSingleDevice) (string, bool) { @@ -71,6 +73,19 @@ func DecodeMigAllocations(raw string) ([]MigAllocation, error) { if allocation.ContainerIndex < 0 || allocation.DeviceIndex < 0 || allocation.GPUUUID == "" || allocation.Profile == "" || allocation.Placement.Size == 0 { return nil, fmt.Errorf("MIG allocation %d is incomplete", i) } + runtimeFields := 0 + if allocation.MigUUID != "" { + runtimeFields++ + } + if allocation.GPUInstanceID != nil { + runtimeFields++ + } + if allocation.ComputeInstanceID != nil { + runtimeFields++ + } + if runtimeFields != 0 && runtimeFields != 3 { + return nil, fmt.Errorf("MIG allocation %d has partial runtime identity", i) + } key := [2]int{allocation.ContainerIndex, allocation.DeviceIndex} if _, ok := seen[key]; ok { return nil, fmt.Errorf("duplicate MIG allocation for container %d device %d", allocation.ContainerIndex, allocation.DeviceIndex) diff --git a/pkg/device/nvidia/mig_profiles_test.go b/pkg/device/nvidia/mig_profiles_test.go index a7b51d7576..aa9d6c7025 100644 --- a/pkg/device/nvidia/mig_profiles_test.go +++ b/pkg/device/nvidia/mig_profiles_test.go @@ -48,3 +48,24 @@ func TestDecodeMigAllocationsRejectsDuplicateDeviceIndex(t *testing.T) { t.Fatal("duplicate container/device allocation index should fail") } } + +func TestDecodeMigAllocationsPreservesZeroRuntimeInstanceIDs(t *testing.T) { + raw := `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"2g.10gb","placement":{"start":0,"size":2},"migUUID":"MIG-a","gpuInstanceID":4,"computeInstanceID":0}]` + allocations, err := DecodeMigAllocations(raw) + if err != nil { + t.Fatalf("decode allocation: %v", err) + } + if len(allocations) != 1 || allocations[0].GPUInstanceID == nil || allocations[0].ComputeInstanceID == nil { + t.Fatalf("runtime identity missing: %+v", allocations) + } + if *allocations[0].GPUInstanceID != 4 || *allocations[0].ComputeInstanceID != 0 { + t.Fatalf("runtime identity mismatch: %+v", allocations[0]) + } +} + +func TestDecodeMigAllocationsRejectsPartialRuntimeIdentity(t *testing.T) { + raw := `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"2g.10gb","placement":{"start":0,"size":2},"migUUID":"MIG-a"}]` + if _, err := DecodeMigAllocations(raw); err == nil { + t.Fatal("partial runtime identity should fail") + } +} diff --git a/pkg/scheduler/mig_allocation_test.go b/pkg/scheduler/mig_allocation_test.go index e680b3b7a5..87d7955061 100644 --- a/pkg/scheduler/mig_allocation_test.go +++ b/pkg/scheduler/mig_allocation_test.go @@ -22,7 +22,7 @@ func TestGetNodesUsageRestoresMigAllocationByProfileAndPlacement(t *testing.T) { pods := device.NewPodManager() pods.AddPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{ UID: "pod-1", Name: "pod-1", Namespace: "default", - Annotations: map[string]string{nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":6,"size":1},"migUUID":"MIG-a"}]`}, + Annotations: map[string]string{nvidia.MigAllocationsAnnotation: `[{"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":6,"size":1},"migUUID":"MIG-a","gpuInstanceID":7,"computeInstanceID":0}]`}, }}, "node1", device.PodDevices{nvidia.NvidiaGPUDevice: {{{UUID: "GPU-a", Usedmem: 5120, Usedcores: 14}}}}) s := Scheduler{nodeManager: nodes, podManager: pods} nodeNames := []string{"node1"} @@ -31,7 +31,7 @@ func TestGetNodesUsageRestoresMigAllocationByProfileAndPlacement(t *testing.T) { t.Fatal(err) } allocations := (*usage)["node1"].Devices.DeviceLists[0].Device.MigAllocationsInUse - if len(allocations) != 1 || allocations[0].Profile != "1g.5gb" || allocations[0].Placement != (device.MigPlacement{Start: 6, Size: 1}) { + if len(allocations) != 1 || allocations[0].Profile != "1g.5gb" || allocations[0].Placement != (device.MigPlacement{Start: 6, Size: 1}) || !allocations[0].RuntimeReady || allocations[0].GPUInstanceID != 7 || allocations[0].ComputeInstanceID != 0 || allocations[0].MigUUID != "MIG-a" { t.Fatalf("restored allocations: %+v", allocations) } } diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 905ca5ba2c..334bb4bc01 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -750,9 +750,18 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin } allocation := allocations[0] allocationsByGPU[udevice.UUID] = allocations[1:] - d.Device.MigAllocationsInUse = append(d.Device.MigAllocationsInUse, device.MigAllocation{ + migAllocation := device.MigAllocation{ Profile: allocation.Profile, Placement: allocation.Placement, - }) + MigUUID: allocation.MigUUID, + RuntimeReady: allocation.MigUUID != "" && allocation.GPUInstanceID != nil && allocation.ComputeInstanceID != nil, + } + if allocation.GPUInstanceID != nil { + migAllocation.GPUInstanceID = *allocation.GPUInstanceID + } + if allocation.ComputeInstanceID != nil { + migAllocation.ComputeInstanceID = *allocation.ComputeInstanceID + } + d.Device.MigAllocationsInUse = append(d.Device.MigAllocationsInUse, migAllocation) continue } if d.Device.Mode == nvidia.MigMode { From 72278affe8a5f6a376f7ff0ee6a2c17443945d43 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Wed, 5 Aug 2026 19:19:33 +0800 Subject: [PATCH 05/11] add license Signed-off-by: Jifei Wang --- docs/develop/mig-dynamic-deallocate.md | 9 --------- hack/hami-mig-e2e.sh | 14 ++++++++++++++ .../nvinternal/plugin/migmgr_placement_test.go | 16 ++++++++++++++++ pkg/device/nvidia/mig_allocations.go | 16 ++++++++++++++++ pkg/device/nvidia/mig_capability_test.go | 16 ++++++++++++++++ pkg/device/nvidia/mig_profiles.go | 16 ++++++++++++++++ pkg/device/nvidia/mig_profiles_test.go | 16 ++++++++++++++++ pkg/device/nvidia/mig_topology.go | 16 ++++++++++++++++ pkg/device/nvidia/mig_topology_test.go | 16 ++++++++++++++++ pkg/scheduler/mig_allocation_test.go | 16 ++++++++++++++++ 10 files changed, 142 insertions(+), 9 deletions(-) diff --git a/docs/develop/mig-dynamic-deallocate.md b/docs/develop/mig-dynamic-deallocate.md index d654759f9f..b3cc757af9 100644 --- a/docs/develop/mig-dynamic-deallocate.md +++ b/docs/develop/mig-dynamic-deallocate.md @@ -1,14 +1,5 @@ # Dynamic MIG Architecture -## Document Status - -| Item | Value | -| --- | --- | -| Status | Active design | -| Scope | NVIDIA Dynamic MIG in HAMi | -| Audience | Maintainers, reviewers, operators, and contributors | -| Source branch | `feat/mig-dynamic-deallocate` | - ## Context NVIDIA MIG divides a physical GPU into hardware-isolated compute instances. GPU models expose profile capacity and placement rules through NVML, while Kubernetes schedules workloads through declarative resources. diff --git a/hack/hami-mig-e2e.sh b/hack/hami-mig-e2e.sh index 7062cfbe13..baeca40a19 100644 --- a/hack/hami-mig-e2e.sh +++ b/hack/hami-mig-e2e.sh @@ -1,4 +1,18 @@ #!/usr/bin/env bash +# Copyright 2026 The HAMi Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + set -euo pipefail NS=hami-mig-final-e2e diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go index 7e697bff33..a549e7237d 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package plugin import ( diff --git a/pkg/device/nvidia/mig_allocations.go b/pkg/device/nvidia/mig_allocations.go index 85689fc1c0..e476cbe69a 100644 --- a/pkg/device/nvidia/mig_allocations.go +++ b/pkg/device/nvidia/mig_allocations.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package nvidia import ( diff --git a/pkg/device/nvidia/mig_capability_test.go b/pkg/device/nvidia/mig_capability_test.go index a8c3c4e68b..23e2bf80d9 100644 --- a/pkg/device/nvidia/mig_capability_test.go +++ b/pkg/device/nvidia/mig_capability_test.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package nvidia import ( diff --git a/pkg/device/nvidia/mig_profiles.go b/pkg/device/nvidia/mig_profiles.go index cec0ab5eec..35a852df82 100644 --- a/pkg/device/nvidia/mig_profiles.go +++ b/pkg/device/nvidia/mig_profiles.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package nvidia import ( diff --git a/pkg/device/nvidia/mig_profiles_test.go b/pkg/device/nvidia/mig_profiles_test.go index aa9d6c7025..fe70a379ee 100644 --- a/pkg/device/nvidia/mig_profiles_test.go +++ b/pkg/device/nvidia/mig_profiles_test.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package nvidia import ( diff --git a/pkg/device/nvidia/mig_topology.go b/pkg/device/nvidia/mig_topology.go index 4dbc017357..faaa0d42ea 100644 --- a/pkg/device/nvidia/mig_topology.go +++ b/pkg/device/nvidia/mig_topology.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package nvidia import ( diff --git a/pkg/device/nvidia/mig_topology_test.go b/pkg/device/nvidia/mig_topology_test.go index 8363b70ba3..f7a33e1acf 100644 --- a/pkg/device/nvidia/mig_topology_test.go +++ b/pkg/device/nvidia/mig_topology_test.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package nvidia import ( diff --git a/pkg/scheduler/mig_allocation_test.go b/pkg/scheduler/mig_allocation_test.go index 87d7955061..857871ca09 100644 --- a/pkg/scheduler/mig_allocation_test.go +++ b/pkg/scheduler/mig_allocation_test.go @@ -1,3 +1,19 @@ +/* +Copyright 2026 The HAMi Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + package scheduler import ( From dde32fce13e2fa093da3436b05d5ae5512c75d72 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Wed, 5 Aug 2026 19:49:41 +0800 Subject: [PATCH 06/11] add import Signed-off-by: Jifei Wang --- pkg/device/nvidia/device_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/device/nvidia/device_test.go b/pkg/device/nvidia/device_test.go index bc820f607d..3b963e63b3 100644 --- a/pkg/device/nvidia/device_test.go +++ b/pkg/device/nvidia/device_test.go @@ -19,6 +19,7 @@ package nvidia import ( "context" "errors" + "strings" "testing" "time" From c7fcb651e22b3caf970d17d9f043b7d60664a0f3 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Thu, 6 Aug 2026 14:46:56 +0800 Subject: [PATCH 07/11] clean up mig-parted Signed-off-by: Jifei Wang --- docker/Dockerfile | 2 - docker/Dockerfile.hamimaster | 2 - docker/Dockerfile.withlib | 2 - .../nvinternal/plugin/mig_startup.go | 26 -- .../nvinternal/plugin/mig_startup_test.go | 67 ----- .../nvidiadevice/nvinternal/plugin/migmgr.go | 2 +- .../nvinternal/plugin/register_test.go | 235 ++-------------- .../nvidiadevice/nvinternal/plugin/server.go | 57 ---- .../nvinternal/plugin/server_test.go | 253 ++++-------------- .../nvidiadevice/nvinternal/plugin/util.go | 16 -- pkg/device/nvidia/device.go | 16 -- 11 files changed, 79 insertions(+), 599 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index d6f77e0a0f..f3af53cc90 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -8,7 +8,6 @@ ADD . /k8s-vgpu #RUN --mount=type=cache,target=/go/pkg/mod \ # cd /k8s-vgpu && make all RUN cd /k8s-vgpu && make all VERSION=$VERSION -RUN go install github.com/NVIDIA/mig-parted/cmd/nvidia-mig-parted@v0.14.4 FROM $NVIDIA_IMAGE AS nvbuild RUN dnf install -y cmake git && dnf clean all @@ -35,7 +34,6 @@ LABEL version="$VERSION" LABEL maintainer="info@dynamia.ai" COPY ./LICENSE /k8s-vgpu/LICENSE COPY --from=gobuild /k8s-vgpu/bin /k8s-vgpu/bin -COPY --from=gobuild /go/bin/nvidia-mig-parted /k8s-vgpu/bin/ COPY ./docker/entrypoint.sh /k8s-vgpu/bin/entrypoint.sh COPY ./lib /k8s-vgpu/lib COPY --from=nvbuild /libvgpu/build/libvgpu.so /k8s-vgpu/lib/nvidia/libvgpu.so."$VERSION" diff --git a/docker/Dockerfile.hamimaster b/docker/Dockerfile.hamimaster index ca83173c2a..18f898b8fe 100644 --- a/docker/Dockerfile.hamimaster +++ b/docker/Dockerfile.hamimaster @@ -8,7 +8,6 @@ ADD . /k8s-vgpu ARG VERSION RUN go env -w GO111MODULE=on RUN cd /k8s-vgpu && make all VERSION=$VERSION -RUN go install github.com/NVIDIA/mig-parted/cmd/nvidia-mig-parted@v0.14.4 FROM nvidia/cuda:13.3.1-base-ubuntu22.04@sha256:ea2a388e99273207c6389b3098a61ebca5ab41d31168dfc4350d0cfa614703e4 RUN apt-get update && \ @@ -26,7 +25,6 @@ LABEL version="$VERSION" LABEL maintainer="opensource@4paradigm.com" COPY ./LICENSE /k8s-vgpu/LICENSE COPY --from=GOBUILD /k8s-vgpu/bin /k8s-vgpu/bin -COPY --from=GOBUILD /go/bin/nvidia-mig-parted /k8s-vgpu/bin/ COPY ./docker/entrypoint.sh /k8s-vgpu/bin/entrypoint.sh COPY ./docker/vgpu-init.sh /k8s-vgpu/bin/vgpu-init.sh COPY ./lib /k8s-vgpu/lib diff --git a/docker/Dockerfile.withlib b/docker/Dockerfile.withlib index 829b6d9ede..fd7ae94140 100644 --- a/docker/Dockerfile.withlib +++ b/docker/Dockerfile.withlib @@ -8,7 +8,6 @@ ARG GOPROXY=https://goproxy.cn,direct ARG VERSION RUN go env -w GO111MODULE=on RUN cd /k8s-vgpu && make all VERSION=$VERSION -RUN go install github.com/NVIDIA/mig-parted/cmd/nvidia-mig-parted@v0.14.4 FROM nvidia/cuda:13.3.1-base-ubuntu22.04@sha256:ea2a388e99273207c6389b3098a61ebca5ab41d31168dfc4350d0cfa614703e4 RUN apt-get update && \ @@ -26,7 +25,6 @@ LABEL version="$VERSION" LABEL maintainer="info@dynamia.ai" COPY ./LICENSE /k8s-vgpu/LICENSE COPY --from=GOBUILD /k8s-vgpu/bin /k8s-vgpu/bin -COPY --from=GOBUILD /go/bin/nvidia-mig-parted /k8s-vgpu/bin/ COPY ./docker/entrypoint.sh /k8s-vgpu/bin/entrypoint.sh COPY ./docker/vgpu-init.sh /k8s-vgpu/bin/vgpu-init.sh COPY ./lib /k8s-vgpu/lib diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go index 519f96a66b..7ae55ec486 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup.go @@ -36,32 +36,6 @@ func sortedIntSetKeys(s map[int]struct{}) []int { return out } -// resetIdleMigGPUs edits the per-device MIG spec in place: GPUs that show no -// sign of in-use compute are returned to "MIG-on with no partitions" so the -// on-demand migMgr path can reshape them per request without destroying live -// GIs on busy cards. Returns the set of GPU indexes that were reset. -// -// cfg is expected to already be in per-device form (one MigConfigSpec per -// Devices=[i]); this is how processMigConfigs arranges it. -func resetIdleMigGPUs(cfg nvidia.MigConfigSpecSlice, inUse map[int]struct{}) []int { - reset := []int{} - for i := range cfg { - devs := cfg[i].Devices - if len(devs) == 0 { - continue - } - gpu := int(devs[0]) - if _, busy := inUse[gpu]; busy { - continue - } - cfg[i].MigEnabled = true - cfg[i].MigDevices = map[string]int32{} - reset = append(reset, gpu) - } - sort.Ints(reset) - return reset -} - // collectInUseGPUs returns the set of GPU indexes that have at least one // in-use MIG instance, unioned from two sources: // - live Pod allocation annotations (authoritative for HAMi allocations). diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go index d22ce48db4..81aa044f0f 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go @@ -20,73 +20,6 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) -func TestResetIdleMigGPUs(t *testing.T) { - tests := []struct { - name string - input nvidia.MigConfigSpecSlice - inUse map[int]struct{} - wantReset []int - wantLayout nvidia.MigConfigSpecSlice - }{ - { - name: "all idle gets reset to MIG-on empty partitions", - input: nvidia.MigConfigSpecSlice{ - {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{"1g.5gb": 4}}, - {MigEnabled: false, Devices: []int32{1}, MigDevices: map[string]int32{}}, - }, - inUse: map[int]struct{}{}, - wantReset: []int{0, 1}, - wantLayout: nvidia.MigConfigSpecSlice{ - {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{}}, - {MigEnabled: true, Devices: []int32{1}, MigDevices: map[string]int32{}}, - }, - }, - { - name: "busy gpu is preserved, idle is reset", - input: nvidia.MigConfigSpecSlice{ - {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{"1g.5gb": 4}}, - {MigEnabled: true, Devices: []int32{1}, MigDevices: map[string]int32{"3g.20gb": 2}}, - }, - inUse: map[int]struct{}{0: {}}, - wantReset: []int{1}, - wantLayout: nvidia.MigConfigSpecSlice{ - {MigEnabled: true, Devices: []int32{0}, MigDevices: map[string]int32{"1g.5gb": 4}}, - {MigEnabled: true, Devices: []int32{1}, MigDevices: map[string]int32{}}, - }, - }, - { - name: "spec entry without devices is left alone", - input: nvidia.MigConfigSpecSlice{ - {MigEnabled: true, Devices: nil, MigDevices: map[string]int32{"1g.5gb": 4}}, - }, - inUse: map[int]struct{}{}, - wantReset: []int{}, - wantLayout: nvidia.MigConfigSpecSlice{ - {MigEnabled: true, Devices: nil, MigDevices: map[string]int32{"1g.5gb": 4}}, - }, - }, - { - name: "empty input", - input: nvidia.MigConfigSpecSlice{}, - inUse: map[int]struct{}{}, - wantReset: []int{}, - wantLayout: nvidia.MigConfigSpecSlice{}, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got := resetIdleMigGPUs(tc.input, tc.inUse) - if !reflect.DeepEqual(got, tc.wantReset) { - t.Errorf("resetIdleMigGPUs reset set = %v, want %v", got, tc.wantReset) - } - if !reflect.DeepEqual(tc.input, tc.wantLayout) { - t.Errorf("resetIdleMigGPUs mutated layout = %+v, want %+v", tc.input, tc.wantLayout) - } - }) - } -} - func TestSortedIntSetKeys(t *testing.T) { got := sortedIntSetKeys(map[int]struct{}{3: {}, 0: {}, 1: {}}) want := []int{0, 1, 3} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go index 0ed811c455..f749122859 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go @@ -96,7 +96,7 @@ func profileSliceKey(profile string) string { } // ResetIdleGPUs prepares idle MIG-capable GPUs for on-demand slot creation -// without calling nvidia-mig-parted. Busy GPUs are left untouched; idle GPUs +// through NVML. Busy GPUs are left untouched; idle GPUs // have MIG mode enabled and all existing GI/CI instances destroyed. func (m *MigInstanceManager) ResetIdleGPUs(deviceCount int, inUse map[int]struct{}) ([]int, error) { reset := []int{} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register_test.go index 496b477aaa..1e37fa9c41 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register_test.go @@ -21,10 +21,9 @@ import ( "time" "github.com/NVIDIA/go-nvml/pkg/nvml" - mock "github.com/NVIDIA/go-nvml/pkg/nvml/mock" + nvmlmock "github.com/NVIDIA/go-nvml/pkg/nvml/mock" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/rm" - "github.com/Project-HAMi/HAMi/pkg/device/nvidia" ) func TestInt8SliceString(t *testing.T) { @@ -87,230 +86,54 @@ func TestInt8SliceString(t *testing.T) { func TestGetNumaNode(t *testing.T) { t.Run("error getting PCI info", func(t *testing.T) { - d := &mock.Device{ - GetPciInfoFunc: func() (nvml.PciInfo, nvml.Return) { - return nvml.PciInfo{}, nvml.ERROR_UNKNOWN - }, - } - hasNode, node, err := GetNumaNode(d) - if err == nil { - t.Fatal("expected an error getting PCI info") - } - if hasNode || node != 0 { - t.Errorf("got hasNode=%v node=%v, want false/0", hasNode, node) + dev := &nvmlmock.Device{GetPciInfoFunc: func() (nvml.PciInfo, nvml.Return) { + return nvml.PciInfo{}, nvml.ERROR_UNKNOWN + }} + hasNode, node, err := GetNumaNode(dev) + if err == nil || hasNode || node != 0 { + t.Fatalf("GetNumaNode() = (%v, %v, %v), want (false, 0, error)", hasNode, node, err) } }) - t.Run("numa_node file not present for this bus ID", func(t *testing.T) { - // GetPciInfo succeeds but no /sys/bus/pci/devices entry exists for this bus ID on the test host. - d := &mock.Device{ - GetPciInfoFunc: func() (nvml.PciInfo, nvml.Return) { - return nvml.PciInfo{BusId: [32]int8{'0', '0', '0', '0', '0', '0', '0', '0', ':', '0', '2', ':', '0', '0', '.', '0'}}, nvml.SUCCESS - }, - } - hasNode, node, err := GetNumaNode(d) - if err == nil { - t.Fatal("expected an error reading the (nonexistent) numa_node file") - } - if hasNode || node != 0 { - t.Errorf("got hasNode=%v node=%v, want false/0", hasNode, node) + t.Run("numa node file is absent", func(t *testing.T) { + dev := &nvmlmock.Device{GetPciInfoFunc: func() (nvml.PciInfo, nvml.Return) { + return nvml.PciInfo{BusId: [32]int8{'0', '0', '0', '0', '0', '0', '0', '0', ':', '0', '2', ':', '0', '0', '.', '0'}}, nvml.SUCCESS + }} + hasNode, node, err := GetNumaNode(dev) + if err == nil || hasNode || node != 0 { + t.Fatalf("GetNumaNode() = (%v, %v, %v), want (false, 0, error)", hasNode, node, err) } }) } -// TestGetAPIDevices_PanicsOnNVMLInitFailure verifies getAPIDevices panics when nvml.Init fails, using a stubbed nvmlInit so the test doesn't need a real NVIDIA driver. -func TestGetAPIDevices_PanicsOnNVMLInitFailure(t *testing.T) { - origInit := nvmlInit +func TestGetAPIDevicesPanicsOnNVMLInitFailure(t *testing.T) { + originalInit := nvmlInit nvmlInit = func() nvml.Return { return nvml.ERROR_LIBRARY_NOT_FOUND } - defer func() { nvmlInit = origInit }() - - mockRM := &rm.ResourceManagerMock{ - DevicesFunc: func() rm.Devices { return rm.Devices{} }, - } - plugin := &NvidiaDevicePlugin{rm: mockRM} + defer func() { nvmlInit = originalInit }() + plugin := &NvidiaDevicePlugin{rm: &rm.ResourceManagerMock{DevicesFunc: func() rm.Devices { return rm.Devices{} }}} defer func() { - if r := recover(); r == nil { - t.Fatal("expected getAPIDevices to panic when nvml.Init fails") + if recover() == nil { + t.Fatal("getAPIDevices did not panic when NVML initialization failed") } }() plugin.getAPIDevices() } -// TestGetAPIDevices_ShutsDownOnNVMLInitSuccess verifies nvml.Shutdown is deferred only after nvmlInit succeeds, using stubbed seams so the real NVML library is never touched. -func TestGetAPIDevices_ShutsDownOnNVMLInitSuccess(t *testing.T) { - origInit := nvmlInit - origShutdown := nvml.Shutdown +func TestGetAPIDevicesShutsDownAfterNVMLInit(t *testing.T) { + originalInit := nvmlInit + originalShutdown := nvml.Shutdown nvmlInit = func() nvml.Return { return nvml.SUCCESS } nvml.Shutdown = func() nvml.Return { return nvml.SUCCESS } defer func() { - nvmlInit = origInit - nvml.Shutdown = origShutdown + nvmlInit = originalInit + nvml.Shutdown = originalShutdown }() - mockRM := &rm.ResourceManagerMock{ - DevicesFunc: func() rm.Devices { return rm.Devices{} }, - } - plugin := &NvidiaDevicePlugin{rm: mockRM} - - got := plugin.getAPIDevices() - if got == nil || len(*got) != 0 { - t.Fatalf("got %v, want a non-nil empty slice", got) - } -} - -func TestProcessMigConfigs(t *testing.T) { - plugin := &NvidiaDevicePlugin{} - - tests := []struct { - name string - migConfigs map[string]nvidia.MigConfigSpecSlice - deviceCount int - expectErr bool - expectedLen int - validateResult func(t *testing.T, result nvidia.MigConfigSpecSlice) - }{ - { - name: "nil migConfigs returns error", - migConfigs: nil, - deviceCount: 2, - expectErr: true, - }, - { - name: "zero deviceCount returns error", - migConfigs: map[string]nvidia.MigConfigSpecSlice{"current": {}}, - deviceCount: 0, - expectErr: true, - }, - { - name: "negative deviceCount returns error", - migConfigs: map[string]nvidia.MigConfigSpecSlice{"current": {}}, - deviceCount: -1, - expectErr: true, - }, - { - name: "single config with empty devices expands to all devices", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{}, - MigEnabled: true, - MigDevices: map[string]int32{"1g.5gb": 7}, - }, - }, - }, - deviceCount: 3, - expectErr: false, - expectedLen: 3, - validateResult: func(t *testing.T, result nvidia.MigConfigSpecSlice) { - for i, cfg := range result { - if len(cfg.Devices) != 1 || cfg.Devices[0] != int32(i) { - t.Errorf("config[%d].Devices = %v, want [%d]", i, cfg.Devices, i) - } - if !cfg.MigEnabled { - t.Errorf("config[%d].MigEnabled = false, want true", i) - } - if cfg.MigDevices["1g.5gb"] != 7 { - t.Errorf("config[%d].MigDevices[1g.5gb] = %d, want 7", i, cfg.MigDevices["1g.5gb"]) - } - } - }, - }, - { - name: "multiple configs with explicit device mapping", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{0, 1}, - MigEnabled: true, - MigDevices: map[string]int32{"1g.5gb": 7}, - }, - nvidia.MigConfigSpec{ - Devices: []int32{2}, - MigEnabled: true, - MigDevices: map[string]int32{"2g.10gb": 3}, - }, - }, - }, - deviceCount: 3, - expectErr: false, - expectedLen: 3, - validateResult: func(t *testing.T, result nvidia.MigConfigSpecSlice) { - // Device 0 should get 1g.5gb config - if result[0].MigDevices["1g.5gb"] != 7 { - t.Errorf("device 0: MigDevices[1g.5gb] = %d, want 7", result[0].MigDevices["1g.5gb"]) - } - if len(result[0].Devices) != 1 || result[0].Devices[0] != 0 { - t.Errorf("device 0: Devices = %v, want [0]", result[0].Devices) - } - // Device 1 should get 1g.5gb config - if result[1].MigDevices["1g.5gb"] != 7 { - t.Errorf("device 1: MigDevices[1g.5gb] = %d, want 7", result[1].MigDevices["1g.5gb"]) - } - // Device 2 should get 2g.10gb config - if result[2].MigDevices["2g.10gb"] != 3 { - t.Errorf("device 2: MigDevices[2g.10gb] = %d, want 3", result[2].MigDevices["2g.10gb"]) - } - }, - }, - { - name: "device not found in config returns error", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{0}, - MigEnabled: true, - MigDevices: map[string]int32{"1g.5gb": 7}, - }, - }, - }, - deviceCount: 3, - expectErr: true, - }, - { - name: "single device single config", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{0}, - MigEnabled: false, - MigDevices: map[string]int32{}, - }, - }, - }, - deviceCount: 1, - expectErr: false, - expectedLen: 1, - validateResult: func(t *testing.T, result nvidia.MigConfigSpecSlice) { - if result[0].MigEnabled { - t.Error("config[0].MigEnabled = true, want false") - } - if len(result[0].Devices) != 1 || result[0].Devices[0] != 0 { - t.Errorf("config[0].Devices = %v, want [0]", result[0].Devices) - } - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result, err := plugin.processMigConfigs(tt.migConfigs, tt.deviceCount) - if tt.expectErr { - if err == nil { - t.Fatal("expected error, got nil") - } - return - } - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if len(result) != tt.expectedLen { - t.Fatalf("result length = %d, want %d", len(result), tt.expectedLen) - } - if tt.validateResult != nil { - tt.validateResult(t, result) - } - }) + plugin := &NvidiaDevicePlugin{rm: &rm.ResourceManagerMock{DevicesFunc: func() rm.Devices { return rm.Devices{} }}} + devices := plugin.getAPIDevices() + if devices == nil || len(*devices) != 0 { + t.Fatalf("getAPIDevices() = %v, want non-nil empty slice", devices) } } diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index b6874b98ec..c2f3fc2b8b 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -110,7 +110,6 @@ type NvidiaDevicePlugin struct { cdiAnnotationPrefix string operatingMode string - migCurrent nvidia.MigPartedSpec deviceCache string // migMgr tracks live MIG GI+CI instances so we can destroy and recreate @@ -206,7 +205,6 @@ func (o *options) devicePluginForResource(ctx context.Context, nvconfig *nvidia. cdiAnnotationPrefix: *o.config.Flags.Plugin.CDIAnnotationPrefix, schedulerConfig: sConfig.NvidiaConfig, operatingMode: mode, - migCurrent: nvidia.MigPartedSpec{}, deviceCache: "", // These will be reinitialized every @@ -243,17 +241,6 @@ func (plugin *NvidiaDevicePlugin) Devices() rm.Devices { return plugin.rm.Devices() } -// BuildFallbackMigConfig - fallback to non-MIG mode -func (plugin *NvidiaDevicePlugin) buildFallbackMigConfig(deviceNumbers int) { - plugin.migCurrent.MigConfigs = make(map[string]nvidia.MigConfigSpecSlice) - configSlice := nvidia.MigConfigSpecSlice{} - for i := 0; i < deviceNumbers; i++ { - conf := nvidia.MigConfigSpec{MigEnabled: false, Devices: []int32{int32(i)}} - configSlice = append(configSlice, conf) - } - plugin.migCurrent.MigConfigs["current"] = configSlice -} - // Start starts the gRPC server, registers the device plugin with the Kubelet,and starts the device healthchecks. func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { plugin.initialize() @@ -1120,47 +1107,3 @@ func (plugin *NvidiaDevicePlugin) apiDevices() []*kubeletdevicepluginv1beta1.Dev numaTopology := plugin.schedulerConfig.EnableNUMATopology != nil && *plugin.schedulerConfig.EnableNUMATopology return plugin.Devices().GetPluginDevices(*plugin.schedulerConfig.DeviceSplitCount, numaTopology) } - -func (plugin *NvidiaDevicePlugin) processMigConfigs(migConfigs map[string]nvidia.MigConfigSpecSlice, deviceCount int) (nvidia.MigConfigSpecSlice, error) { - if migConfigs == nil { - return nil, fmt.Errorf("migConfigs cannot be nil") - } - if deviceCount <= 0 { - return nil, fmt.Errorf("deviceCount must be positive") - } - - transformConfigs := func() (nvidia.MigConfigSpecSlice, error) { - var result nvidia.MigConfigSpecSlice - - if len(migConfigs["current"]) == 1 && len(migConfigs["current"][0].Devices) == 0 { - for i := 0; i < deviceCount; i++ { - config := deepCopyMigConfig(migConfigs["current"][0]) - config.Devices = []int32{int32(i)} - result = append(result, config) - } - return result, nil - } - - deviceToConfig := make(map[int32]*nvidia.MigConfigSpec) - for i := range migConfigs["current"] { - for _, device := range migConfigs["current"][i].Devices { - deviceToConfig[device] = &migConfigs["current"][i] - } - } - - for i := 0; i < deviceCount; i++ { - deviceIndex := int32(i) - config, exists := deviceToConfig[deviceIndex] - if !exists { - return nil, fmt.Errorf("device %d does not match any MIG configuration", i) - } - newConfig := deepCopyMigConfig(*config) - newConfig.Devices = []int32{deviceIndex} - result = append(result, newConfig) - - } - return result, nil - } - - return transformConfigs() -} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go index 0adcadd006..0f208a9ea1 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go @@ -49,6 +49,7 @@ import ( "github.com/Project-HAMi/HAMi/pkg/util" "github.com/Project-HAMi/HAMi/pkg/util/client" "github.com/stretchr/testify/require" + "google.golang.org/grpc" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes/fake" @@ -200,205 +201,6 @@ func TestCDIAllocateResponse(t *testing.T) { } } -func ptr[T any](x T) *T { - return &x -} - -type MigDeviceConfigs struct { - Configs []map[string]int32 -} - -func Test_processMigConfigs(t *testing.T) { - type testCase struct { - name string - migConfigs map[string]nvidia.MigConfigSpecSlice - deviceCount int - expectError bool - validate func(t *testing.T, result nvidia.MigConfigSpecSlice) - } - - testConfigs := MigDeviceConfigs{ - Configs: []map[string]int32{ - { - "1g.10gb": 4, - "2g.20gb": 1, - }, - { - "3g.30gb": 2, - }, - {}, - }, - } - - testCases := []testCase{ - { - name: "SingleConfigForAllDevices", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{}, - MigEnabled: true, - MigDevices: testConfigs.Configs[1], - }, - }, - }, - deviceCount: 3, - expectError: false, - validate: func(t *testing.T, result nvidia.MigConfigSpecSlice) { - if len(result) != 3 { - t.Errorf("Expected 3 configs, got %d", len(result)) - } - for i, config := range result { - if len(config.Devices) != 1 || config.Devices[0] != int32(i) { - t.Errorf("Config for device %d is incorrect: %v", i, config) - } - if !config.MigEnabled { - t.Error("MigEnabled should be true") - } - if len(config.MigDevices) != 1 || config.MigDevices["3g.30gb"] != 2 { - t.Error("MigDevices not preserved correctly") - } - } - }, - }, - { - name: "MultipleConfigsForSpecificDevicesWithNoEnabled", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{0, 1}, - MigEnabled: true, - MigDevices: testConfigs.Configs[0], - }, - nvidia.MigConfigSpec{ - Devices: []int32{2}, - MigEnabled: false, - MigDevices: testConfigs.Configs[1], - }, - }, - }, - deviceCount: 3, - expectError: false, - validate: func(t *testing.T, result nvidia.MigConfigSpecSlice) { - if len(result) != 3 { - t.Errorf("Expected 3 configs, got %d", len(result)) - } - for i := 0; i < 2; i++ { - if len(result[i].Devices) != 1 || result[i].Devices[0] != int32(i) { - t.Errorf("Config for device %d is incorrect: %v", i, result[i]) - } - if !result[i].MigEnabled { - t.Error("MigEnabled should be true for device", i) - } - if len(result[i].MigDevices) != 2 || (result[i].MigDevices["1g.10gb"] != 4 || result[i].MigDevices["2g.20gb"] != 1) { - t.Error("MigDevices not preserved correctly for device", i) - } - } - if len(result[2].Devices) != 1 || result[2].Devices[0] != 2 { - t.Errorf("Config for device 2 is incorrect: %v", result[2]) - } - if result[2].MigEnabled { - t.Error("MigEnabled should be false for device 2") - } - if len(result[2].MigDevices) != 1 || result[2].MigDevices["3g.30gb"] != 2 { - t.Error("MigDevices not preserved correctly for device 2") - } - }, - }, - { - name: "MultipleConfigsForSpecificDevicesWithAllEnabled", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{0, 1}, - MigEnabled: true, - MigDevices: testConfigs.Configs[0], - }, - nvidia.MigConfigSpec{ - Devices: []int32{2}, - MigEnabled: true, - MigDevices: testConfigs.Configs[1], - }, - }, - }, - deviceCount: 3, - expectError: false, - validate: func(t *testing.T, result nvidia.MigConfigSpecSlice) { - if len(result) != 3 { - t.Errorf("Expected 3 configs, got %d", len(result)) - } - for i := 0; i < 2; i++ { - if len(result[i].Devices) != 1 || result[i].Devices[0] != int32(i) { - t.Errorf("Config for device %d is incorrect: %v", i, result[i]) - } - if !result[i].MigEnabled { - t.Error("MigEnabled should be true for device", i) - } - if len(result[i].MigDevices) != 2 || (result[i].MigDevices["1g.10gb"] != 4 || result[i].MigDevices["2g.20gb"] != 1) { - t.Error("MigDevices not preserved correctly for device", i) - } - } - if len(result[2].Devices) != 1 || result[2].Devices[0] != 2 { - t.Errorf("Config for device 2 is incorrect: %v", result[2]) - } - if !result[2].MigEnabled { - t.Error("MigEnabled should be false for device 2") - } - if len(result[2].MigDevices) != 1 || result[2].MigDevices["3g.30gb"] != 2 { - t.Error("MigDevices not preserved correctly for device 2") - } - t.Log(result) - }, - }, - { - name: "DeviceNotMatched", - migConfigs: map[string]nvidia.MigConfigSpecSlice{ - "current": { - nvidia.MigConfigSpec{ - Devices: []int32{0, 1}, - MigEnabled: true, - }, - }, - }, - deviceCount: 3, - expectError: true, - validate: nil, - }, - } - - plugin := NvidiaDevicePlugin{ - config: &nvidia.DeviceConfig{ - Config: &v1.Config{ - Flags: v1.Flags{ - CommandLineFlags: v1.CommandLineFlags{}, - }, - }, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - result, err := plugin.processMigConfigs(tc.migConfigs, tc.deviceCount) - - if tc.expectError { - if err == nil { - t.Error("Expected error but got nil") - } - t.Log(err) - return - } - - if err != nil { - t.Fatalf("Unexpected error: %v", err) - } - - if tc.validate != nil { - tc.validate(t, result) - } - }) - } -} - func TestSelectPreferredDeviceIDsFromAnnotatedDevices(t *testing.T) { plugin := &NvidiaDevicePlugin{} // Use real NVIDIA GPU UUID format: GPU-xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx @@ -1020,11 +822,6 @@ func TestAllocatePreservesContainerOrderWhenOneContainerFallsBack(t *testing.T) getPendingPod = func(context.Context, string) (*corev1.Pod, error) { return pod, nil } defer func() { getPendingPod = previousGetPendingPod }() - // Provide a fake K8s client so the real patchErasedAnnotation can patch - previousKubeClient := client.KubeClient - client.KubeClient = fake.NewSimpleClientset(pod) - defer func() { client.KubeClient = previousKubeClient }() - previousPodAllocationFailed := podAllocationFailed podAllocationFailed = func(string, *corev1.Pod, string) {} defer func() { podAllocationFailed = previousPodAllocationFailed }() @@ -1047,3 +844,51 @@ func TestAllocatePreservesContainerOrderWhenOneContainerFallsBack(t *testing.T) require.Equal(t, "3000m", response.ContainerResponses[0].Envs["CUDA_DEVICE_MEMORY_LIMIT_0"]) require.Equal(t, "4000m", response.ContainerResponses[1].Envs["CUDA_DEVICE_MEMORY_LIMIT_0"]) } + +type mockListAndWatchServer struct { + grpc.ServerStream + sendErrs []error + sent []*kubeletdevicepluginv1beta1.ListAndWatchResponse +} + +func (m *mockListAndWatchServer) Send(response *kubeletdevicepluginv1beta1.ListAndWatchResponse) error { + m.sent = append(m.sent, response) + if len(m.sendErrs) == 0 { + return nil + } + err := m.sendErrs[0] + m.sendErrs = m.sendErrs[1:] + return err +} + +func TestListAndWatchSendError(t *testing.T) { + mockRM := &rm.ResourceManagerMock{ + DevicesFunc: func() rm.Devices { return rm.Devices{} }, + ResourceFunc: func() v1.ResourceName { return v1.ResourceName("nvidia.com/gpu") }, + } + + t.Run("initial send fails", func(t *testing.T) { + expectedErr := fmt.Errorf("initial send failed") + server := &mockListAndWatchServer{sendErrs: []error{expectedErr}} + plugin := &NvidiaDevicePlugin{ + rm: mockRM, stop: make(chan any), health: make(chan *rm.Device, 1), + schedulerConfig: nvidia.NvidiaConfig{NodeDefaultConfig: nvidia.NodeDefaultConfig{DeviceSplitCount: ptr[uint](1)}}, + } + err := plugin.ListAndWatch(&kubeletdevicepluginv1beta1.Empty{}, server) + require.ErrorIs(t, err, expectedErr) + require.Len(t, server.sent, 1) + }) + + t.Run("update send fails", func(t *testing.T) { + expectedErr := fmt.Errorf("update send failed") + server := &mockListAndWatchServer{sendErrs: []error{nil, expectedErr}} + plugin := &NvidiaDevicePlugin{ + rm: mockRM, stop: make(chan any), health: make(chan *rm.Device, 1), + schedulerConfig: nvidia.NvidiaConfig{NodeDefaultConfig: nvidia.NodeDefaultConfig{DeviceSplitCount: ptr[uint](1)}}, + } + plugin.health <- &rm.Device{Device: kubeletdevicepluginv1beta1.Device{ID: "gpu-1"}} + err := plugin.ListAndWatch(&kubeletdevicepluginv1beta1.Empty{}, server) + require.NoError(t, err) + require.Len(t, server.sent, 2) + }) +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go index a4b2a2e988..48630000a7 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go @@ -288,22 +288,6 @@ func containsModel(target string, models []string) bool { return false } -// Helper function to deepcopy new mig spec -func deepCopyMigConfig(src nvidia.MigConfigSpec) nvidia.MigConfigSpec { - dst := src - if src.Devices != nil { - dst.Devices = make([]int32, len(src.Devices)) - copy(dst.Devices, src.Devices) - } - if src.MigDevices != nil { - dst.MigDevices = make(map[string]int32) - for k, v := range src.MigDevices { - dst.MigDevices[k] = v - } - } - return dst -} - func (nv *NvidiaDevicePlugin) GetContainerDeviceStrArray(c device.ContainerDevices, pod *corev1.Pod, containerName string) ([]string, error) { if nv.operatingMode != "mig" { out := make([]string, 0, len(c)) diff --git a/pkg/device/nvidia/device.go b/pkg/device/nvidia/device.go index 311113dafe..1b495f81ed 100644 --- a/pkg/device/nvidia/device.go +++ b/pkg/device/nvidia/device.go @@ -66,22 +66,6 @@ var ( MemoryFactor int32 = 1 ) -type MigPartedSpec struct { - Version string `json:"version" yaml:"version"` - MigConfigs map[string]MigConfigSpecSlice `json:"mig-configs,omitempty" yaml:"mig-configs,omitempty"` -} - -// MigConfigSpec defines the spec to declare the desired MIG configuration for a set of GPUs. -type MigConfigSpec struct { - DeviceFilter any `json:"device-filter,omitempty" yaml:"device-filter,flow,omitempty"` - Devices []int32 `json:"devices" yaml:"devices,flow"` - MigEnabled bool `json:"mig-enabled" yaml:"mig-enabled"` - MigDevices map[string]int32 `json:"mig-devices" yaml:"mig-devices"` -} - -// MigConfigSpecSlice represents a slice of 'MigConfigSpec'. -type MigConfigSpecSlice []MigConfigSpec - // GPUCoreUtilizationPolicy is set nvidia gpu core isolation policy. type GPUCoreUtilizationPolicy string From 9542ac7edd5a772b21af7f3681ed72caf37a9068 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Fri, 7 Aug 2026 12:15:42 +0800 Subject: [PATCH 08/11] resolve commnet Signed-off-by: Jifei Wang --- hack/hami-mig-e2e.sh | 64 +++++++++++++++---- .../nvinternal/plugin/mig_startup_test.go | 3 +- .../nvidiadevice/nvinternal/plugin/migmgr.go | 55 +++++++++------- .../nvinternal/plugin/register.go | 4 ++ .../nvidiadevice/nvinternal/plugin/server.go | 22 ++++++- .../nvidiadevice/nvinternal/plugin/util.go | 18 +++++- pkg/device/nvidia/device.go | 3 + pkg/device/nvidia/mig_capability_test.go | 43 +++++++++++++ 8 files changed, 171 insertions(+), 41 deletions(-) diff --git a/hack/hami-mig-e2e.sh b/hack/hami-mig-e2e.sh index baeca40a19..dc7580621a 100644 --- a/hack/hami-mig-e2e.sh +++ b/hack/hami-mig-e2e.sh @@ -15,8 +15,11 @@ set -euo pipefail -NS=hami-mig-final-e2e -HAMI_NS=hami-system +NS=${MIG_E2E_NAMESPACE:-hami-mig-final-e2e} +HAMI_NS=${HAMI_NAMESPACE:-hami-system} +TARGET_NODE=${TARGET_NODE:-} +DEVICE_PLUGIN_DAEMONSET=${DEVICE_PLUGIN_DAEMONSET:-hami-device-plugin} +DEVICE_PLUGIN_CONTAINER=${DEVICE_PLUGIN_CONTAINER:-device-plugin} IMAGE=${GPU_TEST_IMAGE:-nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda12.5.0-ubuntu22.04} GPU_PROGRESS_WAIT=${GPU_PROGRESS_WAIT:-5} @@ -29,17 +32,49 @@ cleanup() { kubectl get namespace "$NS" >/dev/null 2>&1 || return 0 sleep 2 done - fail "namespace cleanup timed out" + echo "FAIL: namespace cleanup timed out" >&2 + return 1 +} + +trap 'status=$?; cleanup || true; exit "$status"' EXIT + +device_plugin_pod() { + kubectl get pods -n "$HAMI_NS" --field-selector "spec.nodeName=${TARGET_NODE}" -o json | + jq -r --arg daemonset "$DEVICE_PLUGIN_DAEMONSET" '.items[] | select([.metadata.ownerReferences[]? | select(.kind == "DaemonSet" and .name == $daemonset)] | length > 0) | .metadata.name' | + sed -n '1p' +} + +node_nvidia_smi() { + local pod + pod=$(device_plugin_pod) + [[ -n "$pod" ]] || fail "no ${DEVICE_PLUGIN_DAEMONSET} Pod found on ${TARGET_NODE}" + kubectl exec -n "$HAMI_NS" "$pod" -c "$DEVICE_PLUGIN_CONTAINER" -- nvidia-smi "$@" +} + +restart_target_device_plugin() { + local old_pod new_pod ready + old_pod=$(device_plugin_pod) + [[ -n "$old_pod" ]] || fail "no ${DEVICE_PLUGIN_DAEMONSET} Pod found on ${TARGET_NODE}" + kubectl delete pod -n "$HAMI_NS" "$old_pod" --wait=false + for _ in $(seq 1 90); do + new_pod=$(device_plugin_pod || true) + if [[ -n "$new_pod" && "$new_pod" != "$old_pod" ]]; then + ready=$(kubectl get pod -n "$HAMI_NS" "$new_pod" -o json 2>/dev/null | jq -r --arg container "$DEVICE_PLUGIN_CONTAINER" '.status.containerStatuses[]? | select(.name == $container) | .ready' || true) + [[ "$ready" == true ]] && return 0 + fi + sleep 2 + done + fail "device plugin on ${TARGET_NODE} did not restart" } create_pod() { local name=$1 memory=$2 - printf '%s\n' "{\"apiVersion\":\"v1\",\"kind\":\"Pod\",\"metadata\":{\"name\":\"${name}\",\"namespace\":\"${NS}\",\"annotations\":{\"nvidia.com/vgpu-mode\":\"mig\"}},\"spec\":{\"schedulerName\":\"hami-scheduler\",\"restartPolicy\":\"Never\",\"containers\":[{\"name\":\"cuda\",\"image\":\"${IMAGE}\",\"imagePullPolicy\":\"IfNotPresent\",\"command\":[\"bash\",\"-lc\",\"set -euo pipefail; nvidia-smi -L; n=0; echo 0 > /tmp/gpu-progress; while true; do if ! /cuda-samples/vectorAdd > /tmp/vectoradd.last 2>&1; then cat /tmp/vectoradd.last >&2; exit 1; fi; n=\$((n + 1)); echo \\\"\$n\\\" > /tmp/gpu-progress.next; mv /tmp/gpu-progress.next /tmp/gpu-progress; done\"],\"resources\":{\"limits\":{\"nvidia.com/gpu\":1,\"nvidia.com/gpumem\":${memory}}}}]}}" | kubectl apply -f - + printf '%s\n' "{\"apiVersion\":\"v1\",\"kind\":\"Pod\",\"metadata\":{\"name\":\"${name}\",\"namespace\":\"${NS}\",\"annotations\":{\"nvidia.com/vgpu-mode\":\"mig\"}},\"spec\":{\"schedulerName\":\"hami-scheduler\",\"nodeSelector\":{\"kubernetes.io/hostname\":\"${TARGET_NODE}\"},\"restartPolicy\":\"Never\",\"containers\":[{\"name\":\"cuda\",\"image\":\"${IMAGE}\",\"imagePullPolicy\":\"IfNotPresent\",\"command\":[\"bash\",\"-lc\",\"set -euo pipefail; nvidia-smi -L; n=0; echo 0 > /tmp/gpu-progress; while true; do if ! /cuda-samples/vectorAdd > /tmp/vectoradd.last 2>&1; then cat /tmp/vectoradd.last >&2; exit 1; fi; n=\$((n + 1)); echo \\\"\$n\\\" > /tmp/gpu-progress.next; mv /tmp/gpu-progress.next /tmp/gpu-progress; done\"],\"resources\":{\"limits\":{\"nvidia.com/gpu\":1,\"nvidia.com/gpumem\":${memory}}}}]}}" | kubectl apply -f - } wait_ready() { kubectl wait -n "$NS" --for=condition=Ready "pod/$1" --timeout=180s; } -mig_count() { nvidia-smi -L | grep -c 'MIG ' || true; } -profile_count() { nvidia-smi -L | grep -c "MIG $1" || true; } +mig_count() { node_nvidia_smi -L | grep -c 'MIG ' || true; } +profile_count() { node_nvidia_smi -L | grep -c "MIG $1" || true; } pod_uuid() { kubectl exec -n "$NS" "$1" -- nvidia-smi -L | sed -n 's/.*UUID: \(MIG-[^)]*\)).*/\1/p' | head -1; } gpu_progress() { kubectl exec -n "$NS" "$1" -- cat /tmp/gpu-progress; } @@ -87,7 +122,7 @@ wait_count() { while true; do count=$(mig_count) [[ "$count" == "$want" ]] && { echo "MIG_COUNT=${want}"; return; } - (( $(date +%s) - start < timeout )) || { nvidia-smi -L; fail "MIG count=${count}, want=${want}"; } + (( $(date +%s) - start < timeout )) || { node_nvidia_smi -L; fail "MIG count=${count}, want=${want}"; } sleep 2 done } @@ -112,7 +147,13 @@ assert_pending_unbound() { [[ "$phase" == Pending && -z "$node" ]] || fail "$pod expected Pending/unbound, got phase=${phase} node=${node}" } -log "baseline cleanup" +[[ -n "$TARGET_NODE" ]] || fail "set TARGET_NODE to the Kubernetes node under test" +kubectl get node "$TARGET_NODE" >/dev/null || fail "target node ${TARGET_NODE} does not exist" +node_nvidia_smi -L >/dev/null +physical_gpu_count=$(node_nvidia_smi --query-gpu=index --format=csv,noheader | wc -l | tr -d ' ') +[[ "$physical_gpu_count" == 1 ]] || fail "target node ${TARGET_NODE} must expose exactly one physical GPU, found ${physical_gpu_count}" + +log "baseline cleanup on ${TARGET_NODE}" cleanup kubectl create namespace "$NS" wait_count 0 @@ -165,8 +206,7 @@ echo "PASS CASE 2 capacity enforced" log "CASE 3: busy device-plugin restart preserves every MIG UUID" progress_before_restart=$(snapshot_gpu_progress one-a two-a two-b two-c) -kubectl rollout restart daemonset/hami-device-plugin -n "$HAMI_NS" -kubectl rollout status daemonset/hami-device-plugin -n "$HAMI_NS" --timeout=180s +restart_target_device_plugin sleep 15 assert_gpu_progress_since "$progress_before_restart" one-a two-a two-b two-c assert_uuid one-a "$uuid_one_a" @@ -175,7 +215,7 @@ assert_uuid two-b "$uuid_two_b" assert_uuid two-c "$uuid_two_c" wait_count 4 [[ "$(profile_count 1g.5gb)" == 1 && "$(profile_count 2g.10gb)" == 3 ]] || fail "layout changed during restart" -kubectl logs -n "$HAMI_NS" daemonset/hami-device-plugin -c device-plugin --since=3m | grep 'inUseGPUs=\[0\]' >/dev/null || fail "startup did not detect busy GPU" +kubectl logs -n "$HAMI_NS" "$(device_plugin_pod)" -c "$DEVICE_PLUGIN_CONTAINER" --since=3m | grep 'inUseGPUs=\[0\]' >/dev/null || fail "startup did not detect busy GPU" echo "PASS CASE 3 all UUIDs preserved" log "CASE 4: immediate delete/replacement after restart" @@ -283,5 +323,5 @@ cleanup wait_count 0 kubectl get nodes kubectl get pods -n "$HAMI_NS" -nvidia-smi -q | grep -A3 'MIG Mode' +node_nvidia_smi -q | grep -A3 'MIG Mode' echo "ALL_FIXED_MIG_E2E_TESTS_PASSED" diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go index 81aa044f0f..7ca734d184 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/mig_startup_test.go @@ -15,9 +15,10 @@ import ( "testing" "time" - "github.com/Project-HAMi/HAMi/pkg/device/nvidia" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" ) func TestSortedIntSetKeys(t *testing.T) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go index f749122859..b31c7c91b7 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go @@ -252,13 +252,18 @@ func destroyAllMigInstances(dev nvml.Device) error { continue } for _, gi := range gis { - ciInfoRet := profileIDToCIProfileID(giProfileID) - if ciInfo, r := gi.GetComputeInstanceProfileInfo(ciInfoRet, nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED); r == nvml.SUCCESS { - if cis, r2 := gi.GetComputeInstances(&ciInfo); r2 == nvml.SUCCESS { - for _, ci := range cis { - if d := ci.Destroy(); d != nvml.SUCCESS { - return fmt.Errorf("destroy compute instance: %s", nvml.ErrorString(d)) - } + for ciProfileID := 0; ciProfileID < nvml.COMPUTE_INSTANCE_PROFILE_COUNT; ciProfileID++ { + ciInfo, r := gi.GetComputeInstanceProfileInfo(ciProfileID, nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED) + if r != nvml.SUCCESS { + continue + } + cis, r := gi.GetComputeInstances(&ciInfo) + if r != nvml.SUCCESS { + continue + } + for _, ci := range cis { + if d := ci.Destroy(); d != nvml.SUCCESS { + return fmt.Errorf("destroy compute instance: %s", nvml.ErrorString(d)) } } } @@ -303,8 +308,10 @@ func allocationKey(gpuIndex int, profile string, placement nvml.GpuInstancePlace } // EnsureAllocation realizes exactly the scheduler-reserved profile and -// placement. It never retries another placement. -func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, placement nvml.GpuInstancePlacement) (string, error) { +// placement. It returns whether this call created the instance, allowing the +// caller to roll back only its own partial allocation. It never retries +// another placement. +func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, placement nvml.GpuInstancePlacement) (string, bool, error) { key := allocationKey(gpuIndex, profile, placement) lk := m.gpuLock(gpuIndex) lk.Lock() @@ -314,33 +321,33 @@ func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, plac if inst := m.byAllocation[key]; inst != nil && inst.Present { uuid := inst.MigUUID m.mu.Unlock() - return uuid, nil + return uuid, false, nil } m.mu.Unlock() if err := ensureMigModeEnabled(gpuIndex); err != nil { - return "", err + return "", false, err } profileKey := profileSliceKey(profile) giProfileID, ok := profileNameToGIProfileID[profileKey] if !ok { - return "", fmt.Errorf("unsupported MIG profile %q", profile) + return "", false, fmt.Errorf("unsupported MIG profile %q", profile) } ciProfileID, ok := profileNameToCIProfileID[profileKey] if !ok { - return "", fmt.Errorf("unsupported MIG compute profile %q", profile) + return "", false, fmt.Errorf("unsupported MIG compute profile %q", profile) } dev, err := deviceHandleByIndex(gpuIndex) if err != nil { - return "", err + return "", false, err } giInfo, ret := dev.GetGpuInstanceProfileInfo(giProfileID) if ret != nvml.SUCCESS { - return "", fmt.Errorf("get GI profile %s: %s", profile, nvml.ErrorString(ret)) + return "", false, fmt.Errorf("get GI profile %s: %s", profile, nvml.ErrorString(ret)) } possible, ret := dev.GetGpuInstancePossiblePlacements(&giInfo) if ret != nvml.SUCCESS { - return "", fmt.Errorf("get placements for %s: %s", profile, nvml.ErrorString(ret)) + return "", false, fmt.Errorf("get placements for %s: %s", profile, nvml.ErrorString(ret)) } valid := false for _, candidate := range possible { @@ -350,38 +357,38 @@ func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, plac } } if !valid { - return "", fmt.Errorf("scheduler selected invalid placement %+v for profile %s", placement, profile) + return "", false, fmt.Errorf("scheduler selected invalid placement %+v for profile %s", placement, profile) } gi, ret := dev.CreateGpuInstanceWithPlacement(&giInfo, &placement) if ret != nvml.SUCCESS { - return "", fmt.Errorf("create GI profile=%s placement=%+v: %s", profile, placement, nvml.ErrorString(ret)) + return "", false, fmt.Errorf("create GI profile=%s placement=%+v: %s", profile, placement, nvml.ErrorString(ret)) } giData, ret := gi.GetInfo() if ret != nvml.SUCCESS { gi.Destroy() - return "", fmt.Errorf("get GI info: %s", nvml.ErrorString(ret)) + return "", false, fmt.Errorf("get GI info: %s", nvml.ErrorString(ret)) } ciInfo, ret := gi.GetComputeInstanceProfileInfo(ciProfileID, nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED) if ret != nvml.SUCCESS { gi.Destroy() - return "", fmt.Errorf("get CI profile info: %s", nvml.ErrorString(ret)) + return "", false, fmt.Errorf("get CI profile info: %s", nvml.ErrorString(ret)) } ci, ret := gi.CreateComputeInstance(&ciInfo) if ret != nvml.SUCCESS { gi.Destroy() - return "", fmt.Errorf("create CI: %s", nvml.ErrorString(ret)) + return "", false, fmt.Errorf("create CI: %s", nvml.ErrorString(ret)) } ciData, ret := ci.GetInfo() if ret != nvml.SUCCESS { ci.Destroy() gi.Destroy() - return "", fmt.Errorf("get CI info: %s", nvml.ErrorString(ret)) + return "", false, fmt.Errorf("get CI info: %s", nvml.ErrorString(ret)) } migUUID, err := findMigUUIDForGI(dev, giData.Id) if err != nil { ci.Destroy() gi.Destroy() - return "", err + return "", false, err } inst := &migInstance{Profile: profile, Placement: placement, Present: true, GIID: giData.Id, CIID: ciData.Id, MigUUID: migUUID} m.mu.Lock() @@ -389,7 +396,7 @@ func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, plac m.byAllocationMigUUID[migUUID] = key m.mu.Unlock() klog.InfoS("created scheduler-reserved MIG allocation", "uuid", migUUID, "gpu", gpuIndex, "profile", profile, "start", placement.Start, "size", placement.Size, "gpuInstanceID", giData.Id, "computeInstanceID", ciData.Id) - return migUUID, nil + return migUUID, true, nil } func (m *MigInstanceManager) AllocationRuntimeInfo(gpuIndex int, profile string, placement nvml.GpuInstancePlacement) (migAllocationRuntimeInfo, bool) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go index 83104e1566..fdb4fa2b7d 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go @@ -189,6 +189,10 @@ func (plugin *NvidiaDevicePlugin) getAPIDevices() *[]*device.DeviceInfo { } if isMigMode { info.MIGProfiles = plugin.discoverMigProfiles(ndev, Model) + if len(info.MIGProfiles) == 0 { + klog.InfoS("skip MIG device with no discovered profile capacity", "id", UUID, "model", Model) + continue + } info.Count = 0 for _, profile := range info.MIGProfiles { if int32(profile.InstanceCount) > info.Count { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index c2f3fc2b8b..9304203f18 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -54,6 +54,7 @@ import ( "google.golang.org/grpc/credentials/insecure" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + k8stypes "k8s.io/apimachinery/pkg/types" "k8s.io/klog/v2" kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" cdiapi "tags.cncf.io/container-device-interface/pkg/cdi" @@ -189,6 +190,10 @@ func (o *options) devicePluginForResource(ctx context.Context, nvconfig *nvidia. if err := nvidia.ValidateMigProfileAllowlist(sConfig.NvidiaConfig.MigProfileAllowlist); err != nil { return nil, fmt.Errorf("validate MIG profile allowlist: %w", err) } + var migMgr *MigInstanceManager + if mode == "mig" { + migMgr = NewMigInstanceManager() + } return &NvidiaDevicePlugin{ ctx: ctx, rm: resourceManager, @@ -205,6 +210,7 @@ func (o *options) devicePluginForResource(ctx context.Context, nvconfig *nvidia. cdiAnnotationPrefix: *o.config.Flags.Plugin.CDIAnnotationPrefix, schedulerConfig: sConfig.NvidiaConfig, operatingMode: mode, + migMgr: migMgr, deviceCache: "", // These will be reinitialized every @@ -280,7 +286,7 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { klog.Fatalf("RemoveMigApplyLock failed: %v", err) } - deviceSupportMig := true + deviceSupportMig := len(deviceNames) > 0 for _, name := range deviceNames { supported := false for _, allowlist := range plugin.schedulerConfig.MigProfileAllowlist { @@ -305,7 +311,6 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { plugin.WatchAndRegister(plugin.disableWatchAndRegister, plugin.ackDisableWatchAndRegister) }() if plugin.operatingMode == "mig" { - plugin.migMgr = NewMigInstanceManager() if deviceSupportMig { inUse, detectErr := collectInUseGPUs(plugin.ctx, os.Getenv(util.NodeNameEnvName)) if detectErr != nil { @@ -386,7 +391,18 @@ func (plugin *NvidiaDevicePlugin) annotateMigRuntimeInfo(pod *corev1.Pod) error if err != nil { return err } - if err := util.PatchPodAnnotations(pod, map[string]string{nvidia.MigAllocationsAnnotation: string(raw)}); err != nil { + patch, err := json.Marshal(map[string]any{ + "metadata": map[string]any{ + "resourceVersion": pod.ResourceVersion, + "annotations": map[string]string{ + nvidia.MigAllocationsAnnotation: string(raw), + }, + }, + }) + if err != nil { + return err + } + if _, err := client.GetClient().CoreV1().Pods(pod.Namespace).Patch(plugin.ctx, pod.Name, k8stypes.MergePatchType, patch, metav1.PatchOptions{}); err != nil { return err } pod.Annotations[nvidia.MigAllocationsAnnotation] = string(raw) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go index 48630000a7..bae38a5eb5 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go @@ -334,6 +334,18 @@ func (nv *NvidiaDevicePlugin) GetContainerDeviceStrArray(c device.ContainerDevic if err := nv.reconcileActiveMigAllocations(); err != nil { return nil, fmt.Errorf("reconcile MIG allocations before allocation: %w", err) } + createdMigUUIDs := make([]string, 0, len(c)) + allocationCompleted := false + defer func() { + if allocationCompleted { + return + } + for i := len(createdMigUUIDs) - 1; i >= 0; i-- { + if err := nv.migMgr.Release(createdMigUUIDs[i]); err != nil { + klog.ErrorS(err, "failed to roll back partial MIG allocation", "uuid", createdMigUUIDs[i]) + } + } + }() out := make([]string, 0, len(c)) for i, reservation := range containerAllocations { if reservation.GPUUUID != c[i].UUID { @@ -343,12 +355,16 @@ func (nv *NvidiaDevicePlugin) GetContainerDeviceStrArray(c device.ContainerDevic if !ok { return nil, fmt.Errorf("resolve parent GPU %s", reservation.GPUUUID) } - migUUID, err := nv.migMgr.EnsureAllocation(gpuIndex, reservation.Profile, nvml.GpuInstancePlacement{Start: reservation.Placement.Start, Size: reservation.Placement.Size}) + migUUID, created, err := nv.migMgr.EnsureAllocation(gpuIndex, reservation.Profile, nvml.GpuInstancePlacement{Start: reservation.Placement.Start, Size: reservation.Placement.Size}) if err != nil { return nil, err } + if created { + createdMigUUIDs = append(createdMigUUIDs, migUUID) + } out = append(out, migUUID) } + allocationCompleted = true return out, nil } diff --git a/pkg/device/nvidia/device.go b/pkg/device/nvidia/device.go index 1b495f81ed..a888239f7e 100644 --- a/pkg/device/nvidia/device.go +++ b/pkg/device/nvidia/device.go @@ -580,6 +580,9 @@ func (dev *NvidiaGPUDevices) CustomFilterRule(allocated *device.PodDevices, requ if devusage.Mode == MigMode { occupied := occupiedMigPlacements(devusage.MigAllocationsInUse) for _, existing := range toAllocate { + if existing.UUID != devusage.ID { + continue + } _, placement, ok := selectMigCandidate(devusage.MigProfiles, occupied, existing.Usedmem) if !ok { return false diff --git a/pkg/device/nvidia/mig_capability_test.go b/pkg/device/nvidia/mig_capability_test.go index 23e2bf80d9..fab88d76d3 100644 --- a/pkg/device/nvidia/mig_capability_test.go +++ b/pkg/device/nvidia/mig_capability_test.go @@ -19,6 +19,8 @@ package nvidia import ( "testing" + corev1 "k8s.io/api/core/v1" + "github.com/Project-HAMi/HAMi/pkg/device" ) @@ -65,3 +67,44 @@ func TestCustomFilterUsesReportedPlacementCapacity(t *testing.T) { t.Fatal("1g request should fit the remaining placement") } } + +func TestCustomFilterIgnoresQueuedAllocationsForOtherGPUs(t *testing.T) { + dev := &NvidiaGPUDevices{} + usage := &device.DeviceUsage{ + ID: "GPU-a", Mode: MigMode, MigProfiles: a100MigProfiles(), + MigAllocationsInUse: []device.MigAllocation{ + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 0, Size: 2}}, + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 2, Size: 2}}, + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 4, Size: 2}}, + }, + } + queued := device.ContainerDevices{{UUID: "GPU-b", Usedmem: 5000}} + if !dev.CustomFilterRule(nil, device.ContainerDeviceRequest{Memreq: 5000}, queued, usage) { + t.Fatal("allocation queued on another GPU consumed this GPU's remaining placement") + } + queued[0].UUID = "GPU-a" + if dev.CustomFilterRule(nil, device.ContainerDeviceRequest{Memreq: 5000}, queued, usage) { + t.Fatal("two allocations on this GPU should not fit its single remaining placement") + } +} + +func TestFitUsesEffectivePercentageMemoryForMigPlacement(t *testing.T) { + dev := InitNvidiaDevice(NvidiaConfig{}) + usage := &device.DeviceUsage{ + ID: "GPU-a", Type: NvidiaGPUDevice, Mode: MigMode, Health: true, + Count: 7, Used: 3, Totalmem: 40960, Usedmem: 30720, Totalcore: 100, + MigProfiles: a100MigProfiles(), + MigAllocationsInUse: []device.MigAllocation{ + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 0, Size: 2}}, + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 2, Size: 2}}, + {Profile: "2g.10gb", Placement: device.MigPlacement{Start: 4, Size: 2}}, + }, + } + request := device.ContainerDeviceRequest{ + Nums: 1, Type: NvidiaGPUDevice, MemPercentagereq: 25, + } + fit, _, _ := dev.Fit([]*device.DeviceUsage{usage}, request, &corev1.Pod{}, &device.NodeInfo{}, &device.PodDevices{}) + if fit { + t.Fatal("25 percent memory request should not fit when only a 1g placement remains") + } +} From cf45c6bc831adee9cc963f9f402512b54a6b94ea Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Fri, 7 Aug 2026 13:31:00 +0800 Subject: [PATCH 09/11] fix(nvidia): serialize MIG startup and restore reservations Signed-off-by: Jifei Wang --- .../nvidiadevice/nvinternal/plugin/server.go | 55 +++++++++--------- .../nvinternal/plugin/server_test.go | 4 ++ pkg/scheduler/mig_allocation_test.go | 34 +++++++++++ pkg/scheduler/scheduler.go | 58 +++++++++++++++---- 4 files changed, 113 insertions(+), 38 deletions(-) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index 9304203f18..2bc7636226 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -261,22 +261,6 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { return err } - err = plugin.Serve() - if err != nil { - klog.Infof("Could not start device plugin for '%s': %s", plugin.rm.Resource(), err) - plugin.cleanup() - return err - } - klog.Infof("Starting to serve '%s' on %s", plugin.rm.Resource(), plugin.socket) - - err = plugin.Register(kubeletSocket) - if err != nil { - klog.Infof("Could not register device plugin: %s", err) - plugin.Stop() - return err - } - klog.Infof("Registered device plugin for '%s' with Kubelet", plugin.rm.Resource()) - // Prepare the lock directory before any dynamic MIG operation. A stale // lock can be left behind when the previous plugin process exits midway. if err = CreateMigApplyLockDir(); err != nil { @@ -300,16 +284,6 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { break } } - go func() { - err := plugin.rm.CheckHealth(plugin.stop, plugin.health, plugin.disableHealthChecks, plugin.ackDisableHealthChecks) - if err != nil { - klog.Infof("Failed to start health check: %v; continuing with health checks disabled", err) - } - }() - - go func() { - plugin.WatchAndRegister(plugin.disableWatchAndRegister, plugin.ackDisableWatchAndRegister) - }() if plugin.operatingMode == "mig" { if deviceSupportMig { inUse, detectErr := collectInUseGPUs(plugin.ctx, os.Getenv(util.NodeNameEnvName)) @@ -333,6 +307,35 @@ func (plugin *NvidiaDevicePlugin) Start(kubeletSocket string) error { klog.InfoS("mig init: failed to adopt active MIG allocations", "err", err) } } + } + + err = plugin.Serve() + if err != nil { + klog.Infof("Could not start device plugin for '%s': %s", plugin.rm.Resource(), err) + plugin.cleanup() + return err + } + klog.Infof("Starting to serve '%s' on %s", plugin.rm.Resource(), plugin.socket) + + err = plugin.Register(kubeletSocket) + if err != nil { + klog.Infof("Could not register device plugin: %s", err) + plugin.Stop() + return err + } + klog.Infof("Registered device plugin for '%s' with Kubelet", plugin.rm.Resource()) + + go func() { + err := plugin.rm.CheckHealth(plugin.stop, plugin.health, plugin.disableHealthChecks, plugin.ackDisableHealthChecks) + if err != nil { + klog.Infof("Failed to start health check: %v; continuing with health checks disabled", err) + } + }() + + go func() { + plugin.WatchAndRegister(plugin.disableWatchAndRegister, plugin.ackDisableWatchAndRegister) + }() + if plugin.operatingMode == "mig" { // Pod annotations are the allocation source of truth. Periodically // reconcile the manager with live Pods so completed or deleted Pods // release their exact profile+placement allocation. diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go index 0f208a9ea1..21bda35763 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go @@ -56,6 +56,10 @@ import ( kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" ) +func ptr[T any](value T) *T { + return &value +} + func TestCDIAllocateResponse(t *testing.T) { testCases := []struct { description string diff --git a/pkg/scheduler/mig_allocation_test.go b/pkg/scheduler/mig_allocation_test.go index 857871ca09..5854d376ce 100644 --- a/pkg/scheduler/mig_allocation_test.go +++ b/pkg/scheduler/mig_allocation_test.go @@ -74,3 +74,37 @@ func TestGetNodesUsageFailsClosedWithoutMigAllocation(t *testing.T) { t.Fatal("MIG device must be fail-closed when an allocated Pod lacks profile/placement") } } + +func TestGetNodesUsageRestoresUnconsumedMigReservations(t *testing.T) { + nodes := newNodeManager() + nodes.addNode("node1", &device.NodeInfo{ + ID: "node1", Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: {{ + ID: "GPU-a", Count: 7, Devmem: 40960, Devcore: 100, Mode: nvidia.MigMode, Health: true, + MIGProfiles: []device.MigProfile{{Name: "1g.5gb", MemoryMB: 5120, Core: 14, Placements: []device.MigPlacement{{Start: 5, Size: 1}, {Start: 6, Size: 1}}}}, + }}}, + }) + pods := device.NewPodManager() + pods.AddPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + UID: "pod-1", Name: "pod-1", Namespace: "default", + Annotations: map[string]string{nvidia.MigAllocationsAnnotation: `[ + {"containerIndex":0,"deviceIndex":0,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":6,"size":1}}, + {"containerIndex":0,"deviceIndex":1,"gpuUUID":"GPU-a","profile":"1g.5gb","placement":{"start":5,"size":1}} + ]`}, + }}, "node1", device.PodDevices{nvidia.NvidiaGPUDevice: {{{UUID: "GPU-a", Usedmem: 5120, Usedcores: 14}}}}) + + s := Scheduler{nodeManager: nodes, podManager: pods} + nodeNames := []string{"node1"} + usage, _, _, err := s.getNodesUsage(&nodeNames, nil) + if err != nil { + t.Fatal(err) + } + deviceUsage := (*usage)["node1"].Devices.DeviceLists[0].Device + if !deviceUsage.Health { + t.Fatal("MIG device should remain healthy when every reservation can be restored") + } + allocations := deviceUsage.MigAllocationsInUse + if len(allocations) != 2 || allocations[0].Placement != (device.MigPlacement{Start: 6, Size: 1}) || allocations[1].Placement != (device.MigPlacement{Start: 5, Size: 1}) { + t.Fatalf("restored allocations: %+v", allocations) + } +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 334bb4bc01..2b65789459 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -700,6 +700,21 @@ func nodeListLen(nodes *corev1.NodeList) int { return len(nodes.Items) } +func migAllocationUsage(allocation nvidia.MigAllocation) device.MigAllocation { + usage := device.MigAllocation{ + Profile: allocation.Profile, Placement: allocation.Placement, + MigUUID: allocation.MigUUID, + RuntimeReady: allocation.MigUUID != "" && allocation.GPUInstanceID != nil && allocation.ComputeInstanceID != nil, + } + if allocation.GPUInstanceID != nil { + usage.GPUInstanceID = *allocation.GPUInstanceID + } + if allocation.ComputeInstanceID != nil { + usage.ComputeInstanceID = *allocation.ComputeInstanceID + } + return usage +} + // returns all nodes and its device memory usage, and we filter it with nodeSelector, taints, nodeAffinity // unschedulerable and nodeName. func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[string]*NodeUsage, *map[string]*NodeUsage, map[string]string, error) { @@ -750,18 +765,7 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin } allocation := allocations[0] allocationsByGPU[udevice.UUID] = allocations[1:] - migAllocation := device.MigAllocation{ - Profile: allocation.Profile, Placement: allocation.Placement, - MigUUID: allocation.MigUUID, - RuntimeReady: allocation.MigUUID != "" && allocation.GPUInstanceID != nil && allocation.ComputeInstanceID != nil, - } - if allocation.GPUInstanceID != nil { - migAllocation.GPUInstanceID = *allocation.GPUInstanceID - } - if allocation.ComputeInstanceID != nil { - migAllocation.ComputeInstanceID = *allocation.ComputeInstanceID - } - d.Device.MigAllocationsInUse = append(d.Device.MigAllocationsInUse, migAllocation) + d.Device.MigAllocationsInUse = append(d.Device.MigAllocationsInUse, migAllocationUsage(allocation)) continue } if d.Device.Mode == nvidia.MigMode { @@ -773,6 +777,36 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin } } } + for gpuUUID, allocations := range allocationsByGPU { + if len(allocations) == 0 { + continue + } + matched := false + for _, d := range node.Devices.DeviceLists { + if d.Device.ID != gpuUUID { + continue + } + matched = true + if d.Device.Mode != nvidia.MigMode { + klog.ErrorS(nil, "unconsumed MIG reservations reference a non-MIG device", "pod", klog.KRef(p.Namespace, p.Name), "gpuUUID", gpuUUID, "reservations", len(allocations)) + d.Device.Health = false + break + } + for _, allocation := range allocations { + d.Device.MigAllocationsInUse = append(d.Device.MigAllocationsInUse, migAllocationUsage(allocation)) + } + klog.InfoS("restored MIG reservations missing from cached Pod devices", "pod", klog.KRef(p.Namespace, p.Name), "gpuUUID", gpuUUID, "reservations", len(allocations)) + break + } + if !matched { + klog.ErrorS(nil, "unconsumed MIG reservations reference an unknown device", "pod", klog.KRef(p.Namespace, p.Name), "gpuUUID", gpuUUID, "reservations", len(allocations)) + for _, d := range node.Devices.DeviceLists { + if d.Device.Mode == nvidia.MigMode { + d.Device.Health = false + } + } + } + } klog.V(5).Infof("usage: pod %v assigned %v %v", p.Name, p.NodeID, p.Devices) } if nodes == nil { From 079cb7ca42fdcadca39a37e87a8f5d69bacc43f9 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Fri, 7 Aug 2026 13:55:42 +0800 Subject: [PATCH 10/11] clenup migmgr Signed-off-by: Jifei Wang --- .../nvidiadevice/nvinternal/plugin/migmgr.go | 147 +++--------------- .../plugin/migmgr_placement_test.go | 75 --------- 2 files changed, 23 insertions(+), 199 deletions(-) delete mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go index b31c7c91b7..b06dfe3dff 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr.go @@ -47,14 +47,11 @@ type migAllocationKey struct { Size uint32 } -// migInstance tracks the nvml-level identity of a MIG GI+CI pair bound to a -// slot. Absent means the slot's GI+CI have been destroyed (e.g. on task end) -// but we remember the profile and placement so we can recreate the instance -// at the same physical slice when the next task claims this slot. +// migInstance tracks the NVML-level identity of a live MIG GI+CI pair bound to +// a scheduler-reserved profile and physical placement. type migInstance struct { Profile string // slice group, e.g. "1g" Placement nvml.GpuInstancePlacement - Present bool GIID uint32 CIID uint32 MigUUID string @@ -95,7 +92,7 @@ func profileSliceKey(profile string) string { return profile } -// ResetIdleGPUs prepares idle MIG-capable GPUs for on-demand slot creation +// ResetIdleGPUs prepares idle MIG-capable GPUs for on-demand instance creation // through NVML. Busy GPUs are left untouched; idle GPUs // have MIG mode enabled and all existing GI/CI instances destroyed. func (m *MigInstanceManager) ResetIdleGPUs(deviceCount int, inUse map[int]struct{}) ([]int, error) { @@ -201,10 +198,10 @@ func ensureMigModeEnabled(gpuIndex int) error { return fmt.Errorf("gpu %d mig mode is not enabled after set (current=%d pending=%d)", gpuIndex, curMode, pendingMode) } -// destroyPresentMigInstance destroys the tracked GI+CI on hardware. Returns +// destroyMigInstance destroys the tracked GI+CI on hardware. Returns // nil when the instance is already gone or was destroyed successfully. -func destroyPresentMigInstance(gpuIndex int, inst *migInstance) error { - if inst == nil || !inst.Present { +func destroyMigInstance(gpuIndex int, inst *migInstance) error { + if inst == nil { return nil } dev, err := deviceHandleByIndex(gpuIndex) @@ -232,7 +229,7 @@ func destroyPresentMigInstance(gpuIndex int, inst *migInstance) error { } // destroyAllMigInstances enumerates and destroys every GI+CI on the device. -// Used on template switches when no scheduler-allocated slot is in use. +// It is used to reset idle GPUs before accepting scheduler allocations. func destroyAllMigInstances(dev nvml.Device) error { for _, giProfileID := range []int{ nvml.GPU_INSTANCE_PROFILE_1_SLICE, @@ -290,10 +287,10 @@ func (m *MigInstanceManager) Release(migUUID string) error { m.mu.Lock() inst := m.byAllocation[key] m.mu.Unlock() - if inst == nil || !inst.Present { + if inst == nil { return nil } - if err := destroyPresentMigInstance(key.GPUIndex, inst); err != nil { + if err := destroyMigInstance(key.GPUIndex, inst); err != nil { return err } m.mu.Lock() @@ -318,7 +315,7 @@ func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, plac defer lk.Unlock() m.mu.Lock() - if inst := m.byAllocation[key]; inst != nil && inst.Present { + if inst := m.byAllocation[key]; inst != nil { uuid := inst.MigUUID m.mu.Unlock() return uuid, false, nil @@ -390,7 +387,7 @@ func (m *MigInstanceManager) EnsureAllocation(gpuIndex int, profile string, plac gi.Destroy() return "", false, err } - inst := &migInstance{Profile: profile, Placement: placement, Present: true, GIID: giData.Id, CIID: ciData.Id, MigUUID: migUUID} + inst := &migInstance{Profile: profile, Placement: placement, GIID: giData.Id, CIID: ciData.Id, MigUUID: migUUID} m.mu.Lock() m.byAllocation[key] = inst m.byAllocationMigUUID[migUUID] = key @@ -404,7 +401,7 @@ func (m *MigInstanceManager) AllocationRuntimeInfo(gpuIndex int, profile string, m.mu.Lock() defer m.mu.Unlock() inst := m.byAllocation[key] - if inst == nil || !inst.Present { + if inst == nil { return migAllocationRuntimeInfo{}, false } return migAllocationRuntimeInfo{ @@ -424,10 +421,15 @@ func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID stri if err != nil { return err } - giProfileID, ok := profileNameToGIProfileID[profileSliceKey(profile)] + profileKey := profileSliceKey(profile) + giProfileID, ok := profileNameToGIProfileID[profileKey] if !ok { return fmt.Errorf("unsupported MIG profile %q", profile) } + ciProfileID, ok := profileNameToCIProfileID[profileKey] + if !ok { + return fmt.Errorf("unsupported MIG compute profile %q", profile) + } profileInfo, ret := dev.GetGpuInstanceProfileInfo(giProfileID) if ret != nvml.SUCCESS { return fmt.Errorf("get GI profile %s: %s", profile, nvml.ErrorString(ret)) @@ -436,7 +438,6 @@ func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID stri if ret != nvml.SUCCESS { return fmt.Errorf("list GI profile %s: %s", profile, nvml.ErrorString(ret)) } - ciProfileID := profileIDToCIProfileID(giProfileID) for _, gi := range instances { giInfo, r := gi.GetInfo() if r != nvml.SUCCESS || giInfo.Placement != placement || giInfo.Id != gpuInstanceID { @@ -460,7 +461,7 @@ func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID stri } key := allocationKey(gpuIndex, profile, placement) m.mu.Lock() - m.byAllocation[key] = &migInstance{Profile: profile, Placement: placement, Present: true, GIID: giInfo.Id, CIID: ciData.Id, MigUUID: migUUID} + m.byAllocation[key] = &migInstance{Profile: profile, Placement: placement, GIID: giInfo.Id, CIID: ciData.Id, MigUUID: migUUID} m.byAllocationMigUUID[migUUID] = key m.mu.Unlock() return nil @@ -471,10 +472,8 @@ func (m *MigInstanceManager) AdoptAllocation(gpuIndex int, profile, migUUID stri func (m *MigInstanceManager) ReconcileActiveAllocations(active map[migAllocationKey]struct{}) error { m.mu.Lock() keys := make([]migAllocationKey, 0, len(m.byAllocation)) - for key, inst := range m.byAllocation { - if inst.Present { - keys = append(keys, key) - } + for key := range m.byAllocation { + keys = append(keys, key) } m.mu.Unlock() for _, key := range keys { @@ -486,9 +485,9 @@ func (m *MigInstanceManager) ReconcileActiveAllocations(active map[migAllocation m.mu.Lock() inst := m.byAllocation[key] m.mu.Unlock() - if inst != nil && inst.Present { + if inst != nil { oldUUID := inst.MigUUID - if err := destroyPresentMigInstance(key.GPUIndex, inst); err != nil { + if err := destroyMigInstance(key.GPUIndex, inst); err != nil { lk.Unlock() return err } @@ -534,103 +533,3 @@ func findMigUUIDForGI(dev nvml.Device, giID uint32) (string, error) { } return "", fmt.Errorf("no MIG device found for GI %d", giID) } - -// pickFreePlacement returns a placement for the given GI profile that does -// not overlap with any of the placements already in use on this GPU. -func pickFreePlacement(dev nvml.Device, info *nvml.GpuInstanceProfileInfo, inUse map[uint32]uint32) (nvml.GpuInstancePlacement, error) { - possible, ret := dev.GetGpuInstancePossiblePlacements(info) - if ret != nvml.SUCCESS { - return nvml.GpuInstancePlacement{}, fmt.Errorf("get possible placements: %s", nvml.ErrorString(ret)) - } - return chooseFreePlacement(possible, inUse, preferHighPlacement(info.SliceCount)) -} - -func preferHighPlacement(sliceCount uint32) bool { - return sliceCount == 1 || sliceCount == 3 -} - -func sortPlacements(possible []nvml.GpuInstancePlacement, preferHigh bool) { - sort.SliceStable(possible, func(i, j int) bool { - if preferHigh { - return possible[i].Start > possible[j].Start - } - return possible[i].Start < possible[j].Start - }) -} - -func placementCandidates(previous nvml.GpuInstancePlacement, possible []nvml.GpuInstancePlacement) []nvml.GpuInstancePlacement { - candidates := make([]nvml.GpuInstancePlacement, 0, len(possible)+1) - if previous.Size != 0 { - candidates = append(candidates, previous) - } - for _, candidate := range possible { - if previous.Size != 0 && candidate == previous { - continue - } - candidates = append(candidates, candidate) - } - return candidates -} - -// chooseFreePlacement packs 1g and 3g instances from high addresses while 2g -// instances use low addresses. This matches NVIDIA's A100 balanced placement -// (2g at 0:2, 1g at 2:1 and 3:1, 3g at 4:4) while retaining the canonical -// 1x1g + 3x2g layout. -func chooseFreePlacement(possible []nvml.GpuInstancePlacement, inUse map[uint32]uint32, preferHigh bool) (nvml.GpuInstancePlacement, error) { - sortPlacements(possible, preferHigh) - for _, p := range possible { - if !placementOverlaps(p, inUse) { - return p, nil - } - } - return nvml.GpuInstancePlacement{}, fmt.Errorf("no free placement for profile") -} - -func placementOverlaps(p nvml.GpuInstancePlacement, inUse map[uint32]uint32) bool { - for start, size := range inUse { - if p.Start < start+size && start < p.Start+p.Size { - return true - } - } - return false -} - -func profileIDToCIProfileID(giProfileID int) int { - switch giProfileID { - case nvml.GPU_INSTANCE_PROFILE_1_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE - case nvml.GPU_INSTANCE_PROFILE_2_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_2_SLICE - case nvml.GPU_INSTANCE_PROFILE_3_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_3_SLICE - case nvml.GPU_INSTANCE_PROFILE_4_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_4_SLICE - case nvml.GPU_INSTANCE_PROFILE_6_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_6_SLICE - case nvml.GPU_INSTANCE_PROFILE_7_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_7_SLICE - case nvml.GPU_INSTANCE_PROFILE_8_SLICE: - return nvml.COMPUTE_INSTANCE_PROFILE_8_SLICE - } - return nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE -} - -func giProfileIDToSliceKey(giProfileID int) string { - switch giProfileID { - case nvml.GPU_INSTANCE_PROFILE_1_SLICE: - return "1g" - case nvml.GPU_INSTANCE_PROFILE_2_SLICE: - return "2g" - case nvml.GPU_INSTANCE_PROFILE_3_SLICE: - return "3g" - case nvml.GPU_INSTANCE_PROFILE_4_SLICE: - return "4g" - case nvml.GPU_INSTANCE_PROFILE_6_SLICE: - return "6g" - case nvml.GPU_INSTANCE_PROFILE_7_SLICE: - return "7g" - case nvml.GPU_INSTANCE_PROFILE_8_SLICE: - return "8g" - } - return "" -} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go deleted file mode 100644 index a549e7237d..0000000000 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/migmgr_placement_test.go +++ /dev/null @@ -1,75 +0,0 @@ -/* -Copyright 2026 The HAMi Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package plugin - -import ( - "testing" - - "github.com/NVIDIA/go-nvml/pkg/nvml" -) - -func TestChooseFreePlacementPackingDirection(t *testing.T) { - possible := []nvml.GpuInstancePlacement{ - {Start: 0, Size: 1}, - {Start: 6, Size: 1}, - {Start: 3, Size: 1}, - } - high, err := chooseFreePlacement(append([]nvml.GpuInstancePlacement(nil), possible...), nil, true) - if err != nil || high.Start != 6 { - t.Fatalf("high packing = %+v, %v; want start 6", high, err) - } - low, err := chooseFreePlacement(append([]nvml.GpuInstancePlacement(nil), possible...), nil, false) - if err != nil || low.Start != 0 { - t.Fatalf("low packing = %+v, %v; want start 0", low, err) - } - free, err := chooseFreePlacement(append([]nvml.GpuInstancePlacement(nil), possible...), map[uint32]uint32{6: 1}, true) - if err != nil || free.Start != 3 { - t.Fatalf("overlap-aware high packing = %+v, %v; want start 3", free, err) - } -} - -func TestPreferHighPlacementMatchesA100BalancedLayout(t *testing.T) { - if !preferHighPlacement(1) || !preferHighPlacement(3) { - t.Fatal("1g and 3g profiles must pack from high placements") - } - if preferHighPlacement(2) { - t.Fatal("2g profiles must pack from low placements") - } -} - -func TestPlacementCandidatesPreferPreviousThenFallback(t *testing.T) { - previous := nvml.GpuInstancePlacement{Start: 6, Size: 1} - possible := []nvml.GpuInstancePlacement{ - {Start: 5, Size: 1}, - {Start: 3, Size: 1}, - {Start: 6, Size: 1}, - } - want := []nvml.GpuInstancePlacement{ - {Start: 6, Size: 1}, - {Start: 5, Size: 1}, - {Start: 3, Size: 1}, - } - got := placementCandidates(previous, possible) - if len(got) != len(want) { - t.Fatalf("candidate count = %d, want %d: %+v", len(got), len(want), got) - } - for i := range want { - if got[i] != want[i] { - t.Fatalf("candidate %d = %+v, want %+v", i, got[i], want[i]) - } - } -} From 9c5464e6413e842b578d9c6346b1c49fa301a142 Mon Sep 17 00:00:00 2001 From: Jifei Wang Date: Fri, 7 Aug 2026 17:54:35 +0800 Subject: [PATCH 11/11] update MIG configuration and documentation Signed-off-by: Jifei Wang --- charts/hami/values.yaml | 8 +- docs/develop/dynamic-mig-migration.md | 300 +++++++++++++++++++++++++ docs/develop/dynamic-mig.md | 3 +- docs/develop/mig-dynamic-deallocate.md | 2 + 4 files changed, 306 insertions(+), 7 deletions(-) create mode 100644 docs/develop/dynamic-mig-migration.md diff --git a/charts/hami/values.yaml b/charts/hami/values.yaml index a7e13d1b11..593b2fedcb 100644 --- a/charts/hami/values.yaml +++ b/charts/hami/values.yaml @@ -550,12 +550,8 @@ prometheus: # device-config: # content: | # nvidia: -# knownMigGeometries: +# migProfileAllowlist: # - models: [ "RTX 6000 Ada" ] -# allowedGeometries: -# - - name: "1g.6gb" -# core: 25 -# memory: 6144 -# count: 4 +# profiles: [ "1g.6gb", "2g.12gb", "4g.24gb" ] device-config: content: "" diff --git a/docs/develop/dynamic-mig-migration.md b/docs/develop/dynamic-mig-migration.md new file mode 100644 index 0000000000..8e9aa7e7dd --- /dev/null +++ b/docs/develop/dynamic-mig-migration.md @@ -0,0 +1,300 @@ +# Migrating to HAMi Dynamic MIG + +This guide is intended for two groups of users: + +- users of the MIG Geometry/Template implementation on the HAMi `master` branch; and +- users of NVIDIA GPU Operator MIG Manager who manage fixed MIG geometries. + +The goal of this migration is not to promise that nodes will never need to be drained again. It is to remove draining from the routine profile-switching path. The scheduler reserves a specific MIG profile and physical placement for each Pod, the device plugin creates the corresponding GI/CI on demand, and the instance is reclaimed when the Pod terminates. + +> The current implementation does not support a seamless rolling migration that preserves legacy MIG Pods. For the initial handover, cordon, drain, upgrade, and validate nodes one at a time. After migration, routine mixed-profile scheduling usually no longer requires draining a node merely to switch the geometry of an entire GPU. + +## Why migrate + +Operating a fixed geometry typically starts by selecting a whole-GPU layout such as `all-1g`, `all-3g`, or a mixed configuration. When the workload mix changes and the current layout cannot satisfy a request, operators must clear the GPU, destroy the existing GI/CI instances, and apply a different layout. + +NVIDIA MIG Manager can trigger reconfiguration by changing `nvidia.com/mig.config`, but NVIDIA still requires that no user workloads are running on GPUs being reconfigured. Enabling or disabling MIG mode can also require a GPU reset or node reboot in some environments. Production procedures therefore commonly cordon or drain the node first. See the [NVIDIA GPU Operator MIG documentation](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-operator-mig.html). + +The implementation on the HAMi `master` branch is also centered on predefined geometries. When a request cannot fit the current geometry, the whole GPU must be switched to another template. This model works well for stable, long-lived resource pools, but mixed inference workloads, bursty profile demand, and frequently created short-lived jobs expose several costs: + +- operators must maintain memory, compute, instance counts, and geometry combinations for every GPU model; +- changing the layout affects the entire GPU, not just the instance required by a new request; +- running instances prevent geometry reconfiguration; +- unused instances in a fixed layout continue to occupy slices; and +- draining nodes and recreating workloads become part of capacity management. + +The current dynamic MIG implementation uses a reservation-first model: + +```text +device plugin publishes profiles and legal placements discovered through NVML + ↓ +scheduler selects GPU + profile + placement for a Pod + ↓ +Pod annotation persists the logical reservation + ↓ +device plugin creates the GI/CI at that placement + ↓ +device plugin records MIG UUID, GI ID, and CI ID + ↓ +the exact CI/GI is destroyed when the Pod terminates +``` + +The main differences are: + +| Area | Fixed Geometry / MIG Manager | HAMi Dynamic MIG | +| --- | --- | --- | +| Layout scope | Node or whole GPU | Per-Pod profile and placement | +| Profile capability | Manually configured geometry | Allowlist defines policy; NVML supplies actual capability | +| Instance creation | Pre-created fixed instance pool | Created from the reservation during `Allocate` | +| Instance reclamation | Usually retained until reconfiguration | Exact GI/CI reclaimed after the Pod terminates | +| Workload mix changes | May require a whole-GPU layout switch | Scheduled directly when a legal free placement exists | +| Restart recovery | Depends on the existing layout | Pod annotations are verified against NVML and adopted | + +Dynamic MIG does not remove MIG hardware constraints. A slice occupied by a GI cannot be converted in place into an overlapping layout. Fragmentation can temporarily prevent placement of a large profile. Enabling or disabling MIG mode, driver maintenance, and rollback can still require draining or rebooting a node. + +## Protocol changes to understand before migration + +### Configuration changes from geometries to a profile allowlist + +HAMi `master` configures complete geometries: + +```yaml +nvidia: + knownMigGeometries: + - models: ["A100-SXM4-40GB"] + allowedGeometries: + - - name: 1g.5gb + core: 14 + memory: 5120 + count: 7 + - - name: 2g.10gb + core: 28 + memory: 10240 + count: 3 + - name: 1g.5gb + core: 14 + memory: 5120 + count: 1 +``` + +The current implementation configures only the profiles that the cluster permits: + +```yaml +nvidia: + migProfileAllowlist: + - models: ["A100-SXM4-40GB"] + profiles: ["1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb"] +``` + +Operators no longer duplicate `core`, `memory`, `count`, or legal placement data. The node that owns the GPU discovers those values through NVML. The allowlist remains important: it defines which profiles the scheduler may use instead of automatically exposing every capability reported by the driver. + +When a legacy configuration contains several geometries, migration normally takes the union of their profile names. For example: + +```text +7 × 1g +3 × 2g + 1 × 1g +2 × 3g +1 × 7g +``` + +becomes: + +```yaml +profiles: ["1g.5gb", "2g.10gb", "3g.20gb", "7g.40gb"] +``` + +Verify profile names for each actual GPU model. Do not infer them solely from nominal memory capacity. Start with the model mappings in the current Chart and the device plugin discovery logs, then validate them with the target driver and hardware. + +### Allocation identity moves from a UUID suffix to a Pod annotation + +The legacy implementation encodes the template and slot in the device identifier, for example: + +```text +GPU-xxxxxxxx[1-2] +``` + +The current implementation stores the complete allocation identity in `hami.io/vgpu-mig-allocations`: + +```json +[ + { + "containerIndex": 0, + "deviceIndex": 0, + "gpuUUID": "GPU-xxxxxxxx", + "profile": "2g.10gb", + "placement": {"start": 2, "size": 2}, + "migUUID": "MIG-xxxxxxxx", + "gpuInstanceID": 4, + "computeInstanceID": 0 + } +] +``` + +The scheduler records the parent GPU, profile, and placement. After creating the instance, the device plugin adds the MIG UUID, GI ID, and CI ID. This annotation is the durable contract used to rebuild scheduler occupancy, recover after a device plugin restart, and reclaim instances. Users must not create or modify it manually. + +Legacy Pods do not contain this complete identity. A legacy template/slot index alone cannot reliably determine the physical placement for every GPU model and existing hardware layout. The current implementation therefore fails safely instead of guessing and risking overlapping slice allocations. This is the primary reason legacy MIG Pods must be drained during the initial upgrade. + +## Support boundaries + +### Operations that normally no longer require draining after migration + +- creating new Pods with different allowlisted profiles; +- deleting a Pod and reclaiming its MIG instance; +- reusing a legal free placement; and +- adopting active instances after a device plugin restart when their complete runtime annotations can be verified through NVML. + +### Operations that can still require draining or rebooting + +- the initial migration from the legacy geometry model; +- transferring hardware mutation ownership from NVIDIA MIG Manager; +- enabling or disabling MIG mode on a physical GPU; +- driver upgrades, GPU resets, or platform-required node reboots; +- rolling back to a version that understands only legacy geometries and UUID encoding; +- satisfying a new layout that would require moving a running GI/CI; and +- repairing a state in which HAMi Pod annotations cannot be correlated with NVML hardware state. + +## Migrating from HAMi master + +### Migration principle + +Do not allow a legacy scheduler and a current device plugin to serve MIG requests together: + +- the legacy scheduler produces template/slot encoding and does not create the new MIG reservation annotation; +- the current device plugin requires an explicit profile and placement in that reservation; and +- the current scheduler reads `migProfiles` capability, while legacy nodes publish `migtemplate`. + +A mixed-version deployment can conservatively report no capacity or fail during `Allocate`. Stop new MIG scheduling, upgrade the control plane, and then upgrade device plugins one node at a time. + +### Recommended procedure + +1. **Inventory and back up the current state.** Save the scheduler device ConfigMap, MIG node registration annotations, the list of active MIG Pods, and `nvidia-smi -L` output. Confirm that application Pods can be recreated. +2. **Stop new scheduling.** Cordon the MIG nodes being migrated so that the legacy scheduler cannot create new legacy-format allocations during the migration window. +3. **Drain legacy MIG Pods.** Wait for workloads to finish or migrate them elsewhere. Confirm that no user GPU process must be preserved. Do not restart only the device plugin and assume that legacy Pods can be adopted automatically. +4. **Migrate configuration.** Convert `knownMigGeometries` to `migProfileAllowlist`. Retain the profiles that administrators want to expose and remove manually maintained `core`, `memory`, `count`, and geometry combinations. +5. **Upgrade the scheduler.** Upgrade the scheduler and its configuration before upgrading device plugins, preventing a legacy scheduler from sending incompatible allocations to current nodes. +6. **Upgrade device plugins one node at a time.** Use a small number of nodes as canaries. On startup, idle GPUs are prepared in a clean MIG-ready state, which can destroy existing GI/CI instances on those idle GPUs. +7. **Validate node capability.** Confirm that device plugin logs show profile and placement discovery and that the Node registration annotation contains non-empty `migProfiles`. +8. **Validate the full lifecycle.** Create a MIG Pod and inspect its reservation annotation, the NVML-visible instance, and the MIG UUID visible to the container. Delete the Pod, wait for reconciliation, and confirm that the instance is released. +9. **Resume scheduling.** Uncordon nodes gradually after the canary succeeds, then restore production workloads. + +The project Helm Chart now provides `migProfileAllowlist` in its default configuration. If `device-config.content` or an external ConfigMap overrides the defaults, update that custom content as well. Legacy fields are not automatically converted to the new allowlist. + +## Migrating from NVIDIA MIG Manager + +### Establish a single owner first + +NVIDIA MIG Manager and HAMi Dynamic MIG both mutate GI/CI state. They must not manage the same physical GPU at the same time. MIG Manager can reapply a whole-GPU geometry in response to a Node label, while HAMi creates and destroys instances on demand from Pod reservations. + +GPU Operator can continue to provide the driver, Container Toolkit, DCGM, and other components, but MIG Manager must no longer apply a geometry to target nodes. The exact method for disabling that reconciliation depends on the GPU Operator version and deployment policy. Before migration, verify that MIG Manager on a target node will not continue to react to changes in `nvidia.com/mig.config`. + +### Recommended procedure + +1. **Record the current state.** Save `nvidia.com/mig.config`, `nvidia.com/mig.config.state`, the MIG Manager ConfigMap, custom geometries, and `nvidia-smi -L` output. +2. **Cordon target nodes and migrate GPU workloads.** NVIDIA requires that no user GPU workload be running during reconfiguration. HAMi's initial handover also needs an explicitly empty and verifiable baseline. +3. **Stop MIG Manager reconciliation on target nodes.** Ensure that it cannot reapply the previous geometry after HAMi creates GI/CI instances. Deleting a MIG Manager Pod once is not sufficient if its controller configuration immediately recreates it. +4. **Keep the required GPU Operator infrastructure.** The driver and container runtime remain prerequisites for HAMi GPU access. Stopping MIG Manager does not necessarily mean uninstalling GPU Operator. +5. **Configure the HAMi node with `mig` operating mode and set `migProfileAllowlist`.** The allowlist can be built from the profiles actually required in the previous MIG Manager configuration. +6. **Start the HAMi scheduler and device plugin.** The device plugin validates profiles and placements through NVML and removes old instances from idle GPUs to establish a predictable hardware baseline. +7. **Run canaries.** Begin with one profile and one Pod, then validate mixed profiles, capacity saturation, reclamation after Pod deletion, and recovery after a device plugin restart. +8. **Expand one node at a time.** Keep an unmigrated static MIG pool as short-term fallback capacity until the dynamic pool has passed production validation. + +## Validation checklist + +### Node capability + +- The registered GPU `mode` is `mig`. +- `migProfiles` is non-empty for every target GPU. +- Profile memory, slice count, and placements match NVML capability. +- Unsupported or non-allowlisted GPU models are not exposed accidentally. + +### Scheduling and realization + +- The Pod uses `nvidia.com/vgpu-mode: "mig"`. +- The scheduler writes `hami.io/vgpu-mig-allocations`. +- The selected profile satisfies the memory request and its placement does not overlap an active reservation. +- After a successful `Allocate`, the annotation contains a MIG UUID, GI ID, and CI ID. +- The MIG UUID visible in the container matches the annotation and NVML. + +Example workload: + +```yaml +apiVersion: v1 +kind: Pod +metadata: + name: mig-canary + annotations: + nvidia.com/vgpu-mode: "mig" +spec: + restartPolicy: Never + containers: + - name: workload + image: ubuntu:22.04 + command: ["bash", "-c", "sleep 3600"] + resources: + limits: + nvidia.com/gpu: 1 + nvidia.com/gpumem: 8000 +``` + +This example validates resource allocation and device injection only. A production canary should use a trusted image with CUDA or NVML tools and run an actual GPU workload. + +### Reclamation and recovery + +- Deleting the canary Pod releases its CI/GI without affecting instances owned by other Pods. +- A subsequent Pod can reuse the same slice. +- Active GPUs are not reset by startup cleanup during a device plugin restart. +- After restart, active instances with complete annotations are verified through NVML and adopted into the manager. +- A Kubernetes API or annotation read failure skips destructive reconciliation instead of guessing and deleting instances. + +### Recommended scenario coverage + +Test at least: + +1. creation and deletion of one `1g` Pod; +2. multiple non-overlapping `1g` instances on one GPU; +3. mixed placement of `1g`, `2g`, and `3g` instances; +4. a Pod remaining Pending when capacity is exhausted; +5. reuse of a placement after deleting a small instance; +6. a device plugin restart while a CUDA workload remains active; +7. safe failure when an annotation is missing or has only part of its runtime identity; and +8. no destructive reclamation while the Kubernetes API is temporarily unavailable. + +## Rollback + +### Before production workloads resume + +If canary validation fails: + +1. keep the node cordoned; +2. stop the current scheduler and device plugin from providing MIG service on the target node; +3. restore the legacy `knownMigGeometries` or MIG Manager configuration; +4. reapply a previously validated fixed geometry; and +5. verify device resource registration before uncordoning the node. + +### After new Dynamic MIG Pods have run + +Do not roll component binaries directly back to a legacy version. The legacy implementation does not understand the new reservation and placement protocol and cannot safely inherit the current manager state. Drain Dynamic MIG Pods again, stop HAMi from mutating GI/CI state, and then restore the legacy controller and fixed geometry. + +## Frequently asked questions + +### Does migration eliminate draining entirely? + +No. Routine creation and deletion of allowlisted profiles on free slices normally does not require draining. Initial handover, MIG mode changes, driver maintenance, re-layout that requires moving active instances, and rollback still can. + +### Can NVIDIA MIG Manager and HAMi manage different GPUs on the same node? + +Consider this only when both systems provide explicit, stable, and verified device-level ownership isolation. This migration guide does not rely on such a setup. By default, assign MIG hardware mutation on a target node to one controller so that whole-GPU geometry reapplication cannot conflict with Pod-level creation and deletion. + +### Why cannot a legacy Pod be migrated automatically from `GPU-UUID[template-slot]`? + +The legacy index describes a logical position in a scheduler template. Across GPU models, driver versions, and actual hardware state, it does not uniquely prove the GI placement, MIG UUID, GI ID, and CI ID. A conversion that is not verified through NVML could map two reservations to overlapping slices. The current implementation prioritizes safety and therefore establishes the new protocol only after legacy workloads have been drained. + +### Must users change their workload YAML? + +Usually not. Users continue to request `nvidia.com/gpu` and `nvidia.com/gpumem` and select `nvidia.com/vgpu-mode: "mig"`. The scheduler and device plugin manage `hami.io/vgpu-mig-allocations`; it is not a user-facing API. + +## Conclusion + +If a cluster's MIG demand is stable for long periods, pre-partitioned node pools remain a simple and reliable option. Dynamic MIG is most valuable when the profile mix changes with Pod lifecycles, static instance pools have low utilization, or geometry switching has become a routine operational burden. + +The migration itself requires one controlled drain because the legacy protocol does not contain enough information to prove the physical identity of existing instances. After migration, HAMi connects profile selection, placement reservation, GI/CI creation, and lifecycle reclamation through one convergent workflow. It reduces the frequency of routine reconfiguration and the scope of whole-GPU layout changes; it does not bypass NVIDIA MIG hardware or driver constraints. diff --git a/docs/develop/dynamic-mig.md b/docs/develop/dynamic-mig.md index d95d7ad310..927c989660 100644 --- a/docs/develop/dynamic-mig.md +++ b/docs/develop/dynamic-mig.md @@ -1,5 +1,7 @@ # NVIDIA GPU MPS and MIG dynamic slice plugin +> **Legacy design:** This document describes the `knownMigGeometries` implementation on the `master` branch. The current topology-aware implementation discovers MIG capability through NVML and uses `migProfileAllowlist`. See [Migrating to HAMi Dynamic MIG](./dynamic-mig-migration.md) for operational guidance and [Dynamic MIG Architecture](./mig-dynamic-deallocate.md) for the current design. + ## Special Thanks This feature will not be implemented without the help of @sailorvii. @@ -169,4 +171,3 @@ If you submit the example on an empty A100-PCIE-40GB node, then it will select a ``` Then start the container with 2g.10gb instances * 2 - diff --git a/docs/develop/mig-dynamic-deallocate.md b/docs/develop/mig-dynamic-deallocate.md index b3cc757af9..037281f1c2 100644 --- a/docs/develop/mig-dynamic-deallocate.md +++ b/docs/develop/mig-dynamic-deallocate.md @@ -1,5 +1,7 @@ # Dynamic MIG Architecture +For an operator-focused upgrade path from the `master` Geometry implementation or NVIDIA MIG Manager, see [Migrating to HAMi Dynamic MIG](./dynamic-mig-migration.md). + ## Context NVIDIA MIG divides a physical GPU into hardware-isolated compute instances. GPU models expose profile capacity and placement rules through NVML, while Kubernetes schedules workloads through declarative resources.