阿里云Kubernetes 1.9上利用Helm运行TensorFlow 分布式模型训练

阿里云Kubernetes 1.9上利用Helm运行TensorFlow 分布式模型训练

TensorFlow是业界最流行的深度学习框架， 但是如何将TensorFlow真正运用于生产环境却并不简单，它面临着资源隔离，应用调度和部署，GPU资源分配，训练生命周期管理等挑战。特别是大规模的分布式训练场景， 单靠手动部署和人力运维已经无法有效处理。特别启动每个模块都需要指定好分布式集群的clusterSpec, 更是让人挠头。

在Kubernetes集群上运行分布式TensorFlow模型训练，可以依靠Kubernetes本身在应用调度，GPU资源分配，共享存储等方面的能力，实现训练任务和参数服务器的调度以及生命周期的管理。同时利用共享存储查看训练的收敛程度，调整超参。

但是手动写部署Yaml对于最终用户来说还是非常酸爽的，阿里云容器服务提供了基于Helm的TensorFlow分布式训练解决方案：

• 同时支持GPU和非GPU集群
• 不再需要手动配置clusterspec信息，只需要指定worker和ps的数目，能自动生成clusterspec
• 内置Tensorboard可以有效监控训练的收敛性，方便快速调整参数epoch，batchsize， learning rate

以下就是一个利用Helm运行端到端的分布式模型训练示例：

1. 准备数据

1.1 创建NAS文件存储，并且设置vpc内挂载点。可以参考阿里云NAS文档。并且查看挂载点，这里假设挂载点为aliyunxxxx.cn-shanghai.nas.aliyuncs.com

1.2 准备名字为/data的数据文件夹

mkdir /nfs
mount -t nfs -o vers=4.0 aliyunxxxx.cn-shanghai.nas.aliyuncs.com:/ /nfs
mkdir -p /nfs/data
umount /nfs

2. 创建persistent volume

以下为创建NAS的nas.yaml样例，实际上也可以创建云盘或者OSS等持久化存储

---
apiVersion: v1
kind: PersistentVolume
metadata:
  labels:
    train: mnist
  name: pv-nas-train
spec:
  persistentVolumeReclaimPolicy: Retain
  accessModes:
    - ReadWriteMany
  capacity:
    storage: 5Gi
  flexVolume:
    driver: alicloud/nas
    options:
      mode: "755"
      path: /data
      server: aliyunxxxx.cn-shanghai.nas.aliyuncs.com
      vers: "4.0"

注意这里需要指定label为model: mnist, storageClassName需要为nas, 这两个标签对于pvc选择pv绑定非常重要。
另外和NAS相关的具体配置可以参考Kubernetes使用阿里云NAS

运行kubectl命令创建

kubectl create -f nas.yaml
persistentvolume "pv-nas" created

部署完成后，可以通过dashboard检查运行状态：

3. 通过Helm部署TensorFlow分布式训练的应用

3.1 可以通过应用目录，点击acs-tensorflow-training

以下为支持GPU的自定义配置参数的training.yaml文件

---
# Default values for acs-dl-distributed-training.
# This is a YAML-formatted file.
# Declare variables to be passed into your templates.
worker:
  number: 2
  gpuCount: 1
  image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/tf-mnist-k8s:gpu
  imagePullPolicy: IfNotPresent
  # if you'd like to choose the cusomtized docker image, 
  #image: ""
  port: 8000

ps:
  number: 2
  image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/tf-mnist-k8s:cpu
  imagePullPolicy: IfNotPresent
   # if you'd like to choose the cusomtized docker image, 
  #image: ""
  port: 9000

tensorboard:
  image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/tensorboard:1.1.0
  serviceType: LoadBalancer

persistence: 
  mountPath: /data
  pvc: 
    matchLabels: 
      train: mnist
    storage: 5Gi

如果你运行的Kubernetes集群不含有GPU可以使用一下配置

---
worker:
  number: 2
  # if you'd like to choose the cusomtized docker image
  image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/tf-mnist-k8s:cpu
  imagePullPolicy: IfNotPresent

ps:
  number: 2
  # if you'd like to choose the cusomtized docker image
  image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/tf-mnist-k8s:cpu
  imagePullPolicy: IfNotPresent

tensorboard:
  image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/tensorboard:1.1.0
  serviceType: LoadBalancer

hyperparams:
  epochs: 100
  batchsize: 20
  learningrate: 0.001

persistence:
  mountPath: /data
  pvc:
    matchLabels:
      train: mnist
    storage: 5Gi

这里镜像的参考代码来自于；https://github.com/cheyang/tensorflow-sample-code

3.2 点击参数, 就可以通过修改参数配置点击部署

也可运行helm命令部署

helm install --values values.yaml --name mnist incubator/acs-tensorflow-tarining
helm install --debug --dry-run --values values.yaml --name mnist incubator/acs-tensorflow-tarining

3.3 部署完成后，可以查看应用运行状态

4. 利用helm命令查看部署的信息

4.1 登录到Kubernetes的master上利用helm命令查看部署应用的列表

# helm list
NAME              REVISION  UPDATED                   STATUS    CHART                         NAMESPACE
mnist-dist-train  1         Mon Mar 19 15:23:51 2018  DEPLOYED  acs-tensorflow-training-0.1.0 default

4.2 利用helm status命令检查具体应用的配置

# helm status mnist-dist-train
LAST DEPLOYED: Mon Mar 19 15:23:51 2018
NAMESPACE: default
STATUS: DEPLOYED

RESOURCES:
==> v1/ConfigMap
NAME             DATA  AGE
tf-cluster-spec  1     7m

==> v1/Service
NAME         TYPE       CLUSTER-IP     EXTERNAL-IP   PORT(S)   AGE
worker-0     ClusterIP  None           <none>        8000/TCP  7m
ps-1         ClusterIP  None           <none>        9000/TCP  7m
tensorboard  ClusterIP  172.19.13.242  106.1.1.1 80/TCP    7m
ps-0         ClusterIP  None           <none>        9000/TCP  7m
worker-1     ClusterIP  None           <none>        8000/TCP  7m

==> v1beta1/Deployment
NAME         DESIRED  CURRENT  UP-TO-DATE  AVAILABLE  AGE
tensorboard  1        1        1           1          7m

==> v1/Job
NAME      DESIRED  SUCCESSFUL  AGE
ps-1      1        0           7m
worker-0  1        0           7m
ps-0      1        0           7m
worker-1  1        0           7m

==> v1/Pod(related)
NAME                          READY  STATUS   RESTARTS  AGE
tensorboard-5c785fbd97-7cwk2  1/1    Running  0         7m
ps-1-lkbtb                    1/1    Running  0         7m
worker-0-2mpmb                1/1    Running  0         7m
ps-0-ncxch                    1/1    Running  0         7m
worker-1-4hngw                1/1    Running  0         7m

这里可以看到Tensorboard的对外IP是106.1.1.1，可以在训练过程中查看cost的收敛程度

4.3 检查任务运行状况， 此时worker都是出于运行中的状态

# kubectl get job
NAME       DESIRED   SUCCESSFUL   AGE
ps-0       1         0            5m
ps-1       1         0            5m
worker-0   1         0            5m
worker-1   1         0            5m

# kubectl get po
NAME                          READY     STATUS    RESTARTS   AGE
ps-0-jndpd                    1/1       Running   0          6m
ps-1-b8zgz                    1/1       Running   0          6m
tensorboard-f78b4d57b-pm2nf   1/1       Running   0          6m
worker-0-rqmvl                1/1       Running   0          6m
worker-1-7pgx6                1/1       Running   0          6m

4.4 检查训练日志

# kubectl logs --tail=10 worker-0-rqmvl
Step: 124607,  Epoch: 24,  Batch: 1600 of 2750,  Cost: 0.8027,  AvgTime: 6.79ms
Step: 124800,  Epoch: 24,  Batch: 1700 of 2750,  Cost: 0.7805,  AvgTime: 6.10ms
Step: 124989,  Epoch: 24,  Batch: 1800 of 2750,  Cost: 1.4159,  AvgTime: 5.98ms
Step: 125184,  Epoch: 24,  Batch: 1900 of 2750,  Cost: 0.6790,  AvgTime: 6.33ms
Step: 125376,  Epoch: 24,  Batch: 2000 of 2750,  Cost: 1.3145,  AvgTime: 6.35ms
Step: 125565,  Epoch: 24,  Batch: 2100 of 2750,  Cost: 0.6310,  AvgTime: 6.13ms
Step: 125759,  Epoch: 24,  Batch: 2200 of 2750,  Cost: 1.1366,  AvgTime: 6.36ms
Step: 125948,  Epoch: 24,  Batch: 2300 of 2750,  Cost: 0.5678,  AvgTime: 6.02ms
Step: 126143,  Epoch: 24,  Batch: 2400 of 2750,  Cost: 0.6040,  AvgTime: 6.84ms
Step: 126310,  Epoch: 24,  Batch: 2500 of 2750,  Cost: 0.7697,  AvgTime: 6.01ms

4.5 可以通过watch job状态，可以监视到job已经完成

# kubectl get job
NAME       DESIRED   SUCCESSFUL   AGE
ps-0       1         0            1h
ps-1       1         0            1h
worker-0   1         1            1h
worker-1   1         1            1h

4.6 此时再查看训练日志，发现训练已经完成

# kubectl logs --tail=10 -f worker-0-rqmvl
Step: 519757,  Epoch: 100,  Batch: 2300 of 2750,  Cost: 0.1770,  AvgTime: 6.45ms
Step: 519950,  Epoch: 100,  Batch: 2400 of 2750,  Cost: 0.2142,  AvgTime: 6.33ms
Step: 520142,  Epoch: 100,  Batch: 2500 of 2750,  Cost: 0.1940,  AvgTime: 6.02ms
Step: 520333,  Epoch: 100,  Batch: 2600 of 2750,  Cost: 0.5144,  AvgTime: 6.21ms
Step: 520521,  Epoch: 100,  Batch: 2700 of 2750,  Cost: 0.5694,  AvgTime: 5.80ms
Step: 520616,  Epoch: 100,  Batch: 2750 of 2750,  Cost: 0.5333,  AvgTime: 2.94ms
Test-Accuracy: 0.89
Total Time: 1664.68s
Final Cost: 0.5333
done

5. 通过Tensorboad查看训练效果，前面已经获得了Tensorboard的外部ip 106.1.1.1， 直接登录链接 http://106.1.1.1/, 就可以观测到训练的效果

总结

TensorFlow和Kubernetes分别作为深度学习和容器编排领域的领航者，二者的强强联合可以真正释放分布式训练的洪荒之力。而阿里云的Helm解决方案降低了部署的难度，降低了这把屠龙刀的使用难度。欢迎大家尝试阿里云Kubernetes容器服务，利用分布式TensorFLow运行自己的模型训练。我们也会持续优化，增加日志和监控，GPU亲和性调度等能力。