Switch to go-nvml instead of gpu-monitoring-tools

Makefile

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-MODULE := github.com/nvidia/go-gpuallocator
+MODULE := github.com/NVIDIA/go-gpuallocator
 
 DOCKER ?= docker
 
-GOLANG_VERSION := 1.15
+GOLANG_VERSION := 1.20
 
 ifeq ($(IMAGE),)
 REGISTRY ?= nvidia
@@ -94,11 +94,10 @@ coverage: test
 $(DOCKER_TARGETS): docker-%: .build-image
 	@echo "Running 'make $(*)' in docker container $(BUILDIMAGE)"
 	$(DOCKER) run \
-		--rm \
+		--rm --privileged \
 		-e GOCACHE=/tmp/.cache \
-		-v $(PWD):$(PWD) \
+		-v $(PWD):$(PWD):z \
 		-w $(PWD) \
-		--user $$(id -u):$$(id -g) \
 		$(BUILDIMAGE) \
 			make $(*)
 

docker/Dockerfile.devel

@@ -11,7 +11,8 @@
 # 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.
-ARG GOLANG_VERSION=1.15
+ARG GOLANG_VERSION=1.20
 FROM golang:${GOLANG_VERSION}
 
-RUN go get -u golang.org/x/lint/golint
+RUN go install golang.org/x/lint/golint@latest
+RUN go install github.com/matryer/moq@latest

go.mod

@@ -1,8 +1,11 @@
 module github.com/NVIDIA/go-gpuallocator
 
-go 1.15
+go 1.20
 
-require github.com/NVIDIA/gpu-monitoring-tools v0.0.0-20201109160820-d08ea3cdcce4
+require (
+	github.com/NVIDIA/go-nvml v0.12.0-1
+	gitlab.com/nvidia/cloud-native/go-nvlib v0.0.0-20230711114409-e6fcecfa6790
+)
 
 replace (
 	k8s.io/api => k8s.io/api v0.18.2

gpuallocator/allocator.go

@@ -1,12 +1,26 @@
-// Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+/**
+# Copyright (c) NVIDIA CORPORATION.  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
+#
+# 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 gpuallocator
 
 import (
 	"fmt"
 	"runtime"
 
-	"github.com/NVIDIA/gpu-monitoring-tools/bindings/go/nvml"
+	"github.com/NVIDIA/go-gpuallocator/internal/gpulib"
 )
 
 // Allocator defines the primary object for allocating and freeing the
@@ -45,9 +59,9 @@ func NewBestEffortAllocator() (*Allocator, error) {
 
 // NewAllocator creates a new Allocator using the given allocation policy
 func NewAllocator(policy Policy) (*Allocator, error) {
-	err := nvml.Init()
-	if err != nil {
-		return nil, fmt.Errorf("error initializing NVML: %v", err)
+	ret := gpulib.Init()
+	if ret.Value() != gpulib.SUCCESS {
+		return nil, fmt.Errorf("error initializing NVML: %v", ret.Error())
 	}
 
 	devices, err := NewDevices()
@@ -58,8 +72,8 @@ func NewAllocator(policy Policy) (*Allocator, error) {
 	allocator := newAllocatorFrom(devices, policy)
 
 	runtime.SetFinalizer(allocator, func(allocator *Allocator) {
-		// Explicitly ignore any errors from nvml.Shutdown().
-		_ = nvml.Shutdown()
+		// Explicitly ignore any errors from gpulib.Shutdown().
+		_ = gpulib.Shutdown()
 	})
 
 	return allocator, nil

gpuallocator/besteffort_policy.go

@@ -1,11 +1,25 @@
-// Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+/**
+# Copyright (c) NVIDIA CORPORATION.  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
+#
+# 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 gpuallocator
 
 import (
 	"fmt"
 
-	"github.com/NVIDIA/gpu-monitoring-tools/bindings/go/nvml"
+	"github.com/NVIDIA/go-gpuallocator/internal/gpulib"
 )
 
 type bestEffortPolicy struct{}
@@ -313,41 +327,41 @@ func calculateGPUPairScore(gpu0 *Device, gpu1 *Device) int {
 
 	for _, link := range gpu0.Links[gpu1.Index] {
 		switch link.Type {
-		case nvml.P2PLinkCrossCPU:
+		case gpulib.P2PLinkCrossCPU:
 			score += 10
-		case nvml.P2PLinkSameCPU:
+		case gpulib.P2PLinkSameCPU:
 			score += 20
-		case nvml.P2PLinkHostBridge:
+		case gpulib.P2PLinkHostBridge:
 			score += 30
-		case nvml.P2PLinkMultiSwitch:
+		case gpulib.P2PLinkMultiSwitch:
 			score += 40
-		case nvml.P2PLinkSingleSwitch:
+		case gpulib.P2PLinkSingleSwitch:
 			score += 50
-		case nvml.P2PLinkSameBoard:
+		case gpulib.P2PLinkSameBoard:
 			score += 60
-		case nvml.SingleNVLINKLink:
+		case gpulib.SingleNVLINKLink:
 			score += 100
-		case nvml.TwoNVLINKLinks:
+		case gpulib.TwoNVLINKLinks:
 			score += 200
-		case nvml.ThreeNVLINKLinks:
+		case gpulib.ThreeNVLINKLinks:
 			score += 300
-		case nvml.FourNVLINKLinks:
+		case gpulib.FourNVLINKLinks:
 			score += 400
-		case nvml.FiveNVLINKLinks:
+		case gpulib.FiveNVLINKLinks:
 			score += 500
-		case nvml.SixNVLINKLinks:
+		case gpulib.SixNVLINKLinks:
 			score += 600
-		case nvml.SevenNVLINKLinks:
+		case gpulib.SevenNVLINKLinks:
 			score += 700
-		case nvml.EightNVLINKLinks:
+		case gpulib.EightNVLINKLinks:
 			score += 800
-		case nvml.NineNVLINKLinks:
+		case gpulib.NineNVLINKLinks:
 			score += 900
-		case nvml.TenNVLINKLinks:
+		case gpulib.TenNVLINKLinks:
 			score += 1000
-		case nvml.ElevenNVLINKLinks:
+		case gpulib.ElevenNVLINKLinks:
 			score += 1100
-		case nvml.TwelveNVLINKLinks:
+		case gpulib.TwelveNVLINKLinks:
 			score += 1200
 		}
 	}