Resource and Memory Management

[karllessard]

This issue is used to collect most of our ideas or proposals to improve the memory management scheme present in the current TF Java client.

Actually, TF Java mostly relies on the try-with-resources approach to release native resources, particularly in graph execution mode.

Eager execution mode releases all resources allocated within a EagerSession when this one is closed, and is also listening to objects discarded by the garbage collector at runtime to free their native counterpart as well while the session is valid.

Can we do better?

[saudet]

Can we do better?

Yes :) In any case, one thing I consider important is interoperability with other libraries. A lot of users use, for example, OpenCV alongside TF in ways similar to what is exemplified in this blog post:

• http://bytedeco.org/news/2018/07/17/bytedeco-as-distribution/

JavaCPP already supports PointerScope that way with the C API of TensorFlow as well, and if we could implement other schemes like reference counting in there that users might like, they would automatically become available to all other libraries using JavaCPP.

[saudet]

FWIW, MXNet decided to basically reimplement PointerScope, but in Scala as they are not using the Java language:
https://cwiki.apache.org/confluence/display/MXNET/JVM+Memory+Management
https://github.com/apache/incubator-mxnet/blob/master/scala-package/memory-management.md
Obviously, they've missed the point about interoperability...

[karllessard]

I agree with you @saudet, could make sense if JavaCPP abstract most of the memory management details for layers on top.

Also, for comparison, the PointerScope seems to accomplish pretty much the same job as the EagerSession, except that the later also releases resources as they are discarded by the GC, while still in the « scope ». Is it something we can add also to JavaCPP?

This allows recycling the memory used by unreferenced objects before the scope (or session) is closed.

[Craigacp]

I think binding the API to JavaCPP by surfacing it's types in the public API (by exposing a PointerScope) is going to make it difficult to switch to whatever the next library is which makes native interop simpler.

With eager mode tensors that are returned are likely to be flagged by IDEs as objects which should be closed, so what happens if they are closed? If they are cleaned up by the EagerSession when it goes out of scope then that's fine but will lead to a lot of noise while writing programs using it.

[karllessard]

I think binding the API to JavaCPP by surfacing it's types in the public API (by exposing a PointerScope) is going to make it difficult to switch to whatever the next library is which makes native interop simpler.

Maybe EagerSession, Graph or whatever TF class we expose in the public API can extends from PointerScope to inherit that functionality?

With eager mode tensors that are returned are likely to be flagged by IDEs as objects which should be closed, so what happens if they are closed?

I think graph and eager mode should both use the same memory management scheme. Right now they are distinct because I didn't want to touch existing mechanisms for graphs when I added eager environment support. That being said, in eager mode if you explicitly close a Tensor, then the tensor will be released immediately as expected and the session won't attempt to discard it later.

If they are cleaned up by the EagerSession when it goes out of scope then that's fine but will lead to a lot of noise while writing programs using it.

I'm a bit neutral about this but just an idea like that, if we decide to apply the eager memory management scheme to graphs, then we won't use try-with-resources anymore so Tensor won't need to be (Auto)Closeable. Still, we should let the user release the memory explicitly if they really want to, maybe via another method (e.g. release()) to make IDEs happy?

[Craigacp]

Yeah, if we remove the AutoCloseable interface from it then it will stop being flagged in the IDE (we can leave the close method to allow the user to clean it up manually). The issue is if we're still supporting graph mode then those will need a containing scope (maybe the standard Session object?) which can release them when it goes out of scope. I've been working on an ONNX-runtime Java wrapper which only creates tensors out of an ONNX session object, so that can track them, but we don't create TF Tensors like that at the moment.

We should probably do some testing though as relying upon the GC to free resources when they go out of scope is might cause some issues. Any ML code will probably have very little allocation in the main loop, and the input and output Tensors at inference time are very small objects (because they are basically a pointer), so it's might not create enough GC pressure to trigger a GC. Without a GC it won't clean up the TF Tensors on the native heap, so we could allocate way too much native memory. Reference counting would be significantly more deterministic, but I'm not sure how that would be implemented without the user of the TF API having to understand that the memory was tracked with reference counting.

One further question, @saudet how does the JavaCPP wrapper interact with the ByteBuffer representation of the Tensors? Can JavaCPP convert it's Pointers into a ByteBuffer (or appropriate subtyped Buffer).

[karllessard]

We should probably do some testing though as relying upon the GC to free resources when they go out of scope is might cause some issues.

That’s what I think too, relying on the GC can be unpredictable but we don’t know yet if it’s a real issue in our case. At the time of adding EagerSession, I did a few tests that seems to provide positive results but more extensive testing would be required if we want to apply that scheme overall the library (and to see if object scoping is required at all).

For reference count, @EronWright can you explain a bit your solution? I know it’s not ready and I might be fast-forwarding here, but when I’m looking at this link you provided me, it looks like allocated objects must be unreferenced manually, via unref(), and that forgetting to do it will lead to a memory leak?

IMHO, the Java API should abstract as much as possible that users are dealing with a native library and we should stick to Java idioms, where users don’t have to worry much about memory deallocation. Unfortunately, I don’t know how we can mimic C++ smart pointers in Java in that way.

[karllessard]

One further question, @saudet how does the JavaCPP wrapper interact with the ByteBuffer representation of the Tensors? Can JavaCPP convert it's Pointers into a ByteBuffer (or appropriate subtyped Buffer).

@Craigacp, I know that the question is not addressed to me but I don’t think we have to map JavaCPP pointers to ByteBuffer, as the later only support 32-bits indexation. My plan was to map those pointers directly to the DataBuffer interface, that is used by the upcoming NdArray API.

[Craigacp]

@karllessard sure that would work but will trigger a jni call for every put or get as it's a wrapper around a pointer to the native heap. Which leads to the complexity that ND4J has where it strongly prefers you never pull data into Java for operations, you only work on it with the provided native methods (i.e. you can't map a lambda element-wise over a tensor). From what I understand of a ByteBuffer that doesn't happen as the JVM understands the ByteBuffer so it doesn't trigger a call to jni.

[EronWright]

@karllessard when we talk about reference counting and the need to manually manage the reference count, it is true but there's a few points to make:

1. An object based on AutoCloseable must also be manually dereferenced (closed), but is limited to a reference count of one.
2. A reference-counted object should have a finalizer to avoid a leak.
3. Netty-style reference counting may actually simplify the user code in certain access patterns, e.g. feeding tensors to a call to run(). In this example, observe that the feeds are simpler to work with, and the fetches are properly cleaned up.

[karllessard]

@karllessard sure that would work but will trigger a jni call for every put or get as it's a wrapper around a pointer to the native heap.

Oh I am not aware of the internal details of JavaCPP yet but if it does trigger a JNI call for every I/O operation then yes, I agree, we should map it to a ByteBuffer to avoid it

[EronWright]

@karllessard regarding the use of direct buffers in the NdArray API, I'm concerned about relying only on GC. I have a doubt that the GC does a good job in accounting for direct buffers (e.g. whether they're considered a 'large' object). One often resorts to the use of sun.misc.Unsafe to explicitly clean up.

Check out the Netty project's efforts along those lines, the ByteBuf class which is a reference-counted alternative to wrapper for ByteBuffer. See UnpooledDirectByteBuf::freeDirect. More info on the use of Netty's buffer library (in a standalone manner), here. Unfortunately it uses 32-bit indexes.