dry run onnx models on setup

11 files changed, 161 insertions, 10 deletions

diff --git a/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp b/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp
index 3f56610dcaa..ce7070d6b2b 100644
--- a/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp
+++ b/eval/src/apps/analyze_onnx_model/analyze_onnx_model.cpp
@@ -172,7 +172,7 @@ int usage(const char *self) {
     return 1;
 }
 
-int main(int argc, char **argv) {
+int my_main(int argc, char **argv) {
     if (argc < 2) {
         return usage(argv[0]);
     }
@@ -206,3 +206,12 @@ int main(int argc, char **argv) {
     fprintf(stderr, "estimated model evaluation time: %g ms\n", min_time_s * 1000.0);
     return 0;
 }
+
+int main(int argc, char **argv) {
+    try {
+        return my_main(argc, argv);
+    } catch (const std::exception &ex) {
+        fprintf(stderr, "got exception: %s\n", ex.what());
+        return 2;
+    }
+}
diff --git a/searchcore/src/apps/verify_ranksetup/verify_ranksetup.cpp b/searchcore/src/apps/verify_ranksetup/verify_ranksetup.cpp
index 853395be3e1..7a3803626c8 100644
--- a/searchcore/src/apps/verify_ranksetup/verify_ranksetup.cpp
+++ b/searchcore/src/apps/verify_ranksetup/verify_ranksetup.cpp
@@ -67,7 +67,7 @@ RankingExpressions make_expressions(const RankingExpressionsConfig &expressionsC
         if (auto file = get_file(entry.fileref, myCfg)) {
             expressions.add(entry.name, file.value());
         } else {
-            LOG(warning, "could not find file for ranking expression '%s' (ref:'%s')",
+            LOG(warning, "could not find file name for ranking expression '%s' (ref:'%s')",
                 entry.name.c_str(), entry.fileref.c_str());
         }
     }
@@ -81,7 +81,7 @@ OnnxModels make_models(const OnnxModelsConfig &modelsCfg, const VerifyRanksetupC
             model_list.emplace_back(entry.name, file.value());
             OnnxModels::configure(entry, model_list.back());
         } else {
-            LOG(warning, "could not find file for onnx model '%s' (ref:'%s')",
+            LOG(warning, "could not find file name for onnx model '%s' (ref:'%s')",
                 entry.name.c_str(), entry.fileref.c_str());
         }
     }
diff --git a/searchcore/src/tests/proton/verify_ranksetup/verify_ranksetup_test.cpp b/searchcore/src/tests/proton/verify_ranksetup/verify_ranksetup_test.cpp
index fc70bafed7f..c13440de4d5 100644
--- a/searchcore/src/tests/proton/verify_ranksetup/verify_ranksetup_test.cpp
+++ b/searchcore/src/tests/proton/verify_ranksetup/verify_ranksetup_test.cpp
@@ -188,6 +188,7 @@ struct Setup {
                 out.fmt("model[%zu].output[%zu].as \"%s\"\n", idx, idx2, output.second.c_str());
                 ++idx2;
             }
+            out.fmt("model[%zu].dry_run_on_setup %s\n", idx, entry.second.dry_run_on_setup() ? "true" : "false");
             ++idx;
         }
     }
@@ -269,6 +270,10 @@ struct OnnxSetup : Setup {
                        .input_feature("attribute_tensor", "rankingExpression(at)")
                        .input_feature("bias_tensor", "rankingExpression(bt)")
                        .output_name("output", "result"));
+        add_onnx_model(OnnxModel("fragile", TEST_PATH("../../../../../searchlib/src/tests/features/onnx_feature/fragile.onnx"))
+                       .dry_run_on_setup(true));
+        add_onnx_model(OnnxModel("unfragile", TEST_PATH("../../../../../searchlib/src/tests/features/onnx_feature/fragile.onnx"))
+                       .dry_run_on_setup(false));
     }
 };
 
@@ -417,6 +422,24 @@ TEST_F("require that onnx model can have inputs and outputs mapped", OnnxSetup()
     f.verify_valid({"onnxModel(mapped).result"});
 }
 
+TEST_F("require that fragile model can pass verification", OnnxSetup()) {
+    f.rank_expr("in1", "tensor<float>(a[2]):[1,2]");
+    f.rank_expr("in2", "tensor<float>(a[2]):[3,4]");
+    f.verify_valid({"onnxModel(fragile)"});
+}
+
+TEST_F("require that broken fragile model fails verification", OnnxSetup()) {
+    f.rank_expr("in1", "tensor<float>(a[2]):[1,2]");
+    f.rank_expr("in2", "tensor<float>(a[3]):[3,4,31515]");
+    f.verify_invalid({"onnxModel(fragile)"});
+}
+
+TEST_F("require that broken fragile model without dry-run passes verification", OnnxSetup()) {
+    f.rank_expr("in1", "tensor<float>(a[2]):[1,2]");
+    f.rank_expr("in2", "tensor<float>(a[3]):[3,4,31515]");
+    f.verify_valid({"onnxModel(unfragile)"});
+}
+
 //-----------------------------------------------------------------------------
 
 TEST_F("cleanup files", Setup()) {
diff --git a/searchcore/src/vespa/searchcore/proton/matching/onnx_models.cpp b/searchcore/src/vespa/searchcore/proton/matching/onnx_models.cpp
index ed80ca28bd6..43e29e37b62 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/onnx_models.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/onnx_models.cpp
@@ -46,6 +46,7 @@ OnnxModels::configure(const ModelConfig &config, Model &model)
     for (const auto &output: config.output) {
         model.output_name(output.name, output.as);
     }
+    model.dry_run_on_setup(config.dryRunOnSetup);
 }
 
 }
diff --git a/searchlib/src/tests/features/onnx_feature/fragile.onnx b/searchlib/src/tests/features/onnx_feature/fragile.onnx
new file mode 100644
index 00000000000..2a05500e95b
--- /dev/null
+++ b/searchlib/src/tests/features/onnx_feature/fragile.onnx
@@ -0,0 +1,15 @@
+
+fragile.py:b
+
+in1
+in2out"AddfragileZ
+in1
+
+

+Z
+in2
+

	
+batchb
+out
+

	
+batchB
\ No newline at end of file
diff --git a/searchlib/src/tests/features/onnx_feature/fragile.py b/searchlib/src/tests/features/onnx_feature/fragile.py
new file mode 100755
index 00000000000..e4eaf168e14
--- /dev/null
+++ b/searchlib/src/tests/features/onnx_feature/fragile.py
@@ -0,0 +1,30 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+import onnx
+from onnx import helper, TensorProto
+
+INPUT1 = helper.make_tensor_value_info('in1', TensorProto.FLOAT, [2])
+INPUT2 = helper.make_tensor_value_info('in2', TensorProto.FLOAT, ['batch'])
+
+OUTPUT = helper.make_tensor_value_info('out', TensorProto.FLOAT, ['batch'])
+
+nodes = [
+    helper.make_node(
+        'Add',
+        ['in1', 'in2'],
+        ['out'],
+    ),
+]
+graph_def = helper.make_graph(
+    nodes,
+    'fragile',
+    [
+        INPUT1,
+        INPUT2,
+    ],
+    [
+        OUTPUT,
+    ],
+)
+model_def = helper.make_model(graph_def, producer_name='fragile.py', opset_imports=[onnx.OperatorSetIdProto(version=12)])
+onnx.save(model_def, 'fragile.onnx')
diff --git a/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp b/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp
index 6a1e4ef9fa1..c07ebc48604 100644
--- a/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp
+++ b/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp
@@ -28,6 +28,7 @@ std::string vespa_dir = source_dir + "/" + "../../../../..";
 std::string simple_model = vespa_dir + "/" + "eval/src/tests/tensor/onnx_wrapper/simple.onnx";
 std::string dynamic_model = vespa_dir + "/" + "eval/src/tests/tensor/onnx_wrapper/dynamic.onnx";
 std::string strange_names_model = source_dir + "/" + "strange_names.onnx";
+std::string fragile_model = source_dir + "/" + "fragile.onnx";
 
 uint32_t default_docid = 1;
 
@@ -61,13 +62,19 @@ struct OnnxFeatureTest : ::testing::Test {
     void add_onnx(const OnnxModel &model) {
         indexEnv.addOnnxModel(model);
     }
-    void compile(const vespalib::string &seed) {
+    bool try_compile(const vespalib::string &seed) {
         resolver->addSeed(seed);
-        ASSERT_TRUE(resolver->compile());
+        if (!resolver->compile()) {
+            return false;
+        }
         MatchDataLayout mdl;
         QueryEnvironment queryEnv(&indexEnv);
         match_data = mdl.createMatchData();
         program.setup(*match_data, queryEnv, overrides);
+        return true;
+    }
+    void compile(const vespalib::string &seed) {
+        ASSERT_TRUE(try_compile(seed));
     }
     TensorSpec get(const vespalib::string &feature, uint32_t docid) {
         auto result = program.get_all_features(false);
@@ -137,4 +144,27 @@ TEST_F(OnnxFeatureTest, input_features_and_output_names_can_be_specified) {
     EXPECT_EQ(get("onnxModel(custom_names).my_second_output", 1), expect_sub);
 }
 
+TEST_F(OnnxFeatureTest, fragile_model_can_be_evaluated) {
+    add_expr("in1", "tensor<float>(x[2]):[docid,5]");
+    add_expr("in2", "tensor<float>(x[2]):[docid,10]");
+    add_onnx(OnnxModel("fragile", fragile_model));
+    EXPECT_TRUE(try_compile(onnx_feature("fragile")));
+    EXPECT_EQ(get(1), TensorSpec::from_expr("tensor<float>(d0[2]):[2,15]"));
+    EXPECT_EQ(get(3), TensorSpec::from_expr("tensor<float>(d0[2]):[6,15]"));
+}
+
+TEST_F(OnnxFeatureTest, runtime_broken_model_can_be_set_up_without_dry_run) {
+    add_expr("in1", "tensor<float>(x[2]):[docid,5]");
+    add_expr("in2", "tensor<float>(x[3]):[docid,10,31515]");
+    add_onnx(OnnxModel("fragile", fragile_model).dry_run_on_setup(false));
+    EXPECT_TRUE(try_compile(onnx_feature("fragile")));
+}
+
+TEST_F(OnnxFeatureTest, runtime_broken_model_fails_with_dry_run) {
+    add_expr("in1", "tensor<float>(x[2]):[docid,5]");
+    add_expr("in2", "tensor<float>(x[3]):[docid,10,31515]");
+    add_onnx(OnnxModel("fragile", fragile_model).dry_run_on_setup(true));
+    EXPECT_FALSE(try_compile(onnx_feature("fragile")));
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchlib/src/tests/features/onnx_feature/strange_names.py b/searchlib/src/tests/features/onnx_feature/strange_names.py
index 681da641264..2e3d3fe4dd1 100755
--- a/searchlib/src/tests/features/onnx_feature/strange_names.py
+++ b/searchlib/src/tests/features/onnx_feature/strange_names.py
@@ -32,5 +32,5 @@ graph_def = helper.make_graph(
         OUTPUT2,
     ],
 )
-model_def = helper.make_model(graph_def, producer_name='strange_names.py')
+model_def = helper.make_model(graph_def, producer_name='strange_names.py', opset_imports=[onnx.OperatorSetIdProto(version=12)])
 onnx.save(model_def, 'strange_names.onnx')
diff --git a/searchlib/src/vespa/searchlib/features/onnx_feature.cpp b/searchlib/src/vespa/searchlib/features/onnx_feature.cpp
index e9fecb3578e..0af03f3aa86 100644
--- a/searchlib/src/vespa/searchlib/features/onnx_feature.cpp
+++ b/searchlib/src/vespa/searchlib/features/onnx_feature.cpp
@@ -5,6 +5,9 @@
 #include <vespa/searchlib/fef/onnx_model.h>
 #include <vespa/searchlib/fef/featureexecutor.h>
 #include <vespa/eval/eval/value.h>
+#include <vespa/eval/eval/tensor_spec.h>
+#include <vespa/eval/eval/fast_value.h>
+#include <vespa/eval/eval/value_codec.h>
 #include <vespa/vespalib/util/stringfmt.h>
 #include <vespa/vespalib/util/stash.h>
 
@@ -18,7 +21,11 @@ using search::fef::IIndexEnvironment;
 using search::fef::IQueryEnvironment;
 using search::fef::ParameterList;
 using vespalib::Stash;
+using vespalib::eval::Value;
 using vespalib::eval::ValueType;
+using vespalib::eval::TensorSpec;
+using vespalib::eval::FastValueBuilderFactory;
+using vespalib::eval::value_from_spec;
 using vespalib::make_string_short::fmt;
 using vespalib::eval::Onnx;
 
@@ -41,8 +48,27 @@ vespalib::string normalize_name(const vespalib::string &name, const char *contex
     return result;
 }
 
+vespalib::string my_dry_run(const Onnx &model, const Onnx::WireInfo &wire) {
+    vespalib::string error_msg;
+    try {
+        Onnx::EvalContext context(model, wire);
+        std::vector<Value::UP> inputs;
+        for (const auto &input_type: wire.vespa_inputs) {
+            TensorSpec spec(input_type.to_spec());
+            inputs.push_back(value_from_spec(spec, FastValueBuilderFactory::get()));
+        }
+        for (size_t i = 0; i < inputs.size(); ++i) {
+            context.bind_param(i, *inputs[i]);
+        }
+        context.eval();
+    } catch (const Ort::Exception &ex) {
+        error_msg = ex.what();
+    }
+    return error_msg;
 }
 
+} // <unnamed>
+
 /**
  * Feature executor that evaluates an onnx model
  */
@@ -94,7 +120,7 @@ OnnxBlueprint::setup(const IIndexEnvironment &env,
             _cache_token = OnnxModelCache::load(model_cfg->file_path());
             _model = &(_cache_token->get());
         }
-    } catch (std::exception &ex) {
+    } catch (const Ort::Exception &ex) {
         return fail("model setup failed: %s", ex.what());
     }
     Onnx::WirePlanner planner;
@@ -131,6 +157,12 @@ OnnxBlueprint::setup(const IIndexEnvironment &env,
         describeOutput(output_name.value(), "output from onnx model", FeatureType::object(output_type));
     }
     _wire_info = planner.get_wire_info(*_model);
+    if (model_cfg->dry_run_on_setup()) {
+        auto error_msg = my_dry_run(*_model, _wire_info);
+        if (!error_msg.empty()) {
+            return fail("onnx model dry-run failed: %s", error_msg.c_str());
+        }
+    }
     return true;
 }
 
diff --git a/searchlib/src/vespa/searchlib/fef/onnx_model.cpp b/searchlib/src/vespa/searchlib/fef/onnx_model.cpp
index ba5adaae857..c301fcf6dca 100644
--- a/searchlib/src/vespa/searchlib/fef/onnx_model.cpp
+++ b/searchlib/src/vespa/searchlib/fef/onnx_model.cpp
@@ -10,7 +10,8 @@ OnnxModel::OnnxModel(const vespalib::string &name_in,
     : _name(name_in),
       _file_path(file_path_in),
       _input_features(),
-      _output_names()
+      _output_names(),
+      _dry_run_on_setup(false)
 {
 }
 
@@ -26,6 +27,13 @@ OnnxModel::output_name(const vespalib::string &model_output_name, const vespalib
     return *this;
 }
 
+OnnxModel &
+OnnxModel::dry_run_on_setup(bool value)
+{
+    _dry_run_on_setup = value;
+    return *this;
+}
+
 std::optional<vespalib::string>
 OnnxModel::input_feature(const vespalib::string &model_input_name) const {
     auto pos = _input_features.find(model_input_name);
@@ -46,8 +54,8 @@ OnnxModel::output_name(const vespalib::string &model_output_name) const {
 
 bool
 OnnxModel::operator==(const OnnxModel &rhs) const {
-    return (std::tie(_name, _file_path, _input_features, _output_names) ==
-            std::tie(rhs._name, rhs._file_path, rhs._input_features, rhs._output_names));
+    return (std::tie(_name, _file_path, _input_features, _output_names, _dry_run_on_setup) ==
+            std::tie(rhs._name, rhs._file_path, rhs._input_features, rhs._output_names, rhs._dry_run_on_setup));
 }
 
 OnnxModel::~OnnxModel() = default;
diff --git a/searchlib/src/vespa/searchlib/fef/onnx_model.h b/searchlib/src/vespa/searchlib/fef/onnx_model.h
index 2195a50600d..2d22dd8eb8d 100644
--- a/searchlib/src/vespa/searchlib/fef/onnx_model.h
+++ b/searchlib/src/vespa/searchlib/fef/onnx_model.h
@@ -19,6 +19,7 @@ private:
     vespalib::string _file_path;
     std::map<vespalib::string,vespalib::string> _input_features;
     std::map<vespalib::string,vespalib::string> _output_names;
+    bool _dry_run_on_setup;
 
 public:
     OnnxModel(const vespalib::string &name_in,
@@ -29,8 +30,10 @@ public:
     const vespalib::string &file_path() const { return _file_path; }
     OnnxModel &input_feature(const vespalib::string &model_input_name, const vespalib::string &input_feature);
     OnnxModel &output_name(const vespalib::string &model_output_name, const vespalib::string &output_name);
+    OnnxModel &dry_run_on_setup(bool value);
     std::optional<vespalib::string> input_feature(const vespalib::string &model_input_name) const;
     std::optional<vespalib::string> output_name(const vespalib::string &model_output_name) const;
+    bool dry_run_on_setup() const { return _dry_run_on_setup; }
     bool operator==(const OnnxModel &rhs) const;
     const std::map<vespalib::string,vespalib::string> &inspect_input_features() const { return _input_features; }
     const std::map<vespalib::string,vespalib::string> &inspect_output_names() const { return _output_names; }