From bb086d4b44556f90f6bdd3711557b56888570044 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Wed, 6 Apr 2022 14:30:49 +0200
Subject: [PATCH] Add auto discover mode and use this as the default

This change adds an 'auto' discover mode that attempts to select the correct mode
for a given platform. This currently attempts to detect whether the platform is a
Tegra-based system in which case the 'csv' discover mode is used. The 'legacy'
discover mode is used as the fallback.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 .../modifier/experimental.go                  | 49 ++++++++++++++++++-
 .../modifier/experimental_test.go             | 29 +++++++++++
 internal/config/runtime.go                    |  2 +-
 internal/config/runtime_test.go               |  2 +-
 4 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/cmd/nvidia-container-runtime/modifier/experimental.go b/cmd/nvidia-container-runtime/modifier/experimental.go
index 8ce5cf12..ead4ce19 100644
--- a/cmd/nvidia-container-runtime/modifier/experimental.go
+++ b/cmd/nvidia-container-runtime/modifier/experimental.go
@@ -18,6 +18,8 @@ package modifier
 
 import (
 	"fmt"
+	"os"
+	"strings"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
@@ -60,7 +62,8 @@ func NewExperimentalModifier(logger *logrus.Logger, cfg *config.Config, ociSpec
 	root := cfg.NVIDIAContainerCLIConfig.Root
 
 	var d discover.Discover
-	switch cfg.NVIDIAContainerRuntimeConfig.DiscoverMode {
+
+	switch resolveAutoDiscoverMode(logger, cfg.NVIDIAContainerRuntimeConfig.DiscoverMode) {
 	case "legacy":
 		legacyDiscoverer, err := discover.NewLegacyDiscoverer(logger, root)
 		if err != nil {
@@ -115,3 +118,47 @@ func (m experimental) Modify(spec *specs.Spec) error {
 
 	return specEdits.Modify(spec)
 }
+
+// resolveAutoDiscoverMode determines the correct discover mode for the specified platform if set to "auto"
+func resolveAutoDiscoverMode(logger *logrus.Logger, mode string) (rmode string) {
+	if mode != "auto" {
+		return mode
+	}
+	defer func() {
+		logger.Infof("Auto-detected discover mode as '%v'", rmode)
+	}()
+
+	isTegra, reason := isTegraSystem()
+	logger.Debugf("Is Tegra-based system? %v: %v", isTegra, reason)
+
+	if isTegra {
+		return "csv"
+	}
+
+	return "legacy"
+}
+
+// isTegraSystem returns true if the system is detected as a Tegra-based system
+func isTegraSystem() (bool, string) {
+	const tegraReleaseFile = "/etc/nv_tegra_release"
+	const tegraFamilyFile = "/sys/devices/soc0/family"
+
+	if info, err := os.Stat(tegraReleaseFile); err == nil && !info.IsDir() {
+		return true, fmt.Sprintf("%v found", tegraReleaseFile)
+	}
+
+	if info, err := os.Stat(tegraFamilyFile); err != nil || !info.IsDir() {
+		return false, fmt.Sprintf("%v not found", tegraFamilyFile)
+	}
+
+	contents, err := os.ReadFile(tegraFamilyFile)
+	if err != nil {
+		return false, fmt.Sprintf("could not read %v", tegraFamilyFile)
+	}
+
+	if strings.HasPrefix(strings.ToLower(string(contents)), "tegra") {
+		return true, fmt.Sprintf("%v has 'tegra' prefix", tegraFamilyFile)
+	}
+
+	return false, fmt.Sprintf("%v has no 'tegra' prefix", tegraFamilyFile)
+}
diff --git a/cmd/nvidia-container-runtime/modifier/experimental_test.go b/cmd/nvidia-container-runtime/modifier/experimental_test.go
index 271a2a87..9e09873a 100644
--- a/cmd/nvidia-container-runtime/modifier/experimental_test.go
+++ b/cmd/nvidia-container-runtime/modifier/experimental_test.go
@@ -320,3 +320,32 @@ func TestExperimentalModifier(t *testing.T) {
 		})
 	}
 }
+
+func TestResolveDiscoverMode(t *testing.T) {
+	logger, _ := testlog.NewNullLogger()
+
+	testCases := []struct {
+		description  string
+		mode         string
+		expectedMode string
+	}{
+		{
+			description:  "non-auto resolves to input",
+			mode:         "not-auto",
+			expectedMode: "not-auto",
+		},
+		// TODO: The following test is brittle in that it will break on Tegra-based systems.
+		// {
+		// 	description:  "auto resolves to legacy",
+		// 	mode:         "auto",
+		// 	expectedMode: "legacy",
+		// },
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			mode := resolveAutoDiscoverMode(logger, tc.mode)
+			require.EqualValues(t, tc.expectedMode, mode)
+		})
+	}
+}
diff --git a/internal/config/runtime.go b/internal/config/runtime.go
index c0c43011..9ae8de48 100644
--- a/internal/config/runtime.go
+++ b/internal/config/runtime.go
@@ -47,7 +47,7 @@ func getDefaultRuntimeConfig() *RuntimeConfig {
 	c := RuntimeConfig{
 		DebugFilePath: "/dev/null",
 		Experimental:  false,
-		DiscoverMode:  "legacy",
+		DiscoverMode:  "auto",
 	}
 
 	return &c
diff --git a/internal/config/runtime_test.go b/internal/config/runtime_test.go
index 32ea7a0b..fe78b529 100644
--- a/internal/config/runtime_test.go
+++ b/internal/config/runtime_test.go
@@ -63,7 +63,7 @@ func TestGetConfig(t *testing.T) {
 				NVIDIAContainerRuntimeConfig: RuntimeConfig{
 					DebugFilePath: "/dev/null",
 					Experimental:  false,
-					DiscoverMode:  "legacy",
+					DiscoverMode:  "auto",
 				},
 			},
 		},