feat(tools/mcp): MCP server for ModelOpt launcher (OMNIML-5123) (#1701)

## Summary

`tools/mcp/` — a new MCP server exposing the existing
`tools/launcher/core.py` orchestration as **typed MCP tools** that codex
/ Claude Code agents can call directly, instead of shelling out to `uv
run launch.py --yaml ...` and parsing prose output.

Tracked under
[OMNIML-5123](https://jirasw.nvidia.com/browse/OMNIML-5123) (Epic).
Ships **Phase 1 + Phase 1.5** together: the core launcher surface plus
the four highest-leverage helpers from the `cell.md` simplification loop
([OMNIML-5128](https://jirasw.nvidia.com/browse/OMNIML-5128) partial,
[OMNIML-5132](https://jirasw.nvidia.com/browse/OMNIML-5132) full).

## Nine tools

**Phase 1 — core launcher surface:**

| Tool | Description |
|---|---|
| `list_examples` | Enumerate `tools/launcher/examples/` with model +
description metadata extracted from each YAML |
| `verify_setup` | Fail-fast probe for the named executor. Docker:
`docker info` (daemon up) + `docker info --format` runtime-registry
check for the `nvidia` runtime — no image pull, daemon-fast. Slurm: `ssh
-o BatchMode=yes -o ConnectTimeout=5` to the cluster login node. ~1 s
probe saves 30+ s of wasted submission on bad config |
| `submit_job` | Submit a launcher YAML. Mode is determined by
mutually-exclusive args: `hf_local` → Docker (local GPU), `cluster_host`
→ Slurm (remote SSH). Returns experiment_id immediately; the actual job
runs detached |
| `job_status` | Filesystem-based status from nemo_run's experiment dir
(`_DONE`, `status_*.out`) — no in-memory registry, survives MCP server
restarts |
| `job_logs` | Read `log_<task>.out` from experiment dir; per-task
filtering + optional tail |

**Phase 1.5 — `cell.md` simplification (OMNIML-5128 / 5132):**

| Tool | Description |
|---|---|
| `wait_for_experiment` | Replaces the agent's `while True: status;
sleep` poll loop with one tool call. Reuses `job_status_impl` so
terminal-state semantics stay identical. Returns final status plus
`waited_seconds`; on timeout returns structured `{ok: False, reason:
"wait_timeout", last_status: …}` |
| `provision_passwordless_ssh_dry_run` | No-side-effect inspection of
`~/.ssh/` that emits the exact `ssh-keygen` / `ssh-copy-id` commands the
operator should run to make `verify_setup(executor='slurm')` pass.
Closes the verify_setup "ssh_auth_failed → now what?" gap |
| `read_cluster_artifact` | Uses nemo_run's tunnel primitives, not a
reinvented SSH layer. `path=None` wraps `nemo experiment logs <id>
<job_idx>` (built-in log fetch); with a `path`, uses the experiment's
`Tunnel` to read the file. Structured failure on subprocess error /
timeout |
| `open_draft_pr` | `git push -u origin HEAD` + `gh pr create --draft
…`. Validates cwd is a git repo first; on gh failure after push
succeeds, reports `branch_pushed=True` so the operator can retry just
the PR-open step |

## Design constants

1. **Single `submit_job` with mode by args** (not separate
`submit_docker` / `submit_slurm` tools). Keeps the LLM tool catalog
compact; mutual-exclusion is a runtime check.
2. **Filesystem is the source of truth** for status + logs. No in-memory
registry. Survives MCP server restarts cleanly — important because
operators / agents kill + restart their hosts often.
3. **`verify_setup` is auto-called by `submit_job`** by default
(skippable when caller just probed). The probe is ~1 s; the cost of a
misconfigured submission is 30+ s of cluster timeout or container-pull.
Always-on verify pays back immediately.
4. **Delegate to nemo_run for tunnels.** `read_cluster_artifact` and
`wait_for_experiment` use nemo_run's existing `Experiment` / `Tunnel` /
`nemo experiment logs` primitives rather than reinventing SSH/rsync. One
source of truth for cluster I/O.

## Layout

```
tools/mcp/
├── pyproject.toml          # name: modelopt-mcp, console_script
├── modelopt_mcp/
│   ├── __init__.py
│   ├── server.py           # FastMCP entry; 9 tool definitions
│   └── bridge.py           # thin wrapper over launcher's core.py
│                           #   + filesystem status/log helpers
│                           #   + tunnel/PR helpers (Phase 1.5)
└── tests/
    └── test_bridge.py      # 34 unit tests, fully hermetic
                            # (mocked subprocess + tmp_path fixtures)
```

## Install

Two paths, both **from source via uv**. No PyPI wheel exists;
OMNIML-5123 opted for the uvx-from-git pattern to skip publication
overhead.

### End-user install (recommended)

`uvx` from the git subdirectory — single command, no manual clone:

```bash
# Claude Code
claude mcp add modelopt -- uvx --from \
  "git+https://github.com/NVIDIA/Model-Optimizer.git#subdirectory=tools/mcp" \
  modelopt-mcp

# Codex
codex mcp add modelopt -- uvx --from \
  "git+https://github.com/NVIDIA/Model-Optimizer.git#subdirectory=tools/mcp" \
  modelopt-mcp
```

Under the hood `uvx` clones the whole repo to its cache, installs
`tools/mcp/` as the entry, and resolves the sibling `modelopt-launcher`
dep via `[tool.uv.sources]` (`path = "../launcher"`) inside the cloned
tree.

### Dev install (local checkout)

```bash
uv pip install -e tools/launcher    # sibling dep first
uv pip install -e tools/mcp         # then this package
modelopt-mcp                         # entry on PATH
```

### Why no plain `pip install` today

Two specific reasons, worth flagging so reviewers know what's
intentional vs missing:

1. **Nothing on PyPI yet.** Neither `modelopt-mcp` nor
`modelopt-launcher` are published — this PR introduces the package but
doesn't add release machinery.
2. **`pip` doesn't read `[tool.uv.sources]`.** Even from a local
checkout, plain `pip install -e tools/mcp` fails because
`modelopt-launcher` is a bare name (no URL) and pip can't find it.
Sticking with `uv` / `uvx` is the practical path while we're git-only.

If we later want plain-pip support: publish to PyPI, or switch to a
PEP-440 direct URL (`"modelopt-launcher @
git+…#subdirectory=tools/launcher"`). Out of scope for this PR.

## Post-review changes

Addressed all CodeRabbit + claude[bot] review findings on the original
Phase-1 surface. See the inline replies for details; the substantive
bug-fix highlights:

* **Slurm `cluster_host`** — propagate via `env=child_env` (launch.py
reads SLURM_HOST, not a CLI arg)
* **`shlex.quote`** removed from nemo-run k=v overrides (subprocess
list-form doesn't shell-quote)
* **Docker `Popen`** now uses `stdout=DEVNULL, stderr=DEVNULL,
start_new_session=True` to avoid pipe-buffer blocking
* **`NEMORUN_HOME`** pinned in subprocess env so submit + status sides
agree
* **GPU verify** swapped from `docker run --gpus all` image-pull (slow +
flaky) to `docker info --format` runtime-registry check (daemon-fast)
* **Task-status word match** anchors on first word against a fixed
failure-word set (no more `"fail" in "succeeded after retry; previous
attempt failed"` false-positive)
* **`experiment_id` regex** generalized for non-NVIDIA cluster paths
* **`pyproject.toml`** dropped the unsatisfiable `modelopt-launcher`
bare-name dep (launcher is a file-layout sibling, not a Python import
dep)
* **`Field(ge=1)`** on `job_logs.tail`
* **Docstring contract** clarified (Docker returns `pid`, Slurm returns
`experiment_id`)

## Validation

- [x] `uv pip install -e .` succeeds (modelopt-launcher resolved
transitively)
- [x] 34/34 unit tests pass (`uv run python -m pytest tests/`)
- [x] stdio handshake works end-to-end; `tools/list` returns all 9 with
full schemas + descriptions
- [x] Mode-resolution: `submit_job` correctly rejects no-executor +
both-executors with structured `reason`
- [x] Filesystem status: correctly classifies `done` / `failed` /
`running` from `_DONE` + `status_*.out`
- [x] `wait_for_experiment` short-circuits on already-terminal
experiments; honors timeout without raising
- [x] `provision_passwordless_ssh_dry_run` distinguishes no-key /
key-only / key+pubkey cases
- [x] `read_cluster_artifact` handles subprocess timeout + non-zero exit
with structured reasons
- [x] `open_draft_pr` reports `branch_pushed=True` on
gh-failure-after-push so retries are cheap
- [x] Pre-commit clean: ruff, ruff-format, mypy, bandit, license-headers

## Acceptance criteria

**OMNIML-5123 (Phase 1):**

- [x] `list_examples` returns all bundled YAMLs with path and model name
- [x] `submit_job` with `hf_local` runs via Docker executor and returns
immediately (Phase 1: returns PID; experiment_id capture in Phase 2)
- [x] `submit_job` with `cluster_host`/`user` runs via Slurm executor
(`detach=True`) and returns experiment_id
- [x] `job_status` correctly reflects running / done / failed from
nemo_run filesystem
- [x] `job_logs` returns stdout for a completed job
- [x] `uvx --from git+...#subdirectory=tools/mcp modelopt-mcp --help`
resolves and starts
- [x] Existing launcher tests unaffected (no changes to
`tools/launcher/`)

**OMNIML-5128 (Phase 1.5, partial):**

- [x] `wait_for_experiment` blocks until terminal or timeout
- [x] `read_cluster_artifact` pulls remote artifacts via nemo_run tunnel
- [x] `open_draft_pr` opens a draft PR against a target repo
- [ ] Capture `experiment_id` from Docker subprocess output — deferred
to Phase 2

**OMNIML-5132 (Phase 1.5, full):**

- [x] `provision_passwordless_ssh_dry_run` emits operator-facing
commands without side effects

## Phase 2 (separate PR)

* Capture `experiment_id` from Docker subprocess output (tail until
nemo_run logs the id).
* Extract the verify + submit helpers into a shared lib that
[`nmm-sandbox-mcp`](https://gitlab-master.nvidia.com/omniml/integration/nmm-sandbox/-/tree/main/tools/mcp)
(companion server, separate repo) can consume for internal-ergonomics
tools — cluster short-name → factory lookup + GitLab CI dispatch.
* NEL integration
([OMNIML-5133](https://jirasw.nvidia.com/browse/OMNIML-5133)) +
checkpoint introspection
([OMNIML-5134](https://jirasw.nvidia.com/browse/OMNIML-5134)).


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

* **New Features**
* Added ModelOpt MCP server and console entrypoint; tools:
list_examples, verify_setup, submit_job, job_status, job_logs,
wait_for_experiment, provision_passwordless_ssh_dry_run,
read_cluster_artifact, open_draft_pr; Docker and Slurm support.

* **Documentation**
* Expanded README with install steps, design notes, end-to-end agent
example, roadmap, and repo layout.

* **Tests**
* Expanded unit tests covering bridge helpers, polling, SSH flows,
artifact reads, and PR automation.

* **Chores**
  * CI updated to run MCP tests; package/meta config added.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Chenhan D. Yu
2026-06-12 18:56:18 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 2640551515
commit 9f37fe1969
8 changed files with 2809 additions and 1 deletions
+27 -1
View File
@@ -12,6 +12,7 @@ on:
- "pyproject.toml"
- "tests/unit/**"
- "tools/launcher/**"
- "tools/mcp/**"
- ".agents/skills/**"
schedule:
- cron: "0 0 * * *" # Nightly
@@ -55,6 +56,7 @@ jobs:
pyproject.toml
tests/unit/**
tools/launcher/**
tools/mcp/**
.agents/skills/**
linux:
needs: [check-dco]
@@ -147,6 +149,29 @@ jobs:
uv venv .venv
uv pip install -e . pytest
uv run python3 -m pytest -v
mcp:
if: needs.check-file-changes.outputs.any_changed == 'true'
needs: [linux, check-file-changes]
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@v6
with:
submodules: recursive
- name: Run modelopt-mcp tests
working-directory: tools/mcp
run: |
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
uv venv .venv
# Install the sibling launcher package first; it's a runtime
# dep declared in tools/mcp/pyproject.toml as `modelopt-launcher`
# but uv resolves the source via [tool.uv.sources] to a local
# editable path. -e on both packages keeps the install cheap
# and matches the dev-mode install in tools/mcp/README.md.
uv pip install -e ../launcher
uv pip install -e . pytest
uv run python3 -m pytest -v
skills:
if: needs.check-file-changes.outputs.any_changed == 'true'
needs: [linux, check-file-changes]
@@ -167,7 +192,7 @@ jobs:
unit-pr-required-check:
# Run even if some jobs are skipped
if: ${{ github.event_name == 'pull_request' && always() }}
needs: [check-file-changes, linux, windows, multi-version, partial-install, launcher, skills]
needs: [check-file-changes, linux, windows, multi-version, partial-install, launcher, mcp, skills]
runs-on: ubuntu-latest
steps:
- name: Required unit tests did not succeed
@@ -177,6 +202,7 @@ jobs:
needs.multi-version.result != 'success' ||
needs.partial-install.result != 'success' ||
needs.launcher.result != 'success' ||
needs.mcp.result != 'success' ||
needs.skills.result != 'success'
)) }}
run: exit 1