Merge pull request #25063 from vespa-engine/mpolden/add-gpu-fields

Add GPU fields to node repository client

2 files changed, 25 insertions, 12 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java
index a4a047b524c..76ee3a52d64 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java
@@ -195,7 +195,13 @@ public class RealNodeRepository implements NodeRepository {
                 nodeResources.bandwidthGbps,
                 diskSpeedFromString(nodeResources.diskSpeed),
                 storageTypeFromString(nodeResources.storageType),
-                architectureFromString(nodeResources.architecture));
+                architectureFromString(nodeResources.architecture),
+                gpuResourcesFrom(nodeResources));
+    }
+
+    private static NodeResources.GpuResources gpuResourcesFrom(NodeRepositoryNode.NodeResources nodeResources) {
+        if (nodeResources.gpuCount == null || nodeResources.gpuMemoryGb == null) return NodeResources.GpuResources.zero();
+        return new NodeResources.GpuResources(nodeResources.gpuCount, nodeResources.gpuMemoryGb);
     }
 
     private static NodeResources.DiskSpeed diskSpeedFromString(String diskSpeed) {
@@ -233,7 +239,6 @@ public class RealNodeRepository implements NodeRepository {
             case fast -> "fast";
             case slow -> "slow";
             case any -> "any";
-            default -> throw new IllegalArgumentException("Unknown disk speed '" + diskSpeed.name() + "'");
         };
     }
 
@@ -242,7 +247,6 @@ public class RealNodeRepository implements NodeRepository {
             case remote -> "remote";
             case local -> "local";
             case any -> "any";
-            default -> throw new IllegalArgumentException("Unknown storage type '" + storageType.name() + "'");
         };
     }
 
@@ -251,7 +255,6 @@ public class RealNodeRepository implements NodeRepository {
             case arm64 -> "arm64";
             case x86_64 -> "x86_64";
             case any -> "any";
-            default -> throw new IllegalArgumentException("Unknown architecture '" + architecture.name() + "'");
         };
     }
 
@@ -274,6 +277,10 @@ public class RealNodeRepository implements NodeRepository {
             node.resources.diskSpeed = toString(resources.diskSpeed());
             node.resources.storageType = toString(resources.storageType());
             node.resources.architecture = toString(resources.architecture());
+            if (!resources.gpuResources().isZero()) {
+                node.resources.gpuCount = resources.gpuResources().count();
+                node.resources.gpuMemoryGb = resources.gpuResources().memoryGb();
+            }
         });
         node.type = addNode.nodeType.name();
         node.ipAddresses = addNode.ipAddresses;
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java
index 956795e09e9..434acca767a 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java
@@ -197,18 +197,24 @@ public class NodeRepositoryNode {
         public String storageType;
         @JsonProperty
         public String architecture;
+        @JsonProperty
+        public Integer gpuCount;
+        @JsonProperty
+        public Double gpuMemoryGb;
 
         @Override
         public String toString() {
             return "NodeResources{" +
-                    "vcpu=" + vcpu +
-                    ", memoryGb=" + memoryGb +
-                    ", diskGb=" + diskGb +
-                    ", bandwidthGbps=" + bandwidthGbps +
-                    ", diskSpeed='" + diskSpeed + '\'' +
-                    ", storageType='" + storageType + '\'' +
-                    ", architecture='" + architecture + '\'' +
-                    '}';
+                   "vcpu=" + vcpu +
+                   ", memoryGb=" + memoryGb +
+                   ", diskGb=" + diskGb +
+                   ", bandwidthGbps=" + bandwidthGbps +
+                   ", diskSpeed='" + diskSpeed + '\'' +
+                   ", storageType='" + storageType + '\'' +
+                   ", architecture='" + architecture + '\'' +
+                   ", gpuCount=" + gpuCount +
+                   ", gpuMemoryGb=" + gpuMemoryGb +
+                   '}';
         }
     }