diff --git a/internal/info/proc/devices/devices.go b/internal/info/proc/devices/devices.go new file mode 100644 index 00000000..377f20d1 --- /dev/null +++ b/internal/info/proc/devices/devices.go @@ -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) +} diff --git a/internal/info/proc/devices/devices_mock.go b/internal/info/proc/devices/devices_mock.go new file mode 100644 index 00000000..315541c2 --- /dev/null +++ b/internal/info/proc/devices/devices_mock.go @@ -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 +} diff --git a/internal/info/proc/devices/devices_test.go b/internal/info/proc/devices/devices_test.go new file mode 100644 index 00000000..78bce52f --- /dev/null +++ b/internal/info/proc/devices/devices_test.go @@ -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) +}