Snorkel AI が Amazon EKS を使用して機械学習ワークロードをスケーリングすることで 40% 以上のコスト削減を達成した方法

機械学習 (ML) のスタートアップは、ハイエンド GPU を使用して大規模モデルをトレーニングし、推論用に大規模にデプロイするため、多くの場合、コンピューティングのヘビーユーザーです。AWS スタートアップは創業から IPO までスタートアップと提携し、何千人もの創設者や人工知能 (AI) イノベーターが Amazon Elastic Kubernetes Service (Amazon EKS) でビジネスを構築するのを支援してきました。Amazon EKS は、Kubernetes の柔軟性と、可用性の高いコンテナ化されたワークロードの構築に最適化された AWS マネージドサービスとしてのセキュリティと耐障害性を備えているため、ML モデルの構築とホストにおける人気の選択肢となっています。

Snorkel AI は、Amazon EKS の恩恵を受けている企業の 1 つです。Snorkel AI は、Fortune 500 企業、連邦政府機関、AI イノベーターが、基盤モデル (FM) や大規模言語モデル (LLM) を構築、適応、抽出して、ドメイン固有のデータセット上で高精度で実行できるようにしています。AI 開発に向けた Snorkel のデータ中心型アプローチにより、組織は保険金請求処理、資金の拡散、臨床試験分析、海洋掘削の積極的な井戸管理の促進などのユースケース向けに、本番環境対応の AI サービスを構築してきました。

過去数か月にわたり、Snorkel チームは、インフラストラクチャの費用を増やしたり、開発者の作業速度を低下させたり、ユーザーエクスペリエンスを阻害したりすることなく、ML 開発ワークロードをサポートする効率的なインフラストラクチャを設計するという固有の課題に熱心に取り組んできました。チームの最終的な目標は、エンドツーエンドの ML プラットフォームである Snorkel Flow のクラスターコンピューティング支出を 40% 以上削減することでした。

Snorkel Flow の概要

Snorkel の Snorkel Flow AI データ開発プラットフォームにより、データチームはプログラムによるラベリング、迅速なモデルトレーニング、およびエラー分析の反復ループを採用することで、AI ア プリケーションを迅速に構築できます。各プロジェクトは、少数のラベリング関数をユーザーが作成すると始まります。

ラベリング関数は、シンプルなヒューリスティック、外部データベース、レガシーモデル、さらには大規模な言語モデルの呼び出しを使用して、エンコードされた専門的直感に基づいて、ラベル付けされていない大量のデータにラベルを適用します。プラットフォームの弱い監視アルゴリズムは、これらのルールベースの関数を組み合わせて、各レコードで最も可能性の高いラベルを決定します。次に、ユーザーはこれらの確率的データポイントに基づいて簡単なモデルをトレーニングし、各ラベリング関数の影響を評価します。分析段階では、ユーザーはモデルのパフォーマンスが低いデータの一部を調査します。その後、ラベリング関数を構築または変更し、別のクイックモデルをトレーニングして、ループを続けます。ユーザーがラベルの品質に満足したら、ロジスティック回帰から FM まで、モデル動物園のアーキテクチャに基づいて最終モデルを構築し、デプロイ用にエクスポートします。

このワークフローの性質上、Snorkel Flow のインフラストラクチャでは、さまざまな期間にコンピューティング使用率が高くなります。Snorkel Flow の顧客ベースと ML 製品の機能が拡大するにつれて、運用コストは自然に増加しました。Snorkel は、効率的な成長を実現するために、最先端の ML ソフトウェアを運用しながら利益を拡大する方法を理解することを目指しました。Snorkel は以下のプラクティスを実装し、クラスターコンピューティングコストを 40% 以上削減しました。

クラウドコスト最適化のためのソリューション 

Software-as-a-Service (SaaS) のスタートアップは、多くの場合、クラウド支出を最適化する機会があります。こうしたコストを左右する独自の要因を理解することが不可欠です。

Snorkel には、2 つの重要な要因がありました。

1. ML 開発ワークロードには、多くの場合、GPU などの特殊で高価なハードウェアが必要です。通常、これらのワークロードは本質的に「バースト性」です。
2. フォーチュン 500 の企業や大規模な連邦機関は、コンテナ化されたプラットフォームを使用して Snorkel を使用しています。それらの機関には、特定のデプロイ要件やデータプライバシー要件を持つ高度な IT 部門がある大手金融機関も含まれます。

Snorkel のチームは、インフラストラクチャのコストを直線的に増加させることなく効率的なスケーリングを可能にするシステムの開発に取り組んでいます。その結果、Snorkel はクラウド費用の問題に対処するために、Amazon EKS での ML ワークロードに合わせた包括的な自動スケーリングソリューションを開発しました。このソリューションは、バーストコンピューティングを必要とするワークロードを促進しただけでなく、コスト削減の目標も達成しました。

自動スケーリングソリューションに加えて、クラウド費用の 40% 以上の削減に貢献した主な戦略には以下が含まれます。

• エンジニアリングリーダーや AWS チームと協力して、クラウド構成の最適化を通じて Savings Plans を採用する。
• Prometheus でノード使用率を監視し、バックエンドエンジニアと相談してプラットフォームコンポーネントのニーズを測定することで、リソースを適切にサイジングする。
• Amazon EKS でコスト効率の高い仮想マシン (VM) タイプに切り替え、マルチ GPU の Amazon Elastic Compute Cloud (Amazon EC2) インスタンスを利用してコストパフォーマンスを向上させる。
• エンジニアが顧客対応チームと協力して内部プロセスの変更を実施し、アイドル状態のコンピューティングを最小限に抑える。

この投稿では、Snorkel が、ML システムのより良いインフラストラクチャの設計を促進するために、スケーリングに関するこれらの課題に対処するプロセスを紹介します。Kubernetes を初めて使用する場合は、Snorkel の「Kubernetes 入門」の投稿を読んで、Kubernetes の基本と機械学習について詳しく学んでください。Snorkel の投稿で学んだ知識は、これまでの Kubernetes との道のりについてさらに学ぶためのものです。

AWS での Snorkel Flow の様子

実際には、Snorkel Flow と AWS のやり取りは以下の順序で行われます。Snorkel Flow プラットフォームはコンテナに大きく依存しているため、AWS への移行はほぼシームレスでした。

• ユーザーは、Amazon Route 53 のルールにマップされるウェブブラウザを介して Snorkel Flow インスタンスにアクセスします。
• Route 53 はリクエストを Application Load Balancer に転送します。
• 次に、Application Load Balancer は、共有 EKS クラスター上で実行されている Snorkel Flow ポッドにリクエストを転送します。Snorkel は、コストを最適化するために EC2 インスタンスタイプを m5 から m6a に切り替えました。その結果、計算量を 10% 節約できるようになり、同じ CPU と RAM の 1 時間あたりのコストに基づくパフォーマンスへの影響は無視できるものでした。
• さらに、単一の g4dn.8xlarge GPU インスタンスから g4dn.12xlarge マルチ GPU インスタンスにアップグレードしました。これにより、4 倍の GPU ポッドを提供できるようになりました。
• 各 Snorkel Flow インスタンスは、Amazon Elastic File System (Amazon EFS) ボリュームを使用して、ファイルをディスクに保存します。
• EC2 インスタンスのポッドにあるセルフホストの Redis キューは、受信したジョブを保持し、ワーカーポッドがそれらをピックアップするのを待機します。
• EKS メトリクスは Amazon CloudWatch へプッシュされ、カスタムスクリプトがクラスターパフォーマンスの異常がないかログを監視します。

このアーキテクチャは、Snorkel Flow ユーザーに安定した快適な体験をもたらしました。

スケーリング方法を再検討する

図 1 で説明したアーキテクチャが登場する前は、Snorkel のインフラストラクチャの初期段階では固定リソースを使用していました。Snorkel のユーザーは、このような大量のワークロードは完了するまでに時間がかかりすぎて、エクスペリエンスに悪影響を及ぼす可能性があるという意見を共有していました。

コンピューティングリソースを手動でスケーリングすると、スケーリングできず、エラーが発生しやすくなり、使用率が低いときでもクラウドのコストは上昇し続けました。これは、クラウドのコスト効率が低く、パフォーマンスが必要以上に遅いという両方の面で最悪でした。

これらの課題に対処するために、Snorkel は次のセクションで説明するように、インフラストラクチャの複数のレベルで自動スケーリングを実装しました。

コスト効率を考慮したスケーラブルなインフラストラクチャを設計する

Snorkel Flow の Kubernetes ディストリビューションには、プラットフォームのさまざまなコンポーネントを実行するポッドを含む EKS クラスターで実行される一連のデプロイが含まれます。

図 2 に示すように、バースト性のコンピューティングワークロードを扱う際の固有の課題に対処するために、Snorkel のチームは Kubernetes ポッドの新しい概念を導入しました。それは、ポッドをセマンティックに「固定」または「柔軟」のいずれかに分類することでした。

• 固定ポッドは、ノードからノードへと安全に移動できません。これは、重要なメモリ内の状態 (チェックポイントなしの進行中のコンピューティングジョブなど) を失うという理由で、または基盤となるプラットフォームコンポーネント (例えば、Ray クラスターのオーケストレーター) の回避可能なダウンタイムを最小限に抑えるためです。
• 柔軟ポッドは、新しいノードに安全に移動できます。この区別は自動スケーリングのコンテキストでは意味があります。これは、ノードをダウンスケーリングすると、ポッドが停止したときに、使用率の低いノードからポッドを移動させることがあるからです。

この固定/柔軟のフレームワークにより、Snorkel はドメイン固有の方法でクラスターの自動ダウンスケーリングが可能になり、財務部門から 1 時間ごとに通知を受けることなく、Amazon EKS でクラスター自動スケーラーを有効にできます。

Snorkel の最初のアプローチは、podDisruptionBudgets を EKS クラスターにデプロイして、クラスター自動スケーラーが日中に柔軟ポッドを移動したり、固定ポッドをまったく移動させたりしないようにすることでした。これは効果的ではありましたが、理論的に最適なノード数よりもはるかに少ない数のノードしかダウンスケールしていなかったため、Snorkel チームは満足しませんでした。

これに対処するため、Snorkel は固定ポッドを小さな固定ノードグループに分離するポッドスケジューリング最適化を重点的に導入しました。残りの柔軟ノードグループに柔軟ポッドとワーカーポッド (固定ポッドと見なされるが、ワーカーノードの自動スケーリングによって一時的になっている) をスケジュールしました。

これらの変更により、Snorkel は、柔軟ポッド間を移動してワーカーポッドの大部分をスケールダウンすることが安全になった夜間に柔軟ノードを効率的にダウンスケールすることができました。

クラスターの大半のノード (つまり、柔軟ノード) の効率的なダウンスケーリングを可能にすることで、Snorkel は Snorkel Flow をホストするためのクラウドコストを 40% 以上削減するという目標を達成することができました。

Snorkel の自動スケーリングソリューションの詳細

Snorkel は、前のセクションで説明したソリューションの実装を、次の 3 つの連続した取り組みに分けています。

1. まず、Snorkel は 「ワーカー自動スケーリング」 を実装しました。これは Redis ベースのカスタム自動スケーリングサービスで、ワーカーのキュー内のジョブに基づいてワーカーポッドをスケールアップおよびスケールダウンできます。
2. 次に、Kubernetes クラスター自動スケーラーがノードのスケールアップに加えてノードをスケールダウンできるように Kubernetes のデプロイを再設定することで、「クラスター自動スケーリング」を実装しました。
3. 第 3 に、Snorkel は、固定ポッドが残りのノードのダウンスケーリングを妨げないように、固定ポッドを固定ノードの小さなグループにグループ化することで、「ノードダウンスケーリングの最適化」を実装しました。

ワーカー自動スケーリング

Snorkel Flow プラットフォームは、コンピューティングを抽象化して、ジョブが Redis キューで待機し、ワーカーがワーカーポッド内のプロセスとして実行されるというパラダイムにします。

Snorkel は、Snorkel Flow のバックエンド API で繰り返し関数を実行することにより、ワーカーポッド用のワーカー自動スケーリングソリューション (図 3) を実装しました。この関数は数秒ごとに Kubernetes クラスターと Redis をチェックして、アップスケーリングとダウンスケーリングの両方の適格性を確認します。

1 つまたは複数の関連する Redis ベースのキューで待機中のジョブがある場合、この関数は Kubernetes API に、これらのジョブを処理するための追加のワーカーポッドをプロビジョニングするように要求します。Redis キューが空で、ジョブレジストリに実行中のジョブがない場合、この関数は Kubernetes API に、ワーカーポッドを破棄して予約済み CPU と RAM リソースを解放するように要求します。

図 4 に示すように、このワーカー自動スケーリングの実装が公開されると、Snorkel Flow のワーカーポッドは一時的になり、ジョブを処理する必要がある場合にのみクラスターに表示されるようになりました。

クラスター自動スケーリング

PodDisruptionBudget リソースは、特定の時点で利用できなくなる可能性のあるポッドレプリカの最大数を指定できるようにすることで、特定のポッドを中断 (例えば、自発的な再起動) に対して保護します。図 5 に示すように、デプロイでこの値を明示的に 0 に設定すると、クラスター自動スケーラーがデプロイのポッドを実行しているノードをダウンスケールしなくなります。

ホストされている Snorkel Flow インスタンスにこのリソースを実装することで、クラスター自動スケーラーは、使用率の低いノードを安全にダウンスケールできました。しかし、Snorkel が実現したコスト削減はごくわずかでした。Snorkel Flow のすべてのポッドが関連付けられた podDisruptionBudget によって保護されていたため、ノードの大部分をダウンスケールすることはできませんでした。

詳しく調べてみると、Snorkel のチームは、この保護が常に存在する必要はないことに気付きました。ワークロードはバースト性であるため、Snorkel Flow とのユーザーインタラクションのほとんどはお客様の営業日に行われます。つまり、営業時間外にはこの保護を緩和しても安全です。ワーカー自動スケーリングと同様に、Snorkel は、使用できないポッドレプリカの最大数を 0 から 1 に増やして、インスタンスの柔軟ポッドの podDisruptionBudgets を一晩中「オフ」に切り替える繰り返し関数を実装しました (図 6)。以前のワーカー自動スケーリングソリューションを ClusterAutoscaler および PodDisruptionBudget 機能と組み合わせると、使用率の低いワーカーノードを以前よりも大幅にダウンスケールできました。Snorkel Flow をクラウドにデプロイしているお客様は、必要に応じてこれを設定できます。

ノードダウンスケーリングの最適化

これらの改善があっても、Snorkel は、使用率の低いノードの大部分がまったくダウンスケールされていないことに気付きました。

さらに調査した結果、Snorkel は、問題が同じノードを占有している固定ポッドと柔軟ポッドに起因していることに気付きました。これは問題でした。なぜなら、柔軟ポッドを含むノードに疑似ランダムに割り当てられた固定ポッドは、そのノードが「ピン」され、使用率が低い場合でもダウンスケールされないからです。固定ポッドのスケジューリングを制御できないため、クラスターのノードの大半は、その時点で必要とされていたよりもはるかに高いコンピューティング能力を備えていたにもかかわらず、ダウンスケールできない期間が生じました。

Snorkel は、podAffinities Kubernetes リソースを活用してこれに対処しました。これにより、特定のノードですでに実行されている他のポッドのラベルに基づいて、ポッドが実行できるノードを制限できるようになりました。固定ポッドと柔軟ポッドを区別するためにポッドにラベルを追加し、デプロイ設定に podAntiAffinity スタンザを追加して、柔軟ポッドを実行しているノードで固定ポッドがスケジュールされないようにしました。その逆も同様です。

この podAffinities の実装により、Snorkel AI はノードを 2 つの関数グループに分割することができました。1 つは、ノード間で安全に移動できない固定ポッドを含む固定ノードグループ (例えば、キャッシュが原因で Redis)、もう 1 つは、一時的 (ワーカーポッドなど) か、営業時間外 (夜間など) に移動しても安全な「柔軟」ポッドを含む柔軟なノードグループです。

プラットフォームのメンテナンス中に手動で介入することも可能ですが、Snorkel は、固定ノードを自動的にダウンスケールすることはできません。しかし、このソリューションでは、移動できないポッドを固定ノードに分離しているため、固定ノードを自動的にダウンスケールできます。

まとめ

Snorkel と AWS スタートアップのチームは、この思考プロセスとこれらのソリューションを共有することで、他のスタートアップが ML ワークロードのより良いインフラストラクチャを構築できるようになることを期待しています。ML、大規模言語モデル、その他の FM が世界中の組織の本番環境に導入されるにつれて、その重要性は急速に高まっています。

AWS re:Invent 2023 に参加していますか? このデプロイは、Snorkel のセッション AI の未来を切り拓く: Amazon EKS での生成モデルのデプロイ (セッション CON312) の一環として取り上げられます。 ぜひチェックしてください。

Snorkel のこの技術的ビジョンを実現するために尽力してくれた David Hao 氏、Edmond Liu 氏、Alec Xiang 氏に感謝いたします。また特に、この記事について思慮深いフィードバックを寄せてくれた Matt Casey 氏、Henry Ehrenberg 氏、Anthony Bishopric 氏、そして Snorkel のインフラストラクチャーエンジニアリングチーム全員に感謝いたします。