Only attempt to set the numba backend if the utils.set_numba_threading()

function is actually called.  TOAST does not use numba since it conflicts
with our existing explicit parallelism.  The set_numba_threading() function
is designed to be used before calling external packages that use numba,
in order to ensure consistency between toast and numba.

src/toast/mpi.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from .utils import Environment, Logger, set_numba_threading
+from .utils import Environment, Logger
 
 from .pshmem import MPIShared, MPILock
 
@@ -28,21 +28,6 @@
             "your python search path?"
         )
 
-# We set the numba threading here, **after** importing MPI.  The reasons are:
-#
-#  1. The import of MPI is time critical.  MPI_Init must be called quickly
-#     before the scheduling system thinks that the job is hung and kills it.
-#
-#  2. This source file is loaded by the top level module import, so placing
-#     code here will ensure that it is run once and only once when toast itself
-#     is imported and before any other dependencies are imported which might
-#     use numba internally.
-
-_have_set_numba_threading = False
-if not _have_set_numba_threading:
-    set_numba_threading()
-    _have_set_numba_threading = True
-
 
 def get_world():
     """Retrieve the default world communicator and its properties.

src/toast/utils.py

@@ -36,7 +36,14 @@
     vfast_erfinv,
 )
 
-numba_threading_layer = "NA"
+
+# This function sets the numba threading layer to (hopefully) be compatible with TOAST.
+# The TOAST threading concurrency is used to attempt to set the numba threading.  We
+# try to use the OpenMP backend for numba and then TBB.  The "workqueue" backend (which
+# is process based).  May not be compatible with all systems, so we use that as a
+# last resort.  This function should be called by any operators that use numba.
+
+numba_threading_layer = None
 
 
 def set_numba_threading():
@@ -58,6 +65,10 @@ def set_numba_threading():
 
     """
     global numba_threading_layer
+    if numba_threading_layer is not None:
+        # Already set.
+        return
+
     # Get the number of threads used by TOAST at runtime.
     env = Environment.get()
     log = Logger.get()
@@ -75,13 +86,15 @@ def set_numba_threading():
     try:
         # New style package layout
         from numba.np.ufunc import omppool
+
         have_numba_omp = True
         if rank == 0:
             log.debug("Numba has OpenMP threading support")
     except ImportError:
         try:
             # Old style
             from numba.npyufunc import omppool
+
             have_numba_omp = True
             if rank == 0:
                 log.debug("Numba has OpenMP threading support")
@@ -93,13 +106,15 @@ def set_numba_threading():
     try:
         # New style package layout
         from numba.np.ufunc import tbbpool
+
         have_numba_tbb = True
         if rank == 0:
             log.debug("Numba has TBB threading support")
     except ImportError:
         try:
             # Old style
             from numba.npyufunc import tbbpool
+
             have_numba_tbb = True
             if rank == 0:
                 log.debug("Numba has TBB threading support")