在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文件内容如下：
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"
执行命令开始部署：
kubectl apply -f pytorch.yaml -n namespace
正常运行结束后可以看到两个pod均为complete状态，查看log输出：
training-operator的一些设计
当我们通过training-operator创建对应的pod资源时，describe worker可以看到如下的env信息：
可以看到master默认端口为23456，world_size即总共的运行节点为2，该worker对应的rank为1，master的地址为master的pod的name，这其实是training-operator为它创建了对应的svc，查看svc可以看到与pod同名的svc：
通过svc配置worker到master pod的tcp连接，以便在distributed.init_process_group以及接下来参数同步等需要网络传输的操作能够正常进行。
而training-operator通过解析yaml文件来为部署的pod配置环境变量以便init_process_group时不用自己配置分布式训练的参数，而是直接读取env配置好的参数，可以简化开发者的开发难度。不过这也能看出来，training-operator对pytorch分布式的支持实际上也是基于pytorch原有的pytorch分布式框架进行开发，在operator上添加相关代码给予支持。
总结
到此这篇关于在k8s上部署pytorch分布式程序的文章就介绍到这了,更多相关k8s部署pytorch分布式程序内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！