Add type for machine-learned models and ensure correct file reference is passed to `OnnxModelCost`

1 files changed, 6 insertions, 2 deletions

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 595cd97e6b6..33ed55ecaef 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
@@ -7,6 +7,8 @@ import com.yahoo.config.application.api.ApplicationFile;
 import com.yahoo.config.application.api.ApplicationPackage;
 import com.yahoo.config.application.api.DeployLogger;
 
+import java.net.URI;
+
 /**
  * @author bjorncs
  */
@@ -17,14 +19,16 @@ public interface OnnxModelCost {
     interface Calculator {
         long aggregatedModelCostInBytes();
         void registerModel(ApplicationFile path);
-        void registerModel(ModelReference ref);
+        @Deprecated(forRemoval = true) void registerModel(ModelReference ref); // TODO(bjorncs): remove once no longer in use by old config models
+        void registerModel(URI uri);
     }
 
     static OnnxModelCost disabled() {
         return (__, ___) -> new Calculator() {
             @Override public long aggregatedModelCostInBytes() { return 0; }
             @Override public void registerModel(ApplicationFile path) {}
-            @Override public void registerModel(ModelReference ref) {}
+            @SuppressWarnings("removal") @Override public void registerModel(ModelReference ref) {}
+            @Override public void registerModel(URI uri) {}
         };
     }
 }