Cleanup names of constants to better match ENVVARs

Signed-off-by: Kevin Klues <kklues@nvidia.com>
This commit is contained in:
Kevin Klues 2019-12-20 16:02:51 +00:00
parent 8da6d14a20
commit b3de846f66

View File

@ -13,14 +13,17 @@ 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 (
allCapabilities = "compute,compat32,graphics,utility,video,display"
defaultCapability = "utility"
) )
type nvidiaConfig struct { type nvidiaConfig 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...)
@ -177,7 +180,7 @@ func getNvidiaConfigLegacy(env map[string]string) *nvidiaConfig {
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)
@ -193,7 +196,7 @@ func getNvidiaConfigLegacy(env map[string]string) *nvidiaConfig {
} }
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.