WP-hpc-tools/lmod/ipython/startup/lmod_python_fix.py

# -*- python -*-
# modified "module" command for Wahab

#from __future__ import print_function
#from subprocess import PIPE, Popen
import os, sys 

try:
    with open("/etc/cluster", "r") as _cluster_F:
        _cluster = next(_cluster_F).strip()
except:
    sys.stderr.write("Unable to determine cluster, 'module' command will not be available.\n")
else:
    if _cluster == "wahab":
        sys.path.append('/shared/apps/common/lmod/current/init')
    elif _cluster == "turing":
        sys.path.append('/cm/shared/applications/lmod/lmod/init')
    else:
        sys.stderr.write("Unknown cluster, 'module' command will not be available.\n")
    del _cluster

try:
    from env_modules_python import module as lmod_module
except:
    pass
else:
    def module(command, *arguments):
        split_path_str = lambda P: P.split(':') if P is not None else []

        PYTHONPATH_ORIG = os.environ.get('PYTHONPATH')
        lmod_module(command, *arguments)
        PYTHONPATH = os.environ.get('PYTHONPATH')

        # update sys.path
        # Must account the *changes* in PYTHONPATH and reflect that to sys.path
        #
        # WARNING: This is not perfect!
        # Due to the algorithm used below, the updated PYTHONPATH
        # will not be reflected 100% in its order in sys.path
        #
        # 1) It assumes that additional paths are prepended.
        # 2) It will completely delete a path that is remove
        #    as a result of module command
        # 3) If there are simultaneous deletion(s) and adition(s),
        #    then the added paths will take precedence (unfortunately),
        #    which may result in unexpected behavior in corner cases,
        #    at least in principle.
        #    This should not be the case for "well behaved" HPC packages
        #    where order of loaded modules should NOT cause change in behavior.

        # Save the original in case needed for debugging
        global sys_path_orig
        sys_path_orig = sys.path

        python_path_orig = split_path_str(PYTHONPATH_ORIG)
        python_path = split_path_str(PYTHONPATH)

        sys_path_added = set(python_path) - set(python_path_orig)
        sys_path_deleted = set(python_path_orig) - set(python_path)

        # delete removed paths
        sys_path_trimmed = [ p for p in sys_path_orig if p not in sys_path_deleted ]

        # prepend new paths
        sys_path_prepended = [ p for p in python_path if p in sys_path_added ]

        sys.path = sys_path_prepended + sys_path_trimmed

        # Extra: if LMOD_MANUAL_LD_LIBRARY_PATH is specified, then
        # the shared library symlinks are updated.
        manual_ld_library_dir = os.environ.get('LMOD_MANUAL_LD_LIBRARY_PATH')
        if manual_ld_library_dir is None:
            return

        for search_path in os.getenv('LD_LIBRARY_PATH').split(':')[::-1]:
            if not os.path.isdir(search_path):
                continue
            if search_path == manual_ld_library_dir:
                continue

            for f in os.listdir(search_path):
                if '.so' in f:
                    src = f'{search_path}/{f}'
                    tgt = f'{manual_ld_library_dir}/{f}'
                    if os.path.islink(tgt):
                        os.unlink(tgt)

                    os.symlink(src, tgt)