[RFC] Switching to JNI for Tensorflow engine

[stu1130]

Background

In the following doc, TensorFlow Java or tf java refers to Java language binding on TensorFlow repo and sig-jvm refers to new java sub-repo outside TensorFlow that uses JavaCpp. There are two executors supported by TensorFlow. On TensorFlow 1.x, it only supports symbolic graph executor called GraphExecutor. When TensorFlow 2.0 was launched, it is recommended to use EagerExecutor that runs the operation in imperative way. In terms of bridge library, TensorFlow Java uses JNI while sig-jvm uses JavaCpp.

Problem

When more and more customers get onboarded with DJL TensorFlow, we ran into several issues. The first one is build problem of TensorFlow native binaries. When upgrading to latest TensorFlow version, it is hard for us to adjust native code such as C APIs to make it work with limited TensorFlow expertise. JavaCpp is yet another black-box for us. It has to build Java and cpp altogether due to its limitation. Similarly, now libtensorflow built by sig-jvm is top of old MKLDNN which has performance degradation by over 50% (see reference issue). We worked around the issue by removing the MKLDNN. The other is GC issue. Our customer Stan from Netflix reported the performance problem caused by JavaCpp. We don’t have control over JavaCpp and don’t deeply understand how it works. All factors make us difficult to provide best optimal TensorFlow libraries for customers. In the doc, we will revisit TensorFlow java and propose several solutions.

Overview

TensorFlow Java

TensorFlow Java package includes two parts: operators and other sources. All operators are generated by Javapoet and op_gen tool. I will dive deeper into it in section TensorFlow Java Operators Generator.

tensorflow/java/src - gen (operator generator tool)
                    |
                    - main -- java/org/tensorflow (java classes & example)
                            |
                            - native (JNI C++ code)

The src/main/java/org/tensorflow has basic Java and utilities classes except for operators themselves. The example classes are Graph, Session, EagerOperation, Tensor, DataType and NativeLibrary that loads JNI. The entry points of C++ code (function signature with native) are also included. There is the other native directory under main. It is mainly for JNI C++ code that interacts with TensorFlow C API. There is an example that demos a simple image classification tf java 1.x graph executor.

Sig-JVM

Sig-JVM have three packages. tensorflow-core provides low-level libraries similar to TensorFlow Java. tensorflow-framework offers more high-level APIs like DJL API package. ndarray is a utility for tensor and data I/O. Currently DJL only depends on tensorflow-core. So we will focus on it. In tensorflow-core, sig-jvm also copied an amount of Java code from original TensorFlow repo. They replaced the JNI layer with JavaCpp mapping code. There are also adding new stuff like more datatype and integration code with their ndarray package.

tf-core-api/src - bazel --- api_def (operator definition prototxt)
                |         |
                |         - op_generator (operator generator tool)
                |
                - gen (gen operators/ javacpp bridge code/ protobuf gen code)      
                |
                - main/java/org/tensorflow (the code copied from tf java/
                                            new dtype/ javacpp mapping code)

Similarly, they use TensorFlow Java Operators Generator and put tool code in tensorflow-core-api/src/bazel/op_generator/. The generated operator are also all check-in tensorflow-core-api/src/gen/java/org/tensorflow/op/.

TensorFlow Java Operators Generator

Both TensorFlow Java and sig-jvm use generated operator. The source code is in tensorflow/java/src/gen. A java class called OperatorProcessor.java that uses Javapoet and a binary gen_op written by C++ are involved. The operator is produced during the compilation time. It first builds gen_op and then packed into libtensorflow jar along with other components. The operator definition could be found by in tensorflow/core/api_def/java_api.

Take concat for instances.

Here is the original proto txt file.

op {
  graph_op_name: "ConcatV2"
  endpoint {
    name: "Concat"
  }
}

The generated ops on Java side are actually operator functions that wraps around OperationBuilder.

public final class Concat<T> extends PrimitiveOp implements Operand<T> {
    private Output<T> output;

    public static <T, U extends Number> Concat<T> create(Scope scope, Iterable<Operand<T>> values, Operand<U> axis) {
        OperationBuilder opBuilder = scope.env().opBuilder("ConcatV2", scope.makeOpName("Concat"));
        opBuilder.addInputList(Operands.asOutputs(values));
        opBuilder.addInput(axis.asOutput());
        opBuilder = scope.applyControlDependencies(opBuilder);
        return new Concat(opBuilder.build());
    }

    public Output<T> output() {
        return this.output;
    }

    public Output<T> asOutput() {
        return this.output;
    }

    private Concat(Operation operation) {
        super(operation);
        int outputIdx = 0;
        byte var10002 = outputIdx;
        int var3 = outputIdx + 1;
        this.output = operation.output(var10002);
    }
}

We can actually get rid of java layer and implement the logic on TfNDArray.java.
The real JNI call is

// for tf 1.x Graph Session in Session.java
private static native byte[] run(
      long handle,
      byte[] runOptions,
      long[] inputTensorHandles,
      long[] inputOpHandles,
      int[] inputOpIndices,
      long[] outputOpHandles,
      int[] outputOpIndices,
      long[] targetOpHandles,
      boolean wantRunMetadata,
      long[] outputTensorHandles);
      
// for tf 2.x Eager Session in EagerOperationBuilder.java
# There are more native function call here, just name one for instance
private static native long[] execute(long opHandle); 

Proposed Solutions

1. Switch to TensorFlow Java

We copy JNI source code and whatever Java classes we need for DJL, own and maintain it. As a result, we have full control over the source code. We can solve any memory issue on JNI. Building custom libtensorflow should be also straightforward following the official doc. We might consider to get rid of operators and generator and do it in MXNet op builder way.

pros:

• We have control over the source code. We have flexibility to add new feature requested by customers, optimize the performance and clean up the Java operators we don’t need. Beside Interaction between Tensorflow Java and JavaCPP Pointer deallocation tensorflow/java#208, another example is that https://github.com/alibaba/flink-ai-extended is still using TensorFlow 1.x, they are ran into crash issue due to instability of Java APIs, if we control the tesorflow java source code, we might solve the issue and convince them to use DJL.
• We have better debuggability for GC issues and apply work around solution if it can’t be properly fixed. In worse scenario where no solution can be applied, we should be able to give customers recommendation to avoid the issue.
• Customers can bring their own libtensorflow on the fly as long as C APIs are not changed.

cons:

• We now own the build and publish pipeline for new JNI layer. It will be painful especially on windows platform probably.
• Although we can refer to existing TensorFlow Java and sig-jvm, there is still an amount of effort to implement JNI and Java classes
• We will have to redo multithreading benchmark in terms of throughput and memory usage
• less support from sig-jvm community in terms of build problem and code contribution

2. Stick with sig-jvm but workaround or fork JavaCpp

We study JavaCpp source code and understand what mechanism it uses to release native resource. And work around on our side like what we did to reuse objects on JNA for performance ads team. If it doesn’t work out, we can fork TensorFlow Java or even JavaCpp and adjust the source code to meet our customers’ need. The benefit is we are not far away from sig-jvm and can leverage their power. But we might end up doing lots of changes and the time spent on it is more than approach 1.

pros:

• Nothing critical issue so far. The APIs are working fine. In particular, there are customers deploying the Tensorflow model on production with multi-threading.
• We can get support from sig-jvm community.
• We don’t own build and publish pipeline

cons:

• We heavily rely on sig-jvm to build latest version of TensorFlow
• We may not be able to solve high GC issue and can’t give customers recommendation before we really understand JavaCpp

Q & A

Q: How do we deal with swig?

The swig is on the way to be deprecated. It was used to act as a bridge connecting Python and C++. Now they switch to pybind11. For TensorFlow java, I can’t find any *.i file. So the refactoring effort should be done.

Reference Issues

1. • Sig-JVM Java 0.2.0 with MKLDNN is 3 times slower than TF Java 1.14 on same model
   1. 
   2. JavaCPP MKLDNN version remains at 0.21: https://github.com/bytedeco/javacpp-presets/tree/master/mkl-dnn
2. GPU performance
   1. The Java Tensorflow library does not seem to be using GPU tensorflow/java#140
   2. Performance regression on GPU tensorflow/java#134
3. JavaCPP GC: Interaction between Tensorflow Java and JavaCPP Pointer deallocation tensorflow/java#208
4. Build issues:
   1. Sig-JVM GitHub Action time out if enable XLA or more CUDA arch support
      1. Update GPU Compute Capacity support to match tensorflow tensorflow/java#200

[saudet]

Hello, JavaCPP author here! I'm sorry you feel of JavaCPP as a "black box", but I would be happy to support your needs in any way that is required. JavaCPP is basically to Java what cython or pybind11 with setuptools is to Python. There is no reason to hack JNI manually in your build system. For reference, a minimally wrapped C API, for example, in the case of MXNet, looks like what I did for this pull request: apache/mxnet#19797.

1. • Sig-JVM Java 0.2.0 with MKLDNN is 3 times slower than TF Java 1.14 on same model
   1. 
   2. JavaCPP MKLDNN version remains at 0.21: https://github.com/bytedeco/javacpp-presets/tree/master/mkl-dnn

2. GPU performance

   1. tensorflow/java#140
   2. tensorflow/java#134

That's unrelated to JNI or JavaCPP and should be fixed with TF 2.4.x, see this pull request: tensorflow/java#212.

3. JavaCPP GC: tensorflow/java#208

To disable using the GC, we can set the "org.bytedeco.javacpp.nopointergc" system property to "true", and that's it!

4. Build issues:
   1. Sig-JVM GitHub Action time out if enable XLA or more CUDA arch support
      1. tensorflow/java#200

We're only having problems securing compute resources. Neither Google nor Oracle are being too helpful. Our requests have gone nowhere for over a year now. If Amazon can help us get access to hardware more quickly, we would be glad to use it!

/cc @karllessard

[Craigacp]

It's also worth noting that SIG-JVM has spent it's time focusing on adding features and improving compatibility with TF 2 in python. Things like training, SavedModel & function support. We've also been working on making it more typesafe and improving the usability from Java. We've not got around to the performance work yet, but as we start to reach feature parity then we can move on to looking at performance.

[karllessard]

TF Java in the main TF repository (the 1.x series) is no longer maintained and you should not rely on it. The new version maintained by SIG JVM is still in a pre-release mode and, like @Craigacp pointed out, current efforts are more aiming the API and functionalities than the performance issues right now. But we know as a certainty that they will need to be addressed before we can do any official release of the library.

Since both TF Java and JavaCPP are projects open to contributions, it would be very valuable if we can merge our efforts and try to identify the problems at the source instead of writing workarounds in libraries consuming them, like DJL (I personally do not recommend to start writing your own JNI layer as TensorFlow's C ABI is becoming way more complex than it used to be in the 1.x days).

@saudet already mentioned that we are currently upgrading to TF 2.4 and that could be a good starting point to start investigating more deeply into these core issues, and we would certainly benefit from the help of external contributors as well. Still, I suspect there won't be a lot of work required to reach the performances we are looking for.

[stu1130]

Thanks @saudet for being active and @karllessard to point out EagerSession in tensorflow/java#208 (comment).
@Craigacp we will probably focus on low-level TensorFlow APIs. Ideally we only use javacpp layer and interact with our internal APIs to reduce intermediate java object creation. For example, use raw OpBuilder instead of generated operators. But I agreed we will lose contribution and effort from tensorflow/java.

Here are issues we are facing.

1. To reduce bad performance caused by GC and avoid OOM
   1. In terms of releasing native resource, we would need to figure out a way to integrate with javacpp PointerScope, so we can release native memory as soon as possible. As such, we can get rid of GC and set org.bytedeco.javacpp.nopointergc=true without running out of memory.
   2. We were asked by internal amazon customers to reuse small objects like Pointer class in JNA to reduce the GC pressure. They have a quite strict throughput and latency requirement. Can we do this kind of optimization in javacpp or TF java? If we use JNI, Pointer would be pure long primitive type. No object, no problem.
2. We are happy to offer our instances. CodeBuild is good candidate. We can create a web hook from CodeBuild for you once we get the admin in TensorFlow/java repo. If the machine problem is solved, could we expect the SLA of new version like 1 month? Btw, I looked at the javacpp tutorial but still have a question. Do we have to build tensorflow from source everytime to use javacpp?
3. Could we make the libtensorflow plugable? User can switch to their custom libtensorflow on the fly without building from source.
4. In terms of API stability, we can expect high-level APIs changes as it is still pre-release but is it safe to say low-level javacpp interaction part is stable and will be barely changed in the future?

@saudet we are also interested in how javacpp do performance optimization so it is better than hand-written JNI. Our PyTorch JNI might be benefited from it.

Again, Thanks for everyone's feedback. Those are valuable to DJL team.

We can also schedule a meeting. Let me know when is best time for you.

[karllessard]

I think that is a good observation, @Craigacp. Basically, both the operation and any of its outputs is attached to the EagerSession in scope and will remain in memory until the session is closed (or until the GC collects them eventually but as we know, we cannot entirely rely on it, especially when there is a very large number of resources allocated).

So it's all about managing the lifetime of the session returned by this.manager.getEagerSession() in this case. If the session remains active during the whole process (like the default EagerSession does), it might lead to memory issues.

[roywei]

Hi @Craigacp @karllessard thanks for your inputs!

Yes, we closed the input tensor after TfNDArray creation, only keep the constant operand for future eager operations(usually preprocess before prediction). The preprocess output will be converted back to Tensor and feed into Classic Session(loaded from SavedModel) for inference. We use try catch blocks around NDManagers similar to using eager sessions.

DJL is using NDManagers to manage NDArray lifecycles. When NDManager closes, it will close all NDArrays it created. In TF engine case, we attach one eagerSession with each NDManager, and it will close all constants/tensors attached to it. see: https://github.com/awslabs/djl/blob/master/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDManager.java#L428

There could be multiple NDManagers and sub managers created during end to end inference. For each batch of input, there will be a TranslatorContext Manager that handle preprocess input and postprocess output, all intermediate NDArray/Tensors/Constants will be closed after each predict call.
https://github.com/awslabs/djl/blob/master/api/src/main/java/ai/djl/inference/Predictor.java#L136

Only SystemManager(like the default eagersession) persist the entire process but we don't use it to create any NDArrays, it usually does not hold data.

On side note, unrelated to memory, if we use multiple eager sessions like current setup. Each eager session will try to create its own thread pool resulting too many threads leading to performance issues. We end up using environment variables like TF_NUM_INTEROP_THREADS and TF_NUM_INTRAOP_THREADS to control the threads created. I experimented with preserving eagerSessions (try to reuse it) but it may lead to memory issue and it's not thread safe.(only classic Session is thread safe)

If you have any suggestions on how to improve this please let us know. Thanks!

[saudet]

Thanks @saudet for being active and @karllessard to point out EagerSession in tensorflow/java#208 (comment).
@Craigacp we will probably focus on low-level TensorFlow APIs. Ideally we only use javacpp layer and interact with our internal APIs to reduce intermediate java object creation. For example, use raw OpBuilder instead of generated operators. But I agreed we will lose contribution and effort from tensorflow/java.

If it's something that has enough demand, I think it could still make sense to have a "low-latency inference module" of sorts in https://github.com/tensorflow/java, but in extreme cases, I think we'd be better served with something else like, for example, TVM:
http://bytedeco.org/news/2020/12/12/deploy-models-with-javacpp-and-tvm/

Here are issues we are facing.

1. To reduce bad performance caused by GC and avoid OOM

   1. In terms of releasing native resource, we would need to figure out a way to integrate with javacpp PointerScope, so we can release native memory as soon as possible. As such, we can get rid of GC and set org.bytedeco.javacpp.nopointergc=true without running out of memory.
   2. We were asked by internal amazon customers to reuse small objects like Pointer class in JNA to reduce the GC pressure. They have a quite strict throughput and latency requirement. Can we do this kind of optimization in javacpp or TF java? If we use JNI, Pointer would be pure long primitive type. No object, no problem.

Wow, now that is starting to sound like something I want to work on :)

We can of course reuse Pointer objects, although it may lead to awkward code, and it's also possible to map pointer types to long with JavaCPP, but it only really makes sense to do that for opaque types. For normal data structures, we need to map them into some class to be able to access them from Java. In any case, for opaque types, it's really easy to do. @lanking520 already asked about mapping opaque types from MXNet to Pointer, but we can just as easily map them to long by changing only a single line like I've done in commit saudet/incubator-mxnet@cde6f97 just to show you what that looks like. Mapping opaque types that way though makes it impossible to use PointerScope, so I don't think it's a good idea.

2. We are happy to offer our instances. CodeBuild is good candidate. We can create a web hook from CodeBuild for you once we get the admin in TensorFlow/java repo. If the machine problem is solved, could we expect the SLA of new version like 1 month? Btw, I looked at the javacpp tutorial but still have a question. Do we have to build tensorflow from source everytime to use javacpp?

It looks like we can kind of use CodeBuild from GitHub Actions with this https://github.com/aws-actions/aws-codebuild-run-build? That sounds awesome! @karllessard What do you think? I'm not sure who would be able to provide an SLA at this point though. I think we'd have to discuss this with our respective employers see what we can agree on, which is probably @agibsonccc and @gsw85 in my case.

Once TF Java is built, we can deploy the JAR files, and we don't need to build from source after that, no. Those binaries can be reused as is by other higher-level modules.

3. Could we make the libtensorflow plugable? User can switch to their custom libtensorflow on the fly without building from source.

Sure, it already works that way. We just need to provide a different libtensorflow_cc.so.2 either in the class path, possibly bundled in a JAR file, or as part of the java.library.path. The JNI bindings compiled in libjnitensorflow.so don't need to be updated as long as the C/C++ API mapped remains compatible.

4. In terms of API stability, we can expect high-level APIs changes as it is still pre-release but is it safe to say low-level javacpp interaction part is stable and will be barely changed in the future?

We can generally assume that it's stable, but TF Core has released breaking changes for string tensors with 2.4.x, so we should probably expect things to change bit from time to time.

@saudet we are also interested in how javacpp do performance optimization so it is better than hand-written JNI. Our PyTorch JNI might be benefited from it.

I'm sure it could! We got a performance boost by getting rid for JNI and using JavaCPP for TF with ~50% less code, see tensorflow/java#18 (comment).

Again, Thanks for everyone's feedback. Those are valuable to DJL team.

We can also schedule a meeting. Let me know when is best time for you.

Thanks to you for sharing! I live in Japan and generally available between 9 am and 5 pm JST on weekdays, although yesterday was a holiday. I can usually do after 10 pm as well after the kids are sleeping, but that's not ideal... Do you already have a Slack channel or something to organize meetings like this? If not, let's start an email thread?

[Craigacp]

On side note, unrelated to memory, if we use multiple eager sessions like current setup. Each eager session will try to create its own thread pool resulting too many threads leading to performance issues. We end up using environment variables like TF_NUM_INTEROP_THREADS and TF_NUM_INTRAOP_THREADS to control the threads created. I experimented with preserving eagerSessions (try to reuse it) but it may lead to memory issue and it's not thread safe.(only classic Session is thread safe)

If you have any suggestions on how to improve this please let us know. Thanks!

@roywei You can use the config protobuf to set those without having to use an environment variable, but yes the eager session does hold on to a few too many things at the moment. There is work on tensor scoping which should make it easier to release resources, but that's not finished yet. I'm not sure that the current approach of folding tensors into eager operations rather than keeping the tensors around will mesh well with the work that we are doing in TF-Java to try and make it easier to clean things up.

Here are issues we are facing.

1. To reduce bad performance caused by GC and avoid OOM

   1. In terms of releasing native resource, we would need to figure out a way to integrate with javacpp PointerScope, so we can release native memory as soon as possible. As such, we can get rid of GC and set org.bytedeco.javacpp.nopointergc=true without running out of memory.
   2. We were asked by internal amazon customers to reuse small objects like Pointer class in JNA to reduce the GC pressure. They have a quite strict throughput and latency requirement. Can we do this kind of optimization in javacpp or TF java? If we use JNI, Pointer would be pure long primitive type. No object, no problem.

Wow, now that is starting to sound like something I want to work on :)

We can of course reuse Pointer objects, although it may lead to awkward code, and it's also possible to map pointer types to long with JavaCPP, but it only really makes sense to do that for opaque types. For normal data structures, we need to map them into some class to be able to access them from Java. In any case, for opaque types, it's really easy to do. @lanking520 already asked about mapping opaque types from MXNet to Pointer, but we can just as easily map them to long by changing only a single line like I've done in commit saudet/incubator-mxnet@cde6f97 just to show you what that looks like. Mapping opaque types that way though makes it impossible to use PointerScope, so I don't think it's a good idea.

@saudet @stu1130 It might be better to go the other way and make Pointer completely immutable. Then the C2 and Graal JIT compilers can perform escape analysis, scalarise the objects away and they never touch the heap. Unfortunately that only happens when the JIT compilers get hold of the relevant code, so it's hard to figure out the benefit without proper benchmarking or long-ish running programs (as they will still be allocated by the interpreter and C1 so anything short lived won't see any effect). This trade off might not be worth it for DJL or TF-Java, but it's hard to know without prototyping.

Plus complete immutability will make it easier to migrate to primitive classes (https://openjdk.java.net/jeps/8251554) which will also reduce allocation pressure as they may be inlined into host objects or arrays.

[saudet]

@saudet @stu1130 It might be better to go the other way and make Pointer completely immutable. Then the C2 and Graal JIT compilers can perform escape analysis, scalarise the objects away and they never touch the heap. Unfortunately that only happens when the JIT compilers get hold of the relevant code, so it's hard to figure out the benefit without proper benchmarking or long-ish running programs (as they will still be allocated by the interpreter and C1 so anything short lived won't see any effect). This trade off might not be worth it for DJL or TF-Java, but it's hard to know without prototyping.

Plus complete immutability will make it easier to migrate to primitive classes (https://openjdk.java.net/jeps/8251554) which will also reduce allocation pressure as they may be inlined into host objects or arrays.

Right, that's the general direction that Panama is taking. That's probably the kind of thing I'll start thinking about when Panama starts becoming useful, including on platforms such as Android. At this point, ART still has a pretty hard time optimizing for such cases, so we often still have to make do with good old recycling of objects anyway.

[skirdey]

@saudet
Does this table around JavaCPP and its GC activites look accurate?

BLOCKING?	noPointerGC=true	maxPhysicalBytes>0	maxBytes>0	maxPhysicalBytes=0	maxBytes=0
noPointerGC=true	not blocking	?		?	?
maxPhysicalBytes>0	?	blocking	blocking	?	?
maxBytes>0	?	blocking	?	?	?
maxPhysicalBytes=0	?	?	?	?	?
maxBytes=0	?	?	?	?	?

[saudet]

I'm not sure I understand what the pairing refers to, but essentially it should never block when noPointerGC is true or when both maxBytes and maxPhysicalBytes are 0. If one of them are greater than 0, blocking (waiting after the GC) may occur when at least one of those thresholds is exceeded, but it's not supposed to be blocking otherwise. If that's happening though, please let me know.

[skirdey]

the pairing is mostly when you use both params at the same time.

[karllessard]

We have identified a potential issue in TF Java that could result in OOM when eager sessions were remaining alive for a relatively long time. There is a PR actually opened to fix it and I think this solution could also resolve some issues that was observed previously when using DJL with TensorFlow.

Now, TF Java is offering very lightweight bindings to interact with the TensorFlow runtime in a Java idiomatic way. While nothing prevents users to call directly the JavaCPP wrappers inside the library to access directly the raw TensorFlow's ABI, TF Java is not designed nor maintained with this purpose and code generated by JavaCPP, always subject to changes, is reserved for internal usage.

I would like to reiterate my previous suggestion that we should try to solve any problems with the integration of TF Java in DJL by investigating (and potentially fixing) the issues in both libraries first, before attempting a drastic change of design as described in the current proposal.

[saudet]

Please note that the "fix" only reverts it to the previous behavior from TF 1.x. It still relies on GC to prevent running out of memory, but since GC does not track off-heap memory, there are cases where it will still run out of memory anyway. If DJL needs something more reliable, please let us know. I agree that we should be working together to make something useful for all the community.