Serialize GPU resources

3 files changed, 61 insertions, 42 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeResourcesSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeResourcesSerializer.java
index f7be2b510c7..8843ce69994 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeResourcesSerializer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeResourcesSerializer.java
@@ -19,6 +19,9 @@ public class NodeResourcesSerializer {
     private static final String diskSpeedKey = "diskSpeed";
     private static final String storageTypeKey = "storageType";
     private static final String architectureKey = "architecture";
+    private static final String gpuKey = "gpu";
+    private static final String gpuCountKey = "gpuCount";
+    private static final String gpuMemoryKey = "gpuMemory";
 
     static void toSlime(NodeResources resources, Cursor resourcesObject) {
         if (resources.isUnspecified()) return;
@@ -29,6 +32,11 @@ public class NodeResourcesSerializer {
         resourcesObject.setString(diskSpeedKey, diskSpeedToString(resources.diskSpeed()));
         resourcesObject.setString(storageTypeKey, storageTypeToString(resources.storageType()));
         resourcesObject.setString(architectureKey, architectureToString(resources.architecture()));
+        if (!resources.gpuResources().isDefault()) {
+            Cursor gpuObject = resourcesObject.setObject(gpuKey);
+            gpuObject.setLong(gpuCountKey, resources.gpuResources().count());
+            gpuObject.setDouble(gpuMemoryKey, resources.gpuResources().memoryGb());
+        }
     }
 
     static NodeResources resourcesFromSlime(Inspector resources) {
@@ -39,7 +47,8 @@ public class NodeResourcesSerializer {
                                  resources.field(bandwidthKey).asDouble(),
                                  diskSpeedFromSlime(resources.field(diskSpeedKey)),
                                  storageTypeFromSlime(resources.field(storageTypeKey)),
-                                 architectureFromSlime(resources.field(architectureKey)));
+                                 architectureFromSlime(resources.field(architectureKey)),
+                                 gpuResourcesFromSlime(resources.field(gpuKey)));
     }
 
     static Optional<NodeResources> optionalResourcesFromSlime(Inspector resources) {
@@ -47,58 +56,61 @@ public class NodeResourcesSerializer {
     }
 
     private static NodeResources.DiskSpeed diskSpeedFromSlime(Inspector diskSpeed) {
-        switch (diskSpeed.asString()) {
-            case "fast" : return NodeResources.DiskSpeed.fast;
-            case "slow" : return NodeResources.DiskSpeed.slow;
-            case "any" : return NodeResources.DiskSpeed.any;
-            default: throw new IllegalStateException("Illegal disk-speed value '" + diskSpeed.asString() + "'");
-        }
+        return switch (diskSpeed.asString()) {
+            case "fast" -> NodeResources.DiskSpeed.fast;
+            case "slow" -> NodeResources.DiskSpeed.slow;
+            case "any" -> NodeResources.DiskSpeed.any;
+            default -> throw new IllegalStateException("Illegal disk-speed value '" + diskSpeed.asString() + "'");
+        };
     }
 
     private static String diskSpeedToString(NodeResources.DiskSpeed diskSpeed) {
-        switch (diskSpeed) {
-            case fast : return "fast";
-            case slow : return "slow";
-            case any : return "any";
-            default: throw new IllegalStateException("Illegal disk-speed value '" + diskSpeed + "'");
-        }
+        return switch (diskSpeed) {
+            case fast -> "fast";
+            case slow -> "slow";
+            case any -> "any";
+        };
     }
 
     private static NodeResources.StorageType storageTypeFromSlime(Inspector storageType) {
-        switch (storageType.asString()) {
-            case "remote" : return NodeResources.StorageType.remote;
-            case "local" : return NodeResources.StorageType.local;
-            case "any" : return NodeResources.StorageType.any;
-            default: throw new IllegalStateException("Illegal storage-type value '" + storageType.asString() + "'");
-        }
+        return switch (storageType.asString()) {
+            case "remote" -> NodeResources.StorageType.remote;
+            case "local" -> NodeResources.StorageType.local;
+            case "any" -> NodeResources.StorageType.any;
+            default -> throw new IllegalStateException("Illegal storage-type value '" + storageType.asString() + "'");
+        };
     }
 
     private static String storageTypeToString(NodeResources.StorageType storageType) {
-        switch (storageType) {
-            case remote : return "remote";
-            case local : return "local";
-            case any : return "any";
-            default: throw new IllegalStateException("Illegal storage-type value '" + storageType + "'");
-        }
+        return switch (storageType) {
+            case remote -> "remote";
+            case local -> "local";
+            case any -> "any";
+        };
     }
 
     private static NodeResources.Architecture architectureFromSlime(Inspector architecture) {
         if ( ! architecture.valid()) return NodeResources.Architecture.getDefault(); // TODO: Remove this line after March 2022
-        switch (architecture.asString()) {
-            case "arm64" : return NodeResources.Architecture.arm64;
-            case "x86_64" : return NodeResources.Architecture.x86_64;
-            case "any" : return NodeResources.Architecture.any;
-            default: throw new IllegalStateException("Illegal architecture value '" + architecture.asString() + "'");
-        }
+        return switch (architecture.asString()) {
+            case "arm64" -> NodeResources.Architecture.arm64;
+            case "x86_64" -> NodeResources.Architecture.x86_64;
+            case "any" -> NodeResources.Architecture.any;
+            default -> throw new IllegalStateException("Illegal architecture value '" + architecture.asString() + "'");
+        };
     }
 
     private static String architectureToString(NodeResources.Architecture architecture) {
-        switch (architecture) {
-            case arm64 : return "arm64";
-            case x86_64 : return "x86_64";
-            case any : return "any";
-            default: throw new IllegalStateException("Illegal architecture value '" + architecture + "'");
-        }
+        return switch (architecture) {
+            case arm64 -> "arm64";
+            case x86_64 -> "x86_64";
+            case any -> "any";
+        };
+    }
+
+    private static NodeResources.GpuResources gpuResourcesFromSlime(Inspector gpu) {
+        if (!gpu.valid()) return NodeResources.GpuResources.getDefault();
+        return new NodeResources.GpuResources((int) gpu.field(gpuCountKey).asLong(),
+                                              gpu.field(gpuMemoryKey).asDouble());
     }
 
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/FlavorConfigBuilder.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/FlavorConfigBuilder.java
index 4567d596c35..2bc5a0719d9 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/FlavorConfigBuilder.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/FlavorConfigBuilder.java
@@ -27,7 +27,7 @@ public class FlavorConfigBuilder {
                                                   double disk,
                                                   double bandwidth,
                                                   Flavor.Type type) {
-        return addFlavor(flavorName, cpu, mem, disk, bandwidth, true, true, type, Architecture.x86_64);
+        return addFlavor(flavorName, cpu, mem, disk, bandwidth, true, true, type, Architecture.x86_64, 0, 0);
     }
 
     public FlavorsConfig.Flavor.Builder addFlavor(String flavorName,
@@ -37,7 +37,7 @@ public class FlavorConfigBuilder {
                                                   double bandwidth,
                                                   Flavor.Type type,
                                                   Architecture architecture) {
-        return addFlavor(flavorName, cpu, mem, disk, bandwidth, true, true, type, architecture);
+        return addFlavor(flavorName, cpu, mem, disk, bandwidth, true, true, type, architecture, 0, 0);
     }
 
     public FlavorsConfig.Flavor.Builder addFlavor(String flavorName,
@@ -48,7 +48,9 @@ public class FlavorConfigBuilder {
                                                   boolean fastDisk,
                                                   boolean remoteStorage,
                                                   Flavor.Type type,
-                                                  Architecture architecture) {
+                                                  Architecture architecture,
+                                                  int gpuCount,
+                                                  double gpuMemory) {
         FlavorsConfig.Flavor.Builder flavor = new FlavorsConfig.Flavor.Builder();
         flavor.name(flavorName);
         flavor.minDiskAvailableGb(disk);
@@ -59,6 +61,8 @@ public class FlavorConfigBuilder {
         flavor.fastDisk(fastDisk);
         flavor.remoteStorage(remoteStorage);
         flavor.architecture(architecture.name());
+        flavor.gpuCount(gpuCount);
+        flavor.gpuMemoryGb(gpuMemory);
         builder.flavor(flavor);
         return flavor;
     }
@@ -84,6 +88,8 @@ public class FlavorConfigBuilder {
                 flavorConfigBuilder.addFlavor(flavorName, 4.,  80., 100, 10, Flavor.Type.BARE_METAL);
             else if (flavorName.equals("arm64"))
                 flavorConfigBuilder.addFlavor(flavorName,2.,  30., 20., 3, Flavor.Type.BARE_METAL, Architecture.arm64);
+            else if (flavorName.equals("gpu"))
+                flavorConfigBuilder.addFlavor(flavorName,4,  16, 125, 10, true, false, Flavor.Type.BARE_METAL, Architecture.x86_64, 1, 16);
             else
                 flavorConfigBuilder.addFlavor(flavorName, 1.,  30., 20., 3, Flavor.Type.BARE_METAL);
         }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializerTest.java
index 2536d6b97ca..e8155234989 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializerTest.java
@@ -59,7 +59,7 @@ import static org.junit.Assert.assertTrue;
  */
 public class NodeSerializerTest {
 
-    private final NodeFlavors nodeFlavors = FlavorConfigBuilder.createDummies("default", "large", "ugccloud-container", "arm64");
+    private final NodeFlavors nodeFlavors = FlavorConfigBuilder.createDummies("default", "large", "ugccloud-container", "arm64", "gpu");
     private final NodeSerializer nodeSerializer = new NodeSerializer(nodeFlavors, 1000);
     private final ManualClock clock = new ManualClock();
 
@@ -79,7 +79,8 @@ public class NodeSerializerTest {
     public void reserved_node_serialization() {
         Node node = createNode();
         NodeResources requestedResources = new NodeResources(1.2, 3.4, 5.6, 7.8,
-                                                             DiskSpeed.any, StorageType.any, Architecture.arm64);
+                                                             DiskSpeed.any, StorageType.any, Architecture.arm64,
+                                                             new NodeResources.GpuResources(1, 16));
 
         assertEquals(0, node.history().events().size());
         node = node.allocate(ApplicationId.from(TenantName.from("myTenant"),