Use memory statistics from model probing in calculation

3 files changed, 29 insertions, 6 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/DefaultOnnxModelCost.java b/config-model/src/main/java/com/yahoo/vespa/model/DefaultOnnxModelCost.java
index 76733872882..9794cfe4ad7 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/DefaultOnnxModelCost.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/DefaultOnnxModelCost.java
@@ -4,8 +4,10 @@ package com.yahoo.vespa.model;
 
 import com.yahoo.config.ModelReference;
 import com.yahoo.config.application.api.ApplicationFile;
+import com.yahoo.config.application.api.ApplicationPackage;
 import com.yahoo.config.application.api.DeployLogger;
 import com.yahoo.config.model.api.OnnxModelCost;
+import com.yahoo.vespa.model.ml.OnnxModelProbe;
 
 import java.io.IOException;
 import java.net.URI;
@@ -29,16 +31,18 @@ import static com.yahoo.yolean.Exceptions.uncheck;
 public class DefaultOnnxModelCost implements OnnxModelCost {
 
     @Override
-    public Calculator newCalculator(DeployLogger logger) {
-        return new CalculatorImpl(logger);
+    public Calculator newCalculator(ApplicationPackage appPkg, DeployLogger logger) {
+        return new CalculatorImpl(appPkg, logger);
     }
 
     private static class CalculatorImpl implements Calculator {
         private final DeployLogger log;
+        private final ApplicationPackage appPkg;
 
         private final ConcurrentMap<String, Long> modelCost = new ConcurrentHashMap<>();
 
-        private CalculatorImpl(DeployLogger log) {
+        private CalculatorImpl(ApplicationPackage appPkg, DeployLogger log) {
+            this.appPkg = appPkg;
             this.log = log;
         }
 
@@ -52,7 +56,17 @@ public class DefaultOnnxModelCost implements OnnxModelCost {
             String path = f.getPath().getRelative();
             if (alreadyAnalyzed(path)) return;
             log.log(Level.FINE, () -> "Register model '%s'".formatted(path));
-            deductJvmHeapSizeWithModelCost(f.exists() ? f.getSize() : 0, path);
+            if (f.exists()) {
+                var memoryStats = OnnxModelProbe.probeMemoryStats(appPkg, f.getPath()).orElse(null);
+                if (memoryStats != null) {
+                    log.log(Level.FINE, () -> "Register model '%s' with memory stats: %s".formatted(path, memoryStats));
+                    deductJvmHeapSizeWithModelCost(f.getSize(), memoryStats, path);
+                } else {
+                    deductJvmHeapSizeWithModelCost(f.getSize(), path);
+                }
+            } else {
+                deductJvmHeapSizeWithModelCost(0, path);
+            }
         }
 
         @Override
@@ -92,6 +106,13 @@ public class DefaultOnnxModelCost implements OnnxModelCost {
             modelCost.put(source, estimatedCost);
         }
 
+        private void deductJvmHeapSizeWithModelCost(long size, OnnxModelProbe.MemoryStats stats, String source) {
+            long estimatedCost = (long)(1.1D * stats.vmSize());
+            log.log(Level.FINE, () ->
+                    "Estimated %s footprint for model of size %s ('%s')".formatted(mb(estimatedCost), mb(size), source));
+            modelCost.put(source, estimatedCost);
+        }
+
         private boolean alreadyAnalyzed(String source) { return modelCost.containsKey(source); }
 
         private static String mb(long bytes) { return "%dMB".formatted(bytes / (1024*1024)); }
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 da6e3387d6a..d49276457b0 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
@@ -130,7 +130,8 @@ public final class ApplicationContainerCluster extends ContainerCluster<Applicat
         heapSizePercentageOfAvailableMemory = deployState.featureFlags().heapSizePercentage() > 0
                 ? Math.min(99, deployState.featureFlags().heapSizePercentage())
                 : defaultHeapSizePercentageOfAvailableMemory;
-        onnxModelCost = deployState.onnxModelCost().newCalculator(deployState.getDeployLogger());
+        onnxModelCost = deployState.onnxModelCost().newCalculator(
+                deployState.getApplicationPackage(), deployState.getDeployLogger());
         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 086f2fe778f..9ce1bb0d02a 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
@@ -4,6 +4,7 @@ package com.yahoo.vespa.model.application.validation;
 
 import com.yahoo.config.ModelReference;
 import com.yahoo.config.application.api.ApplicationFile;
+import com.yahoo.config.application.api.ApplicationPackage;
 import com.yahoo.config.application.api.DeployLogger;
 import com.yahoo.config.model.NullConfigModelRegistry;
 import com.yahoo.config.model.api.OnnxModelCost;
@@ -112,7 +113,7 @@ class JvmHeapSizeValidatorTest {
 
         ModelCostDummy(long modelCost) { this.modelCost = modelCost; }
 
-        @Override public Calculator newCalculator(DeployLogger logger) { return this; }
+        @Override public Calculator newCalculator(ApplicationPackage appPkg, DeployLogger logger) { return this; }
         @Override public long aggregatedModelCostInBytes() { return totalCost.get(); }
         @Override public void registerModel(ApplicationFile path) {}