Amazon Web Services ブログ

Amazon EC2でのAlphaFold v2.0の実行

この記事は、”Run AlphaFold v2.0 on Amazon EC2” を翻訳したものです。

DeepMindによるGitHub上のAlphaFold v2.0のオープンソースに関するNatureの論文の後、科学および研究コミュニティの多くは、DeepMindのAlphaFold実装を直接試してみたいと考えていました。Amazon Elastic Compute Cloud (Amazon EC2)とNVIDIA GPUを介したコンピューティングリソースを使用すると、AlphaFoldをすばやく実行して自分で試すことができます。

この記事では、NVIDIA GPUを使用してEC2インスタンスにAlphaFoldをインストールする方法について、順を追って説明します。

ソリューション概要

このプロセスは、Deep Learning Amazon Machine Image (DLAMI) から始まります。インストール後、インスタンスで CASP14サンプルを持つ AlphaFold モデルを使用して予測を実行します。また、Amazon Elastic Block Store (Amazon EBS) スナップショットを将来使用できるように作成して、再設定する手間を減らし、コストを節約する方法も示します。

新しいEC2インスタンスをゼロからセットアップせずにAlphaFoldを実行するには、この記事の最後のセクションに進んでください。提供されたパブリックEBSスナップショットを使用して、短時間で新しいEC2インスタンスを作成できます。

すべてのステップを完了し、24時間以内にすべてのリソースをシャットダウンした場合、この投稿で使用したAWSリソースの総コストは100ドル未満になります。EBSスナップショットを作成してAWSアカウント内に保存した場合、EBSスナップショットのストレージコストは月額約150ドルになります。

DLAMIを使用してEC2インスタンスを起動

このセクションでは、AWSからDLAMIを使用してEC2インスタンスをセットアップする方法を示します。AlphaFoldの依存関係がすでに多数プリインストールされており、セットアップにかかる時間を節約できます。

  1. Amazon EC2コンソールで、要件にあったAWS リージョンを選択します。
  2. Deep Learning AMIを検索して、DLAMIを使用して新しいEC2インスタンスを起動します。ここでは、この記事の執筆時点での最新バージョンUbuntu 18.04をベースにしたDLAMI version 48.0を使っています。EC2 AMI
  3. インスタンスタイプとしてGPUが1つあるp3.2xlargeインスタンスを選択します。p3.2xlargeインスタンスに十分なクォータがない場合は、AWSアカウントのAmazon EC2クォータを増やすことができます。P3 2xlarge
  4. AWS環境の要件に基づいて、適切なAmazon Virtual Private Cloud (Amazon VPC)を設定します。Amazon VPCを初めて設定する場合は、デフォルトのAmazon VPCの使用を検討し、「Amazon VPCの使用を開始する」を確認してください。
  5. システムボリュームを200 GiBに設定し、サイズが3 TB (3072 GiB)の新しいデータボリュームを1つ追加します。EBS Volume
  6. セキュリティグループの設定で SSH を使用してEC2インスタンスにアクセスできること、およびEC2インスタンスがインターネットに接続してAlphaFoldやその他のパッケージをインストールできることを確認します。
  7. EC2インスタンスを起動します。
  8. EC2インスタンスの準備が整うのを待ち、SSHを使用してAmazon EC2ターミナルにアクセスします。
  9. 必要に応じて、新しいEC2インスタンスに必要なソフトウェアがほかにある場合は、ここでインストールします。

AlphaFoldをインストール

これでAlphaFoldをインストールする準備ができました。

  1. SSHを使用してAmazon EC2ターミナルにアクセスしたら、まずすべてのパッケージを更新します。
    sudo apt update
  2. データボリュームを/dataフォルダにマウントします。詳細についてはLinuxでAmazon EBSボリュームを使用できるようにするを参照してください。
  3. lsblkコマンドを使用して、使用可能なディスクデバイスとそのマウントポイント (該当する場合) を表示します。使用する正しいデバイス名を判断するのに役立ちます。
    lsblk

    lsblk command

  4. ボリュームにファイルシステムがあるかどうかを確認します。新しいボリュームはrawブロックデバイスであり、マウントして使用する前に、そのボリューム上にファイルシステムを作成する必要があります。デバイスは空のボリュームです。
    sudo file -s /dev/xvdb

    sudo file

  5. ボリューム上にファイルシステムを作成し、ボリュームを/dataフォルダにマウントします。
    sudo mkfs.xfs /dev/xvdb
    sudo mkdir /data
    sudo mount /dev/xvdb /data
    sudo chown ubuntu:ubuntu -R /data
    df -h /data
  6. AlphaFoldとその他の必要なツールをインストールします。
    sudo apt install aria2 rsync git vim wget tmux tree -y
  7. 作業フォルダを作成し、GitHubリポジトリからAlphaFoldコードのクローンを作成します。
    cd /data
    mkdir -p /data/af_download_data
    mkdir -p /data/output/alphafold
    mkdir -p /data/input
    git clone https://github.com/deepmind/alphafold.git

    新しいボリュームは排他的に使用でき、後で作成するスナップショットには必要なデータがすべて揃っています。

  8. 提供されたスクリプトをバックグラウンドで使用して、データをダウンロードします。AlphaFoldには、複数の遺伝子 (配列) データベースとモデルパラメーターが必要です。
    nohup /data/alphafold/scripts/download_all_data.sh /data/af_download_data &

    ダウンロードプロセス全体に10時間以上かかることがありますので、完了するまで待ちます。次のコマンドを使用して、ダウンロードと解凍のプロセスを監視できます。

    du -sh /data/af_download_data/*

    ダウンロードプロセスが完了すると、/data/af_download_dataフォルダに次のファイルがあるはずです。

    $DOWNLOAD_DIR/                             # Total: ~ 2.2 TB (download: 438 GB)
        bfd/                                   # ~ 1.7 TB (download: 271.6 GB)
            # 6 files.
        mgnify/                                # ~ 64 GB (download: 32.9 GB)
            mgy_clusters_2018_12.fa
        params/                                # ~ 3.5 GB (download: 3.5 GB)
            # 5 CASP14 models,
            # 5 pTM models,
            # LICENSE,
            # = 11 files.
        pdb70/                                 # ~ 56 GB (download: 19.5 GB)
            # 9 files.
        pdb_mmcif/                             # ~ 206 GB (download: 46 GB)
            mmcif_files/
                # About 180,000 .cif files.
            obsolete.dat
        small_bfd/                             # ~ 17 GB (download: 9.6 GB)
            bfd-first_non_consensus_sequences.fasta
        uniclust30/                            # ~ 86 GB (download: 24.9 GB)
            uniclust30_2018_08/
                # 13 files.
        uniref90/                              # ~ 58 GB (download: 29.7 GB)
            uniref90.fasta
  9. /data/alphafold/docker/run_docker.pyを更新して、構成がローカルパスに反映されるようにします。
    vim /data/alphafold/docker/run_docker.py

    作成したフォルダでは、構成は次のようになります。EC2インスタンスで別のフォルダ構造を設定する場合は、それに応じて設定します。

    #### USER CONFIGURATION ####
    
    # Set to target of scripts/download_all_databases.sh
    DOWNLOAD_DIR = '/data/af_download_data'
    
    # Name of the AlphaFold Docker image.
    docker_image_name = 'alphafold'
    
    # Path to a directory that will store the results.
    output_dir = '/data/output/alphafold'

    User Config

  10. NVIDIAコンテナキットがインストールされていることを確認します。
    sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

    次のスクリーンショットのような出力が表示されるはずです。
    NVidia docker

  11. AlphaFold Dockerイメージをビルドします。ローカルパスが/data/alphafoldであることを確認してください。.dockerignoreファイルはそのフォルダの下にあります。
    cd /data/alphafold
    docker build -f docker/Dockerfile -t alphafold .
    docker images

    ビルドが完了すると、新しいDockerイメージが表示されます。
    alpha folder docker image

  12. pipを使用して、AlphaFoldに必要なすべてのPython依存関係をインストールします。
    pip3 install -r /data/alphafold/docker/requirements.txt
  13. CASP14ターゲットリストに移動し、T1050のプレーンテキストリンクからシーケンスをコピーします。
    CASP14 T1050
  14. コンテンツを新しいT1050.fastaファイルにコピーし、/data/inputフォルダーに保存します。
    T1050 fasta file
  15. これと同じプロセスを使用して、/data/inputフォルダーの下にテスト用の.fastaファイルをさらにいくつか作成できます。
    fasta files

GPU用のCloudWatchモニタリングをインストール (オプション)

必要に応じて、GPU用のAmazon CloudWatchモニタリングをインストールします。これにはAWS Identity and Access Management (IAM)ロールが必要です。

  1. CloudWatch用のAmazon EC2 IAMロールを作成し、EC2インスタンスにアタッチします。
    IAM Role
  2. gpumon.pyを開いて、EC2 インスタンスを起動しているリージョンを設定し、CloudWatch名前空間としてAlphaFoldのような新しい名前空間を指定します。
    vim ~/tools/GPUCloudWatchMonitor/gpumon.py

    gpumon config

  3. gpumonを起動し、CloudWatchへのGPUメトリックスの送信を開始します。
    source activate python3
    python ~/tools/GPUCloudWatchMonitor/gpumon.py &

AlphaFoldで予測

AlphaFoldで予測を実行する準備が整いました。

  1. 次のコマンドを使用して、/data/input/T1050.fastaからのタンパク質配列の予測を実行します。
    nohup python3 /data/alphafold/docker/run_docker.py --fasta_paths=/data/input/T1050.fasta --max_template_date=2020-05-14 &
  2. tailコマンドを使用して、予測の進行状況を確認します。
    tail -F /data/nohup.out

    全体の予測が完了するまでに数時間かかります。予測が完了すると、出力フォルダーに次の内容が表示されます。この場合、<target_name>T1050です。

    <target_name>/
        features.pkl
        ranked_{0,1,2,3,4}.pdb
        ranking_debug.json
        relaxed_model_{1,2,3,4,5}.pdb
        result_model_{1,2,3,4,5}.pkl
        timings.json
        unrelaxed_model_{1,2,3,4,5}.pdb
        msas/
            bfd_uniclust_hits.a3m
            mgnify_hits.sto
            pdb70_hits.hhr
            uniref90_hits.sto
    
  3. 出力フォルダの所有者をrootから変更して、コピーできるようにします。
    sudo chown ubuntu:ubuntu /data/output/alphafold/ -R
  4. scpを使用して、予測出力フォルダからローカルフォルダに出力をコピーします。
    scp -i <ec2-key-path>.pem -r ubuntu@<ec2-ip>:/data/output/alphafold/T1050 ~/Downloads/
  5. NIHのタンパク質3Dビューアを使用して、結果フォルダーから予測された3D構造を表示します。
  6. 信頼度が最も高い予測を含むranked_0.pdbを選択します。
    Open PDB file
    以下は、AlphaFoldによるT1050の予測構造の3Dビューです。
    T1050 3D view

データボリュームからスナップショットを作成

AlphaFoldをEC2インスタンスにインストールするには時間がかかります。ただし、P3インスタンスとEBSボリュームを常に実行し続けると、コストが高くなる可能性があります。EC2インスタンスをすぐに準備したいが、必要なたびに環境の再構築に時間をかけたくない場合もあります。EBSスナップショットは、時間とコストの両方を節約するのに役立ちます。

  1. Amazon EC2コンソールで、Elastic Block Storeの下のナビゲーションペインで[ボリューム]を選択します。
  2. EC2インスタンスIDでフィルタリングします。2つのボリュームが表示されます。
  3. サイズが3072 GiBのデータボリュームを選択します。
  4. [アクション]メニューの[スナップショットの作成]を選択します。
    Create snapshot
    スナップショットの作成には数時間かかります。
  5. スナップショットが完了したら、[スナップショット]を選択すると、新しいスナップショットがリストに表示されます。

これで EC2インスタンスを安全にシャットダウンできます。この時点で、データボリューム内のすべてのデータは、将来の使用に備えてスナップショットに安全に保存されます。

スナップショットを使用してEC2インスタンスを再作成

AlphaFoldを使用して新しいEC2インスタンスを再作成するには、最初にEC2インスタンスをゼロから作成したときに行った手順と似ています。ただし、データボリュームをゼロから作成する代わりに、Amazon EBSスナップショットから復元した新しいボリュームをアタッチします。

  1. お好みのAWSリージョンでAmazon EC2コンソールを開き、Deep Learning AMIを検索してDLAMIを使用して新しいEC2インスタンスを起動します。
  2. Deep Learning AMI (Ubuntu 18.04)を選択します。
    EC2 AMI
  3. インスタンスタイプとして1つの GPU を持つp3.2xlargeを選択します。p3.2xlargeインスタンスに十分なクォータがない場合は、AWSアカウントのAmazon EC2クォータを増やすことができます。
    P3 2xlarge
  4. AWS環境の要件に基づいて、適切なAmazon VPC設定を構成します。Amazon VPCを初めて設定する場合は、デフォルトのAmazon VPCの使用を検討し、「Amazon VPC の使用を開始する」を確認してください。
  5. システムボリュームを200 GiBに設定しますが、新しいデータボリュームを追加しないでください。
    EBS volume
  6. セキュリティグループ設定でEC2インスタンスへのアクセスが許可され、EC2インスタンスがインターネットに接続してPythonおよびDockerパッケージをインストールできることを確認します。
  7. EC2インスタンスを起動します。
  8. 後のステップで使用するため、インスタンスのアベイラビリティーゾーンとインスタンスID を書き留めておきます。
    restored ec2 AZ
  9. Amazon EC2コンソールで、[スナップショット]を選択します。
  10. 以前に作成したスナップショットを選択するか、以下の手順11で示されている公開スナップショットを使用します。
  11. [アクション]メニューの[ボリュームの作成]を選択します。
    この記事では、リージョンus-east-1us-west-2、およびeu-west-1で公開スナップショットを提供しています。パブリックスナップショットは、スナップショットIDで検索できます。us-east-1の場合はsnap-0d736c6e22d0110d0us-west-2の場合はsnap-080e5bbdfe190ee7eeu-west-1ではsnap-08d06a7c7c3295567です。
  12. 新しいデータボリューム設定を適宜セットアップします。アベイラビリティーゾーンが新しく作成されたEC2インスタンスと同じであることを確認します。そうしないと、新しいEC2インスタンスにボリュームをマウントできません。
  13. [ボリュームの作成]を選択して、新しいデータボリュームを作成します。
    restore snapshot
  14. [ボリューム]を選択すると、新しく作成されたデータボリュームが表示されます。その状態は[使用可能]であるはずです。
  15. ボリュームを選択し、[アクション]メニューの[ボリュームのアタッチ]を選択します。
  16. 新しく作成したEC2インスタンスを選択し、ボリュームをアタッチします。
    Attached volume
  17. SSHを使用してAmazon EC2ターミナルにアクセスし、lsblkを実行します。新しいデータボリュームがアンマウントされているのがわかります。この場合は/dev/xvdfです。
    lsblk
  18. ボリュームにファイルシステムがあるかどうかを確認します。スナップショットから作成されたデータボリュームには、すでにXFSファイルシステムがあります。
    sudo file -s /dev/xvdf

    sudo file

  19. 新しいデータボリュームを/dataフォルダにマウントします。
    sudo mkdir /data
    sudo mount /dev/xvdf /data
    sudo chown ubuntu:ubuntu -R /data
  20. システム上のすべてのパッケージを更新し、依存関係をインストールします。AlphaFold Dockerイメージを再構築する必要があります。
    sudo apt update
    sudo apt install aria2 rsync git vim wget tmux tree -y
    pip3 install -r /data/alphafold/docker/requirements.txt
    
    cd /data/alphafold
    docker build -f docker/Dockerfile -t alphafold .
    docker images

    alpha folder docker image

  21. NVIDIAコンテナキットがインストールされていることを確認します。
    sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

    次のスクリーンショットのような出力が表示されるはずです。
    NVidia docker

  22. 次のコマンドを使用して、/data/input/T1024.fastaのタンパク質配列の予測を実行します。
    cd /data
    nohup python3 /data/alphafold/docker/run_docker.py --fasta_paths=/data/input/T1024.fasta --max_template_date=2020-05-14 &

    スナップショットにはすでにT1050の結果が含まれているため、別のタンパク質配列を使用しています。T1050の予測をもう一度実行する場合は、新しい予測を実行する前に、まず既存のT1050結果フォルダーを削除するか名前を変更します。

  23. tailコマンドを使用して、予測の進行状況を監視します。
    tail -F /data/nohup.out

    全体の予測が完了するまでに数時間かかります。

クリーンアップ

すべての予測を完了し、結果をローカルにコピーしたら、AWSリソースをクリーンアップしてコストを削減します。EC2インスタンスがシャットダウンされたときにEBSデータボリュームが削除されなかった場合は、EC2インスタンスを安全にシャットダウンし、EBSデータボリュームを削除できます。AlphaFoldを再度使用する必要がある場合は、同じプロセスに従って新しいEC2インスタンスを起動し、新しい予測をほんの数分で開始できます。また、EBSスナップショットストレージ以外の追加コストは発生しません。

まとめ

Amazon EC2とNVIDIA GPU、Deep Learning AMIを使用すると、DeepMindの新しいAlphaFold実装をインストールし、CASP14サンプルで予測を実行できます。データボリューム上のデータをその時点でスナップショットにバックアップするため、EC2インスタンスやEBSボリュームが不要なときに料金を支払う必要がありません。前のスナップショットに基づいてEBSボリュームを作成すると、AlphaFoldを使用してEC2インスタンスを再作成するのに必要な時間が大幅に短縮されます。したがって、短時間で予測の実行を開始できます。


著者について

Qi WangはAWSのグローバルヘルスケアおよびライフサイエンスチームに所属するシニアソリューションアーキテクトです。彼はヘルスケア・ライフサイエンス分野におけるビジネスイノベーションとデジタルトランスフォーメーションにおいて10年以上の経験があります。AWSでは、創薬研究、臨床試験、および商業化においてライフサイエンスのお客様の変革に貢献しています。

翻訳はIndustry Solutions Architectの松永が担当しました。原文はこちらです。