AlphaFoldをスパコン（名古屋大学の「不老」）で動かしてみた

• メモ
  • AlphaFold v2.1.1向けに色々修正・加筆しました。GPUを使ってTensorFlowを実行したときに出てくる「いつものログ」を確認したので、おそらく大丈夫なはず。
  • 2021年11月12日、「不老」では公式にインストール済AlphaFoldが提供されるようになりました。わざわざ個別にインストールとかしなくても簡単に使えます。（コンテナではなく直接利用できるような形式でインストールされています。）
  • 自分で環境構築してみたい方や、同様にSingularityな環境で動かしたい方は、この記事を参考にしてやってください。
  • 注意：この記事を書いている人はスパコンの専門家であり分子生物学の専門家ではなく、AlphaFoldも今回初めて触ったので、AlphaFoldの細かい使い方や用語については正確ではない可能性があります。（脳内補完よろしくお願いします、ということで……。）

本編

AlphaFoldをスパコンで動かしたいという話があったので試してみた。無事に動作したようだが、躓いた点もあったので記録しておく。
（そもそもコンテナを使わずにネイティブ環境（？）で動かせば良いじゃんという気もしなくもないが、コンテナを使いたいという需要はあるだろうし、「不老」以外でもDockerではなくSingularityを採用している環境を使う人にとって有益な情報になる可能性もあると思うので。）

AlphaFold本家と@Ag_smithによる資料を元に、名古屋大学情報基盤センターの「不老」Type IIサブシステム（1ノードあたりV100を4基搭載している、いわゆるGPUスパコン）で動作確認した。

リポジトリのチェックアウト（クローン）

何はともあれ、AlphaFoldを本家リポジトリから拾ってくるところから。ブランチ（タグ）を明示するかどうかは使い方にあわせて。今回はv2.1.1を使うことにした。（v2.0からv2.1へのアップデートでデータベースや実行時のオプションの互換性が少し欠けてしまったそうなので注意。）

# masterブランチ（alphafoldディレクトリが作られる）
$ git clone https://github.com/deepmind/alphafold.git
# v2.1.1を明示する例（本記事はv2.1.1で動作を確認しました）
$ git clone https://github.com/deepmind/alphafold.git -v v2.1.1

GPU利用のための準備

本家リポジトリで公開されているものはGPUの利用が考慮されていない、具体的には内部で利用しているTensorFlowがGPU向けになっていないため、コンテナを作成する前に一部のファイルを書き換えておく。
書き換えポイントは以下の2つ。いずれも、AlphaFoldに限らずTensorFlowでGPUを利用する際に注意が必要なポイント。これらに加えてさらにコンテナ作成時や利用時に気を付けるべきポイントもあるのだが、それは「コンテナの準備」の中で。

書き換えポイント1

requirements.txtに書かれているtensorflow-cpu==2.5.0をtensorflow-gpu==2.5.0に書き換える。

手っ取り早く直すならこんな感じ。（sedコマンドで文字列の置換をしているだけ。）

$ sed -i -e "s/cpu/gpu/" requirements.txt

書き換えポイント2

alphafold/relax/amber_minimize.pyに書かれているopenmm.Platform.getPlatformByName("CPU")の部分のCPUをCUDAに書き換える。同じ記述が2箇所ある。

手っ取り早く直すならこんな感じ。（同上）

$ sed -i -e "s/CPU/CUDA/" alphafold/relax/amber_minimize.py

コンテナの構築

本家ではDockerを用いた環境構築方法が案内されているが、「不老」など共用のスパコンではDockerの代わりにSingularityを提供していることが多いはず。Dockerコンテナ自体を持ち込んで変換するという手段もあるが、今回はDockerfileからSingularity Recipe（defファイル）を生成してbuildしてみた。
実は変換自体はSingularity Python (spython)を使うと簡単に行えるのだが、そのままでは正しく動作しなかったため一部修正した。

「不老」にはデフォルトでspythonは入っていないが、condaなどで簡単に導入可能。以下、変換例。

$ spython recipe docker/Dockerfile singularity.def

以下はspythonにより生成されたファイルを修正した後のファイルの例。修正ポイントについては日本語でコメントを付加してあるので参考にして欲しい。
（/tmp/hh-suiteまわりで躓いたので/tmp/hh-suite_buildディレクトリを作る構成にしていたが、おそらく同じノードで複数のビルドを走らせてしまったために生じた問題だと思われるため、元に戻した。）

singularity.def

Bootstrap: docker
From: nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04
Stage: spython-base

%files
. /app/alphafold
%post
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

CUDA=11.0
# FROM directive resets ARGS, so we specify again (the value is retained if
# previously set).

# Use bash to support string substitution.
# Dockerfileではシェルを明示できるが、spythonではできないようなので単純にコメントアウト
# SHELL ["/bin/bash", "-c"]

# 11-0の部分は${CUDA/./-}というシェルによる文字列置き換えがなされていたが、エラーしていたためべた書きに変更
apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
build-essential \
cmake \
cuda-command-line-tools-11-0 \
git \
hmmer \
kalign \
tzdata \
wget \
&& rm -rf /var/lib/apt/lists/*

# Compile HHsuite from source.
# ディレクトリの移動はpushd/popdで行っていたが、シェルの都合か効かなかったため同等の効果が得られるように修正
git clone --branch v3.3.0 https://github.com/soedinglab/hh-suite.git /tmp/hh-suite \
&& mkdir /tmp/hh-suite/build \
&& cd /tmp/hh-suite/build \
&& cmake -DCMAKE_INSTALL_PREFIX=/opt/hhsuite .. \
&& make -j 4 && make install \
&& ln -s /opt/hhsuite/bin/* /usr/bin \
&& cd ${OLDPWD} \
&& rm -rf /tmp/hh-suite

# Install Miniconda package manager.
wget -q -P /tmp \
https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
&& bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \
&& rm /tmp/Miniconda3-latest-Linux-x86_64.sh

# Install conda packages.
# CUDA_VERSIONの参照は大丈夫そう（もし失敗するなら明記も検討）
PATH="/opt/conda/bin:$PATH"
echo "CUDA_VERSION=${CUDA_VERSION}"
conda update -qy conda \
&& conda install -y -c conda-forge \
openmm=7.5.1 \
cudatoolkit==${CUDA_VERSION} \
pdbfixer \
pip \
python=3.7

wget -q -P /app/alphafold/alphafold/common/ \
https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt

# Install pip packages.
# ここでも${CUDA/./}が効かなかったため手動で110に変更
pip3 install --upgrade pip \
&& pip3 install -r /app/alphafold/requirements.txt \
&& pip3 install --upgrade jax jaxlib==0.1.69+cuda110 -f \
https://storage.googleapis.com/jax-releases/jax_releases.html

# Apply OpenMM patch.
cd /opt/conda/lib/python3.7/site-packages
patch -p0 < /app/alphafold/docker/openmm.patch

# We need to run `ldconfig` first to ensure GPUs are visible, due to some quirk
# with Debian. See https://github.com/NVIDIA/nvidia-docker/issues/1399 for
# details.
# ENTRYPOINT does not support easily running multiple commands, so instead we
# write a shell script to wrap them up.
cd /app/alphafold
# 「'」の直前の余計な$を消した
echo '#!/bin/bash\n\
ldconfig\n\
python /app/alphafold/run_alphafold.py "$@"' > /app/run_alphafold.sh \
&& chmod +x /app/run_alphafold.sh
%environment
export PATH="/opt/conda/bin:$PATH"
# コンテナ（の中のTensorFlow）実行時にlibcusolverが見つけられずGPUが使えなくなるのを防ぐためにライブラリパスを追加しておく
export LD_LIBRARY_PATH="/home/center/opt/x86_64/cores/hpc_sdk/Linux_x86_64/21.2/math_libs/11.2/lib64:$LD_LIBRARY_PATH"
%runscript
cd /app/alphafold
exec /app/run_alphafold.sh "$@"
%startscript
cd /app/alphafold
exec /app/run_alphafold.sh "$@"

あとはsingularity buildするだけ。「不老」では-fオプション（fakeroot）を使えるのは計算ノードだけであるため、バッチジョブで実行。それなりに時間がかかる。
ちなみにインタラクティブジョブで実行しても良いが、消費ポイント的にcx-shareを使う方が安い。さらに言えば一度構築したことのあるコンテナを再度構築する際はキャッシュが効くため、一度ログインノードでnvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04をビルドしてキャッシュを得ておくと時間短縮になる、はず。（ちゃんと比較したわけではない。）

build_singularity.sh

#!/bin/bash -x
#PJM -L rscgrp=cx-share
#PJM -L jobenv=singularity
#PJM -j
#PJM -S
module load singularity
singularity build -f alphafold.sif ./singularity.def

なおレシピファイルには以下の記述があり、このコンテナ作成時にはリポジトリからチェックアウトしてきたディレクトリの中身全てがコンテナの中に押し込まれることになる。
（DockerfileにあるCOPY . /app/alphafoldが変換されてこうなっている。）
singularity shellでコンテナに入り、/app/alphafoldディレクトリの中を覗いてみると確認することができる。

%files
. /app/alphafold

そのため、alphafoldの挙動を調整するために改造などしたい場合は、あらかじめチェックアウトしてきたディレクトリ内のファイルの更新をしておいてからsingularity buildによりコンテナ化をせねばならない。

また、コンテナをテスト用に作り直したりしていると（作成したコンテナファイルが残っていたりすると）、作成するコンテナの容量が肥大化してしまうため、作業用のディレクトリは分けた方が良さそうである。例えば、チェックアウトしたalphafoldと同じ階層にalphafold_workディレクトリを作成しておき、alphafold_workディレクトリ内で以下を実行するなどするとスムーズである。

build_singularity2.sh

#!/bin/bash -x
#PJM -L rscgrp=cx-share
#PJM -L jobenv=singularity
#PJM -j
#PJM -S
module load singularity
cd ../alphafold
singularity build -f ../alphafold_work/alphafold.sif ./singularity.def

必要なデータの準備（データベースの準備）

ダウンロードに使われているaria(aria2c)は自分で用意しておくこと。それ以外には特別な配慮は不要。参考資料を真似する。ただし容量が大きい点には注意。
（フルでDLすると3TBあり「不老」の無料ストレージ領域が溢れる。不便なので、運用サイドで誰でも参照できるデータベースディレクトリを用意している。）

コンテナの実行（v2.1向けのオプション指定に更新した）

本家ではrun_docker.pyを使っているが、もちろんsingularity環境ではこれをそのまま使うことはできないため、自分で適切に実行する必要がある。

run_docker.pyには様々なオプション（引数）があるようだが、とりあえずこんな感じにすれば動くよ、という例を以下に示す。

主なポイントはコンテナ内外のディレクトリの対応付け（マウント）を行う--bindオプション。
必要そうなものを真似したが、各自の実行方法にてさらに必要なものがあれば真似して追加してほしい。

環境変数NVIDIA_VISIBLE_DEVICESにallを設定してノード内の4GPUを使ってくれるようにしている。

TensorFlowがGPU利用時にlibcusolver.so.11を参照するのにコンテナ内にはlibcusolver.so.10しか入っていないという問題があるため、/home/center/optをbindしてLD_LIBRARY_PATHで指定したパスから適切なライブラリを参照できるようにしている。
（本当はレシピファイルの時点で適切なcuda関係のライブラリをインストールすると良いのかもしれないが、とりあえずこのLD_LIBRARY_PATHで解消できるので、またあとで考える。CUDA 11.0と一緒に入ってくるcusolverが10.xなのに、tensorflow-gpu==2.5.0がcusolver 11.0を要求している、というのが問題のような気はする。）
コンテナをビルドする際にLD_LIBRARY_PATHを対策してあるため、このタイミングで必要なのはそのディレクトリがコンテナ内から見えるようにすること（ --bind /home/center/opt:/home/center/opt を書くこと）だけである。（当初はLD_LIBRARY_PATHを--envで渡すなどの対応が必要かと思ったが、どうやら不要のようだ。）

fasta_paths など計算対象を指定するオプションは各自で適宜調整して欲しい。

ちなみに --bind /data/group1/$USER:/data/group1/$USER はホームディレクトリ以下ではなく /data/group1/$USER 以下でコンテナを実行する際に必要なオプション。（Singularityの仕様。）

run_alphafold.sh

#!/bin/bash -x
#PJM -L rscgrp=cx-single
#PJM -L elapse=24:00:00
#PJM -L jobenv=singularity
#PJM -j
#PJM -S

export DB_DIR=データベースが格納されているディレクトリのパス
export OUT_DIR=結果が格納されるディレクトリ。あらかじめ作っておく必要がある。（もちろん、このスクリプト内で作っても良い。）

date
module load singularity

singularity run --nv \
--bind /data/group1/$USER:/data/group1/$USER \
--bind ${DB_DIR}/uniref90:/mnt/uniref90_database_path \
--bind ${DB_DIR}/mgnify:/mnt/mgnify_database_path \
--bind ${DB_DIR}/pdb70:/mnt/pdb70_database_path \
--bind ${DB_DIR}/pdb_mmcif:/mnt/template_mmcif_dir \
--bind ${DB_DIR}/pdb_mmcif:/mnt/obsolete_pdbs_path \
--bind ${DB_DIR}/uniclust30/uniclust30_2018_08:/mnt/uniclust30_database_path \
--bind ${DB_DIR}/bfd:/mnt/bfd_database_path \
--bind ${OUT_DIR}:/mnt/output \
--bind /home/center/opt:/home/center/opt \
--env NVIDIA_VISIBLE_DEVICES=all \
./alphafold.sif \
--uniref90_database_path=/mnt/uniref90_database_path/uniref90.fasta \
--mgnify_database_path=/mnt/mgnify_database_path/mgy_clusters_2018_12.fa \
--pdb70_database_path=/mnt/pdb70_database_path/pdb70 \
--template_mmcif_dir=/mnt/template_mmcif_dir/mmcif_files \
--obsolete_pdbs_path=/mnt/obsolete_pdbs_path/obsolete.dat \
--uniclust30_database_path=/mnt/uniclust30_database_path/uniclust30_2018_08 \
--bfd_database_path=/mnt/bfd_database_path/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
--fasta_paths=`pwd`/Q9KVZ7.fasta --model_preset=monomer \
--data_dir=${DB_DIR} --output_dir=${OUT_DIR} \
--max_template_date=2021-07-14 --db_preset=full_dbs --benchmark=False --logtostderr

date

実行してみて気がついたこと

Q9KVZ7という問題を実行してみたが、実行時間の殆どがGPUを使っていないHHblitsという処理だった。ここに11時間、GPUを使えるTensorFlowに10分というレベル。Type IIサブシステムに搭載されているSSDを使ってみたところ、11時間が10分程度に短縮できるというトンデモな結果に。しかし3TBのデータコピーには数時間かかる。
ユーザが誰でも共通して使うデータなので、NVMESHで誰でもすぐに参照できるようにするのが良さそうである。（対策中。→NVMESHにデータが配置され、誰でも参照できる状態になった。継続的なデータの更新という問題はあるが……。）