SageMaker Studioから直接Amazon SageMaker HyperPod Spacesを管理する

このたび私たちは、Amazon SageMaker Studio UIから直接、Amazon SageMaker HyperPod EKSクラスター上でAmazon SageMaker Spacesを作成・管理できる機能を発表しました。データサイエンティストおよび機械学習(ML)エンジニアは、ブラウザーを離れたりコマンドラインツールを使ったりすることなく、HyperPodクラスター上でJupyterLabやCode Editor環境を起動できるようになり、クラスターへのアクセスから生産的な開発開始までの時間が数回のクリックに短縮されます。

背景

Amazon SageMaker HyperPod は、大規模な基盤モデル(FM)のトレーニングと推論のために特別に構築されたインフラストラクチャを提供します。 Amazon Elastic Kubernetes Service (Amazon EKS) によるオーケストレーションにより、チームは組み込みの耐障害性と自動的な障害復旧機能を備えた、数百のアクセラレーターにまたがる分散トレーニングジョブを実行できます。トレーニングに加えて、HyperPodはこのEKSオーケストレーションされたインフラストラクチャを、数十億パラメーターの基盤モデルに対する低レイテンシーでスケーラブルな推論の提供にも拡張しています。

本年初頭、私たちは Amazon SageMaker Spaces for HyperPodを発表しました。これは、ML開発者がHyperPod EKSクラスター上に直接インタラクティブな開発環境を作成できるアドオンです。これにより、組織はGPUのフラクショナル割り当てに対応しながら、同じインフラストラクチャ上でトレーニングジョブやモデルデプロイと並行してインタラクティブなワークロードを実行することで、GPUへの投資を最大化できました。

以前は、Spacesの作成と管理は主に HyperPod CLI または kubectl コマンドに依存していました。このアプローチはインフラストラクチャ管理者にとって強力で詳細な制御を提供しますが、ビジュアルインターフェースを好むデータサイエンティストは、この新しいSageMaker Studioの機能を利用してコマンドラインツールを回避し、モデル開発に完全に集中できるようになりました。

新機能

この新機能により、データサイエンティストはSageMaker Studioから直接Spacesの作成、設定、開始、停止、オープンができるようになりました。HyperPodクラスターの詳細ページにある新しいIDE and Notebooksタブは、Space管理のための完全なユーザーインターフェースを提供し、日常的なSpaceの操作にCLIツールを使用する必要をなくなります。

Studioから利用可能な主な機能は次のとおりです:

• ガイド付きフォームを通じて、設定可能なコンピュート、ネームスペース、ストレージ、コンピュートクォータ管理のための HyperPod Task Governance 、およびイメージ設定を備えたSpacesを作成する。
• 名前、アプリケーションタイプ、ステータス、アクセスタイプ、ストレージ、GPU、vCPUの割り当てを表示する検索可能なテーブルですべてのSpacesを表示する。
• ワンクリックでSpacesを開始・停止し、使用されていないSpacesのコンピュートリソースを解放する。
• Spacesをブラウザーで直接開く(JupyterLabまたはCode Editor)、または選択したリモートIDE(例:VS Code)から接続する。

図1:HyperPodクラスターの詳細ページにあるIDE and Notebooksタブには、すべてのSpacesのステータス、コンピュート割り当て、および停止、オープン、リモートIDEで開くためのクイックアクションが表示されます

はじめに

セットアップには2つのロールが関わります: 管理者 がクラスターを準備し、 データサイエンティスト がSpacesを作成して開きます。以下のセクションでそれぞれを説明します。

管理者向け

管理者は SageMaker Spaces アドオンをインストールし HyperPod EKS クラスターに対して、SageMaker HyperPod EKS クラスターの IDE and Notebooks タブから、 Quick install (最適化されたデフォルト設定でのワンクリック) または Custom install オプション (Web UI アクセスを設定するために必要) のいずれかを使用してインストールします。インストール後、管理者は名前空間の設定、Space テンプレートの作成、および EKS アクセスエントリを通じたアクセス管理を行うことができます。

以下は、管理者が実施する必要のある 1 回限りのセットアップです。

• Spaces アドオンをインストールします: Amazon SageMaker AI コンソールから HyperPod クラスターを開き、 IDE and Notebooks タブに移動して、 Quick install または Custom install (Web ブラウザアクセスを有効にするには Custom install が必要です) を選択します。詳細な手順については、 AWS ドキュメント を参照してください。
• EKS アクセスエントリを設定します: 3 つの管理ポリシー AmazonSagemakerHyperpodSpacePolicy, AmazonSagemakerHyperpodUserClusterPolicyと AmazonSagemakerHyperpodSpaceTemplatePolicy を、データサイエンティストが使用する AWS Identity and Access Management (IAM) ロールにアタッチします。
• StudioドメインでユーザーごとのID伝播を有効にします: お使いのStudioドメインがSageMaker StudioとHyperPod Spacesの統合が展開される前に作成されたものである場合は、HyperPod EKSクラスターへのユーザーごとのID伝播を有効にする必要があります。これにより、クラスターに対する各Studioユーザーの操作(Spaceの作成、停止、削除)が、EKSアクセスエントリーおよびAWS CloudTrailにおいて該当ユーザーのユーザープロファイルに紐付けられることが保証されます。さらに、このIDマッピングにより、厳格なSpaceの所有権が強制されます。どのユーザーが環境を作成したかが追跡され、SpaceがSageMaker Studioドメイン内でプライベートか共有かが判定されます。

Studioドメインごとに以下のコマンドを一度実行します:

aws sagemaker update-domain \
    --domain-id $DOMAIN_ID \
    --default-user-settings '{
  "StudioWebPortalSettings": {
    "ExecutionRoleSessionNameMode": "USER_IDENTITY"
  }
}'

ドメインの更新後、以下のコマンドが「USER_IDENTITY」を返すことを確認します:

aws sagemaker describe-domain --domain-id $DOMAIN_ID \
    --query 'DefaultUserSettings.StudioWebPortalSettings.ExecutionRoleSessionNameMode'

既に実行中のアプリには影響ありません。ユーザーは次回のサインイン時に新しい設定が適用されます。

• 必要に応じて追加の機能を有効にします: 以下の「オプション機能」の表から、任意の機能を有効にしてください。

データサイエンティスト向け

アドオンがインストールされ、アクセスが設定された後、データサイエンティストはSageMaker StudioのComputeの下にあるHyperPodクラスターに移動し → 「IDE and Notebooks」タブを選択してSpaces管理インターフェースを表示します(図1を参照)。Spaceの作成と管理の詳細については、以下を参照してください: HyperPodでのSpaceの作成と管理.

Spaceのステータスが Running と表示されたら(コールドクラスターでは通常数分、オーバープロビジョニングがあれば約30～40秒)、「 Open 」を選択してブラウザーでJupyterLabまたはCode Editorを起動するか(図2および図3)、「 Open in VS Code 」を選択してSSH-over-SSM経由でローカルエディターから接続します(図4)。

図2: WebブラウザーからアクセスしたJupyterLab Space。利用可能なノートブックカーネル、コンソール、ターミナルアクセスを備えたランチャーが表示されています

JupyterLab Spaceでの作業

JupyterLab Spaceを開くと、以下にアクセスできる完全に構成された開発環境が提供されます:

• ipykernelを備えたPython 3ノートブック。
• Glue PySparkおよびGlue Sparkコンソール。
• SparkMagicのPySparkおよびSparkカーネル。
• コマンド実行用のターミナルアクセス。
• 永続ストレージを備えたファイルブラウザー。
• 組み込みチャットとコンテキストヘルプ。

作業内容は接続されたAmazon Elastic Block Store(Amazon EBS)ボリュームに保持されるため、進捗を失うことなくSpaceを停止および再起動できます。

図3: HyperPod上で実行されているCode Editor Space。ファイルエクスプローラー、エディター、統合ターミナルを備えたVS CodeスタイルのWebインターフェースが表示されています

Code Editor Spaceでの作業

ブラウザーでVS Codeスタイルのエクスペリエンスを好む開発者のために、Code Editor Spacesは以下を備えた軽量のWebベースIDEを提供します:

• 構文ハイライトとIntelliSenseを備えた完全なファイル編集。
• シェルコマンドの実行、トレーニングジョブの送信、クラスターリソースとのやり取りのための統合ターミナル。
• 言語サーバー、リンター、フォーマッターの拡張機能サポート。
• バージョン管理ワークフローのためのGit統合。
• クラスターファイルシステムおよびマウントされたAmazon FSxボリュームへの直接アクセス。

Code Editor Spacesは、トレーニングスクリプトの作成とデバッグ、実験設定の管理、コードリポジトリの作業に最適で、これらすべてをブラウザを離れることなく行えます。

図4: HyperPod Space にリモート接続されたローカル VS Code インスタンス。リモート接続インジケーターと、クラスターコンピュート上で動作する完全な IDE 機能が表示されています

リモート IDE 接続 (VS Code)

Spaces テーブルから「Open in VS Code」を選択すると、ローカルの Visual Studio Code を HyperPod 上で実行されている Space に接続できます。これは内部的に SSH-over-SSM トンネリングを使用しており、SSH キーの管理やポート 22 の公開を行うことなく、安全な接続を提供します。

ローカルの VS Code 環境のフル機能(拡張機能、テーマ、キーバインディングを含む)をそのままに、HyperPod クラスターコンピュート上でコードを実行できます。

また、AWS Toolkit for Visual Studio Code を使用して接続することもできます。このツールキットでは、SageMaker AI > HyperPod の下に Spaces が一覧表示され、ツールキットパネルから直接 Space の開始、停止、接続を行えます。

オプションの機能

以下に挙げる機能はすべてオプションであり、組み合わせて利用できます。チームのニーズに合った任意の組み合わせを有効にしてください。

Pro tip: ノードのオーバープロビジョニングで Space の起動時間を短縮

デフォルトでは、Karpenter オートスケーリングを使用する HyperPod EKS クラスター上の SageMaker Spaces には、 コールドスタート遅延として 5〜7 分 が発生します。これは、スケールトゥーゼロのクラスターで Space を初めて作成する際に生じるもので、主に以下が原因です:

• Amazon Elastic Compute Cloud (Amazon EC2) インスタンスの起動。
• Kubernetes ノードの登録。
• SageMaker Distribution (SMD) イメージのプル。

レイテンシーの影響を受けやすいインタラクティブワークロード (JupyterLab、Code Editor) の場合は、 事前にウォームアップされ、イメージがキャッシュされたノードのプールを維持することができます 標準的なKubernetesのオーバープロビジョニングパターンを使用しています。これにより、Spaceの起動時間が数分からおよそ 30〜40秒に短縮されます。

仕組み

1. 低優先度の(-1000)プレースホルダーDeploymentが、ウォーム状態の各ノード上に1つのKubernetes Podを保持します。これらのPodは、実際のSpaceと同じCPU/メモリを要求します。
2. 各プレースホルダー上の initContainer が、Karpenterがノードをプロビジョニングする際に、SageMaker Distributionイメージをそのノードに事前プルします。
3. ユーザーがSpaceを作成すると(デフォルトの優先度 0)、Kubernetesスケジューラーはプレースホルダー(優先度 -1000)をプリエンプトします。これにより、Spaceはイメージプルやノード起動の待ち時間なしで、すでにウォーム状態でイメージがキャッシュされたノードに数秒で配置されます。
4. Karpenter は、置き換えられたプレースホルダー用の代替ノードをバックグラウンドでプロビジョニングします。

Pro tip のデプロイの詳細については、以下をご覧ください: Overprovisioning for HyperPod Spaces.

検証済みの起動レイテンシー on ml.m5.12xlarge (24 vCPU の割り当て可能、2 vCPU のプレースホルダー、8 GiB のプレースホルダーメモリ) に、 sagemaker-distribution:latest-cpu SMD イメージ (約 3.5 GB) を使用した場合:

Note: これらの数値は CPU のみの場合です。GPU Space では、GPU イメージを事前プルした上で、 nvidia.com/gpu を要求する専用のプレースホルダー Deployment が必要です。そうしないと、GPU ノードはコールドのままとなり、GPU イメージは約 10 GB あるため、3.5 GB の CPU イメージよりもはるかに大きなプルコストがかかります。

Note: ウォームノードはそれぞれ、Running 状態の EC2 インスタンスを 1 つ保持します。オンデマンドノードの場合、これはウォームノードを起動し続けるための追加コストとなります。

HyperPod Spaces アドオンのインストールと使い始め方の詳細については、以下を参照してください: AWS ドキュメント.

料金

SageMaker Spaces アドオンの設定に追加料金はかかりません。お支払いいただくのは、Space が消費する基盤となる HyperPod クラスターコンピュートの料金と、SSH-over-SSM リモート接続に使用される AWS Systems Manager Advanced On-Premises Instance の 1 時間あたりの料金です。詳細については、 AWS Systems Manager の料金 を参照してください。

前述のオーバープロビジョニングを使用する場合は、ウォームノードの追加コストが発生することに注意してください。インスタンスタイプとサイズに応じて、これらのノードは Space のプロビジョニングを待つために Running 状態のまま保持されます。

Conclusion

SageMaker Studio による HyperPod Spaces の管理により、データサイエンティストと高性能コンピュートインフラストラクチャの間のギャップが埋まります。チームは、CLI ツールや Kubernetes の概念を学ぶことなく、クラスターへのアクセスから、数分で稼働中の JupyterLab や Code Editor 環境を利用できるようになりました。HyperPod Task Governance、フラクショナル GPU サポート、アイドルシャットダウンなどの機能と組み合わせることで、組織はコスト管理とリソースの公平性を維持しながら、共有クラスターへのセルフサービスアクセスを提供できます。

始めるには、SageMaker AI コンソールで HyperPod EKS クラスターに移動し、IDE and Notebooks タブを選択してください。詳細については、SageMaker HyperPod Spaces のドキュメントを参照してください。

About the authors