Skip to content

Commit b893505

Browse files
feat: enhance NVIDIA device plugin with MIG support and add unit tests
1 parent 4639608 commit b893505

File tree

4 files changed

+345
-0
lines changed

4 files changed

+345
-0
lines changed

cmd/device-plugin/nvidia/main.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import (
3939
"github.com/Project-HAMi/HAMi/pkg/nvidia-plugin/pkg/plugin"
4040
"github.com/Project-HAMi/HAMi/pkg/nvidia-plugin/pkg/rm"
4141
"github.com/Project-HAMi/HAMi/pkg/nvidia-plugin/pkg/watch"
42+
"github.com/Project-HAMi/HAMi/pkg/util"
4243
)
4344

4445
type options struct {
@@ -236,6 +237,13 @@ func loadConfig(c *cli.Context, flags []cli.Flag) (*spec.Config, error) {
236237
func start(c *cli.Context, o *options) error {
237238
klog.InfoS(fmt.Sprintf("Starting %s", c.App.Name), "version", c.App.Version)
238239

240+
util.NodeName = os.Getenv(util.NodeNameEnvName)
241+
// watcher, err := newFSWatcher(kubeletdevicepluginv1beta1.DevicePluginPath)
242+
// if err != nil {
243+
// return fmt.Errorf("failed to create FS watcher: %v", err)
244+
// }
245+
// defer watcher.Close()
246+
239247
kubeletSocketDir := filepath.Dir(o.kubeletSocket)
240248
klog.Infof("Starting FS watcher for %v", kubeletSocketDir)
241249
watcher, err := watch.Files(kubeletSocketDir)
@@ -244,6 +252,7 @@ func start(c *cli.Context, o *options) error {
244252
}
245253
defer watcher.Close()
246254

255+
klog.Infof("Start working on node %s", util.NodeName)
247256
klog.Info("Starting OS watcher.")
248257
sigs := watch.Signals(syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
249258

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The HAMi Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
/*
10+
* Licensed to NVIDIA CORPORATION under one or more contributor
11+
* license agreements. See the NOTICE file distributed with
12+
* this work for additional information regarding copyright
13+
* ownership. NVIDIA CORPORATION licenses this file to you under
14+
* the Apache License, Version 2.0 (the "License"); you may
15+
* not use this file except in compliance with the License.
16+
* You may obtain a copy of the License at
17+
*
18+
* http://www.apache.org/licenses/LICENSE-2.0
19+
*
20+
* Unless required by applicable law or agreed to in writing,
21+
* software distributed under the License is distributed on an
22+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
23+
* KIND, either express or implied. See the License for the
24+
* specific language governing permissions and limitations
25+
* under the License.
26+
*/
27+
28+
/*
29+
* Modifications Copyright The HAMi Authors. See
30+
* GitHub history for details.
31+
*/
32+
33+
package plugin
34+
35+
import (
36+
"fmt"
37+
"os/exec"
38+
"strconv"
39+
"strings"
40+
"time"
41+
42+
"github.com/NVIDIA/go-nvml/pkg/nvml"
43+
"k8s.io/klog/v2"
44+
45+
"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
46+
"github.com/Project-HAMi/HAMi/pkg/util"
47+
)
48+
49+
func (plugin *NvidiaDevicePlugin) getNumaInformation(idx int) (int, error) {
50+
cmd := exec.Command("nvidia-smi", "topo", "-m")
51+
out, err := cmd.CombinedOutput()
52+
if err != nil {
53+
return 0, err
54+
}
55+
klog.V(5).InfoS("nvidia-smi topo -m output", "result", string(out))
56+
return parseNvidiaNumaInfo(idx, string(out))
57+
}
58+
59+
func parseNvidiaNumaInfo(idx int, nvidiaTopoStr string) (int, error) {
60+
result := 0
61+
numaAffinityColumnIndex := 0
62+
for index, val := range strings.Split(nvidiaTopoStr, "\n") {
63+
if !strings.Contains(val, "GPU") {
64+
continue
65+
}
66+
// Example: GPU0 X 0-7 N/A N/A
67+
// Many values are separated by two tabs, but this actually represents 5 values instead of 7
68+
// So add logic to remove multiple tabs
69+
words := strings.Split(strings.ReplaceAll(val, "\t\t", "\t"), "\t")
70+
klog.V(5).InfoS("parseNumaInfo", "words", words)
71+
// get numa affinity column number
72+
if index == 0 {
73+
for columnIndex, headerVal := range words {
74+
// The topology output of a single card is as follows:
75+
// GPU0 CPU Affinity NUMA Affinity GPU NUMA ID
76+
// GPU0 X 0-7 N/A N/A
77+
//Legend: Other content omitted
78+
79+
// The topology output in the case of multiple cards is as follows:
80+
// GPU0 GPU1 CPU Affinity NUMA Affinity
81+
// GPU0 X PHB 0-31 N/A
82+
// GPU1 PHB X 0-31 N/A
83+
// Legend: Other content omitted
84+
85+
// We need to get the value of the NUMA Affinity column, but their column indexes are inconsistent,
86+
// so we need to get the index first and then get the value.
87+
if strings.Contains(headerVal, "NUMA Affinity") {
88+
// The header is one column less than the actual row.
89+
numaAffinityColumnIndex = columnIndex
90+
continue
91+
}
92+
}
93+
continue
94+
}
95+
klog.V(5).InfoS("nvidia-smi topo -m row output", "row output", words, "length", len(words))
96+
if strings.Contains(words[0], fmt.Sprint(idx)) {
97+
if len(words) <= numaAffinityColumnIndex || words[numaAffinityColumnIndex] == "N/A" {
98+
klog.InfoS("current card has not established numa topology", "gpu row info", words, "index", idx)
99+
return 0, nil
100+
}
101+
result, err := strconv.Atoi(words[numaAffinityColumnIndex])
102+
if err != nil {
103+
return result, err
104+
}
105+
}
106+
}
107+
return result, nil
108+
}
109+
110+
func (plugin *NvidiaDevicePlugin) getAPIDevices() *[]*util.DeviceInfo {
111+
devs := plugin.Devices()
112+
klog.V(5).InfoS("getAPIDevices", "devices", devs)
113+
nvml.Init()
114+
res := make([]*util.DeviceInfo, 0, len(devs))
115+
for UUID := range devs {
116+
ndev, ret := nvml.DeviceGetHandleByUUID(UUID)
117+
if ret != nvml.SUCCESS {
118+
klog.Errorln("nvml new device by index error uuid=", UUID, "err=", ret)
119+
panic(0)
120+
}
121+
idx, ret := ndev.GetIndex()
122+
if ret != nvml.SUCCESS {
123+
klog.Errorln("nvml get index error ret=", ret)
124+
panic(0)
125+
}
126+
memoryTotal := 0
127+
memory, ret := ndev.GetMemoryInfo()
128+
if ret == nvml.SUCCESS {
129+
memoryTotal = int(memory.Total)
130+
} else {
131+
klog.Error("nvml get memory error ret=", ret)
132+
panic(0)
133+
}
134+
Model, ret := ndev.GetName()
135+
if ret != nvml.SUCCESS {
136+
klog.Error("nvml get name error ret=", ret)
137+
panic(0)
138+
}
139+
140+
registeredmem := int32(memoryTotal / 1024 / 1024)
141+
if plugin.schedulerConfig.DeviceMemoryScaling != 1 {
142+
registeredmem = int32(float64(registeredmem) * plugin.schedulerConfig.DeviceMemoryScaling)
143+
}
144+
klog.Infoln("MemoryScaling=", plugin.schedulerConfig.DeviceMemoryScaling, "registeredmem=", registeredmem)
145+
health := true
146+
for _, val := range devs {
147+
if strings.Compare(val.ID, UUID) == 0 {
148+
// when NVIDIA-Tesla P4, the device info is : ID:GPU-e290caca-2f0c-9582-acab-67a142b61ffa,Health:Healthy,Topology:nil,
149+
// it is more reasonable to think of healthy as case-insensitive
150+
if strings.EqualFold(val.Health, "healthy") {
151+
health = true
152+
} else {
153+
health = false
154+
}
155+
break
156+
}
157+
}
158+
numa, err := plugin.getNumaInformation(idx)
159+
if err != nil {
160+
klog.ErrorS(err, "failed to get numa information", "idx", idx)
161+
}
162+
res = append(res, &util.DeviceInfo{
163+
ID: UUID,
164+
Index: uint(idx),
165+
Count: int32(plugin.schedulerConfig.DeviceSplitCount),
166+
Devmem: registeredmem,
167+
Devcore: int32(plugin.schedulerConfig.DeviceCoreScaling * 100),
168+
Type: fmt.Sprintf("%v-%v", "NVIDIA", Model),
169+
Numa: numa,
170+
Mode: plugin.operatingMode,
171+
Health: health,
172+
})
173+
klog.Infof("nvml registered device id=%v, memory=%v, type=%v, numa=%v", idx, registeredmem, Model, numa)
174+
}
175+
return &res
176+
}
177+
178+
func (plugin *NvidiaDevicePlugin) RegistrInAnnotation() error {
179+
devices := plugin.getAPIDevices()
180+
klog.InfoS("start working on the devices", "devices", devices)
181+
annos := make(map[string]string)
182+
node, err := util.GetNode(util.NodeName)
183+
if err != nil {
184+
klog.Errorln("get node error", err.Error())
185+
return err
186+
}
187+
encodeddevices := util.EncodeNodeDevices(*devices)
188+
annos[nvidia.HandshakeAnnos] = "Reported " + time.Now().String()
189+
annos[nvidia.RegisterAnnos] = encodeddevices
190+
klog.Infof("patch node with the following annos %v", fmt.Sprintf("%v", annos))
191+
err = util.PatchNodeAnnotations(node, annos)
192+
193+
if err != nil {
194+
klog.Errorln("patch node error", err.Error())
195+
}
196+
return err
197+
}
198+
199+
func (plugin *NvidiaDevicePlugin) WatchAndRegister() {
200+
klog.Info("Starting WatchAndRegister")
201+
errorSleepInterval := time.Second * 5
202+
successSleepInterval := time.Second * 30
203+
for {
204+
err := plugin.RegistrInAnnotation()
205+
if err != nil {
206+
klog.Errorf("Failed to register annotation: %v", err)
207+
klog.Infof("Retrying in %v seconds...", errorSleepInterval)
208+
time.Sleep(errorSleepInterval)
209+
} else {
210+
klog.Infof("Successfully registered annotation. Next check in %v seconds...", successSleepInterval)
211+
time.Sleep(successSleepInterval)
212+
}
213+
}
214+
}
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The HAMi Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
/*
10+
* Licensed to NVIDIA CORPORATION under one or more contributor
11+
* license agreements. See the NOTICE file distributed with
12+
* this work for additional information regarding copyright
13+
* ownership. NVIDIA CORPORATION licenses this file to you under
14+
* the Apache License, Version 2.0 (the "License"); you may
15+
* not use this file except in compliance with the License.
16+
* You may obtain a copy of the License at
17+
*
18+
* http://www.apache.org/licenses/LICENSE-2.0
19+
*
20+
* Unless required by applicable law or agreed to in writing,
21+
* software distributed under the License is distributed on an
22+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
23+
* KIND, either express or implied. See the License for the
24+
* specific language governing permissions and limitations
25+
* under the License.
26+
*/
27+
28+
/*
29+
* Modifications Copyright The HAMi Authors. See
30+
* GitHub history for details.
31+
*/
32+
33+
package plugin
34+
35+
import "testing"
36+
37+
func Test_parseNvidiaNumaInfo(t *testing.T) {
38+
39+
tests := []struct {
40+
name string
41+
idx int
42+
nvidiaTopoStr string
43+
want int
44+
wantErr bool
45+
}{
46+
{
47+
name: "single Tesla P4 NUMA",
48+
idx: 0,
49+
nvidiaTopoStr: `GPU0 CPU Affinity NUMA Affinity ...
50+
...`,
51+
want: 0,
52+
wantErr: false,
53+
},
54+
{
55+
name: "two Tesla P4 NUMA topo with index 0",
56+
idx: 0,
57+
nvidiaTopoStr: `GPU0 GPU1 CPU Affinity NUMA Affinity ...
58+
...`,
59+
want: 0,
60+
wantErr: false,
61+
},
62+
{
63+
name: "two Tesla P4 NUMA topo with index 1",
64+
idx: 1,
65+
nvidiaTopoStr: `GPU0 GPU1 CPU Affinity NUMA Affinity ...
66+
...`,
67+
want: 0,
68+
wantErr: false,
69+
},
70+
{
71+
name: "NUMA Affinity is empty",
72+
idx: 0,
73+
nvidiaTopoStr: `GPU0 CPU Affinity NUMA Affinity GPU NUMA ID
74+
GPU0 X`,
75+
want: 0,
76+
wantErr: false,
77+
},
78+
}
79+
80+
for _, tt := range tests {
81+
t.Run(tt.name, func(t *testing.T) {
82+
got, err := parseNvidiaNumaInfo(tt.idx, tt.nvidiaTopoStr)
83+
if (err != nil) != tt.wantErr {
84+
t.Errorf("parseNvidiaNumaInfo() error = %v, wantErr %v", err, tt.wantErr)
85+
return
86+
}
87+
if got != tt.want {
88+
t.Errorf("parseNvidiaNumaInfo() got = %v, want %v", got, tt.want)
89+
}
90+
})
91+
}
92+
}

0 commit comments

Comments
 (0)