diff --git a/conda_env_setup.sh b/conda_env_setup.sh new file mode 100755 index 000000000..35d175576 --- /dev/null +++ b/conda_env_setup.sh @@ -0,0 +1,249 @@ +#!/bin/bash + +# Helper script to create a conda environment with all of the +# dependencies needed for developing SO software. This includes +# packages needed when building some tools from source. +# +# This script is for anyone who does not want to use the "soconda" +# tools and / or does not need things like the libactpol / moby2 +# packages. This script installs as much as possible as conda +# packages, and then pip-installs the remainder. +# +# IMPORTANT NOTES: +# +# 1. You should use a conda-forge "base" environment. For +# example downloaded from: https://conda-forge.org/download/ +# Make sure "conda" is in your path before running this script. +# +# 2. If you are running on a cluster / HPC center you +# should set the MPICC environment variable to the +# name of your MPI C compiler. For example, at NERSC +# you would do: +# +# MPICC=cc ./conda_dev_env.sh -e /path/to/my/env -p 3.13 +# +# OTHER NOTES: +# +# If you would like to create your environments in an alternate +# directory, first load the base environment and add the +# alternate location to your conda config: +# +# conda config --append envs_dirs path/to/envs +# +# Then run this script with the full location: +# +# ./conda_dev_env.sh path/to/envs/my_env +# +# and then when activating the environment you can reference +# it with just "my_env". +# +# NOTE: by default, the full path to the conda env will be +# displayed in the shell prompt. To avoid this, edit ~/.condarc +# and add this line: +# +# env_prompt: '({name}) ' +# +# Now only the basename of the environment will show up. + +# Location of this script +pushd $(dirname $0) >/dev/null 2>&1 +scriptdir=$(pwd) +popd >/dev/null 2>&1 + +show_help () { + echo "" >&2 + echo "Usage: $0" >&2 + echo " [-e ]" >&2 + echo " [-p ]" >&2 + echo "" >&2 + echo " Create a conda environment for SO development." >&2 + echo "" >&2 + echo "" >&2 +} + +envname="" +pyversion=$(python3 --version 2>&1 | awk '{print $2}' | sed -e "s#\(.*\)\.\(.*\)\..*#\1.\2#") +extra="no" + +while getopts ":e:p:x" opt; do + case $opt in + e) + envname="${OPTARG}" + ;; + p) + pyversion="${OPTARG}" + ;; + \?) + show_help + exit 1 + ;; + :) + echo "Option -$OPTARG requires an argument." >&2 + show_help + exit 1 + ;; + esac +done + +shift $((OPTIND-1)) + +if [ -z "${envname}" ]; then + echo "" + echo "You must specify the environment name or path!" + show_help + exit 1 +fi + +conda_exe=$(which conda) +if [ -z "${conda_exe}" ]; then + echo "No conda executable found- is your base environment activated?" + exit 1 +fi + +# Activate the base environment +basedir=$(dirname $(dirname "${conda_exe}")) +if [[ -e ${basedir}/etc/profile.d/conda.sh ]]; then + source ${basedir}/etc/profile.d/conda.sh +else + source /etc/profile.d/conda.sh +fi +conda deactivate +conda activate base + +# Determine whether the environment is a name or a +# full path. +env_noslash=$(echo "${envname}" | sed -e 's/\///g') + +is_path=no +if [ "${env_noslash}" != "${envname}" ]; then + # This was a path + is_path=yes + env_check="" + if [ -e "${envname}/bin/conda" ]; then + # It already exists + env_check="${envname}" + fi +else + env_check=$(conda env list | { grep "${envname} " || true; }) +fi + +# Determine whether we are installing MPI with conda. This matters due +# to a bug in the ucx package (required by the MPI packages). The ucx package +# will only install correctly if listed as an initial package during the +# "create" command. + +conda_mpi="yes" +mpi_pkgs="mpich mpi4py" +if [ -n "${MPICC}" ]; then + conda_mpi="no" + mpi_pkgs="" +fi + +if [ -z "${mpi_pkgs}" ]; then + initial_pkgs="python=${pyversion}" +else + if [ "$(uname)" = "Linux" ]; then + initial_pkgs="ucx python=${pyversion}" + else + initial_pkgs="python=${pyversion}" + fi +fi + +if [ -z "${env_check}" ]; then + # Environment does not yet exist. Create it. + echo "Creating new environment \"${envname}\"" + if [ ${is_path} = "no" ]; then + conda create --yes -n "${envname}" + else + conda create --yes -p "${envname}" + fi + echo "Activating environment \"${envname}\"" + conda activate "${envname}" + echo "Setting default channel in this env to conda-forge" + conda config --env --add channels conda-forge + conda config --env --set channel_priority strict +else + echo "Activating environment \"${envname}\"" + conda activate "${envname}" +fi + +# Install conda packages. First the python we are using. +echo "Installing python-${pyversion}..." +conda install --yes --update-all ${initial_pkgs} + +# Next, install build tools +echo "Installing conda packages for development..." +conda install --yes --update-all \ + compilers \ + pybind11 \ + cmake \ + ninja \ + scikit-build-core \ + setuptools \ + setuptools-scm + +# Reload the environment to pick up compiler environment variables +conda deactivate +conda activate "${envname}" + +# The conda compiler packages make a symlink "cc", which conflicts +# with the Cray MPI compiler needed for mpi4py. Remove this symlink. +rm -f "${CONDA_PREFIX}/bin/cc" + +# Install dependencies that we have on conda-forge +echo "Installing SO packages that are on conda-forge" +# FIXME: add so3g here eventually (spt3g is already there) +conda install --yes pixell toast qpoint + +installed_pkgs="$(conda list | awk '{print $1}')" + +# Some of our dependencies are only available from pypi. Install +# pipgrip so that we can try to install as many dependencies as +# possible from conda-forge +python -m pip install pipgrip + +pushd "${scriptdir}" 2>&1 >/dev/null + +# FIXME: remove so3g (and the loop) once that package is on conda-forge +for pkg in "so3g" "."; do + deps=$(pipgrip --threads 4 --pipe ${pkg} | sed -e 's/\(==[^[:space:]]\+\)//g') + for dep in ${deps}; do + if [ ${dep} != ${pkg} ]; then + # Special handling of ruamel.yaml / ruamel-yaml. This is a longstanding + # mess across PyPI and conda ecosystems... + if [ "${dep}" = "ruamel-yaml" ] || [ "${dep}" = "ruamel-yaml-clib" ]; then + dep="ruamel.yaml" + fi + depcheck=$(echo "$installed_pkgs" | grep -E "^${dep}\$") + if [ -z "${depcheck}" ]; then + # It is not already installed, try to install it with conda + echo "Attempt to install conda package for dependency \"${dep}\"..." + conda install --yes --quiet ${dep} + # A failure of the above command is NOT AN ERROR. If the + # conda package does not exist, then the dependency will be + # installed by pip below. + if [[ $? = 0 ]]; then + installed_pkgs="${installed_pkgs}"$'\n'"${dep}" + fi + else + echo " Package for dependency \"${dep}\" already installed" + fi + fi + done + echo "Installing package ${pkg}" + python3 -m pip install ${pkg} + [[ $? != 0 ]] && exit 1 + installed_pkgs="${installed_pkgs}"$'\n'"${pkg}" +done + +popd 2>&1 >/dev/null + +# Install mpi4py from source if using an external MPI +if [ "${conda_mpi}" = "no" ]; then + echo "Building mpi4py with MPICC=\"${MPICC}\"" + pip install --force-reinstall --no-cache-dir --no-binary=mpi4py mpi4py + [[ $? != 0 ]] && exit 1 +else + conda install --yes ${mpi_pkgs} +fi +