chore: import upstream snapshot with attribution

python/mlc_llm/testing/pytest_utils.py

@@ -0,0 +1,86 @@
+"""Extra utilities to mark tests"""
+
+import functools
+import inspect
+from pathlib import Path
+from typing import Callable
+
+import pytest
+
+from mlc_llm.support.constants import MLC_TEST_MODEL_PATH
+
+
+def require_test_model(*models: str):
+    """Testcase decorator to require a model
+
+    Examples
+    --------
+    .. code::
+
+        @require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+        def test_reload_reset_unload(model):
+            # model now points to the right path
+            # specified by MLC_TEST_MODEL_PATH
+            engine = mlc_llm.MLCEngine(model)
+            # test code follows
+
+    Parameters
+    ----------
+    models : List[str]
+        The model directories or URLs.
+    """
+    model_paths = []
+    missing_models = []
+
+    for model in models:
+        model_path = None
+        for base_path in MLC_TEST_MODEL_PATH:
+            if (base_path / model / "mlc-chat-config.json").is_file():
+                model_path = base_path / model
+                break
+        if model_path is None and (Path(model) / "mlc-chat-config.json").is_file():
+            model_path = Path(model)
+
+        if model_path is None:
+            missing_models.append(model)
+        else:
+            model_paths.append(str(model_path))
+
+    message = (
+        f"Model {', '.join(missing_models)} not found in candidate paths "
+        f"{[str(p) for p in MLC_TEST_MODEL_PATH]},"
+        " if you set MLC_TEST_MODEL_PATH, please ensure model paths are in the right location,"
+        " by default we reuse cache, try to run mlc_llm chat to download right set of models."
+    )
+
+    def _decorator(func: Callable[..., None]):
+        wrapped = functools.partial(func, *model_paths)
+        wrapped.__name__ = func.__name__
+
+        if inspect.iscoroutinefunction(wrapped):
+            # The function is a coroutine function ("async def func(...)")
+            @functools.wraps(wrapped)
+            async def wrapper(*args, **kwargs):
+                if len(missing_models) > 0:
+                    print(f"{message} skipping...")
+                    return
+                await wrapped(*args, **kwargs)
+
+        else:
+            # The function is a normal function ("def func(...)")
+            @functools.wraps(wrapped)
+            def wrapper(*args, **kwargs):
+                if len(missing_models) > 0:
+                    print(f"{message} skipping...")
+                    return
+                wrapped(*args, **kwargs)
+
+        return pytest.mark.skipif(len(missing_models) > 0, reason=message)(wrapper)
+
+    return _decorator
+
+
+def require_test_tokenizers(*models: str):
+    """Testcase decorator to require a path to tokenizers"""
+    # redirect to require models for now
+    return require_test_model(*models)