kokkos_teardown.hpp

Go to the documentation of this file.

// core — uniform Kokkos teardown for every peclet Python extension module (nanobind).
//
// THE PROBLEM. Kokkos::finalize() has to run from a Python `atexit` hook: the alternative —
// Kokkos's own static destructors, which run after the CUDA runtime has unloaded — aborts every
// CUDA process at exit with cudaErrorCudartUnloading. But Python runs atexit hooks BEFORE it tears
// down module globals, so any bound object (a Solver / Simulation / Tessellation at script or
// notebook scope) or any zero-copy array (a capsule holding a View) that is still referenced when
// the hook runs destroys its Kokkos View AFTER finalize, and
// Kokkos::Impl::SharedAllocationRecord::decrement calls Kokkos::abort — `Kokkos allocation "x" is
// being deallocated after Kokkos::finalize was called` + a backtrace, SIGABRT / exit 134 — on
// OpenMP exactly as on CUDA. That is the normal situation in a plain script, in `python -c`, and in
// every Jupyter / Quarto kernel; `del obj; gc.collect()` before exit was the only way out, and it
// does not exist for a notebook.
//
// THE PATTERN (the same in flow, dem, voro, pnm, coupling and core.amr):
//   * every binding-side owner of Kokkos Views derives from Releasable: its constructor registers
//     it in this module's live registry, its destructor unregisters it, release() drops its Views.
//     A wrapper that cannot drop its state piecemeal releases by destroying its own bound instance
//     (destruct_bound_instance): nanobind runs the C++ destructor now and marks the Python object
//     uninitialized, so a later method call raises a TypeError instead of touching freed memory,
//     and Python's final dealloc skips the destructor.
//   * every zero-copy capsule made by view_to_ndarray (ndarray_interop.hpp) is a Releasable too: it
//     keeps a reference to the exported View and release() drops it. The array's memory is gone
//     from then on — after shutdown it is only ever touched by Python's own dealloc, never read.
//   * a C++ class that keeps its own registry (dem's Simulation::releaseAll) plugs in via add_hook.
//   * install(m) initializes Kokkos — bounding the host thread pool by the process's real CPU
//     budget when a cgroup quota makes it narrower than the visible machine (cpu_budget.hpp)
//     — registers ONE atexit hook per module that runs release_all()
//     and THEN Kokkos::finalize(), exposes it as m.finalize() for deterministic teardown
//     (idempotent; after it, every bound object and zero-copy array of the module is dead), and
//     publishes m.execution_space.
// Each extension module statically links its own Kokkos (libkokkoscore.a), so the registry, the
// hook and the finalize are per module: modules tear down independently, in whatever order.
//
// Included by binding translation units only (they link nanobind + Kokkos); never by kernels.
#ifndef PECLET_CORE_PYTHON_KOKKOS_TEARDOWN_HPP
#define PECLET_CORE_PYTHON_KOKKOS_TEARDOWN_HPP
 
#include <nanobind/nanobind.h>
 
#include <functional>
#include <Kokkos_Core.hpp>
#include <set>
#include <type_traits>
#include <vector>
 
#include "peclet/core/common/cpu_budget.hpp"
 
namespace peclet::core::python {
 
namespace nb = nanobind;
 
class Releasable {
 public:
  Releasable() noexcept { registry().insert(this); }
  Releasable(const Releasable&) noexcept { registry().insert(this); }
  Releasable(Releasable&&) noexcept { registry().insert(this); }
  Releasable& operator=(const Releasable&) noexcept { return *this; }
  Releasable& operator=(Releasable&&) noexcept { return *this; }
  virtual ~Releasable() { registry().erase(this); }
  virtual void release() noexcept = 0;
 
  static std::set<Releasable*>& registry() {
    static auto* s = new std::set<Releasable*>;
    return *s;
  }
};
 
inline std::vector<std::function<void()>>& release_hooks() {
  static auto* v = new std::vector<std::function<void()>>;
  return *v;
}
inline void add_release_hook(std::function<void()> f) {
  release_hooks().push_back(std::move(f));
}
 
inline void release_all() noexcept {
  auto& reg = Releasable::registry();
  std::vector<Releasable*> snapshot(reg.begin(), reg.end());
  for (Releasable* r : snapshot)
    if (reg.count(r))
      r->release();
  for (auto& h : release_hooks())
    h();
}
 
inline void shutdown() noexcept {
  release_all();
  if (Kokkos::is_initialized() && !Kokkos::is_finalized())
    Kokkos::finalize();
}
 
template <class T>
inline void destruct_bound_instance(T* self) noexcept {
  nb::handle h = nb::find(self);
  if (h.is_valid())
    nb::inst_destruct(h);
}
 
inline void install(nb::module_& m) {
  if (!Kokkos::is_initialized()) {
    // Size the host backend by the CPUs this process may USE, not the ones it can see: inside a
    // cgroup CPU quota (Colab, Binder, Docker, Slurm) those differ, and the oversubscribed pool
    // spin-waits itself to a standstill — >35x, silent, on a user's first run. defaultHostThreads()
    // returns 0, and so says nothing at all, whenever OMP_NUM_THREADS / KOKKOS_NUM_THREADS is set
    // or the budget is the whole machine, which makes this inert on an ordinary workstation.
    // See peclet/core/common/cpu_budget.hpp and suite docs/SCALING_ISSUES.md issue 7.
    // Only the OpenMP backend sizes itself: its runtime reads OMP_NUM_THREADS (Kokkos reads only
    // KOKKOS_NUM_THREADS) and its default is the whole machine. Kokkos::Threads -- what a Windows
    // or macOS wheel carries -- does neither: without hwloc it defaults to ONE thread. So on that
    // backend defaultHostThreads() is the difference between a multicore wheel and a 7x cut.
    Kokkos::InitializationSettings settings;
#if defined(KOKKOS_ENABLE_OPENMP)
    constexpr bool kHostSizesItself =
        std::is_same_v<Kokkos::DefaultHostExecutionSpace, Kokkos::OpenMP>;
#else
    constexpr bool kHostSizesItself = false;
#endif
    if (const int threads = defaultHostThreads(kHostSizesItself); threads > 0)
      settings.set_num_threads(threads);
    Kokkos::initialize(settings);
  }
  m.def("finalize", &shutdown,
        "Release every live object and zero-copy array of this module, then Kokkos::finalize() "
        "(deterministic teardown; also run automatically at interpreter exit). Idempotent. After "
        "it, the module's solver objects and *_view arrays must not be used.");
  m.attr("execution_space") = nb::str(Kokkos::DefaultExecutionSpace::name());
  nb::module_::import_("atexit").attr("register")(nb::cpp_function(&shutdown));
}
 
}  // namespace peclet::core::python
 
#endif  // PECLET_CORE_PYTHON_KOKKOS_TEARDOWN_HPP