Skip to content

Commit 483a59b

Browse files
ShourenShouren
committed
refactor: simplify code by using modern constructs
Signed-off-by: Shouren Yang <[email protected]>
1 parent 38bed6f commit 483a59b

File tree

27 files changed

+90
-162
lines changed

27 files changed

+90
-162
lines changed

Diff for: ‎cmd/vGPUmonitor/feedback.go

+6-6
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ func getUsedGPUPid() ([]uint, nvml.Return) {
6262
if err != nvml.SUCCESS {
6363
return []uint{}, err
6464
}
65-
for i := 0; i < count; i++ {
65+
for i := range count {
6666
device, err := nvml.DeviceGetHandleByIndex(i)
6767
if err != nvml.SUCCESS {
6868
return []uint{}, err
@@ -167,11 +167,11 @@ func getUsedGPUPid() ([]uint, nvml.Return) {
167167
//}
168168

169169
func CheckBlocking(utSwitchOn map[string]UtilizationPerDevice, p int, c *nvidia.ContainerUsage) bool {
170-
for i := 0; i < c.Info.DeviceMax(); i++ {
170+
for i := range c.Info.DeviceMax() {
171171
uuid := c.Info.DeviceUUID(i)
172172
_, ok := utSwitchOn[uuid]
173173
if ok {
174-
for i := 0; i < p; i++ {
174+
for i := range p {
175175
if utSwitchOn[uuid][i] > 0 {
176176
return true
177177
}
@@ -184,11 +184,11 @@ func CheckBlocking(utSwitchOn map[string]UtilizationPerDevice, p int, c *nvidia.
184184

185185
// Check whether task with higher priority use GPU or there are other tasks with the same priority.
186186
func CheckPriority(utSwitchOn map[string]UtilizationPerDevice, p int, c *nvidia.ContainerUsage) bool {
187-
for i := 0; i < c.Info.DeviceMax(); i++ {
187+
for i := range c.Info.DeviceMax() {
188188
uuid := c.Info.DeviceUUID(i)
189189
_, ok := utSwitchOn[uuid]
190190
if ok {
191-
for i := 0; i < p; i++ {
191+
for i := range p {
192192
if utSwitchOn[uuid][i] > 0 {
193193
return true
194194
}
@@ -210,7 +210,7 @@ func Observe(lister *nvidia.ContainerLister) {
210210
if recentKernel > 0 {
211211
recentKernel--
212212
if recentKernel > 0 {
213-
for i := 0; i < c.Info.DeviceMax(); i++ {
213+
for i := range c.Info.DeviceMax() {
214214
//for _, devuuid := range val.sr.uuids {
215215
// Null device condition
216216
if !c.Info.IsValidUUID(i) {

Diff for: ‎cmd/vGPUmonitor/metrics.go

+3-6
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ func (cc ClusterManagerCollector) collectGPUInfo(ch chan<- prometheus.Metric) er
203203
return err
204204
}
205205

206-
for ii := 0; ii < devnum; ii++ {
206+
for ii := range devnum {
207207
if err := cc.collectGPUDeviceMetrics(ch, ii); err != nil {
208208
klog.Error("Failed to collect metrics for GPU device ", ii, ": ", err)
209209
}
@@ -353,7 +353,7 @@ func (cc ClusterManagerCollector) collectContainerMetrics(ch chan<- prometheus.M
353353
}
354354

355355
// Iterate through each device
356-
for i := 0; i < c.Info.DeviceNum(); i++ {
356+
for i := range c.Info.DeviceNum() {
357357
uuid := c.Info.DeviceUUID(i)
358358
if len(uuid) < 40 {
359359
klog.Errorf("Invalid UUID length for device %d in Pod %s/%s, Container %s", i, pod.Namespace, pod.Name, ctr.Name)
@@ -398,10 +398,7 @@ func (cc ClusterManagerCollector) collectContainerMetrics(ch chan<- prometheus.M
398398

399399
// Send last kernel time metric if valid
400400
if lastKernelTime > 0 {
401-
lastSec := nowSec - lastKernelTime
402-
if lastSec < 0 {
403-
lastSec = 0
404-
}
401+
lastSec := max(nowSec-lastKernelTime, 0)
405402
if err := sendMetric(ch, ctrDeviceLastKernelDesc, prometheus.GaugeValue, float64(lastSec), labels...); err != nil {
406403
klog.Errorf("Failed to send last kernel time metric for device %d in Pod %s/%s, Container %s: %v", i, pod.Namespace, pod.Name, ctr.Name, err)
407404
return err

Diff for: ‎hack/tools/preferredimports/preferredimports.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ type analyzer struct {
5454
fset *token.FileSet // positions are relative to fset
5555
ctx build.Context
5656
failed bool
57-
donePaths map[string]interface{}
57+
donePaths map[string]any
5858
}
5959

6060
func newAnalyzer() *analyzer {
@@ -64,7 +64,7 @@ func newAnalyzer() *analyzer {
6464
a := &analyzer{
6565
fset: token.NewFileSet(),
6666
ctx: ctx,
67-
donePaths: make(map[string]interface{}),
67+
donePaths: make(map[string]any),
6868
}
6969

7070
return a

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ type NvidiaDevicePlugin struct {
9696

9797
server *grpc.Server
9898
health chan *rm.Device
99-
stop chan interface{}
99+
stop chan any
100100
}
101101

102102
func readFromConfigFile(sConfig *nvidia.NvidiaConfig) (string, error) {
@@ -183,7 +183,7 @@ func NewNvidiaDevicePlugin(config *nvidia.DeviceConfig, resourceManager rm.Resou
183183
func (plugin *NvidiaDevicePlugin) initialize() {
184184
plugin.server = grpc.NewServer([]grpc.ServerOption{}...)
185185
plugin.health = make(chan *rm.Device)
186-
plugin.stop = make(chan interface{})
186+
plugin.stop = make(chan any)
187187
}
188188

189189
func (plugin *NvidiaDevicePlugin) cleanup() {

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/allocate.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ func (r *resourceManager) distributedAlloc(available, required []string, size in
116116
// to the list of devices to allocate, remove it from the candidate list,
117117
// down its available count in the replicas map, and repeat.
118118
var devices []string
119-
for i := 0; i < needed; i++ {
119+
for range needed {
120120
sort.Slice(candidates, func(i, j int) bool {
121121
iid := AnnotatedID(candidates[i]).GetID()
122122
jid := AnnotatedID(candidates[j]).GetID()

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/device_map.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -321,7 +321,7 @@ func updateDeviceMapWithReplicas(config *nvidia.DeviceConfig, oDevices DeviceMap
321321
name = r.Rename
322322
}
323323
for _, id := range ids {
324-
for i := 0; i < r.Replicas; i++ {
324+
for i := range r.Replicas {
325325
annotatedID := string(NewAnnotatedID(id, i))
326326
replicatedDevice := *(oDevices[r.Name][id])
327327
replicatedDevice.ID = annotatedID

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/devices.go

+3-7
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ import (
3737
"strconv"
3838
"strings"
3939

40+
"slices"
41+
4042
kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
4143
)
4244

@@ -214,13 +216,7 @@ func (d Device) AlignedAllocationSupported() bool {
214216
return false
215217
}
216218

217-
for _, p := range d.Paths {
218-
if p == "/dev/dxg" {
219-
return false
220-
}
221-
}
222-
223-
return true
219+
return !slices.Contains(d.Paths, "/dev/dxg")
224220
}
225221

226222
// IsMigDevice returns checks whether d is a MIG device or not.

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/health.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ const (
5555
)
5656

5757
// CheckHealth performs health checks on a set of devices, writing to the 'unhealthy' channel with any unhealthy devices
58-
func (r *nvmlResourceManager) checkHealth(stop <-chan interface{}, devices Devices, unhealthy chan<- *Device) error {
58+
func (r *nvmlResourceManager) checkHealth(stop <-chan any, devices Devices, unhealthy chan<- *Device) error {
5959
disableHealthChecks := strings.ToLower(os.Getenv(envDisableHealthChecks))
6060
if disableHealthChecks == "all" {
6161
disableHealthChecks = allHealthChecks

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/nvml_manager.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,6 @@ func (r *nvmlResourceManager) GetDevicePaths(ids []string) []string {
115115
}
116116

117117
// CheckHealth performs health checks on a set of devices, writing to the 'unhealthy' channel with any unhealthy devices
118-
func (r *nvmlResourceManager) CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device) error {
118+
func (r *nvmlResourceManager) CheckHealth(stop <-chan any, unhealthy chan<- *Device) error {
119119
return r.checkHealth(stop, r.devices, unhealthy)
120120
}

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/rm.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ type ResourceManager interface {
5757
Devices() Devices
5858
GetDevicePaths([]string) []string
5959
GetPreferredAllocation(available, required []string, size int) ([]string, error)
60-
CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device) error
60+
CheckHealth(stop <-chan any, unhealthy chan<- *Device) error
6161
}
6262

6363
// NewResourceManagers returns a []ResourceManager, one for each resource in 'config'.

Diff for: ‎pkg/device-plugin/nvidiadevice/nvinternal/rm/tegra_manager.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,6 @@ func (r *tegraResourceManager) GetDevicePaths(ids []string) []string {
8787
}
8888

8989
// CheckHealth is disabled for the tegraResourceManager
90-
func (r *tegraResourceManager) CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device) error {
90+
func (r *tegraResourceManager) CheckHealth(stop <-chan any, unhealthy chan<- *Device) error {
9191
return nil
9292
}

Diff for: ‎pkg/device/ascend/device.go

+5-13
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ import (
2929
"github.com/Project-HAMi/HAMi/pkg/util"
3030
"github.com/Project-HAMi/HAMi/pkg/util/nodelock"
3131

32+
"slices"
33+
3234
corev1 "k8s.io/api/core/v1"
3335
"k8s.io/apimachinery/pkg/api/resource"
3436
"k8s.io/klog/v2"
@@ -57,7 +59,7 @@ var (
5759
)
5860

5961
func (dev *Devices) trimMemory(m int64) (int64, string) {
60-
for i := 0; i < len(dev.config.Templates); i++ {
62+
for i := range dev.config.Templates {
6163
if m <= dev.config.Templates[i].Memory {
6264
return dev.config.Templates[i].Memory, dev.config.Templates[i].Name
6365
}
@@ -216,25 +218,15 @@ func (dev *Devices) CheckUUID(annos map[string]string, d util.DeviceUsage) bool
216218
klog.V(5).Infof("check uuid for ascend user uuid [%s], device id is %s", userUUID, d.ID)
217219
// use , symbol to connect multiple uuid
218220
userUUIDs := strings.Split(userUUID, ",")
219-
for _, uuid := range userUUIDs {
220-
if d.ID == uuid {
221-
return true
222-
}
223-
}
224-
return false
221+
return slices.Contains(userUUIDs, d.ID)
225222
}
226223

227224
noUserUUID, ok := annos[dev.noUseUUIDAnno]
228225
if ok {
229226
klog.V(5).Infof("check uuid for ascend not user uuid [%s], device id is %s", noUserUUID, d.ID)
230227
// use , symbol to connect multiple uuid
231228
noUserUUIDs := strings.Split(noUserUUID, ",")
232-
for _, uuid := range noUserUUIDs {
233-
if d.ID == uuid {
234-
return false
235-
}
236-
}
237-
return true
229+
return !slices.Contains(noUserUUIDs, d.ID)
238230
}
239231
return true
240232
}

Diff for: ‎pkg/device/cambricon/device.go

+5-13
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ import (
2828
"github.com/Project-HAMi/HAMi/pkg/util"
2929
"github.com/Project-HAMi/HAMi/pkg/util/client"
3030

31+
"slices"
32+
3133
corev1 "k8s.io/api/core/v1"
3234
"k8s.io/apimachinery/pkg/api/resource"
3335
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -94,7 +96,7 @@ func (dev *CambriconDevices) setNodeLock(node *corev1.Node) error {
9496
}
9597

9698
patchedAnnotation, err := json.Marshal(
97-
map[string]interface{}{
99+
map[string]any{
98100
"metadata": map[string]map[string]string{"annotations": {
99101
DsmluLockTime: time.Now().Format(time.RFC3339),
100102
}}})
@@ -222,25 +224,15 @@ func (dev *CambriconDevices) CheckUUID(annos map[string]string, d util.DeviceUsa
222224
klog.V(5).Infof("check uuid for mlu user uuid [%s], device id is %s", userUUID, d.ID)
223225
// use , symbol to connect multiple uuid
224226
userUUIDs := strings.Split(userUUID, ",")
225-
for _, uuid := range userUUIDs {
226-
if d.ID == uuid {
227-
return true
228-
}
229-
}
230-
return false
227+
return slices.Contains(userUUIDs, d.ID)
231228
}
232229

233230
noUserUUID, ok := annos[MLUNoUseUUID]
234231
if ok {
235232
klog.V(5).Infof("check uuid for mlu not user uuid [%s], device id is %s", noUserUUID, d.ID)
236233
// use , symbol to connect multiple uuid
237234
noUserUUIDs := strings.Split(noUserUUID, ",")
238-
for _, uuid := range noUserUUIDs {
239-
if d.ID == uuid {
240-
return false
241-
}
242-
}
243-
return true
235+
return !slices.Contains(noUserUUIDs, d.ID)
244236
}
245237
return true
246238
}

Diff for: ‎pkg/device/devices.go

+10-10
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ func InitDevicesWithConfig(config *Config) error {
9696
var initErrors []error
9797

9898
// Helper function to initialize devices and handle errors
99-
initializeDevice := func(deviceType string, commonWord string, initFunc func(interface{}) (Devices, error), config interface{}) {
99+
initializeDevice := func(deviceType string, commonWord string, initFunc func(any) (Devices, error), config any) {
100100
klog.Infof("Initializing %s device", commonWord)
101101
device, err := initFunc(config)
102102
if err != nil {
@@ -113,52 +113,52 @@ func InitDevicesWithConfig(config *Config) error {
113113
deviceInitializers := []struct {
114114
deviceType string
115115
commonWord string
116-
initFunc func(interface{}) (Devices, error)
117-
config interface{}
116+
initFunc func(any) (Devices, error)
117+
config any
118118
}{
119-
{nvidia.NvidiaGPUDevice, nvidia.NvidiaGPUCommonWord, func(cfg interface{}) (Devices, error) {
119+
{nvidia.NvidiaGPUDevice, nvidia.NvidiaGPUCommonWord, func(cfg any) (Devices, error) {
120120
nvidiaConfig, ok := cfg.(nvidia.NvidiaConfig)
121121
if !ok {
122122
return nil, fmt.Errorf("invalid configuration for %s", nvidia.NvidiaGPUCommonWord)
123123
}
124124
return nvidia.InitNvidiaDevice(nvidiaConfig), nil
125125
}, config.NvidiaConfig},
126-
{cambricon.CambriconMLUDevice, cambricon.CambriconMLUCommonWord, func(cfg interface{}) (Devices, error) {
126+
{cambricon.CambriconMLUDevice, cambricon.CambriconMLUCommonWord, func(cfg any) (Devices, error) {
127127
cambriconConfig, ok := cfg.(cambricon.CambriconConfig)
128128
if !ok {
129129
return nil, fmt.Errorf("invalid configuration for %s", cambricon.CambriconMLUCommonWord)
130130
}
131131
return cambricon.InitMLUDevice(cambriconConfig), nil
132132
}, config.CambriconConfig},
133-
{hygon.HygonDCUDevice, hygon.HygonDCUCommonWord, func(cfg interface{}) (Devices, error) {
133+
{hygon.HygonDCUDevice, hygon.HygonDCUCommonWord, func(cfg any) (Devices, error) {
134134
hygonConfig, ok := cfg.(hygon.HygonConfig)
135135
if !ok {
136136
return nil, fmt.Errorf("invalid configuration for %s", hygon.HygonDCUCommonWord)
137137
}
138138
return hygon.InitDCUDevice(hygonConfig), nil
139139
}, config.HygonConfig},
140-
{iluvatar.IluvatarGPUDevice, iluvatar.IluvatarGPUCommonWord, func(cfg interface{}) (Devices, error) {
140+
{iluvatar.IluvatarGPUDevice, iluvatar.IluvatarGPUCommonWord, func(cfg any) (Devices, error) {
141141
iluvatarConfig, ok := cfg.(iluvatar.IluvatarConfig)
142142
if !ok {
143143
return nil, fmt.Errorf("invalid configuration for %s", iluvatar.IluvatarGPUCommonWord)
144144
}
145145
return iluvatar.InitIluvatarDevice(iluvatarConfig), nil
146146
}, config.IluvatarConfig},
147-
{mthreads.MthreadsGPUDevice, mthreads.MthreadsGPUCommonWord, func(cfg interface{}) (Devices, error) {
147+
{mthreads.MthreadsGPUDevice, mthreads.MthreadsGPUCommonWord, func(cfg any) (Devices, error) {
148148
mthreadsConfig, ok := cfg.(mthreads.MthreadsConfig)
149149
if !ok {
150150
return nil, fmt.Errorf("invalid configuration for %s", mthreads.MthreadsGPUCommonWord)
151151
}
152152
return mthreads.InitMthreadsDevice(mthreadsConfig), nil
153153
}, config.MthreadsConfig},
154-
{metax.MetaxGPUDevice, metax.MetaxGPUCommonWord, func(cfg interface{}) (Devices, error) {
154+
{metax.MetaxGPUDevice, metax.MetaxGPUCommonWord, func(cfg any) (Devices, error) {
155155
metaxConfig, ok := cfg.(metax.MetaxConfig)
156156
if !ok {
157157
return nil, fmt.Errorf("invalid configuration for %s", metax.MetaxGPUCommonWord)
158158
}
159159
return metax.InitMetaxDevice(metaxConfig), nil
160160
}, config.MetaxConfig},
161-
{metax.MetaxSGPUDevice, metax.MetaxSGPUCommonWord, func(cfg interface{}) (Devices, error) {
161+
{metax.MetaxSGPUDevice, metax.MetaxSGPUCommonWord, func(cfg any) (Devices, error) {
162162
metaxConfig, ok := cfg.(metax.MetaxConfig)
163163
if !ok {
164164
return nil, fmt.Errorf("invalid configuration for %s", metax.MetaxGPUCommonWord)

Diff for: ‎pkg/device/devices_test.go

+5-8
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ import (
2727
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2828
"k8s.io/client-go/kubernetes/fake"
2929

30+
"slices"
31+
3032
"github.com/Project-HAMi/HAMi/pkg/device/ascend"
3133
"github.com/Project-HAMi/HAMi/pkg/device/cambricon"
3234
"github.com/Project-HAMi/HAMi/pkg/device/hygon"
@@ -174,8 +176,8 @@ func Test_LoadConfig(t *testing.T) {
174176

175177
dataDrivenTests := []struct {
176178
name string
177-
expectedConfig interface{}
178-
actualConfig interface{}
179+
expectedConfig any
180+
actualConfig any
179181
}{
180182
{"NVIDIA Config", createNvidiaConfig(), configData.NvidiaConfig},
181183
{"Cambricon Config", createCambriconConfig(), configData.CambriconConfig},
@@ -351,12 +353,7 @@ func setupTest(t *testing.T) (map[string]string, map[string]Devices) {
351353
}
352354

353355
func containsString(slice []string, str string) bool {
354-
for _, v := range slice {
355-
if v == str {
356-
return true
357-
}
358-
}
359-
return false
356+
return slices.Contains(slice, str)
360357
}
361358

362359
// Test_InitDevicesWithConfig_Success tests the initialization of devices with the provided config.

0 commit comments

Comments
 (0)