test: deduplicate suites and prune low-value cases

- 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

tests/trainer/test_train_strategy.py

@@ -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():