Merge branch 'bartgol/eamxx/pysession-header-only' (PR #7704)

Avoids recompilation of several copies of the class

[BFB]

Author: bartgol

components/eamxx/src/share/core/CMakeLists.txt

@@ -73,12 +73,11 @@ if (EAMXX_ENABLE_PYTHON)
 
   find_package(pybind11 REQUIRED HINTS ${pybind11_ROOT})
 
-  # Add py sources and link pybind11
-  target_sources (eamxx_core PUBLIC
-    eamxx_pysession.cpp)
-
   target_link_libraries(eamxx_core PUBLIC pybind11::embed)
   target_compile_definitions (eamxx_core PUBLIC EAMXX_HAS_PYTHON)
+
+  # Used in eamxx_pysession.hpp to get current path
+  target_link_libraries(eamxx_core PUBLIC stdc++fs)
 endif()
 
 if (NOT SCREAM_LIB_ONLY)

components/eamxx/src/share/core/eamxx_pysession.cpp

@@ -2,7 +2,12 @@
 #define EAMXX_PY_SESSION_HPP
 
 #include <pybind11/pybind11.h>
+#include <pybind11/embed.h>
 
+#include <ekat_assert.hpp>
+#include <ekat_fpe.hpp>
+
+#include <filesystem>
 #include <string>
 #include <any>
 
@@ -67,6 +72,80 @@ class PySession {
   std::any py_guard;
 };
 
+// ==================== IMPLEMENTATION ===================== //
+
+inline void PySession::initialize () {
+  if (num_customers==0) {
+    // Note: if Py interpreter is already inited, we ASSUME someone else
+    // is handling the interpreter initialization/finalization
+    if (not Py_IsInitialized()) {
+      py_guard = std::make_shared<pybind11::scoped_interpreter>();
+    }
+  }
+  ++num_customers;
+}
+
+inline void PySession::finalize () {
+  EKAT_REQUIRE_MSG (num_customers>0,
+      "Error! Invalid number of customers.\n"
+      "  Did you call PySession::finalize() without calling PySession::initialize()?\n");
+
+  --num_customers;
+  if (num_customers==0) {
+    py_guard.reset();
+  }
+}
+
+inline void PySession::add_path (const std::string& path)
+{
+  EKAT_REQUIRE_MSG (is_initialized(),
+      "Error! Cannot modify python's sys.path, since PySession was not initialized yet.\n");
+
+  try {
+    // Import the sys module
+    pybind11::module sysModule = pybind11::module::import("sys");
+
+    // Get the sys.path list
+    pybind11::list sysPath = sysModule.attr("path");
+
+    // Append the new path to sys.path
+    sysPath.append(path);
+  } catch (const pybind11::error_already_set& e) {
+    std::cerr << "Error: " << e.what() << std::endl;
+    throw std::runtime_error("Could not modify sys.path. Aborting.");
+  }
+}
+
+inline void PySession::add_curr_path ()
+{
+  auto curr_path = std::filesystem::current_path();
+  add_path(curr_path.string());
+}
+
+inline pybind11::module PySession::safe_import (const std::string& module_name) const
+{
+  pybind11::module m;
+
+  // Disable FPEs while loading the module, then immediately re-enable them
+  auto fpes = ekat::get_enabled_fpes();
+  ekat::disable_all_fpes();
+  try {
+    m = pybind11::module::import(module_name.c_str());
+  } catch (const pybind11::error_already_set& e) {
+    std::cout << "[PySession::safe_import] Error! Python module import failed.\n"
+                 " - module name: " + module_name + "\n"
+                 " - pybind11 error: " + std::string(e.what()) + "\n"
+                 "Did you forget to call PySession::add_path to add the module location to sys.path?\n";
+    throw e;
+  }
+  ekat::enable_fpes(fpes);
+
+  EKAT_REQUIRE_MSG (not m.is_none(),
+      "Error! Could not import module '" + module_name + "'.\n");
+
+  return m;
+}
+
 } // namespace scream
 
 #endif // EAMXX_PY_SESSION_HPP