Repository: pfnet/pytorch-pfn-extras

Issue:
Embed PPE version to ExtensionsManager snapshot and warn if version mismatch
A snapshot format may not be compatible with the one saved with the earlier PPE releases. Users will experience hard-to-debug error messages in that case. To ease the diagnosis, I think it's better to add a PPE version number to ExtensionsManager snapshot and warn if it does not match with the current version on loading state dict.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/training/manager.py b/pytorch_pfn_extras/training/manager.py
--- a/pytorch_pfn_extras/training/manager.py
+++ b/pytorch_pfn_extras/training/manager.py
@@ -10,6 +10,7 @@
 
 import torch
 
+import pytorch_pfn_extras
 from pytorch_pfn_extras import writing
 from pytorch_pfn_extras import reporting
 from pytorch_pfn_extras.training import extension as extension_module
@@ -519,12 +520,30 @@ def state_dict(
                                  for name in self._optimizers}
         to_save['extensions'] = {name: self._extensions[name].state_dict()
                                  for name in self._extensions}
+        to_save['ppe_version'] = pytorch_pfn_extras.__version__
         return to_save
 
+    def _check_snapshot_version(self, ppe_version: Optional[str]) -> None:
+        must_warn = ppe_version is None or (
+            ppe_version != pytorch_pfn_extras.__version__)
+
+        if not must_warn:
+            return
+
+        msg = ('You are trying to load a snapshot file taken using a different '
+               'PPE version.\n')
+
+        if ppe_version is not None:
+            msg += (f'Snapshot taken with PPE {ppe_version} but '
+                    f'currently using PPE {pytorch_pfn_extras.__version__}')
+
+        warnings.warn(msg)
+
     def load_state_dict(
             self,
             to_load: Dict[str, Any],
     ) -> None:
+        self._check_snapshot_version(to_load.get('ppe_version', None))
         self._start_iteration = to_load['_start_iteration']
         self.iteration = self._start_iteration
         self._start_execution = to_load.get('_start_execution', self.iteration)


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py b/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
@@ -3,6 +3,7 @@
 import torch
 from torch import nn
 
+import pytorch_pfn_extras as ppe
 from pytorch_pfn_extras import training
 
 
@@ -186,6 +187,7 @@ def test_extensions_manager_state_dict():
     state_dict = manager.state_dict()
 
     assert state_dict == {
+        'ppe_version': ppe.__version__,
         '_start_execution': passed_iteration,
         '_start_iteration': passed_iteration,
         'models': {'model_name': model_state_dict},
@@ -214,6 +216,105 @@ def test_extensions_manager_state_dict():
     assert new_optimizer.called_load_state_dict == 1
 
 
+def test_extensions_manager_state_dict_old_ppe_no_version():
+    model_state_dict = object()
+    optimizer_state_dict = object()
+    max_epochs = 5
+    iters_per_epoch = 4
+    passed_iteration = 11
+
+    manager = training.ExtensionsManager(
+        {'model_name': _StateDictModel(state_dict=model_state_dict)},
+        {'optimizer_name': _StateDictObj(state_dict=optimizer_state_dict)},
+        max_epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+
+    for _ in range(passed_iteration):
+        with manager.run_iteration():
+            pass
+
+    new_model = _StateDictModel(state_dict_to_be_loaded=model_state_dict)
+    new_optimizer = _StateDictObj(state_dict_to_be_loaded=optimizer_state_dict)
+    manager_2 = training.ExtensionsManager(
+        {'model_name': new_model},
+        {'optimizer_name': new_optimizer},
+        max_epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+
+    state_dict = manager.state_dict()
+    del state_dict['ppe_version']
+    with pytest.warns(UserWarning, match='version'):
+        manager_2.load_state_dict(state_dict)
+
+
+def test_extensions_manager_state_dict_old_ppe_version():
+    model_state_dict = object()
+    optimizer_state_dict = object()
+    max_epochs = 5
+    iters_per_epoch = 4
+    passed_iteration = 11
+
+    manager = training.ExtensionsManager(
+        {'model_name': _StateDictModel(state_dict=model_state_dict)},
+        {'optimizer_name': _StateDictObj(state_dict=optimizer_state_dict)},
+        max_epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+
+    for _ in range(passed_iteration):
+        with manager.run_iteration():
+            pass
+
+    new_model = _StateDictModel(state_dict_to_be_loaded=model_state_dict)
+    new_optimizer = _StateDictObj(state_dict_to_be_loaded=optimizer_state_dict)
+    manager_2 = training.ExtensionsManager(
+        {'model_name': new_model},
+        {'optimizer_name': new_optimizer},
+        max_epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+
+    state_dict = manager.state_dict()
+    state_dict['ppe_version'] = '0.4.0'
+    with pytest.warns(UserWarning, match='version'):
+        manager_2.load_state_dict(state_dict)
+
+
+def test_extensions_manager_state_dict_future_ppe_version():
+    model_state_dict = object()
+    optimizer_state_dict = object()
+    max_epochs = 5
+    iters_per_epoch = 4
+    passed_iteration = 11
+
+    manager = training.ExtensionsManager(
+        {'model_name': _StateDictModel(state_dict=model_state_dict)},
+        {'optimizer_name': _StateDictObj(state_dict=optimizer_state_dict)},
+        max_epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+
+    for _ in range(passed_iteration):
+        with manager.run_iteration():
+            pass
+
+    new_model = _StateDictModel(state_dict_to_be_loaded=model_state_dict)
+    new_optimizer = _StateDictObj(state_dict_to_be_loaded=optimizer_state_dict)
+    manager_2 = training.ExtensionsManager(
+        {'model_name': new_model},
+        {'optimizer_name': new_optimizer},
+        max_epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+
+    state_dict = manager.state_dict()
+    state_dict['ppe_version'] = '23.0.0'
+    with pytest.warns(UserWarning, match='version'):
+        manager_2.load_state_dict(state_dict)
+
+
 def test_ignite_extensions_manager_state_dict():
 
     try:
@@ -251,6 +352,7 @@ def test_ignite_extensions_manager_state_dict():
     state_dict = manager.state_dict()
 
     assert state_dict == {
+        'ppe_version': ppe.__version__,
         '_start_execution': passed_iteration,
         '_start_iteration': passed_iteration,
         '_epoch_length': iters_per_epoch,
@@ -299,6 +401,7 @@ def test_extensions_manager_with_plain_model_and_optimizer():
     state_dict = manager.state_dict()
 
     assert state_dict == {
+        'ppe_version': ppe.__version__,
         '_start_execution': 0,
         '_start_iteration': 0,
         'models': {'main': model_state_dict},