add fail fast actions

Author: lanluo-nvidia

.github/workflows/build-test-linux-x86_64.yml

@@ -75,9 +75,50 @@ jobs:
       smoke-test-script: ${{ matrix.smoke-test-script }}
       trigger-event: ${{ github.event_name }}
 
+  tests-py-fail-fast:
+    name: Test fail fast [Python]
+    needs: [filter-matrix, build]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/tensorrt
+            package-name: torch_tensorrt
+            pre-script: packaging/pre_build_script.sh
+            post-script: packaging/post_build_script.sh
+            smoke-test-script: packaging/smoke_test_script.sh
+    uses: ./.github/workflows/linux-test.yml
+    with:
+      job-name: tests-py-critical-fail-fast
+      repository: "pytorch/tensorrt"
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.filter-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      script: |
+        export USE_HOST_DEPS=1
+        export CI_BUILD=1
+        export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH
+        pushd .
+        cd tests/py
+        python -m pip install -r requirements.txt
+
+        # test dynamo
+        python -m pytest -m critical --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_backend_test_results.xml -n 4 dynamo/backend/
+        python -m pytest -m critical -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_models_result.xml --ir dynamo dynamo/models/
+        python -m pytest -m critical --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_automatic_plugin_results.xml dynamo/automatic_plugin/
+        python -m pytest -m critical --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_partitioning_results.xml dynamo/partitioning/
+        python -m pytest -m critical --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_lowering_results.xml dynamo/lowering/
+        python -m pytest -m critical --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_runtime_results.xml dynamo/runtime/
+
+        # test core
+        python -m pytest -m critical --junitxml=${RUNNER_TEST_RESULTS_DIR}/core_test_results.xml core/
+        popd
+
   tests-py-torchscript-fe:
     name: Test torchscript frontend [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:
@@ -113,7 +154,7 @@ jobs:
 
   tests-py-dynamo-converters:
     name: Test dynamo converters [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:
@@ -148,7 +189,7 @@ jobs:
 
   tests-py-dynamo-fe:
     name: Test dynamo frontend [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:
@@ -180,7 +221,7 @@ jobs:
 
   tests-py-dynamo-serde:
     name: Test dynamo export serde [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:
@@ -245,7 +286,7 @@ jobs:
 
   tests-py-dynamo-core:
     name: Test dynamo core [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:
@@ -278,7 +319,7 @@ jobs:
 
   tests-py-dynamo-cudagraphs:
     name: Test dynamo cudagraphs [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:
@@ -311,7 +352,7 @@ jobs:
 
   tests-py-core:
     name: Test core [Python]
-    needs: [filter-matrix, build]
+    needs: [filter-matrix, build, tests-py-fail-fast]
     strategy:
       fail-fast: false
       matrix:

tests/py/core/test_classes.py

@@ -2,16 +2,17 @@
 import unittest
 from typing import Dict
 
+import pytest
+import tensorrt as trt
 import torch
 import torch_tensorrt
 import torch_tensorrt as torchtrt
 import torchvision.models as models
 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule
 
-import tensorrt as trt
-
 
 class TestDevice(unittest.TestCase):
+    @pytest.mark.critical
     def test_from_string_constructor(self):
         device = torchtrt.Device("cuda:0")
         self.assertEqual(device.device_type, torchtrt.DeviceType.GPU)

tests/py/dynamo/automatic_plugin/test_automatic_plugin.py

@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import pytest
 import torch
 import torch.nn as nn
 import torch_tensorrt
@@ -66,6 +67,7 @@ class TestAutomaticPlugin(DispatchTestCase):
             ((256, 256), torch.int),
         ]
     )
+    @pytest.mark.critical
     def test_mul_plugin_float(self, input_shape, dtype):
         class elementwise_mul(nn.Module):
             def forward(self, lhs, rhs):

tests/py/dynamo/backend/test_backend_compiler.py

@@ -1,6 +1,7 @@
 # type: ignore
 from copy import deepcopy
 
+import pytest
 import torch
 import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase, run_tests
@@ -10,6 +11,7 @@
 
 
 class TestTRTModuleNextCompilation(TestCase):
+    @pytest.mark.critical
     def test_trt_module_next_full_support(self):
         class FullySupportedMultiOp(torch.nn.Module):
             def forward(self, x, y):

tests/py/dynamo/distributed/test_nccl_ops.py

@@ -1,5 +1,6 @@
 import os
 
+import pytest
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -17,6 +18,7 @@
 
 
 class TestGatherNcclOpsConverter(DispatchTestCase):
+    @pytest.mark.critical
     @parameterized.expand([8])
     def test_nccl_ops(self, linear_layer_dim):
         class DistributedGatherModel(nn.Module):

tests/py/dynamo/lowering/test_aten_lowering_passes.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase, run_tests
@@ -6,6 +7,7 @@
 
 
 class TestInputAsOutput(TestCase):
+    @pytest.mark.critical
     def test_input_as_output(self):
         class InputAsOutput(torch.nn.Module):
             def forward(self, x, y):
@@ -56,6 +58,7 @@ def forward(self, x, y):
 
 
 class TestLoweringPassMembership(TestCase):
+    @pytest.mark.critical
     def insert_at_end(self):
         from torch_tensorrt.dynamo.lowering.passes import (
             ATEN_LOWERING_PASSES,

tests/py/dynamo/lowering/test_decompositions.py

@@ -14,6 +14,7 @@
 
 
 class TestLowering(TestCase):
+    @pytest.mark.critical
     def test_lowering_inplace_op(self):
         class InPlace(torch.nn.Module):
             def __init__(self, *args, **kwargs) -> None:

tests/py/dynamo/models/test_dyn_models.py

@@ -11,6 +11,7 @@
 assertions = unittest.TestCase()
 
 
+@pytest.mark.critical
 @pytest.mark.unit
 def test_base_dynamic(ir):
     """
@@ -175,6 +176,7 @@ def forward(self, x):
     )
 
 
+@pytest.mark.critical
 @pytest.mark.unit
 def test_resnet_dynamic(ir):
     """

tests/py/dynamo/models/test_engine_cache.py

@@ -57,6 +57,7 @@ def load(self, hash: str, prefix: str = "blob") -> Optional[bytes]:
 
 
 class TestHashFunction(TestCase):
+    @pytest.mark.critical
     def test_reexport_is_equal(self):
         pyt_model = models.resnet18(pretrained=True).eval().to("cuda")
         example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -23,6 +23,7 @@
 
 
 @pytest.mark.unit
+@pytest.mark.critical
 def test_custom_model():
     class net(nn.Module):
         def __init__(self):
@@ -83,6 +84,7 @@ def forward(self, x, b=5, c=None, d=None):
 
 
 @pytest.mark.unit
+@pytest.mark.critical
 def test_custom_model_with_dynamo_trace():
     class net(nn.Module):
         def __init__(self):

tests/py/dynamo/models/test_export_serde.py

@@ -17,6 +17,7 @@
 trt_ep_path = os.path.join(tempfile.gettempdir(), "trt.ep")
 
 
+@pytest.mark.critical
 @pytest.mark.unit
 def test_base_full_compile(ir):
     """

tests/py/dynamo/models/test_models.py

@@ -85,6 +85,7 @@ def test_resnet18_cpu_offload(ir):
 
 
 @pytest.mark.unit
+@pytest.mark.critical
 def test_mobilenet_v2(ir):
     model = models.mobilenet_v2(pretrained=True).eval().to("cuda")
     input = torch.randn((1, 3, 224, 224)).to("cuda")

tests/py/dynamo/models/test_models_export.py

@@ -81,6 +81,7 @@ def test_mobilenet_v2(ir):
 
 
 @pytest.mark.unit
+@pytest.mark.critical
 def test_efficientnet_b0(ir):
     model = timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda")
     input = torch.randn((1, 3, 224, 224)).to("cuda")

tests/py/dynamo/partitioning/test_fast_partitioning.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 
 import numpy as np
+import pytest
 import torch
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch_tensorrt.dynamo import partitioning
@@ -55,6 +56,7 @@ def forward(self, x, y):
             "Single operators can be segmented if full compilation is required",
         )
 
+    @pytest.mark.critical
     def test_partition_fully_supported_multi_op(self):
         class FullySupportedMultiOp(torch.nn.Module):
             def __init__(self, *args, **kwargs) -> None:

tests/py/dynamo/partitioning/test_global_partitioning.py

@@ -95,6 +95,7 @@ def forward(self, x, y):
             "Single operators can be segmented if full compilation is required",
         )
 
+    @pytest.mark.critical
     def test_partition_fully_supported_multi_op(self):
         class FullySupportedMultiOp(torch.nn.Module):
             def __init__(self, *args, **kwargs) -> None:

tests/py/dynamo/runtime/test_000_convert_module_to_trt_engine.py

@@ -1,14 +1,15 @@
 import unittest
 
+import pytest
+import tensorrt as trt
 import torch
 import torch_tensorrt
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule
 from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
 
-import tensorrt as trt
-
 
 class TestConvertModuleToTrtEngine(unittest.TestCase):
+    @pytest.mark.critical
     def test_convert_module(self):
         class Test(torch.nn.Module):
             def forward(self, a, b):

tests/py/dynamo/runtime/test_output_allocator.py

@@ -48,6 +48,7 @@ class TestOutputAllocatorStaticModel(TestCase):
             ("cpp_runtime", False),
         ]
     )
+    @pytest.mark.critical
     def test_cudagraphs_and_output_allocator(self, _, use_python_runtime):
         model = StaticModel().eval().cuda()
         inputs = [torch.randn((2, 3), dtype=torch.float).cuda()]
@@ -157,6 +158,7 @@ class TestOutputAllocatorDDSModel(TestCase):
             ("cpp_runtime", False),
         ]
     )
+    @pytest.mark.critical
     def test_cudagraphs_and_output_allocator(self, _, use_python_runtime):
         model = DDSModel().eval().cuda()
         inputs = (torch.randint(low=0, high=3, size=(10,), dtype=torch.int).to("cuda"),)

tests/py/dynamo/runtime/test_pre_allocated_outputs.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 import torch_tensorrt as torchtrt
 from parameterized import parameterized
@@ -14,6 +15,7 @@ class TestPreAllocatedOutputs(TestCase):
             ("cpp_runtime", False),
         ]
     )
+    @pytest.mark.critical
     def test_pre_allocated_outputs_default(self, _, use_python_runtime):
         class SampleModel(torch.nn.Module):
             def forward(self, x):