mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Chenhan D. YuandClaude Opus 4.7
9f37fe1969
## 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>
53 lines
2.3 KiB
Python
53 lines
2.3 KiB
Python
# 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.
|
|
|
|
"""ModelOpt launcher MCP server.
|
|
|
|
Anchor design: OMNIML-5123. Exposes the launcher's submit / status / logs
|
|
operations as typed MCP tools that codex / Claude Code agents can call
|
|
directly, instead of shelling out to ``uv run launch.py --yaml ...``.
|
|
|
|
Tool surface (Phase 1):
|
|
|
|
* ``list_examples`` — discover bundled example YAMLs under
|
|
``tools/launcher/examples/`` with their model + description metadata.
|
|
* ``verify_setup`` — fail-fast probe for the named executor (docker or
|
|
slurm) BEFORE the user burns minutes on a misconfigured submission.
|
|
* ``submit_job`` — submit a launcher YAML; mode determined by args
|
|
(``hf_local`` → Docker; ``cluster_host`` → Slurm). Returns
|
|
immediately: Slurm returns ``experiment_id`` (parsed from launch.py's
|
|
detach-mode stdout), Docker returns the background subprocess
|
|
``pid`` (Phase 2 will tail launcher output to capture the nemo_run
|
|
experiment_id for the Docker path too).
|
|
* ``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_*.out`` per task, with optional tail.
|
|
|
|
Two design constants:
|
|
|
|
1. **Mode determined by args, not by tool choice.** Single
|
|
``submit_job`` rather than separate ``submit_docker`` / ``submit_slurm``.
|
|
The mutually-exclusive arg shape is a runtime check; the LLM catalog
|
|
stays compact.
|
|
2. **Filesystem is the source of truth.** Status + logs both read from
|
|
nemo_run's experiment dir. The MCP server carries no per-job state
|
|
across calls — survives restarts cleanly.
|
|
"""
|
|
|
|
from modelopt_mcp.server import main
|
|
|
|
__all__ = ["main"]
|