Added inf support config and neuron runtime override.

Added an agent config InferentiaSupportEnabled populated by ECS_ENABLE_INF_SUPPORT env. For a container that has AWS_NEURON_VISIBLE_DEVICES specified, if InferentiaSupportEnabled is on, the agent will override its runtime to the neuron docker runtime needed for using the inferentia devices.

This enables us to only use the neuron runtime for container that needs the inf device, and only do so when such runtime is installed on the AMI (which is indicated by the ECS_ENABLE_INF_SUPPORT config that we will add together with installing the neuron runtime).

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{