From 7ca2c3d68d8ad842477fb19325a0aac4f8f2d6bd Mon Sep 17 00:00:00 2001
From: Tibor Vass <tibor@docker.com>
Date: Tue, 12 Mar 2019 14:48:59 -0700
Subject: [PATCH] contrib/nvidia: export helper binary path and list of Nvidia
 capabilities

Signed-off-by: Tibor Vass <tibor@docker.com>
---
 contrib/nvidia/nvidia.go | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/contrib/nvidia/nvidia.go b/contrib/nvidia/nvidia.go
index 5a1d3c7f1..a48594b00 100644
--- a/contrib/nvidia/nvidia.go
+++ b/contrib/nvidia/nvidia.go
@@ -29,7 +29,8 @@ import (
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 )
 
-const nvidiaCLI = "nvidia-container-cli"
+// NvidiaCLI is the path to the Nvidia helper binary
+const NvidiaCLI = "nvidia-container-cli"
 
 // Capability specifies capabilities for the gpu inside the container
 // Detailed explanation of options can be found:
@@ -51,13 +52,16 @@ const (
 	Display Capability = "display"
 )
 
-var allCaps = []Capability{
-	Compute,
-	Compat32,
-	Graphics,
-	Utility,
-	Video,
-	Display,
+// AllCaps returns the complete list of supported Nvidia capabilties.
+func AllCaps() []Capability {
+	return []Capability{
+		Compute,
+		Compat32,
+		Graphics,
+		Utility,
+		Video,
+		Display,
+	}
 }
 
 // WithGPUs adds NVIDIA gpu support to a container
@@ -76,7 +80,7 @@ func WithGPUs(opts ...Opts) oci.SpecOpts {
 			}
 			c.OCIHookPath = path
 		}
-		nvidiaPath, err := exec.LookPath(nvidiaCLI)
+		nvidiaPath, err := exec.LookPath(NvidiaCLI)
 		if err != nil {
 			return err
 		}
@@ -166,7 +170,7 @@ func WithAllDevices(c *config) error {
 
 // WithAllCapabilities adds all capabilities to the container for the gpus
 func WithAllCapabilities(c *config) error {
-	c.Capabilities = allCaps
+	c.Capabilities = AllCaps()
 	return nil
 }