計(jì)實(shí)戰(zhàn):在 GKE 上以 NVIDIA A100 80GB GPU 運(yùn)行 JAX 多節(jié)點(diǎn)分布式「Hello World」)
Meshery 設(shè)計(jì)實(shí)戰(zhàn)在 GKE 上以 NVIDIA A100 80GB GPU 運(yùn)行 JAX 多節(jié)點(diǎn)分布式「Hello World」【免費(fèi)下載鏈接】mesheryMeshery, the cloud native manager項(xiàng)目地址: https://gitcode.com/GitHub_Trending/me/meshery本篇基于 Meshery 官方 Catalog 中的部署設(shè)計(jì)Deployment Design「JAX Hello World using NVIDIA GPUs A100-80GB on GKE」講解如何把一套 JAX 多機(jī)多卡分布式訓(xùn)練示例以可視化 Design 的形式托管在 Meshery 中并在 Google Kubernetes EngineGKE的 A2 系列節(jié)點(diǎn)NVIDIA A100 80GB Tensor Core GPU上完成部署。讀完本文你將掌握該 Design 中每個(gè) Kubernetes 組件Indexed Job、GPU 資源聲明、Headless Service的作用與協(xié)作方式并能使用mesheryctl design命令將其導(dǎo)入、查看并應(yīng)用到自己的集群。背景為什么要在 GKE 上跑 JAX 多節(jié)點(diǎn)任務(wù)JAX 是一個(gè)快速發(fā)展的 Python 高性能數(shù)值計(jì)算與機(jī)器學(xué)習(xí)ML研究庫廣泛應(yīng)用于大語言模型、藥物發(fā)現(xiàn)、物理 ML、強(qiáng)化學(xué)習(xí)和神經(jīng)圖形學(xué)等領(lǐng)域。它對(duì)開發(fā)者與研究者極具吸引力一方面提供易用的 NumPy 風(fēng)格 API內(nèi)置自動(dòng)微分auto differentiation與優(yōu)化能力另一方面僅需少量代碼即可在多節(jié)點(diǎn)、多 GPU 系統(tǒng)上發(fā)起分布式處理并通過 NVIDIA GPU 上 XLA 優(yōu)化的內(nèi)核獲得加速性能。然而「少量代碼就能分布式」的背后是集群層面繁瑣的準(zhǔn)備工作需要為每個(gè)訓(xùn)練進(jìn)程分配 GPU 資源、讓各進(jìn)程通過 DNS 互相發(fā)現(xiàn)、指定協(xié)調(diào)者coordinator端口并保證進(jìn)程間網(wǎng)絡(luò)可達(dá)。本目錄條目所描述的 Design正是把這一類集群級(jí)配置以聲明式、可復(fù)用的方式固化下來它在 4 個(gè)節(jié)點(diǎn)上運(yùn)行一個(gè)簡單的 Hello World 應(yīng)用每個(gè)節(jié)點(diǎn) 8 個(gè)進(jìn)程、使用 8 張 GPU即共 32 個(gè)訓(xùn)練進(jìn)程、32 張 A100 GPU。該設(shè)計(jì)的目錄條目文檔位于 docs/catalog/deployment/8b041687-3c09-4cfd-8613-cf326a54e1b2.md可執(zhí)行的完整設(shè)計(jì)定義位于 design.yml。設(shè)計(jì)文件整體結(jié)構(gòu)該設(shè)計(jì)遵循 Meshery 設(shè)計(jì)格式schemaVersion: designs.meshery.io/v1beta1由Components組件與Relationships關(guān)系兩部分構(gòu)成。設(shè)計(jì)中的組件全部來自kubernetes模型類型為 deployment兼容目標(biāo)為kubernetes。整體組成如下組件Kubernetes 類型作用defaultNamespacev1部署的目標(biāo)命名空間job-name主 JobJobbatch/v1定義jax-worker容器與分布式協(xié)調(diào)參數(shù)completionMode: Indexedjob-nameGPU 覆蓋Jobbatch/v1追加 GPU 資源請求nvidia.com/gpu: 1與 GPU 污點(diǎn)容忍service-nameServicev1Headless ServiceclusterIP: None為進(jìn)程間 DNS 發(fā)現(xiàn)提供穩(wěn)定地址其配套的包元數(shù)據(jù)文件 artifacthub-pkg.yml 記錄了該設(shè)計(jì)的版本0.0.1、描述、許可證Apache-2.0以及安裝命令提示mesheryctl design import -f。逐個(gè)組件拆解1. Namespace部署邊界設(shè)計(jì)中第一個(gè)組件是名為default的 Namespacekind: Namespace,version: v1。它界定了后續(xù) Job 與 Service 的部署范圍。該組件帶有 Meshery 為 Kubernetes 組件注冊的通用能力capabilities例如 Performance Test性能測試、Workload Configuration工作負(fù)載配置、Labels and Annotations Configuration 以及 Relationships 查看等說明在 Meshery UI 中你可以直接對(duì)該組件發(fā)起這些操作。2. 主 Job定義 JAX 分布式協(xié)調(diào)邏輯核心組件是一個(gè)batch/v1類型的 Job它定義了訓(xùn)練容器jax-workerspec: template: spec: subdomain: headless-svc containers: - name: jax-worker image: gcr.io/PROJECT/jax/hello:latest command: [python, train.py] args: - --num_processes - WILL_BE_REPLACED - --job_name - WILL_BE_REPLACED - --sub_domain - WILL_BE_REPLACED - --coordinator_port - WILL_BE_REPLACED ports: - containerPort: 1234 restartPolicy: Never completions: 1 parallelism: 1 backoffLimit: 1 completionMode: Indexed這段配置對(duì)應(yīng) JAX 分布式初始化jax.distributed.initialize所需的關(guān)鍵參數(shù)理解它們是把任務(wù)跑起來的前提--num_processes參與訓(xùn)練的總進(jìn)程數(shù)。按目錄條目的說明4 節(jié)點(diǎn) × 8 進(jìn)程 32此處應(yīng)替換為32。--job_name作業(yè)名用于在集群內(nèi)區(qū)分不同的訓(xùn)練作業(yè)。--sub_domain子域名與 Pod 模板中的subdomain: headless-svc對(duì)應(yīng)是進(jìn)程間通過 DNS 互相發(fā)現(xiàn)的依據(jù)。--coordinator_port協(xié)調(diào)者監(jiān)聽端口與容器聲明的containerPort: 1234一致train.py內(nèi)部會(huì)據(jù)此建立 gRPC 通信。WILL_BE_REPLACED是設(shè)計(jì)作者留下的占位符實(shí)際部署前必須替換為真實(shí)值——這也正是「模板化設(shè)計(jì)」的典型用法。鏡像gcr.io/PROJECT/jax/hello:latest同樣包含占位符PROJECT需要替換為你自己的 GCR 項(xiàng)目名該鏡像用于在每臺(tái) A100 節(jié)點(diǎn)上啟動(dòng) JAX 訓(xùn)練進(jìn)程。3. GPU 覆蓋 Job請求與容忍 A100 GPU第二個(gè) Job 組件是對(duì)上一步的增強(qiáng)/覆蓋在 Meshery 中以獨(dú)立組件形式存在通過 Relationships 關(guān)聯(lián)它向 Pod 模板追加 GPU 資源聲明與污點(diǎn)容忍containers: - name: jax-worker resources: limits: nvidia.com/gpu: 1 tolerations: - key: nvidia.com/gpu effect: NoSchedule operator: Existsnvidia.com/gpu: 1通過 NVIDIA Device Plugin 暴露的擴(kuò)展資源聲明每個(gè)jax-worker容器需要 1 張 GPU。調(diào)度器會(huì)據(jù)此將 Pod 放置到具備可用 A100 GPU 的 A2 節(jié)點(diǎn)上。容忍tolerationsGKE 的 GPU 節(jié)點(diǎn)通常帶有nvidia.com/gpu: NoSchedule污點(diǎn)只有聲明了對(duì)應(yīng)容忍的 Pod 才允許被調(diào)度上去operator: Exists表示只要污點(diǎn)鍵存在即容忍無需匹配具體值。4. Headless Service進(jìn)程間 DNS 發(fā)現(xiàn)最后是v1類型的 Service其配置極其精簡卻至關(guān)重要spec: selector: job-name: job-name clusterIP: NoneclusterIP: None表示這是一個(gè) Headless Service。Kubernetes 不會(huì)為其分配 ClusterIP而是為每個(gè)匹配的 Pod 生成獨(dú)立的 DNS 記錄。selector: {job-name: job-name}選中該 Job 產(chǎn)生的所有 Pod。配合主 Job 中 Pod 模板的subdomain: headless-svc每個(gè)jax-workerPod 都可以通過pod-name.headless-svc.namespace.svc.cluster.local形式的穩(wěn)定域名被其他節(jié)點(diǎn)上的進(jìn)程解析到。這正是 JAX 多節(jié)點(diǎn)協(xié)調(diào)中「進(jìn)程之間如何找到彼此」的實(shí)現(xiàn)機(jī)制--sub_domain指向的就是這個(gè) headless-svc。Relationships組件之間如何被串起來與三個(gè)組件同級(jí)的還有一組 RelationshipsschemaVersion: relationships.meshery.io/v1alpha3它們描述了組件間的關(guān)聯(lián)語義Meshery 在部署與清理時(shí)會(huì)據(jù)此推導(dǎo)正確的處理順序與依賴關(guān)系hierarchical parentinventoryJob/PodTemplate 與 Namespace 之間的父子歸屬關(guān)系——組件所屬的命名空間由 Namespace 組件決定通過mutatorRef/mutatedRef把命名空間注入各組件配置。hierarchical siblingmatchlabelsJob 與 Service 之間基于標(biāo)簽的兄弟關(guān)系Service 通過job-name: job-name選擇器與 Job 產(chǎn)生的 Pod 關(guān)聯(lián)。關(guān)系還帶有approved/deleted狀態(tài)說明設(shè)計(jì)在保存與迭代過程中關(guān)系會(huì)被重新評(píng)估。這些關(guān)系體現(xiàn)了 Meshery 設(shè)計(jì)的核心思想設(shè)計(jì)是描述性、聲明式的組件是構(gòu)建塊關(guān)系是它們之間的語義連接詳見 Designs 概念文檔。部署到 GKE A2 節(jié)點(diǎn)前置條件與注意事項(xiàng)目錄條目的 patternCaveats 明確提醒「確保網(wǎng)絡(luò)配置正確并且每個(gè)資源都應(yīng)用了正確的注解」Ensure networking is setup properly and correct annotation are applied to each resource。結(jié)合設(shè)計(jì)內(nèi)容部署前應(yīng)確認(rèn)以下幾點(diǎn)節(jié)點(diǎn)池集群需包含使用 A2 超算系列機(jī)型、配備 A100 80GB GPU 的節(jié)點(diǎn)池GKE 中 A2 機(jī)型即面向 A100 的節(jié)點(diǎn)池。GPU 驅(qū)動(dòng)與 Device PluginGKE 默認(rèn)節(jié)點(diǎn)池可開啟 GPU 驅(qū)動(dòng)安裝Driver/Device Plugin 就緒后nvidia.com/gpu擴(kuò)展資源才會(huì)在節(jié)點(diǎn)上可見Job 才能被調(diào)度。網(wǎng)絡(luò)與注解headless-svc 依賴集群內(nèi) DNSCoreDNS正常工作若啟用了網(wǎng)絡(luò)策略或 Service Mesh需要保證 1234 等協(xié)調(diào)端口在 Pod 之間可達(dá)并按設(shè)計(jì)提示為資源補(bǔ)充正確的注解。替換占位符將PROJECT鏡像倉庫項(xiàng)目與四個(gè)WILL_BE_REPLACED參數(shù)替換為真實(shí)值后再行部署。從倉庫中的 Deployment Engine 概念文檔 可以推斷該設(shè)計(jì)的履約路徑其全部組件都來自kubernetes模型注冊者registrant是已連接的 Kubernetes 集群本身不對(duì)外暴露可供 Meshery 調(diào)用的網(wǎng)絡(luò)端口因此這些組件會(huì)由 Meshery Server 通過進(jìn)程內(nèi)的 Kubernetes 客戶端直接應(yīng)用到所選集群Path A無需 Meshery Adapter 參與。這是大多數(shù) Catalog 中純 Kubernetes 設(shè)計(jì)的典型履約方式。使用 mesheryctl 導(dǎo)入并管理該設(shè)計(jì)該目錄條目的設(shè)計(jì)文件以標(biāo)準(zhǔn) YAML 格式存放你可以用 Meshery CLI 將其導(dǎo)入自己的 Meshery 實(shí)例。mesheryctl design import支持從本地文件系統(tǒng)路徑或遠(yuǎn)程 URL 導(dǎo)入 Helm Chart、Kubernetes Manifest、Docker Compose 或 Meshery 設(shè)計(jì)詳見 mesheryctl design import 命令參考# 導(dǎo)入 Meshery 設(shè)計(jì)source-type 為 design可省略 mesheryctl design import -f design.yml -s design -n jax-hello-gke # 從本地文件導(dǎo)入并自定義名稱 mesheryctl design import -f design.yml -n design-name導(dǎo)入后按 Catalog 概念文檔 所列的 CLI 用法還可以這樣管理設(shè)計(jì)# 查看/列出設(shè)計(jì) mesheryctl design list mesheryctl design view [design name | ID] # 應(yīng)用部署設(shè)計(jì)到當(dāng)前連接的集群 mesheryctl design apply --file design.yml # 刪除設(shè)計(jì) mesheryctl design delete --file design.yml如果你從 design.yml 下載文件到本地即可直接執(zhí)行上面的導(dǎo)入命令。設(shè)計(jì)被導(dǎo)入后即成為你賬號(hào)下的可部署單元可在 Meshery UI 中以可視化畫布形式查看組件與關(guān)系再一鍵部署到所選環(huán)境。延伸從設(shè)計(jì)到模式Patterns值得說明的是本目錄條目被標(biāo)記為type: deployment其內(nèi)容是完整的、可直接部署的設(shè)計(jì)。而 Meshery 生態(tài)中還區(qū)分了更高一層的抽象——Patterns模式Pattern 是模板化的設(shè)計(jì)把復(fù)雜設(shè)計(jì)封裝為單一可復(fù)用組件隱藏底層復(fù)雜度便于在團(tuán)隊(duì)內(nèi)傳播最佳實(shí)踐。根據(jù)倉庫中的 Patterns 概念文檔Patterns 是規(guī)劃中的路線圖特性目標(biāo)版本 v0.9.0。在此之前Catalog 中的設(shè)計(jì)如本文介紹的 JAX 部署設(shè)計(jì)正是共享與復(fù)用基礎(chǔ)設(shè)施配置的主要載體你可以在 Catalog 概念文檔 中了解將設(shè)計(jì)發(fā)布到 Catalog 的完整審批流程提交 → 管理員審閱 → 校驗(yàn) → 發(fā)布并觸發(fā) GitHub Workflow 同步到 Meshery.io。小結(jié)「JAX Hello World using NVIDIA GPUs A100-80GB on GKE」這一 Catalog 設(shè)計(jì)以四個(gè) Kubernetes 組件加一組關(guān)系完整描述了在 GKE A2 節(jié)點(diǎn)上運(yùn)行 32 進(jìn)程 JAX 分布式任務(wù)所需的一切Indexed Job 定義分布式協(xié)調(diào)參數(shù)GPU 資源聲明與污點(diǎn)容忍保證 Pod 落到 A100 節(jié)點(diǎn)Headless Service 提供進(jìn)程間 DNS 發(fā)現(xiàn)。通過mesheryctl design import將其導(dǎo)入 Meshery你可以在可視化畫布上審視其結(jié)構(gòu)替換占位符后一鍵部署——這套「以設(shè)計(jì)為單位的聲明式部署」工作流也正是 Meshery 作為云原生管理器管理 GPU 訓(xùn)練工作負(fù)載的典型范式?!久赓M(fèi)下載鏈接】mesheryMeshery, the cloud native manager項(xiàng)目地址: https://gitcode.com/GitHub_Trending/me/meshery創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考