Copy files from nvidia-container-toolkit

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/proc/devices.go

@@ -0,0 +1,116 @@
+/*
+# 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 proc
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"strings"
+)
+
+const (
+	procDevicesPath    = "/proc/devices"
+	nvidiaDevicePrefix = "nvidia"
+)
+
+// Device represents a device as specified under /proc/devices
+type Device struct {
+	Name  string
+	Major int
+}
+
+// NvidiaDevices represents the set of nvidia owned devices under /proc/devices
+type NvidiaDevices interface {
+	Exists(name string) bool
+	Get(name string) (Device, bool)
+}
+
+type nvidiaDevices map[string]Device
+
+var _ NvidiaDevices = nvidiaDevices(nil)
+
+// Exists checks if a Device with a given name exists or not
+func (d nvidiaDevices) Exists(name string) bool {
+	_, exists := d[name]
+	return exists
+}
+
+// Get a Device from NvidiaDevices
+func (d nvidiaDevices) Get(name string) (Device, bool) {
+	device, exists := d[name]
+	return device, exists
+}
+
+func (d nvidiaDevices) add(devices ...Device) {
+	for _, device := range devices {
+		d[device.Name] = device
+	}
+}
+
+// NewMockNvidiaDevices returns NvidiaDevices populated from the devices passed in
+func NewMockNvidiaDevices(devices ...Device) NvidiaDevices {
+	nvds := make(nvidiaDevices)
+	nvds.add(devices...)
+	return nvds
+}
+
+// GetNvidiaDevices returns the set of NvidiaDevices on the machine
+func GetNvidiaDevices() (NvidiaDevices, 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 processDeviceFile(devicesFile), nil
+}
+
+func processDeviceFile(devicesFile io.Reader) NvidiaDevices {
+	nvidiaDevices := make(nvidiaDevices)
+	scanner := bufio.NewScanner(devicesFile)
+	for scanner.Scan() {
+		device, major, err := processProcDeviceLine(scanner.Text())
+		if err != nil {
+			log.Printf("Skipping line in devices file: %v", err)
+			continue
+		}
+		if strings.HasPrefix(device, nvidiaDevicePrefix) {
+			nvidiaDevices.add(Device{device, major})
+		}
+	}
+	return nvidiaDevices
+}
+
+func processProcDeviceLine(line string) (string, int, error) {
+	trimmed := strings.TrimSpace(line)
+
+	var name string
+	var major int
+
+	n, _ := fmt.Sscanf(trimmed, "%d %s", &major, &name)
+	if n == 2 {
+		return name, major, nil
+	}
+
+	return "", 0, fmt.Errorf("unparsable line: %v", line)
+}

internal/proc/devices_test.go

@@ -0,0 +1,92 @@
+/*
+# 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 proc
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNvidiaDevices(t *testing.T) {
+	devices := []Device{
+		{"nvidia-frontend", 195},
+		{"nvidia-nvlink", 234},
+		{"nvidia-caps", 235},
+		{"nvidia-uvm", 510},
+		{"nvidia-nvswitch", 511},
+	}
+
+	nvidiaDevices := NewMockNvidiaDevices(devices...)
+	for _, d := range devices {
+		device, exists := nvidiaDevices.Get(d.Name)
+		require.True(t, exists, "Unexpected missing device")
+		require.Equal(t, device.Name, d.Name, "Unexpected device name")
+		require.Equal(t, device.Major, d.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 []Device
+	}{
+		{[]string{}, []Device{}},
+		{[]string{"Not a valid line:"}, []Device{}},
+		{[]string{"195 nvidia-frontend"}, []Device{{"nvidia-frontend", 195}}},
+		{[]string{"195 nvidia-frontend", "235 nvidia-caps"}, []Device{{"nvidia-frontend", 195}, {"nvidia-caps", 235}}},
+		{[]string{"  195 nvidia-frontend"}, []Device{{"nvidia-frontend", 195}}},
+		{[]string{"Not a valid line:", "", "195 nvidia-frontend"}, []Device{{"nvidia-frontend", 195}}},
+		{[]string{"195 not-nvidia-frontend"}, []Device{}},
+	}
+	for _, tc := range testCases {
+		contents := strings.NewReader(strings.Join(tc.lines, "\n"))
+		d := processDeviceFile(contents)
+		require.Equalf(t, NewMockNvidiaDevices(tc.expected...), d, "testCase: %v", tc)
+	}
+}
+
+func TestProcessDeviceFileLine(t *testing.T) {
+	testCases := []struct {
+		line  string
+		name  string
+		major int
+		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 _, tc := range testCases {
+		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)
+		}
+
+	}
+}