Implement billing for GPU resources

4 files changed, 41 insertions, 20 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanRegistryMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanRegistryMock.java
index ecce9b7ccfa..2c496b1e3e5 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanRegistryMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanRegistryMock.java
@@ -12,9 +12,9 @@ import java.util.stream.Stream;
 
 public class PlanRegistryMock implements PlanRegistry {
 
-    public static final Plan freeTrial = new MockPlan("trial", false, false, 0, 0, 0, 200, "Free Trial - for testing purposes");
-    public static final Plan paidPlan  = new MockPlan("paid", true, true, "0.09", "0.009", "0.0003", 500, "Paid Plan - for testing purposes");
-    public static final Plan nonePlan  = new MockPlan("none", false, false, 0, 0, 0, 0, "None Plan - for testing purposes");
+    public static final Plan freeTrial = new MockPlan("trial", false, false, 0, 0, 0, 0, 200, "Free Trial - for testing purposes");
+    public static final Plan paidPlan  = new MockPlan("paid", true, true, "0.09", "0.009", "0.0003", "0.075", 500, "Paid Plan - for testing purposes");
+    public static final Plan nonePlan  = new MockPlan("none", false, false, 0, 0, 0, 0, 0, "None Plan - for testing purposes");
 
     @Override
     public Plan defaultPlan() {
@@ -41,12 +41,12 @@ public class PlanRegistryMock implements PlanRegistry {
         private final boolean billed;
         private final boolean supported;
 
-        public MockPlan(String planId, boolean billed, boolean supported, double cpuPrice, double memPrice, double dgbPrice, int quota, String description) {
-            this(PlanId.from(planId), billed, supported, new MockCostCalculator(cpuPrice, memPrice, dgbPrice), () -> Quota.unlimited().withBudget(quota), description);
+        public MockPlan(String planId, boolean billed, boolean supported, double cpuPrice, double memPrice, double dgbPrice, double gpuPrice, int quota, String description) {
+            this(PlanId.from(planId), billed, supported, new MockCostCalculator(cpuPrice, memPrice, dgbPrice, gpuPrice), () -> Quota.unlimited().withBudget(quota), description);
         }
 
-        public MockPlan(String planId, boolean billed, boolean supported, String cpuPrice, String memPrice, String dgbPrice, int quota, String description) {
-            this(PlanId.from(planId), billed, supported, new MockCostCalculator(cpuPrice, memPrice, dgbPrice), () -> Quota.unlimited().withBudget(quota), description);
+        public MockPlan(String planId, boolean billed, boolean supported, String cpuPrice, String memPrice, String dgbPrice, String gpuPrice, int quota, String description) {
+            this(PlanId.from(planId), billed, supported, new MockCostCalculator(cpuPrice, memPrice, dgbPrice, gpuPrice), () -> Quota.unlimited().withBudget(quota), description);
         }
 
         public MockPlan(PlanId planId, boolean billed, boolean supported, MockCostCalculator calculator, QuotaCalculator quota, String description) {
@@ -94,19 +94,21 @@ public class PlanRegistryMock implements PlanRegistry {
         private final BigDecimal cpuHourCost;
         private final BigDecimal memHourCost;
         private final BigDecimal dgbHourCost;
+        private final BigDecimal gpuHourCost;
 
-        public MockCostCalculator(String cpuPrice, String memPrice, String dgbPrice) {
-            this(new BigDecimal(cpuPrice), new BigDecimal(memPrice), new BigDecimal(dgbPrice));
+        public MockCostCalculator(String cpuPrice, String memPrice, String dgbPrice, String gpuPrice) {
+            this(new BigDecimal(cpuPrice), new BigDecimal(memPrice), new BigDecimal(dgbPrice), new BigDecimal(gpuPrice));
         }
 
-        public MockCostCalculator(double cpuPrice, double memPrice, double dgbPrice) {
-            this(BigDecimal.valueOf(cpuPrice), BigDecimal.valueOf(memPrice), BigDecimal.valueOf(dgbPrice));
+        public MockCostCalculator(double cpuPrice, double memPrice, double dgbPrice, double gpuPrice) {
+            this(BigDecimal.valueOf(cpuPrice), BigDecimal.valueOf(memPrice), BigDecimal.valueOf(dgbPrice), BigDecimal.valueOf(gpuPrice));
         }
 
-        public MockCostCalculator(BigDecimal cpuPrice, BigDecimal memPrice, BigDecimal dgbPrice) {
+        public MockCostCalculator(BigDecimal cpuPrice, BigDecimal memPrice, BigDecimal dgbPrice, BigDecimal gpuPrice) {
             this.cpuHourCost = cpuPrice;
             this.memHourCost = memPrice;
             this.dgbHourCost = dgbPrice;
+            this.gpuHourCost = gpuPrice;
         }
 
         @Override
@@ -114,6 +116,7 @@ public class PlanRegistryMock implements PlanRegistry {
             var cpuCost = usage.getCpuMillis().multiply(cpuHourCost).divide(millisPerHour, RoundingMode.HALF_UP).setScale(2, RoundingMode.HALF_UP);
             var memCost = usage.getMemoryMillis().multiply(memHourCost).divide(millisPerHour, RoundingMode.HALF_UP).setScale(2, RoundingMode.HALF_UP);
             var dgbCost = usage.getDiskMillis().multiply(dgbHourCost).divide(millisPerHour, RoundingMode.HALF_UP).setScale(2, RoundingMode.HALF_UP);
+            var gpuCost = usage.getGpuMillis().multiply(gpuHourCost).divide(millisPerHour, RoundingMode.HALF_UP).setScale(2, RoundingMode.HALF_UP);
 
             return new CostInfo(
                     usage.getApplicationId(),
@@ -121,9 +124,11 @@ public class PlanRegistryMock implements PlanRegistry {
                     usage.getCpuMillis().divide(millisPerHour, RoundingMode.HALF_UP),
                     usage.getMemoryMillis().divide(millisPerHour, RoundingMode.HALF_UP),
                     usage.getDiskMillis().divide(millisPerHour, RoundingMode.HALF_UP),
+                    usage.getGpuMillis().divide(millisPerHour, RoundingMode.HALF_UP),
                     cpuCost,
                     memCost,
                     dgbCost,
+                    gpuCost,
                     usage.getArchitecture(),
                     usage.getMajorVersion()
             );
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/CostInfo.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/CostInfo.java
index 571a7b5a58e..a05d8c9d52b 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/CostInfo.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/CostInfo.java
@@ -17,24 +17,28 @@ public class CostInfo {
     private final BigDecimal cpuHours;
     private final BigDecimal memoryHours;
     private final BigDecimal diskHours;
+    private final BigDecimal gpuHours;
     private final BigDecimal cpuCost;
     private final BigDecimal memoryCost;
     private final BigDecimal diskCost;
+    private final BigDecimal gpuCost;
     private final NodeResources.Architecture architecture;
     private final int majorVersion;
 
 
     public CostInfo(ApplicationId applicationId, ZoneId zoneId,
-                    BigDecimal cpuHours, BigDecimal memoryHours, BigDecimal diskHours,
-                    BigDecimal cpuCost, BigDecimal memoryCost, BigDecimal diskCost, NodeResources.Architecture architecture, int majorVersion) {
+                    BigDecimal cpuHours, BigDecimal memoryHours, BigDecimal diskHours, BigDecimal gpuHours,
+                    BigDecimal cpuCost, BigDecimal memoryCost, BigDecimal diskCost, BigDecimal gpuCost, NodeResources.Architecture architecture, int majorVersion) {
         this.applicationId = applicationId;
         this.zoneId = zoneId;
         this.cpuHours = cpuHours;
         this.memoryHours = memoryHours;
         this.diskHours = diskHours;
+        this.gpuHours = gpuHours;
         this.cpuCost = cpuCost;
         this.memoryCost = memoryCost;
         this.diskCost = diskCost;
+        this.gpuCost = gpuCost;
         this.architecture = architecture;
         this.majorVersion = majorVersion;
     }
@@ -59,6 +63,10 @@ public class CostInfo {
         return diskHours;
     }
 
+    public BigDecimal getGpuHours() {
+        return gpuHours;
+    }
+
     public BigDecimal getCpuCost() {
         return cpuCost;
     }
@@ -71,8 +79,12 @@ public class CostInfo {
         return diskCost;
     }
 
+    public BigDecimal getGpuCost() {
+        return gpuCost;
+    }
+
     public BigDecimal getTotalCost() {
-        return cpuCost.add(memoryCost).add(diskCost);
+        return cpuCost.add(memoryCost).add(diskCost).add(gpuCost);
     }
 
     public NodeResources.Architecture getArchitecture() {
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java
index c1d210e990d..c0d006dcb53 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java
@@ -82,8 +82,8 @@ public class ResourceDatabaseClientMock implements ResourceDatabaseClient {
                             a.getVersion().getMajor(),
                             BigDecimal.valueOf(a.resources().vcpu()).multiply(d),
                             BigDecimal.valueOf(a.resources().memoryGb()).multiply(d),
-                            BigDecimal.valueOf(a.resources().diskGb()).multiply(d)
-                    );
+                            BigDecimal.valueOf(a.resources().diskGb()).multiply(d),
+                            BigDecimal.valueOf(a.resources().gpuResources().count() * a.resources().gpuResources().memoryGb()).multiply(d));
                 })
                 .toList();
     }
@@ -101,8 +101,8 @@ public class ResourceDatabaseClientMock implements ResourceDatabaseClient {
                 a.getMajorVersion(),
                 a.getCpuMillis().add(b.getCpuMillis()),
                 a.getMemoryMillis().add(b.getMemoryMillis()),
-                a.getDiskMillis().add(b.getDiskMillis())
-        );
+                a.getDiskMillis().add(b.getDiskMillis()),
+                a.getGpuMillis().add(b.getGpuMillis()));
     }
 
     @Override
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceUsage.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceUsage.java
index bd1252d0879..0834cf52ec1 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceUsage.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceUsage.java
@@ -17,16 +17,18 @@ public class ResourceUsage {
     private final BigDecimal cpuMillis;
     private final BigDecimal memoryMillis;
     private final BigDecimal diskMillis;
+    private final BigDecimal gpuMillis;
     private final NodeResources.Architecture architecture;
     private final int majorVersion;
 
     public ResourceUsage(ApplicationId applicationId, ZoneId zoneId, Plan plan, NodeResources.Architecture architecture,
-                         int majorVersion, BigDecimal cpuMillis, BigDecimal memoryMillis, BigDecimal diskMillis) {
+                         int majorVersion, BigDecimal cpuMillis, BigDecimal memoryMillis, BigDecimal diskMillis, BigDecimal gpuMillis) {
         this.applicationId = applicationId;
         this.zoneId = zoneId;
         this.cpuMillis = cpuMillis;
         this.memoryMillis = memoryMillis;
         this.diskMillis = diskMillis;
+        this.gpuMillis = gpuMillis;
         this.plan = plan;
         this.architecture = architecture;
         this.majorVersion = majorVersion;
@@ -52,6 +54,8 @@ public class ResourceUsage {
         return diskMillis;
     }
 
+    public BigDecimal getGpuMillis() { return gpuMillis; }
+
     public Plan getPlan() {
         return plan;
     }