UM-Bridge: Kubernetes を活用してクラウドでのスケーラブルな不確実性の定量化を実現
Google Cloud Japan Team
※この投稿は米国時間 2024 年 2 月 7 日に、Google Cloud blog に投稿されたものの抄訳です。
シミュレーションと不確実性
数値シミュレーションは、科学や工学の多くの分野でイノベーションを推進するツールとして確立されており、数学モデルを使って実世界のプロセスを予測するのに役立っています。たとえば、エンジニアは CAD ソフトウェアを使用して機械的シミュレーションの結果を得ることができますし、医療画像処理では推論が重要な役割を果たしています。このようなシナリオでは、データの不確実性に直面するのが一般的であり、モデリングへの決定論的手法だけでは、研究対象のシステムの複雑さを完全に捉えきれないことがあります。
不確実性は、測定誤差、不完全な情報、特定のプロセスに固有の確率的性質など、さまざまな要因によって生じる可能性があります。物理プロセスをシミュレートする際、不確実なデータは、紛れもなく不確実な予測や推論を意味します。不確実性を踏まえた分析を行うことで、研究者はモデルの限界を理解し、異なる結果に関連するリスクを定量化できます。特定のシナリオが起こる可能性を理解することは重要です。その例として、放射性廃棄物処分の評価、津波の早期警報システム、航空機の翼の設計時に重大な欠陥が生じうる確率の考慮などが挙げられます。
複雑な問題、複雑なソフトウェア
したがって、不確実性の定量化（UQ）には、数学的側面と技術的側面のそれぞれに多くの課題が存在します。非常に単純な UQ の問題であっても、膨大な計算量を必要とする場合があります。そのため、比較的少ないシミュレーション実行回数で良好な結果が得られる効率的な UQ アルゴリズム、最先端の数値ソルバー、ハイ パフォーマンス コンピューティング（HPC）リソースを、すべて 1 つのアプリケーションにまとめる必要があります。
さらに、UQ コミュニティと数値シミュレーション コミュニティでは根本的に要件が異なるため、まったく違うツールを使用する傾向があります。その結果、プログラミング言語、ビルドシステム、並列化スキームなどには互換性がありません。また、大規模なシミュレーション ソフトウェアを拡張して不確実性を組み込む場合、通常モノリシックに行われるため、非常に複雑なソフトウェアになります。たとえば、一般的な気候モデルのコードは、50 万行をゆうに超える場合があります。つまり、高度な UQ 手法を高度なモデルと組み合わせることができるのは、非常に大規模で学際的なチームのみということになってしまいます。
では、全体的な複雑さを管理しやすくし、UQ 手法をより身近なものにするにはどうすればよいのでしょうか。
UM-Bridge: UQ の複雑さを解消
私たちは、あらゆるプログラミング言語やフレームワークからアクセスできる数値モデルの統合インターフェースとして、UM-Bridge1（UQ and Modeling Bridge）を開発し、数学モデルのサービスベースの抽象化を導入しました。UM-Bridge では、UQ とモデルコードは 2 つの独立したアプリケーションであり、HTTP ベースのプロトコルを使用してモデル評価リクエストを渡します。多くの UQ アルゴリズムでは、モデルを、パラメータ ベクトルをモデルの結果ベクトルにマッピングする数学関数として扱うため（場合によっては導関数を伴う）、このようなアーキテクチャが可能になります（これは、数学モデルで動作する多くの最適化アルゴリズムや ML アルゴリズムにも当てはまります）。
多くの言語やフレームワークに UM-Bridge がネイティブに統合されているため、それぞれの言語で通常の関数呼び出しを行うように簡単にモデルを呼び出すことができます。サービスベースのアーキテクチャと統合インターフェースにより、UM-Bridge クライアントは、それぞれのプログラミング言語やフレームワークに関係なく、任意の UM-Bridge サポートモデルに接続できます。
UQ のエコシステムとクラウドをつなぐ
UQ アルゴリズム自体の計算費用は低くなる傾向があります。一般的に、費用を増加させるのは多くの数値モデル評価です。そのため、並列化された UQ アルゴリズムを単一のマシンで実行し、負荷の高いモデルの計算を UM-Bridge でクラウド クラスタにオフロードすることができます。クラスタ側のロードバランサは、これらの並列モデル評価リクエストを受け取り、（潜在的に多数の）独立したモデル インスタンス間で分散します。
すでに多くの UQ パッケージが単一ノードでの並列処理をサポートしているものの、マルチノードはサポートしていません。マルチノードをサポートしているものは通常、MPI や Ray などの特定のテクノロジーに限定されているため、任意のモデルコードとの互換性は制限されることになります。以下に示す新しいアーキテクチャは、これらの UQ コードのいずれであってもモデルの計算をクラスタに透過的にオフロードできるため、特に魅力的といえるでしょう。
私たちは、UM-Bridge フレームワークを実装するための科学計算プラットフォームとして、Google Cloud 上の Google Kubernetes Engine を選択しました。ワークロードをオーケストレートするためのコンテナ化された環境や、タスク分散のための高度なロード バランシング機能など、私たちの要件に必要な機能を提供してくれるためです。
ロードバランサには、一度に 1 つのリクエストだけが 1 つのモデル インスタンスに渡されるように構成された HAProxy を採用しています。各モデル インスタンス自体は、MPI を介して複数のノード（つまりコンテナ）にまたがって並列化できますが、これは Kubeflow の mpi-operator を使用する特別なケースとしてサポートしています。最後に、一部の GKE 固有の最適化は、コンパクトに配置されたクラスタで実行されるため、ノード間のレイテンシが短縮されます。加えて、一貫したパフォーマンスを実現するために同時マルチスレッディング（SMT）を無効にしています。
私たちが行った設定は、再現可能な構成として公開しています。これは、任意の UM-Bridge モデルコンテナをサポートするあらゆる Kubernetes クラスタに応用できます。UM-Bridge は UQ パッケージとクラウド クラスタの橋渡しを行い、UQ とモデル間のインテグレーションの問題を解決するだけでなく、専用の HPC UQ コードも不要にします。
応用
私たちは、この Kubernetes の設定を実行して UQ の多くの問題を解決することに成功し、そのうちの 3 つの設定を公開2しています。クライアント側は、MATLAB の「parfor」ループによって簡単に並列化されたスパース グリッド コードから、Ray で並列化された Python UQ コードまで幅広く対応しています。同様に、Fortran やコード生成された C++ など、いくつかの異なるツールを使用してモデルを構築しています。UM-Bridge と Kubernetes のおかげで、簡単にツールを組み合わせることができました。
この手法の最大の応用例が、利用可能なブイのデータに基づいた、2011 年東北地方太平洋沖地震の津波の震源に対するベイズ推定です。私たちは、高速近似モデルによって高速化された階層的 UQ 手法を使用しました。それでも、さまざまな震源で 3,000 回以上のシミュレーションを行い、その結果を測定データと比較する必要がありました。そこで、100 個の c2d-highcpu-56 VM を使用して、負荷の高いシミュレーションの実行を 5,600 個の vCPU を擁する GKE クラスタにオフロードしました。この応用ケースでは、1 つのノードに割り当てられた 56 個の vCPU をすべて使用して、各シミュレーションと並行して 100 件のシミュレーションを実行しました。
MPI 並列 UQ コードを使用し、ベアメタル HPC クラスタ3でモデルを実行する従来のアプローチと比較して、開発期間を数か月からわずか数日に短縮することができました。UM-Bridge、コンテナ化、Kubernetes によって、以前直面していた技術的な問題が解決されたため、私たちは UQ の問題に完全に集中できるようになりました。
まとめ
UM-Bridge は、マイクロサービスを活用するアーキテクチャを導入することで、複雑な応用ケースにおいても UQ を可能にします。大規模な GKE Standard クラスタで、私たちはそのスケーラビリティを実証することができました。コンテナの完全なサポートと私たちが公開している汎用的な Kubernetes の設定により、Kubernetes が大規模な UQ を実行するうえでの魅力的な選択肢となることは間違いありません。
ドキュメントやチュートリアル、Kubernetes の設定については、https://um-bridge-benchmarks.readthedocs.io をご覧ください。ここには、UM-Bridge ベンチマーク ライブラリの一部として、上述の津波モデルも用意されています。
UM-Bridge および UQ の応用に関するコミュニティを積極的に構築していますので、お気軽にお問い合わせください。現在の決定論的数値シミュレーションのように、UQ をユビキタスなものにしていきましょう。
1.UM-Bridge
2. Scaling up on GKE
3. Large-scale UQ on conventional HPC cluster
-ダラム大学、助教授 Anne Reinarz 博士
-ハイデルベルク大学、博士研究員 Linus Seelinger 博士
