From e07c7f0fa24f467b613b9136e781ff570357c071 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 24 May 2022 09:34:35 +0200
Subject: [PATCH] Ignore NVIDIA_REQUIRE_JETPACK* for image requirements

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 CHANGELOG.md                             |  1 +
 internal/config/image/cuda_image.go      |  3 +-
 internal/config/image/cuda_image_test.go | 52 ++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b37b9890..06f92232 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## v1.10.0-rc.3
 
+* Ignore NVIDIA_REQUIRE_JETPACK* environment variables for requirement checks
 * Fix bug in detection of Tegra systems where `/sys/devices/soc0/family` is ignored
 * Fix bug where links to devices were detected as devices
 * [libnvida-container] Fix bug introduced when adding libcudadebugger.so to list of libraries
diff --git a/internal/config/image/cuda_image.go b/internal/config/image/cuda_image.go
index c92b75a5..2315d4d0 100644
--- a/internal/config/image/cuda_image.go
+++ b/internal/config/image/cuda_image.go
@@ -29,6 +29,7 @@ const (
 	envCUDAVersion      = "CUDA_VERSION"
 	envNVRequirePrefix  = "NVIDIA_REQUIRE_"
 	envNVRequireCUDA    = envNVRequirePrefix + "CUDA"
+	envNVRequireJetpack = envNVRequirePrefix + "JETPACK"
 	envNVDisableRequire = "NVIDIA_DISABLE_REQUIRE"
 )
 
@@ -84,7 +85,7 @@ func (i CUDA) GetRequirements() ([]string, error) {
 	// All variables with the "NVIDIA_REQUIRE_" prefix are passed to nvidia-container-cli
 	var requirements []string
 	for name, value := range i {
-		if strings.HasPrefix(name, envNVRequirePrefix) {
+		if strings.HasPrefix(name, envNVRequirePrefix) && !strings.HasPrefix(name, envNVRequireJetpack) {
 			requirements = append(requirements, value)
 		}
 	}
diff --git a/internal/config/image/cuda_image_test.go b/internal/config/image/cuda_image_test.go
index 71d11cd2..a330b70c 100644
--- a/internal/config/image/cuda_image_test.go
+++ b/internal/config/image/cuda_image_test.go
@@ -69,3 +69,55 @@ func TestParseMajorMinorVersionInvalid(t *testing.T) {
 		})
 	}
 }
+
+func TestGetRequirements(t *testing.T) {
+	testCases := []struct {
+		description  string
+		env          []string
+		requirements []string
+	}{
+		{
+			description:  "NVIDIA_REQUIRE_JETPACK is ignored",
+			env:          []string{"NVIDIA_REQUIRE_JETPACK=csv-mounts=all"},
+			requirements: nil,
+		},
+		{
+			description:  "NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS is ignored",
+			env:          []string{"NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS=base-only"},
+			requirements: nil,
+		},
+		{
+			description:  "single requirement set",
+			env:          []string{"NVIDIA_REQUIRE_CUDA=cuda>=11.6"},
+			requirements: []string{"cuda>=11.6"},
+		},
+		{
+			description:  "requirements are concatenated requirement set",
+			env:          []string{"NVIDIA_REQUIRE_CUDA=cuda>=11.6", "NVIDIA_REQUIRE_BRAND=brand=tesla"},
+			requirements: []string{"cuda>=11.6", "brand=tesla"},
+		},
+		{
+			description:  "legacy image",
+			env:          []string{"CUDA_VERSION=11.6"},
+			requirements: []string{"cuda>=11.6"},
+		},
+		{
+			description:  "legacy image with additional requirement",
+			env:          []string{"CUDA_VERSION=11.6", "NVIDIA_REQUIRE_BRAND=brand=tesla"},
+			requirements: []string{"cuda>=11.6", "brand=tesla"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			image, err := NewCUDAImageFromEnv(tc.env)
+			require.NoError(t, err)
+
+			requirements, err := image.GetRequirements()
+			require.NoError(t, err)
+			require.ElementsMatch(t, tc.requirements, requirements)
+
+		})
+
+	}
+}