[jax] [package list]

jax/0.4.13-rocm-5.6.1-python-3.10-singularity-20240207 (jax-0.4.13-rocm-5.6.1-python-3.10-singularity-20240207.eb)

Install with the EasyBuild-user module in partition/container:

module load LUMI partition/container EasyBuild-user
eb jax-0.4.13-rocm-5.6.1-python-3.10-singularity-20240207.eb

The module will be available in all versions of the LUMI stack and in the CrayEnv stack.

To access module help after installation use module spider jax/0.4.13-rocm-5.6.1-python-3.10-singularity-20240207.

EasyConfig:

# Developed by Kurt Lust for LUMI
#DOC Contains JAX 0.4.13 with NumPy 1.26.3 and SciPy 1.12.0.
easyblock = 'MakeCp'

local_NumPy_version = '1.26.3'
local_SciPy_version = '1.12.0'

name =          'jax'
version =       '0.4.13'
versionsuffix = '-rocm-5.6.1-python-3.10-singularity-20240207'

local_sif =    'lumi-jax-rocm-5.6.1-python-3.10-jax-0.4.13-dockerhash-825cf0044cce.sif'
#local_docker = 'lumi-pytorch-rocm-5.6.1-python-3.10-pytorch-v2.1.0.docker'

homepage = 'https://jax.readthedocs.io/en/latest/'

whatis = [
    'Description: JAX is Autograd and XLA, brought together for high-performance numerical computing.',
    'Keywords: JAX'
]

description = f"""
This module provides a container with JAX %(version)s with NumPy {local_NumPy_version} and SciPy {local_SciPy_version}.

The module defines a number of environment variables:
*   SIF and SIFJAX: The full path and name of the Singularity SIF file 
    to use with singularity exec etc.
*   SINGULAIRTY_BINDPATH: Mounts the necessary directories from the system,
    including /users, /project, /scratch and /flash so that you should be
    able to use your regular directories in the container.
*   RUNSCRIPTS and RUNSCRIPTSJAX: The directory with some sample
    runscripts.

Note that this container uses a Conda environment (jax) internally. When in
the container, the command to activate the container is contained in the
environment variable WITH_CONDA.
"""

docurls = [
    'JAX online docs: https://jax.readthedocs.io/en/latest/'    
]

toolchain = SYSTEM

sources = [
    {
        'filename':    local_sif,
        'extract_cmd': '/bin/cp -L %s .'
    },
]

skipsteps = ['build']

files_to_copy = [
    ([local_sif], '.'),
]

local_bin_start_shell="""
#!/bin/bash -e

# Run application
if [[ -f "/.singularity.d/Singularity" ]] 
then
    # In a singularity container, just in case a user would add this to the path.
    if [ "\$#" -eq 0 ]
    then
        \$WITH_CONDA ; exec bash
    else
        \$WITH_CONDA ; exec bash "\$@"
    fi
else
    # Not yet in the container
    if [ "\$#" -eq 0 ]
    then
        singularity exec \$SIFJAX bash -c '\\$WITH_CONDA ; exec bash'
    else
        singularity exec \$SIFJAX bash "\$@"
    fi
fi

"""

local_runscript_python_simple="""
#!/bin/bash -e

# Start conda environment inside the container
\$WITH_CONDA

# Run application
python "\$@"

"""

postinstallcmds = [
    'mkdir -p %(installdir)s/bin',
    f'cat >%(installdir)s/bin/start-shell <<EOF {local_bin_start_shell}EOF',
    'chmod a+x %(installdir)s/bin/start-shell',
    'mkdir -p %(installdir)s/runscripts',
    f'cat >%(installdir)s/runscripts/conda-python-simple <<EOF {local_runscript_python_simple}EOF',
    'chmod a+x %(installdir)s/runscripts/conda-python-simple',
]

sanity_check_paths = {
    # We deliberately don't check for local_sif as the user is allowed to remove that file
    # but may still want to regenerate the module which would then fail in the sanity check.
    #'files': [f'share/docker-defs/{local_docker}'],
    'files': [],
    'dirs':  ['runscripts'],
}

sanity_check_commands = [
    """start-shell -c '/runscripts/conda-python-simple -c "import numpy ; import scipy ; import jax ; print( f'"'JAX {jax.__version__}, NumPy {numpy.__version__}, SciPy {scipy.__version__}.'"' )"' """   
]

modextravars = {
    # SIF variables currently set by a function via modluafooter.
    #'SIF':             '%(installdir)s/' + local_sif,
    #'SIFJAX':          '%(installdir)s/' + local_sif,
    'RUNSCRIPTS':       '%(installdir)s/runscripts',
    'RUNSCRIPTSJAX':    '%(installdir)s/runscripts',
    'SINGULARITY_BIND': '/var/spool/slurmd,/opt/cray,/usr/lib64/libcxi.so.1,/usr/lib64/libjansson.so.4,' +
                        '%(installdir)s/runscripts:/runscripts,' + 
                        '/pfs,/scratch,/projappl,/project,/flash,/appl'
}

modluafooter = f"""
-- Call a routine to set the various environment variables.
create_container_vars( '{local_sif}', '%(name)s', '%(installdir)s' )
"""

moduleclass = 'devel'

[jax] [package list]