在 SageMaker HyperPod 中實施登錄節點負載平衡以增強多用戶體驗

Amazon SageMaker HyperPod 是為了支持大規模機器學習 (ML) 操作而設計的，提供一個穩定的環境來進行基礎模型 (FMs) 的長時間訓練。多位使用者，例如 ML 研究人員、軟體工程師、數據科學家和叢集管理員，可以同時在同一叢集上工作，每個人都能管理自己的工作和檔案，而不會干擾到其他人。

使用 HyperPod 時，可以使用熟悉的協調選項，例如 Slurm 或 Amazon Elastic Kubernetes Service (Amazon EKS)。這篇文章特別針對使用 Slurm 作為協調者的 HyperPod 叢集。在這些叢集中，登錄節點的概念可供叢集管理員添加，以便利用戶訪問。這些登錄節點作為用戶與叢集計算資源互動的入口。通過使用登錄節點，用戶可以將他們的互動活動（例如瀏覽檔案、提交工作和編譯代碼）與叢集的主節點分開。這種分離有助於防止任何單一用戶的活動影響整個叢集的性能。

然而，雖然 HyperPod 提供使用登錄節點的能力，但並不提供一個整合的機制來平衡這些節點上的用戶活動。因此，用戶需要手動選擇登錄節點，這可能導致某些節點過度使用，而其他節點則未被充分利用。這不僅影響資源使用的效率，還可能導致不同用戶的性能體驗不均衡。

在這篇文章中，我們探討了一個解決方案，用於在基於 Slurm 的 HyperPod 叢集中實施登錄節點的負載平衡。通過在所有可用節點之間均勻分配用戶活動，這種方法提供了更一致的性能、更好的資源利用率，以及更順暢的用戶體驗。我們將指導您完成設置過程，提供實用步驟以實現有效的負載平衡。

解決方案概述

在 HyperPod 中，登錄節點作為用戶與叢集計算資源互動的訪問點，讓他們能夠管理任務而不影響主節點。雖然訪問這些登錄節點的默認方法是通過 AWS Systems Manager，但在某些情況下，直接的安全外殼 (SSH) 訪問更為合適。SSH 提供了一種更傳統和靈活的管理互動方式，特別是對於需要特定網絡配置或需要 TCP 負載平衡等功能的用戶，這是 Systems Manager 不支持的。

考慮到 HyperPod 通常部署在使用私有子網的虛擬私有雲 (VPC) 中，直接 SSH 訪問登錄節點需要安全的網絡連接進入私有子網。有幾種選擇可以實現這一點：

AWS Site-to-Site VPN – 在您的本地網絡和 VPC 之間建立安全連接，適合企業環境
AWS Direct Connect – 提供高吞吐量和低延遲需求的專用網絡連接
AWS VPN Client – 一種基於軟體的解決方案，遠程用戶可以使用它安全地連接到 VPC，提供靈活且簡單的登錄節點訪問

這篇文章展示了如何使用 AWS VPN Client 建立到 VPC 的安全連接。我們在私有子網中設置了一個網絡負載平衡器 (NLB)，以均勻分配 SSH 流量到可用的登錄節點，並使用 VPN 連接到 VPC 中的 NLB。NLB 確保用戶會話在節點之間平衡，防止任何單一節點成為瓶頸，從而改善整體性能和資源利用率。

對於 VPN 連接可能不可行的環境，另一種選擇是在公共子網中部署 NLB，以允許從互聯網直接 SSH 訪問。在這種配置中，NLB 可以通過限制安全組的訪問來加以保護，只允許來自指定的受信 IP 地址的 SSH 流量。因此，授權用戶可以直接連接到登錄節點，同時保持對叢集訪問的某種控制。然而，這種面向公共的方式不在本篇文章的範疇內，並且不建議用於生產環境，因為將 SSH 訪問暴露於互聯網可能會引入額外的安全風險。

以下圖表提供了解決方案架構的概述。

前提條件

在遵循本篇文章中的步驟之前，請確保您已經設置了 HyperPod 叢集的基本組件。這包括 HyperPod 叢集的核心基礎設施和安全訪問所需的網絡配置。具體來說，您需要：

HyperPod 叢集 – 本篇文章假設您已經部署了 HyperPod 叢集。如果沒有，請參考《開始使用 SageMaker HyperPod》和 HyperPod 工作坊以獲取創建和配置叢集的指導。
VPC、子網和安全組 – 您的 HyperPod 叢集應該位於一個具有相關子網的 VPC 中。要部署新的 VPC 和子網，請按照 HyperPod 工作坊的《擁有帳戶》部分中的說明進行。此過程包括部署 AWS CloudFormation 堆疊以創建基本資源，例如 VPC、子網、安全組和 Amazon FSx for Lustre 共享存儲卷。

設置登錄節點以訪問叢集

登錄節點是專用的訪問點，用戶可以通過它們與 HyperPod 叢集的計算資源互動，而不影響主節點。通過登錄節點連接，用戶可以獨立瀏覽檔案、提交工作和編譯代碼，促進叢集資源的更有序和高效使用。

如果您尚未設置登錄節點，請參考 HyperPod 工作坊中的登錄節點部分，該部分提供了將這些節點添加到叢集配置的詳細說明。

每個 HyperPod 叢集中的登錄節點都有一個與您的 VPC 相關聯的網絡介面。網絡介面，也稱為彈性網絡介面，代表一個虛擬網絡卡，將每個登錄節點連接到您的 VPC，使其能夠在網絡上進行通信。這些介面有分配的 IPv4 地址，這對於將流量從 NLB 路由到登錄節點至關重要。

要繼續進行負載平衡器的設置，您需要獲取每個登錄節點的 IPv4 地址。您可以從 AWS 管理控制台獲取這些地址，或在 HyperPod 叢集的主節點上執行命令。

使用 AWS 管理控制台

要使用 AWS 管理控制台設置登錄節點以訪問叢集，請按照以下步驟操作：

在 Amazon EC2 控制台中，選擇導航窗格中的網絡介面
在搜索欄中，選擇 VPC ID = (等於) 並選擇包含 HyperPod 叢集的 VPC ID
在搜索欄中，選擇描述 : (包含) 並輸入包含您的登錄節點的實例組的名稱（通常是 login-group）

對於每個登錄節點，您將在列表中找到一個條目，如下圖所示。記下您叢集所有登錄節點的 IPv4 地址。

使用 HyperPod 主節點

另外，您也可以通過在 HyperPod 叢集的主節點上輸入以下命令來檢索 IPv4 地址：

sudo cat /opt/ml/config/resource_config.json \ | jq ‘.InstanceGroups[] | select(.Name==”login-group”).Instances[].CustomerIpAddress’

創建網絡負載平衡器

下一步是創建 NLB 以管理叢集登錄節點之間的流量。

對於 NLB 部署，您需要之前收集的登錄節點的 IPv4 地址和適當的安全組配置。如果您按照 HyperPod 工作坊的說明部署了叢集，則應已經存在允許所有叢集節點之間通信的安全組。

這個安全組可以應用於負載平衡器，如以下說明所示。或者，您可以選擇創建一個專用的安全組，專門授予對登錄節點的訪問。

創建目標組

首先，我們創建將由 NLB 使用的目標組。

在 Amazon EC2 控制台中，選擇導航窗格中的目標組
選擇創建目標組
創建一個目標組，參數如下：

對於選擇目標類型，選擇 IP 地址
對於目標組名稱，輸入 smhp-login-node-tg
對於協議 : 端口，選擇 TCP 並輸入 22
對於 IP 地址類型，選擇 IPv4
對於 VPC，選擇 SageMaker HyperPod VPC（這是使用 CloudFormation 模板創建的）
對於健康檢查協議，選擇 TCP

選擇下一步，如下圖所示

在註冊目標部分，將登錄節點的 IP 地址註冊為目標
對於端口，輸入 22 並選擇包含待處理的選項，如下圖所示

登錄節點的 IP 將顯示為待處理健康狀態的目標。選擇創建目標組，如下圖所示

創建負載平衡器

要創建負載平衡器，請按照以下步驟操作：

在 Amazon EC2 控制台中，選擇導航窗格中的負載平衡器
選擇創建負載平衡器
選擇網絡負載平衡器並選擇創建，如下圖所示

提供一個名稱（例如，smhp-login-node-lb）並選擇內部作為方案

對於網絡映射，選擇包含您的 HyperPod 叢集的 VPC 和相關的私有子網，如下圖所示

選擇允許對登錄節點的端口 22 訪問的安全組。如果您按照 HyperPod 工作坊的說明部署了叢集，則可以使用此部署的安全組。

選擇您之前創建的目標組，並選擇 TCP 作為協議，端口為 22，如下圖所示

選擇創建負載平衡器

在負載平衡器創建後，您可以在負載平衡器的詳細頁面上找到其 DNS 名稱，如下圖所示。

確保登錄節點之間的主機密鑰一致

在負載平衡環境中使用多個登錄節點時，保持所有節點之間的 SSH 主機密鑰一致非常重要。SSH 主機密鑰是每個伺服器用來向連接的客戶端證明其身份的唯一標識符。如果每個登錄節點都有不同的主機密鑰，則用戶在連接到不同的節點時會遇到“警告：SSH 主機密鑰已更改”的消息，這會造成困惑，並可能導致用戶質疑連接的安全性。

為了避免這些警告，請在負載平衡輪換中的所有登錄節點上配置相同的 SSH 主機密鑰。這樣的設置確保用戶在被負載平衡器路由到不同的節點時不會收到主機密鑰不匹配的警報。

您可以在叢集的主節點上輸入以下腳本，將第一個登錄節點的 SSH 主機密鑰複製到 HyperPod 叢集中的其他登錄節點：

#!/bin/bash

SUDOER_USER=”ubuntu”

login_nodes=($(sudo cat /opt/ml/config/resource_config.json | jq ‘.InstanceGroups[] | select(.Name==”login-group”).Instances[].CustomerIpAddress’ | tr ‘\n’ ‘ ‘ | tr -d ‘”‘))
source_node=”$login_nodes[0]”
key_paths=(“/etc/ssh/ssh_host_rsa_key”
“/etc/ssh/ssh_host_rsa_key.pub”
“/etc/ssh/ssh_host_ecdsa_key”
“/etc/ssh/ssh_host_ecdsa_key.pub”
“/etc/ssh/ssh_host_ed25519_key”
“/etc/ssh/ssh_host_ed25519_key.pub”)

tmp_dir=”/tmp/ssh_host_keys_$(uuidgen)”

copy_cmd=””
for key_path in “$key_paths[@]”; do
copy_cmd=”sudo cp $key_path $tmp_dir/;$copy_cmd”
done

ssh $source_node “mkdir -p $tmp_dir;$copy_cmd sudo chown -R $SUDOER_USER $tmp_dir;”

for node in “$login_nodes[@]:1”; do
echo “從 $source_node 複製 SSH 主機密鑰到 $node…”
scp -r $source_node:$tmp_dir $node:$tmp_dir
ssh $node “sudo chown -R root:root $tmp_dir; sudo mv $tmp_dir/ssh_host_* /etc/ssh/;”
done

for node in “$login_nodes[@]”; do
echo “在 $node 上清理臨時目錄 $tmp_dir…”
ssh $node “sudo rm -r $tmp_dir”
done

創建 AWS 客戶端 VPN 端點

由於 NLB 是以內部方案創建的，因此只能從 HyperPod VPC 內部訪問。要訪問 VPC 並向 NLB 發送請求，我們在這篇文章中使用 AWS 客戶端 VPN。

AWS 客戶端 VPN 是一種管理的基於客戶端的 VPN 服務，使您能夠安全訪問 AWS 資源和您本地網絡中的資源。

我們將設置一個 AWS 客戶端 VPN 端點，為客戶提供訪問 HyperPod VPC 的權限，並使用雙向身份驗證。通過雙向身份驗證，客戶端 VPN 使用證書在客戶端和客戶端 VPN 端點之間進行身份驗證。

要使用雙向身份驗證部署客戶端 VPN 端點，您可以遵循《開始使用 AWS 客戶端 VPN》中概述的步驟。在配置客戶端 VPN 以訪問 HyperPod VPC 和登錄節點時，請記住對以下步驟進行這些調整：

步驟 2（創建客戶端 VPN 端點） – 默認情況下，所有客戶端流量都通過客戶端 VPN 隧道路由。要允許在不通過 VPN 路由流量的情況下訪問互聯網，您可以在創建端點時啟用選項“啟用分隧道”。啟用此選項後，只有目標為客戶端 VPN 端點路由表中匹配路由的流量才會通過 VPN 隧道路由。詳細信息，請參見《客戶端 VPN 端點的分隧道》。
步驟 3（目標網絡關聯） – 選擇 HyperPod 叢集使用的 VPC 和私有子網，其中包含叢集登錄節點。
步驟 4（授權規則） – 選擇與 HyperPod VPC 相關的無類別域間路由 (CIDR) 範圍。如果您遵循了 HyperPod 工作坊的說明，則 CIDR 範圍為 10.0.0.0/16。
步驟 6（安全組） – 選擇您之前在創建 NLB 時使用的安全組。

連接到登錄節點

在 AWS 客戶端 VPN 配置完成後，客戶端可以建立到 HyperPod VPC 的 VPN 連接。建立 VPN 連接後，客戶端可以使用 SSH 連接到 NLB，該 NLB 將把他們路由到其中一個登錄節點。

ssh -i /path/to/your/private-key.pem user@<NLB-IP-or-DNS>

要允許對登錄節點的 SSH 訪問，您必須在叢集中創建用戶帳戶，並將他們的公鑰添加到每個登錄節點的 authorized_keys 檔案中（或在所有節點上，如果需要）。有關管理多用戶訪問的詳細說明，請參見 HyperPod 工作坊的多用戶部分。

除了使用 AWS 客戶端 VPN，您還可以從其他 AWS 服務（例如 Amazon Elastic Compute Cloud (Amazon EC2) 實例）訪問 NLB，只要它們滿足以下要求：

VPC 連接 – EC2 實例必須與 NLB 位於同一 VPC 中，或能夠通過對等連接或類似的網絡設置訪問 HyperPod VPC。
安全組配置 – EC2 實例的安全組必須允許對 NLB 安全組的端口 22 進行出站連接。同樣，NLB 安全組應配置為接受來自 EC2 實例安全組的端口 22 的入站 SSH 流量。

清理

要刪除已部署的資源，您可以按照以下順序進行清理：

刪除客戶端 VPN 端點
刪除網絡負載平衡器
刪除與負載平衡器相關的目標組

如果您還想刪除 HyperPod 叢集，請遵循以下額外步驟：

刪除 HyperPod 叢集
刪除 CloudFormation 堆疊，其中包括 VPC、子網、安全組和 FSx for Lustre 卷

結論

在這篇文章中，我們探討了如何為 SageMaker HyperPod 叢集實施登錄節點的負載平衡。通過使用網絡負載平衡器將用戶流量分配到登錄節點，您可以優化資源利用率，並增強整體多用戶體驗，為每個用戶提供無縫訪問叢集資源的能力。

這種方法僅代表自定義 HyperPod 叢集的一種方式。由於 SageMaker HyperPod 的靈活性，您可以根據自己的獨特需求調整配置，同時受益於管理的、彈性的環境。無論您需要擴展基礎模型工作負載、在不同任務之間共享計算資源，還是支持長時間運行的訓練工作，SageMaker HyperPod 都提供了一個可以隨著需求發展的多功能解決方案。

有關如何充分利用 SageMaker HyperPod 的更多詳細信息，請深入了解 HyperPod 工作坊，並探索更多涵蓋 HyperPod 的博客文章。

關於作者

Janosch Woschitz 是 AWS 的高級解決方案架構師，專注於 AI/ML。擁有超過 15 年的經驗，他在全球支持客戶利用 AI 和 ML 進行創新解決方案，並在 AWS 上構建 ML 平台。他的專業知識涵蓋機器學習、數據工程和可擴展的分佈式系統，並擁有強大的軟體工程背景和自動駕駛等領域的行業專業知識。

Giuseppe Angelo Porcelli 是 Amazon Web Services 的首席機器學習專家解決方案架構師。擁有幾年的軟體工程和 ML 背景，他與各種規模的客戶合作，了解他們的業務和技術需求，並設計最佳利用 AWS 雲和 Amazon 機器學習堆疊的 AI 和 ML 解決方案。他曾參與不同領域的項目，包括 MLOps、計算機視覺和自然語言處理，涉及廣泛的 AWS 服務。在空閒時間，Giuseppe 喜歡踢足球。