在阿里 PPU 环境中使用 onnxruntime-gpu 的 Java 库

核心思路就是利用环境变量让 Java 应用加载 pixi 中安装的 onnxruntime 的 so。

前置操作请参阅 使用 pixi 配置阿里 PPU 环境

我这里项目使用的是 Gradle，在项目根目录的 build.gradle.kts 中添加如下内容：

subprojects {
    val pixiEnvs = File(rootDir, "torrenkt-tslm-model/.pixi/envs")
    if (!pixiEnvs.exists()) {
        logger.error("please run 'pixi install' in ./torrenkt-tslm-model")
        return@subprojects
    }
    val env = when (providers.gradleProperty("torch.gpu.env").getOrElse("cpu")) {
        "ppu" -> "ppu"
        "cuda" -> "default"
        else -> "cpu"
    }
    val path = File(pixiEnvs, "$env/bin").canonicalPath
    val libs = File(pixiEnvs, "$env/lib").canonicalPath
    val onnxruntime = File(libs, "./python3.12/site-packages/onnxruntime/capi").canonicalPath // 注意 python 版本换成你实际使用的版本
    tasks.withType<javaexec>().configureEach {
        environment("PATH", "$path:${environment["PATH"]}")
        environment("LD_LIBRARY_PATH", "$libs:$onnxruntime:${environment["LD_LIBRARY_PATH"]}")
    }
}

然后在 ~/.gradle/gradle.properties 中添加内容：

torch.gpu.env=ppu

这样不仅能让 ppu 环境使用 python 包提供的 onnxruntime.so，其他环境也能加载 pixi 中的 cuda-toolkit。

打包后的成品运行直接添加环境变量再运行即可，这里不再赘述。