Java's OnnxSequence Does Not Free Resources on Close

[sbonebrake]

Describe the issue

It seems several years ago the behavior of OnnxSequence.java changed to no longer free up memory when it closes. This can lead to a memory leak that can will eventually eat up all your memory an crash your container. The change appears to be intentional in this PR #13012 by @Craigacp for performance reasons. Since OnnxValue implements AutoCloseable it feels like the default behavior should be to close each OnnxTensor it creates. I think this is the behavior of the other classes that implement OnnxValue as well.

I think a reasonable solution would be to have OnnxSequence close its OnnxTensor's by default but allow for the existing behavior by adding a method getValueUnmanaged(). I'm happy to work on a pull request, but thought I would gather some feedback first. Thanks!

To reproduce

public float[] predict(float[][] featureVectorArray) throws OrtException {
  try (OnnxTensor onnxTensor = OnnxTensor.createTensor(ortEnvironment, featureVectorArray);
      Result result = session.run(Collections.singletonMap(inputName, onnxTensor))) {
    OnnxValue onnxValue = result.get(1);
 
    List<OnnxMap> onnxProbabilities = (List<OnnxMap>) onnxValue.getValue();
    float[] probabilities = new float[onnxProbabilities.size()];
    for (int i = 0; i < onnxProbabilities.size(); i++) {
      // THESE VALUES ARE NOT CLOSED
      probabilities[i] = (float) onnxProbabilities.get(i).getValue().get(1L);
    }
    return probabilities;
  }

Urgency

No response

Platform

Linux

OS Version

Docker with azurelinux 2.0

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.20.0

ONNX Runtime API

Java

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[Craigacp]

The semantics here are a little complicated. If you don't pull the values out of the OnnxSequence then there is nothing to close as they are allocated in the getValue call. If you do pull the values out, then the sequence would need to track them forever, and if getValue is called multiple times it will make multiple wrapper objects which need tracking (as they would be closed independently). The similar case is OrtSession.Result which does close the tensors it contains, but you always get the same Java object out of each call to result.get(), whereas at the moment you get different ones out of OnnxSequence.getValue (this is something that might be worth fixing, because it is weird, though it's how the C API behaves which is why it's like that). It is noted in the documentation for OnnxSequence.getValue that you need to close the values it returns separately, and there is a close utility method for doing that on OnnxValue.

It is annoying how it interacts with the Result object, so I'll have a think more about what the semantics should be.

[sbonebrake]

Thanks for the reply @Craigacp. I'm still trying familiarize myself with JNI and these API calls. What is the reasoning behind getting different nativeHandles for repeated calls to OnnxSequence.getValue()? Can the underlying values change for repeated calls? If the values cannot change, then perhaps OnnxSequence can lazily initialize and store references to maps and tensors in getValue() and reuse them for repeated calls. Once we have that, then close() them when it sequence close() is called.

[Craigacp]

The reasoning is that it follows what the C API's GetValue call does, which gives back the map or tensor specified at that location in the sequence, where it's documented that you must close that value separately from the sequence. It doesn't have any other constraints on the method behaviour. Consequently I don't know if the underlying values can change as I'm not sure if the ORT C API actually does give out different values (which would induce copying and memory allocation) or if you get a reference to the same memory. I've tried to use the C API as documented without letting undocumented behaviour leak into the way I'm using it, but this is an area that could use clarity from the MS team.

We don't allow users to write into tensors directly, so that means things shouldn't change from the Java side, but it's feasible that you could pull a tensor out and use it as a pinned output for another call to session.run() where it would change the value of the tensor.

I can write some tests to see what the behaviour of the C API is.

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[snnn]

This issue has been automatically closed as 'not planned' because it has been marked as 'stale' for more than 30 days without activity. If you believe this is still an issue, please feel free to reopen it.