Merge pull request #2510 from crosbymichael/gpucaps2

Change gpu Capability type to string
This commit is contained in:
Derek McGowan 2018-07-30 14:35:59 -07:00 committed by GitHub
commit 45d0df8fe3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

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