概要

Amazon SageMaker HyperPodは、大規模言語モデルや生成AIモデルのトレーニング向けに設計された高性能Cluster Serviceです。数千GPUへのスケーラビリティ、自動障害検出・回復機能により長時間トレーニングの信頼性を高め、事前最適化されたMLインフラで開発効率を向上させます。大規模AIモデル開発において、コスト効率と安定性を両立させる理想的なインフラストラクチャソリューションです。

本記事では、Terraformを活用してHyperPod環境を構築する手順を解説します。

以下は、今回構築するAWS Architectureです。
SageMaker HyperPodの構築にはAWS提供のLifecycle Scriptsを活用し、基盤となるインフラ部分はTerraformを用いて構築します。

Amazon Sagemaker Hyperpod環境を構築するための前提条件

HyperPod環境を構築する前に、以下の前提条件を満たしていることを確認します。

前提条件

AWS Accountがあり、適切なIAM権限を持っていること
AWS CLIが設定済みであること
- Sagemaker Cluster CLIを利用するには、2.14.3以上が必要
Terraformを活用するための基本知識があること
Terraformがローカル環境にインストールされていること
- required_version = "~> 1.8.0"
SLURMを活用するための基本知識があること
Service Quotas(※Accountごとに異なるため、確認が必要)
- Service: Amazon SageMaker
  - Maximum size of EBS volume in GB for a SageMaker HyperPod cluster instance
    - Applied account-level quota value: 50 以上
  - Total number of instances allowed across SageMaker HyperPod clusters
    - Applied account-level quota value: 10 以上
  - ml.g5.xlarge for cluster usage
    - Requested quota value:1 以上

Terraformを用いた環境のコード化

AWSではHyperPod環境をCloudFormationで提供していますが、本記事ではTerraformを活用して同様の環境を構築します。CloudFormationの利用も可能ですが、Terraformを用いることでコードの再利用性や管理の柔軟性が向上します。
AWS Official: Amazon Sagemaker Hyperpod Cloudformation

Terraformのコードは個人のGitHubに公開しているため、必要に応じて取得し、terraform applyを実行してください。
Amazon Sagemaker Hyperpod Terraform

Terraform Dir構成

.
├── README.md
├── files
├── fsx.tf
├── iam.tf
├── nw.tf
├── output.tf
├── provider.tf
├── s3.tf
├── sg.tf
├── tfplan.txt
└── vars.tf

Terraformの適用が完了すると、AWS上にHyperPodのインフラが構築され、以下のようなOutputsが出力されます。

これらの情報は、HyperPodの立ち上げや運用時に必要となるため、メモしておくか、ターミナルの出力を適宜整理して作業を進めることを推奨します。

Outputs:

efa_sg_id = "sg-0e4b9d4777c5d81a9"
fsx_dns_name = "fs-07feee9b0e55d7c5d.fsx.ap-northeast-1.amazonaws.com"
fsx_mount_name = "4lpqtbev"
primary_private_subnet_id = "subnet-0e30f0724f1463291"
s3_lifecycle_bucket_Uri = "s3://leo-stg-hyperpod-lifecycle-7d460cc5/lifecycle-script-directory/src/"
sagemaker_execution_role_arn = "arn:aws:iam::xxxxxxx:role/leo-stg-SagemakerHyperpod-Role"
second_private_subnet_id = "subnet-0861f28ede3762119"

SLURMクラスタのセットアップ（Head Node、Compute Node）

セットアップする前に、簡単に構成について説明します。

AWS HyperPodにおけるSlurmクラスターの構成と主要要素

AWS HyperPodのSlurmクラスターは、計算リソースを効率的に管理し、ジョブスケジューリングを最適化するために構築されます。
HyperPodの構成は、大きく Instance Groups、Lifecycle Scripts、Service の3つの要素で成り立っています。

Instance Groups: 計算ノード（Compute Nodes）をグループ化し、ジョブの並列処理を可能にします。
Lifecycle Scripts: .sh や .py スクリプトを活用し、ノードのセットアップや環境設定を自動化します。
Service: Slurmのヘッドノード（Head Nodes）を通じて、ジョブを管理・スケジュールし、計算ノードに最適に配分します。

AWS HyperPodにおけるLifecycle Scriptsの役割と起動プロセス

AWS HyperPodのSlurmクラスターでは、Lifecycle Scripts がノードのセットアップと初期設定を自動化し、起動プロセスを管理します。

起動プロセスの流れ
- on_create.sh が create_cluster.json を基にクラスタ作成を開始
- lifecycle_script.py が provisioning_params.json を参照し、ノードごとの設定を適用
- カスタムスクリプト（install_x.sh など）を実行し、パッケージのインストールや設定を実施
主な役割
- Slurmの構成 (slurm.conf の生成、slurmctld / slurmd の起動)
- 認証設定 (munge.key の作成と配布)
- FSx for Lustreのマウント
- スケーラビリティの確保（ノード追加・削除の自動対応）

Terraformを使用してAWSインフラ環境の準備が整ったので、次に、以下のイメージの赤いボックス部分にあたるSLURMクラスタのセットアップを行います。

1. provisioning_parameters.json設定

provisioning_params.jsonは、Slurmクラスタの定義やPartitionの設定、FSx for Lustreの指定を行う構成ファイルです。

まず、Lifecycle Scriptsを利用するために、AWS公式のGitHubリポジトリをクローンします。


git clone --depth=1 https://github.com/aws-samples/awsome-distributed-training/

次に、/base-config フォルダ内のファイルを作業しやすいディレクトリに移動します。（そのまま /base-config で作業しても問題ありません。）
私の場合、以下のようにフォルダを整理しています。
クローンしたリポジトリには多くのファイルが含まれており、すべてを確認するのは大変でした。しかし、何度も試行した結果、以下のようなファイル構成であれば、最低限のHyperPodクラスター構築には問題なく対応できることを確認しました。

.
├── README.md
├── cluster
│   ├── create_cluster.json     # 個別に作成するJSONFile
│   ├── easy-ssh.sh             # awsome-distributed-training/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh
│   ├── lifecycle_files         # awsome-distributed-training/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/
│   ├── update_cluster.json     # 個別に作成するJSONFile
│   └── validate-config.py      # awsome-distributed-training/1.architectures/5.sagemaker-hyperpod/validate-config.py
└── terraform                   # Terraform Code
    ├── README.md
    ├── files
    ├── fsx.tf
    ├── iam.tf
    ├── nw.tf
    ├── output.tf
    ├── provider.tf
    ├── s3.tf
    ├── sg.tf
    ├── tfplan.txt
    └── vars.tf

provisioning_parameters.jsonは個別に作成する必要がありますが、
配置場所が
私のDir構成の場合、/cluster/lifecycle_filesとなり、
AWS公式のGitHubの場合、awsome-distributed-training/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/となります。
provisioning_parameters.jsonの内容は次の通りです。


{
    "version": "1.0.0",
    "workload_manager": "slurm",
    "controller_group": "your-head-node-name", # Head Node名
    "worker_groups": [
        {
            "instance_group_name": "your-compute-node-name", # Compute Node名
            "partition_name": "partition-1" # Partition名
        }
    ],
    "fsx_dns_name": "fs-07feee9b0e55d7c5d.fsx.ap-northeast-1.amazonaws.com", # Terraform outputsのfsx_dns_name
    "fsx_mountname": "4lpqtbev" # Terraform outputsのfsx_mount_name
}

上記のLifecycle Scriptsで説明したように、ノードのセットアップと初期設定を自動化し、起動プロセスを管理するためのスクリプトをS3に配置します。このファイルをlifecycle_script.pyが参照し、各ノードに適用します。そのため、Lifecycle Scriptsを適切に管理し、provisioning_params.jsonを含めた設定をS3にアップロードする必要があります。

以下のコマンドでS3にアップロードします。


# 以下 /cluster上での実行コマンド

# aws s3 sync "lifecycle_files/" {Terraform outputsのsyncs3_lifecycle_bucket_Uri}

aws s3 sync "lifecycle_files/" s3://leo-stg-hyperpod-lifecycle-7d460cc5/lifecycle-script-directory/src

2. create_cluster.json設定

Lifecycle Scriptsが準備できましたので、Clusterを起動するJSON Fileを作成します。
create_cluster.jsonは次のように作成します。
今回はHead NodeとCompute Nodeは1台構成になります。
細かいOptionsの説明は公式Docsを参考にしてください。
Create a SageMaker HyperPod cluster


{
    "ClusterName": "your-cluster-name", # HyperPod Cluster名
    
    "InstanceGroups": [
        {
            
            "InstanceGroupName": "your-head-node-name", # provisioning_params.jsonと同じHead Node名
            "InstanceType": "ml.m5.large",
            "InstanceCount": 1,
            "LifeCycleConfig": {
              "SourceS3Uri": "s3://leo-stg-hyperpod-lifecycle-7d460cc5/lifecycle-script-directory/src", # Terraform outputsのsyncs3_lifecycle_bucket_Uri
              "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "arn:aws:iam::xxxxxxx:role/leo-stg-SagemakerHyperpod-Role", # Terraform outputsのsagemaker_execution_role_arn
            "InstanceStorageConfigs": [
                {
                   "EbsVolumeConfig":{     
                      "VolumeSizeInGB": 30
                   }
                }
            ]
        },
        {
            "InstanceGroupName": "your-compute-node-name", # provisioning_params.jsonと同じCompute Node名
            "InstanceType": "ml.g5.xlarge",
            "InstanceCount": 1,
            "LifeCycleConfig": {
              "SourceS3Uri": "s3://leo-stg-hyperpod-lifecycle-7d460cc5/lifecycle-script-directory/src", # Terraform outputsのsyncs3_lifecycle_bucket_Uri
              "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "arn:aws:iam::xxxxxxx:role/leo-stg-SagemakerHyperpod-Role", # Terraform outputsのsagemaker_execution_role_arn
            "InstanceStorageConfigs": [
                {
                   "EbsVolumeConfig":{     
                      "VolumeSizeInGB": 30
                   }
                }
            ]
        }
    ],
    "VpcConfig": {
      "SecurityGroupIds": ["sg-0e4b9d4777c5d81a9"], # Terraform outputsのefa_sg_id
      "Subnets":["subnet-0e30f0724f1463291"] # Terraform outputsのprimary_private_subnet_id
    }
}

3. Create Clusterコマンド実行

クラスター作成用のJSONが準備できたので、実行前に設定に問題がないかを確認します。
この確認には、validate-config.pyを使用できます。このスクリプトを実行することで、設定ファイルの整合性や構成の適切性を事前に検証し、エラーや不備がないか確認することが可能です。

例えば、次のようcreate_cluster.jsonとprovisioning_parameters.jsonを指定することで、確認が可能です。エラーが起きた場合、エラーメッセージに従って対応します。


 python3 ../validate-config.py --cluster-config create_cluster.json --provisioning-parameters lifecycle_files/provisioning_parameters.json
✔️  Validated instance group name leo-stg-compute-ig is correct ...
✔️  Validated subnet subnet-0e30f0724f1463291 ...
✔️  Validated security group sg-0e4b9d4777c5d81a9 ingress rules ...
✔️  Validated security group sg-0e4b9d4777c5d81a9 egress rules ...
✔️  Validated FSx Lustre DNS name fs-07feee9b0e55d7c5d.fsx.ap-northeast-1.amazonaws.com ...
✔️  Validated FSx Lustre mount name 4lpqtbev ...
✔️  Validated provisioning_parameter.json is valid json ...
✅ Cluster Validation succeeded

validate-config.pyで設定に問題がないことを確認したら、次のコマンドを実行してクラスターを構築します。


aws sagemaker create-cluster --cli-input-json file://create_cluster.json

AWSマネジメントコンソールで確認すると、Cluster StatusがCreatingになっていることが分かります。

今回はCompute NodeにGPUを利用しているため、クラスターのプロビジョニングには一定時間がかかりますが、しばらくするとCluster StatusがInServiceになり、HyperPodクラスターの構築が完了したことを確認できます。

4. HyperPod Head Nodeに接続する方法

HyperPodクラスターにアクセスするために、easy-ssh.sh スクリプトを使用してノードに接続する方法を紹介します。本手順では、AWS Systems Manager（HyperPodクラスターの仕様）を活用し、SSHを簡単に設定・利用できるようにします。

HyperPodのHead Nodeに接続するには、以下のコマンドを実行します。


# ./easy-ssh.sh -c {YOUR_HEADNODE_NAME} {YOUR_CLUSTER_NAME}

./easy-ssh.sh -c leo-stg-head-ig leo-stg-hyperpod-cluster

# コマンドを実行すると、以下の情報が出力されます。

Cluster id: rw1dauane3qr
Instance id: i-0067cfc34c41a5c56
Node Group: leo-stg-head-ig
aws ssm start-session  --target sagemaker-cluster:rw1dauane3qr_leo-stg-head-ig-i-0067cfc34c41a5c56
Add the following to your ~/.ssh/config to easily connect:

cat <<EOF >> ~/.ssh/config
Host leo-stg-hyperpod-cluster
  User ubuntu
  ProxyCommand sh -c "aws ssm start-session  --target sagemaker-cluster:rw1dauane3qr_leo-stg-head-ig-i-0067cfc34c41a5c56 --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"
EOF

Add your ssh keypair and then you can do:

$ ssh leo-stg-hyperpod-cluster


Starting session with SessionId: test_hyperpod_leo-hgi52c5cp77rf9oukfteqxcq7u
# sudo su - ubuntu  >> Default UserはUbuntu
ubuntu@ip-10-1-23-70:~$

接続が成功すると、次のように sinfo コマンドでSlurmクラスターの状態を確認できます。

ubuntu@ip-10-1-23-70:~$ sinfo
PARTITION   AVAIL  TIMELIMIT  NODES  STATE NODELIST
dev*           up   infinite      1   idle ip-10-1-44-196
partition-1    up   infinite      1   idle ip-10-1-44-196

これで、AI/MLモデルの開発環境が整いました。

FSx for LustreやS3との統合

HyperPodクラスターの設定に関する説明の中で触れていない部分として、FSx for LustreとS3の統合についてはすでにTerraformで対応済みです。

AWS FSx for Lustreは、高性能な並列ファイルシステムを提供し、大規模なHPCやAI/MLワークロードに最適化されています。特に、S3との統合により、ストレージ管理の柔軟性と効率性が向上します。

S3連携
- FSx for Lustreは、S3バケットをバックエンドとして使用でき、データの取り込み (fsx lustre data repository association) や、S3への書き戻し (export path) を簡単に実行可能です。
高速アクセス
- S3データをキャッシュし、POSIX互換の低レイテンシアクセスを提供。頻繁にアクセスするデータをFSxで高速処理し、S3でコストを最適化できます。
スケーラビリティ
- TB〜PB級のデータ処理をサポートし、大規模並列処理が可能。AI/MLやシミュレーション用途で有効です。

この統合により、HyperPod環境では、S3のスケーラブルなストレージとFSx for Lustreの高速性を活かしたデータ処理基盤を構築できます。

HyperPodクラスターの削除方法

AWS HyperPodクラスターは、用途を考慮せずに構築すると不要なコストが発生するため、不要な場合は適切に削除することが重要です。

以下のコマンドを実行すると、一定時間後にAWSマネジメントコンソール上からもリソースが削除されます。


# aws sagemaker delete-cluster --cluster-name {YOUR_CLUSTER_NAME}

aws sagemaker delete-cluster --cluster-name leo-stg-hyperpod-cluste

また、FSx for LustreやNAT Gatewayなどのサービスも継続的にコストが発生するため、Terraformを利用して一括で削除することを推奨します。
terraform destroyを実行することで、関連するすべてのリソースを適切に削除できます。

最後に

Amazon SageMaker HyperPodは新しいサービスであり、触れる機会が限られていますが、仕事の一環として数ヶ月間取り組むことができました。

当初は、構築に必要な準備項目が多く、全体を整理するだけでも難しく感じました。しかし、何度も試行錯誤を繰り返し、自分に合った構成を確立することで、最適な形でHyperPodの運用を行えるようになりました。

この経験を活かし、本ブログでは可能な限りシンプルな構成を紹介しています。少しでも役立てば幸いです。Amazon SageMaker HyperPodの特長や運用に関して、今後も整理しながら情報をまとめていく予定です。

以上、「Amazon SageMaker HyperPodの構築手順 – Terraformを活用した構築」に関する記事でした。ご覧いただき、ありがとうございました！