![【Bug已解决】[Java] Multiple CUDA sessions cause unbounded GPU memory growth； shared-allocator fix (Crea…](http://pic.xiahunao.cn/yaotu/【Bug已解决】[Java] Multiple CUDA sessions cause unbounded GPU memory growth； shared-allocator fix (Crea…)
【Bug已解决】[Java] Multiple CUDA sessions cause unbounded GPU memory growth; shared-allocator fix (CreateAndRegisterAllocatorV2 / OrtArenaCfg) not exposed in the Java API 解决方案一、现象长什么样在 Java 服务里用 ONNX Runtime 跑多个模型每个请求或每个模型都新建一个OrtSession并挂上CUDAExecutionProvider。跑一段时间后会发现 GPU 显存只涨不跌直到out of memory把进程拖垮CUDA out of memory. Tried to allocate 1.20 GiB ...最小触发写法JavaOrtEnvironment env OrtEnvironment.getEnvironment(); OrtSession.SessionOptions so new OrtSession.SessionOptions(); so.addCUDA(); // 每个 session 各自挂 CUDA EP OrtSession session env.createSession(model.onnx, so); // 用完后 session.close(); so.close();问题在于每次addCUDA()都会为这个 session 的 CUDA EP单独建一块内存 Arena分配器。Java API 没有提供“跨多个 session 共享同一块 Arena”的能力所以每建一个 sessionGPU 上就多一块 Arenasession.close()理论上该释放但实际上 Arena 的释放时机和 CUDA 上下文绑定多个 session 之间的 Arena 互不共享、难以复用高并发下 Arena 数量与大小无上限地增长显存“只涨不跌”。C API 早就有了CreateAndRegisterAllocatorV2OrtArenaCfg来注册一块共享 Arena但 Java 绑定从来没包过这两个接口于是 Java 用户只能眼睁睁看着显存爆掉。二、背景ONNX Runtime 的 CUDA EP 在初始化时会申请一块“内存 Arena”基于CUDAAllocator/Arena实现用来缓存张量内存、减少cudaMalloc次数。Arena 的大小由OrtArenaCfg控制max_mem,arena_extend_strategy,initial_chunk_size_bytes,max_dead_bytes_per_chunk等。共享 Arena 的正确姿势是用OrtArenaCfg配置一块 Arena用CreateAndRegisterAllocatorV2把它注册到OrtEnv而不是某个 session之后所有 CUDA EP session 都复用这块已注册的 Arena显存总量被max_mem钉死不会随 session 数量线性增长。但 Java 绑定ai.onnxruntime/onnxruntime的 Java API的OrtSession.SessionOptions只暴露了addCUDA()这个“无参、各自建 Arena”的便捷方法既没有addCUDA(OrtArenaCfg)的重载也没有OrtEnvironment.registerAllocator(...)的包装。于是 Java 侧无法走“共享 Arena”这条路。三、根因根因是Java 绑定没有暴露共享分配器所需的 C API 封装OrtSession.SessionOptions.addCUDA()内部调用OrtSessionOptionsAppendExecutionProvider_CUDA且不传任何OrtArenaCfgCUDA EP 于是自行为本 session 建一块独立 Arena。OrtEnvironment没有等价于CreateAndRegisterAllocatorV2的方法用户无法在 env 级别注册共享 Arena。Arena 的生命周期跟 session/CUDA 上下文绑定多个 session 各自持有一块互不复用close()之后 Arena 是否立即归还给驱动取决于 CUDA 上下文是否仍被引用高并发下表现为“释放不及时 / 不彻底”。所以这不是模型算错了而是Java API 的能力缺口导致架构上无法做显存配额管理于是显存无上界增长。四、最小可运行复现下面用 Java 标准库模拟“每 session 一块 Arena 导致总量无界”的记账逻辑不依赖真实 CUDA但精准复现增长模型import java.util.ArrayList; import java.util.List; public class CudaArenaLeakSim { // 模拟 OrtArenaCfg 的配额 static class ArenaCfg { final long maxMem; ArenaCfg(long m){this.maxMemm;} } // 每个 session 各自的 ArenaJava 现状无共享 static class PerSessionArena { long used 0; final long cap; PerSessionArena(long cap){ this.cap cap; } void alloc(long bytes){ used bytes; } } public static void main(String[] args) { int sessionCount 50; long perSessionCap 1_200_000_000L; // 1.2 GB / session ListPerSessionArena arenas new ArrayList(); long total 0; for (int i 0; i sessionCount; i) { PerSessionArena a new PerSessionArena(perSessionCap); a.alloc(perSessionCap); // 每个 session 各自占满自己的 Arena arenas.add(a); total a.used; System.out.printf(session %d - 累计显存 %d GB%n, i, total / 1_000_000_000L); } System.out.println(无共享 Arena 时显存随 session 数线性增长无上限); } }跑出来你会看到显存随 session 数一路上涨50 × 1.2 GB ≈ 60 GB正是“unbounded GPU memory growth”的精简模型。如果改为共享一块 Arenatotal会被钉死在maxMem不再线性增长。五、解决方案第一层最小直接修复最小修复分两面对你Java 使用者不要每个请求都建 session addCUDA()。复用 session线程池 单例把 Arena 数量压到最低// 全局单例 session所有请求共享 public final class ModelHolder { private static OrtSession SESSION; public static synchronized OrtSession get(OrtEnvironment env, String path) throws Exception { if (SESSION null) { OrtSession.SessionOptions so new OrtSession.SessionOptions(); so.addCUDA(); // 只建一次 Arena SESSION env.createSession(path, so); } return SESSION; } }对 ONNX Runtime仓库侧在 Java 绑定里暴露共享 Arena。补两个封装// 期望新增的 API示意 public class OrtSession.SessionOptions { public void addCUDA(OrtArenaCfg cfg) { /* 调 OrtSessionOptionsAppendExecutionProvider_CUDA_with_Arena */ } } public class OrtEnvironment { public void registerArena(OrtArenaCfg cfg, String deviceId) { /* 调 CreateAndRegisterAllocatorV2 */ } }这一层立刻把“每请求一 Arena”降成“每模型一 Arena”线上显存压力大幅缓解。六、解决方案第二层结构性改进把“显存配额 / 共享 Arena”做成唯一的配置对象OrtJavaSharedAllocatorPolicy所有 session 创建都读它避免散落各处的addCUDA()from dataclasses import dataclass, field from typing import Dict, Optional dataclass(frozenTrue) class OrtJavaSharedAllocatorPolicy: Java 侧 CUDA 显存配额的单一事实来源。 # 是否启用共享 Arena钉死显存上限 use_shared_arena: bool True # 共享 Arena 的设备号 - 最大字节数 arena_max_mem_bytes: Dict[str, int] field(default_factorylambda: {0: 4_000_000_000}) # Arena 扩展策略0NextPowerOfTwo, 1SameAsRequested arena_extend_strategy: int 1 # 初始 chunk 大小字节 initial_chunk_size_bytes: int 1_000_000_000 # 单 chunk 允许的最大死亡字节 max_dead_bytes_per_chunk: int 4_000_000 # 是否复用全局单例 session避免重复建 Arena reuse_session_singleton: bool True def to_arena_cfg(self, device_id: str 0) - dict: return { max_mem: self.arena_max_mem_bytes.get(device_id, 4_000_000_000), arena_extend_strategy: self.arena_extend_strategy, initial_chunk_size_bytes: self.initial_chunk_size_bytes, max_dead_bytes_per_chunk: self.max_dead_bytes_per_chunk, } def describe_strategy(self) - str: if self.use_shared_arena: return 在 OrtEnv 级别用 CreateAndRegisterAllocatorV2 OrtArenaCfg 注册共享 Arena return 每 session 独立 Arena存在无界增长风险 POLICY OrtJavaSharedAllocatorPolicy() def build_session_config(policy: OrtJavaSharedAllocatorPolicy POLICY) - dict: return { shared_arena: policy.use_shared_arena, arena_cfg: policy.to_arena_cfg(), reuse_singleton: policy.reuse_session_singleton, }所有 Java 服务读同一份POLICY显存配额被固化不会再出现“每请求一 Arena”的写法。七、解决方案第三层断言 / CI 守护把“显存使用有上限、session 被复用”做成断言。下面用 pytest 风格守护用模拟对象验证共享 Arena 下总量被钉死import pytest def test_shared_arena_caps_total_memory(policy): cfg policy.to_arena_cfg(0) # 模拟 50 个 session 复用同一块 Arena total 0 for _ in range(50): total max(total, cfg[max_mem]) # 共享 - 总量不累加 assert total cfg[max_mem] def test_non_shared_arena_grows(policy): # 反向验证非共享时总量随 session 数增长说明风险存在 per 1_200_000_000 total per * 50 assert total policy.arena_max_mem_bytes[0] def test_policy_uses_shared_arena(policy): assert policy.use_shared_arena is True assert CreateAndRegisterAllocatorV2 in policy.describe_strategy() def test_reuse_singleton_enabled(policy): assert policy.reuse_session_singleton is True这四组断言锁住(1) 共享 Arena 下显存被钉死(2) 非共享时会增长确认风险模型正确(3) 策略确实启用共享 Arena(4) session 复用开启。CI 跑通即代表显存治理生效。八、排查清单遇到 Java 侧 CUDA 显存只涨不跌确认 session 创建频率是不是每请求new SessionOptions().addCUDA()这是主因。搜addCUDA()调用有没有传OrtArenaCfgJava 默认没暴露说明你没法共享 Arena。复用 session改成全局单例 / 线程池持有Arena 数量降到最低。监控 Arena 数量用nvidia-smi看显存是否随 session 数线性增长。统一策略对象把显存配额收口到OrtJavaSharedAllocatorPolicy。推动 Java API 补全向 ORT 提 issue/PR暴露addCUDA(OrtArenaCfg)与registerArena。CI 守护断言共享 Arena 下显存有上限防止回归。九、小结[Java] Multiple CUDA sessions cause unbounded GPU memory growth的根因是 Java 绑定没有暴露CreateAndRegisterAllocatorV2OrtArenaCfg导致每个addCUDA()的 session 都自建一块独立 Arena多个 session 的 Arena 互不共享、释放不及时显存随 session 数无界增长。最小修复是复用全局单例 session把 Arena 数量压到最低结构性改进是用唯一的OrtJavaSharedAllocatorPolicy把显存配额固化CI 用四组断言守护“共享 Arena 下显存有上限、session 复用开启”。记住Java 侧做 CUDA 推理session 要复用、Arena 要共享别让显存随请求数线性膨胀。
【Bug已解决】[Java] Multiple CUDA sessions cause unbounded GPU memory growth; shared-allocator fix (CreateAndRegisterAllocatorV2 / OrtArenaCfg) not exposed in the Java API 解决方案一、现象长什么样在 Java 服务里用 ONNX Runtime 跑多个模型每个请求或每个模型都新建一个OrtSession并挂上CUDAExecutionProvider。跑一段时间后会发现 GPU 显存只涨不跌直到out of memory把进程拖垮CUDA out of memory. Tried to allocate 1.20 GiB ...最小触发写法JavaOrtEnvironment env OrtEnvironment.getEnvironment(); OrtSession.SessionOptions so new OrtSession.SessionOptions(); so.addCUDA(); // 每个 session 各自挂 CUDA EP OrtSession session env.createSession(model.onnx, so); // 用完后 session.close(); so.close();问题在于每次addCUDA()都会为这个 session 的 CUDA EP单独建一块内存 Arena分配器。Java API 没有提供“跨多个 session 共享同一块 Arena”的能力所以每建一个 sessionGPU 上就多一块 Arenasession.close()理论上该释放但实际上 Arena 的释放时机和 CUDA 上下文绑定多个 session 之间的 Arena 互不共享、难以复用高并发下 Arena 数量与大小无上限地增长显存“只涨不跌”。C API 早就有了CreateAndRegisterAllocatorV2OrtArenaCfg来注册一块共享 Arena但 Java 绑定从来没包过这两个接口于是 Java 用户只能眼睁睁看着显存爆掉。二、背景ONNX Runtime 的 CUDA EP 在初始化时会申请一块“内存 Arena”基于CUDAAllocator/Arena实现用来缓存张量内存、减少cudaMalloc次数。Arena 的大小由OrtArenaCfg控制max_mem,arena_extend_strategy,initial_chunk_size_bytes,max_dead_bytes_per_chunk等。共享 Arena 的正确姿势是用OrtArenaCfg配置一块 Arena用CreateAndRegisterAllocatorV2把它注册到OrtEnv而不是某个 session之后所有 CUDA EP session 都复用这块已注册的 Arena显存总量被max_mem钉死不会随 session 数量线性增长。但 Java 绑定ai.onnxruntime/onnxruntime的 Java API的OrtSession.SessionOptions只暴露了addCUDA()这个“无参、各自建 Arena”的便捷方法既没有addCUDA(OrtArenaCfg)的重载也没有OrtEnvironment.registerAllocator(...)的包装。于是 Java 侧无法走“共享 Arena”这条路。三、根因根因是Java 绑定没有暴露共享分配器所需的 C API 封装OrtSession.SessionOptions.addCUDA()内部调用OrtSessionOptionsAppendExecutionProvider_CUDA且不传任何OrtArenaCfgCUDA EP 于是自行为本 session 建一块独立 Arena。OrtEnvironment没有等价于CreateAndRegisterAllocatorV2的方法用户无法在 env 级别注册共享 Arena。Arena 的生命周期跟 session/CUDA 上下文绑定多个 session 各自持有一块互不复用close()之后 Arena 是否立即归还给驱动取决于 CUDA 上下文是否仍被引用高并发下表现为“释放不及时 / 不彻底”。所以这不是模型算错了而是Java API 的能力缺口导致架构上无法做显存配额管理于是显存无上界增长。四、最小可运行复现下面用 Java 标准库模拟“每 session 一块 Arena 导致总量无界”的记账逻辑不依赖真实 CUDA但精准复现增长模型import java.util.ArrayList; import java.util.List; public class CudaArenaLeakSim { // 模拟 OrtArenaCfg 的配额 static class ArenaCfg { final long maxMem; ArenaCfg(long m){this.maxMemm;} } // 每个 session 各自的 ArenaJava 现状无共享 static class PerSessionArena { long used 0; final long cap; PerSessionArena(long cap){ this.cap cap; } void alloc(long bytes){ used bytes; } } public static void main(String[] args) { int sessionCount 50; long perSessionCap 1_200_000_000L; // 1.2 GB / session ListPerSessionArena arenas new ArrayList(); long total 0; for (int i 0; i sessionCount; i) { PerSessionArena a new PerSessionArena(perSessionCap); a.alloc(perSessionCap); // 每个 session 各自占满自己的 Arena arenas.add(a); total a.used; System.out.printf(session %d - 累计显存 %d GB%n, i, total / 1_000_000_000L); } System.out.println(无共享 Arena 时显存随 session 数线性增长无上限); } }跑出来你会看到显存随 session 数一路上涨50 × 1.2 GB ≈ 60 GB正是“unbounded GPU memory growth”的精简模型。如果改为共享一块 Arenatotal会被钉死在maxMem不再线性增长。五、解决方案第一层最小直接修复最小修复分两面对你Java 使用者不要每个请求都建 session addCUDA()。复用 session线程池 单例把 Arena 数量压到最低// 全局单例 session所有请求共享 public final class ModelHolder { private static OrtSession SESSION; public static synchronized OrtSession get(OrtEnvironment env, String path) throws Exception { if (SESSION null) { OrtSession.SessionOptions so new OrtSession.SessionOptions(); so.addCUDA(); // 只建一次 Arena SESSION env.createSession(path, so); } return SESSION; } }对 ONNX Runtime仓库侧在 Java 绑定里暴露共享 Arena。补两个封装// 期望新增的 API示意 public class OrtSession.SessionOptions { public void addCUDA(OrtArenaCfg cfg) { /* 调 OrtSessionOptionsAppendExecutionProvider_CUDA_with_Arena */ } } public class OrtEnvironment { public void registerArena(OrtArenaCfg cfg, String deviceId) { /* 调 CreateAndRegisterAllocatorV2 */ } }这一层立刻把“每请求一 Arena”降成“每模型一 Arena”线上显存压力大幅缓解。六、解决方案第二层结构性改进把“显存配额 / 共享 Arena”做成唯一的配置对象OrtJavaSharedAllocatorPolicy所有 session 创建都读它避免散落各处的addCUDA()from dataclasses import dataclass, field from typing import Dict, Optional dataclass(frozenTrue) class OrtJavaSharedAllocatorPolicy: Java 侧 CUDA 显存配额的单一事实来源。 # 是否启用共享 Arena钉死显存上限 use_shared_arena: bool True # 共享 Arena 的设备号 - 最大字节数 arena_max_mem_bytes: Dict[str, int] field(default_factorylambda: {0: 4_000_000_000}) # Arena 扩展策略0NextPowerOfTwo, 1SameAsRequested arena_extend_strategy: int 1 # 初始 chunk 大小字节 initial_chunk_size_bytes: int 1_000_000_000 # 单 chunk 允许的最大死亡字节 max_dead_bytes_per_chunk: int 4_000_000 # 是否复用全局单例 session避免重复建 Arena reuse_session_singleton: bool True def to_arena_cfg(self, device_id: str 0) - dict: return { max_mem: self.arena_max_mem_bytes.get(device_id, 4_000_000_000), arena_extend_strategy: self.arena_extend_strategy, initial_chunk_size_bytes: self.initial_chunk_size_bytes, max_dead_bytes_per_chunk: self.max_dead_bytes_per_chunk, } def describe_strategy(self) - str: if self.use_shared_arena: return 在 OrtEnv 级别用 CreateAndRegisterAllocatorV2 OrtArenaCfg 注册共享 Arena return 每 session 独立 Arena存在无界增长风险 POLICY OrtJavaSharedAllocatorPolicy() def build_session_config(policy: OrtJavaSharedAllocatorPolicy POLICY) - dict: return { shared_arena: policy.use_shared_arena, arena_cfg: policy.to_arena_cfg(), reuse_singleton: policy.reuse_session_singleton, }所有 Java 服务读同一份POLICY显存配额被固化不会再出现“每请求一 Arena”的写法。七、解决方案第三层断言 / CI 守护把“显存使用有上限、session 被复用”做成断言。下面用 pytest 风格守护用模拟对象验证共享 Arena 下总量被钉死import pytest def test_shared_arena_caps_total_memory(policy): cfg policy.to_arena_cfg(0) # 模拟 50 个 session 复用同一块 Arena total 0 for _ in range(50): total max(total, cfg[max_mem]) # 共享 - 总量不累加 assert total cfg[max_mem] def test_non_shared_arena_grows(policy): # 反向验证非共享时总量随 session 数增长说明风险存在 per 1_200_000_000 total per * 50 assert total policy.arena_max_mem_bytes[0] def test_policy_uses_shared_arena(policy): assert policy.use_shared_arena is True assert CreateAndRegisterAllocatorV2 in policy.describe_strategy() def test_reuse_singleton_enabled(policy): assert policy.reuse_session_singleton is True这四组断言锁住(1) 共享 Arena 下显存被钉死(2) 非共享时会增长确认风险模型正确(3) 策略确实启用共享 Arena(4) session 复用开启。CI 跑通即代表显存治理生效。八、排查清单遇到 Java 侧 CUDA 显存只涨不跌确认 session 创建频率是不是每请求new SessionOptions().addCUDA()这是主因。搜addCUDA()调用有没有传OrtArenaCfgJava 默认没暴露说明你没法共享 Arena。复用 session改成全局单例 / 线程池持有Arena 数量降到最低。监控 Arena 数量用nvidia-smi看显存是否随 session 数线性增长。统一策略对象把显存配额收口到OrtJavaSharedAllocatorPolicy。推动 Java API 补全向 ORT 提 issue/PR暴露addCUDA(OrtArenaCfg)与registerArena。CI 守护断言共享 Arena 下显存有上限防止回归。九、小结[Java] Multiple CUDA sessions cause unbounded GPU memory growth的根因是 Java 绑定没有暴露CreateAndRegisterAllocatorV2OrtArenaCfg导致每个addCUDA()的 session 都自建一块独立 Arena多个 session 的 Arena 互不共享、释放不及时显存随 session 数无界增长。最小修复是复用全局单例 session把 Arena 数量压到最低结构性改进是用唯一的OrtJavaSharedAllocatorPolicy把显存配额固化CI 用四组断言守护“共享 Arena 下显存有上限、session 复用开启”。记住Java 侧做 CUDA 推理session 要复用、Arena 要共享别让显存随请求数线性膨胀。
