[code-reflection] RFR: Model lifetimes of onnx session-related objects more explicitly [v5]
Paul Sandoz
psandoz at openjdk.org
Mon Mar 3 22:11:17 UTC 2025
On Mon, 3 Mar 2025 22:07:46 GMT, Maurizio Cimadamore <mcimadamore at openjdk.org> wrote:
>> The class representing an onnx session is auto closeable. But, in the current code, a session is closed immediately after its `run` method is called. This is problematic because a session returns some ORTValues (tensors) which also need to be freed, but that cannot be freed immediately after calling `run` (as they need to be used by clients).
>>
>> To address this problem, I tweaked the session code to accept an external arena. All the allocation of session-related data structures now happens using that external arena. This means that the client can now be in charge of managing the lifetime of a session (see changes to MNIST demo).
>>
>> To test, I tweaked the MNIST code to do 10K iterations on each button pressed. Predictably, a single button pressed resulted in over 3g of memory being leaked. With these changes the memory arrives at ~400K (there is still some minor leak, but not sure worth pushing more).
>>
>> If the changes to the demo are not deemed good, I can withdraw this PR -- I mostly wanted to capture the result of my exploration somewhere.
>
> Maurizio Cimadamore has updated the pull request incrementally with one additional commit since the last revision:
>
> Fix missing arena parameter
Marked as reviewed by psandoz (Lead).
cr-examples/onnx/src/main/java/oracle/code/onnx/OnnxRuntime.java line 138:
> 136: }
> 137:
> 138: public static <T> Tensor<T> execute(MethodHandles.Lookup l, OnnxFunction<Tensor<T>> codeLambda, Arena sessionArena) {
Can we make the `codeLambda` be the last parameter? Then the arguments are more easily identifiable, with the "configuration" stuff prefixed.
-------------
PR Review: https://git.openjdk.org/babylon/pull/332#pullrequestreview-2655464038
PR Review Comment: https://git.openjdk.org/babylon/pull/332#discussion_r1978311695
More information about the babylon-dev
mailing list