diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index 90f5e07..a6f1c2e 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -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,26 +492,26 @@ 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"]) - 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"]) - data_dir = _write_json_dataset( - test_dir, - tokenizer_path, - [{"text": "hello world"}, {"text": "foo bar baz qux"}], - ) + if use_jsonl: + data_dir = _write_jsonl_dataset( + test_dir, + tokenizer_path, + [{"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" @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 2475176..8edb6cb 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -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) diff --git a/tests/inference/test_task.py b/tests/inference/test_task.py index b969c74..b74886f 100644 --- a/tests/inference/test_task.py +++ b/tests/inference/test_task.py @@ -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(): diff --git a/tests/inference/test_tool_parser.py b/tests/inference/test_tool_parser.py index 2a6b945..5d27f52 100644 --- a/tests/inference/test_tool_parser.py +++ b/tests/inference/test_tool_parser.py @@ -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(): +@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]"), + ], +) +def test_streaming_args_variants(text, arg_substr): 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"}}' - ) - _, args_chunks = _simulate_streaming(parser, text) - joined = "".join(args_chunks) - assert "\u4f60\u597d" in joined - - -def test_streaming_args_with_array(): - 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() diff --git a/tests/module/test_lora.py b/tests/module/test_lora.py index d95a397..39d7b27 100644 --- a/tests/module/test_lora.py +++ b/tests/module/test_lora.py @@ -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(): diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index 9b62434..d7c531b 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/tests/trainer/test_grpo_strategy.py @@ -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() diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index 27e9b94..5d38080 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -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))) diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index 94d5630..a71e365 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -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(): diff --git a/tests/trainer/test_signal_handler.py b/tests/trainer/test_signal_handler.py index ece946b..e95cd28 100644 --- a/tests/trainer/test_signal_handler.py +++ b/tests/trainer/test_signal_handler.py @@ -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"]) diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py index d941dcd..3a59baa 100644 --- a/tests/trainer/test_trainer.py +++ b/tests/trainer/test_trainer.py @@ -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 - - -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: - 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, - ) - 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"] +@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=4, + ) + trainer = Trainer(train_config) + trainer.train()