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

refactor: eliminate test duplication via shared helpers

Browse Source 
- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers
- Replace 15 copies of device one-liner with session-scoped fixture
- Collapse 5 near-identical Dataset subclasses into RandomTokenDataset
- Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions
- Make test_callbacks and test_early_stopping use existing train_config_factory
- Replace 6 duplicate meta.json read blocks with load_shard_meta
- Fix mkdtemp leaks in test_lora.py with TemporaryDirectory
This commit is contained in:
ViperEkura
2026-07-27 22:34:53 +08:00
parent c26a47b0df
commit 5ba21f4eb3
16 changed files with 364 additions and 661 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/trainer/conftest.py
+5 -46
View File
@@ -2,33 +2,15 @@ import os
import pytest
import torch
from torch.utils.data import Dataset
from astrai.config import TrainConfig
from astrai.trainer.schedule import SchedulerFactory
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,)),
}
from tests.helpers import RandomTokenDataset
def create_train_config(
model_fn,
dataset: Dataset,
dataset,
test_dir: str,
device: str,
strategy: str = "seq",
@@ -40,25 +22,7 @@ def create_train_config(
random_seed: int = 42,
**kwargs,
):
"""Factory function to create common TrainConfig for tests.
Args:
model_fn: Model factory (callable returning nn.Module)
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_per_device: Batch size per device (default: 2)
grad_accum_steps: Gradient accumulation steps (default: 1)
max_grad_norm: Maximum gradient norm for clipping (default: 1.0)
ckpt_interval: Checkpoint save interval in optimizer steps (default: 5)
random_seed: Random seed for reproducibility (default: 42)
**kwargs: Additional arguments passed to TrainConfig
Returns:
TrainConfig instance configured for testing
"""
"""Factory function to create common TrainConfig for tests."""
def optimizer_fn(m):
return torch.optim.AdamW(m.parameters(), lr=0.001)
@@ -89,16 +53,11 @@ def create_train_config(
@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.
"""
"""Fixture providing the ``create_train_config`` factory function."""
return create_train_config
@pytest.fixture
def trainer_dataset():
"""Fixture providing a dataset for trainer tests."""
dataset = TrainerDataset()
yield dataset
return RandomTokenDataset()
tests/trainer/test_callbacks.py
+12 -52
View File
@@ -1,10 +1,6 @@
import os
import torch
from astrai.config.train_config import TrainConfig
from astrai.model.components.decoder_block import DecoderBlock
from astrai.trainer.schedule import SchedulerFactory
from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
from astrai.trainer.trainer import Trainer
@@ -94,69 +90,35 @@ def test_gradient_checkpointing_backward(test_model):
assert p.grad is None or p.grad.sum().item() == 0, f"{name} grad not zeroed"
def test_gradient_checkpointing_trainer_integration(base_test_env, random_dataset):
def test_gradient_checkpointing_trainer_integration(
base_test_env, random_dataset, train_config_factory, device
):
"""Gradient checkpointing runs end-to-end via Trainer."""
def optimizer_fn(model):
return torch.optim.AdamW(model.parameters())
def scheduler_fn(optim):
return SchedulerFactory.create(
"cosine", optim, warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
train_config = TrainConfig(
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
strategy="seq",
dataset=random_dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
ckpt_dir=base_test_env["test_dir"],
log_dir=os.path.join(base_test_env["test_dir"], "logs"),
n_epoch=1,
batch_per_device=2,
test_dir=base_test_env["test_dir"],
device=device,
ckpt_interval=3,
grad_accum_steps=1,
max_grad_norm=1.0,
random_seed=42,
device_type=base_test_env["device"],
gradient_checkpointing_modules=[DecoderBlock],
)
trainer = Trainer(train_config)
trainer.train()
# no crash = callback correctly enabled/disabled
def test_callback_integration(base_test_env, random_dataset):
def test_callback_integration(
base_test_env, random_dataset, train_config_factory, device
):
"""Test that all callbacks are properly integrated"""
def optimizer_fn(model):
return torch.optim.AdamW(model.parameters())
def scheduler_fn(optim):
return SchedulerFactory.create(
"cosine", optim, warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
train_config = TrainConfig(
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
strategy="seq",
dataset=random_dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
ckpt_dir=base_test_env["test_dir"],
log_dir=os.path.join(base_test_env["test_dir"], "logs"),
n_epoch=1,
batch_per_device=2,
test_dir=base_test_env["test_dir"],
device=device,
ckpt_interval=3,
grad_accum_steps=1,
max_grad_norm=1.0,
random_seed=42,
device_type=base_test_env["device"],
)
# Create custom callbacks to track calls
callback_calls = []
class TrackingCallback(TrainCallback):
@@ -170,10 +132,8 @@ def test_callback_integration(base_test_env, random_dataset):
callback_calls.append("on_epoch_end")
trainer = Trainer(train_config, callbacks=[TrackingCallback()])
trainer.train()
# Verify callbacks were called
assert "on_train_begin" in callback_calls
assert "on_batch_end" in callback_calls
assert "on_epoch_end" in callback_calls
tests/trainer/test_early_stopping.py
+8 -29
View File
@@ -1,43 +1,25 @@
import os
import numpy as np
import torch
from astrai.config.train_config import TrainConfig
from astrai.trainer.schedule import SchedulerFactory
from astrai.trainer.trainer import Trainer
from tests.helpers import load_checkpoint_meta
def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
def test_early_stopping_simulation(
base_test_env, early_stopping_dataset, train_config_factory, device
):
"""Simulate early stopping behavior"""
def optimizer_fn(model):
return torch.optim.AdamW(model.parameters())
def scheduler_fn(optim):
return SchedulerFactory.create(
"cosine", optim, warmup_steps=10, lr_decay_steps=10, min_rate=0.05
)
train_config = TrainConfig(
strategy="seq",
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
train_config = train_config_factory(
model_fn=lambda: base_test_env["model"],
dataset=early_stopping_dataset,
ckpt_dir=base_test_env["test_dir"],
log_dir=os.path.join(base_test_env["test_dir"], "logs"),
test_dir=base_test_env["test_dir"],
device=device,
n_epoch=2,
batch_per_device=2,
ckpt_interval=1,
grad_accum_steps=2,
random_seed=np.random.randint(1e4),
device_type=base_test_env["device"],
)
trainer = Trainer(train_config)
# Should handle early stopping gracefully
try:
trainer.train()
except Exception:
@@ -50,8 +32,5 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
# Verify checkpoint was saved at expected step
load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_step_5")
import json
with open(os.path.join(load_dir, "meta.json")) as f:
meta = json.load(f)
meta = load_checkpoint_meta(load_dir)
assert meta["consumed_samples"] == 20
tests/trainer/test_grpo_strategy.py
+7 -54
View File
@@ -1,35 +1,9 @@
import pytest
import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.strategy import GRPOStrategy
class _FakeExecutor:
"""Minimal executor stub providing ``unwrap_model`` for ref model creation."""
def unwrap_model(self, model):
return model.state_dict()
def _make_config(vocab_size=200, max_position_embeddings=64):
return AutoRegressiveLMConfig(
vocab_size=vocab_size,
hidden_size=16,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=32,
max_position_embeddings=max_position_embeddings,
num_hidden_layers=2,
rms_norm_eps=1e-5,
)
def _make_model(device):
config = _make_config()
model = AutoRegressiveLM(config).to(device=device)
return model, config
from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config
def _make_batch(
@@ -44,9 +18,7 @@ def _make_batch(
responses = torch.randint(
0, 200, (batch_size, group_size, response_len), device=device
)
# All response tokens valid.
masks = torch.ones(batch_size, group_size, response_len, device=device)
# Distinct rewards per group member so std > 0.
rewards = torch.randn(batch_size, group_size, device=device)
return {
"prompts": prompts,
@@ -56,23 +28,12 @@ def _make_batch(
}
def _make_frozen_copy(model, device):
"""Create a frozen copy of ``model`` with independent weights loaded."""
config = _make_config()
copy = AutoRegressiveLM(config).to(device=device)
copy.load_state_dict(model.state_dict())
copy.requires_grad_(False)
copy.eval()
return copy
@pytest.fixture
def grpo_strategy():
def grpo_strategy(device):
"""Build a GRPOStrategy with a small real model and fake executor."""
device = "cuda" if torch.cuda.is_available() else "cpu"
model, config = _make_model(device)
old_model = _make_frozen_copy(model, device)
ref_model = _make_frozen_copy(model, device)
model, config = make_model(device)
old_model = make_frozen(model, device)
ref_model = make_frozen(model, device)
strategy = GRPOStrategy(
model=model,
@@ -83,7 +44,7 @@ def grpo_strategy():
kl_coef=0.01,
group_size=4,
model_fn=lambda c=config: AutoRegressiveLM(c).to(device=device),
executor=_FakeExecutor(),
executor=FakeExecutor(),
)
return strategy, device
@@ -103,7 +64,6 @@ def test_grpo_loss_backward(grpo_strategy):
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
loss.backward()
# At least some parameter should receive a gradient.
has_grad = any(
p.grad is not None and p.grad.abs().sum().item() > 0
for p in strategy.model.parameters()
@@ -136,32 +96,27 @@ def test_grpo_prompt_tokens_masked(grpo_strategy):
the policy loss should be zero (no valid tokens contribute)."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
# Zero out all response masks → no response token contributes.
batch["masks"] = torch.zeros_like(batch["masks"])
loss = strategy.compute_loss(batch)
# With no valid tokens, policy_loss term is 0 and KL term is 0.
assert loss.item() == pytest.approx(0.0, abs=1e-6)
def test_grpo_identical_rewards_zero_advantage(grpo_strategy):
"""When all group rewards are identical, advantage is 0 → policy_loss is 0.
"""When all group rewards are identical, advantage is 0 -> policy_loss is 0.
Only the KL term remains (which is 0 when policy == ref at init)."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
batch["rewards"] = torch.ones(batch["rewards"].shape, device=device)
loss = strategy.compute_loss(batch)
# At init policy == old == ref, so ratio == 1, KL == 0; advantage == 0.
assert loss.item() == pytest.approx(0.0, abs=1e-5)
def test_grpo_sync_old_model(grpo_strategy):
"""sync_old_model copies current policy weights into old_model."""
strategy, device = grpo_strategy
# Perturb policy model so it differs from old.
with torch.no_grad():
for p in strategy.model.parameters():
p.add_(0.05)
# old_model should still hold original weights (differ from policy).
policy_sd = strategy.model.state_dict()
old_sd = strategy.old_model.state_dict()
differs_before = any(
@@ -195,14 +150,12 @@ 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
# Diverge policy from ref.
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()
# With distinct rewards and diverged policy, loss should be non-trivial.
assert loss.abs().item() > 1e-4
tests/trainer/test_online_e2e.py
+2 -14
View File
@@ -12,19 +12,7 @@ from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.rollout import BaseRewardModel
from astrai.trainer.schedule import SchedulerFactory
from astrai.trainer.trainer import Trainer
_CHAT_TEMPLATE = (
"{% for message in messages %}"
"{% if message['role'] == 'system' %}"
"SYSTEM: {{ message['content'] }}\n"
"{% elif message['role'] == 'user' %}"
"USER: {{ message['content'] }}\n"
"{% elif message['role'] == 'assistant' %}"
"ASSISTANT: {{ message['content'] }}\n"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}ASSISTANT: {% endif %}"
)
from tests.helpers import CHAT_TEMPLATE
class InstructionDataset(Dataset):
@@ -97,7 +85,7 @@ def test_online_dpo_end_to_end(base_test_env):
# Equip tokenizer with a chat template so RolloutGenerator can
# render instruction/input via apply_chat_template.
tokenizer.set_chat_template(_CHAT_TEMPLATE)
tokenizer.set_chat_template(CHAT_TEMPLATE)
tokenizer.save_pretrained(test_dir)
model_fn = partial(_model_fn, model_config)
tests/trainer/test_online_strategy.py
+13 -59
View File
@@ -9,7 +9,6 @@ the per-strategy ``prepare_from_rollout`` mappings for both
import pytest
import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.rollout import RolloutResult
from astrai.trainer.strategy import (
@@ -17,47 +16,7 @@ from astrai.trainer.strategy import (
GRPOStrategy,
StrategyFactory,
)
class _FakeExecutor:
"""Executor stub tracking ``sync_gradients`` and providing unwrap_model."""
def __init__(self, sync_gradients=True):
self._sync_gradients = sync_gradients
@property
def sync_gradients(self):
return self._sync_gradients
def unwrap_model(self, model):
return model.state_dict()
def _make_config(vocab_size=200, max_position_embeddings=64):
return AutoRegressiveLMConfig(
vocab_size=vocab_size,
hidden_size=16,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=32,
max_position_embeddings=max_position_embeddings,
num_hidden_layers=2,
rms_norm_eps=1e-5,
)
def _make_model(device):
cfg = _make_config()
return AutoRegressiveLM(cfg).to(device=device), cfg
def _make_frozen(model, device):
cfg = _make_config()
copy = AutoRegressiveLM(cfg).to(device=device)
copy.load_state_dict(model.state_dict())
copy.requires_grad_(False)
copy.eval()
return copy
from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config
def _make_rollout_result(B=2, G=4, P=6, R=8, device="cpu"):
@@ -75,7 +34,7 @@ class _RecordingRunner:
"""Fake RolloutRunner returning a fixed result with freshness tracking.
Freshness is ``True`` on the first call after construction or after
:meth:`swap_result`; ``False`` on subsequent cached calls — mirroring
:meth:`swap_result`; ``False`` on subsequent cached calls -- mirroring
the real ``RolloutRunner`` contract without invoking generation.
"""
@@ -99,15 +58,10 @@ class _RecordingRunner:
self._fresh = True
@pytest.fixture
def device():
return "cuda" if torch.cuda.is_available() else "cpu"
def _make_grpo(device, executor=None):
model, _ = _make_model(device)
old_model = _make_frozen(model, device)
ref_model = _make_frozen(model, device)
model, _ = make_model(device)
old_model = make_frozen(model, device)
ref_model = make_frozen(model, device)
return GRPOStrategy(
model=model,
device=device,
@@ -116,22 +70,22 @@ def _make_grpo(device, executor=None):
clip_eps=0.2,
kl_coef=0.01,
group_size=4,
model_fn=lambda c=_make_config(): AutoRegressiveLM(c).to(device=device),
executor=executor or _FakeExecutor(),
model_fn=lambda c=make_rollout_config(): AutoRegressiveLM(c).to(device=device),
executor=executor or FakeExecutor(),
)
def _make_dpo(device, executor=None):
model, _ = _make_model(device)
ref_model = _make_frozen(model, device)
model, _ = make_model(device)
ref_model = make_frozen(model, device)
return DPOStrategy(
model=model,
device=device,
ref_model=ref_model,
beta=0.1,
reduction="sum",
model_fn=lambda c=_make_config(): AutoRegressiveLM(c).to(device=device),
executor=executor or _FakeExecutor(),
model_fn=lambda c=make_rollout_config(): AutoRegressiveLM(c).to(device=device),
executor=executor or FakeExecutor(),
)
@@ -295,7 +249,7 @@ def test_dpo_no_sync_hook_when_new_rollout_result(device):
def test_step_not_called_when_sync_gradients_false(device):
executor = _FakeExecutor(sync_gradients=False)
executor = FakeExecutor(sync_gradients=False)
strat = _make_grpo(device, executor=executor)
runner = _RecordingRunner(_make_rollout_result(device=device))
strat.set_rollout_runner(runner)
@@ -304,7 +258,7 @@ def test_step_not_called_when_sync_gradients_false(device):
def test_step_called_when_sync_gradients_true(device):
executor = _FakeExecutor(sync_gradients=True)
executor = FakeExecutor(sync_gradients=True)
strat = _make_grpo(device, executor=executor)
runner = _RecordingRunner(_make_rollout_result(device=device))
strat.set_rollout_runner(runner)
tests/trainer/test_rollout.py
+5 -75
View File
@@ -3,9 +3,7 @@
import pytest
import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.inference.core.scheduler import InferenceScheduler
from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.rollout import (
BaseRewardModel,
RawRollout,
@@ -13,50 +11,7 @@ from astrai.trainer.rollout import (
RolloutResult,
RolloutRunner,
)
_CHAT_TEMPLATE = (
"{% for message in messages %}"
"{% if message['role'] == 'system' %}SYSTEM: {{ message['content'] }}\n{% endif %}"
"{% if message['role'] == 'user' %}USER: {{ message['content'] }}\n{% endif %}"
"{% if message['role'] == 'assistant' %}ASSISTANT: {{ message['content'] }}\n{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}ASSISTANT: {% endif %}"
)
class FakeTokenizer:
"""Minimal stub tokenizer with a chat template for rollout tests."""
stop_ids = [2]
def __init__(self):
from astrai.tokenize.chat_template import ChatTemplate
self._chat_template = ChatTemplate.from_string(_CHAT_TEMPLATE)
def encode(self, texts, **_):
if isinstance(texts, str):
texts = [texts]
return [[b for b in t.encode("utf-8")] for t in texts]
def decode(self, ids, skip_special_tokens=True):
if isinstance(ids, list):
return bytes(b for b in ids if b > 2).decode("utf-8", errors="ignore")
return str(ids)
def apply_chat_template(
self, messages, tokenize=True, add_generation_prompt=True, **_
):
rendered = self._chat_template.render(
messages=messages, add_generation_prompt=add_generation_prompt
)
if tokenize:
return (
self.encode(rendered)[0]
if isinstance(rendered, str)
else [self.encode(t)[0] for t in rendered]
)
return rendered
from tests.helpers import FakeTokenizer, make_model
class ConstantRewardModel(BaseRewardModel):
@@ -83,26 +38,6 @@ class NonFiniteRewardModel(BaseRewardModel):
return torch.full((B, G), float("nan"))
def _make_config(vocab_size=200, max_position_embeddings=128):
return AutoRegressiveLMConfig(
vocab_size=vocab_size,
hidden_size=16,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=32,
max_position_embeddings=max_position_embeddings,
num_hidden_layers=2,
rms_norm_eps=1e-5,
)
def _make_model(device):
cfg = _make_config()
m = AutoRegressiveLM(cfg).to(device=device)
m.eval()
return m, cfg
def _make_scheduler(model, tokenizer, max_batch_size=8, max_len=128):
return InferenceScheduler(
model=model,
@@ -160,14 +95,9 @@ def test_constant_reward_model_shape():
assert torch.all(out == 0.5)
@pytest.fixture
def device():
return "cuda" if torch.cuda.is_available() else "cpu"
def _make_generator(device, **kw):
model, _ = _make_model(device)
tokenizer = FakeTokenizer()
model, _ = make_model(device, max_position_embeddings=128)
tokenizer = FakeTokenizer(with_chat_template=True)
scheduler = _make_scheduler(
model,
tokenizer,
@@ -229,7 +159,7 @@ def test_rollout_generator_mask_matches_responses(device):
def test_rollout_generator_logprobs_are_nonpositive(device):
"""Behaviour-policy logprobs of sampled tokens should be ≤ 0."""
"""Behaviour-policy logprobs of sampled tokens should be <= 0."""
gen, _ = _make_generator(device, group_size=2, max_tokens=4)
batch = _make_instruction_batch(n=1)
r = gen.generate(batch)
@@ -241,7 +171,7 @@ def test_rollout_generator_logprobs_are_nonpositive(device):
def test_rollout_generator_instruction_role_mapping(device):
"""instruction → system, input → user, output → assistant."""
"""instruction -> system, input -> user, output -> assistant."""
gen, _ = _make_generator(device, group_size=1, max_tokens=4)
batch = {
"instruction": ["Be helpful"],
tests/trainer/test_signal_handler.py
+8 -33
View File
@@ -1,4 +1,3 @@
import json
import multiprocessing as mp
import os
import signal
@@ -10,15 +9,15 @@ import torch.optim as optim
from torch.utils.data import Dataset
from astrai.config import TrainConfig
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.model.transformer import AutoRegressiveLM
from astrai.parallel.signal_handler import register_signal_handlers
from astrai.trainer import Trainer
from astrai.trainer.schedule import SchedulerFactory
from astrai.trainer.train_context import TrainContext
from tests.helpers import load_checkpoint_meta, make_tiny_config
class _PicklableDataset(Dataset):
class PicklableDataset(Dataset):
def __init__(self, length=200, max_length=64, vocab_size=1000):
self.length = length
self.max_length = max_length
@@ -35,16 +34,7 @@ class _PicklableDataset(Dataset):
def _build_model():
config = AutoRegressiveLMConfig(
vocab_size=1000,
hidden_size=8,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=16,
max_position_embeddings=64,
num_hidden_layers=2,
rms_norm_eps=1e-5,
)
config = make_tiny_config()
device = "cuda" if torch.cuda.is_available() else "cpu"
return AutoRegressiveLM(config).to(device=device)
@@ -61,7 +51,7 @@ class _ReadyCallback:
def _inner_run(batch_per_device, ckpt_interval, ckpt_dir, log_dir, ready_file):
dataset = _PicklableDataset()
dataset = PicklableDataset()
def model_fn():
return _build_model()
@@ -147,17 +137,7 @@ 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)"
ckpt_dir = base_test_env["test_dir"]
meta_files = []
for root, dirs, files in os.walk(ckpt_dir):
for f in files:
if f == "meta.json":
meta_files.append(os.path.join(root, f))
assert len(meta_files) > 0, f"No checkpoint meta.json found in {ckpt_dir}"
with open(meta_files[-1]) as f:
meta = json.load(f)
meta = load_checkpoint_meta(base_test_env["test_dir"])
assert "consumed_samples" in meta
assert meta["consumed_samples"] >= 0
@@ -167,11 +147,6 @@ def test_sigint_triggers_checkpoint_save(base_test_env):
exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGINT)
assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
ckpt_dir = base_test_env["test_dir"]
meta_files = []
for root, dirs, files in os.walk(ckpt_dir):
for f in files:
if f == "meta.json":
meta_files.append(os.path.join(root, f))
assert len(meta_files) > 0, f"No checkpoint meta.json found in {ckpt_dir}"
meta = load_checkpoint_meta(base_test_env["test_dir"])
assert "consumed_samples" in meta
assert meta["consumed_samples"] >= 0