Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
lightning/tests/utilities/test_auto_restart.py

1199 lines
44 KiB
Python
Raw Blame History

# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import os
import random
import random as python_random
from collections import defaultdict
from collections.abc import Iterable
from contextlib import suppress
from copy import deepcopy
from typing import List, Optional
from unittest import mock
from unittest.mock import ANY
import numpy as np
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.utils.data import BatchSampler, DistributedSampler, RandomSampler, SequentialSampler
from torch.utils.data._utils.worker import get_worker_info
from torch.utils.data.dataloader import DataLoader, default_collate
from torch.utils.data.dataset import Dataset, IterableDataset
import tests.helpers.utils as tutils
from pytorch_lightning import Callback, LightningModule, seed_everything, Trainer
from pytorch_lightning.utilities.auto_restart import (
_add_capture_metadata_collate,
_dataloader_load_state_dict,
_dataloader_to_state_dict,
CaptureIterableDataset,
CaptureMapDataset,
FastForwardSampler,
MergedIteratorState,
)
from pytorch_lightning.utilities.enums import AutoRestartBatchKeys
from pytorch_lightning.utilities.exceptions import ExitGracefullyException, MisconfigurationException
from pytorch_lightning.utilities.fetching import DataFetcher
from pytorch_lightning.utilities.imports import _fault_tolerant_training
from tests.helpers.boring_model import BoringModel, RandomDataset
from tests.helpers.runif import RunIf
# Credit to PyTorch Team.
# Taken from:
# https://github.com/pytorch/pytorch/blob/3b977a0d2834d300c0301a0c6af98c8e939019ce/torch/utils/data/_utils/worker.py#L151
# Not available until torch 1.9.0
def _generate_state(base_seed, worker_id):
INIT_A = 0x43B0D7E5
MULT_A = 0x931E8875
INIT_B = 0x8B51F9DD
MULT_B = 0x58F38DED
MIX_MULT_L = 0xCA01F9DD
MIX_MULT_R = 0x4973F715
XSHIFT = 4 * 8 // 2
MASK32 = 0xFFFFFFFF
entropy = [worker_id, base_seed & MASK32, base_seed >> 32, 0]
pool = [0] * 4
hash_const_A = INIT_A
def hash(value):
nonlocal hash_const_A
value = (value ^ hash_const_A) & MASK32
hash_const_A = (hash_const_A * MULT_A) & MASK32
value = (value * hash_const_A) & MASK32
value = (value ^ (value >> XSHIFT)) & MASK32
return value
def mix(x, y):
result_x = (MIX_MULT_L * x) & MASK32
result_y = (MIX_MULT_R * y) & MASK32
result = (result_x - result_y) & MASK32
result = (result ^ (result >> XSHIFT)) & MASK32
return result
# Add in the entropy to the pool.
for i in range(len(pool)):
pool[i] = hash(entropy[i])
# Mix all bits together so late bits can affect earlier bits.
for i_src in range(len(pool)):
for i_dst in range(len(pool)):
if i_src != i_dst:
pool[i_dst] = mix(pool[i_dst], hash(pool[i_src]))
hash_const_B = INIT_B
state = []
for i_dst in range(4):
data_val = pool[i_dst]
data_val = (data_val ^ hash_const_B) & MASK32
hash_const_B = (hash_const_B * MULT_B) & MASK32
data_val = (data_val * hash_const_B) & MASK32
data_val = (data_val ^ (data_val >> XSHIFT)) & MASK32
state.append(data_val)
return state
def test_fast_forward_getattr():
dataset = range(15)
sampler = SequentialSampler(dataset)
batch_sampler = BatchSampler(sampler, 3, False)
index_batch_sampler = FastForwardSampler(batch_sampler)
assert index_batch_sampler.batch_size == 3
assert index_batch_sampler.sampler == sampler
def test_fast_forward_on_batch_sampler():
"""This test ensures ``FastForwardSampler`` applied to ``BatchSampler`` correctly retrived the right next batch
on restart."""
dataset = range(15)
sampler = SequentialSampler(dataset)
batch_sampler = BatchSampler(sampler, 3, False)
index_batch_sampler = FastForwardSampler(batch_sampler)
assert isinstance(index_batch_sampler, Iterable)
index_batch_sampler_iter = iter(index_batch_sampler)
assert next(index_batch_sampler_iter) == [0, 1, 2]
assert next(index_batch_sampler_iter) == [3, 4, 5]
state_dict = index_batch_sampler.state_dict(2)
index_batch_sampler = FastForwardSampler(batch_sampler)
index_batch_sampler.load_state_dict(state_dict)
index_batch_sampler_iter = iter(index_batch_sampler)
assert next(index_batch_sampler_iter) == [6, 7, 8]
def test_fast_forward_on_sequential_sampler():
"""This test ensures ``FastForwardSampler`` applied to ``SequentialSampler`` correctly retrived the right next
batch on restart."""
dataset = range(15)
sequential_sampler = SequentialSampler(dataset)
sampler = FastForwardSampler(sequential_sampler)
sampler.setup(3)
batch_sampler = BatchSampler(sampler, 3, False)
batch_sampler_iter = iter(batch_sampler)
assert next(batch_sampler_iter) == [0, 1, 2]
assert next(batch_sampler_iter) == [3, 4, 5]
state_dict = sampler.state_dict(2)
assert state_dict[0]["current_iteration"] == 6
sampler.load_state_dict(state_dict)
batch_sampler_iter = iter(batch_sampler)
assert next(batch_sampler_iter) == [6, 7, 8]
@pytest.mark.skipif(torch.cuda.is_available(), reason="todo (tchaton) Need more investigation")
def test_fast_forward_on_random_sampler():
"""This test ensures ``FastForwardSampler`` applied to ``RandomSampler`` correctly retrived the right next
batch on restart."""
seed = 42
seed_everything(42)
dataset = range(15)
generator = torch.Generator().manual_seed(seed)
values = list(RandomSampler(dataset, generator=generator))
generator = torch.Generator().manual_seed(seed)
random_sampler = RandomSampler(dataset, generator=generator)
sampler = FastForwardSampler(random_sampler)
sampler.setup(3)
batch_sampler = BatchSampler(sampler, 3, False)
batch_sampler_iter = iter(batch_sampler)
assert next(batch_sampler_iter) == values[:3]
assert next(batch_sampler_iter) == values[3:6]
assert next(batch_sampler_iter) == values[6:9]
state_dict = sampler.state_dict(3)
assert state_dict[0]["current_iteration"] == 9
state_dict[0]["current_iteration"] = 6
seed_everything(42)
generator = torch.Generator().manual_seed(seed)
random_sampler = RandomSampler(dataset, generator=generator)
sampler = FastForwardSampler(random_sampler)
sampler.setup(3)
batch_sampler = BatchSampler(sampler, 3, False)
sampler.load_state_dict(state_dict)
batch_sampler_iter = iter(batch_sampler)
assert next(batch_sampler_iter) == values[6:9]
has_raised = False
try:
for _ in range(5):
next(batch_sampler_iter)
except StopIteration:
has_raised = True
assert sampler._current_iteration == 0
sampler.load_state_dict(sampler.state_dict(0))
assert has_raised
class RangeIterableDataset(IterableDataset):
def __init__(self, data, num_workers: int, batch_size: int, state_dict=None, attr_name: str = "iter_sampler"):
self.data = list(data)
self.batch_size = batch_size
self.num_workers = num_workers
self.state_dict = state_dict
self.attr_name = attr_name
def __iter__(self):
worker_info = get_worker_info()
if worker_info and self.num_workers == 2:
id = worker_info.id
num_samples = len(self.data)
if id == 0:
self.data = list(self.data)[: num_samples // 2]
else:
self.data = list(self.data)[num_samples // 2 :]
self.user_sampler = RandomSampler(self.data)
else:
self.user_sampler = RandomSampler(self.data)
setattr(self, self.attr_name, iter(self.user_sampler))
return self
def __next__(self):
iter_sampler = getattr(self, self.attr_name)
return self.data[next(iter_sampler)]
@pytest.mark.skipif(torch.cuda.is_available(), reason="This test takes around 30 sec and should be skipped in Azure CI")
@pytest.mark.parametrize("num_workers", [0, 1, 2])
def test_fast_forward_sampler_over_iterable_dataset(num_workers):
"""This test ensures ``FastForwardSampler`` and ``CaptureIterableDataset`` are properly being used to capture
workers states."""
batch_size = 3
initial_seed = seed_everything(42)
generator = torch.Generator()
generator.manual_seed(initial_seed)
dataset = RangeIterableDataset(range(20), num_workers, batch_size, True)
dataset = CaptureIterableDataset(dataset)
dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, generator=generator)
_add_capture_metadata_collate(dataloader)
iter_dataloader = iter(dataloader)
batches = []
for _ in range(5):
batches.append(next(iter_dataloader))
# restarting on batch_1 and getting 3 extra batches
state_dict = {"iter_sampler": {}}
for batch in batches[:2]:
batch, _state_dict = batch["data"], batch[AutoRestartBatchKeys.PL_RESTART_META]
for k, v in _state_dict.items():
state_dict[k].update(v)
assert len(state_dict["iter_sampler"]) == (num_workers if num_workers > 1 else 1)
initial_seed = seed_everything(42)
generator.manual_seed(initial_seed)
dataset = RangeIterableDataset(range(20), num_workers, batch_size, state_dict=state_dict)
dataset = CaptureIterableDataset(dataset)
dataset.load_state_dict(state_dict)
dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, generator=generator)
_add_capture_metadata_collate(dataloader)
iter_dataloader = iter(dataloader)
batches_restart = []
for _ in range(3):
batches_restart.append(next(iter_dataloader))
assert torch.equal(batches_restart[0]["data"], batches[2]["data"])
assert torch.equal(batches_restart[1]["data"], batches[3]["data"])
assert torch.equal(batches_restart[2]["data"], batches[4]["data"])
def _setup_ddp(rank, worldsize):
os.environ["MASTER_ADDR"] = "localhost"
# initialize the process group
dist.init_process_group("gloo", rank=rank, world_size=worldsize)
def _test_fast_forward_sampler_with_distributed_sampler(rank, worldsize):
_setup_ddp(rank, worldsize)
initial_seed = seed_everything(42)
generator = torch.Generator()
generator.manual_seed(initial_seed)
num_workers = 2
batch_size = 4
dataset = range(30)
sampler = FastForwardSampler(DistributedSampler(dataset, num_replicas=worldsize, rank=rank, seed=initial_seed))
sampler.setup(batch_size)
dataloader = DataLoader(
dataset, batch_size=batch_size, num_workers=num_workers, generator=generator, sampler=sampler
)
iter_dataloader = iter(dataloader)
num_yielded = 0
batches = []
while True:
try:
batches.append(next(iter_dataloader))
num_yielded += 1
except StopIteration:
break
expected = torch.tensor([17, 27, 24]) if rank == 0 else torch.tensor([19, 5, 3])
assert torch.equal(batches[-1], expected)
assert sampler.state_dict(num_yielded)[0]["current_iteration"] == 16
reload_state_dict = sampler.state_dict(num_yielded - 1)
assert reload_state_dict[0]["current_iteration"] == 12
sampler = FastForwardSampler(DistributedSampler(dataset, num_replicas=worldsize, rank=rank, seed=initial_seed))
sampler.setup(batch_size)
sampler.load_state_dict(reload_state_dict)
dataloader = DataLoader(
dataset, batch_size=batch_size, num_workers=num_workers, generator=generator, sampler=sampler
)
iter_dataloader = iter(dataloader)
batches = []
while True:
try:
batches.append(next(iter_dataloader))
except StopIteration:
break
assert torch.equal(batches[-1], expected)
assert sampler.state_dict(num_yielded)[0]["current_iteration"] == 16
@pytest.mark.skipif(torch.cuda.is_available(), reason="This test takes around 25 sec and should be skipped in Azure CI")
@RunIf(skip_windows=True)
def test_fast_forward_sampler_with_distributed_sampler():
"""Make sure result logging works with DDP."""
tutils.set_random_main_port()
worldsize = 2
mp.spawn(_test_fast_forward_sampler_with_distributed_sampler, args=(worldsize,), nprocs=worldsize)
class MetaLearningDataset(IterableDataset):
def __init__(
self,
dataset: Dataset,
batch_size: int,
drop_last: bool,
task_num_classes: int = 5,
num_workers: Optional[int] = None,
global_rank: Optional[int] = None,
world_size: Optional[int] = None,
initial_seed: Optional[int] = None,
shuffle: bool = True,
debugging: bool = False,
):
self.dataset = dataset
self.batch_size = batch_size
self.drop_last = drop_last
self.num_workers = num_workers or 1
self.global_rank = global_rank
self.world_size = world_size
self.task_num_classes = task_num_classes
self.labels = labels = getattr(dataset, "labels")
self.initial_seed = initial_seed
self.generator: Optional[torch.Generator] = None
self.current_task_iteration = 0
self.shuffle = shuffle
self.debugging = debugging
if labels is None:
raise MisconfigurationException(f"Provided {self.dataset} should have an attribute labels.")
if len(labels) != len(dataset):
raise MisconfigurationException("Found provided ``labels`` don't match the dataset length.")
if (isinstance(global_rank, int) and world_size is None) or (
isinstance(world_size, int) and global_rank is None
):
raise MisconfigurationException("Both ``world_size`` and ``global_rank`` should be provided !")
self.unique_labels = np.unique(self.labels)
@property
def worker_id(self) -> int:
worker_info = get_worker_info()
return worker_info.id if worker_info else 0
@property
def is_distributed(self) -> bool:
return self.world_size is not None and self.world_size > 1
def set_seed(self, shared: bool = False):
initial_seed = self.initial_seed + self.current_task_iteration
if shared:
seed = initial_seed
np_seed = _generate_state(initial_seed, 0)
else:
seed = initial_seed + self.worker_id + self.global_rank + self.current_task_iteration
np_seed = _generate_state(initial_seed, self.worker_id + self.global_rank)
random.seed(seed)
torch.manual_seed(seed)
np.random.seed(np_seed)
def sample_task_indices(self):
self.set_seed(shared=True)
self.selected_indexes = np.random.choice(self.unique_labels, self.task_num_classes, replace=False)
self.selected_indexes.sort()
# subset of indices from the entire dataset where the labels are actually among the
# task_num_classes selected_indexes
self.task_indices = np.arange(len(self.dataset))[np.in1d(self.labels, self.selected_indexes)]
self.task_length = len(self.task_indices)
self.set_seed(shared=False)
@property
def worker_rank(self) -> int:
worker_id = self.worker_id
is_global_zero = self.global_rank == 0
return self.global_rank + worker_id + int(not is_global_zero)
def create_sampler(self):
data = range(self.task_length)
if self.world_size == 1 and self.num_workers in (0, 1):
if self.shuffle:
self.sampler = RandomSampler(data, generator=self.generator)
else:
self.sampler = SequentialSampler(data)
else:
num_workers = 1 if self.num_workers in (None, 0) else self.num_workers
num_replicas = num_workers * self.world_size
current_seed = self.initial_seed + self.current_task_iteration
self.sampler = DistributedSampler(
data, num_replicas=num_replicas, rank=self.worker_rank, shuffle=self.shuffle, seed=current_seed
)
def __iter__(self):
if self.generator is None:
self.generator = torch.Generator().manual_seed(self.initial_seed)
self.sample_task_indices()
self.create_sampler()
self.batch_sampler = BatchSampler(self.sampler, batch_size=self.batch_size, drop_last=self.drop_last)
self.iter_sampler = iter(self.batch_sampler)
self.is_first_batch = True
self.current_task_iteration += 1
return self
def increment_iteration(self):
self.current_task_iteration += 1
def __next__(self):
# this is optional, but useful to accumulate gradient over the entire task.
is_first_batch = self.is_first_batch if self.debugging else (self.is_first_batch and self.worker_id == 0)
if is_first_batch:
self.is_first_batch = False
return {"task_length": len(self.batch_sampler), "selected_indexes": self.selected_indexes}
random_indices = next(self.iter_sampler)
task_indices = [self.task_indices[idx] for idx in random_indices]
return default_collate([self.dataset[idx] for idx in task_indices])
class ClassificationDataset(Dataset):
def __init__(self, inputs, labels):
self.inputs = inputs
self.labels = labels
assert len(self.inputs) == len(self.labels)
def __getitem__(self, index):
return (self.inputs[index], self.labels[index])
def __len__(self):
return len(self.inputs)
def _test_fast_forward_sampler_with_distributed_sampler_and_iterative_dataset(rank, worldsize):
if worldsize > 1:
_setup_ddp(rank, worldsize)
def all_gather(tensor, world_size):
tensor_list = [torch.zeros_like(tensor, dtype=torch.int64) for _ in range(world_size)]
torch.distributed.all_gather(tensor_list, tensor)
return tensor_list
initial_seed = seed_everything(42)
generator = torch.Generator()
generator.manual_seed(initial_seed)
num_workers = 2
batch_size = 4
dataset_length = 60
num_classes = 10
labels = np.random.randint(0, num_classes, dataset_length)
dataset = ClassificationDataset(range(dataset_length), labels)
dataset = MetaLearningDataset(
dataset,
batch_size=batch_size,
drop_last=True,
num_workers=num_workers,
global_rank=rank,
world_size=worldsize,
initial_seed=initial_seed,
debugging=True,
shuffle=True,
)
dataset = CaptureIterableDataset(dataset)
dataloader = DataLoader(dataset, num_workers=num_workers, batch_size=1, generator=generator)
_add_capture_metadata_collate(dataloader)
epoch_results = []
for _ in range(2):
iter_dataloader = iter(dataloader)
batches = []
while True:
try:
batches.append(next(iter_dataloader))
except StopIteration:
break
epoch_results.append(batches)
dataloader.dataset.dataset.current_task_iteration += 1
assert len(epoch_results) == 2
assert len(epoch_results[0]) == math.ceil((dataset_length / (num_workers * worldsize)) / batch_size) + 2
if worldsize == 1:
assert epoch_results[0][0]["data"]["task_length"] == epoch_results[0][1]["data"]["task_length"]
assert torch.equal(
epoch_results[0][0]["data"]["selected_indexes"], epoch_results[0][1]["data"]["selected_indexes"]
)
assert 0 in epoch_results[0][2][AutoRestartBatchKeys.PL_RESTART_META]["iter_sampler"] # worker id 0
assert 1 in epoch_results[0][3][AutoRestartBatchKeys.PL_RESTART_META]["iter_sampler"] # worker id 1
assert not torch.equal(epoch_results[0][2]["data"][0], epoch_results[0][3]["data"][0])
else:
first_task_metadata = all_gather(epoch_results[0][0]["data"]["task_length"], worldsize)
second_task_metadata = all_gather(epoch_results[0][1]["data"]["task_length"], worldsize)
assert torch.equal(first_task_metadata[0], first_task_metadata[1])
assert torch.equal(second_task_metadata[0], second_task_metadata[1])
assert torch.equal(first_task_metadata[0], second_task_metadata[1])
first_batch_list = all_gather(epoch_results[0][2]["data"][0], worldsize)
assert not torch.equal(first_batch_list[0], first_batch_list[1])
second_batch_list = all_gather(epoch_results[0][3]["data"][0], worldsize)
assert not torch.equal(second_batch_list[0], second_batch_list[1])
# restarting on epoch 0 / real batch 2
state_dict = {"iter_sampler": {}}
for batch in epoch_results[0][2:4]:
batch, _state_dict = batch["data"], batch[AutoRestartBatchKeys.PL_RESTART_META]
for k, v in _state_dict.items():
state_dict[k].update(v)
dataset = ClassificationDataset(range(dataset_length), labels)
dataset = MetaLearningDataset(
dataset,
batch_size=batch_size,
drop_last=True,
num_workers=num_workers,
global_rank=rank,
world_size=worldsize,
initial_seed=initial_seed,
debugging=True,
shuffle=True,
)
dataset = CaptureIterableDataset(dataset)
dataset.load_state_dict(state_dict)
dataloader = DataLoader(dataset, num_workers=num_workers, batch_size=1, generator=generator)
_add_capture_metadata_collate(dataloader)
epoch_results_restart = []
for _ in range(2):
iter_dataloader = iter(dataloader)
batches = []
while True:
try:
batches.append(next(iter_dataloader))
except StopIteration:
break
epoch_results_restart.append(batches)
dataloader.dataset.dataset.increment_iteration()
dataloader.dataset.reset_on_epoch()
assert len(epoch_results_restart[0]) + 2 == len(epoch_results[0])
epoch_tensors = [e["data"][0] for e in epoch_results[0][4:]]
epoch_tensors_restart = [e["data"][0] for e in epoch_results_restart[0][2:]]
for t, tr in zip(epoch_tensors, epoch_tensors_restart):
assert torch.equal(t, tr)
epoch_tensors = [e["data"][0] for e in epoch_results[1][2:]]
epoch_tensors_restart = [e["data"][0] for e in epoch_results_restart[1][2:]]
for t, tr in zip(epoch_tensors, epoch_tensors_restart):
assert torch.equal(t, tr)
@pytest.mark.skipif(torch.cuda.is_available(), reason="This test takes around 45 sec and should be skipped in Azure CI")
def test_fast_forward_sampler_iterative_dataset():
_test_fast_forward_sampler_with_distributed_sampler_and_iterative_dataset(0, 1)
@pytest.mark.skipif(torch.cuda.is_available(), reason="This test takes around 55 sec and should be skipped in Azure CI")
@RunIf(skip_windows=True)
def test_fast_forward_sampler_with_distributed_sampler_and_iterative_dataset():
"""Make sure result logging works with DDP."""
tutils.set_random_main_port()
worldsize = 2
mp.spawn(
_test_fast_forward_sampler_with_distributed_sampler_and_iterative_dataset, args=(worldsize,), nprocs=worldsize
)
@mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "1"})
@RunIf(max_torch="1.7")
def test_fault_tolerant_not_supported():
assert not _fault_tolerant_training()
def create_iterable_dataset(batch_size, num_workers, attr_name="iter_sampler", wrap: bool = True):
dataset = RangeIterableDataset(range(50), num_workers=num_workers, batch_size=batch_size, attr_name=attr_name)
if wrap:
dataset = CaptureIterableDataset(dataset)
return dataset
def test_dataloader_to_state_dict_and_reload():
"""
Note: Those utilities are used only with DataLoader wrapping a ``mapping`` based dataset.
"""
def create_dataloader():
dataset = range(50)
batch_size = 8
sampler = FastForwardSampler(SequentialSampler(dataset))
sampler.setup(batch_size)
return DataLoader(dataset, sampler=sampler, batch_size=batch_size)
dataloader = create_dataloader()
iter_dataloader = iter(dataloader)
_ = next(iter_dataloader)
_ = next(iter_dataloader)
state_dict = _dataloader_to_state_dict(dataloader, iter_dataloader)
assert state_dict == {
"num_workers": 0,
"previous_worker": None,
0: {"current_iteration": 16},
}
dataloader = create_dataloader()
dataloader = _dataloader_load_state_dict(dataloader, state_dict)
iter_dataloader = iter(dataloader)
_ = next(iter_dataloader)
state_dict = _dataloader_to_state_dict(dataloader, iter_dataloader)
assert state_dict == {
"num_workers": 0,
"previous_worker": None,
0: {"current_iteration": 24},
}
@RunIf(min_torch="1.7.0")
@pytest.mark.parametrize("use_fault_tolerant", ["0", "1"])
def test_data_loading_wraps_dataset_and_samplers(use_fault_tolerant, tmpdir):
"""This test ensures the dataset and sampler are properly wrapped when fault tolerant is enabled."""
class CustomBatchSampler(BatchSampler):
pass
dataset = range(50)
class TestModel(BoringModel):
def train_dataloader(self):
return {
"a": [
DataLoader(create_iterable_dataset(3, 1, wrap=False), num_workers=0, batch_size=3),
DataLoader(dataset, batch_size=8),
DataLoader(
dataset,
batch_sampler=CustomBatchSampler(SequentialSampler(dataset), batch_size=8, drop_last=False),
),
],
"b": DataLoader(
create_iterable_dataset(2, num_workers=1, attr_name="custom_sampler", wrap=False),
num_workers=0,
batch_size=2,
),
}
def training_step(self, batch, batch_idx):
assert batch == {
"a": [ANY, ANY, ANY],
"b": ANY,
}
def validation_step(self, batch, batch_idx):
assert isinstance(batch, torch.Tensor)
validation_epoch_end = None
class Check(Callback):
def on_train_batch_start(self, trainer, *_) -> None:
loaders = trainer.train_dataloader.loaders
if use_fault_tolerant == "1":
assert isinstance(loaders["a"][0].loader.dataset, CaptureIterableDataset)
assert isinstance(loaders["a"][1].loader.sampler, FastForwardSampler)
assert isinstance(loaders["a"][1].loader.dataset, CaptureMapDataset)
assert isinstance(loaders["a"][2].loader.batch_sampler, FastForwardSampler)
assert isinstance(loaders["a"][2].loader.dataset, CaptureMapDataset)
assert isinstance(loaders["b"].loader.dataset, CaptureIterableDataset)
else:
assert isinstance(loaders["a"][0].loader.dataset, RangeIterableDataset)
assert isinstance(loaders["a"][1].loader.sampler, SequentialSampler)
assert not isinstance(loaders["a"][1].loader.dataset, CaptureMapDataset)
assert isinstance(loaders["a"][2].loader.batch_sampler, CustomBatchSampler)
assert not isinstance(loaders["a"][2].loader.dataset, CaptureMapDataset)
assert isinstance(loaders["b"].loader.dataset, RangeIterableDataset)
with mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": use_fault_tolerant}):
model = TestModel()
model.training_epoch_end = None
trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, limit_train_batches=1, callbacks=Check())
trainer.fit(model)
class SequentialGetItemDataset(Dataset):
def __init__(self, length, *_):
self.len = length
def __getitem__(self, index):
return torch.tensor([index]).float()
def __len__(self):
return self.len
class RandomGetItemDataset(Dataset):
"""A dataset with random elements generated using global rng from torch, numpy and python."""
def __init__(self, length, size):
self.size = size
self.len = length
def __getitem__(self, index):
t = torch.rand(self.size)
n = torch.from_numpy(np.random.rand(self.size))
p = torch.tensor([python_random.random() for _ in range(self.size)])
sample = (index + (t + n + p) / 10).float()
return sample
def __len__(self):
return self.len
# TODO: test with `RandomGeneratorGetItemDataset`
@mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "1"})
@RunIf(min_torch="1.7.0")
@pytest.mark.parametrize(
"dataset_class",
[
SequentialGetItemDataset,
RandomGetItemDataset,
# RandomGeneratorGetItemDataset,
],
)
@pytest.mark.parametrize("num_workers", [0])
@pytest.mark.parametrize("batch_size", [1, 2, 3])
def test_dataset_rng_states_restart(dataset_class, num_workers, batch_size):
"""Test that the sequence of batches coming from a random number generator continues with the correct sequence
after reloading the state."""
def create_dataset_sampler():
dset = CaptureMapDataset(dataset_class(16, 8))
random_sampler = RandomSampler(dset, generator=torch.Generator())
return dset, random_sampler
def create_dataloader_sampler(dset, sampler):
sampler = FastForwardSampler(sampler)
sampler.setup(batch_size)
dl = DataLoader(dset, num_workers=num_workers, sampler=sampler, batch_size=batch_size)
_add_capture_metadata_collate(dl)
return dl, sampler
def fetch(fetcher, prefetch_iter, num_batches_fetched):
batch, _ = next(prefetch_iter)
state: List[MergedIteratorState] = fetcher.state
assert len(state) == 1
assert isinstance(state[0], MergedIteratorState)
assert len(fetcher.dataloader_iter.cache_states) == 1
if num_workers == 0:
assert state[0].state[0].num_batches_fetched == num_batches_fetched
return state
dataset, random_sampler = create_dataset_sampler()
dataloader, ff_sampler = create_dataloader_sampler(dataset, random_sampler)
fetcher = DataFetcher()
fetcher.setup(dataloader)
prefetch_iter = iter(fetcher)
# fetch 4 batches
fetch(fetcher, prefetch_iter, 1)
fetch(fetcher, prefetch_iter, 2)
fetch(fetcher, prefetch_iter, 3)
# (A) capture the state after fetching 4 batches
state = fetch(fetcher, prefetch_iter, 4)
state = deepcopy(state[0])
# (B) simulate 2 additional batches
batch05, _ = next(prefetch_iter)
batch06, _ = next(prefetch_iter)
# start reloading
dataset, random_sampler = create_dataset_sampler()
dataloader, ff_sampler = create_dataloader_sampler(dataset, random_sampler)
# load the state dict saved at (A)
ff_sampler.load_state_dict(state.sampler_states)
dataset.load_state_dict(state.dataset_states, latest_worker_id=state.latest_worker_id, num_workers=num_workers)
prefetcher = DataFetcher()
prefetcher.setup(dataloader)
prefetch_iter = iter(prefetcher)
# fetch 2 random batches, these should match exactly the batches seen at (B)
batch05_restart, _ = next(prefetch_iter)
batch06_restart, _ = next(prefetch_iter)
assert torch.equal(batch05, batch05_restart)
assert torch.equal(batch06, batch06_restart)
class CustomException(Exception):
pass
class SequentialIterableDataset(IterableDataset):
def __init__(self, length, *_):
self.len = length
self.sampler = SequentialSampler(range(self.len))
def __iter__(self):
self.sampler_iter = iter(self.sampler)
return self
def __next__(self):
indices = next(self.sampler_iter)
return torch.tensor([indices]).float()
class SequentialDictIterableDataset(SequentialIterableDataset):
def __next__(self):
indices = next(self.sampler_iter)
return {"data": torch.tensor([indices]).float()}
class TestModel(LightningModule):
def __init__(self, fail_on_step: int = -1):
super().__init__()
self.layer = torch.nn.Linear(1, 2)
self.seen_batches = []
self.fail_on_step = fail_on_step
def training_step(self, batch, batch_idx):
if self.global_step == self.fail_on_step:
raise CustomException()
batch = batch["data"] if isinstance(batch, dict) else batch
self.seen_batches.append(torch.stack(batch) if isinstance(batch, list) else batch)
loss = sum(self.layer(b).sum() for b in batch)
return loss
def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)
def _run_training(trainer_kwargs, dataset_classes, fail_on_step: int = -1, ckpt_path=None):
seed_everything(1)
train_dataloader = [
DataLoader(dataset_class(3, 1), batch_size=1, num_workers=0) for dataset_class in dataset_classes
]
train_dataloader = train_dataloader[0] if len(train_dataloader) == 1 else train_dataloader
model = TestModel(fail_on_step=fail_on_step)
trainer = Trainer(**trainer_kwargs)
with suppress(CustomException):
trainer.fit(model, train_dataloader=train_dataloader, ckpt_path=ckpt_path)
return model.seen_batches, model.parameters()
@mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "1"})
@RunIf(min_torch="1.7.0")
@pytest.mark.parametrize(
"dataset_classes",
[
# single training dataset
[RandomGetItemDataset],
[SequentialIterableDataset],
[SequentialDictIterableDataset],
# multiple training datasets (combinded dataloader)
[SequentialGetItemDataset, SequentialIterableDataset],
[SequentialIterableDataset, SequentialIterableDataset],
# [RandomGetItemDataset, RandomGetItemDataset], # TODO: support in the future
],
)
@pytest.mark.parametrize("multiple_trainloader_mode", ["min_size", "max_size_cycle"])
def test_dataset_rng_states_restart_with_lightning(tmpdir, dataset_classes, multiple_trainloader_mode):
"""Test that the Trainer can resume from a failed run in the case of several types of datasets."""
trainer_kwargs = dict(
default_root_dir=tmpdir,
max_epochs=3,
enable_progress_bar=False,
enable_model_summary=False,
multiple_trainloader_mode=multiple_trainloader_mode,
)
all_batches, weights0 = _run_training(trainer_kwargs, dataset_classes)
all_batches = torch.stack(all_batches)
assert len(all_batches) == 9
# Simulate 1st failure
complete_batches, _ = _run_training(trainer_kwargs, dataset_classes, fail_on_step=4)
assert len(complete_batches) == 4
checkpoint_path = os.path.join(tmpdir, ".pl_auto_save.ckpt")
assert os.path.exists(checkpoint_path)
# Resume after failure
resumed_batches, weights1 = _run_training(
trainer_kwargs, dataset_classes, fail_on_step=-1, ckpt_path=checkpoint_path
)
assert len(resumed_batches) == 5
# the resumed batches should match the batches of the successful training
all_batches_resumed = torch.stack(complete_batches + resumed_batches)
assert len(all_batches_resumed) == 9
assert torch.equal(all_batches, all_batches_resumed)
# the final weights of a resumed training should equal the weights of an uninterrupted training
for w0, w1 in zip(weights0, weights1):
assert w0 is not w1
assert torch.allclose(w0, w1)
@mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "1"})
@RunIf(min_torch="1.7.0")
@pytest.mark.parametrize(
["train_datasets", "val_datasets"],
[
([RandomGetItemDataset], [RandomGetItemDataset]),
([RandomGetItemDataset], [RandomGetItemDataset, RandomGetItemDataset]),
],
)
@pytest.mark.parametrize(
"val_check_interval",
[
pytest.param(
0.5,
marks=pytest.mark.xfail(
reason=(
"TODO: the `train_dataloader` random state overrides the validation state when restarting training"
)
),
),
1.0,
],
)
def test_auto_restart_within_validation_loop(train_datasets, val_datasets, val_check_interval, tmpdir):
n_val_dataloaders = len(val_datasets)
stop_dataloader = n_val_dataloaders - 1
stop_batch = 1
class ValidationLoopTestModel(LightningModule):
def __init__(self, should_fail):
super().__init__()
self.layer = torch.nn.Linear(1, 2)
self.should_fail = should_fail
self.training_batches = []
self.validation_batches = defaultdict(list)
def step(self, batch):
return sum(self.layer(b).sum() for b in batch)
def training_step(self, batch, batch_idx):
self.training_batches.append(batch)
return self.step(batch)
def validation_step(self, batch, batch_idx, dataloader_idx=0):
if self.should_fail and stop_dataloader == dataloader_idx and batch_idx == stop_batch:
raise CustomException
self.validation_batches[dataloader_idx].append(batch)
return self.step(batch)
def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)
def train_dataloader(self):
return [DataLoader(cls(4, 1)) for cls in train_datasets]
def val_dataloader(self):
return [DataLoader(cls(4, 1)) for cls in val_datasets]
def run(should_fail, resume):
if not resume:
seed_everything(42)
model = ValidationLoopTestModel(should_fail)
ckpt_path = str(tmpdir / ".pl_auto_save.ckpt") if resume else None
trainer = Trainer(
default_root_dir=tmpdir,
max_epochs=1,
val_check_interval=val_check_interval,
num_sanity_val_steps=0,
)
if should_fail:
with pytest.raises(CustomException):
trainer.fit(model, ckpt_path=ckpt_path)
else:
trainer.fit(model, ckpt_path=ckpt_path)
return model.training_batches, model.validation_batches
total_train_batches, total_val_batches = run(should_fail=False, resume=False)
pre_fail_train_batches, pre_fail_val_batches = run(should_fail=True, resume=False)
post_fail_train_batches, post_fail_val_batches = run(should_fail=False, resume=True)
torch.testing.assert_allclose(total_train_batches, pre_fail_train_batches + post_fail_train_batches)
for k in total_val_batches:
torch.testing.assert_allclose(total_val_batches[k], pre_fail_val_batches[k] + post_fail_val_batches[k])
class TestAutoRestartModelUnderSignal(BoringModel):
def __init__(self, should_signal: bool, failure_on_step: bool, failure_on_training: bool, on_last_batch: bool):
super().__init__()
self.should_signal = should_signal
self.failure_on_step = failure_on_step
self.failure_on_training = failure_on_training
self.on_last_batch = on_last_batch
self.seen_train_batches = []
def _signal(self):
if self.should_signal:
# simulate `os.kill(os.getpid(), signal.SIGUSR1)`
self.trainer._terminate_gracefully = True
def training_step(self, batch, batch_idx):
self.seen_train_batches.append(batch)
should_signal = self.trainer.fit_loop.epoch_loop._is_training_done if self.on_last_batch else batch_idx == 2
if self.failure_on_step and self.failure_on_training and should_signal:
self._signal()
return super().training_step(batch, batch_idx)
def validation_step(self, batch, batch_idx):
should_signal = (
self.trainer.fit_loop.epoch_loop.val_loop.epoch_loop.batch_progress.is_last_batch
if self.on_last_batch
else batch_idx == 2
)
if self.failure_on_step and not self.failure_on_training and should_signal:
self._signal()
return super().validation_step(batch, batch_idx)
def training_epoch_end(self, outputs) -> None:
if not self.failure_on_step and self.failure_on_training:
self._signal()
def validation_epoch_end(self, outputs) -> None:
if not self.failure_on_step and not self.failure_on_training:
self._signal()
def train_dataloader(self):
return DataLoader(RandomDataset(32, 4))
def val_dataloader(self):
return DataLoader(RandomDataset(32, 4))
def _fit_model(
tmpdir, should_signal, val_check_interval, failure_on_step, failure_on_training, on_last_batch, status=None
):
seed_everything(42)
model = TestAutoRestartModelUnderSignal(should_signal, failure_on_step, failure_on_training, on_last_batch)
trainer_kwargs = dict(
default_root_dir=tmpdir,
max_epochs=1,
limit_train_batches=4,
limit_val_batches=4,
val_check_interval=val_check_interval,
num_sanity_val_steps=0,
)
trainer = Trainer(**trainer_kwargs)
if should_signal:
with pytest.raises(ExitGracefullyException, match=status):
trainer.fit(model)
else:
trainer.fit(model)
assert trainer._terminate_gracefully == should_signal
return model
@pytest.mark.parametrize("on_last_batch", [False, True])
@pytest.mark.parametrize("val_check_interval", [0.5, 1.0])
@pytest.mark.parametrize("failure_on_training", [False, True])
@pytest.mark.parametrize("failure_on_step", [False, True])
@mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "1"})
@RunIf(min_torch="1.7.0", skip_windows=True)
def test_auto_restart_under_signal(on_last_batch, val_check_interval, failure_on_training, failure_on_step, tmpdir):
"""This test asserts that if a signal is being sent during the training / validation phase, the model should
restart in a reproducible way."""
model_total = _fit_model(tmpdir, False, val_check_interval, failure_on_step, failure_on_training, on_last_batch)
if failure_on_step:
if on_last_batch:
if failure_on_training:
# Breaking on first validation batch.
# This is done to capture the random state of the validation dataloader.
status = "EvaluationEpochLoop:advance"
else:
# when breaking on last batch of validation, we should exist on `run_end` val_check_interval == 1.0
status = (
"TrainingEpochLoop:on_run_end" if val_check_interval == 1.0 else "TrainingEpochLoop:on_advance_end"
)
else:
status = "TrainingEpochLoop:on_advance_end" if failure_on_training else "EvaluationEpochLoop:advance"
else:
if val_check_interval == 1.0:
status = "TrainingEpochLoop:on_run_end"
else:
# `training_epoch_end` happens after `validation_epoch_end` since Lightning v1.4
status = "TrainingEpochLoop:on_run_end" if failure_on_training else "TrainingEpochLoop:on_advance_end"
model_signaled = _fit_model(
tmpdir, True, val_check_interval, failure_on_step, failure_on_training, on_last_batch, status=status
)
checkpoint_path = str(tmpdir / ".pl_auto_save.ckpt")
assert os.path.exists(checkpoint_path)
model_restarted = _fit_model(tmpdir, False, val_check_interval, failure_on_step, failure_on_training, on_last_batch)
# check the batches
actual = torch.cat(model_signaled.seen_train_batches + model_restarted.seen_train_batches)
expected = torch.cat(model_total.seen_train_batches)
assert torch.equal(actual, expected)
# FIXME: why `on_last_batch` doesn't work ?
if failure_on_step and failure_on_training and not on_last_batch:
assert not torch.equal(model_total.layer.weight, model_signaled.layer.weight)
assert torch.equal(model_restarted.layer.weight, model_total.layer.weight)
checkpoint = torch.load(checkpoint_path)["loops"]["fit_loop"]
p = checkpoint["epoch_loop.batch_progress"]
if p["is_last_batch"] and p["current"]["completed"] == 4:
assert "dataloader_state_dict" not in checkpoint["epoch_loop.state_dict"]
else:
assert "dataloader_state_dict" in checkpoint["epoch_loop.state_dict"]
state_dict = checkpoint["epoch_loop.val_loop.epoch_loop.state_dict"]
p = checkpoint["epoch_loop.val_loop.epoch_loop.batch_progress"]
if (p["is_last_batch"] and p["current"]["completed"] == 4) or p["current"]["ready"] == 0:
assert "dataloader_state_dict" not in state_dict
else:
assert "dataloader_state_dict" in state_dict
Reference in New Issue View Git Blame Copy Permalink