wehub-resource-sync
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
"""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)
|