ViperEkura
dcc96de12a
- Convert protocol and MoE test classes to plain functions - Add real server/engine integration and generate_async tests - Isolate test model per test and use pytest tmp_path - Reset FastAPI engine state after inference tests - Fix generate_async StopIteration handling on Python 3.12 - Fix HF adapter MoE dense/shared and Gemma qk_norm mapping - Correct dev dependency httpx2 to httpx
89 lines
2.5 KiB
Python
89 lines
2.5 KiB
Python
import json
|
|
import os
|
|
|
|
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
|
|
def test_model(device):
|
|
"""Function-scoped small AutoRegressiveLM model, isolated per test."""
|
|
config = make_tiny_config()
|
|
model = AutoRegressiveLM(config).to(device=device)
|
|
return {"model": model, "device": device, "config": config}
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir(tmp_path):
|
|
"""Function-scoped temporary directory, cleaned up by pytest."""
|
|
return str(tmp_path)
|
|
|
|
|
|
@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)
|