Fix --trust_calibration_data being mutually exclusive with calibration data paths#1540
Open
adityasingh2400 wants to merge 1 commit into
Open
Conversation
…n data paths
The ONNX PTQ CLI added --trust_calibration_data to the same mutually
exclusive group as --calibration_data_path and --calibration_cache_path.
That flag only gates allow_pickle when loading the data those paths point
to, so it is meant to be combined with them, not used as an alternative.
With the current grouping, the secure opt-in the loader's error message
points users to ("enable pickle deserialization by setting the
--trust_calibration_data flag") is unreachable: argparse rejects
--calibration_data_path X --trust_calibration_data before main() runs.
Move --trust_calibration_data out of the group so it is an independent
flag while --calibration_data_path and --calibration_cache_path stay
mutually exclusive. Add regression tests for both behaviors.
Signed-off-by: Aditya Singh <adisin650@gmail.com>
Contributor
📝 WalkthroughWalkthroughThis PR fixes a CLI argument mutual-exclusion bug in the ONNX quantization module. The ChangesONNX Quantization CLI Flag Fix
🎯 2 (Simple) | ⏱️ ~12 minutes 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 `@tests/unit/onnx/quantization/test_autotune_quantization_integration.py`:
- Line 21: The test imports get_parser at module scope so re-import inside
test_quantization_cli_parser_imports_without_tensorrt doesn't see the
monkeypatched absence of TensorRT; modify the test to remove or reload
modelopt.onnx.quantization.__main__ after monkeypatching sys.modules["tensorrt"]
(e.g., pop sys.modules["modelopt.onnx.quantization.__main__"] or use
importlib.reload) and then import or call get_parser so the module-level
import-time behavior in __main__ runs under the patched tensorrt context.
🪄 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: 408776cf-db90-410d-80cb-cafd670138a0
📒 Files selected for processing (3)
CHANGELOG.rstmodelopt/onnx/quantization/__main__.pytests/unit/onnx/quantization/test_autotune_quantization_integration.py
|
|
||
| import pytest | ||
|
|
||
| from modelopt.onnx.quantization.__main__ import get_parser |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify the test currently reloads autotune but not __main__, while __main__.get_parser is imported at module scope.
rg -n -C2 'from modelopt\.onnx\.quantization\.__main__ import get_parser$|get_parser as get_parser_no_trt|importlib\.reload\(modelopt\.onnx\.quantization\.autotune\)|importlib\.reload\(.*__main__' tests/unit/onnx/quantization/test_autotune_quantization_integration.pyRepository: NVIDIA/Model-Optimizer
Length of output: 422
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="tests/unit/onnx/quantization/test_autotune_quantization_integration.py"
# Show relevant ranges around the module import and the no-TensorRT test.
sed -n '1,120p' "$FILE" | nl -ba | sed -n '1,120p'Repository: NVIDIA/Model-Optimizer
Length of output: 109
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="tests/unit/onnx/quantization/test_autotune_quantization_integration.py"
wc -l "$FILE"
echo "----"
cat -n "$FILE" | sed -n '1,140p'Repository: NVIDIA/Model-Optimizer
Length of output: 4014
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and print the __main__.py implementation of get_parser and any tensorrt-related code.
FILE="$(fd -t f "__main__.py" modelopt/onnx/quantization -a 2>/dev/null | head -n 5 | head -n 1)"
echo "Using: $FILE"
wc -l "$FILE"
echo "---- get_parser excerpt ----"
python3 - <<'PY'
import re,sys
path=sys.argv[1]
with open(path,'r',encoding='utf-8') as f:
s=f.read().splitlines()
# print around def get_parser
for i,line in enumerate(s,1):
if re.match(r"\s*def\s+get_parser\s*\(", line):
start=max(1,i-20); end=min(len(s), i+120)
for j in range(start,end+1):
print(f"{j:4d}\t{s[j-1]}")
break
else:
print("get_parser not found")
PY "$FILE"
echo "---- tensorrt references in __main__.py ----"
rg -n "tensorrt" "$FILE" || true
echo "---- get_parser callers imports ----"
rg -n "get_parser" "$FILE" || trueRepository: NVIDIA/Model-Optimizer
Length of output: 416
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="$(fd -t f "__main__.py" modelopt/onnx/quantization -a 2>/dev/null | head -n 1)"
echo "Using: $FILE"
wc -l "$FILE"
echo "---- get_parser excerpt ----"
python3 - "$FILE" <<'PY'
import re,sys
path=sys.argv[1]
with open(path,'r',encoding='utf-8') as f:
s=f.read().splitlines()
for i,line in enumerate(s,1):
if re.match(r"\s*def\s+get_parser\s*\(", line):
start=max(1,i-20); end=min(len(s), i+120)
for j in range(start,end+1):
print(f"{j:4d}\t{s[j-1]}")
break
else:
print("get_parser not found")
PY
echo "---- tensorrt references in __main__.py ----"
rg -n "tensorrt" "$FILE" || true
echo "---- get_parser callers imports in __main__.py ----"
rg -n "get_parser" "$FILE" || trueRepository: NVIDIA/Model-Optimizer
Length of output: 6122
🏁 Script executed:
#!/bin/bash
set -euo pipefail
MAIN="modelopt/onnx/quantization/__main__.py"
echo "== Top of __main__.py =="
cat -n "$MAIN" | sed -n '1,140p'
echo
echo "== References to autotune in __main__.py =="
rg -n "autotune" "$MAIN" || true
echo
echo "== TensorRT references in autotune module/package =="
# autotune may be a module or a package; search both by locating files under its path.
AUTOTUNE_DIR="$(dirname "$(fd -t f "autotune.py" modelopt/onnx/quantization -a 2>/dev/null | head -n 1)")/autotune" || true
# Fallback: just search under modelopt/onnx/quantization/autotune
if [ -d "modelopt/onnx/quantization/autotune" ]; then
rg -n "tensorrt" "modelopt/onnx/quantization/autotune" || true
else
echo "No directory modelopt/onnx/quantization/autotune found; searching any autotune.py"
fd -t f "autotune*.py" modelopt/onnx/quantization -a 2>/dev/null | xargs -r rg -n "tensorrt" || true
fi
echo
echo "== Does importing __main__ pull in autotune at module scope? =="
# Show first 80 lines of imports in __main__.py (already printed), plus any non-local imports.
rg -n "^from |^import " "$MAIN" | sed -n '1,120p'Repository: NVIDIA/Model-Optimizer
Length of output: 10289
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "tensorrt|TRT" modelopt/onnx/quantization/autotune/utils* modelopt/onnx/quantization/quantize* || true
echo "---- autotune/utils.py excerpt (imports) ----"
UTILS="$(fd -t f "utils.py" modelopt/onnx/quantization/autotune -a 2>/dev/null | head -n 1)"
echo "Using $UTILS"
wc -l "$UTILS"
sed -n '1,120p' "$UTILS" | cat -n
echo "---- quantize.py excerpt (imports) ----"
QUANT="$(fd -t f "quantize.py" modelopt/onnx/quantization -a 2>/dev/null | head -n 1)"
echo "Using $QUANT"
wc -l "$QUANT"
sed -n '1,160p' "$QUANT" | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 13735
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== tensorrt references in modelopt/onnx/trt_utils.py =="
TRT_UTIL="$(fd -t f "trt_utils.py" modelopt/onnx -a 2>/dev/null | head -n 1)"
echo "Using: $TRT_UTIL"
wc -l "$TRT_UTIL"
rg -n "tensorrt|TRT_AVAILABLE" "$TRT_UTIL" || true
echo
echo "== Top of trt_utils.py =="
sed -n '1,220p' "$TRT_UTIL" | cat -n
echo
echo "== tensorrt references in modelopt/onnx/quantization/ort_utils.py =="
ORT_UTIL="$(fd -t f "ort_utils.py" modelopt/onnx/quantization -a 2>/dev/null | head -n 1)"
echo "Using: $ORT_UTIL"
wc -l "$ORT_UTIL"
rg -n "tensorrt|TRT_AVAILABLE" "$ORT_UTIL" || true
echo
echo "== Top of ort_utils.py =="
sed -n '1,220p' "$ORT_UTIL" | cat -n
echo
echo "== What does reloading autotune do at package import time? =="
AUTOTUNE_INIT="$(fd -t f "__init__.py" modelopt/onnx/quantization/autotune -a 2>/dev/null | head -n 1)"
echo "Using: $AUTOTUNE_INIT"
wc -l "$AUTOTUNE_INIT"
sed -n '1,260p' "$AUTOTUNE_INIT" | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 23680
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="modelopt/onnx/trt_utils.py"
echo "== trt_utils.py lines 1-70 (imports/try-except) =="
sed -n '1,90p' "$FILE" | cat -n
echo
echo "== trt_utils.py around infer_types_shapes_tensorrt (260-420) =="
sed -n '240,430p' "$FILE" | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 12888
Fix no-TensorRT CLI parser test to reload __main__ under the monkeypatched tensorrt context
get_parser is imported at module scope (line 21), so inside test_quantization_cli_parser_imports_without_tensorrt (lines 35-37) the re-import can reuse the already-loaded modelopt.onnx.quantization.__main__ from before sys.modules["tensorrt"] is patched. Reload __main__ after the monkeypatch so the import-time behavior actually reflects “no TensorRT”.
Suggested fix
- from modelopt.onnx.quantization.__main__ import get_parser as get_parser_no_trt
-
- parser = get_parser_no_trt()
+ import modelopt.onnx.quantization.__main__ as quant_main
+ importlib.reload(quant_main)
+ parser = quant_main.get_parser()🤖 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/unit/onnx/quantization/test_autotune_quantization_integration.py` at
line 21, The test imports get_parser at module scope so re-import inside
test_quantization_cli_parser_imports_without_tensorrt doesn't see the
monkeypatched absence of TensorRT; modify the test to remove or reload
modelopt.onnx.quantization.__main__ after monkeypatching sys.modules["tensorrt"]
(e.g., pop sys.modules["modelopt.onnx.quantization.__main__"] or use
importlib.reload) and then import or call get_parser so the module-level
import-time behavior in __main__ runs under the patched tensorrt context.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
In
python -m modelopt.onnx.quantization, the--trust_calibration_dataflag is declared inside the sameadd_mutually_exclusive_group()as--calibration_data_pathand--calibration_cache_path. That flag only controlsallow_picklewhen loading the data those paths point to, so it is meant to be combined with them rather than used as an alternative.Because of the grouping, the secure pickle opt-in that the loader points users to is unreachable from the CLI.
main()raises:but
--calibration_data_path X --trust_calibration_datais rejected by argparse before that code runs:This moves
--trust_calibration_dataout of the mutually exclusive group so it is an independent flag, while--calibration_data_pathand--calibration_cache_pathstay mutually exclusive. The flag was added in #626.Usage
# Previously errored at argparse; now works: python -m modelopt.onnx.quantization \ --onnx_path model.onnx \ --calibration_data_path calib.npy \ --trust_calibration_dataTesting
Added unit tests in
tests/unit/onnx/quantization/test_autotune_quantization_integration.pythat assert--trust_calibration_dataparses together with either calibration path, that it defaults toFalse, and that--calibration_data_path/--calibration_cache_pathremain mutually exclusive. Verified the new combine test fails on the pre-fix parser and passes after the change.ruff,mypy, andbanditpre-commit hooks pass on the changed files.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
The change is confined to the CLI argument grouping; the existing
allow_picklesecurity logic inmain()is unchanged.Summary by CodeRabbit
Bug Fixes
--trust_calibration_dataflag to be used simultaneously with calibration data path or cache path options. Previously, a mutually-exclusive argument group incorrectly rejected this valid command combination.