Add MLflow tracking to modelopt MCP#1986
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
📝 WalkthroughWalkthroughChangesMLflow support is added to MCP job submission and experiment completion. The bridge creates runs, propagates identifiers, records submission failures, archives logs, and finalizes terminal status. Server tools, tests, dependency metadata, and documentation are updated. MLflow Job Tracking
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant MCPClient
participant submit_job
participant MLflow
participant Launcher
participant wait_for_experiment
participant finalize_mlflow_run
MCPClient->>submit_job: enable_mlflow=True
submit_job->>MLflow: create pre-submit run
submit_job->>Launcher: pass MLflow identifiers
Launcher-->>submit_job: experiment id and submission result
MCPClient->>wait_for_experiment: finalize_mlflow=True
wait_for_experiment->>finalize_mlflow_run: terminal experiment and run id
finalize_mlflow_run->>MLflow: upload logs and record final status
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
5ed2f88 to
6f11a8e
Compare
6f11a8e to
00b6ac8
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1986 +/- ##
==========================================
- Coverage 76.42% 72.83% -3.60%
==========================================
Files 520 525 +5
Lines 58339 63093 +4754
==========================================
+ Hits 44587 45951 +1364
- Misses 13752 17142 +3390
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:
|
00b6ac8 to
e5a9c6e
Compare
|
/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
Focused, well-structured PR adding optional MLflow tracking to the modelopt-mcp launcher bridge. Scope: 5 files (bridge.py, server.py, README.md, pyproject.toml, test_bridge.py). I reviewed all source files and the tests; skipped only the doc/pyproject text beyond a scan.
Findings — CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 0
What's solid
- Lazy import of
mlflowvia_setup_mlflow()correctly gates the optional dependency — users who don't enable tracking never import it, and a missing package returns a structuredmlflow_not_installedresult rather than raising. - Run lifecycle is coherent: pre-submit run creation,
MLFLOW_RUN_ID/MLFLOW_TRACKING_URIinjected into both Docker and Slurm child envs, terminal finalization uploads launcher logs and records a pass/failresultmetric. The env injection sits after_apply_launcher_env, so it isn't clobbered. - Failure marking is applied on nearly every post-creation bail-out (Slurm FileNotFound/timeout/returncode/unparsed, Docker FileNotFound), so runs generally don't leak in the
submittedstate. - stdout redirect via
contextlib.redirect_stdoutkeeps MLflow chatter out of the MCP stdio channel — important for an MCP server. - Test coverage is thorough, exercising env injection, run-name construction, artifact upload, Nemo-log fetch, and the
wait_for_experimentdelegation path.
Blocking-ish issue (1 IMPORTANT)
The Docker docker_submit_log_unavailable (OSError) early-return is the one post-run-creation failure path that omits _mark_mlflow_submission_failed. When enable_mlflow=True and the submit-log dir can't be created, the MLflow run is orphaned in lifecycle=submitted with no terminal result — the exact leak the other paths avoid. Inline comment has the one-line fix.
Risk assessment
Low risk. Tooling-only change (tools/mcp/), no modelopt core, no mode/state/export surface, feature is opt-in and defaults off. The single finding is a consistency gap on a rare error path, easily addressed.
🤖 Generated with Claude Code
![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: 6
🤖 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 799-833: Update the Nemo log-fetch loop around subprocess.run so
stdout and stderr stream directly to each target artifact instead of using
capture_output and duplicating the full body in memory. Open target for writing
before invoking the process, retain only a bounded stderr/output diagnostic tail
for error reporting, and preserve timeout, missing-CLI, and fetched-file
handling in the existing loop.
- Around line 779-787: Update _infer_nemo_log_task_count so an explicit_count is
only used when positive; treat zero or negative values as unknown and continue
to infer the count from status, falling back to 1 when the task_statuses map is
empty or unavailable. Preserve the existing nonnegative handling for positive
explicit counts and ensure callers passing 0 still fetch one terminal task log.
- Around line 1525-1544: Update the post-creation failure paths in the enclosing
submission flow so every failure after _create_mlflow_run succeeds calls
_mark_mlflow_submission_failed before returning. Apply this to both the Docker
log-creation failure path and the additional failure path around the referenced
lines, while preserving each existing response payload and handling only runs
that were successfully created.
- Around line 803-808: Remove the # nosec B603 B607 suppression from the
subprocess.run call in the bridge execution flow, then satisfy the Bandit check
by validating the argv contents before invocation or using the repository’s
approved subprocess abstraction. Preserve the fixed-argument, no-shell behavior
and existing timeout and output handling.
- Around line 874-907: Update the finalization flow before log fetching and
MLflow setup to reject any non-terminal status: after resolving final_status,
return an experiment_not_terminal failure for statuses other than done or
failed, including the experiment identifier and existing response fields as
appropriate. Allow the existing finalized lifecycle, result metric, and
log-upload path to run only for terminal statuses.
- Around line 848-864: Update finalize_mlflow_run_impl to validate experiment_id
through _validate_experiment_id before calling _resolve_experiment_dir or
performing any artifact/log upload; preserve the existing missing-mlflow-run
handling and ensure invalid external identifiers are rejected at this interface
boundary.
🪄 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: 5fc75483-6a5f-4934-bbf1-5c490f3c354d
📒 Files selected for processing (5)
tools/mcp/README.mdtools/mcp/modelopt_mcp/bridge.pytools/mcp/modelopt_mcp/server.pytools/mcp/pyproject.tomltools/mcp/tests/test_bridge.py
| proc = subprocess.run( # nosec B603 B607 - fixed nemo CLI argv; no shell. | ||
| argv, | ||
| capture_output=True, | ||
| text=True, | ||
| timeout=120, | ||
| check=False, |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Remove the prohibited Bandit bypass.
Line 803 adds # nosec B603 B607. Remove the suppression and address the check through validated arguments or the repository’s approved subprocess abstraction.
As per coding guidelines, “# nosec comments are not allowed as a bypass for security checks.”
🧰 Tools
🪛 ast-grep (0.44.1)
[error] 802-808: Command coming from incoming request
Context: subprocess.run( # nosec B603 B607 - fixed nemo CLI argv; no shell.
argv,
capture_output=True,
text=True,
timeout=120,
check=False,
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
🤖 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 803 - 808, Remove the # nosec
B603 B607 suppression from the subprocess.run call in the bridge execution flow,
then satisfy the Bandit check by validating the argv contents before invocation
or using the repository’s approved subprocess abstraction. Preserve the
fixed-argument, no-shell behavior and existing timeout and output handling.
Source: Coding guidelines
|
/ok to test e5a9c6e |
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
e5a9c6e to
d256b07
Compare
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: 4
🤖 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 895-901: Update the MLflow finalization loop around
_infer_nemo_log_task_count and _fetch_single_nemo_log_to_file to enforce one
overall deadline for the operation rather than allowing each sequential fetch
its own 120-second timeout. Use bounded concurrency or stop scheduling work when
the deadline is exhausted, while preserving per-task log retrieval and result
handling within the available budget.
- Around line 813-824: Update the parent result-handling logic around
result_queue so an absent child message is treated as failure, not success: do
not default result.get("ok", True) to true, and require an explicit successful
result before archiving. Preserve the existing error reporting for messages with
ok set to false.
- Around line 762-778: Harden _mlflow_artifact_files and the related
finalization log-capture/upload paths so job-created symlinks are never
followed: open log targets with no-follow semantics, reject symlinks and
non-regular files, and upload only securely snapshotted regular files whose
resolved location remains within exp_dir. Apply the same validation to the
referenced log-writing and MLflow upload call sites, preserving the existing
artifact patterns and deduplication behavior.
- Around line 699-720: The failure path in _mark_mlflow_submission_failed() must
explicitly terminate the reopened MLflow run as FAILED instead of allowing the
start_run context to finish as FINISHED. Update that method to use
MlflowClient.set_terminated with status="FAILED" or call
mlflow.end_run(status="FAILED") while preserving the existing failure-recording
behavior.
🪄 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: 8547b75b-3097-4d1c-8bc2-c98747aa9757
📒 Files selected for processing (5)
tools/mcp/README.mdtools/mcp/modelopt_mcp/bridge.pytools/mcp/modelopt_mcp/server.pytools/mcp/pyproject.tomltools/mcp/tests/test_bridge.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tools/mcp/pyproject.toml
- tools/mcp/README.md
- tools/mcp/modelopt_mcp/server.py
| with contextlib.redirect_stdout(stdout): | ||
| mlflow.set_tracking_uri(uri) | ||
| mlflow.set_experiment(exp_name) | ||
| with mlflow.start_run(run_name=run_name) as run: | ||
| run_id = run.info.run_id | ||
| experiment_id = run.info.experiment_id | ||
| mlflow.log_params( | ||
| { | ||
| "target_yaml": yaml_path, | ||
| "executor": executor, | ||
| "cluster_host": cluster_host or "", | ||
| "source_ref": source.ref if source else "", | ||
| "source_sha": source.resolved_sha if source else "", | ||
| } | ||
| ) | ||
| mlflow.set_tags( | ||
| { | ||
| "lifecycle": "submitted", | ||
| "managed_by": "modelopt-mcp", | ||
| } | ||
| ) | ||
| except Exception as exc: # pragma: no cover - exercised through unit-level monkeypatches. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C3 '\b(start_run|create_run|set_terminated|end_run)\s*\(' \
tools/mcp/modelopt_mcp/bridge.py tools/mcp/tests/test_bridge.py
rg -n -C2 'mlflow' tools/mcp/pyproject.tomlRepository: NVIDIA/Model-Optimizer
Length of output: 3518
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('tools/mcp/modelopt_mcp/bridge.py')
lines = p.read_text().splitlines()
for start,end in [(660,760),(940,1010)]:
print(f"\n--- {p} lines {start}-{end} ---")
for i in range(start-1, min(end, len(lines))):
print(f"{i+1:4d}: {lines[i]}")
PY
printf '\n--- AST outline ---\n'
ast-grep outline tools/mcp/modelopt_mcp/bridge.py --view expandedRepository: NVIDIA/Model-Optimizer
Length of output: 12123
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('tools/mcp/tests/test_bridge.py')
lines = p.read_text().splitlines()
for start,end in [(760,860),(1540,1665)]:
print(f"\n--- {p} lines {start}-{end} ---")
for i in range(start-1, min(end, len(lines))):
print(f"{i+1:4d}: {lines[i]}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 9210
🌐 Web query:
MLflow start_run context manager status FINISHED on exit run_id resume run create_run set_terminated documentation
💡 Result:
In MLflow, the mlflow.start_run function serves as the primary entry point for managing experiment runs, providing a context manager interface to ensure runs are automatically terminated [1][2][3]. Context Manager Behavior: When used as a context manager (e.g., with mlflow.start_run:), the run remains active throughout the duration of the code block [1]. Upon exiting the block—whether successfully or due to an exception—the run is automatically closed [1][4]. The automatic termination behavior defaults to a status of FINISHED if no error occurs, though the exact status logic is managed internally by the run's lifecycle [1][5][3]. Resuming a Run: You can resume an existing run by passing its unique run_id to mlflow.start_run(run_id=...) [6][7][8]. When a run is resumed, MLflow sets its status to RUNNING [7][8]. If you are experiencing issues where a resumed run does not properly transition to a running state, it is common to use the MlflowClient to explicitly update the status [9][6][10]. Manual Run Management: If you need finer control over the run lifecycle, you can use the MLflow Fluent API or the MlflowClient: - mlflow.end_run(status='FINISHED'): Manually terminates the currently active run with the specified status (e.g., FINISHED, FAILED, or KILLED) [2][5][3]. - MlflowClient.create_run(...): Creates a new run object without automatically setting it as the active run, which is useful when you want to manage runs without the implicit context manager behavior [10]. - MlflowClient.set_terminated(run_id, status='FINISHED'): Explicitly updates an existing run's status to a terminal state [10]. For complex scenarios where standard context management is insufficient, the MlflowClient provides the most robust methods for programmatically creating, updating, and terminating runs [10].
Citations:
- 1: https://mlflow.org/docs/2.13.0/tracking/tracking-api.html
- 2: https://mlflow.org/docs/3.3.0/api_reference/python_api/mlflow.html
- 3: https://mlflow.org/docs/latest/api_reference/python_api/mlflow.html
- 4: https://mlflow.org/docs/2.12.1/tracking/tracking-api.html
- 5: https://mlflow.org/docs/2.20.3/python_api/mlflow.html
- 6: https://stackoverflow.com/questions/67062145/continue-stopped-run-in-mlflow
- 7: https://github.com/mlflow/mlflow/blob/14edcfa9/mlflow/tracking/fluent.py
- 8: https://github.com/mlflow/mlflow/blob/4bf54ce9ae22248d2692d2f981e776719a767acf/mlflow/tracking/fluent.py
- 9: [BUG] Resuming runs with mlflow.start_run(run_id=) does not set state properly. mlflow/mlflow#3228
- 10: https://mlflow.org/docs/latest/api_reference/python_api/mlflow.client.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C3 '_mark_mlflow_submission_failed|_create_mlflow_run\(' tools/mcp/modelopt_mcp/bridge.pyRepository: NVIDIA/Model-Optimizer
Length of output: 2758
Mark failed submissions as FAILED.
_mark_mlflow_submission_failed() reopens the run with start_run(run_id=...), but that context still exits as FINISHED, so launcher failures are recorded as successful runs. Use MlflowClient.set_terminated(..., status="FAILED") or mlflow.end_run(status="FAILED") here.
🤖 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 699 - 720, The failure path in
_mark_mlflow_submission_failed() must explicitly terminate the reopened MLflow
run as FAILED instead of allowing the start_run context to finish as FINISHED.
Update that method to use MlflowClient.set_terminated with status="FAILED" or
call mlflow.end_run(status="FAILED") while preserving the existing
failure-recording behavior.
| def _mlflow_artifact_files(exp_dir: Path) -> list[Path]: | ||
| """Local launcher files worth preserving durably in MLflow.""" | ||
| patterns = [ | ||
| "log_*.out", | ||
| "nemo_logs_task_*.log", | ||
| "status_*.out", | ||
| "_DONE", | ||
| _SLURM_STATUS_META, | ||
| ] | ||
| seen: set[Path] = set() | ||
| files: list[Path] = [] | ||
| for pattern in patterns: | ||
| for path in sorted(exp_dir.glob(pattern)): | ||
| if path.is_file() and path not in seen: | ||
| seen.add(path) | ||
| files.append(path) | ||
| return files |
There was a problem hiding this comment.
🔒 Security & Privacy | 🔴 Critical | 🏗️ Heavy lift
Do not follow job-controlled symlinks during log capture or upload.
A workload can pre-create nemo_logs_task_0.log or log_*.out as a symlink. Finalization then either truncates a host file via open("w") or uploads an arbitrary readable host file to MLflow. Open targets with no-follow semantics and upload only securely snapshotted regular files contained within exp_dir.
Also applies to: 807-812, 895-900, 983-985
🤖 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 762 - 778, Harden
_mlflow_artifact_files and the related finalization log-capture/upload paths so
job-created symlinks are never followed: open log targets with no-follow
semantics, reject symlinks and non-regular files, and upload only securely
snapshotted regular files whose resolved location remains within exp_dir. Apply
the same validation to the referenced log-writing and MLflow upload call sites,
preserving the existing artifact patterns and deduplication behavior.
| except BaseException as exc: | ||
| with contextlib.suppress(Exception): | ||
| result_queue.put( | ||
| { | ||
| "ok": False, | ||
| "error_type": type(exc).__name__, | ||
| "diagnostic": str(exc), | ||
| } | ||
| ) | ||
| return | ||
| else: | ||
| result_queue.put({"ok": True}) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not treat a missing child result as success.
The child catches failures and exits with code 0; if its queue message is unavailable, result.get("ok", True) silently reports success and archives a partial file.
Proposed fix
try:
- result = result_queue.get_nowait()
+ result = result_queue.get(timeout=1)
except queue.Empty:
- result = {}
- if proc.exitcode == 0 and result.get("ok", True):
+ return {
+ "ok": False,
+ "reason": "logs_fetch_failed",
+ "diagnostic": "Nemo log fetch exited without returning a result.",
+ "exit_code": proc.exitcode,
+ }
+ if proc.exitcode == 0 and result.get("ok") is True:
return {"ok": True}Also applies to: 861-866
🤖 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 813 - 824, Update the parent
result-handling logic around result_queue so an absent child message is treated
as failure, not success: do not default result.get("ok", True) to true, and
require an explicit successful result before archiving. Preserve the existing
error reporting for messages with ok set to false.
| for job_idx in range(_infer_nemo_log_task_count(experiment_id, task_count)): | ||
| target = exp_dir / f"nemo_logs_task_{job_idx}.log" | ||
| result = _fetch_single_nemo_log_to_file( | ||
| experiment_id=experiment_id, | ||
| job_idx=job_idx, | ||
| target=target, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Bound total MLflow finalization time.
Sequential 120-second fetches make worst-case latency task_count × 120s; a terminal job with many tasks can block the MCP request for hours. Enforce an overall deadline and use bounded concurrency or stop once the budget is exhausted.
🤖 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 895 - 901, Update the MLflow
finalization loop around _infer_nemo_log_task_count and
_fetch_single_nemo_log_to_file to enforce one overall deadline for the operation
rather than allowing each sequential fetch its own 120-second timeout. Use
bounded concurrency or stop scheduling work when the deadline is exhausted,
while preserving per-task log retrieval and result handling within the available
budget.
1 participant
Summary
Testing
Summary by CodeRabbit