Thermo Fisher ScientificがAWS ParallelClusterを使ってクライオ電子顕微鏡を高速化した方法

Thermo Fisher Scientific は、計算機によるクライオ電子顕微鏡（Cryogenic electron microscopy, Cryo-EM）の画像や動画の処理のワークロードを含む、科学サービスを提供する世界的なリーダーです。科学者はクライオ電子顕微鏡を、生体分子の3次元構造を原子レベルに近い分解能で決定するために使用し、Thermo Fisher の装置を使用してテラバイトの画像と動画が生成されるプロセスで利用します。これらのデータセットを繰り返し処理することによって、結果を得るためのトータルコスト、費やされる時間、および運用上のオーバーヘッドを削減することは、医薬品化学の結果のスピードと質を向上させるために重要な技術革新となります。

このブログでは、AWS ParallelCluster、Amazon FSx for Lustre、およびcryoSPARC（Structura Biotechnology社製）を使用したCryo-EM ベンチマークを成功させるためのプロセスを紹介し、その過程における設計を説明します。

課題

Cryo-EM は2017年のノーベル化学賞を受賞しており、SARS-COV-2 のスパイクタンパク質を含む重要な創薬ターゲットの価値の高い3D 構造の生成にも利用されました。Cryo-EM 透過型電子顕微鏡は、タンパク質のサンプルをガラス質の氷へ瞬間冷凍し、ミクロの構造をネイティブに近い状態で3D デジタル表現するものです。この装置は、1 サンプルあたり数テラバイトのデータを生成し、データの解像度は新しい装置の世代が進むごとに向上しています。このデータは、ハイパフォーマンスコンピューティング（HPC）リソースを使用して処理する必要があり、1 つまたは複数のGPU で加速された複雑なパイプラインを使用します。またこれは、知識のある科学者が実施する必要があるインタラクティブな作業であり、結果には3D の可視化も必要です。

AWS Parallel Cluster を利用した複雑な計算要件の簡素化

Cryo-EM のHPC ニーズは、オンプレミスでは解決が困難です。これは、Cryo-EM の解析にはコンピューティングリソースに対する柔軟なアクセスが必要だからです。処理パイプラインの仕様によっては、各計算ステージの速度をスケールさせるために、それぞれのステージにおいて適切なGPU数が必要となりますが、これを予測することは困難です。この問題に対して、オンプレミスデータセンターの老朽化した高価なハードウェアに依存するのではなく、AWS 上のCryo-EM のワークロードでは、Amazon Elastic Compute Cloud（EC2）リソースの弾力性を利用することができます。

AWS ParallelCluster は、オープンソースのクラスタ管理ツールで、infrastructure as code (IaC) を使用して、ワークフローの各ステップのニーズに合わせて設定したAmazon EC2 インスタンスのクラスタをプロビジョニングします。科学者は、ワーカーノードあたりのGPU の数や種類をベンチマークして、コストと性能の最適なバランスを見つけるために実験することができます。ParallelCluster の利用には追加費用がかかりません。フレームワークによってプロビジョニングされたリソースに対してのみ支払いが発生します。必要な分だけリソースを増減できるので、従来のオンプレミス型クラスタが抱えていた、処理のピーク時におけるリソース競合の悩みが解消されます。

FSx for Lustre による高パフォーマンスなコラボレーションとコスト効率に優れたファイルストレージ

計算の柔軟性に加えて、科学者にはデータセットや結果を地理的に離れた同僚と共有する必要性が高まっています。帯域幅が限られたネットワーク上で大きなデータセットを転送すると、数日から数週間の遅延が発生し、転送後のデータへのアクセス者を制御することができません。Amazon Simple Storage Service (Amazon S3) にデータを保存すれば、数回のクリックでグローバルな共有が可能になり、ビジネスとサイエンスの側面におけるマネジメント、セキュリティ、ガバナンスの要件に合わせたきめ細かいアクセス制御が可能になります。

GPU ノードのパフォーマンスを可能な限り高速化するには、GPU に供給し続けるための高速キャッシングファイルストレージが重要となります。Cryo-EM 解析では、大きなパーティクルスタックが GPU  のメモリに出入りする必要があるため、シーケンシャル I/Oパフォーマンスが特に重要です。Amazon FSx for Lustre は、HPCワークロードに最適化された高性能な並列ファイルシステムを提供するフルマネージドなAWS サービスで、ParallelCluster とAmazon S3 の両方と統合されています。このS3とFSx for Lustre のシームレスな組み合わせを使用することで、ユーザーはコストと性能の両面で最高のものを手に入れることができます。FSx for Lustre はすべての計算ノードで共有されるため、ローカルストレージを搭載した計算ノードを必要とせず、さらにコストを削減することができます。つまりFSx for Lustre は、一時的なCryo-EM ファイルへの迅速な読み書きのアクセスにより、Cryo-EM の再構成を高速化するためのコスト効率と性能に優れた共有ドライブなのです。

CryoSPARC を用いたベンチマーク

CryoSPARC は、Structura Biotechnology Inc. によって開発された、Cryo-EM データ処理のための完全なソリューションで、研究および創薬の分野で使用されています。生データを高速に処理し、3次元再構成を行うことにより、計算量を小さく抑えることができます。柔軟性の問題に対処する独自のアルゴリズムが特徴で、膜タンパク質のような治療に関連するターゲットに対して特に優れた性能を発揮します。Thermo Fisher 社は、Cryo-EM の処理環境とインフラをセットアップするポストインストールスクリプトの実行を含む設定ファイルを使用して、AWS 上でのCryoSPARC の展開を自動化しました。コマンドラインオプション1つで、ハードウェアとソフトウェアがCryo-EM の処理に対応できるようになります。

ソリューション

このブログ記事では、技術的な考慮事項のウォークスルーとソリューションの構成要素の背後にある「理由」を説明します。追加の前提条件、クリーンアップなど、自分のアカウントでこのソリューションを再現するためのステップバイステップの手順については、AWS Samples GitHub  リポジトリを参照してください。

図1 – Architecture diagram of cryoSPARC deployed on AWS ParallelCluster

ネットワーク、セキュリティ、およびコンピュートアベイラビリティーの前提条件

デフォルトVPC では、パブリックとプライベートのサブネットが複数のAvailability Zone (AZ) にまたがってバランスよく配置されるのが一般的な使い方です。しかし、ParallelCluster のようなHPC クラスタは通常、通信のレイテンシーを低く抑えクラスタプレイスメントグループを使用できるように単一のAZ利用が好まれます。また計算ノード群のために、比較的多くのIPアドレスを持つ大規模なプライベートサブネットを作成することができます。そして、ヘッドノードを動かすための、最小限のIPアドレスレンジのパブリックサブネットを作成します。

P4dインスタンスファミリーのようなHPC EC2インスタンスは、すべてのAZで利用できるわけではありません。つまり、あるリージョン内のどのAZ に、必要なコンピュートファミリーが揃っているかを判断する必要があります。これは、AWS CLI のdescribe-instance-type-offerings コマンドを実行することで調べることができます。これを行う最も簡単な方法は、CloudShellを利用することです。このサービスは数分でAWS CLI コマンドを発行できるシェル環境を提供します。CloudShell 環境が用意されたら、テキストをコピーしてシェルにペーストし、括弧内のプレースホルダーに希望のリージョンを指定します。

aws ec2 describe-instance-type-offerings \
--location-type availability-zone \
--region <region> \
--filters Name=instance-type,Values=p4d.24xlarge \
--query "InstanceTypeOfferings[*].Location" \
--output text

入力パラメータに記述したインスタンスがどのAZ にあるかが出力されます。

また、クラスタの準備ができたらヘッドノードにSSH アクセスできるように、ParallelCluster の設定の入力パラメータとしてEC2 のSSH キーペアが必要です。

データ転送

実験機器からAmazon S3 にデータを移動するための適切なデータ転送の仕組みは、研究所の接続性と転送するデータ量に依存します。まず最初にお勧めするのは、オンプレミスからクラウドへの安全なデータ転送を最小限の開発工数で簡単に自動化できるAWS DataSync です。AWS Storage Gateway, Amazon S3 File Gateway は、特に研究所の接続性が限られている場合や、オンプレミスから転送されたデータセットへの継続的な双方向アクセスが必要な場合、もう一つの有効な選択肢となります。DataSyncとStorage Gateway の両方は、HPC 以外のビジネスクリティカルなニーズを保護するために帯域幅を制限することができます。

またこの他にも、AWS CLI を使用して個々のファイルを転送したり、パートナーソリューションを使用して素早く開始することもできます。

AWS Identity and Access Management (IAM)  の権限

ParallelCluster は、デフォルトで独自の最小特権のロールとポリシーを作成しますが、多くのエンタープライズ企業は AWS アカウントユーザーの IAM アクションへのアクセスを制限しています。ParallelCluster は、事前に作成されたIAMリソースの使用または追加もサポートしているため、IT 管理者に事前に作成されるようにリクエストして利用することができます。その際に必要な権限とロールは、ParallelCluster のドキュメントに記載されているので、すぐに使い始めることができます。

FSx ファイルシステムは、永続的またはスクラッチとしてプロビジョニングすることができます。永続的ファイルシステムは、Amazon S3 にデータを自動的にエクスポートして戻すことができますが、スクラッチファイルシステムはそうではありません。このGitHub リポジトリにあるサンプル実装では、本番環境ではなくベンチマーク環境をプロビジョニングしているため、スクラッチファイルシステムを使用しています。ParallelCluster のジョブスケジューラにデータエクスポートタスクを組み込んで、ジョブが完了するたびにバックグラウンドで透過的にデータエクスポートを実行したい場合は、ヘッドノードのインスタンスプロファイルにIAM Policy ステートメントを追加する必要があります。このポリシーは、この GitHub リポジトリの FSxLustreDataRepoTasksPolicy.yml ファイルにあります。エクスポートを実行する場合は、ParallelCluster のプロビジョニングを実行するために使用しているロールに、このポリシーが含まれていることを確認してください。

コンピュートクラスタの構成とエンドユーザーアクセス

このブログで紹介したアーキテクチャでは、複数のEC2 コンピュートファミリーで繰り返しテストを実行することが可能で、ワークフローをsmall, medium, large のキューに分割することで、各計算ステージでコストとパフォーマンスを最適化した構成を作成することができるようになっています。図1を参照すると、Queue1 にはg4dn GPU ノード、Queue2 にはp4dn ノード、Queue3 にはユーティリティコンピュートに最適化されたc5n ノードが含まれています。

CryoSPARC はヘッドノード上でWeb アプリケーションとMongoDB データベースを稼働させており、ヘッドノード上のリバースSSH トンネル、またはAWS Systems Manager Session Manager を使ってアクセスすることができます。ウェブアプリケーションでは、ワークフロー中にパーティクルをインタラクティブに選択したり、処理中（および処理後）のビジュアライゼーションにアクセスしたりすることができます。

ヘッドノードには、cryoSPARC アプリケーションをサポートするのに十分な大きさを持ち、コストを最小限に抑えるインスタンスを選択する必要があります。この例ではc5.4xlarge を使用していますが、使用パターンに合わせて他のインスタンスタイプを選択することができます。

インストールとストレージに関する考察

ParallelCluster のポストインストールスクリプトは、提供されたライセンスキーでcryoSPARC のインストールを実行し、1-500GB の範囲の共有ボリュームにインストールします。私たちのユースケースでは、100GB で十分であることがわかりました。AWS ParallelCluster は、EBS ボリュームをヘッドノードにアタッチし、それをNFS マウントとしてコンピュートノードに公開します。クラスタが繰り返しデプロイされる場合は、AWS ParallelCluster のEC2 Image Builder サポートを利用して、ソフトウェアスタックがすでにインストールされた独自のカスタムAMI を作成することができます。これにより、ポストインストールアクションを繰り返さないようにし、プロビジョニング時間を短縮することができます。

また、一時的なライフスパンのベンチマークワークロード用に12TB のAmazon FSx for Lustre のスクラッチファイルシステムをスピンアップしています。

実行

このブログ記事で説明したParallelCluster のセットアップを再現するには、前に説明したネットワークとセキュリティの前提条件を使用して、サブネットID、EC2 キーペア名、cryoSPARC ライセンス、S3 バケット名で構成ファイルのプレースホルダを更新してください。その後、README ファイルに記載されているデプロイメントの指示に従います。

クラスタのプロビジョニングが完了したら、Structura のドキュメントに記載されているように、cryoSPARC のジョブを実行する準備が整った状態です。ベンチマーク結果は、Thermo Fisher 社のホワイトペーパー、Cryo-EM processing at the pace of medicinal chemistry on AWSで公開されています。チームの成果についてのライブトークを聞くには、YouTube上にある Lab Roots Webinar: Keeping up with the chemists – Cloud processing for pharma cryo-EMをご覧ください。

まとめ

Thermo Fisher 社は、AWS ParallelCluster を使った再現性がありセキュアで柔軟があり、繰り返しデプロイ可能なソリューションにより、公開データセットの処理時間をオンプレミスの4週間から4日間に短縮しました。ベンチマーク実行にかかる総費用は、オンプレミス上の場合はハードウェアを何十万ドルもかけて購入し継続的な管理することが必要ですが、AWS上では800ドル以下でした。最終的に、AWS 上で利用可能な最新のハードウェアを使用できたため、得られた粒子画像の解像度が向上しました。

提供されたリソースを使用して独自のデータセットをベンチマークし、AWS における柔軟性、パフォーマンス、およびコストの最適化を活用して、研究を加速してください。またAWS アカウントでの意図しない課金を避けるために、GitHub のREADME ファイルにある手順を使用してクラスタをクリーンアップすることを忘れないでください。

著者について

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