test: prune low-value and duplicate tests
- Remove tautological test_trainer assertions that never trained - Drop grpo isfinite-only smokes and merge frozen-model checks via parametrize - Merge duplicate tool_parser cases (find/streaming/factory) with parametrize - Collapse duplicate dataset store/detect_format tests - Remove misleading scheduler/task tests that asserted the opposite of their names - Merge signal-handler SIGTERM/SIGINT into one parametrized case - Drop cross-file grpo strategy duplication kept in online_strategy
This commit is contained in:
@@ -217,14 +217,8 @@ def test_unloaded_sample_window_raises():
|
||||
store.sample_window(0)
|
||||
|
||||
|
||||
def test_unloaded_dataset_len():
|
||||
"""__len__ on a store with no data returns 0."""
|
||||
store = MmapStore(window_size=64, stride=64)
|
||||
assert len(store) == 0
|
||||
|
||||
|
||||
def test_store_unloaded_len():
|
||||
"""Unloaded Store has __len__ == 0"""
|
||||
"""Unloaded Store has __len__ == 0."""
|
||||
store = MmapStore()
|
||||
assert len(store) == 0
|
||||
assert store.keys == []
|
||||
@@ -498,21 +492,21 @@ def _write_json_dataset(test_dir, tokenizer_path, records, config_overrides=None
|
||||
return data_dir
|
||||
|
||||
|
||||
def test_detect_format_jsonl_dir(base_test_env):
|
||||
@pytest.mark.parametrize(
|
||||
"use_jsonl",
|
||||
[True, False],
|
||||
)
|
||||
def test_detect_format_data_dir(base_test_env, use_jsonl):
|
||||
"""detect_format returns 'jsonl' for dirs of .jsonl or .json files."""
|
||||
test_dir = base_test_env["test_dir"]
|
||||
tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
|
||||
if use_jsonl:
|
||||
data_dir = _write_jsonl_dataset(
|
||||
test_dir,
|
||||
tokenizer_path,
|
||||
[{"text": "hello world"}, {"text": "foo bar baz"}],
|
||||
)
|
||||
assert detect_format(data_dir) == "jsonl"
|
||||
|
||||
|
||||
def test_detect_format_json_dir(base_test_env):
|
||||
"""detect_format returns 'jsonl' for directory with .json files."""
|
||||
test_dir = base_test_env["test_dir"]
|
||||
tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
|
||||
else:
|
||||
data_dir = _write_json_dataset(
|
||||
test_dir,
|
||||
tokenizer_path,
|
||||
@@ -745,32 +739,6 @@ def test_sft_jsonl_explicit_config_takes_priority(base_test_env):
|
||||
assert "loss_mask" in dataset.keys
|
||||
|
||||
|
||||
def test_jsonl_store_pipeline_config_roundtrip(base_test_env):
|
||||
test_dir = base_test_env["test_dir"]
|
||||
config_path = os.path.join(test_dir, "dataset_config.json")
|
||||
with open(config_path, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"tokenizer_path": os.path.join(test_dir, "tokenizer"),
|
||||
"version": 1,
|
||||
"input": {"sections": [{"field": "text", "action": "train"}]},
|
||||
"mask": {"assistant": "train"},
|
||||
"preprocessing": {"max_seq_len": 64},
|
||||
"output": {"position_ids_mode": "doc_reset"},
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
raw.pop("tokenizer_path")
|
||||
config = PipelineConfig.from_dict(raw)
|
||||
assert config.output.position_ids_mode == "doc_reset"
|
||||
assert config.preprocessing.max_seq_len == 64
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GRPO end-to-end: builder → JsonlStore → GRPODataset → collate_fn
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -177,24 +177,6 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
|
||||
assert stats["total_tasks"] >= 0
|
||||
|
||||
|
||||
def test_prefill_skips_fully_cached_tasks(mock_model_and_tokenizer):
|
||||
"""Tasks whose entire prompt is cached skip the prefill phase."""
|
||||
mock_model, mock_tokenizer = mock_model_and_tokenizer
|
||||
|
||||
with patch("astrai.inference.core.scheduler.AutoModel"):
|
||||
with patch("astrai.inference.core.scheduler.AutoTokenizer"):
|
||||
scheduler = InferenceScheduler(
|
||||
model=mock_model,
|
||||
tokenizer=mock_tokenizer,
|
||||
max_batch_size=4,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
task_id = scheduler.add_task("short prompt", stream_callback=lambda t: None)
|
||||
scheduler.stop()
|
||||
assert task_id.startswith("task_")
|
||||
|
||||
|
||||
def _make_real_scheduler(device):
|
||||
"""Build a scheduler backed by a tiny real model for run_batch tests."""
|
||||
cfg = make_rollout_config(max_position_embeddings=64)
|
||||
|
||||
@@ -51,7 +51,7 @@ def test_task_manager_add_task():
|
||||
assert len(tm.waiting_queue) == 1
|
||||
|
||||
|
||||
def test_task_manager_add_task_too_long_immediate_stop():
|
||||
def test_task_manager_long_prompt_truncated_not_stopped():
|
||||
t = _make_mock_tokenizer()
|
||||
t.encode.return_value = list(range(9000))
|
||||
cb_calls = []
|
||||
@@ -60,6 +60,7 @@ def test_task_manager_add_task_too_long_immediate_stop():
|
||||
tm.add_task("long", stream_callback=lambda tok: cb_calls.append(tok))
|
||||
assert len(cb_calls) == 0
|
||||
assert len(tm.waiting_queue) == 1
|
||||
assert len(tm.waiting_queue[0].prompt_ids) == 16
|
||||
|
||||
|
||||
def test_task_manager_remove_task():
|
||||
|
||||
@@ -59,14 +59,15 @@ def test_find_multiple_tool_calls():
|
||||
assert results[1]["name"] == "f2"
|
||||
|
||||
|
||||
def test_find_no_tool_call():
|
||||
results = _find_tool_calls("Hello, how are you?")
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_find_non_tool_json_skipped():
|
||||
results = _find_tool_calls('{"not_a_tool": true}')
|
||||
assert len(results) == 0
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected_count",
|
||||
[
|
||||
("Hello, how are you?", 0),
|
||||
('{"not_a_tool": true}', 0),
|
||||
],
|
||||
)
|
||||
def test_find_no_tool_call(text, expected_count):
|
||||
assert len(_find_tool_calls(text)) == expected_count
|
||||
|
||||
|
||||
def test_find_no_arguments_field():
|
||||
@@ -76,79 +77,6 @@ def test_find_no_arguments_field():
|
||||
assert results[0]["args"] == ""
|
||||
|
||||
|
||||
def test_find_deeply_nested_arguments():
|
||||
text = '{"name": "deep", "arguments": {"a": {"b": {"c": {"d": 4}}}}}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "deep"
|
||||
assert '"d": 4' in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_arguments_with_boolean_and_null():
|
||||
text = '{"name": "flags", "arguments": {"active": true, "count": 0, "nick": null}}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "flags"
|
||||
assert "true" in results[0]["args"]
|
||||
assert "null" in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_arguments_with_array():
|
||||
text = '{"name": "add_items", "arguments": {"items": [1, 2, 3], "name": "list"}}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "add_items"
|
||||
assert "[1, 2, 3]" in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_arguments_with_nested_array_of_objects():
|
||||
text = '{"name": "batch", "arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert '"rows"' in results[0]["args"]
|
||||
assert '"id": 1' in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_arguments_as_string_not_object():
|
||||
text = '{"name": "echo", "arguments": "just a string"}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "echo"
|
||||
assert "just a string" in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_arguments_with_unicode():
|
||||
text = (
|
||||
'{"name": "translate", "arguments": {"text": "\u4f60\u597d\uff0c\u4e16\u754c"}}'
|
||||
)
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "translate"
|
||||
|
||||
|
||||
def test_find_arguments_with_escaped_quotes():
|
||||
text = '{"name": "format", "arguments": {"template": "he said \\"hello\\""}}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert 'he said \\"hello\\"' in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_arguments_with_braces_in_string():
|
||||
text = '{"name": "eval", "arguments": {"code": "function(x) { return x + 1; }"}}'
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "eval"
|
||||
assert "function(x) { return x + 1; }" in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_many_properties():
|
||||
args = ",".join(f'"{chr(97 + i % 26)}" : {i}' for i in range(20))
|
||||
text = '{"name": "many", "arguments": {' + args + "}}"
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "many"
|
||||
|
||||
|
||||
def test_find_empty_arguments():
|
||||
results = _find_tool_calls('{"name": "ping", "arguments": {}}')
|
||||
assert len(results) == 1
|
||||
@@ -164,6 +92,62 @@ def test_find_extracts_correct_arg_start_position():
|
||||
assert json_str == text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected_name,arg_substr",
|
||||
[
|
||||
(
|
||||
'{"name": "deep", "arguments": {"a": {"b": {"c": {"d": 4}}}}}',
|
||||
"deep",
|
||||
'"d": 4',
|
||||
),
|
||||
(
|
||||
'{"name": "flags", "arguments": {"active": true, "count": 0, "nick": null}}',
|
||||
"flags",
|
||||
"null",
|
||||
),
|
||||
(
|
||||
'{"name": "add_items", "arguments": {"items": [1, 2, 3], "name": "list"}}',
|
||||
"add_items",
|
||||
"[1, 2, 3]",
|
||||
),
|
||||
(
|
||||
'{"name": "batch", "arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}',
|
||||
"batch",
|
||||
'"id": 1',
|
||||
),
|
||||
('{"name": "echo", "arguments": "just a string"}', "echo", "just a string"),
|
||||
(
|
||||
'{"name": "translate", "arguments": {"text": "\u4f60\u597d\uff0c\u4e16\u754c"}}',
|
||||
"translate",
|
||||
"\u4f60\u597d",
|
||||
),
|
||||
(
|
||||
'{"name": "format", "arguments": {"template": "he said \\"hello\\""}}',
|
||||
"format",
|
||||
'he said \\"hello\\"',
|
||||
),
|
||||
(
|
||||
'{"name": "eval", "arguments": {"code": "function(x) { return x + 1; }"}}',
|
||||
"eval",
|
||||
"function(x) { return x + 1; }",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_find_arguments_variants(text, expected_name, arg_substr):
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == expected_name
|
||||
assert arg_substr in results[0]["args"]
|
||||
|
||||
|
||||
def test_find_many_properties():
|
||||
args = ",".join(f'"{chr(97 + i % 26)}" : {i}' for i in range(20))
|
||||
text = '{"name": "many", "arguments": {' + args + "}}"
|
||||
results = _find_tool_calls(text)
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "many"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected_name,expected_complete",
|
||||
[
|
||||
@@ -340,30 +324,21 @@ def test_streaming_multiple_tool_calls_incremental():
|
||||
assert "f2" in names
|
||||
|
||||
|
||||
def test_streaming_deeply_nested_args():
|
||||
parser = SimpleJsonToolParser()
|
||||
text = '{"name": "deep", "arguments": {"a": {"b": {"c": 42}}}}'
|
||||
_, args_chunks = _simulate_streaming(parser, text)
|
||||
joined = "".join(args_chunks)
|
||||
assert '"c": 42' in joined
|
||||
|
||||
|
||||
def test_streaming_args_with_unicode():
|
||||
parser = SimpleJsonToolParser()
|
||||
text = (
|
||||
'{"name": "translate", "arguments": {"text": "\u4f60\u597d\uff0c\u4e16\u754c"}}'
|
||||
@pytest.mark.parametrize(
|
||||
"text,arg_substr",
|
||||
[
|
||||
('{"name": "deep", "arguments": {"a": {"b": {"c": 42}}}}', '"c": 42'),
|
||||
(
|
||||
'{"name": "translate", "arguments": {"text": "\u4f60\u597d\uff0c\u4e16\u754c"}}',
|
||||
"\u4f60\u597d",
|
||||
),
|
||||
('{"name": "add", "arguments": {"items": [1, 2, 3]}}', "[1, 2, 3]"),
|
||||
],
|
||||
)
|
||||
_, args_chunks = _simulate_streaming(parser, text)
|
||||
joined = "".join(args_chunks)
|
||||
assert "\u4f60\u597d" in joined
|
||||
|
||||
|
||||
def test_streaming_args_with_array():
|
||||
def test_streaming_args_variants(text, arg_substr):
|
||||
parser = SimpleJsonToolParser()
|
||||
text = '{"name": "add", "arguments": {"items": [1, 2, 3]}}'
|
||||
_, args_chunks = _simulate_streaming(parser, text)
|
||||
joined = "".join(args_chunks)
|
||||
assert "[1, 2, 3]" in joined
|
||||
assert arg_substr in "".join(args_chunks)
|
||||
|
||||
|
||||
def test_streaming_empty_arguments():
|
||||
@@ -514,7 +489,6 @@ def test_feed_then_parse_complete_same_instance():
|
||||
('{ "name" : "f"}', True),
|
||||
('{"other": 1}', False),
|
||||
('prefix {"name": "f", "args": {}}', True),
|
||||
('{"name": "f"}', True), # match at start
|
||||
(' {"name": "f"}', True),
|
||||
],
|
||||
)
|
||||
@@ -526,10 +500,6 @@ def test_pattern_regex(text, matches):
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_pattern_name_at_start():
|
||||
assert _TOOL_CALL_HEAD_RE.match('{"name": "f"}')
|
||||
|
||||
|
||||
def test_factory_register_and_create():
|
||||
parser = ToolParserFactory.create("simple_json")
|
||||
assert isinstance(parser, BaseToolParser)
|
||||
@@ -547,10 +517,6 @@ def test_factory_list_registered():
|
||||
assert "simple_json" in ToolParserFactory.list_registered()
|
||||
|
||||
|
||||
def test_factory_create_with_no_extra_kwargs():
|
||||
assert isinstance(ToolParserFactory.create("simple_json"), BaseToolParser)
|
||||
|
||||
|
||||
def test_factory_create_with_tools_only():
|
||||
tools = [
|
||||
{
|
||||
@@ -563,13 +529,6 @@ def test_factory_create_with_tools_only():
|
||||
assert parser.tool_choice == "auto"
|
||||
|
||||
|
||||
def test_feed_accepts_token_ids_and_ignores_them():
|
||||
parser = SimpleJsonToolParser()
|
||||
text = '{"name": "get_weather", "arguments": {"city": "Beijing"}}'
|
||||
deltas_with = parser.feed(text, current_token_ids=[123, 456], delta_token_ids=[456])
|
||||
assert len(deltas_with) > 0
|
||||
|
||||
|
||||
def test_feed_token_ids_do_not_affect_parsing():
|
||||
parser_no_ids = SimpleJsonToolParser()
|
||||
parser_with_ids = SimpleJsonToolParser()
|
||||
|
||||
@@ -98,18 +98,9 @@ def test_loralinear_merge():
|
||||
assert lora._merged
|
||||
assert not hasattr(lora, "lora_A")
|
||||
|
||||
|
||||
def test_loralinear_merge_is_idempotent():
|
||||
base = Linear(4, 4)
|
||||
with torch.no_grad():
|
||||
base.weight.zero_()
|
||||
|
||||
lora = LoRALinear(base, r=2, alpha=2)
|
||||
with torch.no_grad():
|
||||
lora.lora_B.fill_(1.0)
|
||||
|
||||
lora.merge()
|
||||
# merge is guarded by _merged — a second call is a no-op.
|
||||
lora.merge()
|
||||
assert lora._merged
|
||||
|
||||
|
||||
def test_inject_lora_default_target():
|
||||
|
||||
@@ -71,23 +71,14 @@ def test_grpo_loss_backward(grpo_strategy):
|
||||
assert has_grad
|
||||
|
||||
|
||||
def test_grpo_ref_model_not_updated(grpo_strategy):
|
||||
"""Backward should not populate gradients on ref_model."""
|
||||
@pytest.mark.parametrize("model_name", ["ref_model", "old_model"])
|
||||
def test_grpo_frozen_models_not_updated(grpo_strategy, model_name):
|
||||
"""Backward should not populate gradients on ref_model or old_model."""
|
||||
strategy, device = grpo_strategy
|
||||
batch = _make_batch(device=device)
|
||||
loss = strategy.compute_loss(batch)
|
||||
loss.backward()
|
||||
for p in strategy.ref_model.parameters():
|
||||
assert p.grad is None
|
||||
|
||||
|
||||
def test_grpo_old_model_not_updated(grpo_strategy):
|
||||
"""Backward should not populate gradients on old_model."""
|
||||
strategy, device = grpo_strategy
|
||||
batch = _make_batch(device=device)
|
||||
loss = strategy.compute_loss(batch)
|
||||
loss.backward()
|
||||
for p in strategy.old_model.parameters():
|
||||
for p in getattr(strategy, model_name).parameters():
|
||||
assert p.grad is None
|
||||
|
||||
|
||||
@@ -133,45 +124,3 @@ def test_grpo_sync_old_model(grpo_strategy):
|
||||
if k in old_sd_after
|
||||
)
|
||||
assert matches
|
||||
|
||||
|
||||
def test_grpo_partial_mask(grpo_strategy):
|
||||
"""Only the first half of response tokens are valid."""
|
||||
strategy, device = grpo_strategy
|
||||
batch = _make_batch(device=device)
|
||||
B, G, R = batch["masks"].shape
|
||||
half = R // 2
|
||||
batch["masks"][:, :, half:] = 0.0
|
||||
loss = strategy.compute_loss(batch)
|
||||
assert torch.isfinite(loss).item()
|
||||
|
||||
|
||||
def test_grpo_clipping_effect(grpo_strategy):
|
||||
"""After diverging policy from ref, ratio should be clipped to [1-eps, 1+eps]
|
||||
on the surrogate. Verify loss is finite and non-zero for distinct rewards."""
|
||||
strategy, device = grpo_strategy
|
||||
with torch.no_grad():
|
||||
for p in strategy.model.parameters():
|
||||
p.add_(0.3)
|
||||
batch = _make_batch(device=device)
|
||||
loss = strategy.compute_loss(batch)
|
||||
assert torch.isfinite(loss).item()
|
||||
assert loss.abs().item() > 1e-4
|
||||
|
||||
|
||||
def test_grpo_no_reduction_param():
|
||||
"""GRPOStrategy.__init__ must not accept ``reduction`` (removed)."""
|
||||
import inspect
|
||||
|
||||
sig = inspect.signature(GRPOStrategy.__init__)
|
||||
assert "reduction" not in sig.parameters
|
||||
|
||||
|
||||
def test_grpo_shapes_3d_batch(grpo_strategy):
|
||||
"""Verify compute_loss handles non-square prompt/response lengths."""
|
||||
strategy, device = grpo_strategy
|
||||
batch = _make_batch(
|
||||
batch_size=3, group_size=4, prompt_len=10, response_len=8, device=device
|
||||
)
|
||||
loss = strategy.compute_loss(batch)
|
||||
assert torch.isfinite(loss).item()
|
||||
|
||||
@@ -96,12 +96,10 @@ def test_factory_registers_online_aliases():
|
||||
assert StrategyFactory.get_component_class("online_dpo") is DPOStrategy
|
||||
|
||||
|
||||
def test_grpo_supports_online(device):
|
||||
assert _make_grpo(device).supports_online() is True
|
||||
|
||||
|
||||
def test_dpo_supports_online(device):
|
||||
assert _make_dpo(device).supports_online() is True
|
||||
@pytest.mark.parametrize("make_fn", ["_make_grpo", "_make_dpo"])
|
||||
def test_online_strategies_support_online(device, make_fn):
|
||||
maker = {"_make_grpo": _make_grpo, "_make_dpo": _make_dpo}[make_fn]
|
||||
assert maker(device).supports_online() is True
|
||||
|
||||
|
||||
def test_base_strategy_prepare_from_rollout_raises_by_default(device):
|
||||
@@ -267,17 +265,6 @@ def test_step_called_when_sync_gradients_true(device):
|
||||
assert runner.step_calls == 1
|
||||
|
||||
|
||||
def test_loss_is_differentiable_grpo(device):
|
||||
strat = _make_grpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
loss.backward()
|
||||
has_grad = any(
|
||||
p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
|
||||
)
|
||||
assert has_grad
|
||||
|
||||
|
||||
def test_loss_is_differentiable_dpo(device):
|
||||
strat = _make_dpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
@@ -289,17 +276,6 @@ def test_loss_is_differentiable_dpo(device):
|
||||
assert has_grad
|
||||
|
||||
|
||||
def test_ref_and_old_model_not_updated_by_backward_grpo(device):
|
||||
strat = _make_grpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
loss.backward()
|
||||
for p in strat.ref_model.parameters():
|
||||
assert p.grad is None
|
||||
for p in strat.old_model.parameters():
|
||||
assert p.grad is None
|
||||
|
||||
|
||||
def test_ref_model_not_updated_by_backward_dpo(device):
|
||||
strat = _make_dpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
|
||||
@@ -81,6 +81,9 @@ def test_rollout_result_inherits_raw_rollout_fields():
|
||||
assert r.prompts.shape == (2, 4)
|
||||
assert r.responses.shape == (2, 3, 5)
|
||||
assert r.prompt_mask.shape == (2, 4)
|
||||
# RolloutResult must carry every RawRollout field.
|
||||
raw_fields = {f for f in RawRollout.__dataclass_fields__}
|
||||
assert raw_fields.issubset(set(RolloutResult.__dataclass_fields__))
|
||||
|
||||
|
||||
def test_base_reward_model_is_abstract():
|
||||
|
||||
@@ -131,18 +131,10 @@ def test_register_signal_handlers():
|
||||
assert ctx.stop_requested
|
||||
|
||||
|
||||
def test_sigterm_triggers_checkpoint_save(base_test_env):
|
||||
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGTERM)
|
||||
assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
|
||||
|
||||
meta = load_checkpoint_meta(base_test_env["test_dir"])
|
||||
assert "consumed_samples" in meta
|
||||
assert meta["consumed_samples"] >= 0
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_sigint_triggers_checkpoint_save(base_test_env):
|
||||
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGINT)
|
||||
@pytest.mark.parametrize("sig", [signal.SIGTERM, signal.SIGINT])
|
||||
def test_signal_triggers_checkpoint_save(base_test_env, sig):
|
||||
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], sig)
|
||||
assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
|
||||
|
||||
meta = load_checkpoint_meta(base_test_env["test_dir"])
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
import pytest
|
||||
|
||||
from astrai.trainer import Trainer
|
||||
|
||||
# train_config_factory is injected via fixture
|
||||
|
||||
|
||||
def test_different_batch_sizes(base_test_env, random_dataset, train_config_factory):
|
||||
"""Test training with different batch sizes"""
|
||||
batch_sizes = [1, 2, 4, 8]
|
||||
|
||||
for batch_per_device in batch_sizes:
|
||||
def test_training_runs_with_various_batch_sizes(
|
||||
base_test_env, random_dataset, train_config_factory
|
||||
):
|
||||
"""Training should complete for a range of batch sizes without error."""
|
||||
for batch_per_device in [1, 2, 4]:
|
||||
train_config = train_config_factory(
|
||||
model_fn=lambda: base_test_env["model"],
|
||||
dataset=random_dataset,
|
||||
@@ -15,48 +15,22 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto
|
||||
device=base_test_env["device"],
|
||||
batch_per_device=batch_per_device,
|
||||
)
|
||||
|
||||
assert train_config.batch_per_device == batch_per_device
|
||||
trainer = Trainer(train_config)
|
||||
trainer.train()
|
||||
|
||||
|
||||
def test_gradient_accumulation(base_test_env, random_dataset, train_config_factory):
|
||||
"""Test training with different gradient accumulation steps"""
|
||||
grad_accum_steps_list = [1, 2, 4]
|
||||
|
||||
for grad_accum_steps in grad_accum_steps_list:
|
||||
@pytest.mark.slow
|
||||
def test_gradient_accumulation_runs(
|
||||
base_test_env, random_dataset, train_config_factory
|
||||
):
|
||||
"""Training with gradient accumulation should complete."""
|
||||
train_config = train_config_factory(
|
||||
model_fn=lambda: base_test_env["model"],
|
||||
dataset=random_dataset,
|
||||
test_dir=base_test_env["test_dir"],
|
||||
device=base_test_env["device"],
|
||||
batch_per_device=2,
|
||||
grad_accum_steps=grad_accum_steps,
|
||||
grad_accum_steps=4,
|
||||
)
|
||||
|
||||
trainer = Trainer(train_config)
|
||||
trainer.train()
|
||||
|
||||
assert train_config.grad_accum_steps == grad_accum_steps
|
||||
|
||||
|
||||
def test_memory_efficient_training(base_test_env, random_dataset, train_config_factory):
|
||||
"""Test training with memory-efficient configurations"""
|
||||
# Test with smaller batch sizes and gradient checkpointing
|
||||
small_batch_configs = [
|
||||
{"batch_per_device": 1, "grad_accum_steps": 8},
|
||||
{"batch_per_device": 2, "grad_accum_steps": 4},
|
||||
{"batch_per_device": 4, "grad_accum_steps": 2},
|
||||
]
|
||||
|
||||
for config in small_batch_configs:
|
||||
train_config = train_config_factory(
|
||||
model_fn=lambda: base_test_env["model"],
|
||||
dataset=random_dataset,
|
||||
test_dir=base_test_env["test_dir"],
|
||||
device=base_test_env["device"],
|
||||
batch_per_device=config["batch_per_device"],
|
||||
grad_accum_steps=config["grad_accum_steps"],
|
||||
)
|
||||
|
||||
assert train_config.grad_accum_steps == config["grad_accum_steps"]
|
||||
assert train_config.batch_per_device == config["batch_per_device"]
|
||||
|
||||
Reference in New Issue
Block a user