Device Name parsing

api/src/main/java/ai/djl/Device.java

@@ -15,6 +15,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /**
  * The {@code Device} class provides the specified assignment for CPU/GPU processing on the {@code
@@ -33,6 +35,8 @@ public final class Device {
     private static final Device CPU = new Device(Type.CPU, -1);
     private static final Device GPU = Device.of(Type.GPU, 0);
 
+    private static final Pattern DEVICE_NAME = Pattern.compile("([a-z]+)([0-9]*)");
+
     private String deviceType;
     private int deviceId;
 
@@ -63,6 +67,46 @@ public static Device of(String deviceType, int deviceId) {
         return CACHE.computeIfAbsent(key, k -> new Device(deviceType, deviceId));
     }
 
+    /**
+     * Parses a deviceName string into a device.
+     *
+     * <p>The main format of a device name string is "cpu", "gpu0", or "nc1". This is simply
+     * deviceType concatenated with the deviceId. If no deviceId is used, -1 will be assumed.
+     *
+     * <p>There are also several simplified formats. "-1", "", and null deviceNames correspond to
+     * cpu. Non-negative integer deviceNames such as "0", "1", or "2" correspond to gpus with those
+     * deviceIds.
+     *
+     * @param deviceName deviceName string
+     * @return the device
+     */
+    public static Device fromName(String deviceName) {
+        if (deviceName == null || deviceName.isEmpty()) {
+            return Device.cpu();
+        }
+
+        Matcher matcher = DEVICE_NAME.matcher(deviceName);
+        if (matcher.matches()) {
+            String deviceType = matcher.group(1);
+            int deviceId = -1;
+            if (!matcher.group(2).isEmpty()) {
+                deviceId = Integer.parseInt(matcher.group(2));
+            }
+            return Device.of(deviceType, deviceId);
+        }
+
+        try {
+            int deviceId = Integer.parseInt(deviceName);
+
+            if (deviceId < 0) {
+                return Device.cpu();
+            }
+            return Device.gpu(deviceId);
+        } catch (NumberFormatException ignored) {
+        }
+        throw new IllegalArgumentException("Failed to parse device name: " + deviceName);
+    }
+
     /**
      * Returns the device type of the Device.
      *

api/src/test/java/ai/djl/DeviceTest.java

@@ -36,4 +36,19 @@ public void testDevice() {
 
         Engine.debugEnvironment();
     }
+
+    @Test
+    public void testDeviceName() {
+        Assert.assertEquals(Device.fromName("cpu"), Device.cpu());
+        Assert.assertEquals(Device.fromName(""), Device.cpu());
+        Assert.assertEquals(Device.fromName("-1"), Device.cpu());
+        Assert.assertEquals(Device.fromName(null), Device.cpu());
+
+        Assert.assertEquals(Device.fromName("gpu0"), Device.gpu());
+        Assert.assertEquals(Device.fromName("0"), Device.gpu());
+        Assert.assertEquals(Device.fromName("1"), Device.gpu(1));
+
+        Assert.assertEquals(Device.fromName("nc1"), Device.of("nc", 1));
+        Assert.assertEquals(Device.fromName("a999"), Device.of("a", 999));
+    }
 }