Enable loading legacy checkpoints that pickled the _FaultToleranceMode enum (#18094)

Author: awaelchli

src/lightning/pytorch/CHANGELOG.md

@@ -172,6 +172,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed validation of non-PyTorch LR schedulers in manual optimization mode ([#18092](https://github.com/Lightning-AI/lightning/pull/18092))
 
 
+- Fixed an attribute error for `_FaultTolerantMode` when loading an old checkpoint that pickled the enum ([#18094](https://github.com/Lightning-AI/lightning/pull/18094))
+
+
 ## [2.0.5] - 2023-07-07
 
 ### Fixed

src/lightning/pytorch/utilities/migration/utils.py

@@ -21,6 +21,7 @@
 from packaging.version import Version
 
 import lightning.pytorch as pl
+from lightning.fabric.utilities.enums import LightningEnum
 from lightning.fabric.utilities.imports import _IS_WINDOWS
 from lightning.fabric.utilities.types import _PATH
 from lightning.fabric.utilities.warnings import PossibleUserWarning
@@ -79,6 +80,8 @@ class pl_legacy_patch:
            version 1.2.8. See: https://github.com/Lightning-AI/lightning/pull/6898
         2. ``lightning.pytorch.utilities.argparse_utils``: A module that was deprecated in 1.2 and removed in 1.4,
            but still needs to be available for import for legacy checkpoints.
+        3. ``lightning.pytorch.utilities.enums._FaultTolerantMode``: This enum was removed in 2.0 but was pickled
+           into older checkpoints.
 
     Example:
 
@@ -95,6 +98,14 @@ def __enter__(self) -> "pl_legacy_patch":
         # `_gpus_arg_default` used to be imported from these locations
         legacy_argparse_module._gpus_arg_default = lambda x: x
         pl.utilities.argparse._gpus_arg_default = lambda x: x
+
+        # `_FaultTolerantMode` was removed from the enums
+        class _FaultTolerantMode(LightningEnum):
+            DISABLED = "disabled"
+            AUTOMATIC = "automatic"
+            MANUAL = "manual"
+
+        pl.utilities.enums._FaultTolerantMode = _FaultTolerantMode
         return self
 
     def __exit__(
@@ -106,6 +117,8 @@ def __exit__(
         if hasattr(pl.utilities.argparse, "_gpus_arg_default"):
             delattr(pl.utilities.argparse, "_gpus_arg_default")
         del sys.modules["lightning.pytorch.utilities.argparse_utils"]
+        if hasattr(pl.utilities.enums, "_FaultTolerantMode"):
+            delattr(pl.utilities.enums, "_FaultTolerantMode")
         _lock.release()
 
 

tests/tests_pytorch/utilities/migration/test_utils.py

@@ -42,6 +42,14 @@ def test_patch_legacy_gpus_arg_default():
     assert not hasattr(pl.utilities.argparse, "_gpus_arg_default")
 
 
+def test_patch_legacy_fault_tolerant_mode():
+    with pl_legacy_patch():
+        from lightning.pytorch.utilities.enums import _FaultTolerantMode
+
+        assert _FaultTolerantMode.AUTOMATIC.value == "automatic"
+    assert not hasattr(pl.utilities.enums, "_FaultTolerantMode")
+
+
 def test_migrate_checkpoint(monkeypatch):
     """Test that the correct migration function gets executed given the current version of the checkpoint."""
     # A checkpoint that is older than any migration point in the index