跳转至

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 runmambaw execmambaw shell 命令来运行环境中的程序。

使用用例

创建环境

例如,要创建一个包含 openmmmace 的基础环境,可以使用 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