[Lmod-users] Conda/venv modules

[Shelly J. (she/her) <She...@uv...>]

Hi All,

Following a discussion during the Lmod-Users meeting on Tuesday about modules that activate Conda environments, I promised to share a function that I created to check for active environments (Conda or venv) before proceeding with loading a module that activates it’s own environment.

The function checks for active Conda or venv environments.  If an active environment is detected, an appropriate message is printed and loading of the module is aborted.

The function is defined and registered to the Lmod sandbox in the SitePackage.lua file as follows:

-- check_python_environment(): Here we check for active conda or venv environments

function check_python_environment(module_name)
    -- This check only needs to run when loading a module
    if mode() ~= "load" then return end

    -- Fetch the environment variables
    local condaLevel = tonumber(os.getenv("CONDA_SHLVL"))
    local virtualEnv = os.getenv("VIRTUAL_ENV")

    -- Check for active Conda environments
    if (condaLevel ~= nil and condaLevel > 0) then
        local condaWarn = string.format([[

    WARNING: %i active Conda environment(s) detected.

    If this virtual environment is active because of a module that you loaded,
    please unload that module with the command 'module unload <module_name>'
    and then load the '%s' module again.

    If you activated this conda environment directly, enter 'conda deactivate'
    until all environments are deactivated and then load the '%s'
    module again.
    ]], condaLevel, module_name, module_name)
        LmodMessage(condaWarn)
        os.exit(1)
    end

    -- Check for active venv environments
    if (virtualEnv ~= nil) then
        local venvWarn = string.format([[

    WARNING: The following active virtual environment has been detected:
        %s

    If this virtual environment is active because of a module that you loaded,
    please unload that module with the command 'module unload <module_name>'
    and then load the '%s' module again.

    If you activated this virtual environment directly, then please enter
    'deactivate' so that no conflicting virtual environments are loaded and
    then load the '%s' module again.
        ]], virtualEnv, module_name, module_name)
        LmodMessage(venvWarn)
        os.exit(1)
    end
end

-- Register function to the Lmod sandbox
sandbox_registration({
    check_python_environment = check_python_environment
})

The function takes an argument that provides the name of the current module or application being loaded that is incorporated into printed message when an active environment is found.  To use the exact module name, perhaps the most obvious choice, then use the first option.  If you might want to specify the app differently in the output, there are other suggested approaches.

To use this function, include the function name and argument, using one of the following 3 approaches, in your module file:

-- Option a: Using the myModuleName() function
check_python_environment(myModuleName())

-- Option b: Specify a string
check_python_environment("my_pkg_name")

-- Option c: Use a local variable for the string
local app = "gubbins"
check_python_environment(app)

The message can be customized as you wish.

Regarding how to activate and deactivate the environment using Lmod, I've tried different approaches to varying success.

For venv environments, I've modified the init scripts to export the PS1 variable, using:

# edit the module function to export the PS1 variable
sed -i.orig 's/$LMOD_CMD/PS1=$PS1 $LMOD_CMD/g' bash

So that, within my module file, I can do something like:

local oldPS1 = os.getenv("PS1")
source_sh("bash", pathJoin(installDir, "turbomole-venv/bin/activate"))
pushenv("PS1","(turbomole-venv) " .. oldPS1)

This seems to work better than using execute {cmd="source /path/to/venv/bin/activate",modeA={"load"}} with execute {cmd="deactivate",modeA={"unload"}}, which did not reliably restore paths and unset variables.


While writing this up, I realized I likely can use a similar approach for Conda modules.  I had been, until now, mostly using the "execute {cmd=" approach to (1) source conda.sh, (2) conda activate, and (3) conda deactivate.  But that required me to subsequently "unset" variables and use "remove_path" to fully return the environment to it's pre-load state.

While having only tested on a single module prior to sending this message, it seems that the following works well for activating a Conda environment and successfully returning the environment to pre-load state upon unloading:

local oldPS1 = os.getenv("PS1")
source_sh("bash", "/path/to/pkgs/miniforge/25.11.0-1/bin/activate /path/to/pkgs/dRep/3.6.2-1/dRep-3.6.2")
pushenv("PS1","(dRep-3.6.2) " .. oldPS1)

Again, this is not thoroughly tested, but I'm cautiously optimistic that it will be generally reliable for Conda environments.  Hopefully, the above description/examples make sense.  Happy to elaborate or clarify further, if needed. Has anyone found a better or more efficient solution that they recommend?

Best,
Shelly