Author: xidianwangtao@gmail.com
更多关于Kubernetes的深度文章,请到我oschina/WalonWang的博客主页。
2016年4月TensorFlow发布了0.8版本宣布支持分布式计算,这个特性,我们称之为Distributed TensorFlow。
这是非常重要的一个特性,因为在AI的世界里,训练数据的size通常会大到让人瞠目结舌。比如Google Brain实验室今年发表的论文OUTRAGEOUSLY LARGE NEURAL NETWORKS: THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER中提到,下图中MOE Layer Model可以达到680亿个Parameters的规模,如此复杂的模型,如果只能单机训练,那耗时难于接受。通过Distributed TensorFlow,可以利用众多服务器构建TensorFlow Cluster来提高训练效率。
关于Distributed TensorFlow的更多内容,请参考官方内容www.tensorflow.org/deplopy/distributed,这里给出Distributed TensorFlow结构图:
Distributed TensorFlow虽然提供了分布式能力,可以利用服务器集群加快训练,但是还有许多缺点,比如资源无法隔离、PS进程遗留问题等等,而这些正是Kubernetes所擅长的地方。下图是总结的你需要将TensorFlow运行在Kubernetes上的理由:
对于我们来说,前期最大的用户痛点就是算法团队使用的HDFS Read性能不及预期,经过网上查找资料及我们自己简单的对比测试,发现***GlusterFS***可能是最适合我们的分布式存储了。因此在我们的TensorFlow on Kubernetes项目中使用GlusterFS来存放训练数据,worker将从GlusterFS中读取训练数据进行计算。
关于PS进程遗留问题,TensorFlow社区有很多讨论,但至今没有官方的实现方案,在Kubernetes中,这将比较好解决,在后面的Thinking小节中会单独讨论。
说明:
整个系统涉及以下核心Components:
网络方案:contiv netplugin + ovs + vlan. 日志方案:fluentd + Kafka + ES + Kibana. 监控方案:cadvisor + prometheus + Grafana.
CaaS的细节不在这里讨论,其实也是大家非常熟悉的方案了。
大家可以参考Kyle Bai的https://github.com/kairen/workshop413,他这里时候用NFS作为后端存储,需要改成你们自己的存储,大家自己去尝试吧,我这就不一步一步来了,好无聊。
这个Demo,我改成NodePort方式暴露Jupyter Nodebook,登录时输入正确的token即可:
这是一个In-Graph集群,点击master_client.ipynb,可以看到具体的训练算法内容:
点击执行,可以在下面看到输出:
这只是个简单的Demo,实际使用上,自动化生成各个ps, worker, pvc对应的kubernetes yaml,使用域名进行服务发现,不然如果你使用IP的话,可能就需要利用Pod的ProStart Hook来反馈各个Task的IP了,这将比较麻烦。
tf.train.Saver
提供能力在worker上进行checkpoint,大概原理就是逐个从PS task中get Parameters,并进行save持久化。
{%- set name = "imagenet" -%}
{%- set worker_replicas = 3 -%}
{%- set ps_replicas = 2 -%}
{%- set script = "http://xxx.xx.xx.xxx:80/imagenet/imagenet.py" -%}
{%- set image = "tensorflow/tensorflow:1.3.0" -%}
{%- set data_dir = "/data" -%}
{%- set log_dir = "/log" -%}
{%- set port = 2222 -%}
{%- set replicas = {"worker": worker_replicas, "ps": ps_replicas} -%}
{%- macro worker_hosts() -%}
{%- for i in range(worker_replicas) -%}
{%- if not loop.first -%},{%- endif -%}
{{ name }}-worker-{{ i }}:{{ port }}
{%- endfor -%}
{%- endmacro -%}
{%- macro ps_hosts() -%}
{%- for i in range(ps_replicas) -%}
{%- if not loop.first -%},{%- endif -%}
{{ name }}-ps-{{ i }}:{{ port }}
{%- endfor -%}
{%- endmacro -%}
{%- for job in ["worker", "ps"] -%}
{%- for i in range(replicas[job]) -%}
kind: Service
apiVersion: v1
metadata:
name: {{ name }}-{{ job }}-{{ i }}
spec:
selector:
name: {{ name }}
job: {{ job }}
task: "{{ i }}"
ports:
- port: {{ port }}
targetPort: 2222
{% if job == "worker" %}
---
kind: Job
apiVersion: batch/v1
metadata:
name: {{ name }}-{{ job }}-{{ i }}
spec:
replicas: 1
template:
metadata:
labels:
name: {{ name }}
job: {{ job }}
task: "{{ i }}"
spec:
containers:
- name: {{ name }}-{{ job }}-{{ i }}
image: {{ image }}
ports:
- containerPort: 2222
command: ["/bin/sh", "-c"]
args:["
curl {{ script }} -o /opt/{{ name }}.py;
python /opt/{{ name }}.py \
--ps_hosts={{ ps_hosts() }} \
--worker_hosts={{ worker_hosts() }} \
--job_name={{ job }} \
--task_index={{ i }} \
--log_path={{ log_dir }} \
--data_dir={{ data_dir }} ;"]
volumeMounts:
- name: data
mountPath: {{ data_dir }}
- name: log
mountPath: {{ log_dir }}
restartPolicy: Never
volumes:
- name: data
persistentVolumeClaim:
claimName: {{ name }}-data-pvc
- name: log
persistentVolumeClaim:
claimName: {{ name }}-log-pvc
{% endif %}
{% if job == "ps" %}
---
kind: Deployment
apiVersion: extensions/v1beta1
metadata:
name: {{ name }}-{{ job }}-{{ i }}
spec:
replicas: 1
template:
metadata:
labels:
name: {{ name }}
job: {{ job }}
task: "{{ i }}"
spec:
containers:
- name: {{ name }}-{{ job }}-{{ i }}
image: {{ image }}
ports:
- containerPort: 2222
command: ["/bin/sh", "-c"]
args:["
curl {{ script }} -o /opt/{{ name }}.py;
python /opt/{{ name }}.py \
--ps_hosts={{ ps_hosts() }} \
--worker_hosts={{ worker_hosts() }} \
--job_name={{ job }} \
--task_index={{ i }} \
--log_path={{ log_dir }} ;"]
volumeMounts:
- name: log
mountPath: {{ log_dir }}
restartPolicy: Never
volumes:
- name: log
persistentVolumeClaim:
claimName: {{ name }}-log-pvc
{% endif %}
---
{% endfor %}
{%- endfor -%}
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
name: {{ name }}-log-pvc
annotations:
volume.beta.kubernetes.io/storage-class: glusterfs
spec:
accessModes:
- ReadWriteMany
resources:
requests:
storage: 10Gi
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
name: {{ name }}-data-pvc
annotations:
volume.beta.kubernetes.io/storage-class: glusterfs
spec:
accessModes:
- ReadWriteMany
resources:
requests:
storage: 10Gi
---
TensorFlow和Kubernetes分别作为深度学习和容器编排领域的王者,两者的合理整合可以将释放Distributed TensorFlow的能力,本文只是我对TensorFlow on Kubernetes的浅尝,未来还有很多工作需要做,比如给某些算法定制特殊的调度策略、网络IO性能调优、在DevOps上开发TaaS,提升易用性、TensorFlow Serving的快速部署等等,欢迎对这方面有浓厚兴趣的同学加我微信xidianwangtao交流。