fix: 修复 dataset 和 checkpoint 的 bug

tests/data/test_checkpoint.py

@@ -1,12 +1,12 @@
-import os
 import torch
 import tempfile
+import torch.distributed as dist
 
 from pathlib import Path
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from khaosz.data.checkpoint import Checkpoint
-from khaosz.parallel.setup import spawn_parallel_fn
+from khaosz.parallel.setup import get_rank, spawn_parallel_fn
 
 def test_single_process():
     model = torch.nn.Linear(10, 5)
@@ -26,8 +26,7 @@ def test_single_process():
         scheduler.step()
     
     checkpoint = Checkpoint(
-        optimizer_state_dict=optimizer.state_dict(),
-        scheduler_state_dict=scheduler.state_dict(),
+        state_dict=model.state_dict(),
         epoch=3,
         iteration=30,
         metrics={
@@ -45,21 +44,14 @@ def test_single_process():
         assert loaded_checkpoint.iteration == 30
         assert loaded_checkpoint.metrics["loss"] == [0.5, 0.4, 0.3, 0.2, 0.1]
         
-        assert 'param_groups' in loaded_checkpoint.optimizer_state_dict
-        assert 'state' in loaded_checkpoint.optimizer_state_dict
-        
         png_files = list(Path(tmpdir).glob("*.png"))
         assert png_files
 
 def simple_training():
-    rank = int(os.environ.get('LOCAL_RANK', 0))
-    
-    # 简单的训练逻辑
     model = torch.nn.Linear(10, 5)
     optimizer = AdamW(model.parameters(), lr=1e-3)
     scheduler = CosineAnnealingLR(optimizer, T_max=10)
     
-    # 训练步骤
     for epoch in range(2):
         for iteration in range(5):
             x = torch.randn(16, 10)
@@ -71,18 +63,29 @@ def simple_training():
         scheduler.step()
 
     checkpoint = Checkpoint(
-        optimizer_state_dict=optimizer.state_dict(),
-        scheduler_state_dict=scheduler.state_dict(),
+        state_dict=model.state_dict(),
         epoch=2,
         iteration=10,
         metrics={"loss": [0.3, 0.2, 0.1]}
     )
     
-    with tempfile.TemporaryDirectory() as tmpdir:
-        checkpoint.save(tmpdir)
-        loaded = Checkpoint.load(tmpdir)
-        assert loaded.epoch == 2
-        print(f"Rank {rank}: Checkpoint test passed")
+    rank = get_rank()
+    
+    if rank == 0:
+        shared_dir = tempfile.mkdtemp()
+        checkpoint.save(shared_dir)
+    else:
+        shared_dir = None
+
+    
+    if dist.is_initialized():
+        dir_list = [shared_dir]
+        dist.broadcast_object_list(dir_list, src=0)
+        shared_dir = dir_list[0]
+    
+ 
+    loaded = Checkpoint.load(shared_dir)
+    assert loaded.epoch == 2
 
 def test_multi_process():
     spawn_parallel_fn(