diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bf8e2639e76..afaec6bc805 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -10,6 +10,11 @@ Changelog - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers. Weight scales are an orthogonal axis, selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian`` with that algorithm's own options), so one recipe can combine a weight calibration with this activation policy in a single pass. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. +*Misc* + +- 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``, and honour MLflow's own ``MLFLOW_TRACKING_URI`` so a shell can opt in without changing the command. A tracked run records the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries, with every command-line argument as a searchable param. The run is opened before the model loads, so an unusable server fails there rather than after calibration, and a failed run is still recorded with its traceback 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..e80926bea91 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -19,6 +19,7 @@ This section focuses on Post-training quantization, a technique that reduces mod | Evaluate Accuracy | Evaluate your model's accuracy! | \[[Link](#evaluate-accuracy)\] | | | Exporting Checkpoints | Export to Hugging Face Unified Checkpoint and deploy on TRT-LLM/vLLM/SGLang | \[[Link](#exporting-checkpoints)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/deployment/3_unified_hf.html)\] | | Pre-Quantized Checkpoints | Ready to deploy Hugging Face pre-quantized checkpoints | \[[Link](#pre-quantized-checkpoints)\] | | +| Tracking runs with MLflow | Record a PTQ run on an MLflow server so it can be reproduced from its entry alone | \[[Link](#tracking-runs-with-mlflow)\] | | | Resources | Extra links to relevant resources | \[[Link](#resources)\] | | @@ -638,6 +639,61 @@ After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build c - Deployable on [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM), [vLLM](https://github.com/vllm-project/vllm) and [SGLang](https://github.com/sgl-project/sglang) - More models coming soon! +## Tracking runs with MLflow + +Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow `, to record a PTQ +run on an MLflow server so it 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. + +
+Uploaded artifacts + +| Artifact | Contents | +| --- | --- | +| `command.txt` | The full invocation, copy-pasteable, with credentials masked | +| `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` | The run's Python stdout/stderr, 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, when the run produces them | + +
+ +Every command-line argument is also logged as a searchable param, alongside +`user` / `hostname` / `modelopt_version` / `git_sha` tags. A run that fails is +still recorded, with status `FAILED` and its log attached. + +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`. +- `$MLFLOW_TRACKING_URI` enables tracking on its own; `--mlflow` overrides it. A URI taken + from the environment is best-effort — if the client is missing or the server is + unreachable the run warns and continues untracked, since the variable is often exported + for other tooling. An explicit `--mlflow` fails loudly instead. + +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. + ## Resources - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d74ffb34efb..dc5c6f5d932 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import argparse import copy import glob import hashlib @@ -23,6 +24,7 @@ import shutil import warnings from collections.abc import Callable, Iterable +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from datetime import timedelta from pathlib import Path @@ -30,6 +32,7 @@ import torch import transformers +import yaml from accelerate import infer_auto_device_map, init_empty_weights from accelerate.utils import get_max_memory from safetensors import safe_open @@ -43,6 +46,7 @@ ProcessorMixin, ) +from modelopt.recipe import load_recipe from modelopt.torch.export.model_utils import is_multimodal_model try: @@ -51,6 +55,11 @@ snapshot_download = None from modelopt.torch.utils import distributed as dist_utils +from modelopt.torch.utils.mlflow import ( + MlflowRunLogger, + default_experiment_name, + validate_tracking_uri, +) logger = logging.getLogger(__name__) @@ -1070,3 +1079,129 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str] if isinstance(algo.get("layerwise"), dict) and "checkpoint_dir" in algo["layerwise"]: algo["layerwise"]["checkpoint_dir"] = resolved return quant_cfg, resolved + + +def add_mlflow_args(parser: argparse.ArgumentParser) -> None: + """Add the MLflow tracking flags.""" + parser.add_argument( + "--mlflow", + 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. MLflow's own $MLFLOW_TRACKING_URI enables tracking without this " + "flag, which overrides it. A URI taken from the environment is best-effort: if " + "it is unusable the run warns and continues untracked." + ), + ) + parser.add_argument( + "--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.", + ) + + +def resolve_mlflow_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None: + """Settle where tracking is configured from, and name the experiment.""" + # MLflow's own variable enables tracking on its own; --mlflow overrides it. Only the + # flag is a deliberate request, so only the flag is fatal when the URI is unusable: the + # variable is commonly exported for unrelated tooling and must not fail a quantization. + args.mlflow_required = args.mlflow is not None + args.mlflow = args.mlflow or os.environ.get("MLFLOW_TRACKING_URI") or None + if args.mlflow: + try: + args.mlflow = validate_tracking_uri(args.mlflow) + except ValueError as e: + if args.mlflow_required: + parser.error(f"--mlflow: {e}") + warnings.warn(f"Ignoring MLFLOW_TRACKING_URI, continuing untracked: {e}") + args.mlflow = None + else: + 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, + ) + + +_MLFLOW_NON_PARAM_ARGS = frozenset( + {"dist_state", "mlflow", "mlflow_experiment", "mlflow_required", "mlflow_run_name"} +) + + +def _mlflow_run_inputs(args: argparse.Namespace) -> tuple[dict, dict]: + """Params and start-time artifacts describing this PTQ run.""" + params = {k: v for k, v in vars(args).items() if k not in _MLFLOW_NON_PARAM_ARGS} + # dist_state is an object, so record the one field worth searching on. + params["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_logger(args: argparse.Namespace) -> MlflowRunLogger: + """Build this run's logger; inert unless --mlflow was given and this is the main rank.""" + return MlflowRunLogger( + args.mlflow, + args.mlflow_experiment, + run_name=args.mlflow_run_name, + enabled=bool(args.mlflow) and args.dist_state.is_main, + required=args.mlflow_required, + ) + + +def mlflow_run(args: argparse.Namespace) -> AbstractContextManager: + """Track this invocation for the duration of the block, or do nothing if untracked.""" + logger = _mlflow_logger(args) + if not logger.enabled: + # Gathering the inputs re-reads the recipe, so keep it off the untracked path. + return nullcontext() + params, texts = _mlflow_run_inputs(args) + return logger.track( + params=params, + tags=_mlflow_run_tags(args), + texts=texts, + files=_mlflow_run_outputs(args), + ) + + +def _mlflow_run_tags(args: argparse.Namespace) -> dict[str, str]: + """Tags shared with the evaluation side, so a PTQ run and the evaluations of the + checkpoint it produced can be found together on one tracking server. + + ``checkpoint_path`` is the checkpoint this run *writes*, because that is what an + evaluation is later pointed at (NEL takes ``deployment.checkpoint_path``); the input is + kept separately. It is resolved because ``--export_path`` defaults to a relative path, + which is useless as a join key. + """ + return { + "model": Path(args.pyt_ckpt_path).name, + "checkpoint_path": str(Path(args.export_path).resolve()), + "source_checkpoint_path": args.pyt_ckpt_path, + } + + +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", + } diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6a4bd476984..d4eebc97ddf 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -29,6 +29,7 @@ from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static from example_utils import ( _resolve_model_path, + add_mlflow_args, build_quant_cfg, cleanup_distributed, copy_custom_model_files, @@ -39,9 +40,11 @@ is_enc_dec, is_nemotron_vl, load_mtp_weights, + mlflow_run, mtp_layer_prefixes_from_checkpoint, needs_checkpoint_path_update, resolve_checkpoint_dir, + resolve_mlflow_args, run_nemotron_vl_preview, setup_distributed_args, validate_fsdp2_supported, @@ -1622,7 +1625,11 @@ def parse_args() -> argparse.Namespace: ), ) + add_mlflow_args(parser) + args = parser.parse_args() + resolve_mlflow_args(args, parser) + if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") @@ -1658,6 +1665,8 @@ def parse_args() -> argparse.Namespace: return args +# Derived state and the tracking settings themselves; everything else argparse parsed is a +# parameter of the run. Deriving the list means a new flag is tracked without touching this. def main(args: argparse.Namespace): if not torch.cuda.is_available(): raise OSError("GPU is required for inference.") @@ -1668,31 +1677,17 @@ def main(args: argparse.Namespace): setup_distributed_args(args) try: - # launch a memory monitor to read the currently used GPU memory. - launch_memory_monitor() + # Entered inside the try: opening the run is fatal by design, and skipping + # cleanup_distributed would leave the other ranks blocked on the first collective + # until the NCCL timeout. + with mlflow_run(args): + # launch a memory monitor to read the currently used GPU memory. + launch_memory_monitor() - # Force eager execution for all model types. - torch.compiler.set_stance("force_eager") - - ( - full_model, - language_model, - model_type, - calibration_only, - processor, - tokenizer, - default_padding_side, - default_pad_token, - device, - ) = load_model(args) + # Force eager execution for all model types. + torch.compiler.set_stance("force_eager") - if args.sparsity_fmt != "dense": - # Sparse - sparsity_main(args, full_model, tokenizer, device) - else: - # Quantize - quantize_main( - args, + ( full_model, language_model, model_type, @@ -1702,7 +1697,25 @@ def main(args: argparse.Namespace): default_padding_side, default_pad_token, device, - ) + ) = load_model(args) + + if args.sparsity_fmt != "dense": + # Sparse + sparsity_main(args, full_model, tokenizer, device) + else: + # Quantize + quantize_main( + args, + full_model, + language_model, + model_type, + calibration_only, + processor, + tokenizer, + default_padding_side, + default_pad_token, + device, + ) finally: cleanup_distributed(args) diff --git a/examples/hf_ptq/requirements.txt b/examples/hf_ptq/requirements.txt index e4756ce34c7..42f39af4c04 100644 --- a/examples/hf_ptq/requirements.txt +++ b/examples/hf_ptq/requirements.txt @@ -1,6 +1,7 @@ compressed-tensors fire flash-attn>=2.6.0 +mlflow-skinny>=2.9 psutil transformers_stream_generator zstandard diff --git a/modelopt/torch/utils/logging.py b/modelopt/torch/utils/logging.py index d1a9fc1aef9..85d3b9df18f 100644 --- a/modelopt/torch/utils/logging.py +++ b/modelopt/torch/utils/logging.py @@ -34,6 +34,7 @@ __all__ = [ "DeprecatedError", + "TeeStream", "atomic_print", "capture_io", "no_stdout", @@ -219,5 +220,40 @@ def custom_showwarning(message, category, filename, lineno, file=None, line=None warnings.showwarning = original_showwarning +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): + # Guard the wrapped attributes themselves: __getattr__ runs whenever they are absent + # (during unpickling, or on a copy), and delegating then would recurse forever. + if name in ("_stream", "_sink"): + raise AttributeError(name) + return getattr(self._stream, name) + + class DeprecatedError(NotImplementedError): """Error for deprecated functions.""" diff --git a/modelopt/torch/utils/mlflow.py b/modelopt/torch/utils/mlflow.py new file mode 100644 index 00000000000..07850371297 --- /dev/null +++ b/modelopt/torch/utils/mlflow.py @@ -0,0 +1,524 @@ +# 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 contextlib +import getpass +import logging +import os +import re +import shlex +import shutil +import socket +import sys +import tempfile +import time +import traceback +from collections.abc import 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 modelopt +from modelopt.torch.utils.logging import TeeStream + +__all__ = ["MlflowRunLogger", "current_user", "default_experiment_name", "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 = "***" + + +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 ``//-``. + + 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() -> str: + """The invocation, as a copy-pasteable line.""" + lines = [shlex.join([sys.executable, *_redact_argv(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 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/