快捷導(dǎo)航

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

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

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

集群配置

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

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

下載Kubeflow training-operator對應(yīng)的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

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

training-operator的一些設(shè)計

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

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

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

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