mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Keval MorabiaandClaude Opus 5
4eb86524f0
### What does this PR do? Type of change: refactor (no functional change) **[1/2] of a split. Merge this first; #2514 is [2/2] and is based on this branch.** Three example scripts had each reimplemented the same MLflow wiring: the flags, the `$USER/<tool>/<model>-<variant>` experiment convention, the params/tags/artifacts a run uploads, and the open/close dance with its status. The copies had already drifted — only `hf_ptq` wrote a provenance pointer, only `vllm_serve` republished the resolved URI — and every new tracked script meant another copy. What a script records is now one declarative `Tool` record, **declared in the script itself, beside the flags it reads**: ```python # examples/megatron_bridge/quantize.py QUANTIZE = Tool( name="megatron_bridge_quantize", tracks="Track this run on an MLflow server, uploading the command, the resolved recipe, ...", variant_help="recipe name, or --quant_cfg if no --recipe", variant=lambda args: Path(args.recipe).stem if args.recipe else (args.quant_cfg or "none"), model=lambda args: args.hf_model_name_or_path, checkpoint=lambda args: args.export_megatron_path, texts=lambda args: resolved_recipe_texts(args.recipe), outputs=lambda args: {"summary/quant_summary.txt": Path(args.export_megatron_path) / ".quant_summary.txt"}, ) ``` `tracked_run` takes that record and runs the whole thing, so a script adds tracking in three lines: `add_mlflow_args(parser, TOOL)`, `resolve_mlflow_args(args, parser, TOOL)`, and `with mlflow_run(args, TOOL):`. The shared module knows no script's flags. `examples/hf_ptq`, `examples/vllm_serve` and `examples/megatron_bridge/quantize.py` move onto it. Three helpers fall away as redundant (`track_run`, `checkpoint_run_tags`, and `hf_ptq`'s two flag pass-throughs). ### Usage No user-facing change. The flags, their spellings and the experiment naming are exactly as before; a script author now writes a `Tool` instead of four functions. ### Testing - `tests/unit/torch/utils/test_mlflow.py`, `tests/examples/hf_ptq/test_hf_ptq_args.py`, `tests/examples/vllm_serve/test_vllm_mlflow_utils.py` — **179 pass**. - `tests/examples/megatron_bridge` in `nvcr.io/nvidia/nemo:26.08` (the only lane that runs it), which drives `quantize.py` for real: **34 passed**, locally and in this PR's `megatron` lane. - `pre-commit run --files <changed>`: all hooks pass. - The four suites shared four copies of a stand-in for the `mlflow` module, which had drifted — one recorded artifacts as a list, another as a dict, a third made `log_artifact` a no-op, so a test asserting on an upload asserted nothing. They now share one `tests/_test_utils/mlflow.py`, which also emulates the fluent API's habit of opening a run when none is active. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — `track_run` and `checkpoint_run_tags` are removed, but neither shipped in a release (0.47.0's `__all__` is `MlflowRunLogger`, `command_text`, `current_user`, `default_experiment_name`, `validate_tracking_uri`, all unchanged here). Three deliberate behaviour changes, each in shared code and each tested: - The `source_checkpoint_path` tag resolves to an absolute path where it recorded the raw argument, which a chain of runs needs to join on the pair. `run_tags` is shared, so this applies to every script that writes the tag — `hf_ptq` **and** `megatron_bridge/quantize.py`, for a local `--hf_model_name_or_path`. A source that names no directory, such as a Hub `org/name` id, is still recorded as given. - `MlflowRunLogger.track()` — which *did* ship in 0.47.0 — records a block ending in `SystemExit(0)` as `FINISHED` where it recorded `FAILED`, since a script that ends by calling `sys.exit()` rather than returning has still finished. - `.experiment.json`'s `tracking_uri` and the `run_url` built from it drop a trailing `/` from the tracking URI, so the link is `https://host/#/...` rather than `https://host//#/...`. Only reachable by constructing `MlflowRunLogger` directly; every CLI path already stripped the slash in `resolve_tracking_uri`. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — no user-visible change; the entry is in [2/2]. - Did you get Claude approval on this PR?: several rounds; re-requested on this head. ### Additional Information Split out of #2514. This half is the enabling refactor with no behaviour change; #2514 is the feature it unlocks and is based on this branch. At ~605 changed lines of core logic it is over the ~500 guideline; the owner accepted a two-PR split rather than three, and everything #2514 alone consumes — `split_tracking_credentials`, `log_active_run_experiment_json`, `MlflowRunLogger._reattach` — lands there rather than here. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>