Skip to content

fix(tf): close EwaldRecp TensorFlow sessions #5708

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wanghan-iapcm wants to merge 1 commit into
base: master
Choose a base branch
from
+57 −0
Open

fix(tf): close EwaldRecp TensorFlow sessions #5708

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions deepmd/tf/infer/ewald_recp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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()
6 changes: 6 additions & 0 deletions deepmd/tf/modifier/dipole_charge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Comment on lines +108 to +113

@coderabbitai coderabbitai Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check whether DeepDipole/DeepEval/BaseModifier define close() or otherwise
# manage self.sess/self.graph lifecycle.
rg -n -C3 'class DeepDipole\b|class DeepEval\b|class BaseModifier\b' --type=py
rg -n -C5 '\bdef close\b' deepmd/tf

Repository: deepmodeling/deepmd-kit

Length of output: 161

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate the relevant files and inspect the class definitions and close-related methods.
git ls-files 'deepmd/tf/modifier/dipole_charge.py' 'deepmd/tf/**/*.py' | sed -n '1,200p'
echo '--- dipole_charge.py ---'
cat -n deepmd/tf/modifier/dipole_charge.py | sed -n '1,220p'
echo '--- search close/session lifecycle ---'
rg -n -C4 '\bdef close\b|self\.sess|self\.graph|class DeepDipole\b|class DeepEval\b|class BaseModifier\b' deepmd/tf

Repository: deepmodeling/deepmd-kit

Length of output: 49678

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the base classes that provide the TF session/graph and whether they already
# define close() or other cleanup hooks.
sed -n '1,260p' deepmd/tf/infer/deep_dipole.py
echo '---'
sed -n '1,260p' deepmd/tf/modifier/base_modifier.py
echo '---'
sed -n '1,260p' deepmd/tf/infer/deep_eval.py | sed -n '1,220p'
echo '---'
rg -n -C3 '\bdef close\b|`@cached_property`\s+def sess\b|self\.sess\s*=|self\.graph\s*=' deepmd/tf/infer/deep_dipole.py deepmd/tf/infer/deep_eval.py deepmd/tf/modifier/base_modifier.py

Repository: deepmodeling/deepmd-kit

Length of output: 11583

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Check whether DeepPot/DeepTensor already provide cleanup hooks that DipoleChargeModifier
# would inherit through BaseModifier -> DeepPot -> DeepEval/DeepTensor.
rg -n -C4 '\bdef close\b|class DeepPot\b|class DeepTensor\b|class DeepEvalBackend\b|class DeepEvalOld\b' deepmd/tf/infer/deep_pot.py deepmd/tf/infer/deep_tensor.py deepmd/tf/infer/deep_eval.py

Repository: deepmodeling/deepmd-kit

Length of output: 1079

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Find any cleanup hooks on the inherited TF evaluation classes.
rg -n -C3 '\bdef close\b|class DeepPot\b|class DeepTensor\b|class DeepEvalOld\b|class DeepEvalBackend\b' \
  deepmd/tf/infer/deep_pot.py \
  deepmd/tf/infer/deep_tensor.py \
  deepmd/tf/infer/deep_eval.py \
  deepmd/tf/infer/deep_dipole.py

Repository: deepmodeling/deepmd-kit

Length of output: 900

Close the modifier's own TF session as well. close() currently releases only self.er; DipoleChargeModifier also keeps self.sess open from DeepDipoleOld, so repeated create/discard cycles still leak the main TensorFlow session. Add a guarded self.sess.close() here too.

🤖 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/modifier/dipole_charge.py` around lines 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.

def serialize(self) -> dict:
"""Serialize the modifier.

Expand Down
10 changes: 10 additions & 0 deletions source/tests/tf/test_dipolecharge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
import os
import unittest
from unittest import (
mock,
)

import numpy as np

Expand Down Expand Up @@ -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)
Expand Down
23 changes: 23 additions & 0 deletions source/tests/tf/test_ewald.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading