Add triggering compile tests on diff

Author: qubvel

requirements/test.txt

@@ -1,5 +1,6 @@
+gitpython==3.1.44
 packaging==24.2
 pytest==8.3.4
 pytest-xdist==3.6.1
 pytest-cov==6.0.0
-ruff==0.9.1
+ruff==0.9.1

tests/conftest.py

@@ -3,13 +3,14 @@ def pytest_addoption(parser):
         "--non-marked-only", action="store_true", help="Run only non-marked tests"
     )
 
+
 def pytest_collection_modifyitems(config, items):
     if config.getoption("--non-marked-only"):
         non_marked_items = []
         for item in items:
             # Check if the test has no marks
             if not item.own_markers:
                 non_marked_items.append(item)
-        
+
         # Update the test collection to only include non-marked tests
         items[:] = non_marked_items

tests/encoders/base.py

@@ -4,12 +4,15 @@
 import segmentation_models_pytorch as smp
 
 from functools import lru_cache
-from tests.utils import default_device
+from tests.utils import default_device, check_run_test_on_diff_or_main
 
 
 class BaseEncoderTester(unittest.TestCase):
     encoder_names = []
 
+    # some tests might be slow, running them only on diff
+    files_for_diff = []
+
     # standard encoder configuration
     num_output_features = 6
     output_strides = [1, 2, 4, 8, 16, 32]
@@ -213,6 +216,9 @@ def test_dilated(self):
 
     @pytest.mark.compile
     def test_compile(self):
+        if not check_run_test_on_diff_or_main(self.files_for_diff):
+            self.skipTest("No diff and not on `main`.")
+
         sample = self._get_sample(
             batch_size=self.default_batch_size,
             num_channels=self.default_num_channels,

tests/encoders/test_pretrainedmodels_encoders.py

@@ -10,6 +10,7 @@ class TestDPNEncoder(base.BaseEncoderTester):
         if not RUN_ALL_ENCODERS
         else ["dpn68", "dpn68b", "dpn92", "dpn98", "dpn107", "dpn131"]
     )
+    files_for_diff = ["encoders/dpn.py"]
 
     def get_tiny_encoder(self):
         params = {
@@ -32,11 +33,13 @@ class TestInceptionResNetV2Encoder(base.BaseEncoderTester):
     encoder_names = (
         ["inceptionresnetv2"] if not RUN_ALL_ENCODERS else ["inceptionresnetv2"]
     )
+    files_for_diff = ["encoders/inceptionresnetv2.py"]
 
 
 class TestInceptionV4Encoder(base.BaseEncoderTester):
     supports_dilated = False
     encoder_names = ["inceptionv4"] if not RUN_ALL_ENCODERS else ["inceptionv4"]
+    files_for_diff = ["encoders/inceptionv4.py"]
 
 
 class TestSeNetEncoder(base.BaseEncoderTester):
@@ -52,6 +55,7 @@ class TestSeNetEncoder(base.BaseEncoderTester):
             # "senet154",  # extra large model
         ]
     )
+    files_for_diff = ["encoders/senet.py"]
 
     def get_tiny_encoder(self):
         params = {
@@ -73,3 +77,4 @@ def get_tiny_encoder(self):
 class TestXceptionEncoder(base.BaseEncoderTester):
     supports_dilated = False
     encoder_names = ["xception"] if not RUN_ALL_ENCODERS else ["xception"]
+    files_for_diff = ["encoders/xception.py"]

tests/encoders/test_smp_encoders.py

@@ -17,6 +17,7 @@ class TestMobileoneEncoder(base.BaseEncoderTester):
             "mobileone_s4",
         ]
     )
+    files_for_diff = ["encoders/mobileone.py"]
 
 
 class TestMixTransformerEncoder(base.BaseEncoderTester):
@@ -25,6 +26,7 @@ class TestMixTransformerEncoder(base.BaseEncoderTester):
         if not RUN_ALL_ENCODERS
         else ["mit_b0", "mit_b1", "mit_b2", "mit_b3", "mit_b4", "mit_b5"]
     )
+    files_for_diff = ["encoders/mix_transformer.py"]
 
     def get_tiny_encoder(self):
         params = {
@@ -59,6 +61,7 @@ class TestEfficientNetEncoder(base.BaseEncoderTester):
             # "efficientnet-b7",  # extra large model
         ]
     )
+    files_for_diff = ["encoders/efficientnet.py"]
 
     def test_compile(self):
         self.skipTest("compile fullgraph is not supported for efficientnet encoders")

tests/encoders/test_timm_ported_encoders.py

@@ -24,6 +24,7 @@ class TestTimmEfficientNetEncoder(base.BaseEncoderTester):
             "timm-tf_efficientnet_lite4",
         ]
     )
+    files_for_diff = ["encoders/timm_efficientnet.py"]
 
 
 class TestTimmGERNetEncoder(base.BaseEncoderTester):
@@ -144,3 +145,4 @@ class TestTimmSkNetEncoder(base.BaseEncoderTester):
             "timm-skresnext50_32x4d",
         ]
     )
+    files_for_diff = ["encoders/timm_sknet.py"]

tests/encoders/test_timm_universal.py

@@ -14,3 +14,4 @@
 
 class TestTimmUniversalEncoder(base.BaseEncoderTester):
     encoder_names = timm_encoders
+    files_for_diff = ["encoders/timm_universal.py"]

tests/encoders/test_torchvision_encoders.py

@@ -22,6 +22,7 @@ class TestResNetEncoder(base.BaseEncoderTester):
             "resnext101_32x48d",
         ]
     )
+    files_for_diff = ["encoders/resnet.py"]
 
     def get_tiny_encoder(self):
         params = {
@@ -39,6 +40,7 @@ class TestDenseNetEncoder(base.BaseEncoderTester):
         if not RUN_ALL_ENCODERS
         else ["densenet121", "densenet169", "densenet161"]
     )
+    files_for_diff = ["encoders/densenet.py"]
 
     def get_tiny_encoder(self):
         params = {
@@ -52,6 +54,7 @@ def get_tiny_encoder(self):
 
 class TestMobileNetEncoder(base.BaseEncoderTester):
     encoder_names = ["mobilenet_v2"] if not RUN_ALL_ENCODERS else ["mobilenet_v2"]
+    files_for_diff = ["encoders/mobilenet.py"]
 
 
 class TestVggEncoder(base.BaseEncoderTester):
@@ -70,6 +73,7 @@ class TestVggEncoder(base.BaseEncoderTester):
             "vgg19_bn",
         ]
     )
+    files_for_diff = ["encoders/vgg.py"]
 
     def get_tiny_encoder(self):
         params = {

tests/utils.py

@@ -1,31 +1,21 @@
 import os
+import re
 import timm
 import torch
 import unittest
 
+from git import Repo
+from typing import List
 from packaging.version import Version
 
 
 has_timm_test_models = Version(timm.__version__) >= Version("1.0.12")
 default_device = "cuda" if torch.cuda.is_available() else "cpu"
 
-
-def get_commit_message():
-    commit_msg = os.getenv("COMMIT_MESSAGE", "")
-    return commit_msg.lower()
-
-
-# Check both environment variables and commit message
-commit_message = get_commit_message()
-RUN_ALL_ENCODERS = (
-    os.getenv("RUN_ALL_ENCODERS", "false").lower() in ["true", "1", "y", "yes"]
-    or "run-all-encoders" in commit_message
-)
-
-RUN_SLOW = (
-    os.getenv("RUN_SLOW", "false").lower() in ["true", "1", "y", "yes"]
-    or "run-slow" in commit_message
-)
+YES_LIST = ["true", "1", "y", "yes"]
+RUN_ALL_ENCODERS = os.getenv("RUN_ALL_ENCODERS", "false").lower() in YES_LIST
+RUN_SLOW = os.getenv("RUN_SLOW", "false").lower() in YES_LIST
+RUN_ALL = os.getenv("RUN_ALL", "false").lower() in YES_LIST
 
 
 def slow_test(test_case):
@@ -45,3 +35,26 @@ def requires_torch_greater_or_equal(version: str):
         torch_version >= provided_version,
         f"torch version {torch_version} is less than {provided_version}",
     )
+
+
+def check_run_test_on_diff_or_main(filepath_patterns: List[str]):
+    if RUN_ALL:
+        return True
+
+    try:
+        repo = Repo(".")
+        current_branch = repo.active_branch.name
+        diff_files = repo.git.diff("main", name_only=True).splitlines()
+
+    except Exception:
+        return True
+
+    if current_branch == "main":
+        return True
+
+    for pattern in filepath_patterns:
+        for file_path in diff_files:
+            if re.search(pattern, file_path):
+                return True
+
+    return False