在k8s上部署pytorch分布式程序的完整步驟記錄

更新時間：2024年08月23日 11:17:42 作者：道阻&且長

Kubernetes的核心優(yōu)勢在于其能夠提供一個可擴展、靈活且高度可配置的平臺,使得應用程序的部署、擴展和管理變得前所未有的簡單下面這篇文章主要給大家介紹了關于在k8s上部署pytorch分布式程序的完整步驟,需要的朋友可以參考下

集群配置

在k8s集群安裝stable版training-operator：

kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=v1.7.0"

下載Kubeflow training-operator對應的Python SDK：

pip3 install kubeflow-training

簡單的pytorch CPU分布式測試

demo.py文件內(nèi)容如下：

import datetime
import torch
torch.distributed.init_process_group(init_method="env://",timeout=datetime.timedelta(minutes=1))
rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()
print(f"rank {rank} world_size {world_size}")
a = torch.tensor([1])
torch.distributed.all_reduce(a)
print(f"rank {rank} world_size {world_size} result {a}")
torch.distributed.barrier()
print(f"rank {rank} world_size {world_size}")

Dockerfile文件如下：

FROM python:3.8
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
WORKDIR /
COPY demo.py demo.py

打包好鏡像上傳后，便可以在集群中跑pytorchjob，pytorch.yaml部署文件如下：

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: torchrun-cpu
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch-demo
              command:
                - "python3"
                - "demo.py"
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch-demo
              command:
                - "python3"
                - "demo.py"

執(zhí)行命令開始部署：

kubectl apply -f pytorch.yaml -n namespace

正常運行結束后可以看到兩個pod均為complete狀態(tài)，查看log輸出：

training-operator的一些設計

當我們通過training-operator創(chuàng)建對應的pod資源時，describe worker可以看到如下的env信息：

可以看到master默認端口為23456，world_size即總共的運行節(jié)點為2，該worker對應的rank為1，master的地址為master的pod的name，這其實是training-operator為它創(chuàng)建了對應的svc，查看svc可以看到與pod同名的svc：

通過svc配置worker到master pod的tcp連接，以便在distributed.init_process_group以及接下來參數(shù)同步等需要網(wǎng)絡傳輸?shù)牟僮髂軌蛘＿M行。

而training-operator通過解析yaml文件來為部署的pod配置環(huán)境變量以便init_process_group時不用自己配置分布式訓練的參數(shù)，而是直接讀取env配置好的參數(shù)，可以簡化開發(fā)者的開發(fā)難度。不過這也能看出來，training-operator對pytorch分布式的支持實際上也是基于pytorch原有的pytorch分布式框架進行開發(fā)，在operator上添加相關代碼給予支持。