[code-reflection] RFR: Model lifetimes of onnx session-related objects more explicitly

[Maurizio Cimadamore]

On Fri, 28 Feb 2025 12:42:24 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.

A possible way forward:
* create a class called OnnxArena which implements FFM's Arena
* move all the model-related execute methods from OnnxRuntime to OnnxArena
* OnnxArena will create a confined arena internally -- to deal with the lifetime of an underlying onnx session

Then the MNIST demo code will look something like:


try (OnnxArena onnxArena = OnnxArena.of()) {
      var prediction = onnxArena.execute(MethodHandles.lookup(),
                           () -> cnn(imageTensor));
      var result = prediction.data().toArray(ValueLayout.JAVA_FLOAT);
      ...
} // close the onnx session, and release all memory allocated within the arena

-------------

PR Comment: https://git.openjdk.org/babylon/pull/332#issuecomment-2693812580