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 9943a02a2f28..e8b12f22b200 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 @@ 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"; 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 4196af18fb64..e3c697e32628 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 @@ * * @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 models; public FileDistributedOnnxModels(FileRegistry fileRegistry, Collection 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 c254385a96e2..a50ddd4aeea4 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 @@ -192,6 +192,12 @@ private static void addOnnxModels(Collection 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() ) { diff --git a/config-model/src/main/javacc/SchemaParser.jj b/config-model/src/main/javacc/SchemaParser.jj index fa9d34139eae..61e8574bc872 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); } | + num = integer() { onnxModel.setGpuDevice(num, false); } | + num = integer() { onnxModel.setStatelessIntraOpThreads(num); } | + num = integer() { onnxModel.setStatelessInterOpThreads(num); } | () { 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 000000000000..f705f8965404 --- /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 3e48f912931c..cd3dfcea3d04 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 000000000000..6283159c324c --- /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 000000000000..4c52b72b5190 --- /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