Refactor to standardize pytest usage (#1844)

Author: fynnsu

setup.py

@@ -147,7 +147,6 @@ def localversion_func(version: ScmVersion) -> str:
             "pytest>=6.0.0",
             "pytest-mock>=3.6.0",
             "pytest-rerunfailures>=13.0",
-            "parameterized",
             "lm_eval==0.4.5",
             # test dependencies
             "beautifulsoup4~=4.12.3",

tests/custom_test.py

@@ -10,21 +10,14 @@
 from bs4 import BeautifulSoup, ResultSet, Tag
 from cmarkgfm import github_flavored_markdown_to_html as gfm_to_html
 
-from tests.testing_utils import run_cli_command
+from tests.testing_utils import requires_gpu, run_cli_command
 
 _T = TypeVar("_T")
 
 
 def requires_gpu_count(num_required_gpus: int) -> pytest.MarkDecorator:
-    """
-    Pytest decorator to skip based on number of available GPUs. This plays nicely with
-    the CUDA_VISIBLE_DEVICES environment variable.
-    """
-    import torch
-
-    num_gpus = torch.cuda.device_count()
-    reason = f"{num_required_gpus} GPUs required, {num_gpus} GPUs detected"
-    return pytest.mark.skipif(num_required_gpus > num_gpus, reason=reason)
+    # Remove after #1801
+    return requires_gpu(num_required_gpus)
 
 
 def requires_gpu_mem(required_amount: Union[int, float]) -> pytest.MarkDecorator:

tests/data.py

@@ -5,6 +5,8 @@
 
 import pytest
 
+from llmcompressor.modifiers.factory import ModifierFactory
+
 try:
     import wandb
 except Exception:
@@ -15,6 +17,12 @@
 os.environ["NM_TEST_LOG_DIR"] = "nm_temp_test_logs"
 
 
+@pytest.fixture
+def setup_modifier_factory():
+    ModifierFactory.refresh()
+    assert ModifierFactory._loaded, "ModifierFactory not loaded"
+
+
 def _get_files(directory: str) -> List[str]:
     list_filepaths = []
     for root, dirs, files in os.walk(directory):

tests/examples/utils.py

@@ -6,15 +6,12 @@
 from llmcompressor.modifiers.awq import AWQMapping, AWQModifier
 from llmcompressor.modifiers.awq.base import get_lowest_common_parent
 from llmcompressor.modifiers.factory import ModifierFactory
-from tests.llmcompressor.modifiers.conf import setup_modifier_factory
 
 
 @pytest.mark.unit
+@pytest.mark.usefixtures("setup_modifier_factory")
 def test_awq_is_registered():
     """Ensure AWQModifier is registered in ModifierFactory"""
-
-    setup_modifier_factory()
-
     modifier = ModifierFactory.create(
         type_="AWQModifier",
         allow_experimental=False,

tests/llmcompressor/conftest.py

@@ -3,12 +3,6 @@
 from torch.utils.data import DataLoader
 
 from llmcompressor.core import Event, EventType, State
-from llmcompressor.modifiers.factory import ModifierFactory
-
-
-def setup_modifier_factory():
-    ModifierFactory.refresh()
-    assert ModifierFactory._loaded, "ModifierFactory not loaded"
 
 
 class LifecyleTestingHarness:

tests/llmcompressor/modifiers/awq/test_base.py

@@ -1,38 +1,28 @@
-import unittest
-
 import pytest
 
 from llmcompressor.modifiers.factory import ModifierFactory
 from llmcompressor.modifiers.logarithmic_equalization.base import (
     LogarithmicEqualizationModifier,
 )
 from llmcompressor.modifiers.smoothquant.base import SmoothQuantModifier
-from tests.llmcompressor.modifiers.conf import setup_modifier_factory
 
 
 @pytest.mark.unit
-class TestLogarithmicEqualizationIsRegistered(unittest.TestCase):
-    def setUp(self):
-        self.kwargs = dict(
-            smoothing_strength=0.3,
-            mappings=[(["layer1", "layer2"], "layer3")],
-        )
-        setup_modifier_factory()
-
-    def test_log_equalization_is_registered(self):
-        modifier = ModifierFactory.create(
-            type_="LogarithmicEqualizationModifier",
-            allow_experimental=False,
-            allow_registered=True,
-            **self.kwargs,
-        )
-
-        self.assertIsInstance(
-            modifier,
-            LogarithmicEqualizationModifier,
-            "PyTorch LogarithmicEqualizationModifier not registered",
-        )
+@pytest.mark.usefixtures("setup_modifier_factory")
+def test_logarithmic_equalization_is_registered():
+    smoothing_strength = 0.3
+    mappings = [(["layer1", "layer2"], "layer3")]
+    modifier = ModifierFactory.create(
+        type_="LogarithmicEqualizationModifier",
+        allow_experimental=False,
+        allow_registered=True,
+        smoothing_strength=smoothing_strength,
+        mappings=mappings,
+    )
 
-        self.assertIsInstance(modifier, SmoothQuantModifier)
-        self.assertEqual(modifier.smoothing_strength, self.kwargs["smoothing_strength"])
-        self.assertEqual(modifier.mappings, self.kwargs["mappings"])
+    assert isinstance(
+        modifier, LogarithmicEqualizationModifier
+    ), "PyTorch LogarithmicEqualizationModifier not registered"
+    assert isinstance(modifier, SmoothQuantModifier)
+    assert modifier.smoothing_strength == smoothing_strength
+    assert modifier.mappings == mappings

tests/llmcompressor/modifiers/conf.py

@@ -1,31 +1,22 @@
-import unittest
-
 import pytest
 
 from llmcompressor.modifiers.factory import ModifierFactory
 from llmcompressor.modifiers.obcq.base import SparseGPTModifier
-from tests.llmcompressor.modifiers.conf import setup_modifier_factory
 
 
 @pytest.mark.unit
-class TestSparseGPTIsRegistered(unittest.TestCase):
-    def setUp(self):
-        self.kwargs = dict(
-            sparsity=0.5,
-            targets="__ALL_PRUNABLE__",
-        )
-        setup_modifier_factory()
-
-    def test_wanda_is_registered(self):
-        type_ = ModifierFactory.create(
-            type_="SparseGPTModifier",
-            allow_experimental=False,
-            allow_registered=True,
-            **self.kwargs,
-        )
+@pytest.mark.usefixtures("setup_modifier_factory")
+def test_sparse_gpt_is_registered():
+    sparsity = 0.5
+    targets = "__ALL_PRUNABLE__"
+    type_ = ModifierFactory.create(
+        type_="SparseGPTModifier",
+        allow_experimental=False,
+        allow_registered=True,
+        sparsity=sparsity,
+        targets=targets,
+    )
 
-        self.assertIsInstance(
-            type_,
-            SparseGPTModifier,
-            "PyTorch SparseGPTModifier not registered",
-        )
+    assert isinstance(
+        type_, SparseGPTModifier
+    ), "PyTorch SparseGPTModifier not registered"

tests/llmcompressor/modifiers/logarithmic_equalization/test_base.py

@@ -1,31 +1,22 @@
-import unittest
-
 import pytest
 
 from llmcompressor.modifiers.factory import ModifierFactory
 from llmcompressor.modifiers.pruning.wanda.base import WandaPruningModifier
-from tests.llmcompressor.modifiers.conf import setup_modifier_factory
 
 
 @pytest.mark.unit
-class TestWandaIsRegistered(unittest.TestCase):
-    def setUp(self):
-        self.kwargs = dict(
-            sparsity=0.5,
-            targets="__ALL_PRUNABLE__",
-        )
-        setup_modifier_factory()
-
-    def test_wanda_is_registered(self):
-        type_ = ModifierFactory.create(
-            type_="WandaPruningModifier",
-            allow_experimental=False,
-            allow_registered=True,
-            **self.kwargs,
-        )
+@pytest.mark.usefixtures("setup_modifier_factory")
+def test_wanda_is_registered():
+    sparsity = 0.5
+    targets = "__ALL_PRUNABLE__"
+    type_ = ModifierFactory.create(
+        type_="WandaPruningModifier",
+        allow_experimental=False,
+        allow_registered=True,
+        sparsity=sparsity,
+        targets=targets,
+    )
 
-        self.assertIsInstance(
-            type_,
-            WandaPruningModifier,
-            "PyTorch WandaPruningModifier not registered",
-        )
+    assert isinstance(
+        type_, WandaPruningModifier
+    ), "PyTorch WandaPruningModifier not registered"