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

test: deduplicate suites and prune low-value cases

Browse Source 
- extract shared helpers for dataset writers, scheduler construction, thread interleaving, hf roundtrips, and moe configs
- remove about 20 cases whose only assertions were format checks, restated declarations, fake-taxonomy duplicates, or test-local scaffolding
- strengthen weak cases into exact reference comparisons, positional mask checks, and deterministic outcomes
- replace two schedule factory smoke tests with cosine/sgdr formula assertions
- delete root-level CLI tests whose merge-priority facts are covered by tests/config/test_cli.py
- suite shrinks from 857 to 826 items; ruff format, import order, and pytest all green
This commit is contained in:
ViperEkura
2026-09-03 21:54:14 +08:00
parent 28d11f1610
commit 9d3ae76683
22 changed files with 407 additions and 1027 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/data/test_dataset.py
+87 -189
View File
@@ -26,11 +26,22 @@ from astrai.serialization import (
) )
from tests.data.factories import make_grpo_config from tests.data.factories import make_grpo_config
SIMPLE_CHAT_TEMPLATE = (
"{% for message in messages %}"
"{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}"
)
def _rand_seq(length, vocab=1000): def _rand_seq(length, vocab=1000):
return torch.randint(0, vocab, (length,), dtype=torch.int64) return torch.randint(0, vocab, (length,), dtype=torch.int64)
def _dump_jsonl(path, records):
with open(path, "w", encoding="utf-8") as f:
for record in records:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
def _save_test_tokenizer(test_dir, tokenizer): def _save_test_tokenizer(test_dir, tokenizer):
tokenizer_path = os.path.join(test_dir, "tokenizer") tokenizer_path = os.path.join(test_dir, "tokenizer")
os.makedirs(tokenizer_path, exist_ok=True) os.makedirs(tokenizer_path, exist_ok=True)
@@ -38,19 +49,19 @@ def _save_test_tokenizer(test_dir, tokenizer):
return tokenizer_path return tokenizer_path
def _write_jsonl_dataset(test_dir, tokenizer_path, records, config_overrides=None): def _write_text_dataset(
data_dir = os.path.join(test_dir, "jsonl_data") test_dir, dirname, tokenizer_path, records, config_overrides=None
):
"""Write a JSONL dataset directory with a text-section default config."""
data_dir = os.path.join(test_dir, dirname)
os.makedirs(data_dir, exist_ok=True) os.makedirs(data_dir, exist_ok=True)
_dump_jsonl(os.path.join(data_dir, "data.jsonl"), records)
with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
for record in records:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
config = { config = {
"tokenizer_path": tokenizer_path, "tokenizer_path": tokenizer_path,
"version": 1, "version": 1,
"input": {"sections": [{"field": "text", "action": "train"}]}, "input": {"sections": [{"field": "text", "action": "train"}]},
"preprocessing": {"max_seq_len": 128}, "preprocessing": {"max_seq_len": 128, "min_chars": 0},
"output": {"position_ids_mode": "continuous"}, "output": {"position_ids_mode": "continuous"},
} }
if config_overrides: if config_overrides:
@@ -71,6 +82,27 @@ def _fake_fetch_record(self, idx, keys):
return {k: self._data[k][idx] for k in keys} return {k: self._data[k][idx] for k in keys}
def _grpo_fake_store(prompts, responses, masks, rewards):
"""Fake GRPO record store matching real Store semantics."""
return type(
"FakeStore",
(),
{
"keys": ["prompts", "responses", "masks", "rewards"],
"num_records": len(prompts),
"token_count": 0,
"_data": {
"prompts": prompts,
"responses": responses,
"masks": masks,
"rewards": rewards,
},
"fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
},
)()
def _make_seq_dataset( def _make_seq_dataset(
test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs
): ):
@@ -85,32 +117,6 @@ def _make_seq_dataset(
) )
def test_dataset_loader_random_paths(base_test_env):
"""Test dataset loader with multiple random paths"""
test_dir = base_test_env["test_dir"]
loaded_dataset = None
num_files = np.random.randint(2, 5)
for i in range(num_files):
seq_length = np.random.randint(200, 400)
dummy_data = {"sequence": [_rand_seq(seq_length) for _ in range(10)]}
sub_dir = os.path.join(test_dir, f"sub_{i}")
os.makedirs(sub_dir, exist_ok=True)
loaded_dataset = _make_seq_dataset(
sub_dir, f"data_{i}", seq_length, data=dummy_data
)
assert loaded_dataset is not None
assert len(loaded_dataset) > 0
# Test that we can get items without errors
for i in range(len(loaded_dataset)):
item = loaded_dataset[i]
assert "input_ids" in item
assert "target_ids" in item
assert item["input_ids"].shape == item["target_ids"].shape
assert item["input_ids"].shape[0] == 64
def test_dpo_strategy_with_random_data(base_test_env): def test_dpo_strategy_with_random_data(base_test_env):
"""Test DPO strategy with randomized preference data""" """Test DPO strategy with randomized preference data"""
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
@@ -345,25 +351,12 @@ def test_normalize_mixed_empty_key():
def test_grpo_dataset_dtype(base_test_env): def test_grpo_dataset_dtype(base_test_env):
"""GRPO dataset returns correct dtypes for per-record structured data.""" """GRPO dataset returns correct dtypes for per-record structured data."""
G = 4 G = 4
store = type( store = _grpo_fake_store(
"FakeStore", prompts=[torch.randint(0, 100, (10,), dtype=torch.int32)],
(), responses=[[torch.randint(0, 100, (5,), dtype=torch.int32) for _ in range(G)]],
{ masks=[[torch.ones(5, dtype=torch.int32) for _ in range(G)]],
"keys": ["prompts", "responses", "masks", "rewards"], rewards=[torch.rand(G, dtype=torch.float32)],
"num_records": 1, )
"token_count": 0,
"_data": {
"prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)],
"responses": [
[torch.randint(0, 100, (5,), dtype=torch.int32) for _ in range(G)]
],
"masks": [[torch.ones(5, dtype=torch.int32) for _ in range(G)]],
"rewards": [torch.rand(G, dtype=torch.float32)],
},
"fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
},
)()
dataset = GRPODataset(store=store) dataset = GRPODataset(store=store)
item = dataset[0] item = dataset[0]
@@ -378,23 +371,12 @@ def test_grpo_dataset_load(base_test_env):
G = 3 G = 3
prompt_len = 8 prompt_len = 8
resp_lens = [5, 7, 4] resp_lens = [5, 7, 4]
store = type( store = _grpo_fake_store(
"FakeStore", prompts=[torch.randint(0, 100, (prompt_len,))],
(), responses=[[torch.randint(0, 100, (rl,)) for rl in resp_lens]],
{ masks=[[torch.ones(rl, dtype=torch.int64) for rl in resp_lens]],
"keys": ["prompts", "responses", "masks", "rewards"], rewards=[torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)],
"num_records": 1, )
"token_count": 0,
"_data": {
"prompts": [torch.randint(0, 100, (prompt_len,))],
"responses": [[torch.randint(0, 100, (rl,)) for rl in resp_lens]],
"masks": [[torch.ones(rl, dtype=torch.int64) for rl in resp_lens]],
"rewards": [torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)],
},
"fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
},
)()
dataset = GRPODataset(store=store) dataset = GRPODataset(store=store)
assert len(dataset) == 1 assert len(dataset) == 1
@@ -465,62 +447,28 @@ def test_dataset_load_explicit_storage_type(base_test_env):
assert dataset.token_count == 200 assert dataset.token_count == 200
def _write_json_dataset(test_dir, tokenizer_path, records, config_overrides=None): @pytest.mark.parametrize("use_jsonl", [True, False])
"""Write JSONL dataset — one JSON object per line."""
data_dir = os.path.join(test_dir, "json_data")
os.makedirs(data_dir, exist_ok=True)
with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
for rec in records:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
config = {
"tokenizer_path": tokenizer_path,
"version": 1,
"input": {"sections": [{"field": "text", "action": "train"}]},
"preprocessing": {"max_seq_len": 128, "min_chars": 0},
"output": {"position_ids_mode": "continuous"},
}
if config_overrides:
config.update(config_overrides)
with open(
os.path.join(data_dir, "dataset_config.json"), "w", encoding="utf-8"
) as f:
json.dump(config, f, ensure_ascii=False, indent=2)
return data_dir
@pytest.mark.parametrize(
"use_jsonl",
[True, False],
)
def test_detect_format_data_dir(base_test_env, use_jsonl): def test_detect_format_data_dir(base_test_env, use_jsonl):
"""detect_format returns 'jsonl' for dirs of .jsonl or .json files.""" """detect_format returns 'jsonl' for dirs of .jsonl or .json files."""
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
if use_jsonl: data_dir = _write_text_dataset(
data_dir = _write_jsonl_dataset( test_dir,
test_dir, "jsonl_data" if use_jsonl else "json_data",
tokenizer_path, tokenizer_path,
[{"text": "hello world"}, {"text": "foo bar baz"}], [{"text": "hello world"}, {"text": "foo bar baz"}],
) )
else:
data_dir = _write_json_dataset(
test_dir,
tokenizer_path,
[{"text": "hello world"}, {"text": "foo bar baz qux"}],
)
assert detect_format(data_dir) == "jsonl" assert detect_format(data_dir) == "jsonl"
def test_json_store_seq(base_test_env): @pytest.mark.parametrize("dirname", ["json_data", "jsonl_data"])
"""JsonlStore loads .json array correctly.""" def test_json_store_seq(base_test_env, dirname):
"""JsonlStore loads a text JSONL dataset and feeds seq training."""
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
data_dir = _write_json_dataset( data_dir = _write_text_dataset(
test_dir, test_dir,
dirname,
tokenizer_path, tokenizer_path,
[{"text": "hello world"}, {"text": "foo bar baz qux"}], [{"text": "hello world"}, {"text": "foo bar baz qux"}],
) )
@@ -535,15 +483,14 @@ def test_json_store_seq(base_test_env):
item = dataset[0] item = dataset[0]
assert "input_ids" in item assert "input_ids" in item
assert "target_ids" in item assert "target_ids" in item
assert item["input_ids"].dtype == torch.long
def test_json_store_no_tokenizer_path(base_test_env): def test_json_store_no_tokenizer_path(base_test_env):
"""JsonlStore uses dataset dir as tokenizer_path when omitted.""" """JsonlStore uses dataset dir as tokenizer_path when omitted."""
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer = base_test_env["tokenizer"] tokenizer = base_test_env["tokenizer"]
tokenizer.set_chat_template( tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE)
"{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}"
)
data_dir = os.path.join(test_dir, "self_contained") data_dir = os.path.join(test_dir, "self_contained")
os.makedirs(data_dir, exist_ok=True) os.makedirs(data_dir, exist_ok=True)
@@ -560,9 +507,7 @@ def test_json_store_no_tokenizer_path(base_test_env):
] ]
} }
] ]
with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: _dump_jsonl(os.path.join(data_dir, "data.jsonl"), records)
for rec in records:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
# dataset_config.json WITHOUT tokenizer_path # dataset_config.json WITHOUT tokenizer_path
config = { config = {
@@ -587,38 +532,14 @@ def test_json_store_no_tokenizer_path(base_test_env):
assert "loss_mask" in store.keys assert "loss_mask" in store.keys
def test_jsonl_store_seq(base_test_env):
test_dir = base_test_env["test_dir"]
tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
data_dir = _write_jsonl_dataset(
test_dir,
tokenizer_path,
[{"text": "hello world"}, {"text": "foo bar baz qux"}],
config_overrides={"preprocessing": {"max_seq_len": 128, "min_chars": 0}},
)
store = StoreFactory.create("jsonl")
store.load(data_dir, transform=_build_jsonl_transform(data_dir))
assert len(store) > 0
assert "sequence" in store.keys
dataset = DatasetFactory.load("seq", data_dir, window_size=8)
assert len(dataset) > 0
item = dataset[0]
assert "input_ids" in item
assert "target_ids" in item
assert item["input_ids"].dtype == torch.long
def test_jsonl_store_sft(base_test_env): def test_jsonl_store_sft(base_test_env):
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer = base_test_env["tokenizer"] tokenizer = base_test_env["tokenizer"]
tokenizer.set_chat_template( tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE)
"{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}"
)
tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
data_dir = _write_jsonl_dataset( data_dir = _write_text_dataset(
test_dir, test_dir,
"sft_jsonl",
tokenizer_path, tokenizer_path,
[ [
{ {
@@ -661,9 +582,7 @@ def test_sft_jsonl_default_messages_config(base_test_env):
""" """
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer = base_test_env["tokenizer"] tokenizer = base_test_env["tokenizer"]
tokenizer.set_chat_template( tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE)
"{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}"
)
tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
data_dir = os.path.join(test_dir, "jsonl_data") data_dir = os.path.join(test_dir, "jsonl_data")
@@ -683,9 +602,7 @@ def test_sft_jsonl_default_messages_config(base_test_env):
] ]
}, },
] ]
with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f: _dump_jsonl(os.path.join(data_dir, "data.jsonl"), records)
for record in records:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
dataset = DatasetFactory.load( dataset = DatasetFactory.load(
"sft", data_dir, window_size=8, tokenizer_path=tokenizer_path "sft", data_dir, window_size=8, tokenizer_path=tokenizer_path
@@ -706,13 +623,12 @@ def test_sft_jsonl_explicit_config_takes_priority(base_test_env):
"""When dataset_config.json exists, it overrides the default messages config.""" """When dataset_config.json exists, it overrides the default messages config."""
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer = base_test_env["tokenizer"] tokenizer = base_test_env["tokenizer"]
tokenizer.set_chat_template( tokenizer.set_chat_template(SIMPLE_CHAT_TEMPLATE)
"{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}"
)
tokenizer_path = _save_test_tokenizer(test_dir, tokenizer) tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
data_dir = _write_jsonl_dataset( data_dir = _write_text_dataset(
test_dir, test_dir,
"sft_explicit",
tokenizer_path, tokenizer_path,
[ [
{ {
@@ -748,10 +664,7 @@ def _write_grpo_jsonl(test_dir, tokenizer_path, records):
"""Write a GRPO JSONL dataset directory with config.""" """Write a GRPO JSONL dataset directory with config."""
data_dir = os.path.join(test_dir, "grpo_jsonl") data_dir = os.path.join(test_dir, "grpo_jsonl")
os.makedirs(data_dir, exist_ok=True) os.makedirs(data_dir, exist_ok=True)
_dump_jsonl(os.path.join(data_dir, "data.jsonl"), records)
with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
for rec in records:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
config = { config = {
"tokenizer_path": tokenizer_path, "tokenizer_path": tokenizer_path,
@@ -927,28 +840,15 @@ def test_grpo_multiple_records(base_test_env):
[torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)] [torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)]
for _ in range(n_records) for _ in range(n_records)
] ]
store = type( store = _grpo_fake_store(
"FakeStore", prompts=[torch.randint(0, 100, (10,)) for _ in range(n_records)],
(), responses=dummy_responses,
{ masks=[
"keys": ["prompts", "responses", "masks", "rewards"], [torch.ones(r.shape[0], dtype=torch.int64) for r in resps]
"num_records": n_records, for resps in dummy_responses
"token_count": 0, ],
"_data": { rewards=[torch.rand(G, dtype=torch.float32) for _ in range(n_records)],
"prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)], )
"responses": dummy_responses,
"masks": [
[torch.ones(r.shape[0], dtype=torch.int64) for r in resps]
for resps in dummy_responses
],
"rewards": [
torch.rand(G, dtype=torch.float32) for _ in range(n_records)
],
},
"fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
},
)()
dataset = GRPODataset(store=store) dataset = GRPODataset(store=store)
assert len(dataset) == n_records assert len(dataset) == n_records
@@ -965,9 +865,7 @@ def test_grpo_multiple_records(base_test_env):
def _write_dpo_jsonl(test_dir, records): def _write_dpo_jsonl(test_dir, records):
"""Write a raw DPO JSONL file (no dataset_config.json).""" """Write a raw DPO JSONL file (no dataset_config.json)."""
path = os.path.join(test_dir, "dpo.jsonl") path = os.path.join(test_dir, "dpo.jsonl")
with open(path, "w", encoding="utf-8") as f: _dump_jsonl(path, records)
for rec in records:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
return path return path
@@ -1075,12 +973,12 @@ def test_jsonl_store_eager_len_returns_token_count(base_test_env):
"""JsonlStore in eager mode: num_records reflects per-record count.""" """JsonlStore in eager mode: num_records reflects per-record count."""
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
data_dir = _write_jsonl_dataset( data_dir = _write_text_dataset(
test_dir, test_dir,
"jsonl_data",
tokenizer_path, tokenizer_path,
[{"text": "hello world"}, {"text": "foo bar"}], [{"text": "hello world"}, {"text": "foo bar"}],
config_overrides={ config_overrides={
"preprocessing": {"max_seq_len": 128, "min_chars": 0},
"output": {"position_ids_mode": "none"}, "output": {"position_ids_mode": "none"},
}, },
) )
tests/data/test_preprocess_builder.py
+29 -92
View File
@@ -14,7 +14,6 @@ from astrai.preprocessing.builder import (
) )
from tests.data.factories import ( from tests.data.factories import (
CHAT_SECTIONS, CHAT_SECTIONS,
INSTRUCTION_SECTIONS,
TEXT_SECTIONS, TEXT_SECTIONS,
make_chat_config, make_chat_config,
make_dpo_chat_config, make_dpo_chat_config,
@@ -35,9 +34,6 @@ def test_chat_simple(chat_tokenizer, builder):
} }
result = builder.build(item, config, chat_tokenizer) result = builder.build(item, config, chat_tokenizer)
assert result is not None assert result is not None
assert "sequence" in result
assert "loss_mask" in result
assert len(result["sequence"]) == len(result["loss_mask"])
ids = chat_tokenizer.decode(result["sequence"], skip_special_tokens=False) ids = chat_tokenizer.decode(result["sequence"], skip_special_tokens=False)
assert "system" in ids.lower() or "<|im_start|>system" in ids assert "system" in ids.lower() or "<|im_start|>system" in ids
@@ -51,21 +47,26 @@ def test_chat_simple(chat_tokenizer, builder):
def test_chat_mask_only_assistant(chat_tokenizer, builder): def test_chat_mask_only_assistant(chat_tokenizer, builder):
config = make_chat_config() config = make_chat_config()
item = { messages = [
"messages": [ {"role": "user", "content": "What is 2+2?"},
{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"},
{"role": "assistant", "content": "4"}, ]
] result = builder.build({"messages": messages}, config, chat_tokenizer)
}
result = builder.build(item, config, chat_tokenizer)
mask = result["loss_mask"] mask = result["loss_mask"]
ids = result["sequence"]
assert len(ids) == len(mask)
trained = [i for i, m in enumerate(mask) if m == 1] def template_ids(message):
masked = [i for i, m in enumerate(mask) if m == 0] rendered = chat_tokenizer.apply_chat_template(
assert len(trained) > 0 [message], tokenize=False, add_generation_prompt=False
assert len(masked) > 0 )
return chat_tokenizer.encode(rendered, add_special_tokens=False)
user_len = len(template_ids(messages[0]))
assistant_len = len(template_ids(messages[1]))
bos = 1 if chat_tokenizer.bos_token_id is not None else 0
assert len(mask) == bos + user_len + assistant_len
assert all(m == 0 for m in mask[: bos + user_len])
assert all(m == 1 for m in mask[bos + user_len :])
def test_chat_batch_matches_single(chat_tokenizer, builder): def test_chat_batch_matches_single(chat_tokenizer, builder):
@@ -166,14 +167,6 @@ def test_chat_truncation(chat_tokenizer, builder):
assert len(result["loss_mask"]) == len(result["sequence"]) assert len(result["loss_mask"]) == len(result["sequence"])
def test_instruction_basic(test_tokenizer, builder):
config = make_instruction_config()
item = {"prompt": "Translate to French: Hello", "response": "Bonjour"}
result = builder.build(item, config, test_tokenizer)
assert result is not None
assert len(result["sequence"]) == len(result["loss_mask"])
def test_instruction_batch_matches_single(test_tokenizer, builder): def test_instruction_batch_matches_single(test_tokenizer, builder):
config = make_instruction_config() config = make_instruction_config()
items = [ items = [
@@ -224,7 +217,6 @@ def test_text_basic(test_tokenizer, builder):
item = {"text": "Hello world. This is a test document."} item = {"text": "Hello world. This is a test document."}
result = builder.build(item, config, test_tokenizer) result = builder.build(item, config, test_tokenizer)
assert result is not None assert result is not None
assert "sequence" in result
assert len(result["sequence"]) > 0 assert len(result["sequence"]) > 0
assert "loss_mask" not in result assert "loss_mask" not in result
@@ -253,72 +245,17 @@ def test_text_truncation(test_tokenizer, builder):
assert len(result["sequence"]) <= 3 assert len(result["sequence"]) <= 3
def test_sectioned_chat(chat_tokenizer, builder): @pytest.mark.parametrize(
config = PipelineConfig( ("name", "builder_cls"),
input=InputConfig(sections=CHAT_SECTIONS), [
mask={"system": "mask", "user": "mask", "assistant": "train"}, ("single", SingleOutputMaskBuilder),
mask_default="mask", ("multi", MultiOutputMaskBuilder),
preprocessing=ProcessingConfig(max_seq_len=2048), ("sectioned", SectionedMaskBuilder),
) ],
item = { )
"messages": [ def test_factory_create(name, builder_cls):
{"role": "user", "content": "What is 2+2?"}, assert name in MaskBuilderFactory.list_registered()
{"role": "assistant", "content": "4"}, assert isinstance(MaskBuilderFactory.create(name), builder_cls)
]
}
result = builder.build(item, config, chat_tokenizer)
assert result is not None
assert len(result["sequence"]) == len(result["loss_mask"])
assert sum(result["loss_mask"]) > 0
assert 0 in result["loss_mask"]
def test_sectioned_instruction(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=INSTRUCTION_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
)
item = {"prompt": "Q: Why?", "response": "A: Because."}
result = builder.build(item, config, test_tokenizer)
assert result is not None
mask = result["loss_mask"]
assert mask[0] == 0
assert mask[-1] == 1
def test_sectioned_text(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
)
item = {"text": "Hello world, this is a test."}
result = builder.build(item, config, test_tokenizer)
assert result is not None
assert "loss_mask" not in result
def test_sectioned_text_too_short(test_tokenizer, builder):
config = PipelineConfig(
input=InputConfig(sections=TEXT_SECTIONS),
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
)
assert builder.build({"text": "short"}, config, test_tokenizer) is None
def test_factory_registered():
names = MaskBuilderFactory.list_registered()
assert "single" in names
assert "multi" in names
assert "sectioned" in names
def test_factory_create():
single = MaskBuilderFactory.create("single")
assert isinstance(single, SingleOutputMaskBuilder)
multi = MaskBuilderFactory.create("multi")
assert isinstance(multi, MultiOutputMaskBuilder)
sectioned = MaskBuilderFactory.create("sectioned")
assert isinstance(sectioned, SectionedMaskBuilder)
def test_dpo_chat_basic(chat_tokenizer, builder): def test_dpo_chat_basic(chat_tokenizer, builder):
tests/data/test_preprocess_pipeline.py
-43
View File
@@ -115,49 +115,6 @@ def test_full_text_pipeline(temp_dir, tokenizer_dir):
assert "loss_mask" not in meta assert "loss_mask" not in meta
def test_full_instruction_pipeline(temp_dir, tokenizer_dir):
jsonl_path = os.path.join(temp_dir, "instruct.jsonl")
with open(jsonl_path, "w", encoding="utf-8") as f:
f.write(
json.dumps(
{
"prompt": "Tell me a joke",
"response": "Why did the chicken cross the road?",
}
)
+ "\n"
)
f.write(
json.dumps(
{
"prompt": "What is AI?",
"response": "Artificial Intelligence is a field of computer science.",
}
)
+ "\n"
)
config = PipelineConfig(
input=InputConfig(sections=INSTRUCTION_SECTIONS),
mask={"prompt": "mask", "response": "train"},
mask_default="mask",
preprocessing=ProcessingConfig(max_seq_len=2048),
output=OutputConfig(storage_format="bin"),
)
out_dir = os.path.join(temp_dir, "output")
Pipeline(
config=config,
input_paths=[jsonl_path],
output_dir=out_dir,
tokenizer_path=tokenizer_dir,
).run()
meta = load_shard_meta(out_dir)
assert "sequence" in meta
assert "loss_mask" in meta
def test_dtype_override(temp_dir, tokenizer_dir): def test_dtype_override(temp_dir, tokenizer_dir):
jsonl_path = os.path.join(temp_dir, "data.jsonl") jsonl_path = os.path.join(temp_dir, "data.jsonl")
with open(jsonl_path, "w", encoding="utf-8") as f: with open(jsonl_path, "w", encoding="utf-8") as f:
tests/extension/test_kernel_mask.py
+30 -3
View File
@@ -1,5 +1,7 @@
"""Kernel-level mask dimension support (2D, 3D, 4D).""" """Kernel-level mask dimension support (2D, 3D, 4D)."""
import math
import torch import torch
from astrai.extension.ops.attention import attn_prefill from astrai.extension.ops.attention import attn_prefill
@@ -7,9 +9,27 @@ from tests.conftest import skip_no_kernel
from tests.extension.conftest import D from tests.extension.conftest import D
def _reference(q, k, v, mask):
"""fp32 masked GQA attention reference (True=keep)."""
b, s_q, h, d = q.shape
rep = h // k.shape[2]
qf = q.float().transpose(1, 2)
kf = k.float().repeat_interleave(rep, dim=2).transpose(1, 2)
vf = v.float().repeat_interleave(rep, dim=2).transpose(1, 2)
scores = qf @ kf.transpose(-1, -2) / math.sqrt(d)
if mask.dim() == 2:
mask = mask[:, None, None, :]
elif mask.dim() == 3:
mask = mask[:, None, :, :]
# 4D [batch, 1, q_len, kv_len] broadcasts over heads as-is
scores = scores.masked_fill(~mask, float("-inf"))
return (scores.softmax(dim=-1) @ vf).transpose(1, 2).to(q.dtype)
@skip_no_kernel @skip_no_kernel
def test_kernel_accepts_2d_mask(): def test_kernel_accepts_2d_mask():
"""Kernel should accept 2D mask [batch, kv_len].""" """2D mask [batch, kv_len] gates the softmax, not just parses."""
torch.manual_seed(11)
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1 batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8 kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16) q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
@@ -20,25 +40,31 @@ def test_kernel_accepts_2d_mask():
out = attn_prefill(q, k, v, mask=mask, is_causal=False) out = attn_prefill(q, k, v, mask=mask, is_causal=False)
assert out.shape == (batch, q_len, n_heads, D) assert out.shape == (batch, q_len, n_heads, D)
torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05)
@skip_no_kernel @skip_no_kernel
def test_kernel_accepts_3d_mask(): def test_kernel_accepts_3d_mask():
"""Kernel should accept 3D mask [batch, q_len, kv_len].""" """3D mask [batch, q_len, kv_len] applies per-query-row gating."""
torch.manual_seed(12)
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1 batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8 kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16) q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16) k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16) v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
mask = torch.ones(batch, q_len, kv_len, dtype=torch.bool, device="cuda") mask = torch.ones(batch, q_len, kv_len, dtype=torch.bool, device="cuda")
mask[:, 0, 5:] = False # differs per query row: only the 3D path can apply it
mask[:, 1, 6:] = False
out = attn_prefill(q, k, v, mask=mask, is_causal=False) out = attn_prefill(q, k, v, mask=mask, is_causal=False)
assert out.shape == (batch, q_len, n_heads, D) assert out.shape == (batch, q_len, n_heads, D)
torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05)
@skip_no_kernel @skip_no_kernel
def test_kernel_accepts_4d_mask(): def test_kernel_accepts_4d_mask():
"""Kernel should accept 4D mask [batch, n_heads, q_len, kv_len].""" """4D mask [batch, 1, q_len, kv_len] broadcasts over heads and gates."""
torch.manual_seed(13)
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1 batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8 kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16) q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
@@ -49,6 +75,7 @@ def test_kernel_accepts_4d_mask():
out = attn_prefill(q, k, v, mask=mask, is_causal=False) out = attn_prefill(q, k, v, mask=mask, is_causal=False)
assert out.shape == (batch, q_len, n_heads, D) assert out.shape == (batch, q_len, n_heads, D)
torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05)
@skip_no_kernel @skip_no_kernel
tests/extension/test_linear_dispatch.py
-5
View File
@@ -6,7 +6,6 @@ import torch
import torch.nn.functional as F import torch.nn.functional as F
from astrai.extension import is_available, linear from astrai.extension import is_available, linear
from astrai.extension.backend import linear as public_linear
from astrai.extension.dispatch import explain, op_backend, resolve from astrai.extension.dispatch import explain, op_backend, resolve
# The package attribute ``linear`` is the dispatched function; reach the # The package attribute ``linear`` is the dispatched function; reach the
@@ -36,10 +35,6 @@ def _routes_to_gemv(monkeypatch, x, weight, bias=None) -> bool:
return linear(x, weight, bias) is sentinel return linear(x, weight, bias) is sentinel
def test_linear_backend_is_public():
assert linear is public_linear
def test_model_linear_routes_through_backend(monkeypatch): def test_model_linear_routes_through_backend(monkeypatch):
sentinel = torch.randn(2, 4) sentinel = torch.randn(2, 4)
tests/extension/test_swiglu_dispatch.py
+5 -2
View File
@@ -38,10 +38,13 @@ def test_cpu_and_training_calls_fall_back_with_gradients(monkeypatch):
def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog): def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode") monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode")
x = torch.randn(2, 8)
up_weight = torch.randn(4, 8)
gate_weight = torch.randn(4, 8)
with caplog.at_level(logging.WARNING): with caplog.at_level(logging.WARNING):
actual = swiglu(torch.randn(2, 8), torch.randn(4, 8), torch.randn(4, 8)) actual = swiglu(x, up_weight, gate_weight)
assert actual.shape == (2, 4)
assert "using auto" in caplog.text assert "using auto" in caplog.text
torch.testing.assert_close(actual, reference_swiglu(x, up_weight, gate_weight))
def test_mlp_routes_through_swiglu_backend(monkeypatch): def test_mlp_routes_through_swiglu_backend(monkeypatch):
tests/inference/test_cache.py
+1 -14
View File
@@ -34,13 +34,6 @@ def _make_task_cache(pool: PagePool) -> TaskCacheManager:
# ---- page_hash ---- # ---- page_hash ----
def test_page_hash_full_page():
token_ids = list(range(256))
h = page_hash(token_ids, 0, 64)
assert isinstance(h, int)
assert h >= 0
def test_page_hash_different_page_differs(): def test_page_hash_different_page_differs():
token_ids = list(range(256)) token_ids = list(range(256))
assert page_hash(token_ids, 0, 64) != page_hash(token_ids, 1, 64) assert page_hash(token_ids, 0, 64) != page_hash(token_ids, 1, 64)
@@ -121,19 +114,13 @@ def test_prefix_cache_ignores_partial_last_page():
def test_prefix_cache_on_evict_clears_mappings(): def test_prefix_cache_on_evict_clears_mappings():
prefix = RadixCache(64) prefix = RadixCache(64)
assert not prefix.has_page(0)
prefix.record(0, list(range(64)), 0) prefix.record(0, list(range(64)), 0)
assert prefix.has_page(0) assert prefix.has_page(0)
prefix.evict(0) prefix.evict(0)
assert not prefix.has_page(0) assert not prefix.has_page(0)
def test_prefix_cache_has_page():
prefix = RadixCache(64)
assert not prefix.has_page(0)
prefix.record(0, list(range(64)), 0)
assert prefix.has_page(0)
def test_prefix_cache_does_not_reuse_page_without_parent_prefix(): def test_prefix_cache_does_not_reuse_page_without_parent_prefix():
prefix = RadixCache(2) prefix = RadixCache(2)
prefix.record(0, [1, 2, 3, 4], 0) prefix.record(0, [1, 2, 3, 4], 0)
tests/inference/test_engine.py
-6
View File
@@ -20,12 +20,6 @@ def _make_engine_mocks(decode=None):
return mock_model, mock_tokenizer return mock_model, mock_tokenizer
def test_result_append_single():
r = GenerateResult(count=1)
r.append("hello", 0)
assert r.results[0] == "hello"
def test_result_append_multiple_tasks(): def test_result_append_multiple_tasks():
r = GenerateResult(count=3) r = GenerateResult(count=3)
r.append("a", 0) r.append("a", 0)
tests/inference/test_protocol.py
-25
View File
@@ -71,31 +71,6 @@ def test_check_empty_sequences():
assert sc.check("hello") is None assert sc.check("hello") is None
def test_gen_context_defaults():
ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
assert ctx.completion_tokens == 0
def test_gen_context_fields_mutable():
ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
ctx.completion_tokens = 42
assert ctx.completion_tokens == 42
def test_stop_info_defaults():
s = StopInfo()
assert s.matched is None
assert s.body == ""
assert s.yielded == ""
def test_stop_info_with_values():
s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
assert s.matched == "stop"
assert s.body == "hello stop"
assert s.yielded == "hello "
def test_openai_prepare_returns_prompt_ctx_stops(): def test_openai_prepare_returns_prompt_ctx_stops():
builder = _make_openai_builder() builder = _make_openai_builder()
req = MagicMock() req = MagicMock()
tests/inference/test_sample.py
+2 -1
View File
@@ -62,8 +62,9 @@ def test_top_p_nucleus_filtering():
logits = torch.tensor([[10.0, 1.0, 1.0, 1.0, 1.0]]) logits = torch.tensor([[10.0, 1.0, 1.0, 1.0, 1.0]])
s = TopPStrategy(top_p=0.5) s = TopPStrategy(top_p=0.5)
result = s.apply(logits.clone(), filter_value=-1e9) result = s.apply(logits.clone(), filter_value=-1e9)
# The dominant logit alone exceeds the nucleus mass; the rest are filtered.
kept = (result > -1e9).sum().item() kept = (result > -1e9).sum().item()
assert kept >= 1 assert kept == 1
def test_top_p_skip_when_one(): def test_top_p_skip_when_one():
tests/inference/test_scheduler.py
+34 -69
View File
@@ -40,18 +40,32 @@ def mock_model_and_tokenizer():
return mock_model, mock_tokenizer return mock_model, mock_tokenizer
def _make_mock_scheduler(mock_model_and_tokenizer):
"""Build a CPU scheduler over mocks, patching scheduler-internal imports."""
mock_model, mock_tokenizer = mock_model_and_tokenizer
with (
patch("astrai.inference.scheduler.AutoModel"),
patch("astrai.inference.scheduler.AutoTokenizer"),
):
return InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
max_batch_size=4,
device="cpu",
)
def _run_threads(*workers, timeout=10.0):
threads = [threading.Thread(target=worker) for worker in workers]
for t in threads:
t.start()
for t in threads:
t.join(timeout=timeout)
def test_scheduler_concurrent_add_task(mock_model_and_tokenizer): def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
"""Test concurrent add_task operations.""" """Test concurrent add_task operations."""
mock_model, mock_tokenizer = mock_model_and_tokenizer scheduler = _make_mock_scheduler(mock_model_and_tokenizer)
with patch("astrai.inference.scheduler.AutoModel"):
with patch("astrai.inference.scheduler.AutoTokenizer"):
scheduler = InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
max_batch_size=4,
device="cpu",
)
results = {"task_ids": [], "errors": []} results = {"task_ids": [], "errors": []}
lock = threading.Lock() lock = threading.Lock()
@@ -65,13 +79,7 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
except Exception as e: except Exception as e:
results["errors"].append(str(e)) results["errors"].append(str(e))
threads = [threading.Thread(target=add_task_worker, args=(i,)) for i in range(5)] _run_threads(*(lambda wid=i: add_task_worker(wid) for i in range(5)))
for t in threads:
t.start()
for t in threads:
t.join()
scheduler.stop() scheduler.stop()
@@ -205,16 +213,7 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer): def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
"""Test concurrent add and remove task operations.""" """Test concurrent add and remove task operations."""
mock_model, mock_tokenizer = mock_model_and_tokenizer scheduler = _make_mock_scheduler(mock_model_and_tokenizer)
with patch("astrai.inference.scheduler.AutoModel"):
with patch("astrai.inference.scheduler.AutoTokenizer"):
scheduler = InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
max_batch_size=4,
device="cpu",
)
results = {"added": [], "removed": [], "errors": []} results = {"added": [], "removed": [], "errors": []}
add_ready = threading.Event() add_ready = threading.Event()
@@ -238,14 +237,7 @@ def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
except Exception as e: except Exception as e:
results["errors"].append(f"Remove: {str(e)}") results["errors"].append(f"Remove: {str(e)}")
add_thread = threading.Thread(target=add_worker) _run_threads(add_worker, remove_worker)
remove_thread = threading.Thread(target=remove_worker)
add_thread.start()
remove_thread.start()
add_thread.join()
remove_thread.join()
scheduler.stop() scheduler.stop()
assert len(results["errors"]) == 0, f"Errors: {results['errors']}" assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
@@ -254,16 +246,7 @@ def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer): def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
"""Test concurrent get_stats operations.""" """Test concurrent get_stats operations."""
mock_model, mock_tokenizer = mock_model_and_tokenizer scheduler = _make_mock_scheduler(mock_model_and_tokenizer)
with patch("astrai.inference.scheduler.AutoModel"):
with patch("astrai.inference.scheduler.AutoTokenizer"):
scheduler = InferenceScheduler(
model=mock_model,
tokenizer=mock_tokenizer,
max_batch_size=4,
device="cpu",
)
results = {"stats": [], "errors": []} results = {"stats": [], "errors": []}
started = threading.Event() started = threading.Event()
@@ -287,17 +270,9 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
except Exception as e: except Exception as e:
results["errors"].append(f"Get stats: {str(e)}") results["errors"].append(f"Get stats: {str(e)}")
add_thread = threading.Thread(target=add_tasks) _run_threads(add_tasks, get_stats)
stats_thread = threading.Thread(target=get_stats)
add_thread.start()
stats_thread.start()
add_thread.join()
stats_done.wait(timeout=5.0)
scheduler.stop() scheduler.stop()
stats_done.wait(timeout=5.0)
stats_thread.join()
assert len(results["errors"]) == 0, f"Errors: {results['errors']}" assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
assert len(results["stats"]) == 50 assert len(results["stats"]) == 50
@@ -504,16 +479,6 @@ def test_ragged_prefill_matches_sequential_greedy_tokens_and_logprobs(device):
scheduler.stop() scheduler.stop()
def test_run_batch_respects_max_tokens(device):
scheduler, _tok, _model = _make_real_scheduler(device)
try:
prompts = [[10, 20, 30]]
results = scheduler.run_batch(prompts, max_tokens=3, temperature=1.0)
assert len(results[0]) <= 3
finally:
scheduler.stop()
def test_run_batch_zero_max_tokens_returns_empty(device): def test_run_batch_zero_max_tokens_returns_empty(device):
scheduler, _tok, _model = _make_real_scheduler(device) scheduler, _tok, _model = _make_real_scheduler(device)
try: try:
@@ -526,12 +491,12 @@ def test_run_batch_stop_id_terminates(device):
"""A token matching stop_ids terminates generation for that prompt.""" """A token matching stop_ids terminates generation for that prompt."""
scheduler, _tok, _model = _make_real_scheduler(device) scheduler, _tok, _model = _make_real_scheduler(device)
try: try:
# Make every token a stop id: generation must end after exactly
# one token (the stop token itself) instead of running to max_tokens.
scheduler._task_mgr.tokenizer.stop_ids = list(range(200))
prompts = [[10, 20, 30]] prompts = [[10, 20, 30]]
results = scheduler.run_batch(prompts, max_tokens=32, temperature=1.0) results = scheduler.run_batch(prompts, max_tokens=32, temperature=1.0)
# If stop token 2 was produced, it is the last token assert len(results[0]) == 1
if results[0] and results[0][-1] == 2:
# No tokens after stop should exist (since we terminate)
assert 2 not in results[0][:-1]
finally: finally:
scheduler.stop() scheduler.stop()
tests/inference/test_tool_parser.py
+4 -29
View File
@@ -252,13 +252,15 @@ def test_feed_with_tools_constructor():
tools = [{"type": "function", "function": {"name": "get_weather"}}] tools = [{"type": "function", "function": {"name": "get_weather"}}]
parser = SimpleJsonToolParser(tools=tools, tool_choice="auto") parser = SimpleJsonToolParser(tools=tools, tool_choice="auto")
deltas = parser.feed('{"name": "get_weather", "arguments": {"city": "BJ"}}') deltas = parser.feed('{"name": "get_weather", "arguments": {"city": "BJ"}}')
assert len(deltas) > 0 tc_deltas = [d for d in deltas if "tool_calls" in d]
assert tc_deltas[0]["tool_calls"][0]["function"]["name"] == "get_weather"
def test_feed_content_after_tool_call_is_not_emitted(): def test_feed_content_after_tool_call_is_not_emitted():
parser = SimpleJsonToolParser() parser = SimpleJsonToolParser()
parser.feed('{"name": "f", "arguments": {}} trailing text') deltas = parser.feed('{"name": "f", "arguments": {}} trailing text')
assert parser.has_tool_calls assert parser.has_tool_calls
assert not any("trailing" in d.get("content", "") for d in deltas)
def _simulate_streaming(parser, text): def _simulate_streaming(parser, text):
@@ -513,10 +515,6 @@ def test_factory_create_passes_tools():
assert parser.tool_choice == "required" assert parser.tool_choice == "required"
def test_factory_list_registered():
assert "simple_json" in ToolParserFactory.list_registered()
def test_factory_create_with_tools_only(): def test_factory_create_with_tools_only():
tools = [ tools = [
{ {
@@ -544,29 +542,6 @@ def test_feed_token_ids_do_not_affect_parsing():
) )
def test_parser_uses_token_ids_for_detection():
class TokenIdParser(BaseToolParser):
def __init__(self, tools=None, tool_choice="auto"):
super().__init__(tools, tool_choice)
self._detections = 0
def feed(self, body, current_token_ids=None, delta_token_ids=None):
if current_token_ids and 999 in current_token_ids:
self._detections += 1
return []
def parse_complete(self, body):
return None
@property
def has_tool_calls(self):
return self._detections > 0
parser = TokenIdParser()
parser.feed("hello", current_token_ids=[1, 999, 3])
assert parser.has_tool_calls
def test_streaming_partial_name_prefix_never_leaks_into_content(): def test_streaming_partial_name_prefix_never_leaks_into_content():
parser = SimpleJsonToolParser() parser = SimpleJsonToolParser()
parts = ["Hello ", '{"', '{"n', '{"na', '{"name"'] parts = ["Hello ", '{"', '{"n', '{"na', '{"name"']
tests/module/test_forward_configs.py
+3 -1
View File
@@ -369,7 +369,9 @@ def test_moe_component_forward_returns_ffn_output():
output = moe(torch.randn(2, 8, 8)) output = moe(torch.randn(2, 8, 8))
assert output["hidden_states"].shape == (2, 8, 8) assert output["hidden_states"].shape == (2, 8, 8)
assert output["aux_loss"] is not None assert output["aux_loss"].ndim == 0
assert output["aux_loss"].requires_grad
assert torch.isfinite(output["aux_loss"])
@pytest.mark.parametrize("decoder_sparse_step", [0, -1]) @pytest.mark.parametrize("decoder_sparse_step", [0, -1])
tests/module/test_lora.py
+3 -1
View File
@@ -332,4 +332,6 @@ def test_collect_lora_info():
info = _collect_lora_info(model) info = _collect_lora_info(model)
assert "q_proj" in info assert "q_proj" in info
assert "o_proj" in info assert "o_proj" in info
assert "q_proj" in info # each layer has one # every decoder layer contributes one of each attention projection
assert len(info["q_proj"]) == len(info["o_proj"]) == 2
assert all(name.endswith("attention.q_proj") for name in info["q_proj"])
tests/optim/test_mano_adamw.py
+21 -7
View File
@@ -34,18 +34,32 @@ def test_mano_one_step_projects_to_tangent_space():
def test_mano_alternates_projection_axis(): def test_mano_alternates_projection_axis():
param = torch.nn.Parameter(torch.eye(4) * 3.0) """Step 0 projects along dim 0, step 1 along dim 1 (steps % 2)."""
param.grad = torch.ones(4, 4) original = torch.tensor([[3.0, 4.0], [0.0, 2.0]])
param = torch.nn.Parameter(original.clone())
eps = 1e-8
optimizer = Mano(
[param], lr=0.1, momentum=0.0, nesterov=False, eps=eps, weight_decay=0.0
)
optimizer = Mano([param], lr=0.1, momentum=0.0, nesterov=False) param.grad = torch.ones(2, 2)
optimizer.step() optimizer.step()
dim_step0 = 0 after_step0 = param.detach().clone()
param.grad = torch.ones(4, 4) param.grad = torch.ones(2, 2)
optimizer.step() optimizer.step()
dim_step1 = 1
assert dim_step0 != dim_step1 grad = torch.ones(2, 2)
def projected(after, dim):
tangent = grad - (torch.sum(grad * after, dim=dim, keepdim=True) * after)
direction = tangent / (torch.norm(tangent, p=2, dim=dim, keepdim=True) + eps)
adjusted_lr = 0.1 * 0.2 * math.sqrt(direction.shape[dim])
return after - adjusted_lr * direction
torch.testing.assert_close(param.detach(), projected(after_step0, dim=1))
# The dim=1 result is distinct, so the assertion above pins the axis.
assert not torch.allclose(param.detach(), projected(after_step0, dim=0))
def test_mano_rejects_non_2d_parameters(): def test_mano_rejects_non_2d_parameters():
tests/serialization/test_hf_adapter.py
+51 -91
View File
@@ -113,6 +113,44 @@ def to_hf_keys(state_dict, head_dim=None):
return out return out
MOE_KWARGS = {
"ffn_type": "moe",
"n_routed_experts": 2,
"n_shared_experts": 1,
"n_activated_experts": 1,
"moe_intermediate_size": 16,
"shared_expert_intermediate_size": 16,
}
def _hf_keyed_state_dict(model, cfg):
return to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads)
def _assert_hf_roundtrip(cfg, convert_cfg=None):
"""HF-keyed weights of a fresh model must convert back exactly."""
model = AutoRegressiveLM(cfg)
sd = model.state_dict()
converted = convert_hf_weights(_hf_keyed_state_dict(model, cfg), convert_cfg or cfg)
assert_state_dicts_equal(converted, sd)
def _assert_hf_directory_load(tmp_path, cfg, raw_config):
"""Save HF-format weights and check from_pretrained reproduces logits."""
model = AutoRegressiveLM(cfg).eval()
save_model(
config=raw_config,
state_dict=_hf_keyed_state_dict(model, cfg),
save_directory=str(tmp_path),
)
loaded = AutoModel.from_pretrained(tmp_path).eval()
input_ids = torch.randint(0, cfg.vocab_size, (1, 8))
with torch.no_grad():
torch.testing.assert_close(
loaded(input_ids)["logits"], model(input_ids)["logits"]
)
def test_convert_hf_config_llama(): def test_convert_hf_config_llama():
cfg = convert_hf_config(LLAMA_RAW) cfg = convert_hf_config(LLAMA_RAW)
assert cfg["model_type"] == "autoregressive_lm" assert cfg["model_type"] == "autoregressive_lm"
@@ -168,31 +206,14 @@ def test_adapt_config_passthrough():
def test_convert_hf_weights_dense_roundtrip(): def test_convert_hf_weights_dense_roundtrip():
cfg = make_tiny_config() _assert_hf_roundtrip(make_tiny_config())
model = AutoRegressiveLM(cfg)
converted = convert_hf_weights(
to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
)
assert_state_dicts_equal(converted, model.state_dict())
def test_convert_hf_weights_moe_roundtrip(): def test_convert_hf_weights_moe_roundtrip():
cfg = make_tiny_config( cfg = make_tiny_config(**MOE_KWARGS)
ffn_type="moe",
n_routed_experts=2,
n_shared_experts=1,
n_activated_experts=1,
moe_intermediate_size=16,
shared_expert_intermediate_size=16,
)
model = AutoRegressiveLM(cfg)
hf_raw = convert_hf_config(MOE_RAW) hf_raw = convert_hf_config(MOE_RAW)
hf_cfg = ConfigFactory.load(hf_raw) hf_cfg = ConfigFactory.load(hf_raw)
converted = convert_hf_weights( _assert_hf_roundtrip(cfg, convert_cfg=hf_cfg)
to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads),
hf_cfg,
)
assert_state_dicts_equal(converted, model.state_dict())
def test_convert_hf_weights_keeps_astrai_keys(): def test_convert_hf_weights_keeps_astrai_keys():
@@ -239,66 +260,28 @@ def test_convert_hf_config_gemma_enables_qk_norm():
def test_convert_hf_weights_moe_with_dense_layers_roundtrip(): def test_convert_hf_weights_moe_with_dense_layers_roundtrip():
cfg = make_tiny_config( _assert_hf_roundtrip(
ffn_type="moe", make_tiny_config(**MOE_KWARGS, mlp_only_layers=[0], decoder_sparse_step=1)
n_routed_experts=2,
n_shared_experts=1,
n_activated_experts=1,
moe_intermediate_size=16,
shared_expert_intermediate_size=16,
mlp_only_layers=[0],
decoder_sparse_step=1,
) )
model = AutoRegressiveLM(cfg)
converted = convert_hf_weights(
to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
)
assert_state_dicts_equal(converted, model.state_dict())
def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip(): def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
cfg = make_tiny_config( cfg = make_tiny_config(**MOE_KWARGS)
ffn_type="moe",
n_routed_experts=2,
n_shared_experts=1,
n_activated_experts=1,
moe_intermediate_size=16,
shared_expert_intermediate_size=16,
)
model = AutoRegressiveLM(cfg) model = AutoRegressiveLM(cfg)
hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads)
hf_sd = { hf_sd = {
k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items() k.replace("shared_experts.", "shared_expert.", 1): v
for k, v in _hf_keyed_state_dict(model, cfg).items()
} }
converted = convert_hf_weights(hf_sd, cfg) converted = convert_hf_weights(hf_sd, cfg)
assert_state_dicts_equal(converted, model.state_dict()) assert_state_dicts_equal(converted, model.state_dict())
def test_convert_hf_weights_gemma_qk_norm_roundtrip(): def test_convert_hf_weights_gemma_qk_norm_roundtrip():
cfg = make_tiny_config(use_qk_norm=True) _assert_hf_roundtrip(make_tiny_config(use_qk_norm=True))
model = AutoRegressiveLM(cfg)
converted = convert_hf_weights(
to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
)
assert_state_dicts_equal(converted, model.state_dict())
def test_from_pretrained_hf_directory(tmp_path): def test_from_pretrained_hf_directory(tmp_path):
cfg = make_tiny_config() _assert_hf_directory_load(tmp_path, make_tiny_config(), LLAMA_RAW)
model = AutoRegressiveLM(cfg).eval()
save_model(
config=LLAMA_RAW,
state_dict=to_hf_keys(
model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
),
save_directory=str(tmp_path),
)
loaded = AutoModel.from_pretrained(tmp_path).eval()
input_ids = torch.randint(0, cfg.vocab_size, (1, 8))
with torch.no_grad():
torch.testing.assert_close(
loaded(input_ids)["logits"], model(input_ids)["logits"]
)
def test_from_pretrained_astrai_directory(tmp_path): def test_from_pretrained_astrai_directory(tmp_path):
@@ -332,9 +315,7 @@ def test_from_pretrained_weights_format_astrai_rejects_hf(tmp_path):
model = AutoRegressiveLM(cfg) model = AutoRegressiveLM(cfg)
save_model( save_model(
config=LLAMA_RAW, config=LLAMA_RAW,
state_dict=to_hf_keys( state_dict=_hf_keyed_state_dict(model, cfg),
model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
),
save_directory=str(tmp_path), save_directory=str(tmp_path),
) )
with pytest.raises(ValueError): with pytest.raises(ValueError):
@@ -355,7 +336,7 @@ def test_from_pretrained_invalid_weights_format(tmp_path):
def test_from_pretrained_hf_directory_sharded(tmp_path): def test_from_pretrained_hf_directory_sharded(tmp_path):
cfg = make_tiny_config() cfg = make_tiny_config()
model = AutoRegressiveLM(cfg).eval() model = AutoRegressiveLM(cfg).eval()
hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads) hf_sd = _hf_keyed_state_dict(model, cfg)
keys = sorted(hf_sd) keys = sorted(hf_sd)
split = len(keys) // 2 split = len(keys) // 2
shard_a = {k: hf_sd[k] for k in keys[:split]} shard_a = {k: hf_sd[k] for k in keys[:split]}
@@ -507,25 +488,4 @@ def test_hf_import_qk_norm_matches_norm_before_rope_reference():
def test_from_pretrained_hf_directory_with_moe(tmp_path): def test_from_pretrained_hf_directory_with_moe(tmp_path):
cfg = make_tiny_config( _assert_hf_directory_load(tmp_path, make_tiny_config(**MOE_KWARGS), MOE_RAW)
ffn_type="moe",
n_routed_experts=2,
n_shared_experts=1,
n_activated_experts=1,
moe_intermediate_size=16,
shared_expert_intermediate_size=16,
)
model = AutoRegressiveLM(cfg).eval()
save_model(
config=MOE_RAW,
state_dict=to_hf_keys(
model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
),
save_directory=str(tmp_path),
)
loaded = AutoModel.from_pretrained(tmp_path).eval()
input_ids = torch.randint(0, cfg.vocab_size, (1, 8))
with torch.no_grad():
torch.testing.assert_close(
loaded(input_ids)["logits"], model(input_ids)["logits"]
)
tests/test_serve_cli.py
-105
View File
@@ -1,105 +0,0 @@
"""Unit tests for the serving CLI YAML merge logic."""
import click
import pytest
import torch
from click.testing import CliRunner
from scripts.tools.server import (
_merge_yaml_into_kwargs,
_resolve_server_config,
server_command,
)
def _passed() -> dict:
return {
"host": "0.0.0.0",
"port": 8000,
"reload": False,
"param_path": None,
"device": "cuda",
"dtype": "bfloat16",
"max_batch_size": 16,
"max_seq_len": None,
}
def test_yaml_overrides_click_defaults_but_not_explicit_cli(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text(
"server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n",
encoding="utf-8",
)
merged = _merge_yaml_into_kwargs(
str(config_path), _passed(), explicit_keys={"device"}
)
assert merged["device"] == "cuda"
assert merged["dtype"] == "float16"
assert merged["max_batch_size"] == 8
def test_resolve_config_yaml_wins_by_default(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text(
"server:\n port: 9000\n max_seq_len: 2048\n",
encoding="utf-8",
)
resolved = _resolve_server_config(str(config_path), _passed())
assert resolved["port"] == 9000
assert resolved["max_seq_len"] == 2048
assert resolved["device"] == "cuda"
assert resolved["dtype"] == "bfloat16"
def test_resolve_config_rejects_bad_dtype(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8")
with pytest.raises(click.UsageError, match="server.dtype"):
_resolve_server_config(str(config_path), _passed())
def test_server_command_rejects_bad_yaml_dtype(tmp_path):
config_path = tmp_path / "serve.yaml"
config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8")
result = CliRunner().invoke(server_command, ["--config", str(config_path)])
assert result.exit_code == 2
assert "server.dtype" in result.output
def test_server_command_merges_yaml_and_cli(tmp_path, monkeypatch):
"""Full CLI path: YAML values apply, explicit CLI flags override, args reach run_server."""
config_path = tmp_path / "serve.yaml"
config_path.write_text(
"server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n",
encoding="utf-8",
)
captured = {}
def fake_run_server(**kwargs):
captured.update(kwargs)
monkeypatch.setattr("scripts.tools.server.run_server", fake_run_server)
result = CliRunner().invoke(
server_command,
["--config", str(config_path), "--max_batch_size", "32"],
)
assert result.exit_code == 0, result.output
assert captured["device"] == "cpu"
assert captured["dtype"] == torch.float16
assert captured["max_batch_size"] == 32
assert captured["port"] == 8000
def test_config_option_rejects_missing_file(tmp_path):
result = CliRunner().invoke(
server_command, ["--config", str(tmp_path / "nope.yaml")]
)
assert result.exit_code == 2
tests/test_serve_runtime.py
-82
View File
@@ -1,82 +0,0 @@
"""Unit tests for the serving runtime configuration parser."""
import pytest
from scripts.docker.serve_runtime import load_runtime
def _write(tmp_path, body: str) -> str:
config_path = tmp_path / "serve.yaml"
config_path.write_text(body, encoding="utf-8")
return str(config_path)
def test_runtime_exports_defaults(tmp_path):
config_path = _write(
tmp_path,
"runtime:\n"
" port: 8000\n"
" paths:\n"
" param: ./params\n"
"server:\n"
" device: cuda\n",
)
runtime = load_runtime(config_path)
assert runtime["SERVE_PORT"] == "8000"
assert runtime["SERVE_CONTAINER_PORT"] == "8000"
assert runtime["SERVE_PARAM_DIR"] == str((tmp_path / "params").resolve())
assert runtime["SERVE_GPU_ENABLED"] == "true"
assert "CUDA_VISIBLE_DEVICES" not in runtime
assert runtime["SERVE_DEVICE"] == "cuda"
assert runtime["CUDA_TAG"] == "cu128"
assert runtime["SERVE_JOB_NAME"] == ""
def test_runtime_gpu_disabled_requires_cpu(tmp_path):
config_path = _write(
tmp_path,
"runtime:\n gpu:\n enabled: false\nserver:\n device: cuda\n",
)
with pytest.raises(ValueError, match="server.device must be 'cpu'"):
load_runtime(config_path)
def test_runtime_gpu_devices_single_and_ports(tmp_path):
config_path = _write(
tmp_path,
"runtime:\n"
" gpu:\n"
" devices: [1]\n"
" port: 8080\n"
"server:\n"
" port: 9000\n"
" device: cuda\n",
)
runtime = load_runtime(config_path)
assert runtime["SERVE_PORT"] == "8080"
assert runtime["SERVE_CONTAINER_PORT"] == "9000"
assert runtime["CUDA_VISIBLE_DEVICES"] == "1"
def test_runtime_gpu_devices_rejects_multi(tmp_path):
config_path = _write(
tmp_path,
"runtime:\n gpu:\n devices: [0, 1]\n",
)
with pytest.raises(ValueError, match="single-device"):
load_runtime(config_path)
def test_runtime_port_out_of_range(tmp_path):
config_path = _write(tmp_path, "runtime:\n port: 70000\n")
with pytest.raises(ValueError, match="between 1 and 65535"):
load_runtime(config_path)
def test_runtime_environment_export(tmp_path):
config_path = _write(
tmp_path,
"runtime:\n environment:\n TOKENIZERS_PARALLELISM: 'false'\n",
)
runtime = load_runtime(config_path)
assert runtime["environment"] == {"TOKENIZERS_PARALLELISM": "false"}
tests/test_train_cli.py
-62
View File
@@ -1,62 +0,0 @@
import re
from click.testing import CliRunner
from scripts.tools.train import _merge_yaml_into_kwargs, train_command
def test_yaml_overrides_click_defaults_but_not_explicit_cli(tmp_path):
config_path = tmp_path / "train.yaml"
config_path.write_text(
"training:\n"
" optimizer: nora_nadamw\n"
" max_lr: 0.0002\n"
" nora_lr: 0.004\n"
" batch_per_device: 8\n",
encoding="utf-8",
)
click_values = {
"optimizer": "nora_nadamw",
"max_lr": 3e-4,
"nora_lr": 5e-3,
"batch_per_device": 16,
}
merged = _merge_yaml_into_kwargs(
str(config_path), click_values, explicit_keys={"batch_per_device"}
)
assert merged["max_lr"] == 2e-4
assert merged["nora_lr"] == 4e-3
assert merged["batch_per_device"] == 16
def test_train_dry_run_uses_yaml_then_explicit_cli(tmp_path):
data_path = tmp_path / "data"
model_path = tmp_path / "model"
data_path.mkdir()
model_path.mkdir()
config_path = tmp_path / "train.yaml"
config_path.write_text(
"data:\n"
f" data_root_path: {data_path}\n"
"model:\n"
f" param_path: {model_path}\n"
"training:\n"
" train_type: seq\n"
" optimizer: nora_nadamw\n"
" max_lr: 0.0002\n"
" nora_lr: 0.004\n"
" batch_per_device: 8\n",
encoding="utf-8",
)
result = CliRunner().invoke(
train_command,
["--config", str(config_path), "--dry-run", "--batch_per_device", "16"],
)
assert result.exit_code == 0, result.output
assert re.search(r"Optimizer\s+: nora_nadamw", result.output)
assert re.search(r"Batch/device\s+: 16", result.output)
assert re.search(r"Max LR\s+: 0.0002", result.output)
tests/trainer/test_callbacks.py
+6 -2
View File
@@ -23,9 +23,9 @@ def test_gradient_checkpointing_enable_disable(test_model):
for layer in model.layers: for layer in model.layers:
callback._enable(layer) callback._enable(layer)
for layer in model.layers: for i, layer in enumerate(model.layers):
assert hasattr(layer, "_original_forward") assert hasattr(layer, "_original_forward")
assert layer.forward is not originals[0] assert layer.forward is not originals[i]
for layer in model.layers: for layer in model.layers:
callback._disable(layer) callback._disable(layer)
@@ -110,6 +110,10 @@ def test_gradient_checkpointing_trainer_integration(
) )
trainer = Trainer(train_config) trainer = Trainer(train_config)
gc_callbacks = [
c for c in trainer.callbacks if isinstance(c, GradientCheckpointingCallback)
]
assert gc_callbacks and gc_callbacks[0].modules == (DecoderBlock,)
trainer.train() trainer.train()
tests/trainer/test_rollout.py
+71 -96
View File
@@ -58,6 +58,46 @@ def _make_instruction_batch(n=2):
return {"instruction": instructions, "input": inputs} return {"instruction": instructions, "input": inputs}
def _blocking_hook(original, started, release):
"""Wrap a hook so it signals ``started`` then blocks until ``release``."""
def hook(*args, **kwargs):
started.set()
assert release.wait(timeout=5)
return original(*args, **kwargs)
return hook
def _assert_interleaved(first, second, *, started, release, finished):
"""Run ``first`` until it blocks, then assert ``second`` cannot finish
while ``first`` holds the lock; release, join both, and surface errors."""
errors = []
def run_safely(fn):
def run():
try:
fn()
except BaseException as exc:
errors.append(exc)
return run
first_thread = threading.Thread(target=run_safely(first))
second_thread = threading.Thread(target=run_safely(second))
first_thread.start()
assert started.wait(timeout=5)
second_thread.start()
assert not finished.wait(timeout=0.1)
release.set()
first_thread.join(timeout=5)
second_thread.join(timeout=5)
assert not first_thread.is_alive()
assert not second_thread.is_alive()
assert errors == []
def test_raw_rollout_fields(): def test_raw_rollout_fields():
r = RawRollout( r = RawRollout(
prompts=torch.zeros(2, 4, dtype=torch.long), prompts=torch.zeros(2, 4, dtype=torch.long),
@@ -96,13 +136,6 @@ def test_base_reward_model_is_abstract():
BaseRewardModel() BaseRewardModel()
def test_constant_reward_model_shape():
rm = ConstantRewardModel(0.5)
out = rm.score(["a", "b"], [["x", "y", "z"], ["p", "q", "r"]])
assert out.shape == (2, 3)
assert torch.all(out == 0.5)
def _make_generator(device, **kw): def _make_generator(device, **kw):
model, _ = make_model(device, max_position_embeddings=128) model, _ = make_model(device, max_position_embeddings=128)
tokenizer = FakeTokenizer(with_chat_template=True) tokenizer = FakeTokenizer(with_chat_template=True)
@@ -159,41 +192,18 @@ def test_rollout_generator_serializes_generation_and_policy_update(device):
generation_started = threading.Event() generation_started = threading.Event()
allow_generation_to_finish = threading.Event() allow_generation_to_finish = threading.Event()
update_finished = threading.Event() update_finished = threading.Event()
thread_errors = [] gen._generate_eval = _blocking_hook(
original = gen._generate_eval gen._generate_eval, generation_started, allow_generation_to_finish
)
def blocking_generate(batch, generation_version): _assert_interleaved(
generation_started.set() lambda: gen.generate(_make_instruction_batch(n=1)),
assert allow_generation_to_finish.wait(timeout=5) lambda: gen.apply_weight_update(1, update_finished.set),
return original(batch, generation_version) started=generation_started,
release=allow_generation_to_finish,
finished=update_finished,
)
gen._generate_eval = blocking_generate
def generate():
try:
gen.generate(_make_instruction_batch(n=1))
except BaseException as exc:
thread_errors.append(exc)
def apply_update():
try:
gen.apply_weight_update(1, update_finished.set)
except BaseException as exc:
thread_errors.append(exc)
generation_thread = threading.Thread(target=generate)
update_thread = threading.Thread(target=apply_update)
generation_thread.start()
assert generation_started.wait(timeout=5)
update_thread.start()
assert not update_finished.wait(timeout=0.1)
allow_generation_to_finish.set()
generation_thread.join(timeout=5)
update_thread.join(timeout=5)
assert not generation_thread.is_alive()
assert not update_thread.is_alive()
assert thread_errors == []
assert update_finished.is_set() assert update_finished.is_set()
assert gen.policy_version == 1 assert gen.policy_version == 1
@@ -203,43 +213,23 @@ def test_rollout_generator_serializes_direct_scheduler_update(device):
generation_started = threading.Event() generation_started = threading.Event()
allow_generation_to_finish = threading.Event() allow_generation_to_finish = threading.Event()
update_finished = threading.Event() update_finished = threading.Event()
thread_errors = [] gen._generate_eval = _blocking_hook(
original = gen._generate_eval gen._generate_eval, generation_started, allow_generation_to_finish
)
def blocking_generate(batch, generation_version):
generation_started.set()
assert allow_generation_to_finish.wait(timeout=5)
return original(batch, generation_version)
gen._generate_eval = blocking_generate
rollout = [] rollout = []
def generate():
try:
rollout.append(gen.generate(_make_instruction_batch(n=1)))
except BaseException as exc:
thread_errors.append(exc)
def update_scheduler_directly(): def update_scheduler_directly():
try: gen.scheduler.update_weights(1)
gen.scheduler.update_weights(1) update_finished.set()
update_finished.set()
except BaseException as exc:
thread_errors.append(exc)
generation_thread = threading.Thread(target=generate) _assert_interleaved(
update_thread = threading.Thread(target=update_scheduler_directly) lambda: rollout.append(gen.generate(_make_instruction_batch(n=1))),
generation_thread.start() update_scheduler_directly,
assert generation_started.wait(timeout=5) started=generation_started,
update_thread.start() release=allow_generation_to_finish,
assert not update_finished.wait(timeout=0.1) finished=update_finished,
)
allow_generation_to_finish.set()
generation_thread.join(timeout=5)
update_thread.join(timeout=5)
assert not generation_thread.is_alive()
assert not update_thread.is_alive()
assert thread_errors == []
assert rollout[0].policy_version == 0 assert rollout[0].policy_version == 0
assert gen.policy_version == 1 assert gen.policy_version == 1
@@ -466,7 +456,6 @@ def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device):
allow_final_validation_to_finish = threading.Event() allow_final_validation_to_finish = threading.Event()
update_finished = threading.Event() update_finished = threading.Event()
rollout_finished = threading.Event() rollout_finished = threading.Event()
thread_errors = []
validation_calls = 0 validation_calls = 0
original_validate = runner._validate_policy_version original_validate = runner._validate_policy_version
@@ -481,31 +470,17 @@ def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device):
runner._validate_policy_version = blocking_validate runner._validate_policy_version = blocking_validate
def produce_rollout(): def produce_rollout():
try: runner(_make_instruction_batch(n=1))
runner(_make_instruction_batch(n=1)) rollout_finished.set()
rollout_finished.set()
except BaseException as exc:
thread_errors.append(exc)
def apply_update(): _assert_interleaved(
try: produce_rollout,
runner.apply_weight_update(1, update_finished.set) lambda: runner.apply_weight_update(1, update_finished.set),
except BaseException as exc: started=final_validation_started,
thread_errors.append(exc) release=allow_final_validation_to_finish,
finished=update_finished,
)
rollout_thread = threading.Thread(target=produce_rollout)
update_thread = threading.Thread(target=apply_update)
rollout_thread.start()
assert final_validation_started.wait(timeout=5)
update_thread.start()
assert not update_finished.wait(timeout=0.1)
allow_final_validation_to_finish.set()
rollout_thread.join(timeout=5)
update_thread.join(timeout=5)
assert not rollout_thread.is_alive()
assert not update_thread.is_alive()
assert thread_errors == []
assert rollout_finished.is_set() assert rollout_finished.is_set()
assert update_finished.is_set() assert update_finished.is_set()
assert runner._cache is not None assert runner._cache is not None
tests/trainer/test_train_strategy.py
+60 -102
View File
@@ -1,120 +1,78 @@
import numpy as np import math
import pytest
import torch import torch
from astrai.trainer.schedule import CosineScheduler, SchedulerFactory, SGDRScheduler from astrai.trainer.schedule import CosineScheduler, SchedulerFactory, SGDRScheduler
def test_schedule_factory_random_configs(): def _stepped_lrs(scheduler, optimizer, n_steps):
"""Test scheduler factory with random configurations""" """Return the lr after construction plus each of *n_steps* steps."""
lrs = list(scheduler.get_last_lr())
for _ in range(n_steps):
optimizer.step()
scheduler.step()
lrs.append(scheduler.get_last_lr()[0])
return lrs
# Create a simple model and optimizer for testing
def test_cosine_scheduler_warms_up_then_decays_to_floor():
"""lr ramps linearly to base_lr during warmup, cosine-decays after it,
and never drops below min_rate * base_lr."""
base_lr = 0.001
model = torch.nn.Linear(10, 2) model = torch.nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr)
scheduler = SchedulerFactory.create(
"cosine", optimizer, warmup_steps=2, lr_decay_steps=4, min_rate=0.1
)
# Test multiple random configurations assert isinstance(scheduler, CosineScheduler)
for _ in range(5): # Test 5 random configurations lrs = _stepped_lrs(scheduler, optimizer, n_steps=7)
# Test multiple random configurations
cosine_params = {
"schedule_type": "cosine",
"warmup_steps": np.random.randint(50, 200),
"total_steps": np.random.randint(1000, 5000),
"min_rate": np.random.uniform(0.01, 0.1),
}
sgdr_params = {
"schedule_type": "sgdr",
"warmup_steps": np.random.randint(50, 200),
"cycle_length": np.random.randint(500, 2000),
"t_mult": np.random.randint(1, 3),
"min_rate": np.random.uniform(0.01, 0.1),
}
for params in [cosine_params, sgdr_params]:
schedule_type = params["schedule_type"]
# Convert parameters for scheduler constructor
if schedule_type == "cosine":
warmup_steps = params["warmup_steps"]
total_steps = params["total_steps"]
min_rate = params["min_rate"]
lr_decay_steps = total_steps - warmup_steps
scheduler = SchedulerFactory.create(
schedule_type,
optimizer,
warmup_steps=warmup_steps,
lr_decay_steps=lr_decay_steps,
min_rate=min_rate,
)
assert isinstance(scheduler, CosineScheduler)
assert scheduler.warmup_steps == warmup_steps
assert scheduler.lr_decay_steps == lr_decay_steps
assert scheduler.min_rate == min_rate
elif schedule_type == "sgdr":
warmup_steps = params["warmup_steps"]
cycle_length = params["cycle_length"]
t_mult = params["t_mult"]
min_rate = params["min_rate"]
scheduler = SchedulerFactory.create(
schedule_type,
optimizer,
warmup_steps=warmup_steps,
cycle_length=cycle_length,
t_mult=t_mult,
min_rate=min_rate,
)
assert isinstance(scheduler, SGDRScheduler)
assert scheduler.warmup_steps == warmup_steps
assert scheduler.cycle_length == cycle_length
assert scheduler.t_mult == t_mult
assert scheduler.min_rate == min_rate
# Test scheduler state dict functionality assert lrs[0] == pytest.approx(0.1 * base_lr) # warmup starts at the floor
state_dict = scheduler.state_dict() assert lrs[1] == pytest.approx(0.5 * base_lr) # halfway through warmup
assert "warmup_steps" in state_dict assert lrs[2] == pytest.approx(base_lr) # warmup complete
assert "min_rate" in state_dict expected_mid = base_lr * 0.5 * (1.0 + math.cos(math.pi * 0.25))
assert lrs[3] == pytest.approx(expected_mid) # quarter into decay
# Test scheduler step functionality assert lrs[5] > 0.1 * base_lr # 3/4 into decay: not clamped yet
initial_lr = scheduler.get_last_lr() assert lrs[6] == pytest.approx(0.1 * base_lr) # clamped at min_rate floor
optimizer.step() assert lrs[7] == pytest.approx(0.1 * base_lr) # stays at the floor
scheduler.step() assert all(lr >= 0.1 * base_lr - 1e-12 for lr in lrs)
new_lr = scheduler.get_last_lr()
# Learning rate should change after step, or if it's the first step,
# the epoch counter should increment
assert initial_lr != new_lr or scheduler.last_epoch > -1
def test_schedule_factory_edge_cases(): def test_cosine_scheduler_decays_to_zero_with_min_rate_zero():
"""Test scheduler factory with edge cases and boundary conditions""" """min_rate=0 must reach exactly 0.0 at the end of decay, not NaN."""
base_lr = 0.001
model = torch.nn.Linear(10, 2) model = torch.nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr)
scheduler = SchedulerFactory.create(
"cosine", optimizer, warmup_steps=1, lr_decay_steps=9, min_rate=0.0
)
# Test edge cases for CosineScheduleConfig lrs = _stepped_lrs(scheduler, optimizer, n_steps=11)
edge_cases = [
# Minimal warmup and steps
{"warmup_steps": 1, "total_steps": 10, "min_rate": 0.01},
# Large values
{"warmup_steps": 1000, "total_steps": 10000, "min_rate": 0.5},
# Zero min_rate (edge case)
{"warmup_steps": 100, "total_steps": 1000, "min_rate": 0.0},
]
for params in edge_cases: assert lrs[10] == 0.0
warmup_steps = params["warmup_steps"] assert lrs[11] == 0.0
total_steps = params["total_steps"] assert all(math.isfinite(lr) for lr in lrs)
min_rate = params["min_rate"]
lr_decay_steps = total_steps - warmup_steps
scheduler = SchedulerFactory.create(
"cosine",
optimizer,
warmup_steps=warmup_steps,
lr_decay_steps=lr_decay_steps,
min_rate=min_rate,
)
assert scheduler is not None
# Test multiple steps
for _ in range(10): def test_sgdr_scheduler_restarts_each_cycle():
optimizer.step() """lr anneals within a cycle, then jumps back to base_lr on restart."""
scheduler.step() base_lr = 0.001
model = torch.nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr)
scheduler = SchedulerFactory.create(
"sgdr", optimizer, warmup_steps=2, cycle_length=4, t_mult=1, min_rate=0.1
)
assert isinstance(scheduler, SGDRScheduler)
lrs = _stepped_lrs(scheduler, optimizer, n_steps=7)
assert lrs[2] == pytest.approx(base_lr) # cycle start
expected_mid = base_lr * (0.1 + 0.9 * 0.5) # halfway through the cycle
assert lrs[4] == pytest.approx(expected_mid)
assert lrs[5] < lrs[4] # still annealing at the cycle end
assert lrs[6] == pytest.approx(base_lr) # restart: back to full lr
def test_schedule_factory_state_persistence(): def test_schedule_factory_state_persistence():