Skip to content

Add MPS control-daemon health probe - #5095

Merged
harishxr merged 1 commit into
aws:devfrom
harishxr:mps-ec2-daemon-probe
Aug 18, 2026
Merged

Add MPS control-daemon health probe#5095
harishxr merged 1 commit into
aws:devfrom
harishxr:mps-ec2-daemon-probe

Conversation

@harishxr

@harishxr harishxr commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds a functional health probe for the NVIDIA MPS control daemon to the shared lib. ProbeControlDaemon execs the MPS control utility with a control command and reports whether the daemon is serving, distinguishing three outcomes: serving (exit 0), not serving (positive exit code from the utility), and wedged (the utility hangs and is caught by a timeout).

This is the first PR in a stack that adds an MPS daemon health gate to the agent:

  1. This PR: the probe primitive.
  2. An MPSDaemonResource TaskResource that wraps the probe in a bounded retry and gates task launch.
  3. Wiring that attaches the resource to MPS tasks only.

ProbeControlDaemon is exported but has no caller in this PR by design.

Implementation details

  • ProbeControlDaemon(exec, command) runs /usr/bin/nvidia-cuda-mps-control under a 3s context timeout, feeding the control command on stdin, and returns a ProbeResult{ExitCode, Stdout, Latency, TimedOut, Err}.
  • A wedged daemon accepts the socket connection but never replies, so the utility hangs. The context timeout converts that hang into TimedOut=true, which is the failure mode systemd liveness cannot detect. Exit code, output, and latency are captured so a caller can surface them as a task stop reason.

Testing

  • The probe was tested end to end as part of the health-gate integration testing: a healthy daemon passed, a stopped daemon was refused via the exit-1 path, and a SIGSTOP-wedged daemon was refused via the timeout path (probe returned TimedOut=true at ~3s per attempt while systemd still reported the unit active).
    New tests cover the changes: yes

Description for the changelog

Enhancement - Add MPS control-daemon health probe

Additional Information

Does this PR include breaking model changes? If so, Have you added transformation functions? No

Does this PR include the addition of new environment variables in the README? No

Licensing

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@harishxr
harishxr force-pushed the mps-ec2-daemon-probe branch from 3911962 to 7b40d6c Compare August 17, 2026 22:12
@harishxr
harishxr force-pushed the mps-ec2-daemon-probe branch 2 times, most recently from 2f02cfd to b873769 Compare August 17, 2026 22:44
@harishxr harishxr changed the title Add MPS control-daemon health probe to ecs-agent/utils/mps Add MPS control-daemon health probe Aug 17, 2026
@harishxr
harishxr marked this pull request as ready for review August 17, 2026 22:52
@harishxr
harishxr requested a review from a team as a code owner August 17, 2026 22:52
amogh09
amogh09 previously approved these changes Aug 17, 2026
@harishxr
harishxr enabled auto-merge (rebase) August 18, 2026 00:22
@harishxr
harishxr merged commit 21e7d47 into aws:dev Aug 18, 2026
45 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants