Merge pull request #26537 from vespa-engine/arnej/add-stateless-settings-in-schemav8.144.19

Arnej/add stateless settings in schema

8 files changed, 117 insertions, 2 deletions

diff --git a/config-model/src/main/java/com/yahoo/schema/derived/Derived.java b/config-model/src/main/java/com/yahoo/schema/derived/Derived.java
index 9943a02a2f2..e8b12f22b20 100644
--- a/config-model/src/main/java/com/yahoo/schema/derived/Derived.java
+++ b/config-model/src/main/java/com/yahoo/schema/derived/Derived.java
@@ -95,7 +95,7 @@ public abstract class Derived implements Exportable {
      * @param toDirectory the directory to export to, or null
      *
      */
-    public final void export(String toDirectory) throws IOException {
+    public void export(String toDirectory) throws IOException {
         Writer writer = null;
         try {
             String fileName = getDerivedName() + ".cfg";
diff --git a/config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java b/config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java
index 4196af18fb6..e3c697e3262 100644
--- a/config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java
+++ b/config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java
@@ -16,10 +16,13 @@ import java.util.logging.Logger;
  *
  * @author bratseth
  */
-public class FileDistributedOnnxModels {
+public class FileDistributedOnnxModels extends Derived implements OnnxModelsConfig.Producer {
 
     private static final Logger log = Logger.getLogger(FileDistributedOnnxModels.class.getName());
 
+    @Override
+    public String getDerivedName() { return "onnx-models"; }
+
     private final Map<String, OnnxModel> models;
 
     public FileDistributedOnnxModels(FileRegistry fileRegistry, Collection<OnnxModel> models) {
diff --git a/config-model/src/main/java/com/yahoo/schema/derived/RankProfileList.java b/config-model/src/main/java/com/yahoo/schema/derived/RankProfileList.java
index c254385a96e..a50ddd4aeea 100644
--- a/config-model/src/main/java/com/yahoo/schema/derived/RankProfileList.java
+++ b/config-model/src/main/java/com/yahoo/schema/derived/RankProfileList.java
@@ -193,6 +193,12 @@ public class RankProfileList extends Derived implements RankProfilesConfig.Produ
     public String getDerivedName() { return "rank-profiles"; }
 
     @Override
+    public void export(String toDirectory) throws java.io.IOException {
+        super.export(toDirectory);
+        onnxModels.export(toDirectory);
+    }
+
+    @Override
     public void getConfig(RankProfilesConfig.Builder builder) {
         for (RawRankProfile rank : rankProfiles.values() ) {
             rank.getConfig(builder);
diff --git a/config-model/src/main/javacc/SchemaParser.jj b/config-model/src/main/javacc/SchemaParser.jj
index fa9d34139ea..61e8574bc87 100644
--- a/config-model/src/main/javacc/SchemaParser.jj
+++ b/config-model/src/main/javacc/SchemaParser.jj
@@ -186,6 +186,9 @@ TOKEN :
 | < SUFFIX: "suffix" >
 | < CONSTANT: "constant">
 | < ONNXMODEL: "onnx-model">
+| < INTRAOPTHREADS: "intraop-threads">
+| < INTEROPTHREADS: "interop-threads">
+| < GPUDEVICE: "gpu-device">
 | < MODEL: "model" >
 | < MUTATE: "mutate" >
 | < QUERY: "query" >
@@ -1594,11 +1597,15 @@ OnnxModel onnxModel() :
 void onnxModelItem(OnnxModel onnxModel) :
 {
     String path = null;
+    int num;
 }
 {
     (
         (path = fileItem()) { onnxModel.setFileName(path); } |
         (path = uriItem()) { onnxModel.setUri(path); } |
+        <GPUDEVICE> <COLON> num = integer() { onnxModel.setGpuDevice(num, false); } |
+        <INTRAOPTHREADS> <COLON> num = integer() { onnxModel.setStatelessIntraOpThreads(num); } |
+        <INTEROPTHREADS> <COLON> num = integer() { onnxModel.setStatelessInterOpThreads(num); } |
         (<ONNX_INPUT_SL>) {
             String name = token.image.substring(5, token.image.lastIndexOf(":")).trim();
             if (name.startsWith("\"")) { name = name.substring(1, name.length() - 1); }
diff --git a/config-model/src/test/derived/globalphase_onnx_inside/onnx-models.cfg b/config-model/src/test/derived/globalphase_onnx_inside/onnx-models.cfg
new file mode 100644
index 00000000000..f705f896540
--- /dev/null
+++ b/config-model/src/test/derived/globalphase_onnx_inside/onnx-models.cfg
@@ -0,0 +1,64 @@
+model[].name "direct"
+model[].fileref "files/ax_plus_b.onnx"
+model[].input[].name "vector_B"
+model[].input[].source "query(bb)"
+model[].input[].name "vector_A"
+model[].input[].source "attribute(aa)"
+model[].input[].name "matrix_X"
+model[].input[].source "constant(xx)"
+model[].output[].name "vector_Y"
+model[].output[].as "out"
+model[].dry_run_on_setup true
+model[].stateless_execution_mode ""
+model[].stateless_interop_threads -1
+model[].stateless_intraop_threads -1
+model[].gpu_device -1
+model[].gpu_device_required false
+model[].name "inside"
+model[].fileref "files/ax_plus_b.onnx"
+model[].input[].name "vector_B"
+model[].input[].source "rankingExpression(indirect_b)"
+model[].input[].name "vector_A"
+model[].input[].source "rankingExpression(indirect_a)"
+model[].input[].name "matrix_X"
+model[].input[].source "rankingExpression(indirect_x)"
+model[].output[].name "vector_Y"
+model[].output[].as "foobar"
+model[].dry_run_on_setup true
+model[].stateless_execution_mode ""
+model[].stateless_interop_threads 5
+model[].stateless_intraop_threads 3
+model[].gpu_device 2
+model[].gpu_device_required false
+model[].name "twoside"
+model[].fileref "files/ax_plus_b.onnx"
+model[].input[].name "vector_B"
+model[].input[].source "rankingExpression(indirect_b)"
+model[].input[].name "vector_A"
+model[].input[].source "rankingExpression(indirect_a)"
+model[].input[].name "matrix_X"
+model[].input[].source "rankingExpression(indirect_x)"
+model[].output[].name "vector_Y"
+model[].output[].as "foobar"
+model[].dry_run_on_setup true
+model[].stateless_execution_mode ""
+model[].stateless_interop_threads -1
+model[].stateless_intraop_threads -1
+model[].gpu_device -1
+model[].gpu_device_required false
+model[].name "another"
+model[].fileref "files/ax_plus_b.onnx"
+model[].input[].name "vector_B"
+model[].input[].source "rankingExpression(indirect_b)"
+model[].input[].name "vector_A"
+model[].input[].source "rankingExpression(indirect_a)"
+model[].input[].name "matrix_X"
+model[].input[].source "rankingExpression(indirect_x)"
+model[].output[].name "vector_Y"
+model[].output[].as "foobar"
+model[].dry_run_on_setup true
+model[].stateless_execution_mode ""
+model[].stateless_interop_threads -1
+model[].stateless_intraop_threads -1
+model[].gpu_device -1
+model[].gpu_device_required false
diff --git a/config-model/src/test/derived/globalphase_onnx_inside/test.sd b/config-model/src/test/derived/globalphase_onnx_inside/test.sd
index 3e48f912931..cd3dfcea3d0 100644
--- a/config-model/src/test/derived/globalphase_onnx_inside/test.sd
+++ b/config-model/src/test/derived/globalphase_onnx_inside/test.sd
@@ -49,6 +49,9 @@ schema test {
             input matrix_X: indirect_x
             input vector_B: indirect_b
             output vector_Y: foobar
+            intraop-threads: 3
+            interop-threads: 5
+            gpu-device: 2
         }
         first-phase {
             expression: sum(attribute(aa))
diff --git a/config-model/src/test/derived/globalphase_token_functions/onnx-models.cfg b/config-model/src/test/derived/globalphase_token_functions/onnx-models.cfg
new file mode 100644
index 00000000000..6283159c324
--- /dev/null
+++ b/config-model/src/test/derived/globalphase_token_functions/onnx-models.cfg
@@ -0,0 +1,16 @@
+model[].name "my_ranking_model"
+model[].fileref "files/ranking_model.onnx"
+model[].input[].name "input_ids"
+model[].input[].source "rankingExpression(input_ids)"
+model[].input[].name "attention_mask"
+model[].input[].source "rankingExpression(attention_mask)"
+model[].input[].name "token_type_ids"
+model[].input[].source "rankingExpression(token_type_ids)"
+model[].output[].name "score"
+model[].output[].as "score"
+model[].dry_run_on_setup true
+model[].stateless_execution_mode ""
+model[].stateless_interop_threads -1
+model[].stateless_intraop_threads -1
+model[].gpu_device -1
+model[].gpu_device_required false
diff --git a/config-model/src/test/derived/vector_constant/onnx-models.cfg b/config-model/src/test/derived/vector_constant/onnx-models.cfg
new file mode 100644
index 00000000000..4c52b72b519
--- /dev/null
+++ b/config-model/src/test/derived/vector_constant/onnx-models.cfg
@@ -0,0 +1,16 @@
+model[].name "inside"
+model[].fileref "ax_plus_b.onnx"
+model[].input[].name "vector_B"
+model[].input[].source "constant(bb)"
+model[].input[].name "vector_A"
+model[].input[].source "rankingExpression(indirect_a)"
+model[].input[].name "matrix_X"
+model[].input[].source "constant(xx)"
+model[].output[].name "vector_Y"
+model[].output[].as "foobar"
+model[].dry_run_on_setup true
+model[].stateless_execution_mode ""
+model[].stateless_interop_threads -1
+model[].stateless_intraop_threads -1
+model[].gpu_device -1
+model[].gpu_device_required false