Add optional MLflow tracking to the vLLM fake-quant server - #2120
Conversation
📝 WalkthroughWalkthroughThe vLLM fake-quant example now supports optional MLflow tracking. The launcher validates and forwards configuration, rank-zero workers record quantization and serving data, artifacts are uploaded, failures are marked, and shared MLflow utilities support command and text logging. ChangesMLflow tracking integration
Estimated code review effort: 3 (Moderate) | ~25 minutes Mergeability Score: 🔵 Low · up to The optional tracking feature is mergeable with owner follow-up: test coverage should accurately exercise the documented worker-reload path, and the test import should be clarified or moved to module scope. These are bounded test-maintenance risks rather than a demonstrated production behavior problem. Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Launcher
participant RayWorker
participant FakeQuantMlflowTracker
participant MLflow
Launcher->>Launcher: validate and resolve MLflow arguments
Launcher->>RayWorker: forward tracking settings and credentials
RayWorker->>FakeQuantMlflowTracker: start before model loading
FakeQuantMlflowTracker->>MLflow: create run and log configuration
RayWorker->>FakeQuantMlflowTracker: log quantization data and summary
FakeQuantMlflowTracker->>MLflow: upload artifacts
RayWorker->>FakeQuantMlflowTracker: finish after warmup
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
![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: 3
🤖 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/vllm_serve/vllm_mlflow_utils.py`:
- Line 51: Define an explicit module-level __all__ in vllm_mlflow_utils
containing FakeQuantMlflowTracker, MLFLOW_ENV_VARS, add_mlflow_args, and
resolve_mlflow_args as the module’s public API.
- Line 256: Move the local imports of modelopt.torch.quantization (and the
import at the other referenced location) to module scope with the other imports.
If either import must remain deferred due to an optional, circular, or unusually
heavy dependency, keep it local and add a brief comment stating that
justification.
- Line 154: In examples/vllm_serve/vllm_mlflow_utils.py lines 154-154, update
the tracking log in the MLflow launcher to avoid printing the raw URI, using a
redacted URI or generic tracking-enabled message while preserving the experiment
context. In examples/vllm_serve/fakequant_worker.py lines 164-170, sanitize
captured worker stdout/stderr before MLflow upload, removing path-bearing
sensitive output such as MODELOPT_STATE_PATH and other credentials, tokens,
sensitive paths, or proprietary model details.
🪄 Autofix
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: 6d633ebd-878b-423e-8237-1c440172e06d
📒 Files selected for processing (9)
CHANGELOG.rstexamples/vllm_serve/Dockerfileexamples/vllm_serve/README.mdexamples/vllm_serve/fakequant_worker.pyexamples/vllm_serve/vllm_mlflow_utils.pyexamples/vllm_serve/vllm_serve_fakequant.pymodelopt/torch/utils/mlflow.pytests/examples/vllm_serve/test_vllm_mlflow_utils.pytests/unit/torch/utils/test_mlflow.py
| ) | ||
| if args.mlflow_run_name: | ||
| os.environ[RUN_NAME_ENV] = args.mlflow_run_name | ||
| print(f"[mlflow] tracking to {uri}, experiment {os.environ[EXPERIMENT_ENV]}") |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift
Remove sensitive values from MLflow-related logs.
The launcher prints the raw tracking URI. A URI can contain basic-auth credentials or query tokens. The worker also starts raw stdout/stderr capture for MLflow upload. examples/vllm_serve/fakequant_worker.py already prints the full MODELOPT_STATE_PATH at line 68, so tracking uploads that path as an artifact.
examples/vllm_serve/vllm_mlflow_utils.py#L154-L154: do not print the raw URI. Print a redacted URI or a generic tracking-enabled message.examples/vllm_serve/fakequant_worker.py#L164-L170: redact sensitive values before uploading captured logs. Remove or sanitize path-bearing worker output, includingMODELOPT_STATE_PATH.
As per path instructions, SECURITY.md says to never log credentials, tokens, sensitive paths, or proprietary model details.
📍 Affects 2 files
examples/vllm_serve/vllm_mlflow_utils.py#L154-L154(this comment)examples/vllm_serve/fakequant_worker.py#L164-L170
🤖 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/vllm_serve/vllm_mlflow_utils.py` at line 154, In
examples/vllm_serve/vllm_mlflow_utils.py lines 154-154, update the tracking log
in the MLflow launcher to avoid printing the raw URI, using a redacted URI or
generic tracking-enabled message while preserving the experiment context. In
examples/vllm_serve/fakequant_worker.py lines 164-170, sanitize captured worker
stdout/stderr before MLflow upload, removing path-bearing sensitive output such
as MODELOPT_STATE_PATH and other credentials, tokens, sensitive paths, or
proprietary model details.
Source: Path instructions
There was a problem hiding this comment.
The first point is fixed in d30483c. resolve_mlflow_args now prints the URI through a _without_credentials() helper that strips any user:token@, matching what MlflowRunLogger already does for every URI it prints or uploads. It mattered more than it looks: that line is captured into the worker log which the run then uploads as an artifact. Regression test asserts the credentials are absent from stdout while the unmasked URI still reaches the workers, which need it to authenticate.
The second point (sanitizing captured worker stdout, including MODELOPT_STATE_PATH) is not changed by this commit.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- claimed commit ---'
git show --no-ext-diff --format=fuller --stat d30483c14 -- \
examples/vllm_serve/vllm_mlflow_utils.py \
tests/examples/vllm_serve/test_vllm_mlflow_utils.py
printf '%s\n' '--- relevant implementation outline ---'
ast-grep outline examples/vllm_serve/vllm_mlflow_utils.py \
--items all --type function --match 'without_credentials|resolve_mlflow_args'
printf '%s\n' '--- implementation and test references ---'
rg -n -C 5 \
'_without_credentials|resolve_mlflow_args|credential|password|token|tracking to' \
examples/vllm_serve/vllm_mlflow_utils.py \
tests/examples/vllm_serve/test_vllm_mlflow_utils.py
printf '%s\n' '--- remaining worker capture and state-path output ---'
rg -n -C 5 \
'MODELOPT_STATE_PATH|redirect_stdout|redirect_stderr|capture|log_text|log_artifact' \
examples/vllm_serve/fakequant_worker.py \
examples/vllm_serve/vllm_mlflow_utils.pyLength of output: 15084
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- URI redaction helper ---'
sed -n '157,178p' examples/vllm_serve/vllm_mlflow_utils.py
printf '%s\n' '--- URI redaction test cases ---'
rg -n -C 4 \
'printed_uri|without_credentials|[?&](token|access_token|api_key|sig|signature|password)=' \
tests/examples/vllm_serve/test_vllm_mlflow_utils.py \
examples/vllm_serve/vllm_mlflow_utils.pyLength of output: 3525
@cjluo-nv, the basic-auth credential fix is valid. The regression test confirms that user:token@ does not reach stdout. The worker environment keeps the original URI, as required.
The review finding remains open for two reasons:
_without_credentials()preserves URI query parameters and fragments. For example,https://mlflow.example.com?access_token=secretstill printsaccess_token=secret. The test covers only user-info credentials.- The captured worker stdout/stderr remains unsanitized. This can still upload
MODELOPT_STATE_PATHand other sensitive output.
Please redact query and fragment components before printing the URI, or print a generic tracking-enabled message. Also sanitize the captured worker log before MLflow uploads it.
You are interacting with an AI system.
There was a problem hiding this comment.
On the second point — sanitizing captured worker stdout, including MODELOPT_STATE_PATH — I am not planning to change this, for three reasons:
- It is already documented behaviour, not an oversight.
MlflowRunLogger's class docstring carries aNote:saying the captured log "is whatever the script printed, so a secret echoed to stdout still reaches the server. Prefer passing credentials via the environment." That is the contract this example relies on, and it points users at the right mitigation. MODELOPT_STATE_PATHis a filesystem path the user passes on their own command line, to a tracking server they name themselves with--mlflow. It is not a credential, and the artifact inherits that server's access control. Theprintatfakequant_worker.py:68also predates this PR — it came in with vllm fakequant reload with modelopt state for HF #805 (48da281, 2026-03-30).- A scrubbing pass over all worker stdout would strip exactly what the log exists for. The reason to upload it is to debug a calibration that went wrong, and calibration diagnostics are mostly module names and paths.
Happy to reconsider if there is a concrete value here that is secret rather than merely path-shaped.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- relevant symbols ---'
ast-grep outline modelopt/torch/utils/mlflow.py examples/vllm_serve/fakequant_worker.py examples/vllm_serve/vllm_mlflow_utils.py \
--items all --type class,function --match 'MlflowRunLogger|FakeQuantWorker|FakeQuantMlflowTracker'
printf '%s\n' '--- documentation and log capture/upload sites ---'
rg -n -C 6 \
'whatever the script printed|Prefer passing credentials|MODELOPT_STATE_PATH|redirect_stdout|redirect_stderr|log_text|log_artifact|captur' \
modelopt/torch/utils/mlflow.py \
examples/vllm_serve/fakequant_worker.py \
examples/vllm_serve/vllm_mlflow_utils.py
printf '%s\n' '--- state-path history ---'
git show --no-ext-diff --format=fuller --stat 48da2812c -- examples/vllm_serve/fakequant_worker.py
git show --no-ext-diff --format= -- examples/vllm_serve/fakequant_worker.py 48da2812c -- 2>/dev/null || trueLength of output: 26980
@cjluo-nv, I verified the capture contract and the history.
MlflowRunLogger explicitly documents that captured stdout and stderr are uploaded unchanged. The documentation tells users to pass credentials through the environment. Commit 48da2812c added MODELOPT_STATE_PATH before this PR.
MODELOPT_STATE_PATH is a user-provided filesystem path. It is not a credential. I do not have a concrete secret value that requires sanitization here. Broad worker-log sanitization would also reduce the diagnostic value of the uploaded log.
I withdraw the worker-log sanitization part of this finding. The URI query and fragment redaction concern remains separate.
✏️ Learnings added
Learnt from: cjluo-nv
Repo: NVIDIA/Model-Optimizer PR: 2120
File: examples/vllm_serve/vllm_mlflow_utils.py:0-0
Timestamp: 2026-08-12T20:24:38.874Z
Learning: In `modelopt/torch/utils/mlflow.py`, `MlflowRunLogger` intentionally uploads captured stdout and stderr unchanged as an MLflow artifact. Its public docstring warns that secrets printed by a script will reach the tracking server and directs users to pass credentials through the environment. In `examples/vllm_serve/fakequant_worker.py`, `MODELOPT_STATE_PATH` is a user-provided filesystem path for fake-quant state restoration, not a credential. Do not request broad sanitization of this diagnostic log solely because it contains filesystem paths.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
You are interacting with an AI system.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2120 +/- ##
==========================================
- Coverage 78.65% 78.18% -0.48%
==========================================
Files 522 522
Lines 60454 60461 +7
==========================================
- Hits 47552 47272 -280
- Misses 12902 13189 +287
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:
|
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
Design review: the PR solves the launcher/worker split needed to associate vLLM fake-quantization with an MLflow run. The main alternatives are reusing the existing MlflowRunLogger/HF PTQ integration (which this does), calling the already-installed optional MLflow client directly, or keeping only local logs/artifacts. The PR body explicitly explains why tracking must be opened in the worker and why the existing logger is extended, so the architectural choice is reasonably justified. The large diff is cohesive and well tested, and the new-file license headers match LICENSE_HEADER. Two small implementation issues remain: a best-effort connection failure leaks its staging directory, and two new ModelOpt imports are unnecessarily function-local.
| raise | ||
|
|
||
| def log_quant_config(self, quant_cfg: Any) -> None: | ||
| """Upload the merged ``QUANT_CFG``/``KV_QUANT_CFG`` config, when that is what ran. |
There was a problem hiding this comment.
Bot comment.
When tracking comes from $MLFLOW_TRACKING_URI, MlflowRunLogger.start() handles an open failure by setting enabled = False and returning rather than raising. In that path this tracker has already created _staging, but every later method (including finish) returns immediately because the logger is disabled, so the temp directory is never removed. Please detect that post-start() state and call _discard_staging() (and add a regression test for a best-effort start failure).
There was a problem hiding this comment.
Confirmed and fixed in d30483c. start() now calls _discard_staging() when MlflowRunLogger.start() returns with enabled false, which is the best-effort path you described — the existing except BaseException only covered the explicit --mlflow case that raises.
Two regression tests in tests/examples/vllm_serve/test_vllm_mlflow_utils.py, one per path. I checked the best-effort one fails without the fix rather than passing vacuously.
| Only the first call has an effect: vLLM drives the worker through several guarded | ||
| steps, and a run already closed as ``FINISHED`` must not be reopened or downgraded | ||
| by a failure in whatever the server does next. | ||
| """ |
There was a problem hiding this comment.
Bot comment.
Please move this ModelOpt import to the top of the file. The same applies to from modelopt.recipe import load_recipe in _start_texts; neither has a stated circular-import, optional-dependency, platform, or startup-cost reason for being local. (The local vllm import is justified by this module's explicit no-vLLM design.)
There was a problem hiding this comment.
Fixed in d30483c — both modelopt.torch.quantization and from modelopt.recipe import load_recipe are now at module scope.
You were right that there was no stated reason, and there was no real one either: the module already imports modelopt.torch.utils.mlflow at the top, which pulls both torch and modelopt.torch.quantization into sys.modules, so deferring them saved nothing. vllm_ptq_utils.py and fakequant_worker.py already import them at the top too. The local vllm import stays, as you noted.
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
Re-review: the design question is adequately addressed. The problem is carrying fake-quant tracking across the vLLM launcher/worker boundary; alternatives include the existing HF-PTQ MlflowRunLogger pattern, direct use of the already-optional mlflow-skinny client, or local-only artifacts. This PR reuses and minimally extends the in-repo logger, and the PR body explains why the worker owns the run and why environment handoff is needed. The 1,133-line diff is large but cohesive, with substantial unit coverage and a reported hardware run. The new source/test headers match LICENSE_HEADER.
Previous review status: the staging-directory leak (correctness) is fixed with both best-effort and required-path regression tests; the function-local ModelOpt imports (minor) are moved to module scope; the __all__ and broad worker-log-scrubbing suggestions were withdrawn. One security issue remains from the URI-redaction thread: query/fragment secrets are still emitted and can also enter command.txt.
|
|
||
| def quant_variant() -> str: | ||
| """What distinguishes this serve of the model, for the default experiment name. | ||
|
|
There was a problem hiding this comment.
Bot comment.
The remaining URI-redaction concern is still present: this removes only user:token@, while preserving query parameters and fragments. For example, --mlflow 'https://host/path?access_token=secret' prints the secret here; because command_text() relies on the same user-info-only redaction in modelopt/torch/utils/mlflow.py, it also stores that URI in command.txt, and MlflowRunLogger.run_url can print it again. Please redact/drop query and fragment data (preferably centrally in the MLflow utility and reuse it here), and add a regression test showing a query token is absent from launcher output and the uploaded command/run URL while the unmodified URI still reaches MLflow.
Wires examples/vllm_serve/vllm_serve_fakequant.py up to modelopt.torch.utils.mlflow via --mlflow <tracking-uri>, so a fake-quant serve records what it actually quantized and an evaluation of that endpoint can be traced back to a recipe. Without the flag, behavior is unchanged. Quantization happens in the vLLM worker, not the launcher: the launcher is the API-server frontend, and the engine and its workers are separate processes whose stdout it never sees. So the launcher only settles the tracking configuration -- validating the URI, naming the experiment, recording the command the user actually typed -- and publishes it through the environment, which is how every other setting in this example reaches the workers. Global rank 0 opens the run. The run opens before the weights load, so an unreachable server or a missing token fails in seconds rather than after a load and a full calibration, and it closes FINISHED once the model is quantized and warmed up; serving itself is not tracked. Uploaded artifacts: command.txt (the launcher's invocation, not the worker's spawn argv), version.txt, recipe/resolved_recipe.yaml, recipe/quant_cfg.yaml (only on the QUANT_CFG/KV_QUANT_CFG path, where the merged and MLA-adjusted config is recorded nowhere else), logs/<script>.log and summary/quant_summary.txt. Both the quantization and the serving settings are logged as searchable params. The checkpoint_path tag matches the one hf_ptq.py sets, so a checkpoint's PTQ run and the serves of it join up. Library additions, both used by the new example module: command_text() now takes an optional argv, so a worker can record the launcher's invocation instead of its own; MlflowRunLogger.log_text() uploads a value settled midway through a run, so a crash during calibration still keeps the config that caused it. The multi-word flags are registered under both spellings: vLLM's FlexibleArgumentParser rewrites every --foo_bar to --foo-bar before matching, so a flag registered only under the underscored spelling is unreachable from its CLI. The example Dockerfile installs the mlflow extra; the client stays an optional dependency and is imported only once tracking is enabled. Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Three findings from review, all confirmed against the code first. MlflowRunLogger.start() reports an unusable server by disabling itself rather than raising when tracking was inferred from the environment. The tracker had already created its staging directory by then, and every later method -- finish() included -- returns on the disabled logger before reaching the cleanup, so the temp directory outlived the process. Discard it when start() comes back disabled. Covered for both the best-effort and the explicit --mlflow paths; the best-effort test fails without the fix. Move `modelopt.torch.quantization` and `modelopt.recipe` to module scope. The justification given for deferring them was wrong: the existing top-level import of modelopt.torch.utils.mlflow already pulls torch and modelopt.torch.quantization into sys.modules, so the local imports saved nothing, and both neighbouring modules in this example import them at the top. The local vllm import stays -- that one is justified by this module's no-vLLM design and says so. Mask any user:token@ in the tracking URI the launcher prints. MlflowRunLogger masks the same thing in every URI it prints or uploads, and this line is captured into the worker log that the run then uploads as an artifact, so leaving it raw contradicted the library's own policy. Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
d30483c to
ec26101
Compare
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/examples/vllm_serve/test_vllm_mlflow_utils.py`:
- Around line 481-488: Align test_a_run_with_no_summary_uploads_none with its
stated scenario: configure QUANT_CONFIG["modelopt_state_path"] and
MODELOPT_STATE_PATH as needed, and use a non-zero-rank _worker() so the test
exercises the reload path; alternatively, revise the docstring to accurately
describe the current rank-0 no-summary behavior.
- Around line 467-471: Move the import of modelopt.torch.quantization to module
scope and reuse that module in the monkeypatch within the affected test,
avoiding importlib.import_module during test execution. If the local import is
required, retain it only with a brief comment documenting the optional or
unusually heavy dependency rationale.
🪄 Autofix
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: 4d3122e7-5b76-484d-bf4f-4b68c8333f2f
📒 Files selected for processing (3)
CHANGELOG.rsttests/examples/vllm_serve/test_vllm_mlflow_utils.pytests/unit/torch/utils/test_mlflow.py
🚧 Files skipped from review as they are similar to previous changes (2)
- CHANGELOG.rst
- tests/unit/torch/utils/test_mlflow.py
| monkeypatch.setattr( | ||
| importlib.import_module("modelopt.torch.quantization"), | ||
| "print_quant_summary", | ||
| write_summary, | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Move or document the local import.
Line 468 imports modelopt.torch.quantization during test execution. This can defer an import error until the test runs. Move the import to module scope. If it must remain local, add a brief comment that names the optional or unusually heavy dependency reason.
As per path instructions, “Imports inside functions or test methods without explicit justification” are IMPORTANT issues.
Proposed fix
import pytest
import yaml
+import modelopt.torch.quantization as mtq
@@
monkeypatch.setattr(
- importlib.import_module("modelopt.torch.quantization"),
+ mtq,
"print_quant_summary",
write_summary,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| monkeypatch.setattr( | |
| importlib.import_module("modelopt.torch.quantization"), | |
| "print_quant_summary", | |
| write_summary, | |
| ) | |
| import pytest | |
| import yaml | |
| import modelopt.torch.quantization as mtq | |
| monkeypatch.setattr( | |
| mtq, | |
| "print_quant_summary", | |
| write_summary, | |
| ) |
🤖 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 `@tests/examples/vllm_serve/test_vllm_mlflow_utils.py` around lines 467 - 471,
Move the import of modelopt.torch.quantization to module scope and reuse that
module in the monkeypatch within the affected test, avoiding
importlib.import_module during test execution. If the local import is required,
retain it only with a brief comment documenting the optional or unusually heavy
dependency rationale.
Sources: Coding guidelines, Path instructions
| def test_a_run_with_no_summary_uploads_none(mlflow_utils, monkeypatch, fake_mlflow): | ||
| """A reload from MODELOPT_STATE_PATH on a non-zero rank writes no summary.""" | ||
| monkeypatch.setenv("MLFLOW_TRACKING_URI", URI) | ||
| tracker = mlflow_utils.FakeQuantMlflowTracker(_worker(), QUANT_CONFIG) | ||
| tracker.start() | ||
| tracker.finish("FINISHED") | ||
|
|
||
| assert "quant_summary.txt" not in fake_mlflow.artifacts |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Exercise the scenario described by the test.
The docstring claims a MODELOPT_STATE_PATH reload on a non-zero rank. The test uses _worker() at rank 0, leaves QUANT_CONFIG["modelopt_state_path"] as None, and does not set MODELOPT_STATE_PATH. It only verifies that a normal run with no generated summary uploads no summary.
Configure the stated worker path through the real worker flow, or change the docstring to describe the behavior this test actually covers.
As per coding guidelines, “Tests must exercise the behavior they claim to validate.”
🤖 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 `@tests/examples/vllm_serve/test_vllm_mlflow_utils.py` around lines 481 - 488,
Align test_a_run_with_no_summary_uploads_none with its stated scenario:
configure QUANT_CONFIG["modelopt_state_path"] and MODELOPT_STATE_PATH as needed,
and use a non-zero-rank _worker() so the test exercises the reload path;
alternatively, revise the docstring to accurately describe the current rank-0
no-summary behavior.
Sources: Coding guidelines, Path instructions
4 participants
What does this PR do?
Type of change: new feature
Wires
examples/vllm_serve/vllm_serve_fakequant.pyup tomodelopt.torch.utils.mlflowvia--mlflow <tracking-uri>, the same way #2023 did forhf_ptq.py, so a fake-quant serve records what it actually quantized and an evaluation of that endpoint can be traced back to a recipe. Without the flag, behavior is unchanged — every hook is gated on it.Three design points worth review:
The run is recorded in the vLLM worker, not the launcher.
vllm_serve_fakequant.pyis the API-server frontend; the engine and its workers are separate processes whose stdout it never sees, so a run opened there would capture none of the calibration. The launcher instead only settles the tracking configuration — validating the URI, naming the experiment, recording the command the user actually typed — and publishes it through the environment, which is how every other setting in this example (QUANT_CFG,RECIPE_PATH, …) already reaches the workers. Global rank 0 opens the run, so a TP-8 serve produces one run.The run covers load-through-warm-up, not the server's lifetime. It opens before the weights load, so an unreachable server or a missing token fails in seconds rather than after a load and a full calibration, and it closes
FINISHEDonce the model is quantized and warmed up. A run that stayed open for the serving lifetime would never close cleanly on SIGTERM.recipe/quant_cfg.yamlis only written on the preset path. WithRECIPE_PATH,get_quant_configreturns the recipe'squantizesection unchanged andresolved_recipe.yamlalready carries it. WithQUANT_CFG/KV_QUANT_CFGit is the only record of what ran: the params carry the preset names, while the config reachingmtq.quantizeis those two deep-copied, merged, and — for an MLA model — extended at runtime with*kv_c_bmm_quantizer/*k_pe_bmm_quantizerby inspecting the loaded model.Uploaded artifacts:
command.txtversion.txtrecipe/resolved_recipe.yamlRECIPE_PATHwith its$imports expandedrecipe/quant_cfg.yamlQUANT_CFG/KV_QUANT_CFG+ MLA fixup (preset path only)logs/<script>.logsummary/quant_summary.txtPlus the quantization and serving settings as searchable params, and
user/hostname/modelopt_version/git_sha/vllm_versiontags. Thecheckpoint_pathtag matches the onehf_ptq.pysets, so a checkpoint's PTQ run and every serve of it join up.Two small library additions, both consumed by the new example module:
command_text(argv=None)— records another process's invocation, since a spawned worker's ownsys.argvis vLLM plumbing rather than anything a user typed.MlflowRunLogger.log_text()— uploads a value settled midway through a run, so a crash during calibration still keeps the config that caused it.The example
Dockerfileinstalls themlflowextra; the client remains optional and is imported only once tracking is enabled.Usage
--mlflow-experiment/--mlflow-run-nameoverride the defaults.$MLFLOW_TRACKING_URIenables tracking on its own and is best-effort; an explicit--mlflowoverrides it and fails loudly.Testing
Unit — 87 passing (
tests/examples/vllm_serve/test_vllm_mlflow_utils.py, 33 new;tests/unit/torch/utils/test_mlflow.py, +5).vllm_mlflow_utilsdeliberately imports no vLLM, so the whole launcher→worker handover is covered without a GPU, a server, or the mlflow client.End to end on aws-cmh (4× GB300,
simple_evals.gpqa_diamond, Nemotron-3.5-Lightning-30B-A3B-BF16 fake-quantized withgeneral/ptq/nvfp4_mlp_only-kv_fp8_cast): runFINISHEDin 261.5 s, opened byWorker_TP0only, all artifacts present and verified by content —command.txtheld the launcher's invocation rather than the worker's spawn argv, andresolved_recipe.yamlwas 6797 B against 1845 B of source. 104 quantizers enabled (92 NVFP4 dynamic block-16 expert weight/input with calibrated amax, 12 FP8 KV bmm). The eval then ran to completion against the served endpoint, 22/22 requests HTTP 200.Two bugs the hardware run caught, both fixed here with regression tests:
--mlflow_run_namewas rejected. vLLM'sFlexibleArgumentParser.parse_argsrewrites every--foo_barto--foo-barbefore matching, so a flag registered only under the underscored spelling is unreachable from its CLI. Both spellings are now registered. A unit test on a plainArgumentParsercould not have caught this.recipe/quant_cfg.yamluploaded a Pythonreprblob under a.yamlname: a recipe'squantizeis aQuantizeConfig,yaml.safe_dumpraisesRepresenterErroron it, and the old JSON fallback stringified the object._dump_yamlnow unwraps pydantic viamodel_dump(mode="json")and raises otherwise, with the caller downgrading that to a warning so a bad config cannot take down a serve.Known coverage gap: the preset (
QUANT_CFG/KV_QUANT_CFG) path — the only one that now writesrecipe/quant_cfg.yaml— is covered by unit test but has not been exercised on hardware; the canary usedRECIPE_PATH. Likewise the case where$MLFLOW_TRACKING_URIis present inside the deployment container and--mlflowoverrides it is unit-tested only: NeMo Evaluator Launcher forwards only declared env vars, so the eval server's URI never entered the container in the canary.Before your PR is "Ready for review"
--mlflowmeans no behavior change.CONTRIBUTING.md: ✅ — no new dependency. Uses the existing optionalnvidia-modelopt[mlflow]extra (mlflow-skinny, Apache-2.0) added in Add optional MLflow tracking to hf_ptq.py #2023; the exampleDockerfilenow installs it. No code copied from other sources./claude reviewnot yet run.Additional Information
Follows #2023, which added
MlflowRunLoggerand thehf_ptq.pyintegration.Note for anyone tracking from an OCI cluster:
mlflow-modelopt.nvidia.comis unreachable from oci-nrt and oci-hsg. TCP 443 completes and the connection is then reset on the first application byte, regardless of SNI or protocol, one RTT away — the PDX PaaS ingress appears to apply a source-IP policy, and the OCI clusters egress from Oracle-owned addresses (155.248.190.0,168.110.199.1) rather than NVIDIA's. gcp-nrt, aws-cmh and cw-dfw all reach it. This is an infrastructure matter, not a property of this change, but it determines where the feature is usable today.🤖 Generated with Claude Code
Summary by CodeRabbit