Repository: deepjavalibrary/djl
Pull request #3256 (2024-06-16 05:11:40)

[examples] Enabled training unit tests on macOS M1
## 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/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);