-
Notifications
You must be signed in to change notification settings - Fork 3.3k
/
auto_restart.py
573 lines (454 loc) · 23 KB
/
auto_restart.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
# 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.
from copy import deepcopy
from dataclasses import dataclass, field
from functools import partial, wraps
from random import getstate as python_get_rng_state
from random import setstate as python_set_rng_state
from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, Tuple, Union
import numpy as np
import torch
from torch.utils.data import Dataset, get_worker_info, Sampler
from torch.utils.data.dataloader import _MultiProcessingDataLoaderIter, DataLoader, IterableDataset
import pytorch_lightning as pl
from pytorch_lightning.utilities.enums import AutoRestartBatchKeys
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.imports import _fault_tolerant_training
class FastForwardSampler(Sampler):
"""This FastForwardSampler wraps a :class:`torch.utils.data.Sampler` and records the number of iterations
performed during an epoch.
It maintains a state, saved with :meth:`state_dict`, that can be reloaded with
:meth:`load_state_dict`. If the sampler is used in a multiprocessing context, the ``FastForwardSampler`` will record
the state of the current worker.
When reloading, the ``FastForwardSampler`` will "fast-forward" the wrapped sampler by iterating through all the
samples seen in the last iterations (for the current worker).
"""
def __init__(self, sampler: Union[Sampler, Generator], attr_name: Optional[str] = None) -> None:
super().__init__(data_source=None)
self._sampler = sampler
self.restarting: bool = False
self._current_iteration = 0
self._counter = 0
self._dataloader_batch_size: Optional[int] = None
self._cached_state_dict: Optional[Dict[int, Any]] = None
self._attr_name = attr_name
def __getattr__(self, key: str) -> Any:
if key in self.__dict__:
return self.__dict__[key]
return getattr(self._sampler, key, None)
def setup(self, dataloader_batch_size: Optional[int] = None) -> None:
"""Setup the ``FastForwardSampler``.
This is required only when the provided dataset subclassed
:class:`torch.utils.data.Dataset`.
"""
self._dataloader_batch_size = dataloader_batch_size
@property
def worker_id(self) -> int:
worker_info = get_worker_info()
return worker_info.id if worker_info else 0
def __iter__(self) -> Iterator[Any]:
self.sampler_iter = iter(self._sampler)
self._current_iteration = 0
self._counter = 0
return self
def __next__(self):
# the `state dict` was cached as workers were unavailable before.
if self._cached_state_dict is not None:
self._load_non_random_state(self._cached_state_dict)
while self._counter < self._current_iteration:
next(self.sampler_iter)
self._counter += 1
# here: i == self._current_iteration
if self._cached_state_dict is not None:
self._cached_state_dict = None
# recreate iterator to be sure loading is reflected there as well
self._current_iteration += 1
self._counter += 1
has_raised = False
try:
return next(self.sampler_iter)
except StopIteration:
has_raised = True
self._current_iteration = 0
self._counter = 0
self._cached_state_dict = None
self.restarting = False
if has_raised:
raise StopIteration
def __len__(self) -> int:
return len(self._sampler)
def state_dict(self, num_batches_processed: Optional[int] = None) -> Dict[int, Dict[str, int]]:
"""Returns the state of the sampler in the current worker.
The worker id indexes the state dict.
"""
return {self.worker_id: {"current_iteration": self._compute_current_iteration(num_batches_processed)}}
def load_state_dict(self, state_dict: Dict[int, Any]) -> None:
"""Loads the saved state for the wrapped sampler.
If the ``state_dict`` contains multiple states, it means there were multiple workers. The state will be cached
and fully reloaded (fast-forward) the first time :meth:`__iter__` is called.
"""
# as workers aren't available, the ``state_dict``` is cached until workers are made available.
state_dict = deepcopy(state_dict)
self._cached_state_dict = state_dict
self.restarting = True
def _compute_current_iteration(self, num_batches_processed: Optional[int] = None) -> int:
"""This function is used to compute the effective iteration.
As DataLoader can perform ``prefecthing`` or training can fail while processing a batch, the current iteration
needs to be computed using the ``num_batches_processed`` processed information.
"""
if num_batches_processed is not None:
current_iteration = num_batches_processed
else:
current_iteration = self._current_iteration
if self._dataloader_batch_size and num_batches_processed is not None:
current_iteration *= self._dataloader_batch_size
return current_iteration
def _load_non_random_state(self, state_dict: Dict[int, Dict[str, Any]]) -> None:
self._current_iteration = state_dict[self.worker_id]["current_iteration"]
@dataclass(frozen=True, unsafe_hash=True)
class IteratorState:
"""The state of an iterator in a single worker process."""
dataset_state: Dict[int, Any] = field(default_factory=dict)
sampler_state: Dict[int, Any] = field(default_factory=dict)
worker_id: int = 0
num_workers: int = 0
num_batches_fetched: int = 0
name: Optional[str] = None
@classmethod
def from_state_dict(cls, state_dict) -> "IteratorState":
return cls(**state_dict)
@dataclass
class MergedIteratorState:
"""This class is used to hold the current iterator state and lives on the iterator.
It holds the current merged states from all worker processes. Once an iterator advances, it can store updates of the
worker states in this merged iterator state.
"""
state: Union[Dict[Union[int, str], Union[Dict[str, IteratorState], IteratorState]]] = field(default_factory=dict)
latest_worker_id: int = 0
represent_map_dataset: Optional[bool] = None
def update(self, generator_name: Optional[str], new_state: IteratorState) -> None:
# a map based dataset doesn't own a generator and therefore `generator_name` should be None.
self.represent_map_dataset = generator_name is None
if self.represent_map_dataset:
state = self.state
else:
if generator_name not in self.state:
self.state[generator_name] = {}
state = self.state[generator_name]
latest_worker_id = new_state.worker_id
state[latest_worker_id] = new_state
self.latest_worker_id = latest_worker_id
@property
def sampler_states(self) -> Dict[int, Any]:
"""Returns the merged sampler states for all worker processes."""
return {0: self.state[k].sampler_state[0] for k in self.state.keys()}
@property
def dataset_states(self) -> Dict[int, Any]:
"""Returns the merged dataset states for all worker processes."""
return {k: self.state[k].dataset_state[k] for k in self.state.keys()}
@classmethod
def from_state_dict(cls, state_dict) -> "MergedIteratorState":
if state_dict["represent_map_dataset"]:
state_dict["state"] = {
worker_id: IteratorState.from_state_dict(state) for worker_id, state in state_dict["state"].items()
}
else:
state_dict["state"] = {
sampler_name: {
worker_id: IteratorState.from_state_dict(state) for worker_id, state in worker_state.items()
}
for sampler_name, worker_state in state_dict["state"].items()
}
return cls(**state_dict)
def __len__(self) -> int:
return len(self.state)
class CaptureMapDataset(Dataset):
"""This class is used to capture the state from the map-based state dataset.
Note:
We currently don't support restoring if we fail during the first `N = num_workers` batches, where
`num_workers` is the number of workers spawned by the dataloader.
"""
def __init__(self, dataset: Dataset) -> None:
self.dataset = dataset
self._cached_state_dict = None
@property
def worker_id(self) -> int:
worker_info = get_worker_info()
return worker_info.id if worker_info else 0
def __getitem__(self, item) -> Tuple[Any, Dict[int, Dict]]:
if self._cached_state_dict is not None:
if self.worker_id in self._cached_state_dict:
set_rng_states(self._cached_state_dict[self.worker_id]["rng_states"])
self._cached_state_dict = None
return self.dataset[item]
def __len__(self) -> int:
return len(self.dataset)
def load_state_dict(self, state_dict: Dict[int, Any], latest_worker_id: int, num_workers: int) -> None:
# as workers aren't available, the ``state_dict``` is cached until workers are made available.
state_dict = deepcopy(state_dict)
if num_workers > 0:
# remap states to worker ids starting at 0
next_worker_id = latest_worker_id + 1
old_to_new_worker_id_map = [((next_worker_id + i) % num_workers, i) for i in range(num_workers)]
state_dict = {
new_id: state_dict[old_id] for old_id, new_id in old_to_new_worker_id_map if old_id in state_dict
}
self._cached_state_dict = state_dict
def state_dict(self) -> Dict[int, Dict[str, Any]]:
return {self.worker_id: {"rng_states": collect_rng_states()}}
def collect_rng_states() -> Dict[str, Any]:
"""Collect the global random state of :mod:`torch`, :mod:`numpy` and Python."""
return {"torch": torch.get_rng_state(), "numpy": np.random.get_state(), "python": python_get_rng_state()}
def set_rng_states(rng_state_dict: Dict[str, Any]) -> None:
"""Set the global random state of :mod:`torch`, :mod:`numpy` and Python in the current process."""
torch.set_rng_state(rng_state_dict.get("torch"))
np.random.set_state(rng_state_dict.get("numpy"))
version, state, gauss = rng_state_dict.get("python")
python_set_rng_state((version, tuple(state), gauss))
class CaptureIterableDataset(IterableDataset):
"""The ``CaptureIterableDataset`` is used to wrap an :class:`torch.utils.data.IterableDataset`.
On ``__iter__`` function call, the ``CaptureIterableDataset`` will wrap the wrapped dataset generators into
``FastForwardSampler`` to keep track of progress. On ``__next__`` function call, the ``CaptureIterableDataset`` will
return a dictionary containing user data and metadata containing the ``FastForwardSampler`` samplers state_dict.
"""
def __init__(self, dataset: IterableDataset) -> None:
super().__init__()
self.dataset = deepcopy(dataset)
self.samplers: Optional[Dict[str, FastForwardSampler]] = None
self._state_dict: Optional[Dict[int, Any]] = None
self._has_wrapped: bool = False
@property
def sampler(self) -> Sampler:
return self.dataset.sampler
def state_dict(self) -> Dict[str, Any]:
return {k: v.state_dict() for k, v in self.samplers.items()}
def load_state_dict(self, state_dict: Dict[int, Any]) -> None:
self._state_dict = deepcopy(state_dict)
def _wrap_generator_samplers(self) -> None:
self.samplers = {}
# access wrapped dataset attributes
dataset_dict = self.dataset.__dict__
# create a dictionary of generator present within the dataset attributes
dataset_sampler_generators = {k: v for k, v in dataset_dict.items() if isinstance(v, (Generator, Iterator))}
# iterate over the generator. If a generator was created from a `Sampler`,
# it will be wrapped into a `FastForwardSampler`.
for (generator_attr_name, generator) in dataset_sampler_generators.items():
if isinstance(generator, Sampler):
continue
# wrap the generator into a `FastForwardSampler`
sampler = FastForwardSampler(generator, attr_name=generator_attr_name)
# if `CaptureIterableDataset` was available, the sampler should reload its own state.
if self._state_dict is not None:
sampler.load_state_dict(self._state_dict[generator_attr_name])
# store the samplers
self.samplers[generator_attr_name] = sampler
# replace generator with the generator from the `FastForwardSampler`.
dataset_dict[generator_attr_name] = iter(sampler)
self.reset_on_epoch()
def reset_on_epoch(self):
self._state_dict = None
def __iter__(self) -> Iterator:
# create a generator from the wrapped Iterative Dataset
# if the dataset contained samplers, they will be transformed into generators
self.iter_data = iter(self.dataset)
# wrap any generator associated to a Sampler into a `FastForwardSampler`.
if isinstance(self.iter_data, Generator):
raise MisconfigurationException(
"PyTorch Lightning Fault-Tolerant feature does not support `__iter__` returning a generator."
" Please use the `__next__` function to fetch the next batch and use a sampler for"
" doing your iterations."
)
self._wrap_generator_samplers()
return self
def __next__(self) -> Dict[str, Any]:
return next(self.iter_data)
def _find_fast_forward_samplers(dataloader: DataLoader) -> Optional[FastForwardSampler]:
"""If the ``DataLoader`` is wrapping a mapping based Dataset, return the ``FastForwardSampler``."""
if isinstance(dataloader.sampler, FastForwardSampler):
return dataloader.sampler
if isinstance(dataloader.batch_sampler, FastForwardSampler):
return dataloader.batch_sampler
def _cycle_to_next_worker_and_reset(dataloader: DataLoader, state_dict: Dict[str, Any]) -> Iterator:
"""This function is used to cycle back the DataLoader ``_MultiProcessingDataLoaderIter`` workers and call the
reset function.
Returns:
iterator: Return the iterator generated from the provided ``DataLoader``.
"""
# create iterator from dataloader
iter_dataloader = iter(dataloader)
# get current num workers
num_workers = getattr(iter_dataloader, "_num_workers", 0)
# as `state_dict` are workers dependent, Lightning doesn't support changing
# the `num_workers` for fault tolerant training
if state_dict["num_workers"] != num_workers:
raise MisconfigurationException(
f"The provided `num_workers` {num_workers} doesn't match the one used "
f"while generating the checkpoint: {state_dict['num_workers']}"
)
# when using multiple workers, we will cycle back the worker queue idx to
# start back on the failed worker.
if isinstance(iter_dataloader, _MultiProcessingDataLoaderIter):
# move back to 0
while next(iter_dataloader._worker_queue_idx_cycle) != 0:
pass
# increment previous worker
if isinstance(state_dict["previous_worker"], int):
for _ in range(state_dict["previous_worker"] - 1):
next(iter_dataloader._worker_queue_idx_cycle)
# we can finally call reset and apply prefecthing.
iter_dataloader._reset = iter_dataloader._original_reset
iter_dataloader._reset(dataloader, first_iter=True)
# return the iterator
return iter_dataloader
def _dataloader_to_state_dict(
dataloader: DataLoader, iterator: Iterator, num_batches_processed: int = None
) -> List[Dict[str, Any]]:
"""Convert a dataloader to its associated state dict."""
out = {}
if iterator is not None:
out.update(_find_current_worker(iterator))
if not isinstance(dataloader.dataset, CaptureIterableDataset):
fast_forward_sampler = _find_fast_forward_samplers(dataloader)
if fast_forward_sampler is not None:
out.update(fast_forward_sampler.state_dict(num_batches_processed=num_batches_processed))
return out
def _dataloader_load_state_dict(dataloader: DataLoader, state_dict: List[Dict[str, Any]]) -> DataLoader:
"""Reload ``DataLoader`` fast-forward sampler state dict."""
fast_forward_sampler = _find_fast_forward_samplers(dataloader)
if isinstance(fast_forward_sampler, Sampler):
state_dict = {k: v for k, v in state_dict.items() if k not in ("num_workers", "previous_worker")}
fast_forward_sampler.load_state_dict(state_dict)
return dataloader
def _find_current_worker(iterator: Iterator) -> Dict[str, Optional[int]]:
"""Find the current DataLoader Iterator worker if multiple workers were used."""
# get the current number of workers
num_workers = getattr(iterator, "_num_workers", 0)
if isinstance(iterator, _MultiProcessingDataLoaderIter):
# fetch next worker
next_worker = (next(iterator._worker_queue_idx_cycle)) % num_workers
# get the current worker from next one
previous_worker = (next_worker - 1) % num_workers
# reset back the `worker_queue_idx` to current one, so we can keep
# going without perturbation.
while next(iterator._worker_queue_idx_cycle) != previous_worker:
pass
else:
previous_worker = None
# return the captured metadata.
return {"num_workers": num_workers, "previous_worker": previous_worker}
def _capture_metadata_collate(samples: List, dataset: Dataset, default_collate: Callable) -> Dict:
"""A collate function that adds the state dict of a :class:`CaptureIterableDataset` or
:class:`CaptureMapDataset` used in the worker processes. This function gets executed within the worker
processes. The structure will be:
.. code-block:: python
{
"data": ..., # data returned by Dataset
"__pl_restart_meta": {"sampler_name0": state_dict0, "sampler_name1": state_dict1},
}
"""
data = default_collate(samples)
if not isinstance(dataset, (CaptureIterableDataset, CaptureMapDataset)):
return data
metadata = dataset.state_dict()
return {"data": data, AutoRestartBatchKeys.PL_RESTART_META: metadata}
def patch_dataloader_iterator(
dataloader: DataLoader,
iterator: Iterator,
data_fetcher: "pl.utilities.fetching.DataFetcher",
num_batches_fetched: int = 0,
) -> None:
"""Patches the iterator of a PyTorch dataloader by injecting logic for fault-tolerant training when it is
necessary to remove the sampler state dict from provided data batch.
The custom data has this format:
.. code-block:: python
{
"batch": ..., # data returned by DataLoader
"__pl_restart_meta": {
"sampler0": {
0: {"current_iteration": ...},
1: {"current_iteration": ...},
},
"sampler1": ...,
},
}
Each sampler in the worker process tracks the current iteration. We return all of them to the main process
as part of the sample and then a special collate function :func:`_capture_metadata_collate`
will extract the current iteration as part of the metadata returned by a custom batch.
"""
assert isinstance(dataloader.dataset, (CaptureMapDataset, CaptureIterableDataset))
def _next_data_wrapper(fn, it, dl, num_batches_fetched) -> Callable:
@wraps(fn)
def wrapper():
nonlocal num_batches_fetched
nonlocal it
nonlocal dl
dataset = dl.dataset
combined_batch = fn()
batch, state = combined_batch["data"], combined_batch[AutoRestartBatchKeys.PL_RESTART_META]
num_batches_fetched += 1
if isinstance(dataset, CaptureIterableDataset):
state = [
IteratorState(
num_workers=dataloader.num_workers,
sampler_state=iterator_state,
num_batches_fetched=num_batches_fetched,
worker_id=list(iterator_state.keys())[0],
name=sampler_iter_name,
)
for sampler_iter_name, iterator_state in state.items()
]
elif isinstance(dataset, CaptureMapDataset):
ff_sampler = _find_fast_forward_samplers(dl)
state = [
IteratorState(
num_workers=dataloader.num_workers,
sampler_state=ff_sampler.state_dict(num_batches_fetched),
dataset_state=state,
worker_id=list(state.keys())[0],
num_batches_fetched=num_batches_fetched,
)
]
data_fetcher._store_dataloader_iter_state(it, state)
return batch
return wrapper
iterator._next_data = _next_data_wrapper(iterator._next_data, iterator, dataloader, num_batches_fetched)
def _add_capture_metadata_collate(dataloader: DataLoader) -> None:
"""Wrap default collate function to retrive captured dataset state dict when fault tolerant is enabled."""
dataloader.collate_fn = partial(
_capture_metadata_collate, dataset=dataloader.dataset, default_collate=dataloader.collate_fn
)
def reload_dataloader_state_dict(dataloader: DataLoader, state_dict: Dict[str, Any]) -> None:
"""Utility to reload state_dict within dataloader for fault tolerance."""
if not _fault_tolerant_training():
return
dataset = dataloader.dataset
if isinstance(dataset, CaptureMapDataset):
iterator_state = state_dict["state"][0]
if not isinstance(iterator_state, IteratorState):
iterator_state = IteratorState.from_state_dict(iterator_state)
# reload sampler state
ff_sampler = _find_fast_forward_samplers(dataloader)
ff_sampler.load_state_dict(iterator_state.sampler_state)
# reload dataset state
dataset.load_state_dict(
iterator_state.dataset_state,
latest_worker_id=state_dict["latest_worker_id"],
num_workers=iterator_state.num_workers,
)
elif isinstance(dataset, CaptureIterableDataset):
dataset.load_state_dict(
{sampler_name: state[0]["sampler_state"] for sampler_name, state in state_dict["state"].items()}
)
else:
raise MisconfigurationException("This shouldn't happen. Please, open an issue on PyTorch Lightning Github.")