UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-03-31 15:47:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge branch 'cleanup-driver-caps' into 'master'

Browse Source 
Cleanup naming of constants and functions

See merge request nvidia/container-toolkit/toolkit!2
This commit is contained in:
Jon Mayo 2019-12-20 17:24:26 +00:00
commit 2155c2d587
2 changed files with 38 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
nvidia-container-toolkit/container_config.go
View File

@ -13,21 +13,24 @@ import (
var envSwarmGPU *string var envSwarmGPU *string
const ( const (
envCUDAVersion = "CUDA_VERSION"
envNVRequirePrefix = "NVIDIA_REQUIRE_" envNVRequirePrefix = "NVIDIA_REQUIRE_"
envLegacyCUDAVersion = "CUDA_VERSION"
envNVRequireCUDA = envNVRequirePrefix + "CUDA" envNVRequireCUDA = envNVRequirePrefix + "CUDA"
envNVGPU = "NVIDIA_VISIBLE_DEVICES"
envNVDriverCapabilities = "NVIDIA_DRIVER_CAPABILITIES"
defaultCapability = "utility"
allCapabilities = "compute,compat32,graphics,utility,video,display"
envNVDisableRequire = "NVIDIA_DISABLE_REQUIRE" envNVDisableRequire = "NVIDIA_DISABLE_REQUIRE"
envNVVisibleDevices = "NVIDIA_VISIBLE_DEVICES"
envNVDriverCapabilities = "NVIDIA_DRIVER_CAPABILITIES"
)
const (
allDriverCapabilities = "compute,compat32,graphics,utility,video,display"
defaultDriverCapabilities = "utility"
) )
type nvidiaConfig struct { type nvidiaConfig struct {
Devices string Devices string
Capabilities string DriverCapabilities string
Requirements []string Requirements []string
DisableRequire bool DisableRequire bool
} }
type containerConfig struct { type containerConfig struct {
@ -111,7 +114,7 @@ func loadSpec(path string) (spec *Spec) {
} }
func getDevices(env map[string]string) *string { func getDevices(env map[string]string) *string {
gpuVars := []string{envNVGPU} gpuVars := []string{envNVVisibleDevices}
if envSwarmGPU != nil { if envSwarmGPU != nil {
// The Swarm resource has higher precedence. // The Swarm resource has higher precedence.
gpuVars = append([]string{*envSwarmGPU}, gpuVars...) gpuVars = append([]string{*envSwarmGPU}, gpuVars...)
@ -125,7 +128,7 @@ func getDevices(env map[string]string) *string {
return nil return nil
} }
func getCapabilities(env map[string]string) *string { func getDriverCapabilities(env map[string]string) *string {
if capabilities, ok := env[envNVDriverCapabilities]; ok { if capabilities, ok := env[envNVDriverCapabilities]; ok {
return &capabilities return &capabilities
} }
@ -160,24 +163,24 @@ func getNvidiaConfigLegacy(env map[string]string) *nvidiaConfig {
devices = "" devices = ""
} }
var capabilities string var driverCapabilities string
if c := getCapabilities(env); c == nil { if c := getDriverCapabilities(env); c == nil {
// Environment variable unset: default to "all". // Environment variable unset: default to "all".
capabilities = allCapabilities driverCapabilities = allDriverCapabilities
} else if len(*c) == 0 { } else if len(*c) == 0 {
// Environment variable empty: use default capability. // Environment variable empty: use default capability.
capabilities = defaultCapability driverCapabilities = defaultDriverCapabilities
} else { } else {
// Environment variable non-empty. // Environment variable non-empty.
capabilities = *c driverCapabilities = *c
} }
if capabilities == "all" { if driverCapabilities == "all" {
capabilities = allCapabilities driverCapabilities = allDriverCapabilities
} }
requirements := getRequirements(env) requirements := getRequirements(env)
vmaj, vmin, _ := parseCudaVersion(env[envLegacyCUDAVersion]) vmaj, vmin, _ := parseCudaVersion(env[envCUDAVersion])
cudaRequire := fmt.Sprintf("cuda>=%d.%d", vmaj, vmin) cudaRequire := fmt.Sprintf("cuda>=%d.%d", vmaj, vmin)
requirements = append(requirements, cudaRequire) requirements = append(requirements, cudaRequire)
@ -185,15 +188,15 @@ func getNvidiaConfigLegacy(env map[string]string) *nvidiaConfig {
disableRequire, _ := strconv.ParseBool(env[envNVDisableRequire]) disableRequire, _ := strconv.ParseBool(env[envNVDisableRequire])
return &nvidiaConfig{ return &nvidiaConfig{
Devices: devices, Devices: devices,
Capabilities: capabilities, DriverCapabilities: driverCapabilities,
Requirements: requirements, Requirements: requirements,
DisableRequire: disableRequire, DisableRequire: disableRequire,
} }
} }
func getNvidiaConfig(env map[string]string) *nvidiaConfig { func getNvidiaConfig(env map[string]string) *nvidiaConfig {
legacyCudaVersion := env[envLegacyCUDAVersion] legacyCudaVersion := env[envCUDAVersion]
cudaRequire := env[envNVRequireCUDA] cudaRequire := env[envNVRequireCUDA]
if len(legacyCudaVersion) > 0 && len(cudaRequire) == 0 { if len(legacyCudaVersion) > 0 && len(cudaRequire) == 0 {
// Legacy CUDA image detected. // Legacy CUDA image detected.
@ -212,16 +215,16 @@ func getNvidiaConfig(env map[string]string) *nvidiaConfig {
devices = "" devices = ""
} }
var capabilities string var driverCapabilities string
if c := getCapabilities(env); c == nil || len(*c) == 0 { if c := getDriverCapabilities(env); c == nil || len(*c) == 0 {
// Environment variable unset or set but empty: use default capability. // Environment variable unset or set but empty: use default capability.
capabilities = defaultCapability driverCapabilities = defaultDriverCapabilities
} else { } else {
// Environment variable set and non-empty. // Environment variable set and non-empty.
capabilities = *c driverCapabilities = *c
} }
if capabilities == "all" { if driverCapabilities == "all" {
capabilities = allCapabilities driverCapabilities = allDriverCapabilities
} }
requirements := getRequirements(env) requirements := getRequirements(env)
@ -230,10 +233,10 @@ func getNvidiaConfig(env map[string]string) *nvidiaConfig {
disableRequire, _ := strconv.ParseBool(env[envNVDisableRequire]) disableRequire, _ := strconv.ParseBool(env[envNVDisableRequire])
return &nvidiaConfig{ return &nvidiaConfig{
Devices: devices, Devices: devices,
Capabilities: capabilities, DriverCapabilities: driverCapabilities,
Requirements: requirements, Requirements: requirements,
DisableRequire: disableRequire, DisableRequire: disableRequire,
} }
} }

2
nvidia-container-toolkit/main.go
View File

@ -124,7 +124,7 @@ func doPrestart() {
args = append(args, fmt.Sprintf("--device=%s", nvidia.Devices)) args = append(args, fmt.Sprintf("--device=%s", nvidia.Devices))
} }
for _, cap := range strings.Split(nvidia.Capabilities, ",") { for _, cap := range strings.Split(nvidia.DriverCapabilities, ",") {
if len(cap) == 0 { if len(cap) == 0 {
break break
} }