Repository: deepjavalibrary/djl
Pull request #2819 (2023-10-24 22:40:13)

Creates MultiDevice
This creates an abstraction for combining devices into a single device. The main use case for now is in DJL Serving TP_parallel. It will allow us to create a WorkerGroup and a PyPredictor for a set of devices and then track the usage of devices properly. It could also be used later for multi-gpu training or other multi-device cases.

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1365 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`6ef8895`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.20%.
> Report is 899 commits behind head on master.

: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    #2819      +/-   ##
============================================
+ Coverage     72.08%   72.20%   +0.11%     
- Complexity     5126     7140    +2014     
============================================
  Files           473      707     +234     
  Lines         21970    31849    +9879     
  Branches       2351     3305     +954     
============================================
+ Hits          15838    22997    +7159     
- Misses         4925     7280    +2355     
- Partials       1207     1572     +365     
```


| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [...ava/ai/djl/inference/streaming/StreamingBlock.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL2luZmVyZW5jZS9zdHJlYW1pbmcvU3RyZWFtaW5nQmxvY2suamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Dimension.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9EaW1lbnNpb24uamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Unit.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9Vbml0LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/audio/Audio.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2F1ZGlvL0F1ZGlvLmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...ava/ai/djl/modality/cv/output/DetectedObjects.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9EZXRlY3RlZE9iamVjdHMuamF2YQ==) | `96.29% <100.00%> (+1.29%)` | :arrow_up: |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <100.00%> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...i/djl/modality/cv/transform/RandomResizedCrop.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21SZXNpemVkQ3JvcC5qYXZh) | `94.11% <100.00%> (+5.22%)` | :arrow_up: |
| ... and [225 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [378 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2819?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/Device.java b/api/src/main/java/ai/djl/Device.java
index ce9b29ae5..b99e5a648 100644
--- a/api/src/main/java/ai/djl/Device.java
+++ b/api/src/main/java/ai/djl/Device.java
@@ -14,11 +14,16 @@ package ai.djl;
 
 import ai.djl.engine.Engine;
 
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 /**
  * The {@code Device} class provides the specified assignment for CPU/GPU processing on the {@code
@@ -30,7 +35,7 @@ import java.util.regex.Pattern;
  * @see <a href="https://d2l.djl.ai/chapter_deep-learning-computation/use-gpu.html">The D2L chapter
  *     on GPU devices</a>
  */
-public final class Device {
+public class Device {
 
     private static final Map<String, Device> CACHE = new ConcurrentHashMap<>();
 
@@ -39,8 +44,8 @@ public final class Device {
 
     private static final Pattern DEVICE_NAME = Pattern.compile("([a-z]+)([0-9]*)");
 
-    private String deviceType;
-    private int deviceId;
+    protected String deviceType;
+    protected int deviceId;
 
     /**
      * Creates a {@code Device} with basic information.
@@ -101,6 +106,13 @@ public final class Device {
             return engine.defaultDevice();
         }
 
+        if (deviceName.contains("+")) {
+            String[] split = deviceName.split("\\+");
+            List<Device> subDevices =
+                    Arrays.stream(split).map(n -> fromName(n, engine)).collect(Collectors.toList());
+            return new MultiDevice(subDevices);
+        }
+
         Matcher matcher = DEVICE_NAME.matcher(deviceName);
         if (matcher.matches()) {
             String deviceType = matcher.group(1);
@@ -214,4 +226,91 @@ public final class Device {
         String CPU = "cpu";
         String GPU = "gpu";
     }
+
+    /** A combined {@link Device} representing the composition of multiple other devices. */
+    public static class MultiDevice extends Device {
+
+        List<Device> devices;
+
+        /**
+         * Constructs a {@link MultiDevice} with a range of new devices.
+         *
+         * @param deviceType the type of the sub-devices
+         * @param startInclusive the start (inclusive) of the devices range
+         * @param endExclusive the end (exclusive) of the devices range
+         */
+        public MultiDevice(String deviceType, int startInclusive, int endExclusive) {
+            this(
+                    IntStream.range(startInclusive, endExclusive)
+                            .mapToObj(i -> Device.of(deviceType, i))
+                            .collect(Collectors.toList()));
+        }
+
+        /**
+         * Constructs a {@link MultiDevice} from sub devices.
+         *
+         * @param devices the sub devices
+         */
+        public MultiDevice(Device... devices) {
+            this(Arrays.asList(devices));
+        }
+
+        /**
+         * Constructs a {@link MultiDevice} from sub devices.
+         *
+         * @param devices the sub devices
+         */
+        public MultiDevice(List<Device> devices) {
+            super(null, -1);
+            devices.sort(
+                    Comparator.comparing(Device::getDeviceType, String.CASE_INSENSITIVE_ORDER)
+                            .thenComparingInt(Device::getDeviceId));
+            this.deviceType =
+                    String.join(
+                            "+",
+                            (Iterable<String>)
+                                    () ->
+                                            devices.stream()
+                                                    .map(d -> d.getDeviceType() + d.getDeviceId())
+                                                    .iterator());
+            this.devices = devices;
+        }
+
+        /**
+         * Returns the sub devices.
+         *
+         * @return the sub devices
+         */
+        public List<Device> getDevices() {
+            return devices;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            if (!super.equals(o)) {
+                return false;
+            }
+            MultiDevice that = (MultiDevice) o;
+            return Objects.equals(devices, that.devices);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public int hashCode() {
+            return Objects.hash(super.hashCode(), devices);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public String toString() {
+            return deviceType + "()";
+        }
+    }
 }
diff --git a/api/src/main/java/ai/djl/training/ParameterStore.java b/api/src/main/java/ai/djl/training/ParameterStore.java
index 7029282c4..15c83bde8 100644
--- a/api/src/main/java/ai/djl/training/ParameterStore.java
+++ b/api/src/main/java/ai/djl/training/ParameterStore.java
@@ -14,6 +14,7 @@
 package ai.djl.training;
 
 import ai.djl.Device;
+import ai.djl.Device.MultiDevice;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.nn.Parameter;
@@ -64,6 +65,10 @@ public class ParameterStore {
         this.parameterServer = parameterServer;
         deviceMap.clear();
         for (int i = 0; i < devices.length; ++i) {
+            if (devices[i] instanceof MultiDevice) {
+                throw new IllegalArgumentException(
+                        "The parameter store does not support MultiDevices");
+            }
             if (deviceMap.put(devices[i], i) != null) {
                 throw new IllegalArgumentException("Duplicated devices are not allowed.");
             }
diff --git a/api/src/main/java/ai/djl/training/listener/EarlyStoppingListener.java b/api/src/main/java/ai/djl/training/listener/EarlyStoppingListener.java
new file mode 100644
index 000000000..6c013c377
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/listener/EarlyStoppingListener.java
@@ -0,0 +1,281 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.training.listener;
+
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingResult;
+
+import java.time.Duration;
+
+/**
+ * Listener that allows the training to be stopped early if the validation loss is not improving, or
+ * if time has expired. <br>
+ *
+ * <p>Usage: Add this listener to the training config, and add it as the last one.
+ *
+ * <pre>
+ *  new DefaultTrainingConfig(...)
+ *        .addTrainingListeners(EarlyStoppingListener.builder()
+ *                .setEpochPatience(1)
+ *                .setEarlyStopPctImprovement(1)
+ *                .setMaxDuration(Duration.ofMinutes(42))
+ *                .setMinEpochs(1)
+ *                .build()
+ *        );
+ * </pre>
+ *
+ * <p>Then surround the fit with a try catch that catches the {@link
+ * EarlyStoppingListener.EarlyStoppedException}. <br>
+ * Example:
+ *
+ * <pre>
+ * try {
+ *   EasyTrain.fit(trainer, 5, trainDataset, testDataset);
+ * } catch (EarlyStoppingListener.EarlyStoppedException e) {
+ *   // handle early stopping
+ *   log.info("Stopped early at epoch {} because: {}", e.getEpoch(), e.getMessage());
+ * }
+ * </pre>
+ *
+ * <br>
+ * Note: Ensure that Metrics are set on the trainer.
+ */
+public final class EarlyStoppingListener implements TrainingListener {
+    private final double objectiveSuccess;
+
+    private final int minEpochs;
+    private final long maxMillis;
+    private final double earlyStopPctImprovement;
+    private final int epochPatience;
+
+    private long startTimeMills;
+    private double prevLoss;
+    private int numberOfEpochsWithoutImprovements;
+
+    private EarlyStoppingListener(
+            double objectiveSuccess,
+            int minEpochs,
+            long maxMillis,
+            double earlyStopPctImprovement,
+            int earlyStopPatience) {
+        this.objectiveSuccess = objectiveSuccess;
+        this.minEpochs = minEpochs;
+        this.maxMillis = maxMillis;
+        this.earlyStopPctImprovement = earlyStopPctImprovement;
+        this.epochPatience = earlyStopPatience;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void onEpoch(Trainer trainer) {
+        int currentEpoch = trainer.getTrainingResult().getEpoch();
+        // stopping criteria
+        final double loss = getLoss(trainer.getTrainingResult());
+        if (currentEpoch >= minEpochs) {
+            if (loss < objectiveSuccess) {
+                throw new EarlyStoppedException(
+                        currentEpoch,
+                        String.format(
+                                "validation loss %s < objectiveSuccess %s",
+                                loss, objectiveSuccess));
+            }
+            long elapsedMillis = System.currentTimeMillis() - startTimeMills;
+            if (elapsedMillis >= maxMillis) {
+                throw new EarlyStoppedException(
+                        currentEpoch,
+                        String.format("%s ms elapsed >= %s maxMillis", elapsedMillis, maxMillis));
+            }
+            // consider early stopping?
+            if (Double.isFinite(prevLoss)) {
+                double goalImprovement = prevLoss * (100 - earlyStopPctImprovement) / 100.0;
+                boolean improved = loss <= goalImprovement; // false if any NANs
+                if (improved) {
+                    numberOfEpochsWithoutImprovements = 0;
+                } else {
+                    numberOfEpochsWithoutImprovements++;
+                    if (numberOfEpochsWithoutImprovements >= epochPatience) {
+                        throw new EarlyStoppedException(
+                                currentEpoch,
+                                String.format(
+                                        "failed to achieve %s%% improvement %s times in a row",
+                                        earlyStopPctImprovement, epochPatience));
+                    }
+                }
+            }
+        }
+        if (Double.isFinite(loss)) {
+            prevLoss = loss;
+        }
+    }
+
+    private static double getLoss(TrainingResult trainingResult) {
+        Float vLoss = trainingResult.getValidateLoss();
+        if (vLoss != null) {
+            return vLoss;
+        }
+        Float tLoss = trainingResult.getTrainLoss();
+        if (tLoss == null) {
+            return Double.NaN;
+        }
+        return tLoss;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void onTrainingBatch(Trainer trainer, BatchData batchData) {
+        // do nothing
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void onValidationBatch(Trainer trainer, BatchData batchData) {
+        // do nothing
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void onTrainingBegin(Trainer trainer) {
+        this.startTimeMills = System.currentTimeMillis();
+        this.prevLoss = Double.NaN;
+        this.numberOfEpochsWithoutImprovements = 0;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void onTrainingEnd(Trainer trainer) {
+        // do nothing
+    }
+
+    /**
+     * Creates a builder to build a {@link EarlyStoppingListener}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** A builder for a {@link EarlyStoppingListener}. */
+    public static final class Builder {
+        private final double objectiveSuccess;
+        private int minEpochs;
+        private long maxMillis;
+        private double earlyStopPctImprovement;
+        private int epochPatience;
+
+        /** Constructs a {@link Builder} with default values. */
+        public Builder() {
+            this.objectiveSuccess = 0;
+            this.minEpochs = 0;
+            this.maxMillis = Long.MAX_VALUE;
+            this.earlyStopPctImprovement = 0;
+            this.epochPatience = 0;
+        }
+
+        /**
+         * Set the minimum # epochs, defaults to 0.
+         *
+         * @param minEpochs the minimum # epochs
+         * @return this builder
+         */
+        public Builder optMinEpochs(int minEpochs) {
+            this.minEpochs = minEpochs;
+            return this;
+        }
+
+        /**
+         * Set the maximum duration a training run should take, defaults to Long.MAX_VALUE in ms.
+         *
+         * @param duration the maximum duration a training run should take
+         * @return this builder
+         */
+        public Builder optMaxDuration(Duration duration) {
+            this.maxMillis = duration.toMillis();
+            return this;
+        }
+
+        /**
+         * Set the maximum # milliseconds a training run should take, defaults to Long.MAX_VALUE.
+         *
+         * @param maxMillis the maximum # milliseconds a training run should take
+         * @return this builder
+         */
+        public Builder optMaxMillis(int maxMillis) {
+            this.maxMillis = maxMillis;
+            return this;
+        }
+
+        /**
+         * Consider early stopping if not x% improvement, defaults to 0.
+         *
+         * @param earlyStopPctImprovement the percentage improvement to consider early stopping,
+         *     must be between 0 and 100.
+         * @return this builder
+         */
+        public Builder optEarlyStopPctImprovement(double earlyStopPctImprovement) {
+            this.earlyStopPctImprovement = earlyStopPctImprovement;
+            return this;
+        }
+
+        /**
+         * Stop if insufficient improvement for x epochs in a row, defaults to 0.
+         *
+         * @param epochPatience the number of epochs without improvement to consider stopping, must
+         *     be greater than 0.
+         * @return this builder
+         */
+        public Builder optEpochPatience(int epochPatience) {
+            this.epochPatience = epochPatience;
+            return this;
+        }
+
+        /**
+         * Builds a {@link EarlyStoppingListener} with the specified values.
+         *
+         * @return a new {@link EarlyStoppingListener}
+         */
+        public EarlyStoppingListener build() {
+            return new EarlyStoppingListener(
+                    objectiveSuccess, minEpochs, maxMillis, earlyStopPctImprovement, epochPatience);
+        }
+    }
+
+    /**
+     * Thrown when training is stopped early, the message will contain the reason why it is stopped
+     * early.
+     */
+    public static class EarlyStoppedException extends RuntimeException {
+        private static final long serialVersionUID = 1L;
+        private final int stopEpoch;
+
+        /**
+         * Constructs an {@link EarlyStoppedException} with the specified message and epoch.
+         *
+         * @param stopEpoch the epoch at which training was stopped early
+         * @param message the message/reason why training was stopped early
+         */
+        public EarlyStoppedException(int stopEpoch, String message) {
+            super(message);
+            this.stopEpoch = stopEpoch;
+        }
+
+        /**
+         * Gets the epoch at which training was stopped early.
+         *
+         * @return the epoch at which training was stopped early.
+         */
+        public int getStopEpoch() {
+            return stopEpoch;
+        }
+    }
+}


Test patch:
diff --git a/api/src/test/java/ai/djl/DeviceTest.java b/api/src/test/java/ai/djl/DeviceTest.java
index 92a0474c6..635728108 100644
--- a/api/src/test/java/ai/djl/DeviceTest.java
+++ b/api/src/test/java/ai/djl/DeviceTest.java
@@ -13,6 +13,7 @@
 
 package ai.djl;
 
+import ai.djl.Device.MultiDevice;
 import ai.djl.engine.Engine;
 
 import org.testng.Assert;
@@ -37,6 +38,8 @@ public class DeviceTest {
 
         System.setProperty("test_key", "test");
         Engine.debugEnvironment();
+
+        Assert.assertEquals(2, new MultiDevice(Device.gpu(1), Device.gpu(2)).getDevices().size());
     }
 
     @Test
@@ -54,5 +57,9 @@ public class DeviceTest {
         Device defaultDevice = Engine.getInstance().defaultDevice();
         Assert.assertEquals(Device.fromName(""), defaultDevice);
         Assert.assertEquals(Device.fromName(null), defaultDevice);
+
+        Assert.assertEquals(
+                Device.fromName("gpu1+gpu2"), new MultiDevice(Device.gpu(2), Device.gpu(1)));
+        Assert.assertEquals(Device.fromName("gpu1+gpu2"), new MultiDevice("gpu", 1, 3));
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
new file mode 100644
index 000000000..c3dd34d03
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
@@ -0,0 +1,189 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.integration.tests.training.listener;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.cv.classification.Mnist;
+import ai.djl.basicmodelzoo.basic.Mlp;
+import ai.djl.integration.util.TestUtils;
+import ai.djl.metric.Metrics;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Activation;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingResult;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.listener.EarlyStoppingListener;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.Loss;
+import ai.djl.training.optimizer.Optimizer;
+import ai.djl.training.tracker.Tracker;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.AfterTest;
+import org.testng.annotations.BeforeTest;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.time.Duration;
+
+public class EarlyStoppingListenerTest {
+
+    private final Optimizer sgd =
+            Optimizer.sgd().setLearningRateTracker(Tracker.fixed(0.1f)).build();
+
+    private NDManager manager;
+    private Mnist testMnistDataset;
+    private Mnist trainMnistDataset;
+
+    @BeforeTest
+    public void setUp() throws IOException, TranslateException {
+        manager = NDManager.newBaseManager(TestUtils.getEngine());
+        testMnistDataset =
+                Mnist.builder()
+                        .optUsage(Dataset.Usage.TEST)
+                        .optManager(manager)
+                        .optLimit(8)
+                        .setSampling(8, false)
+                        .build();
+        testMnistDataset.prepare();
+
+        trainMnistDataset =
+                Mnist.builder()
+                        .optUsage(Dataset.Usage.TRAIN)
+                        .optManager(manager)
+                        .optLimit(16)
+                        .setSampling(8, false)
+                        .build();
+        trainMnistDataset.prepare();
+    }
+
+    @AfterTest
+    public void closeResources() {
+        manager.close();
+    }
+
+    @Test
+    public void testEarlyStoppingStopsOnEpoch2() throws Exception {
+        Mlp mlpModel = new Mlp(784, 1, new int[] {256}, Activation::relu);
+
+        try (Model model = Model.newInstance("lin-reg", TestUtils.getEngine())) {
+            model.setBlock(mlpModel);
+
+            DefaultTrainingConfig config =
+                    new DefaultTrainingConfig(Loss.l2Loss())
+                            .optOptimizer(sgd)
+                            .addTrainingListeners(TrainingListener.Defaults.logging())
+                            .addTrainingListeners(
+                                    EarlyStoppingListener.builder()
+                                            .optEpochPatience(1)
+                                            .optEarlyStopPctImprovement(99)
+                                            .optMaxDuration(Duration.ofMinutes(1))
+                                            .optMinEpochs(1)
+                                            .build());
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.initialize(new Shape(1, 784));
+                Metrics metrics = new Metrics();
+                trainer.setMetrics(metrics);
+
+                try {
+                    // Set epoch to 5 as we expect the early stopping to stop after the second epoch
+                    EasyTrain.fit(trainer, 5, trainMnistDataset, testMnistDataset);
+                } catch (EarlyStoppingListener.EarlyStoppedException e) {
+                    Assert.assertEquals(
+                            e.getMessage(), "failed to achieve 99.0% improvement 1 times in a row");
+                    Assert.assertEquals(e.getStopEpoch(), 2);
+                }
+
+                TrainingResult trainingResult = trainer.getTrainingResult();
+                Assert.assertEquals(trainingResult.getEpoch(), 2);
+            }
+        }
+    }
+
+    @Test
+    public void testEarlyStoppingStopsOnEpoch3AsMinEpochsIs3() throws Exception {
+        Mlp mlpModel = new Mlp(784, 1, new int[] {256}, Activation::relu);
+
+        try (Model model = Model.newInstance("lin-reg", TestUtils.getEngine())) {
+            model.setBlock(mlpModel);
+
+            DefaultTrainingConfig config =
+                    new DefaultTrainingConfig(Loss.l2Loss())
+                            .optOptimizer(sgd)
+                            .addTrainingListeners(TrainingListener.Defaults.logging())
+                            .addTrainingListeners(
+                                    EarlyStoppingListener.builder()
+                                            .optEpochPatience(1)
+                                            .optEarlyStopPctImprovement(50)
+                                            .optMaxMillis(60_000)
+                                            .optMinEpochs(3)
+                                            .build());
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.initialize(new Shape(1, 784));
+                Metrics metrics = new Metrics();
+                trainer.setMetrics(metrics);
+
+                try {
+                    // Set epoch to 5 as we expect the early stopping to stop after the second epoch
+                    EasyTrain.fit(trainer, 5, trainMnistDataset, testMnistDataset);
+                } catch (EarlyStoppingListener.EarlyStoppedException e) {
+                    Assert.assertEquals(
+                            e.getMessage(), "failed to achieve 50.0% improvement 1 times in a row");
+                    Assert.assertEquals(e.getStopEpoch(), 3);
+                }
+
+                TrainingResult trainingResult = trainer.getTrainingResult();
+                Assert.assertEquals(trainingResult.getEpoch(), 3);
+            }
+        }
+    }
+
+    @Test
+    public void testEarlyStoppingStopsOnEpoch1AsMaxDurationIs1ms() throws Exception {
+        Mlp mlpModel = new Mlp(784, 1, new int[] {256}, Activation::relu);
+
+        try (Model model = Model.newInstance("lin-reg", TestUtils.getEngine())) {
+            model.setBlock(mlpModel);
+
+            DefaultTrainingConfig config =
+                    new DefaultTrainingConfig(Loss.l2Loss())
+                            .optOptimizer(sgd)
+                            .addTrainingListeners(TrainingListener.Defaults.logging())
+                            .addTrainingListeners(
+                                    EarlyStoppingListener.builder().optMaxMillis(1).build());
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.initialize(new Shape(1, 784));
+                Metrics metrics = new Metrics();
+                trainer.setMetrics(metrics);
+
+                try {
+                    // Set epoch to 5 as we expect the early stopping to stop after the second epoch
+                    EasyTrain.fit(trainer, 5, trainMnistDataset, testMnistDataset);
+                } catch (EarlyStoppingListener.EarlyStoppedException e) {
+                    Assert.assertTrue(e.getMessage().contains("ms elapsed >= 1 maxMillis"));
+                    Assert.assertEquals(e.getStopEpoch(), 1);
+                }
+
+                TrainingResult trainingResult = trainer.getTrainingResult();
+                Assert.assertEquals(trainingResult.getEpoch(), 1);
+            }
+        }
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/listener/package-info.java b/integration/src/main/java/ai/djl/integration/tests/training/listener/package-info.java
new file mode 100644
index 000000000..88680e5fe
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/training/listener/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains tests using the listeners {@link ai.djl.training}. */
+package ai.djl.integration.tests.training.listener;