Add tests for examples (#149)

* Add tests for examples

* Ignore examples tests by default

* Trailing comma

* Add test for "quantizing_moe_fp8" example

* Update "quantizing_moe" example tests

* Add test for "compressed_inference" example folder

* Remove unused import

* Add new dependency

* Different approach for "flash_attn"

* Add comment about flash_attn requirement

* Test additional quantizing_moe example script

* Add decorator to skip based on available VRAM

* Limit GPU usage in 'cpu_offloading' example

* Add optional pytest-xdist parallelization

* Reduce persistent /tmp usage

* Fix parametrization in big_models

* Add pytest mark for GPU count requirement

* Add 'multi_gpu' pytest marker

* Skip 'deepseek_moe_w4a16.py' by default

* Fix skip mark

* Mark 'ex_trl_distillation.py' as multi_gpu

* Abstract command copy/run to helper functions

* Update for MoE examples PR #192

* Reduce test marker/decorator redundancy

* style fixes

* Rip out unused run parallelization

* Exclude 'deepseek_moe_w8a8_fp8' from multi-GPU

* Use variable for repeated string literal

* Use `requires_gpu_count` over `requires_gpu`

Makefile

@@ -12,6 +12,9 @@ endif
 ifneq ($(findstring pytorch,$(TARGETS)),pytorch)
     PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/llmcompressor/pytorch
 endif
+ifneq ($(findstring examples,$(TARGETS)),examples)
+    PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/examples
+endif
 
 # run checks on all files for the repo
 # leaving out mypy src for now

pyproject.toml

@@ -27,4 +27,7 @@ markers = [
     "regression: detailed tests to ensure major functions work correctly",
     "integration: tests which integrate with a third party service such as HF",
     "unit: tests to ensure code correctness and regression test functionality",
+    "example: tests for content in the 'examples' folder",
+    "multi_gpu: tests that require multiple GPUs",
 ]
+tmp_path_retention_policy = "failed"

setup.py

@@ -66,11 +66,15 @@
     ],
     extras_require={
         "dev": [
-            # testing
+            # testing framework
             "pytest>=6.0.0",
             "pytest-mock>=3.6.0",
             "pytest-rerunfailures>=13.0",
             "parameterized",
+            # example test dependencies
+            "beautifulsoup4~=4.12.3",
+            "cmarkgfm~=2024.1.14",
+            "trl>=0.10.1",
             # linting, formatting, and type checking
             "black~=24.4.2",
             "isort~=5.13.2",

tests/examples/test_big_models_with_accelerate.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    ReadMe,
+    copy_and_run_script,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+    requires_gpu_mem,
+)
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/big_models_with_accelerate"
+
+
+@pytest.mark.example
+class TestBigModelsWithAccelerate:
+    """
+    Tests for examples in the "big_models_with_accelerate" example folder.
+    """
+
+    def test_readme_has_install_command(self, example_dir: str):
+        """
+        Test that the README has a valid install command.
+        """
+        readme_path = Path.cwd() / example_dir / "README.md"
+        readme = ReadMe(readme_path)
+
+        code = readme.get_code_block_content(position=1, lang="shell")
+        assert "pip install" in code
+
+        assert code.startswith("pip install llmcompressor")
+
+    @pytest.mark.parametrize(
+        ("script_filename", "visible_gpus"),
+        [
+            pytest.param("cpu_offloading_fp8.py", "0", id="cpu_offloading"),
+            pytest.param(
+                "multi_gpu_int8.py",
+                "",
+                id="multi_gpu_int8",
+                marks=[
+                    requires_gpu_mem(630),
+                    requires_gpu_count(2),
+                    pytest.mark.multi_gpu,
+                ],
+            ),
+            pytest.param(
+                "multi_gpu_int8_sequential_update.py",
+                "",
+                id="multi_gpu_int8_sequential_update",
+                marks=[requires_gpu_count(2), pytest.mark.multi_gpu],
+            ),
+        ],
+    )
+    @requires_gpu_count(1)
+    def test_example_scripts(
+        self,
+        example_dir: str,
+        visible_gpus: str,
+        script_filename: str,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """
+        Test for the example scripts in the folder.
+        """
+
+        if visible_gpus:
+            monkeypatch.setenv("CUDA_VISIBLE_DEVICES", visible_gpus)
+
+        command, result = copy_and_run_script(tmp_path, example_dir, script_filename)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)

tests/examples/test_compressed_inference.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    copy_and_run_script,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+)
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/compressed_inference"
+
+
+@pytest.mark.example
+@requires_gpu_count(1)
+class TestCompressedInference:
+    """
+    Tests for examples in the "compressed_inference" example folder.
+    """
+
+    def test_fp8_example_script(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the "fp8_compressed_inference.py" script in the folder.
+        """
+        script_filename = "fp8_compressed_inference.py"
+        command, result = copy_and_run_script(tmp_path, example_dir, script_filename)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)

tests/examples/test_quantization_24_sparse_w4a16.py

@@ -0,0 +1,62 @@
+import shlex
+import shutil
+import sys
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    ReadMe,
+    copy_and_run_command,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+)
+from tests.testing_utils import run_cli_command
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/quantization_24_sparse_w4a16"
+
+
+@pytest.mark.example
+@requires_gpu_count(1)
+class TestQuantization24SparseW4A16:
+    """
+    Tests for examples in the "quantization_24_sparse_w4a16" example folder.
+    """
+
+    def test_doc_example_command(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the example command in the README.
+        """
+        readme_path = Path.cwd() / example_dir / "README.md"
+        readme = ReadMe(readme_path)
+
+        command = readme.get_code_block_content(position=2, lang="shell")
+        assert command.startswith("python")
+
+        command = shlex.split(command)
+        result = copy_and_run_command(tmp_path, example_dir, command)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)
+
+    def test_alternative_recipe(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the example command in the README with the alternative recipe file.
+        """
+        shutil.copytree(str(Path.cwd() / example_dir), tmp_path / example_dir)
+
+        # replace recipe file with alternative
+        script_filename = "llama7b_sparse_w4a16.py"
+        script_path = tmp_path / example_dir / script_filename
+        content = script_path.read_text(encoding="utf-8")
+        content = content.replace(
+            "2:4_w4a16_recipe.yaml", "2:4_w4a16_group-128_recipe.yaml"
+        )
+        script_path.write_text(content, encoding="utf-8")
+
+        command = [sys.executable, script_filename]
+        result = run_cli_command(command, cwd=tmp_path / example_dir)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)

tests/examples/test_quantization_kv_cache.py

@@ -0,0 +1,39 @@
+import shlex
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    ReadMe,
+    copy_and_run_command,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+)
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/quantization_kv_cache"
+
+
+@pytest.mark.example
+@requires_gpu_count(1)
+class TestQuantizationKVCache:
+    """
+    Tests for examples in the "quantization_kv_cache" example folder.
+    """
+
+    def test_doc_example_command(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the example command in the README.
+        """
+        readme_path = Path.cwd() / example_dir / "README.md"
+        readme = ReadMe(readme_path)
+
+        command = readme.get_code_block_content(position=2, lang="shell")
+        assert command.startswith("python")
+
+        command = shlex.split(command)
+        result = copy_and_run_command(tmp_path, example_dir, command)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)

tests/examples/test_quantization_w4a16.py

@@ -0,0 +1,39 @@
+import shlex
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    ReadMe,
+    copy_and_run_command,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+)
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/quantization_w4a16"
+
+
+@pytest.mark.example
+@requires_gpu_count(1)
+class TestQuantizationW4A16:
+    """
+    Tests for examples in the "quantization_w4a16" example folder.
+    """
+
+    def test_doc_example_command(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the example command in the README.
+        """
+        readme_path = Path.cwd() / example_dir / "README.md"
+        readme = ReadMe(readme_path)
+
+        command = readme.get_code_block_content(position=2, lang="shell")
+        assert command.startswith("python")
+
+        command = shlex.split(command)
+        result = copy_and_run_command(tmp_path, example_dir, command)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)

tests/examples/test_quantization_w8a8_fp8.py

@@ -0,0 +1,49 @@
+import shlex
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    ReadMe,
+    copy_and_run_command,
+    copy_and_run_script,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+)
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/quantization_w8a8_fp8"
+
+
+@pytest.mark.example
+@requires_gpu_count(1)
+class TestQuantizationW8A8_FP8:
+    """
+    Tests for examples in the "quantization_w8a8_fp8" example folder.
+    """
+
+    def test_doc_example_command(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the example command in the README.
+        """
+        readme_path = Path.cwd() / example_dir / "README.md"
+        readme = ReadMe(readme_path)
+
+        command = readme.get_code_block_content(position=2, lang="shell")
+        assert command.startswith("python")
+
+        command = shlex.split(command)
+        result = copy_and_run_command(tmp_path, example_dir, command)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)
+
+    def test_gemma2_example_script(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the "gemma2_example.py" script in the folder.
+        """
+        script_filename = "gemma2_example.py"
+        command, result = copy_and_run_script(tmp_path, example_dir, script_filename)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)

tests/examples/test_quantization_w8a8_int8.py

@@ -0,0 +1,49 @@
+import shlex
+from pathlib import Path
+
+import pytest
+
+from tests.examples.utils import (
+    ReadMe,
+    copy_and_run_command,
+    copy_and_run_script,
+    gen_cmd_fail_message,
+    requires_gpu_count,
+)
+
+
+@pytest.fixture
+def example_dir() -> str:
+    return "examples/quantization_w8a8_int8"
+
+
+@pytest.mark.example
+@requires_gpu_count(1)
+class TestQuantizationW8A8_Int8:
+    """
+    Tests for examples in the "quantization_w8a8_int8" example folder.
+    """
+
+    def test_doc_example_command(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the example command in the README.
+        """
+        readme_path = Path.cwd() / example_dir / "README.md"
+        readme = ReadMe(readme_path)
+
+        command = readme.get_code_block_content(position=2, lang="shell")
+        assert command.startswith("python")
+
+        command = shlex.split(command)
+        result = copy_and_run_command(tmp_path, example_dir, command)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)
+
+    def test_gemma2_example_script(self, example_dir: str, tmp_path: Path):
+        """
+        Test for the "gemma2_example.py" script in the folder.
+        """
+        script_filename = "gemma2_example.py"
+        command, result = copy_and_run_script(tmp_path, example_dir, script_filename)
+
+        assert result.returncode == 0, gen_cmd_fail_message(command, result)