mirror of
https://github.com/NVIDIA/nvidia-container-toolkit
synced 2024-11-24 21:14:00 +00:00
Add internal/info/proc/devices package to read device majors
This change adds basic functionality to process the /proc/devices file to extract device majors. Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
parent
f9330a4c2c
commit
95394e0fc8
125
internal/info/proc/devices/devices.go
Normal file
125
internal/info/proc/devices/devices.go
Normal file
@ -0,0 +1,125 @@
|
||||
/*
|
||||
# Copyright (c) 2021, 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 devices
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Device major numbers and device names for NVIDIA devices
|
||||
const (
|
||||
NVIDIAUVMMinor = 0
|
||||
NVIDIAUVMToolsMinor = 1
|
||||
NVIDIACTLMinor = 255
|
||||
NVIDIAModesetMinor = 254
|
||||
|
||||
NVIDIAFrontend = Name("nvidia-frontend")
|
||||
NVIDIAGPU = NVIDIAFrontend
|
||||
NVIDIACaps = Name("nvidia-caps")
|
||||
NVIDIAUVM = Name("nvidia-uvm")
|
||||
|
||||
procDevicesPath = "/proc/devices"
|
||||
nvidiaDevicePrefix = "nvidia"
|
||||
)
|
||||
|
||||
// Name represents the name of a device as specified under /proc/devices
|
||||
type Name string
|
||||
|
||||
// Major represents a device major as specified under /proc/devices
|
||||
type Major int
|
||||
|
||||
// Devices represents the set of devices under /proc/devices
|
||||
//
|
||||
//go:generate moq -stub -out devices_mock.go . Devices
|
||||
type Devices interface {
|
||||
Exists(Name) bool
|
||||
Get(Name) (Major, bool)
|
||||
}
|
||||
|
||||
type devices map[Name]Major
|
||||
|
||||
var _ Devices = devices(nil)
|
||||
|
||||
// Exists checks if a Device with a given name exists or not
|
||||
func (d devices) Exists(name Name) bool {
|
||||
_, exists := d[name]
|
||||
return exists
|
||||
}
|
||||
|
||||
// Get a Device from Devices
|
||||
func (d devices) Get(name Name) (Major, bool) {
|
||||
device, exists := d[name]
|
||||
return device, exists
|
||||
}
|
||||
|
||||
// GetNVIDIADevices returns the set of NVIDIA Devices on the machine
|
||||
func GetNVIDIADevices() (Devices, error) {
|
||||
devicesFile, err := os.Open(procDevicesPath)
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error opening devices file: %v", err)
|
||||
}
|
||||
defer devicesFile.Close()
|
||||
|
||||
return nvidiaDeviceFrom(devicesFile), nil
|
||||
}
|
||||
|
||||
func nvidiaDeviceFrom(reader io.Reader) devices {
|
||||
allDevices := devicesFrom(reader)
|
||||
nvidiaDevices := make(devices)
|
||||
for n, d := range allDevices {
|
||||
if !strings.HasPrefix(string(n), nvidiaDevicePrefix) {
|
||||
continue
|
||||
}
|
||||
nvidiaDevices[n] = d
|
||||
}
|
||||
|
||||
return nvidiaDevices
|
||||
}
|
||||
|
||||
func devicesFrom(reader io.Reader) devices {
|
||||
allDevices := make(devices)
|
||||
scanner := bufio.NewScanner(reader)
|
||||
for scanner.Scan() {
|
||||
device, major, err := processProcDeviceLine(scanner.Text())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
allDevices[device] = major
|
||||
}
|
||||
return allDevices
|
||||
}
|
||||
|
||||
func processProcDeviceLine(line string) (Name, Major, error) {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
|
||||
var name Name
|
||||
var major Major
|
||||
|
||||
n, _ := fmt.Sscanf(trimmed, "%d %s", &major, &name)
|
||||
if n == 2 {
|
||||
return name, major, nil
|
||||
}
|
||||
|
||||
return "", 0, fmt.Errorf("unparsable line: %v", line)
|
||||
}
|
125
internal/info/proc/devices/devices_mock.go
Normal file
125
internal/info/proc/devices/devices_mock.go
Normal file
@ -0,0 +1,125 @@
|
||||
// Code generated by moq; DO NOT EDIT.
|
||||
// github.com/matryer/moq
|
||||
|
||||
package devices
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Ensure, that DevicesMock does implement Devices.
|
||||
// If this is not the case, regenerate this file with moq.
|
||||
var _ Devices = &DevicesMock{}
|
||||
|
||||
// DevicesMock is a mock implementation of Devices.
|
||||
//
|
||||
// func TestSomethingThatUsesDevices(t *testing.T) {
|
||||
//
|
||||
// // make and configure a mocked Devices
|
||||
// mockedDevices := &DevicesMock{
|
||||
// ExistsFunc: func(name Name) bool {
|
||||
// panic("mock out the Exists method")
|
||||
// },
|
||||
// GetFunc: func(name Name) (Major, bool) {
|
||||
// panic("mock out the Get method")
|
||||
// },
|
||||
// }
|
||||
//
|
||||
// // use mockedDevices in code that requires Devices
|
||||
// // and then make assertions.
|
||||
//
|
||||
// }
|
||||
type DevicesMock struct {
|
||||
// ExistsFunc mocks the Exists method.
|
||||
ExistsFunc func(name Name) bool
|
||||
|
||||
// GetFunc mocks the Get method.
|
||||
GetFunc func(name Name) (Major, bool)
|
||||
|
||||
// calls tracks calls to the methods.
|
||||
calls struct {
|
||||
// Exists holds details about calls to the Exists method.
|
||||
Exists []struct {
|
||||
// Name is the name argument value.
|
||||
Name Name
|
||||
}
|
||||
// Get holds details about calls to the Get method.
|
||||
Get []struct {
|
||||
// Name is the name argument value.
|
||||
Name Name
|
||||
}
|
||||
}
|
||||
lockExists sync.RWMutex
|
||||
lockGet sync.RWMutex
|
||||
}
|
||||
|
||||
// Exists calls ExistsFunc.
|
||||
func (mock *DevicesMock) Exists(name Name) bool {
|
||||
callInfo := struct {
|
||||
Name Name
|
||||
}{
|
||||
Name: name,
|
||||
}
|
||||
mock.lockExists.Lock()
|
||||
mock.calls.Exists = append(mock.calls.Exists, callInfo)
|
||||
mock.lockExists.Unlock()
|
||||
if mock.ExistsFunc == nil {
|
||||
var (
|
||||
bOut bool
|
||||
)
|
||||
return bOut
|
||||
}
|
||||
return mock.ExistsFunc(name)
|
||||
}
|
||||
|
||||
// ExistsCalls gets all the calls that were made to Exists.
|
||||
// Check the length with:
|
||||
//
|
||||
// len(mockedDevices.ExistsCalls())
|
||||
func (mock *DevicesMock) ExistsCalls() []struct {
|
||||
Name Name
|
||||
} {
|
||||
var calls []struct {
|
||||
Name Name
|
||||
}
|
||||
mock.lockExists.RLock()
|
||||
calls = mock.calls.Exists
|
||||
mock.lockExists.RUnlock()
|
||||
return calls
|
||||
}
|
||||
|
||||
// Get calls GetFunc.
|
||||
func (mock *DevicesMock) Get(name Name) (Major, bool) {
|
||||
callInfo := struct {
|
||||
Name Name
|
||||
}{
|
||||
Name: name,
|
||||
}
|
||||
mock.lockGet.Lock()
|
||||
mock.calls.Get = append(mock.calls.Get, callInfo)
|
||||
mock.lockGet.Unlock()
|
||||
if mock.GetFunc == nil {
|
||||
var (
|
||||
majorOut Major
|
||||
bOut bool
|
||||
)
|
||||
return majorOut, bOut
|
||||
}
|
||||
return mock.GetFunc(name)
|
||||
}
|
||||
|
||||
// GetCalls gets all the calls that were made to Get.
|
||||
// Check the length with:
|
||||
//
|
||||
// len(mockedDevices.GetCalls())
|
||||
func (mock *DevicesMock) GetCalls() []struct {
|
||||
Name Name
|
||||
} {
|
||||
var calls []struct {
|
||||
Name Name
|
||||
}
|
||||
mock.lockGet.RLock()
|
||||
calls = mock.calls.Get
|
||||
mock.lockGet.RUnlock()
|
||||
return calls
|
||||
}
|
102
internal/info/proc/devices/devices_test.go
Normal file
102
internal/info/proc/devices/devices_test.go
Normal file
@ -0,0 +1,102 @@
|
||||
/*
|
||||
# Copyright (c) 2021, 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 devices
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNvidiaDevices(t *testing.T) {
|
||||
devices := map[Name]Major{
|
||||
"nvidia-frontend": 195,
|
||||
"nvidia-nvlink": 234,
|
||||
"nvidia-caps": 235,
|
||||
"nvidia-uvm": 510,
|
||||
"nvidia-nvswitch": 511,
|
||||
}
|
||||
|
||||
nvidiaDevices := testDevices(devices)
|
||||
for name, major := range devices {
|
||||
device, exists := nvidiaDevices.Get(name)
|
||||
require.True(t, exists, "Unexpected missing device")
|
||||
require.Equal(t, device, major, "Unexpected device major")
|
||||
}
|
||||
_, exists := nvidiaDevices.Get("bogus")
|
||||
require.False(t, exists, "Unexpected 'bogus' device found")
|
||||
}
|
||||
|
||||
func TestProcessDeviceFile(t *testing.T) {
|
||||
testCases := []struct {
|
||||
lines []string
|
||||
expected devices
|
||||
}{
|
||||
{[]string{}, make(devices)},
|
||||
{[]string{"Not a valid line:"}, make(devices)},
|
||||
{[]string{"195 nvidia-frontend"}, devices{"nvidia-frontend": 195}},
|
||||
{[]string{"195 nvidia-frontend", "235 nvidia-caps"}, devices{"nvidia-frontend": 195, "nvidia-caps": 235}},
|
||||
{[]string{" 195 nvidia-frontend"}, devices{"nvidia-frontend": 195}},
|
||||
{[]string{"Not a valid line:", "", "195 nvidia-frontend"}, devices{"nvidia-frontend": 195}},
|
||||
{[]string{"195 not-nvidia-frontend"}, make(devices)},
|
||||
}
|
||||
for i, tc := range testCases {
|
||||
t.Run(fmt.Sprintf("testcase %d", i), func(t *testing.T) {
|
||||
contents := strings.NewReader(strings.Join(tc.lines, "\n"))
|
||||
d := nvidiaDeviceFrom(contents)
|
||||
|
||||
require.EqualValues(t, tc.expected, d)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDeviceFileLine(t *testing.T) {
|
||||
testCases := []struct {
|
||||
line string
|
||||
name Name
|
||||
major Major
|
||||
err bool
|
||||
}{
|
||||
{"", "", 0, true},
|
||||
{"0", "", 0, true},
|
||||
{"notint nvidia-frontend", "", 0, true},
|
||||
{"195 nvidia-frontend", "nvidia-frontend", 195, false},
|
||||
{" 195 nvidia-frontend", "nvidia-frontend", 195, false},
|
||||
}
|
||||
|
||||
for i, tc := range testCases {
|
||||
t.Run(fmt.Sprintf("testcase %d", i), func(t *testing.T) {
|
||||
name, major, err := processProcDeviceLine(tc.line)
|
||||
|
||||
require.Equal(t, tc.name, name)
|
||||
require.Equal(t, tc.major, major)
|
||||
if tc.err {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// testDevices creates a set of test NVIDIA devices
|
||||
func testDevices(d map[Name]Major) Devices {
|
||||
return devices(d)
|
||||
}
|
Loading…
Reference in New Issue
Block a user