Fix model imports (triton-inference-server#25)

Tabrizian · web-flow · commit c4fb7130ac92 · 2020-11-24T09:47:51.000-08:00
* Fix model imports

* Update documentation
diff --git a/README.md b/README.md
@@ -55,31 +55,25 @@ $ cp examples/add_sub.py models/add_sub/1/model.py
 $ cp examples/config.pbtxt models/add_sub/config.pbtxt
 ```
 
-4. Copy `triton_python_backend_utils.py`
-
-```
-$ cp src/resources/triton_python_backend_utils.py models/add_sub/1/
-```
-
-5. Start the Triton server.
+4. Start the Triton server.
 
 ```
 $ tritonserver --model-repository `pwd`/models
 ```
 
-6. In the host machine, start the client container.
+5. In the host machine, start the client container.
 
 ```
- docker run -ti --net host nvcr.io/nvidia/tritonserver:20.10-py3-clientsdk /bin/bash
+ docker run -ti --net host nvcr.io/nvidia/tritonserver:20.10-py3-sdk /bin/bash
 ```
 
-7. In the client container, clone the Python backend repository.
+6. In the client container, clone the Python backend repository.
 
 ```
 $ git clone https://github.com/triton-inference-server/python_backend -b r20.10
 ```
 
-8. Run the example client.
+7. Run the example client.
 ```
 $ python3 python_backend/examples/add_sub_client.py
 ```
@@ -122,20 +116,14 @@ but the listed CMake argument can be used to override.
 Set `DCMAKE_INSTALL_PREFIX` to the location where the Triton Server is installed. In the released containers,
 this location is `/opt/tritonserver`.
 
-2. Copy example model and configuration
+3. Copy example model and configuration
 
 ```
 $ mkdir -p models/add_sub/1/
 $ cp examples/add_sub/model.py models/add_sub/1/model.py
 $ cp examples/add_sub/config.pbtxt models/add_sub/config.pbtxt
 ```
 
-3. Copy `triton_python_backend_utils.py`
-
-```
-$ cp src/resources/triton_python_backend_utils.py models/add_sub/1/
-```
-
 4. Start the Triton Server
 
 ```
@@ -287,23 +275,15 @@ the model configuration. In order to use this backend you must set the `backend`
 field of your model `config.pbtxt` file to `python`. You shouldn't set
 `platform` field of the configuration.
 
-Also, you need to make a copy of
-[triton_python_backend_utils.py](src/resources/triton_python_backend_utils.py)
-available to your `model.py`.
-
 Your models directory should look like below:
 ```
 models
 └── add_sub
     ├── 1
-    │   ├── model.py
-    │   └── triton_python_backend_utils.py
+    │   └── model.py
     └── config.pbtxt
 ```
 
-It is recommended to have one copy of `triton_python_backend_utils.py` along
-with every `model.py` file like the tree structure shown above.
-
 ## Changing Python Runtime Path
 
 Python backend by default uses `python3` available inside `PATH`. In order to change
diff --git a/examples/add_sub/model.py b/examples/add_sub/model.py
@@ -28,12 +28,10 @@
 import sys
 import json
 
-# You need to include triton_python_backend_utils here to be able to work with
-# inference requests and responses. It also contains some utility functions for
-# extracting information from model_config and converting Triton input/output
-# types to numpy types. You must copy
-# src/resources/triton_python_backend_utils.py in the appropriate location so
-# that the import below works properly.
+# triton_python_backend_utils is available in every Triton Python model. You
+# need to use this module to create inference requests and responses. It also
+# contains some utility functions for extracting information from model_config
+# and converting Triton input/output types to numpy types.
 import triton_python_backend_utils as pb_utils
 
 
diff --git a/examples/pytorch/model.py b/examples/pytorch/model.py
@@ -29,12 +29,10 @@
 import json
 from torch import nn
 
-# You need to include triton_python_backend_utils here to be able to work with
-# inference requests and responses. It also contains some utility functions for
-# extracting information from model_config and converting Triton input/output
-# types to numpy types. You must copy
-# src/resources/triton_python_backend_utils.py in the appropriate location so
-# that the import below works properly.
+# triton_python_backend_utils is available in every Triton Python model. You
+# need to use this module to create inference requests and responses. It also
+# contains some utility functions for extracting information from model_config
+# and converting Triton input/output types to numpy types.
 import triton_python_backend_utils as pb_utils
 
 
@@ -53,7 +51,6 @@ def forward(self, input0, input1):
         return (input0 + input1), (input0 - input1)
 
 
-
 class TritonPythonModel:
     """Your Python model must use the same class name. Every Python model
     that is created must have "TritonPythonModel" as the class name.
@@ -92,7 +89,7 @@ def initialize(self, args):
             output0_config['data_type'])
         self.output1_dtype = pb_utils.triton_string_to_numpy(
             output1_config['data_type'])
-        
+
         # Instantiate the PyTorch model
         self.add_sub_model = AddSubNet()
 
diff --git a/src/resources/startup.py b/src/resources/startup.py
@@ -35,6 +35,7 @@
 import time
 import struct
 import traceback
+from pathlib import Path
 
 from grpc_channelz.v1 import channelz
 from grpc_channelz.v1 import channelz_pb2
@@ -155,10 +156,19 @@ class PythonHost(PythonInterpreterServicer):
 
     def __init__(self, module_path, *args, **kwargs):
         super(PythonInterpreterServicer, self).__init__(*args, **kwargs)
-        spec = importlib.util.spec_from_file_location('TritonPythonModel',
-                                                      module_path)
+
+        module_path = Path(module_path).resolve()
+        # Add model parent directories so that relative and absolute import work
+        sys.path.append(str(module_path.parent))
+        sys.path.append(str(module_path.parent.parent))
+
+        # We need to import the parent directory of the module
+        # so that the relative imports work too.
+        spec = importlib.util.spec_from_file_location(
+            f'{module_path.parent.name}.{module_path.name}', str(module_path))
         module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(module)
+
         self.module_path = module_path
 
         if hasattr(module, 'TritonPythonModel'):
