Record the MLflow run that produced a checkpoint in .experiment.json (#2374)

### What does this PR do?

Type of change: new feature

A tracked `hf_ptq` run already tags itself with the checkpoint it writes
(`checkpoint_path`), so a run can be followed to its output. The reverse
was missing: given a checkpoint on disk, there was no way to find the
run that quantized it without searching the tracking server by path.

A tracked run now writes `.experiment.json` into `--export_path` naming
the experiment, the MLflow run id and the run URL, and uploads the same
bytes as the `experiment.json` artifact so a downloaded artifact set is
self-describing. `MlflowRunLogger` gains a `run_info` property carrying
that identity, with the tracking URI credential-masked the way `run_url`
already was.

Two deliberate behaviours:

- **Written from a `finally`**, so a run that crashes after export still
leaves the pointer behind.
- **Skipped when the export directory is absent** — a run that exported
nothing has nowhere to put it, and creating the directory would suggest
a checkpoint that does not exist. The artifact is still uploaded in that
case, so a failed run is traceable from the server side.

A failed local write warns and continues rather than failing the job,
consistent with the rest of the MLflow path. Only the main rank writes,
since the logger is inert on other ranks.

### Usage

```bash
python hf_ptq.py --pyt_ckpt_path Qwen/Qwen3.5-0.8B --qformat fp8 \
    --export_path /tmp/qwen35-fp8 --mlflow https://<your-mlflow-server>
```

```console
$ cat /tmp/qwen35-fp8/.experiment.json
{
  "tracking_uri": "https://<your-mlflow-server>",
  "experiment_name": "alice/hf_ptq/Qwen3.5-0.8B-fp8",
  "experiment_id": "36",
  "run_id": "7bec239a3a154970b062f3024a5ff20e",
  "run_name": "20260910-175422",
  "run_url": "https://<your-mlflow-server>/#/experiments/36/runs/7bec239a3a154970b062f3024a5ff20e"
}
```

```python
# checkpoint -> run
import json, mlflow
info = json.load(open("/tmp/qwen35-fp8/.experiment.json"))
mlflow.set_tracking_uri(info["tracking_uri"])
run = mlflow.get_run(info["run_id"])
```

### Testing

**Unit** — `tests/unit/torch/utils/test_mlflow.py` (61 passed):
`run_info` contents before/after the run opens, the defaulted run name
being reported rather than left blank, and credential masking of the
tracking URI.

**Example** — `tests/examples/hf_ptq/test_hf_ptq_args.py` (27 passed):
the file landing in the checkpoint and on the server with identical
content, the failed-run path, the no-export path, and untracked runs
writing nothing.

**Real runs**, 1x H200, `Qwen3.5-0.8B` FP8 PTQ,
`tensorrt-llm/release:1.3.0rc26`:

- Against a local MLflow server — checkpoint copy and uploaded artifact
byte-identical; artifacts on the run were `command.txt`,
`experiment.json`, `logs/hf_ptq.log`, `summary/quant_summary.txt`,
`version.txt`.
- Against the internal `mlflow-modelopt` server (experiment
`chenjiel/hf_ptq/Qwen3.5-0.8B-fp8`, run
`7bec239a3a154970b062f3024a5ff20e`) — same result, confirming artifact
upload against a real backend. Reading `.experiment.json` back and
calling `mlflow.get_run(run_id)` resolved to `FINISHED` with
`checkpoint_path` pointing at the export directory.
- Crash path exercised for real when a first attempt died on a gated
calibration dataset: no export directory created, `experiment.json`
still uploaded, run closed `FAILED`.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- 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)?:
✅ — new entry under `*Misc*` in the open 0.48.0 section, matching where
the MLflow entries sit in 0.47.0.
- Did you get Claude approval on this PR?: ❌

🤖 Generated with [Claude Code](https://claude.com/claude-code)


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

* **New Features**
* Exported checkpoints now record experiment and run traceability
metadata in `.experiment.json`.
* Checkpoint metadata is uploaded with opened MLflow runs, including
runs where export fails.
* Active MLflow run details—including identifiers, resolved run name,
URL, and tracking server—are available with credentials redacted.
* **Bug Fixes**
* Improved handling of failed, untracked, and pre-existing exports to
prevent inherited metadata pointers.
* **Documentation**
* Updated MLflow integration guidance and changelog information for
checkpoint metadata.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

examples/hf_ptq/README.md

@@ -700,6 +700,7 @@ seconds rather than after a full calibration.
 | --- | --- |
 | `command.txt` | The full invocation, copy-pasteable, with credentials masked |
 | `version.txt` | The ModelOpt version that ran |
+| `experiment.json` | The experiment name, run id and run URL — the same file written into `--export_path` |
 | `recipe/resolved_recipe.yaml` | The `--recipe` with its `$import`s expanded, so it stands alone |
 | `logs/hf_ptq.log` | The run's Python stdout/stderr, including the traceback if it crashed |
 | `summary/quant_summary.txt` | The per-quantizer summary (unless `--no-verbose`) |
@@ -711,6 +712,22 @@ Every command-line argument is also logged as a searchable param, alongside
 `user` / `hostname` / `modelopt_version` / `git_sha` tags. A run that fails is
 still recorded, with status `FAILED` and its log attached.
 
+A tracked run also drops `.experiment.json` into `--export_path`, so a checkpoint found on
+disk names the run that produced it:
+
+```bash
+cat <quantized_ckpt_path>/.experiment.json
+# {"tracking_uri": ..., "experiment_name": ..., "experiment_id": ..., "run_id": ...,
+#  "run_name": ..., "run_url": ...}
+```
+
+The local file is written only once the export itself completes, so a run that fails
+earlier leaves whatever checkpoint is already in `--export_path` — and its pointer —
+untouched. The `experiment.json` artifact is uploaded for every run that opened, so a
+failed run stays traceable from the server. An export that is *not* tracked removes any
+pointer it would otherwise inherit, from a reused `--export_path` or from a tracked source
+checkpoint.
+
 Other flags:
 
 - `--mlflow_experiment` — defaults to `$USER/hf_ptq/<checkpoint basename>-<recipe name>`,

examples/hf_ptq/example_utils.py

@@ -22,8 +22,8 @@ import json
 import logging
 import os
 import warnings
-from collections.abc import Callable, Iterable
-from contextlib import AbstractContextManager, nullcontext
+from collections.abc import Callable, Iterable, Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import timedelta
 from pathlib import Path
@@ -101,7 +101,11 @@ _HF_PTQ_WEIGHT_FILE_PATTERNS = (
     "*.tgz",
     "*.zip",
 )
+# Dotted like the other sidecars hf_ptq drops in the export directory, so it is ignored by
+# from_pretrained and does not look like part of the model.
+_EXPERIMENT_JSON = ".experiment.json"
 _HF_PTQ_EXPORT_OWNED_FILES = {
+    _EXPERIMENT_JSON,
     "config.json",
     "hf_quant_config.json",
     "quant_config.json",
@@ -1309,9 +1313,11 @@ def add_mlflow_args(parser: argparse.ArgumentParser) -> None:
         help=(
             "Track this run on an MLflow server (e.g. https://<your-mlflow-server>/), "
             "uploading the command, the resolved recipe, the run log and the quantization "
-            "summaries. MLflow's own $MLFLOW_TRACKING_URI enables tracking without this "
-            "flag, which overrides it. A URI taken from the environment is best-effort: if "
-            "it is unusable the run warns and continues untracked."
+            "summaries, and writing .experiment.json into --export_path so the checkpoint "
+            "names the run that produced it. MLflow's own $MLFLOW_TRACKING_URI enables "
+            "tracking without this flag, which overrides it. A URI taken from the "
+            "environment is best-effort: if it is unusable the run warns and continues "
+            "untracked."
         ),
     )
     parser.add_argument(
@@ -1353,7 +1359,14 @@ def resolve_mlflow_args(args: argparse.Namespace, parser: argparse.ArgumentParse
 
 
 _MLFLOW_NON_PARAM_ARGS = frozenset(
-    {"dist_state", "mlflow", "mlflow_experiment", "mlflow_required", "mlflow_run_name"}
+    {
+        "checkpoint_exported",
+        "dist_state",
+        "mlflow",
+        "mlflow_experiment",
+        "mlflow_required",
+        "mlflow_run_name",
+    }
 )
 
 
@@ -1382,19 +1395,81 @@ def _mlflow_logger(args: argparse.Namespace) -> MlflowRunLogger:
     )
 
 
-def mlflow_run(args: argparse.Namespace) -> AbstractContextManager:
-    """Track this invocation for the duration of the block, or do nothing if untracked."""
+@contextmanager
+def mlflow_run(args: argparse.Namespace) -> Iterator[None]:
+    """Track this invocation for the duration of the block, and keep the checkpoint's
+    provenance pointer honest whether or not the run is tracked."""
     logger = _mlflow_logger(args)
+    export_path = Path(args.export_path)
     if not logger.enabled:
         # Gathering the inputs re-reads the recipe, so keep it off the untracked path.
-        return nullcontext()
+        try:
+            yield
+        finally:
+            _drop_inherited_experiment_json(args, export_path)
+        return
     params, texts = _mlflow_run_inputs(args)
-    return logger.track(
+    with logger.track(
         params=params,
         tags=_mlflow_run_tags(args),
         texts=texts,
         files=_mlflow_run_outputs(args),
-    )
+    ):
+        try:
+            yield
+        finally:
+            _log_experiment_json(logger, args, export_path)
+
+
+def _log_experiment_json(
+    logger: MlflowRunLogger, args: argparse.Namespace, export_path: Path
+) -> None:
+    """Record which MLflow run produced this checkpoint, in the checkpoint and on the server.
+
+    The tags point from the run to the checkpoint it wrote; this file is the reverse, so a
+    checkpoint found on disk can be traced back to the run that quantized it without
+    searching the server.
+
+    The artifact goes up for any run that opened, so a failure is traceable from the server
+    side. The local copy is written only once ``export_quantized`` has returned, because the
+    file claims authorship of the checkpoint sitting next to it: ``--export_path`` existing
+    proves nothing, since ``print_quant_summary`` creates it before quantization and the
+    directory may hold a valid checkpoint from an earlier attempt whose weights this run
+    never touched.
+
+    There is nothing to record at all when the run never opened, which a URI taken from the
+    environment reaches by design: it disables tracking from inside the block rather than
+    failing the quantization.
+    """
+    info = logger.run_info
+    if not info:
+        return
+    text = json.dumps(info, indent=2) + "\n"
+    logger.log_text(_EXPERIMENT_JSON.removeprefix("."), text)
+    if not args.checkpoint_exported:
+        return
+    try:
+        (export_path / _EXPERIMENT_JSON).write_text(text)
+    except OSError as e:
+        print(f"[mlflow] WARNING: could not write {export_path / _EXPERIMENT_JSON}: {e}")
+
+
+def _drop_inherited_experiment_json(args: argparse.Namespace, export_path: Path) -> None:
+    """Remove a pointer an untracked export would otherwise inherit.
+
+    A fresh checkpoint written into a reused ``--export_path`` would keep the previous run's
+    pointer, and one quantized from a tracked source checkpoint could be handed that
+    source's pointer. Either way the file would name a run that did not produce these
+    weights. Only a completed export clears it; a failed run leaves whatever checkpoint was
+    already there, pointer included.
+    """
+    if not args.checkpoint_exported or not args.dist_state.is_main:
+        return
+    stale = export_path / _EXPERIMENT_JSON
+    try:
+        stale.unlink(missing_ok=True)
+    except OSError as e:
+        print(f"Warning: could not remove stale {stale}: {e}")
 
 
 def _mlflow_run_tags(args: argparse.Namespace) -> dict[str, str]:

examples/hf_ptq/hf_ptq.py

@@ -880,6 +880,7 @@ def export_quantized(
         # No tokenizer saving needed for spec ckpts
         if has_spec_opt(full_model):
             export_speculative_decoding(full_model, export_dir=export_path)
+            args.checkpoint_exported = True
             print(f"Quantized speculative decoding checkpoint exported to: {export_path}")
             return
 
@@ -976,6 +977,7 @@ def export_quantized(
                 exclude_files=exclude_files,
             )
 
+        args.checkpoint_exported = True
         end_time = time.time()
         print_rank_0(
             f"Quantized model exported to: {export_path}. Total time used {end_time - start_time}s"
@@ -1689,6 +1691,9 @@ def parse_args() -> argparse.Namespace:
     add_mlflow_args(parser)
 
     args = parser.parse_args()
+    # Flipped by export_quantized once a checkpoint is actually on disk. The MLflow pointer
+    # is gated on it rather than on --export_path existing, which proves nothing.
+    args.checkpoint_exported = False
     resolve_mlflow_args(args, parser)
 
     if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0):