microsoft / microsoft/onnxruntime
[Java] Multiple CUDA sessions cause unbounded GPU memory growth; shared-allocator fix (CreateAndRegisterAllocatorV2 / OrtArenaCfg) not exposed in the Java API
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
When two CUDA `OrtSession`s are kept resident on one GPU and used **alternately** within a workload — in our case an embedding model (stage 1) then a cross-encoder reranker (stage 3) of a retrieval pipeline, repeated per request — GPU memory (`nvidia-smi memory.used`) grows by a roughly fixed **~2 GB on every iteration and never plateaus**, ending in a `bfc_arena.cc` allocation failure (OOM) after a handful of iterations.
Key observations (all measured):
- Each model used **alone** plateaus fine (no growth after warm‑up).
- Two **embedding** models alternating grow mildly, then **plateau**. Only **embedding ↔ cross-encoder** alternation fails to plateau.
- Independent of **batch size** (batch=1 leaks the same ~2 GB as batch=32), of sequence length/shape, and of **`arena_extend_strategy`** (both `kNextPowerOfTwo` and `kSameAsRequested` grow).
- All `OnnxTensor` / `OrtSession.Result` are closed (try-with-resources); session count stays constant (no reloads). An in-process probe shows the ~2 GB is allocated inside the reranker session's `run()` and never reused on the next request.
This is consistent with the documented behavior that each session owns its own CUDA arena (`cudaMalloc`, never returned to the OS) and the arena re-extends instead of reusing freed blocks when two arenas interleave on the device.
### The Java-API gap (the actual ask)
The canonical fix is a **shared arena registered on the environment** — `CreateAndRegisterAllocator` / `...V2` + `OrtArenaCfg` + the `session.use_env_allocators=1` config — so multiple sessions draw from one arena. **This is not reachable from the Java binding (1.26.0):**
- `OrtEnvironment` exposes no `createAndRegisterAllocator` / `...V2`.
- There is no `OrtArenaCfg` class in `ai.onnxruntime`.
- `SessionOptions.addConfigEntry("session.use_env_allocators", "1")` exists, but is inert with no env allocator to bind to.
**Request:** expose the shared CUDA allocator (`CreateAndRegisterAllocatorV2` + `OrtArenaCfg`) in the Java API, or document a supported Java-only remedy for multi-session CUDA memory growth.
### Java workarounds we validated (for others hitting this)
1. **Per-session `gpu_mem_limit`** — `OrtCUDAProviderOptions.add("gpu_mem_limit", )`. Forces arena reuse within a ceiling; **fully stops** the growth in our tests. Downside: a hard cap blocks any single model larger than the cap, and per-session caps sum toward total VRAM.
2. **`memory.enable_memory_arena_shrinkage` run option** — `RunOptions.addRunConfigEntry("memory.enable_memory_arena_shrinkage", "gpu:0")`, applied per run. Returns the arena after a run, so no hard cap (large models still load). In our tests it slows the growth ~10× but does not fully eliminate a residual periodic bump.
### Environment
- onnxruntime_gpu **1.26.0**, **Java** binding, CUDA EP, single GPU.
### Related
Appears related to #12920, #14474, #25785, #21577.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the Java binding's OrtEnvironment, SessionOptions, and the native CreateAndRegisterAllocatorV2/OrtArenaCfg APIs named in the issue. Trace how session.use_env_allocators is handled and determine whether the shared CUDA allocator can be exposed safely; done means a supported Java API or documented Java-only remedy for multiple CUDA sessions.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- java
- Domain
- api, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100