mirror of
				https://github.com/NVIDIA/nvidia-container-toolkit
				synced 2025-06-26 18:18:24 +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