Merge pull request #26537 from vespa-engine/arnej/add-stateless-setti…

…ngs-in-schema

Arnej/add stateless settings in schema

config-model/src/main/java/com/yahoo/schema/derived/Derived.java

@@ -95,7 +95,7 @@ protected String getIndex(int number, boolean labels) {
      * @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";

config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java

@@ -16,10 +16,13 @@
  *
  * @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) {

config-model/src/main/java/com/yahoo/schema/derived/RankProfileList.java

@@ -192,6 +192,12 @@ private static void addOnnxModels(Collection<OnnxModel> source,
     @Override
     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() ) {

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); }

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

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))

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

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