fix(tf): close C++ API sessions

[njzjz]

Closes #5657.

Summary

• initialize the TensorFlow Session* member to nullptr in DeepPotTF, DeepSpinTF, DeepTensorTF, and DipoleChargeModifierTF;
• centralize exclusive session teardown as Close() followed by delete and pointer reset;
• use an initialization rollback guard so public init() calls are failure-atomic and can be retried without overwriting a session left by an earlier failed attempt;
• clean sessions from both initializing-constructor exception paths and normal destructors;
• make DeepPotJAX initialization failure-atomic by unconditionally releasing all partially created TensorFlow C API handles (TF_Session, TF_Graph, TF_SessionOptions, TF_Status, TFE_Context, TFE_ContextOptions, and TF_Function references);
• add direct lifecycle regressions for failed initialization/retry and successful initialization/normal destruction.

Root cause

The TensorFlow wrappers owned raw sessions created by NewSession, but their destructors and initializing-constructor catch blocks only deleted the GraphDef. Default construction followed by a failing public init() could also retain a partially initialized session, allowing a retry to overwrite the pointer.

DeepPotJAX had the same failure mode with TensorFlow C API handles: its destructor only released resources when inited == true, so an exception before the final initialization commit leaked every handle created up to that point.

The TensorFlow rollback guard now owns cleanup until initialization has completely committed. DeepPotJAX::clear_tf_resources() is null-safe for partial initialization and is called both from the exception path and the destructor. Successful initialization commits only after all metadata has been loaded.

Test coverage

The lifecycle test intentionally white-box-links deepmd_backend_tf and deepmd_backend_jax into runUnitTests_cc, rather than going through the dlopen plugin facade. This makes the concrete backend destructors directly visible to the LSan test process.

The tests cover:

• repeated failed initialization followed by safe destruction for all four TensorFlow C++ wrappers;
• successful initialization followed by normal destruction for all four TensorFlow C++ wrappers;
• repeated failed DeepPotJAX initialization;
• successful DeepPotJAX initialization followed by normal destruction when the generated SavedModel test artifact is available.

Ordinary builds verify retry safety and destructor behavior; the CI leak-sanitizer leg provides the direct leak regression signal.

Validation

• built deepmd_backend_tf, deepmd_backend_jax, runUnitTests_cc, and deepmd_op in a TensorFlow/JAX C++ test configuration with LAMMPS plugin mode disabled;
• TestTensorFlowSessionLifecycle.*: 3 passed, 1 skipped locally because the generated JAX SavedModel artifact was unavailable;
• ruff check .;
• ruff format .;
• clang-format check;
• CMake formatting hook;
• git diff --check.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

📝 Walkthrough

Walkthrough

TensorFlow C++ and JAX API wrappers now centralize resource cleanup, handle partial initialization, and validate retry and destruction lifecycles through expanded backend tests.

Changes

TensorFlow session lifecycle

Layer / File(s)	Summary
Session cleanup contract source/api_cc/include/tf_private.h	Adds a non-throwing session cleanup helper and RAII guard.
C++ wrapper lifecycle integration source/api_cc/src/DeepPotTF.cc, source/api_cc/src/DeepSpinTF.cc, source/api_cc/src/DeepTensorTF.cc, source/api_cc/src/DataModifierTF.cc	Cleans sessions during construction failure, destruction, and failed initialization, releasing guards after success.
JAX resource lifecycle cleanup source/api_cc/include/DeepPotJAX.h, source/api_cc/src/DeepPotJAX.cc	Initializes TensorFlow C API handles to null, centralizes teardown, and cleans partial initialization failures.
Lifecycle test wiring and validation source/api_cc/tests/CMakeLists.txt, source/api_cc/tests/test_tf_session_lifecycle.cc	Links TensorFlow and JAX backends and tests failed-init retries plus normal destruction paths.

Estimated code review effort: 3 (Moderate) | ~30 minutes

Sequence Diagram(s)

sequenceDiagram
  participant APIWrapper
  participant SessionCleanupGuard
  participant TensorFlowSession
  APIWrapper->>SessionCleanupGuard: guard session during init
  APIWrapper->>TensorFlowSession: create and initialize session
  TensorFlowSession-->>APIWrapper: success or exception
  alt initialization fails
    SessionCleanupGuard->>TensorFlowSession: Close()
    SessionCleanupGuard->>TensorFlowSession: delete and null pointer
  else initialization succeeds
    APIWrapper->>SessionCleanupGuard: release()
  end

Loading

Suggested reviewers: wanghan-iapcm

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The PR also changes DeepPotJAX and test CMake/lifecycle coverage, which are outside the TensorFlow session leak fix.	Move the JAX cleanup and test linkage changes to a separate PR, or extend the issue scope to cover them.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The four TensorFlow wrappers now close/delete sessions in destructors and failure paths, covering the leak fix.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: fixing TensorFlow C++ API session cleanup.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[wanghan-iapcm]

Reviewed via the code-review skill. Correct and complete: all four tensorflow::Session* owners now close+delete the session; the SessionCleanupGuard binds session by reference (so it cleans the pointer NewSession later assigns), releases only after inited=true, and its member-declaration order matches the reordered init lists (no -Wreorder). No double-free on the throwing-constructor path (the in-init guard nulls session before the catch block runs), and Close().IgnoreError() is the right cleanup-path call. Nicely, unlike a pure-UB fix, this genuinely fails pre-fix under CI's -fsanitize=leak leg, which does run runUnitTests_cc, and all four wrappers create a session before the forced failure so the leak is LSan-visible. Long-overdue fix (PR #3062 fixed the analogous graph_def leak in 2023 but never the session).

Two non-blocking notes:

1. Same leak family, still open elsewhere: DeepPotJAX (source/api_cc/src/DeepPotJAX.cc) has the identical leak-on-throwing-init pattern with TF C-API handles (TF_Session/TF_Graph/TFE_Context/TF_Function) -- its multi-arg constructor has no try/catch and ~DeepPotJAX only frees resources if inited==true, so a throw after TF_LoadSessionFromSavedModel but before inited=true leaks everything. Out of #5657's scope (different handle type), but worth a follow-up issue.
2. The new lifecycle test covers only the failed-init/retry path; the successful-init-then-normal-destroy path (the headline #5657 scenario, which exercises the destructor's Close() on a LIVE session) is only covered implicitly by the pre-existing functional tests under the same LSan leg -- a direct case would make that explicit. Also worth calling out in the PR description that this is the first test to white-box-link deepmd_backend_tf into runUnitTests_cc rather than going through the dlopen plugin facade. LGTM.

[njzjz-bot]

Thanks for the non-blocking follow-up notes. I addressed both of them in 142db5c55:

• DeepPotJAX now releases every partially created TensorFlow C API handle when init() throws, resets the object to a retryable state, and uses the same null-safe cleanup path during normal destruction.
• The lifecycle suite now directly covers successful initialization followed by normal destruction for all four TensorFlow C++ wrappers. It also covers failed/retried JAX initialization and successful JAX destruction when the generated SavedModel artifact is available.
• The PR description now explicitly documents that the lifecycle test white-box-links deepmd_backend_tf and deepmd_backend_jax into runUnitTests_cc instead of using the dlopen facade.

I rebuilt deepmd_backend_tf, deepmd_backend_jax, runUnitTests_cc, and deepmd_op. The focused lifecycle run passed 3 tests; the JAX successful-initialization case was skipped locally because its generated SavedModel artifact was not present.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

[codecov]

Codecov Report

❌ Patch coverage is 92.11823% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.73%. Comparing base (7d5ad38) to head (142db5c).
⚠️ Report is 8 commits behind head on master.

Files with missing lines	Patch %	Lines
source/api_cc/src/DeepPotJAX.cc	89.32%	9 Missing and 2 partials ⚠️
source/api_cc/tests/test_tf_session_lifecycle.cc	91.89%	2 Missing and 1 partial ⚠️
source/api_cc/include/tf_private.h	91.66%	0 Missing and 1 partial ⚠️
source/api_cc/src/DeepSpinTF.cc	92.85%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5800      +/-   ##
==========================================
- Coverage   79.85%   77.73%   -2.12%     
==========================================
  Files        1022     1052      +30     
  Lines      117351   120707    +3356     
  Branches     4313     4363      +50     
==========================================
+ Hits        93706    93835     +129     
- Misses      22101    25319    +3218     
- Partials     1544     1553       +9     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.