如何让 Apache Spark 在 GPU 上运行？

Question

我想将 apache spark 与 GPU 集成，但 spark 可以在 java 上运行，而 gpu 使用 CUDA/OpenCL，那么我们如何合并它们。

Answer 1

这取决于你想做什么。如果您想使用 spark 使用 GPU 分配计算，则不必使用 java。您可以将 python (pyspark) 与具有 cuda 模块的 numba 一起使用。

例如，如果您希望工作节点在 RDD 的每个块上计算操作（此处为 gpu_function），则可以应用此代码。

rdd = rdd.mapPartition(gpu_function)

和 ：

def gpu_function(x):
    ...
    input = f(x)
    output = ...
    gpu_cuda[grid_size,block_size](input,output)
    return output

和 ：

from numba import cuda
@cuda.jit("(float32[:],float32[:])")
def gpu_cuda(input,output)
    output = g(input)

我建议您查看 幻灯片共享网址：https ://fr.slideshare.net/continuumio/gpu-computing-with-apache-spark-and-python ，特别是第 34 张幻灯片。

您只需要在每个工作节点上安装 numba 和 cuda 驱动程序。

Answer 2

有一些库可以帮助解决这个困境。

Databricks 正在开发一个适用于 Spark 和 TensorFlow 的解决方案，该解决方案将允许您使用集群或机器的 GPU。

如果您想了解更多相关信息，请参阅Spark Summit Europe 2016的演示文稿。该演示文稿将稍微介绍 TensorFrames 的工作原理。

其他这是 DataBricks 博客中关于TensoFrames的帖子。

有关更多代码信息，请参阅Tensorframes 的 Git。