mlflow.py

# 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 argparse
import contextlib
import getpass
import json
import logging
import os
import re
import shlex
import shutil
import socket
import sys
import tempfile
import time
import traceback
import warnings
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from typing import Any
from urllib.parse import urlparse

import yaml

import modelopt
from modelopt.torch.utils.logging import TeeStream

__all__ = [
    "EXPERIMENT_JSON",
    "TRACKING_URI_ENV",
    "MlflowRunLogger",
    "add_mlflow_args",
    "checkpoint_run_tags",
    "command_text",
    "current_user",
    "default_experiment_name",
    "drop_experiment_json",
    "mask_tracking_uri",
    "masked_args",
    "resolve_mlflow_args",
    "resolve_tracking_uri",
    "resolved_recipe_texts",
    "track_run",
    "validate_tracking_uri",
]

# MLflow experiment names are stored in a VARCHAR(256) column by the SQL-backed stores. The
# per-component cap stops one pathological component from crowding out the others; the name
# cap is what actually keeps the result storable.
_MAX_COMPONENT_LEN = 100
_MAX_NAME_LEN = 250
_UNSAFE_CHARS = re.compile(r"[^A-Za-z0-9._-]+")

# Anything uploaded or printed passes through _redact first: a tracking URI may carry
# ``user:token@`` and a caller's own flags may carry a secret.
_SECRET_NAME = re.compile(r"token|api[-_]?key|password|passwd|secret|credential", re.IGNORECASE)
_URI_USERINFO = re.compile(r"(?<=://)[^/\s@]+(?=@)")
_MASK = "***"

# Provenance pointer a checkpoint carries, dotted like the other sidecars a quantization
# drops next to the weights so a checkpoint loader ignores it and it does not look like part
# of the model.
EXPERIMENT_JSON = ".experiment.json"

# MLflow's own variable, so a shell that already exports it opts in without a flag. Public
# because the vLLM example republishes the resolved URI under it for its worker processes.
TRACKING_URI_ENV = "MLFLOW_TRACKING_URI"


def _stat_key(path: Path) -> tuple[int, int] | None:
    """Identity of a file's contents-in-time, or ``None`` when it does not exist."""
    try:
        stat = path.stat()
    except OSError:
        return None
    return (stat.st_mtime_ns, stat.st_size)


def _redact(value: Any) -> Any:
    """Mask credentials embedded in a URI, leaving non-strings untouched."""
    return _URI_USERINFO.sub(_MASK, value) if isinstance(value, str) else value


def _redact_argv(argv: list[str]) -> list[str]:
    """Mask the value of any ``--*token*`` style option, and credentials in any URI."""
    redacted: list[str] = []
    mask_next = False
    for token in argv:
        if mask_next:
            # Unconditionally, since a secret may itself start with "-"; an option there
            # instead would mean the caller passed no value, which argparse rejects anyway.
            redacted.append(_MASK)
        elif token.startswith("-") and _SECRET_NAME.search(token):
            option, sep, _ = token.partition("=")
            redacted.append(option + sep + _MASK if sep else option)
        else:
            redacted.append(_redact(token))
        mask_next = (
            token.startswith("-") and _SECRET_NAME.search(token) is not None and "=" not in token
        )
    return redacted


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.

    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"):
        message = f"MLflow tracking URI must be http(s), got {uri!r}."
        if not parsed.scheme:
            # Only a bare host is plausibly a forgotten scheme; suggesting https://sqlite:///...
            # for a URI that already has one would be nonsense.
            message += f" Did you mean https://{uri.lstrip('/')}?"
        raise ValueError(message)
    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 ``<user>/<tool>/<model>-<variant>``.

    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; *variant* is whatever
    distinguishes this run of *tool* on *model*, such as a recipe name or a quantization
    format. Each component is reduced to ``[A-Za-z0-9._-]`` so the ``/`` separators stay
    meaningful, and *user* defaults to the current user.

    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()
    name = (
        f"{_sanitize(owner)}/{_sanitize(tool)}/{_sanitize(Path(model).name)}-{_sanitize(variant)}"
    )
    return name[:_MAX_NAME_LEN]


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 git checkout.

    Read out of ``.git`` rather than by shelling out to ``git``, which keeps the library
    free of subprocess use. Handles worktrees, where ``.git`` is a file pointing at the
    real git directory and refs live in the main checkout alongside it.
    """
    try:
        git_path = Path(__file__).resolve().parents[3] / ".git"
        if git_path.is_file():
            git_dir = Path(git_path.read_text().split("gitdir:", 1)[1].strip())
        else:
            git_dir = git_path
        head = (git_dir / "HEAD").read_text().strip()
        if not head.startswith("ref: "):
            return head[:9]  # detached HEAD
        ref = head.removeprefix("ref: ")
        # A worktree keeps HEAD locally but shares refs with the checkout named by commondir.
        bases = [git_dir]
        commondir = git_dir / "commondir"
        if commondir.is_file():
            bases.append((git_dir / commondir.read_text().strip()).resolve())
        for base in bases:
            if (base / ref).is_file():
                return (base / ref).read_text().strip()[:9]
            packed = base / "packed-refs"
            if packed.is_file():
                for line in packed.read_text().splitlines():
                    sha, _, name = line.partition(" ")
                    if name.strip() == ref:
                        return sha[:9]
    except (OSError, IndexError):
        pass
    return "unknown"


def command_text(argv: list[str] | None = None) -> str:
    """The invocation, as a copy-pasteable line, with credentials masked.

    *argv* defaults to this process's own ``sys.argv``. Pass another process's argv when the
    run is opened somewhere the user never typed a command -- a worker subprocess, whose own
    ``sys.argv`` is an implementation detail rather than a reproducible invocation.
    """
    lines = [shlex.join([sys.executable, *_redact_argv(sys.argv if argv is None else 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 MlflowRunLogger:
    """Record one script invocation as an MLflow run.

    :meth:`start` opens the run *before* the expensive work begins, so a bad URI, a missing
    token or an unreachable server fails there 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/<script>.log``. Logging handlers that libraries bound to ``sys.stderr`` at import
    time are re-pointed at the tee for the duration and handed back afterwards.

    Failures after the run is open are reported as warnings and never raised: losing a
    tracking server must not turn a successful job into a failed one.

    Note:
        ``command.txt`` masks ``--*token*``-style option values and credentials embedded in
        a URI, but the captured log is whatever the script printed, so a secret echoed to
        stdout still reaches the server. Prefer passing credentials via the environment.

    *tracking_uri* must already be validated (see :func:`validate_tracking_uri`),
    *experiment_name* is created if absent, *run_name* defaults to the UTC start time
    ``YYYYmmdd-HHMMSS``, and ``enabled=False`` makes every method a no-op -- which is how
    callers skip non-main ranks or an absent flag. ``required=False`` additionally downgrades
    a failure to open the run into a warning: use it when tracking was inferred from the
    environment rather than asked for, so an uninstalled client or an unreachable server
    cannot take the job down with it.

    Example:
        >>> logger = MlflowRunLogger(uri, "alice/hf_ptq/Qwen3-0.6B-nvfp4")
        >>> logger.start(params={"model": ckpt}, texts={"config.yaml": config_yaml})
        >>> status = "FAILED"
        >>> try:
        ...     quantize_and_export()
        ...     status = "FINISHED"
        ... finally:
        ...     logger.finish(status, files={"summary/report.txt": report_path})
    """

    def __init__(
        self,
        tracking_uri: str,
        experiment_name: str,
        run_name: str | None = None,
        enabled: bool = True,
        required: bool = True,
    ):
        """Configure the run without contacting the server; see the class docstring."""
        self.tracking_uri = tracking_uri
        self.experiment_name = experiment_name
        self.run_name = run_name
        self.enabled = enabled
        self.required = required
        self._mlflow: Any = None
        self._run: Any = None
        self._log_path: Path | None = None
        self._saved_streams: tuple | None = None
        self._tees: tuple | None = None
        self._file_stats: dict[str, tuple[int, int] | None] = {}
        self._start_time = 0.0

    @property
    def run_url(self) -> str:
        """Link to this run in the MLflow UI, or ``""`` before the run is open."""
        if self._run is None:
            return ""
        info = self._run.info
        uri = _redact(self.tracking_uri)
        return f"{uri}/#/experiments/{info.experiment_id}/runs/{info.run_id}"

    @property
    def run_info(self) -> dict[str, str]:
        """Identity of this run on the server, or ``{}`` before it is open.

        Enough for a consumer holding only this run's outputs to find it again: ``run_id``
        is MLflow's own identifier for the run, a uuid4 hex, unique across experiments.
        Every field is read back off the run the server returned rather than off what was
        requested, so a run MLflow resolved differently is reported as it really is.
        """
        if self._run is None:
            return {}
        info = self._run.info
        return {
            "tracking_uri": _redact(self.tracking_uri),
            "experiment_name": self.experiment_name,
            "experiment_id": str(info.experiment_id),
            "run_id": str(info.run_id),
            "run_name": getattr(info, "run_name", None) or self.run_name or "",
            "run_url": self.run_url,
        }

    def start(
        self,
        params: dict[str, Any] | None = None,
        tags: dict[str, Any] | None = None,
        texts: dict[str, str] | None = None,
        files: Mapping[str, Path | str] | None = None,
    ) -> None:
        """Open the run: capture output, verify the server, upload the inputs.

        *params* are searchable; *tags* merge over the defaults (user, hostname, ModelOpt
        version and commit); *texts* maps artifact path to content, uploaded here rather
        than at the end so it survives a crash. *files* names the outputs the run is
        expected to produce, so :meth:`finish` can tell them from files that were already
        there -- pass the same mapping to both.

        Opening the run is the readiness check: it is MLflow's own first request, so it
        honours the client's TLS and retry configuration rather than second-guessing it.
        Set ``MLFLOW_HTTP_REQUEST_MAX_RETRIES`` to shorten the wait on a dead host.

        Raises:
            ImportError: If ``mlflow`` is not installed and ``required``.
            Exception: Whatever MLflow raises for an unusable server, if ``required``.
        """
        if not self.enabled or self._run is not None:
            return
        self._start_time = time.time()
        # Keyed through Path on both sides: a caller may pass strings, and "./out/x" and
        # "out/x" are the same file but not the same string.
        self._file_stats = {str(p): _stat_key(p) for p in map(Path, (files or {}).values())}
        self._start_capture()
        try:
            self._open_run()
            self._log_inputs(params, tags, texts)
        except Exception as e:
            # start_run() may already have succeeded, and the caller gets an exception
            # rather than a logger to call finish() on, so close the run here.
            self._abort_run()
            self._stop_capture()
            if self.required:
                raise
            self.enabled = False
            print(f"[mlflow] WARNING: tracking disabled, continuing without it ({e})")

    @contextmanager
    def track(
        self,
        params: dict[str, Any] | None = None,
        tags: dict[str, Any] | None = None,
        texts: dict[str, str] | None = None,
        files: Mapping[str, Path | str] | None = None,
        metrics: dict[str, float] | None = None,
    ) -> Iterator["MlflowRunLogger"]:
        """Open the run for the duration of the block, closing it with the right status.

        Mirrors ``mlflow.start_run()``. *files* and *metrics* are uploaded when the block
        exits; naming the paths upfront is fine because only files this run actually wrote
        are uploaded (see :meth:`finish`).

        Example:
            >>> with logger.track(params={"model": ckpt}, files={"summary.txt": report}):
            ...     quantize_and_export()
        """
        self.start(params=params, tags=tags, texts=texts, files=files)
        status = "FAILED"
        try:
            yield self
            status = "FINISHED"
        finally:
            self.finish(status, files=files, metrics=metrics)

    def log_text(self, artifact_path: str, text: str) -> None:
        """Upload *text* as an artifact while the run is open, best-effort.

        For a value that is only settled midway through the run and is worth having even if
        the run later crashes -- the quantization config a calibration is about to apply,
        say. :meth:`start` and :meth:`finish` cover everything known at the two ends.
        """
        if not self.enabled or self._run is None:
            return
        try:
            self._log_texts({artifact_path: text})
        except Exception as e:
            print(f"[mlflow] WARNING: could not upload {artifact_path}: {e}")

    def log_experiment_json(self, checkpoint_dir: Path | str | None = None) -> None:
        """Record which MLflow run produced a checkpoint, on the server and in the checkpoint.

        Tags point from a run to the checkpoint it wrote; this is the reverse, so a checkpoint
        found on disk can be traced back to the run that produced it without searching the
        server. The artifact goes up for any run that opened, so a failure is traceable from
        the server side too.

        *checkpoint_dir* also writes the JSON there as :data:`EXPERIMENT_JSON`. Pass it only
        once the checkpoint is really on disk, since the file claims authorship of the weights
        sitting next to it: an output directory existing proves nothing, as it may hold a
        checkpoint from an earlier attempt whose weights this run never touched.

        After the checkpoint is written, the pointer beside it is this run's or absent --
        never a previous run's. So a run that never opened *removes* the pointer rather than
        leaving one: tracking can disable itself mid-flight (an unreachable server or an
        uninstalled client, which a URI inherited from the environment tolerates by design),
        and the caller's untracked cleanup was skipped because tracking looked configured.
        """
        info = self.run_info
        if not info:
            if checkpoint_dir is not None:
                drop_experiment_json(checkpoint_dir)
            return
        text = json.dumps(info, indent=2) + "\n"
        self.log_text(EXPERIMENT_JSON.removeprefix("."), text)
        if checkpoint_dir is None:
            return
        target = Path(checkpoint_dir) / EXPERIMENT_JSON
        try:
            target.write_text(text)
        except OSError as e:
            print(f"[mlflow] WARNING: could not write {target}: {e}")

    def _abort_run(self) -> None:
        """End a run that failed before :meth:`start` returned, so it is not left RUNNING."""
        if self._run is None:
            return
        try:
            self._mlflow.end_run(status="FAILED")
        except Exception as e:
            print(f"[mlflow] WARNING: could not close the interrupted run: {e}")
        self._run = None

    def finish(
        self,
        status: str,
        texts: dict[str, str] | None = None,
        files: Mapping[str, Path | str] | None = None,
        metrics: dict[str, float] | None = None,
    ) -> None:
        """Upload the run's outputs and close it with *status*, e.g. ``"FINISHED"``.

        *texts* and *files* both map artifact path to content, from memory and from disk
        respectively. A *files* entry is skipped when its file is absent, or was last
        modified before the run started -- so callers can list optional outputs, and a run
        that produced none of them does not upload a previous run's leftovers.
        *metrics* merges over the default ``total_time_s``.
        """
        if not self.enabled or self._run is None:
            self._stop_capture()
            return
        if status != "FINISHED":
            self._note_active_exception()
        try:
            self._log_outputs(texts, files, metrics)
        except Exception as e:
            print(f"[mlflow] WARNING: could not upload run outputs: {e}")
        self._stop_capture()
        try:
            self._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 _note_active_exception(self) -> None:
        """Append the exception being handled to the captured log.

        :meth:`finish` runs from the caller's ``finally``, which is *before* the interpreter
        prints the traceback to ``sys.stderr`` -- no longer teed by then -- so the log would
        otherwise stop at the last line the script printed. Written to the file only, so the
        console still shows the traceback exactly once.
        """
        if sys.exc_info()[0] is None or self._saved_streams is None:
            return
        sink = self._saved_streams[2]
        if not sink.closed:
            sink.write("\n" + traceback.format_exc())

    def _open_run(self) -> None:
        try:
            import mlflow  # optional dependency: only needed once tracking is enabled
        except ImportError as e:
            raise ImportError(
                "MLflow tracking requires the 'mlflow' package: pip install nvidia-modelopt[mlflow]"
            ) from e

        self._mlflow = mlflow
        mlflow.set_tracking_uri(self.tracking_uri)
        mlflow.set_experiment(self.experiment_name)
        # Settled here rather than passed straight through, so run_info reports the name the
        # run actually carries.
        self.run_name = self.run_name or datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
        self._run = mlflow.start_run(run_name=self.run_name)
        print(f"[mlflow] experiment: {self.experiment_name}\n[mlflow] run: {self.run_url}")

    def _log_inputs(self, params, tags, texts) -> None:
        if params:
            self._mlflow.log_params(
                {k: _MASK if _SECRET_NAME.search(k) else _redact(v) for k, v in params.items()}
            )
        self._mlflow.set_tags(
            {
                "user": current_user(),
                "hostname": socket.gethostname(),
                "modelopt_version": modelopt.__version__,
                "git_sha": _git_sha(),
                **(tags or {}),
            }
        )
        # The version is a tag as well, for searching; the artifact travels with the run.
        self._log_texts(
            {
                "command.txt": command_text(),
                "version.txt": f"{modelopt.__version__}\n",
                **(texts or {}),
            }
        )

    def _log_outputs(self, texts, files, metrics) -> None:
        elapsed = time.time() - self._start_time
        self._mlflow.log_metrics({"total_time_s": elapsed, **(metrics or {})})
        self._log_texts(texts)
        sys.stdout.flush()
        sys.stderr.flush()
        if self._log_path is not None:
            self._log_file(f"logs/{self._log_path.name}", self._log_path)
        for artifact_path, local in (files or {}).items():
            path = Path(local)
            if not path.is_file():
                continue
            # Only what this run produced: an export directory is commonly reused across
            # attempts, so a run that crashes early would otherwise upload the previous
            # run's summary as its own. Compared against the stat taken when the run opened
            # rather than against the wall clock, whose resolution outruns the filesystem's.
            if str(path) in self._file_stats and self._file_stats[str(path)] == _stat_key(path):
                continue
            self._log_file(artifact_path, path)

    def _log_texts(self, texts) -> None:
        for artifact_path, text in (texts or {}).items():
            self._mlflow.log_text(text, artifact_path)

    def _log_file(self, artifact_path: str, local: Path) -> None:
        """Upload *local* to *artifact_path*, staging a copy when it must be renamed."""
        target = PurePosixPath(artifact_path)
        directory = str(target.parent) if str(target.parent) != "." else None
        if local.name == target.name:
            self._mlflow.log_artifact(str(local), artifact_path=directory)
            return
        # log_artifact keeps the local basename, so rename via a staged copy rather than
        # reading the file into memory -- these can be hundreds of MB.
        with tempfile.TemporaryDirectory() as staging:
            staged = Path(staging) / target.name
            shutil.copy2(local, staged)
            self._mlflow.log_artifact(str(staged), artifact_path=directory)

    def _start_capture(self) -> None:
        script = Path(sys.argv[0]).stem or "run"
        self._log_path = Path(tempfile.mkdtemp(prefix="modelopt-mlflow-")) / f"{script}.log"
        sink = open(self._log_path, "w", buffering=1, encoding="utf-8")
        stdout, stderr = sys.stdout, sys.stderr
        self._saved_streams = (stdout, stderr, sink)
        self._tees = (TeeStream(stdout, sink), TeeStream(stderr, sink))
        sys.stdout, sys.stderr = self._tees
        self._repoint_handlers({stdout: self._tees[0], stderr: self._tees[1]})
        print(f"[mlflow] capturing this run's log to {self._log_path}")

    @staticmethod
    def _repoint_handlers(replacements: dict) -> None:
        """Move already-configured logging handlers from one stream to another.

        transformers and huggingface_hub bind ``sys.stderr`` into a ``StreamHandler`` when
        they are imported, long before the capture starts; without this their warnings reach
        the console but never the log. Scanning again on the way out -- rather than replaying
        a list captured on the way in -- also hands back handlers a library bound *during*
        the run, so nothing is left pointing at the tee once its file is closed.
        """
        loggers = [logging.getLogger(), *list(logging.Logger.manager.loggerDict.values())]
        for logger in loggers:
            for handler in list(getattr(logger, "handlers", [])):
                if not isinstance(handler, logging.StreamHandler):
                    continue
                if handler.stream not in replacements:
                    continue
                # logging._StderrHandler exposes ``stream`` as a read-only property that
                # already resolves to whatever sys.stderr currently is, so it follows the tee
                # on its own and cannot -- and must not -- be repointed.
                with contextlib.suppress(AttributeError):
                    handler.setStream(replacements[handler.stream])

    def _stop_capture(self) -> None:
        if self._saved_streams is None:
            return
        stdout, stderr, _ = self._saved_streams
        if self._tees is not None:
            self._repoint_handlers({self._tees[0]: stdout, self._tees[1]: stderr})
            self._tees = None
        sys.stdout, sys.stderr, sink = self._saved_streams
        sink.close()
        self._saved_streams = None
        if self._log_path is not None:
            shutil.rmtree(self._log_path.parent, ignore_errors=True)
            self._log_path = None


# The CLI surface below is shared by the example scripts that offer tracking, so a run is
# configured the same way and named by the same convention whichever script opened it.

_ENV_HELP = (
    f"MLflow's own ${TRACKING_URI_ENV} 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."
)

_TRACKS_HELP = (
    "Track this run on an MLflow server (e.g. https://<your-mlflow-server>/), uploading the "
    "command, the resolved configuration, the run log and the run's summaries."
)


def add_mlflow_args(
    parser: argparse.ArgumentParser,
    tool: str,
    tracks: str = _TRACKS_HELP,
    variant_help: str = "recipe name, or the quantization format",
) -> None:
    """Add ``--mlflow``, ``--mlflow_experiment`` and ``--mlflow_run_name`` to *parser*.

    *tool* names the script in the default experiment ``<user>/<tool>/<model>-<variant>`` (see
    :func:`default_experiment_name`), *tracks* is the leading description of ``--mlflow`` --
    what this particular script uploads -- and *variant_help* says what the script derives the
    variant from. Pair with :func:`resolve_mlflow_args`.

    The multi-word flags are registered under both the underscored and the dashed spelling:
    vLLM's ``FlexibleArgumentParser`` rewrites every ``--foo_bar`` on the command line to
    ``--foo-bar`` before matching, so the dashed spelling has to exist for the flag to be
    reachable there at all, and a user moving between the example scripts should not have to
    remember which spelling each one took.
    """
    parser.add_argument("--mlflow", default=None, help=f"{tracks} {_ENV_HELP}")
    parser.add_argument(
        "--mlflow_experiment",
        "--mlflow-experiment",
        default=None,
        help=f"MLflow experiment name. Default: $USER/{tool}/<model basename>-<{variant_help}>.",
    )
    parser.add_argument(
        "--mlflow_run_name",
        "--mlflow-run-name",
        default=None,
        help="MLflow run name. Default: the UTC start time as YYYYmmdd-HHMMSS.",
    )


def mask_tracking_uri(uri: str | None) -> str | None:
    """Mask any ``user:token@`` a tracking URI carries, for printing.

    Credentials in the URI are a supported form, so everything this module prints or uploads
    masks them -- ``command.txt``, the logged params, :attr:`MlflowRunLogger.run_url`. A
    caller that prints the URI itself (a script echoing its parsed arguments, say) has to do
    the same, or the secret reaches a console log that is routinely archived.
    """
    return _redact(uri)


def resolve_tracking_uri(
    uri: str | None, parser: argparse.ArgumentParser
) -> tuple[str | None, bool]:
    """Settle the tracking URI from ``--mlflow`` and the environment.

    Returns ``(uri or None, required)``, where *required* records that the flag was passed.
    Only the flag is a deliberate request, so only the flag is fatal when the URI is unusable:
    the environment variable is commonly exported for unrelated tooling and must not fail a
    job that would otherwise have worked.
    """
    # An empty value is not a deliberate request: ``--mlflow "$UNSET_VAR"`` is a wrapper
    # script whose variable did not resolve, so it neither names a server nor should take a
    # job down. It falls back like an absent flag -- but loudly, since the caller did ask
    # for tracking and what it gets is whatever the environment names, or nothing.
    required = bool(uri)
    if not uri:
        if uri is not None:
            warnings.warn(
                f"--mlflow was given an empty value; falling back to ${TRACKING_URI_ENV} "
                "if it is set, and running untracked otherwise."
            )
        uri = os.environ.get(TRACKING_URI_ENV) or None
        if uri is None:
            return None, required
    try:
        return validate_tracking_uri(uri), required
    except ValueError as e:
        if required:
            parser.error(f"--mlflow: {e}")  # exits
        warnings.warn(f"Ignoring ${TRACKING_URI_ENV}, continuing untracked: {e}")
        return None, required


def resolve_mlflow_args(
    args: argparse.Namespace,
    parser: argparse.ArgumentParser,
    tool: str,
    model: str,
    variant: str,
) -> None:
    """Settle where tracking is configured from, and name the experiment, in place.

    Sets ``args.mlflow`` to the validated URI or ``None``, ``args.mlflow_required`` to whether
    the flag asked for it, and defaults ``args.mlflow_experiment`` from *tool*, *model* and
    *variant*. Pair with :func:`add_mlflow_args`.
    """
    args.mlflow, args.mlflow_required = resolve_tracking_uri(args.mlflow, parser)
    if args.mlflow:
        args.mlflow_experiment = args.mlflow_experiment or default_experiment_name(
            tool, model, variant
        )


def drop_experiment_json(checkpoint_dir: Path | str) -> None:
    """Remove a provenance pointer an untracked export would otherwise inherit.

    A fresh checkpoint written into a reused output directory would keep the previous run's
    pointer, and one produced 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. Call it
    only for a completed export; a failed run leaves whatever checkpoint was already there,
    pointer included.
    """
    stale = Path(checkpoint_dir) / EXPERIMENT_JSON
    try:
        stale.unlink(missing_ok=True)
    except OSError as e:
        print(f"[mlflow] WARNING: could not remove stale {stale}: {e}")


def masked_args(args: argparse.Namespace, attr: str = "mlflow") -> argparse.Namespace:
    """A copy of *args* whose tracking URI cannot leak credentials into a printed namespace.

    For a script that echoes its parsed arguments: a ``user:token@`` in the URI is a supported
    form that this module masks wherever it prints or uploads one, and a job log is routinely
    archived. Uploaded artifacts are unaffected -- :func:`command_text` redacts, and the URI
    is not worth logging as a param.
    """
    return argparse.Namespace(**{**vars(args), attr: mask_tracking_uri(getattr(args, attr, None))})


def checkpoint_run_tags(source_model: str, checkpoint_dir: Path | str) -> dict[str, str]:
    """Tags a quantization run and whatever is later done with the checkpoint it wrote.

    Shared so the two can be found together on one tracking server. ``checkpoint_path`` is
    the checkpoint the run *writes*, because that is what an export or an evaluation is later
    pointed at (NEL takes ``deployment.checkpoint_path``); the input is kept separately. It is
    resolved because a relative path is useless as a join key.
    """
    return {
        "model": Path(source_model).name,
        "checkpoint_path": str(Path(checkpoint_dir).resolve()),
        "source_checkpoint_path": source_model,
    }


def resolved_recipe_texts(recipe: str | None) -> dict[str, str]:
    r"""``{artifact path: content}`` for *recipe*, or ``{}`` when the run used none.

    The resolved recipe, not the source file: a recipe may be a directory or use ``$import``\ s,
    and only the resolved form stands alone.
    """
    if not recipe:
        return {}
    # Lazy import: modelopt.recipe imports modelopt.torch.quantization, which imports this
    # package at top level (circular), as with the other upward imports in modelopt/torch/utils.
    from modelopt.recipe import load_recipe

    resolved = load_recipe(recipe).model_dump(mode="json")
    return {"recipe/resolved_recipe.yaml": yaml.safe_dump(resolved, sort_keys=False)}


@contextmanager
def track_run(
    logger: MlflowRunLogger,
    checkpoint_dir: Path | str,
    is_main: bool,
    exported: Callable[[], bool],
    describe: Callable[[], Mapping[str, Any]] | None = None,
) -> Iterator[MlflowRunLogger]:
    """Track a checkpoint-producing run, keeping its provenance pointer honest either way.

    *logger* is inert unless tracking was configured *and* this is the rank that records it,
    so the caller needs no branching. *checkpoint_dir* is where the run writes its checkpoint
    and *is_main* gates writes every rank would otherwise race on. *exported* is read on the
    way out, not on the way in: only a completed export may claim the checkpoint the pointer
    sits next to, since the directory usually exists before the weights do.

    *describe* returns the keyword arguments for :meth:`MlflowRunLogger.track` (``params``,
    ``tags``, ``texts``, ``files``) and is called only when the run is tracked, so an
    untracked run does not pay for gathering them -- re-reading a recipe, say.

    Example:
        >>> with track_run(logger, args.export_path, is_main, lambda: args.exported, describe):
        ...     quantize_and_export(args)
    """
    path = Path(checkpoint_dir)
    if not logger.enabled:
        try:
            yield logger
        finally:
            if exported() and is_main:
                drop_experiment_json(path)
        return
    with logger.track(**(describe() if describe is not None else {})):
        try:
            yield logger
        finally:
            logger.log_experiment_json(path if exported() else None)