Fault Tolerant Manual: Add support for DDP

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+
 - Added a flag `SLURMEnvironment(auto_requeue=True|False)` to control whether Lightning handles the requeuing ([#10601](https://github.com/PyTorchLightning/pytorch-lightning/issues/10601))
 
 
@@ -21,6 +22,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Add logic to reload the states across data loading components ([#10699](https://github.com/PyTorchLightning/pytorch-lightning/issues/10699))
     * Cleanup some fault tolerant utilities ([#10703](https://github.com/PyTorchLightning/pytorch-lightning/issues/10703))
     * Enable Fault Tolerant Manual Training ([#10707](https://github.com/PyTorchLightning/pytorch-lightning/issues/10707))
+    * Broadcast the `_terminate_gracefully` to all processes and add support for DDP ([#10638](https://github.com/PyTorchLightning/pytorch-lightning/issues/10638))
+
 
 - Added support for re-instantiation of custom (subclasses of) `DataLoaders` returned in the `*_dataloader()` methods, i.e., automatic replacement of samplers now works with custom types of `DataLoader` ([#10680](https://github.com/PyTorchLightning/pytorch-lightning/issues/10639))
 

pytorch_lightning/loops/epoch/evaluation_epoch_loop.py

@@ -17,13 +17,20 @@
 from functools import lru_cache
 from typing import Any, Dict, Iterator, Optional, Union
 
+import torch
 from deprecate import void
 
 from pytorch_lightning.loops.base import Loop
 from pytorch_lightning.loops.utilities import _update_dataloader_iter
 from pytorch_lightning.trainer.progress import BatchProgress
-from pytorch_lightning.utilities.auto_restart import _reload_dataloader_state_dict, MergedIteratorState
+from pytorch_lightning.utilities.auto_restart import (
+    _collect_states_on_rank_zero_over_collection,
+    _reload_dataloader_state_dict,
+    MergedIteratorState,
+)
+from pytorch_lightning.utilities.distributed import distributed_available
 from pytorch_lightning.utilities.fetching import AbstractDataFetcher, DataFetcher
+from pytorch_lightning.utilities.imports import _fault_tolerant_training
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.types import EPOCH_OUTPUT, STEP_OUTPUT
 
@@ -174,11 +181,20 @@ def on_save_checkpoint(self) -> Dict:
         state: Optional[MergedIteratorState] = getattr(self._data_fetcher.dataloader_iter, state_to_save, None)
         if state:
             state_dict["dataloader_state_dict"] = asdict(state)
+        state_dict["dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
+            state_dict["dataloader_state_dict"], device=self.trainer.training_type_plugin.root_device
+        )
         return state_dict
 
     def on_load_checkpoint(self, state_dict: Dict) -> None:
         # cache the dataloader state dict until the dataloader objects are available
-        self._dataloader_state_dict = state_dict.get("dataloader_state_dict")
+        # dataset states are collected across all ranks
+        if _fault_tolerant_training():
+            dataloader_state_dict = state_dict.get("dataloader_state_dict", None)
+            if not dataloader_state_dict:
+                return
+            rank = torch.distributed.get_rank() if distributed_available() else 0
+            self._dataloader_state_dict = dataloader_state_dict[rank]
 
     def _reload_dataloader_state_dict(self, data_fetcher: AbstractDataFetcher):
         if not self.trainer.sanity_checking and self._dataloader_state_dict:

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -25,6 +25,7 @@
 from pytorch_lightning.trainer.progress import BatchProgress, SchedulerProgress
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection
+from pytorch_lightning.utilities.auto_restart import _collect_states_on_rank_zero_over_collection
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.fetching import AbstractDataFetcher
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -320,9 +321,14 @@ def on_save_checkpoint(self) -> Dict:
             or self.batch_progress.current.ready == 0  # did not start
         ):
             return state_dict
+
         state_dict["dataloader_state_dict"] = self.trainer.train_dataloader.state_dict(
             has_completed=self._has_completed()
         )
+
+        state_dict["dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
+            state_dict["dataloader_state_dict"], device=self.trainer.training_type_plugin.root_device
+        )
         return state_dict
 
     def on_load_checkpoint(self, state_dict: Dict) -> None:

pytorch_lightning/trainer/connectors/data_connector.py

@@ -19,6 +19,7 @@
 
 import pytorch_lightning as pl
 from pytorch_lightning.utilities import rank_zero_deprecation
+from pytorch_lightning.utilities.auto_restart import _teardown_dataloader_get_iterators
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.fetching import (
     AbstractDataFetcher,
@@ -254,6 +255,7 @@ def teardown(self) -> None:
         if self.sanity_check_data_fetcher:
             self.sanity_check_data_fetcher.teardown()
             self.sanity_check_data_fetcher = None
+        _teardown_dataloader_get_iterators()
 
 
 @dataclass

pytorch_lightning/trainer/supporters.py

@@ -29,6 +29,7 @@
     patch_dataloader_iterator,
 )
 from pytorch_lightning.utilities.data import get_len
+from pytorch_lightning.utilities.distributed import distributed_available
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _fault_tolerant_training
 
@@ -403,6 +404,10 @@ def create_loader_iters(dataloader: DataLoader, state_dict: Dict) -> Iterator:
             if isinstance(dataloader, CycleIterator):
                 dataloader = dataloader_to_iter_on.loader
 
+            # dataset states are collected across all ranks
+            rank = torch.distributed.get_rank() if distributed_available() else 0
+            state_dict = state_dict[rank]
+
             _reload_dataloader_state_dict(dataloader, state_dict)
 
             # We finally spawned the workers if any.

pytorch_lightning/trainer/trainer.py

@@ -2094,7 +2094,11 @@ def _results(self) -> Optional[ResultCollection]:
             return active_loop._results
 
     def _exit_gracefully_on_signal(self) -> None:
-        if _fault_tolerant_training() and self._terminate_gracefully:
+        if _fault_tolerant_training():
+            # the signal should be sent to rank 0
+            should_terminate_gracefully = self.training_type_plugin.broadcast(self._terminate_gracefully)
+            if not should_terminate_gracefully:
+                return
             caller = inspect.stack()[1]
             class_name = caller[0].f_locals["self"].__class__.__name__
             raise ExitGracefullyException(f"Exiting gracefully on {class_name}:{caller.function}")

pytorch_lightning/utilities/auto_restart.py

@@ -31,6 +31,8 @@
 from typing_extensions import Protocol, runtime_checkable
 
 import pytorch_lightning as pl
+from pytorch_lightning.utilities.apply_func import apply_to_collection
+from pytorch_lightning.utilities.distributed import _collect_states_on_rank_zero
 from pytorch_lightning.utilities.enums import _FaultTolerantMode, AutoRestartBatchKeys
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
@@ -737,3 +739,15 @@ def _teardown_dataloader_get_iterators() -> None:
     if get_iterator:
         DataLoader._get_iterator = get_iterator
         del DataLoader._ori_get_iterator
+
+
+def _collect_states_on_rank_zero_over_collection(state_dict: Any, device: torch.device) -> Any:
+    """This utility collects the state across processes for a collection of state."""
+
+    def fn(state: Dict):
+        nonlocal device
+        if state.get("state", None) is not None:
+            return _collect_states_on_rank_zero(state, device=device)
+        return {k: apply_to_collection(v, Dict, fn) for k, v in state.items()}
+
+    return apply_to_collection(state_dict, Dict, fn)

pytorch_lightning/utilities/imports.py

@@ -14,7 +14,6 @@
 """General utilities."""
 import importlib
 import operator
-import os
 import platform
 import sys
 from importlib.util import find_spec
@@ -111,4 +110,6 @@ def _compare_version(package: str, op: Callable, version: str, use_base_version:
 
 # experimental feature within PyTorch Lightning.
 def _fault_tolerant_training() -> bool:
-    return bool(int(os.getenv("PL_FAULT_TOLERANT_TRAINING", 0)))
+    from pytorch_lightning.utilities.enums import _FaultTolerantMode
+
+    return _FaultTolerantMode.detect_current_mode().is_enabled

tests/utilities/test_auto_restart.py

@@ -39,6 +39,7 @@
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.utilities.auto_restart import (
     _add_capture_metadata_collate,
+    _collect_states_on_rank_zero_over_collection,
     _MultiProcessingDataLoaderIterStateful,
     _patch_dataloader_get_iterators,
     _reload_dataloader_state_dict,
@@ -1254,6 +1255,13 @@ def load_state_dict(self, state_dict):
         self.counter = state_dict[0]["counter"]
 
 
+def test_collect_states_with_collection():
+    state = {"state": 0}
+    collection = [{"a": state, "b": [{"a": state}]}]
+    generated = _collect_states_on_rank_zero_over_collection(collection, torch.device("cpu"))
+    assert generated == [{"a": {0: state}, "b": [{"a": {0: state}}]}]
+
+
 @pytest.mark.parametrize("num_workers", [0])
 @mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "2"})
 def test_stateful_workers(num_workers):