Supply cluster id to OnnxModelCost

5 files changed, 15 insertions, 6 deletions

diff --git a/config-model-api/abi-spec.json b/config-model-api/abi-spec.json
index b4bf45d655a..e1081ebec99 100644
--- a/config-model-api/abi-spec.json
+++ b/config-model-api/abi-spec.json
@@ -1477,7 +1477,7 @@
     ],
     "methods" : [
       "public void <init>()",
-      "public com.yahoo.config.model.api.OnnxModelCost$Calculator newCalculator(com.yahoo.config.application.api.ApplicationPackage, com.yahoo.config.provision.ApplicationId)",
+      "public com.yahoo.config.model.api.OnnxModelCost$Calculator newCalculator(com.yahoo.config.application.api.ApplicationPackage, com.yahoo.config.provision.ApplicationId, com.yahoo.config.provision.ClusterSpec$Id)",
       "public long aggregatedModelCostInBytes()",
       "public void registerModel(com.yahoo.config.application.api.ApplicationFile, com.yahoo.config.model.api.OnnxModelOptions)",
       "public void registerModel(java.net.URI, com.yahoo.config.model.api.OnnxModelOptions)",
@@ -1519,7 +1519,8 @@
       "abstract"
     ],
     "methods" : [
-      "public abstract com.yahoo.config.model.api.OnnxModelCost$Calculator newCalculator(com.yahoo.config.application.api.ApplicationPackage, com.yahoo.config.provision.ApplicationId)",
+      "public com.yahoo.config.model.api.OnnxModelCost$Calculator newCalculator(com.yahoo.config.application.api.ApplicationPackage, com.yahoo.config.provision.ApplicationId)",
+      "public abstract com.yahoo.config.model.api.OnnxModelCost$Calculator newCalculator(com.yahoo.config.application.api.ApplicationPackage, com.yahoo.config.provision.ApplicationId, com.yahoo.config.provision.ClusterSpec$Id)",
       "public static com.yahoo.config.model.api.OnnxModelCost disabled()"
     ],
     "fields" : [ ]
diff --git a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java
index 17a0d64aac9..4764d501865 100644
--- a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java
+++ b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java
@@ -4,6 +4,7 @@ package com.yahoo.config.model.api;
 import com.yahoo.config.application.api.ApplicationFile;
 import com.yahoo.config.application.api.ApplicationPackage;
 import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.ClusterSpec;
 
 import java.net.URI;
 import java.util.Map;
@@ -14,7 +15,11 @@ import java.util.Optional;
  */
 public interface OnnxModelCost {
 
-    Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId);
+    // TODO: Remove when no longer in use (oldest model version is 8.283)
+    default Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId) {
+        return newCalculator(appPkg, applicationId, null);
+    }
+    Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId, ClusterSpec.Id clusterId);
 
     interface Calculator {
         long aggregatedModelCostInBytes();
@@ -41,7 +46,7 @@ public interface OnnxModelCost {
     static OnnxModelCost disabled() { return new DisabledOnnxModelCost(); }
 
     class DisabledOnnxModelCost implements OnnxModelCost, Calculator {
-        @Override public Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId) { return this; }
+        @Override public Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId, ClusterSpec.Id clusterId) { return this; }
         @Override public long aggregatedModelCostInBytes() {return 0;}
         @Override public void registerModel(ApplicationFile path, OnnxModelOptions onnxModelOptions) {}
         @Override public void registerModel(URI uri, OnnxModelOptions onnxModelOptions) {}
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainerCluster.java b/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainerCluster.java
index 8c4adfb96cb..9896ca95e97 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainerCluster.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/container/ApplicationContainerCluster.java
@@ -16,6 +16,7 @@ import com.yahoo.config.model.api.OnnxModelCost;
 import com.yahoo.config.model.deploy.DeployState;
 import com.yahoo.config.model.producer.TreeConfigProducer;
 import com.yahoo.config.provision.AllocatedHosts;
+import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.HostSpec;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.container.bundle.BundleInstantiationSpecification;
@@ -139,7 +140,7 @@ public final class ApplicationContainerCluster extends ContainerCluster<Applicat
                 : defaultHeapSizePercentageOfAvailableMemory;
         onnxModelCost = deployState.onnxModelCost();
         onnxModelCostCalculator = deployState.onnxModelCost().newCalculator(
-                deployState.getApplicationPackage(), deployState.getProperties().applicationId());
+                deployState.getApplicationPackage(), deployState.getProperties().applicationId(), ClusterSpec.Id.from(clusterId));
         logger = deployState.getDeployLogger();
     }
 
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/application/validation/JvmHeapSizeValidatorTest.java b/config-model/src/test/java/com/yahoo/vespa/model/application/validation/JvmHeapSizeValidatorTest.java
index 31e4c661151..e4ba003df64 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/application/validation/JvmHeapSizeValidatorTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/application/validation/JvmHeapSizeValidatorTest.java
@@ -13,6 +13,7 @@ import com.yahoo.config.model.deploy.TestProperties;
 import com.yahoo.config.model.provision.InMemoryProvisioner;
 import com.yahoo.config.model.test.MockApplicationPackage;
 import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.text.Text;
 import com.yahoo.vespa.model.VespaModel;
@@ -121,6 +122,7 @@ class JvmHeapSizeValidatorTest {
         ModelCostDummy(long modelCost) { this.modelCost = modelCost; }
 
         @Override public Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId) { return this; }
+        @Override public Calculator newCalculator(ApplicationPackage appPkg, ApplicationId applicationId, ClusterSpec.Id clusterId) { return this; }
         @Override public Map<String, ModelInfo> models() { return Map.of(); }
         @Override public void setRestartOnDeploy() {}
         @Override public boolean restartOnDeploy() { return false;}
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/application/validation/change/RestartOnDeployForOnnxModelChangesValidatorTest.java b/config-model/src/test/java/com/yahoo/vespa/model/application/validation/change/RestartOnDeployForOnnxModelChangesValidatorTest.java
index 13389689de5..f2013f0f26b 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/application/validation/change/RestartOnDeployForOnnxModelChangesValidatorTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/application/validation/change/RestartOnDeployForOnnxModelChangesValidatorTest.java
@@ -79,7 +79,7 @@ public class RestartOnDeployForOnnxModelChangesValidatorTest {
     }
 
     private static OnnxModelCost onnxModelCost(long estimatedCost, long hash) {
-        return (appPkg, applicationId) -> new OnnxModelCost.Calculator() {
+        return (appPkg, applicationId, clusterId) -> new OnnxModelCost.Calculator() {
 
             private final Map<String, OnnxModelCost.ModelInfo> models = new HashMap<>();
             private boolean restartOnDeploy = false;