AlphaFold を Google Cloud Life Sciences で実行しよう
AlphaFold
AlphaFold は、DeepMind 社が開発したタンパク質の 3D 構造を予測するための AI モデルです。2018年に最初のAlphaFold (version 1)が公開され、その後 2020年に version 2 が発表されました。これは、生物学の分野では 50 年来の課題を解決したとも言われるほどの革新的な発明として世界中から関心を集めています。また、DeepMind 社は、2021年にこの AlphaFold を GitHub 上に公開しました。これにより、誰もが AlphaFold を手元で動かせるようになり、近い将来には未知のタンパク質の機能を解明し、新薬の開発を加速し、多くの人の病気を治療できるようになることが期待されています。
AlphaFold が GitHub に公開された直後から、研究機関や多くの個人がクラウドやオンプレミスのサーバ上で AlphaFold を実行するための手順をブログ等で公開しています。ただ、 AlphaFold を動かすには、2.2 TB以上にもなるデータベースファイルをインターネットからダウンロードして、ディスク上に格納しておく必要があります。また、CPUやメモリ、もしくは GPU もそれなりに必要になるため、AlphaFold を使うためにクラウドのサーバを構築・設定し、長時間稼働し続けておくのは気が引けます。
この記事では、Google Cloud 上で仮想サーバ を立てずに、高速かつ安価に AlphaFold を実行するための手順を説明します。
準備作業(初回のみ)
仮想サーバを立てずにと言っておきながら、初回の準備だけは VM インスタンスを作成します。これは、最初にディスクイメージを作成するために一時的に VM インスタンスを使う必要があるためです。
作業は、以下のような流れになります。
- AlphaFold コンテナイメージを作成しコンテナレジストリに登録
- データベースファイルをダウンロードし、ディスクをイメージに変換
- Cloud Life Sciences API の有効化
- dsub をインストール
- コンテナ内で実行するスクリプトを作成
VM インスタンスが必要になるのは上記の手順の 2 までです。それ以降の手順は、Cloud Shell 上での操作を想定しています。
それでは、手順を順番に説明していきます。
1. AlphFold コンテナイメージを作成しコンテナレジストリに登録
Compute Engine VM インスタンスを作成します。この VM には 2.2 TB以上のディスクをアタッチして、VM インスタンス上でマウントします。コマンドラインで us-central1-a ゾーンに VM インスタンスを作成する場合のコマンド例は、以下のようになります。boot-diskのサイズを 100 GB、create-disk で 3000 GB のディスクを作成している点に注意してください。
gcloud compute instances create <INSTANCE_NAME> \
--zone=us-central1-a \
--machine-type=e2-standard-8 \
--boot-disk-size=100GB \
--create-disk=mode=rw,size=3000,type=projects/<PROJECT_ID>/zones/us-central1-a/diskTypes/pd-balanced,name=alphafold-data,device-name=alphafold-data
gcloud コマンドで実行する以外に、Google Cloud 管理コンソール上でも同じことができます。
Google Cloud では ”イメージはグローバルリソース” なので、任意のリージョンで一度イメージを作成すれば、世界中のどのリージョン・ゾーンからでもそのイメージを利用することができます。またその際、リージョン間のデータ転送費用もかかりません。そのため、今回 VM インスタンスを作成するリージョンはどこでも構いませんが、比較的単価の安い米国のリージョンを推奨します。
次に、作成した VM インスタンスに SSH ログインして、GitHub から AlphaFold をクローンし、コンテナイメージをビルドしてコンテナレジストリに push します。具体的には以下のコマンドを実行します。なお、dockerのインストール手順は省略しています。
sudo apt-get update
sudo apt-get install git
git clone https://github.com/deepmind/alphafold.git
cd alphafold
docker build -f docker/Dockerfile -t alphafold .
docker tag alphafold gcr.io/<PROJECT_ID>/alphafold
docker push gcr.io/<PROJECT_ID>/alphafold
2. データベースファイルをダウンロードし、ディスクをイメージに変換
アタッチしたディスクをマウントしたパスを指定して、データベースファイルをダウンロードするためのスクリプトを実行します。ここで、MOUNT_POINTは、追加ディスクをマウントしたパスを指します。非ブートディスクをフォーマットしてマウントする方法も参照してください。
sudo apt-get install rsync
sudo apt-get install aria2
cd scripts
./download_all_data.sh <MOUNT_POINT>
このダウンロードが完了するには、 2〜3 時間かかります。なお、日本リージョンでこのダウンロードを実行する場合はさらに長時間かかるため、ミラーサイト等からダウンロードできるようにダウンロードスクリプトを編集することをおすすめします(詳細は割愛)。
このダウンロードの実行が完了すると、アタッチしたディスク内は以下のような構成になります。
/ # Total: ~ 2.2 TB (download:438GB)
bfd/ #~1.7TB (download:271.6GB)
# 6 files.
mgnify/ # ~64GB (download:32.9GB)
mgy_clusters_2018_12.fa
params/ # ~3.5GB (download:3.5GB)
# 5 CASP14 models,
# 5 pTM models,
# LICENSE,
# = 11 files.
pdb70/ # ~56GB (download:19.5GB)
# 9 files.
pdb_mmcif/ # ~206GB (download:46GB)
mmcif_files/
# About 180,000 .cif files.
obsolete.dat
small_bfd/ # ~17GB (download:9.6GB)
bfd-first_non_consensus_sequences.fasta
uniclust30/ # ~86GB (download:24.9GB)
uniclust30_2018_08/
# 13 files.
uniref90/ # ~58GB (download:29.7 GB)
uniref90.fasta
ダウンロードが終わったら、ディスクをデタッチし、そのディスクからイメージを作成します。イメージの作成は、Google Cloud の管理コンソールやコマンドラインでも実行できますが、コマンドラインで実行する場合は、以下のようになります。
gcloud compute images create <IMAGE_NAME> --source-disk <DISK_NAME>
イメージを作成できたら、もはや VM インスタンスは不要なので、削除してください。VM インスタンスを削除するのもとても簡単です。コマンドで削除する場合は以下のようになります。
gcloud compute instances delete <INSTANCE_NAME>
ディスクイメージを作成できたので、ここからあとの操作は Cloud Shell 上で実行します。Cloud Shell の利用方法は、以下のドキュメントを参照してください。
https://cloud.google.com/shell/docs/launching-cloud-shell
3. Cloud Life Sciences APIの有効化
Google Cloud の Cloud Life Sciences は、ライフ サイエンスのデータを管理、処理、変換するための一連のサービスとツールです。このあとの手順で利用する dsub は、この Cloud Life Sciences を利用するので、ここでこの API を有効化しておきましょう。API の有効化も、Google Cloud 管理コンソールやコマンドラインで実行できます。コマンドラインで実行する場合は、以下のようになります。
gcloud services enable lifesciences.googleapis.com
4. dsub をインストール
dsub は、元は Google がバイオインフォマティックス向けの各種ソフトウェアを実行するために開発したコマンドラインのツールです。dsub は指定したコンテナイメージからコンテナを作成し、その中で任意のコマンドやスクリプトを実行できます。また、dsub は、イメージからディスクを作成しそれをマウントしてコンテナからアクセスできるようにしたり、Google Cloud Storage (GCS) に配置したファイルをコンテナに取り込み、コンテナ内で処理された結果を GCS に出力する、などの機能もあります。dsub には他にも様々な機能がありますが、コマンドパラメータを指定するだけでこれらを実行できます。
バイオインフォマティックスのツール群はその多くがコンテナ化されており、また、大量のデータを読み書きする必要があるため、この dsub のようなツールはとても重宝されます。
dsub は、 pip で容易にインストールすることができます。Cloud Shell 上で Python3 の venv を使ってインストールする場合は以下のような手順です。
python3 -m venv env
source env/bin/activate
pip install dsub
dsub の詳細については、 GitHub を参照してください。
5. コンテナ内で実行するスクリプトを作成
コンテナの内部で実行されるスクリプトファイルを作成しておきます。ここでは、“alphafold.sh” というファイル名で、以下の内容を保存しておきます。パラメータが多めですが、/app/run_alphafold.sh を実行しているだけです。また、${}で囲まれた変数は、後述する dsub コマンドのパラメータから引き渡されます。
cd /app/alphafold/app/run_alphafold.sh \
--fasta_paths=${FASTA} \
--uniref90_database_path=${DB}/uniref90/uniref90.fasta \
--mgnify_database_path=${DB}/mgnify/mgy_clusters_2018_12.fa \
--pdb70_database_path=${DB}/pdb70/pdb70 \
--data_dir=${DB} \
--template_mmcif_dir=${DB}/pdb_mmcif/mmcif_files \
--obsolete_pdbs_path=${DB}/pdb_mmcif/obsolete.dat \
--uniclust30_database_path=${DB}/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
-- bfd_database_path=${DB}/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
--output_dir=${OUT_PATH} \
--model_names=model_1,model_2,model_3,model_4,model_5 \
--max_template_date=2020–05–14 \
--preset=full_dbs \
--benchmark=False \
--logtostderr
これで、初回のみに必要な準備作業は完了です。ここまでの作業で、Google Cloud プロジェクト内には、以下のリソースが作成されているはずです。
・コンテナレジストリに AlphaFold のコンテナイメージが登録されている
・ディスクイメージが作成されている
AlphaFold 実行
AlphaFold を実行するための準備はできたので、動かしてみましょう。ここまでの手順は、初回の準備だけに必要な作業なので再度実行することはありません。あとは、以下のような簡単な手順だけで AlphaFold をすぐに何度も実行することができます。
- アミノ酸配列(FASTA)を GCS に保存
- AlphaFold を実行
それぞれの手順を説明していきます。
1. アミノ酸配列(FASTA)を GCS に保存
AlphaFold への入力データとなるアミノ酸配列を用意しましょう。アミノ酸配列のデータは、通常、アルファベットが羅列されたテキスト形式です。
サンプルデータファイルを作成し、それを Google Cloud Storage (GCS) にコピーしましょう。”all0174.fasta ” というファイル名で以下の内容のデータを保存します。最初の ”>” で始まる行はコメント行で、その次の行以降に配列が記述されます。
>all0174
MTEADSSVLQIWGGHPLQGHVKISGAKNSALVIMAGALLCSGDCRIRNVPLLADVERMGEVISALGVRLTRQADIIDINASEIKTSKAPYELVTQLRASFFAIGAILARLGVAQMPLPGGCAIGARPVDLHVRGLQAMGAEVQIEHGICNAYVPGSGGRLKGAKIYLDTPSVGATETLMMAATLADGETILENAAREPEVVDLANFCKAMGANIQGAGTSTITIVGVPKLHSVDYSIIPDRIEAGTFLVAGAITRSEITLSSVVPEHLIPLIAKLRDIGVTIIEESPDCLRILPAEILKATDIDTLPHPGFPTDMQAPFMALLTLAEGDSIINESVFENRLRHASELNRLGADIRVKGNTAFVRGVPLLSGAPVIGTDLRASAALVIAGLAAEGKTTIQGLHHLDRGYDQIDVKLQQLGAKILRVREEPANAEVAVNNNVSPASIST
GCS バケットがなければ、作成しておきます。コマンドラインでバケットを作成する場合は以下のようになります。
gsutil mb gs://<PROJECT_ID>-alphafold
上記で作成した FASTA ファイル(all0174.fasta)を、GCS バケットの下の input フォルダーへコピーします。コマンドラインでコピーする場合は以下の通りです。
gsutil cp all0174.fasta gs://<PROJECT_ID>-alphafold/input/
2. AlphaFold を実行
dsub コマンドで AlphaFold を実行します。コマンドは以下のようになります。
dsub --provider google-cls-v2 \
--project <PROJECT_ID> \
--zones <ZONE_NAME> \
--logging gs://<PROJECT_ID>-alphafold/logs \
--image=gcr.io/<PROJECT_ID>/alphafold:latest \
--script=alphafold.sh \
--input FASTA=gs://<PROJECT_ID>-alphafold/input/all0174.fasta \
--mount DB=”<IMAGE_URL> 3000" \
--output-recursive OUT_PATH=gs://<PROJECT_ID>-alphafold/output \
--machine-type n1-standard-8 \
--boot-disk-size 100 \
--subnetwork <SUBNET_NAME>
パラメータの説明 :
--provider : google-cls-v2 (Cloud Life Sicences API (v2))を指定
--project : PROJECT ID を指定
--zones : ZONE 名を指定(ワイルドカード”*”も可能。例: “us-central1-*”)
--logging : ログの出力先フォルダーを指定
--image : AlphaFold コンテナイメージのURL
--script : 先に作成した script ファイル名を指定
--input : GCS にアップロードした FASTA ファイルのパスを指定
--mount : 先に作成したディスクイメージのURLを指定(*1)
--output-recursive : 成果物を出力する GCS フォルダーのパスを指定
--machine-type : n1-standard-8
--boot-disk-size : 100 GB
--subnetwork : サブネット名を指定する(default以外の場合)*1: 先に作成したディスクイメージのURLは、以下のコマンドで確認できます。
gcloud compute images list --no-standard-images --uri
dsub コマンドを実行するとすぐにジョブIDが表示されます。このジョブIDを使って、以下のコマンドで 実行中ジョブのステータスを確認することができます。
dstat --provider google-cls-v2 --project <PROJECT_ID> --jobs <JOB_ID>
また、ddel コマンドでジョブをキャンセルすることもできます。コマンドは以下のようになります。
ddel --provider google-cls-v2 --project <PROJECT_ID> --jobs <JOB_ID>
なお、上記の dsub コマンドは CPU のみを使って実行しますが、GPU を利用して AlphaFold を高速化することができます。例えば、NVIDIA K80 を 1 つ使用する場合は、以下のパラメータを追加して dsub を実行します。
--accelerator-type nvidia-tesla-k80
--accelerator-count 1
GPU が利用できるリージョンや GPU の種類については、以下の URL で確認できます。
https://cloud.google.com/compute/docs/gpus/gpu-regions-zones
さらに、プリエンプティブル VM を使うことで費用を安価に抑えることができます。プリエンプティブル VM は有効時間が短い分、手頃な料金で利用できる Compute Engine VM インスタンスです。dsub コマンドでプリエンプティブル VM を使う場合は、以下のパラメータを追加して dsub を実行するだけです。
--preemptible
タンパク質 3D 構造を可視化
AlphaFold は、完了するまでに数十分から数時間かかります。また、コマンドが完了すると、dsub の output-recursive パラメータで指定したフォルダーの下に AlphaFold の予測結果であるタンパク質構造を格納したファイル(PDBファイル)などが出力されています。
拡張子が .pdb のファイルを一つダウンロードして、それを可視化してみましょう。
PDB ファイルを扱うためのソフトウェアは多数ありますが、ファイルをアップロードするだけで簡単に3D構造を可視化できるWEBサイトも複数あります。
例えば以下のようなサイトです。
https://www.rcsb.org/3d-view
https://molstar.org/viewer/
https://www.ncbi.nlm.nih.gov/Structure/icn3d/full.html?mmdbid=1TUP
どのWEBサイトも似たような操作で可視化できます。Open File メニューから PDB ファイルを選択し Apply をクリックすると、以下のような3D構造が表示されます。
まとめ
今回紹介した dsub コマンドで AlphaFold を実行するという方法は、仮想サーバを起動し続ける必要がないため費用を安価に抑えることができます。また、GPU + preemptible をセットで使えば、処理を高速化しつつ、さらに安価に実行することができるでしょう。
また、今回は dsub を使いましたが、dsub を使わず Cloud Life Sciences のみでも同じことは実現できます。
Google Cloud では、ライフサイエンスをサポートするための様々なソリューションを提供しています。詳しくは以下のドキュメントを参照してください。
https://cloud.google.com/life-sciences/docs/how-tos