mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-19 05:49:19 +08:00
Chenjie LuoandClaude Opus 5.5
2758c046fb
### What does this PR do?
Type of change: Tests (CI speed-up)
Follow-up to #2724. Each step of the diffusers example tests started a
fresh `python <script>.py` subprocess. On the tiny test models most of
each step was startup: about 20s importing torch/diffusers/modelopt and
initializing CUDA. For example, `quantize.py`'s own timer reports 3.1s
for the SDXL quantize run inside a ~30s step.
- **In-process runner**
(`tests/_test_utils/examples/diffusers_example_runner.py`). It plugs
into the existing `set_in_process_runner` hook, like the megatron_bridge
suite does, and runs the script with `runpy.run_path(...,
run_name="__main__")`. Argument parsing and `if __name__ == "__main__":`
blocks run exactly as on the command line; `diffusion_trt.py` swaps
`torch.nn.RMSNorm` only in that block. After each step it puts back the
process-wide state the script changed:
- sibling modules imported by bare name (`config`, `utils`, `quantize`,
...), with any same-named module imported elsewhere stashed and
restored;
- the `torch.nn.RMSNorm` swap;
- logger levels, plus the handlers `quantize.py` attaches to the
captured stdout;
- warning filters, the environment, the working directory and
`sys.path`.
- **Opt-in per module** with `pytestmark =
pytest.mark.usefixtures("in_process_examples")`: `test_diffusers.py`,
`test_export_diffusers_hf_ckpt.py` and `sparsity/test_sparsity.py`. The
deprecated ONNX/TRT test, which CI doesn't run, keeps using
subprocesses.
- **`torch.compile` steps stay in a subprocess.** They already pass
`env=`, which forces one, and that is where the Inductor settings and
cudagraphs belong. The expected fallback warning is filtered.
- **`examples/diffusers/quantization/{onnx_utils,calib}/__init__.py`.**
As namespace packages these dirs lose to any same-named regular module
anywhere on `sys.path`. Importing `onnxruntime.transformers` appends its
own directory, which holds an `onnx_utils.py`. Once that has happened in
a long-lived process, `from onnx_utils.export import ...` in
`diffusion_trt.py` and in the ONNX/TRT deprecation test resolves to
onnxruntime's module. A fresh subprocess never hit this. Making them
regular packages lets the example's own directory win. Both files
contain only the license header.
### Usage
N/A (tests only).
### Testing
I ran these on an RTX PRO 6000 in `nvcr.io/nvidia/pytorch:26.07-py3`,
set up like the CI lane (editable `.[onnx,hf,dev-test]`,
`diffusers==0.39.0`). Both runs used `pytest tests/examples/diffusers
--ignore=test_diffusers_onnx_trt.py --ignore=fastgen --cov`, with
`COVERAGE_PROCESS_START` set as in CI:
| | `main` (af07754) | this PR |
|---|---:|---:|
| Result | 40 passed | 40 passed |
| Wall time | **22m35s** | **4m07s** |
| Coverage (statements missed / total) | 54899 / 68343 | 54899 / 68343 |
Per-test examples:
- SDXL int8 quantize + HF export + restore: 60s → 11s
- each Wan 2.2 quantize + restore: ~50s → ~3s
- each sparsity script run: ~25s → ~2s
The three `torch.compile` tests remain the bulk. They took 113s locally
on a warm Inductor cache and 190s on #2724's CI run, so I expect the
lane's pytest step to drop from about 21 min to about 6 min.
- `tests/unit/test_diffusers_example_runner.py` (CPU): runs a fake
example that makes all of the above process-wide changes, imports a
bare-name sibling `utils` while another module named `utils` is already
loaded, and exits nonzero on request. It checks the output is captured
and every change is undone, and that a nonzero exit raises with the
output attached.
- `pre-commit` passes.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ (the two `__init__.py` files
only turn namespace packages into regular ones)
- 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?: ✅
`tests/unit/test_diffusers_example_runner.py`
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A (tests only)
- Did you get Claude approval on this PR?: ❌
### Additional Information
- This touches `tests/_test_utils/`, so CI runs every example test lane.
- Example scripts' own warnings now appear in pytest's warnings summary
instead of being lost in subprocess output.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Tests**
* Diffusers example tests can now run examples within the test process,
with output captured for successful and failed runs.
* Test runs restore process state after execution, including environment
settings and logging.
* Added coverage for execution, failure reporting, and cleanup behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>