Making ex18 much much faster (Experimental, Do-Not-Merge)

.github/scripts/validate_nsb_consumer.py

@@ -0,0 +1,167 @@
+"""Run NSB fixtures against a wheel in a separate consumer without build tools."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+import shutil
+import shlex
+import signal
+import subprocess
+import sys
+import venv
+
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def run(command, *, cwd, env, timeout=600):
+    command = [str(value) for value in command]
+    print("$ " + " ".join(command), flush=True)
+    with subprocess.Popen(command, cwd=cwd, env=env, start_new_session=True) as child:
+        try:
+            status = child.wait(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            os.killpg(child.pid, signal.SIGKILL)
+            child.wait()
+            raise
+        if status:
+            # Kill any surviving MPI workers before reporting the failure.
+            try:
+                os.killpg(child.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+            raise subprocess.CalledProcessError(status, command)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        '--backend', choices=('serial', 'openmpi', 'mpich'), required=True
+    )
+    parser.add_argument('--runtime-version', required=True)
+    parser.add_argument('--wheel-dir', type=Path, required=True)
+    parser.add_argument('--work-dir', type=Path, required=True)
+    args = parser.parse_args()
+    work = args.work_dir.resolve()
+    wheels = args.wheel_dir.resolve()
+    if work == ROOT or ROOT in work.parents:
+        parser.error('consumer directory must be outside the checkout')
+    work.mkdir(parents=True, exist_ok=True)
+    if any(work.iterdir()):
+        parser.error('consumer directory must be empty')
+    pymfem_wheels = list(wheels.glob('mfem-*.whl'))
+    if len(pymfem_wheels) != 1:
+        parser.error('expected exactly one PyMFEM wheel')
+    mpi = args.backend != 'serial'
+    env = dict(os.environ)
+    for key in ('PYTHONPATH', 'PYTHONHOME', 'PIP_TARGET', 'PIP_PREFIX', 'PIP_USER'):
+        env.pop(key, None)
+    env.update(PYTHONNOUSERSITE='1', OPENBLAS_NUM_THREADS='1', OMP_NUM_THREADS='1')
+    consumer = work / 'venv'
+    venv.EnvBuilder(with_pip=True).create(consumer)
+    python = consumer / 'bin/python'
+    run([python, '-m', 'pip', 'install', f'numba-swig-bridge-rt=={args.runtime_version}',
+         'numba==0.68.0', 'numpy>=2.4.1', 'scipy'], cwd=work, env=env)
+    if mpi:
+        mpi_wheels = list(wheels.glob('mpi4py-*.whl'))
+        if len(mpi_wheels) != 1:
+            parser.error('expected the source-built mpi4py wheel')
+        run([python, '-m', 'pip', 'install', '--no-deps', *mpi_wheels], cwd=work, env=env)
+    run([python, '-m', 'pip', 'install', '--no-deps', *pymfem_wheels], cwd=work, env=env)
+    run([python, '-m', 'pip', 'check'], cwd=work, env=env)
+
+    # Copy only fixture inputs; never copy the source mfem package.
+    fixtures = work / 'fixtures'
+    for directory in ('examples', 'data'):
+        shutil.copytree(ROOT / directory, fixtures / directory,
+                        ignore=shutil.ignore_patterns('__pycache__', '*.pyc', '*.gf'))
+    tests = fixtures / 'test/bridge'
+    tests.mkdir(parents=True)
+    for source in (ROOT / 'test/bridge').glob('test_*.py'):
+        if source.name != 'test_nsb_state_build.py':
+            shutil.copy2(source, tests / source.name)
+    env['PATH'] = str(python.parent)
+    mpi_command = []
+    if mpi:
+        tools = work / 'mpi-tools'
+        tools.mkdir()
+        original_path = os.environ.get('PATH')
+        launcher = shutil.which(f'mpiexec.{args.backend}', path=original_path)
+        if not launcher:
+            parser.error('selected MPI launcher is unavailable')
+        if args.backend == 'openmpi':
+            # An absolute Open MPI launcher path can implicitly enable prefix
+            # propagation, exposing /usr/bin (and SWIG) inside the workers.
+            # Use the same wrapper for direct and acceptance-driver launches.
+            wrapper = tools / 'mpiexec'
+            wrapper.write_text(
+                '#!/bin/sh\nexec ' + shlex.quote(launcher) + ' --noprefix "$@"\n'
+            )
+            wrapper.chmod(0o755)
+        else:
+            (tools / 'mpiexec').symlink_to(launcher)
+        for name in ('ssh', 'rsh', 'orted', 'prted', 'hydra_pmi_proxy'):
+            tool = shutil.which(name, path=original_path)
+            if tool:
+                (tools / name).symlink_to(tool)
+        env['PATH'] += os.pathsep + str(tools)
+        if args.backend == 'openmpi':
+            env['OMPI_MCA_rmaps_base_oversubscribe'] = '1'
+        mpi_command = [tools / 'mpiexec', '-n', '2']
+
+    check = work / 'check_install.py'
+    check.write_text(
+        "import importlib, importlib.util, pathlib, shutil, sys\n"
+        "assert importlib.util.find_spec('numba_swig_bridge') is None\n"
+        "assert importlib.util.find_spec('swig') is None\n"
+        "assert shutil.which('swig') is None\n"
+        + ("assert importlib.util.find_spec('mpi4py') is None\n" if not mpi else '')
+        + "import mfem, nsb_rt\n"
+        "assert pathlib.Path(mfem.__file__).resolve().is_relative_to(pathlib.Path(sys.prefix))\n"
+        "assert pathlib.Path(nsb_rt.__file__).resolve().is_relative_to(pathlib.Path(sys.prefix))\n"
+        f"assert nsb_rt.__version__ == {args.runtime_version!r}\n"
+        "backend = importlib.import_module('mfem.' + sys.argv[1])\n"
+        "registration = backend.get_bridge_registration()\n"
+        "assert pathlib.Path(registration.__file__).resolve().is_relative_to(pathlib.Path(sys.prefix))\n"
+        "print('Installed wheel and runtime verified:', mfem.__file__, flush=True)\n"
+    )
+    run([python, '-I', check, 'ser'], cwd=work, env=env)
+    completed = []
+    excluded = {'test_ex18p.py', 'test_ex9p.py', 'test_ex19p.py', 'test_parallel_foundation.py',
+                'test_gridfunctionp.py', 'test_native_operatorsp.py'}
+    for test in sorted(tests.glob('test_*.py')):
+        if test.name not in excluded:
+            run([python, '-I', test], cwd=work, env=env, timeout=1200)
+            completed.append(test.name)
+    if mpi:
+        mpi_check = work / 'check_mpi.py'
+        mpi_check.write_text(
+            "from mpi4py import MPI\n"
+            "version = MPI.Get_library_version()\n"
+            f"assert {args.backend!r} in version.lower().replace(' ', ''), version\n"
+            "assert MPI.COMM_WORLD.size == 2\n"
+            "print('MPI rank', MPI.COMM_WORLD.rank, 'size', MPI.COMM_WORLD.size, flush=True)\n"
+        )
+        run([*mpi_command, python, '-I', mpi_check], cwd=work, env=env)
+        run([*mpi_command, python, '-I', check, 'par'], cwd=work, env=env)
+        for name in ('test_parallel_foundation.py', 'test_ex18p.py',
+                     'test_gridfunctionp.py', 'test_native_operatorsp.py'):
+            run([*mpi_command, python, '-I', tests / name], cwd=work, env=env)
+            completed.append(name)
+        # These acceptance drivers launch their own two-rank workers.
+        for name in ('test_ex9p.py', 'test_ex19p.py'):
+            run([python, '-I', tests / name], cwd=work, env=env, timeout=1800)
+            completed.append(name)
+        run([*mpi_command, python, '-I', fixtures / 'examples/bridge/ex4p.py',
+             '--refine', '0', '--compare'], cwd=work, env=env)
+        completed.append('ex4p.py --compare')
+    (work / 'report.json').write_text(json.dumps({
+        'backend': args.backend, 'runtime': args.runtime_version,
+        'wheel': pymfem_wheels[0].name, 'checks': completed, 'status': 'passed',
+    }, indent=2) + '\n')
+    print('NSB installed-consumer validation passed:', args.backend, flush=True)
+
+
+if __name__ == '__main__':
+    main()

.github/workflows/Build-and-test-NSB-callable.yml

@@ -0,0 +1,87 @@
+name: Build and test NSB (callable)
+
+on:
+  workflow_call:
+    inputs:
+      backend:
+        description: 'serial, openmpi or mpich'
+        type: string
+        required: true
+      python-version:
+        type: string
+        default: '3.13'
+      runtime-version:
+        type: string
+        default: '0.15.5'
+
+permissions:
+  contents: read
+
+jobs:
+  wheel-consumer:
+    runs-on: ubuntu-24.04
+    timeout-minutes: 120
+    env:
+      OPENBLAS_NUM_THREADS: '1'
+      OMP_NUM_THREADS: '1'
+      PYMFEM_BUILD_NCPU: '2'
+      NSB_BACKEND: ${{ inputs.backend }}
+      NSB_RUNTIME_VERSION: ${{ inputs.runtime-version }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ inputs.python-version }}
+      - name: Validate backend input
+        run: |
+          case "$NSB_BACKEND" in
+            serial|openmpi|mpich) ;;
+            *) echo "Unsupported backend: $NSB_BACKEND" >&2; exit 2 ;;
+          esac
+      - name: Use local UCX transports for MPICH
+        if: inputs.backend == 'mpich'
+        run: echo 'UCX_TLS=tcp,sm,self' >> "$GITHUB_ENV"
+      - name: Set up MPI (including Ubuntu MPICH hotfix)
+        if: inputs.backend != 'serial'
+        uses: mpi4py/setup-mpi@v1
+        with:
+          mpi: ${{ inputs.backend }}
+      - name: Install native build tools
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends build-essential zlib1g-dev openssh-client
+      - name: Prepare separate build environment
+        run: |
+          python -m venv "$RUNNER_TEMP/nsb-build"
+          "$RUNNER_TEMP/nsb-build/bin/python" -m pip install --upgrade \
+            'pip>=25.1' 'setuptools>=80.0.1' wheel 'cmake>=4.0.0' \
+            'swig==4.5.1' 'numpy>=2.4.1' 'numba==0.68.0' scipy \
+            "numba-swig-bridge-rt==$NSB_RUNTIME_VERSION"
+          mkdir -p "$RUNNER_TEMP/nsb-wheels"
+          if [ "$NSB_BACKEND" != serial ]; then
+            MPI4PY_BUILD_MPICC="mpicc.$NSB_BACKEND" \
+              "$RUNNER_TEMP/nsb-build/bin/python" -m pip wheel \
+              'mpi4py==4.1.2' --no-deps --no-binary=mpi4py --no-cache-dir \
+              --wheel-dir "$RUNNER_TEMP/nsb-wheels"
+            "$RUNNER_TEMP/nsb-build/bin/python" -m pip install "$RUNNER_TEMP"/nsb-wheels/mpi4py-*.whl
+          fi
+      - name: Build PyMFEM wheel from committed artifacts
+        run: |
+          export PATH="$RUNNER_TEMP/nsb-build/bin:$PATH"
+          export SWIG="$RUNNER_TEMP/nsb-build/bin/swig"
+          "$SWIG" -version
+          flags=(-Cwith-numba-swig-bridge=Yes)
+          if [ "$NSB_BACKEND" != serial ]; then
+            flags+=(-Cwith-parallel=Yes "-CMPICC=mpicc.$NSB_BACKEND" "-CMPICXX=mpicxx.$NSB_BACKEND")
+          fi
+          "$RUNNER_TEMP/nsb-build/bin/python" -I -c \
+            "import importlib.util; assert importlib.util.find_spec('numba_swig_bridge') is None"
+          "$RUNNER_TEMP/nsb-build/bin/python" -m pip wheel . \
+            --no-deps --no-build-isolation --verbose \
+            --wheel-dir "$RUNNER_TEMP/nsb-wheels" "${flags[@]}"
+      - name: Test fresh installed wheel consumers
+        run: |
+          "$RUNNER_TEMP/nsb-build/bin/python" .github/scripts/validate_nsb_consumer.py \
+            --backend "$NSB_BACKEND" --runtime-version "$NSB_RUNTIME_VERSION" \
+            --wheel-dir "$RUNNER_TEMP/nsb-wheels" \
+            --work-dir "$RUNNER_TEMP/nsb-consumer"

.github/workflows/Build-and-test-NSB-dispatch.yml

@@ -0,0 +1,22 @@
+name: Build and test NSB (dispatch)
+
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  nsb:
+    if: ${{ !contains(github.event.pull_request.labels.*.name, 'skip-ci') }}
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [serial, openmpi, mpich]
+    name: NSB | ${{ matrix.backend }} | Python 3.13
+    uses: ./.github/workflows/Build-and-test-NSB-callable.yml
+    with:
+      backend: ${{ matrix.backend }}
+

.github/workflows/build-and-test-callable.yml

@@ -103,6 +103,7 @@ jobs:
     - name: Install MPI
       if: inputs.parallel
       run: |
+        sudo apt-get update
         sudo apt-get install openmpi-bin libopenmpi-dev
         export OMPI_MCA_rmaps_base_oversubscribe=1
         # (broken?) sudo apt-get install mpich libmpich-dev
@@ -123,7 +124,8 @@ jobs:
     - name: Download CUDA
       if: inputs.cuda && steps.cache-cuda.outputs.cache-hit == false
       run: |
-        CUDA_URL="https://developer.download.nvidia.com/compute/cuda/${{ inputs.cuda-toolkit-version }}/local_installers/cuda_${{ inputs.cuda-toolkit-version }}_${{ inputs.cuda-driver-version }}_linux.run"
+        ## CUDA_URL="https://developer.download.nvidia.com/compute/cuda/${{ inputs.cuda-toolkit-version }}/local_installers/cuda_${{ inputs.cuda-toolkit-version }}_${{ inputs.cuda-driver-version }}_linux.run"
+        CUDA_URL="https://developer.download.nvidia.com/compute/cuda/${{ inputs.cuda-toolkit-version }}/local_installers/cuda_${{ inputs.cuda-toolkit-version }}_${{ inputs.cuda-driver-version }}_[...]
         curl -o ~/cache/cuda.run --create-dirs $CUDA_URL
 
     - name: Install CUDA

.gitignore

@@ -2,6 +2,7 @@
 *.so
 *.o
 *.a
+/external/
 .DS_store
 #*
 .#*
@@ -12,9 +13,32 @@ build
 dist
 test/*
 !test/*.py
+!test/bridge/
+!test/bridge/*.py
+!test/bridge/README.md
 !test/cpp
 !test/cpp/
 !test/issues/*.py
 !test/issues/data_3960/VisIt_volume
-*.egg-info
-*.py#
+*.egg-info
+*.py#
+mfem/_ser/*_wrap.cxx 
+mfem/_ser/*_wrap.h
+mfem/_par/*_wrap.cxx
+mfem/_par/*_wrap.h
+mfem/_ser/*.py
+mfem/_par/*.py
+!mfem/_par/setup.py
+!mfem/_par/__init__.py
+!mfem/_par/tmop_modules.py
+!mfem/_par/bridge_registration.py
+!mfem/_ser/setup.py
+!mfem/_ser/__init__.py
+!mfem/_ser/tmop_modules.py
+!mfem/_ser/bridge_registration.py
+mfem/__sha__.py
+mfem/common/existing_mfem_headers.i
+mfem/_ser/director_descriptor.json
+mfem/_ser/ordinary_descriptor.json
+mfem/_par/director_descriptor.json
+mfem/_par/ordinary_descriptor.json

INSTALL.md

@@ -47,7 +47,7 @@ and `--with-gslib` flags as follows.
 pip install . -C"with-parallel=Yes" -C"with-gslib=Yes"
 ```
 
-(Warning) The migration to `pip install . ` is on-going effort and Some of the example commands are not tested yet, which are indicated by using old conversion of "python setup.py install XXXX"
+
 
 ## Environmental variable
 `PYMFEM_BUILD_NCPU` :  set number of parallel jobs used when calling SWIG and build-extension.
@@ -182,6 +182,15 @@ pip install . -C"skip-ext=Yes"  -C"skip-swig=Yes" --verbose
 pip install . -C"with-parallel=Yes" -C"skip-ext=Yes"  -C"skip-swig=Yes" --verbose
 ```
 
+### Experimental Numba bridge (prebuilt binaries are not available)
+```shell
+pip install ".[bridge]" -C"with-numba-swig-bridge=Yes"
+```
+
+The `[bridge]` extra installs the bridge runtime in the target environment;
+`-C"with-numba-swig-bridge=Yes"` enables bridge support in the source build.
+Both are needed for a complete source installation.
+
 ### Other options
 `--unverifiedSSL` :
    This addresses error relating SSL certificate. Typical error message is

MANIFEST.in

@@ -13,3 +13,8 @@ include docs/*.txt
 include docs/*.png
 include external/README.txt
 exclude mfem/*/*.so
+
+include mfem/_ser/bridge.toml mfem/_par/bridge.toml
+recursive-include mfem/common/bridges *.i
+recursive-include mfem/_ser/nsb_artifacts *
+recursive-include mfem/_par/nsb_artifacts *