fix: make replace_function / calibrate_with_adapters / enable_fake_quant exception-safe

[arham766]

What does this PR do?

Type of change: Bug fix

Bug 5 in #1902: three context managers yield without try/finally, so an exception in the body leaks patched state (replaced function + cache attribute left installed; LoRA adapters left disabled after a failed calibration; fake-quant flags left set). Siblings in the same module already restore in finally. Success-path semantics unchanged; three exception-path regression tests added.

Usage

N/A

Testing

Regression tests included (added to existing unit test files; each was verified to fail with the fix reverted). Full tests/unit/torch quantization+export+utils+opt battery passes locally with all sibling fixes applied (962 passed). Note: the test-suite PRs #1903-#1907 contain behavior-documenting NOTE tests that pin the OLD behavior fixed here — whichever lands second will be rebased to flip those assertions (happy to do so).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (error-raising on previously-silent invalid input / warning removal / exception-path cleanup only)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A (external contributor)

Additional Information

Issue: #1902

Summary by CodeRabbit

• Bug Fixes

  • Improved cleanup behavior in quantization-related utilities so temporary state is always restored, even if an operation fails.
  • Ensured adapters, function replacements, and fake-quant settings are properly reset after errors, reducing the risk of lingering side effects.

• Tests

  • Added regression coverage for error-handling cases to verify state is restored correctly when exceptions occur.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2ffa12e9-1158-43b6-b680-c722a82a8a0d

📥 Commits

Reviewing files that changed from the base of the PR and between b96a785 and 3119ddb.

📒 Files selected for processing (2)

• modelopt/torch/quantization/utils/core_utils.py
• tests/unit/torch/quantization/test_utils.py

---

📝 Walkthrough

Walkthrough

Three context managers in core_utils.py (calibrate_with_adapters, replace_function, enable_fake_quant) are refactored to move restoration logic into finally blocks, ensuring state cleanup occurs even when exceptions propagate. Corresponding regression tests are added to verify this behavior.

Changes

Exception-safe context manager restoration

Layer / File(s)	Summary
replace_function restoration modelopt/torch/quantization/utils/core_utils.py, tests/unit/torch/quantization/test_utils.py	replace_function restores the original function and removes the cached attribute in a finally block; a new test verifies both restoration and cleanup after an exception, with supporting imports added.
calibrate_with_adapters re-enable modelopt/torch/quantization/utils/core_utils.py, tests/unit/torch/quantization/test_utils.py	calibrate_with_adapters wraps the yield in try/finally to re-enable adapters; a new test confirms adapters remain disabled during a failing calibration and are re-enabled afterward.
enable_fake_quant flag restoration modelopt/torch/quantization/utils/core_utils.py, tests/unit/torch/quantization/test_utils.py	enable_fake_quant restores per-module _fake_quant flags in a finally block; a new test verifies flags are restored to original values after an exception.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

Not applicable; the changes are localized exception-handling refinements without multi-component interaction flow.

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely summarizes the main change: making the three context managers exception-safe.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	Changed files only add try/finally cleanup and regression tests; no torch.load/numpy.load/trust_remote_code/eval/exec/# nosec or dependency issues found.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}