Spark 3.0 对 GPU 做了什么支持

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。

git checkout v3.0.0-preview-rc2 -b v3.0.0-preview-rc2

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

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

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

# 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。

/path/to/spark/core/src/main/scala/org/apache/spark/resource
├── ResourceAllocator.scala
├── ResourceInformation.scala
└── ResourceUtils.scala

ResourceAllocator 是一个 trait，主要是定义了 task 的 acquire 和 release 资源的两个动作。

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

spark.driver.resource.{resourceName}.vendor