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..7f5dbba657 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,22 @@ 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: + # 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. @@ -1230,6 +1249,28 @@ 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: + # 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 4c0b7b7cbc..22f4ee38bd 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,23 @@ 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: + # 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/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_deepeval_close.py b/source/tests/tf/test_deepeval_close.py new file mode 100644 index 0000000000..fa633d556b --- /dev/null +++ b/source/tests/tf/test_deepeval_close.py @@ -0,0 +1,102 @@ +# 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. +""" + +from unittest import ( + TestCase, + main, + 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(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(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(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__": + main() 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)