fix(tf): close EwaldRecp TensorFlow sessions#5708
Conversation
EwaldRecp creates a dedicated tf.Session in its constructor but exposed no way to close it, and DipoleChargeModifier holds an EwaldRecp without releasing it. Repeatedly constructing and discarding these objects leaked TensorFlow sessions and graph resources in long-running processes. Add close(), context-manager support, and a defensive __del__ to EwaldRecp, and have DipoleChargeModifier.close() forward to the EwaldRecp. Add tests that the session is released after close()/context-manager exit and that the modifier forwards close to its evaluator. Fix deepmodeling#5685
📝 WalkthroughWalkthroughAdds session cleanup and context-manager support across TensorFlow evaluator classes, forwards cleanup through ChangesTensorFlow session lifecycle
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant DeepEval
participant DeepEvalBackend
participant TFSession
DeepEval->>DeepEval: close()
DeepEval->>DeepEvalBackend: close()
DeepEvalBackend->>TFSession: release resources if any
DeepEval->>DeepEval: __exit__()
DeepEval->>DeepEvalBackend: close()
Related Issues: Suggested labels: enhancement, tests Suggested reviewers: njzjz 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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}
There was a problem hiding this comment.
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 `@deepmd/tf/modifier/dipole_charge.py`:
- Around line 108-113: The DipoleChargeModifier.close method only releases the
Ewald evaluator stored in self.er and leaves the TensorFlow session inherited
from DeepDipoleOld open. Update close() in DipoleChargeModifier to also guard
and close self.sess after closing er, so repeated create/discard cycles fully
release both sessions.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 91599b1b-6ed5-4ae4-bbea-09f84b087cc3
📒 Files selected for processing (4)
deepmd/tf/infer/ewald_recp.pydeepmd/tf/modifier/dipole_charge.pysource/tests/tf/test_dipolecharge.pysource/tests/tf/test_ewald.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5708 +/- ##
==========================================
- Coverage 81.97% 79.48% -2.49%
==========================================
Files 959 1014 +55
Lines 105748 115412 +9664
Branches 4102 4275 +173
==========================================
+ Hits 86684 91739 +5055
- Misses 17573 22130 +4557
- Partials 1491 1543 +52 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz
left a comment
njzjz
left a comment
There was a problem hiding this comment.
It seems to me that all TensorFlow DeepEval has a sess attribute. So I am not sure if it is useful to fix only one class.
Generalizes the EwaldRecp session-close fix to every TensorFlow DeepEval evaluator, which each cache a tf.Session via the sess cached_property and previously leaked it. Adds a no-op close() plus context-manager support on the DeepEvalBackend base, overrides close() on the TF DeepEval and legacy DeepEvalOld backends to close and drop the cached session (guarded so it is never materialized just to be closed), and forwards close()/__enter__/__exit__ from the high-level DeepEval wrapper. Addresses the review note that the leak affects all TF DeepEval classes, not only EwaldRecp.
|
Good point — addressed in f54d49a. You are right that every TF |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
deepmd/tf/infer/deep_eval.py (1)
1247-1263: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueConsider guarding
__del__against exceptions during interpreter shutdown.
__del__unconditionally callsclose(), which callssess.close(). If module-level TF/Python state is already torn down at interpreter shutdown, this can raise noisy "Exception ignored in..." messages rather than propagating cleanly. Wrapping the shutdown path in atry/except Exception: passis a common mitigation for__del__cleanup methods, distinct from the explicitclose()call which should still raise normally.🛡️ Optional guard for `__del__`
def __del__(self) -> None: - self.close() + try: + self.close() + except Exception: + pass🤖 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 `@deepmd/tf/infer/deep_eval.py` around lines 1247 - 1263, Guard the cleanup path in the DeepEval destructor: keep `close()` unchanged so explicit calls still surface errors, but update `__del__` in `deep_eval.py` to wrap `self.close()` in a broad `try/except Exception` and ignore failures that can occur during interpreter shutdown when TensorFlow/Python state may already be torn down. Use the existing `close`, `__enter__`, `__exit__`, and `__del__` methods to locate the change.
🤖 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 `@deepmd/infer/deep_eval.py`:
- Around line 96-108: DeepEvalBackend.close is intentionally a no-op default,
but Ruff B027 will flag the empty method and fail checks. Update the close
method in DeepEvalBackend to explicitly suppress the no-op warning (for example
with a targeted lint ignore or equivalent) while keeping __enter__ and __exit__
behavior unchanged. Make sure the suppression is scoped to this default
implementation so subclasses can still override close normally.
---
Nitpick comments:
In `@deepmd/tf/infer/deep_eval.py`:
- Around line 1247-1263: Guard the cleanup path in the DeepEval destructor: keep
`close()` unchanged so explicit calls still surface errors, but update `__del__`
in `deep_eval.py` to wrap `self.close()` in a broad `try/except Exception` and
ignore failures that can occur during interpreter shutdown when
TensorFlow/Python state may already be torn down. Use the existing `close`,
`__enter__`, `__exit__`, and `__del__` methods to locate the change.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 4cba7939-241f-43c0-9a9b-a6db31e66a5d
📒 Files selected for processing (3)
deepmd/infer/deep_eval.pydeepmd/tf/infer/deep_eval.pysource/tests/tf/test_deepeval_close.py
Wrap the session cleanup in the EwaldRecp/DeepEval/DeepEvalOld __del__ methods in try/except so garbage collection during interpreter shutdown, when TF/Python state may already be torn down, does not emit noisy "Exception ignored in" messages; the explicit close() path still raises normally. Also consolidate the duplicate unittest import in the new test into a single from-import to satisfy the CodeQL py/import-and-import-from warning.
|
Pushed 7901e91 with the review follow-ups:
The earlier |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/tf/infer/deep_eval.py (1)
1252-1273: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
close()/__del__logic vs.DeepEval(lines 312-327).The body of
close()and__del__inDeepEvalOldis byte-for-byte identical to the version inDeepEvalabove. SinceDeepEvalOlddoesn't share a base class withDeepEvalhere, consider extracting a small mixin (e.g.,_CachedSessionCloseMixin) withclose()/__del__to avoid maintaining two copies of the same cleanup logic.♻️ Suggested mixin extraction
class _CachedSessionCloseMixin: """Shared close()/__del__ for classes with a cached_property `sess`.""" def close(self) -> None: sess = self.__dict__.pop("sess", None) if sess is not None: sess.close() def __del__(self) -> None: try: self.close() except Exception: pass🤖 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 `@deepmd/tf/infer/deep_eval.py` around lines 1252 - 1273, The cleanup logic in DeepEvalOld duplicates the same close()/__del__ behavior already used by DeepEval, so extract that shared cached-session teardown into a small mixin such as _CachedSessionCloseMixin and have both DeepEval and DeepEvalOld inherit from it. Move the session cache pop-and-close logic from close() plus the guarded __del__ shutdown handling into the shared mixin, then remove the duplicated implementations from DeepEvalOld so the two classes use one maintained copy of the cleanup code.
🤖 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.
Nitpick comments:
In `@deepmd/tf/infer/deep_eval.py`:
- Around line 1252-1273: The cleanup logic in DeepEvalOld duplicates the same
close()/__del__ behavior already used by DeepEval, so extract that shared
cached-session teardown into a small mixin such as _CachedSessionCloseMixin and
have both DeepEval and DeepEvalOld inherit from it. Move the session cache
pop-and-close logic from close() plus the guarded __del__ shutdown handling into
the shared mixin, then remove the duplicated implementations from DeepEvalOld so
the two classes use one maintained copy of the cleanup code.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 82778837-297a-4353-954f-c36c2f6b1fd5
📒 Files selected for processing (3)
deepmd/tf/infer/deep_eval.pydeepmd/tf/infer/ewald_recp.pysource/tests/tf/test_deepeval_close.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/tests/tf/test_deepeval_close.py
aff66bd
3 participants
Problem
EwaldRecpcreates a dedicatedtf.Sessionin its constructor but exposes no way to close it (noclose(), context-manager path, or destructor).DipoleChargeModifiercreates and keeps anEwaldRecpinstance without releasing it either. Creating and discarding these objects therefore leaks TensorFlow sessions and graph resources; long-running processes that repeatedly construct modifiers accumulate them until process exit.Fix
Add
close()(idempotent, guarded), context-manager support (__enter__/__exit__), and a defensive__del__toEwaldRecp, and haveDipoleChargeModifier.close()forward to the heldEwaldRecp.Test
test_ewald.pypreviously exercised onlyeval, never lifecycle. New tests assert thatEwaldRecp.close()and context-manager exit release the session (a subsequentevalraises on the closed session), and thatDipoleChargeModifier.close()forwards to itsEwaldRecp.close(). All three fail on the current code withAttributeError(no such methods) and pass after the fix.Fix #5685
Summary by CodeRabbit
close()forwarding for the dipole-charge modifier to release its associated evaluator resources.close()behavior, context-manager exit, and error handling after sessions are released for Ewald reciprocal and DeepEval (including wrapper forwarding and no-op behavior when nothing is cached).