SageMaker HyperPod のタスクガバナンス機能を試してみた
2025.01.12こんにちは！クラウド事業本部コンサルティング部のたかくに（@takakuni_）です。
re:Invent 2024 で Amazon SageMaker HyperPod にタスクガバナンス機能が登場しました。
当時は実際に試せてなかったため、試してみたいと思います。
タスクガバナンス
タスクガバナンスはざっくり説明すると、 HyperPod クラスターに所属するコンピュート（CPU や GPU）を、チームごとに管理する機能です。
タスクガバナンスを利用することで、以下の管理をチームごとに設定できるようになります。
- タスクの優先度
- 各チームのコンピューティング割り当て
- 各チームがアイドル状態のコンピューティングリソースを貸し借りする方法
- チームが自らのタスクを先取りした場合
前提条件
タスクガバナンスの機能は EKS オーケストレータのみサポートしている状況になります。Kubernetes v1.30.0 以上でサポートです。
Kueue を利用してジョブの管理を行うため、すでに Kueue をインストールしている場合はアンインストールが必要になります。加えて、HyperPod タスクガバナンスはアドオンを利用してセットアップを行うのですが、 HyperPod クラスター（ノード）作成後にインストールが必要になります。
また、利用できるリージョンは。執筆時点でバージニア北部、オハイオ、オレゴンの 3 リージョンです。
Amazon SageMaker HyperPod タスクガバナンスは、米国東部 (バージニア北部)、米国東部 (オハイオ)、米国西部 (オレゴン) の AWS リージョンでご利用いただけるようになりました。HyperPod タスクガバナンスは追加コストなしでご利用いただけます。詳細については、SageMaker HyperPod の製品ページにアクセスしてください。
やってみる
それでは早速、タスクガバナンス機能をセットアップします。
GPU が 4 枚搭載された ml.g5.12xlarge を 4 台起動し、 John チームは 3 台、 Smith チームは 1 台使えるような状況とします。
HashiCorp Terraform でセットアップしますが、要所要所ポイントをかいつまんでいきます。
EKS アドオンのインストール
それでは、タスクガバナンスの EKS アドオンをインストールします。HyperPod ノードがセットアップ完了したのちにインストールが必要のため、 depends_on で依存関係を追加しています。
resource "aws_eks_addon" "amazon_sagemaker_hyperpod_taskgovernance" {
cluster_name = aws_eks_cluster.this.name
addon_name = "amazon-sagemaker-hyperpod-taskgovernance"
resolve_conflicts_on_create = "OVERWRITE"
addon_version = "v1.0.0-eksbuild.4"
depends_on = [
awscc_sagemaker_cluster.this,
]
}
いくつか kueue-system 名前空間でリソースが作成されていますね。
takakuni@ sagemaker_hyperpod_eks_task_governance % kubectl get all -n kueue-system
NAME READY STATUS RESTARTS AGE
pod/kueue-controller-manager-685d6d4664-j8gq5 2/2 Running 0 10m
NAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGE
service/kueue-controller-manager-metrics-service ClusterIP 172.20.250.28 <none> 8443/TCP 10m
service/kueue-webhook-service ClusterIP 172.20.232.68 <none> 443/TCP 10m
NAME READY UP-TO-DATE AVAILABLE AGE
deployment.apps/kueue-controller-manager 1/1 1 1 10m
NAME DESIRED CURRENT READY AGE
replicaset.apps/kueue-controller-manager-685d6d4664 1 1 1 10m
クラスターポリシーの作成
アドオンが成功したため、クラスターポリシーを設定します。ここはコンソールで行います。
クラスターポリシーは、キューに投げるタスクやアイドル状態のコンピュートをどのように共有するのかを定義します。1 つのポリシーで タスクの優先順位付け と アイドル状態のコンピューティング割り当て の 2 軸を設定し、このポリシーはクラスター全体で共有します。
タスクの優先順位付け
タスクの優先順位付けは、キューで待機している PyTorch などのタスクの優先順位を定義します。
先着順でタスクを制御するのか、優先クラスを定義し、特定のクラスごとに優先させるのか定義します。
アイドル状態のコンピューティング割り当て
アイドル状態のコンピューティング割り当ては、使われていないコンピューティングを、チーム全体にどのように割り当てるかを定義します。定義は先着順と適正な共有に分かれます。
適正な共有は、チーム間でアイドル状態のコンピュートを借りられる共通の値（優先度）を設けて共有する仕組みです。この共通の値は後ほど設定します。
今回、タスクの優先順位付けは先着順、アイドル状態のコンピューティング割り当ては適正な共有で設定しました。
チームの作成
それではチームの作成を行います。コンピューティング割り当てから、チーム名や説明等を設定します。コンピューティング割り当ては、チームごとにどれくらいコンピュート（CPU,GPU,Memory）を利用できるのか定義します。
チームを作成すると、チーム名ごとに名前空間が自動作成されます。
適正な共有の加重が、先ほどのアイドル状態のコンピュートを借りられる優先度を意味します。
プリエンプションは、他チームにアイドル状態のコンピュートを貸している状況で、チーム内の割り込みタスクが入った場合に、アイドルで貸していたコンピュートの処理を中断させるかどうかを決定します。
つづいて、コンピューティングです。
クォータはチームに割り当てるインスタンスの数を設定します。
画面の場合、 ml.g5.12xlarge が 4 台利用可能ですが、インスタンスグループで作成していないインスタンスタイプも指定可能です。
貸借はチームがコンピュートを貸し借りできるかどうかの設定値です。
- 貸借
- 他のチームと貸し借り可能
- 貸す
- 他のチームに貸すことだけできる。借りることはできない
- 貸さないでください
- 他のチームに貸すことができない。借りることもできない
なお、貸借を選んだ場合は借用限度を 1 ～ 500%、デフォルトは 50% で決めることができます。
最終的な設定値
今回は john チームと smith チームに以下を設定しました。
john
- 適正な共有の加重：0
- プリエンプション：有効
- コンピュート：ml.g5.12xlarge 3 台
- 貸借モード
- 借用限度：50%
smith
- 適正な共有の加重：100
- プリエンプション：有効
- コンピュート：ml.g5.12xlarge 1 台
- 貸借モード
- 借用限度：100%
smith の方が優先して、借りられるようになっています。
ジョブの投入
John チーム
8 GPU を利用する PyTorch FSBP ジョブを John チームで実行します。
ワークショップの内容では ml.g5.8xlarge x 8 ですが、少し内容変えて ml.g5.12xlarge x 4 でやって見ました。
（メモリ不足にならなくて良かったです）
ラベルや名前空間をチームに割り当てられた内容に変更します。
apiVersion: v1
kind: Service
metadata:
name: etcd
+ namespace: hyperpod-ns-john
spec:
ports:
- name: etcd-client-port
port: 2379
protocol: TCP
targetPort: 2379
selector:
app: etcd
---
apiVersion: apps/v1
kind: Deployment
metadata:
+ namespace: hyperpod-ns-john
labels:
app: etcd
name: etcd
spec:
replicas: 1
selector:
matchLabels:
app: etcd
template:
metadata:
labels:
app: etcd
+ kueue.x-k8s.io/queue-name: hyperpod-ns-john-localqueue
spec:
containers:
- name: etcd
command: ['/usr/local/bin/etcd']
args:
- '--data-dir'
- '/var/lib/etcd'
- '--enable-v2'
- '--listen-client-urls'
- 'http://0.0.0.0:2379'
- '--advertise-client-urls'
- 'http://0.0.0.0:2379'
- '--initial-cluster-state'
- 'new'
image: quay.io/coreos/etcd:latest
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
restartPolicy: Always
---
apiVersion: 'kubeflow.org/v1'
kind: PyTorchJob
metadata:
name: fsdp
+ namespace: hyperpod-ns-john
labels:
+ kueue.x-k8s.io/queue-name: hyperpod-ns-john-localqueue
spec:
elasticPolicy:
rdzvBackend: etcd
rdzvHost: etcd
rdzvPort: 2379
minReplicas: 1
maxReplicas: 64
maxRestarts: 100
metrics:
- type: Resource
resource:
name: cpu
target:
type: Utilization
averageUtilization: 90
pytorchReplicaSpecs:
Worker:
replicas: 2
restartPolicy: OnFailure
template:
metadata:
labels:
app: fsdp
+ kueue.x-k8s.io/queue-name: hyperpod-ns-john-localqueue
spec:
volumes:
- name: shmem
hostPath:
path: /dev/shm
- name: local
hostPath:
path: /mnt/k8s-disks/0
#nodeSelector:
# node.kubernetes.io/instance-type: "ml.g5.12xlarge"
containers:
- name: pytorch
image: 123456789012.dkr.ecr.us-west-2.amazonaws.com/hyprpd-tskgvrn-fsdp:pytorch2.2
imagePullPolicy: Always
resources:
requests:
nvidia.com/gpu: 4
vpc.amazonaws.com/efa: 1
limits:
nvidia.com/gpu: 4
vpc.amazonaws.com/efa: 1
env:
# for P5 FI_* should be commented out
- name: LOGLEVEL
value: 'DEBUG'
#- name: FI_PROVIDER
# value: efa
#- name: FI_EFA_USE_DEVICE_RDMA
# value: "1"
#- name: FI_EFA_FORK_SAFE
# value: "1"
#- name: FI_LOG_LEVEL
# value: "1"
#- name: FI_EFA_ENABLE_SHM_TRANSFER
# value: "1"
- name: TORCH_DISTRIBUTED_DEBUG
value: 'DETAIL'
- name: TORCH_NCCL_ENABLE_MONITORING
value: '1'
- name: TORCH_NCCL_TRACE_BUFFER_SIZE
value: '20000'
- name: TORCH_NCCL_DUMP_ON_TIMEOUT
value: '1'
- name: TORCH_NCCL_DEBUG_INFO_TEMP_FILE
value: '/local/nccl_trace_rank_'
- name: PYTORCH_CUDA_ALLOC_CONF
value: 'expandable_segments:True'
- name: NCCL_DEBUG
value: 'INFO'
- name: NCCL_SOCKET_IFNAME
value: '^lo'
- name: TORCH_NCCL_ASYNC_ERROR_HANDLING
value: '1'
#- name: TORCH_DIST_INIT_BARRIER
# value: "1"
#- name: NCCL_IGNORE_DISABLED_P2P
# value: "1"
#- name: NCCL_NVLS_ENABLE
# value: "0"
command:
- /usr/local/bin/torchrun
- --nproc_per_node=4
- --nnodes=2
- /fsdp/train.py
- --max_context_width=4096
- --num_key_value_heads=32
- --intermediate_size=11008
- --hidden_width=4096
- --num_layers=32
- --num_heads=32
- --model_type=llama_v2
- --tokenizer=hf-internal-testing/llama-tokenizer
- --checkpoint_freq=5000
- --validation_freq=500
- --max_steps=5000
- --checkpoint_dir=/checkpoints
- --dataset=allenai/c4
- --dataset_config_name=en
- --resume_from_checkpoint=/checkpoints
- --train_batch_size=1
- --val_batch_size=1
- --sharding_strategy=full
- --offload_activation=1
volumeMounts:
- name: shmem
mountPath: /dev/shm
- name: local
mountPath: /local
ジョブを投げると、ダッシュボードの GPU 利用率がちょうど 50 % になりました。（4 枚使っているので想定通りですね）
ちょうど 50 % ほどになっていますね。ただ、50%しか使えていないため、もったいないですね。
Smith
Smith チームにも、同じようにジョブを投入してみます。
apiVersion: v1
kind: Service
metadata:
name: etcd
namespace: hyperpod-ns-smith
spec:
ports:
- name: etcd-client-port
port: 2379
protocol: TCP
targetPort: 2379
selector:
app: etcd
---
apiVersion: apps/v1
kind: Deployment
metadata:
namespace: hyperpod-ns-smith
labels:
app: etcd
name: etcd
spec:
replicas: 1
selector:
matchLabels:
app: etcd
template:
metadata:
labels:
app: etcd
kueue.x-k8s.io/queue-name: hyperpod-ns-smith-localqueue
spec:
containers:
- name: etcd
command: ['/usr/local/bin/etcd']
args:
- '--data-dir'
- '/var/lib/etcd'
- '--enable-v2'
- '--listen-client-urls'
- 'http://0.0.0.0:2379'
- '--advertise-client-urls'
- 'http://0.0.0.0:2379'
- '--initial-cluster-state'
- 'new'
image: quay.io/coreos/etcd:latest
ports:
- containerPort: 2379
name: client
protocol: TCP
- containerPort: 2380
name: server
protocol: TCP
restartPolicy: Always
---
apiVersion: 'kubeflow.org/v1'
kind: PyTorchJob
metadata:
name: fsdp
namespace: hyperpod-ns-smith
labels:
kueue.x-k8s.io/queue-name: hyperpod-ns-smith-localqueue
spec:
elasticPolicy:
rdzvBackend: etcd
rdzvHost: etcd
rdzvPort: 2379
minReplicas: 1
maxReplicas: 64
maxRestarts: 100
metrics:
- type: Resource
resource:
name: cpu
target:
type: Utilization
averageUtilization: 90
pytorchReplicaSpecs:
Worker:
replicas: 2
restartPolicy: OnFailure
template:
metadata:
labels:
app: fsdp
kueue.x-k8s.io/queue-name: hyperpod-ns-smith-localqueue
spec:
volumes:
- name: shmem
hostPath:
path: /dev/shm
- name: local
hostPath:
path: /mnt/k8s-disks/0
#nodeSelector:
# node.kubernetes.io/instance-type: "ml.g5.12xlarge"
containers:
- name: pytorch
image: 123456789012.dkr.ecr.us-west-2.amazonaws.com/hyprpd-tskgvrn-fsdp:pytorch2.2
imagePullPolicy: Always
resources:
requests:
nvidia.com/gpu: 4
vpc.amazonaws.com/efa: 1
limits:
nvidia.com/gpu: 4
vpc.amazonaws.com/efa: 1
env:
# for P5 FI_* should be commented out
- name: LOGLEVEL
value: 'DEBUG'
#- name: FI_PROVIDER
# value: efa
#- name: FI_EFA_USE_DEVICE_RDMA
# value: "1"
#- name: FI_EFA_FORK_SAFE
# value: "1"
#- name: FI_LOG_LEVEL
# value: "1"
#- name: FI_EFA_ENABLE_SHM_TRANSFER
# value: "1"
- name: TORCH_DISTRIBUTED_DEBUG
value: 'DETAIL'
- name: TORCH_NCCL_ENABLE_MONITORING
value: '1'
- name: TORCH_NCCL_TRACE_BUFFER_SIZE
value: '20000'
- name: TORCH_NCCL_DUMP_ON_TIMEOUT
value: '1'
- name: TORCH_NCCL_DEBUG_INFO_TEMP_FILE
value: '/local/nccl_trace_rank_'
- name: PYTORCH_CUDA_ALLOC_CONF
value: 'expandable_segments:True'
- name: NCCL_DEBUG
value: 'INFO'
- name: NCCL_SOCKET_IFNAME
value: '^lo'
- name: TORCH_NCCL_ASYNC_ERROR_HANDLING
value: '1'
#- name: TORCH_DIST_INIT_BARRIER
# value: "1"
#- name: NCCL_IGNORE_DISABLED_P2P
# value: "1"
#- name: NCCL_NVLS_ENABLE
# value: "0"
command:
- /usr/local/bin/torchrun
- --nproc_per_node=4
- --nnodes=2
- /fsdp/train.py
- --max_context_width=4096
- --num_key_value_heads=32
- --intermediate_size=11008
- --hidden_width=4096
- --num_layers=32
- --num_heads=32
- --model_type=llama_v2
- --tokenizer=hf-internal-testing/llama-tokenizer
- --checkpoint_freq=5000
- --validation_freq=500
- --max_steps=5000
- --checkpoint_dir=/checkpoints
- --dataset=allenai/c4
- --dataset_config_name=en
- --resume_from_checkpoint=/checkpoints
- --train_batch_size=1
- --val_batch_size=1
- --sharding_strategy=full
- --offload_activation=1
volumeMounts:
- name: shmem
mountPath: /dev/shm
- name: local
mountPath: /local
余していた GPU が使いきれました。
グラフも 50% から 100% に変化しています。
チームのサマリも Smith チームは John チームの余していた GPU を借りていますね。無事コンピュートを貸し借りできました。
まとめ
以上、「SageMaker HyperPod のタスクガバナンス機能を試してみた」でした。
各チームに一定割合のコンピュートを割り当てつつ、アイドル状態のリソースがあれば、貸し借りできるルールを設けられるのは非常に便利ですね。
今回はアイドル状態のコンピュートに重きを置いていましたが、ガバナンス周りはまた別のブログで書ければと思います。
このブログがどなたかの参考になれば幸いです。クラウド事業本部コンサルティング部のたかくに（@takakuni_）でした！
