fix nvidia runtime assignment and check

aws · Sep 28, 2021 · 0d16d11 · 0d16d11
1 parent 2652f69
commit 0d16d11
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/agent/api/task/task.go b/agent/api/task/task.go
@@ -462,8 +462,8 @@ func (task *Task) addGPUResource(cfg *config.Config) error {
 		// For internal instances, GPU IDs are handled by env var
 		if !cfg.External.Enabled() {
 			task.populateGPUEnvironmentVariables()
+			task.NvidiaRuntime = cfg.NvidiaRuntime
 		}
-		task.NvidiaRuntime = cfg.NvidiaRuntime
 	}
 	return nil
 }
@@ -1535,10 +1535,10 @@ func (task *Task) dockerHostConfig(container *apicontainer.Container, dockerCont
 func (task *Task) overrideContainerRuntime(container *apicontainer.Container, hostCfg *dockercontainer.HostConfig,
 	cfg *config.Config) *apierrors.HostConfigError {
 	if task.isGPUEnabled() && task.shouldRequireNvidiaRuntime(container) {
-		if task.NvidiaRuntime == "" {
-			return &apierrors.HostConfigError{Msg: "Runtime is not set for GPU containers"}
-		}
 		if !cfg.External.Enabled() {
+			if task.NvidiaRuntime == "" {
+				return &apierrors.HostConfigError{Msg: "Runtime is not set for GPU containers"}
+			}
 			seelog.Debugf("Setting runtime as %s for container %s", task.NvidiaRuntime, container.Name)
 			hostCfg.Runtime = task.NvidiaRuntime
 		}