Add ONNX and ORT support + Docs for TensorRT

[msaroufim]

EDIT 11/4: To make this easier to merge I've cut the scope to only ONNX and ORT, will revisit the rest after pytorch/tensorrt when it gets official pypi binaries but this is now ready for review

EDIT 11/6: I'm going to write a brief doc page on using different optimization runtimes

EDIT 11/8: Addressed most feedback, rest will be addressed in future work

This PR

• a --serialized-file that's in .onnx format, which will be correctly loaded by the base handler using an ort.InferenceSession() https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html
• TensorRT format is a .ts extension which can be loaded via torch.jit.load() https://pytorch.org/TensorRT/getting_started/getting_started_with_python_api.html#getting-started-with-python-api
• A hardcoded model_config.json similar to index_to_name.json which would have model specific information which you can get access via `model_config.get("property"). For now I'm not using any special configs
• Starting with torch 1.12 nvfuser is now the default backend for torchscript so that's a free performance enhancement we got where we don't need to do anything

Open question

• Should we do the conversion from torch to ONNX in the initialize()? My gut is no, we don't ask users to train models in initialize() they should prepare their models and once prepared used serve but I could be convinced otherwise - EDIT: NO
• Should I put a schema on model_config.json? I'd like to wait and see what schema gets chosen for Modularize ipex optimization in base_handler.py into ts/utils/ipex_optimization.py.  #1664 and then decide
• There are other ways of designing the same thing but I felt like this was the simplest, I thought about the here https://gist.github.com/msaroufim/c5aa505fdd7053ca184da921536c26f1

Future work

• Add AITemplate support https://ai.facebook.com/blog/gpu-inference-engine-nvidia-amd-open-source/?utm_source=twitter&utm_medium=organic_social&utm_campaign=blog
• nvfuser is actually already supported and starting 1.12 has become the default fusing backend for torch.jit.script()

[codecov]

Codecov Report

Merging #1857 (0733614) into master (a5fc61f) will increase coverage by 2.99%.
The diff coverage is 25.67%.

❗ Current head 0733614 differs from pull request most recent head 6c9c7c5. Consider uploading reports for the commit 6c9c7c5 to get more accurate results

@@            Coverage Diff             @@
##           master    #1857      +/-   ##
==========================================
+ Coverage   41.67%   44.66%   +2.99%     
==========================================
  Files          55       63       +8     
  Lines        2282     2624     +342     
  Branches        1       56      +55     
==========================================
+ Hits          951     1172     +221     
- Misses       1331     1452     +121     

Impacted Files	Coverage Δ
setup.py	0.00% <0.00%> (ø)
ts/torch_handler/base_handler.py	0.00% <0.00%> (ø)
ts/utils/util.py	37.50% <ø> (ø)
...l-archiver/model_archiver/model_packaging_utils.py	54.66% <54.28%> (ø)
...el-archiver/model_archiver/model_archiver_error.py	100.00% <0.00%> (ø)
workflow-archiver/workflow_archiver/version.py	100.00% <0.00%> (ø)
model-archiver/model_archiver/model_packaging.py	90.00% <0.00%> (ø)
...w-archiver/workflow_archiver/workflow_packaging.py	89.65% <0.00%> (ø)
model-archiver/model_archiver/version.py	100.00% <0.00%> (ø)
...hiver/workflow_archiver/workflow_archiver_error.py	100.00% <0.00%> (ø)
... and 2 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[HamidShojanazeri]

Beside the inline comments, we also would need to think of preprocessing for onnx (to avoid custom handlers), as the input need to be converted to numpy. We might be able to have separate util files for onnx, trt or future backends/dynamo or alternatively have one backend utils to keep all the related helper functions there. This may help to keep the base_handler cleaner and more maintainable. cc @lxning on this as well.

[lxning]

it should be dynamically loaded based on model runtime type.

[msaroufim]

I think we should do that for our common dependencies but for our developer dependencies we should install everything so we can run tests in CI

[msaroufim]

Just a quick update, the main challenge for this PR will be correctly packaging setup.py with all the correct package versions and I'm seeing all sorts of basic issues even with basic models to setup some E2E test. Some relevant issues

• ❓ [Question] How to install "tensorrt" package?  TensorRT#1233 (comment)
• Warning: Checker does not support models with experimental ops: ATen ort#150
• Other alternatives like AITemplate still don't have an official release and only support Ampere (A100+) Any plans for official releases? facebookincubator/AITemplate#36

Is linking to specific versions enough? Should I use git submodules? When do I update them? How will docker support work?

EDIT:

The installation experience is now pip install torchserve[tensorrt] if remote or pip install .[tensorrt] which will install the necessary packages and the full set of options is onnx, tensorrt, ipex

[msaroufim]

And here's an example of an inference running from my logs. Repro is in test/pytest/test_onnx.py

(base) ubuntu@ip-172-31-17-70:~$ curl -X POST http://127.0.0.1:8080/predictions/onnx --data-binary '1'
[
  -0.28003010153770447
](base) ubuntu@ip-172-31-17-70:~$ 

[lxning]

this would potentially cause docker image (> 10g) too big if we want to support multi-platform. We need figure out a way to package for multi-platform support

[msaroufim]

That's true, I guess for dev dependencies that seems fine though. I can remove these if we disable the ONNX tests by default but I don't believe we should either

[HamidShojanazeri]

Thanks @msaroufim , I added some comments, but more generally I wonder if we have plans to move all the backend initialization logics + utilities into a backend specific files / utils. This would make the base_handler easier to maintain and more readable. Maybe its a good boot camp task?

[HamidShojanazeri]

are we going to replace the above Lines 120-137 with this function?

[msaroufim]

Yeah this is a typo, will fix

[msaroufim]

Discussed offline with Li and Hamid

To merge this PR

• Disable test by default in CI
• Remove requirements from developer dependencies
• Add documentation to describe how people should pre and post process with ONNX
• Update test

Future

• automation of pre and post processing
• Refactor docker image to pull in necessary runtime dependencies optionally
• Refactor base handler

[msaroufim]

@HamidShojanazeri @lxning I made all the changes we discussed

Logs

ubuntu@ip-172-31-17-70:~/serve/test/pytest$ curl -X POST http://127.0.0.1:8080/predictions/onnx --data-binary '1'
Prediction Succeeded

Pytest

(ort) ubuntu@ip-172-31-17-70:~/serve/test/pytest$ pytest test_onnx.py 
==================================== test session starts =====================================
platform linux -- Python 3.8.13, pytest-7.2.0, pluggy-1.0.0
rootdir: /home/ubuntu/serve
plugins: mock-3.10.0, cov-4.0.0
collected 5 items                                                                            

test_onnx.py .....                                                                     [100%]

===================================== 5 passed in 1.17s ======================================

Link check is failing because I'm linking to a test file that doesn't exist on master yet https://github.com/pytorch/serve/actions/runs/3425428486/jobs/5706209933#step:5:867 - this wont be a problem after we merge

[HamidShojanazeri]

Thanks @msaroufim , LGTM. I also suggest for next set of PRs maybe having linting/ formatting as a separate PR makes it easier to focus on the changes.