Add ACCELERATED_COMPUTE health check for the MPS control daemon

agent/app/agent.go

@@ -725,6 +725,9 @@ func (agent *ecsAgent) newDoctorWithHealthchecks(cluster, containerInstanceARN s
 		runtimeHealthCheck,
 	}
 
+	// register the MPS control daemon health check on MPS-capable instances
+	healthcheckList = agent.appendMpsDaemonHealthcheck(healthcheckList)
+
 	// set up the doctor and return it
 	return doctor.NewDoctor(healthcheckList, cluster, containerInstanceARN)
 }

agent/app/agent_capability_unix.go

@@ -126,22 +126,33 @@ func (agent *ecsAgent) appendNvidiaDriverVersionAttribute(capabilities []types.A
 	return capabilities
 }
 
-// appendGpuSharingMpsCapability advertises ecs.capability.gpu-sharing-mps when the
-// instance can run MPS: a GPU is present, the MPS control binary and its systemd unit
-// are installed and enabled, the GPU is not a vGPU slice, and every discovered GPU has
-// a usable-memory value. The facts are gathered by ecs-init and read from the NvidiaGPUManager;
-// the decision itself lives in the shared gpu package so the MI agent reaches the same verdict from the same inputs.
-func (agent *ecsAgent) appendGpuSharingMpsCapability(capabilities []types.Attribute) []types.Attribute {
+// mpsCapabilityInputs gathers the host facts the MPS gate decision is made from,
+// reading them off the NvidiaGPUManager. ok is false when there is no GPU manager, in
+// which case MPS cannot be evaluated on this instance. It is the single source of truth
+// shared by the gpu-sharing-mps capability and the MPS daemon health check.
+func (agent *ecsAgent) mpsCapabilityInputs() (gpu.MpsCapabilityInputs, bool) {
 	if agent.resourceFields == nil || agent.resourceFields.NvidiaGPUManager == nil {
-		return capabilities
+		return gpu.MpsCapabilityInputs{}, false
 	}
 	mgr := agent.resourceFields.NvidiaGPUManager
-	inputs := gpu.MpsCapabilityInputs{
+	return gpu.MpsCapabilityInputs{
 		GPUPresent:        len(mgr.GetDevices()) > 0,
 		MpsBinaryPresent:  mgr.GetMpsControlBinaryPresent(),
 		MpsServiceEnabled: mgr.GetMpsServiceEnabled(),
 		IsVGPU:            mgr.GetHasVGPU(),
 		AllGPUsHaveMemory: gpu.AllGPUMemoryReported(mgr.GetGPUIDsUnsafe(), mgr.GetGPUMemoryMiBUnsafe()),
+	}, true
+}
+
+// appendGpuSharingMpsCapability advertises ecs.capability.gpu-sharing-mps when the
+// instance can run MPS: a GPU is present, the MPS control binary and its systemd unit
+// are installed and enabled, the GPU is not a vGPU slice, and every discovered GPU has
+// a usable-memory value. The facts are gathered by ecs-init and read from the NvidiaGPUManager;
+// the decision itself lives in the shared gpu package so the MI agent reaches the same verdict from the same inputs.
+func (agent *ecsAgent) appendGpuSharingMpsCapability(capabilities []types.Attribute) []types.Attribute {
+	inputs, ok := agent.mpsCapabilityInputs()
+	if !ok {
+		return capabilities
 	}
 	advertise, conditions := gpu.ShouldAdvertiseMpsCapability(inputs)
 	if !advertise {

agent/app/agent_mps_healthcheck_unix.go

@@ -0,0 +1,39 @@
+//go:build linux
+// +build linux
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package app
+
+import (
+	agentdoctor "github.com/aws/amazon-ecs-agent/agent/doctor"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/utils/gpu"
+)
+
+// appendMpsDaemonHealthcheck registers the MPS control-daemon health check on exactly
+// the instances that advertise ecs.capability.gpu-sharing-mps, reusing the same
+// predicate so the instance health signal and the capability agree on what
+// MPS-capable means. A plain GPU box with no nvidia-mps.service would otherwise fail
+// every probe and report a false ACCELERATED_COMPUTE=IMPAIRED.
+func (agent *ecsAgent) appendMpsDaemonHealthcheck(list []doctor.Healthcheck) []doctor.Healthcheck {
+	inputs, ok := agent.mpsCapabilityInputs()
+	if !ok {
+		return list
+	}
+	if advertise, _ := gpu.ShouldAdvertiseMpsCapability(inputs); !advertise {
+		return list
+	}
+	return append(list, agentdoctor.NewMpsDaemonHealthcheck())
+}

agent/app/agent_mps_healthcheck_unix_test.go

@@ -0,0 +1,83 @@
+//go:build linux && unit
+// +build linux,unit
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package app
+
+import (
+	"testing"
+
+	"github.com/aws/amazon-ecs-agent/agent/gpu"
+	"github.com/aws/amazon-ecs-agent/agent/taskresource"
+	"github.com/stretchr/testify/assert"
+)
+
+// TestAppendMpsDaemonHealthcheck asserts that the MPS control-daemon health check is
+// registered on exactly the instances that would advertise gpu-sharing-mps: it must be
+// appended when ShouldAdvertiseMpsCapability is satisfied and withheld otherwise, so a
+// plain GPU box never reports a false ACCELERATED_COMPUTE=IMPAIRED. It reuses the same
+// NvidiaGPUManager setup as the capability test to keep the two decisions in lockstep.
+func TestAppendMpsDaemonHealthcheck(t *testing.T) {
+	// newMpsManager returns a GPU manager carrying the given MPS facts and one device
+	// (with usable memory) unless gpuPresent is false. haveMemory controls whether that
+	// device reports a memory value, exercising the AllGPUsHaveMemory gate.
+	newMpsManager := func(gpuPresent, binary, service, vgpu, haveMemory bool) *gpu.NvidiaGPUManager {
+		m := &gpu.NvidiaGPUManager{
+			MpsControlBinaryPresent: binary,
+			MpsServiceEnabled:       service,
+			HasVGPU:                 vgpu,
+		}
+		if gpuPresent {
+			m.SetGPUIDs([]string{"gpu-0"})
+			if haveMemory {
+				m.SetGPUMemoryMiB(map[string]uint64{"gpu-0": 22563})
+			}
+			m.SetDevices()
+		}
+		return m
+	}
+
+	cases := []struct {
+		name     string
+		mgr      *gpu.NvidiaGPUManager // nil means no GPU manager on the instance
+		register bool
+	}{
+		{"all conditions met", newMpsManager(true, true, true, false, true), true},
+		{"no gpu manager", nil, false},
+		{"no gpu present", newMpsManager(false, true, true, false, true), false},
+		{"mps binary absent", newMpsManager(true, false, true, false, true), false},
+		{"mps service disabled", newMpsManager(true, true, false, false, true), false},
+		{"is vgpu", newMpsManager(true, true, true, true, true), false},
+		{"gpu present but no memory reported", newMpsManager(true, true, true, false, false), false},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			// NvidiaGPUManager is an interface, so a typed-nil pointer would read as a
+			// non-nil interface; leave the field unset to model "no GPU manager".
+			rf := &taskresource.ResourceFields{}
+			if tc.mgr != nil {
+				rf.NvidiaGPUManager = tc.mgr
+			}
+			agent := &ecsAgent{resourceFields: rf}
+			got := agent.appendMpsDaemonHealthcheck(nil)
+			if tc.register {
+				assert.Len(t, got, 1, "the MPS health check must be registered")
+			} else {
+				assert.Empty(t, got, "the MPS health check must not be registered")
+			}
+		})
+	}
+}

agent/app/agent_mps_healthcheck_unsupported.go

@@ -0,0 +1,27 @@
+//go:build !linux
+// +build !linux
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package app
+
+import (
+	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
+)
+
+// appendMpsDaemonHealthcheck is a no-op off Linux, where the MPS control daemon and its
+// pipe directory do not exist.
+func (agent *ecsAgent) appendMpsDaemonHealthcheck(list []doctor.Healthcheck) []doctor.Healthcheck {
+	return list
+}

agent/doctor/mps_daemon_healthcheck.go

@@ -0,0 +1,122 @@
+//go:build linux
+// +build linux
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package doctor
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/aws/amazon-ecs-agent/agent/doctor/statustracker"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger/field"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/utils/execwrapper"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/utils/mps"
+)
+
+// mpsDaemonImpairedThreshold is the number of consecutive probe failures before the
+// instance is reported ACCELERATED_COMPUTE=IMPAIRED
+const mpsDaemonImpairedThreshold = 3
+
+type mpsDaemonHealthcheck struct {
+	*statustracker.HealthCheckStatusTracker
+
+	exec        execwrapper.Exec
+	statPipeDir func(string) (os.FileInfo, error)
+	probeCmd    string
+
+	consecutiveFailures int
+	threshold           int
+}
+
+// NewMpsDaemonHealthcheck is the constructor for the MPS control-daemon health check.
+func NewMpsDaemonHealthcheck() doctor.Healthcheck {
+	return newMpsDaemonHealthcheck(execwrapper.NewExec(), os.Stat)
+}
+
+func newMpsDaemonHealthcheck(exec execwrapper.Exec,
+	statPipeDir func(string) (os.FileInfo, error)) *mpsDaemonHealthcheck {
+	return &mpsDaemonHealthcheck{
+		HealthCheckStatusTracker: statustracker.NewHealthCheckStatusTracker(),
+		exec:                     exec,
+		statPipeDir:              statPipeDir,
+		probeCmd:                 mps.ProbeCommand,
+		threshold:                mpsDaemonImpairedThreshold,
+	}
+}
+
+// RunCheck runs one probe per tick and folds the result into the consecutive-failure
+// counter. It reports Impaired only after threshold failures in a row; any success
+// resets the counter, returning the status to Ok.
+func (m *mpsDaemonHealthcheck) RunCheck() ecstcs.InstanceHealthCheckStatus {
+	res := m.probe()
+	serving := res.Err == nil
+
+	if serving {
+		m.consecutiveFailures = 0
+	} else {
+		m.consecutiveFailures++
+	}
+
+	status := ecstcs.InstanceHealthCheckStatusOk
+	if m.consecutiveFailures >= m.threshold {
+		status = ecstcs.InstanceHealthCheckStatusImpaired
+	}
+
+	switch {
+	case status == ecstcs.InstanceHealthCheckStatusImpaired:
+		logger.Error("MPS control daemon health check impaired", logger.Fields{
+			"consecutiveFailures": m.consecutiveFailures,
+			"timedOut":            res.TimedOut,
+			field.Error:           res.Err,
+		})
+	case !serving:
+		logger.Warn("MPS control daemon probe failed, below impairment threshold", logger.Fields{
+			"consecutiveFailures": m.consecutiveFailures,
+			"timedOut":            res.TimedOut,
+			field.Error:           res.Err,
+		})
+	default:
+		logger.Debug("MPS control daemon is serving")
+	}
+
+	m.SetHealthcheckStatus(status)
+	return m.GetHealthcheckStatus()
+}
+
+// probe runs the pipe-directory pre-check and, if it passes, a single control-daemon
+// probe. The daemon is serving when the returned result has a nil Err. An unusable
+// pipe directory counts as this tick's failure and skips the exec; it can recover on a
+// later tick.
+func (m *mpsDaemonHealthcheck) probe() mps.ProbeResult {
+	fi, err := m.statPipeDir(mps.PipeDirectory)
+	if err != nil {
+		return mps.ProbeResult{Err: err}
+	}
+	if !fi.IsDir() {
+		return mps.ProbeResult{
+			Err: fmt.Errorf("mps pipe directory %s is not a directory", mps.PipeDirectory),
+		}
+	}
+	return mps.ProbeControlDaemon(m.exec, m.probeCmd)
+}
+
+// GetHealthcheckType returns the type of this health check.
+func (m *mpsDaemonHealthcheck) GetHealthcheckType() string {
+	return ecstcs.InstanceHealthCheckTypeAcceleratedCompute
+}