feat: 增加server, 并且修改测试单元

tests/trainer/conftest.py

@@ -0,0 +1,97 @@
+import torch
+from torch.utils.data import Dataset
+import pytest
+
+
+class TrainerDataset(Dataset):
+    """Base dataset for trainer tests with consistent interface."""
+
+    def __init__(self, length=100, max_length=64, vocab_size=1000):
+        self.length = length
+        self.max_length = max_length
+        self.vocab_size = vocab_size
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx):
+        return {
+            "input_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
+            "target_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
+        }
+
+
+def create_train_config(
+    model: torch.nn.Module,
+    dataset: Dataset,
+    test_dir: str,
+    device: str,
+    strategy: str = "seq",
+    n_epoch: int = 1,
+    batch_size: int = 2,
+    accumulation_steps: int = 1,
+    max_grad_norm: float = 1.0,
+    ckpt_interval: int = 5,
+    random_seed: int = 42,
+    **kwargs,
+):
+    """Factory function to create common TrainConfig for tests.
+
+    Args:
+        model: The model to train
+        dataset: Training dataset
+        test_dir: Checkpoint directory
+        device: Device type ("cuda" or "cpu")
+        strategy: Training strategy type (default: "seq")
+        n_epoch: Number of epochs (default: 1)
+        batch_size: Batch size (default: 2)
+        accumulation_steps: Gradient accumulation steps (default: 1)
+        max_grad_norm: Maximum gradient norm for clipping (default: 1.0)
+        ckpt_interval: Checkpoint save interval in iterations (default: 5)
+        random_seed: Random seed for reproducibility (default: 42)
+        **kwargs: Additional arguments passed to TrainConfig
+
+    Returns:
+        TrainConfig instance configured for testing
+    """
+    from astrai.config import TrainConfig
+    from astrai.config.schedule_config import CosineScheduleConfig
+    from astrai.trainer.schedule import SchedulerFactory
+
+    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
+    optimizer_fn = lambda m: torch.optim.AdamW(m.parameters(), lr=0.001)
+    scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
+
+    return TrainConfig(
+        strategy=strategy,
+        model=model,
+        dataset=dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=test_dir,
+        n_epoch=n_epoch,
+        batch_size=batch_size,
+        ckpt_interval=ckpt_interval,
+        accumulation_steps=accumulation_steps,
+        max_grad_norm=max_grad_norm,
+        random_seed=random_seed,
+        device_type=device,
+        **kwargs,
+    )
+
+
+@pytest.fixture
+def train_config_factory():
+    """Fixture that provides the create_train_config factory function.
+
+    This fixture can be used by tests to create consistent TrainConfig
+    instances with sensible defaults for testing.
+    """
+    return create_train_config
+
+
+@pytest.fixture
+def trainer_dataset():
+    """Fixture providing a dataset for trainer tests."""
+    dataset = TrainerDataset()
+    yield dataset

tests/trainer/test_trainer.py

@@ -1,63 +1,39 @@
 import torch
-import numpy as np
 
-
-from astrai.config import *
-from astrai.trainer import *
 from astrai.data.dataset import *
+from astrai.trainer import Trainer
+
+# train_config_factory is injected via fixture
 
 
-def test_different_batch_sizes(base_test_env, random_dataset):
+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_size in batch_sizes:
-        schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-        optimizer_fn = lambda model: torch.optim.AdamW(model.parameters())
-        scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
-
-        train_config = TrainConfig(
-            strategy="seq",
+        train_config = train_config_factory(
             model=base_test_env["model"],
             dataset=random_dataset,
-            optimizer_fn=optimizer_fn,
-            scheduler_fn=scheduler_fn,
-            ckpt_dir=base_test_env["test_dir"],
-            n_epoch=1,
+            test_dir=base_test_env["test_dir"],
+            device=base_test_env["device"],
             batch_size=batch_size,
-            ckpt_interval=5,
-            accumulation_steps=1,
-            max_grad_norm=1.0,
-            random_seed=np.random.randint(1000),
-            device_type=base_test_env["device"],
         )
 
         assert train_config.batch_size == batch_size
 
 
-def test_gradient_accumulation(base_test_env, random_dataset):
+def test_gradient_accumulation(base_test_env, random_dataset, train_config_factory):
     """Test training with different gradient accumulation steps"""
     accumulation_steps_list = [1, 2, 4]
 
     for accumulation_steps in accumulation_steps_list:
-        schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-        optimizer_fn = lambda model: torch.optim.AdamW(model.parameters())
-        scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
-
-        train_config = TrainConfig(
-            strategy="seq",
+        train_config = train_config_factory(
             model=base_test_env["model"],
-            optimizer_fn=optimizer_fn,
-            scheduler_fn=scheduler_fn,
             dataset=random_dataset,
-            ckpt_dir=base_test_env["test_dir"],
-            n_epoch=1,
+            test_dir=base_test_env["test_dir"],
+            device=base_test_env["device"],
             batch_size=2,
-            ckpt_interval=10,
             accumulation_steps=accumulation_steps,
-            max_grad_norm=1.0,
-            random_seed=42,
-            device_type=base_test_env["device"],
         )
 
         trainer = Trainer(train_config)
@@ -66,7 +42,7 @@ def test_gradient_accumulation(base_test_env, random_dataset):
         assert train_config.accumulation_steps == accumulation_steps
 
 
-def test_memory_efficient_training(base_test_env, random_dataset):
+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 = [
@@ -76,24 +52,13 @@ def test_memory_efficient_training(base_test_env, random_dataset):
     ]
 
     for config in small_batch_configs:
-        schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-        optimizer_fn = lambda model: torch.optim.AdamW(model.parameters())
-        scheduler_fn = lambda optim: SchedulerFactory.load(optim, schedule_config)
-
-        train_config = TrainConfig(
-            strategy="seq",
+        train_config = train_config_factory(
             model=base_test_env["model"],
             dataset=random_dataset,
-            optimizer_fn=optimizer_fn,
-            scheduler_fn=scheduler_fn,
-            ckpt_dir=base_test_env["test_dir"],
-            n_epoch=1,
+            test_dir=base_test_env["test_dir"],
+            device=base_test_env["device"],
             batch_size=config["batch_size"],
-            ckpt_interval=5,
             accumulation_steps=config["accumulation_steps"],
-            max_grad_norm=1.0,
-            random_seed=42,
-            device_type=base_test_env["device"],
         )
 
         assert train_config.accumulation_steps == config["accumulation_steps"]