AKS上での大規模Ray実行

Microsoft の Azure Kubernetes Service (AKS) チームは、Anyscale が提供するマネージド Ray サービスを大規模に実行するためのガイダンスを共有しました。彼らは GPU 容量の制限、散在する ML ストレージ、および認証情報の有効期限切れという 3 つの主要な課題に焦点を当てています。

この投稿では、AKS 上のオープンソース KubeRay に関する以前の概要について補足しています。今回は、以前 RayTurbo と呼ばれていた Anyscale の改善されたランタイム（runtime）を取り上げています。このランタイムは、スマートな自動スケーリング、監視機能の向上、および耐障害性のあるトレーニング機能を備えています。これらすべてはオープンソースの Ray フレームワークに基づいています。

Ray は、単一のラップトップから数千ノードにわたるクラスターまで AI および ML ワークロードをスケールできるように設計された、Python ネイティブの分散コンピューティングフレームワークです。Anyscale のマネージドプラットフォームは、本番環境での使用に向けた機能を Ray に追加しています。新しいガイダンスは、Microsoft と Anyscale が Azure 統合を改善するためのパートナーシップを示すものです。

GPU の不足は、大規模な ML における最も重要な運用課題の一つです。NVIDIA GPU のような需要の高いアクセラレーターは、Azure リージョンにおいてクォータや利用可能性の問題に直面することが多く、これによりクラスターのセットアップやジョブのスケジューリングが遅れる可能性があります。

Microsoft が提案する解決策は、マルチクラスター・マルチリージョン構成です。異なる Azure リージョン内のさまざまな AKS インスタンスに Ray クラスターを分散配置することで、チームは以下のことが可能になります：リージョンごとの制限を超えて GPU のクォータを集約する、障害や容量の問題が発生した際にワークロードを自動的にリルートする、Azure Arc を使用して AKS と連携し、オンプレミスシステムや他のクラウドプロバイダーへ計算プールを拡張する。

Anyscale コンソールでは、登録されたクラスターを一つのビューで確認できます。Anyscale Workspaces は、利用可能なキャパシティに基づいて、手動または自動でワークロードのスケジューリングを管理します。新しいリージョンを追加するには、cloud_resource.yaml マニフェストを作成し、その後 Anyscale CLI を使用して適用します。この設定ファーストのアプローチにより、マルチリージョンへの拡張が容易に管理できます。

ML 運用における一般的な課題は、パイプラインの各ステージ間でトレーニングデータ、モデルチェックポイント、およびアーティファクトを転送することです。これには、事前学習から微調整、そして推論へとデータを移動させることが含まれます。この課題への対応策として Azure BlobFuse2 が提案されており、これは Ray ワーカーポッドに POSIX 互換のファイルシステムとして Azure Blob Storage をマウントします。

Ray の観点からは、マウントポイントは単なるローカルディレクトリです。タスクとアクターは標準的なファイル I/O を使用してデータセットを読み込み、チェックポイントを書き出します。BlobFuse2 はその後、データを Azure Blob Storage に保存します。これにより、ポッド間およびノードプール間でデータが利用可能になります。ローカルキャッシングは大規模なトレーニング実行中の GPU のストールを防ぎ、データと計算リソースが分離されているため、Ray クラスターはデータ損失を伴うことなくスケールアップ・ダウンできます。

セットアップには、クラスター作成時に Blob CSI ドライバーを有効化します。次に、ワークロードアイデンティティを使用して認証を行う StorageClass を定義してください。最後に、ReadWriteMany アクセス権を持つ PersistentVolumeClaim を作成します。これにより、異なるノード上の複数の Ray ワーカーが同時に共有データにアクセスできるようになります。このアプローチにより、Ray コードの移植性が向上し、インフラストラクチャ層に Azure ネイティブストレージの耐久性とスケーラビリティが付加されます。

もう一つの重要なトピックは認証の信頼性です。Anyscale と Azure は以前、30 日ごとに期限切れとなる CLI トークンまたは API キーとの統合を行っており、これにより手動でのローテーションが必要となり、サービス停止のリスクがありました。

新しい手法では、Microsoft Entra サービスプリンシパルと AKS ワークロードアイデンティティを使用し、短期間のトークンを自動的に発行します。Anyscale Kubernetes Operator ポッドはユーザー割り当てマネージドアイデンティティを使用しており、このアイデンティティが Entra ID から Anyscale サービスプリンシパルのアクセストークンを要求します。Azure はトークンの更新を透過的に処理するため、クラスター内に長期有効な認証情報を保存する必要がなく、手動でのローテーションも不要です。

著者らは、これはマルチクラスター環境において特に重要であると述べています。ここでは多数のクラスター間で認証情報を手動で管理することは運用負荷を増大させます。ワークロードアイデンティティモデルは、Azure リソースアクセスに対する細粒度の RBAC を提供し、副産物として Azure Activity Logs を通じて完全な監査証跡を生成します。

Anyscale の AKS 統合機能は現在、クローズドベータ版として提供されています。アクセスを希望するチームは、担当の Microsoft アカウントチームまでお問い合わせください。また、AKS の GitHub リポジトリにてリクエストを提出することも可能です。その際は、Ray ワークロードの詳細と対象リージョンを含めて記載してください。GitHub 上の Azure-Samples/aks-anyscale リポジトリでは、DeepSpeed や LLaMA-Factory を用いたファインチューニングのための例セットアップやワークロードを確認できます。これには大規模言語モデル（LLM）の推論エンドポイントも含まれています。

Microsoft がこの賭けを単独で行っているわけではありません。AWS は Ray Summit 2024 で Anyscale とのパートナーシップを発表しました。これにより、E クラスターが RayTurbo ランタイムと接続されます。これは、NVIDIA GPU に AWS の Trainium や Inferentia アクセラレータを組み合わせてハードウェアの柔軟性を強調するものです。さらに、SageMaker HyperPod は、ノードレベルでの耐障害性が必要な長時間実行型のトレーニングジョブに対するデプロイ先として新たに追加されました。Google Cloud はオープンソースへの貢献においてリードしています。

GKE チームは Anyscale のエンジニアと協力し、ラベルベースのスケジューリングを Ray v2.49 へアップストリームしました。また、マルチチップ TPU セットアップにおけるリソースの断片化を削減するため、ray.util.tpu レイヤーも作成されました。さらに、新しい GB200 ベースのインスタンス向けに動的リソース割り当て（Dynamic Resource Allocation）が追加されています。

3 つの主要クラウドプロバイダーはすべて同じマネージド Ray オペレーターを採用し、それぞれが独自のインフラストラクチャを追加しました。これは、業界が AI ワークロードに対して Kubernetes と Ray の組み合わせを好んでいることを示しています。現在、競争の焦点はランタイムそのものではなく、どのクラウドが周辺インフラストラクチャを最も効率的に最適化できるかという点に移っています。

著者について

Claudio Masolo

Claudio は Nearform のシニア DevOps エンジニアです。

余暇には、ランニング、読書、そして昔のビデオゲームを遊ぶことが好きです。