diff --git a/README.md b/README.md index 2585d915..5f1682cc 100644 --- a/README.md +++ b/README.md @@ -249,7 +249,9 @@ class TritonPythonModel: inputs = [{ 'name': 'INPUT0', 'data_type': 'TYPE_FP32', - 'dims': [4] + 'dims': [4], + # this parameter will set `INPUT0 as an optional input` + 'optional': True }, { 'name': 'INPUT1', 'data_type': 'TYPE_FP32', @@ -394,6 +396,23 @@ function to gain read-only access to the `pb_utils.ModelConfig` object. The `pb_utils.ModelConfig` object being returned from here will be used as the final configuration for the model. +In addition to minimal properties, you can also set [model_transaction_policy]( + https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_configuration.md#model-transaction-policy) +through `auto_complete_config` using `set_model_transaction_policy`. +For example, +```python +import triton_python_backend_utils as pb_utils + + +class TritonPythonModel: + @staticmethod + def auto_complete_config(auto_complete_model_config): + ... + transaction_policy = {"decoupled": True} + auto_complete_model_config.set_model_transaction_policy(transaction_policy) + ... +``` + Note: The Python interpreter used to invoke this function will be destroyed upon returning from this function and as a result none of the objects created here will be available in the `initialize`, `execute`, or `finalize` diff --git a/src/resources/triton_python_backend_utils.py b/src/resources/triton_python_backend_utils.py index 560a3198..9828ab5b 100644 --- a/src/resources/triton_python_backend_utils.py +++ b/src/resources/triton_python_backend_utils.py @@ -381,12 +381,12 @@ def add_input(self, input): Raises ------ ValueError - If input contains property other than 'name', 'data_type' - and 'dims' or any of the properties are not set, or if an - input with the same name already exists in the configuration - but has different data_type or dims property + If input contains property other than 'name', 'data_type', + 'dims', 'optional' or any of the non-optional properties + are not set, or if an input with the same name already exists + in the configuration but has different data_type or dims property """ - valid_properties = ["name", "data_type", "dims"] + valid_properties = ["name", "data_type", "dims", "optional"] for current_property in input: if current_property not in valid_properties: raise ValueError( @@ -447,9 +447,26 @@ def add_input(self, input): + " but the model configuration specifies dims " + str(current_input["dims"]) ) + elif ( + "optional" in current_input + and "optional" in input + and current_input["optional"] != input["optional"] + ): + raise ValueError( + "model '" + + self._model_config["name"] + + "', tensor '" + + input["name"] + + "': the model expects optional " + + str(input["optional"]) + + " but the model configuration specifies optional " + + str(current_input["optional"]) + ) else: current_input["data_type"] = input["data_type"] current_input["dims"] = input["dims"] + if "optional" in input: + current_input["optional"] = input["optional"] return self._model_config["input"].append(input) @@ -538,6 +555,53 @@ def add_output(self, output): self._model_config["output"].append(output) + def set_model_transaction_policy(self, transaction_policy_dict): + """ + Set model transaction policy for the model. + Parameters + ---------- + transaction_policy_dict : dict + The dict, containing all properties to be set as a part + of `model_transaction_policy` field. + Raises + ------ + ValueError + If transaction_policy_dict contains property other + than 'decoupled', or if `model_transaction_policy` already exists + in the configuration, but has different `decoupled` property. + """ + valid_properties = ["decoupled"] + for current_property in transaction_policy_dict.keys(): + if current_property not in valid_properties: + raise ValueError( + "model transaction property in auto-complete-config " + + "function for model '" + + self._model_config["name"] + + "' contains property other than 'decoupled'." + ) + + if "model_transaction_policy" not in self._model_config: + self._model_config["model_transaction_policy"] = {} + + if "decoupled" in transaction_policy_dict.keys(): + if ( + "decoupled" in self._model_config["model_transaction_policy"] + and self._model_config["model_transaction_policy"]["decoupled"] + != transaction_policy_dict["decoupled"] + ): + raise ValueError( + "trying to change decoupled property in auto-complete-config " + + "for model '" + + self._model_config["name"] + + "', which is already set to '" + + str(self._model_config["model_transaction_policy"]["decoupled"]) + + "'." + ) + + self._model_config["model_transaction_policy"][ + "decoupled" + ] = transaction_policy_dict["decoupled"] + TRITONSERVER_REQUEST_FLAG_SEQUENCE_START = 1 TRITONSERVER_REQUEST_FLAG_SEQUENCE_END = 2