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:
Evan Lezar 2023-01-19 15:24:17 +01:00
parent f9330a4c2c
commit 95394e0fc8
3 changed files with 352 additions and 0 deletions

View 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)
}

View 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
}

View 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)
}