diff --git a/.github/actions/install/psydac-req/action.yml b/.github/actions/install/psydac-req/action.yml index 2905f49fa..b5314240a 100644 --- a/.github/actions/install/psydac-req/action.yml +++ b/.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)" \ No newline at end of file + # print(f)" diff --git a/.github/workflows/test-PR-psydac.yml b/.github/workflows/test-PR-psydac.yml index 62a14bddf..92caad818 100644 --- a/.github/workflows/test-PR-psydac.yml +++ b/.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 - \ No newline at end of file diff --git a/feectools/__init__.py b/feectools/__init__.py index e69de29bb..fabe3d391 100644 --- a/feectools/__init__.py +++ b/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 diff --git a/feectools/api/settings.py b/feectools/api/settings.py index d96c420c7..5fbd4095f 100644 --- a/feectools/api/settings.py +++ b/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 diff --git a/feectools/ddm/mpi.py b/feectools/ddm/mpi.py index 9b6caf23d..242bd4ea3 100644 --- a/feectools/ddm/mpi.py +++ b/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 diff --git a/feectools/ddm/partition.py b/feectools/ddm/partition.py index 8b2b0d3b7..538a58cc9 100644 --- a/feectools/ddm/partition.py +++ b/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): diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 1b6096d12..da932c149 100644 --- a/feectools/linalg/direct_solvers.py +++ b/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