Skip to content

Commit

Permalink
Enhanced python_backend autocomplete (#317)
Browse files Browse the repository at this point in the history
* Added to python_backend autocomplete: optional input and model_transaction_policy
  • Loading branch information
oandreeva-nv authored Nov 2, 2023
1 parent cba7ed3 commit 0f12211
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 6 deletions.
21 changes: 20 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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`
Expand Down
74 changes: 69 additions & 5 deletions src/resources/triton_python_backend_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 0f12211

Please sign in to comment.