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