diff --git a/tools/mcp/README.md b/tools/mcp/README.md index 3715612d155..231bfd0a63b 100644 --- a/tools/mcp/README.md +++ b/tools/mcp/README.md @@ -21,10 +21,11 @@ Mode is determined by which args you pass, not by which tool you call. One tool, |---|---| | `list_examples` | Enumerate bundled launcher YAMLs under `tools/launcher/examples/` with model + description metadata extracted from each YAML. Discovery primitive — call this first when you don't know which YAML to launch. | | `verify_setup(executor, ...)` | Fail-fast probe for the named executor. Docker: `docker info` (daemon up) + `docker info --format` runtime-registry check (looks for `"nvidia"` runtime registered by the NVIDIA Container Toolkit — no image pull, daemon-fast). Slurm: `ssh -o BatchMode=yes -o ConnectTimeout=5` to the cluster login node. Returns structured failure on auth / network / daemon issues — no exception. | -| `submit_job(yaml_path, hf_local? \| cluster_host?, ..., dry_run?, source_ref?, source_repo?)` | Submit a launcher YAML. Mode resolved from mutually-exclusive args. Before launching, materializes a managed Model-Optimizer checkout at `source_ref` (branch, tag, or SHA; default `main`) and initializes recursive submodules, then runs that checkout's launcher. Returns `experiment_id` (Slurm) or PID plus captured `experiment_id` when available (Docker) immediately; if Docker's short tail times out, it still returns PID with `experiment_id=None` and a persistent `stdout_log` under `$NEMORUN_HOME/.modelopt-mcp/docker-submit-logs/` for diagnostics. The actual job runs detached. Auto-runs `verify_setup` first by default (skippable). **Pass `dry_run=True`** to validate the YAML via `launch.py --dryrun --yes` without contacting the cluster / spawning a container / running sbatch — returns `{ok, dry_run: True, validated: bool, diagnostic?, exit_code, stdout_tail, stderr_tail, argv, source_sha, source_root}` instead of `experiment_id`. Used by verify-task workflow stages (deployment_support, hidden_state_dump_support, mlm_eval, ...). | +| `submit_job(yaml_path, hf_local? \| cluster_host?, ..., dry_run?, source_ref?, source_repo?, enable_mlflow?)` | Submit a launcher YAML. Mode resolved from mutually-exclusive args. Before launching, materializes a managed Model-Optimizer checkout at `source_ref` (branch, tag, or SHA; default `main`) and initializes recursive submodules, then runs that checkout's launcher. Returns `experiment_id` (Slurm) or PID plus captured `experiment_id` when available (Docker) immediately; if Docker's short tail times out, it still returns PID with `experiment_id=None` and a persistent `stdout_log` under `$NEMORUN_HOME/.modelopt-mcp/docker-submit-logs/` for diagnostics. The actual job runs detached. Auto-runs `verify_setup` first by default (skippable). When `enable_mlflow=True`, creates an MLflow run before launch and passes standard `MLFLOW_RUN_ID` / `MLFLOW_TRACKING_URI` into the job. **Pass `dry_run=True`** to validate the YAML via `launch.py --dryrun --yes` without contacting the cluster / spawning a container / running sbatch — returns `{ok, dry_run: True, validated: bool, diagnostic?, exit_code, stdout_tail, stderr_tail, argv, source_sha, source_root}` instead of `experiment_id`. | | `job_status(experiment_id)` | Filesystem-based status from nemo_run's experiment dir (`_DONE`, `status_*.out`). Returns `done` / `failed` / `running` plus per-task statuses. No in-memory registry; survives MCP server restarts. | | `job_logs(experiment_id, task?, tail?)` | Read `log_.out` from the experiment dir. Per-task filtering + optional tail to truncate. | -| `wait_for_experiment(experiment_id, timeout_sec?, poll_interval_sec?)` | Block until `job_status` returns `done` / `failed`, or until `timeout_sec` elapses. Single tool call replaces the agent's `while True: status; sleep` loop — saves tool-call turns and avoids overshooting the poll interval. Returns the final status plus `waited_seconds`. | +| `wait_for_experiment(experiment_id, timeout_sec?, poll_interval_sec?, finalize_mlflow?, mlflow_run_id?)` | Block until `job_status` returns `done` / `failed`, or until `timeout_sec` elapses. Single tool call replaces the agent's `while True: status; sleep` loop — saves tool-call turns and avoids overshooting the poll interval. When `finalize_mlflow=True`, fetches terminal Nemo/Slurm logs with `nemo experiment logs`, then attaches final pass/fail status and launcher logs to the MLflow run created by `submit_job(enable_mlflow=True)`. | +| `finalize_mlflow_run(experiment_id, mlflow_run_id, mlflow_tracking_uri?)` | Fetch terminal Nemo/Slurm logs and attach final status plus launcher logs to an already-created MLflow run. Usually use `wait_for_experiment(..., finalize_mlflow=True)` instead. | | `provision_passwordless_ssh_dry_run(cluster_host, cluster_user, identity?)` | Operator UX helper. Inspects `~/.ssh/` and emits the exact `ssh-keygen` / `ssh-copy-id` commands the user should run to make `verify_setup(executor='slurm')` pass. No side effects — pure inspection + shell-command formatting. | | `read_cluster_artifact(experiment_id, path?, job_idx?)` | Pull artifact content from the remote cluster via nemo_run's tunnel primitives. With `path=None`, wraps `nemo experiment logs ` (built-in log fetch). With a `path`, uses the experiment's `Tunnel` to `cat` the file. No reinvented SSH. | | `open_draft_pr(target_repo, title, body, base_branch?, cwd?)` | Push the current branch and open a draft PR via `gh pr create --draft`. Validates `cwd` is a git repo before doing anything. On `gh` failure after a successful push, returns `branch_pushed=True` so the operator can retry just the PR-open step. | @@ -115,14 +116,16 @@ result = mcp__modelopt__submit_job( identity="/home/alice/.ssh/id_ed25519", skip_verify=True, # we just probed source_ref="main", # optional; omit to use main + enable_mlflow=True, # optional; passes MLFLOW_RUN_ID into the job ) -# {"ok": True, "experiment_id": "cicd_1781240000", "slurm_job_id": "12345", "source_sha": "...", ...} +# {"ok": True, "experiment_id": "cicd_1781240000", "mlflow": {"mlflow_run_id": "..."}, ...} -# 4. Poll until done -while True: - status = mcp__modelopt__job_status(experiment_id="cicd_1781240000") - if status["status"] in ("done", "failed"): - break +# 4. Poll until done and attach durable logs to MLflow +status = mcp__modelopt__wait_for_experiment( + experiment_id=result["experiment_id"], + finalize_mlflow=True, + mlflow_run_id=result["mlflow"]["mlflow_run_id"], +) # 5. Fetch logs logs = mcp__modelopt__job_logs( @@ -147,6 +150,12 @@ For local Docker execution, drop `cluster_host`/`cluster_user`/`identity` and pa | `MODELOPT_MCP_DISABLE_MANAGED_SOURCE` | (optional) local dev | Set to `1` to skip managed checkout and invoke the installed `modelopt-launcher` entrypoint directly. | | `MODELOPT_MCP_UV` | (optional) `submit_job` | Override the `uv` binary used for `uv run --project /tools/launcher ...`. | | `MODELOPT_LAUNCHER_EXAMPLES_DIR` | (optional) `list_examples` | Override the examples directory location. Defaults to `../launcher/examples/` relative to this package. | +| `MLFLOW_TRACKING_URI` / `MODELOPT_MCP_MLFLOW_TRACKING_URI` | (optional) `submit_job(enable_mlflow=True)` | Tracking server for MCP-created runs. Defaults to `https://mlflow-modelopt.nvidia.com/`. | +| `MODELOPT_MCP_MLFLOW_EXPERIMENT` | (optional) `submit_job(enable_mlflow=True)` | Experiment name for MCP-created runs. Defaults to `ci-tests`. | + +MCP-created MLflow runs default to the run name +`///`. Passing +`job_name` to `submit_job` overrides that run name. ## Design principles diff --git a/tools/mcp/modelopt_mcp/bridge.py b/tools/mcp/modelopt_mcp/bridge.py index 5ac251e83b1..70c92f17897 100644 --- a/tools/mcp/modelopt_mcp/bridge.py +++ b/tools/mcp/modelopt_mcp/bridge.py @@ -33,9 +33,14 @@ from __future__ import annotations +import contextlib +import getpass import hashlib +import io import json +import multiprocessing import os +import queue import re import shutil import subprocess # nosec B404 - fixed-argv CLI probes are required; shell=True is not used. @@ -43,6 +48,7 @@ import time from dataclasses import dataclass from pathlib import Path +from typing import Any import yaml @@ -53,6 +59,8 @@ _DEFAULT_SOURCE_REPO = "https://github.com/NVIDIA/Model-Optimizer.git" _DEFAULT_SOURCE_REF = "main" +_DEFAULT_MLFLOW_TRACKING_URI = "https://mlflow-modelopt.nvidia.com/" +_DEFAULT_MLFLOW_EXPERIMENT = "ci-tests" # Canonical task-status failure tokens — matched against the FIRST word # of each ``status_.out`` file by ``job_status_impl``. @@ -580,6 +588,425 @@ def _launcher_argv(abs_yaml: Path, checkout: SourceCheckout | None, *flags: str) ] +# --------------------------------------------------------------------------- +# Optional MLflow tracking +# --------------------------------------------------------------------------- + + +def _mlflow_tracking_uri(tracking_uri: str | None) -> str: + """Resolve the tracking URI used for MCP-managed MLflow runs.""" + return ( + tracking_uri + or os.environ.get("MLFLOW_TRACKING_URI") + or os.environ.get("MODELOPT_MCP_MLFLOW_TRACKING_URI") + or _DEFAULT_MLFLOW_TRACKING_URI + ) + + +def _mlflow_experiment_name(experiment_name: str | None) -> str: + """Resolve the experiment name for MCP-managed MLflow runs.""" + return ( + experiment_name + or os.environ.get("MODELOPT_MCP_MLFLOW_EXPERIMENT") + or _DEFAULT_MLFLOW_EXPERIMENT + ) + + +def _setup_mlflow() -> tuple[Any | None, dict | None]: + """Import MLflow lazily so users who do not enable tracking avoid the dependency.""" + try: + import mlflow + except ImportError: + return None, { + "ok": False, + "reason": "mlflow_not_installed", + "diagnostic": ( + "MLflow tracking was requested, but the `mlflow` Python package " + "is not installed in the modelopt-mcp environment." + ), + } + return mlflow, None + + +def _mlflow_run_url(tracking_uri: str, experiment_id: str | None, run_id: str) -> str | None: + """Best-effort browser URL for MLflow UIs that use the standard hash route.""" + if not tracking_uri.startswith(("http://", "https://")) or not experiment_id: + return None + return f"{tracking_uri.rstrip('/')}/#/experiments/{experiment_id}/runs/{run_id}" + + +def _mlflow_basename(yaml_path: str) -> str: + """Human-readable run-name base from a launcher YAML path.""" + path = Path(yaml_path) + parent = path.parent.name + stem = path.stem + return f"{parent}/{stem}" if parent else stem + + +def _mlflow_user_label() -> str: + """Stable user component for MCP-created MLflow run names.""" + return os.environ.get("USER") or os.environ.get("USERNAME") or getpass.getuser() + + +def _mlflow_cluster_label(cluster_host: str | None, executor: str) -> str: + """Stable cluster component for MCP-created MLflow run names.""" + return cluster_host or executor + + +def _mlflow_run_name( + *, + yaml_path: str, + executor: str, + cluster_host: str | None, + job_name: str | None, +) -> str: + """Default MCP MLflow run name: user/cluster/yaml.""" + if job_name: + return job_name + return ( + f"{_mlflow_user_label()}/" + f"{_mlflow_cluster_label(cluster_host, executor)}/" + f"{_mlflow_basename(yaml_path)}" + ) + + +def _create_mlflow_run( + *, + yaml_path: str, + executor: str, + tracking_uri: str | None, + experiment_name: str | None, + cluster_host: str | None, + job_name: str | None, + source: SourceCheckout | None, +) -> dict: + """Create the pre-submit MLflow run and return fields for submit_job.""" + mlflow, error = _setup_mlflow() + if error: + return error + assert mlflow is not None + + uri = _mlflow_tracking_uri(tracking_uri) + exp_name = _mlflow_experiment_name(experiment_name) + run_name = _mlflow_run_name( + yaml_path=yaml_path, + executor=executor, + cluster_host=cluster_host, + job_name=job_name, + ) + stdout = io.StringIO() + try: + 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. + return { + "ok": False, + "reason": "mlflow_create_failed", + "diagnostic": f"Failed to create MLflow run: {exc}", + "mlflow_tracking_uri": uri, + "mlflow_experiment": exp_name, + } + + return { + "ok": True, + "mlflow_run_id": run_id, + "mlflow_tracking_uri": uri, + "mlflow_experiment": exp_name, + "mlflow_experiment_id": experiment_id, + "mlflow_run_url": _mlflow_run_url(uri, experiment_id, run_id), + } + + +def _mark_mlflow_submission_failed(mlflow_run_id: str, tracking_uri: str, reason: str) -> None: + """Best-effort update when submission fails after a run was created.""" + mlflow, error = _setup_mlflow() + if error: + return + assert mlflow is not None + with contextlib.suppress(Exception), contextlib.redirect_stdout(io.StringIO()): + mlflow.set_tracking_uri(tracking_uri) + with mlflow.start_run(run_id=mlflow_run_id): + mlflow.set_tags({"lifecycle": "submit_failed", "result_label": "SUBMIT_FAILED"}) + mlflow.log_param("submit_failure_reason", reason) + mlflow.log_metric("result", 0.0) + + +def _mlflow_result_metric(status: str) -> float: + """Convert MCP terminal status to a compact numeric result.""" + if status == "done": + return 1.0 + if status == "failed": + return 0.0 + return -1.0 + + +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 + + +def _infer_nemo_log_task_count(experiment_id: str, explicit_count: int | None = None) -> int: + """Best-effort task count for `nemo experiment logs ` fetches.""" + if explicit_count is not None and explicit_count > 0: + return explicit_count + status = job_status_impl(experiment_id) + task_statuses = status.get("task_statuses") if status.get("ok") else None + if isinstance(task_statuses, dict) and task_statuses: + return len(task_statuses) + return 1 + + +def _nemo_experiment_logs(experiment_id: str, job_idx: int) -> None: + """Run Nemo Run's log fetch command in-process.""" + from nemo_run.cli.experiment import logs + + logs(experiment_id, job_idx) + + +def _run_nemo_logs_process( + experiment_id: str, + job_idx: int, + target: str, + result_queue: multiprocessing.Queue, +) -> None: + """Child-process entry point that streams Nemo logs into target.""" + try: + with ( + Path(target).open("w", encoding="utf-8", errors="replace") as out, + contextlib.redirect_stdout(out), + contextlib.redirect_stderr(out), + ): + _nemo_experiment_logs(experiment_id, job_idx) + 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}) + + +def _read_text_tail(path: Path) -> str: + """Best-effort diagnostic tail from a possibly partial artifact.""" + try: + return _tail(path.read_text(encoding="utf-8", errors="replace")) + except OSError: + return "" + + +def _fetch_single_nemo_log_to_file( + *, + experiment_id: str, + job_idx: int, + target: Path, + timeout_sec: int = 120, +) -> dict: + """Fetch one Nemo task log to target without buffering it in the MCP process.""" + ctx = multiprocessing.get_context() + result_queue = ctx.Queue() + proc = ctx.Process( + target=_run_nemo_logs_process, + args=(experiment_id, job_idx, str(target), result_queue), + ) + proc.daemon = True + proc.start() + proc.join(timeout_sec) + if proc.is_alive(): + proc.terminate() + proc.join(5) + return { + "ok": False, + "reason": "logs_fetch_timeout", + "diagnostic": f"`nemo experiment logs {experiment_id} {job_idx}` timed out.", + } + + try: + result = result_queue.get_nowait() + except queue.Empty: + result = {} + if proc.exitcode == 0 and result.get("ok", True): + return {"ok": True} + + error_type = result.get("error_type") + reason = ( + "nemo_run_not_installed" + if error_type in {"ImportError", "ModuleNotFoundError"} + else "logs_fetch_failed" + ) + diagnostic = result.get("diagnostic") or f"Nemo log fetch exited with code {proc.exitcode}." + tail = _read_text_tail(target) + if tail: + diagnostic = f"{diagnostic} Output tail: {tail}" + return { + "ok": False, + "reason": reason, + "diagnostic": diagnostic, + "exit_code": proc.exitcode, + } + + +def _fetch_nemo_logs_to_experiment_dir( + *, + experiment_id: str, + exp_dir: Path, + task_count: int | None = None, +) -> dict: + """Fetch terminal Nemo/Slurm logs into exp_dir so MLflow can archive them.""" + fetched: list[str] = [] + errors: list[dict] = [] + 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, + ) + if result.get("ok"): + fetched.append(target.name) + else: + errors.append({"job_idx": job_idx, "artifact": target.name, **result}) + + return {"fetched_artifacts": fetched, "errors": errors} + + +def finalize_mlflow_run_impl( + *, + experiment_id: str, + mlflow_run_id: str, + tracking_uri: str | None = None, + status: str | None = None, + fetch_remote_logs: bool = True, + task_count: int | None = None, +) -> dict: + """Attach final MCP status and local launcher logs to a pre-created MLflow run.""" + if not mlflow_run_id: + return { + "ok": False, + "reason": "missing_mlflow_run_id", + "diagnostic": "finalize_mlflow_run requires the mlflow_run_id returned by submit_job.", + } + invalid = _validate_experiment_id(experiment_id) + if invalid: + return {**invalid, "mlflow_run_id": mlflow_run_id} + + exp_dir = _resolve_experiment_dir(experiment_id) + if exp_dir is None: + return { + "ok": False, + "experiment_id": experiment_id, + "mlflow_run_id": mlflow_run_id, + "reason": "experiment_dir_not_found", + "diagnostic": _experiment_not_found_diagnostic(), + } + + final_status = status or job_status_impl(experiment_id).get("status", "unknown") + if final_status not in {"done", "failed"}: + return { + "ok": False, + "experiment_id": experiment_id, + "experiment_dir": str(exp_dir), + "mlflow_run_id": mlflow_run_id, + "reason": "experiment_not_terminal", + "status": final_status, + } + log_fetch = {"fetched_artifacts": [], "errors": []} + if fetch_remote_logs: + log_fetch = _fetch_nemo_logs_to_experiment_dir( + experiment_id=experiment_id, + exp_dir=exp_dir, + task_count=task_count, + ) + uri = _mlflow_tracking_uri(tracking_uri) + mlflow, error = _setup_mlflow() + if error: + return { + **error, + "experiment_id": experiment_id, + "mlflow_run_id": mlflow_run_id, + "log_fetch": log_fetch, + } + assert mlflow is not None + + uploaded: list[str] = [] + stdout = io.StringIO() + try: + with contextlib.redirect_stdout(stdout): + mlflow.set_tracking_uri(uri) + with mlflow.start_run(run_id=mlflow_run_id): + with contextlib.suppress(Exception): + mlflow.log_param("experiment_id", experiment_id) + mlflow.log_metric("result", _mlflow_result_metric(final_status)) + mlflow.set_tags( + { + "lifecycle": "finalized", + "result_label": final_status.upper(), + } + ) + for path in _mlflow_artifact_files(exp_dir): + mlflow.log_artifact(str(path), artifact_path="logs") + uploaded.append(path.name) + except Exception as exc: + return { + "ok": False, + "reason": "mlflow_finalize_failed", + "diagnostic": f"Failed to finalize MLflow run: {exc}", + "experiment_id": experiment_id, + "mlflow_run_id": mlflow_run_id, + "mlflow_tracking_uri": uri, + "uploaded_artifacts": uploaded, + "log_fetch": log_fetch, + } + + return { + "ok": True, + "experiment_id": experiment_id, + "experiment_dir": str(exp_dir), + "mlflow_run_id": mlflow_run_id, + "mlflow_tracking_uri": uri, + "status": final_status, + "uploaded_artifacts": uploaded, + "log_fetch": log_fetch, + } + + # --------------------------------------------------------------------------- # list_examples # --------------------------------------------------------------------------- @@ -1039,6 +1466,9 @@ def submit_job_impl( dry_run: bool = False, source_ref: str | None = None, source_repo: str | None = None, + enable_mlflow: bool = False, + mlflow_tracking_uri: str | None = None, + mlflow_experiment: str | None = None, ) -> dict: """Submit a launcher YAML. @@ -1167,6 +1597,27 @@ def submit_job_impl( ), } + mlflow_info: dict = {"ok": True, "enabled": False} + if enable_mlflow: + mlflow_info = _create_mlflow_run( + yaml_path=str(abs_yaml), + executor=executor, + tracking_uri=mlflow_tracking_uri, + experiment_name=mlflow_experiment, + cluster_host=cluster_host, + job_name=job_name, + source=checkout, + ) + if not mlflow_info.get("ok"): + return { + **mlflow_info, + "executor": executor, + "yaml_path": yaml_path, + "resolved_path": str(abs_yaml), + **_source_result_fields(checkout), + } + mlflow_info["enabled"] = True + # ---- Dispatch to the launcher --------------------------------- # Subprocess `uv run launch.py --yaml --yes ...` rather # than calling core.run_jobs directly in-process. Why subprocess: @@ -1223,6 +1674,9 @@ def submit_job_impl( control_socket=control_socket, reconnect_command=reconnect_command, ) + if mlflow_info.get("enabled"): + child_env["MLFLOW_RUN_ID"] = mlflow_info["mlflow_run_id"] + child_env["MLFLOW_TRACKING_URI"] = mlflow_info["mlflow_tracking_uri"] if executor == "docker": # Docker mode: spawn detached. Redirect stdout/stderr to a side-channel @@ -1244,6 +1698,12 @@ def submit_job_impl( mode="w+b", ) except OSError as e: + if mlflow_info.get("enabled"): + _mark_mlflow_submission_failed( + mlflow_info["mlflow_run_id"], + mlflow_info["mlflow_tracking_uri"], + "docker_submit_log_unavailable", + ) return { "ok": False, "executor": "docker", @@ -1264,6 +1724,12 @@ def submit_job_impl( except FileNotFoundError: log_file.close() log_path.unlink(missing_ok=True) + if mlflow_info.get("enabled"): + _mark_mlflow_submission_failed( + mlflow_info["mlflow_run_id"], + mlflow_info["mlflow_tracking_uri"], + "launcher_not_installed", + ) return _launcher_not_installed(argv) finally: log_file.close() @@ -1278,6 +1744,7 @@ def submit_job_impl( "experiment_id": experiment_id, "stdout_log": str(log_path), "stdout_tail": stdout_tail, + "mlflow": mlflow_info, **_source_result_fields(checkout), "diagnostic": ( "Docker mode launched detached and experiment_id was captured from launcher output." @@ -1304,8 +1771,20 @@ def submit_job_impl( check=False, ) except FileNotFoundError: + if mlflow_info.get("enabled"): + _mark_mlflow_submission_failed( + mlflow_info["mlflow_run_id"], + mlflow_info["mlflow_tracking_uri"], + "launcher_not_installed", + ) return _launcher_not_installed(argv) except subprocess.TimeoutExpired as e: + if mlflow_info.get("enabled"): + _mark_mlflow_submission_failed( + mlflow_info["mlflow_run_id"], + mlflow_info["mlflow_tracking_uri"], + "submission_timeout", + ) return { "ok": False, "executor": "slurm", @@ -1326,6 +1805,12 @@ def submit_job_impl( stderr_tail = str(proc.stderr or "")[-2000:] if proc.returncode != 0 or _launcher_reported_error(stdout_tail, stderr_tail): + if mlflow_info.get("enabled"): + _mark_mlflow_submission_failed( + mlflow_info["mlflow_run_id"], + mlflow_info["mlflow_tracking_uri"], + "launch_py_failed", + ) return { "ok": False, "executor": "slurm", @@ -1346,6 +1831,12 @@ def submit_job_impl( experiment_id, experiment_dir, slurm_job_id = _parse_launcher_submission(stdout_tail) if not experiment_id: + if mlflow_info.get("enabled"): + _mark_mlflow_submission_failed( + mlflow_info["mlflow_run_id"], + mlflow_info["mlflow_tracking_uri"], + "launch_result_unparsed", + ) return { "ok": False, "executor": "slurm", @@ -1382,6 +1873,7 @@ def submit_job_impl( "exit_code": 0, "stdout_tail": stdout_tail, "argv": argv, + "mlflow": mlflow_info, **_source_result_fields(checkout), } @@ -1923,6 +2415,10 @@ def wait_for_experiment_impl( experiment_id: str, timeout_sec: int, poll_interval_sec: int, + *, + finalize_mlflow: bool = False, + mlflow_run_id: str | None = None, + mlflow_tracking_uri: str | None = None, ) -> dict: """Block until ``experiment_id`` reaches a terminal status or the timeout elapses. @@ -1947,7 +2443,26 @@ def wait_for_experiment_impl( # dir doesn't exist. return {**status, "waited_seconds": time.monotonic() - started} if status["status"] in ("done", "failed"): - return {**status, "waited_seconds": time.monotonic() - started} + result = {**status, "waited_seconds": time.monotonic() - started} + if finalize_mlflow: + if not mlflow_run_id: + result["mlflow_finalize"] = { + "ok": False, + "reason": "missing_mlflow_run_id", + "diagnostic": ( + "finalize_mlflow=True requires the mlflow_run_id " + "returned by submit_job(enable_mlflow=True)." + ), + } + else: + result["mlflow_finalize"] = finalize_mlflow_run_impl( + experiment_id=experiment_id, + mlflow_run_id=mlflow_run_id, + tracking_uri=mlflow_tracking_uri, + status=status["status"], + task_count=len(status.get("task_statuses") or {}), + ) + return result if time.monotonic() - started > timeout_sec: return { "ok": False, diff --git a/tools/mcp/modelopt_mcp/server.py b/tools/mcp/modelopt_mcp/server.py index 2506b18a33c..ee99bf36ec6 100644 --- a/tools/mcp/modelopt_mcp/server.py +++ b/tools/mcp/modelopt_mcp/server.py @@ -319,6 +319,38 @@ def submit_job( ) ), ] = False, + enable_mlflow: Annotated[ + bool, + Field( + description=( + "If True, create an MLflow run before submission and pass " + "standard MLFLOW_RUN_ID plus MLFLOW_TRACKING_URI into the " + "launched job environment. User code may log metrics/config " + "to that run; call wait_for_experiment(..., " + "finalize_mlflow=True, mlflow_run_id=) to attach final " + "status and launcher logs." + ) + ), + ] = False, + mlflow_tracking_uri: Annotated[ + str | None, + Field( + description=( + "Optional MLflow tracking URI override. Defaults to " + "MLFLOW_TRACKING_URI, MODELOPT_MCP_MLFLOW_TRACKING_URI, " + "or https://mlflow-modelopt.nvidia.com/." + ) + ), + ] = None, + mlflow_experiment: Annotated[ + str | None, + Field( + description=( + "Optional MLflow experiment name. Defaults to " + "MODELOPT_MCP_MLFLOW_EXPERIMENT or ci-tests." + ) + ), + ] = None, ) -> dict: return bridge.submit_job_impl( yaml_path=yaml_path, @@ -343,6 +375,9 @@ def submit_job( dry_run=dry_run, source_ref=source_ref, source_repo=source_repo, + enable_mlflow=enable_mlflow, + mlflow_tracking_uri=mlflow_tracking_uri, + mlflow_experiment=mlflow_experiment, ) @mcp.tool( @@ -432,11 +467,64 @@ def wait_for_experiment( ), ), ] = 30, + finalize_mlflow: Annotated[ + bool, + Field( + description=( + "If True, finalize a pre-created MLflow run when the " + "experiment reaches done/failed by logging final status " + "and fetching/uploading Nemo/Slurm logs as artifacts. Requires " + "mlflow_run_id." + ) + ), + ] = False, + mlflow_run_id: Annotated[ + str | None, + Field(description="MLflow run id returned by submit_job when enable_mlflow=True."), + ] = None, + mlflow_tracking_uri: Annotated[ + str | None, + Field(description="Optional MLflow tracking URI override for finalization."), + ] = None, ) -> dict: return bridge.wait_for_experiment_impl( experiment_id, timeout_sec, poll_interval_sec, + finalize_mlflow=finalize_mlflow, + mlflow_run_id=mlflow_run_id, + mlflow_tracking_uri=mlflow_tracking_uri, + ) + + @mcp.tool( + name="finalize_mlflow_run", + description=( + "Fetch terminal Nemo/Slurm logs and attach final status plus " + "launcher logs to an MLflow run created by " + "submit_job(enable_mlflow=True). Usually you can use " + "wait_for_experiment(..., finalize_mlflow=True) instead; this " + "tool is useful when the experiment has already reached a " + "terminal state." + ), + ) + def finalize_mlflow_run( + experiment_id: Annotated[ + str, + Field(description="The experiment id returned by submit_job."), + ], + mlflow_run_id: Annotated[ + str, + Field(description="The MLflow run id returned by submit_job."), + ], + mlflow_tracking_uri: Annotated[ + str | None, + Field(description="Optional MLflow tracking URI override."), + ] = None, + ) -> dict: + return bridge.finalize_mlflow_run_impl( + experiment_id=experiment_id, + mlflow_run_id=mlflow_run_id, + tracking_uri=mlflow_tracking_uri, ) @mcp.tool( diff --git a/tools/mcp/pyproject.toml b/tools/mcp/pyproject.toml index 3e30411baa1..6c90e008f36 100644 --- a/tools/mcp/pyproject.toml +++ b/tools/mcp/pyproject.toml @@ -4,6 +4,7 @@ version = "0.1.0" description = "MCP server exposing ModelOpt launcher operations (submit, status, logs) as typed tools for codex / Claude Code agents" requires-python = ">=3.10" dependencies = [ + "mlflow", "mcp>=1.0", "modelopt-launcher", "pyyaml", diff --git a/tools/mcp/tests/test_bridge.py b/tools/mcp/tests/test_bridge.py index 69d57822119..963a1f3b261 100644 --- a/tools/mcp/tests/test_bridge.py +++ b/tools/mcp/tests/test_bridge.py @@ -19,6 +19,8 @@ import json import subprocess +from pathlib import Path +from typing import Any import pytest @@ -618,6 +620,59 @@ def fail_named_temporary_file(*args, **kwargs): assert "disk full" in result["diagnostic"] +def test_submit_job_docker_log_creation_failure_marks_mlflow_failed(monkeypatch, tmp_path): + """Post-MLflow Docker submit failures should not leave submitted runs orphaned.""" + yaml_dir = tmp_path / "examples" + yaml_dir.mkdir() + yaml_path = yaml_dir / "config.yaml" + yaml_path.write_text("job_name: t\npipeline: []\n") + monkeypatch.setenv("MODELOPT_LAUNCHER_EXAMPLES_DIR", str(yaml_dir)) + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path / "nemo")) + monkeypatch.setattr(bridge, "verify_docker_setup_impl", lambda: {"ok": True}) + monkeypatch.setattr( + bridge, + "_create_mlflow_run", + lambda **kwargs: { + "ok": True, + "enabled": True, + "mlflow_run_id": "run-123", + "mlflow_tracking_uri": "https://mlflow.example", + }, + ) + + marked = {} + + def fake_mark(run_id, tracking_uri, reason): + marked.update({"run_id": run_id, "tracking_uri": tracking_uri, "reason": reason}) + + def fail_named_temporary_file(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(bridge, "_mark_mlflow_submission_failed", fake_mark) + monkeypatch.setattr(bridge.tempfile, "NamedTemporaryFile", fail_named_temporary_file) + + result = bridge.submit_job_impl( + yaml_path="config.yaml", + hf_local="/tmp/hf", + cluster_host=None, + cluster_user=None, + identity=None, + job_dir=None, + job_name=None, + extra_overrides=None, + skip_verify=False, + enable_mlflow=True, + ) + + assert result["ok"] is False + assert result["reason"] == "docker_submit_log_unavailable" + assert marked == { + "run_id": "run-123", + "tracking_uri": "https://mlflow.example", + "reason": "docker_submit_log_unavailable", + } + + def test_submit_job_slurm_zero_exit_without_ids_is_failure(monkeypatch, tmp_path): """Slurm submit must not report success when launcher emits no ids.""" yaml_dir = tmp_path / "examples" @@ -706,6 +761,96 @@ def fake_run(argv, **kwargs): assert meta["cluster_user"] == "user" +def test_submit_job_slurm_creates_mlflow_run_and_injects_env(monkeypatch, tmp_path): + """enable_mlflow=True creates a run and passes standard env vars to launcher.""" + yaml_dir = tmp_path / "examples" + yaml_dir.mkdir() + yaml_path = yaml_dir / "config.yaml" + yaml_path.write_text("job_name: t\npipeline: []\n") + monkeypatch.setenv("MODELOPT_LAUNCHER_EXAMPLES_DIR", str(yaml_dir)) + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + monkeypatch.setenv("USER", "alice") + (tmp_path / "experiments" / "cicd" / "cicd_1782173197").mkdir(parents=True) + monkeypatch.setattr(bridge, "verify_slurm_setup_impl", lambda **_: {"ok": True}) + + calls = [] + + class _Info: + run_id = "run-123" + experiment_id = "7" + + class _Run: + info = _Info() + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + class _MLflow: + def set_tracking_uri(self, uri): + calls.append(("tracking_uri", uri)) + + def set_experiment(self, name): + calls.append(("experiment", name)) + + def start_run(self, **kwargs): + calls.append(("start_run", kwargs)) + return _Run() + + def log_params(self, params): + calls.append(("params", params)) + + def set_tags(self, tags): + calls.append(("tags", tags)) + + monkeypatch.setattr(bridge, "_setup_mlflow", lambda: (_MLflow(), None)) + captured = {} + + def fake_run(argv, **kwargs): + captured["env"] = kwargs["env"] + return subprocess.CompletedProcess( + args=argv, + returncode=0, + stdout=( + "Experiment Status for cicd_1782173197\n" + "- Job id: 13049989\n" + 'experiment = run.Experiment.from_id("cicd_1782173197")\n' + ), + stderr="", + ) + + monkeypatch.setattr(subprocess, "run", fake_run) + + result = bridge.submit_job_impl( + yaml_path="config.yaml", + cluster_host="cluster.example.com", + cluster_user="user", + skip_verify=False, + enable_mlflow=True, + mlflow_tracking_uri="https://mlflow.example", + mlflow_experiment="mcp-tests", + ) + + assert result["ok"] is True + assert result["mlflow"]["mlflow_run_id"] == "run-123" + assert ( + result["mlflow"]["mlflow_run_url"] == "https://mlflow.example/#/experiments/7/runs/run-123" + ) + assert captured["env"]["MLFLOW_RUN_ID"] == "run-123" + assert captured["env"]["MLFLOW_TRACKING_URI"] == "https://mlflow.example" + assert ("tracking_uri", "https://mlflow.example") in calls + assert ("experiment", "mcp-tests") in calls + assert ( + "start_run", + {"run_name": "alice/cluster.example.com/examples/config"}, + ) in calls + tags = next(item[1] for item in calls if item[0] == "tags") + assert tags["lifecycle"] == "submitted" + assert tags["managed_by"] == "modelopt-mcp" + + def test_submit_job_slurm_accepts_nmm_cluster_fields(monkeypatch, tmp_path): """nmm-sandbox resolved cluster config maps to launcher overrides and env.""" yaml_dir = tmp_path / "examples" @@ -1392,6 +1537,229 @@ def test_wait_for_experiment_returns_terminal_immediately(tmp_path, monkeypatch) assert result["waited_seconds"] < 1 # didn't actually wait +def test_finalize_mlflow_run_uploads_launcher_logs(tmp_path, monkeypatch): + """Finalization logs final status and uploads local launcher artifacts.""" + exp = tmp_path / "experiments" / "exp_done" + exp.mkdir(parents=True) + (exp / "_DONE").touch() + (exp / "status_task_0.out").write_text("succeeded\n") + (exp / "log_task_0.out").write_text("hello from slurm\n") + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + monkeypatch.setattr( + bridge, + "_fetch_nemo_logs_to_experiment_dir", + lambda **kwargs: {"fetched_artifacts": [], "errors": []}, + ) + + calls: list[tuple[Any, ...]] = [] + + class _Run: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + class _MLflow: + def set_tracking_uri(self, uri): + calls.append(("tracking_uri", uri)) + + def start_run(self, **kwargs): + calls.append(("start_run", kwargs)) + return _Run() + + def log_param(self, key, value): + calls.append(("param", key, value)) + + def log_metric(self, key, value): + calls.append(("metric", key, value)) + + def set_tags(self, tags): + calls.append(("tags", tags)) + + def log_artifact(self, path, artifact_path=None): + calls.append(("artifact", path, artifact_path)) + + monkeypatch.setattr(bridge, "_setup_mlflow", lambda: (_MLflow(), None)) + + result = bridge.finalize_mlflow_run_impl( + experiment_id="exp_done", + mlflow_run_id="run-123", + tracking_uri="https://mlflow.example", + ) + + assert result["ok"] is True + assert result["status"] == "done" + assert set(result["uploaded_artifacts"]) == {"_DONE", "log_task_0.out", "status_task_0.out"} + assert ("start_run", {"run_id": "run-123"}) in calls + assert ("metric", "result", 1.0) in calls + assert result["log_fetch"] == {"fetched_artifacts": [], "errors": []} + artifacts = [item for item in calls if item[0] == "artifact"] + assert all(item[2] == "logs" for item in artifacts) + + +def test_finalize_mlflow_run_fetches_nemo_logs_before_upload(tmp_path, monkeypatch): + """Terminal finalization fetches Nemo logs and uploads the fetched files.""" + exp = tmp_path / "experiments" / "exp_done" + exp.mkdir(parents=True) + (exp / "_DONE").touch() + (exp / "status_task_0.out").write_text("succeeded\n") + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + + fetch_calls = [] + + def fake_fetch_single(*, experiment_id, job_idx, target, timeout_sec=120): + fetch_calls.append( + { + "experiment_id": experiment_id, + "job_idx": job_idx, + "target": target, + "timeout_sec": timeout_sec, + } + ) + target.write_text(f"log for task {job_idx}\n") + return {"ok": True} + + calls: list[tuple[Any, ...]] = [] + + class _Run: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + class _MLflow: + def set_tracking_uri(self, uri): + calls.append(("tracking_uri", uri)) + + def start_run(self, **kwargs): + calls.append(("start_run", kwargs)) + return _Run() + + def log_param(self, key, value): + calls.append(("param", key, value)) + + def log_metric(self, key, value): + calls.append(("metric", key, value)) + + def set_tags(self, tags): + calls.append(("tags", tags)) + + def log_artifact(self, path, artifact_path=None): + calls.append(("artifact", path, artifact_path)) + + monkeypatch.setattr(bridge, "_fetch_single_nemo_log_to_file", fake_fetch_single) + monkeypatch.setattr(bridge, "_setup_mlflow", lambda: (_MLflow(), None)) + + result = bridge.finalize_mlflow_run_impl( + experiment_id="exp_done", + mlflow_run_id="run-123", + tracking_uri="https://mlflow.example", + task_count=2, + ) + + assert result["ok"] is True + assert result["log_fetch"] == { + "fetched_artifacts": ["nemo_logs_task_0.log", "nemo_logs_task_1.log"], + "errors": [], + } + assert (exp / "nemo_logs_task_0.log").read_text() == "log for task 0\n" + assert (exp / "nemo_logs_task_1.log").read_text() == "log for task 1\n" + assert [call["job_idx"] for call in fetch_calls] == [0, 1] + uploaded = {Path(item[1]).name for item in calls if item[0] == "artifact"} + assert {"nemo_logs_task_0.log", "nemo_logs_task_1.log"} <= uploaded + + +def test_fetch_nemo_logs_treats_zero_task_count_as_unknown(tmp_path, monkeypatch): + """An empty local status map should still fetch task 0 after terminal status.""" + exp = tmp_path / "experiments" / "exp_done" + exp.mkdir(parents=True) + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + + calls = [] + + def fake_fetch_single(*, experiment_id, job_idx, target, timeout_sec=120): + calls.append(job_idx) + target.write_text("task 0 log\n") + return {"ok": True} + + monkeypatch.setattr(bridge, "_fetch_single_nemo_log_to_file", fake_fetch_single) + + result = bridge._fetch_nemo_logs_to_experiment_dir( + experiment_id="exp_done", + exp_dir=exp, + task_count=0, + ) + + assert result == {"fetched_artifacts": ["nemo_logs_task_0.log"], "errors": []} + assert calls == [0] + + +def test_finalize_mlflow_run_rejects_invalid_experiment_id(): + """finalize_mlflow_run is an external boundary and validates ids itself.""" + result = bridge.finalize_mlflow_run_impl( + experiment_id="../secret", + mlflow_run_id="run-123", + ) + + assert result["ok"] is False + assert result["reason"] == "invalid_experiment_id" + assert result["mlflow_run_id"] == "run-123" + + +def test_finalize_mlflow_run_rejects_non_terminal_experiment(tmp_path, monkeypatch): + """Standalone finalization should not mark running experiments as finalized.""" + exp = tmp_path / "experiments" / "exp_running" + exp.mkdir(parents=True) + (exp / "status_task_0.out").write_text("running\n") + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + + result = bridge.finalize_mlflow_run_impl( + experiment_id="exp_running", + mlflow_run_id="run-123", + ) + + assert result["ok"] is False + assert result["reason"] == "experiment_not_terminal" + assert result["status"] == "running" + + +def test_wait_for_experiment_can_finalize_mlflow(tmp_path, monkeypatch): + """wait_for_experiment(..., finalize_mlflow=True) delegates finalization at terminal.""" + exp = tmp_path / "experiments" / "exp_done" + exp.mkdir(parents=True) + (exp / "_DONE").touch() + (exp / "status_task_0.out").write_text("succeeded\n") + monkeypatch.setenv("NEMORUN_HOME", str(tmp_path)) + + seen = {} + + def fake_finalize(**kwargs): + seen.update(kwargs) + return {"ok": True, "mlflow_run_id": kwargs["mlflow_run_id"]} + + monkeypatch.setattr(bridge, "finalize_mlflow_run_impl", fake_finalize) + result = bridge.wait_for_experiment_impl( + "exp_done", + timeout_sec=10, + poll_interval_sec=1, + finalize_mlflow=True, + mlflow_run_id="run-123", + mlflow_tracking_uri="https://mlflow.example", + ) + + assert result["ok"] is True + assert result["mlflow_finalize"]["ok"] is True + assert seen == { + "experiment_id": "exp_done", + "mlflow_run_id": "run-123", + "tracking_uri": "https://mlflow.example", + "status": "done", + "task_count": 1, + } + + def test_wait_for_experiment_polls_until_done(tmp_path, monkeypatch): """Spin through running → done.""" exp = tmp_path / "experiments" / "exp_in_flight"