UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-26 18:18:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Replace go-nvlib/pkg/nvml with go-nvml/pkg/nvml

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2024-04-18 15:01:20 +02:00
parent 2019cd6f0a
commit b6be911eaa
12 changed files with 30 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/info/additional_info.go
View File

@@ -22,7 +22,7 @@ import (
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
"github.com/NVIDIA/go-nvlib/pkg/nvml"
"github.com/NVIDIA/go-nvml/pkg/nvml"
)
// additionalInfo allows for the info.Interface to be extened to implement the infoInterface.

37
internal/info/additional_info_test.go
View File

@@ -20,7 +20,8 @@ import (
"testing"
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
"github.com/NVIDIA/go-nvlib/pkg/nvml"
"github.com/NVIDIA/go-nvml/pkg/nvml"
"github.com/NVIDIA/go-nvml/pkg/nvml/mock"
"github.com/stretchr/testify/require"
)
@@ -32,7 +33,7 @@ func TestUsesNVGPUModule(t *testing.T) {
}{
{
description: "init failure returns false",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.ERROR_LIBRARY_NOT_FOUND
},
@@ -41,7 +42,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "no devices returns false",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -56,7 +57,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "DeviceGetCount error returns false",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -71,7 +72,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "Failure to get device name returns false",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -82,7 +83,7 @@ func TestUsesNVGPUModule(t *testing.T) {
return 1, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) {
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
return "", nvml.ERROR_UNKNOWN
},
@@ -94,7 +95,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "nested panic returns false",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -105,7 +106,7 @@ func TestUsesNVGPUModule(t *testing.T) {
return 1, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) {
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
panic("deep panic")
},
@@ -117,7 +118,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "Single device name with no nvgpu",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -128,7 +129,7 @@ func TestUsesNVGPUModule(t *testing.T) {
return 1, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) {
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
return "NVIDIA A100-SXM4-40GB", nvml.SUCCESS
},
@@ -140,7 +141,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "Single device name with nvgpu",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -151,7 +152,7 @@ func TestUsesNVGPUModule(t *testing.T) {
return 1, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) {
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
return "Orin (nvgpu)", nvml.SUCCESS
},
@@ -163,7 +164,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "Multiple device names with no nvgpu",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -174,7 +175,7 @@ func TestUsesNVGPUModule(t *testing.T) {
return 2, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) {
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
return "NVIDIA A100-SXM4-40GB", nvml.SUCCESS
},
@@ -186,7 +187,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "Multiple device names with nvgpu",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -197,7 +198,7 @@ func TestUsesNVGPUModule(t *testing.T) {
return 2, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) {
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
return "Orin (nvgpu)", nvml.SUCCESS
},
@@ -209,7 +210,7 @@ func TestUsesNVGPUModule(t *testing.T) {
},
{
description: "Mixed device names",
nvmllib: &nvml.InterfaceMock{
nvmllib: &mock.Interface{
InitFunc: func() nvml.Return {
return nvml.SUCCESS
},
@@ -226,7 +227,7 @@ func TestUsesNVGPUModule(t *testing.T) {
} else {
deviceName = "Orin (nvgpu)"
}
device := &nvml.DeviceMock{
device := &mock.Device{
GetNameFunc: func() (string, nvml.Return) {
return deviceName, nvml.SUCCESS
},

2
internal/info/auto.go
View File

@@ -19,7 +19,7 @@ package info
import (
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
"github.com/NVIDIA/go-nvlib/pkg/nvml"
"github.com/NVIDIA/go-nvml/pkg/nvml"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"