Fix ensemble check in requiring multiple versions of the same model c…

…ase (triton-inference-server#595)

* Fix ensemble check in requiring multiple versions of the same model

* Improve error message and remoe redundant code

* Add corresponding test in L0_model_config

* Deploy identity model on CPU

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version/expected

@@ -0,0 +1 @@
+ensemble 'no_required_version' depends on 'simple' whose required version 2 is not loaded

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version/no_required_version/config.pbtxt

@@ -0,0 +1,51 @@
+name: "no_required_version"
+platform: "ensemble"
+max_batch_size: 8
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+ensemble_scheduling {
+  step [
+    {
+      model_name: "simple"
+      model_version: 2
+      input_map {
+        key: "INPUT0"
+        value: "INPUT0"
+      }
+      input_map {
+        key: "INPUT1"
+        value: "INPUT1"
+      }
+      output_map {
+        key: "OUTPUT0"
+        value: "OUTPUT0"
+      }
+      output_map {
+        key: "OUTPUT1"
+        value: "OUTPUT1"
+      }
+    }
+  ]
+}

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version/simple/config.pbtxt

@@ -0,0 +1,33 @@
+name: "simple"
+platform: "custom"
+max_batch_size: 8
+version_policy : { all {} }
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+instance_group [
+  {
+    kind: KIND_CPU
+  }
+]

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version_2/expected

@@ -0,0 +1 @@
+ensemble 'no_required_version_2' depends on 'simple' whose required version 2 is not loaded

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version_2/no_required_version_2/config.pbtxt

@@ -0,0 +1,71 @@
+name: "no_required_version_2"
+platform: "ensemble"
+max_batch_size: 8
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+ensemble_scheduling {
+  step [
+    {
+      model_name: "simple"
+      model_version: -1
+      input_map {
+        key: "INPUT0"
+        value: "INPUT0"
+      }
+      input_map {
+        key: "INPUT1"
+        value: "INPUT1"
+      }
+      output_map {
+        key: "OUTPUT0"
+        value: "temp0"
+      }
+      output_map {
+        key: "OUTPUT1"
+        value: "temp1"
+      }
+    },
+    {
+      model_name: "simple"
+      model_version: 2
+      input_map {
+        key: "INPUT0"
+        value: "temp0"
+      }
+      input_map {
+        key: "INPUT1"
+        value: "temp1"
+      }
+      output_map {
+        key: "OUTPUT0"
+        value: "OUTPUT0"
+      }
+      output_map {
+        key: "OUTPUT1"
+        value: "OUTPUT1"
+      }
+    }
+  ]
+}

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version_2/simple/config.pbtxt

@@ -0,0 +1,33 @@
+name: "simple"
+platform: "custom"
+max_batch_size: 8
+version_policy : { all {} }
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+instance_group [
+  {
+    kind: KIND_CPU
+  }
+]

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version_3/expected

@@ -0,0 +1 @@
+ensemble 'no_required_version_3' depends on 'simple' whose required version 2 is not loaded

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version_3/no_required_version_3/config.pbtxt

@@ -0,0 +1,71 @@
+name: "no_required_version_3"
+platform: "ensemble"
+max_batch_size: 8
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+ensemble_scheduling {
+  step [
+    {
+      model_name: "simple"
+      model_version: 2
+      input_map {
+        key: "INPUT0"
+        value: "INPUT0"
+      }
+      input_map {
+        key: "INPUT1"
+        value: "INPUT1"
+      }
+      output_map {
+        key: "OUTPUT0"
+        value: "temp0"
+      }
+      output_map {
+        key: "OUTPUT1"
+        value: "temp1"
+      }
+    },
+    {
+      model_name: "simple"
+      model_version: -1
+      input_map {
+        key: "INPUT0"
+        value: "temp0"
+      }
+      input_map {
+        key: "INPUT1"
+        value: "temp1"
+      }
+      output_map {
+        key: "OUTPUT0"
+        value: "OUTPUT0"
+      }
+      output_map {
+        key: "OUTPUT1"
+        value: "OUTPUT1"
+      }
+    }
+  ]
+}

qa/L0_model_config/autofill_noplatform/ensemble/no_required_version_3/simple/config.pbtxt

@@ -0,0 +1,33 @@
+name: "simple"
+platform: "custom"
+max_batch_size: 8
+version_policy : { all {} }
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+output [
+  {
+    name: "OUTPUT0"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  },
+  {
+    name: "OUTPUT1"
+    data_type: TYPE_INT32
+    dims: [ 16 ]
+  }
+]
+instance_group [
+  {
+    kind: KIND_CPU
+  }
+]

qa/L0_model_config/test.sh

@@ -88,6 +88,12 @@ for modelpath in \
        autofill_noplatform/ensemble/unmapped_input/fp32_dim1_batch4_input4/1 \
        autofill_noplatform/ensemble/circular_dependency/circular_dependency/1 \
        autofill_noplatform/ensemble/circular_dependency/circular_dependency_2/1 \
+       autofill_noplatform/ensemble/no_required_version/no_required_version/1 \
+       autofill_noplatform/ensemble/no_required_version/simple/1 \
+       autofill_noplatform/ensemble/no_required_version_2/no_required_version_2/1 \
+       autofill_noplatform/ensemble/no_required_version_2/simple/1 \
+       autofill_noplatform/ensemble/no_required_version_3/no_required_version_3/1 \
+       autofill_noplatform/ensemble/no_required_version_3/simple/1 \
        autofill_noplatform_success/ensemble/embedded_ensemble/embedded_ensemble/1 \
        autofill_noplatform_success/ensemble/embedded_ensemble/fp32_dim1_batch4/1 \
        autofill_noplatform_success/ensemble/embedded_ensemble/inner_ensemble/1 \

src/core/model_repository_manager.cc

@@ -1437,6 +1437,7 @@ ModelRepositoryManager::ConnectDependencyGraph(DependencyNode* updated_node)
   if (updated_node->model_config_.has_ensemble_scheduling()) {
     for (const auto& step :
          updated_node->model_config_.ensemble_scheduling().step()) {
+      DependencyNode* upstream_node = nullptr;
       const auto& model_name = step.model_name();
       auto dit = dependency_graph_.find(model_name);
       if (dit == dependency_graph_.end()) {
@@ -1449,12 +1450,17 @@ ModelRepositoryManager::ConnectDependencyGraph(DependencyNode* updated_node)
         // Add the node to missing node's downstream so that when the missing
         // node is added, the downstreams can be found easily.
         mit->second->downstreams_.emplace(updated_node);
-        updated_node->upstreams_.emplace(
-            mit->second.get(), step.model_version());
+        upstream_node = mit->second.get();
       } else {
         dit->second->downstreams_.emplace(updated_node);
-        updated_node->upstreams_.emplace(
-            dit->second.get(), step.model_version());
+        upstream_node = dit->second.get();
+      }
+      auto res = updated_node->upstreams_.emplace(
+          upstream_node, std::set<int64_t>({step.model_version()}));
+      // If map insertion doesn't happen, the same model is required in
+      // different step, insert the version to existing required version set.
+      if (!res.second) {
+        res.first->second.insert(step.model_version());
       }
     }
     return true;
@@ -1546,14 +1552,20 @@ ModelRepositoryManager::CheckNode(DependencyNode* node)
             RequestStatusCode::INVALID_ARG,
             "ensemble '" + node->model_name_ + "' depends on '" +
                 upstream.first->model_name_ + "' which has no loaded version");
-      } else if (upstream.second != -1) {
-        auto it = upstream.first->loaded_versions_.find(upstream.second);
-        if (it == upstream.first->loaded_versions_.end()) {
-          node->status_ = Status(
-              RequestStatusCode::INVALID_ARG,
-              "ensemble '" + node->model_name_ + "' depends on '" +
-                  upstream.first->model_name_ +
-                  "' whose required version is not loaded");
+      } else {
+        for (const auto& required_version : upstream.second) {
+          if (required_version == -1) {
+            continue;
+          }
+
+          auto it = upstream.first->loaded_versions_.find(required_version);
+          if (it == upstream.first->loaded_versions_.end()) {
+            node->status_ = Status(
+                RequestStatusCode::INVALID_ARG,
+                "ensemble '" + node->model_name_ + "' depends on '" +
+                    upstream.first->model_name_ + "' whose required version " +
+                    std::to_string(required_version) + " is not loaded");
+          }
         }
       }
       if (!node->status_.IsOk()) {

src/core/model_repository_manager.h

@@ -62,7 +62,7 @@ class ModelRepositoryManager {
     ModelConfig model_config_;
     std::set<int64_t> loaded_versions_;
     std::set<DependencyNode*> missing_upstreams_;
-    std::unordered_map<DependencyNode*, int64_t> upstreams_;
+    std::unordered_map<DependencyNode*, std::set<int64_t>> upstreams_;
     std::set<DependencyNode*> downstreams_;
   };