Repository: pytorch/ignite

Issue:
Add meta-data as optional arg for BaseSaveHandler
## 🚀 Feature

According to https://github.com/pytorch/ignite/issues/1056#issuecomment-632919883
it could be helpful to provide checkpoint metadata to a save handler. 

Idea is to have
```python
class BaseSaveHandler(metaclass=ABCMeta):
    """Base class for save handlers"""

    @abstractmethod
    def __call__(self, checkpoint: Mapping, filename: str, metadata: Mapping =None) -> None:
        pass
```

We also need to make sure to update all existing save handlers like `DiskSaver`, `NeptuneSaver` etc.


Discussion:
erip: LGTM!

Patch:
diff --git a/docs/source/handlers.rst b/docs/source/handlers.rst
index 196ce8ea..96b18175 100644
--- a/docs/source/handlers.rst
+++ b/docs/source/handlers.rst
@@ -5,6 +5,7 @@ Complete list of handlers
 -------------------------
 
     - :class:`~ignite.handlers.Checkpoint`
+    - :class:`~ignite.handlers.checkpoint.BaseSaveHandler`
     - :class:`~ignite.handlers.DiskSaver`
     - :class:`~ignite.handlers.ModelCheckpoint`
     - :class:`~ignite.handlers.EarlyStopping`
@@ -17,6 +18,9 @@ Complete list of handlers
 .. autoclass:: Checkpoint
     :members: load_objects
 
+.. autoclass:: ignite.handlers.checkpoint.BaseSaveHandler
+    :members: __call__, remove
+
 .. autoclass:: DiskSaver
 
 .. autoclass:: ModelCheckpoint
diff --git a/ignite/contrib/handlers/neptune_logger.py b/ignite/contrib/handlers/neptune_logger.py
index e342cf30..434d6182 100644
--- a/ignite/contrib/handlers/neptune_logger.py
+++ b/ignite/contrib/handlers/neptune_logger.py
@@ -1,7 +1,7 @@
 import numbers
 import tempfile
 import warnings
-from typing import Mapping
+from typing import Mapping, Optional
 
 import torch
 
@@ -574,7 +574,7 @@ class NeptuneSaver(BaseSaveHandler):
     def __init__(self, neptune_logger: NeptuneLogger):
         self._logger = neptune_logger
 
-    def __call__(self, checkpoint: Mapping, filename: str) -> None:
+    def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
 
         with tempfile.NamedTemporaryFile() as tmp:
             torch.save(checkpoint, tmp.name)
diff --git a/ignite/contrib/handlers/trains_logger.py b/ignite/contrib/handlers/trains_logger.py
index 28e520a4..7cae3d24 100644
--- a/ignite/contrib/handlers/trains_logger.py
+++ b/ignite/contrib/handlers/trains_logger.py
@@ -643,8 +643,8 @@ class TrainsSaver(DiskSaver):
         if output_uri:
             self.task.output_uri = output_uri
 
-    def __call__(self, checkpoint: Mapping, filename: str) -> None:
-        super(TrainsSaver, self).__call__(checkpoint, filename)
+    def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
+        super(TrainsSaver, self).__call__(checkpoint, filename, metadata)
 
         try:
             import trains
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
index 80342f52..dd62d70f 100644
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -15,14 +15,39 @@ __all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler"]
 
 
 class BaseSaveHandler(metaclass=ABCMeta):
-    """Base class for save handlers"""
+    """Base class for save handlers
+
+    Methods to override:
+
+    - :meth:`~ignite.handlers.checkpoint.BaseSaveHandler.__call__`
+    - :meth:`~ignite.handlers.checkpoint.BaseSaveHandler.remove`
+    """
 
     @abstractmethod
-    def __call__(self, checkpoint: Mapping, filename: str) -> None:
+    def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
+        """Method to save `checkpoint` with `filename`. Additionally, metadata dictionary is provided.
+
+        Metadata contains:
+
+        - `basename`: file prefix (if provided) with checkpoint name, e.g. `epoch_checkpoint`.
+        - `score_name`: score name if provided, e.g `val_acc`.
+        - `priority`: checkpoint priority value (higher is better), e.g. `12` or `0.6554435`
+
+        Args:
+            checkpoint (Mapping): checkpoint dictionary to save.
+            filename (str): filename associated with checkpoint.
+            metadata (Mapping, optional): metadata on checkpoint to save.
+
+        """
         pass
 
     @abstractmethod
     def remove(self, filename: str) -> None:
+        """Method to remove saved checkpoint.
+
+        Args:
+            filename (str): filename associated with checkpoint.
+        """
         pass
 
 
@@ -34,9 +59,9 @@ class Checkpoint:
     Args:
         to_save (Mapping): Dictionary with the objects to save. Objects should have implemented `state_dict` and `
             load_state_dict` methods.
-        save_handler (callable or `BaseSaveHandler`): Method or callable class to use to save engine and other provided
-            objects. Function receives two objects: checkpoint as a dictionary and filename. If `save_handler` is
-            callable class, it can
+        save_handler (callable or :class:`~ignite.handlers.checkpoint.BaseSaveHandler`): Method or callable class to
+            use to save engine and other provided objects. Function receives two objects: checkpoint as a dictionary
+            and filename. If `save_handler` is callable class, it can
             inherit of :class:`~ignite.handlers.checkpoint.BaseSaveHandler` and optionally implement `remove` method to
             keep a fixed number of saved checkpoints. In case if user needs to save engine's checkpoint on a disk,
             `save_handler` can be defined with :class:`~ignite.handlers.DiskSaver`.
@@ -184,7 +209,7 @@ class Checkpoint:
         self.global_step_transform = global_step_transform
 
     @property
-    def last_checkpoint(self) -> str:
+    def last_checkpoint(self) -> Optional[str]:
         if len(self._saved) < 1:
             return None
         return self._saved[-1].filename
@@ -237,7 +262,16 @@ class Checkpoint:
             if any(item.filename == filename for item in self._saved):
                 return
 
-            self.save_handler(checkpoint, filename)
+            metadata = {
+                "basename": "{}{}".format(self._fname_prefix, name),
+                "score_name": self._score_name,
+                "priority": priority,
+            }
+
+            try:
+                self.save_handler(checkpoint, filename, metadata)
+            except TypeError:
+                self.save_handler(checkpoint, filename)
 
             self._saved.append(Checkpoint.Item(priority, filename))
             self._saved.sort(key=lambda item: item[0])
@@ -350,7 +384,7 @@ class DiskSaver(BaseSaveHandler):
                     "".format(matched, dirname)
                 )
 
-    def __call__(self, checkpoint: Mapping, filename: str) -> None:
+    def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
         path = os.path.join(self.dirname, filename)
 
         if not self._atomic:


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
index c646bc85..495d9272 100644
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -87,13 +87,15 @@ def test_checkpoint_default():
         checkpointer(trainer)
         assert save_handler.call_count == 1
 
-        save_handler.assert_called_with(obj, "{}_0.pt".format(name))
+        metadata = {"basename": name, "score_name": None, "priority": 0}
+        save_handler.assert_called_with(obj, "{}_0.pt".format(name), metadata)
 
         trainer.state.epoch = 12
         trainer.state.iteration = 1234
         checkpointer(trainer)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}_1234.pt".format(name))
+        metadata["priority"] = 1234
+        save_handler.assert_called_with(obj, "{}_1234.pt".format(name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}_0.pt".format(name))
         assert checkpointer.last_checkpoint == "{}_1234.pt".format(name)
@@ -128,13 +130,15 @@ def test_checkpoint_with_global_step_transform():
         if len(filename_prefix) > 0:
             filename_prefix += "_"
 
-        save_handler.assert_called_with(obj, "{}{}_1.pt".format(filename_prefix, name))
+        metadata = {"basename": "{}{}".format(filename_prefix, name), "score_name": None, "priority": 1}
+        save_handler.assert_called_with(obj, "{}{}_1.pt".format(filename_prefix, name), metadata)
 
         trainer.state.epoch = 12
         trainer.state.iteration = 1234
         checkpointer(trainer)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}{}_12.pt".format(filename_prefix, name))
+        metadata["priority"] = 1234
+        save_handler.assert_called_with(obj, "{}{}_12.pt".format(filename_prefix, name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}{}_1.pt".format(filename_prefix, name))
         assert checkpointer.last_checkpoint == "{}{}_12.pt".format(filename_prefix, name)
@@ -162,7 +166,8 @@ def test_checkpoint_with_score_function():
         checkpointer(trainer)
         assert save_handler.call_count == 1
 
-        save_handler.assert_called_with(obj, "{}_0.7700.pt".format(name))
+        metadata = {"basename": name, "score_name": None, "priority": 0.77}
+        save_handler.assert_called_with(obj, "{}_0.7700.pt".format(name), metadata)
 
         trainer.state.epoch = 12
         trainer.state.iteration = 1234
@@ -170,7 +175,8 @@ def test_checkpoint_with_score_function():
 
         checkpointer(trainer)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}_0.7800.pt".format(name))
+        metadata["priority"] = 0.78
+        save_handler.assert_called_with(obj, "{}_0.7800.pt".format(name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}_0.7700.pt".format(name))
         assert checkpointer.last_checkpoint == "{}_0.7800.pt".format(name)
@@ -199,7 +205,8 @@ def test_checkpoint_with_score_name_and_function():
         checkpointer(trainer)
         assert save_handler.call_count == 1
 
-        save_handler.assert_called_with(obj, "{}_loss=-0.7700.pt".format(name))
+        metadata = {"basename": name, "score_name": "loss", "priority": -0.77}
+        save_handler.assert_called_with(obj, "{}_loss=-0.7700.pt".format(name), metadata)
 
         trainer.state.epoch = 12
         trainer.state.iteration = 1234
@@ -207,7 +214,8 @@ def test_checkpoint_with_score_name_and_function():
 
         checkpointer(trainer)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}_loss=-0.7600.pt".format(name))
+        metadata["priority"] = -0.76
+        save_handler.assert_called_with(obj, "{}_loss=-0.7600.pt".format(name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}_loss=-0.7700.pt".format(name))
         assert checkpointer.last_checkpoint == "{}_loss=-0.7600.pt".format(name)
@@ -241,14 +249,16 @@ def test_checkpoint_with_int_score():
         checkpointer(trainer)
         assert save_handler.call_count == 1
 
-        save_handler.assert_called_with(obj, "{}_{}1.pt".format(name, score_name))
+        metadata = {"basename": name, "score_name": score_name[:-1] if len(score_name) > 0 else None, "priority": 1}
+        save_handler.assert_called_with(obj, "{}_{}1.pt".format(name, score_name), metadata)
 
         trainer.state.epoch = 12
         trainer.state.iteration = 1234
 
         checkpointer(trainer)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}_{}12.pt".format(name, score_name))
+        metadata["priority"] = 12
+        save_handler.assert_called_with(obj, "{}_{}12.pt".format(name, score_name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}_{}1.pt".format(name, score_name))
         assert checkpointer.last_checkpoint == "{}_{}12.pt".format(name, score_name)
@@ -284,14 +294,16 @@ def test_checkpoint_with_score_function_and_trainer_epoch():
         checkpointer(evaluator)
         assert save_handler.call_count == 1
 
-        save_handler.assert_called_with(obj, "{}_11_0.7700.pt".format(name))
+        metadata = {"basename": name, "score_name": None, "priority": 0.77}
+        save_handler.assert_called_with(obj, "{}_11_0.7700.pt".format(name), metadata)
 
         trainer.state.epoch = 12
         evaluator.state.metrics["val_acc"] = 0.78
 
         checkpointer(evaluator)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}_12_0.7800.pt".format(name))
+        metadata["priority"] = 0.78
+        save_handler.assert_called_with(obj, "{}_12_0.7800.pt".format(name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}_11_0.7700.pt".format(name))
         assert checkpointer.last_checkpoint == "{}_12_0.7800.pt".format(name)
@@ -322,14 +334,16 @@ def test_checkpoint_with_score_name_and_function_and_trainer_epoch():
         checkpointer(evaluator)
         assert save_handler.call_count == 1
 
-        save_handler.assert_called_with(obj, "{}_11_val_acc=0.7700.pt".format(name))
+        metadata = {"basename": name, "score_name": "val_acc", "priority": 0.77}
+        save_handler.assert_called_with(obj, "{}_11_val_acc=0.7700.pt".format(name), metadata)
 
         trainer.state.epoch = 12
         evaluator.state.metrics["val_acc"] = 0.78
 
         checkpointer(evaluator)
         assert save_handler.call_count == 2
-        save_handler.assert_called_with(obj, "{}_12_val_acc=0.7800.pt".format(name))
+        metadata["priority"] = 0.78
+        save_handler.assert_called_with(obj, "{}_12_val_acc=0.7800.pt".format(name), metadata)
         assert save_handler.remove.call_count == 1
         save_handler.remove.assert_called_with("{}_11_val_acc=0.7700.pt".format(name))
         assert checkpointer.last_checkpoint == "{}_12_val_acc=0.7800.pt".format(name)
@@ -379,6 +393,20 @@ def test_checkpoint_last_checkpoint_on_score():
     assert checkpointer.last_checkpoint == "{}_val_acc=0.9000.pt".format("model")
 
 
+def test_checkpoint_save_handler_callable():
+    def save_handler(c, f):
+        assert f == "model_12.pt"
+
+    to_save = {"model": DummyModel()}
+
+    checkpointer = Checkpoint(to_save, save_handler=save_handler,)
+
+    trainer = Engine(lambda e, b: None)
+
+    trainer.state = State(epoch=1, iteration=12)
+    checkpointer(trainer)
+
+
 def test_model_checkpoint_args_validation(dirname):
     existing = os.path.join(dirname, "existing_dir")
     nonempty = os.path.join(dirname, "nonempty")