Merge pull request #120 from underworldcode/feature/pixi-petsc-arch-static

Build cache coherence across PETSc version switches

Author: lmoresi

petsc-custom/activate-petsc-arch.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# activate-petsc-arch.sh — sourced by pixi on environment activation.
+#
+# Dynamically sets PETSC_ARCH for AMR environments (custom PETSc builds)
+# by combining the active pixi environment name with the PETSc version
+# pinned in petsc-custom/.petsc-version.
+#
+# This replaces hardcoded PETSC_ARCH entries in pixi.toml so that
+# switching PETSc versions (via ./uw petsc switch <tag>) does not require
+# editing tracked configuration.
+
+_uw_mpi=""
+_uw_suffix=""
+
+# Default MPI for AMR envs is platform-dependent and mirrors the conda deps
+# pinned in pixi.toml: openmpi on macOS (feature.amr.target.osx-arm64),
+# mpich on Linux (feature.amr.target.linux-64).
+_uw_default_mpi="mpich"
+[ "$(uname -s)" = "Darwin" ] && _uw_default_mpi="openmpi"
+
+case "${PIXI_ENVIRONMENT_NAME:-}" in
+    amr|amr-runtime|amr-dev)     _uw_mpi="$_uw_default_mpi" ;;
+    amr-mpich|amr-mpich-dev)     _uw_mpi="mpich" ;;
+    amr-openmpi|amr-openmpi-dev) _uw_mpi="openmpi" ;;
+    amr-debug)                   _uw_mpi="$_uw_default_mpi"; _uw_suffix="-debug" ;;
+esac
+
+if [ -n "$_uw_mpi" ]; then
+    # Match the version-discovery logic in ./uw (petsc_version_short):
+    # prefer .petsc-version, strip whitespace, fall back to 324.
+    _uw_ver="324"
+    _uw_ver_file="${PIXI_PROJECT_ROOT:-.}/petsc-custom/.petsc-version"
+    if [ -f "$_uw_ver_file" ]; then
+        _uw_read=$(tr -d '[:space:]' < "$_uw_ver_file" 2>/dev/null)
+        [ -n "$_uw_read" ] && _uw_ver="$_uw_read"
+    fi
+
+    export PETSC_ARCH="petsc-${_uw_ver}-uw-${_uw_mpi}${_uw_suffix}"
+fi
+
+unset _uw_mpi _uw_suffix _uw_default_mpi _uw_ver _uw_ver_file _uw_read

pixi.toml

@@ -165,6 +165,9 @@ cmake = ">=3.31,<4"
 make = ">=4.4,<5"
 mpi4py = ">=4,<5"
 
+[feature.amr.activation]
+scripts = ["petsc-custom/activate-petsc-arch.sh"]
+
 [feature.amr.activation.env]
 PETSC_DIR = "$PIXI_PROJECT_ROOT/petsc-custom/petsc"
 
@@ -173,16 +176,10 @@ openmpi = ">=5.0,<6"
 hdf5 = { version = ">=1.14,<2", build = "*openmpi*" }
 h5py = { version = ">=3.12,<4", build = "*openmpi*" }
 
-[feature.amr.target.osx-arm64.activation.env]
-PETSC_ARCH = "petsc-4-uw-openmpi"
-
 [feature.amr.target.linux-64.dependencies]
 hdf5 = { version = ">=1.14,<2", build = "*mpich*" }
 h5py = { version = ">=3.12,<4", build = "*mpich*" }
 
-[feature.amr.target.linux-64.activation.env]
-PETSC_ARCH = "petsc-4-uw-mpich"
-
 [feature.amr.tasks]
 petsc-local-build = { cmd = "./build-petsc.sh", cwd = "petsc-custom" }
 petsc-local-clean = { cmd = "./build-petsc.sh clean", cwd = "petsc-custom" }
@@ -202,9 +199,11 @@ mpi4py = ">=4,<5"
 hdf5 = { version = ">=1.14,<2", build = "*mpich*" }
 h5py = { version = ">=3.12,<4", build = "*mpich*" }
 
+[feature.amr-mpich.activation]
+scripts = ["petsc-custom/activate-petsc-arch.sh"]
+
 [feature.amr-mpich.activation.env]
 PETSC_DIR = "$PIXI_PROJECT_ROOT/petsc-custom/petsc"
-PETSC_ARCH = "petsc-4-uw-mpich"
 
 [feature.amr-mpich.tasks]
 petsc-local-build = { cmd = "./build-petsc.sh", cwd = "petsc-custom" }
@@ -221,9 +220,11 @@ mpi4py = ">=4,<5"
 hdf5 = { version = ">=1.14,<2", build = "*openmpi*" }
 h5py = { version = ">=3.12,<4", build = "*openmpi*" }
 
+[feature.amr-openmpi.activation]
+scripts = ["petsc-custom/activate-petsc-arch.sh"]
+
 [feature.amr-openmpi.activation.env]
 PETSC_DIR = "$PIXI_PROJECT_ROOT/petsc-custom/petsc"
-PETSC_ARCH = "petsc-4-uw-openmpi"
 
 [feature.amr-openmpi.tasks]
 petsc-local-build = { cmd = "./build-petsc.sh", cwd = "petsc-custom" }
@@ -236,19 +237,26 @@ petsc-local-clean = { cmd = "./build-petsc.sh clean", cwd = "petsc-custom" }
 # Uses the same conda deps as amr but points to a separate PETSC_ARCH
 # so the optimised build is untouched.
 #
-# Setup:
+# PETSC_ARCH is set dynamically by petsc-custom/activate-petsc-arch.sh
+# and resolves to:  petsc-<ver>-uw-<mpi>-debug
+# where <ver> comes from petsc-custom/.petsc-version and <mpi> follows
+# the platform default (openmpi on macOS, mpich on Linux).
+#
+# Setup (inside a pixi shell -e amr-debug, so $PETSC_ARCH is already set):
 #   cd petsc-custom/petsc
 #   pixi run -e amr-debug python3 ./configure \
-#       --with-petsc-arch=petsc-4-uw-openmpi-debug --with-debugging=1 \
+#       --with-petsc-arch="$PETSC_ARCH" --with-debugging=1 \
 #       --with-mpi-dir="$CONDA_PREFIX" --with-hdf5=0 \
 #       --download-mpich=0 --download-openmpi=0 --download-mpi4py=0 \
 #       --with-petsc4py=0 --with-x=0 --with-pragmatic=0 --with-slepc=0 \
 #       "--COPTFLAGS=-g -O0" "--CXXOPTFLAGS=-g -O0" "--FOPTFLAGS=-g -O0"
-#   make PETSC_DIR=$(pwd) PETSC_ARCH=petsc-4-uw-openmpi-debug all
+#   make PETSC_DIR=$(pwd) PETSC_ARCH="$PETSC_ARCH" all
+
+[feature.amr-debug.activation]
+scripts = ["petsc-custom/activate-petsc-arch.sh"]
 
 [feature.amr-debug.activation.env]
 PETSC_DIR = "$PIXI_PROJECT_ROOT/petsc-custom/petsc"
-PETSC_ARCH = "petsc-4-uw-openmpi-debug"
 
 # ============================================
 # HPC CLUSTER FEATURE

setup.py

@@ -76,46 +76,42 @@ def configure():
     LIBRARY_DIRS = []
     LIBRARIES = []
 
-    PETSC_DIR = ""
-    PETSC_ARCH = ""
-
-    # try get PETSC_DIR from petsc pip installation
-    try:
-        import petsc
-
-        PETSC_DIR = petsc.get_petsc_dir()
-    except:
-        pass
-
-    # PETSc
     import os
 
-    if not os.path.exists(PETSC_DIR):
-        print(f"PETSC_INFO from petsc4py - {petsc4py.get_config()}")
-        PETSC_DIR = petsc4py.get_config()["PETSC_DIR"]
-        PETSC_ARCH = petsc4py.get_config()["PETSC_ARCH"]
-
-    # It is preferable to use the petsc4py paths to the
-    # petsc libraries for consistency but the pip installation
-    # of PETSc sometimes points to the temporary setup up path
-
-    if not os.path.exists(PETSC_DIR):
-        print(f"PETSC_DIR {PETSC_DIR} is bad - trying another ...")
+    PETSC_DIR = ""
+    PETSC_ARCH = ""
 
-        if os.environ.get("CONDA_PREFIX") and not os.environ.get("PETSC_DIR"):
+    # Priority 1: Environment variables (set by pixi activation for custom builds)
+    if os.environ.get("PETSC_DIR") and os.path.exists(os.environ["PETSC_DIR"]):
+        PETSC_DIR = os.environ["PETSC_DIR"]
+        PETSC_ARCH = os.environ.get("PETSC_ARCH", "")
+
+    # Priority 2: petsc4py configuration (matches the installed petsc4py)
+    if not PETSC_DIR or not os.path.exists(PETSC_DIR):
+        config = petsc4py.get_config()
+        PETSC_DIR = config["PETSC_DIR"]
+        PETSC_ARCH = config.get("PETSC_ARCH", "")
+
+    # Priority 3: conda petsc package
+    if not PETSC_DIR or not os.path.exists(PETSC_DIR):
+        try:
+            import petsc
+            PETSC_DIR = petsc.get_petsc_dir()
+        except ImportError:
+            pass
+
+    # Priority 4: conda prefix fallback
+    if not PETSC_DIR or not os.path.exists(PETSC_DIR):
+        if os.environ.get("CONDA_PREFIX"):
             import sys
-
             py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
             PETSC_DIR = os.path.join(
                 os.environ["CONDA_PREFIX"],
                 "lib",
                 "python" + py_version,
                 "site-packages",
                 "petsc",
-            )  # symlink to latest python
-            PETSC_ARCH = os.environ.get("PETSC_ARCH", "")
-        else:
-            PETSC_DIR = os.environ["PETSC_DIR"]
+            )
             PETSC_ARCH = os.environ.get("PETSC_ARCH", "")
 
     print(f"Using PETSc:")

uw

@@ -159,7 +159,7 @@ run_build() {
         # Step 2: Check/build petsc4py
         if ! petsc4py_installed "$env"; then
             echo "  Installing petsc4py for $env..."
-            $PIXI run -e "$env" pip install "$PETSC_CUSTOM/src/binding/petsc4py" --no-build-isolation || {
+            $PIXI run -e "$env" pip install "$PETSC_CUSTOM/src/binding/petsc4py" --no-build-isolation --no-cache-dir || {
                 echo -e "${YELLOW}petsc4py build failed${NC}"
                 exit 1
             }
@@ -293,8 +293,25 @@ run_petsc_cmd() {
                 exit 1
             fi
             $PIXI run -e "$(get_env)" "$SCRIPT_DIR/petsc-custom/build-petsc.sh" checkout "$arg"
+
+            # Nuke stale build artifacts and petsc4py. The next ./uw build's
+            # target-change detection can't fire because the still-installed
+            # petsc4py reports the OLD arch, so we force a full rebuild here.
+            # (run_build already passes --no-cache-dir to pip install . so we
+            # don't need a global pip cache purge — only petsc4py's local-path
+            # install needs its own --no-cache-dir.)
+            echo ""
+            echo "Clearing stale build artifacts for clean rebuild..."
+            rm -rf "$SCRIPT_DIR"/build/lib.* "$SCRIPT_DIR"/build/temp.* "$SCRIPT_DIR"/build/bdist.* 2>/dev/null
+            rm -f "$SCRIPT_DIR/build/.petsc_target"
+            if is_amr_env "$(get_env)"; then
+                echo "Uninstalling old petsc4py (linked to previous PETSc build)..."
+                $PIXI run -e "$(get_env)" pip uninstall -y petsc4py 2>/dev/null || true
+            fi
+
             echo ""
-            echo "To complete the switch, rebuild:"
+            echo "To complete the switch, rebuild PETSc then underworld3:"
+            echo "  ./uw petsc build"
             echo "  ./uw build"
             ;;
         active)