Repository: deepjavalibrary/djl (java)

Issue:
ONNX Extensions
I tried to get around the limitations I encountered when trying to load tf-text (https://github.com/deepjavalibrary/djl/issues/1803) by converting [the model](https://tfhub.dev/google/universal-sentence-encoder-multilingual-large/3) to ONNX and trying to run it via the ONNX Runtime. This however also failed with `Fatal error: SentencepieceTokenizer is not a registered function/op`. After some googling, it appears that one would need to also load the extensions from https://github.com/microsoft/onnxruntime-extensions. Is there a way to do that in DJL?

Pull request description:
[onnx] Fixes OrtNDArray double close issue
1. Fixes #1804
2. Fixes crash caused by double close OrtNDArray
3. NDArray.get() should use alternativeManager if available
4. getNDArrayInternal() should fail with alternativeManager is disabled

Change-Id: I37767394f3c68dc8a48ba21ae589dfa81407bdb8

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 60d1d02a9..541e93d75 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -74,6 +74,7 @@ public abstract class NDArrayAdapter implements NDArray {
         manager.attachInternal(getUid(), this);
         alternativeManager = ((BaseNDManager) manager).getAlternativeManager();
         if (alternativeManager == null) {
+            // to prevent hybrid engine memory leak
             alternativeManager = manager;
         }
     }
@@ -195,6 +196,12 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray get(NDIndex index) {
+        return get(alternativeManager, index);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void set(Buffer data) {
@@ -1075,7 +1082,11 @@ public abstract class NDArrayAdapter implements NDArray {
     /** {@inheritDoc} */
     @Override
     public NDArrayEx getNDArrayInternal() {
-        return getAlternativeArray().getNDArrayInternal();
+        NDArray array = getAlternativeArray();
+        if (array instanceof NDArrayAdapter) {
+            throw new UnsupportedOperationException("Operation not supported.");
+        }
+        return array.getNDArrayInternal();
     }
 
     /** {@inheritDoc} */
@@ -1083,11 +1094,11 @@ public abstract class NDArrayAdapter implements NDArray {
     public void close() {
         if (!isClosed) {
             manager.detachInternal(getUid());
+            isClosed = true;
             if (alternativeArray != null) {
                 alternativeArray.close();
                 alternativeArray = null;
             }
-            isClosed = true;
         }
     }
 
diff --git a/engines/onnxruntime/onnxruntime-engine/build.gradle b/engines/onnxruntime/onnxruntime-engine/build.gradle
index c961d5691..162663344 100644
--- a/engines/onnxruntime/onnxruntime-engine/build.gradle
+++ b/engines/onnxruntime/onnxruntime-engine/build.gradle
@@ -7,7 +7,8 @@ dependencies {
     testImplementation("org.testng:testng:${testng_version}") {
         exclude group: "junit", module: "junit"
     }
-    testImplementation(project(":testing"))
+    testImplementation project(":testing")
+    testImplementation project(":engines:pytorch:pytorch-engine")
 
     testRuntimeOnly "org.slf4j:slf4j-simple:${slf4j_version}"
 }
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
index 0a1f3758e..07088cdd2 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
@@ -25,11 +25,12 @@ import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.charset.Charset;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
 
 /** {@code OrtNDArray} is the ONNX Runtime implementation of {@link NDArray}. */
 public class OrtNDArray extends NDArrayAdapter {
 
-    private OnnxTensor tensor;
+    private AtomicReference<OnnxTensor> tensor;
 
     /**
      * Constructs an ONNX Runtime NDArray from a {@link OnnxTensor} (internal. Use {@link NDManager}
@@ -41,19 +42,24 @@ public class OrtNDArray extends NDArrayAdapter {
      */
     OrtNDArray(OrtNDManager manager, NDManager alternativeManager, OnnxTensor tensor) {
         super(manager, alternativeManager, null, null, UUID.randomUUID().toString());
-        this.tensor = tensor;
+        this.tensor = new AtomicReference<>(tensor);
         manager.attachInternal(uid, this);
     }
 
-    OnnxTensor getTensor() {
-        return tensor;
+    /**
+     * Returns the {@code OnnxTensor} representation of this OrtNDArray.
+     *
+     * @return the {@code OnnxTensor} representation of this OrtNDArray
+     */
+    public OnnxTensor getTensor() {
+        return tensor.get();
     }
 
     /** {@inheritDoc} */
     @Override
     public DataType getDataType() {
         if (dataType == null) {
-            dataType = OrtUtils.toDataType(tensor.getInfo().type);
+            dataType = OrtUtils.toDataType(tensor.get().getInfo().type);
         }
         return dataType;
     }
@@ -62,7 +68,7 @@ public class OrtNDArray extends NDArrayAdapter {
     @Override
     public Shape getShape() {
         if (shape == null) {
-            shape = new Shape(tensor.getInfo().getShape());
+            shape = new Shape(tensor.get().getInfo().getShape());
         }
         return shape;
     }
@@ -70,8 +76,7 @@ public class OrtNDArray extends NDArrayAdapter {
     /** {@inheritDoc} */
     @Override
     public void intern(NDArray replaced) {
-        tensor.close();
-        tensor = ((OrtNDArray) replaced).tensor;
+        tensor.getAndSet(((OrtNDArray) replaced).getTensor()).close();
     }
 
     /** {@inheritDoc} */
@@ -85,7 +90,7 @@ public class OrtNDArray extends NDArrayAdapter {
     @Override
     public String[] toStringArray(Charset charset) {
         try {
-            Object obj = tensor.getValue();
+            Object obj = tensor.get().getValue();
             if (obj instanceof String) {
                 // Scalar type;
                 return new String[] {(String) obj};
@@ -102,15 +107,16 @@ public class OrtNDArray extends NDArrayAdapter {
         if (getDataType() == DataType.STRING) {
             throw new IllegalArgumentException("Please use toStringArray() for String NDArray.");
         }
-        return tensor.getByteBuffer().order(ByteOrder.nativeOrder());
+        return tensor.get().getByteBuffer().order(ByteOrder.nativeOrder());
     }
 
     /** {@inheritDoc} */
     @Override
     public void close() {
-        if (!isClosed) {
-            tensor.close();
-            super.close();
+        super.close();
+        OnnxTensor ortTensor = tensor.getAndSet(null);
+        if (ortTensor != null) {
+            ortTensor.close();
         }
     }
 }


Test patch:
diff --git a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
index c5cba2a56..a07f59672 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
@@ -14,6 +14,7 @@ package ai.djl.onnxruntime.engine;
 
 import ai.djl.Model;
 import ai.djl.ModelException;
+import ai.djl.engine.Engine;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Classifications;
 import ai.djl.ndarray.NDArray;
@@ -25,6 +26,7 @@ import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
+import ai.onnxruntime.OrtException;
 
 import org.testng.Assert;
 import org.testng.SkipException;
@@ -32,6 +34,7 @@ import org.testng.annotations.Test;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.reflect.Field;
 import java.nio.file.Files;
 import java.nio.file.Path;
 
@@ -88,7 +91,7 @@ public class OrtTest {
     }
 
     @Test
-    public void testNDArray() {
+    public void testNDArray() throws OrtException {
         TestRequirements.notArm();
 
         try (NDManager manager = OrtNDManager.getSystemManager().newSubManager()) {
@@ -111,6 +114,9 @@ public class OrtTest {
 
             array = manager.create(new String[] {"string1", "string2"});
             Assert.assertEquals(array.toStringArray()[1], "string2");
+
+            float[][] value = (float[][]) ((OrtNDArray) ones).getTensor().getValue();
+            Assert.assertEquals(value[0], new float[] {1, 1});
         }
     }
 
@@ -118,7 +124,7 @@ public class OrtTest {
     public void testStringTensor() throws ModelException, IOException, TranslateException {
         TestRequirements.notArm();
 
-        System.setProperty("ai.djl.onnx.disable_alternative", "true");
+        setAlternativeEngineDisabled(true);
         Criteria<NDList, NDList> criteria =
                 Criteria.builder()
                         .setTypes(NDList.class, NDList.class)
@@ -137,6 +143,38 @@ public class OrtTest {
             Assert.assertEquals(result.size(), 2);
             Assert.assertEquals(result.get(0).toLongArray(), new long[] {1});
         }
-        System.clearProperty("ai.djl.onnx.disable_alternative");
+        setAlternativeEngineDisabled(false);
+    }
+
+    @Test
+    public void testAlternativeArray() {
+        TestRequirements.notArm();
+
+        try (NDManager manager = OrtNDManager.getSystemManager().newSubManager()) {
+            NDArray array = manager.zeros(new Shape(1, 2));
+            Assert.assertEquals(array.get(0).toFloatArray(), new float[] {0, 0});
+        }
+
+        setAlternativeEngineDisabled(true);
+        try (NDManager manager = OrtNDManager.getSystemManager().newSubManager()) {
+            NDArray array = manager.zeros(new Shape(1, 2));
+            Assert.expectThrows(UnsupportedOperationException.class, () -> array.get(0));
+        }
+        setAlternativeEngineDisabled(false);
+    }
+
+    private void setAlternativeEngineDisabled(boolean enable) {
+        System.setProperty("ai.djl.onnx.disable_alternative", String.valueOf(enable));
+        Engine engine = Engine.getEngine(OrtEngine.ENGINE_NAME);
+        try {
+            Field field = OrtEngine.class.getDeclaredField("initialized");
+            field.setAccessible(true);
+            field.setBoolean(engine, false);
+            field = OrtEngine.class.getDeclaredField("alternativeEngine");
+            field.setAccessible(true);
+            field.set(engine, null);
+        } catch (ReflectiveOperationException ignore) {
+            // ignore
+        }
     }
 }