OpenFold を AWS ParallelCluster で動かしてタンパク質の立体構造を予測してみた

大村保貴
2025.02.28
タンパク質の立体構造予測には強力な計算リソースが必要です。本記事では、AWS ParallelCluster と Apptainer の組み合わせ、ストレージは FSx for Lusrete を使用したタンパク質の立体構造予測環境を紹介します。この組み合わせには以下のメリットがあります。
スケーラブルな GPU 計算環境の構築
GPU リソースの占有利用
コンテナ利用による再現性の高い実行環境の確保
実際の環境構築から実行、性能比較までの一連の流れを解説します。
 前提条件と参考資料本記事は以下の過去記事を参考にしています。
Apptainerを使ったOpenFoldの実行
AWS ParallelClusterでApptainerを実行する方法
 システム構成今回は以下のシステム構成を採用しました。
AWS ParallelCluster で HPC 環境を構築
既存の FSx for Lustre をマウントし、S3 バケットとリンク
OpenFold のデータベースと Apptainer イメージを S3 に保存
コンピュートノードから Lustre 経由でデータにアクセス
今回の検証では、ParallelCluster + Apptainer の組み合わせで OpenFold を実行することと、FSx for Lustre のスループットが全体の処理時間与える影響を検証しました。
 検証環境

コンポーネント
バージョン/仕様


ParallelCluster
3.7.0

CUDA
12.4

OS
Ubuntu 22.04

FSx for Lustre
2.15

ヘッドノード
m7i-flex.large

計算ノード
g6.4xlarge (NVIDIA L4 GPU)

FSx for Lustre (1回目)
スループット: 125MB/s

FSx for Lustre (2回目)
スループット: 1000MB/s

 環境構築 ParallelClusterの設定と作成以下のコンフィグファイルを使用して ParallelCluster を作成します。

ヘッドノードはm7i-flex.large、GPU 計算用ノードはg6.4xlargeを使用します。
コンフィグ全文 config.yaml
Region: us-east-1
Image:
  Os: ubuntu2204
Tags:
  - Key: Name
    Value: apptainer-cluster

# ----------------------------------------------------------------
# Head Node Settings
# ----------------------------------------------------------------
HeadNode:
  InstanceType: m7i-flex.large
  Networking:
    ElasticIp: true
    SubnetId: subnet-01a6ed0305dc2dd6f
  LocalStorage:
    RootVolume:
      Size: 40
      Encrypted: true
      VolumeType: gp3
      Iops: 3000
      Throughput: 125
  Iam:
    AdditionalIamPolicies:
      - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore

# ----------------------------------------------------------------
# Compute Node Settings
# ----------------------------------------------------------------
Scheduling:
  Scheduler: slurm
  SlurmSettings:
    ScaledownIdletime: 5
  SlurmQueues:
  # ------ Compute ------
    - Name: test
      ComputeResources:
        - Name: test
          Instances:
            - InstanceType: c7i.large
          MinCount: 0
          MaxCount: 10
          DisableSimultaneousMultithreading: true
      ComputeSettings:
        LocalStorage:
          RootVolume:
            Size: 40
            Encrypted: true
            VolumeType: gp3
            Iops: 3000
            Throughput: 125
      CapacityType: SPOT
      AllocationStrategy: price-capacity-optimized
      Networking:
        SubnetIds:
          - subnet-097a4d57a0ba52d7f
        AssignPublicIp: true
        PlacementGroup:
          Enabled: false
      Iam:
        AdditionalIamPolicies:
          - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
    - Name: gpu1
      ComputeResources:
        - Name: gpu1
          Instances:
            - InstanceType: g6.4xlarge
          MinCount: 0
          MaxCount: 10
          DisableSimultaneousMultithreading: true
      ComputeSettings:
        LocalStorage:
          RootVolume:
            Size: 40
            Encrypted: true
            VolumeType: gp3
            Iops: 3000
            Throughput: 125
      CapacityType: ONDEMAND
      Networking:
        SubnetIds:
          - subnet-097a4d57a0ba52d7f
        AssignPublicIp: true
        PlacementGroup:
          Enabled: false
      CustomActions:
        OnNodeConfigured:
          Sequence:
            - Script: s3://ohmura-parallelcsluter-scripts/install.sh
      Iam:
        AdditionalIamPolicies:
          - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
        S3Access:
          - BucketName: ohmura-parallelcsluter-scripts
            EnableWriteAccess: false

# ----------------------------------------------------------------
# Shared Storage Settings
# ----------------------------------------------------------------
SharedStorage:
  - MountDir: /mnt/lustre
    Name: lustre
    StorageType: FsxLustre
    FsxLustreSettings:
      FileSystemId: fs-07ef5bd10f711dc1d

# ----------------------------------------------------------------
#  Other Settings
# ----------------------------------------------------------------
Monitoring:
  Logs:
    CloudWatch:
      Enabled: true
      RetentionInDays: 180
      DeletionPolicy: "Delete"
  Dashboards:
    CloudWatch:
      Enabled: false
以下のコマンドでクラスターを作成します。
pcluster create-cluster --cluster-name openfold-cluster --cluster-configuration config.yaml
 Apptainer のインストールコンピュートノードに Apptainer をインストールするスクリプトを作成します。

このスクリプトは ParallelCluster のOnNodeConfiguredアクションで実行されます。
#! /bin/bash

# 必要なパッケージのインストール
sudo apt update
sudo apt install -y wget

# 最新の Apptainer Debian パッケージをダウンロード
cd /tmp
wget https://github.com/apptainer/apptainer/releases/download/v1.3.6/apptainer_1.3.6_amd64.deb

# パッケージをインストール
sudo apt install -y ./apptainer_1.3.6_amd64.deb
このスクリプトを S3 バケット（例：s3://ohmura-parallelcsluter-scripts/install.sh）にアップロードします。
 データベースの準備OpenFold の実行には大量のデータベースファイルが必要です。

データベースのダウンロードは、OpenFold を使用してタンパク質の立体構造予測をやってみたの記事で説明した手順でダウンロードします。
データベースは以下の手順で準備します。
OpenFold のデータベースファイルをダウンロード
S3 バケットにアップロード
FSx for Lustre とリンク
!FSx for Lustre を使うことで、S3 に保存された大量のデータに高速アクセスできます。

データベースは約 3TB あるため、高速な大容量ストレージには FSx for Lustre + S3 が最適です。
この構成により、複数のコンピュートノードから同じデータベースにアクセスできます。
 OpenFold の Apptainer イメージ準備OpenFold の Apptainer イメージは、Apptainer を使った OpenFold の実行の記事で説明した手順で作成します。

作成したイメージファイル（openfold.sif）は以下の手順で配置します。
作成した Apptainer イメージ（openfold.sif）を S3 バケットにアップロード
FSx for Lustre を通じてコンピュートノードからアクセス
この構成により、複数のコンピュートノードから同じイメージファイルを共有して使用できます。
 ヘッドノードで準備作業ヘッドノードにログインし、テストデータを準備します。

OpenFold のリポジトリをクローンし、CUDA 12 対応ブランチを使用します。
cd ~
git clone -b pl_upgrades https://github.com/aqlaboratory/openfold.git
テスト用のアミノ酸配列ファイルは /home/ubuntu/openfold/examples/monomer/6kwc.fasta を利用しました。
 OpenFoldの実行 ジョブスクリプトの作成以下の Slurm ジョブスクリプトを作成して OpenFold を実行します。
#!/bin/bash
#SBATCH --job-name=openfold_opt
#SBATCH --output=openfold_%j.log
#SBATCH --error=openfold_%j.log
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8       # 全物理コアを使用
#SBATCH --gres=gpu:1            # 1つのGPUを使用

# エラー時に停止するように設定
set -e

# 環境変数の設定
IMAGE_PATH="/mnt/lustre/s3/images"
DATABASE_PATH="/mnt/lustre/s3/data"
WORK_DIR="/home/ubuntu/openfold/examples/monomer"
OUTPUT_DIR="/mnt/lustre/s3/openfold_results"

# タイムスタンプ付きの出力ディレクトリを作成
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
RUN_OUTPUT_DIR="${OUTPUT_DIR}/${TIMESTAMP}"
sudo mkdir -p $RUN_OUTPUT_DIR
sudo chown -R $USER:$USER $RUN_OUTPUT_DIR

echo "Job started at $(date)"
echo "Running on host: $(hostname)"
echo "Working directory: $WORK_DIR"
echo "Output directory: $RUN_OUTPUT_DIR"

# 作業ディレクトリに移動
cd $WORK_DIR

# システム情報を表示
echo "=== SYSTEM RESOURCES ==="
echo "CPU Information:"
lscpu | grep -E 'Model name|Socket|Thread|CPU\(s\)'
echo "Memory Information:"
free -h
echo "GPU Information:"
nvidia-smi

export OMP_NUM_THREADS=8        # OpenMPスレッド数を物理コア数に設定
export MKL_NUM_THREADS=8        # Intel MKLスレッド数を設定
export NUMEXPR_NUM_THREADS=8    # NumExprスレッド数を設定

# Apptainer 実行コマンド
echo "Starting OpenFold with Apptainer..."
apptainer exec --nv \
    --bind ${WORK_DIR}/fasta_dir:/local/fasta_dir,${RUN_OUTPUT_DIR}:/local/results,${DATABASE_PATH}:/database \
    --env OMP_NUM_THREADS=8,MKL_NUM_THREADS=8,NUMEXPR_NUM_THREADS=8 \
    ${IMAGE_PATH}/openfold.sif \
    python3 /opt/openfold/run_pretrained_openfold.py \
    /local/fasta_dir \
    /database/pdb_mmcif/mmcif_files/ \
    --uniref90_database_path /database/uniref90/uniref90.fasta \
    --mgnify_database_path /database/mgnify/mgy_clusters_2022_05.fa \
    --pdb70_database_path /database/pdb70/pdb70 \
    --uniclust30_database_path /database/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --bfd_database_path /database/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
    --jackhmmer_binary_path /opt/conda/bin/jackhmmer \
    --hhblits_binary_path /opt/conda/bin/hhblits \
    --hhsearch_binary_path /opt/conda/bin/hhsearch \
    --kalign_binary_path /opt/conda/bin/kalign \
    --openfold_checkpoint_path /database/openfold_params/finetuning_ptm_2.pt \
    --config_preset "model_1_ptm" \
    --model_device "cuda:0" \
    --output_dir /local/results \
    --cpus 8

# 実行後のGPU状態を記録
echo "GPU state after execution:"
nvidia-smi

echo "Job completed at $(date)"
echo "Results are available at: $RUN_OUTPUT_DIR"
 ジョブの実行作成したジョブスクリプトを以下のコマンドで実行します。
sbatch openfold_job.sh
ジョブが終わるまで待ちます。
$ squeue
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
                22      gpu1 openfold   ubuntu  R    2:18:43      1 gpu1-dy-gpu1-4
 実行結果と性能比較 FSx for Lustreスループットの影響今回の検証では同じジョブを異なる FSx for Lustre スループット設定で実行し、処理時間を比較しました。


スループット設定
実行時間
アライメント生成
推論時間
リラクゼーション


125MB/s
2時間23分
約2時間20分
34.03秒
9.00秒

1000MB/s
54分38秒
約54分
33.80秒
7.93秒

ジョブの実行結果から以下の情報が得られました。
構造予測の推論処理は非常に高速なのですが、データベースを検索する処理がストレージのスループット依存でした。
実行ログ(1 回目)全文Job started at Thu Feb 27 04:49:23 UTC 2025
Running on host: gpu1-dy-gpu1-4
Working directory: /home/ubuntu/openfold/examples/monomer
Output directory: /mnt/lustre/s3/openfold_results/20250227_044923
=== SYSTEM RESOURCES ===
CPU Information:
CPU(s):                             16
On-line CPU(s) list:                0-7
Off-line CPU(s) list:               8-15
Model name:                         AMD EPYC 7R13 Processor
Thread(s) per core:                 1
Socket(s):                          1
NUMA node0 CPU(s):                  0-7
Memory Information:
               total        used        free      shared  buff/cache   available
Mem:            60Gi       518Mi        58Gi       1.0Mi       1.5Gi        59Gi
Swap:             0B          0B          0B
GPU Information:
Thu Feb 27 04:49:23 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.127.08             Driver Version: 550.127.08     CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA L4                      On  |   00000000:35:00.0 Off |                    0 |
| N/A   31C    P8             11W /   72W |       1MiB /  23034MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+
Starting OpenFold with Apptainer...
[2025-02-27 04:49:38,512] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)
INFO:/opt/openfold/openfold/utils/script_utils.py:Loaded OpenFold parameters at /database/openfold_params/finetuning_ptm_2.pt...
INFO:/opt/openfold/run_pretrained_openfold.py:Generating alignments for 6KWC_1...
INFO:/opt/openfold/openfold/utils/script_utils.py:Running inference for 6KWC_1...
INFO:/opt/openfold/openfold/utils/script_utils.py:Inference time: 34.03283635799926
INFO:/opt/openfold/run_pretrained_openfold.py:Output written to /local/results/predictions/6KWC_1_model_1_ptm_unrelaxed.pdb...
INFO:/opt/openfold/run_pretrained_openfold.py:Running relaxation on /local/results/predictions/6KWC_1_model_1_ptm_unrelaxed.pdb...
WARNING:root:Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.
INFO:/opt/openfold/openfold/utils/script_utils.py:Relaxation time: 8.998917716000506
INFO:/opt/openfold/openfold/utils/script_utils.py:Relaxed output written to /local/results/predictions/6KWC_1_model_1_ptm_relaxed.pdb...
GPU state after execution:
Thu Feb 27 07:12:19 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.127.08             Driver Version: 550.127.08     CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA L4                      On  |   00000000:35:00.0 Off |                    0 |
| N/A   51C    P0             29W /   72W |       1MiB /  23034MiB |     73%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+
Job completed at Thu Feb 27 07:12:19 UTC 2025
Results are available at: /mnt/lustre/s3/openfold_results/20250227_044923
 まとめ今回の検証では、AWS ParallelCluster、Apptainer、FSx for Lustre を組み合わせた OpenFold 実行環境を構築し、FSx for Lustre のスループット設定が処理時間に与える影響を検証しました。
ストレージスループットの重要性
FSx for Lustre のスループットを 125MB/s から 1000MB/s に向上させることで、全体の処理時間が約 2 時間 23 分から 54 分 38 秒へと大幅に短縮
データベース（ファイル）アクセスが処理時間の大部分を占めていた
計算ノードの性能だけでなく、ストレージのスループットも考慮する必要がある

コスト最適化
処理時間の短縮により、高額な GPU インスタンスの利用時間が減少
結果として、コスト効率が向上した

 おわりにAWS ParallelCluster を使うと(きっと)研究を加速させる強力な計算環境を自前で構築できます。ぜひ皆さんも試してみてください。
OpenFold を AWS ParallelCluster で動かしてタンパク質の立体構造を予測してみた

前提条件と参考資料

システム構成

検証環境

環境構築

ParallelClusterの設定と作成

Apptainer のインストール

データベースの準備

OpenFold の Apptainer イメージ準備

ヘッドノードで準備作業

OpenFoldの実行

ジョブスクリプトの作成

ジョブの実行

実行結果と性能比較

FSx for Lustreスループットの影響

まとめ

おわりに

関連記事

主なカテゴリ

AWS

おすすめ

プロダクト

コンテンツ

お問い合わせ

運営会社

コンポーネント	バージョン/仕様
ParallelCluster	3.7.0
CUDA	12.4
OS	Ubuntu 22.04
FSx for Lustre	2.15
ヘッドノード	m7i-flex.large
計算ノード	g6.4xlarge (NVIDIA L4 GPU)
FSx for Lustre (1回目)	スループット: 125MB/s
FSx for Lustre (2回目)	スループット: 1000MB/s
スループット設定	実行時間	アライメント生成	推論時間	リラクゼーション
125MB/s	2時間23分	約2時間20分	34.03秒	9.00秒
1000MB/s	54分38秒	約54分	33.80秒	7.93秒