ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

style: 使用ruff 工具优化代码风格

Browse Source
This commit is contained in:
ViperEkura
2026-03-30 23:32:28 +08:00
parent 345fd2f091
commit 426af2d75f
52 changed files with 1836 additions and 1493 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/data/test_checkpoint.py
+15 -22
View File
@@ -7,6 +7,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
from khaosz.data.serialization import Checkpoint
from khaosz.parallel.setup import get_rank, spawn_parallel_fn
def test_single_process():
model = torch.nn.Linear(10, 5)
optimizer = AdamW(model.parameters(), lr=1e-3)
@@ -14,34 +15,31 @@ def test_single_process():
for epoch in range(3):
for iteration in range(10):
x = torch.randn(32, 10)
y = torch.randn(32, 5)
loss = model(x).mean()
loss.backward()
optimizer.step()
optimizer.zero_grad()
scheduler.step()
checkpoint = Checkpoint(
state_dict=model.state_dict(),
epoch=3,
iteration=30
)
checkpoint = Checkpoint(state_dict=model.state_dict(), epoch=3, iteration=30)
with tempfile.TemporaryDirectory() as tmpdir:
checkpoint.save(tmpdir)
loaded_checkpoint = Checkpoint.load(tmpdir)
assert loaded_checkpoint.epoch == 3
assert loaded_checkpoint.iteration == 30
def simple_training():
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)
@@ -57,28 +55,23 @@ def simple_training():
epoch=2,
iteration=10,
)
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(
simple_training,
world_size=2,
backend="gloo"
)
spawn_parallel_fn(simple_training, world_size=2, backend="gloo")
tests/data/test_dataset.py
+44 -45
View File
@@ -5,30 +5,32 @@ from khaosz.data.serialization import save_h5
from khaosz.data.dataset import *
def test_dataset_loader_random_paths(base_test_env):
"""Test dataset loader with multiple random paths"""
test_dir = base_test_env["test_dir"]
# Create multiple mmap dataset directories with random data
num_files = np.random.randint(2, 5)
for i in range(num_files):
seq_length = np.random.randint(200, 400)
dummy_data = {
"sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64) for _ in range(10)],
"sequence": [
torch.randint(0, 1000, (seq_length,), dtype=torch.int64)
for _ in range(10)
],
}
save_h5(test_dir, f"data_{i}", dummy_data)
# Test loading with multiple paths
loaded_dataset = DatasetLoader.load(
train_type="seq",
load_path=test_dir,
window_size=64,
train_type="seq",
load_path=test_dir,
window_size=64,
)
assert loaded_dataset is not None
assert len(loaded_dataset) > 0
# Test that we can get items without errors
for i in range(len(loaded_dataset)):
item = loaded_dataset[i]
@@ -41,30 +43,30 @@ def test_dataset_loader_random_paths(base_test_env):
def test_dpo_strategy_with_random_data(base_test_env):
"""Test DPO strategy with randomized preference data"""
test_dir = base_test_env["test_dir"]
# Create DPO-style data with memory mapping format
seq_length = np.random.randint(100, 200)
dummy_data = {
"chosen": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
"rejected": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
"chosen_mask": [torch.ones(seq_length, dtype=torch.bool)],
"rejected_mask": [torch.ones(seq_length, dtype=torch.bool)]
"rejected_mask": [torch.ones(seq_length, dtype=torch.bool)],
}
save_h5(test_dir, "dpo_data", dummy_data)
# Load DPO dataset
dpo_dataset = DatasetLoader.load(
train_type="dpo",
load_path=test_dir,
window_size=64,
train_type="dpo",
load_path=test_dir,
window_size=64,
)
assert dpo_dataset is not None
assert hasattr(dpo_dataset, 'fetcher')
assert hasattr(dpo_dataset, "fetcher")
assert len(dpo_dataset) > 0
# Test that we can get DPO items without errors
for i in range(min(3, len(dpo_dataset))):
item = dpo_dataset[i]
@@ -79,28 +81,28 @@ def test_dpo_strategy_with_random_data(base_test_env):
def test_sft_dataset_with_random_data(base_test_env):
"""Test SFT dataset with random data"""
test_dir = base_test_env["test_dir"]
# Create SFT-style data with memory mapping format
seq_length = np.random.randint(100, 200)
dummy_data = {
"sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
"loss_mask": [torch.ones(seq_length, dtype=torch.bool)]
"loss_mask": [torch.ones(seq_length, dtype=torch.bool)],
}
save_h5(test_dir, "sft_data", dummy_data)
# Load SFT dataset
sft_dataset = DatasetLoader.load(
train_type="sft",
load_path=test_dir,
window_size=64,
train_type="sft",
load_path=test_dir,
window_size=64,
)
assert sft_dataset is not None
assert hasattr(sft_dataset, 'fetcher')
assert hasattr(sft_dataset, "fetcher")
assert len(sft_dataset) > 0
# Test that we can get SFT items without errors
for i in range(min(3, len(sft_dataset))):
item = sft_dataset[i]
@@ -114,33 +116,30 @@ def test_sft_dataset_with_random_data(base_test_env):
def test_dataset_with_custom_stride(base_test_env):
"""Test dataset with custom stride parameter"""
test_dir = base_test_env["test_dir"]
# Create test data
seq_length = 200
dummy_data = {
"sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
}
save_h5(test_dir,"stride_test_data", dummy_data)
save_h5(test_dir, "stride_test_data", dummy_data)
# Test with custom stride
custom_stride = 32
dataset = DatasetLoader.load(
train_type="seq",
load_path=test_dir,
window_size=64,
stride=custom_stride
train_type="seq", load_path=test_dir, window_size=64, stride=custom_stride
)
assert dataset is not None
assert len(dataset) > 0
# With stride 32 and window 64 on 200 length data, we should get more samples
# than with default stride (which equals window size)
default_stride_dataset = DatasetLoader.load(
train_type="seq",
load_path=test_dir,
window_size=64,
train_type="seq",
load_path=test_dir,
window_size=64,
)
assert len(dataset) > len(default_stride_dataset)
tests/data/test_sampler.py
+14 -12
View File
@@ -1,30 +1,32 @@
from khaosz.trainer import *
from khaosz.data import *
def test_random_sampler_consistency(random_dataset):
"""Test RandomSampler produces consistent results with same seed"""
dataset = random_dataset
# Create two samplers with same seed
sampler1 = ResumableDistributedSampler(dataset, seed=42)
sampler2 = ResumableDistributedSampler(dataset, seed=42)
indices1 = list(iter(sampler1))
indices2 = list(iter(sampler2))
assert indices1 == indices2
def test_random_sampler_different_seeds(random_dataset):
"""Test RandomSampler produces different results with different seeds"""
dataset = random_dataset
# Create two samplers with different seeds
sampler1 = ResumableDistributedSampler(dataset, seed=42)
sampler2 = ResumableDistributedSampler(dataset, seed=123)
indices1 = list(iter(sampler1))
indices2 = list(iter(sampler2))
# Very high probability they should be different
assert indices1 != indices2
@@ -33,20 +35,20 @@ def test_sampler_across_epochs(random_dataset):
"""Test sampler behavior across multiple epochs"""
dataset = random_dataset
n = len(dataset)
sampler = ResumableDistributedSampler(dataset, seed=42)
# Get indices for first epoch
epoch1_indices = list(iter(sampler))
assert len(epoch1_indices) == n
# Get indices for second epoch
epoch2_indices = list(iter(sampler))
assert len(epoch2_indices) == n
# Check that epochs have different order (should be random)
assert epoch1_indices != epoch2_indices
# Check that all indices are present in each epoch
assert set(epoch1_indices) == set(range(n))
assert set(epoch2_indices) == set(range(n))
assert set(epoch2_indices) == set(range(n))