データベースファイルを S3 に保存して Mountpoint for Amazon S3 を利用した OpenFold の実行は実用的か検証してみた
この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。
はじめに
OpenFold はタンパク質の立体構造を予測するツールです。本記事では、OpenFold の実行に必要なデータベースの保存先として、Amazon S3 が利用可能か検証します。具体的には、Mountpoint for Amazon S3 を使用した実行結果と課題を紹介します。
Icons made by Freepik from Flaticon
検証の目的と概要
OpenFold は大量のデータベースファイルを必要とします。ストレージの選択が実行性能に大きく影響します。
検証目的
- データベースを S3 保存によるランニングコスト削減
- Mountpoint for Amazon S3 で実行可能か確認
- 必要なシステムリソースの特定
検証内容
- データベースを S3 へのアップロード
- 異なるメモリサイズでの実行テスト
- パフォーマンスとコストの評価
検証結果
OpenFold のデータベースを Mountpoint for Amazon S3 を利用した実行は推奨できません。
主な理由
- 実行のたびに S3 からのデータロードに時間がかかる
- 処理が途中で停止する
前提知識
Apptainer での OpenFold 実行環境構築については、以下の記事を参照してください。
検証環境
推論用の GPUインスタンス
OpenFold の実行には GPU インスタンスが必要です。
CUDA 11 と 12 をサポートしている OpenFold に対し、CUDA 12 の実行環境がセットアップ済みの AMI を利用しました。
インスタンス構成
| 項目 | 仕様 |
|---|---|
| インスタンスタイプ | g6.xlarge / g6.4xlarge / g6.8xlarge |
| GPU | NVIDIA L4 Tensor Core GPU × 1 |
| メモリ | 16GB / 64GB / 128GB |
| vCPU | 4コア / 16コア / 32コア |
ソフトウェア環境
| 項目 | バージョン |
|---|---|
| AMI | Deep Learning Base OSS Nvidia Driver GPU AMI |
| OS | Ubuntu 22.04 LTS |
| CUDA | 12.4 |
| Apptainer | 1.3.6 |
S3 にデータベースをアップロード
VPC エンドポイント（ゲートウェイタイプ）経由で、EC2 から S3 へデータベースをアップロードしました。約 21 万個の小さいファイルの転送に 5 時間要しました。
S3 Storage Lens でのファイル数確認した様子。
S3 Put リクエストの料金は $1 もかかりませんでした。
Mountpoint for Amazon S3 の設定
S3 バケットのマウント方法は、以下の記事の手順で実施しました。
検証プロセス
OpenFold の実行
以下のコマンドで検証を実施しました。
# 環境変数の設定
IMAGE_PATH="/mnt/s3/images"
DATABASE_PATH2="/mnt/s3/data"
CURRENT_DIR=$PWD
cd /home/ubuntu/openfold/examples/monomer
mkdir results-s3
# Apptainer 実行コマンド
nohup apptainer exec --nv \
--bind ${CURRENT_DIR}:/local,${DATABASE_PATH2}:/database \
${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 \
[コマンド省略] \
--output_dir /local/results-s3 > s3.log &`
実行結果と課題
最初の検証（16GB メモリ）では、HHblits の実行で失敗しました。
[2025-01-03 08:48:37,146] [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...
ERROR:root:HHblits failed. HHblits stderr begin:
ERROR:root:- 09:20:09.396 INFO: Searching 65983866 column state sequences.
ERROR:root:HHblits stderr end
Traceback (most recent call last):
File "/opt/openfold/run_pretrained_openfold.py", line 493, in <module>
main(args)
File "/opt/openfold/run_pretrained_openfold.py", line 291, in main
precompute_alignments(tags, seqs, alignment_dir, args)
File "/opt/openfold/run_pretrained_openfold.py", line 111, in precompute_alignments
alignment_runner.run(
File "/opt/openfold/openfold/data/data_pipeline.py", line 534, in run
hhblits_bfd_unirefclust_result = run_msa_tool(
File "/opt/openfold/openfold/data/data_pipeline.py", line 262, in run_msa_tool
result = msa_runner.query(fasta_path)[0]
File "/opt/openfold/openfold/data/tools/hhblits.py", line 160, in query
raise RuntimeError(
RuntimeError: HHblits failed
stdout:
stderr:
- 09:20:09.396 INFO: Searching 65983866 column state sequences.
HHBlits はメモリを大きく消費するために、64GB メモリの g6.4xlarge で実行し直しました。htop の出力結果から、HHBlits と Mountpoint for Amazon S3 でメモリを消費していますが、まだメモリには余裕ありそうです。
top - 22:06:23 up 2:34, 3 users, load average: 4.07, 4.00, 3.38
Tasks: 311 total, 1 running, 310 sleeping, 0 stopped, 0 zombie
%Cpu(s): 0.2 us, 0.2 sy, 0.0 ni, 87.1 id, 12.4 wa, 0.0 hi, 0.0 si, 0.0 st
MiB Mem : 26.1/61912.1 [|||||||||||||||||||| ]
MiB Swap: 0.0/0.0 [ ]
PID USER PR NI VIRT RES SHR S %CPU %MEM TIME+ COMMAND
1392 ubuntu 20 0 5234596 1.5g 9600 S 4.7 2.5 12:07.91 /usr/bin/mount-s3 --allow-delet+
2322 ubuntu 20 0 1866.3g 16.3g 4.8g D 1.0 27.0 0:21.50 /opt/conda/bin/hhblits -i /loca+
632 root 20 0 548336 23168 8704 S 0.3 0.0 0:34.06 /usr/bin/nv-hostengine -n --ser+
651 root 20 0 2317660 44008 31488 S 0.3 0.1 0:05.09 /usr/bin/containerd
nohup でコマンドでバックグランド実行していたのでjobsの結果を確認し、バックグランド実行がないことを確認しました。実行ログを確認すると、実行完了のログを確認できず処理は途中のまま終わっていました。
[2025-01-04 03:32:58,129] [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...
インスタンスタイプをg6.8xlargeに変更してメモリを 128GB にしましたが、同様の結果となりました。メモリの問題ではないようです。そして、CloudWatch メトリクスから以下の問題を確認しました。
- データベースの読み込みに 40 分以上を要する
- 小さいファイルの大量アクセスでスループットが低下する
これらの結果から、S3 をデータベースの保存先として使用することは適切ではないと判断し検証を終了しました。
考察
パフォーマンスの観点
- データベースのアクセス速度の課題
- S3 からのデータロード時間が処理のボトルネック
- 小さいファイルの大量アクセスによる遅延があるのではないかと思われる
OpenFold 実行開始してから 40 分かけて合計約 1.7TB のデータ読み込みが発生しました。
- システムリソースの利用状況
- HHblits はメモリ消費が大きい
- Mountpoint for Amazon S3 自体のメモリを消費する
コストの観点
S3 に約 2.7TB のデータベースファイルを保存して実行できた場合、OpenFold 実行時に GET リクエストが如何ほどか気にしていました。検証のために 4-5 回実行したのですけど、リクエスト課金は$1 以下でした。
推奨構成
-
データベースファイルのストレージ構成
- FSx for シリーズまたは、gp3, io2 などの高速な IOPS 確保できる EBS を推奨
- gp3 の EBS で動作実績はあります
-
インスタンスタイプ選択
- 最小 64GB 以上のメモリ
- GPU 搭載インスタンス
おわりに
OpenFold の実運用では、パフォーマンスを優先したインフラ構成を推奨します。高速なストレージとして、FSx ファミリーや EBS gp3 の使用を検討してください。S3 へのアクセスは、FSx for Lustre を介した構成がベターでした。
