From e1b798f7cc8ad87f976b2c6ba54b7507f59c02c7 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 28 Jul 2026 07:37:58 +0000 Subject: [PATCH 01/10] Add optional MLflow tracking to hf_ptq.py Records a PTQ run on an MLflow server so it can be reproduced from its MLflow entry alone. Enabled by --mlflow ; without the flag nothing changes. The run is opened before the model loads, so an unreachable server or a bad URI fails in seconds rather than after hours of calibration. The invocation and the recipe are uploaded at that point too, which keeps a crashed run useful: it is still recorded, with status FAILED and its log attached. Uploads command.txt, recipe/resolved_recipe.yaml, logs/hf_ptq.log and the quantization summaries, plus the model / format / calibration settings as searchable params. The recipe is uploaded resolved rather than verbatim because a recipe may be a directory or use $imports -- for the Qwen3.6 MoE AutoQuantize recipe the source file is 2.2 KB against 7.6 KB resolved, so only the resolved form describes what actually ran. hf_ptq.py has no logging framework, so the log is produced by teeing stdout/stderr. Handlers that libraries bound to sys.stderr at import time are re-pointed at the tee for the duration, otherwise transformers and huggingface_hub warnings reach the console but never the log. Experiment defaults to $USER/hf_ptq/-; run name defaults to the UTC start time. Only the main rank uploads. MLflow failures never fail the quantization. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 + examples/hf_ptq/README.md | 42 +++ examples/hf_ptq/hf_ptq.py | 41 +++ examples/hf_ptq/mlflow_utils.py | 342 ++++++++++++++++++ examples/hf_ptq/requirements.txt | 1 + tests/examples/hf_ptq/test_mlflow_utils.py | 401 +++++++++++++++++++++ 6 files changed, 829 insertions(+) create mode 100644 examples/hf_ptq/mlflow_utils.py create mode 100644 tests/examples/hf_ptq/test_mlflow_utils.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..d31c30d049e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,8 @@ Changelog **New Features** +- Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py`` to record a PTQ run on an MLflow server: the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries are uploaded as artifacts, with the model / format / calibration settings as searchable params. The run is opened before the model loads so an unreachable server fails in seconds instead of after calibration, and a failed run is still recorded with its log attached. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. + **Backward Breaking Changes** **Deprecations** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index a3967565ae0..8adf755488b 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -293,6 +293,48 @@ scripts/huggingface_example.sh --model --quant nvfp4 --vlm --calib_with_ > Note: when `--calib_with_images` is set, `--calib_size` must be a single value, and the calibration dataset is nvidia/nemotron_vlm_dataset_v2. This functionality is currently in beta and has been tested on `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16`. +### Tracking runs with MLflow + +Pass `--mlflow ` to record a PTQ run on an MLflow server, so the run can be +reproduced later from its MLflow entry alone: + +```bash +python hf_ptq.py \ + --pyt_ckpt_path \ + --recipe general/ptq/nvfp4_default-kv_fp8_cast \ + --export_path \ + --mlflow https:/// +``` + +The run is opened *before* the model loads, so a bad URI or a missing token fails within +seconds rather than after a full calibration. These artifacts are uploaded: + +| Artifact | Contents | +| --- | --- | +| `command.txt` | The full invocation, copy-pasteable | +| `recipe/resolved_recipe.yaml` | The `--recipe` with its `$import`s expanded, so it stands alone | +| `logs/hf_ptq.log` | Everything the run printed, including the traceback if it crashed | +| `summary/quant_summary.txt` | The per-quantizer summary (unless `--no-verbose`) | +| `summary/moe.html` | Per-expert calibration token counts, for MoE models | + +The model, format, recipe and calibration settings are also logged as searchable params, +alongside `user` / `hostname` / `modelopt_version` / `git_sha` tags. A run that fails is +still recorded, with status `FAILED` and its log attached. + +Other flags: + +- `--mlflow_experiment` — defaults to `$USER/hf_ptq/-`, + falling back to `--qformat` when no `--recipe` is used. +- `--mlflow_run_name` — defaults to the UTC start time, `YYYYmmdd-HHMMSS`. +- Passing `--mlflow` with no value uses `$MLFLOW_TRACKING_URI`. + +Authentication uses MLflow's own environment variables (`MLFLOW_TRACKING_TOKEN`, or +`MLFLOW_TRACKING_USERNAME` / `MLFLOW_TRACKING_PASSWORD`). + +> Note: only the main rank uploads, so `--use_fsdp2` runs produce a single run. The log +> captures Python output; output written directly by native libraries (NCCL, CUDA) goes to +> the terminal only. On SLURM, keep the job's own `.out` file for those. + ### Megatron-Bridge Example Script Please refer to [examples/megatron_bridge/README.md](../megatron_bridge/README.md) for example scripts for PTQ / QAD with Megatron-Bridge which is generally more performant than the Hugging Face scripts. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6a4bd476984..2cf68acaf1e 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -46,6 +46,7 @@ setup_distributed_args, validate_fsdp2_supported, ) +from mlflow_utils import MlflowRunLogger, default_experiment_name, validate_tracking_uri from torch.utils.data import DataLoader from transformers import ( AutoConfig, @@ -1622,7 +1623,39 @@ def parse_args() -> argparse.Namespace: ), ) + parser.add_argument( + "--mlflow", + nargs="?", + const=os.environ.get("MLFLOW_TRACKING_URI", ""), + default=None, + help=( + "Track this run on an MLflow server (e.g. https:///), " + "uploading the command, the resolved recipe, the run log and the quantization " + "summaries. Pass the flag without a value to use $MLFLOW_TRACKING_URI." + ), + ) + parser.add_argument( + "--mlflow_experiment", + default=None, + help=( + "MLflow experiment name. Default: " + "$USER/hf_ptq/-." + ), + ) + parser.add_argument( + "--mlflow_run_name", + default=None, + help="MLflow run name. Default: the UTC start time as YYYYmmdd-HHMMSS.", + ) + args = parser.parse_args() + if args.mlflow is not None: + try: + args.mlflow = validate_tracking_uri(args.mlflow) + except ValueError as e: + parser.error(str(e)) + args.mlflow_experiment = args.mlflow_experiment or default_experiment_name(args) + if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") @@ -1667,6 +1700,12 @@ def main(args: argparse.Namespace): setup_distributed_args(args) + # Opened before the model loads so an unreachable server or a bad experiment name + # fails in seconds rather than after a full calibration run. + mlflow_logger = MlflowRunLogger(args) + mlflow_logger.start() + + status = "FAILED" try: # launch a memory monitor to read the currently used GPU memory. launch_memory_monitor() @@ -1703,8 +1742,10 @@ def main(args: argparse.Namespace): default_pad_token, device, ) + status = "FINISHED" finally: cleanup_distributed(args) + mlflow_logger.finish(status) if __name__ == "__main__": diff --git a/examples/hf_ptq/mlflow_utils.py b/examples/hf_ptq/mlflow_utils.py new file mode 100644 index 00000000000..94710873162 --- /dev/null +++ b/examples/hf_ptq/mlflow_utils.py @@ -0,0 +1,342 @@ +# 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. + +"""MLflow tracking for ``hf_ptq.py``, enabled by ``--mlflow ``. + +Uploads the invocation, the resolved recipe, the run log and the quantization +summaries to an MLflow tracking server so a PTQ run can be reproduced from its +MLflow entry alone. Everything here is a no-op unless ``--mlflow`` is given. +""" + +import getpass +import logging +import os +import re +import shlex +import shutil +import socket +import subprocess +import sys +import tempfile +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import yaml + +import modelopt +from modelopt.recipe import load_recipe + +__all__ = ["MlflowRunLogger", "TeeStream", "default_experiment_name", "validate_tracking_uri"] + +# MLflow experiment names are stored in a VARCHAR(256) column by the SQL-backed stores. +_MAX_COMPONENT_LEN = 100 +_UNSAFE_CHARS = re.compile(r"[^A-Za-z0-9._-]+") + + +def validate_tracking_uri(uri: str) -> str: + """Validate an MLflow tracking URI and return it without a trailing slash. + + Only ``http(s)`` servers are accepted; MLflow's local ``file:`` / ``sqlite:`` + backends are not a useful destination for a shared PTQ record. + + Raises: + ValueError: if *uri* is empty, has no host, or is not an http(s) URL. + """ + if not uri: + raise ValueError( + "--mlflow requires a tracking URI (e.g. https:///); " + "pass one explicitly or set MLFLOW_TRACKING_URI." + ) + parsed = urlparse(uri) + if parsed.scheme not in ("http", "https"): + raise ValueError( + f"--mlflow expects an http(s) tracking URI, got {uri!r}. " + "Did you mean https://" + uri.lstrip("/") + "?" + ) + if not parsed.netloc: + raise ValueError(f"--mlflow tracking URI {uri!r} has no host.") + return uri.rstrip("/") + + +def default_experiment_name(args) -> str: + """Build the default experiment name, ``/hf_ptq/-``. + + The model component is the checkpoint's basename, so both a local directory and a + ``org/name`` Hugging Face id collapse to the same readable name. + """ + model = Path(args.pyt_ckpt_path).name + variant = Path(args.recipe).stem if args.recipe else args.qformat + return f"{_sanitize(_user())}/hf_ptq/{_sanitize(model)}-{_sanitize(variant)}" + + +def _sanitize(component: str) -> str: + """Reduce one experiment-name component to ``[A-Za-z0-9._-]``.""" + cleaned = _UNSAFE_CHARS.sub("_", component).strip("._-") + return cleaned[:_MAX_COMPONENT_LEN] or "unknown" + + +def _user() -> str: + try: + return getpass.getuser() + except OSError: # container without a passwd entry for the uid + return "unknown" + + +def _git_sha() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], + cwd=Path(__file__).resolve().parent, + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (subprocess.CalledProcessError, OSError): + return "unknown" + + +def _command_text() -> str: + """The invocation, as a copy-pasteable line.""" + lines = [shlex.join([sys.executable, *sys.argv])] + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size > 1: + lines += [ + "", + f"# Launched under torchrun with WORLD_SIZE={world_size}, " + f"LOCAL_WORLD_SIZE={os.environ.get('LOCAL_WORLD_SIZE', '?')}. The torchrun " + "wrapper is not part of sys.argv and is therefore not shown above.", + ] + return "\n".join(lines) + "\n" + + +class TeeStream: + """Mirror a text stream to *sink* while passing writes through to *stream*. + + ``hf_ptq.py`` reports progress with bare ``print()`` and has no log file; wrapping + ``sys.stdout``/``sys.stderr`` in this is what produces one. Attribute access falls + through to the wrapped stream so ``isatty()`` keeps progress bars behaving. Native + (C-level) writes go straight to the real file descriptor and are *not* captured. + """ + + def __init__(self, stream, sink): + self._stream = stream + self._sink = sink + + def write(self, data: str) -> int: + self._stream.write(data) + if not self._sink.closed: + self._sink.write(data) + return len(data) + + def flush(self) -> None: + self._stream.flush() + if not self._sink.closed: + self._sink.flush() + + def __getattr__(self, name): + return getattr(self._stream, name) + + +class MlflowRunLogger: + """Records one ``hf_ptq.py`` invocation as an MLflow run. + + Disabled -- every method a no-op -- unless ``--mlflow`` was passed and this is the + main rank. :meth:`start` validates the server and opens the run *before* the model + loads, so a bad URI or a missing token fails in seconds rather than after hours of + calibration; :meth:`finish` uploads the run outputs and closes the run. + + Example: + >>> logger = MlflowRunLogger(args) + >>> logger.start() + >>> try: + ... quantize_and_export() + ... status = "FINISHED" + ... finally: + ... logger.finish(status) + """ + + def __init__(self, args): + self.args = args + self.enabled = bool(args.mlflow) and args.dist_state.is_main + self._run: Any = None + self._log_dir: Path | None = None + self._saved_streams: tuple | None = None + self._redirected_handlers: list[tuple[logging.StreamHandler, Any]] = [] + self._start_time = 0.0 + + def start(self) -> None: + """Open the run: capture stdout, verify the server, log the inputs.""" + if not self.enabled: + return + self._start_time = time.time() + self._start_capture() + try: + self._open_run() + self._log_inputs() + except Exception: + self._stop_capture() + raise + + def finish(self, status: str) -> None: + """Upload the run outputs and close the run with *status*. + + Never raises: an MLflow outage must not turn a successful quantization into a + failed exit. + """ + if not self.enabled or self._run is None: + self._stop_capture() + return + try: + self._log_outputs() + except Exception as e: + print(f"[mlflow] WARNING: could not upload run outputs: {e}") + self._stop_capture() + try: + import mlflow + + mlflow.end_run(status=status) + print(f"[mlflow] {status}: {self._run_url()}") + except Exception as e: + print(f"[mlflow] WARNING: could not close the run: {e}") + + def _open_run(self) -> None: + # Optional dependency: only examples using --mlflow need it installed. + try: + import mlflow + except ImportError as e: + raise ImportError("--mlflow requires the 'mlflow' package: pip install mlflow") from e + + self._check_reachable() + mlflow.set_tracking_uri(self.args.mlflow) + mlflow.set_experiment(self.args.mlflow_experiment) + run_name = self.args.mlflow_run_name or datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S") + self._run = mlflow.start_run(run_name=run_name) + print(f"[mlflow] experiment: {self.args.mlflow_experiment}") + print(f"[mlflow] run: {self._run_url()}") + + def _check_reachable(self) -> None: + """Fail fast on an unreachable host. + + Any HTTP response -- including 401 -- means the host is up, so authorization is + left to the first real API call, which reports it precisely. + """ + import requests + + try: + requests.get(f"{self.args.mlflow}/health", timeout=10) + except requests.RequestException as e: + raise ConnectionError(f"MLflow server {self.args.mlflow} is unreachable: {e}") from e + + def _log_inputs(self) -> None: + """Log what the run was given -- upfront, so a crash still leaves a usable record.""" + import mlflow + + args = self.args + mlflow.log_params( + { + "model": args.pyt_ckpt_path, + "qformat": args.qformat, + "kv_cache_qformat": args.kv_cache_qformat, + "recipe": args.recipe or "", + "calib_size": args.calib_size, + "calib_seq": args.calib_seq, + "batch_size": args.batch_size, + "sparsity_fmt": args.sparsity_fmt, + "export_path": args.export_path, + "world_size": args.dist_state.world_size, + } + ) + mlflow.set_tags( + { + "user": _user(), + "hostname": socket.gethostname(), + "modelopt_version": modelopt.__version__, + "git_sha": _git_sha(), + } + ) + mlflow.log_text(_command_text(), "command.txt") + if args.recipe: + # The resolved recipe, not the source file: a recipe may be a directory or + # use $imports, and only the resolved form is self-contained. + resolved = load_recipe(args.recipe).model_dump(mode="json") + mlflow.log_text( + yaml.safe_dump(resolved, sort_keys=False), "recipe/resolved_recipe.yaml" + ) + + def _log_outputs(self) -> None: + """Log what the run produced: the captured log and the quantization summaries.""" + import mlflow + + mlflow.log_metric("total_time_s", time.time() - self._start_time) + sys.stdout.flush() + sys.stderr.flush() + if self._log_dir is not None: + mlflow.log_artifact(str(self._log_dir / "hf_ptq.log"), artifact_path="logs") + export_path = Path(self.args.export_path) + # Written by print_quant_summary / save_expert_token_count_table; hidden names are + # awkward to browse in the MLflow UI, so they are uploaded without the leading dot. + for source, artifact in ( + (".quant_summary.txt", "summary/quant_summary.txt"), + (".moe.html", "summary/moe.html"), + ): + if (export_path / source).is_file(): + text = (export_path / source).read_text(encoding="utf-8", errors="replace") + mlflow.log_text(text, artifact) + + def _start_capture(self) -> None: + self._log_dir = Path(tempfile.mkdtemp(prefix="hf_ptq-mlflow-")) + sink = open(self._log_dir / "hf_ptq.log", "w", buffering=1, encoding="utf-8") + original_stdout, original_stderr = sys.stdout, sys.stderr + self._saved_streams = (original_stdout, original_stderr, sink) + sys.stdout = TeeStream(original_stdout, sink) + sys.stderr = TeeStream(original_stderr, sink) + self._redirect_log_handlers( + {original_stdout: sys.stdout, original_stderr: sys.stderr}, + ) + print(f"[mlflow] capturing this run's log to {self._log_dir / 'hf_ptq.log'}") + + def _redirect_log_handlers(self, replacements: dict) -> None: + """Point already-configured logging handlers at the tee. + + transformers and huggingface_hub bind ``sys.stderr`` into a ``StreamHandler`` when + they are imported, which happens long before the capture starts; without this their + warnings -- rate limits, deprecations -- reach the console but never the log. + """ + self._redirected_handlers = [] + loggers = [logging.getLogger(), *logging.Logger.manager.loggerDict.values()] + for logger in loggers: + for handler in getattr(logger, "handlers", []): + if isinstance(handler, logging.StreamHandler) and handler.stream in replacements: + self._redirected_handlers.append((handler, handler.stream)) + handler.setStream(replacements[handler.stream]) + + def _stop_capture(self) -> None: + if self._saved_streams is None: + return + for handler, stream in self._redirected_handlers: + handler.setStream(stream) + self._redirected_handlers = [] + sys.stdout, sys.stderr, sink = self._saved_streams + sink.close() + self._saved_streams = None + if self._log_dir is not None: + shutil.rmtree(self._log_dir, ignore_errors=True) + self._log_dir = None + + def _run_url(self) -> str: + info = self._run.info + return f"{self.args.mlflow}/#/experiments/{info.experiment_id}/runs/{info.run_id}" diff --git a/examples/hf_ptq/requirements.txt b/examples/hf_ptq/requirements.txt index deb09927544..3e7d4006c94 100644 --- a/examples/hf_ptq/requirements.txt +++ b/examples/hf_ptq/requirements.txt @@ -1,5 +1,6 @@ compressed-tensors fire flash-attn>=2.6.0 +mlflow transformers_stream_generator zstandard diff --git a/tests/examples/hf_ptq/test_mlflow_utils.py b/tests/examples/hf_ptq/test_mlflow_utils.py new file mode 100644 index 00000000000..3dea00f8450 --- /dev/null +++ b/tests/examples/hf_ptq/test_mlflow_utils.py @@ -0,0 +1,401 @@ +# 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. + +import getpass +import importlib +import io +import logging +import sys +from argparse import Namespace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" + + +@pytest.fixture +def mlflow_utils(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + monkeypatch.setattr(getpass, "getuser", lambda: "tester") + return importlib.import_module("mlflow_utils") + + +class FakeMlflow: + """Stand-in for the mlflow module, so these tests need no server and no dependency.""" + + def __init__(self): + self.tracking_uri = None + self.experiment = None + self.run_name = None + self.status = None + self.params = {} + self.tags = {} + self.texts = {} + self.metrics = {} + self.artifacts = [] + + def set_tracking_uri(self, uri): + self.tracking_uri = uri + + def set_experiment(self, name): + self.experiment = name + + def start_run(self, run_name=None): + self.run_name = run_name + return SimpleNamespace(info=SimpleNamespace(experiment_id="7", run_id="deadbeef")) + + def log_params(self, params): + self.params.update(params) + + def set_tags(self, tags): + self.tags.update(tags) + + def log_text(self, text, artifact_file): + self.texts[artifact_file] = text + + def log_artifact(self, local_path, artifact_path=None): + self.artifacts.append((Path(local_path).name, artifact_path)) + + def log_metric(self, key, value): + self.metrics[key] = value + + def end_run(self, status=None): + self.status = status + + +def _args(**overrides): + args = Namespace( + mlflow="https://mlflow.example.com", + mlflow_experiment="tester/hf_ptq/model-nvfp4", + mlflow_run_name=None, + pyt_ckpt_path="/models/Qwen3-0.6B", + recipe=None, + qformat="nvfp4", + kv_cache_qformat="fp8_cast", + calib_size=[1024], + calib_seq=512, + batch_size=0, + sparsity_fmt="dense", + export_path="exported_model", + dist_state=SimpleNamespace(is_main=True, world_size=1), + ) + return Namespace(**{**vars(args), **overrides}) + + +@pytest.fixture +def fake_mlflow(monkeypatch): + fake = FakeMlflow() + monkeypatch.setitem(sys.modules, "mlflow", fake) + import requests + + monkeypatch.setattr(requests, "get", lambda *a, **kw: SimpleNamespace(status_code=200)) + return fake + + +@pytest.mark.parametrize( + ("uri", "expected"), + [ + ("https://mlflow.example.com/", "https://mlflow.example.com"), + ("https://mlflow.example.com", "https://mlflow.example.com"), + ("http://localhost:5000", "http://localhost:5000"), + ("https://host/mlflow/", "https://host/mlflow"), + ], +) +def test_validate_tracking_uri_accepts_http_servers(mlflow_utils, uri, expected): + assert mlflow_utils.validate_tracking_uri(uri) == expected + + +@pytest.mark.parametrize( + "uri", + [ + "", # --mlflow with no value and no MLFLOW_TRACKING_URI + "mlflow.example.com", # missing scheme + "/local/mlruns", # local path + "file:///local/mlruns", # unsupported backend + "sqlite:///mlflow.db", + "https://", # no host + ], +) +def test_validate_tracking_uri_rejects_non_servers(mlflow_utils, uri): + with pytest.raises(ValueError): + mlflow_utils.validate_tracking_uri(uri) + + +@pytest.mark.parametrize( + ("ckpt", "recipe", "qformat", "expected"), + [ + # Local directory, trailing slash, built-in recipe name. + ( + "/models/Llama-3.3-70B-Instruct/", + "general/ptq/nvfp4_default-kv_fp8_cast", + "fp8", + "tester/hf_ptq/Llama-3.3-70B-Instruct-nvfp4_default-kv_fp8_cast", + ), + # Hugging Face id collapses to its basename; recipe file path drops the suffix. + ( + "nvidia/Llama-3.3-70B-Instruct", + "./my/tuned.yaml", + "fp8", + "tester/hf_ptq/Llama-3.3-70B-Instruct-tuned", + ), + # No recipe: the quantization format names the variant. + ("openai/gpt-oss-20b", None, "nvfp4", "tester/hf_ptq/gpt-oss-20b-nvfp4"), + # Comma-separated formats and other unsafe characters are sanitized. + ("/models/my model!", None, "nvfp4,fp8", "tester/hf_ptq/my_model-nvfp4_fp8"), + ], +) +def test_default_experiment_name(mlflow_utils, ckpt, recipe, qformat, expected): + args = _args(pyt_ckpt_path=ckpt, recipe=recipe, qformat=qformat) + assert mlflow_utils.default_experiment_name(args) == expected + + +def test_default_experiment_name_survives_unusable_username(mlflow_utils, monkeypatch): + """A container without a passwd entry for the uid must not break the run.""" + monkeypatch.setattr(getpass, "getuser", lambda: (_ for _ in ()).throw(OSError)) + args = _args(pyt_ckpt_path="/models/Qwen3-0.6B", recipe=None, qformat="nvfp4") + assert mlflow_utils.default_experiment_name(args) == "unknown/hf_ptq/Qwen3-0.6B-nvfp4" + + +def test_tee_stream_writes_to_both_and_delegates(mlflow_utils): + original, sink = io.StringIO(), io.StringIO() + tee = mlflow_utils.TeeStream(original, sink) + + print("hello", file=tee) + tee.flush() + + assert original.getvalue() == "hello\n" + assert sink.getvalue() == "hello\n" + # Progress bars check isatty(); it must report the real stream, not the tee. + assert tee.isatty() == original.isatty() + + +def test_logger_is_inert_without_the_flag(mlflow_utils, monkeypatch): + """Without --mlflow nothing is imported, captured or uploaded.""" + monkeypatch.setitem(sys.modules, "mlflow", None) + logger = mlflow_utils.MlflowRunLogger(_args(mlflow=None)) + stdout = sys.stdout + + logger.start() + assert sys.stdout is stdout + logger.finish("FINISHED") + + +def test_logger_skips_non_main_ranks(mlflow_utils, monkeypatch): + monkeypatch.setitem(sys.modules, "mlflow", None) + args = _args(dist_state=SimpleNamespace(is_main=False, world_size=8)) + logger = mlflow_utils.MlflowRunLogger(args) + + logger.start() + logger.finish("FINISHED") + + +def test_logger_logs_inputs_and_outputs(mlflow_utils, fake_mlflow, tmp_path, monkeypatch): + monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--pyt_ckpt_path", "/models/Qwen3-0.6B"]) + (tmp_path / ".quant_summary.txt").write_text("448 TensorQuantizers found in model\n") + (tmp_path / ".moe.html").write_text("experts") + args = _args( + recipe="general/ptq/nvfp4_default-kv_fp8_cast", + export_path=str(tmp_path), + ) + logger = mlflow_utils.MlflowRunLogger(args) + + logger.start() + try: + assert fake_mlflow.tracking_uri == "https://mlflow.example.com" + assert fake_mlflow.experiment == "tester/hf_ptq/model-nvfp4" + # The default run name is the UTC start time. + assert len(fake_mlflow.run_name) == 15 and fake_mlflow.run_name[8] == "-" + finally: + logger.finish("FINISHED") + + assert fake_mlflow.params["model"] == "/models/Qwen3-0.6B" + assert fake_mlflow.params["qformat"] == "nvfp4" + assert fake_mlflow.params["kv_cache_qformat"] == "fp8_cast" + assert fake_mlflow.tags["user"] == "tester" + + command = fake_mlflow.texts["command.txt"] + assert "hf_ptq.py --pyt_ckpt_path /models/Qwen3-0.6B" in command + assert "torchrun" not in command + + # The recipe is uploaded resolved, so $imports are expanded and it stands alone. + recipe = yaml.safe_load(fake_mlflow.texts["recipe/resolved_recipe.yaml"]) + assert recipe["metadata"]["recipe_type"] == "ptq" + assert recipe["quantize"]["quant_cfg"] + + assert "TensorQuantizers" in fake_mlflow.texts["summary/quant_summary.txt"] + assert fake_mlflow.texts["summary/moe.html"] == "experts" + assert ("hf_ptq.log", "logs") in fake_mlflow.artifacts + assert "total_time_s" in fake_mlflow.metrics + assert fake_mlflow.status == "FINISHED" + + +def test_command_flags_the_invisible_torchrun_wrapper( + mlflow_utils, fake_mlflow, tmp_path, monkeypatch +): + """Under torchrun, sys.argv is the worker's, so the launcher must be called out.""" + monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--use_fsdp2"]) + monkeypatch.setenv("WORLD_SIZE", "8") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") + logger = mlflow_utils.MlflowRunLogger(_args(recipe=None, export_path=str(tmp_path))) + + logger.start() + logger.finish("FINISHED") + + command = fake_mlflow.texts["command.txt"] + assert "hf_ptq.py --use_fsdp2" in command + assert "WORLD_SIZE=8" in command and "not part of sys.argv" in command + + +def test_capture_includes_preconfigured_library_logging(mlflow_utils, fake_mlflow, tmp_path): + """transformers/huggingface_hub bind sys.stderr at import, long before capture starts.""" + library_logger = logging.getLogger("test_preconfigured_library") + handler = logging.StreamHandler(sys.stderr) + library_logger.addHandler(handler) + logger = mlflow_utils.MlflowRunLogger(_args(recipe=None, export_path=str(tmp_path))) + + try: + logger.start() + log_path = logger._log_dir / "hf_ptq.log" + library_logger.warning("Rate limited. Waiting 169.0s before retry") + captured = log_path.read_text() + logger.finish("FINISHED") + finally: + library_logger.removeHandler(handler) + + assert "Rate limited" in captured + # The handler must be handed back its own stream, or later logging writes to a closed file. + assert handler.stream is sys.stderr + + +def test_logger_omits_recipe_artifact_without_a_recipe(mlflow_utils, fake_mlflow, tmp_path): + logger = mlflow_utils.MlflowRunLogger(_args(recipe=None, export_path=str(tmp_path))) + + logger.start() + logger.finish("FINISHED") + + assert "recipe/resolved_recipe.yaml" not in fake_mlflow.texts + assert fake_mlflow.params["recipe"] == "" + + +def test_logger_restores_streams_and_reports_failure(mlflow_utils, fake_mlflow, tmp_path): + """A failed quantization is still recorded, with its log attached.""" + logger = mlflow_utils.MlflowRunLogger(_args(recipe=None, export_path=str(tmp_path))) + stdout, stderr = sys.stdout, sys.stderr + + logger.start() + print("calibrating") + logger.finish("FAILED") + + assert sys.stdout is stdout and sys.stderr is stderr + assert fake_mlflow.status == "FAILED" + assert ("hf_ptq.log", "logs") in fake_mlflow.artifacts + + +def test_logger_never_raises_when_the_server_dies_mid_run(mlflow_utils, fake_mlflow, tmp_path): + logger = mlflow_utils.MlflowRunLogger(_args(recipe=None, export_path=str(tmp_path))) + logger.start() + + def explode(*args, **kwargs): + raise RuntimeError("server gone") + + fake_mlflow.log_artifact = explode + logger.finish("FINISHED") + + assert sys.stdout is not None and not isinstance(sys.stdout, mlflow_utils.TeeStream) + + +def test_unreachable_server_fails_before_the_model_loads(mlflow_utils, monkeypatch, tmp_path): + import requests + + monkeypatch.setitem(sys.modules, "mlflow", FakeMlflow()) + monkeypatch.setattr( + requests, + "get", + lambda *a, **kw: (_ for _ in ()).throw(requests.ConnectionError("no route to host")), + ) + logger = mlflow_utils.MlflowRunLogger(_args(recipe=None, export_path=str(tmp_path))) + stdout = sys.stdout + + with pytest.raises(ConnectionError, match="unreachable"): + logger.start() + + # The capture must be torn down so the failure is readable on the console. + assert sys.stdout is stdout + + +def test_parse_args_defaults_the_experiment_name(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + monkeypatch.setattr(getpass, "getuser", lambda: "tester") + hf_ptq = importlib.import_module("hf_ptq") + monkeypatch.setattr( + sys, + "argv", + [ + "hf_ptq.py", + "--pyt_ckpt_path", + "/models/Qwen3-0.6B", + "--recipe", + "general/ptq/nvfp4_default-kv_fp8_cast", + "--mlflow", + "https://mlflow.example.com/", + ], + ) + + args = hf_ptq.parse_args() + + assert args.mlflow == "https://mlflow.example.com" + assert args.mlflow_experiment == "tester/hf_ptq/Qwen3-0.6B-nvfp4_default-kv_fp8_cast" + assert args.mlflow_run_name is None + + +def test_parse_args_leaves_mlflow_off_by_default(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + hf_ptq = importlib.import_module("hf_ptq") + monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--pyt_ckpt_path", "/models/Qwen3-0.6B"]) + + args = hf_ptq.parse_args() + + assert args.mlflow is None + assert args.mlflow_experiment is None + + +def test_parse_args_rejects_a_bad_tracking_uri(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + hf_ptq = importlib.import_module("hf_ptq") + monkeypatch.setattr( + sys, + "argv", + ["hf_ptq.py", "--pyt_ckpt_path", "/models/Qwen3-0.6B", "--mlflow", "not-a-url"], + ) + + with pytest.raises(SystemExit): + hf_ptq.parse_args() + + +def test_mlflow_flag_falls_back_to_the_environment(monkeypatch): + monkeypatch.setenv("MLFLOW_TRACKING_URI", "https://mlflow.example.com/") + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + monkeypatch.setattr(getpass, "getuser", lambda: "tester") + hf_ptq = importlib.import_module("hf_ptq") + monkeypatch.setattr( + sys, "argv", ["hf_ptq.py", "--pyt_ckpt_path", "/models/Qwen3-0.6B", "--mlflow"] + ) + + args = hf_ptq.parse_args() + + assert args.mlflow == "https://mlflow.example.com" From 79bc5e607d5bc4901f7442bf8113eb5dceb29f32 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 28 Jul 2026 15:09:40 +0000 Subject: [PATCH 02/10] Move MLflow tracking into modelopt.torch.utils and log the version Promotes the tracking helper from examples/hf_ptq to modelopt.torch.utils.mlflow so other example scripts can record runs the same way. The move required decoupling it from hf_ptq: MlflowRunLogger took an argparse Namespace and read ten PTQ-specific attributes off it, which no other caller could supply. It now takes a tracking URI, an experiment name and an explicit enabled flag, with params, tags and artifacts passed in; default_experiment_name takes (tool, model, variant) rather than inspecting args. hf_ptq keeps the PTQ-specific parts in two small helpers. Uploading the recipe moved to the caller as part of that, which also drops the modelopt.recipe import from the library module -- keeping it would have risked a modelopt.torch.utils -> modelopt.recipe -> modelopt.torch.quantization -> modelopt.torch.utils cycle. The captured log now takes its name from sys.argv[0] instead of hardcoding hf_ptq. Also uploads the ModelOpt version as version.txt. It stays a tag as well: the tag is what makes runs filterable, the artifact is what travels with a downloaded run. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 3 +- examples/hf_ptq/README.md | 7 +- examples/hf_ptq/hf_ptq.py | 66 +++- examples/hf_ptq/mlflow_utils.py | 342 ----------------- modelopt/torch/utils/__init__.py | 1 + modelopt/torch/utils/mlflow.py | 423 +++++++++++++++++++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 97 +++++ tests/examples/hf_ptq/test_mlflow_utils.py | 401 ------------------- tests/unit/torch/utils/test_mlflow.py | 324 ++++++++++++++++ 9 files changed, 912 insertions(+), 752 deletions(-) delete mode 100644 examples/hf_ptq/mlflow_utils.py create mode 100644 modelopt/torch/utils/mlflow.py delete mode 100644 tests/examples/hf_ptq/test_mlflow_utils.py create mode 100644 tests/unit/torch/utils/test_mlflow.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d31c30d049e..ef087780448 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,7 +6,8 @@ Changelog **New Features** -- Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py`` to record a PTQ run on an MLflow server: the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries are uploaded as artifacts, with the model / format / calibration settings as searchable params. The run is opened before the model loads so an unreachable server fails in seconds instead of after calibration, and a failed run is still recorded with its log attached. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. +- Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: it uploads the invocation, the ModelOpt version, the run's log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled. +- Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py``, which records a PTQ run through the above and additionally uploads the resolved recipe (``$import``\ s expanded) and the quantization summaries. The run is opened before the model loads so an unreachable server fails in seconds instead of after calibration, and a failed run is still recorded with its log attached. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. **Backward Breaking Changes** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 8adf755488b..a482392a966 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -312,10 +312,11 @@ seconds rather than after a full calibration. These artifacts are uploaded: | Artifact | Contents | | --- | --- | | `command.txt` | The full invocation, copy-pasteable | +| `version.txt` | The ModelOpt version that ran | | `recipe/resolved_recipe.yaml` | The `--recipe` with its `$import`s expanded, so it stands alone | | `logs/hf_ptq.log` | Everything the run printed, including the traceback if it crashed | | `summary/quant_summary.txt` | The per-quantizer summary (unless `--no-verbose`) | -| `summary/moe.html` | Per-expert calibration token counts, for MoE models | +| `summary/moe.html` | Per-expert calibration token counts, when the run produces them | The model, format, recipe and calibration settings are also logged as searchable params, alongside `user` / `hostname` / `modelopt_version` / `git_sha` tags. A run that fails is @@ -331,6 +332,10 @@ Other flags: Authentication uses MLflow's own environment variables (`MLFLOW_TRACKING_TOKEN`, or `MLFLOW_TRACKING_USERNAME` / `MLFLOW_TRACKING_PASSWORD`). +The tracking itself lives in `modelopt.torch.utils.mlflow` +([`MlflowRunLogger`](../../modelopt/torch/utils/mlflow.py)), so other example scripts can +record runs the same way; `hf_ptq.py` only supplies the params and artifacts specific to PTQ. + > Note: only the main rank uploads, so `--use_fsdp2` runs produce a single run. The log > captures Python output; output written directly by native libraries (NCCL, CUDA) goes to > the terminal only. On SLURM, keep the job's own `.out` file for those. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 2cf68acaf1e..2ac7ee525e0 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -24,6 +24,7 @@ import numpy as np import torch +import yaml from accelerate.hooks import remove_hook_from_module from cast_mxfp4_to_nvfp4 import apply_to_model as apply_cast_mxfp4_to_nvfp4 from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static @@ -46,7 +47,6 @@ setup_distributed_args, validate_fsdp2_supported, ) -from mlflow_utils import MlflowRunLogger, default_experiment_name, validate_tracking_uri from torch.utils.data import DataLoader from transformers import ( AutoConfig, @@ -81,7 +81,12 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) -from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils import ( + MlflowRunLogger, + default_experiment_name, + print_rank_0, + validate_tracking_uri, +) from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, @@ -1653,8 +1658,12 @@ def parse_args() -> argparse.Namespace: try: args.mlflow = validate_tracking_uri(args.mlflow) except ValueError as e: - parser.error(str(e)) - args.mlflow_experiment = args.mlflow_experiment or default_experiment_name(args) + parser.error(f"--mlflow: {e}") + args.mlflow_experiment = args.mlflow_experiment or default_experiment_name( + "hf_ptq", + args.pyt_ckpt_path, + Path(args.recipe).stem if args.recipe else args.qformat, + ) if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") @@ -1691,6 +1700,43 @@ def parse_args() -> argparse.Namespace: return args +def _mlflow_run_inputs(args: argparse.Namespace) -> tuple[dict, dict]: + """Params and start-time artifacts describing this PTQ run.""" + params = { + "model": args.pyt_ckpt_path, + "qformat": args.qformat, + "kv_cache_qformat": args.kv_cache_qformat, + "recipe": args.recipe or "", + "calib_size": args.calib_size, + "calib_seq": args.calib_seq, + "batch_size": args.batch_size, + "sparsity_fmt": args.sparsity_fmt, + "export_path": args.export_path, + "world_size": args.dist_state.world_size, + } + texts = {} + if args.recipe: + # The resolved recipe, not the source file: a recipe may be a directory or use + # $imports, and only the resolved form is self-contained. + resolved = load_recipe(args.recipe).model_dump(mode="json") + texts["recipe/resolved_recipe.yaml"] = yaml.safe_dump(resolved, sort_keys=False) + return params, texts + + +def _mlflow_run_outputs(args: argparse.Namespace) -> dict[str, Path]: + """Summaries written by post_quantize, keyed by artifact path. + + Uploaded without the leading dot, which is awkward to browse in the MLflow UI. Missing + entries are skipped: the MoE table only exists for MoE models, and neither file is + written under ``--no-verbose``. + """ + export_path = Path(args.export_path) + return { + "summary/quant_summary.txt": export_path / ".quant_summary.txt", + "summary/moe.html": export_path / ".moe.html", + } + + def main(args: argparse.Namespace): if not torch.cuda.is_available(): raise OSError("GPU is required for inference.") @@ -1702,8 +1748,14 @@ def main(args: argparse.Namespace): # Opened before the model loads so an unreachable server or a bad experiment name # fails in seconds rather than after a full calibration run. - mlflow_logger = MlflowRunLogger(args) - mlflow_logger.start() + mlflow_logger = MlflowRunLogger( + args.mlflow, + args.mlflow_experiment, + run_name=args.mlflow_run_name, + enabled=bool(args.mlflow) and args.dist_state.is_main, + ) + mlflow_params, mlflow_texts = _mlflow_run_inputs(args) + mlflow_logger.start(params=mlflow_params, texts=mlflow_texts) status = "FAILED" try: @@ -1745,7 +1797,7 @@ def main(args: argparse.Namespace): status = "FINISHED" finally: cleanup_distributed(args) - mlflow_logger.finish(status) + mlflow_logger.finish(status, files=_mlflow_run_outputs(args)) if __name__ == "__main__": diff --git a/examples/hf_ptq/mlflow_utils.py b/examples/hf_ptq/mlflow_utils.py deleted file mode 100644 index 94710873162..00000000000 --- a/examples/hf_ptq/mlflow_utils.py +++ /dev/null @@ -1,342 +0,0 @@ -# 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. - -"""MLflow tracking for ``hf_ptq.py``, enabled by ``--mlflow ``. - -Uploads the invocation, the resolved recipe, the run log and the quantization -summaries to an MLflow tracking server so a PTQ run can be reproduced from its -MLflow entry alone. Everything here is a no-op unless ``--mlflow`` is given. -""" - -import getpass -import logging -import os -import re -import shlex -import shutil -import socket -import subprocess -import sys -import tempfile -import time -from datetime import datetime, timezone -from pathlib import Path -from typing import Any -from urllib.parse import urlparse - -import yaml - -import modelopt -from modelopt.recipe import load_recipe - -__all__ = ["MlflowRunLogger", "TeeStream", "default_experiment_name", "validate_tracking_uri"] - -# MLflow experiment names are stored in a VARCHAR(256) column by the SQL-backed stores. -_MAX_COMPONENT_LEN = 100 -_UNSAFE_CHARS = re.compile(r"[^A-Za-z0-9._-]+") - - -def validate_tracking_uri(uri: str) -> str: - """Validate an MLflow tracking URI and return it without a trailing slash. - - Only ``http(s)`` servers are accepted; MLflow's local ``file:`` / ``sqlite:`` - backends are not a useful destination for a shared PTQ record. - - Raises: - ValueError: if *uri* is empty, has no host, or is not an http(s) URL. - """ - if not uri: - raise ValueError( - "--mlflow requires a tracking URI (e.g. https:///); " - "pass one explicitly or set MLFLOW_TRACKING_URI." - ) - parsed = urlparse(uri) - if parsed.scheme not in ("http", "https"): - raise ValueError( - f"--mlflow expects an http(s) tracking URI, got {uri!r}. " - "Did you mean https://" + uri.lstrip("/") + "?" - ) - if not parsed.netloc: - raise ValueError(f"--mlflow tracking URI {uri!r} has no host.") - return uri.rstrip("/") - - -def default_experiment_name(args) -> str: - """Build the default experiment name, ``/hf_ptq/-``. - - The model component is the checkpoint's basename, so both a local directory and a - ``org/name`` Hugging Face id collapse to the same readable name. - """ - model = Path(args.pyt_ckpt_path).name - variant = Path(args.recipe).stem if args.recipe else args.qformat - return f"{_sanitize(_user())}/hf_ptq/{_sanitize(model)}-{_sanitize(variant)}" - - -def _sanitize(component: str) -> str: - """Reduce one experiment-name component to ``[A-Za-z0-9._-]``.""" - cleaned = _UNSAFE_CHARS.sub("_", component).strip("._-") - return cleaned[:_MAX_COMPONENT_LEN] or "unknown" - - -def _user() -> str: - try: - return getpass.getuser() - except OSError: # container without a passwd entry for the uid - return "unknown" - - -def _git_sha() -> str: - try: - return subprocess.check_output( - ["git", "rev-parse", "--short", "HEAD"], - cwd=Path(__file__).resolve().parent, - stderr=subprocess.DEVNULL, - text=True, - ).strip() - except (subprocess.CalledProcessError, OSError): - return "unknown" - - -def _command_text() -> str: - """The invocation, as a copy-pasteable line.""" - lines = [shlex.join([sys.executable, *sys.argv])] - world_size = int(os.environ.get("WORLD_SIZE", "1")) - if world_size > 1: - lines += [ - "", - f"# Launched under torchrun with WORLD_SIZE={world_size}, " - f"LOCAL_WORLD_SIZE={os.environ.get('LOCAL_WORLD_SIZE', '?')}. The torchrun " - "wrapper is not part of sys.argv and is therefore not shown above.", - ] - return "\n".join(lines) + "\n" - - -class TeeStream: - """Mirror a text stream to *sink* while passing writes through to *stream*. - - ``hf_ptq.py`` reports progress with bare ``print()`` and has no log file; wrapping - ``sys.stdout``/``sys.stderr`` in this is what produces one. Attribute access falls - through to the wrapped stream so ``isatty()`` keeps progress bars behaving. Native - (C-level) writes go straight to the real file descriptor and are *not* captured. - """ - - def __init__(self, stream, sink): - self._stream = stream - self._sink = sink - - def write(self, data: str) -> int: - self._stream.write(data) - if not self._sink.closed: - self._sink.write(data) - return len(data) - - def flush(self) -> None: - self._stream.flush() - if not self._sink.closed: - self._sink.flush() - - def __getattr__(self, name): - return getattr(self._stream, name) - - -class MlflowRunLogger: - """Records one ``hf_ptq.py`` invocation as an MLflow run. - - Disabled -- every method a no-op -- unless ``--mlflow`` was passed and this is the - main rank. :meth:`start` validates the server and opens the run *before* the model - loads, so a bad URI or a missing token fails in seconds rather than after hours of - calibration; :meth:`finish` uploads the run outputs and closes the run. - - Example: - >>> logger = MlflowRunLogger(args) - >>> logger.start() - >>> try: - ... quantize_and_export() - ... status = "FINISHED" - ... finally: - ... logger.finish(status) - """ - - def __init__(self, args): - self.args = args - self.enabled = bool(args.mlflow) and args.dist_state.is_main - self._run: Any = None - self._log_dir: Path | None = None - self._saved_streams: tuple | None = None - self._redirected_handlers: list[tuple[logging.StreamHandler, Any]] = [] - self._start_time = 0.0 - - def start(self) -> None: - """Open the run: capture stdout, verify the server, log the inputs.""" - if not self.enabled: - return - self._start_time = time.time() - self._start_capture() - try: - self._open_run() - self._log_inputs() - except Exception: - self._stop_capture() - raise - - def finish(self, status: str) -> None: - """Upload the run outputs and close the run with *status*. - - Never raises: an MLflow outage must not turn a successful quantization into a - failed exit. - """ - if not self.enabled or self._run is None: - self._stop_capture() - return - try: - self._log_outputs() - except Exception as e: - print(f"[mlflow] WARNING: could not upload run outputs: {e}") - self._stop_capture() - try: - import mlflow - - mlflow.end_run(status=status) - print(f"[mlflow] {status}: {self._run_url()}") - except Exception as e: - print(f"[mlflow] WARNING: could not close the run: {e}") - - def _open_run(self) -> None: - # Optional dependency: only examples using --mlflow need it installed. - try: - import mlflow - except ImportError as e: - raise ImportError("--mlflow requires the 'mlflow' package: pip install mlflow") from e - - self._check_reachable() - mlflow.set_tracking_uri(self.args.mlflow) - mlflow.set_experiment(self.args.mlflow_experiment) - run_name = self.args.mlflow_run_name or datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S") - self._run = mlflow.start_run(run_name=run_name) - print(f"[mlflow] experiment: {self.args.mlflow_experiment}") - print(f"[mlflow] run: {self._run_url()}") - - def _check_reachable(self) -> None: - """Fail fast on an unreachable host. - - Any HTTP response -- including 401 -- means the host is up, so authorization is - left to the first real API call, which reports it precisely. - """ - import requests - - try: - requests.get(f"{self.args.mlflow}/health", timeout=10) - except requests.RequestException as e: - raise ConnectionError(f"MLflow server {self.args.mlflow} is unreachable: {e}") from e - - def _log_inputs(self) -> None: - """Log what the run was given -- upfront, so a crash still leaves a usable record.""" - import mlflow - - args = self.args - mlflow.log_params( - { - "model": args.pyt_ckpt_path, - "qformat": args.qformat, - "kv_cache_qformat": args.kv_cache_qformat, - "recipe": args.recipe or "", - "calib_size": args.calib_size, - "calib_seq": args.calib_seq, - "batch_size": args.batch_size, - "sparsity_fmt": args.sparsity_fmt, - "export_path": args.export_path, - "world_size": args.dist_state.world_size, - } - ) - mlflow.set_tags( - { - "user": _user(), - "hostname": socket.gethostname(), - "modelopt_version": modelopt.__version__, - "git_sha": _git_sha(), - } - ) - mlflow.log_text(_command_text(), "command.txt") - if args.recipe: - # The resolved recipe, not the source file: a recipe may be a directory or - # use $imports, and only the resolved form is self-contained. - resolved = load_recipe(args.recipe).model_dump(mode="json") - mlflow.log_text( - yaml.safe_dump(resolved, sort_keys=False), "recipe/resolved_recipe.yaml" - ) - - def _log_outputs(self) -> None: - """Log what the run produced: the captured log and the quantization summaries.""" - import mlflow - - mlflow.log_metric("total_time_s", time.time() - self._start_time) - sys.stdout.flush() - sys.stderr.flush() - if self._log_dir is not None: - mlflow.log_artifact(str(self._log_dir / "hf_ptq.log"), artifact_path="logs") - export_path = Path(self.args.export_path) - # Written by print_quant_summary / save_expert_token_count_table; hidden names are - # awkward to browse in the MLflow UI, so they are uploaded without the leading dot. - for source, artifact in ( - (".quant_summary.txt", "summary/quant_summary.txt"), - (".moe.html", "summary/moe.html"), - ): - if (export_path / source).is_file(): - text = (export_path / source).read_text(encoding="utf-8", errors="replace") - mlflow.log_text(text, artifact) - - def _start_capture(self) -> None: - self._log_dir = Path(tempfile.mkdtemp(prefix="hf_ptq-mlflow-")) - sink = open(self._log_dir / "hf_ptq.log", "w", buffering=1, encoding="utf-8") - original_stdout, original_stderr = sys.stdout, sys.stderr - self._saved_streams = (original_stdout, original_stderr, sink) - sys.stdout = TeeStream(original_stdout, sink) - sys.stderr = TeeStream(original_stderr, sink) - self._redirect_log_handlers( - {original_stdout: sys.stdout, original_stderr: sys.stderr}, - ) - print(f"[mlflow] capturing this run's log to {self._log_dir / 'hf_ptq.log'}") - - def _redirect_log_handlers(self, replacements: dict) -> None: - """Point already-configured logging handlers at the tee. - - transformers and huggingface_hub bind ``sys.stderr`` into a ``StreamHandler`` when - they are imported, which happens long before the capture starts; without this their - warnings -- rate limits, deprecations -- reach the console but never the log. - """ - self._redirected_handlers = [] - loggers = [logging.getLogger(), *logging.Logger.manager.loggerDict.values()] - for logger in loggers: - for handler in getattr(logger, "handlers", []): - if isinstance(handler, logging.StreamHandler) and handler.stream in replacements: - self._redirected_handlers.append((handler, handler.stream)) - handler.setStream(replacements[handler.stream]) - - def _stop_capture(self) -> None: - if self._saved_streams is None: - return - for handler, stream in self._redirected_handlers: - handler.setStream(stream) - self._redirected_handlers = [] - sys.stdout, sys.stderr, sink = self._saved_streams - sink.close() - self._saved_streams = None - if self._log_dir is not None: - shutil.rmtree(self._log_dir, ignore_errors=True) - self._log_dir = None - - def _run_url(self) -> str: - info = self._run.info - return f"{self.args.mlflow}/#/experiments/{info.experiment_id}/runs/{info.run_id}" diff --git a/modelopt/torch/utils/__init__.py b/modelopt/torch/utils/__init__.py index a38c80cac01..1fa04666e2a 100644 --- a/modelopt/torch/utils/__init__.py +++ b/modelopt/torch/utils/__init__.py @@ -23,6 +23,7 @@ from .list import * from .logging import * from .loss_mask import * +from .mlflow import * from .network import * from .perf import * from .regex import * diff --git a/modelopt/torch/utils/mlflow.py b/modelopt/torch/utils/mlflow.py new file mode 100644 index 00000000000..44c909b136b --- /dev/null +++ b/modelopt/torch/utils/mlflow.py @@ -0,0 +1,423 @@ +# 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. + +"""Record a script run on an MLflow tracking server. + +Lets an example script upload its invocation, configuration, log and outputs so the run can +be reproduced from its MLflow entry alone. ``mlflow`` is an optional dependency, imported +only once tracking is actually enabled. +""" + +import getpass +import logging +import os +import re +import shlex +import shutil +import socket +import subprocess # nosec B404 +import sys +import tempfile +import time +from collections.abc import Mapping +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath +from typing import Any +from urllib.parse import urlparse + +import modelopt + +__all__ = [ + "MlflowRunLogger", + "TeeStream", + "current_user", + "default_experiment_name", + "validate_tracking_uri", +] + +# MLflow experiment names are stored in a VARCHAR(256) column by the SQL-backed stores. +_MAX_COMPONENT_LEN = 100 +_UNSAFE_CHARS = re.compile(r"[^A-Za-z0-9._-]+") + + +def validate_tracking_uri(uri: str) -> str: + """Validate an MLflow tracking URI and return it without a trailing slash. + + Only ``http(s)`` servers are accepted; MLflow's local ``file:`` / ``sqlite:`` backends + are not a useful destination for a shared record of a run. + + Args: + uri: The tracking URI to validate, e.g. ``https://mlflow.example.com/``. + + Returns: + The URI with any trailing slash removed. + + Raises: + ValueError: If *uri* is empty, has no host, or is not an http(s) URL. + """ + if not uri: + raise ValueError( + "MLflow tracking URI is empty; pass one explicitly or set MLFLOW_TRACKING_URI." + ) + parsed = urlparse(uri) + if parsed.scheme not in ("http", "https"): + raise ValueError( + f"MLflow tracking URI must be http(s), got {uri!r}. " + "Did you mean https://" + uri.lstrip("/") + "?" + ) + if not parsed.netloc: + raise ValueError(f"MLflow tracking URI {uri!r} has no host.") + return uri.rstrip("/") + + +def default_experiment_name(tool: str, model: str, variant: str, user: str | None = None) -> str: + """Build an experiment name of the form ``//-``. + + Only the basename of *model* is used, so a local checkpoint directory and an + ``org/name`` Hugging Face id collapse to the same readable name. Each component is + reduced to ``[A-Za-z0-9._-]`` so the ``/`` separators stay meaningful. + + Args: + tool: Name of the script producing the run, e.g. ``"hf_ptq"``. + model: Checkpoint path or Hugging Face model id. + variant: What distinguishes this run of *tool* on *model*, e.g. a recipe name + or a quantization format. + user: Owner of the run. Defaults to the current user. + + Returns: + The experiment name. + + Example: + >>> default_experiment_name("hf_ptq", "/models/Qwen3-0.6B/", "nvfp4", user="alice") + 'alice/hf_ptq/Qwen3-0.6B-nvfp4' + """ + owner = user if user is not None else current_user() + return ( + f"{_sanitize(owner)}/{_sanitize(tool)}/{_sanitize(Path(model).name)}-{_sanitize(variant)}" + ) + + +def current_user() -> str: + """Return the current username, or ``"unknown"`` if the uid has no passwd entry.""" + try: + return getpass.getuser() + except OSError: # container without a passwd entry for the uid + return "unknown" + + +def _sanitize(component: str) -> str: + """Reduce one experiment-name component to ``[A-Za-z0-9._-]``.""" + cleaned = _UNSAFE_CHARS.sub("_", component).strip("._-") + return cleaned[:_MAX_COMPONENT_LEN] or "unknown" + + +def _git_sha() -> str: + """Short commit of the ModelOpt source, or ``"unknown"`` outside a checkout.""" + try: + return subprocess.check_output( # nosec B603 B607 + ["git", "rev-parse", "--short", "HEAD"], + cwd=Path(__file__).resolve().parent, + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (subprocess.CalledProcessError, OSError): + return "unknown" + + +def _command_text() -> str: + """The invocation, as a copy-pasteable line.""" + lines = [shlex.join([sys.executable, *sys.argv])] + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size > 1: + lines += [ + "", + f"# Launched under torchrun with WORLD_SIZE={world_size}, " + f"LOCAL_WORLD_SIZE={os.environ.get('LOCAL_WORLD_SIZE', '?')}. The torchrun " + "wrapper is not part of sys.argv and is therefore not shown above.", + ] + return "\n".join(lines) + "\n" + + +class TeeStream: + """Mirror a text stream to *sink* while passing writes through to *stream*. + + Scripts that report progress with bare ``print()`` have no log file; wrapping + ``sys.stdout``/``sys.stderr`` in this is what produces one. Attribute access falls + through to the wrapped stream so ``isatty()`` keeps progress bars behaving. Native + (C-level) writes go straight to the real file descriptor and are *not* captured. + """ + + def __init__(self, stream, sink): + """Wrap *stream*, mirroring everything written to it into the open file *sink*.""" + self._stream = stream + self._sink = sink + + def write(self, data: str) -> int: + """Write to both the original stream and the sink.""" + self._stream.write(data) + if not self._sink.closed: + self._sink.write(data) + return len(data) + + def flush(self) -> None: + """Flush both the original stream and the sink.""" + self._stream.flush() + if not self._sink.closed: + self._sink.flush() + + def __getattr__(self, name): + return getattr(self._stream, name) + + +class MlflowRunLogger: + """Record one script invocation as an MLflow run. + + :meth:`start` verifies the server and opens the run *before* the expensive work begins, + so a bad URI or a missing token fails in seconds rather than after hours; it also + uploads the invocation and any configuration passed to it, which keeps a crashed run + useful. :meth:`finish` uploads the captured log plus any outputs and closes the run. + Everything is a no-op when ``enabled`` is false, so callers need no branching. + + While the run is open, ``stdout``/``stderr`` are teed to a file that is uploaded as + ``logs/