Spark 3.0 对 GPU 做了什么支持

2020-08-06 10:12:22 浏览数 (1)

Spark 本身在数据处理流程里占据非常重要的地位,而在人工智能的战场,传统 Spark 能带来什么呢?

在深度学习里,模型训练一般都被 Tensorflow, PyTorch 等深度学习框架占领了,而 Spark 提供的 GraphX 和 MlLib 可以做一些机器学习的东西,但是在深度学习的战场里,明显没有什么优势,最大的问题就在于硬件加速上,3.0 以前的社区版 Spark 是没有任务调度 GPU 的方法的。

那么 Spark 为什么在深度学习靠拢呢,还有他的价值呢?

答案就是训练模型除了本身的大规模的并行密集计算,从数据到模型,必须有数据处理的过程,这个也就是 Spark 的强项,因为你不太可能用 Pandas 简单清洗汇总 ETL 你的训练数据的,所以 Spark 还是有其一席之地的。

目前 Spark 只支持调度 Nvidia 的 GPU,AMD 等其他厂商的 GPU 还有待测试。

看一下 Spark 源码到底对 GPU 做了什么样的支持。

Checkout 到最新的 v3.0.0-preview-rc2 的 tag。

代码语言:javascript复制
git checkout v3.0.0-preview-rc2 -b v3.0.0-preview-rc2

我们找一下关键词 GPU,全局 搜索一下。

因为我只关注 on Kubernetes 的模式下,如何对 GPU 进行支持,所以全局搜索的时候,可以适当把 Yarn 和 Mesos 的部分剔除,因为这里的代码一般是互斥的,也就是 on Kubernetes 的代码,不会用到 Yarn 和 Mesos 的模块。

这个脚本运行的结果如下图。

代码语言:javascript复制
# ADDRS=`nvidia-smi --query-gpu=index --format=csv,noheader | sed -e ':a' -e 'N' -e'$!ba' -e 's/n/","/g'`
# echo {"name": "gpu", "addresses":["$ADDRS"]}
{"name": "gpu", "addresses":["0","1","2","3","4","5","6","7"]}

query-gpu 选项是用于打印 GPU 相关信息的一个选项。

spark.{driver/executor}.resource.gpu.discoveryScript 这是 Spark 新增内置的一个 config,主要是用于传递脚本,来确认 Pod 里的 Driver/Executor 容器是否已经分配 到了 GPU。当然了,我理解的是,如果你百分百确定可以分配到 GPU,其实这个脚本不传也是 没事的。

core 模块增加了一个 resource 的 pacakge。

代码语言:javascript复制
/path/to/spark/core/src/main/scala/org/apache/spark/resource
├── ResourceAllocator.scala
├── ResourceInformation.scala
└── ResourceUtils.scala

ResourceAllocator 是一个 trait,主要是定义了 task 的 acquirerelease 资源的两个动作。

关于如何申请 GPU,目前 Spark 将所有的资源类型抽象成 resource 这个概念,无论还以后是集成 FPGA 还是 GPU,设置是 IB 网卡等等,都可以通过 resource 的 vendor 和 name 来指定需要给 task 申请的资源。

代码语言:javascript复制
spark.driver.resource.{resourceName}.vendor

0 人点赞