diff --git a/charts/hami/templates/scheduler/device-configmap.yaml b/charts/hami/templates/scheduler/device-configmap.yaml index 7f0b2b254d..1335115bc6 100644 --- a/charts/hami/templates/scheduler/device-configmap.yaml +++ b/charts/hami/templates/scheduler/device-configmap.yaml @@ -230,6 +230,23 @@ data: core: 100 memory: 184320 count: 1 + - 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 cambricon: resourceCountName: {{ .Values.mluResourceName }} resourceMemoryName: {{ .Values.mluResourceMem }} diff --git a/docker/Dockerfile b/docker/Dockerfile index d29cd5ae09..3d26ccb3bb 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -8,7 +8,7 @@ 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.12.2 +RUN go install github.com/NVIDIA/mig-parted/cmd/nvidia-mig-parted@v0.12.3 FROM $NVIDIA_IMAGE AS nvbuild RUN dnf install -y cmake git && dnf clean all diff --git a/docker/Dockerfile.hamimaster b/docker/Dockerfile.hamimaster index fa15934dd7..877204e1fb 100644 --- a/docker/Dockerfile.hamimaster +++ b/docker/Dockerfile.hamimaster @@ -8,7 +8,7 @@ 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.12.2 +RUN go install github.com/NVIDIA/mig-parted/cmd/nvidia-mig-parted@v0.12.3 FROM nvidia/cuda:13.3.0-base-ubuntu22.04 RUN apt-get update && \ diff --git a/docker/Dockerfile.withlib b/docker/Dockerfile.withlib index bdbc865679..2382576633 100644 --- a/docker/Dockerfile.withlib +++ b/docker/Dockerfile.withlib @@ -8,7 +8,7 @@ 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.12.2 +RUN go install github.com/NVIDIA/mig-parted/cmd/nvidia-mig-parted@v0.12.3 FROM nvidia/cuda:13.3.0-base-ubuntu22.04 RUN apt-get update && \ diff --git a/docs/develop/dynamic-mig.md b/docs/develop/dynamic-mig.md index b5805083c6..d95d7ad310 100644 --- a/docs/develop/dynamic-mig.md +++ b/docs/develop/dynamic-mig.md @@ -87,6 +87,20 @@ data: - name: 7g.79gb memory: 80896 count: 1 + - models: [ "RTX PRO 6000 Blackwell Server Edition" ] + allowedGeometries: + - + - name: 1g.24gb + memory: 24576 + count: 4 + - + - name: 2g.48gb + memory: 49152 + count: 2 + - + - name: 4g.96gb + memory: 98304 + count: 1 nodeconfig: - name: nodeA operatingmode: hami-core diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go index e6c2d4f358..48dee199be 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go @@ -60,6 +60,14 @@ func TestGenerateMigTemplate(t *testing.T) { {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}}, + }, + }, }, } @@ -134,6 +142,23 @@ func TestGenerateMigTemplate(t *testing.T) { "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 {