fix(tests): move shared test mixins to tests.common to fix DeviceContext leak#5344
Merged
njzjz merged 2 commits intodeepmodeling:masterfrom Mar 28, 2026
Merged
Conversation
…ext leak
pt_expt tests imported TestCaseSingleFrameWithNlist and get_tols from
tests.pt.model, which triggered tests/pt/__init__.py's
torch.set_default_device("cuda:9999999"). This pushed a DeviceContext
onto the mode stack, causing torch.zeros() (without device=) to target
a fake CUDA device — crashing on CPU-only machines.
Fix: move the shared mixins (pure numpy, no torch dependency) to
tests/common/test_mixins.py. The pt tests re-export from there for
backward compat. pt_expt tests now import directly from tests.common,
avoiding the tests.pt package entirely.
Also simplify the pt_expt conftest.py and remove the manual
_pop_device_contexts() call from test_change_bias.py.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughCentralized shared test utilities into Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt_expt/conftest.py (1)
21-35: Optional cleanup: remove unused return value from_pop_device_contexts.Now that context restore is gone, the returned
poppedlist is dead state and can be dropped for clarity.Proposed simplification
-def _pop_device_contexts() -> list: +def _pop_device_contexts() -> None: """Pop all stale DeviceContext modes from the torch function mode stack.""" - popped = [] while True: modes = _get_current_function_mode_stack() if not modes: break top = modes[-1] if isinstance(top, _device.DeviceContext): top.__exit__(None, None, None) - popped.append(top) else: break - return popped + return NoneAlso applies to: 39-40
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt_expt/conftest.py` around lines 21 - 35, The function _pop_device_contexts currently builds and returns a list named popped but that return value is unused; remove the popped list and its return, turning the function into a void operation that simply iterates the mode stack and calls top.__exit__(None, None, None) for each _device.DeviceContext encountered; update the function body to drop the popped variable and the final "return popped". Apply the same simplification to the analogous code block at the other occurrence noted (lines 39-40) so both places perform side-effect-only context popping without returning unused state.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@source/tests/pt_expt/conftest.py`:
- Around line 21-35: The function _pop_device_contexts currently builds and
returns a list named popped but that return value is unused; remove the popped
list and its return, turning the function into a void operation that simply
iterates the mode stack and calls top.__exit__(None, None, None) for each
_device.DeviceContext encountered; update the function body to drop the popped
variable and the final "return popped". Apply the same simplification to the
analogous code block at the other occurrence noted (lines 39-40) so both places
perform side-effect-only context popping without returning unused state.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 53d4815d-fcdd-4e31-be0b-e8f7196010ed
📒 Files selected for processing (22)
source/tests/common/test_mixins.pysource/tests/pt/model/test_env_mat.pysource/tests/pt/model/test_mlp.pysource/tests/pt_expt/conftest.pysource/tests/pt_expt/descriptor/test_dpa1.pysource/tests/pt_expt/descriptor/test_dpa2.pysource/tests/pt_expt/descriptor/test_dpa3.pysource/tests/pt_expt/descriptor/test_hybrid.pysource/tests/pt_expt/descriptor/test_se_atten_v2.pysource/tests/pt_expt/descriptor/test_se_e2_a.pysource/tests/pt_expt/descriptor/test_se_r.pysource/tests/pt_expt/descriptor/test_se_t.pysource/tests/pt_expt/descriptor/test_se_t_tebd.pysource/tests/pt_expt/fitting/test_dipole_fitting.pysource/tests/pt_expt/fitting/test_dos_fitting.pysource/tests/pt_expt/fitting/test_ener_fitting.pysource/tests/pt_expt/fitting/test_invar_fitting.pysource/tests/pt_expt/fitting/test_polar_fitting.pysource/tests/pt_expt/fitting/test_property_fitting.pysource/tests/pt_expt/loss/test_ener.pysource/tests/pt_expt/test_change_bias.pysource/tests/pt_expt/utils/test_exclusion_mask.py
💤 Files with no reviewable changes (1)
- source/tests/pt_expt/test_change_bias.py
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3ef15274bc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5344 +/- ##
==========================================
- Coverage 82.28% 82.28% -0.01%
==========================================
Files 797 797
Lines 82100 82101 +1
Branches 4003 4004 +1
==========================================
- Hits 67557 67555 -2
- Misses 13336 13337 +1
- Partials 1207 1209 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Reinstate session-scoped fixture to cover mixed test runs where tests/pt/__init__.py may push a DeviceContext before pt_expt tests.
njzjz
approved these changes
Mar 27, 2026
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
2 participants
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.
Summary
TestCaseSingleFrameWithNlistandget_tolsfromtests.pt.modeltotests.common.test_mixinstests.common.test_mixinsdirectlytests/pt_expt/conftest.pyand remove manual_pop_device_contexts()workaroundsRoot cause
tests/pt/__init__.pycallstorch.set_default_device("cuda:9999999")to enforce explicit device usage in pt tests. This pushes aDeviceContextonto the torch mode stack. pt_expt descriptor/fitting tests importedTestCaseSingleFrameWithNlistfromtests.pt.model.test_env_mat, which triggeredtests/pt/__init__.py— leaking theDeviceContextinto pt_expt tests.The leaked
DeviceContextcausedtorch.zeros()calls (without explicitdevice=) inside AOTInductor's lowering pass and PyTorch's Adam optimizer to targetcuda:9999999, crashing on CPU-only CI machines withAssertionError: Torch not compiled with CUDA enabled.Fix
The shared mixins (
TestCaseSingleFrameWithNlist,get_tols) are pure numpy with no torch dependency. Moving them totests/common/test_mixins.pylets pt_expt tests import them without touching thetests.ptpackage. The pt tests re-export from the common location for backward compatibility.Test plan
pytest source/tests/pt_expt/descriptor/test_se_e2_a.py— passes, no DeviceContext leakpytest source/tests/pt/model/test_env_mat.py— passes (backward compat via re-export)pytest source/tests/pt/model/test_mlp.py— passes (backward compat via re-export)import source.tests.pt_expt.descriptor.test_se_e2_ano longer creates DeviceContext🤖 Generated with Claude Code
Summary by CodeRabbit