Create module infrastructure for utilising conda environment

[aidanheerdegen]

conda environments can be activate and deactivated using modules on NCI.

An example of how this is done is the hh5 conda environments:

/g/data/hh5/public/modules/conda/

This could be copied and adapted for use in vk83.

[aidanheerdegen]

Credit: this was all written by @scottwales.

Some explanation:

All the conda/analysis modules are symlinks pointing to the same file

$ ls -l /g/data/hh5/public/modules/conda/
total 16
lrwxrwxrwx  1 saw562   hh5_w  10 Oct  5  2018 analysis27-18.10 -> .common.v2
lrwxrwxrwx  1 saw562   hh5_w  43 Jan 13  2022 analysis3-22.01 -> /g/data/hh5/public/modules/conda/.common.v2
lrwxrwxrwx  1 aph502   hh5_w  43 Apr 11  2022 analysis3-22.04 -> /g/data/hh5/public/modules/conda/.common.v2
lrwxrwxrwx  1 aph502   hh5_w  10 Jul  7  2022 analysis3-22.07 -> .common.v2
lrwxrwxrwx  1 hxw599   hh5_w  43 Dec  6  2022 analysis3-22.10 -> /g/data/hh5/public/modules/conda/.common.v2
lrwxrwxrwx  1 hxw599   hh5_w  43 Feb 28  2023 analysis3-23.01 -> /g/data/hh5/public/modules/conda/.common.v2
lrwxrwxrwx  1 hh5_apps hh5_w  43 Jun  1 19:03 analysis3-23.04 -> /g/data/hh5/public/modules/conda/.common.v2
lrwxrwxrwx  1 hh5_apps hh5_w  43 Aug 24 11:14 analysis3-23.07 -> /g/data/hh5/public/modules/conda/.common.v2
-rwxr-xr-x+ 1 saw562   hh5_w 223 Apr  3  2020 condaenv.sh
lrwxrwxrwx  1 saw562   hh5_w  10 Feb 25  2019 cws-radar -> .common.v2
lrwxrwxrwx  1 aph502   hh5_w  43 Apr 12  2022 esmvaltool-21.08 -> /g/data/hh5/public/modules/conda/.common.v2

This file /g/data/hh5/public/modules/conda/.common.v2 contains this:

#%Module1.0

if {[module-info mode remove]} {
    if {$env(CONDA_SHLVL) > 1} {
        puts stderr "ERROR: Multiple ($::env(CONDA_SHLVL)) conda environments have been loaded, cannot unload with module"
        puts stderr "ERROR: Try 'conda deactivate' first"
        break
    }
}

conflict conda

# Conda install directory
set condahome "/g/data/hh5/public/apps/miniconda3"

# Name of this module's environment
set condaenv [lindex [split [module-info name] {/}] 1]

# Get the environment variables from 'conda activate'
set condaenv [exec /bin/env -i /g/data/hh5/public/modules/conda/condaenv.sh $condaenv]

# Convert the environment into module commands
set lines [split $condaenv '\n']
foreach line $lines {
    regexp {^([^=]*)=(.*)} $line -> key value

    # Exclude $PWD and $_
    if {[lsearch -exact {MODULEPATH PWD _} $key] >= 0} {
        continue
    }

    # Is this some sort of path?
    if {[string match UDUNITS2_XML_PATH $key]} {
        # This is actually a single path
        setenv $key $value
        continue
    } elseif {[string match *?PATH $key]} {
        # A *PATH variable to be prepended with a ':'
        prepend-path $key $value
        continue
    } elseif {[lsearch {_LMFILES_ LOADEDMODULES} $key] >= 0} {
        # Modulefile stuff that works like a path
        prepend-path $key $value
        continue
    } elseif {[string match PATH $key]} {
        # PATH itself (strip out the system paths to keep ordering correct)
        prepend-path $key [regsub {:/usr/bin:/bin} $value {}]
        continue
    }

    # Otherwise set an environment var
    setenv $key $value
}

# Add user environments
if [ file isdirectory "/scratch" ] {
    set userconda /scratch/$::env(PROJECT)/$::env(USER)/conda
    setenv CONDA_ENVS_PATH $userconda/envs
    setenv CONDA_PKGS_DIRS $userconda/pkgs
} elseif [ file isdirectory "/local/$::env(PROJECT)" ] {
    set userconda /local/$::env(PROJECT)/$::env(USER)/conda
    setenv CONDA_ENVS_PATH $userconda/envs
    setenv CONDA_PKGS_DIRS $userconda/pkgs
}

setenv LC_ALL en_AU.utf8

So when the script is called it finds out the name of the module it is supposed to be in this line:

# Name of this module's environment
set condaenv [lindex [split [module-info name] {/}] 1]

The real magic is the call to condaenv.sh using the module name determined above

# Get the environment variables from 'conda activate'
set condaenv [exec /bin/env -i /g/data/hh5/public/modules/conda/condaenv.sh $condaenv]

Where condaenv.sh contains this:

#!/bin/bash
# Prints the environment variables set by 'conda activate $1', for processing by the modules
export PATH=/usr/bin:/bin
source /g/data/hh5/public/apps/miniconda3/etc/profile.d/conda.sh
conda activate $1
/bin/env

So the call activates the conda environment in a subshell, dumps the environment variables and saves them in the condaenv tcl variable.

Then the code loops through all the lines of the dumped environment and converts the environment variables into calls to the tcl module calls setenv.

It's genius.

[ScottWales]

Nowdays with containerised environments the modules just set PATH. Wrapper scripts in $PREFIX/bin for each program in the conda environment then load the container and run that function.

[aidanheerdegen]

Nowdays with containerised environments the modules just set PATH. Wrapper scripts in $PREFIX/bin for each program in the conda environment then load the container and run that function.

Something like Dale's stuff?

https://github.com/coecms/cms-conda-singularity/blob/main/modules/common_v3

Does containerisation overcome the restriction of environment modules on modifying the environment only through the module mechanisms?

I'm wondering if that makes a container a better solution for deploying spack? Can a containerised environment do a . share/spack/setup-env.sh step?

[ScottWales]

Basically yeah - see /g/data/access/ngm/modules/ants/1.1.0//g/data/access/ngm/apps/ants/1.1.0. Dale does some extra tricks there to do with his overlay stuff and to support mpi.

You can set a entry point script on the container which runs before any command where you can activate the conda environment or spack or whatever - here's my spack container's.

[aidanheerdegen]

Awesome, thanks.