跳转至

mambaw: 在容器中使用 mamba

背景

在高性能计算环境中,一些软件对 glibc 版本有较高要求。在编译、安装或运行时,用户可能会遇到如下报错:

/lib64/libm.so.6: version `GLIBC_2.29' not found

针对这种情况,用户可以选择在 Singularity 容器中使用软件。但针对每个软件单独构建容器可能比较繁琐。

为此,平台提供了一个基于 mamba 的解决方案 mambaw。它允许用户借助一个内置高版本 glibc 的通用容器,使用 mamba 来灵活配置和管理自己的 Conda 环境。

使用方法

mambaw 命令的实际路径为 /public/scripts/mambaw。

用户可以在 ~/.bashrc 文件中将该路径所在的目录添加到 PATH 环境变量中:

export PATH=/public/scripts:$PATH

添加后,重新加载配置(source ~/.bashrc)或重新登录即可通过 mambaw 使用。该命令提供的功能与 conda 基本相同,使用时只需要把 conda 替换为 mambaw 即可。

注意

mambaw 不支持 像普通 Conda 环境那样通过 conda activate 激活环境。

用户需要使用 mambaw run、mambaw exec 或 mambaw shell 命令来运行环境中的程序。

使用用例

创建环境

例如,要创建一个包含 openmm 和 mace 的基础环境,可以使用 mambaw create 命令:

mambaw create -n openmm -c conda-forge -y \
    python=3.12.* cuda-version=12.4 cuda-toolkit cuda-libraries cuda-compiler cudnn \
    compilers cmake make cython swig doxygen numpy setuptools

安装软件包

可以使用 mambaw install 在现有环境中安装新软件包(例如 git):

mambaw install -n openmm git -y

如果需要在已创建的环境中使用 pip 安装其他软件包,可以使用 mambaw run:

mambaw run -n openmm \
    pip install mace-torch cuequivariance cuequivariance-torch cuequivariance-ops-torch-cu12 cupy-cuda12x cython \
    openmm openmmml \
    torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --extra-index-url https://download.pytorch.org/whl/cu12

运行环境中的软件

运行该环境下的其他软件,同样可以使用 mambaw run:

mambaw run -n openmm mace_run_train \
    --name="mace_model" \
    --train_file="@TRAIN_FILE" \
    --valid_fraction=0.1 \
    --E0s="average" \
    --model="MACE" \
    --num_interactions=2 \
    --num_channels=128 \
    --energy_key=energy \
    --forces_key=forces \
    --max_L=1 \
    --correlation=3 \
    --r_max=5.0 \
    --batch_size=@BATCH_SIZE \
    --max_num_epochs=@TRAIN_EPOCHS \
    --seed=@SEED \
    --device=cuda \
    --save_cpu

运行脚本

如果要运行自己编写的 Shell 脚本或 Python 脚本,可以使用如下方法:

# 运行 Shell 脚本
mambaw run -n openmm bash .my-script.sh

# 运行 Python 脚本
mambaw run -n openmm python my-script.py