UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2024-11-22 00:08:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge branch 'CNT-4778/bump-gonvlib' into 'main'

Browse Source 
Update to github.com/NVIDIA/go-nvlib@f3264c8a6a7a

See merge request nvidia/container-toolkit/container-toolkit!520
This commit is contained in:
Christopher Desiniotis 2023-12-13 18:11:59 +00:00
commit 9baed635d1
7 changed files with 103 additions and 175 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
go.mod
View File

@ -3,10 +3,9 @@ module github.com/NVIDIA/nvidia-container-toolkit
go 1.20
require (
github.com/NVIDIA/go-nvlib v0.0.0-20231116150931-9fd385bace0d
github.com/NVIDIA/go-nvlib v0.0.0-20231212194527-f3264c8a6a7a
github.com/NVIDIA/go-nvml v0.12.0-1.0.20231020145430-e06766c5e74f
github.com/fsnotify/fsnotify v1.5.4
github.com/google/uuid v1.4.0
github.com/opencontainers/runtime-spec v1.1.0
github.com/pelletier/go-toml v1.9.4
github.com/sirupsen/logrus v1.9.0
@ -21,6 +20,7 @@ require (
require (
github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/google/uuid v1.4.0 // indirect
github.com/hashicorp/errwrap v1.1.0 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/opencontainers/runtime-tools v0.9.1-0.20221107090550-2e043c6bd626 // indirect

4
go.sum
View File

@ -1,6 +1,6 @@
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/NVIDIA/go-nvlib v0.0.0-20231116150931-9fd385bace0d h1:XxRHS7eNkZVcPpZZmUcoT4oO8FEcoYKn06sooQh5niU=
github.com/NVIDIA/go-nvlib v0.0.0-20231116150931-9fd385bace0d/go.mod h1:HPFNPAYqQeoos58MKUboWsdZMu71EzSQrbmd+QBRD40=
github.com/NVIDIA/go-nvlib v0.0.0-20231212194527-f3264c8a6a7a h1:aHaNKihxpWzWnV3yoVkit3bhOF7cg2ScCbzW+gepQ/E=
github.com/NVIDIA/go-nvlib v0.0.0-20231212194527-f3264c8a6a7a/go.mod h1:U82N6/xKp6OnoqpALBH0C5SO59Buu4sX1Z3rQtBsBKQ=
github.com/NVIDIA/go-nvml v0.12.0-1.0.20231020145430-e06766c5e74f h1:FTblgO87K1vPB8tcwM5EOFpFf6UpsrlDpErPm25mFWE=
github.com/NVIDIA/go-nvml v0.12.0-1.0.20231020145430-e06766c5e74f/go.mod h1:7ruy85eOM73muOc/I37euONSwEyFqZsv5ED9AogD4G0=
github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=

76
pkg/nvcdi/identifier.go
View File

@ -1,76 +0,0 @@
package nvcdi
import (
"strconv"
"strings"
"github.com/google/uuid"
)
type identifier string
// isGPUIndex checks if an identifier is a full GPU index
func (i identifier) isGpuIndex() bool {
if _, err := strconv.ParseUint(string(i), 10, 0); err != nil {
return false
}
return true
}
// isMigIndex checks if an identifier is a MIG index
func (i identifier) isMigIndex() bool {
split := strings.SplitN(string(i), ":", 2)
if len(split) != 2 {
return false
}
for _, s := range split {
if _, err := strconv.ParseUint(s, 10, 0); err != nil {
return false
}
}
return true
}
// isUUID checks if an identifier is a UUID
func (i identifier) isUUID() bool {
return i.isGpuUUID() || i.isMigUUID()
}
// isGpuUUID checks if an identifier is a GPU UUID
// A GPU UUID must be of the form GPU-b1028956-cfa2-0990-bf4a-5da9abb51763
func (i identifier) isGpuUUID() bool {
if !strings.HasPrefix(string(i), "GPU-") {
return false
}
_, err := uuid.Parse(strings.TrimPrefix(string(i), "GPU-"))
return err == nil
}
// isMigUUID checks if an identifier is a MIG UUID
// A MIG UUID can be of one of two forms:
// - MIG-b1028956-cfa2-0990-bf4a-5da9abb51763
// - MIG-GPU-b1028956-cfa2-0990-bf4a-5da9abb51763/3/0
func (i identifier) isMigUUID() bool {
if !strings.HasPrefix(string(i), "MIG-") {
return false
}
suffix := strings.TrimPrefix(string(i), "MIG-")
_, err := uuid.Parse(suffix)
if err == nil {
return true
}
split := strings.SplitN(suffix, "/", 3)
if len(split) != 3 {
return false
}
if !identifier(split[0]).isGpuUUID() {
return false
}
for _, s := range split[1:] {
_, err := strconv.ParseUint(s, 10, 0)
if err != nil {
return false
}
}
return true
}

90
pkg/nvcdi/identifier_test.go
View File

@ -1,90 +0,0 @@
package nvcdi
import (
"fmt"
"testing"
"github.com/stretchr/testify/require"
)
func TestIsGpuIndex(t *testing.T) {
testCases := []struct {
id string
expected bool
}{
{"", false},
{"0", true},
{"1", true},
{"not an integer", false},
}
for i, tc := range testCases {
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
actual := identifier(tc.id).isGpuIndex()
require.Equal(t, tc.expected, actual)
})
}
}
func TestIsMigIndex(t *testing.T) {
testCases := []struct {
id string
expected bool
}{
{"", false},
{"0", false},
{"not an integer", false},
{"0:0", true},
{"0:0:0", false},
{"0:0.0", false},
{"0:foo", false},
{"foo:0", false},
}
for i, tc := range testCases {
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
actual := identifier(tc.id).isMigIndex()
require.Equal(t, tc.expected, actual)
})
}
}
func TestIsGpuUUID(t *testing.T) {
testCases := []struct {
id string
expected bool
}{
{"", false},
{"0", false},
{"not an integer", false},
{"GPU-foo", false},
{"GPU-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", true},
{"MIG-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
{"ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
}
for i, tc := range testCases {
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
actual := identifier(tc.id).isGpuUUID()
require.Equal(t, tc.expected, actual)
})
}
}
func TestIsMigUUID(t *testing.T) {
testCases := []struct {
id string
expected bool
}{
{"", false},
{"0", false},
{"not an integer", false},
{"MIG-foo", false},
{"MIG-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", true},
{"GPU-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
{"ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
}
for i, tc := range testCases {
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
actual := identifier(tc.id).isMigUUID()
require.Equal(t, tc.expected, actual)
})
}
}

8
pkg/nvcdi/lib-nvml.go
View File

@ -133,20 +133,20 @@ func (l *nvmllib) getNVMLDevicesByID(identifiers ...string) ([]nvml.Device, erro
func (l *nvmllib) getNVMLDeviceByID(id string) (nvml.Device, error) {
var err error
devID := identifier(id)
devID := device.Identifier(id)
if devID.isUUID() {
if devID.IsUUID() {
return l.nvmllib.DeviceGetHandleByUUID(id)
}
if devID.isGpuIndex() {
if devID.IsGpuIndex() {
if idx, err := strconv.Atoi(id); err == nil {
return l.nvmllib.DeviceGetHandleByIndex(idx)
}
return nil, fmt.Errorf("failed to convert device index to an int: %w", err)
}
if devID.isMigIndex() {
if devID.IsMigIndex() {
var gpuIdx, migIdx int
var parent nvml.Device
split := strings.SplitN(id, ":", 2)

94
vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/identifier.go generated vendored Normal file
View File

@ -0,0 +1,94 @@
/*
* 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 device
import (
"strconv"
"strings"
"github.com/google/uuid"
)
// Identifier can be used to refer to a GPU or MIG device.
// This includes a device index or UUID.
type Identifier string
// IsGpuIndex checks if an identifier is a full GPU index
func (i Identifier) IsGpuIndex() bool {
if _, err := strconv.ParseUint(string(i), 10, 0); err != nil {
return false
}
return true
}
// IsMigIndex checks if an identifier is a MIG index
func (i Identifier) IsMigIndex() bool {
split := strings.Split(string(i), ":")
if len(split) != 2 {
return false
}
for _, s := range split {
if !Identifier(s).IsGpuIndex() {
return false
}
}
return true
}
// IsUUID checks if an identifier is a UUID
func (i Identifier) IsUUID() bool {
return i.IsGpuUUID() || i.IsMigUUID()
}
// IsGpuUUID checks if an identifier is a GPU UUID
// A GPU UUID must be of the form GPU-b1028956-cfa2-0990-bf4a-5da9abb51763
func (i Identifier) IsGpuUUID() bool {
if !strings.HasPrefix(string(i), "GPU-") {
return false
}
_, err := uuid.Parse(strings.TrimPrefix(string(i), "GPU-"))
return err == nil
}
// IsMigUUID checks if an identifier is a MIG UUID
// A MIG UUID can be of one of two forms:
// - MIG-b1028956-cfa2-0990-bf4a-5da9abb51763
// - MIG-GPU-b1028956-cfa2-0990-bf4a-5da9abb51763/3/0
func (i Identifier) IsMigUUID() bool {
if !strings.HasPrefix(string(i), "MIG-") {
return false
}
suffix := strings.TrimPrefix(string(i), "MIG-")
_, err := uuid.Parse(suffix)
if err == nil {
return true
}
split := strings.Split(suffix, "/")
if len(split) != 3 {
return false
}
if !Identifier(split[0]).IsGpuUUID() {
return false
}
for _, s := range split[1:] {
_, err := strconv.ParseUint(s, 10, 0)
if err != nil {
return false
}
}
return true
}

2
vendor/modules.txt vendored
View File

@ -1,4 +1,4 @@
# github.com/NVIDIA/go-nvlib v0.0.0-20231116150931-9fd385bace0d
# github.com/NVIDIA/go-nvlib v0.0.0-20231212194527-f3264c8a6a7a
## explicit; go 1.20
github.com/NVIDIA/go-nvlib/pkg/nvlib/device
github.com/NVIDIA/go-nvlib/pkg/nvlib/info