create_derecho_env.sh

#!/bin/bash

set -e

# ml conda
# ENV_NAME="credit-derecho"
# CURR_DIR=`pwd`
# WHEEL_DIR="/glade/work/dgagne/credit-pytorch-envs/derecho-pytorch-mpi/wheels"
# echo $CURR_DIR
# conda create -n $ENV_NAME python=3.11
# conda init
# conda activate $ENV_NAME
# cd /glade/work/dgagne/credit-pytorch-envs/derecho-pytorch-mpi
# ./embed_nccl_vars_conda.sh
# cd $CURR_DIR
# pip install ${WHEEL_DIR}/torch-2.5.1+derecho.gcc.12.4.0.cray.mpich.8.1.29-cp311-cp311-linux_x86_64.whl
# pip install ${WHEEL_DIR}/torchvision-0.20.1+derecho.gcc.12.4.0-cp311-cp311-linux_x86_64.whl
# pip install -e .


#-----------------------------------------------------------
# set up an initial conda environment at ${CREDIT_ENV_PATH}
# containing Derecho-specific torch & MPI bits.
# (install torchmetrics at this point too, installing it later
# via pip risks an undesirable torch update.)
module load ncarenv/24.12
module load gcc
module load conda
module load mkl
module list
export ESMFMKFILE="/glade/work/dgagne/esmf-8.9.1/lib/libO/Linux.gfortran.64.mpiuni.default/esmf.mk"
topdir=$(git rev-parse --show-toplevel)
CREDIT_ENV_NAME=${CREDIT_ENV_NAME:-"credit-derecho"}
yml=$(mktemp --tmpdir=${topdir} credit-derecho-tmp-XXXXXXXXXX.yml)
echo ${yml}
#yml=derecho.yml

echo "Creating conda env \"${CREDIT_ENV_NAME}\""

cat <<EOF > ${yml}
name: credit
channels:
  - file:///glade/work/benkirk/consulting/conda-recipes/output
  - conda-forge/label/mpi-external
  - conda-forge
dependencies:
  - python=3.11
  - conda-tree
  - mpi4py =*=derecho*
  - pip
  - pytorch ==2.8.0=derecho*2000
  - torchvision =*=derecho*
  - torchmetrics
  - pip:
    - pipdeptree
EOF

# create the environment
export CONDA_VERBOSITY=1
export TIMEFORMAT=$'--> Real time: %3R seconds'
time conda env create \
     --name ${CREDIT_ENV_NAME} \
     --file ${yml} \
    || { cat ${yml}; echo "ERROR Creating Conda env!"; exit 1; }

rm -f ${yml}

#-----------------------------------------------------------
# activate & fix the environment
conda activate ${CREDIT_ENV_NAME}

echo "NCCLs - before cleanup:"
find ${CONDA_PREFIX} -name "libnccl.*"

# remove PIP NCCL, if any.
#  (echo-opt -> xgboost -> nvidia-nccl-cu12 -> problem.)
pip uninstall -y $(pip list | grep nvidia-nccl | awk '{print $1}') || true

#-----------------------------------------------------------
# install credit (editable) with constraints to prevent pip
# from overwriting the conda-installed torch/torchvision/torchmetrics.
constraint_file=$(mktemp --tmpdir=${topdir} credit-constraints-XXXXXXXXXX.txt)
pip list --format=freeze | grep -iE "^(torch|torchvision|torchmetrics)==" > "${constraint_file}"
echo "Using pip constraints:"
cat "${constraint_file}"
pip install --constraint "${constraint_file}" -e .
rm -f "${constraint_file}"

#----------------------------------------------------------------------
# Add activation file with appropriate modules to conda activation path
activate_shell="${CONDA_PREFIX}/etc/conda/activate.d/esmf_activate.sh"

cat << EOF > ${activate_shell}
module load ncarenv/24.12 
module load gcc 
module load mkl
export ESMFMKFILE="/glade/work/dgagne/esmf-8.9.1/lib/libO/Linux.gfortran.64.mpiuni.default/esmf.mk"
EOF

# Install esmpy and xesmf. Built against DJ's install of esmf since module loaded esmf wasn't working properly
pip install /glade/work/dgagne/esmf-8.9.1/src/addon/esmpy/
pip install xesmf

conda-tree deptree --small
pipdeptree --depth 3

echo "NCCLs - after cleanup:"
find ${CONDA_PREFIX} -name "libnccl.*"

python -c "import torch; print('torch version:', torch.__version__); print('CUDA available:', torch.cuda.is_available())"
python -c "import credit"

echo
echo "\"${CREDIT_ENV_NAME}\" conda environment for Derecho successfully installed into CONDA_PREFIX"
echo "use \"conda activate ${CREDIT_ENV_NAME}\" to activate"