fix(tf): close EwaldRecp TensorFlow sessions

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:

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 @@
         # 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 @@
         # 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.
 

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 @@
         )
 
         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

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.
 

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()

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)

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)