Zenn
🥜

LocalColabFold実行環境の構築 | Apptainer (on WSL2 & with GPU)

2023/05/01に公開
WSL 2
#
container
#
bioinformatics
singularity
#
apptainer
tech

※ コードブロック における 🥜 は、プロンプトを表している

🥜 はじめに

いまさらながら LocalColabFold を自宅のマシン環境で使えるようにした。あいかわらず WSL2 上の Apptainer (旧Singularity) を用いて構築。

LocalColabFold についてはAg_smithさんの記事を参照。
Qiita | ColabFoldのローカル版"LocalColabFold"を使ったタンパク質の構造サンプリング

以下、本記事を執筆した際の環境。

specifications
CPU Ryzen 7 5700X (AMD)
グラフィックボード GeForce RTX 2080 Ti VENTUS GP OC (MSI)
メモリ DDR4-3200 16 GB×2 (Crucial)
version
Windows 10 22H2
NVIDIA グラフィクス ドライバー 531.41
WSL2 1.2.5.0
AlmaLinux 8 on WLS2 8.7
Apptainer 1.1.7-1.el8

なお、WSL2 上の AlmaLinux 8 に Apptainer 環境を構築する手順については、 以下の記事をご参照ください。
Zenn | AlmaLinux 8 on WSL2にApptainer環境 (with GPU) を構築する

🥜 LocalColabFold実行コンテナの構築

LocalColabFold 公式の install_colabbatch_linux.sh を参考に、コンテナビルド用の LocalColabFold.def を任意のディレクトリ (例として ~/apps/apptainer_image/def_file ) に作成する。
GitHub | LocalColabFold

~/apps/apptainer_image/def_file/LocalColabFold-v1.5.2.def
BootStrap: docker
From: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

%post
  # localtime関係のWARNINGに対処
  touch /etc/localtime

  # LocalColabFoldに必要なパッケージのインストール (冗長)
  apt update && apt upgrade -y
  apt install -y build-essential gcc-9 g++-9 git curl wget
  
  # localeにen_US.UTF-8を追加する (WSL2だけで使用するなら不要)
  apt install -y locales
  locale-gen en_US.UTF-8
  
  # apt関係のお掃除
  rm -rf /var/lib/apt/lists/* && apt autoremove -y && apt clean

  # gcc-9およびg++-9をデフォルトのgccおよびg++に設定
  update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 99
  gcc --version
  update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-9 99
  g++ --version

  # Pyenvを/usr/local/appsにインストール
  git clone https://github.com/yyuu/pyenv.git /usr/local/apps/pyenv
  export PYENV_ROOT="/usr/local/apps/pyenv"
  export PATH="${PYENV_ROOT}/bin:${PATH}"

  # PyenvでMiniforgeをインストールし、MiniforgeにPATHを通す
  pyenv install --list
  pyenv install miniforge3-22.9.0-3
  pyenv global miniforge3-22.9.0-3
  pyenv versions
  export MINIFORGE3_ROOT="${PYENV_ROOT}/versions/miniforge3-22.9.0-3"
  export PATH="${MINIFORGE3_ROOT}/bin:${PATH}"

  # conda自身のアップデート
  conda update -n base conda
  # conda環境colabfold-condaを作成
  conda create -n colabfold-conda

  # LocalColabFoldに必要なライブラリのインストール
  conda install -n colabfold-conda -c conda-forge \
    python=3.10 cudatoolkit==11.8.0 openmm==7.7.0 pdbfixer
  conda install -n colabfold-conda -c conda-forge -c bioconda \
    kalign2=2.04 hhsuite=3.3.0 mmseqs2=14.7e284
  
  # conda関係のお掃除
  conda clean --all --force-pkgs-dirs --yes
  # conda環境colabfold-condaをアクティベート
  export COLABFOLD_HOME="${MINIFORGE3_ROOT}/envs/colabfold-conda"
  export PATH="${COLABFOLD_HOME}/bin:${PATH}"
  
  # pip自身のアップデート
  python3 -m pip install --no-cache-dir --upgrade pip

  # ColabFoldのインストール
  python3 -m pip install --no-cache-dir --no-warn-conflicts \
    "colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold" tensorflow==2.12.0
  python3 -m pip install --no-cache-dir \
    https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.25+cuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl
  python3 -m pip install --no-cache-dir \
    jax==0.3.25 chex==0.1.6 biopython==1.79
  
  # conda環境colabfold-condaにインストールしたパッケージの由来 (channel) を確認する
  conda list -n colabfold-conda

  # Use 'Agg' for non-GUI backend
  cd ${COLABFOLD_HOME}/lib/python3.10/site-packages/colabfold
  sed -i -e \
    "s#from matplotlib import pyplot as plt#import matplotlib\nmatplotlib.use('Agg')\nimport matplotlib.pyplot as plt#g" ./plot.py
  # modify the default params directory
  sed -i -e \
    "s#appdirs.user_cache_dir(__package__ or \"colabfold\")#\"${MINIFORGE3_ROOT}/envs/colabfold\"#g" ./download.py
  # remove cache directory
  rm -rf ./__pycache__

  # model weightsのダウンロード
  cd ${MINIFORGE3_ROOT}/envs
  python3 -m colabfold.download

%environment
  # LocalColabFold実行のための環境設定
  export MINIFORGE3_ROOT="/usr/local/apps/pyenv/versions/miniforge3-22.9.0-3"
  export COLABFOLD_HOME="${MINIFORGE3_ROOT}/envs/colabfold-conda"
  export PATH="${COLABFOLD_HOME}/bin:${PATH}"
  
  # for WSL2 (?)
  export TF_FORCE_UNIFIED_MEMORY="1"
  export XLA_PYTHON_CLIENT_MEM_FRACTION="4.0"
  export XLA_PYTHON_CLIENT_ALLOCATOR="platform"
  export TF_FORCE_GPU_ALLOW_GROWTH="true"

%runscript
  # コンテナ実行時のスクリプトを定義
  # "$@"は追加の引数を利用するための設定
  colabfold_batch "$@"

コンテナビルドの大まかな流れは以下

  • NVIDIAの CUDA and cuDNN imageをベースにコンテナを構築
  • apt を用いて必要そうなパッケージをインストール (かなり冗長)
  • Pyenv をインストールし、Pyenv 経由で Miniforge をインストール
  • conda 環境 colabfold-conda を作成
  • conda を用いて必要なライブラリを colabfold-conda にインストール
  • pip を用いて必要なライブラリと LocalColabFold を colabfold-conda にインストール
  • モデルパラメータのダウンロード

# for WSL2 (?) に記載の環境変数 ( TF_FORCE_UNIFIED_MEMORY など) は、WSL2 を用いて環境構築する際にのみ必要かもしれない (公式にもそのような記載がある)。

また、公式の install_colabbatch_linux.sh と異なり、conda install の際に cudnn をインストールしていない (ベースのコンテナイメージに含まれているため)。
ちなみに cudatoolkit は openmm の依存ライブラリなので、ベースイメージに含まれる CUDA と重複するがやむを得ず conda install に含めている。

おそらく LocalColabFold 実行時に使用される CUDA 環境は、colabfold-conda ではなくベースイメージに含まれる CUDA 環境である。colabfold-conda 内の CUDA 環境を使うためには、LD_LIBRARY_PATH に ${COLABFOLD_HOME}/lib を通してやる必要がある。

ただし、NVIDIA提供のコンテナイメージを用いずに、素の Ubuntu イメージに conda install で cudatoolkit および cudnn をインストールしても LocalColabFold は動かないようである ( TF_FORCE_UNIFIED_MEMORY なんかの設定が効いていない印象)。

作成した LocalColabFold-v1.5.2.def を基に、LocalColabFold 実行コンテナ LocalColabFold-v1.5.2.def を以下のコマンドでビルドする。

# ~/apps/apptainer_imageに移動
🥜 cd ~/apps/apptainer_image
# LocalColabFold-v1.5.2.defに基づいてコンテナイメージをビルド
🥜 apptainer build ./LocalColabFold-v1.5.2.sif ./def_file/LocalColabFold-v1.5.2.def

🥜 LocalColabFoldの動作テスト

LocalColabFold-v1.5.2.sif を用いてCanis familiaris allergen 1 (Can f 1) の構造予測を実施してみる。

Can f 1のアミノ酸配列 ( ~/hoge/canf1_monomer.fasta )

~/hoge/canf1_monomer.fasta
>canf1
MKTLLLTIGFSLIAILQAQDTPALGKDTVAVSGKWYLKAMTADQEVPEKPDSVTPMILKAQKGGNLEAKI
TMLTNGQCQNITVVLHKTSEPGKYTAYEGQRVVFIQPSPVRDHYILYCEGELHGRQIRMAKLLGRDPEQS
QEALEDFREFSRAKGLNQEILELAQSETCSPGGQ

LocalColabFold実行スクリプト ( ~/hoge/LCF_run.sh )

~/hoge/LCF_run.sh
#!/usr/bin/bash

APPTAINER_IMAGE="~/apps/apptainer_image"

apptainer run --nvccli ${APPTAINER_IMAGE}/LocalColabFold-v1.5.2.sif \
  --amber \
  --use-gpu-relax \
  --templates \
  --num-recycle 3 \
  canf1_monomer.fasta canf1_monomer

以下のように計算を実行する。

# 目的のディレクトリに移動
🥜 cd ~/hoge
# LCF_run.shに実行権限付与
🥜 chmod u+x ./LCF_run.sh
# LCF_run.shを実行
🥜 ./LCF_run.sh

以下のように計算が実行される。

INFO:    Setting 'NVIDIA_VISIBLE_DEVICES=all' to emulate legacy GPU binding.
INFO:    Setting --writable-tmpfs (required by nvidia-container-cli)
2023-05-01 16:33:23,789 Running colabfold 1.5.2 (39a7b60851a7056fe0314cdb0e36741a42f9502f)
2023-05-01 16:33:26.990691: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2023-05-01 16:33:28,803 Running on GPU
2023-05-01 16:33:29,591 Found 7 citations for tools or databases
2023-05-01 16:33:29,592 Query 1/1: canf1 (length 174)
COMPLETE: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 150/150 [elapsed: 00:00 remaining: 00:00]
2023-05-01 16:33:55,106 Sequence 0 found templates: ['5t43_A', '6nre_A', '3zq3_C', '5x7y_B', '2l9c_A', '1gm6_A', '5ngh_A', '1ew3_A', '2rq0_A', '3kff_A', '3kfh_A', '4run_B', '2a2u_A', '4ory_F', '4os3_A', '7dru_C', '4orx_A', '2xst_A', '5wy9_A', '7dru_B']
2023-05-01 16:33:55,153 Setting max_seq=512, max_extra_seq=627
2023-05-01 16:34:44,806 alphafold2_ptm_model_1_seed_000 recycle=0 pLDDT=84.5 pTM=0.764
2023-05-01 16:34:56,388 alphafold2_ptm_model_1_seed_000 recycle=1 pLDDT=85.8 pTM=0.771 tol=4.81
2023-05-01 16:35:07,939 alphafold2_ptm_model_1_seed_000 recycle=2 pLDDT=86.2 pTM=0.774 tol=1.36
2023-05-01 16:35:19,719 alphafold2_ptm_model_1_seed_000 recycle=3 pLDDT=86.6 pTM=0.777 tol=1.84
2023-05-01 16:35:19,720 alphafold2_ptm_model_1_seed_000 took 73.0s (3 recycles)
2023-05-01 16:35:31,374 alphafold2_ptm_model_2_seed_000 recycle=0 pLDDT=85.2 pTM=0.778
2023-05-01 16:35:43,086 alphafold2_ptm_model_2_seed_000 recycle=1 pLDDT=85.5 pTM=0.776 tol=0.789
2023-05-01 16:35:54,895 alphafold2_ptm_model_2_seed_000 recycle=2 pLDDT=86.2 pTM=0.779 tol=0.574
2023-05-01 16:36:06,642 alphafold2_ptm_model_2_seed_000 recycle=3 pLDDT=86.4 pTM=0.778 tol=0.469
2023-05-01 16:36:06,643 alphafold2_ptm_model_2_seed_000 took 46.9s (3 recycles)
2023-05-01 16:36:26,819 alphafold2_ptm_model_3_seed_000 recycle=0 pLDDT=85.4 pTM=0.779
2023-05-01 16:36:38,027 alphafold2_ptm_model_3_seed_000 recycle=1 pLDDT=85.6 pTM=0.775 tol=0.818
2023-05-01 16:36:49,243 alphafold2_ptm_model_3_seed_000 recycle=2 pLDDT=86.2 pTM=0.78 tol=0.573
2023-05-01 16:37:00,487 alphafold2_ptm_model_3_seed_000 recycle=3 pLDDT=86.1 pTM=0.78 tol=0.398
2023-05-01 16:37:00,488 alphafold2_ptm_model_3_seed_000 took 53.8s (3 recycles)
2023-05-01 16:37:11,744 alphafold2_ptm_model_4_seed_000 recycle=0 pLDDT=84.9 pTM=0.771
2023-05-01 16:37:23,048 alphafold2_ptm_model_4_seed_000 recycle=1 pLDDT=85.8 pTM=0.774 tol=0.914
2023-05-01 16:37:34,363 alphafold2_ptm_model_4_seed_000 recycle=2 pLDDT=86 pTM=0.776 tol=0.699
2023-05-01 16:37:45,728 alphafold2_ptm_model_4_seed_000 recycle=3 pLDDT=86.1 pTM=0.778 tol=0.65
2023-05-01 16:37:45,729 alphafold2_ptm_model_4_seed_000 took 45.2s (3 recycles)
2023-05-01 16:37:57,046 alphafold2_ptm_model_5_seed_000 recycle=0 pLDDT=84.9 pTM=0.774
2023-05-01 16:38:08,297 alphafold2_ptm_model_5_seed_000 recycle=1 pLDDT=85.2 pTM=0.774 tol=3.5
2023-05-01 16:38:19,608 alphafold2_ptm_model_5_seed_000 recycle=2 pLDDT=85.9 pTM=0.779 tol=0.814
2023-05-01 16:38:30,865 alphafold2_ptm_model_5_seed_000 recycle=3 pLDDT=85.7 pTM=0.778 tol=1.36
2023-05-01 16:38:30,866 alphafold2_ptm_model_5_seed_000 took 45.1s (3 recycles)
2023-05-01 16:38:30,900 reranking models by 'plddt' metric
2023-05-01 16:38:31,836 Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.
2023-05-01 16:38:35,839 Relaxation took 4.9s
2023-05-01 16:38:35,839 rank_001_alphafold2_ptm_model_1_seed_000 pLDDT=86.6 pTM=0.777
2023-05-01 16:38:38,444 Relaxation took 2.6s
2023-05-01 16:38:38,445 rank_002_alphafold2_ptm_model_2_seed_000 pLDDT=86.4 pTM=0.778
2023-05-01 16:38:41,049 Relaxation took 2.6s
2023-05-01 16:38:41,049 rank_003_alphafold2_ptm_model_4_seed_000 pLDDT=86.1 pTM=0.778
2023-05-01 16:38:43,946 Relaxation took 2.9s
2023-05-01 16:38:43,947 rank_004_alphafold2_ptm_model_3_seed_000 pLDDT=86.1 pTM=0.78
2023-05-01 16:38:46,555 Relaxation took 2.6s
2023-05-01 16:38:46,556 rank_005_alphafold2_ptm_model_5_seed_000 pLDDT=85.7 pTM=0.778
2023-05-01 16:38:47,693 Done

Discussion