From 81e2859e8b5d0a8e272cd1862480413507613beb Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Mon, 30 Jul 2018 15:37:13 -0400
Subject: [PATCH] Change gpu Capability type to string

This helps with mappings so that we are not translating multiple times
from caller to hook.

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 contrib/nvidia/nvidia.go | 38 ++++++++++++++++++--------------------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/contrib/nvidia/nvidia.go b/contrib/nvidia/nvidia.go
index 72c6f63d4..05f2e7e5b 100644
--- a/contrib/nvidia/nvidia.go
+++ b/contrib/nvidia/nvidia.go
@@ -34,23 +34,32 @@ const nvidiaCLI = "nvidia-container-cli"
 // Capability specifies capabilities for the gpu inside the container
 // Detailed explaination of options can be found:
 // https://github.com/nvidia/nvidia-container-runtime#supported-driver-capabilities
-type Capability int
+type Capability string
 
 const (
 	// Compute capability
-	Compute Capability = iota + 1
+	Compute Capability = "compute"
 	// Compat32 capability
-	Compat32
+	Compat32 Capability = "compat32"
 	// Graphics capability
-	Graphics
+	Graphics Capability = "graphics"
 	// Utility capability
-	Utility
+	Utility Capability = "utility"
 	// Video capability
-	Video
+	Video Capability = "video"
 	// Display capability
-	Display
+	Display Capability = "display"
 )
 
+var allCaps = []Capability{
+	Compute,
+	Compat32,
+	Graphics,
+	Utility,
+	Video,
+	Display,
+}
+
 // WithGPUs adds NVIDIA gpu support to a container
 func WithGPUs(opts ...Opts) oci.SpecOpts {
 	return func(_ context.Context, _ oci.Client, _ *containers.Container, s *specs.Spec) error {
@@ -110,7 +119,7 @@ func (c *config) args() []string {
 		args = append(args, fmt.Sprintf("--device=%s", strings.Join(c.Devices, ",")))
 	}
 	for _, c := range c.Capabilities {
-		args = append(args, fmt.Sprintf("--%s", capFlags[c]))
+		args = append(args, fmt.Sprintf("--%s", c))
 	}
 	if c.LDConfig != "" {
 		args = append(args, fmt.Sprintf("--ldconfig=%s", c.LDConfig))
@@ -122,15 +131,6 @@ func (c *config) args() []string {
 	return args
 }
 
-var capFlags = map[Capability]string{
-	Compute:  "compute",
-	Compat32: "compat32",
-	Graphics: "graphics",
-	Utility:  "utility",
-	Video:    "video",
-	Display:  "display",
-}
-
 // Opts are options for configuring gpu support
 type Opts func(*config) error
 
@@ -160,9 +160,7 @@ func WithAllDevices(c *config) error {
 
 // WithAllCapabilities adds all capabilities to the container for the gpus
 func WithAllCapabilities(c *config) error {
-	for k := range capFlags {
-		c.Capabilities = append(c.Capabilities, k)
-	}
+	c.Capabilities = allCaps
 	return nil
 }