global backend instance cache (#1691)

Author: slayoo

.github/workflows/tests.yml

@@ -103,7 +103,7 @@ jobs:
             test-suite: "unit_tests/!(dynamics)"
       fail-fast: false
     runs-on: ${{ matrix.platform }}
-    timeout-minutes: ${{ startsWith(matrix.platform, 'windows-') && 40 || 30 }}
+    timeout-minutes: ${{ startsWith(matrix.platform, 'ubuntu-') && 25 || 40 }}
     steps:
       - uses: actions/[email protected]
         with:
@@ -187,7 +187,7 @@ jobs:
         test-suite: [ "chemistry_extraterrestrial", "freezing", "isotopes", "condensation_a", "condensation_b", "condensation_c", "coagulation", "breakup", "multi-process_a", "multi-process_b", "multi-process_c", "multi-process_d", "multi-process_e"]
       fail-fast: false
     runs-on: ${{ matrix.platform }}
-    timeout-minutes: ${{ startsWith(matrix.platform, 'windows-') && 40 || 30 }}
+    timeout-minutes: ${{ startsWith(matrix.platform, 'ubuntu-') && 20 || 35 }}
     steps:
       - uses: actions/[email protected]
         with:

PySDM/backends/__init__.py

@@ -1,16 +1,22 @@
 """
-Backend classes: CPU=`PySDM.backends.numba.Numba`
-and GPU=`PySDM.backends.thrust_rtc.ThrustRTC`
+Number-crunching backends
 """
 
 import ctypes
+from functools import partial
 import os
 import sys
 import warnings
 
 from numba import cuda
 
-from .numba import Numba
+from . import numba as _numba
+
+# for pdoc
+CPU = None
+GPU = None
+Numba = _numba.Numba
+ThrustRTC = None
 
 
 # https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
@@ -72,12 +78,18 @@ def __call__(self, storage):
 
     ThrustRTC.Random = Random
 
-CPU = Numba
-"""
-alias for Numba
-"""
+_BACKEND_CACHE = {}
 
-GPU = ThrustRTC
-"""
-alias for ThrustRTC
-"""
+
+def _cached_backend(formulae=None, backend_class=None, **kwargs):
+    key = backend_class.__name__ + ":" + str(formulae) + ":" + str(kwargs)
+    if key not in _BACKEND_CACHE:
+        _BACKEND_CACHE[key] = backend_class(formulae=formulae, **kwargs)
+    return _BACKEND_CACHE[key]
+
+
+CPU = partial(_cached_backend, backend_class=Numba)
+""" returns a cached instance of the Numba backend (cache key including formulae parameters) """
+
+GPU = partial(_cached_backend, backend_class=ThrustRTC)
+""" returns a cached instance of the ThrustRTC backend (cache key including formulae parameters) """

PySDM/backends/numba.py

@@ -36,7 +36,9 @@ class Numba(  # pylint: disable=too-many-ancestors,duplicate-code
 
     default_croupier = "local"
 
-    def __init__(self, formulae=None, double_precision=True, override_jit_flags=None):
+    def __init__(
+        self, formulae=None, *, double_precision=True, override_jit_flags=None
+    ):
         if not double_precision:
             raise NotImplementedError()
         self.formulae = formulae or Formulae()

PySDM/backends/thrust_rtc.py

@@ -47,7 +47,7 @@ class ThrustRTC(  # pylint: disable=duplicate-code,too-many-ancestors
     default_croupier = "global"
 
     def __init__(
-        self, formulae=None, double_precision=False, debug=False, verbose=False
+        self, formulae=None, *, double_precision=False, debug=False, verbose=False
     ):
         self.formulae = formulae or Formulae()
 

PySDM/initialisation/init_fall_momenta.py

@@ -6,7 +6,6 @@
 import numpy as np
 
 from PySDM.dynamics.terminal_velocity import GunnKinzer1949
-from PySDM.formulae import Formulae
 from PySDM.particulator import Particulator
 
 
@@ -32,7 +31,7 @@ def init_fall_momenta(
 
     from PySDM.backends import CPU  # pylint: disable=import-outside-toplevel
 
-    particulator = Particulator(0, CPU(Formulae()))  # TODO #1155
+    particulator = Particulator(0, CPU())  # TODO #1155
 
     approximation = terminal_velocity_approx(particulator=particulator)
 

docs/markdown/pysdm_landing.md

@@ -663,8 +663,8 @@ mindmap
         ...
       (...)
     ((backends))
-      CPU
-      GPU
+      Numba
+      ThrustRTC
     ((dynamics))
       AqueousChemistry
       Collision

examples/PySDM_examples/Arabas_and_Shima_2017/simulation.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 import PySDM.products as PySDM_products
-from PySDM.backends import CPU
+from PySDM.backends import Numba
 from PySDM.builder import Builder
 from PySDM.dynamics import AmbientThermodynamics, Condensation
 from PySDM.environments import Parcel
@@ -10,7 +10,7 @@
 
 
 class Simulation:
-    def __init__(self, settings, backend=CPU):
+    def __init__(self, settings, backend=Numba):
         t_half = settings.z_half / settings.w_avg
 
         dt_output = (2 * t_half) / settings.n_output
@@ -23,7 +23,7 @@ def __init__(self, settings, backend=CPU):
                 formulae=settings.formulae,
                 **(
                     {"override_jit_flags": {"parallel": False}}
-                    if backend == CPU
+                    if backend is Numba
                     else {}
                 ),
             ),

examples/PySDM_examples/Arabas_et_al_2015/example_benchmark.py

@@ -6,7 +6,7 @@
 
 import PySDM.backends.impl_numba.conf
 from PySDM import Formulae
-from PySDM.backends import CPU, GPU
+from PySDM.backends import Numba, ThrustRTC
 from PySDM.products import WallTime
 
 
@@ -42,11 +42,11 @@ def main():
     n_sd = range(14, 16, 1)
 
     times = {}
-    backends = [(CPU, "sync"), (CPU, "async")]
-    if GPU.ENABLE:
-        backends.append((GPU, "async"))
+    backends = [(Numba, "sync"), (Numba, "async")]
+    if ThrustRTC.ENABLE:
+        backends.append((ThrustRTC, "async"))
     for backend, mode in backends:
-        if backend is CPU:
+        if backend is Numba:
             PySDM.backends.impl_numba.conf.NUMBA_PARALLEL = mode
             reload_cpu_backend()
         key = f"{backend} (mode={mode})"

examples/PySDM_examples/Berry_1967/figs_5_8_10.ipynb

@@ -51,7 +51,7 @@
    "source": [
     "import os\n",
     "from numpy import errstate\n",
-    "from PySDM.backends import CPU, GPU\n",
+    "from PySDM.backends import CPU, GPU, ThrustRTC\n",
     "from PySDM.dynamics.collisions.collision_kernels import Geometric, Hydrodynamic, Electric\n",
     "from PySDM_examples.Berry_1967.spectrum_plotter import SpectrumPlotter\n",
     "from PySDM_examples.Berry_1967.settings import Settings\n",
@@ -2222,7 +2222,7 @@
     "smooth = widgets.Checkbox(value=True, description='smooth plot')\n",
     "gpu = widgets.Checkbox(value=False, description='GPU')\n",
     "options = [adaptive, smooth]\n",
-    "if GPU.ENABLE:\n",
+    "if ThrustRTC.ENABLE:\n",
     "    options.append(gpu)\n",
     "kernel = widgets.Select(\n",
     "    options=['geometric sweep-out', 'electric field 3000V/cm', 'hydrodynamic capture'],\n",

examples/PySDM_examples/Bulenok_2023_MasterThesis/utils.py

@@ -95,7 +95,7 @@ def go_benchmark(
         backend_configs.append((GPU, None))
 
     for backend_class, n_threads in backend_configs:
-        backend_name = backend_class.__name__
+        backend_name = backend_class().__class__.__name__
         if n_threads:
             numba.set_num_threads(n_threads)
             backend_name += "_" + str(numba.get_num_threads())