[v1] model loader (#9613)

This commit is contained in:
Yaowei Zheng
2025-12-14 11:50:52 +08:00
committed by GitHub
parent fdd24276ed
commit aeda079014
27 changed files with 449 additions and 305 deletions

View File

@@ -18,6 +18,7 @@ Contains shared fixtures, pytest configuration, and custom markers.
"""
import pytest
from pytest import Config, Item
from llamafactory.extras.misc import get_current_device, is_env_enabled
from llamafactory.train.test_utils import patch_valuehead_model
@@ -29,21 +30,15 @@ except Exception:
CURRENT_DEVICE = "cpu"
def pytest_configure(config):
def pytest_configure(config: Config):
"""Register custom pytest markers."""
config.addinivalue_line(
"markers", "slow: marks tests as slow (deselect with '-m \"not slow\"' or set RUN_SLOW=1 to run)"
)
config.addinivalue_line(
"markers", "skip_on_devices: skip test on specified devices, e.g., @pytest.mark.skip_on_devices('npu', 'xpu')"
)
config.addinivalue_line(
"markers", "require_device: test requires specific device, e.g., @pytest.mark.require_device('cuda')"
)
config.addinivalue_line("markers", "runs_on: test requires specific device, e.g., @pytest.mark.runs_on(['cpu'])")
def _handle_runs_on(items):
def _handle_runs_on(items: list[Item]):
"""Skip tests on specified devices based on runs_on marker.
Usage:
@@ -68,7 +63,7 @@ def _handle_runs_on(items):
)
def _handle_slow_tests(items):
def _handle_slow_tests(items: list[Item]):
"""Skip slow tests unless RUN_SLOW environment variable is set.
Usage:
@@ -85,51 +80,9 @@ def _handle_slow_tests(items):
item.add_marker(skip_slow)
def _handle_device_skips(items):
"""Skip tests on specified devices based on skip_on_devices marker.
Usage:
@pytest.mark.skip_on_devices("npu", "xpu")
def test_something():
pass
"""
for item in items:
skip_marker = item.get_closest_marker("skip_on_devices")
if skip_marker:
skip_devices = skip_marker.args
if CURRENT_DEVICE in skip_devices:
item.add_marker(
pytest.mark.skip(
reason=f"test skipped on {CURRENT_DEVICE.upper()} (skip list: {', '.join(skip_devices)})"
)
)
def _handle_device_requirements(items):
"""Skip tests that require a specific device when running on other devices.
Usage:
@pytest.mark.require_device("cuda")
def test_gpu_only():
pass
"""
for item in items:
require_marker = item.get_closest_marker("require_device")
if require_marker:
required_device = require_marker.args[0] if require_marker.args else None
if required_device and CURRENT_DEVICE != required_device:
item.add_marker(
pytest.mark.skip(
reason=f"test requires {required_device.upper()} (current: {CURRENT_DEVICE.upper()})"
)
)
def pytest_collection_modifyitems(config, items):
def pytest_collection_modifyitems(config: Config, items: list[Item]):
"""Modify test collection based on markers and environment."""
_handle_slow_tests(items)
_handle_device_skips(items)
_handle_device_requirements(items)
_handle_runs_on(items)