huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-27 21:12:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(neuron): add helper to batch export models

Browse Source 
Also rename fixture file fro clarity.
This commit is contained in:
David Corvoysier 2025-02-24 17:37:37 +00:00
parent 70e846d53b
commit 53c1226939
2 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
integration-tests/conftest.py
View File

@ -1,4 +1,4 @@
pytest_plugins = ["fixtures.neuron.service", "fixtures.neuron.model"]
pytest_plugins = ["fixtures.neuron.service", "fixtures.neuron.export_models"]
# ruff: noqa: E402
from _pytest.fixtures import SubRequest
import requests

9
integration-tests/fixtures/neuron/model.py → integration-tests/fixtures/neuron/export_models.py
View File

@ -216,6 +216,11 @@ def maybe_export_model(config_name, model_config):
return neuron_model_id
def maybe_export_models():
for config_name, model_config in MODEL_CONFIGURATIONS.items():
maybe_export_model(config_name, model_config)
@pytest.fixture(scope="session", params=MODEL_CONFIGURATIONS.keys())
def neuron_model_config(request):
"""Expose a pre-trained neuron model
@ -262,3 +267,7 @@ def neuron_model_config(request):
@pytest.fixture(scope="module")
def neuron_model_path(neuron_model_config):
yield neuron_model_config["neuron_model_path"]
if __name__ == "__main__":
maybe_export_models()