fix: make MCP Slurm wait poll remote job state#1900
Conversation
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds Slurm submit metadata persistence and SSH-based status resolution in ChangesSlurm Sidecar Status Tracking
Managed-source Launcher Argv Change
Hostname Smoke Test Script
Estimated code review effort: 4 (Complex) | ~40 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tools/mcp/modelopt_mcp/bridge.py (1)
1705-1780: 🔒 Security & Privacy | 🔵 Trivial | 💤 Low valueConsider validating/quoting
slurm_job_idbefore it reaches the remote shell.
job_idis interpolated into thesacct/squeuecommand strings, whichsshexecutes through the remote login shell. Todayslurm_job_idis parsed from trusted launcher output (numeric), so this is not currently exploitable, but a non-numeric value would be interpreted by the remote shell. A cheap defense-in-depth is to assert the id matches an expected pattern (e.g.^\d+(_\d+)?(\.\w+)?$) before building the command.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tools/mcp/modelopt_mcp/bridge.py` around lines 1705 - 1780, The _query_slurm_status helper builds remote sacct/squeue command strings using slurm_job_id, so add a defensive validation step before constructing those commands. In _query_slurm_status, ensure job_id matches the expected Slurm job-id format (for example numeric with optional array/step suffixes) and return a failed diagnostic if it does not, then only pass the validated id into _ssh_argv_for_slurm and the sacct/squeue command strings.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tools/mcp/modelopt_mcp/bridge.py`:
- Around line 1826-1841: In the Slurm status handling branch of the experiment
status lookup, do not always fall back to running when _query_slurm_status
returns a non-ok result. Update the logic around _load_slurm_status_metadata and
_query_slurm_status to treat slurm_status_not_found as a special case and
consult the local completion indicators (_DONE and status_*.out) before
defaulting to running. Keep the returned slurm_status payload intact, but ensure
the overall status can become completed when the local markers show the
experiment has finished so wait_for_experiment_impl stops polling correctly.
---
Nitpick comments:
In `@tools/mcp/modelopt_mcp/bridge.py`:
- Around line 1705-1780: The _query_slurm_status helper builds remote
sacct/squeue command strings using slurm_job_id, so add a defensive validation
step before constructing those commands. In _query_slurm_status, ensure job_id
matches the expected Slurm job-id format (for example numeric with optional
array/step suffixes) and return a failed diagnostic if it does not, then only
pass the validated id into _ssh_argv_for_slurm and the sacct/squeue command
strings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 64116f92-4ae7-4046-b0ef-672474565ecb
📒 Files selected for processing (3)
tools/launcher/common/smoke/hostname.shtools/mcp/modelopt_mcp/bridge.pytools/mcp/tests/test_bridge.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1900 +/- ##
==========================================
+ Coverage 61.21% 61.23% +0.01%
==========================================
Files 515 515
Lines 57245 57245
==========================================
+ Hits 35043 35054 +11
+ Misses 22202 22191 -11
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Review Summary
Reviewed all 3 changed files (tools/mcp/modelopt_mcp/bridge.py, tools/mcp/tests/test_bridge.py, tools/launcher/common/smoke/hostname.sh). This is a focused, well-tested change to the MCP Slurm sidecar status flow. The metadata round-trip, SSH argv construction, sacct/squeue parsing, and state mapping are all sound, and test coverage for the RUNNING/COMPLETED/wait paths is good.
Findings by severity: CRITICAL: 1, IMPORTANT: 0, SUGGESTION: 0
Most impactful finding
[CRITICAL Algorithm] job_status_impl (lines 1826-1841): when the Slurm query returns ok: False, overall is forced to running, discarding the local _DONE marker. This is correct for transient failures (slurm_status_query_failed — SSH down / MFA expired), but it is a regression for slurm_status_not_found: once a completed job's record is purged from sacct (retention is often only hours/days) and has left the queue, the experiment is reported running indefinitely, and wait_for_experiment_impl spins to its full timeout for a job that already finished. Before this PR the _DONE marker returned done immediately. Suggested fix (in the inline comment): special-case slurm_status_not_found to fall back to local completion markers, and add a test for that path.
Assessment
Low-to-moderate risk. The change is scoped to the MCP launcher tooling (no core modelopt/ optimization paths), and the happy paths are well covered. The one CRITICAL is a fallback-path correctness gap that will surface for status queries on aged jobs and for any cluster with short sacct retention; worth addressing before merge. (Note: CodeRabbit's pre-merge gate also flags the new nosec comments on the SSH subprocess calls — that's tracked separately and out of scope for this review.)
| slurm_meta = _load_slurm_status_metadata(exp_dir) | ||
| if slurm_meta is not None: | ||
| slurm_status = _query_slurm_status(slurm_meta) | ||
| if slurm_status.get("ok"): | ||
| overall = slurm_status.get("status", "unknown") | ||
| else: | ||
| overall = "running" | ||
| return { | ||
| "ok": True, | ||
| "experiment_id": experiment_id, | ||
| "experiment_dir": str(exp_dir), | ||
| "status": overall, | ||
| "task_statuses": task_statuses, | ||
| "has_done_marker": done_marker.exists(), | ||
| "slurm_status": slurm_status, | ||
| } |
There was a problem hiding this comment.
[CRITICAL Algorithm] When the sidecar is present but the Slurm query returns ok: False, overall is unconditionally set to "running", discarding the local _DONE/status_*.out signal. This is fine for the transient slurm_status_query_failed case (SSH down, MFA expired), but it is a real bug for slurm_status_not_found.
Why it matters: sacct retention (MinJobAge/purge) is frequently only hours to a few days. Once a completed job's record is purged from the accounting DB and it has left the queue, _query_slurm_status returns slurm_status_not_found. job_status_impl then reports running forever, and wait_for_experiment_impl (which only terminates on done/failed) spins to its full timeout_sec for a job that finished long ago. This is a regression: before this PR the _DONE marker would have returned done immediately.
Suggested fix: treat slurm_status_not_found specially — fall back to the local completion markers instead of forcing running:
slurm_meta = _load_slurm_status_metadata(exp_dir)
if slurm_meta is not None:
slurm_status = _query_slurm_status(slurm_meta)
if slurm_status.get("ok"):
overall = slurm_status.get("status", "unknown")
elif slurm_status.get("reason") == "slurm_status_not_found":
# Job aged out of sacct and is not in the queue: trust local markers.
if done_marker.exists():
overall = "failed" if any_failed else "done"
else:
overall = "running"
else:
# Transient query failure (SSH/MFA): keep polling.
overall = "running"
return {
"ok": True,
...
"slurm_status": slurm_status,
}Recommend adding a test for the slurm_status_not_found + _DONE path, which is currently uncovered.
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
|
/ok to test 1465d79 |
kevalmorabia97
left a comment
kevalmorabia97
left a comment
There was a problem hiding this comment.
LGTM. Please address claude feedback before merging
|
2 participants
Summary
uv run --with-editableso themodelopt-launcherconsole script is availablejob_status/wait_for_experimentprefer remotesacct/squeuestate over local_DONEcommon/smoke/hostname.shused byexamples/smoke/hostname.yamlWhy
Detached Slurm submission creates the local
_DONEmarker when the submit process exits, not when the remote Slurm job completes. This made MCPwait_for_experimentreturndonewhile Slurm jobs were still running. Reproduced on both MFA and non-MFA clusters.Validation
uv run --project tools/mcp ruff check tools/mcp/modelopt_mcp/bridge.py tools/mcp/tests/test_bridge.pyuv run --project tools/mcp pytest tools/mcp/tests/test_bridge.pyuv run --project tools/launcher pytest tools/launcher/tests/test_core_extended.py tools/launcher/tests/test_slurm_config.pyuv run --with-editable tools/launcher modelopt-launcher --yaml tools/launcher/examples/smoke/hostname.yaml --dryrun --yesManual smoke evidence
cicd_1783037181, Slurm job2320415, exit0:02945347was still running; job later completed with exit0:0Summary by CodeRabbit
New Features
Bug Fixes
Tests