AWS、SageMaker HyperPod向けにRay統合機能を発表——Studioからのクラスタ管理や耐障害性向上に対応

要点

• 米Amazon Web Services（AWS）は、機械学習インフラ「Amazon SageMaker HyperPod」において、分散処理フレームワーク「Ray」を統合する新機能を発表した。

• クラスタ作成からダッシュボード閲覧、JupyterLabやCode Editorの接続、ジョブ送信までをSageMaker Studio上で一元管理できるようになる。

• ノード障害時の自動復旧や、階層化ストレージを活用したチェックポイントの高速復旧、KVキャッシュのオフロード機能を提供する。

• 標準的なRay APIおよびオープンソースのKubeRayと互換性を維持しており、既存のスクリプトを改修なしで利用できる。

• 米Amazon Web Services（AWS）は、大規模な機械学習インフラサービス「Amazon SageMaker HyperPod」において、分散Pythonフレームワーク「Ray」を統合する新機能を発表した。この統合により、基盤モデル（大規模なデータで事前学習された汎用AIモデル）の学習や推論に向けた専用インフラ上で、Rayを容易に運用・管理できるようになるという。従来のKubernetes環境で課題となっていた設定や運用の負担を大幅に削減する狙いがある。

Kubernetes環境におけるRay運用の課題を解消

Rayは、GPUクラスタ全体でPythonワークロードを分散処理し、分散学習（Ray Train）からモデル配信（Ray Serve）までを拡張するためのオープンソースフレームワークである。一方、Kubernetes上でRayを運用する際は、オープンソースのオペレーターである「KubeRay」を介してクラスタのライフサイクルを管理するのが一般的となっている。また、AWSが提供するSageMaker HyperPodは、「Amazon Elastic Kubernetes Service（Amazon EKS）」上で動作し、ノードの健全性監視や自動復旧機能を備えた大規模機械学習向けの専用インフラである。

発表によると、これまでKubernetes上でRayを実行するには、データサイエンティスト自身がYAMLマニフェスト（設定ファイル）を作成したり、依存関係を変更するたびにDockerイメージを再ビルドしたりする必要があったという。さらに、運用監視ツールであるRay Dashboardへのアクセスにポート転送（kubectl port-forward）を設定したり、メトリクス可視化のためにPrometheusやGrafanaを手動で構成したりするなど、環境構築や管理に多くの工数がかかっていたと説明している。

SageMaker Studioからの統合的な開発・監視体験

今回の機能提供により、統合開発環境「SageMaker Studio」を通じてRayクラスタの作成から管理、監視までを一貫して行えるようになった。利用者はコンソール上からRayクラスタを立ち上げられるほか、Ray Dashboardや「Amazon Managed Grafana」による可観測性ダッシュボードを直接開くことが可能になる。

開発環境としては、JupyterLabやCode EditorのワークスペースをRayクラスタに接続し、対話的な開発を行える仕組みが整えられた。さらに、分散ジョブの送信や、処理が停止した状態を検出するハングジョブ検知の設定も、すべてSageMaker Studio上から実行できるようになったという。

独自のインフラ連携による耐障害性と推論最適化

基盤モデルの学習や配信を支えるインフラレベルの機能強化も行われている。アプリケーション層においては、SageMaker HyperPodが持つノードヘルス監視と自動復旧機能がRayの学習ジョブに適用され、ノード障害時にも自動的に耐障害性が確保される。また、HyperPodの分散階層化ストレージ（tiered storage）を利用した「階層化チェックポインティング」に対応し、障害発生時などの学習再開を迅速化できるとしている。

モデル推論（Ray Serve）においては、事前学習済みモデルを簡単に利用できる「SageMaker JumpStart」との統合が図られた。これにより、モデルの重みデータをRay Serveのエンドポイントに直接読み込めるようになる。加えて、長文の入力処理（ロングコンテキスト）を効率化するため、中間データであるKVキャッシュを分散階層化ストレージへオフロード（退避）する機能も備えているという。

既存コードとの互換性と利用環境

今回の新機能は、オープンソースのKubeRayおよび標準的なRay APIと互換性を持って動作するよう設計されている。そのため、ユーザーがすでに作成したスクリプトや既存のワークフローに手を加えることなく、そのまま移行して実行できると説明している。

なお、本機能を利用するには、Amazon EKSでオーケストレーションされたSageMaker HyperPodクラスタとSageMaker Studioドメインが必要となる。あわせて、対話型ワークスペースを接続する「SageMaker Spaces EKS add-on」、メトリクスを収集する「HyperPod Observability EKS add-on」、クラスタ管理を行う「KubeRay operator」、認証付きパブリックエンドポイントを生成する「HyperPod Ray Endpoint Operator」の各コンポーネントをクラスタに導入する必要があるとしている。