Amazon SageMaker AI上のBoltzGenでタンパク質設計を加速

Amazon SageMaker AI上でBoltzGenをデプロイすると、GPUコンピュートインフラストラクチャをエンドツーエンドで管理することにより、タンパク質バインダー設計を加速できます。BoltzGenは拡散ベースの生成モデルであり、特定の生体分子標的に結合できるタンパク質やペプチドを設計します。典型的な設計キャンペーンには、バックボーン生成、逆フォールディング、構造検証、候補ランキングといった複数のGPU集約的なステップが含まれます。これらのステップを数百、数千、さらには数百万の設計候補に対して実行するには、インスタンスのプロビジョニング、ステップ間のデータ移動、コスト追跡といった運用上のオーバーヘッドが発生します。SageMaker AIは、インスタンスのプロビジョニングから結果の配信、リソースのクリーンアップまで、このコンピュートライフサイクルを管理するため、インフラストラクチャの運用ではなく設計の反復に集中できます。

この記事では、SageMaker AIにBoltzGenをデプロイし、エンドツーエンドのタンパク質設計実験を実行する方法を説明します。チュートリアルの最後には、迅速な検証実行から本番バッチ処理まで拡張可能な動作セットアップが得られます。この設定は研究のさまざまな段階に最適化された2つの実行モードを提供し、反復ワークフローでの計算コストを削減するためにステップレベルのキャッシュを使用します。

このチュートリアルは、タンパク質バインダー設計、治療用タンパク質工学、またはde novoタンパク質アーキテクチャに取り組むアカデミック研究ラボ、バイオテクノロジー系スタートアップ、製薬R&Dグループ、教育プログラムに適用されます。

SageMaker AIがタンパク質設計のボトルネックにどう対処するか

BoltzGenキャンペーンの各ステップはGPUハードウェア上で実行され、一度に1つの設計仕様を処理します。リポジトリのベンチマークデータによると、4-GPUインスタンス（ml.g5.12xlarge）では、1000サンプルのキャンペーンが完了するまでに約375時間かかります。このインフラストラクチャを運用するには、CUDA環境の構築、GPUインスタンスのライフサイクル調整、ステップ間のデータパイプライン構築、長時間実行ジョブの障害からの回復が必要です。

SageMaker AIはこれらのボトルネックに直接対処します。ジョブを送信すると、SageMaker AIはGPUインスタンスをプロビジョニングし、コンテナ内でBoltzGenを実行します。結果をAmazon Simple Storage Service（Amazon S3）に書き込み、処理が完了するとインスタンスを解放します。課金は秒単位であるため、アイドルGPUコストは発生しません。ml.g4dn.xlargeでの2時間の設計実行は、オンデマンド価格で約1.50ドルです。

この実装は、単一インスタンス内のマルチGPU並列化と、パイプライン全体でのマルチインスタンススケーリングをサポートしています。パイプラインモードでは、各ステップの出力がAmazon S3に7日間の有効期限でキャッシュされるため、フィルタリングパラメータを反復するときに、計算コストの約90％を占める設計生成ステップが再実行されることはありません。

リポジトリ内のセットアップスクリプトはコンテナをビルドしてAmazon Elastic Container Registry（Amazon ECR）にプッシュし、作業例を使用すると数分以内に最初の設計ジョブを送信できます。10個の設計候補を検証するのと同じ構成が、より大規模なキャンペーンにも拡張できます。パラメータ値を変更するだけです。ml.g4dn（最も低コストなT4 GPU）からml.g6e（NVIDIA L40S GPU）までのインスタンスタイプにより、スループットと予算に合わせた柔軟性が得られます。

BoltzGenが提供するもの

BoltzGenは、幅広い生体分子標的に結合できるタンパク質やペプチドを設計するための全原子生成モデルです。このモデルは拡散プロセスを通じてバインダーのバックボーン構造を生成します。次に、BoltzIFと呼ばれる逆フォールディングモデルを使用してアミノ酸配列を生成し、Boltz2フォールディング予測で構造を検証します。

この実装は、GitHubのBoltzgen on SageMakerリポジトリで入手できます。リポジトリには、セットアップスクリプト、両方の実行モード、マルチGPUおよびマルチインスタンススケーリングのサポート、クイックスタートガイドが含まれています。

2つの実行モード

実装では、研究のさまざまな段階に最適化された2つの実行モードを提供しており、迅速な実験から始めて、ニーズの成長に応じて本番ワークフローに移行できます。

SageMaker AI処理ジョブは、迅速な実験のための直接的なバッチ実行を提供します。ジョブを送信すると、SageMaker AIがGPUインスタンスをプロビジョニングし、コンテナを実行し、完了時にインスタンスをシャットダウンします。単一のPythonスクリプトがワークフロー全体をトリガーし、マルチステップのオーケストレーションは不要です。スケールアップ前に設計仕様をテストすることが目標である初期段階の作業では、処理ジョブによりセットアップ時間が最小限に抑えられます。

SageMaker AI Pipelinesは、ステップレベルのキャッシュ、自動スケーリング、SageMaker AI Studioでの視覚的モニタリングを備えた5ステップのオーケストレーションワークフローを提供します。本番ワークロードは、高価な計算ステップを再実行することなくパラメータの反復を実行できるというメリットがあります。各ステップ（設計、逆フォールディング、フォールディング、分析、フィルタリング）は独立してスケーリングでき、入力が変更されていない場合はキャッシュされた結果が再利用されます。

両方のモードは同じワークフローに従います。Amazon ECRのコンテナはプロビジョニングされたGPUインスタンス上で実行され、BoltzGenが設計仕様を処理し、SageMaker AIが結果をAmazon S3に保存してからコンピュートリソースをシャットダウンします。

はじめに

始めるには、AWSアカウント、ローカルにインストールされたいくつかのツール、およびデータ用のAmazon S3バケットが必要です。次のセクションでは、各前提条件を説明し、最初のタンパク質設計実験をガイドします。

前提条件

有効なAWSアカウント（課金が有効）と、インストールおよび設定済みのAWS CLI（aws configure）があることを確認します。ターゲットAWSリージョン（us-east-1など）のGPUインスタンス（ml.g4dn.xlarge以上）に対して十分なサービス割り当て（このチュートリアルでは1など）があることを確認します。

次に、SageMaker AIのAWS Identity and Access Management（IAM）実行ロールを作成します。AmazonSageMakerFullAccess管理ポリシーをアタッチし、データバケットへのAmazon S3読み取り/書き込みアクセス権を付与し、BoltzGenコンテナイメージのAmazon ECRプル権限を付与します。

ローカルマシンに、pipを備えたPython 3.11以上、最新のBoto3とSageMaker SDK（pip install boto3 sagemaker）、およびコンテナイメージをビルドするためのDockerをインストールします。最後に、ターゲットリージョンにAmazon S3バケットを作成します：

aws s3 mb s3://amzn-s3-demo-bucket --region

最初のタンパク質設計を実行する

最も合理化されたセットアップ体験には、Amazon SageMaker StudioまたはAmazon SageMakerノートブックインスタンスをお勧めします。これらの環境には、AWS認証情報、SageMaker SDK、Dockerサポートが事前構成されています。

SageMaker AI Studioを起動

まずSageMaker AI Studioを開き、セットアップと実験のコマンドを実行します。

SageMaker AIコンソールを開きます。

左側のナビゲーションからStudioを選択します。

新しいStudioドメインを作成するか（存在しない場合）、既存のものを開き、AWS CLIを使用してDockerアクセスを有効にします（ローカルモード）。

JupyterLabスペースをインスタンス（例：ml.m5.2xlarge）で起動します。

ml.g4dn.xlargeおよびml.g5.xlarge処理ジョブ使用の割り当て制限を0より大きく引き上げます。

リポジトリをクローン

BoltzGen on SageMaker AIのコードをダウンロードします。これにはセットアップスクリプト、両方の実行モード、サンプル設計仕様が含まれています。

git clone https://github.com/aws-samples/sample-biofm-quickstart.git cd sample-biofm-quickstart/models/boltzgen/inference/amazon-sagemaker/boltzgen-on-sagemaker-processing-job-cli

コンテナイメージをビルドしてプッシュ

コンテナイメージはBoltzGenとその依存関係、GPUドライバをパッケージ化します。一度ビルドして実験全体で再利用します。リージョンを設定し、ビルドスクリプトを実行します：

export AWS_REGION=./sagemaker/build_and_push.sh

AWS認証情報を設定

SageMaker AIはコンピュートリソースをプロビジョニングし、結果をAmazon S3に保存するためにAWSアカウントの詳細を必要とします。環境テンプレートをコピーし、設定を入力します：

cp sagemaker/pipeline/.env.example sagemaker/pipeline/.env vim sagemaker/pipeline/.env

.envの例：

AWS_REGION=us-east-1 AWS_S3_BUCKET=amzn-s3-demo-bucket AWS_ROLE_ARN=arn:aws:iam:::role/SageMakerExecutionRole AWS_IMAGE_URI=.dkr.ecr..amazonaws.com/boltzgen-sagemaker:latest

これで、2つのオプションで最初の設計ジョブを実行する準備が整いました：

オプションA：単一の処理ジョブですべてのステップを迅速に実行：

python sagemaker/run_processing_job.py \ --design-spec example/vanilla_protein/1g13prot.yaml \ --s3-bucket amzn-s3-demo-bucket \ --instance-type ml.g4dn.xlarge \ --num-designs 10 \ --budget 2 \ --wait

拡散モデルによって生成される中間設計候補は10個で、これらは逆フォールディング、再フォールディング、分析、フィルタリングを経る生のバックボーン構造です。実際には、この数は10,000から60,000と大幅に高くする必要があります。最終的な多様性最適化出力セットは2つの設計です。つまり、10の中間設計のうち、品質と配列多様性の両方で最適化された2つだけが出力フォルダに出力されます。

オプションB：異なるステップをパイプラインでオーケストレーションする、より本番向けのワークフロー。

パイプラインディレクトリに移動：

cd sagemaker/pipeline

パイプラインを作成：

python run_pipeline.py --config pipeline_config.yaml create

設計仕様で実行：

python run_pipeline.py --config pipeline_config.yaml run

ステータスを確認（ARNを実行ARNに置き換え）：

python run_pipeline.py --region us-east-1 status --execution-arn

結果をダウンロード

処理ジョブの場合：

aws s3 sync s3://amzn-s3-demo-bucket/boltzgen/output/boltzgen-TIMESTAMP ./results

パイプラインの場合：

aws s3 sync s3://amzn-s3-demo-bucket/boltzgen-pipeline/output/TIMESTAMP ./results

出力ディレクトリにはいくつかのサブディレクトリが含まれています。final_ranked_designs/には、予測結合親和性でランク付けされたトップ設計が保存されます。intermediate_designs/には生成された候補構造が含まれます。metrics/には、二乗平均平方根偏差（RMSD）や接触スコアなどの構造解析データが保存されます。job_metadata.jsonファイルには、完全なジョブパラメータとステータスが記録されています。

次のディレクトリツリーは、--num-designs 100 --budget 10で1G13ターゲットに対して実行した出力構造を示しています：

results/1g13prot/ ├── final_ranked_designs/ │ ├── final_10_designs/ # 多様性と品質で選択されたトップ設計（.cif） │ ├── intermediate_ranked_10_designs/ # 品質スコアでランク付けされたトップ10設計（.cif） │ ├── all_designs_metrics.csv # 全設計の完全なメトリクステーブル │ ├── final_designs_metrics_10.csv # 最終選択設計のメトリクス │ └── results_overview.pdf # フィルタリング基準とプロットを含む要約レポート ├── intermediate_designs/ # 生成されたすべてのバックボーン構造（.cif） ├── intermediate_designs_inverse_folded/ # 配列設計と再フォールディング後の構造 └── config/ # 再現性のためのステップ設定

BoltzGenは自動的に生成されたすべての設計をフィルタリングしランク付けします。各設計は、ランキング段階に入る前に、再フォールディングRMSDが2.5オングストローム未満、バランスの取れたアミノ酸組成などの構造品質しきい値を通過する必要があります。この実行では、100設計のうち3設計がすべてのフィルターを通過しました。その後、パイプラインは複合品質スコアで設計をランク付けし、多様性最適化を適用して最終セットを選択します。

次の表は、主要メトリクスとともに上位5つの設計を示しています。design_ptmおよびdesign_iptmスコアが高いほど、構造的信頼性と結合強度が高いことを示します。filter_rmsdが低いほど、再フォールディング後の構造精度が高いことを示し、delta_sasaが高いほど結合界面が大きいことを示します。

id Rank design_ptm design_iptm filter_rmsd delta_sasa 1g13prot_24 1 0.760 0.304 2.24 584.7 1g13prot_37 2 0.732 0.275 1.98 657.4 1g13prot_21 3 0.748 0.159 2.49 535.7 1g13prot_06 4 0.736 0.426 1.34 1598.0 1g13prot_59 5 0.763 0.527 1.05 2369.0

各最終設計は、ランクプレフィックス（例：rank001_1g13prot_24.cif）付きの.cif構造ファイルとして保存され、ダウンストリーム分析または実験的検証の準備が整います。

最初の実験が完了したので、次のセクションではBoltzGen on SageMaker AIが内部でどのように機能するかを探り、アーキテクチャを最適化できるようにします。