Override container runtime when inferentia support is enabled

agent/api/container/container.go

@@ -77,6 +77,9 @@ const (
 
 	// TargetLogDriver is to show secret target being "LOG_DRIVER", the default will be "CONTAINER"
 	SecretTargetLogDriver = "LOG_DRIVER"
+
+	// neuronVisibleDevicesEnvVar is the env which indicates that the container wants to use inferentia devices.
+	neuronVisibleDevicesEnvVar = "AWS_NEURON_VISIBLE_DEVICES"
 )
 
 // DockerConfig represents additional metadata about a container to run. It's
@@ -1120,3 +1123,11 @@ func (c *Container) GetEnvironmentFiles() []EnvironmentFile {
 
 	return c.EnvironmentFiles
 }
+
+func (c *Container) RequireNeuronRuntime() bool {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+
+	_, ok := c.Environment[neuronVisibleDevicesEnvVar]
+	return ok
+}

agent/api/container/container_test.go

@@ -750,3 +750,10 @@ func TestMergeEnvironmentVariablesFromEnvfiles(t *testing.T) {
 		})
 	}
 }
+
+func TestRequireNeuronRuntime(t *testing.T) {
+	c := &Container{
+		Environment: map[string]string{neuronVisibleDevicesEnvVar: "all"},
+	}
+	assert.True(t, c.RequireNeuronRuntime())
+}

agent/api/task/task.go

@@ -77,6 +77,9 @@ const (
 	NvidiaVisibleDevicesEnvVar = "NVIDIA_VISIBLE_DEVICES"
 	GPUAssociationType         = "gpu"
 
+	// neuronRuntime is the name of the neuron docker runtime.
+	neuronRuntime = "neuron"
+
 	ContainerOrderingCreateCondition = "CREATE"
 	ContainerOrderingStartCondition  = "START"
 
@@ -1430,8 +1433,8 @@ func (task *Task) dockerExposedPorts(container *apicontainer.Container) nat.Port
 }
 
 // DockerHostConfig construct the configuration recognized by docker
-func (task *Task) DockerHostConfig(container *apicontainer.Container, dockerContainerMap map[string]*apicontainer.DockerContainer, apiVersion dockerclient.DockerVersion) (*dockercontainer.HostConfig, *apierrors.HostConfigError) {
-	return task.dockerHostConfig(container, dockerContainerMap, apiVersion)
+func (task *Task) DockerHostConfig(container *apicontainer.Container, dockerContainerMap map[string]*apicontainer.DockerContainer, apiVersion dockerclient.DockerVersion, cfg *config.Config) (*dockercontainer.HostConfig, *apierrors.HostConfigError) {
+	return task.dockerHostConfig(container, dockerContainerMap, apiVersion, cfg)
 }
 
 // ApplyExecutionRoleLogsAuth will check whether the task has execution role
@@ -1459,7 +1462,7 @@ func (task *Task) ApplyExecutionRoleLogsAuth(hostConfig *dockercontainer.HostCon
 	return nil
 }
 
-func (task *Task) dockerHostConfig(container *apicontainer.Container, dockerContainerMap map[string]*apicontainer.DockerContainer, apiVersion dockerclient.DockerVersion) (*dockercontainer.HostConfig, *apierrors.HostConfigError) {
+func (task *Task) dockerHostConfig(container *apicontainer.Container, dockerContainerMap map[string]*apicontainer.DockerContainer, apiVersion dockerclient.DockerVersion, cfg *config.Config) (*dockercontainer.HostConfig, *apierrors.HostConfigError) {
 	dockerLinkArr, err := task.dockerLinks(container, dockerContainerMap)
 	if err != nil {
 		return nil, &apierrors.HostConfigError{Msg: err.Error()}
@@ -1488,12 +1491,8 @@ func (task *Task) dockerHostConfig(container *apicontainer.Container, dockerCont
 		Resources:    resources,
 	}
 
-	if task.isGPUEnabled() && task.shouldRequireNvidiaRuntime(container) {
-		if task.NvidiaRuntime == "" {
-			return nil, &apierrors.HostConfigError{Msg: "Runtime is not set for GPU containers"}
-		}
-		seelog.Debugf("Setting runtime as %s for container %s", task.NvidiaRuntime, container.Name)
-		hostConfig.Runtime = task.NvidiaRuntime
+	if err := task.overrideContainerRuntime(container, hostConfig, cfg); err != nil {
+		return nil, err
 	}
 
 	if container.DockerConfig.HostConfig != nil {
@@ -1537,6 +1536,24 @@ func (task *Task) dockerHostConfig(container *apicontainer.Container, dockerCont
 	return hostConfig, nil
 }
 
+// overrideContainerRuntime overrides the runtime for the container in host config if needed.
+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"}
+		}
+		seelog.Debugf("Setting runtime as %s for container %s", task.NvidiaRuntime, container.Name)
+		hostCfg.Runtime = task.NvidiaRuntime
+	}
+
+	if cfg.InferentiaSupportEnabled && container.RequireNeuronRuntime() {
+		seelog.Debugf("Setting runtime as %s for container %s", neuronRuntime, container.Name)
+		hostCfg.Runtime = neuronRuntime
+	}
+	return nil
+}
+
 // Requires an *apicontainer.Container and returns the Resources for the HostConfig struct
 func (task *Task) getDockerResources(container *apicontainer.Container) dockercontainer.Resources {
 	// Convert MB to B and set Memory

agent/api/task/task_linux_test.go

@@ -422,7 +422,8 @@ func TestPlatformHostConfigOverrideErrorPath(t *testing.T) {
 		},
 	}
 
-	dockerHostConfig, err := task.DockerHostConfig(task.Containers[0], dockerMap(task), defaultDockerClientAPIVersion)
+	dockerHostConfig, err := task.DockerHostConfig(task.Containers[0], dockerMap(task), defaultDockerClientAPIVersion,
+		&config.Config{})
 	assert.Error(t, err)
 	assert.Empty(t, dockerHostConfig)
 }
@@ -464,7 +465,8 @@ func TestDockerHostConfigRawConfigMerging(t *testing.T) {
 		},
 	}
 
-	hostConfig, configErr := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), minDockerClientAPIVersion)
+	hostConfig, configErr := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask),
+		minDockerClientAPIVersion, &config.Config{})
 	assert.Nil(t, configErr)
 
 	expected := dockercontainer.HostConfig{