ViperEkura
84753d3e08
- extract preprocessing config factories into tests/data/factories.py - keep conftest.py fixtures-only; stop importing builders from it - promote temp_dir fixture to root conftest for cross-directory reuse - unify duplicate BPE tokenizer builders into build_test_tokenizer - merge grpo/dpo online e2e tests into one parametrized integration test - extract engine mock factory and shared model batch builders - drop local tempfile usage in favor of shared fixtures No behavior change: 519 tests pass.
93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from astrai.extension import KERNEL_NAMES, is_available
|
|
from astrai.model.transformer import AutoRegressiveLM
|
|
from tests.helpers import (
|
|
TINY_CONFIG,
|
|
RandomTokenDataset,
|
|
build_test_tokenizer,
|
|
make_tiny_config,
|
|
)
|
|
|
|
CUDA_AVAIL = torch.cuda.is_available()
|
|
KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES)
|
|
skip_no_cuda = pytest.mark.skipif(not CUDA_AVAIL, reason="CUDA not available")
|
|
skip_no_kernel = pytest.mark.skipif(not KERNEL_AVAIL, reason="CUDA kernels not built")
|
|
|
|
|
|
def pytest_configure(config):
|
|
config.addinivalue_line("markers", "slow: marks tests as slow")
|
|
config.addinivalue_line("markers", "integration: integration tests")
|
|
config.addinivalue_line("markers", "unit: fast unit tests")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def device():
|
|
"""Session-scoped device string (``"cuda"`` if available, else ``"cpu"``)."""
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def create_test_tokenizer(vocab_size: int = 1000):
|
|
"""Create a simple tokenizer for testing purposes."""
|
|
return build_test_tokenizer(vocab_size)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_tokenizer():
|
|
"""Session-scoped tokenizer, created once for the entire test run."""
|
|
return create_test_tokenizer()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_model(device):
|
|
"""Session-scoped small AutoRegressiveLM model, created once."""
|
|
config = make_tiny_config()
|
|
model = AutoRegressiveLM(config).to(device=device)
|
|
return {"model": model, "device": device, "config": config}
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir():
|
|
"""Function-scoped temporary directory, cleaned up after each test."""
|
|
d = tempfile.mkdtemp()
|
|
yield d
|
|
shutil.rmtree(d, ignore_errors=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def base_test_env(test_model, test_tokenizer, temp_dir):
|
|
"""Function-scoped test environment with isolated temp directory."""
|
|
config_path = os.path.join(temp_dir, "config.json")
|
|
with open(config_path, "w") as f:
|
|
json.dump(TINY_CONFIG, f)
|
|
|
|
return {
|
|
"device": test_model["device"],
|
|
"test_dir": temp_dir,
|
|
"config_path": config_path,
|
|
"transformer_config": test_model["config"],
|
|
"model": test_model["model"],
|
|
"tokenizer": test_tokenizer,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def random_dataset():
|
|
return RandomTokenDataset(length=None)
|
|
|
|
|
|
@pytest.fixture
|
|
def multi_turn_dataset():
|
|
return RandomTokenDataset(length=None, with_loss_mask=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def early_stopping_dataset():
|
|
return RandomTokenDataset(length=10, stop_after=5)
|