ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: deduplicate and restructure test suite

Browse Source 
- 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.
This commit is contained in:
ViperEkura
2026-08-20 01:53:28 +08:00
parent 53a7149577
commit 84753d3e08
13 changed files with 312 additions and 454 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/conftest.py
+20 -21
View File
@@ -5,12 +5,15 @@ import tempfile
import pytest
import torch
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
from astrai.extension import KERNEL_NAMES, is_available
from astrai.model.transformer import AutoRegressiveLM
from astrai.tokenize import AutoTokenizer
from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config
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)
@@ -30,18 +33,9 @@ def device():
return "cuda" if torch.cuda.is_available() else "cpu"
def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
def create_test_tokenizer(vocab_size: int = 1000):
"""Create a simple tokenizer for testing purposes."""
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
trainer = trainers.BpeTrainer(
vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
)
tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
auto_tokenizer = AutoTokenizer()
auto_tokenizer._tokenizer = tokenizer
auto_tokenizer._special_token_map = {"unk_token": "<unk>", "pad_token": "<pad>"}
return auto_tokenizer
return build_test_tokenizer(vocab_size)
@pytest.fixture(scope="session")
@@ -59,24 +53,29 @@ def test_model(device):
@pytest.fixture
def base_test_env(test_model, test_tokenizer):
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."""
test_dir = tempfile.mkdtemp()
config_path = os.path.join(test_dir, "config.json")
config_path = os.path.join(temp_dir, "config.json")
with open(config_path, "w") as f:
json.dump(TINY_CONFIG, f)
yield {
return {
"device": test_model["device"],
"test_dir": str(test_dir),
"test_dir": temp_dir,
"config_path": config_path,
"transformer_config": test_model["config"],
"model": test_model["model"],
"tokenizer": test_tokenizer,
}
shutil.rmtree(test_dir)
@pytest.fixture
def random_dataset():
tests/data/conftest.py
+27 -152
View File
@@ -1,21 +1,15 @@
import json
import os
import tempfile
import pytest
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
from astrai.config.preprocess_config import (
InputConfig,
PipelineConfig,
ProcessingConfig,
)
from astrai.preprocessing.builder import (
MultiOutputMaskBuilder,
SectionedMaskBuilder,
SingleOutputMaskBuilder,
)
from astrai.tokenize import AutoTokenizer
from tests.data.factories import make_grpo_config
from tests.helpers import build_test_tokenizer
_SPECIAL_TOKENS_CONFIG = {
"bos_token": "<|begin_of_sentence|>",
@@ -41,27 +35,8 @@ _CHAT_TEMPLATE = (
"{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
)
_CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
_INSTRUCTION_SECTIONS = [
{"field": "prompt", "action": "mask", "add_special_tokens": True},
{"field": "response", "action": "train"},
]
_TEXT_SECTIONS = [{"field": "text", "action": "train"}]
_GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
def _build_chat_tokenizer():
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tr = trainers.BpeTrainer(
vocab_size=512,
min_frequency=1,
special_tokens=_SPECIAL_TOKENS,
)
train_data = [
_CHAT_TOKENIZER_DATA = [
"hello world",
"Hi there!",
"You are helpful.",
@@ -78,18 +53,24 @@ def _build_chat_tokenizer():
"<|im_end|>",
*[chr(i) for i in range(32, 127)],
]
tok.train_from_iterator(train_data, tr)
auto_tok = AutoTokenizer()
auto_tok._tokenizer = tok
auto_tok._special_token_map = {
_CHAT_TOKENIZER_MAP = {
"bos_token": "<|begin_of_sentence|>",
"eos_token": "<|end_of_sentence|>",
"pad_token": "<|_pad_|>",
"unk_token": "<|_unk_|>",
}
auto_tok.set_chat_template(_CHAT_TEMPLATE)
return auto_tok
def _build_chat_tokenizer():
return build_test_tokenizer(
vocab_size=512,
special_tokens=_SPECIAL_TOKENS,
special_token_map=_CHAT_TOKENIZER_MAP,
add_prefix_space=False,
train_data=_CHAT_TOKENIZER_DATA,
chat_template=_CHAT_TEMPLATE,
)
@pytest.fixture(scope="session")
@@ -97,116 +78,11 @@ def chat_tokenizer():
return _build_chat_tokenizer()
@pytest.fixture
def temp_dir():
d = tempfile.mkdtemp()
yield d
import shutil
shutil.rmtree(d, ignore_errors=True)
def make_chat_config():
return PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
mask={"system": "mask", "user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_instruction_config():
return PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_text_config():
return PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
preprocessing=ProcessingConfig(
max_seq_len=2048, min_chars=1, max_chars=2_000_000
),
)
def make_dpo_chat_config():
return PipelineConfig(
input=InputConfig(
sources={
"chosen": {
"sections": [
{"field": "chosen", "action": "$role", "template": True}
]
},
"rejected": {
"sections": [
{"field": "rejected", "action": "$role", "template": True}
]
},
}
),
mask={"user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_grpo_config():
return PipelineConfig(
input=InputConfig(
sources={
"prompts": {
"sections": [
{"field": "prompt", "action": "mask", "template": True}
]
},
"responses": {
"sections": _GRPO_RESPONSE_SECTIONS,
"list_field": True,
"mask_key": "masks",
},
"rewards": {
"sections": [{"field": "rewards", "action": "value"}],
},
}
),
mask={"user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def make_grpo_no_template_config():
return PipelineConfig(
input=InputConfig(
sources={
"prompts": {
"sections": [
{
"field": "prompt",
"action": "mask",
"add_special_tokens": True,
}
]
},
"responses": {
"sections": _GRPO_RESPONSE_SECTIONS,
"list_field": True,
"mask_key": "masks",
},
"rewards": {
"sections": [{"field": "rewards", "action": "value"}],
},
}
),
mask={"user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
)
def _write_tokenizer_dir(dir_path, tokenizer, tokenizer_config):
"""Persist a tokenizer plus ``tokenizer_config.json`` into *dir_path*."""
tokenizer._tokenizer.save(os.path.join(dir_path, "tokenizer.json"))
with open(os.path.join(dir_path, "tokenizer_config.json"), "w") as f:
json.dump(tokenizer_config, f)
@pytest.fixture
@@ -228,10 +104,10 @@ def multi_builder():
def tokenizer_dir(temp_dir, test_tokenizer):
d = os.path.join(temp_dir, "tok")
os.makedirs(d, exist_ok=True)
test_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
json.dump(
{"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, f
_write_tokenizer_dir(
d,
test_tokenizer,
{"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}},
)
return d
@@ -240,10 +116,9 @@ def tokenizer_dir(temp_dir, test_tokenizer):
def chat_tokenizer_dir(temp_dir, chat_tokenizer):
d = os.path.join(temp_dir, "tok")
os.makedirs(d, exist_ok=True)
chat_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
json.dump(
_write_tokenizer_dir(
d,
chat_tokenizer,
{"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE},
f,
)
return d
tests/data/factories.py
+86
View File
@@ -0,0 +1,86 @@
"""Test data builders for preprocessing and dataset scenarios."""
from astrai.config.preprocess_config import (
InputConfig,
PipelineConfig,
ProcessingConfig,
)
CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
INSTRUCTION_SECTIONS = [
{"field": "prompt", "action": "mask", "add_special_tokens": True},
{"field": "response", "action": "train"},
]
TEXT_SECTIONS = [{"field": "text", "action": "train"}]
GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
def make_pipeline_config(sections, *, mask=None, preprocessing=None, sources=None):
"""Build a pipeline config with the common test defaults."""
return PipelineConfig(
input=InputConfig(sections=sections, sources=sources),
mask={} if mask is None else mask,
mask_default="mask",
preprocessing=preprocessing or ProcessingConfig(max_seq_len=2048),
)
def make_chat_config():
return make_pipeline_config(
CHAT_SECTIONS,
mask={"system": "mask", "user": "mask", "assistant": "train"},
)
def make_instruction_config():
return make_pipeline_config(
INSTRUCTION_SECTIONS,
mask={"prompt": "mask", "response": "train"},
)
def make_text_config():
return make_pipeline_config(
TEXT_SECTIONS,
preprocessing=ProcessingConfig(
max_seq_len=2048, min_chars=1, max_chars=2_000_000
),
)
def make_dpo_chat_config():
sources = {
name: {"sections": [{"field": name, "action": "$role", "template": True}]}
for name in ("chosen", "rejected")
}
return make_pipeline_config(
None,
mask={"user": "mask", "assistant": "train"},
sources=sources,
)
def make_grpo_config(*, template=True):
prompt_section = {"field": "prompt", "action": "mask"}
if template:
prompt_section["template"] = True
else:
prompt_section["add_special_tokens"] = True
sources = {
"prompts": {"sections": [prompt_section]},
"responses": {
"sections": GRPO_RESPONSE_SECTIONS,
"list_field": True,
"mask_key": "masks",
},
"rewards": {"sections": [{"field": "rewards", "action": "value"}]},
}
return make_pipeline_config(
None,
mask={"user": "mask", "assistant": "train"},
sources=sources,
)
def make_grpo_no_template_config():
return make_grpo_config(template=False)
tests/data/test_dataset.py
+2 -2
View File
@@ -24,7 +24,7 @@ from astrai.serialization import (
load_bin,
save_bin,
)
from tests.data.conftest import make_grpo_no_template_config
from tests.data.factories import make_grpo_config
def _rand_seq(length, vocab=1000):
@@ -797,7 +797,7 @@ def test_grpo_builder_preserves_response_boundaries(base_test_env):
_save_test_tokenizer(base_test_env["test_dir"], tokenizer)
builder = SectionedMaskBuilder()
config = make_grpo_no_template_config()
config = make_grpo_config(template=False)
config.preprocessing.max_seq_len = 128
item = {
tests/data/test_preprocess_builder.py
+13 -13
View File
@@ -12,10 +12,10 @@ from astrai.preprocessing.builder import (
SectionedMaskBuilder,
SingleOutputMaskBuilder,
)
from tests.data.conftest import (
_CHAT_SECTIONS,
_INSTRUCTION_SECTIONS,
_TEXT_SECTIONS,
from tests.data.factories import (
CHAT_SECTIONS,
INSTRUCTION_SECTIONS,
TEXT_SECTIONS,
make_chat_config,
make_dpo_chat_config,
make_grpo_config,
@@ -101,7 +101,7 @@ def test_chat_uniform_masking(
mask_rules, mask_default, expect_nonzero, chat_tokenizer, builder
):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask=mask_rules,
mask_default=mask_default,
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -128,7 +128,7 @@ def test_chat_empty_messages(chat_tokenizer, builder):
def test_chat_domain_extraction(chat_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -147,7 +147,7 @@ def test_chat_domain_extraction(chat_tokenizer, builder):
def test_chat_truncation(chat_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=10),
@@ -237,7 +237,7 @@ def test_text_empty(test_tokenizer, builder):
def test_text_too_short(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(min_chars=100),
)
assert builder.build({"text": "short"}, config, test_tokenizer) is None
@@ -245,7 +245,7 @@ def test_text_too_short(test_tokenizer, builder):
def test_text_truncation(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
)
item = {"text": "This is a very long text that should be truncated"}
@@ -255,7 +255,7 @@ def test_text_truncation(test_tokenizer, builder):
def test_sectioned_chat(chat_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"system": "mask", "user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -275,7 +275,7 @@ def test_sectioned_chat(chat_tokenizer, builder):
def test_sectioned_instruction(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
)
item = {"prompt": "Q: Why?", "response": "A: Because."}
@@ -288,7 +288,7 @@ def test_sectioned_instruction(test_tokenizer, builder):
def test_sectioned_text(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
)
item = {"text": "Hello world, this is a test."}
@@ -299,7 +299,7 @@ def test_sectioned_text(test_tokenizer, builder):
def test_sectioned_text_too_short(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
)
assert builder.build({"text": "short"}, config, test_tokenizer) is None
tests/data/test_preprocess_config.py
+7 -7
View File
@@ -4,9 +4,9 @@ from astrai.config.preprocess_config import (
InputConfig,
PipelineConfig,
)
from tests.data.conftest import (
_INSTRUCTION_SECTIONS,
_TEXT_SECTIONS,
from tests.data.factories import (
INSTRUCTION_SECTIONS,
TEXT_SECTIONS,
make_dpo_chat_config,
)
@@ -43,26 +43,26 @@ def test_from_dict_flat():
def test_to_dict_roundtrip():
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
)
d = config.to_dict()
config2 = PipelineConfig.from_dict(d)
assert config2.input.sections == _INSTRUCTION_SECTIONS
assert config2.input.sections == INSTRUCTION_SECTIONS
assert config2.mask == {"prompt": "mask", "response": "train"}
def test_to_file_from_file(temp_dir):
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
mask={"text": "train"},
mask_default="mask",
)
path = os.path.join(temp_dir, "config.json")
config.to_file(path)
loaded = PipelineConfig.from_file(path)
assert loaded.input.sections == _TEXT_SECTIONS
assert loaded.input.sections == TEXT_SECTIONS
assert loaded.mask == {"text": "train"}
tests/data/test_preprocess_pipeline.py
+8 -8
View File
@@ -9,10 +9,10 @@ from astrai.config.preprocess_config import (
)
from astrai.preprocessing.packing import PackingStrategyFactory
from astrai.preprocessing.pipeline import Pipeline, filter_by_length
from tests.data.conftest import (
_CHAT_SECTIONS,
_INSTRUCTION_SECTIONS,
_TEXT_SECTIONS,
from tests.data.factories import (
CHAT_SECTIONS,
INSTRUCTION_SECTIONS,
TEXT_SECTIONS,
make_dpo_chat_config,
make_grpo_no_template_config,
)
@@ -54,7 +54,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir):
)
config = PipelineConfig(
input=InputConfig(sections=_CHAT_SECTIONS),
input=InputConfig(sections=CHAT_SECTIONS),
mask={"system": "mask", "user": "mask", "assistant": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -97,7 +97,7 @@ def test_full_text_pipeline(temp_dir, tokenizer_dir):
)
config = PipelineConfig(
input=InputConfig(sections=_TEXT_SECTIONS),
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=10),
output=OutputConfig(storage_format="bin"),
)
@@ -138,7 +138,7 @@ def test_full_instruction_pipeline(temp_dir, tokenizer_dir):
)
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -164,7 +164,7 @@ def test_dtype_override(temp_dir, tokenizer_dir):
f.write(json.dumps({"prompt": "Q", "response": "A"}) + "\n")
config = PipelineConfig(
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
tests/helpers.py
+40
View File
@@ -4,10 +4,12 @@ import json
import os
import torch
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
from torch.utils.data import Dataset
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
from astrai.tokenize import AutoTokenizer
TINY_CONFIG = dict(
vocab_size=1000,
@@ -57,6 +59,44 @@ def make_model(device, **cfg_overrides):
return model, cfg
def build_test_tokenizer(
vocab_size: int = 1000,
*,
special_tokens=("<unk>", "<pad>"),
special_token_map=None,
add_prefix_space: bool = True,
train_data=None,
chat_template: str | None = None,
) -> AutoTokenizer:
"""Build a lightweight BPE ``AutoTokenizer`` for tests.
``special_token_map`` defaults to ``{"unk_token", "pad_token"}``
pointing at the first two special tokens.
"""
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(
add_prefix_space=add_prefix_space
)
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
min_frequency=1,
special_tokens=list(special_tokens),
)
tokenizer.train_from_iterator(
train_data if train_data is not None else [chr(i) for i in range(256)],
trainer,
)
auto_tokenizer = AutoTokenizer()
auto_tokenizer._tokenizer = tokenizer
auto_tokenizer._special_token_map = special_token_map or {
"unk_token": special_tokens[0],
"pad_token": special_tokens[1],
}
if chat_template is not None:
auto_tokenizer.set_chat_template(chat_template)
return auto_tokenizer
def make_frozen(model, device):
"""Create a frozen, eval-mode copy of *model* with identical weights."""
cfg = make_rollout_config()
tests/inference/test_engine.py
+18 -25
View File
@@ -8,6 +8,17 @@ from astrai.inference import STOP
from astrai.inference.engine import GenerateResult, InferenceEngine
def _make_engine_mocks(decode=None):
"""Build the standard mock model/tokenizer pair used by engine tests."""
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_tokenizer.encode.return_value = [1, 2, 3]
mock_tokenizer.stop_ids = [0]
if decode is not None:
mock_tokenizer.decode.return_value = decode
return mock_model, mock_tokenizer
def test_result_append_single():
r = GenerateResult(count=1)
r.append("hello", 0)
@@ -102,11 +113,7 @@ def test_result_get_results():
def test_engine_generate_non_streaming_single():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_tokenizer.encode.return_value = [1, 2, 3]
mock_tokenizer.decode.return_value = "response"
mock_tokenizer.stop_ids = [0]
mock_model, mock_tokenizer = _make_engine_mocks(decode="response")
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
instance = MockSched.return_value
@@ -125,11 +132,7 @@ def test_engine_generate_non_streaming_single():
def test_engine_generate_streaming_yields_tokens():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_tokenizer.encode.return_value = [1, 2, 3]
mock_tokenizer.decode.return_value = "tok"
mock_tokenizer.stop_ids = [0]
mock_model, mock_tokenizer = _make_engine_mocks(decode="tok")
callbacks_saved = []
@@ -154,11 +157,7 @@ def test_engine_generate_streaming_yields_tokens():
def test_engine_generate_non_streaming_batch():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_tokenizer.encode.return_value = [1, 2, 3]
mock_tokenizer.decode.return_value = "r"
mock_tokenizer.stop_ids = [0]
mock_model, mock_tokenizer = _make_engine_mocks(decode="r")
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
instance = MockSched.return_value
@@ -177,10 +176,7 @@ def test_engine_generate_non_streaming_batch():
def test_engine_generate_zero_max_tokens_returns_empty():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_tokenizer.encode.return_value = [1, 2, 3]
mock_tokenizer.stop_ids = [0]
mock_model, mock_tokenizer = _make_engine_mocks()
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
instance = MockSched.return_value
@@ -192,8 +188,7 @@ def test_engine_generate_zero_max_tokens_returns_empty():
def test_engine_generate_zero_max_tokens_stream_is_empty():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_model, mock_tokenizer = _make_engine_mocks()
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
instance = MockSched.return_value
@@ -203,8 +198,7 @@ def test_engine_generate_zero_max_tokens_stream_is_empty():
def test_engine_passes_backend_to_scheduler():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_model, mock_tokenizer = _make_engine_mocks()
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
InferenceEngine(
@@ -218,8 +212,7 @@ def test_engine_passes_backend_to_scheduler():
def test_generate_captures_calling_backend_context():
mock_model = MagicMock()
mock_tokenizer = MagicMock()
mock_model, mock_tokenizer = _make_engine_mocks()
captured = []
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
tests/module/test_lora.py
+15 -25
View File
@@ -1,9 +1,6 @@
import tempfile
import pytest
import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model import AutoRegressiveLM
from astrai.model.components.linear import Linear
from astrai.model.components.lora import (
@@ -16,22 +13,20 @@ from astrai.model.components.lora import (
merge_lora,
save_lora,
)
from tests.helpers import make_tiny_config
MODEL_KWARGS = dict(
LORA_MODEL_KWARGS = dict(
vocab_size=1000,
hidden_size=64,
num_attention_heads=4,
num_key_value_heads=2,
intermediate_size=128,
num_hidden_layers=2,
max_position_embeddings=32,
rms_norm_eps=1e-5,
)
def _make_model(**kwargs):
kw = {**MODEL_KWARGS, **kwargs}
config = AutoRegressiveLMConfig(**kw)
config = make_tiny_config(**{**LORA_MODEL_KWARGS, **kwargs})
model = AutoRegressiveLM(config)
model.eval()
return model
@@ -227,7 +222,7 @@ def test_state_dict_after_inject_consistent_with_original():
assert len(lora_keys) > 0
def test_save_load_roundtrip():
def test_save_load_roundtrip(temp_dir):
model = _make_model()
cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
@@ -240,12 +235,11 @@ def test_save_load_roundtrip():
with torch.no_grad():
out_src = model(x)["logits"].clone()
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, cfg)
save_lora(model, temp_dir, cfg)
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
load_lora(model2, tmpdir)
load_lora(model2, temp_dir)
with torch.no_grad():
out_dst = model2(x)["logits"]
@@ -253,7 +247,7 @@ def test_save_load_roundtrip():
torch.testing.assert_close(out_src, out_dst)
def test_save_after_merge_raises():
def test_save_after_merge_raises(temp_dir):
model = _make_model()
cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
@@ -262,16 +256,14 @@ def test_save_after_merge_raises():
if isinstance(m, LoRALinear):
m.lora_B.fill_(0.5)
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, cfg)
save_lora(model, temp_dir, cfg)
merge_lora(model)
with tempfile.TemporaryDirectory() as tmpdir2:
with pytest.raises(RuntimeError, match="No LoRA parameters"):
save_lora(model, tmpdir2, cfg)
save_lora(model, temp_dir, cfg)
def test_load_lora_on_already_injected():
def test_load_lora_on_already_injected(temp_dir):
model = _make_model()
inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
@@ -280,18 +272,17 @@ def test_load_lora_on_already_injected():
if isinstance(m, LoRALinear):
m.lora_B.fill_(0.5)
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
save_lora(model, temp_dir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
load_lora(model2, tmpdir)
load_lora(model2, temp_dir)
assert _get_lora_count(model2) > 0
def test_load_lora_mismatched_r_raises():
def test_load_lora_mismatched_r_raises(temp_dir):
model = _make_model()
cfg = inject_lora(model, r=8, alpha=16, target_modules={"q_proj"})
@@ -300,15 +291,14 @@ def test_load_lora_mismatched_r_raises():
if isinstance(m, LoRALinear):
m.lora_B.fill_(0.5)
with tempfile.TemporaryDirectory() as tmpdir:
save_lora(model, tmpdir, cfg)
save_lora(model, temp_dir, cfg)
model2 = _make_model()
model2.load_state_dict(model.state_dict(), strict=False)
inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
with pytest.raises(RuntimeError, match="size mismatch"):
load_lora(model2, tmpdir)
load_lora(model2, temp_dir)
def test_merge_preserves_output():
tests/module/test_model_forward.py
+23 -28
View File
@@ -39,6 +39,21 @@ def _make_model(config=None) -> AutoRegressiveLM:
return AutoRegressiveLM(config)
def _make_batch(config, batch_size=2, seq_len=8, with_extra=False):
"""Build a random token batch, optionally with position ids and loss mask."""
vocab = config.vocab_size
batch = {
"input_ids": torch.randint(0, vocab, (batch_size, seq_len)),
"target_ids": torch.randint(0, vocab, (batch_size, seq_len)),
}
if with_extra:
batch["position_ids"] = (
torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
)
batch["loss_mask"] = torch.ones(batch_size, seq_len, dtype=torch.bool)
return batch
def test_model_forward_contract_uses_dense_training_and_packed_inference():
from astrai.inference.cache import PagePool, TaskCacheManager
from astrai.inference.workspace import InferenceWorkspace
@@ -180,13 +195,6 @@ class TestSEQStrategyMoE:
self.model = _make_model(self.config).to(device)
self.model.train()
def _make_batch(self, batch_size=2, seq_len=8):
vocab = self.config.vocab_size
input_ids = torch.randint(0, vocab, (batch_size, seq_len))
# target = input shifted right
target_ids = torch.randint(0, vocab, (batch_size, seq_len))
return {"input_ids": input_ids, "target_ids": target_ids}
def test_compute_loss_returns_scalar(self):
"""compute_loss should return a scalar tensor."""
strategy = SEQStrategy(
@@ -194,7 +202,7 @@ class TestSEQStrategyMoE:
self.device,
moe_aux_loss_coef=0.01,
)
loss = strategy.compute_loss(self._make_batch())
loss = strategy.compute_loss(_make_batch(self.config))
assert loss.ndim == 0
assert loss.requires_grad
@@ -205,7 +213,7 @@ class TestSEQStrategyMoE:
self.device,
moe_aux_loss_coef=0.01,
)
output = strategy.compute_loss_output(self._make_batch())
output = strategy.compute_loss_output(_make_batch(self.config))
assert "loss" in output
assert "metrics" in output
@@ -225,7 +233,7 @@ class TestSEQStrategyMoE:
self.device,
moe_aux_loss_coef=0.01,
)
strategy.compute_loss_output(self._make_batch())
strategy.compute_loss_output(_make_batch(self.config))
moe_metrics = strategy._moe_metrics
assert moe_metrics, "_moe_metrics should not be empty for MoE model"
@@ -246,7 +254,7 @@ class TestSEQStrategyMoE:
self.device,
moe_aux_loss_coef=0.0,
)
output = strategy.compute_loss_output(self._make_batch())
output = strategy.compute_loss_output(_make_batch(self.config))
metrics = output["metrics"]
# task_loss and loss should be equal (aux weighted by zero)
@@ -268,7 +276,7 @@ class TestSEQStrategyMoE:
self.device,
moe_aux_loss_coef=0.01,
)
output = strategy.compute_loss_output(self._make_batch())
output = strategy.compute_loss_output(_make_batch(self.config))
assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
def test_factory_creates_strategy_with_coef(self):
@@ -294,7 +302,7 @@ class TestSEQStrategyMoE:
self.device,
moe_aux_loss_coef=0.01,
)
output = strategy.compute_loss_output(self._make_batch())
output = strategy.compute_loss_output(_make_batch(self.config))
metrics = output["metrics"]
assert "moe_aux_loss" not in metrics
@@ -313,19 +321,6 @@ class TestSFTStrategyMoE:
self.model = _make_model(self.config).to(device)
self.model.train()
def _make_batch(self, batch_size=2, seq_len=8):
vocab = self.config.vocab_size
input_ids = torch.randint(0, vocab, (batch_size, seq_len))
target_ids = torch.randint(0, vocab, (batch_size, seq_len))
position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
return {
"input_ids": input_ids,
"target_ids": target_ids,
"position_ids": position_ids,
"loss_mask": loss_mask,
}
def test_compute_loss_output_with_aux_loss(self):
"""SFTStrategy produces MoE metrics when coef > 0."""
strategy = SFTStrategy(
@@ -333,7 +328,7 @@ class TestSFTStrategyMoE:
self.device,
moe_aux_loss_coef=0.01,
)
output = strategy.compute_loss_output(self._make_batch())
output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
metrics = output["metrics"]
assert "moe_aux_loss" in metrics
@@ -351,7 +346,7 @@ class TestSFTStrategyMoE:
self.device,
moe_aux_loss_coef=0.0,
)
output = strategy.compute_loss_output(self._make_batch())
output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
metrics = output["metrics"]
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
tests/trainer/test_grpo_e2e.py
-123
View File
@@ -1,123 +0,0 @@
"""End-to-end integration test for online GRPO rollout."""
import os
from functools import partial
import pytest
import torch
from torch.utils.data import Dataset
from astrai.config import TrainConfig
from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.rollout import BaseRewardModel
from astrai.trainer.schedule import SchedulerFactory
from astrai.trainer.trainer import Trainer
from tests.helpers import CHAT_TEMPLATE
class InstructionDataset(Dataset):
"""Toy instruction/input dataset for online RL rollout.
Each sample has an ``instruction`` and an optional ``input``; the
RolloutGenerator renders both through the tokenizer's chat template
so the prompt matches the SFT-trained format.
"""
_SAMPLES = [
{"instruction": "Hello", "input": ""},
{"instruction": "Tell me a story", "input": "about dragons"},
{"instruction": "Summarize", "input": "the article"},
{"instruction": "Translate", "input": "to French: hi"},
]
def __len__(self):
return len(self._SAMPLES)
def __getitem__(self, idx):
return dict(self._SAMPLES[idx])
class LengthRewardModel(BaseRewardModel):
"""Rewards each response by its (non-pad) token count.
Gives the group-normalized advantage a non-degenerate signal.
"""
def score(self, prompts, responses):
B = len(prompts)
G = len(responses[0]) if B else 0
rewards = torch.zeros(B, G)
for i in range(B):
for g in range(G):
rewards[i, g] = float(len(responses[i][g]))
return rewards
def instruction_collate_fn(batch):
"""Stack a list of instruction/input dicts into a batch dict of lists."""
return {
"instruction": [b["instruction"] for b in batch],
"input": [b.get("input", "") for b in batch],
}
def _model_fn(model_config):
return AutoRegressiveLM(model_config).to(dtype=torch.float32)
def _optimizer_fn(m):
return torch.optim.AdamW(m.parameters(), lr=1e-4)
def _scheduler_fn(optim):
return SchedulerFactory.create(
"cosine", optim, warmup_steps=1, lr_decay_steps=4, min_rate=0.05
)
@pytest.mark.integration
def test_online_grpo_end_to_end(base_test_env):
"""Run one epoch of online GRPO with KV-cache-backed rollout."""
test_dir = base_test_env["test_dir"]
device = base_test_env["device"]
tokenizer = base_test_env["tokenizer"]
model_config = base_test_env["transformer_config"]
tokenizer.set_chat_template(CHAT_TEMPLATE)
tokenizer.save_pretrained(test_dir)
model_fn = partial(_model_fn, model_config)
optimizer_fn = _optimizer_fn
scheduler_fn = _scheduler_fn
dataset = InstructionDataset()
train_config = TrainConfig(
strategy="online_grpo",
model_fn=model_fn,
dataset=dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
ckpt_dir=os.path.join(test_dir, "ckpt"),
n_epoch=1,
batch_per_device=2,
ckpt_interval=100,
grad_accum_steps=1,
random_seed=42,
device_type=device,
nprocs=1,
parallel_mode="none",
strategy_kwargs={"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2},
rollout_interval=1,
rollout_temperature=1.0,
rollout_top_k=0,
rollout_top_p=1.0,
rollout_max_tokens=4,
reward_model_fn=LengthRewardModel,
collate_fn=instruction_collate_fn,
)
trainer = Trainer(train_config)
trainer.train(param_path=test_dir)
assert os.path.isdir(os.path.join(test_dir, "ckpt"))
tests/trainer/test_online_e2e.py
+21 -18
View File
@@ -1,4 +1,4 @@
"""End-to-end integration test for online DPO rollout."""
"""End-to-end integration tests for online GRPO/DPO rollout."""
import os
from functools import partial
@@ -40,7 +40,7 @@ class InstructionDataset(Dataset):
class LengthRewardModel(BaseRewardModel):
"""Rewards each response by its (non-pad) token count.
Enough for DPO to distinguish chosen/rejected from the rollout group.
Gives the group-normalized advantage a non-degenerate signal.
"""
def score(self, prompts, responses):
@@ -75,31 +75,34 @@ def _scheduler_fn(optim):
)
_ONLINE_STRATEGIES = [
pytest.param(
"online_grpo",
{"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2},
id="grpo",
),
pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, id="dpo"),
]
@pytest.mark.integration
def test_online_dpo_end_to_end(base_test_env):
"""Run one epoch of online DPO with KV-cache-backed rollout."""
@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES)
def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs):
"""Run one epoch of online RL rollout with KV-cache-backed generation."""
test_dir = base_test_env["test_dir"]
device = base_test_env["device"]
tokenizer = base_test_env["tokenizer"]
model_config = base_test_env["transformer_config"]
# Equip tokenizer with a chat template so RolloutGenerator can
# render instruction/input via apply_chat_template.
tokenizer.set_chat_template(CHAT_TEMPLATE)
tokenizer.save_pretrained(test_dir)
model_fn = partial(_model_fn, model_config)
optimizer_fn = _optimizer_fn
scheduler_fn = _scheduler_fn
dataset = InstructionDataset()
train_config = TrainConfig(
strategy="online_dpo",
model_fn=model_fn,
dataset=dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
strategy=strategy,
model_fn=partial(_model_fn, model_config),
dataset=InstructionDataset(),
optimizer_fn=_optimizer_fn,
scheduler_fn=_scheduler_fn,
ckpt_dir=os.path.join(test_dir, "ckpt"),
n_epoch=1,
batch_per_device=2,
@@ -109,7 +112,7 @@ def test_online_dpo_end_to_end(base_test_env):
device_type=device,
nprocs=1,
parallel_mode="none",
strategy_kwargs={"beta": 0.1, "group_size": 2},
strategy_kwargs=strategy_kwargs,
rollout_interval=1,
rollout_temperature=1.0,
rollout_top_k=0,