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/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:
callback._enable(layer)
for layer in model.layers:
for i, layer in enumerate(model.layers):
assert hasattr(layer, "_original_forward")
assert layer.forward is not originals[0]
assert layer.forward is not originals[i]
for layer in model.layers:
callback._disable(layer)
@@ -110,6 +110,10 @@ def test_gradient_checkpointing_trainer_integration(
)
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()
tests/trainer/test_rollout.py
+71 -96
View File
@@ -58,6 +58,46 @@ def _make_instruction_batch(n=2):
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():
r = RawRollout(
prompts=torch.zeros(2, 4, dtype=torch.long),
@@ -96,13 +136,6 @@ def test_base_reward_model_is_abstract():
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):
model, _ = make_model(device, max_position_embeddings=128)
tokenizer = FakeTokenizer(with_chat_template=True)
@@ -159,41 +192,18 @@ def test_rollout_generator_serializes_generation_and_policy_update(device):
generation_started = threading.Event()
allow_generation_to_finish = threading.Event()
update_finished = threading.Event()
thread_errors = []
original = gen._generate_eval
gen._generate_eval = _blocking_hook(
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)
_assert_interleaved(
lambda: gen.generate(_make_instruction_batch(n=1)),
lambda: gen.apply_weight_update(1, update_finished.set),
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 gen.policy_version == 1
@@ -203,43 +213,23 @@ def test_rollout_generator_serializes_direct_scheduler_update(device):
generation_started = threading.Event()
allow_generation_to_finish = threading.Event()
update_finished = threading.Event()
thread_errors = []
original = gen._generate_eval
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
gen._generate_eval = _blocking_hook(
gen._generate_eval, generation_started, allow_generation_to_finish
)
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():
try:
gen.scheduler.update_weights(1)
update_finished.set()
except BaseException as exc:
thread_errors.append(exc)
gen.scheduler.update_weights(1)
update_finished.set()
generation_thread = threading.Thread(target=generate)
update_thread = threading.Thread(target=update_scheduler_directly)
generation_thread.start()
assert generation_started.wait(timeout=5)
update_thread.start()
assert not update_finished.wait(timeout=0.1)
_assert_interleaved(
lambda: rollout.append(gen.generate(_make_instruction_batch(n=1))),
update_scheduler_directly,
started=generation_started,
release=allow_generation_to_finish,
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 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()
update_finished = threading.Event()
rollout_finished = threading.Event()
thread_errors = []
validation_calls = 0
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
def produce_rollout():
try:
runner(_make_instruction_batch(n=1))
rollout_finished.set()
except BaseException as exc:
thread_errors.append(exc)
runner(_make_instruction_batch(n=1))
rollout_finished.set()
def apply_update():
try:
runner.apply_weight_update(1, update_finished.set)
except BaseException as exc:
thread_errors.append(exc)
_assert_interleaved(
produce_rollout,
lambda: runner.apply_weight_update(1, update_finished.set),
started=final_validation_started,
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 update_finished.is_set()
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
from astrai.trainer.schedule import CosineScheduler, SchedulerFactory, SGDRScheduler
def test_schedule_factory_random_configs():
"""Test scheduler factory with random configurations"""
def _stepped_lrs(scheduler, optimizer, n_steps):
"""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)
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
for _ in range(5): # Test 5 random configurations
# 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
assert isinstance(scheduler, CosineScheduler)
lrs = _stepped_lrs(scheduler, optimizer, n_steps=7)
# Test scheduler state dict functionality
state_dict = scheduler.state_dict()
assert "warmup_steps" in state_dict
assert "min_rate" in state_dict
# Test scheduler step functionality
initial_lr = scheduler.get_last_lr()
optimizer.step()
scheduler.step()
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
assert lrs[0] == pytest.approx(0.1 * base_lr) # warmup starts at the floor
assert lrs[1] == pytest.approx(0.5 * base_lr) # halfway through warmup
assert lrs[2] == pytest.approx(base_lr) # warmup complete
expected_mid = base_lr * 0.5 * (1.0 + math.cos(math.pi * 0.25))
assert lrs[3] == pytest.approx(expected_mid) # quarter into decay
assert lrs[5] > 0.1 * base_lr # 3/4 into decay: not clamped yet
assert lrs[6] == pytest.approx(0.1 * base_lr) # clamped at min_rate floor
assert lrs[7] == pytest.approx(0.1 * base_lr) # stays at the floor
assert all(lr >= 0.1 * base_lr - 1e-12 for lr in lrs)
def test_schedule_factory_edge_cases():
"""Test scheduler factory with edge cases and boundary conditions"""
def test_cosine_scheduler_decays_to_zero_with_min_rate_zero():
"""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)
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
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},
]
lrs = _stepped_lrs(scheduler, optimizer, n_steps=11)
for params in edge_cases:
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(
"cosine",
optimizer,
warmup_steps=warmup_steps,
lr_decay_steps=lr_decay_steps,
min_rate=min_rate,
)
assert scheduler is not None
assert lrs[10] == 0.0
assert lrs[11] == 0.0
assert all(math.isfinite(lr) for lr in lrs)
# Test multiple steps
for _ in range(10):
optimizer.step()
scheduler.step()
def test_sgdr_scheduler_restarts_each_cycle():
"""lr anneals within a cycle, then jumps back to base_lr on restart."""
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():