From 6bc3a4d9c25157df9846d6463aec597559306340 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 2 Jul 2026 09:32:28 +0800 Subject: [PATCH 1/3] fix(tf): close EwaldRecp TensorFlow sessions 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 #5685 --- deepmd/tf/infer/ewald_recp.py | 18 ++++++++++++++++++ deepmd/tf/modifier/dipole_charge.py | 6 ++++++ source/tests/tf/test_dipolecharge.py | 10 ++++++++++ source/tests/tf/test_ewald.py | 23 +++++++++++++++++++++++ 4 files changed, 57 insertions(+) diff --git a/deepmd/tf/infer/ewald_recp.py b/deepmd/tf/infer/ewald_recp.py index 4c0b7b7cbc..4b1a0b27be 100644 --- a/deepmd/tf/infer/ewald_recp.py +++ b/deepmd/tf/infer/ewald_recp.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np +from typing_extensions import ( + Self, +) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, @@ -93,3 +96,18 @@ def eval( ) return energy, force, virial + + def close(self) -> None: + """Close the TensorFlow session held by this object.""" + sess = getattr(self, "sess", None) + if sess is not None: + sess.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + self.close() + + def __del__(self) -> None: + self.close() diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py index a5d1fbf975..c6ec07b234 100644 --- a/deepmd/tf/modifier/dipole_charge.py +++ b/deepmd/tf/modifier/dipole_charge.py @@ -105,6 +105,12 @@ def __init__( self.force = None self.ntypes = len(self.sel_a) + def close(self) -> None: + """Close the TensorFlow session held by the Ewald reciprocal evaluator.""" + er = getattr(self, "er", None) + if er is not None: + er.close() + def serialize(self) -> dict: """Serialize the modifier. diff --git a/source/tests/tf/test_dipolecharge.py b/source/tests/tf/test_dipolecharge.py index 71c46446f6..f677d791f9 100644 --- a/source/tests/tf/test_dipolecharge.py +++ b/source/tests/tf/test_dipolecharge.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import os import unittest +from unittest import ( + mock, +) import numpy as np @@ -125,6 +128,13 @@ def tearDownClass(cls) -> None: os.remove("dipolecharge_d.pb") cls.dp = None + def test_close_forwards_to_ewald(self) -> None: + # closing the modifier must release the EwaldRecp session. Patch the + # evaluator so the shared class-level modifier is not disturbed. + with mock.patch.object(self.dp, "er") as mock_er: + self.dp.close() + mock_er.close.assert_called_once() + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 5) self.assertAlmostEqual(self.dp.get_rcut(), 4.0, places=default_places) diff --git a/source/tests/tf/test_ewald.py b/source/tests/tf/test_ewald.py index 270546fbc8..3b9cee66b2 100644 --- a/source/tests/tf/test_ewald.py +++ b/source/tests/tf/test_ewald.py @@ -224,3 +224,26 @@ def test_virial(self) -> None: np.testing.assert_almost_equal( t_esti.ravel(), virial.ravel(), places, err_msg="virial component failed" ) + + +class TestEwaldRecpClose(tf.test.TestCase): + """EwaldRecp owns a TensorFlow session that must be closeable.""" + + coord = np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]]) + charge = np.array([[1.0, -1.0]]) + box = np.array([[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]]) + + def test_close_releases_session(self) -> None: + er = EwaldRecp(1.0, 1.0) + # a fresh evaluator works + er.eval(self.coord, self.charge, self.box) + er.close() + # after close the underlying session is unusable + with self.assertRaises(RuntimeError): + er.eval(self.coord, self.charge, self.box) + + def test_context_manager_closes_session(self) -> None: + with EwaldRecp(1.0, 1.0) as er: + er.eval(self.coord, self.charge, self.box) + with self.assertRaises(RuntimeError): + er.eval(self.coord, self.charge, self.box) From f54d49a9a42d70d2b7a694beec851ac0f473b9c7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 7 Jul 2026 14:41:18 +0800 Subject: [PATCH 2/3] fix(tf): close TensorFlow sessions for all DeepEval evaluators 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. --- deepmd/infer/deep_eval.py | 23 ++++++ deepmd/tf/infer/deep_eval.py | 31 ++++++++ source/tests/tf/test_deepeval_close.py | 101 +++++++++++++++++++++++++ 3 files changed, 155 insertions(+) create mode 100644 source/tests/tf/test_deepeval_close.py diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 3c6b78f103..70df842745 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -93,6 +93,19 @@ def __init__( ) -> None: pass + def close(self) -> None: + """Release resources held by the backend. + + The base implementation does nothing. Backends that hold persistent + resources (such as a TensorFlow session) should override it. + """ + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + self.close() + def __new__(cls, model_file: str, *args: object, **kwargs: object) -> Self: if cls is DeepEvalBackend: backend = Backend.detect_backend_by_model(model_file) @@ -517,6 +530,16 @@ def __init__( if self.deep_eval.get_has_spin() and hasattr(self, "output_def_mag"): self.deep_eval.output_def = self.output_def_mag + def close(self) -> None: + """Close the underlying backend evaluator, releasing its resources.""" + self.deep_eval.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + self.close() + @property @abstractmethod def output_def(self) -> ModelOutputDef: diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 0ec2f1c74e..32061e3eb2 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -13,6 +13,9 @@ ) import numpy as np +from typing_extensions import ( + Self, +) from deepmd.common import ( make_default_mesh, @@ -306,6 +309,17 @@ def sess(self) -> tf.Session: # start a tf session associated to the graph return tf.Session(graph=self.graph, config=default_tf_session_config) + def close(self) -> None: + """Close the TensorFlow session held by this evaluator.""" + # ``sess`` is a cached_property; only close it if it was materialized, + # and drop the cache so a later access can recreate the session. + sess = self.__dict__.pop("sess", None) + if sess is not None: + sess.close() + + def __del__(self) -> None: + self.close() + def _graph_compatable(self) -> bool: """Check the model compatibility. @@ -1230,6 +1244,23 @@ def sess(self) -> tf.Session: # start a tf session associated to the graph return tf.Session(graph=self.graph, config=default_tf_session_config) + def close(self) -> None: + """Close the TensorFlow session held by this evaluator.""" + # ``sess`` is a cached_property; only close it if it was materialized, + # and drop the cache so a later access can recreate the session. + sess = self.__dict__.pop("sess", None) + if sess is not None: + sess.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + self.close() + + def __del__(self) -> None: + self.close() + def _graph_compatable(self) -> bool: """Check the model compatibility. diff --git a/source/tests/tf/test_deepeval_close.py b/source/tests/tf/test_deepeval_close.py new file mode 100644 index 0000000000..46056d941c --- /dev/null +++ b/source/tests/tf/test_deepeval_close.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Session-close behavior of the TensorFlow ``DeepEval`` evaluators. + +Every TF ``DeepEval`` backend caches a ``tf.Session`` via ``sess`` (a +``cached_property``). These tests verify the session is released by ``close()``, +by the context manager, and that ``close()`` never *materializes* a session that +was not already created. +""" + +import unittest +from unittest import ( + mock, +) + +from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper +from deepmd.infer.deep_eval import ( + DeepEvalBackend, +) +from deepmd.tf.infer.deep_eval import DeepEval as DeepEvalTF +from deepmd.tf.infer.deep_eval import ( + DeepEvalOld, +) + + +class TestTFDeepEvalClose(unittest.TestCase): + """The TF backends own a cached ``tf.Session`` that must be closeable.""" + + # both the modern backend and the legacy one used by DeepDipole/DeepPolar + backends = (DeepEvalTF, DeepEvalOld) + + def _bare(self, cls: type) -> object: + # bypass __init__/__new__ (which need a frozen model); we only exercise + # the close()/context-manager logic, not graph loading. + return object.__new__(cls) + + def test_close_closes_materialized_session(self) -> None: + for cls in self.backends: + with self.subTest(cls=cls.__name__): + obj = self._bare(cls) + fake_sess = mock.Mock() + obj.__dict__["sess"] = fake_sess # simulate the cached_property + obj.close() + fake_sess.close.assert_called_once() + # cache dropped so a later access can recreate the session + self.assertNotIn("sess", obj.__dict__) + + def test_close_without_session_does_not_create_one(self) -> None: + for cls in self.backends: + with self.subTest(cls=cls.__name__): + obj = self._bare(cls) + self.assertNotIn("sess", obj.__dict__) + obj.close() # must be a no-op, not materialize a session + self.assertNotIn("sess", obj.__dict__) + + def test_context_manager_closes_session(self) -> None: + for cls in self.backends: + with self.subTest(cls=cls.__name__): + obj = self._bare(cls) + fake_sess = mock.Mock() + obj.__dict__["sess"] = fake_sess + with obj as entered: + self.assertIs(entered, obj) + fake_sess.close.assert_called_once() + + +class TestDeepEvalWrapperClose(unittest.TestCase): + """The high-level ``DeepEval`` wrapper forwards close() to its backend.""" + + def _bare_wrapper(self) -> DeepEvalWrapper: + class _ConcreteEval(DeepEvalWrapper): + @property + def output_def(self) -> None: # the sole abstract member + return None + + return object.__new__(_ConcreteEval) + + def test_close_forwards_to_backend(self) -> None: + obj = self._bare_wrapper() + obj.deep_eval = mock.Mock() + obj.close() + obj.deep_eval.close.assert_called_once() + + def test_context_manager_forwards_to_backend(self) -> None: + obj = self._bare_wrapper() + obj.deep_eval = mock.Mock() + with obj as entered: + self.assertIs(entered, obj) + obj.deep_eval.close.assert_called_once() + + +class TestDeepEvalBackendBaseClose(unittest.TestCase): + """The backend base close() is a no-op so session-less backends comply.""" + + def test_base_close_is_noop(self) -> None: + # call the unbound base method on an arbitrary object: it must do + # nothing and not raise (backends without a session inherit this). + self.assertIsNone(DeepEvalBackend.close(object())) + + +if __name__ == "__main__": + unittest.main() From 7901e9130dad34bb7f8ccfc21452ba91e5af10e2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 7 Jul 2026 16:18:44 +0800 Subject: [PATCH 3/3] fix(tf): guard DeepEval __del__ during shutdown; dedupe test import 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. --- deepmd/tf/infer/deep_eval.py | 14 ++++++++++++-- deepmd/tf/infer/ewald_recp.py | 7 ++++++- source/tests/tf/test_deepeval_close.py | 11 ++++++----- 3 files changed, 24 insertions(+), 8 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 32061e3eb2..7f5dbba657 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -318,7 +318,12 @@ def close(self) -> None: sess.close() def __del__(self) -> None: - self.close() + # during interpreter shutdown TF/Python state may already be torn down; + # swallow errors so GC does not emit noisy "Exception ignored" messages. + try: + self.close() + except Exception: + pass def _graph_compatable(self) -> bool: """Check the model compatibility. @@ -1259,7 +1264,12 @@ def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> No self.close() def __del__(self) -> None: - self.close() + # during interpreter shutdown TF/Python state may already be torn down; + # swallow errors so GC does not emit noisy "Exception ignored" messages. + try: + self.close() + except Exception: + pass def _graph_compatable(self) -> bool: """Check the model compatibility. diff --git a/deepmd/tf/infer/ewald_recp.py b/deepmd/tf/infer/ewald_recp.py index 4b1a0b27be..22f4ee38bd 100644 --- a/deepmd/tf/infer/ewald_recp.py +++ b/deepmd/tf/infer/ewald_recp.py @@ -110,4 +110,9 @@ def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> No self.close() def __del__(self) -> None: - self.close() + # during interpreter shutdown TF/Python state may already be torn down; + # swallow errors so GC does not emit noisy "Exception ignored" messages. + try: + self.close() + except Exception: + pass diff --git a/source/tests/tf/test_deepeval_close.py b/source/tests/tf/test_deepeval_close.py index 46056d941c..fa633d556b 100644 --- a/source/tests/tf/test_deepeval_close.py +++ b/source/tests/tf/test_deepeval_close.py @@ -7,8 +7,9 @@ was not already created. """ -import unittest from unittest import ( + TestCase, + main, mock, ) @@ -22,7 +23,7 @@ ) -class TestTFDeepEvalClose(unittest.TestCase): +class TestTFDeepEvalClose(TestCase): """The TF backends own a cached ``tf.Session`` that must be closeable.""" # both the modern backend and the legacy one used by DeepDipole/DeepPolar @@ -63,7 +64,7 @@ def test_context_manager_closes_session(self) -> None: fake_sess.close.assert_called_once() -class TestDeepEvalWrapperClose(unittest.TestCase): +class TestDeepEvalWrapperClose(TestCase): """The high-level ``DeepEval`` wrapper forwards close() to its backend.""" def _bare_wrapper(self) -> DeepEvalWrapper: @@ -88,7 +89,7 @@ def test_context_manager_forwards_to_backend(self) -> None: obj.deep_eval.close.assert_called_once() -class TestDeepEvalBackendBaseClose(unittest.TestCase): +class TestDeepEvalBackendBaseClose(TestCase): """The backend base close() is a no-op so session-less backends comply.""" def test_base_close_is_noop(self) -> None: @@ -98,4 +99,4 @@ def test_base_close_is_noop(self) -> None: if __name__ == "__main__": - unittest.main() + main()