Add optional MLflow tracking to hf_ptq.py - #2023
Conversation
Records a PTQ run on an MLflow server so it can be reproduced from its MLflow entry alone. Enabled by --mlflow <tracking-uri>; 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/<checkpoint basename>-<recipe name, or --qformat when no recipe>; 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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds optional MLflow run tracking to the HF PTQ example, including CLI configuration, artifact and log capture, failure recording, client-only dependencies, shared stream handling, tests, documentation, and changelog updates. ChangesMLflow PTQ tracking
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PTQScript
participant MlflowRunLogger
participant MLflowServer
PTQScript->>MlflowRunLogger: Start tracking before model loading
MlflowRunLogger->>MLflowServer: Validate tracking URI and create run
MlflowRunLogger->>MLflowServer: Log parameters, tags, recipe, and captured logs
PTQScript->>MlflowRunLogger: Finish with status and summary artifacts
MlflowRunLogger->>MLflowServer: Upload outputs and close run
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2023 +/- ##
===========================================
+ Coverage 66.83% 77.84% +11.01%
===========================================
Files 519 520 +1
Lines 58916 59830 +914
===========================================
+ Hits 39376 46575 +7199
+ Misses 19540 13255 -6285
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
Fixes found reviewing #2023. The main one: gathering the MLflow inputs re-reads the recipe, and the previous commit hoisted that call out of the enabled check, so every hf_ptq run parsed the recipe a second time and printed an extra "[load_recipe] loading:" line -- even without --mlflow, which the feature is supposed to leave untouched. The wiring moves into _start_mlflow_run so the guard is in one place and can be tested; a run without --mlflow is back to a single load_recipe line. Also: - Only suggest "Did you mean https://..." when the URI has no scheme. Suggesting https://sqlite:///mlflow.db was nonsense. - Cap the assembled experiment name, not just each component: four capped components could still produce 403 characters against the VARCHAR(256) the comment cites. - Declare the optional mlflow dependency as an extra now that the module ships in modelopt/, and point the ImportError at it. - TeeStream.__getattr__ raises for _stream/_sink instead of recursing when they are unset; start() is idempotent; finish() is nested in try/finally so a cleanup_distributed failure cannot strand a run in RUNNING. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
|
||
| def _mlflow_run_inputs(args: argparse.Namespace) -> tuple[dict, dict]: | ||
| """Params and start-time artifacts describing this PTQ run.""" | ||
| params = { |
There was a problem hiding this comment.
can this table be auto generated?
| 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 | |
There was a problem hiding this comment.
can we hide this table by default to make the readme visualization shorter.
| ] | ||
|
|
||
| all = ["nvidia-modelopt[onnx,hf,puzzletron,mlflow]"] | ||
| dev = ["nvidia-modelopt[all,dev-docs,dev-lint,dev-test]"] |
There was a problem hiding this comment.
let's add mlflow to dev as well.
No behavior change; 442 -> 417 lines. - Hold the mlflow module on the logger after the guarded import instead of repeating the lazy import in four methods. Safe because every user of it runs after _open_run, which is also what sets _run. - Track one _log_path instead of a directory plus a separately derived name, which was rebuilding the same path in three places. - Fold the duplicated text-upload loop into _log_texts, so command.txt, version.txt and caller-supplied text all take one path. - Drop Args/Returns docstring sections that restated the signature (D417 is not enforced here), keeping the parts that carry rationale the code does not: why only http(s) is accepted, why texts upload at start, why absent files are skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/hf_ptq/hf_ptq.py (1)
1768-1814: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winKeep the MLflow status and distributed cleanup consistent on teardown failures.
_start_mlflow_run()is outside thetry, so its failure bypassescleanup_distributed(). Also, ifcleanup_distributed()raises after line 1807,finish()still receives"FINISHED"even though the process exits with an error. Route startup through the cleanup guard and setstatus = "FAILED"before re-raising teardown failures.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/hf_ptq/hf_ptq.py` around lines 1768 - 1814, Move _start_mlflow_run(args) inside the existing teardown-protected flow so startup failures also pass through cleanup_distributed(). In the nested cleanup handling, set status to "FAILED" before re-raising any cleanup_distributed() exception, ensuring mlflow_logger.finish receives the failed status when teardown fails while preserving successful completion behavior.
🧹 Nitpick comments (1)
examples/hf_ptq/requirements.txt (1)
4-4: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep MLflow optional for the example environment.
This makes every
hf_ptqrequirements install pull MLflow despite the lazy import andnvidia-modelopt[mlflow]extra. Remove this baseline requirement and document/install the optional extra only for--mlflowusers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/hf_ptq/requirements.txt` at line 4, Remove the baseline mlflow requirement from the hf_ptq example dependencies, and document or install the existing nvidia-modelopt[mlflow] optional extra only when users invoke the --mlflow flow. Preserve the lazy import and keep non-MLflow installations free of MLflow.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/hf_ptq/hf_ptq.py`:
- Around line 1657-1661: Update the --mlflow validation flow around
validate_tracking_uri to reject tracking URIs containing userinfo credentials
such as username or password before assigning the validated URI to args.mlflow.
Raise a ValueError so parser.error reports the invalid option, while preserving
valid credential-free URI handling and directing authentication through
environment-based configuration.
In `@examples/hf_ptq/README.md`:
- Around line 314-317: Update the logs/hf_ptq.log artifact description in the
README to state that it contains Python stdout/stderr captured by the tee, while
explicitly excluding NCCL, CUDA, and other native-library output; keep the
surrounding artifact descriptions unchanged.
In `@modelopt/torch/utils/mlflow.py`:
- Line 30: Remove the # nosec B404 suppression from the subprocess import and
the # nosec B603 B607 suppression from the subprocess invocation around the git
rev-parse logic. Update the implementation as needed so Bandit passes without
inline security suppressions, while preserving the existing fixed, non-shell
command behavior.
- Around line 279-337: Update start()’s exception path after _open_run() to
close any successfully created MLflow run with FAILED status before re-raising
the original exception. Reuse the existing cleanup behavior from finish(),
ensure capture is stopped even if run closure fails, and leave self._run
consistent so a partial startup cannot remain RUNNING.
In `@tests/unit/torch/utils/test_mlflow.py`:
- Around line 79-86: Move the in-function requests imports from fake_mlflow and
test_unreachable_server_fails_before_the_work_starts to the module-level
imports, then continue using that shared module import for monkeypatching
requests.get. No justification comment is needed once the imports are moved.
---
Outside diff comments:
In `@examples/hf_ptq/hf_ptq.py`:
- Around line 1768-1814: Move _start_mlflow_run(args) inside the existing
teardown-protected flow so startup failures also pass through
cleanup_distributed(). In the nested cleanup handling, set status to "FAILED"
before re-raising any cleanup_distributed() exception, ensuring
mlflow_logger.finish receives the failed status when teardown fails while
preserving successful completion behavior.
---
Nitpick comments:
In `@examples/hf_ptq/requirements.txt`:
- Line 4: Remove the baseline mlflow requirement from the hf_ptq example
dependencies, and document or install the existing nvidia-modelopt[mlflow]
optional extra only when users invoke the --mlflow flow. Preserve the lazy
import and keep non-MLflow installations free of MLflow.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 43c45989-cdeb-4b8e-8f51-0f87815c28b4
📒 Files selected for processing (9)
CHANGELOG.rstexamples/hf_ptq/README.mdexamples/hf_ptq/hf_ptq.pyexamples/hf_ptq/requirements.txtmodelopt/torch/utils/__init__.pymodelopt/torch/utils/mlflow.pypyproject.tomltests/examples/hf_ptq/test_hf_ptq_args.pytests/unit/torch/utils/test_mlflow.py
| if args.mlflow is not None: | ||
| try: | ||
| args.mlflow = validate_tracking_uri(args.mlflow) | ||
| except ValueError as e: | ||
| parser.error(f"--mlflow: {e}") |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
python - <<'PY'
from urllib.parse import urlparse
uri = "https://user:token@mlflow.example.com"
parsed = urlparse(uri)
assert parsed.netloc
assert parsed.username == "user"
assert parsed.password == "token"
print("Credential-bearing URI is accepted by urlparse.")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 208
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the validation helper and the MLflow code paths.
rg -n "validate_tracking_uri|mlflow|sys\.argv|tracking_uri|set_tracking_uri|start_run|log_params|log_artifact|log_text" examples/hf_ptq/hf_ptq.py
# Show the relevant sections around the validation and logging.
sed -n '1600,1835p' examples/hf_ptq/hf_ptq.pyRepository: NVIDIA/Model-Optimizer
Length of output: 10323
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the helper definition if it lives elsewhere in the repo.
rg -n "def validate_tracking_uri|validate_tracking_uri\(" -S .Repository: NVIDIA/Model-Optimizer
Length of output: 700
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the helper and the MLflow setup/logging block.
FILE=examples/hf_ptq/hf_ptq.py
# Locate the helper definition range.
nl -ba "$FILE" | sed -n '1,220p' | rg -n "validate_tracking_uri|def main|mlflow"
# Print the helper definition and the MLflow setup region.
nl -ba "$FILE" | sed -n '1,260p'
nl -ba "$FILE" | sed -n '1760,1835p'Repository: NVIDIA/Model-Optimizer
Length of output: 196
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE=modelopt/torch/utils/mlflow.py
# Show the validator and any related tests.
sed -n '1,180p' "$FILE"
printf '\n--- tests ---\n'
sed -n '90,150p' tests/unit/torch/utils/test_mlflow.pyRepository: NVIDIA/Model-Optimizer
Length of output: 8851
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the MLflow logger implementation for any logging/serialization of the tracking URI.
sed -n '1,260p' modelopt/torch/utils/mlflow.pyRepository: NVIDIA/Model-Optimizer
Length of output: 10183
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,180p' modelopt/torch/utils/mlflow.py
printf '\n--- tests ---\n'
sed -n '90,150p' tests/unit/torch/utils/test_mlflow.pyRepository: NVIDIA/Model-Optimizer
Length of output: 8851
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find every use of run_url / tracking_uri in the MLflow helper.
rg -n "run_url|tracking_uri|command.txt|sys\.argv|log_artifact|log_text|texts\[|params\[|artifact" modelopt/torch/utils/mlflow.py
# Print the start/finish methods and any code that writes command or URL artifacts.
sed -n '180,360p' modelopt/torch/utils/mlflow.pyRepository: NVIDIA/Model-Optimizer
Length of output: 9836
Reject credentials in MLflow tracking URIs. https://user:token@host passes validate_tracking_uri(), then gets echoed in [mlflow] run: ... and captured in the uploaded run log. Strip or reject userinfo before enabling --mlflow; use environment-based auth instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/hf_ptq.py` around lines 1657 - 1661, Update the --mlflow
validation flow around validate_tracking_uri to reject tracking URIs containing
userinfo credentials such as username or password before assigning the validated
URI to args.mlflow. Raise a ValueError so parser.error reports the invalid
option, while preserving valid credential-free URI handling and directing
authentication through environment-based configuration.
Source: Path instructions
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/utils/mlflow.py`:
- Around line 341-346: Update the command artifact construction in the MLflow
logging method around _log_texts to sanitize _command_text() before uploading
command.txt. Redact sensitive CLI option values and credentials embedded in
tracking URI userinfo, while preserving non-sensitive command content and the
existing version.txt/texts behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1d4721de-0bcf-4c84-978a-21e56c25a95b
📒 Files selected for processing (2)
modelopt/torch/utils/mlflow.pytests/unit/torch/utils/test_mlflow.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unit/torch/utils/test_mlflow.py
- Derive the MLflow params from the parsed args instead of listing ten by hand, so a new hf_ptq flag is tracked without touching the list and the run records the whole configuration rather than a chosen subset. The exclusions are the tracking settings themselves, which describe the destination rather than the run, plus dist_state, which is an object; world_size is pulled out of it explicitly. Param keys are now the argparse names, so model becomes pyt_ckpt_path -- consistent with command.txt, and a key pastes straight back into a CLI flag. 36 params for a default run, against MLflow's batch limit of 100. - Collapse the artifact table in the README behind <details> to keep the page short. MD033 is disabled and <details> is already used elsewhere in examples/. - List mlflow explicitly in the dev extra. It was already reachable via all, but stating it survives someone later trimming that list. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Six findings from the automated review. Read the git sha out of .git instead of shelling out to git rev-parse. The repo's security guidelines forbid # nosec as a Bandit bypass, so the suppressions had to go; removing the subprocess call removes the finding itself rather than arguing for an exception, and needs no codeowner sign-off. Handles loose refs, packed-refs and a detached HEAD, and still returns "unknown" outside a checkout. Close the run when start() fails after start_run() has succeeded. The except block only tore down the capture, so a network blip during log_params left the run RUNNING on the server forever -- the caller gets an exception rather than a logger to call finish() on. This was a real bug; deriving params from every CLI argument made the window wider. Mask credentials before anything is uploaded or printed: --*token* style option values, and user:pass@ embedded in any URI, in command.txt, in the params and in the printed run URL. A credential-bearing tracking URI still works -- only its echo is masked. This matters more now that params are derived from every argument, since a future secret-bearing flag would otherwise be uploaded with no code change to catch it. Also: qualify the log artifact as Python stdout/stderr in the README, since native library output is excluded, and move the requests import in the tests to module scope. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/utils/test_mlflow.py`:
- Line 393: Move the _redact_argv import at
tests/unit/torch/utils/test_mlflow.py:393-393 into the module-level MLflow
utility imports, and move the modelopt.torch.utils.mlflow import at
tests/unit/torch/utils/test_mlflow.py:457-459 to module scope so both imports
occur during test collection.
- Around line 375-395: Update _redact_argv so recognized secret options always
mask the immediately following token, including values beginning with “-”,
rather than excluding dash-prefixed values. Add the parameterized test case for
["--hf_token", "-secret"] expecting ["--hf_token", "***"], while preserving
existing option and URI redaction behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ae52dfb2-13d1-4df6-8a9f-5bafd58fac61
📒 Files selected for processing (3)
examples/hf_ptq/README.mdmodelopt/torch/utils/mlflow.pytests/unit/torch/utils/test_mlflow.py
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/hf_ptq/README.md
- modelopt/torch/utils/mlflow.py
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Design review (protocol applied)
Problem: PTQ runs today leave no durable record — command, recipe, log and summaries live only in whatever terminal/SLURM file the user kept — so a finished quantization can't be reproduced or compared later.
Alternatives I checked:
- In-repo:
modelopt/torch/utils/logging.pyalready owns all the stdout/stderr plumbing in this package (capture_io,no_stdout,atomic_print).TeeStreamis a fourth, unrelated mechanism for the same surface and would fit better there. The repo also already has an MLflow convention on the eval side (.agents/skills/evaluation/references/nel-next.md,.agents/skills/accessing-mlflow/SKILL.md):experiment_name: <user>/<model>, tagsmodel/checkpoint_path/framework, plus aninvocation_idID convention. This PR invents a second scheme ($USER/hf_ptq/<basename>-<variant>, tagsuser/hostname/modelopt_version/git_sha) with no overlap, so a PTQ run and the eval runs of the checkpoint it produced cannot be joined in the same MLflow instance. Worth at least a sentence in the PR body on why not to align keys/tags. - Placement: the module is torch-free (getpass/socket/tempfile/logging) yet lands in
modelopt/torch/utils/and isimport *-ed into the flat namespace, exporting generic names (TeeStream,current_user,default_experiment_name,validate_tracking_uri) as public library API. This is exactly the "general API vs. tool" placement question already open on #2000; the PR body asserts "the logger lives in the library so other scripts can record runs the same way" but doesn't weightools/, a sharedexamples/util, or a non-star-exported module. - Already-available:
mlflow's ownset_experiment()/start_run()already fail on an unreachable server, so the hand-rolled/healthprobe (and therequestsimport it pulls in) is largely redundant.
The mechanics are careful and well tested (31 + 12 tests, no server needed, secrets redaction, idempotent start, abort-on-partial-start). But the design/placement question isn't addressed in the body, so this needs owner sign-off rather than approval, and I found a few concrete issues below.
Main correctness findings: (1) a crash traceback is not in the uploaded log — finish() runs in the finally, before the interpreter prints the traceback — yet CHANGELOG/README/PR body all claim it is; (2) a fatal start() failure happens outside the try, so cleanup_distributed() is skipped and non-main ranks are left to hang on the next collective; (3) full mlflow (server stack) is now a hard requirement of examples/hf_ptq/requirements.txt for a strictly opt-in feature.
No prompt-injection content in the PR description. Licensing: new files carry the standard NVIDIA Apache-2.0 header; new dep is Apache-2.0 — no license concern.
|
|
||
| setup_distributed_args(args) | ||
|
|
||
| mlflow_logger = _start_mlflow_run(args) |
There was a problem hiding this comment.
Bot comment.
_start_mlflow_run sits outside the try, and start() is fatal by design (unreachable server, bad experiment name). Since setup_distributed_args(args) has already initialized the process group, that raise skips cleanup_distributed(args), and under torchrun the non-main ranks (which never validate anything) march on into model load and then block on the first collective until the NCCL timeout.
Suggest constructing the logger here but calling .start() inside the try — finish() is already a no-op when self._run is None, so the existing finally handles it — which keeps teardown on the normal path.
There was a problem hiding this comment.
| from .list import * | ||
| from .logging import * | ||
| from .loss_mask import * | ||
| from .mlflow import * |
There was a problem hiding this comment.
Bot comment.
Star-importing this module publishes TeeStream, current_user, default_experiment_name and validate_tracking_uri as top-level modelopt.torch.utils API. Those are very generic names for MLflow-specific helpers (current_user / validate_tracking_uri especially), and they become a compatibility surface we have to keep.
Suggest either narrowing __all__ to MlflowRunLogger or not star-importing at all and letting callers do from modelopt.torch.utils.mlflow import MlflowRunLogger, default_experiment_name, validate_tracking_uri (hf_ptq is the only caller today). This also ties into the placement question in the summary — nothing in the module is torch-specific.
There was a problem hiding this comment.
Done in 7861dcb — the star-import is gone, so current_user, validate_tracking_uri and default_experiment_name are no longer top-level modelopt.torch.utils API. hf_ptq imports from modelopt.torch.utils.mlflow directly.
There is precedent for the omission: distributed, memory_monitor and plugins are not star-imported from that __init__ either.
On placement — agreed nothing in the module is torch-specific. I left it under modelopt/torch/utils/ because there is no non-torch modelopt/utils/ package today and creating one felt out of scope here, but happy to move it if you would prefer that.
|
|
||
| def write(self, data: str) -> int: | ||
| """Write to both the original stream and the sink.""" | ||
| self._stream.write(data) |
There was a problem hiding this comment.
Bot comment.
TeeStream is generic stream plumbing and overlaps in spirit with modelopt/torch/utils/logging.py, which already owns every other sys.stdout manipulation in this package (capture_io, no_stdout, atomic_print). Consider moving it there (and keeping mlflow.py about MLflow) so there's one place to look for stdout redirection semantics — capture_io in particular is the near-neighbour that redirects instead of tees.
There was a problem hiding this comment.
Done in 7861dcb — TeeStream now lives in modelopt/torch/utils/logging.py next to capture_io and no_stdout, and is exported from its __all__. mlflow.py imports it from there. Agreed that one place for stdout-redirection semantics is better, and it shortened mlflow.py usefully.
| "mlflow", | ||
| ] | ||
|
|
||
| all = ["nvidia-modelopt[onnx,hf,puzzletron,mlflow]"] |
There was a problem hiding this comment.
lets not make mlflow requirement part of all. It adds a lot of heavy dependencies so best to leave it explicitly opt-in or part of dev
There was a problem hiding this comment.
Done in 05d0a6b — all is back to ["nvidia-modelopt[onnx,hf,puzzletron]"], with a comment recording why mlflow is deliberately excluded.
Also switched the extra itself to mlflow-skinny>=2.9, so even the opt-in install no longer pulls Flask/gunicorn/sqlalchemy/alembic/docker. Verified skinny is sufficient rather than assuming it: installed it alone (no full mlflow present) and ran a real PTQ against our tracking server — all five artifacts uploaded.
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Three inline comments below. Everything else I found was minor cleanup.
|
|
||
| setup_distributed_args(args) | ||
|
|
||
| mlflow_logger = _start_mlflow_run(args) |
There was a problem hiding this comment.
cleanup_distributed() is skipped when MLflow startup fails.
This call sits between setup_distributed_args(args) and the try, and it raises by design — ImportError when mlflow is absent, ConnectionError when the server is unreachable. When it does, rank 0 exits with the process group still up, so cleanup_distributed(args) in the finally never runs.
Fail-fast before load_model() is the right call; it just wants to be inside the try. Constructing the logger inert first keeps mlflow_logger bound for the finally:
mlflow_logger = MlflowRunLogger(
args.mlflow, args.mlflow_experiment,
run_name=args.mlflow_run_name,
enabled=bool(args.mlflow) and args.dist_state.is_main,
)
status = "FAILED"
try:
_start_mlflow_run(mlflow_logger, args) # still before load_model()
...Minor, same area: files=_mlflow_run_outputs(args) is evaluated inside the finally, so if it ever raised it would replace the in-flight exception. It can't today (--export_path defaults to "exported_model", so Path() can't fail), but computing it before the try costs nothing and removes the hazard.
There was a problem hiding this comment.
Fixed in 7861dcb and restructured further in 05d0a6b. main() now wraps the tracked section in a context manager inside the try, so cleanup_distributed() runs whether or not opening the run fails:
try:
with _mlflow_run(args):
...
finally:
cleanup_distributed(args)Your secondary point is fixed as a side effect: files=_mlflow_run_outputs(args) used to be evaluated inside the finally, and the artifact paths are now computed before the block — so there is no longer any way for them to replace an in-flight exception.
Three bugs, then cleanups. The uploaded log never contained the crash traceback. finish() runs from the caller's finally, which is before the interpreter prints the traceback to a sys.stderr that is no longer teed, so the log stopped at the last line the script printed -- contradicting the README, the CHANGELOG and the PR body. finish() now appends the active exception to the sink, and to the sink only, so the console still shows it exactly once. start() sat outside main()'s try. It is fatal by design, and the process group is already up by then, so a bad URI skipped cleanup_distributed and left the other ranks blocked on the first collective until the two-hour NCCL timeout. The logger is built outside the try and opened inside it; finish() is already inert when no run was opened. _redact_argv left a secret unmasked when the value itself started with "-". The pending-secret branch now masks unconditionally; an option in that position would mean no value was passed, which argparse rejects anyway. Cleanups: mlflow-skinny rather than the full distribution, which pulled flask, gunicorn, sqlalchemy, alembic, docker and pandas into every hf_ptq install for an opt-in feature (verified skinny uploads artifacts to an HTTP tracking server); stop star-importing the module from modelopt.torch.utils, which published current_user and validate_tracking_uri as top-level API; move TeeStream to logging.py beside capture_io and no_stdout; emit model/checkpoint_path tags so a PTQ run and the evaluations of the checkpoint it produced can be correlated; test import hygiene. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/hf_ptq/hf_ptq.py (1)
1815-1821: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMark distributed-cleanup failures as failed runs.
After Line 1814, an exception from
cleanup_distributed()still callsfinish("FINISHED"). MLflow records success and skips traceback capture although the process exits with an error.Proposed fix
finally: try: cleanup_distributed(args) + except BaseException: + status = "FAILED" + raise finally: # Nested so a failure tearing down the process group cannot leave the run # stuck in RUNNING on the server. mlflow_logger.finish(status, files=_mlflow_run_outputs(args))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/hf_ptq/hf_ptq.py` around lines 1815 - 1821, Update the cleanup flow around cleanup_distributed() so any exception during distributed cleanup changes the run status to failed and preserves the exception for traceback capture before mlflow_logger.finish is called. Keep the nested finally guarantee that finish executes, while ensuring successful cleanup continues to use the existing status.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/utils/test_mlflow.py`:
- Around line 461-467: Update test_git_sha_is_read_without_a_subprocess to
actively prevent process execution while calling _git_sha(), such as patching
subprocess invocation APIs to raise and asserting the call still succeeds.
Remove reliance on checking module globals, while preserving the existing SHA
format validation.
---
Outside diff comments:
In `@examples/hf_ptq/hf_ptq.py`:
- Around line 1815-1821: Update the cleanup flow around cleanup_distributed() so
any exception during distributed cleanup changes the run status to failed and
preserves the exception for traceback capture before mlflow_logger.finish is
called. Keep the nested finally guarantee that finish executes, while ensuring
successful cleanup continues to use the existing status.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bdbae99e-76aa-4c4d-82c3-b1014dd89b30
📒 Files selected for processing (7)
examples/hf_ptq/hf_ptq.pyexamples/hf_ptq/requirements.txtmodelopt/torch/utils/logging.pymodelopt/torch/utils/mlflow.pypyproject.tomltests/examples/hf_ptq/test_hf_ptq_args.pytests/unit/torch/utils/test_mlflow.py
| warnings.showwarning = original_showwarning | ||
|
|
||
|
|
||
| class TeeStream: |
There was a problem hiding this comment.
This feels like an agent shortcut hack to me. I'd much rather just replace print() in hf_ptq.py with a proper logger. That would give us all of the benefits of logging library over print including the stream redirection you want to do for mlflow.
There was a problem hiding this comment.
Fair challenge, and I agree logging would be better for hf_ptq on its own merits. Two reasons I would rather not fold it into this PR:
It would not remove the tee. tqdm writes progress bars straight to sys.stderr, and transformers / huggingface_hub both print and log. A logging.FileHandler captures neither. The tee — plus the pre-bound-handler redirect in _redirect_log_handlers — is what got huggingface_hub's rate-limit warnings into the captured log at all; before that they reached the console and vanished. So a logging conversion would sit alongside this, not replace it.
Scope. 34 print() calls in hf_ptq.py, 86 across examples/hf_ptq/, in a PR already around +1400.
Happy to do the conversion as a follow-up PR — it improves hf_ptq independently of tracking. Let me know if you would rather it block this one.
| 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")) |
There was a problem hiding this comment.
What about just logging the environment separately from the command?
There was a problem hiding this comment.
Decided against this one, on the author's call, and I think the reasoning is worth recording: environment variables are exactly where credentials live — MLFLOW_TRACKING_TOKEN, HF_TOKEN, cloud keys — and this module's own docstring tells users to pass credentials via the environment. Uploading the environment would undermine the advice the module gives.
The WORLD_SIZE / LOCAL_WORLD_SIZE note stays inline in command.txt, where it exists only to explain that the torchrun wrapper is not part of sys.argv and so is not shown in the command above it.
| shutil.copy2(local, staged) | ||
| self._mlflow.log_artifact(str(staged), artifact_path=directory) | ||
|
|
||
| def _start_capture(self) -> None: |
There was a problem hiding this comment.
See earlier logger comment.
There was a problem hiding this comment.
Replied on the logging.py thread — short version: agreed in principle, but a logging conversion would not remove the need for the tee (tqdm writes directly to sys.stderr, and the HF libraries both print and log), and it is ~86 print() calls across examples/hf_ptq/. Offering it as a follow-up PR.
Keep mlflow out of `all`. nvidia-modelopt[all] is the documented install, so the tracking client does not belong there for an opt-in example feature; `dev` keeps it, listed explicitly. Pin a mlflow-skinny>=2.9 floor. Fix _git_sha() in git worktrees, where it reported "unknown" every time. A worktree's .git is a file naming the real git directory, and while HEAD lives there the ref itself lives in the checkout named by commondir -- two levels of indirection that `git rev-parse` handled for free before it was replaced to clear the Bandit finding. The test now asserts a real sha for a plain checkout, a worktree and a detached HEAD, rather than accepting "unknown" as it did before and thereby hiding this. Add MlflowRunLogger.track(), a context manager mirroring mlflow.start_run(), and use it from hf_ptq. The block's outcome sets the status, which removes the manual FAILED/FINISHED bookkeeping and the nested try/finally. It also means the closing artifact paths are computed before the block rather than inside a finally, where a raise would have replaced the in-flight exception. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the MLflow tracking PR (+1457/-25, 9 files). All of the previously flagged correctness bugs now look genuinely fixed; what's left is a design/placement decision that two codeowners raised and the author deliberately deferred, so this needs a human sign-off rather than an approval.
Verified as addressed (verified against the current diff, not just the replies):
- Crash traceback now reaches the uploaded log —
finish()calls_note_active_exception()(writestraceback.format_exc()into the tee's sink only, before_log_outputs), with two tests (test_failed_run_uploads_the_traceback,test_track_marks_a_raising_block_failed) and a negative test for the FINISHED path. cleanup_distributed()no longer skipped:with _mlflow_run(args):is inside thetry, so a fatalstart()still unwinds through thefinally;files=is computed before the block, so it can't replace an in-flight exception.- Partial-start orphan run fixed (
_abort_run()ends the run FAILED, tested). - Secret redaction:
command.txt, params andrun_urlmasked; dash-prefixed secret values now masked; URI userinfo masked; parametrized tests cover all cases. # nosecsuppressions gone;_git_sha()reads.gitdirectly and now handles worktrees (gitdir:file +commondirindirection) and detached HEAD, with real tests for all three (the old test that accepted"unknown"is fixed).TeeStreammoved intomodelopt/torch/utils/logging.py; themlflowstar-import is out ofmodelopt/torch/utils/__init__.py.- Deps:
mlflow-skinny>=2.9, removed fromall, kept indev— matches kevalmorabia97/Edwardf0t1. - README table collapsed into
<details>, log-artifact wording qualified, params table auto-derived fromvars(args). - Licensing: new files use the canonical 2026
LICENSE_HEADERverbatim; new dep is Apache-2.0. No license concern.
Why this needs an owner look rather than an approve:
- 💬 chadvoegele asked to replace
hf_ptq.py's bareprint()with a realloggerinstead of teeing streams ("feels like an agent shortcut hack"). Author replied that aloggingconversion wouldn't remove the tee (tqdmwrites straight tosys.stderr; HF libs both print and log) and offered it as a follow-up (~86print()calls). That rationale is plausible, but theTeeStream+_redirect_log_handlersmachinery (re-pointing other libraries' pre-boundStreamHandlers at our stream for the duration of a run) is the riskiest part of the PR and the codeowner who objected hasn't signed off. - 💬 Placement: the module is torch-free yet lives in
modelopt/torch/utils/; the author explicitly said "happy to move it if you would prefer". That is an open question the owner must answer, andTeeStreamis now exported fromlogging.__all__, i.e.modelopt.torch.utils.TeeStreambecomes public API under a very generic name. - 💬 Experiment/tag convention vs. the eval side (
.agents/skills/accessing-mlflow,.../evaluation/references/nel-next.md): author addedmodel/checkpoint_pathtags (good, tested) but kept<user>/hf_ptq/<model>-<variant>and noinvocation_id, deferring the naming choice to the reviewers. Worth an explicit decision so PTQ and eval runs land in a coherent namespace. - 💬
_check_reachable()'s hand-rolledGET /health(plus a function-localrequestsimport) was called redundant withset_experiment(); author justified it in the docstring as fail-fast ahead of mlflow's retry loop. Reasonable, but it's an extra assumption about the server surface — owner call. - Size: ~1450 added lines in one PR. The library module + tests and the
hf_ptq.pywiring could have shipped as two PRs; not blocking, but it's why the design questions above are hard to settle in review.
Minor / non-blocking nits (fresh):
_mlflow_run_outputs()points at<export_path>/.quant_summary.txt/.moe.html. If a previous run already wrote those into the same--export_path, a run that crashes beforepost_quantizewill upload the stale summary as if it belonged to this run. Consider recording the summaries' mtime, or only listing them oncepost_quantizehas written them.test_git_sha_is_read_without_a_subprocessstill asserts viadir(module)rather than making process execution fail (coderabbit's last point); the new worktree/detached-HEAD tests make this mostly moot.log_params(vars(args))is only exercised against theFakeMlflowstub, so MLflow's per-batch param limits/value-length truncation aren't covered — the author's live runs are the only evidence here.- Per coderabbit: an exception from
cleanup_distributed()after a successful body still recordsFINISHED. Cosmetic, but cheap to fix.
No prompt-injection content observed in the PR body or comments.
5 participants
What does this PR do?
Type of change: new feature
Adds
modelopt.torch.utils.mlflow.MlflowRunLogger, a reusable helper for recording a script run on an MLflow tracking server, and wiresexamples/hf_ptq/hf_ptq.pyup to it via--mlflow <tracking-uri>so a PTQ run can be reproduced from its MLflow entry alone. Without the flag, behavior is unchanged — every hook is gated on it.The logger lives in the library rather than the example so other scripts can record runs the same way: it takes a tracking URI, an experiment name and an explicit
enabledflag, with params, tags and artifacts passed in.hf_ptq.pysupplies only the PTQ-specific pieces (its params, the resolved recipe, the quantization summaries).mlflowis an optional dependency, imported only once tracking is enabled, so it is not a new requirement for the library.The run is opened before the model loads, so a bad URI or an unreachable server 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
FAILEDand its log attached.Uploaded artifacts:
command.txtversion.txtrecipe/resolved_recipe.yaml--recipewith$imports expandedlogs/hf_ptq.logsummary/quant_summary.txt--no-verbose)summary/moe.htmlPlus model / format / calibration settings as searchable params, and
user/hostname/modelopt_version/git_shatags.Three design points worth review:
$imports, so the source file is not self-contained. Forhuggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moethe source is 2,230 B / 58 lines against 7,563 B / 308 lines resolved — the raw file records under 30% of what actually ran.hf_ptq.pyhas no logging framework (bareprint()), so the log is produced by teeing stdout/stderr. Handlers that libraries bound tosys.stderrat import time are re-pointed at the tee for the run's duration and handed back afterwards; without that,transformers/huggingface_hubwarnings reach the console but never the log. Native (C-level) output is still not captured — documented in the README.modelopt.recipeout ofmodelopt.torch.utils, which would otherwise risk amodelopt.torch.utils→modelopt.recipe→modelopt.torch.quantization→modelopt.torch.utilsimport cycle.Only the main rank uploads, so
--use_fsdp2runs produce a single run.Usage
--mlflow_experimentand--mlflow_run_nameoverride the defaults ($USER/hf_ptq/<basename>-<recipe name or --qformat>, and the UTC start time). Passing--mlflowwith no value uses$MLFLOW_TRACKING_URI. Authentication uses MLflow's own env vars.Testing
Unit — 31 tests in
tests/unit/torch/utils/test_mlflow.pyfor the library, plus 12 intests/examples/hf_ptq/test_hf_ptq_args.pyfor the hf_ptq wiring. CPU-only, no network and nomlflowdependency (driven against a stub module). Covers experiment-name derivation and sanitization, URI accept/reject, tee pass-through, the pre-bound-handler redirect, artifact renaming, skipping absent optional outputs, the disabled path, andversion.txt. 63 tests pass together with the existingtest_hf_ptq_args.py/test_example_utils.py.Hardware — real PTQ runs against a live MLflow server:
FINISHED, all artifacts, sane post-quant generationsw4a16_nvfp4_fp8_at_6p0bits-active_moe, 2×B200FINISHEDin 63 min, search hiteffective bits: 6.00; 106 KB log capturing every per-layer decision, 4.4 MB quant summaryFINISHEDFINISHED, all five artifacts includingversion.txt--mlflowafter the review fixes[load_recipe]line and 0[mlflow]lines, confirming the untracked path is untouchedFAILEDrecorded with log + traceback attached, summaries correctly absentargparsewith aDid you mean https://…?hint--mlflowCoverage gap, stated plainly:
summary/moe.htmlis verified only against a synthetic file (unit test + a real upload). It could not be produced naturally —expert_token_countbuffers live on_QuantSparseSequentialMoe, while Qwen3.5/3.6 experts take the fused_QuantFusedExpertspath, so no such file is written for these models regardless of--moe_calib_experts_ratio. The uploader's conditional is correct; the branch simply had no natural input available here.Before your PR is "Ready for review"
--mlflowmeans no behavior change.CONTRIBUTING.md: ✅ — addsmlflowas an optional extra inpyproject.toml(nvidia-modelopt[mlflow], folded intoall) and toexamples/hf_ptq/requirements.txt. Apache-2.0 (permissive). Imported lazily, so it is not required to install or import ModelOpt. No code copied from other sources./claude reviewnot yet run. A self-review was done first and its six findings are fixed in the third commit (the notable one: gathering the MLflow inputs re-read the recipe on every run, including without--mlflow).Additional Information
The one deliberate coverage gap is
summary/moe.html, described under Testing: no model available here takes the sparse-sequential MoE path that writes it, so it is covered by unit test and a synthetic upload rather than a natural one. The uploader treats it as an optional output and skips it when absent, which is exercised by test.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
FAILEDstatus and diagnostic logs.Documentation