Possible improvement for the OrtEngine

[andreabrduque]

Description

I think it would be interesting to add support for onnx models to disable session threadpools by using a shared global thread pool.

As far as my understanding of DJL goes, when we create one predictor of different models for each thread, we end up running several onnx sessions in parallel.

In my benchmarks being able to control the threadpool when multiple onnx sessions are run in parallel offers slightly better performance and better resource utilisation for some models.

Will this change the current api? How?

I thought about passing the option disablePerSessionThreads as a setter, such as done for the other options from SessionOptions.

However It would also require to pass global thread pool settings through the OrtEnvironment.ThreadingOptions

The part where I got really stuck was because the method OrtEnvironment.getEnvironment() is called both in OrtEngine and OrtNDManager implementations. According to the Onnxruntime java API, there is no way to guarantee that the environment has the appropriate thread pool configuration. So if I pass ThreadingOptions to an environment and then retrieve it again, I will get an IllegalStateException.

I hacked a bit here

Who will benefit from this enhancement?

This benefits the use case in which more than one model instance is being used in parallel (for example, if we load a ZooModel each in one GPU and do that for 4 GPUs).

[frankfliu]

@andreabrduque
OrtEnvironment use a global singleton. If you want to customize ThreadingOptions, you can call the following before DJL is loaded:

OrtEnvironment.getEnvironment(OrtLoggingLevel.ORT_LOGGING_LEVEL_WARNING, "ort-java", threadOptions);

DJL will pick the initialized OrtEnvironment.

[andreabrduque]

Hey @frankfliu , thanks for the reply.

The problem with is that I need to be able to set OrtSession.SessionOptions.disablePerSessionThreads() for the ThreadingOptions to make an effect, and the SessionOptions are passed through DJL.

I can make a PR to add that option, but if users try to use the disablePerSessionThreads() without passing the global threading options, they will get an exception, but I can add a section to the docs to warn that.