Repository: deepjavalibrary/djl
Pull request #3257 (2024-06-16 21:17:32)

[api] Fixes IdEmbedding memory leak
Fixes: #3251

## 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


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3257?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: Patch coverage is `0%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 72.42%. Comparing base [(`bd7f4ca`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bd7f4ca8b7cb88e5e3864814501939884a4d15c4?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) to head [(`6beb95c`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/6beb95c72bf04c626fc4d658632b6bbfc42c8eb8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Report is 12 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3257?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...c/main/java/ai/djl/nn/transformer/IdEmbedding.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3257?src=pr&el=tree&filepath=api%2Fsrc%2Fmain%2Fjava%2Fai%2Fdjl%2Fnn%2Ftransformer%2FIdEmbedding.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25uL3RyYW5zZm9ybWVyL0lkRW1iZWRkaW5nLmphdmE=) | 0.00% | [9 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3257?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #3257      +/-   ##
============================================
+ Coverage     64.26%   72.42%   +8.16%     
- Complexity     6455     7328     +873     
============================================
  Files           676      677       +1     
  Lines         32443    32512      +69     
  Branches       3417     3430      +13     
============================================
+ Hits          20849    23547    +2698     
+ Misses        10016     7350    -2666     
- Partials       1578     1615      +37     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3257?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java b/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
index d135d5c8b..11e1274b6 100644
--- a/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
+++ b/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
@@ -63,15 +63,20 @@ public final class IdEmbedding extends AbstractBlock {
     protected NDList forwardInternal(
             ParameterStore ps, NDList inputs, boolean training, PairList<String, Object> params) {
         NDArray input = inputs.singletonOrThrow();
-        // on info to the right shapes, see: http://beta.mxnet.io/r/api/mx.symbol.gather_nd.html
-        NDArray ids = input.flatten().reshape(1, input.getShape().size());
-        // create the embedding Table
-        NDArray embeddingTable = ps.getValue(embedding, ids.getDevice(), training);
-        // We do not perform a sparse lookup, instead we just project into the table
-        NDArray result = embeddingTable.gatherNd(ids);
-        // we want the original shape of the input + the last dimension of the embedding
-        Shape targetShape = input.getShape().addAll(new Shape(embeddingTable.getShape().get(1)));
-        return new NDList(result.reshape(targetShape));
+        try (NDManager scope = NDManager.subManagerOf(input)) {
+            // on info to the right shapes, see: http://beta.mxnet.io/r/api/mx.symbol.gather_nd.html
+            NDArray ids = input.flatten().reshape(1, input.getShape().size());
+            // create the embedding Table
+            NDArray embeddingTable = ps.getValue(embedding, ids.getDevice(), training);
+            scope.tempAttachAll(embeddingTable);
+            // We do not perform a sparse lookup, instead we just project into the table
+            NDArray result = embeddingTable.gatherNd(ids);
+            result.attach(inputs.getManager());
+            // we want the original shape of the input + the last dimension of the embedding
+            Shape targetShape =
+                    input.getShape().addAll(new Shape(embeddingTable.getShape().get(1)));
+            return new NDList(result.reshape(targetShape));
+        }
     }
 
     /**
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/jni/JniUtils.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/jni/JniUtils.java
index 2af6c6128..4284e7521 100644
--- a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/jni/JniUtils.java
+++ b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/jni/JniUtils.java
@@ -29,6 +29,7 @@ import com.microsoft.ml.lightgbm.lightgbmlib;
 import com.microsoft.ml.lightgbm.lightgbmlibConstants;
 
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.nio.DoubleBuffer;
 import java.nio.FloatBuffer;
 
@@ -104,7 +105,8 @@ public final class JniUtils {
             checkCall(result);
             int length = Math.toIntExact(lightgbmlib.int64_tp_value(outLength));
             if (a.getDataType() == DataType.FLOAT32) {
-                ByteBuffer bb = ByteBuffer.allocateDirect(length * 4);
+                ByteBuffer bb =
+                        ByteBuffer.allocateDirect(length * 4).order(ByteOrder.nativeOrder());
                 FloatBuffer wrapped = bb.asFloatBuffer();
                 for (int i = 0; i < length; i++) {
                     wrapped.put((float) lightgbmlib.doubleArray_getitem(outBuffer, i));
@@ -112,7 +114,8 @@ public final class JniUtils {
                 bb.rewind();
                 return new Pair<>(length, bb);
             } else if (a.getDataType() == DataType.FLOAT64) {
-                ByteBuffer bb = ByteBuffer.allocateDirect(length * 8);
+                ByteBuffer bb =
+                        ByteBuffer.allocateDirect(length * 8).order(ByteOrder.nativeOrder());
                 DoubleBuffer wrapped = bb.asDoubleBuffer();
                 for (int i = 0; i < length; i++) {
                     wrapped.put(lightgbmlib.doubleArray_getitem(outBuffer, i));
diff --git a/engines/ml/xgboost/build.gradle.kts b/engines/ml/xgboost/build.gradle.kts
index 4a6e6cfa2..5dfa5cb0d 100644
--- a/engines/ml/xgboost/build.gradle.kts
+++ b/engines/ml/xgboost/build.gradle.kts
@@ -43,7 +43,7 @@ tasks {
     compileJava { dependsOn(processResources) }
 
     processResources {
-        val jnilibDir = buildDirectory / "resources/java/main/lib/linux/aarch64/"
+        val jnilibDir = buildDirectory / "resources/main/lib/linux/aarch64/"
         inputs.properties(mapOf("xgboostVersion" to libs.versions.xgboost.get()))
         outputs.dir(jnilibDir)
         doLast {
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngineProvider.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngineProvider.java
index 5f45116f6..c7ea55745 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngineProvider.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngineProvider.java
@@ -27,6 +27,12 @@ public class MxEngineProvider implements EngineProvider {
     /** {@inheritDoc} */
     @Override
     public int getEngineRank() {
+        String osName = System.getProperty("os.name");
+        String osArch = System.getProperty("os.arch");
+        if (osName.startsWith("Mac") && "aarch64".equals(osArch)) {
+            // MXNet doesn't support macOS M1
+            return 99;
+        }
         return MxEngine.RANK;
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainBertOnCode.java b/examples/src/main/java/ai/djl/examples/training/TrainBertOnCode.java
index bac7e1873..05e7a0153 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainBertOnCode.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainBertOnCode.java
@@ -54,7 +54,8 @@ public final class TrainBertOnCode {
         BertArguments arguments = (BertArguments) new BertArguments().parseArgs(args);
 
         BertCodeDataset dataset =
-                new BertCodeDataset(arguments.getBatchSize(), arguments.getLimit());
+                new BertCodeDataset(
+                        arguments.getBatchSize(), arguments.getLimit(), arguments.getEngine());
         dataset.prepare();
 
         // Create model & trainer
diff --git a/examples/src/main/java/ai/djl/examples/training/util/BertCodeDataset.java b/examples/src/main/java/ai/djl/examples/training/util/BertCodeDataset.java
index d8d538d57..ce3ed687f 100644
--- a/examples/src/main/java/ai/djl/examples/training/util/BertCodeDataset.java
+++ b/examples/src/main/java/ai/djl/examples/training/util/BertCodeDataset.java
@@ -62,10 +62,10 @@ public class BertCodeDataset implements Dataset {
     private long epochLimit;
     private NDManager manager;
 
-    public BertCodeDataset(int batchSize, long epochLimit) {
+    public BertCodeDataset(int batchSize, long epochLimit, String engine) {
         this.batchSize = batchSize;
         this.epochLimit = epochLimit;
-        this.manager = NDManager.newBaseManager();
+        this.manager = NDManager.newBaseManager(engine);
     }
 
     /** {@inheritDoc} */


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
index 35d831777..b99863a4e 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.examples.training;
 
-import ai.djl.testing.TestRequirements;
 import ai.djl.training.TrainingResult;
 import ai.djl.translate.TranslateException;
 
@@ -25,8 +24,6 @@ public class TrainAirfoilWithTabNetTest {
 
     @Test
     public void testTrainAirfoilWithTabNet() throws TranslateException, IOException {
-        TestRequirements.linux();
-
         String[] args = {"-g", "1", "-e", "20", "-b", "32"};
         if (!Boolean.getBoolean("nightly")) {
             args[3] = "2";
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainBertOnGoemotionsTest.java b/examples/src/test/java/ai/djl/examples/training/TrainBertOnGoemotionsTest.java
index 285a22c67..0707e5c8a 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainBertOnGoemotionsTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainBertOnGoemotionsTest.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.examples.training;
 
-import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
 
 import org.testng.annotations.Test;
@@ -23,9 +22,7 @@ public class TrainBertOnGoemotionsTest {
 
     @Test
     public void testTrainBert() throws IOException, TranslateException {
-        TestRequirements.linux();
-
-        String[] args = new String[] {"-g", "1", "-m", "1", "-e", "1"};
+        String[] args = {"-g", "1", "-m", "1", "-e", "1"};
         TrainBertOnGoemotions.runExample(args);
     }
 }
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainBertTest.java b/examples/src/test/java/ai/djl/examples/training/TrainBertTest.java
index cdb101a29..fc21f70e1 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainBertTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainBertTest.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.examples.training;
 
-import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
 
 import org.testng.annotations.Test;
@@ -23,9 +22,7 @@ public class TrainBertTest {
 
     @Test
     public void testTrainBert() throws IOException, TranslateException {
-        TestRequirements.linux();
-
-        String[] args = new String[] {"-g", "1", "-m", "1", "-e", "1"};
+        String[] args = {"-g", "1", "-m", "1", "-e", "1"};
         TrainBertOnCode.runExample(args);
     }
 }
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
index 1e54b24c1..5d8aa27b9 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
@@ -15,7 +15,6 @@ package ai.djl.examples.training;
 import ai.djl.ModelException;
 import ai.djl.examples.inference.cv.ImageClassification;
 import ai.djl.modality.Classifications;
-import ai.djl.testing.TestRequirements;
 import ai.djl.training.TrainingResult;
 import ai.djl.translate.TranslateException;
 
@@ -28,8 +27,6 @@ public class TrainMnistTest {
 
     @Test
     public void testTrainMnist() throws ModelException, TranslateException, IOException {
-        TestRequirements.linux();
-
         double expectedProb;
         if (Boolean.getBoolean("nightly")) {
             String[] args = new String[] {"-g", "1"};
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainMnistWithLSTMTest.java b/examples/src/test/java/ai/djl/examples/training/TrainMnistWithLSTMTest.java
index 7b9b03804..7693e1954 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainMnistWithLSTMTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainMnistWithLSTMTest.java
@@ -12,7 +12,7 @@
  */
 package ai.djl.examples.training;
 
-import ai.djl.testing.TestRequirements;
+import ai.djl.engine.Engine;
 import ai.djl.training.TrainingResult;
 import ai.djl.translate.TranslateException;
 
@@ -25,10 +25,14 @@ public class TrainMnistWithLSTMTest {
 
     @Test
     public void testTrainMnistWithLSTM() throws IOException, TranslateException {
-        TestRequirements.linux();
-
-        // TODO: PyTorch -- cuDNN error: CUDNN_STATUS_VERSION_MISMATCH
-        String[] args = new String[] {"-g", "1", "-e", "1", "-m", "2", "--engine", "MXNet"};
+        String[] args;
+        Engine engine = Engine.getEngine("PyTorch");
+        if (engine.getGpuCount() > 0) {
+            // TODO: PyTorch -- cuDNN error: CUDNN_STATUS_VERSION_MISMATCH
+            args = new String[] {"-g", "1", "-e", "1", "-m", "2", "--engine", "MXNet"};
+        } else {
+            args = new String[] {"-g", "1", "-e", "1", "-m", "2"};
+        }
         TrainingResult result = TrainMnistWithLSTM.runExample(args);
         Assert.assertNotNull(result);
     }
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java b/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java
index a07d95b5a..d726afdd4 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java
@@ -29,7 +29,7 @@ public class TrainSentimentAnalysisTest {
         TestRequirements.nightly();
         TestRequirements.gpu("MXNet");
 
-        String[] args = new String[] {"-e", "1", "-g", "1", "--engine", "MXNet"};
+        String[] args = {"-e", "1", "-g", "1", "--engine", "MXNet"};
         TrainSentimentAnalysis.runExample(args);
     }
 }
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainSeq2SeqTest.java b/examples/src/test/java/ai/djl/examples/training/TrainSeq2SeqTest.java
index 58c46e214..d8b9a13db 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainSeq2SeqTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainSeq2SeqTest.java
@@ -28,7 +28,7 @@ public class TrainSeq2SeqTest {
     public void testTrainSeq2Seq() throws IOException, TranslateException {
         TestRequirements.linux();
 
-        // TODO: PyTorch -- cuDNN error: CUDNN_STATUS_VERSION_MISMATCH
+        // TODO: PyTorch -- PtNDArray.sequenceMask not implemented
         String[] args = new String[] {"-g", "1", "-e", "1", "-m", "2", "--engine", "MXNet"};
         TrainingResult result = TrainSeq2Seq.runExample(args);
         Assert.assertNotNull(result);
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainTicTacToeTest.java b/examples/src/test/java/ai/djl/examples/training/TrainTicTacToeTest.java
index 6d077552d..c04d2660e 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainTicTacToeTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainTicTacToeTest.java
@@ -13,7 +13,6 @@
 package ai.djl.examples.training;
 
 import ai.djl.engine.Engine;
-import ai.djl.testing.TestRequirements;
 import ai.djl.training.TrainingResult;
 
 import org.testng.Assert;
@@ -25,8 +24,6 @@ public class TrainTicTacToeTest {
 
     @Test
     public void testTrainTicTacToe() throws IOException {
-        TestRequirements.linux();
-
         Engine engine = Engine.getEngine("PyTorch");
         if (Boolean.getBoolean("nightly") && engine.getGpuCount() > 0) {
             String[] args = new String[] {"-g", "1", "-e", "6"};
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainTimeSeriesTest.java b/examples/src/test/java/ai/djl/examples/training/TrainTimeSeriesTest.java
index c1b6ffbca..0fa4cd9bd 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainTimeSeriesTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainTimeSeriesTest.java
@@ -28,7 +28,7 @@ public class TrainTimeSeriesTest {
     public void testTrainTimeSeries() throws TranslateException, IOException {
         TestRequirements.linux();
 
-        // TODO: PyTorch -- cuDNN error: CUDNN_STATUS_VERSION_MISMATCH
+        // TODO: PyTorch -- PtNDArray.gammaln not implemented
         String[] args = new String[] {"-g", "1", "-e", "5", "-b", "32", "--engine", "MXNet"};
         TrainingResult result = TrainTimeSeries.runExample(args);
         Assert.assertNotNull(result);