File tree Expand file tree Collapse file tree 1 file changed +14
-10
lines changed Expand file tree Collapse file tree 1 file changed +14
-10
lines changed Original file line number Diff line number Diff line change @@ -29,7 +29,8 @@ import (
29
29
specs "github.com/opencontainers/runtime-spec/specs-go"
30
30
)
31
31
32
- const nvidiaCLI = "nvidia-container-cli"
32
+ // NvidiaCLI is the path to the Nvidia helper binary
33
+ const NvidiaCLI = "nvidia-container-cli"
33
34
34
35
// Capability specifies capabilities for the gpu inside the container
35
36
// Detailed explanation of options can be found:
@@ -51,13 +52,16 @@ const (
51
52
Display Capability = "display"
52
53
)
53
54
54
- var allCaps = []Capability {
55
- Compute ,
56
- Compat32 ,
57
- Graphics ,
58
- Utility ,
59
- Video ,
60
- Display ,
55
+ // AllCaps returns the complete list of supported Nvidia capabilties.
56
+ func AllCaps () []Capability {
57
+ return []Capability {
58
+ Compute ,
59
+ Compat32 ,
60
+ Graphics ,
61
+ Utility ,
62
+ Video ,
63
+ Display ,
64
+ }
61
65
}
62
66
63
67
// WithGPUs adds NVIDIA gpu support to a container
@@ -76,7 +80,7 @@ func WithGPUs(opts ...Opts) oci.SpecOpts {
76
80
}
77
81
c .OCIHookPath = path
78
82
}
79
- nvidiaPath , err := exec .LookPath (nvidiaCLI )
83
+ nvidiaPath , err := exec .LookPath (NvidiaCLI )
80
84
if err != nil {
81
85
return err
82
86
}
@@ -166,7 +170,7 @@ func WithAllDevices(c *config) error {
166
170
167
171
// WithAllCapabilities adds all capabilities to the container for the gpus
168
172
func WithAllCapabilities (c * config ) error {
169
- c .Capabilities = allCaps
173
+ c .Capabilities = AllCaps ()
170
174
return nil
171
175
}
172
176
You can’t perform that action at this time.
0 commit comments