【Bug已解决】[Training] OrtTrainingSession.evalStep() leaks outputHandlesArr — missing ReleaseLongArrayElements 解决方案一、现象长什么样在 JavaONNX Runtime 的 Java API里用OrtTrainingSession.evalStep()做训练评估跑很多 step 之后发现进程常驻内存持续增长用 MAT / jstat 看是long[]数组不断堆积、GC 回收不掉。进一步用 native 内存工具看是 JNI 层在把long[]返回给 Java 时没有正确释放。现象# 现象 A长时间训练后 Java 堆外/堆内 long[] 泄漏OOM # jstat -gc: FGCT 持续上升最终 java.lang.OutOfMemoryError: Java heap space # 现象 BNative 侧报 JNI 警告 # JNI WARNING: received null array from GetLongArrayElements without release # 或 native 内存随 step 数线性增长 # 现象 C只在 evalStep返回多个输出 handle时泄漏 # trainStep返回单个不泄漏 —— 定位到 outputHandlesArr 这个数组最坑的是现象 A不是立刻崩而是跑几千 step 后慢慢滲满内存训练任务中途被杀极难第一时间联想到 JNI 释放遗漏。二、背景ONNX Runtime 的 Java API 通过 JNI 调用 native C。当evalStep()需要把“多个输出 tensor 的 handle 数组long[] outputHandlesArr”从 native 返回给 Java 时典型的 JNI 写法是jlongArray arr env-NewLongArray(n); jlong* elems env-GetLongArrayElements(arr, nullptr); // 填充 elems[i] handle env-ReleaseLongArrayElements(arr, elems, 0); // ← 必须释放GetLongArrayElements可能返回一个指向 Java 数组内部数据的指针pinned或一份拷贝。无论哪种都必须配对的ReleaseLongArrayElements否则如果是拷贝模式拷贝在 native 侧泄漏如果是 pinned 模式数组被一直 pin 住无法被 GC 移动/回收。ORT 的evalStepJNI 包装在重构时漏写了ReleaseLongArrayElements于是每次调用都泄漏一个long[]的 native 副本。这是 JNI 代码审查里最经典的坑GetXxxArrayElements与ReleaseXxxArrayElements不配对。三、根因GetLongArrayElements没有配对ReleaseLongArrayElementsevalStep填充完outputHandlesArr后直接return忘了 release每次调用泄漏一份long[]native 拷贝。只在多输出路径泄漏evalStep返回多个输出 handle需要数组trainStep返回单个不需要数组所以只有evalStep触发 → 现象 C 的定位线索。缺少 JNI 调用的内存对拍/泄漏检测CI 没跑长时间 evalStep 循环 内存断言泄漏长期存在。本质是JNI 层GetLongArrayElements/ReleaseLongArrayElements不配对且只在多输出的 evalStep 路径暴露缺少泄漏检测。四、最小可运行复现下面用 Python 模拟“JNI 数组元素获取后不释放导致计数泄漏”用引用计数类比 native 拷贝class FakeJniArray: _live_copies 0 def __init__(self, n): self.n n def get_elements(self): FakeJniArray._live_copies 1 # 获取 产生一份 native 拷贝 return [0] * self.n def release_elements(self, elems): FakeJniArray._live_copies - 1 # 释放 拷贝回收 def eval_step_buggy(n_outputs): buggy: 获取后忘了 release。 arr FakeJniArray(n_outputs) elems arr.get_elements() # 填充 elems ... 然后直接返回漏了 arr.release_elements(elems) return elems def eval_step_fixed(n_outputs): arr FakeJniArray(n_outputs) elems arr.get_elements() # ... 填充 arr.release_elements(elems) # 正确释放 return elems FakeJniArray._live_copies 0 for _ in range(1000): eval_step_buggy(4) print(buggy leaked copies after 1000 steps:, FakeJniArray._live_copies) # 1000 FakeJniArray._live_copies 0 for _ in range(1000): eval_step_fixed(4) print(fixed leaked copies after 1000 steps:, FakeJniArray._live_copies) # 0buggy跑 1000 步泄漏 1000 份fixed为 0。五、解决方案第一层最小直接修复最小修复在evalStep的 JNI 包装里填充完outputHandlesArr后必须配对调用ReleaseLongArrayElements// 修正后的 evalStep JNI 片段 jlongArray outputHandlesArr env-NewLongArray(output_count); if (outputHandlesArr nullptr) { /* 异常处理 */ } jlong* handles env-GetLongArrayElements(outputHandlesArr, nullptr); for (int i 0; i output_count; i) { handles[i] static_castjlong(output_tensors[i]-GetHandle()); } // 关键填充完成后立即释放0 表示把改动写回 Java 数组并释放 native 拷贝 env-ReleaseLongArrayElements(outputHandlesArr, handles, 0); return outputHandlesArr;这一层改动最小补一行ReleaseLongArrayElements泄漏消失。但依赖“每个 Get 都记得 Release”下看第二层。六、解决方案第二层结构性改进把“JNI 数组元素的获取/释放必须配对”固化成单一事实来源。下面这个 dataclass 用 RAII 风格的 Python 模型描述一个作用域守卫获取即登记离开作用域必须释放并把这套契约做成可校验的规范。from dataclasses import dataclass, field from typing import Dict, List import contextlib dataclass class OrtTrainingJniLeakPolicy: 单一事实来源JNI 数组元素获取/释放配对的契约。 _pending: Dict[int, str] field(default_factorydict) _counter: int 0 contextlib.contextmanager def array_elements(self, tag: str): 守卫进入即 Get离开必 Release。漏写 Release 会在 __exit__ 报错。 self._counter 1 cid self._counter self._pending[cid] tag try: yield fnative_copy_of_{tag} # 模拟 GetLongArrayElements 返回值 finally: # 对应 ReleaseLongArrayElements self._pending.pop(cid, None) def assert_no_leak(self) - None: if self._pending: raise AssertionError(funreleased JNI arrays: {self._pending}) # 用法正确 policy OrtTrainingJniLeakPolicy() with policy.array_elements(outputHandlesArr): pass # 填充 handles policy.assert_no_leak() # 通过这一层的关键收益RAII 守卫用with确保获取/释放配对漏写 Release 会在退出时报“unreleased”泄漏断言assert_no_leak在每次 evalStep 后检查残留即报错单一事实来源所有 JNI 数组配对约定收口在OrtTrainingJniLeakPolicy。七、解决方案第三层断言 / CI 守护把第二层钉成 pytest挂进 CI确保 JNI 数组不泄漏import contextlib import pytest from your_package.ort_jni_leak import OrtTrainingJniLeakPolicy def test_no_leak_when_released(): # 断言 1正确配对释放后无泄漏 p OrtTrainingJniLeakPolicy() with p.array_elements(outputHandlesArr): pass p.assert_no_leak() def test_leak_detected_without_release(): # 断言 2漏写 Release 必须被守卫抓出 p OrtTrainingJniLeakPolicy() try: with p.array_elements(outputHandlesArr): raise RuntimeError(simulate early return / forgot release path) except RuntimeError: pass # 即使异常退出__exit__ 也应释放若实现正确则无泄漏 p.assert_no_leak() def test_evalstep_loop_no_accumulate(): # 断言 3循环 1000 次 evalStep 后泄漏计数为 0 p OrtTrainingJniLeakPolicy() for _ in range(1000): with p.array_elements(outputHandlesArr): pass p.assert_no_leak()三条断言从“配对无泄漏”“漏释放被抓”“循环不累积”三面把 JNI 泄漏钉死在 CI。八、排查清单Java/ONNX Runtime 训练出现内存随 step 增长时用 jstat/jmap 看long[]是否随 evalStep 调用数线性增长是就怀疑 JNI 数组泄漏。查evalStep的 JNI 包装每个GetLongArrayElements是否都有配对ReleaseLongArrayElementsoutputHandlesArr是重灾区。只有多输出路径泄漏、单输出不泄漏定位到需要数组返回的函数现象 C 线索。用第二层OrtTrainingJniLeakPolicyRAII 守卫 assert_no_leak校验。加第三层 pytest断言“配对无泄漏、漏释放被抓、循环不累积”。JNI 代码必须“获取即释放”理想用 C 的 RAII wrapper 包裹 Get/Release。九、小结OrtTrainingSession.evalStep()的泄漏 bug 本质是JNI 层GetLongArrayElements获取outputHandlesArr后漏写配对的ReleaseLongArrayElements导致每次 evalStep 都泄漏一份long[]native 拷贝训练几千 step 后 OOM且只在多输出的 evalStep 路径暴露单输出 trainStep 正常。修复分三层——第一层补上ReleaseLongArrayElements第二层用OrtTrainingJniLeakPolicy这个 dataclass 把“获取/释放配对”收口成 RAII 守卫并加泄漏断言第三层用三条 pytest 把“配对无泄漏、漏释放被抓、循环不累积”钉死在 CI。核心心法JNI 里GetXxxArrayElements必须与ReleaseXxxArrayElements严格配对多输出返回的数组是最易漏的点必须用守卫或 RAII 兜底。