trending/Model-Optimizer
1
0
Fork 0
mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-02 03:14:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix ModelOpt MCP Slurm launcher submit (#1799)

Browse Source 
## Summary
- fix launcher Slurm task annotation patching so nemo-run CLI resolves
`slurm_factory` correctly for task slots
- harden `modelopt-mcp` submit parsing/status resolution and add
regression coverage for launcher false-positive success cases
- add a minimal `nvidia-smi` smoke YAML/script and fix launcher
packaging so source-backed Slurm jobs package required files recursively

## Validation
- `uv run pytest tests/test_core.py -q`
- `uv run pytest tests/test_bridge.py -q`
- dry-run and live-submit validated through the patched local MCP server
on `cw_dfw`
- interactive smoke job succeeded end-to-end (`nvidia-smi` ran
successfully in-container)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added an NVIDIA SMI GPU smoke test (script and minimal Slurm YAML
example) for launcher integration.
* **Bug Fixes**
* Improved detection of fatal launcher errors, including when the
launcher exits with code 0.
* Strengthened Slurm experiment/job identifier parsing and added early
rejection of unsafe experiment IDs, with clearer “unparsed”/failure
behavior.
* Updated sandbox task Slurm config type handling and improved launcher
packaging so examples/common are included consistently.
* **Tests**
* Expanded unit and filesystem-based coverage for parsing/validation,
dry-run fatal stderr handling, and nested experiment directory layouts.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
Signed-off-by: Chenhan D. Yu <chenhany@nvidia.com>
This commit is contained in:
Chenhan D. Yu
2026-06-24 01:05:59 +05:30
committed by GitHub
parent c81210faec
commit 37dbbdac5a
10 changed files with 1126 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/launcher/common/smoke/nvidia_smi.sh
+22
View File
@@ -0,0 +1,22 @@
#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
set -euo pipefail
echo "MODEL_OPT_SMOKE_START"
hostname
nvidia-smi
echo "MODEL_OPT_SMOKE_DONE"
tools/launcher/core.py
+70 -21
View File
@@ -75,9 +75,19 @@ def set_slurm_config_type(cls):
"""Register the SlurmConfig dataclass type used by SandboxTask."""
global _SLURM_CONFIG_TYPE
_SLURM_CONFIG_TYPE = cls
# Patch SandboxTask's type annotation so nemo-run's CLI parser can resolve factories
SandboxTask.__dataclass_fields__["slurm_config"].type = cls
SandboxTask.__annotations__["slurm_config"] = cls
# Patch every task dataclass so nemo-run's CLI parser sees the concrete
# SlurmConfig type for task_0/task_1/... fields, not the base `object`.
for task_cls in (
SandboxTask,
SandboxTask0,
SandboxTask1,
SandboxTask2,
SandboxTask3,
SandboxTask4,
):
task_cls.__dataclass_fields__["slurm_config"].type = cls
task_cls.__annotations__["slurm_config"] = cls
task_cls.__init__.__annotations__["slurm_config"] = cls
def register_factory(name, fn):
@@ -386,25 +396,64 @@ def build_docker_executor(
def _git_info(path):
"""Get git commit hash and branch for a directory."""
import subprocess # nosec B404
try:
commit = subprocess.run( # nosec B603 B607
["git", "rev-parse", "--short", "HEAD"],
cwd=path,
capture_output=True,
text=True,
timeout=5,
).stdout.strip()
branch = subprocess.run( # nosec B603 B607
["git", "rev-parse", "--abbrev-ref", "HEAD"],
cwd=path,
capture_output=True,
text=True,
timeout=5,
).stdout.strip()
return commit, branch
except Exception:
worktree_dir = os.path.abspath(path)
while True:
git_path = os.path.join(worktree_dir, ".git")
if os.path.isdir(git_path):
git_dir = git_path
break
if os.path.isfile(git_path):
with open(git_path, encoding="utf-8") as file:
marker = file.read().strip()
if not marker.startswith("gitdir:"):
return "unknown", "unknown"
git_dir = marker.removeprefix("gitdir:").strip()
if not os.path.isabs(git_dir):
git_dir = os.path.normpath(os.path.join(worktree_dir, git_dir))
break
parent = os.path.dirname(worktree_dir)
if parent == worktree_dir:
return "unknown", "unknown"
worktree_dir = parent
common_dir = git_dir
commondir_path = os.path.join(git_dir, "commondir")
if os.path.exists(commondir_path):
with open(commondir_path, encoding="utf-8") as file:
common_dir = file.read().strip()
if not os.path.isabs(common_dir):
common_dir = os.path.normpath(os.path.join(git_dir, common_dir))
with open(os.path.join(git_dir, "HEAD"), encoding="utf-8") as file:
head = file.read().strip()
if not head.startswith("ref:"):
return head[:7], "HEAD"
ref = head.removeprefix("ref:").strip()
branch = ref.removeprefix("refs/heads/")
commit = ""
for refs_dir in (git_dir, common_dir):
ref_path = os.path.join(refs_dir, *ref.split("/"))
if os.path.exists(ref_path):
with open(ref_path, encoding="utf-8") as file:
commit = file.read().strip()
break
if not commit:
packed_refs = os.path.join(common_dir, "packed-refs")
if os.path.exists(packed_refs):
with open(packed_refs, encoding="utf-8") as file:
for line in file:
if line.startswith(("#", "^")):
continue
sha, _, packed_ref = line.strip().partition(" ")
if packed_ref == ref:
commit = sha
break
return (commit[:7] if commit else "unknown"), branch
except OSError:
return "unknown", "unknown"
tools/launcher/examples/smoke/nvidia_smi.yaml
+21
View File
@@ -0,0 +1,21 @@
# Minimal Slurm smoke test for launcher/MCP integration.
#
# This intentionally avoids model downloads and HF cache mounts. It verifies:
# MCP submit_job -> launcher YAML parse -> Slurm submit -> container start -> GPU visibility.
job_name: nvidia_smi_smoke
pipeline:
skip: false
allow_to_fail: false
note: "Slurm container GPU smoke test"
task_0:
script: common/smoke/nvidia_smi.sh
slurm_config:
_factory_: "slurm_factory"
nodes: 1
ntasks_per_node: 1
gpus_per_node: 1
time: "00:10:00"
container: nvcr.io/nvidia/cuda:12.4.1-base-ubuntu22.04
container_mounts: []
tools/launcher/launch.py
+32 -14
View File
@@ -30,8 +30,9 @@ Environment variables:
"""
import getpass
import glob
import os
import subprocess # nosec B404
import subprocess # nosec B404 - required for explicit git clean command; no shell is used.
import warnings
import modelopt_launcher as _pkg
@@ -79,20 +80,33 @@ _modelopt_src = os.path.join(LAUNCHER_DIR, "modules", "Model-Optimizer", "modelo
EXPERIMENT_TITLE = "cicd"
DEFAULT_SLURM_ENV, DEFAULT_LOCAL_ENV = get_default_env(EXPERIMENT_TITLE)
_include_pattern = ["examples/*", "common/*"]
_relative_path = [LAUNCHER_DIR, LAUNCHER_DIR]
_include_pattern = []
_relative_path = []
def _add_package_path(path: str) -> None:
"""Add an existing package path using LAUNCHER_DIR as the tar root."""
if os.path.exists(path):
_include_pattern.append(path)
_relative_path.append(LAUNCHER_DIR)
def _add_package_glob(pattern: str) -> None:
"""Expand a glob and add each matching path to the launcher package."""
for path in sorted(glob.glob(pattern)):
_add_package_path(path)
_add_package_path(os.path.join(LAUNCHER_DIR, "examples"))
_add_package_path(os.path.join(LAUNCHER_DIR, "common"))
if _has_modelopt_src:
_include_pattern = [
"modules/Megatron-LM/megatron/*",
"modules/Megatron-LM/examples/*",
"modules/Megatron-LM/*.py",
"modules/Model-Optimizer/modelopt/*",
"modules/Model-Optimizer/modelopt_recipes/*",
"modules/Model-Optimizer/examples/*",
*_include_pattern,
]
_relative_path = [LAUNCHER_DIR] * 6 + _relative_path
_add_package_path(os.path.join(LAUNCHER_DIR, "modules/Megatron-LM/megatron"))
_add_package_path(os.path.join(LAUNCHER_DIR, "modules/Megatron-LM/examples"))
_add_package_glob(os.path.join(LAUNCHER_DIR, "modules/Megatron-LM/*.py"))
_add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/modelopt"))
_add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/modelopt_recipes"))
_add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/examples"))
packager = run.PatternPackager(
include_pattern=_include_pattern,
@@ -127,7 +141,11 @@ def launch(
raise ValueError("--clean requires a dev checkout; modelopt source not found.")
examples_dir = os.path.join(_mo_symlink, "examples")
print(f"Cleaning {examples_dir} with git clean -xdf ...")
subprocess.run(["git", "clean", "-xdf", "."], cwd=examples_dir, check=True) # nosec B603 B607
subprocess.run( # nosec B603 B607 - fixed git CLI argv; no shell.
["git", "clean", "-xdf", "."],
cwd=examples_dir,
check=True,
)
if "NEMORUN_HOME" not in os.environ:
warnings.warn("NEMORUN_HOME is not set. Defaulting to current working directory.")
tools/launcher/tests/test_core.py
+14 -2
View File
@@ -35,6 +35,9 @@ from core import (
SandboxTask,
SandboxTask0,
SandboxTask1,
SandboxTask2,
SandboxTask3,
SandboxTask4,
create_task_from_yaml,
get_default_env,
register_factory,
@@ -185,8 +188,17 @@ class TestSetSlurmConfigType:
host: str = "test"
set_slurm_config_type(MockSlurmConfig)
assert SandboxTask.__annotations__["slurm_config"] is MockSlurmConfig
assert SandboxTask.__dataclass_fields__["slurm_config"].type is MockSlurmConfig
for task_cls in (
SandboxTask,
SandboxTask0,
SandboxTask1,
SandboxTask2,
SandboxTask3,
SandboxTask4,
):
assert task_cls.__annotations__["slurm_config"] is MockSlurmConfig
assert task_cls.__dataclass_fields__["slurm_config"].type is MockSlurmConfig
assert task_cls.__init__.__annotations__["slurm_config"] is MockSlurmConfig
class TestGetDefaultEnv:
tools/launcher/tests/test_core_extended.py
+6
View File
@@ -129,6 +129,12 @@ class TestGitInfo:
assert branch != "unknown"
assert len(commit) >= 7 # short hash
def test_valid_git_repo_from_nested_directory(self):
commit, branch = _git_info(os.path.join(os.getcwd(), "tests"))
assert commit != "unknown"
assert branch != "unknown"
assert len(commit) >= 7 # short hash
def test_nonexistent_directory(self):
commit, branch = _git_info("/tmp/nonexistent_xyz_12345")
assert commit == "unknown"
tools/mcp/README.md
+27 -3
View File
@@ -21,7 +21,7 @@ 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?)` | Submit a launcher YAML. Mode resolved from mutually-exclusive args. Returns `experiment_id` (Slurm) or PID (Docker) immediately; 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}` 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?)` | 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 (Docker) immediately; 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, ...). |
| `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_<task>.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`. |
@@ -47,7 +47,7 @@ codex mcp add modelopt -- uvx --from \
modelopt-mcp
```
`uvx` clones the whole repo to its cache, installs `tools/mcp/` as the entry point, and resolves the sibling `modelopt-launcher` dep via `[tool.uv.sources]` (path → `../launcher`) inside the cloned tree.
`uvx` clones the whole repo to its cache, installs `tools/mcp/` as the entry point, and resolves the sibling `modelopt-launcher` dep via `[tool.uv.sources]` (path → `../launcher`) inside the cloned tree. That install clone is only the server runtime; job submission uses the managed source checkout described below.
### Dev install (local checkout)
@@ -59,6 +59,24 @@ modelopt-mcp # stdio server entry on PATH
Both packages share the launcher's `core.py` orchestrator. The dev path relies on `[tool.uv.sources]` to point `modelopt-launcher` at `../launcher`.
## Managed source checkouts
`submit_job` does not rely on the uvx install clone or on the caller being inside a Model-Optimizer checkout. For each launch it resolves:
1. `source_ref` argument, if provided.
2. `MODELOPT_MCP_SOURCE_REF`, if set.
3. `main`.
It resolves that ref against `source_repo` / `MODELOPT_MCP_SOURCE_REPO` / `https://github.com/NVIDIA/Model-Optimizer.git`, creates a cached checkout under `MODELOPT_MCP_SOURCE_CACHE` (default `$XDG_CACHE_HOME/modelopt-mcp/sources` or `~/.cache/modelopt-mcp/sources`), and runs:
```bash
uv run --project <source_root>/tools/launcher modelopt-launcher --yaml <resolved-yaml> ...
```
The checkout is keyed by resolved commit SHA, so multiple agents using different branches or SHAs get separate source roots. Recursive submodules are initialized in the managed checkout, so launcher packagers can include `tools/launcher/modules/...` content even when MCP was installed outside a repo checkout.
Set `MODELOPT_MCP_DISABLE_MANAGED_SOURCE=1` only for local development when you deliberately want the already-installed `modelopt-launcher` entrypoint.
### Why no plain `pip install` today
`modelopt-mcp` and `modelopt-launcher` are not on PyPI. Plain `pip` doesn't read `[tool.uv.sources]`, so even from a local checkout, `pip install -e tools/mcp` fails to resolve the bare `modelopt-launcher` name. Stick with `uv` / `uvx` while we're git-only.
@@ -96,8 +114,9 @@ result = mcp__modelopt__submit_job(
cluster_user="alice",
identity="/home/alice/.ssh/id_ed25519",
skip_verify=True, # we just probed
source_ref="main", # optional; omit to use main
)
# {"ok": True, "experiment_id": "cicd_1781240000", "slurm_job_id": "12345", ...}
# {"ok": True, "experiment_id": "cicd_1781240000", "slurm_job_id": "12345", "source_sha": "...", ...}
# 4. Poll until done
while True:
@@ -122,6 +141,11 @@ For local Docker execution, drop `cluster_host`/`cluster_user`/`identity` and pa
| `NEMORUN_HOME` | submit + status + logs | Where the launcher writes experiment artifacts. Defaults to cwd if unset. `job_status` / `job_logs` search `$NEMORUN_HOME/experiments/<id>/`. |
| `MODELOPT_MCP_LOG` | (optional) server | Log level. Defaults to `INFO`. Logs go to stderr — stdout is the MCP wire. |
| `MODELOPT_MCP_SKIP_GPU_CHECK` | (optional) `verify_setup(executor='docker')` | Set to skip the `docker info --format` runtime-registry check. Useful for CI hosts where the daemon is up but the NVIDIA Container Toolkit isn't installed. |
| `MODELOPT_MCP_SOURCE_REPO` | (optional) `submit_job` | Default git repository for managed source checkouts. Defaults to `https://github.com/NVIDIA/Model-Optimizer.git`. |
| `MODELOPT_MCP_SOURCE_REF` | (optional) `submit_job` | Default branch, tag, or SHA when `source_ref` is omitted. Defaults to `main`. |
| `MODELOPT_MCP_SOURCE_CACHE` | (optional) `submit_job` | Root for managed source checkouts. Defaults to `$XDG_CACHE_HOME/modelopt-mcp/sources` or `~/.cache/modelopt-mcp/sources`. |
| `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 <source>/tools/launcher ...`. |
| `MODELOPT_LAUNCHER_EXAMPLES_DIR` | (optional) `list_examples` | Override the examples directory location. Defaults to `../launcher/examples/` relative to this package. |
## Design principles
tools/mcp/modelopt_mcp/bridge.py
+535 -55
View File
@@ -33,9 +33,11 @@ need — keeps the surface area auditable.
from __future__ import annotations
import hashlib
import os
import re
import subprocess # nosec B404
import shutil
import subprocess # nosec B404 - fixed-argv CLI probes are required; shell=True is not used.
import time
from dataclasses import dataclass
from pathlib import Path
@@ -47,12 +49,65 @@ import yaml
# uvx-from-git installs (the launcher is a sibling site-packages install).
_THIS_DIR = Path(__file__).resolve().parent
_DEFAULT_SOURCE_REPO = "https://github.com/NVIDIA/Model-Optimizer.git"
_DEFAULT_SOURCE_REF = "main"
# Canonical task-status failure tokens — matched against the FIRST word
# of each ``status_<task>.out`` file by ``job_status_impl``.
_STATUS_FAILURE_WORDS: frozenset[str] = frozenset(
{"failed", "error", "errored", "cancelled", "canceled"}
)
_SAFE_EXPERIMENT_ID_RE = re.compile(r"^[A-Za-z0-9_-]+$")
_LAUNCHER_ERROR_RE = re.compile(
r"(?:^|\n)(?:Unexpected error:|Error processing argument )",
re.IGNORECASE,
)
_GIT_SHA_RE = re.compile(r"^[0-9a-fA-F]{7,40}$")
_SAFE_PATH_TOKEN_RE = re.compile(r"[^A-Za-z0-9_.-]+")
@dataclass
class SourceCheckout:
"""A managed Model-Optimizer checkout used for one launcher invocation."""
repo: str
ref: str
resolved_sha: str
root: Path
@property
def launcher_dir(self) -> Path:
"""Return the launcher package directory inside this checkout."""
return self.root / "tools" / "launcher"
@property
def examples_dir(self) -> Path:
"""Return the launcher examples directory inside this checkout."""
return self.launcher_dir / "examples"
def _launcher_reported_error(stdout: str, stderr: str) -> bool:
"""Return True when launcher text contains a fatal error despite exit 0."""
return bool(_LAUNCHER_ERROR_RE.search(f"{stdout}\n{stderr}"))
def _validate_experiment_id(experiment_id: str) -> dict | None:
"""Reject experiment ids that could escape path joins or alter glob matching."""
if _SAFE_EXPERIMENT_ID_RE.fullmatch(experiment_id):
return None
return {
"ok": False,
"experiment_id": experiment_id,
"reason": "invalid_experiment_id",
"diagnostic": (
"experiment_id must be a single path-safe token containing "
"only letters, numbers, underscores, and hyphens."
),
}
def _find_launcher_examples_dir() -> Path | None:
"""Resolve the launcher examples directory.
@@ -83,8 +138,34 @@ def _find_launcher_examples_dir() -> Path | None:
return None
def _find_launcher_package_dir() -> Path | None:
"""Resolve the installed launcher's package directory."""
try:
import modelopt_launcher
candidate = Path(modelopt_launcher.PACKAGE_DIR)
if candidate.exists():
return candidate
except ImportError:
pass
return None
def _launcher_not_installed(argv: list[str]) -> dict:
"""Structured failure when the ``modelopt-launcher`` binary is not on PATH."""
if argv and argv[0] == _uv_binary():
return {
"ok": False,
"reason": "uv_not_installed",
"diagnostic": (
"`uv` was not found on PATH. Managed Model-Optimizer source "
"checkouts use `uv run --project <checkout>/tools/launcher "
"modelopt-launcher ...`. Install uv or set "
"MODELOPT_MCP_DISABLE_MANAGED_SOURCE=1 to use the installed "
"`modelopt-launcher` entrypoint directly."
),
"argv": argv,
}
return {
"ok": False,
"reason": "launcher_not_installed",
@@ -97,6 +178,301 @@ def _launcher_not_installed(argv: list[str]) -> dict:
}
# ---------------------------------------------------------------------------
# Managed Model-Optimizer source checkouts
# ---------------------------------------------------------------------------
def _uv_binary() -> str:
"""Return the uv executable used for managed-source launcher runs."""
return os.environ.get("MODELOPT_MCP_UV", "uv")
def _source_cache_root() -> Path:
"""Return the root directory for MCP-managed source checkouts."""
env = os.environ.get("MODELOPT_MCP_SOURCE_CACHE")
if env:
return Path(env).expanduser()
xdg_cache = os.environ.get("XDG_CACHE_HOME")
base = Path(xdg_cache).expanduser() if xdg_cache else Path.home() / ".cache"
return base / "modelopt-mcp" / "sources"
def _source_disabled() -> bool:
"""Return True when callers explicitly opt out of managed source checkouts."""
return os.environ.get("MODELOPT_MCP_DISABLE_MANAGED_SOURCE", "").lower() in {
"1",
"true",
"yes",
"on",
}
def _sanitize_path_token(value: str, *, fallback: str) -> str:
"""Make a short, filesystem-safe display token."""
token = _SAFE_PATH_TOKEN_RE.sub("-", value.strip()).strip(".-")
return (token or fallback)[:48]
def _tail(text: str | None, limit: int = 1200) -> str:
"""Return a short tail suitable for structured diagnostics."""
return str(text or "")[-limit:]
def _git_failure(
*,
reason: str,
diagnostic: str,
argv: list[str],
proc: subprocess.CompletedProcess | None = None,
) -> dict:
"""Return a structured managed-source git failure."""
result = {
"ok": False,
"reason": reason,
"diagnostic": diagnostic,
"argv": argv,
}
if proc is not None:
result.update(
{
"exit_code": proc.returncode,
"stdout_tail": _tail(proc.stdout),
"stderr_tail": _tail(proc.stderr),
}
)
return result
def _run_git(argv: list[str], *, cwd: Path | None = None, timeout: int = 300) -> dict:
"""Run a fixed git argv list and return either proc or a structured failure."""
try:
proc = subprocess.run( # nosec B603 B607 - fixed git argv list; no shell.
argv,
cwd=str(cwd) if cwd else None,
capture_output=True,
text=True,
timeout=timeout,
check=False,
)
except FileNotFoundError:
return _git_failure(
reason="git_not_installed",
diagnostic="`git` was not found on PATH; cannot prepare the managed source checkout.",
argv=argv,
)
except subprocess.TimeoutExpired as e:
return {
"ok": False,
"reason": "git_timeout",
"diagnostic": f"`{' '.join(argv)}` did not finish within {timeout}s.",
"argv": argv,
"stdout_tail": (e.stdout or b"").decode(errors="replace")[-1200:]
if isinstance(e.stdout, bytes)
else _tail(e.stdout),
"stderr_tail": (e.stderr or b"").decode(errors="replace")[-1200:]
if isinstance(e.stderr, bytes)
else _tail(e.stderr),
}
if proc.returncode != 0:
return _git_failure(
reason="git_failed",
diagnostic=f"`{' '.join(argv)}` failed while preparing the managed source checkout.",
argv=argv,
proc=proc,
)
return {"ok": True, "proc": proc}
def _resolve_source_ref(repo: str, ref: str) -> dict:
"""Resolve a branch/tag/ref to a commit SHA without mutating local state."""
if _GIT_SHA_RE.fullmatch(ref):
return {"ok": True, "resolved_sha": ref.lower()}
patterns = [ref]
if not ref.startswith("refs/"):
patterns.extend([f"refs/heads/{ref}", f"refs/tags/{ref}", f"refs/tags/{ref}^{{}}"])
argv = ["git", "ls-remote", repo, *patterns]
result = _run_git(argv, timeout=60)
if not result.get("ok"):
return {
**result,
"reason": "source_ref_resolve_failed",
"diagnostic": (
f"Could not resolve Model-Optimizer source ref {ref!r} from {repo}. "
"Check the branch/tag/SHA and network credentials."
),
}
lines = [line.split() for line in result["proc"].stdout.splitlines() if line.strip()]
by_name = {name: sha for sha, name, *_ in lines if len(sha) == 40}
for name in (
f"refs/heads/{ref}",
f"refs/tags/{ref}^{{}}",
f"refs/tags/{ref}",
ref,
):
sha = by_name.get(name)
if sha:
return {"ok": True, "resolved_sha": sha}
for sha, *_ in lines:
if len(sha) == 40:
return {"ok": True, "resolved_sha": sha}
return {
"ok": False,
"reason": "source_ref_not_found",
"diagnostic": f"Model-Optimizer source ref {ref!r} was not found in {repo}.",
"argv": argv,
"stdout_tail": result["proc"].stdout[-1200:],
"stderr_tail": result["proc"].stderr[-1200:],
}
def _checkout_path(repo: str, ref: str, resolved_sha: str) -> Path:
"""Return the immutable checkout path for a resolved source ref."""
repo_hash = hashlib.sha256(repo.encode()).hexdigest()[:12]
ref_token = _sanitize_path_token(ref, fallback="ref")
sha_token = resolved_sha[:12]
return _source_cache_root() / repo_hash / f"{ref_token}-{sha_token}"
def _checkout_ready(path: Path, resolved_sha: str) -> bool:
"""Return True when a managed checkout already exists at the requested SHA."""
if not (path / ".git").exists() or not (path / "tools" / "launcher" / "launch.py").exists():
return False
result = _run_git(["git", "-C", str(path), "rev-parse", "HEAD"], timeout=30)
return bool(result.get("ok") and result["proc"].stdout.strip().startswith(resolved_sha))
def _materialize_checkout(repo: str, ref: str, resolved_sha: str, path: Path) -> dict:
"""Clone Model-Optimizer and initialize submodules for a resolved ref."""
parent = path.parent
parent.mkdir(parents=True, exist_ok=True)
tmp = (
parent / f".tmp-{_sanitize_path_token(ref, fallback='ref')}-{os.getpid()}-{time.time_ns()}"
)
if tmp.exists():
shutil.rmtree(tmp)
clone = ["git", "clone", "--no-checkout", "--filter=blob:none", repo, str(tmp)]
fetch_refs = [resolved_sha] if _GIT_SHA_RE.fullmatch(ref) else [ref, resolved_sha]
post_fetch_steps = [
["git", "-C", str(tmp), "checkout", "--detach", "FETCH_HEAD"],
["git", "-C", str(tmp), "submodule", "sync", "--recursive"],
["git", "-C", str(tmp), "submodule", "update", "--init", "--recursive", "--depth=1"],
]
try:
result = _run_git(clone)
if not result.get("ok"):
return _source_checkout_failure(result, repo, ref, resolved_sha, path)
fetch_result = None
for fetch_ref in fetch_refs:
fetch = ["git", "-C", str(tmp), "fetch", "--depth=1", "origin", fetch_ref]
fetch_result = _run_git(fetch)
if fetch_result.get("ok"):
break
if fetch_result is None or not fetch_result.get("ok"):
return _source_checkout_failure(fetch_result or {}, repo, ref, resolved_sha, path)
for argv in post_fetch_steps:
result = _run_git(argv)
if not result.get("ok"):
return _source_checkout_failure(result, repo, ref, resolved_sha, path)
if path.exists():
if _checkout_ready(path, resolved_sha):
shutil.rmtree(tmp)
else:
shutil.rmtree(path)
tmp.rename(path)
else:
tmp.rename(path)
finally:
if tmp.exists():
shutil.rmtree(tmp)
return {"ok": True}
def _source_checkout_failure(
result: dict,
repo: str,
ref: str,
resolved_sha: str,
path: Path,
) -> dict:
"""Attach source provenance to a failed checkout step."""
return {
**result,
"ok": False,
"reason": "source_checkout_failed",
"diagnostic": (
"Failed to prepare the managed Model-Optimizer source checkout. "
"The launcher was not run."
),
"source_repo": repo,
"source_ref": ref,
"source_sha": resolved_sha,
"source_root": str(path),
}
def _ensure_source_checkout(
source_ref: str | None = None,
source_repo: str | None = None,
) -> dict:
"""Return a managed source checkout, or None when explicitly disabled."""
if _source_disabled():
return {"ok": True, "checkout": None}
repo = source_repo or os.environ.get("MODELOPT_MCP_SOURCE_REPO") or _DEFAULT_SOURCE_REPO
ref = source_ref or os.environ.get("MODELOPT_MCP_SOURCE_REF") or _DEFAULT_SOURCE_REF
resolved = _resolve_source_ref(repo, ref)
if not resolved.get("ok"):
return {**resolved, "source_repo": repo, "source_ref": ref}
resolved_sha = resolved["resolved_sha"]
path = _checkout_path(repo, ref, resolved_sha)
if not _checkout_ready(path, resolved_sha):
materialized = _materialize_checkout(repo, ref, resolved_sha, path)
if not materialized.get("ok"):
return materialized
checkout = SourceCheckout(repo=repo, ref=ref, resolved_sha=resolved_sha, root=path)
return {"ok": True, "checkout": checkout}
def _source_result_fields(checkout: SourceCheckout | None) -> dict:
"""Return source provenance fields for tool results."""
if checkout is None:
return {}
return {
"source_repo": checkout.repo,
"source_ref": checkout.ref,
"source_sha": checkout.resolved_sha,
"source_root": str(checkout.root),
}
def _launcher_argv(abs_yaml: Path, checkout: SourceCheckout | None, *flags: str) -> list[str]:
"""Build the launcher argv for installed or managed-source execution."""
if checkout is None:
return ["modelopt-launcher", "--yaml", str(abs_yaml), *flags]
return [
_uv_binary(),
"run",
"--project",
str(checkout.launcher_dir),
"modelopt-launcher",
"--yaml",
str(abs_yaml),
*flags,
]
# ---------------------------------------------------------------------------
# list_examples
# ---------------------------------------------------------------------------
@@ -196,7 +572,7 @@ def verify_docker_setup_impl() -> dict:
# invoking the docker CLI by name with a fixed argv list, no
# shell-interpretation, no untrusted input.
try:
proc = subprocess.run( # nosec B603 B607
proc = subprocess.run( # nosec B603 B607 - fixed docker CLI argv; no shell.
["docker", "info"],
capture_output=True,
text=True,
@@ -250,7 +626,7 @@ def verify_docker_setup_impl() -> dict:
# runtime when nvidia-ctk runtime configure was last invoked.
# B603/B607 same false-positive shape as daemon check.
try:
gpu = subprocess.run( # nosec B603 B607
gpu = subprocess.run( # nosec B603 B607 - fixed docker CLI argv; no shell.
["docker", "info", "--format", "{{json .}}"],
capture_output=True,
text=True,
@@ -325,7 +701,7 @@ def verify_slurm_setup_impl(
# B603/B607 false positive — `ssh` invoked by name with a controlled
# argv (BatchMode, ConnectTimeout, identity path, target). No shell.
try:
proc = subprocess.run( # nosec B603 B607
proc = subprocess.run( # nosec B603 B607 - fixed ssh CLI argv; no shell.
argv,
capture_output=True,
text=True,
@@ -383,13 +759,14 @@ def verify_slurm_setup_impl(
# ---------------------------------------------------------------------------
def _normalize_yaml_path(yaml_path: str) -> Path:
def _normalize_yaml_path(yaml_path: str, *, examples_dir: Path | None = None) -> Path:
"""Resolve a launcher YAML path to an absolute Path.
Lookup order:
1. Absolute path — use as-is
2. Relative to ``MODELOPT_LAUNCHER_EXAMPLES_DIR`` (or its parent)
3. Relative to cwd
2. Relative to the managed checkout's examples dir, when present
3. Relative to ``MODELOPT_LAUNCHER_EXAMPLES_DIR`` (or its parent)
4. Relative to cwd
The double-fallback lets the agent pass either ``examples/Qwen/.../X.yaml``
or just the absolute path.
@@ -397,13 +774,18 @@ def _normalize_yaml_path(yaml_path: str) -> Path:
p = Path(yaml_path)
if p.is_absolute():
return p
# Look under examples dir
examples_dir = _find_launcher_examples_dir()
# Look under a managed checkout first, then the installed examples dir.
examples_dirs: list[Path] = []
if examples_dir is not None:
candidate = examples_dir / yaml_path
examples_dirs.append(examples_dir)
installed_examples_dir = _find_launcher_examples_dir()
if installed_examples_dir is not None and installed_examples_dir not in examples_dirs:
examples_dirs.append(installed_examples_dir)
for root in examples_dirs:
candidate = root / yaml_path
if candidate.exists():
return candidate
candidate = examples_dir.parent / yaml_path
candidate = root.parent / yaml_path
if candidate.exists():
return candidate
# cwd fallback
@@ -422,6 +804,8 @@ def submit_job_impl(
extra_overrides: dict[str, str] | None,
skip_verify: bool,
dry_run: bool = False,
source_ref: str | None = None,
source_repo: str | None = None,
) -> dict:
"""Submit a launcher YAML.
@@ -454,6 +838,8 @@ def submit_job_impl(
job_dir=job_dir,
job_name=job_name,
extra_overrides=extra_overrides,
source_ref=source_ref,
source_repo=source_repo,
)
# ---- Mode resolution -------------------------------------------
@@ -501,14 +887,23 @@ def submit_job_impl(
"verify_result": check,
}
# ---- Resolve the YAML path ------------------------------------
abs_yaml = _normalize_yaml_path(yaml_path)
# ---- Resolve source + YAML path -------------------------------
source = _ensure_source_checkout(source_ref=source_ref, source_repo=source_repo)
if not source.get("ok"):
return source
checkout: SourceCheckout | None = source["checkout"]
abs_yaml = _normalize_yaml_path(
yaml_path,
examples_dir=checkout.examples_dir if checkout else None,
)
if not abs_yaml.exists():
return {
"ok": False,
"reason": "yaml_not_found",
"yaml_path": yaml_path,
"resolved_path": str(abs_yaml),
**_source_result_fields(checkout),
"diagnostic": (
f"YAML not found at {abs_yaml}. Pass a path under "
f"tools/launcher/examples/ (relative), an absolute path, "
@@ -527,7 +922,7 @@ def submit_job_impl(
# list never goes through a shell, so quoting bakes literal quote chars
# into the values that nemo-run's CLI parser sees. Verbatim values
# carry spaces / special chars safely.
argv = ["modelopt-launcher", "--yaml", str(abs_yaml), "--yes"]
argv = _launcher_argv(abs_yaml, checkout, "--yes")
if hf_local:
argv.append(f"hf_local={hf_local}")
else:
@@ -554,6 +949,10 @@ def submit_job_impl(
# experiment_dir_not_found for jobs that actually succeeded.
child_env = os.environ.copy()
child_env.setdefault("NEMORUN_HOME", os.getcwd())
if checkout is not None:
child_env["MODELOPT_MCP_SOURCE_ROOT"] = str(checkout.root)
child_env["MODELOPT_MCP_SOURCE_REF"] = checkout.ref
child_env["MODELOPT_MCP_SOURCE_SHA"] = checkout.resolved_sha
if executor == "slurm":
# Required for slurm_factory's host default. Verify_setup ran
# against this same host above (when verify_setup=True), so the
@@ -571,7 +970,7 @@ def submit_job_impl(
# in-flight launcher.
# B603 false positive — argv is a controlled list built above.
try:
proc = subprocess.Popen( # nosec B603
proc = subprocess.Popen( # nosec B603 - fixed launcher argv list; no shell.
argv,
env=child_env,
stdout=subprocess.DEVNULL,
@@ -587,6 +986,7 @@ def submit_job_impl(
"argv": argv,
"nemorun_home": child_env["NEMORUN_HOME"],
"experiment_id": None, # Phase 2: tail launcher's output
**_source_result_fields(checkout),
# via a side-channel log file to capture nemo_run's id
"spike_note": (
"Docker mode launched detached. Phase 1: experiment_id "
@@ -599,7 +999,7 @@ def submit_job_impl(
# with detach=true). Capture stdout to parse experiment_id.
# B603 false positive — argv is a controlled list built above.
try:
proc = subprocess.run( # nosec B603
proc = subprocess.run( # nosec B603 - fixed launcher argv list; no shell.
argv,
env=child_env,
capture_output=True,
@@ -620,6 +1020,7 @@ def submit_job_impl(
f"{(e.stdout or b'').decode(errors='replace')[-400:]}"
),
"argv": argv,
**_source_result_fields(checkout),
}
# `proc` here is the CompletedProcess from subprocess.run with
@@ -628,7 +1029,7 @@ def submit_job_impl(
stdout_tail = str(proc.stdout or "")[-2000:]
stderr_tail = str(proc.stderr or "")[-2000:]
if proc.returncode != 0:
if proc.returncode != 0 or _launcher_reported_error(stdout_tail, stderr_tail):
return {
"ok": False,
"executor": "slurm",
@@ -637,11 +1038,13 @@ def submit_job_impl(
"stdout_tail": stdout_tail,
"stderr_tail": stderr_tail,
"diagnostic": (
f"launch.py exited with code {proc.returncode}. Common "
f"causes: SSH publickey rejection, malformed YAML, "
f"NEMORUN_HOME unset. Inspect stderr_tail."
"launch.py failed or printed a fatal launcher error. "
"Common causes: SSH publickey rejection, malformed YAML, "
"factory parsing failure, or NEMORUN_HOME unset. Inspect "
"stdout_tail/stderr_tail."
),
"argv": argv,
**_source_result_fields(checkout),
}
# Best-effort experiment_id + dir + slurm_job_id parse. nemo_run's
@@ -654,14 +1057,26 @@ def submit_job_impl(
# nemo_run prints "Entering Experiment <title>_<id> with id: <id>" —
# match the trailing id directly so we don't have to encode the
# title prefix or hard-code timestamp width.
m = re.search(
r"experiment[_\s-]+id[:\s]+(\S+)",
stdout_tail,
re.IGNORECASE,
)
m = re.search(r'Experiment\.from_id\("([^"]+)"\)', stdout_tail)
if m:
experiment_id = m.group(1)
else:
m = re.search(
r"Experiment Status for\s+(\S+)",
stdout_tail,
re.IGNORECASE,
)
if m:
experiment_id = m.group(1)
if not experiment_id:
m = re.search(
r"experiment[_\s-]+id[:\s]+(\S+)",
stdout_tail,
re.IGNORECASE,
)
if m:
experiment_id = m.group(1)
if not experiment_id:
# Fallback for older nemo_run output that lacked the explicit
# "id:" label. Accepts any path-safe id token following the
# word "experiment" — not just timestamp-style.
@@ -670,7 +1085,7 @@ def submit_job_impl(
stdout_tail,
re.IGNORECASE,
)
if m:
if m and m.group(1).lower() not in {"status"}:
experiment_id = m.group(1)
# Match any path containing `/experiments/<id>/` — don't anchor on
# cluster-specific filesystem roots (NVIDIA's /lustre, partner
@@ -681,6 +1096,28 @@ def submit_job_impl(
m = re.search(r"Submitted batch job (\d+)", stdout_tail)
if m:
slurm_job_id = m.group(1)
else:
m = re.search(r"Job id:\s*(\d+)", stdout_tail, re.IGNORECASE)
if m:
slurm_job_id = m.group(1)
if not experiment_id:
return {
"ok": False,
"executor": "slurm",
"reason": "launch_result_unparsed",
"exit_code": 0,
"slurm_job_id": slurm_job_id,
"stdout_tail": stdout_tail,
"stderr_tail": stderr_tail,
"diagnostic": (
"launch.py exited 0 but did not report an experiment_id "
"that callers can use for job_status/job_logs polling. "
"Treating this as failed even if a Slurm job id was parsed."
),
"argv": argv,
**_source_result_fields(checkout),
}
return {
"ok": True,
@@ -691,6 +1128,7 @@ def submit_job_impl(
"exit_code": 0,
"stdout_tail": stdout_tail,
"argv": argv,
**_source_result_fields(checkout),
}
@@ -704,6 +1142,8 @@ def _submit_job_dry_run(
job_dir: str | None,
job_name: str | None,
extra_overrides: dict[str, str] | None,
source_ref: str | None,
source_repo: str | None,
) -> dict:
"""Validate a launcher YAML by running ``launch.py --dryrun``.
@@ -719,9 +1159,17 @@ def _submit_job_dry_run(
failure / timeout branches (the validated-success branch omits
it since there's nothing to diagnose).
"""
# Same path resolution as the live submit, so dry-run and live use
# exactly the same YAML.
abs_yaml = _normalize_yaml_path(yaml_path)
# Same source + path resolution as the live submit, so dry-run and live
# use exactly the same launcher checkout and YAML.
source = _ensure_source_checkout(source_ref=source_ref, source_repo=source_repo)
if not source.get("ok"):
return {**source, "dry_run": True}
checkout: SourceCheckout | None = source["checkout"]
abs_yaml = _normalize_yaml_path(
yaml_path,
examples_dir=checkout.examples_dir if checkout else None,
)
if not abs_yaml.exists():
return {
"ok": False,
@@ -729,6 +1177,7 @@ def _submit_job_dry_run(
"reason": "yaml_not_found",
"yaml_path": yaml_path,
"resolved_path": str(abs_yaml),
**_source_result_fields(checkout),
"diagnostic": (
f"YAML not found at {abs_yaml}. Pass a path under "
f"tools/launcher/examples/ (relative), an absolute path, "
@@ -745,7 +1194,7 @@ def _submit_job_dry_run(
# blocks on its confirmation prompt — and since we're capturing
# stdout (no TTY), the prompt would hang until the 60-second
# timeout fires.
argv = ["modelopt-launcher", "--yaml", str(abs_yaml), "--dryrun", "--yes"]
argv = _launcher_argv(abs_yaml, checkout, "--dryrun", "--yes")
if hf_local:
argv.append(f"hf_local={hf_local}")
if cluster_user:
@@ -764,6 +1213,10 @@ def _submit_job_dry_run(
# default when cluster_host is set).
child_env = os.environ.copy()
child_env.setdefault("NEMORUN_HOME", os.getcwd())
if checkout is not None:
child_env["MODELOPT_MCP_SOURCE_ROOT"] = str(checkout.root)
child_env["MODELOPT_MCP_SOURCE_REF"] = checkout.ref
child_env["MODELOPT_MCP_SOURCE_SHA"] = checkout.resolved_sha
if cluster_host:
child_env["SLURM_HOST"] = cluster_host
@@ -778,7 +1231,7 @@ def _submit_job_dry_run(
# `submit_job_impl` (Popen at line 563 + run at line 590), the
# verify probes (line 197 + 251), and the SSH probe (line 326).
try:
proc = subprocess.run( # nosec B603
proc = subprocess.run( # nosec B603 - fixed dry-run launcher argv list; no shell.
argv,
env=child_env,
capture_output=True,
@@ -802,12 +1255,13 @@ def _submit_job_dry_run(
"hung."
),
"argv": argv,
**_source_result_fields(checkout),
}
stdout_tail = str(proc.stdout or "")[-2000:]
stderr_tail = str(proc.stderr or "")[-2000:]
if proc.returncode != 0:
if proc.returncode != 0 or _launcher_reported_error(stdout_tail, stderr_tail):
return {
"ok": True, # The tool itself ran cleanly
"dry_run": True,
@@ -816,14 +1270,15 @@ def _submit_job_dry_run(
"stdout_tail": stdout_tail,
"stderr_tail": stderr_tail,
"diagnostic": (
f"launch.py --dryrun rejected the YAML (exit code "
f"{proc.returncode}). Common reasons: invalid YAML "
f"syntax, missing required fields, factory function "
f"not registered, or a referenced file (HF model path, "
f"container tag) doesn't exist. See stderr_tail for the "
f"specific error."
"launch.py --dryrun rejected the YAML or printed a fatal "
"launcher error. Common reasons: invalid YAML syntax, "
"missing required fields, factory function not registered, "
"factory parsing failure, or a referenced file (HF model "
"path, container tag) doesn't exist. See stdout_tail/"
"stderr_tail for the specific error."
),
"argv": argv,
**_source_result_fields(checkout),
}
# Success branch returns the same field set as the failure branch
@@ -837,6 +1292,7 @@ def _submit_job_dry_run(
"stdout_tail": stdout_tail,
"stderr_tail": stderr_tail,
"argv": argv,
**_source_result_fields(checkout),
}
@@ -859,16 +1315,37 @@ def _resolve_experiment_dir(experiment_id: str) -> Path | None:
for the case where the operator didn't set NEMORUN_HOME at all
AND the MCP server's cwd differs from where launch.py ran.
"""
candidates = []
for root in _experiment_search_roots():
direct = root / experiment_id
if direct.exists():
return direct
for nested in root.glob(f"*/{experiment_id}"):
if nested.exists():
return nested
return None
def _experiment_search_roots() -> list[Path]:
"""Return experiment roots searched by status/log tools."""
roots = []
nemorun_home = os.environ.get("NEMORUN_HOME")
if nemorun_home:
candidates.append(Path(nemorun_home) / "experiments" / experiment_id)
candidates.append(Path.cwd() / "experiments" / experiment_id)
candidates.append(Path.cwd() / "local_experiments" / experiment_id)
for c in candidates:
if c.exists():
return c
return None
roots.append(Path(nemorun_home) / "experiments")
roots.append(Path.cwd() / "experiments")
roots.append(Path.cwd() / "local_experiments")
launcher_dir = _find_launcher_package_dir()
if launcher_dir is not None:
roots.append(launcher_dir / "experiments")
return roots
def _experiment_not_found_diagnostic() -> str:
"""Describe all experiment roots used by _resolve_experiment_dir."""
roots = ", ".join(str(root) for root in _experiment_search_roots())
return (
f"Searched experiment roots: {roots}. Either the id is wrong or "
"NEMORUN_HOME isn't set the same as it was at submit time."
)
def job_status_impl(experiment_id: str) -> dict:
@@ -885,18 +1362,17 @@ def job_status_impl(experiment_id: str) -> dict:
Per-task statuses (``status_<task_name>.out``) are also surfaced so
multi-task pipelines can be inspected.
"""
invalid = _validate_experiment_id(experiment_id)
if invalid:
return invalid
exp_dir = _resolve_experiment_dir(experiment_id)
if exp_dir is None:
return {
"ok": False,
"experiment_id": experiment_id,
"reason": "experiment_dir_not_found",
"diagnostic": (
"Searched NEMORUN_HOME/experiments/, ./experiments/, "
"./local_experiments/ — no match. Either the id is "
"wrong or NEMORUN_HOME isn't set the same as it was "
"at submit time."
),
"diagnostic": _experiment_not_found_diagnostic(),
}
done_marker = exp_dir / "_DONE"
@@ -940,6 +1416,10 @@ def job_logs_impl(
If ``task`` is None, returns logs for ALL tasks.
If ``tail`` is set, returns only the last N lines per task.
"""
invalid = _validate_experiment_id(experiment_id)
if invalid:
return invalid
exp_dir = _resolve_experiment_dir(experiment_id)
if exp_dir is None:
return {
@@ -1188,7 +1668,7 @@ def read_cluster_artifact_impl(
str(job_idx),
]
try:
proc = subprocess.run( # nosec B603 B607
proc = subprocess.run( # nosec B603 B607 - fixed nemo CLI argv; no shell.
argv,
capture_output=True,
text=True,
@@ -1367,7 +1847,7 @@ def open_draft_pr_impl(
# Step 1: push
try:
push = subprocess.run( # nosec B603 B607
push = subprocess.run( # nosec B603 B607 - fixed git CLI argv; no shell.
["git", "push", "-u", "origin", "HEAD"],
cwd=str(cwd_path),
capture_output=True,
@@ -1393,7 +1873,7 @@ def open_draft_pr_impl(
# Step 2: gh pr create
try:
gh = subprocess.run( # nosec B603 B607
gh = subprocess.run( # nosec B603 B607 - fixed gh CLI argv; no shell.
[
"gh",
"pr",
tools/mcp/modelopt_mcp/server.py
+22
View File
@@ -197,6 +197,26 @@ def _build_server() -> FastMCP:
)
),
] = None,
source_ref: Annotated[
str | None,
Field(
description=(
"Model-Optimizer branch, tag, or commit SHA to materialize "
"before launching. None resolves the repository default "
"configured by MODELOPT_MCP_SOURCE_REF, falling back to main."
)
),
] = None,
source_repo: Annotated[
str | None,
Field(
description=(
"Git repository URL for the managed Model-Optimizer checkout. "
"None uses MODELOPT_MCP_SOURCE_REPO or the public NVIDIA/"
"Model-Optimizer repository."
)
),
] = None,
skip_verify: Annotated[
bool,
Field(
@@ -241,6 +261,8 @@ def _build_server() -> FastMCP:
extra_overrides=extra_overrides,
skip_verify=skip_verify,
dry_run=dry_run,
source_ref=source_ref,
source_repo=source_repo,
)
@mcp.tool(
tools/mcp/tests/test_bridge.py
+377
View File
@@ -28,6 +28,13 @@ pytest.importorskip("pydantic")
from modelopt_mcp import bridge
@pytest.fixture(autouse=True)
def _disable_managed_source(monkeypatch):
"""Keep legacy subprocess tests offline unless a test opts into source routing."""
monkeypatch.setenv("MODELOPT_MCP_DISABLE_MANAGED_SOURCE", "1")
# ---------------------------------------------------------------------------
# list_examples
# ---------------------------------------------------------------------------
@@ -270,6 +277,280 @@ def test_submit_job_yaml_not_found(monkeypatch, tmp_path):
assert result["reason"] == "yaml_not_found"
def test_ensure_source_checkout_defaults_to_main(monkeypatch, tmp_path):
"""Managed source defaults to Model-Optimizer main and the default repo."""
monkeypatch.delenv("MODELOPT_MCP_DISABLE_MANAGED_SOURCE", raising=False)
monkeypatch.setenv("MODELOPT_MCP_SOURCE_CACHE", str(tmp_path / "cache"))
seen = {}
def fake_resolve(repo, ref):
seen["repo"] = repo
seen["ref"] = ref
return {"ok": True, "resolved_sha": "a" * 40}
monkeypatch.setattr(bridge, "_resolve_source_ref", fake_resolve)
monkeypatch.setattr(bridge, "_checkout_ready", lambda path, sha: True)
result = bridge._ensure_source_checkout()
assert result["ok"] is True
assert seen["repo"] == "https://github.com/NVIDIA/Model-Optimizer.git"
assert seen["ref"] == "main"
assert result["checkout"].resolved_sha == "a" * 40
def test_submit_job_dry_run_uses_managed_source_checkout(monkeypatch, tmp_path):
"""Managed source routes launcher execution through uv --project <checkout>."""
checkout_root = tmp_path / "checkout"
yaml_dir = checkout_root / "tools" / "launcher" / "examples" / "fam" / "model"
yaml_dir.mkdir(parents=True)
yaml_path = yaml_dir / "config.yaml"
yaml_path.write_text("job_name: t\npipeline: []\n")
checkout = bridge.SourceCheckout(
repo="https://example.com/modelopt.git",
ref="feature/ref",
resolved_sha="b" * 40,
root=checkout_root,
)
monkeypatch.setattr(
bridge,
"_ensure_source_checkout",
lambda **_: {"ok": True, "checkout": checkout},
)
captured = {}
def fake_run(argv, **kwargs):
captured["argv"] = argv
captured["env"] = kwargs["env"]
return subprocess.CompletedProcess(
args=argv,
returncode=0,
stdout="Dry-run OK\n",
stderr="",
)
monkeypatch.setattr(subprocess, "run", fake_run)
result = bridge.submit_job_impl(
yaml_path="fam/model/config.yaml",
hf_local=None,
cluster_host=None,
cluster_user=None,
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=True,
dry_run=True,
source_ref="feature/ref",
)
assert result["ok"] is True
assert result["source_ref"] == "feature/ref"
assert result["source_sha"] == "b" * 40
assert captured["argv"][:5] == [
"uv",
"run",
"--project",
str(checkout_root / "tools" / "launcher"),
"modelopt-launcher",
]
assert str(yaml_path) in captured["argv"]
assert captured["env"]["MODELOPT_MCP_SOURCE_ROOT"] == str(checkout_root)
assert captured["env"]["MODELOPT_MCP_SOURCE_SHA"] == "b" * 40
def test_submit_job_source_checkout_failure_short_circuits(monkeypatch):
"""Source checkout failures return structured diagnostics and do not launch."""
monkeypatch.setattr(
bridge,
"_ensure_source_checkout",
lambda **_: {
"ok": False,
"reason": "source_ref_not_found",
"diagnostic": "missing ref",
},
)
result = bridge.submit_job_impl(
yaml_path="examples/test.yaml",
hf_local=None,
cluster_host=None,
cluster_user=None,
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=True,
dry_run=True,
source_ref="ghost",
)
assert result["ok"] is False
assert result["dry_run"] is True
assert result["reason"] == "source_ref_not_found"
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"
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.setattr(
bridge,
"verify_slurm_setup_impl",
lambda **_: {"ok": True},
)
def fake_run(argv, **kwargs):
return subprocess.CompletedProcess(
args=argv,
returncode=0,
stdout="Configuring global options\n",
stderr="",
)
monkeypatch.setattr(subprocess, "run", fake_run)
result = bridge.submit_job_impl(
yaml_path="config.yaml",
hf_local=None,
cluster_host="cluster.example.com",
cluster_user="user",
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=False,
)
assert result["ok"] is False
assert result["reason"] == "launch_result_unparsed"
def test_submit_job_slurm_parses_nemo_job_id(monkeypatch, tmp_path):
"""Parse Slurm job id from Nemo's experiment status output."""
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.setattr(
bridge,
"verify_slurm_setup_impl",
lambda **_: {"ok": True},
)
def fake_run(argv, **kwargs):
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",
hf_local=None,
cluster_host="cluster.example.com",
cluster_user="user",
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=False,
)
assert result["ok"] is True
assert result["slurm_job_id"] == "13049989"
assert result["experiment_id"] == "cicd_1782173197"
def test_submit_job_slurm_job_id_without_experiment_id_is_failure(monkeypatch, tmp_path):
"""A Slurm job id alone is not enough for MCP status/log polling."""
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.setattr(
bridge,
"verify_slurm_setup_impl",
lambda **_: {"ok": True},
)
def fake_run(argv, **kwargs):
return subprocess.CompletedProcess(
args=argv,
returncode=0,
stdout="Task 0\n- Job id: 13049989\n",
stderr="",
)
monkeypatch.setattr(subprocess, "run", fake_run)
result = bridge.submit_job_impl(
yaml_path="config.yaml",
hf_local=None,
cluster_host="cluster.example.com",
cluster_user="user",
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=False,
)
assert result["ok"] is False
assert result["reason"] == "launch_result_unparsed"
assert result["slurm_job_id"] == "13049989"
def test_submit_job_slurm_zero_exit_with_launcher_error_is_failure(monkeypatch, tmp_path):
"""Launcher fatal text must override a misleading zero exit status."""
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.setattr(
bridge,
"verify_slurm_setup_impl",
lambda **_: {"ok": True},
)
def fake_run(argv, **kwargs):
return subprocess.CompletedProcess(
args=argv,
returncode=0,
stdout="Configuring global options\n",
stderr="Unexpected error: Failed to parse slurm_factory\n",
)
monkeypatch.setattr(subprocess, "run", fake_run)
result = bridge.submit_job_impl(
yaml_path="config.yaml",
hf_local=None,
cluster_host="cluster.example.com",
cluster_user="user",
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=False,
)
assert result["ok"] is False
assert result["reason"] == "launch_py_failed"
assert "Unexpected error" in result["stderr_tail"]
# ---------------------------------------------------------------------------
# submit_job dry-run branch
# ---------------------------------------------------------------------------
@@ -353,6 +634,43 @@ def test_submit_job_dry_run_yaml_invalid(monkeypatch, tmp_path):
assert "yaml.YAMLError" in result["stderr_tail"]
def test_submit_job_dry_run_zero_exit_with_launcher_error_is_invalid(monkeypatch, tmp_path):
"""dry-run must treat fatal launcher text as invalid even with exit 0."""
yaml_dir = tmp_path / "examples"
yaml_dir.mkdir()
yaml_path = yaml_dir / "bad.yaml"
yaml_path.write_text("job_name: t\npipeline: []\n")
monkeypatch.setenv("MODELOPT_LAUNCHER_EXAMPLES_DIR", str(yaml_dir))
def fake_run(argv, **kwargs):
return subprocess.CompletedProcess(
args=argv,
returncode=0,
stdout="Configuring global options\n",
stderr="Unexpected error: Failed to parse slurm_factory\n",
)
monkeypatch.setattr(subprocess, "run", fake_run)
result = bridge.submit_job_impl(
yaml_path="bad.yaml",
hf_local=None,
cluster_host=None,
cluster_user=None,
identity=None,
job_dir=None,
job_name=None,
extra_overrides=None,
skip_verify=True,
dry_run=True,
)
assert result["ok"] is True
assert result["dry_run"] is True
assert result["validated"] is False
assert result["exit_code"] == 0
assert "Unexpected error" in result["stderr_tail"]
def test_submit_job_dry_run_yaml_not_found(monkeypatch, tmp_path):
"""dry_run=True + missing yaml → yaml_not_found with dry_run flag set."""
monkeypatch.setenv("MODELOPT_LAUNCHER_EXAMPLES_DIR", str(tmp_path))
@@ -464,6 +782,44 @@ def test_job_status_running(tmp_path, monkeypatch):
assert result["has_done_marker"] is False
def test_job_status_nested_nemo_title_dir(tmp_path, monkeypatch):
"""nemo_run stores experiments under experiments/<title>/<experiment_id>."""
exp = tmp_path / "experiments" / "cicd" / "exp_1781000006"
exp.mkdir(parents=True)
(exp / "status_task_0.out").write_text("running\n")
monkeypatch.setenv("NEMORUN_HOME", str(tmp_path))
result = bridge.job_status_impl("exp_1781000006")
assert result["ok"] is True
assert result["experiment_dir"] == str(exp)
assert result["status"] == "running"
def test_job_status_launcher_experiments_fallback(tmp_path, monkeypatch):
"""Resolve experiments under the installed launcher's package directory."""
launcher_dir = tmp_path / "launcher"
exp = launcher_dir / "experiments" / "cicd" / "exp_1781000007"
exp.mkdir(parents=True)
(exp / "status_task_0.out").write_text("running\n")
monkeypatch.delenv("NEMORUN_HOME", raising=False)
other_cwd = tmp_path / "other"
other_cwd.mkdir()
monkeypatch.chdir(other_cwd)
monkeypatch.setattr(bridge, "_find_launcher_package_dir", lambda: launcher_dir)
result = bridge.job_status_impl("exp_1781000007")
assert result["ok"] is True
assert result["experiment_dir"] == str(exp)
assert result["status"] == "running"
def test_job_status_rejects_unsafe_experiment_id():
"""Experiment ids are path tokens, not filesystem paths or globs."""
result = bridge.job_status_impl("../exp_1781000008")
assert result["ok"] is False
assert result["reason"] == "invalid_experiment_id"
def test_job_status_unknown_id(tmp_path, monkeypatch):
"""No experiment dir matching the id → experiment_dir_not_found."""
monkeypatch.setenv("NEMORUN_HOME", str(tmp_path))
@@ -472,6 +828,20 @@ def test_job_status_unknown_id(tmp_path, monkeypatch):
assert result["reason"] == "experiment_dir_not_found"
def test_job_status_unknown_id_reports_launcher_fallback(tmp_path, monkeypatch):
"""The not-found diagnostic stays in sync with the searched roots."""
launcher_dir = tmp_path / "launcher"
launcher_dir.mkdir()
monkeypatch.delenv("NEMORUN_HOME", raising=False)
monkeypatch.setattr(bridge, "_find_launcher_package_dir", lambda: launcher_dir)
result = bridge.job_status_impl("does_not_exist")
assert result["ok"] is False
assert result["reason"] == "experiment_dir_not_found"
assert str(launcher_dir / "experiments") in result["diagnostic"]
def test_job_logs_all_tasks(tmp_path, monkeypatch):
"""task=None returns logs for every log_*.out under the experiment dir."""
exp = tmp_path / "experiments" / "exp_1781000003"
@@ -511,6 +881,13 @@ def test_job_logs_missing_task(tmp_path, monkeypatch):
assert result["reason"] == "task_log_not_found"
def test_job_logs_rejects_unsafe_experiment_id():
"""job_logs applies the same experiment-id validation as job_status."""
result = bridge.job_logs_impl("exp_*", task=None, tail=None)
assert result["ok"] is False
assert result["reason"] == "invalid_experiment_id"
# ---------------------------------------------------------------------------
# wait_for_experiment
# ---------------------------------------------------------------------------