lazy imports

.github/actions/install/psydac-req/action.yml

@@ -132,7 +132,9 @@ runs:
           # accepts, so its build fails under normal (isolated) pip builds
           # which always fetch the latest setuptools. Pin setuptools<81 and
           # build without isolation so this pinned version is actually used.
-          python -m pip install wheel Cython numpy "setuptools<81"
+          # Cython 3.3 rejects pointer indexing in petsc4py 3.23.3's PC.pyx
+          # with "Invalid index type 'int'"; keep its build on Cython 3.2.
+          python -m pip install wheel "Cython>=3.2,<3.3" numpy "setuptools<81"
           python -m pip install --no-build-isolation src/binding/petsc4py
 
     # - name: Check parallel h5py installation
@@ -142,4 +144,4 @@ runs:
     #       import h5py
     #       # This particular instantiation of h5py.File will fail if parallel h5py isn't installed
     #       f = h5py.File('parallel_test.hdf5', 'w', driver='mpio', comm=MPI.COMM_WORLD)
-    #       print(f)"
+    #       print(f)"

.github/workflows/test-PR-psydac.yml

@@ -14,6 +14,7 @@ defaults:
 
 permissions:
   contents: read
+  packages: read
 
 jobs:
   test_psydac:
@@ -22,8 +23,8 @@ jobs:
     container:
       image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest
       credentials:
-        username: spossann
-        password: ${{ secrets.GHCR_PROJECT_TOKEN }}
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
 
     strategy:
       fail-fast: false
@@ -84,4 +85,3 @@ jobs:
         if: always()
         run: |
           rm -rf pytest
-

feectools/__init__.py

@@ -0,0 +1,4 @@
+# MPI switch, read once when feectools.ddm.mpi is first imported:
+#   None  -> use mpi4py if it is available (default)
+#   False -> do not import mpi4py, use MockMPI (serial runs; saves ~1 s of import time)
+use_mpi = None

feectools/api/settings.py

@@ -43,13 +43,16 @@
                        'openmp'  : False}
 # ...
 
-# Get gfortran version
-gfortran_version_output = subprocess.check_output(['gfortran', '--version']).decode('utf-8') # nosec B603, B607
-gfortran_version_string = re.search(r"(\d+\.\d+\.\d+)", gfortran_version_output).group()
-gfortran_version = Version(gfortran_version_string)
+def _get_gfortran_version():
+    """Version of the installed gfortran. Only queried where it is needed (see below), since it spawns a subprocess."""
+    gfortran_version_output = subprocess.check_output(['gfortran', '--version']).decode('utf-8') # nosec B603, B607
+    gfortran_version_string = re.search(r"(\d+\.\d+\.\d+)", gfortran_version_output).group()
+    return Version(gfortran_version_string)
 
 # Platform-dependent flags
-if platform.system() == "Darwin" and platform.machine() == 'arm64' and gfortran_version >= Version("14"):
+# (the gfortran version is only relevant on Apple silicon; evaluating it lazily avoids running
+#  a subprocess on every import and lets feectools be imported without gfortran on other platforms)
+if platform.system() == "Darwin" and platform.machine() == 'arm64' and _get_gfortran_version() >= Version("14"):
 
     # Apple silicon requires architecture-specific flags (see https://github.com/pyccel/psydac/pull/411)
     # which are only available on GCC version >= 14

feectools/ddm/mpi.py

@@ -1,8 +1,92 @@
+"""Detection of whether the process was launched by an MPI launcher.
+
+Importing ``mpi4py.MPI`` calls ``MPI_Init``, and any collective (``bcast``,
+``Barrier``, ...) issued afterwards costs something even on a single process.
+A plain ``python script.py`` run should therefore never touch MPI at all, even
+when mpi4py happens to be installed. This module answers the only question
+that decides it: was this process started by ``mpirun``/``mpiexec``/``srun``
+(or an equivalent launcher)?
+
+The answer is read from the environment the launcher sets up, so it is
+available before mpi4py is imported.
+"""
+
+import os
+import sys
+
 from dataclasses import dataclass
 from time import time
 from typing import TYPE_CHECKING
 
 
+# Per-process variables exported by the process managers behind the common
+# launchers. Each is set only for processes started *by* the launcher, so the
+# presence of any one of them means "this rank belongs to an MPI job".
+# SLURM_PROCID is deliberately absent: it is also set for the script of a
+# plain `sbatch` job, which is not an MPI launch. `srun` is covered by the
+# PMI/PMIX variables its MPI plugin exports.
+_LAUNCHER_ENV_VARS = (
+    "OMPI_COMM_WORLD_RANK",  # Open MPI (and derivatives: Spectrum, ...)
+    "PMI_RANK",  # MPICH, Intel MPI, MS-MPI, Cray, srun (pmi2)
+    "PMIX_RANK",  # PMIx, used by srun --mpi=pmix and Open MPI 5
+    "MV2_COMM_WORLD_RANK",  # MVAPICH2
+    "MPI_LOCALRANKID",  # Hydra (mpiexec.hydra)
+    "ALPS_APP_PE",  # Cray ALPS aprun
+    "PALS_RANKID",  # Cray PALS palsrun
+)
+
+# Escape hatch: force the decision either way without touching code, e.g. for
+# a launcher whose variables are not listed above.
+_OVERRIDE_ENV_VAR = "STRUPHY_MPI"
+
+_TRUE_VALUES = ("1", "true", "yes", "on")
+_FALSE_VALUES = ("0", "false", "no", "off")
+
+
+def _override() -> bool | None:
+    """Value of ``STRUPHY_MPI``, or None if unset/unrecognized."""
+    value = os.environ.get(_OVERRIDE_ENV_VAR)
+    if value is None:
+        return None
+    value = value.strip().lower()
+    if value in _TRUE_VALUES:
+        return True
+    if value in _FALSE_VALUES:
+        return False
+    return None
+
+
+def launched_under_mpi() -> bool:
+    """Whether this process was started by an MPI launcher.
+
+    Returns
+    -------
+    bool
+        True if a launcher's per-rank environment variable is present, or if
+        the application itself already initialized MPI (in which case using
+        the communicator is free). ``STRUPHY_MPI=0``/``1`` overrides
+        the detection.
+    """
+    override = _override()
+    if override is not None:
+        return override
+
+    if any(var in os.environ for var in _LAUNCHER_ENV_VARS):
+        return True
+
+    # The application may have initialized MPI itself (embedded interpreter,
+    # or an explicit `from mpi4py import MPI`). Only inspect mpi4py if it is
+    # already imported: importing it here is exactly what must be avoided.
+    mpi_module = sys.modules.get("mpi4py.MPI")
+    if mpi_module is not None:
+        try:
+            return bool(mpi_module.Is_initialized())
+        except AttributeError:
+            return False
+
+    return False
+
+
 # Might not be needed
 class MPICommWrapper:
     def __init__(self, use_mpi=True):
@@ -80,23 +164,50 @@ def COMM_WORLD(self):
     #     return 1
 
 
-try:
-    # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
+def _mpi_disabled():
+    """True if the user or the host application opted out of MPI.
+
+    Importing mpi4py initializes MPI, which takes close to a second. Serial runs that
+    never use MPI can skip it entirely, in two ways:
+
+    * export ``FEECTOOLS_MPI=0`` before starting Python, or
+    * set ``feectools.use_mpi = False`` before this module is first imported
+      (in-process, so it is not inherited by subprocesses such as ``mpirun``).
+
+    The MockMPI wrapper below is then used, exactly as if mpi4py were not installed.
+    """
     import os
-    if os.environ.get('ARRAY_BACKEND') == 'cupy':
-        raise ImportError("MPI disabled when using CuPy backend")
-
-    from mpi4py import MPI
 
-    _comm = MPI.COMM_WORLD
-    # rank = _comm.Get_rank()
-    # size = _comm.Get_size()
-    mpi_enabled = True
-except ImportError:
-    # mpi4py not installed
-    mpi_enabled = False
-except Exception:
-    # mpi4py installed but not running under mpirun
+    import feectools
+
+    if getattr(feectools, 'use_mpi', None) is False:
+        return True
+    return os.environ.get('FEECTOOLS_MPI', '').strip().lower() in ('0', 'false', 'no', 'off')
+
+
+if launched_under_mpi():
+    try:
+        # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
+        import os
+        if os.environ.get('ARRAY_BACKEND') == 'cupy':
+            raise ImportError("MPI disabled when using CuPy backend")
+
+        if _mpi_disabled():
+            raise ImportError("MPI disabled (feectools.use_mpi = False or FEECTOOLS_MPI=0)")
+
+        from mpi4py import MPI
+
+        _comm = MPI.COMM_WORLD
+        # rank = _comm.Get_rank()
+        # size = _comm.Get_size()
+        mpi_enabled = True
+    except ImportError:
+        # mpi4py not installed
+        mpi_enabled = False
+    except Exception:
+        # mpi4py installed but not running under mpirun
+        mpi_enabled = False
+else:
     mpi_enabled = False
 
 # TODO: add environment variable for mpi use

feectools/ddm/partition.py

@@ -2,10 +2,49 @@
 import numpy as np
 import numpy.ma as ma
 
-from sympy.ntheory import factorint
+__all__ = ('compute_dims', 'partition_procs_per_patch')
+
+#==============================================================================
+def factorint(n, multiple=False):
+    """
+    Prime factorization of an integer by trial division.
 
+    Drop-in replacement for the subset of ``sympy.ntheory.factorint`` used here.
+    Importing sympy takes ~1 s and is not needed for the small integers
+    (process counts, number of grid points) that are factorized in this module.
 
-__all__ = ('compute_dims', 'partition_procs_per_patch')
+    Parameters
+    ----------
+    n : int
+        Integer to factorize.
+
+    multiple : bool
+        If False (default), return a dict {prime: multiplicity}.
+        If True, return the list of primes in ascending order, repeated
+        according to their multiplicity.
+    """
+    n = int(n)
+    factors = {}
+
+    # same conventions as sympy for non-positive input
+    if n == 0:
+        factors[0] = 1
+    else:
+        if n < 0:
+            factors[-1] = 1
+            n = -n
+        p = 2
+        while p * p <= n:
+            while n % p == 0:
+                factors[p] = factors.get(p, 0) + 1
+                n //= p
+            p += 1 if p == 2 else 2
+        if n > 1:
+            factors[n] = factors.get(n, 0) + 1
+
+    if multiple:
+        return [p for p in sorted(factors) for _ in range(factors[p])]
+    return factors
 
 #==============================================================================
 def partition_procs_per_patch(npts, size):

feectools/linalg/direct_solvers.py

@@ -6,9 +6,7 @@
 from abc                 import abstractmethod
 import cunumpy as xp
 from cunumpy.xp import array_backend
-from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
 from scipy.sparse        import spmatrix, dia_matrix
-from scipy.sparse.linalg import splu
 
 from feectools.linalg.basic    import LinearSolver
 
@@ -53,6 +51,9 @@ def __init__(self, u, l, bmat, transposed=False):
         self._l    = l
         self._transposed = transposed
 
+        # imported here: scipy.linalg costs ~0.3 s to import and is only needed once a solver is built
+        from scipy.linalg.lapack import dgbtrf, dgbtrs, sgbtrf, sgbtrs, cgbtrf, cgbtrs, zgbtrf, zgbtrs
+
         # ... LU factorization
         if bmat.dtype == xp.float32:
             self._factor_function = sgbtrf
@@ -186,6 +187,8 @@ def __init__(self, spmat, transposed=False):
 
         assert isinstance(spmat, spmatrix)
 
+        from scipy.sparse.linalg import splu  # deferred, see BandedSolver
+
         self._space = xp.ndarray
         self._splu  = splu(spmat.tocsc())
         self._transposed = transposed