From db47b5827571511727ce3bca05804b078f7744c1 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Thu, 29 Sep 2022 14:39:56 +0200
Subject: [PATCH] Add utilities for driver capabilities to image packages

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 internal/config/image/capabilities.go | 54 +++++++++++++++++++++++++++
 internal/config/image/cuda_image.go   | 23 +++++++++---
 2 files changed, 72 insertions(+), 5 deletions(-)
 create mode 100644 internal/config/image/capabilities.go

diff --git a/internal/config/image/capabilities.go b/internal/config/image/capabilities.go
new file mode 100644
index 00000000..2ff1f40b
--- /dev/null
+++ b/internal/config/image/capabilities.go
@@ -0,0 +1,54 @@
+/**
+# Copyright (c) 2022, 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 image
+
+// DriverCapability represents the possible values of NVIDIA_DRIVER_CAPABILITIES
+type DriverCapability string
+
+// Constants for the supported driver capabilities
+const (
+	DriverCapabilityAll      DriverCapability = "all"
+	DriverCapabilityCompat32 DriverCapability = "compat32"
+	DriverCapabilityCompute  DriverCapability = "compute"
+	DriverCapabilityDisplay  DriverCapability = "display"
+	DriverCapabilityGraphics DriverCapability = "graphics"
+	DriverCapabilityNgx      DriverCapability = "ngx"
+	DriverCapabilityUtility  DriverCapability = "utility"
+	DriverCapabilityVideo    DriverCapability = "video"
+)
+
+// DriverCapabilities represents the NVIDIA_DRIVER_CAPABILITIES set for the specified image.
+type DriverCapabilities map[DriverCapability]bool
+
+// Has check whether the specified capability is selected.
+func (c DriverCapabilities) Has(capability DriverCapability) bool {
+	if c[DriverCapabilityAll] {
+		return true
+	}
+	return c[capability]
+}
+
+// Any checks whether any of the specified capabilites are set
+func (c DriverCapabilities) Any(capabilities ...DriverCapability) bool {
+	for _, cap := range capabilities {
+		if c.Has(cap) {
+			return true
+		}
+	}
+
+	return false
+}
diff --git a/internal/config/image/cuda_image.go b/internal/config/image/cuda_image.go
index b9ad4c94..097f5d9e 100644
--- a/internal/config/image/cuda_image.go
+++ b/internal/config/image/cuda_image.go
@@ -26,11 +26,12 @@ import (
 )
 
 const (
-	envCUDAVersion      = "CUDA_VERSION"
-	envNVRequirePrefix  = "NVIDIA_REQUIRE_"
-	envNVRequireCUDA    = envNVRequirePrefix + "CUDA"
-	envNVRequireJetpack = envNVRequirePrefix + "JETPACK"
-	envNVDisableRequire = "NVIDIA_DISABLE_REQUIRE"
+	envCUDAVersion          = "CUDA_VERSION"
+	envNVRequirePrefix      = "NVIDIA_REQUIRE_"
+	envNVRequireCUDA        = envNVRequirePrefix + "CUDA"
+	envNVRequireJetpack     = envNVRequirePrefix + "JETPACK"
+	envNVDisableRequire     = "NVIDIA_DISABLE_REQUIRE"
+	envNVDriverCapabilities = "NVIDIA_DRIVER_CAPABILITIES"
 )
 
 // CUDA represents a CUDA image that can be used for GPU computing. This wraps
@@ -142,6 +143,18 @@ func (i CUDA) DevicesFromEnvvars(envVars ...string) []string {
 	return strings.Split(*devices, ",")
 }
 
+// GetDriverCapabilities returns the requested driver capabilities.
+func (i CUDA) GetDriverCapabilities() DriverCapabilities {
+	env := i[envNVDriverCapabilities]
+
+	capabilites := make(DriverCapabilities)
+	for _, c := range strings.Split(env, ",") {
+		capabilites[DriverCapability(c)] = true
+	}
+
+	return capabilites
+}
+
 func (i CUDA) legacyVersion() (string, error) {
 	majorMinor, err := parseMajorMinorVersion(i[envCUDAVersion])
 	if err != nil {