fix(tf): close C++ API sessions (#5800)

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

---------

Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>

Author: 4 people

source/api_cc/include/DeepPotJAX.h

@@ -242,6 +242,13 @@ class DeepPotJAX : public DeepPotBackend {
                            const bool atomic);
 
  private:
+  /** Release every TensorFlow C API handle owned by this backend.
+   *
+   * This helper is intentionally safe for partially initialized objects so an
+   * exception from init() cannot leak handles or make a later retry unsafe.
+   */
+  void clear_tf_resources() noexcept;
+
   bool inited;
   // device
   std::string device;
@@ -284,12 +291,12 @@ class DeepPotJAX : public DeepPotBackend {
   /**  TF C API objects.
    * @{
    */
-  TF_Graph* graph;
-  TF_Status* status;
-  TF_Session* session;
-  TF_SessionOptions* sessionopts;
-  TFE_ContextOptions* ctx_opts;
-  TFE_Context* ctx;
+  TF_Graph* graph = nullptr;
+  TF_Status* status = nullptr;
+  TF_Session* session = nullptr;
+  TF_SessionOptions* sessionopts = nullptr;
+  TFE_ContextOptions* ctx_opts = nullptr;
+  TFE_Context* ctx = nullptr;
   std::vector<TF_Function*> func_vector;
   /**
    * @}

source/api_cc/include/tf_private.h

@@ -18,4 +18,49 @@ typedef tensorflow::tstring STRINGTYPE;
 #else
 typedef std::string STRINGTYPE;
 #endif
+
+/**
+ * @brief Close and release an exclusively owned TensorFlow session.
+ *
+ * TensorFlow permits concurrent Run calls, but requires them all to finish
+ * before the sole Close call.  Owners must therefore stop concurrent use
+ * before destruction; this helper only centralizes the non-throwing cleanup
+ * needed by destructors and partially constructed API objects.
+ *
+ * @param[in,out] session Exclusively owned session, reset to nullptr.
+ */
+inline void close_and_delete_session(tensorflow::Session*& session) noexcept {
+  if (session == nullptr) {
+    return;
+  }
+  session->Close().IgnoreError();
+  delete session;
+  session = nullptr;
+}
+
+/**
+ * @brief Roll back a session assigned during a potentially throwing init call.
+ *
+ * Call release() only after initialization has completed successfully.
+ */
+class SessionCleanupGuard {
+ public:
+  explicit SessionCleanupGuard(tensorflow::Session*& session) noexcept
+      : session_(session), active_(true) {}
+
+  ~SessionCleanupGuard() {
+    if (active_) {
+      close_and_delete_session(session_);
+    }
+  }
+
+  SessionCleanupGuard(const SessionCleanupGuard&) = delete;
+  SessionCleanupGuard& operator=(const SessionCleanupGuard&) = delete;
+
+  void release() noexcept { active_ = false; }
+
+ private:
+  tensorflow::Session*& session_;
+  bool active_;
+};
 }  // namespace deepmd

source/api_cc/src/DataModifierTF.cc

@@ -8,22 +8,29 @@ using namespace deepmd;
 using namespace tensorflow;
 
 DipoleChargeModifierTF::DipoleChargeModifierTF()
-    : inited(false), graph_def(new GraphDef()) {}
+    : session(nullptr), graph_def(new GraphDef()), inited(false) {}
 
 DipoleChargeModifierTF::DipoleChargeModifierTF(const std::string& model,
                                                const int& gpu_rank,
                                                const std::string& name_scope_)
-    : inited(false), name_scope(name_scope_), graph_def(new GraphDef()) {
+    : session(nullptr),
+      name_scope(name_scope_),
+      graph_def(new GraphDef()),
+      inited(false) {
   try {
     init(model, gpu_rank, name_scope_);
   } catch (...) {
     // Clean up and rethrow, as the destructor will not be called
+    deepmd::close_and_delete_session(session);
     delete graph_def;
     throw;
   }
 }
 
-DipoleChargeModifierTF::~DipoleChargeModifierTF() { delete graph_def; };
+DipoleChargeModifierTF::~DipoleChargeModifierTF() {
+  deepmd::close_and_delete_session(session);
+  delete graph_def;
+};
 
 void DipoleChargeModifierTF::init(const std::string& model,
                                   const int& gpu_rank,
@@ -34,6 +41,7 @@ void DipoleChargeModifierTF::init(const std::string& model,
               << std::endl;
     return;
   }
+  deepmd::SessionCleanupGuard session_guard(session);
   name_scope = name_scope_;
   SessionOptions options;
   get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
@@ -72,6 +80,7 @@ void DipoleChargeModifierTF::init(const std::string& model,
   get_vector<int>(sel_type, "model_attr/sel_type");
   sort(sel_type.begin(), sel_type.end());
   inited = true;
+  session_guard.release();
 }
 
 template <class VT>