fix: Upgrade TRT to 8.6.1, parallelize FX tests in CI (#1930)

Author: gs-olive

.circleci/config.yml

@@ -203,10 +203,10 @@ commands:
         default: "8.8.0.121"
       trt-version-short:
         type: string
-        default: "8.6.0"
+        default: "8.6.1"
       trt-version-long:
         type: string
-        default: "8.6.0.12-1"
+        default: "8.6.1.6-1"
       bazel-version:
         type: string
         default: "5.2.0"
@@ -249,7 +249,7 @@ commands:
       parameters:
         trt-version-long:
           type: string
-          default: "8.6.0"
+          default: "8.6.1"
         cudnn-version-long:
           type: string
           default: "8.8.0.121"
@@ -567,10 +567,10 @@ commands:
       - run:
           name: Run FX converter tests
           command: |
+            set -e
             cd py/torch_tensorrt/fx/test
-            pushd converters/acc_op/
-            pytest --junitxml=/tmp/artifacts/test_results/fx/converters/acc_op/test_results.xml
-            popd
+            TESTS_TO_RUN=$(circleci tests glob "converters/acc_op/test_*.py" | circleci tests split --split-by=timings)
+            pytest --junitxml=/tmp/artifacts/test_results/fx/converters/acc_op/test_results.xml $TESTS_TO_RUN
 
       - store_test_results:
           path: /tmp/artifacts
@@ -583,10 +583,10 @@ commands:
       - run:
           name: Run FX converter tests
           command: |
+            set -e
             cd py/torch_tensorrt/fx/test
-            pushd converters/aten_op/
-            pytest --junitxml=/tmp/artifacts/test_results/fx/converters/aten_op/test_results.xml
-            popd
+            TESTS_TO_RUN=$(circleci tests glob "converters/aten_op/test_*.py" | circleci tests split --split-by=timings)
+            pytest --junitxml=/tmp/artifacts/test_results/fx/converters/aten_op/test_results.xml $TESTS_TO_RUN
 
       - store_test_results:
           path: /tmp/artifacts
@@ -954,6 +954,7 @@ jobs:
         type: string
       python-version:
         type: string
+    parallelism: 8
     machine:
       image: linux-cuda-11:2023.02.1
     resource_class: gpu.nvidia.large
@@ -989,6 +990,7 @@ jobs:
         type: string
       python-version:
         type: string
+    parallelism: 8
     machine:
       image: linux-cuda-11:2023.02.1
     resource_class: gpu.nvidia.large
@@ -1374,10 +1376,10 @@ parameters:
     default: "8.8.0.121"
   trt-version-short:
     type: string
-    default: "8.6.0"
+    default: "8.6.1"
   trt-version-long:
     type: string
-    default: "8.6.0"
+    default: "8.6.1"
 
   # Jetson platform config
   torch-jetson-build:

README.md

@@ -119,7 +119,7 @@ These are the following dependencies used to verify the testcases. Torch-TensorR
 - Libtorch 2.1.0.dev20230419 (built with CUDA 11.8)
 - CUDA 11.8
 - cuDNN 8.8.0
-- TensorRT 8.6.0
+- TensorRT 8.6.1
 
 ## Prebuilt Binaries and Wheel files
 

WORKSPACE

@@ -81,10 +81,10 @@ http_archive(
 http_archive(
     name = "tensorrt",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "c1732a1093c57ab79fa0b687f061be369e449c9c17792b660f3663ecd8fa7b63",
-    strip_prefix = "TensorRT-8.6.0.12",
+    sha256 = "15bfe6053d45feec45ecc7123a9106076b0b43fa0435f242d89dca0778337759",
+    strip_prefix = "TensorRT-8.6.1.6",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.0/tars/TensorRT-8.6.0.12.Linux.x86_64-gnu.cuda-11.8.tar.gz",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/tars/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz",
     ],
 )
 

py/requirements.txt

@@ -5,4 +5,4 @@ pybind11==2.6.2
 torch==2.1.0.dev20230419+cu118
 torchvision==0.16.0.dev20230419+cu118
 --extra-index-url https://pypi.ngc.nvidia.com
-tensorrt==8.6.0
+tensorrt==8.6.1

py/torch_tensorrt/fx/README.md

@@ -12,7 +12,7 @@ FX2TRT is merged as FX module in Torch-TensorRT
     $ conda install pytorch torchvision torchtext cudatoolkit=11.8 -c pytorch-nightly
     # Install TensorRT python package
     $ pip3 install nvidia-pyindex
-    $ pip3 install tensorrt==8.6.0
+    $ pip3 install tensorrt==8.6.1
     $ git clone https://github.com/pytorch/TensorRT.git
     $ cd TensorRT/py && python setup.py install --fx-only && cd ..
     $ python -c "import torch_tensorrt.fx"