| |
|
| | '''
|
| | This script fetches all the models used in the server tests.
|
| |
|
| | This is useful for slow tests that use larger models, to avoid them timing out on the model downloads.
|
| |
|
| | It is meant to be run from the root of the repository.
|
| |
|
| | Example:
|
| | python scripts/fetch_server_test_models.py
|
| | ( cd tools/server/tests && ./tests.sh -v -x -m slow )
|
| | '''
|
| | import ast
|
| | import glob
|
| | import logging
|
| | import os
|
| | from typing import Generator
|
| | from pydantic import BaseModel
|
| | from typing import Optional
|
| | import subprocess
|
| |
|
| |
|
| | class HuggingFaceModel(BaseModel):
|
| | hf_repo: str
|
| | hf_file: Optional[str] = None
|
| |
|
| | class Config:
|
| | frozen = True
|
| |
|
| |
|
| | def collect_hf_model_test_parameters(test_file) -> Generator[HuggingFaceModel, None, None]:
|
| | try:
|
| | with open(test_file) as f:
|
| | tree = ast.parse(f.read())
|
| | except Exception as e:
|
| | logging.error(f'collect_hf_model_test_parameters failed on {test_file}: {e}')
|
| | return
|
| |
|
| | for node in ast.walk(tree):
|
| | if isinstance(node, ast.FunctionDef):
|
| | for dec in node.decorator_list:
|
| | if isinstance(dec, ast.Call) and isinstance(dec.func, ast.Attribute) and dec.func.attr == 'parametrize':
|
| | param_names = ast.literal_eval(dec.args[0]).split(",")
|
| | if "hf_repo" not in param_names:
|
| | continue
|
| |
|
| | raw_param_values = dec.args[1]
|
| | if not isinstance(raw_param_values, ast.List):
|
| | logging.warning(f'Skipping non-list parametrize entry at {test_file}:{node.lineno}')
|
| | continue
|
| |
|
| | hf_repo_idx = param_names.index("hf_repo")
|
| | hf_file_idx = param_names.index("hf_file") if "hf_file" in param_names else None
|
| |
|
| | for t in raw_param_values.elts:
|
| | if not isinstance(t, ast.Tuple):
|
| | logging.warning(f'Skipping non-tuple parametrize entry at {test_file}:{node.lineno}')
|
| | continue
|
| | yield HuggingFaceModel(
|
| | hf_repo=ast.literal_eval(t.elts[hf_repo_idx]),
|
| | hf_file=ast.literal_eval(t.elts[hf_file_idx]) if hf_file_idx is not None else None)
|
| |
|
| |
|
| | if __name__ == '__main__':
|
| | logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
|
| |
|
| | models = sorted(list(set([
|
| | model
|
| | for test_file in glob.glob('tools/server/tests/unit/test_*.py')
|
| | for model in collect_hf_model_test_parameters(test_file)
|
| | ])), key=lambda m: (m.hf_repo, m.hf_file))
|
| |
|
| | logging.info(f'Found {len(models)} models in parameterized tests:')
|
| | for m in models:
|
| | logging.info(f' - {m.hf_repo} / {m.hf_file}')
|
| |
|
| | cli_path = os.environ.get(
|
| | 'LLAMA_CLI_BIN_PATH',
|
| | os.path.join(
|
| | os.path.dirname(__file__),
|
| | '../build/bin/Release/llama-cli.exe' if os.name == 'nt' else '../build/bin/llama-cli'))
|
| |
|
| | for m in models:
|
| | if '<' in m.hf_repo or (m.hf_file is not None and '<' in m.hf_file):
|
| | continue
|
| | if m.hf_file is not None and '-of-' in m.hf_file:
|
| | logging.warning(f'Skipping model at {m.hf_repo} / {m.hf_file} because it is a split file')
|
| | continue
|
| | logging.info(f'Using llama-cli to ensure model {m.hf_repo}/{m.hf_file} was fetched')
|
| | cmd = [
|
| | cli_path,
|
| | '-hfr', m.hf_repo,
|
| | *([] if m.hf_file is None else ['-hff', m.hf_file]),
|
| | '-n', '1',
|
| | '-p', 'Hey',
|
| | '--no-warmup',
|
| | '--log-disable',
|
| | '-no-cnv']
|
| | if m.hf_file != 'tinyllamas/stories260K.gguf' and 'Mistral-Nemo' not in m.hf_repo:
|
| | cmd.append('-fa')
|
| | try:
|
| | subprocess.check_call(cmd)
|
| | except subprocess.CalledProcessError:
|
| | logging.error(f'Failed to fetch model at {m.hf_repo} / {m.hf_file} with command:\n {" ".join(cmd)}')
|
| | exit(1)
|
| |
|
