Files
firedrake-conda/setup.sh

578 lines
21 KiB
Bash
Executable File

#!/usr/bin/env bash
# ============================================================================
# setup.sh — Install Firedrake + icepack into a conda environment
#
# This script:
# 1. Creates a conda environment with all native dependencies from conda-forge
# 2. Sets environment variables so Firedrake finds conda's PETSc, MPI, etc.
# 3. pip-installs Firedrake and icepack (skipping binary builds for
# MPI-linked packages so they compile against conda's libraries)
#
# Usage:
# bash setup.sh # env named "firedrake"
# bash setup.sh my-env # env named "my-env"
# bash setup.sh ./envs/firedrake # env at path ./envs/firedrake
# FIREDRAKE_ENV=~/fd FIREDRAKE_REF=main bash setup.sh
#
# Prerequisites:
# - conda or mamba installed
# - internet access (to download packages)
#
# Tested on: Linux x86_64, macOS ARM64
# ============================================================================
set -euo pipefail
# ── Configuration ──────────────────────────────────────────────────────────
# Environment name or path — accepts:
# - A simple name like "firedrake" → conda env -n firedrake
# - An absolute/relative path like "./envs/fd" → conda env -p ./envs/fd
# Set via FIREDRAKE_ENV env var, first positional arg, or default "firedrake".
ENV_TARGET="${FIREDRAKE_ENV:-${1:-firedrake}}"
# Detect if it's a path (contains / or starts with .)
if [[ "$ENV_TARGET" == */* ]] || [[ "$ENV_TARGET" == .* ]]; then
# Resolve to absolute path, creating parent dir if needed
_parent="$(dirname "$ENV_TARGET")"
mkdir -p "$_parent" 2>/dev/null || true
ENV_TARGET="$(cd "$_parent" && pwd)/$(basename "$ENV_TARGET")"
ENV_IS_PATH=true
CONDA_ENV_FLAG="-p"
else
ENV_IS_PATH=false
CONDA_ENV_FLAG="-n"
fi
ENV_DISPLAY="$ENV_TARGET"
ENV_ACTIVATE="$ENV_TARGET"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Firedrake version: set to a release tag (e.g. "2025.10.2") or "main"
FIREDRAKE_REF="${FIREDRAKE_REF:-${2:-2025.10.2}}"
# Prefer mamba if available (much faster solver)
if command -v mamba &> /dev/null; then
CONDA=mamba
elif command -v conda &> /dev/null; then
CONDA=conda
else
echo "ERROR: Neither conda nor mamba found. Please install one first."
echo " See: https://docs.conda.io/en/latest/miniconda.html"
exit 1
fi
# ── Colors ─────────────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
DIM='\033[2m'
NC='\033[0m' # No color
info() { echo -e "${BLUE}[INFO]${NC} $*"; }
ok() { echo -e "${GREEN}[OK]${NC} $*"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
fail() { echo -e "${RED}[FAIL]${NC} $*"; exit 1; }
# ── Live rolling tail ─────────────────────────────────────────────────────
# Runs a command while showing a live, continuously-updating view of the
# last N lines of its combined stdout/stderr. When the command finishes,
# the rolling display is cleared and replaced with a summary.
#
# Usage: run_live <n_lines> <description> <command...>
# Returns the exit code of the command.
TAIL_LINES=20
run_live() {
local nlines="$1"; shift
local desc="$1"; shift
local logfile
logfile=$(mktemp /tmp/firedrake-install-XXXXXX.log)
echo -e "${DIM}─── ${desc} ───${NC}"
# Run the actual command in the background, directing all output to logfile
"$@" > "$logfile" 2>&1 &
local cmd_pid=$!
# Tail the logfile in real-time, showing last N lines with in-place rewriting
python3 -u -c "
import sys, os, time
nlines = int(sys.argv[1])
path = sys.argv[2]
pid = int(sys.argv[3])
cols = int(os.environ.get('COLUMNS', 120))
buf = []
shown = 0
def render():
global shown
# Move up to overwrite previous output
if shown > 0:
sys.stdout.write(f'\033[{shown}A')
shown = 0
for b in buf:
sys.stdout.write(f'\033[2K {b}\n')
shown += 1
sys.stdout.flush()
def pid_alive(p):
try:
os.kill(p, 0)
return True
except ProcessLookupError:
return False
with open(path, 'r') as f:
while True:
line = f.readline()
if line:
line = line.rstrip('\n')[:cols]
buf.append(line)
if len(buf) > nlines:
buf = buf[-nlines:]
render()
else:
# No new data — check if process is still running
if not pid_alive(pid):
# Drain any remaining lines
for line in f:
line = line.rstrip('\n')[:cols]
buf.append(line)
if len(buf) > nlines:
buf = buf[-nlines:]
render()
break
time.sleep(0.05)
# Clear the rolling display
if shown > 0:
sys.stdout.write(f'\033[{shown}A')
for _ in range(shown):
sys.stdout.write('\033[2K\n')
sys.stdout.write(f'\033[{shown}A')
sys.stdout.flush()
" "$nlines" "$logfile" "$cmd_pid"
# Collect the exit code
local exit_code=0
wait "$cmd_pid" || exit_code=$?
if [[ $exit_code -ne 0 ]]; then
warn "${desc} failed (exit code ${exit_code}). Full log: ${logfile}"
echo -e "${DIM}─── Last 30 lines of ${logfile} ───${NC}"
tail -30 "$logfile"
echo -e "${DIM}────────────────────────────────────${NC}"
else
rm -f "$logfile"
fi
return $exit_code
}
# ── Step 1: Create conda environment ──────────────────────────────────────
info "Creating conda environment '${ENV_DISPLAY}' from environment.yml ..."
# Check if environment already exists
ENV_EXISTS=false
if [[ "$ENV_IS_PATH" == true ]]; then
[[ -d "${ENV_TARGET}/conda-meta" ]] && ENV_EXISTS=true
else
$CONDA env list 2>/dev/null | grep -q "^${ENV_TARGET} " && ENV_EXISTS=true
fi
if [[ "$ENV_EXISTS" == true ]]; then
warn "Environment '${ENV_DISPLAY}' already exists."
read -p " Recreate it? [y/N] " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
$CONDA env remove ${CONDA_ENV_FLAG} "${ENV_TARGET}" -y
else
info "Reusing existing environment."
fi
fi
$CONDA env create ${CONDA_ENV_FLAG} "${ENV_TARGET}" -f "${SCRIPT_DIR}/environment.yml" || \
$CONDA env update ${CONDA_ENV_FLAG} "${ENV_TARGET}" -f "${SCRIPT_DIR}/environment.yml"
ok "Conda environment created."
# ── Step 2: Activate and detect paths ─────────────────────────────────────
# We need to source conda's shell hook to use `conda activate` in a script
eval "$(conda shell.bash hook)"
set +u # Allow referencing unset variables for the conda activation script (e.g. GFORTRAN)
conda activate "${ENV_ACTIVATE}"
set -u
CONDA_PREFIX="${CONDA_PREFIX}"
info "CONDA_PREFIX = ${CONDA_PREFIX}"
# Detect platform
OS="$(uname -s)"
ARCH="$(uname -m)"
info "Platform: ${OS} ${ARCH}"
if [[ "$OS" == "Darwin" ]]; then
SOEXT="dylib"
else
SOEXT="so"
fi
# ── Step 3: Verify conda packages ────────────────────────────────────────
info "Verifying critical conda packages ..."
# Check PETSc
if [[ ! -f "${CONDA_PREFIX}/lib/libpetsc.${SOEXT}" ]]; then
fail "PETSc shared library not found at ${CONDA_PREFIX}/lib/libpetsc.${SOEXT}"
fi
ok "PETSc found."
# Check MPI compiler wrappers
if ! command -v mpicc &> /dev/null; then
fail "mpicc not found. MPI compiler wrappers are missing."
fi
ok "MPI compilers found: $(which mpicc)"
# Check HDF5
if [[ ! -f "${CONDA_PREFIX}/lib/libhdf5.${SOEXT}" ]]; then
fail "HDF5 not found."
fi
ok "HDF5 found."
# ── Step 4: Set environment variables ─────────────────────────────────────
info "Setting environment variables for Firedrake ..."
# PETSc: conda installs in "prefix" mode (no PETSC_ARCH)
export PETSC_DIR="${CONDA_PREFIX}"
export PETSC_ARCH=""
# Compilers: use MPI wrappers from conda
export CC=mpicc
export CXX=mpicxx
export FC=mpifort
# HDF5: tell h5py and Firedrake to use MPI-enabled HDF5
export HDF5_MPI="ON"
export HDF5_DIR="${CONDA_PREFIX}"
# Firedrake recommends single-threaded BLAS
export OMP_NUM_THREADS=1
# PyOP2 JIT: conda's mpicc wraps x86_64-conda-linux-gnu-cc which PyOP2
# doesn't recognise → falls back to generic compiler without -shared/-fPIC
export PYOP2_CFLAGS="-fPIC -O2"
export PYOP2_LDFLAGS="-shared -Wl,-rpath,${CONDA_PREFIX}/lib"
# Tell pip not to use cached wheels for MPI-linked packages
# (they may be linked against the wrong MPI/PETSc)
export PIP_NO_CACHE_DIR=1
info " PETSC_DIR = ${PETSC_DIR}"
info " PETSC_ARCH = (empty — prefix install)"
info " CC = ${CC}"
info " HDF5_MPI = ${HDF5_MPI}"
# ── Step 5: Pip install Firedrake ecosystem ───────────────────────────────
info "Installing Firedrake and dependencies via pip ..."
info "This may take several minutes (compiling C extensions) ..."
# Purge pip cache to avoid stale binaries linked against wrong libs
pip cache purge 2>/dev/null || true
# Install packages that have C extensions and must link against conda's libs.
# We use --no-binary to force compilation against the conda environment.
# The order matters: dependencies first, then Firedrake, then icepack.
info " [1/5] Installing build dependencies ..."
# These are needed as build backends / build-time deps by libsupermesh
# and firedrake. We install them now so --no-build-isolation works
# (which ensures C extensions link against conda's MPI/PETSc, not
# some isolated copy).
# Install ALL build backends that Firedrake's dependency tree uses.
# --no-build-isolation means pip won't auto-install these, so we need
# them all present before building anything.
run_live $TAIL_LINES "Installing build backends" \
pip install \
setuptools wheel \
scikit-build-core \
meson-python meson ninja \
hatchling hatch-vcs hatch-fancy-pypi-readme \
flit-core \
cython \
pybind11 \
petsctools
ok "Build dependencies installed."
info " [2/5] Installing mpi4py, petsc4py, h5py (from source against conda libs) ..."
# These should already be installed by conda, but if pip needs to reinstall
# them (e.g. for Firedrake version constraints), force source builds.
# Usually conda's versions are fine and pip will skip them.
info " [3/5] Installing libsupermesh ..."
run_live $TAIL_LINES "libsupermesh" \
pip install --no-build-isolation libsupermesh || {
warn "libsupermesh pip install failed, trying from git ..."
run_live $TAIL_LINES "libsupermesh (git)" \
pip install --no-build-isolation \
"libsupermesh @ git+https://github.com/firedrakeproject/libsupermesh.git"
}
ok "libsupermesh installed."
info " [4/5] Installing Firedrake (${FIREDRAKE_REF}) ..."
# Firedrake's pyproject.toml lists dependencies like petsc4py, mpi4py, h5py
# as git URLs. Conda already provides working versions of these, so we
# generate a pip constraints file that pins them to the conda-installed
# versions. This prevents pip from cloning & rebuilding petsc4py (which
# fails on release tags due to a setuptools dry_run incompatibility).
#
# Clone first so the user can see git progress (pip hides it entirely).
# Clones live in ~/.firedrake-conda/clones/ so they persist across
# environment recreations (conda env remove + re-create).
CLONE_DIR="${HOME}/.firedrake-conda/clones"
mkdir -p "${CLONE_DIR}"
info " Cloning Firedrake (${FIREDRAKE_REF}) ..."
if [[ -d "${CLONE_DIR}/firedrake" ]]; then
info " (updating existing clone)"
# Unshallow if needed (shallow clones can't switch tags)
if [[ -f "${CLONE_DIR}/firedrake/.git/shallow" ]]; then
git -C "${CLONE_DIR}/firedrake" fetch --progress --unshallow --tags origin 2>&1 || \
git -C "${CLONE_DIR}/firedrake" fetch --progress --tags origin 2>&1
else
git -C "${CLONE_DIR}/firedrake" fetch --progress --tags origin 2>&1
fi
git -C "${CLONE_DIR}/firedrake" checkout "${FIREDRAKE_REF}" 2>&1
if [[ "${FIREDRAKE_REF}" == "main" ]]; then
git -C "${CLONE_DIR}/firedrake" reset --hard origin/main
fi
else
# Full clone for tags; shallow clone for main
if [[ "${FIREDRAKE_REF}" == "main" ]]; then
git clone --progress --depth 1 --branch main \
https://github.com/firedrakeproject/firedrake.git \
"${CLONE_DIR}/firedrake" 2>&1
else
git clone --progress --branch "${FIREDRAKE_REF}" \
https://github.com/firedrakeproject/firedrake.git \
"${CLONE_DIR}/firedrake" 2>&1
fi
fi
# ── Generate pip constraints file ──────────────────────────────────────
# Pin every conda-provided package so pip doesn't try to rebuild them
# from git URLs in Firedrake's pyproject.toml.
info " Generating pip constraints for conda packages ..."
CONSTRAINTS="${CONDA_PREFIX}/constraints.txt"
python3 - <<'PYEOF' > "${CONSTRAINTS}"
import importlib.metadata, sys
# Packages that conda provides and pip must not rebuild
CONDA_OWNS = [
"petsc4py", "slepc4py", "mpi4py", "h5py", "numpy", "scipy",
"sympy", "cython", "pybind11", "pkgconfig", "rtree",
"scikit-build-core", "hatchling", "flit-core", "setuptools",
]
for pkg in CONDA_OWNS:
try:
version = importlib.metadata.version(pkg)
print(f"{pkg}=={version}")
except importlib.metadata.PackageNotFoundError:
pass
PYEOF
info " Constraints file:"
cat "${CONSTRAINTS}" | sed 's/^/ /'
# ── Patch pyproject.toml to remove conda-provided deps ─────────────────
# Firedrake's release tags may list petsc4py, slepc4py, mpi4py, h5py as
# git URL deps (@ git+https://...) or version pins. Pip can't satisfy
# direct URL refs from conda, and tries to rebuild them — which fails
# (e.g. petsc4py's setuptools dry_run bug). We strip these deps since
# conda already provides them.
info " Patching pyproject.toml to skip conda-provided deps ..."
PYPROJECT="${CLONE_DIR}/firedrake/pyproject.toml"
python3 - "${PYPROJECT}" <<'PYEOF'
import re, sys
path = sys.argv[1]
text = open(path).read()
# Packages conda provides — remove any line in dependencies[] that
# starts with these names (handles version pins AND @ git+... URLs)
CONDA_PKGS = ["petsc4py", "slepc4py", "mpi4py", "h5py"]
lines = text.split("\n")
patched = []
removed = []
for line in lines:
stripped = line.strip().strip('"').strip("'").strip(",").strip()
skip = False
for pkg in CONDA_PKGS:
# Match "petsc4py...", " petsc4py...", '"petsc4py...', etc.
if re.match(rf"^{re.escape(pkg)}(\s|>|<|=|@|;|$)", stripped):
skip = True
removed.append(stripped)
break
if not skip:
patched.append(line)
if removed:
open(path, "w").write("\n".join(patched))
for r in removed:
print(f" Removed: {r}")
else:
print(" (no conda-provided deps found to remove)")
PYEOF
# ── pip install with constraints ───────────────────────────────────────
info " pip install Firedrake (from local clone) ..."
run_live $TAIL_LINES "Firedrake" \
pip install --no-build-isolation -c "${CONSTRAINTS}" "${CLONE_DIR}/firedrake"
ok "Firedrake installed."
info " [5/5] Installing icepack ..."
info " Cloning icepack ..."
if [[ -d "${CLONE_DIR}/icepack" ]]; then
info " (updating existing clone)"
git -C "${CLONE_DIR}/icepack" pull --progress 2>&1
else
git clone --progress --depth 1 \
https://github.com/icepack/icepack.git \
"${CLONE_DIR}/icepack" 2>&1
fi
info " pip install icepack (from local clone) ..."
run_live $TAIL_LINES "icepack" \
pip install --no-build-isolation -c "${CONSTRAINTS}" "${CLONE_DIR}/icepack"
ok "icepack installed."
# ── Step 6: Install activation script ────────────────────────────────────
info "Installing conda activation/deactivation scripts ..."
ACTIVATE_DIR="${CONDA_PREFIX}/etc/conda/activate.d"
DEACTIVATE_DIR="${CONDA_PREFIX}/etc/conda/deactivate.d"
mkdir -p "${ACTIVATE_DIR}" "${DEACTIVATE_DIR}"
cat > "${ACTIVATE_DIR}/firedrake-env-vars.sh" << 'ACTIVATE_EOF'
#!/usr/bin/env bash
# Set environment variables needed by Firedrake's JIT compiler.
# These are automatically set on `conda activate`.
# Save previous values so we can restore on deactivate
export _FIREDRAKE_OLD_CC="${CC:-}"
export _FIREDRAKE_OLD_CXX="${CXX:-}"
export _FIREDRAKE_OLD_FC="${FC:-}"
export _FIREDRAKE_OLD_PETSC_DIR="${PETSC_DIR:-}"
export _FIREDRAKE_OLD_PETSC_ARCH="${PETSC_ARCH:-}"
export _FIREDRAKE_OLD_HDF5_MPI="${HDF5_MPI:-}"
export _FIREDRAKE_OLD_HDF5_DIR="${HDF5_DIR:-}"
export _FIREDRAKE_OLD_OMP_NUM_THREADS="${OMP_NUM_THREADS:-}"
export _FIREDRAKE_OLD_PYOP2_CFLAGS="${PYOP2_CFLAGS:-}"
export _FIREDRAKE_OLD_PYOP2_LDFLAGS="${PYOP2_LDFLAGS:-}"
# PETSc
export PETSC_DIR="${CONDA_PREFIX}"
export PETSC_ARCH=""
# Compilers — use MPI wrappers so JIT-compiled kernels link against
# the correct MPI and can be loaded alongside PETSc.
export CC=mpicc
export CXX=mpicxx
export FC=mpifort
# HDF5
export HDF5_MPI="ON"
export HDF5_DIR="${CONDA_PREFIX}"
# Firedrake recommends single-threaded BLAS to avoid oversubscription
export OMP_NUM_THREADS=1
# PyOP2 JIT flags — conda's mpicc wraps x86_64-conda-linux-gnu-cc
# which PyOP2 doesn't recognise, so it falls back to a generic compiler
# class that omits -shared and -fPIC. These env vars inject the flags.
export PYOP2_CFLAGS="-fPIC -O2"
export PYOP2_LDFLAGS="-shared -Wl,-rpath,${CONDA_PREFIX}/lib"
ACTIVATE_EOF
cat > "${DEACTIVATE_DIR}/firedrake-env-vars.sh" << 'DEACTIVATE_EOF'
#!/usr/bin/env bash
# Restore environment variables on `conda deactivate`.
export CC="${_FIREDRAKE_OLD_CC}"
export CXX="${_FIREDRAKE_OLD_CXX}"
export FC="${_FIREDRAKE_OLD_FC}"
export PETSC_DIR="${_FIREDRAKE_OLD_PETSC_DIR}"
export PETSC_ARCH="${_FIREDRAKE_OLD_PETSC_ARCH}"
export HDF5_MPI="${_FIREDRAKE_OLD_HDF5_MPI}"
export HDF5_DIR="${_FIREDRAKE_OLD_HDF5_DIR}"
export OMP_NUM_THREADS="${_FIREDRAKE_OLD_OMP_NUM_THREADS}"
export PYOP2_CFLAGS="${_FIREDRAKE_OLD_PYOP2_CFLAGS}"
export PYOP2_LDFLAGS="${_FIREDRAKE_OLD_PYOP2_LDFLAGS}"
unset _FIREDRAKE_OLD_CC
unset _FIREDRAKE_OLD_CXX
unset _FIREDRAKE_OLD_FC
unset _FIREDRAKE_OLD_PETSC_DIR
unset _FIREDRAKE_OLD_PETSC_ARCH
unset _FIREDRAKE_OLD_HDF5_MPI
unset _FIREDRAKE_OLD_HDF5_DIR
unset _FIREDRAKE_OLD_OMP_NUM_THREADS
unset _FIREDRAKE_OLD_PYOP2_CFLAGS
unset _FIREDRAKE_OLD_PYOP2_LDFLAGS
# Clean up empty vars
[[ -z "$CC" ]] && unset CC
[[ -z "$CXX" ]] && unset CXX
[[ -z "$FC" ]] && unset FC
[[ -z "$PETSC_DIR" ]] && unset PETSC_DIR
[[ -z "$PETSC_ARCH" ]] && unset PETSC_ARCH
[[ -z "$HDF5_MPI" ]] && unset HDF5_MPI
[[ -z "$HDF5_DIR" ]] && unset HDF5_DIR
[[ -z "$OMP_NUM_THREADS" ]] && unset OMP_NUM_THREADS
[[ -z "$PYOP2_CFLAGS" ]] && unset PYOP2_CFLAGS
[[ -z "$PYOP2_LDFLAGS" ]] && unset PYOP2_LDFLAGS
DEACTIVATE_EOF
chmod +x "${ACTIVATE_DIR}/firedrake-env-vars.sh"
chmod +x "${DEACTIVATE_DIR}/firedrake-env-vars.sh"
ok "Activation scripts installed."
# ── Step 7: Verify ────────────────────────────────────────────────────────
info "Running verification ..."
python "${SCRIPT_DIR}/verify.py" && ok "All checks passed!" || {
warn "Some checks failed. See output above."
warn "You may need to debug linking issues."
warn "Try: conda activate ${ENV_ACTIVATE} && python verify.py"
}
# ── Done ──────────────────────────────────────────────────────────────────
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
ok "Installation complete!"
echo ""
echo " To use Firedrake + icepack:"
echo ""
echo " conda activate ${ENV_ACTIVATE}"
echo " python -c 'import firedrake; import icepack; print(\"Ready!\")'"
echo ""
echo " Environment variables (PETSC_DIR, CC, etc.) are set automatically"
echo " when you activate the environment."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"