C++ / JNI - Spark 调用 JNI on Yarn

一.引言

前面介绍了如何使用 C++ 与 Java  构建自己的 .so 包并实现了本地的测试，下面介绍如何在 Yarn 集群使用 Spark 调用 JNI 函数。

二.JNI TestOnYarn 实现逻辑

介绍之前，首先回顾一下之前本地测试是如何实现的：

// 加载 .so
System.load(filePath)
// 初始化对应类
TestJNIByJVM testJNIByJVM = new TestJNIByJVM();
// JNI 方法调用
int res = testJNIByJVM.testJNI(input1, input2);

切换到 Spark On Yarn 后，所以我们需要注意三个事情：

A.如何传输文件

由于 Spark On Yarn 运行在分布式环境，所以不能像本地测试一样，在本机存储一份文件即可，需要使用 --files 参数将对应动态库 .so 包分布式传到每一台机器。这一步需要在 Spark Submit 脚本处实现：

--files ${libSoPath} \

如果你的任务本身已经 --files 传了文件，则你需要逗号分隔传递多个文件：

--files ${OtherFile},${LibSoPath} \

如果更多的话，依次类推，就像 --jars 一样传输多个 jar。 

B.如何读取文件

不论是本地测试还是集群测试，这里都推荐使用 Load 方法传递绝对路径读取，而 LoadLibrary 方法读取绝对路径虽然代码量会少一点，但是容易出去出现异常，所以建议前者。前面我们写过 Spark Submit --files 文件读取，大家可以参考其中方法：

SparkFiles.get(fileName)

其中 fileName 为 --files 传输的文件名，通过该方法可以在 Driver 端或者 Executor 端获取该文件在分布式执行机器上的对应存储位置，再使用 System.load 读取绝对路径即可，当然也可以直接使用 System 方法直接本机地址：

System.getenv("PWD")

C.调用方法

由于需要后期传入 Path 再读取，所以不能像上面那样初始化 Object 静态类时直接 System.load，我们需要为当前类提供 init 方法，并在 Driver 端或者 Executor 端调用 init 方法初始化：

- 本地

// 加载 .so
System.load(filePath)
// 初始化对应类
TestJNIByJVM testJNIByJVM = new TestJNIByJVM();

- Spark

这里使用变量初始化 JNI 类为 null，并提供 init 方法供 Dirver、Executor 端调用。

object TestJNI {

  // 初始化默认空值
  var testJNIByJVM: TestJNIByJVM = _

  // 根据 path Load
  def initJNI(filePath: String): Unit = {
    if (testJNIByJVM == null) {
      System.load(filePath);
      testJNIByJVM = new TestJNIByJVM()
      println("初始化 TestJNI!")
    } else {
      println("已初始化 TestJNI!")
    }
  }
  
}

D.Dirver、Executor 调用

就是 JNI 调用是在 Driver 端还是 Executor 端，不管哪一端调用都需要对应的初始化方法，而不是 Driver 端初始化一劳永逸。

- Driver

      val libPath = s"${System.getenv("PWD")}/xxx.so"
      TestJNI.initJNI(libPath)

这一步调用方法是介于 SparkContext 与 RDD 执行逻辑之间，因为你要在 Driver 端调用。

- Executor

Executor 端初始化方法相同，唯一区别是需要将上述在 RDD 逻辑内初始化，建议大家使用 foreachPartition 和 mapPartition 减少初始化次数，当然也可以在 Executor 端加锁或者使用双重检测等等，总之不要初始化太多次浪费时间。

    rdd.foreachPartition(partition => {

      val libPath = s"${System.getenv("PWD")}/xxx.so"
      TestJNI.initJNI(libPath) 
      
      while (partition.hasNext) {
        ...
      }
      
    })

加载成功后，调用静态类变量方法即可，Func 即为你 JNI 实现的方法：

TestJNI.TestJNIByJVM.Func

三.总结

相比本地或者单服务器调用，Spark On Yarn 调用需要多一个步骤，但是整体思路不变，这里还是再次推荐一下使用 load 方法读取绝对路径，如果想要调用绝对路径，需要把 .so 包放到 spark 对应的环境目录下，相对麻烦一些。