specdec_bench: fix run provenance, add MLflow tracking - #2119
Conversation
configuration.json recorded the raw argparse namespace, which is ambiguous about what the engine actually did: - DFLASH takes --block_size (= draft_length + 1) and leaves --draft_length at its default, so a block_size=8 DFLASH run was recorded as draft_length=3. Add num_speculative_tokens, resolved from whichever flag the algorithm uses, for consumers to read instead of re-deriving it from two flags. - draft_model_dir only records a local path, so two published drafters benchmarked against the same verifier are indistinguishable after the job is gone. Add draft_huggingface_model_id via env var, mirroring the existing HUGGINGFACE_MODEL_ID channel. Also stop redacting engine config. _SENSITIVE_SUBSTRINGS matches any key containing 'token', which swallowed num_speculative_tokens, max_num_batched_tokens, skip_tokenizer_init and ten more knobs in the serving config — including the engine's own speculation width, the natural cross-check on the field added above. The existing allowlist can't keep up because the set grows with every engine release, so redaction now also considers the value: a credential is a string, while these knobs are ints, bools and None. hf_token still redacts. Callers holding only a name (argv scanning) keep the name-only verdict. Signed-off-by: Hao Guo <haoguo@nvidia.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
|
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:
📝 WalkthroughWalkthroughThe speculative decoding benchmark now validates Hub model IDs, derives effective speculative-token metadata, applies value-aware redaction, and supports optional MLflow tracking. MLflow records benchmark parameters, tags, metrics, artifacts, and run status. ChangesSpeculative decoding and MLflow updates
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant Benchmark
participant Tracking
participant MLflow
CLI->>Benchmark: provide benchmark and MLflow options
Benchmark->>Tracking: resolve tracking configuration
Benchmark->>Tracking: enter MLflow run context
Tracking->>MLflow: start tagged run
Benchmark->>Benchmark: execute speculative decoding benchmark
Tracking->>MLflow: log parameters, metrics, and artifacts
Tracking->>MLflow: set FINISHED or FAILED status
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 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: 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 `@examples/specdec_bench/specdec_bench/utils.py`:
- Around line 219-240: Update examples/specdec_bench/specdec_bench/utils.py
lines 219-240 so _is_sensitive_key classifies explicit credential names as
sensitive regardless of scalar value, while retaining visibility for engine
configuration knobs; correct the related documentation at lines 35-42. At lines
253-253, apply this classification before _redact_value serializes each field.
At lines 352-357, validate DRAFT_HUGGINGFACE_MODEL_ID as a public model ID and
omit or reject invalid values before writing configuration.json. Add regressions
covering numeric credential values and credential-bearing environment values.
- Around line 307-316: The speculation width recorded by dump_env must match the
runtime value used by vLLM. Update the configuration assignment for
num_speculative_tokens to prefer the runtime serving_config value, or
consistently map block_size without subtracting one, while preserving the
draft_length fallback for other algorithms.
🪄 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: 50081128-77cb-4abb-8b1b-6db4fd4dfbf2
📒 Files selected for processing (3)
examples/specdec_bench/specdec_bench/utils.pytests/examples/specdec_bench/test_dump_env.pytests/examples/specdec_bench/test_redaction.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2119 +/- ##
=======================================
Coverage 78.73% 78.73%
=======================================
Files 522 522
Lines 60342 60349 +7
=======================================
+ Hits 47508 47515 +7
Misses 12834 12834
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:
|
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — specdec_bench: record effective speculation width and drafter id
Scope: full review (trigger comment carried no scoping instructions). 3 files changed / +173 −3; all three reviewed: examples/specdec_bench/specdec_bench/utils.py, tests/examples/specdec_bench/test_dump_env.py, tests/examples/specdec_bench/test_redaction.py. Also traced the width value through examples/specdec_bench/run.py, specdec_bench/models/vllm.py, and specdec_bench/models/sglang.py.
Findings: CRITICAL 1 · IMPORTANT 2 · SUGGESTION 1
The intent here is right on all three counts — the value-shape redaction rule is a genuinely better design than an enumerated allowlist that can't keep up with engine releases, and both new fields close real provenance gaps. The issues are in the derivation, not the premise.
Most impactful
1. block_size - 1 disagrees with the engine by one (CRITICAL). For --block_size 8, both engine paths receive num_speculative_tokens: 8 — vllm.py:115 and sglang.py:80, and run.py:361 documents the flag as being num_speculative_tokens. The new derived field records 7. Since this PR un-redacts serving_config specifically so the engine's own value can serve as the cross-check, the two fields that are supposed to confirm each other will systematically differ. The 7/7 output in the PR description isn't reachable from the current code. If the −1 is deliberate — treating the engines' mapping as the real bug — the correction belongs in the engine branches, not in a second independent derivation that can drift from them.
2. The derivation never reads speculative_algorithm (IMPORTANT). --speculative_algorithm NONE runs no speculation (vllm.py:120-121) but inherits --draft_length's default of 3, so every baseline row now claims a phantom width of 3. And block_size wins unconditionally even on vLLM + EAGLE3, where only draft_length reaches the engine. Keying on the algorithm — the resolution the surrounding comment already describes — fixes both and keeps one source of truth with the engine branches.
3. List-valued knobs still redact (IMPORTANT). _NON_SECRET_VALUE_TYPES is scalars only, so encoder_cudagraph_token_budgets — named in the comment and PR body as one of the twelve knobs restored — is still masked under its real list value. The new test parametrizes it as None, the one value for which it passes, so the suite is green on a case the code doesn't handle. Recursing into containers whose leaves are all non-secret scalars covers it, and leaves the existing hf_token: ["a", "b"] assertion true.
Also one non-blocking note: test_absent_by_default reads ambient env and fails for anyone who exports DRAFT_HUGGINGFACE_MODEL_ID, which the PR presents as the normal way to run.
Risk
Moderate, and confined to examples/specdec_bench/ — no modelopt/ code, no mode registration, config schema, or modelopt_state involvement, so there is no checkpoint or export exposure. The backward-compat claim in the PR body holds: all three changes are additive or widen what gets written.
The risk that matters is provenance integrity rather than runtime breakage. configuration.json is the only surviving record of a run and the PR directs consumers to prefer the new field over the raw flags, so an off-by-one there is worse than the ambiguity it replaces — it is wrong with the appearance of authority, on exactly the DFLASH rows that motivated the fix. Findings 1 and 2 should be settled together, since both concern which flag is authoritative for which algorithm; 3 is independent and self-contained. Once the derived width provably matches serving_config.speculative_config.num_speculative_tokens for a DFLASH run, this is ready.
Skipped as covered by CodeRabbit / gated pre-merge: style and formatting, and the security anti-pattern set in .coderabbit.yaml.
num_speculative_tokens: record what the engine receives, not a derived
draft-token count. models/vllm.py passes block_size straight through as
num_speculative_tokens for DFLASH, so block_size=8 must record 8, not 7.
Confirmed against the runs that motivated this PR: their acceptance
histograms top out at 9 = 8 speculative tokens + 1 bonus.
Redaction: a scalar value is no longer taken as proof of safety for
names that are credentials outright — _is_sensitive_key('hf_token', 1234)
returned False and would have persisted a numeric token. Split the
substrings into always-sensitive (hf_token / api_key / access_key /
secret / password), which redact on the name alone, and the ambiguous
rest, where a scalar still rules the field an engine knob.
Hub ids: validate HUGGINGFACE_MODEL_ID and DRAFT_HUGGINGFACE_MODEL_ID
against <org>/<name> before recording them. They are copied verbatim
into a published configuration.json, so a URL carrying an embedded
credential must not reach it through the environment. A malformed value
warns and is omitted rather than raising, since provenance metadata
should not fail a benchmark that already ran.
Signed-off-by: Hao Guo <haoguo@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Three follow-ups from review, each verified against the code paths in models/vllm.py and against the runs that motivated this PR. Resolve num_speculative_tokens by algorithm instead of by flag precedence, matching the engine branches that forward exactly one flag each. NONE now records 0 rather than inheriting --draft_length's default of 3 — baselines are the rows most often joined against specdec rows, so a phantom width there reads as authoritative. --block_size no longer overrides the width for algorithms that never read it. Clear containers whose leaves are all non-secret, so a knob holding a list survives. A sensitive-key hit replaces the whole subtree before _redact_value can recurse, so encoder_cudagraph_token_budgets — named in this PR as one of the knobs the fix restores — was still redacted under its real list value. The test pinned the one scalar (None) for which it passed; it now uses [256, 512]. A list of strings under a credential name still redacts. test_absent_by_default now clears DRAFT_HUGGINGFACE_MODEL_ID rather than reading the ambient environment, which would fail for anyone who exports it — the documented way to run a benchmark. Signed-off-by: Hao Guo <haoguo@nvidia.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.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 `@examples/specdec_bench/specdec_bench/utils.py`:
- Around line 37-40: Update _is_sensitive_key in
examples/specdec_bench/specdec_bench/utils.py at lines 263-265 to normalize
underscores and hyphens before checking _ALWAYS_SENSITIVE_SUBSTRINGS, so
separator-free camelCase variants such as apiKey, accessKey, and hfToken are
unconditionally redacted. Add numeric-value regression coverage for these
camelCase credential names. The _ALWAYS_SENSITIVE_SUBSTRINGS declaration at
lines 37-40 requires no direct change.
- Around line 228-248: Update _hub_model_id and _HUB_MODEL_ID_RE to use
huggingface_hub.utils.validate_repo_id while retaining the explicit <org>/<name>
requirement, rejecting values such as org/foo..bar, org/foo--bar, and
org/foo.git with the existing warning-and-omit behavior. In
tests/examples/specdec_bench/test_dump_env.py lines 97-120, add these values to
the malformed-ID test cases.
🪄 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: 953cfb8b-b5d2-4152-b6e2-745730e61634
📒 Files selected for processing (3)
examples/specdec_bench/specdec_bench/utils.pytests/examples/specdec_bench/test_dump_env.pytests/examples/specdec_bench/test_redaction.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/examples/specdec_bench/test_redaction.py
_ALWAYS_SENSITIVE_SUBSTRINGS missed hfToken / apiKey / hf-token, so the same numeric-credential hole stayed open for those spellings. Match against the key with separators stripped. Validate Hub ids with huggingface_hub's validate_repo_id in addition to the local pattern. Both are needed: the pattern is the security boundary (it excludes :@?# and whitespace, so a URL with an embedded credential cannot get through) while validate_repo_id alone would accept a single-component name; validate_repo_id adds canonicality, rejecting foo..bar, foo--bar and foo.git. Imported at module scope rather than guarded — utils.py already imports transformers, which hard-requires huggingface-hub, so an ImportError branch would be unreachable. Signed-off-by: Hao Guo <haoguo@nvidia.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
SSH security review
Short answer: this PR contains no SSH code. The diff is 3 files — So if you were asking about SSH in this PR, there is nothing to review. Two things it might have been about: 1. If you meant the launcher's SSH code (
|
A benchmark's results existed only as JSON under --save_dir until somebody uploaded them, so the numbers were invisible until a separate publish step ran. Add --mlflow / --mlflow_experiment / --mlflow_run_name, mirroring examples/hf_ptq, so a run is visible as soon as it finishes: acceptance length, per-category AL and throughput as metrics, the configuration as params, and the result JSONs and run log as artifacts. Experiment defaults to $USER/specdec_bench/<model>-<algorithm>. Metrics are read from --save_dir when the block exits rather than passed to MlflowRunLogger.track upfront, because the benchmark writes them during the run. A malformed or missing result file logs no metric instead of raising: tracking must not be why a finished benchmark reports failure. The drafter is tagged with its full Hub id. The org is what separates two drafters trained for the same verifier, which is the comparison this benchmark exists to make, so only a local path is shortened to its leaf. Also stop MlflowRunLogger masking engine configuration. _SECRET_NAME matches any param containing 'token', which masked the num_speculative_tokens this same PR records — and would mask max_num_batched_tokens and similar knobs for any caller. Redaction now also considers the value: a credential is a non-trivial string, so an ambiguous name holding an int, float, bool or None is kept, while an explicit hf_token / api_key / secret name masks whatever it holds, separators stripped so hfToken and hf-token match too. Signed-off-by: Hao Guo <haoguo@nvidia.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
h-guo18
changed the title
h-guo18
changed the title
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: 6
🤖 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/specdec_bench/specdec_bench/mlflow_tracking.py`:
- Around line 158-166: Update the acceptance metric mapping in the relevant
tracking function to preserve whether values come from AL or AR: emit
`Average_AL` under the AL key and `Average_AR` under a distinct AR key, and
likewise use separate `category_al/` and `category_ar/` prefixes for
`Category_AL` and `Category_AR`. Avoid fallback logic that relabels AR metrics
as AL.
- Line 157: Normalize the values returned by _first for acceptance_rate.json and
perplexity.json to mappings before any .get() access, treating scalar or
list-derived non-mappings as empty mappings. Apply this in the metric-loading
flow around acceptance and perplexity so malformed payloads cannot raise during
the finally argument evaluation and prevent logger.finish() from running.
- Around line 64-68: Update the CLI help text associated with the --mlflow
option to state that the explicit --mlflow value overrides MLflow's
MLFLOW_TRACKING_URI environment variable, matching the precedence implemented
around the argument handling near line 93. Preserve the existing best-effort
warning behavior for unusable environment-derived URIs.
- Around line 36-40: Add a module-level __all__ declaration in
mlflow_tracking.py listing the three intended public helpers: MlflowRunLogger,
default_experiment_name, and validate_tracking_uri, so imported implementation
symbols are not exposed as part of the module API.
- Line 202: Update the logger.start call in the benchmark flow to pass the
expected files mapping through its files parameter, using the same file
definitions later consumed by finish(). Keep the existing params and
_run_tags(args) arguments unchanged so MlflowRunLogger records each file’s
pre-run state and only uploads artifacts generated by this run.
- Around line 96-102: Update the ValueError handling around
validate_tracking_uri in the MLflow argument validation flow so neither
parser.error nor warnings.warn includes the exception text or supplied URI. Emit
a generic tracking-URI validation message while preserving the
required-versus-untracked control flow and args.mlflow reset behavior.
🪄 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: dd423c8a-d936-4c17-a411-21cadc5516be
📒 Files selected for processing (4)
examples/specdec_bench/run.pyexamples/specdec_bench/specdec_bench/mlflow_tracking.pymodelopt/torch/utils/mlflow.pytests/unit/torch/utils/test_mlflow.py
|
Name the expected outputs to logger.start() as well as finish(). start()
stats each path so finish() can tell an output this run wrote from one a
previous attempt left in --save_dir; without it, a run that produced none
of them uploaded the earlier attempt's results as its own.
_headline_metrics now yields {} for any payload that is not a mapping. A
metric file is written by a separate process that may have died mid-write,
and this runs inside the caller's finally — an AttributeError there would
stop the MLflow run from being closed at all, leaving it RUNNING forever.
Record legacy Average_AR / Category_AR as avg_ar / category_ar rather than
folding them into avg_al. The suffix records how the figure was computed,
so mapping both onto one name would silently mix two definitions.
Drop the URI from the validation failure message. validate_tracking_uri
quotes the value it rejected, which may carry credentials in its userinfo
or query, and this lands in terminals and CI logs; naming the source the
caller set is enough to act on.
Also correct the --mlflow help, which had the precedence backwards, and
declare __all__.
Signed-off-by: Hao Guo <haoguo@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
|
Can we support |
|
(AI-generated reply) Both are already supported.
export MLFLOW_TRACKING_URI=https://<your-mlflow-server>/
python3 run.py ... # unchanged
The two differ on failure, matching The other variables on that page are read by the |
2 participants
What does this PR do?
Type of change: Bug fix + new feature
Fix what
configuration.jsonrecords. It is the only record of how a benchmark ran once the job is gone, and three fields were wrong or missing:num_speculative_tokens(new). DFLASH is configured by--block_sizeand ignores--draft_length, so ablock_size=8run was recorded asdraft_length: 3. Resolved by algorithm, matching the engine branches inmodels/vllm.py;NONErecords0instead of inheriting a default.draft_huggingface_model_id(new, via env var).draft_model_diris a local path, so two drafters benchmarked against the same verifier were indistinguishable afterwards. Validated as an<org>/<name>Hub id.tokenmasked twelveserving_configknobs, including the engine's own speculation width. Value shape now decides: credentials are strings, these knobs are ints/bools/None. Explicit names (hf_token,api_key, …) still mask anything,hfTokenandhf-tokenincluded.Add MLflow tracking. Results otherwise sit in
--save_diruntil someone uploads them.--mlflow(plus--mlflow_experiment/--mlflow_run_name, mirroringexamples/hf_ptq) logs acceptance length, per-category AL and throughput as metrics, the configuration as params, and the result JSONs and log as artifacts.MLFLOW_TRACKING_URIopts in without changing the command.This also fixes the same redaction bug in
MlflowRunLogger, which masked thenum_speculative_tokensadded above — andmax_num_batched_tokensand similar for any caller.Usage
{ "draft_length": 3, "block_size": 8, "num_speculative_tokens": 8, "draft_huggingface_model_id": "org/My-Drafter", "serving_config": { "speculative_config": { "num_speculative_tokens": 8 } } }The raw flags stay for provenance;
num_speculative_tokensis what consumers should read, and the nested one is the engine's own value for cross-checking.Testing
test_dump_env.py(new) covers each algorithm's width path,NONE, absent attributes and Hub-id validation.test_redaction.pyandtest_mlflow.pygain engine knobs (including list-valued), credential names under scalar values, and separator-free spellings.The provenance tests were verified to fail against the pre-fix code. The redaction change was replayed over real
configuration.jsonfiles from the results bucket — knobs return,hf_tokenstays masked. MLflow tracking was exercised end-to-end against a live server.69 unit + 83 example tests pass;
pre-commitclean.Before your PR is "Ready for review"
draft_lengthkeep working but should move tonum_speculative_tokens. Already-uploaded files keep their redacted cells until re-uploaded.CONTRIBUTING.md: N/A —mlflowis already optional, imported only when tracking is on.Additional Information
Found while reconciling results on the SpecDec visualizer: DFLASH rows showed
draft_length 3against a histogram topping out at 9, and rows sharing a verifier could not be told apart. The MLflow half is the other end of the same problem — logged runs let the dashboard read the tracking server instead of requiring an upload per result.