Repository: pytorch/ignite

Issue:
TrainsSaver doesn't respect Checkpoint's n_saved
## 🐛 Bug description

<!-- A clear and concise description of what the bug is. -->
As the title says, it seems that `TrainsSaver` bypasses the `Checkpoint` `n_saved` parameter. That means that all models are saved and never updated / deleted.
<!-- Please, add steps on how to reproduce it. -->
Consider this simple example:
```
        task.phases['train'].add_event_handler(
            Events.EPOCH_COMPLETED(every=1),
            Checkpoint(to_save, TrainsSaver(output_uri=output_uri), 'epoch', n_saved=1,
                       global_step_transform=global_step_from_engine(task.phases['train'])))
```
The above saves every checkpoint. You end-up with
```
epoch_checkpoint_1.pt
epoch_checkpoint_2.pt
epoch_checkpoint_3.pt
...
```

Now if we do, the same with `DiskSaver`: 
```
        task.phases['train'].add_event_handler(
            Events.EPOCH_COMPLETED(every=1),
            Checkpoint(to_save, DiskSaver(dirname=dirname), 'epoch', n_saved=1,
                       global_step_transform=global_step_from_engine(task.phases['train'])))
```
We get only:
```
epoch_checkpoint_3.pt
```
as expected.

Same behaviour if we save only best models using `score_function`, i.e. `TrainsSaver` saves every best model.

## Environment

 - PyTorch Version: 1.3.1
 - Ignite Version: 0.4.0.dev20200519  (EDIT: update to latest nightly, issue still exists)
 - OS: Linux
 - How you installed Ignite: pip nightly
 - Python version: 3.6
 - Any other relevant information: trains version: 0.14.3


Discussion:
jkhenning: Hi @vfdev-5,

Please note that `requirements-dev.txt` was updated to require Trains 0.15.1rc0 that supports some new required changes. Once the PR is approved in principal, we'll release Trains 0.15.1 so this PR can be merged.

Thanks 😄 
bmartinn: @vfdev-5 let me add some missing background on the functionality of the callbacks.

The idea is to be able to affect the way Trains logs and uploads any stored model.
These callbacks will be called during from the Trains context, right after `torch.save` stored the mode.
This allows us to add metadata on the Model object created in Trains, change the uploaded file name (i.e. limit according to the n_saved argument), change the Model name in the Trains UI etc.

The `_CallbackContext` makes sure we allocate the correct "upload" file for the current save operation while respecting the `n_saved` limit. Basically it will find an empty "slot" (i.e. filename we can use overwrite) in the `pre_callback` setting the upload filename for Trains, then in the `post_callback` it will rename the Trains Model object with the original filename and add the metadata in the comment section (so it is later searchable)

>Can we test somehow (via mock at least) if we fix the issue ?

I think `examples/contrib/mnist/mnist_with_trains_logger.py` reproduces the issue, as `Checkpoint` is created with `n_saved=1`, or am I missing something ?
vfdev-5: > I think examples/contrib/mnist/mnist_with_trains_logger.py reproduces the issue, as Checkpoint is created with n_saved=1, or am I missing something ?

I wanted to say to unit test it in some way. For example, when remove is called, intended logic is executed ... Probably, I need to play a bit around it to better understand what happens.
vfdev-5: @jkhenning could you please add unit tests on `_CallbacksContext.pre_callback/post_callback` as we discussed above and update the version of trains as 0.15.1 is released :) Thanks 1
jkhenning: @vfdev-5 I've added unit test, updated trains version and used the `CallbackType` enum instead of using plain strings - please let me know what you think.
vfdev-5: @jkhenning looks good, thanks ! I'll try to add some more checks inside the test and I can merge then.

Patch:
diff --git a/examples/contrib/mnist/mnist_with_trains_logger.py b/examples/contrib/mnist/mnist_with_trains_logger.py
index 5d59d95a..f9605f68 100644
--- a/examples/contrib/mnist/mnist_with_trains_logger.py
+++ b/examples/contrib/mnist/mnist_with_trains_logger.py
@@ -122,7 +122,7 @@ def run(train_batch_size, val_batch_size, epochs, lr, momentum):
 
     handler = Checkpoint(
         {"model": model},
-        TrainsSaver(trains_logger, dirname="~/.trains/cache/"),
+        TrainsSaver(trains_logger, dirname="./TRAINS"),
         n_saved=1,
         score_function=lambda e: e.state.metrics["accuracy"],
         score_name="val_acc",
diff --git a/ignite/contrib/handlers/trains_logger.py b/ignite/contrib/handlers/trains_logger.py
index 10833ab5..b0b35e25 100644
--- a/ignite/contrib/handlers/trains_logger.py
+++ b/ignite/contrib/handlers/trains_logger.py
@@ -1,8 +1,10 @@
 import os
 import tempfile
 import warnings
+from collections import defaultdict
 from datetime import datetime
-from typing import Mapping, Optional
+from enum import Enum
+from typing import Any, List, Mapping, Optional, Type
 
 import torch
 
@@ -502,7 +504,8 @@ class TrainsLogger(BaseLogger):
 
     def __init__(self, *_, **kwargs):
         try:
-            import trains
+            from trains import Task
+            from trains.binding.frameworks.tensorflow_bind import WeightsGradientHistHelper
         except ImportError:
             raise RuntimeError(
                 "This contrib module requires trains to be installed. "
@@ -528,18 +531,16 @@ class TrainsLogger(BaseLogger):
 
             self._task = _Stub()
         else:
-            self._task = trains.Task.init(
+            self._task = Task.init(
                 project_name=kwargs.get("project_name"),
                 task_name=kwargs.get("task_name"),
-                task_type=kwargs.get("task_type", trains.Task.TaskTypes.training),
+                task_type=kwargs.get("task_type", Task.TaskTypes.training),
                 **experiment_kwargs,
             )
 
         self.trains_logger = self._task.get_logger()
 
-        self.grad_helper = trains.binding.frameworks.tensorflow_bind.WeightsGradientHistHelper(
-            logger=self.trains_logger,
-        )
+        self.grad_helper = WeightsGradientHistHelper(logger=self.trains_logger,)
 
     @classmethod
     def set_bypass_mode(cls, bypass: bool) -> None:
@@ -634,6 +635,8 @@ class TrainsSaver(DiskSaver):
         if "atomic" not in kwargs:
             kwargs["atomic"] = False
 
+        self._checkpoint_slots = defaultdict(list)
+
         super(TrainsSaver, self).__init__(dirname=dirname, *args, **kwargs)
 
     @idist.one_rank_only()
@@ -659,32 +662,92 @@ class TrainsSaver(DiskSaver):
         if output_uri:
             self._task.output_uri = output_uri
 
+    class _CallbacksContext:
+        def __init__(
+            self,
+            callback_type: Type[Enum],
+            slots: List,
+            checkpoint_key: str,
+            filename: str,
+            basename: str,
+            metadata: Optional[Mapping] = None,
+        ):
+            self._callback_type = callback_type
+            self._slots = slots
+            self._checkpoint_key = str(checkpoint_key)
+            self._filename = filename
+            self._basename = basename
+            self._metadata = metadata
+
+        def pre_callback(self, action: str, model_info: Any):
+            if action != self._callback_type.save:
+                return model_info
+
+            try:
+                slot = self._slots.index(None)
+                self._slots[slot] = model_info.upload_filename
+            except ValueError:
+                self._slots.append(model_info.upload_filename)
+                slot = len(self._slots) - 1
+
+            model_info.upload_filename = "{}_{}{}".format(self._basename, slot, os.path.splitext(self._filename)[1])
+            model_info.local_model_id = "{}:{}".format(self._checkpoint_key, model_info.upload_filename)
+            return model_info
+
+        def post_callback(self, action: str, model_info: Any):
+            if action != self._callback_type.save:
+                return model_info
+
+            model_info.model.name = "{}: {}".format(model_info.task.name, self._filename)
+            prefix = "Checkpoint Metadata: "
+            metadata = "{}{}".format(
+                prefix,
+                ", ".join("{}={}".format(k, v) for k, v in self._metadata.items()) if self._metadata else "none",
+            )
+            comment = "\n".join(
+                metadata if line.startswith(prefix) else line for line in (model_info.model.comment or "").split("\n")
+            )
+            if prefix not in comment:
+                comment += "\n" + metadata
+            model_info.model.comment = comment
+
+            return model_info
+
     def __call__(self, checkpoint: Mapping, filename: str, metadata: Optional[Mapping] = None) -> None:
-        super(TrainsSaver, self).__call__(checkpoint, filename, metadata)
-
-        if idist.get_rank() == 0:
-            # Maybe wont work with XLA
-            if self._atomic:
-                try:
-                    import trains
-                except ImportError:
-                    raise RuntimeError(
-                        "This contrib module requires trains to be installed. "
-                        "You may install trains using: \n pip install trains \n"
-                    )
-
-                # If atomic, DiskSaver's implementation first stores checkpoint into a temporary file
-                # and prohibits trains to automatically detect correct artifact path and name
-                path = os.path.join(self.dirname, filename)
-                if os.path.exists(path):
-                    trains.binding.frameworks.WeightsFileHandler.create_output_model(
-                        model=checkpoint,
-                        saved_path=path,
-                        framework=trains.model.Framework.pytorch,
-                        task=self._task,
-                        singlefile=True,
-                        model_name=os.path.basename(filename),
-                    )
+        try:
+            from trains import Model
+            from trains.binding.frameworks import WeightsFileHandler
+        except ImportError:
+            raise RuntimeError(
+                "This contrib module requires trains to be installed. "
+                "You may install trains using: \n pip install trains \n"
+            )
+
+        try:
+            basename = metadata["basename"]
+        except (TypeError, KeyError):
+            warnings.warn("Checkpoint metadata missing or basename cannot be found")
+            basename = "checkpoint"
+
+        checkpoint_key = (self.dirname, basename)
+
+        cb_context = self._CallbacksContext(
+            callback_type=WeightsFileHandler.CallbackType,
+            slots=self._checkpoint_slots[checkpoint_key],
+            checkpoint_key=str(checkpoint_key),
+            filename=filename,
+            basename=basename,
+            metadata=metadata,
+        )
+
+        pre_cb_id = WeightsFileHandler.add_pre_callback(cb_context.pre_callback)
+        post_cb_id = WeightsFileHandler.add_post_callback(cb_context.post_callback)
+
+        try:
+            super(TrainsSaver, self).__call__(checkpoint, filename, metadata)
+        finally:
+            WeightsFileHandler.remove_pre_callback(pre_cb_id)
+            WeightsFileHandler.remove_post_callback(post_cb_id)
 
     @idist.one_rank_only()
     def get_local_copy(self, filename: str) -> Optional[str]:
@@ -704,3 +767,14 @@ class TrainsSaver(DiskSaver):
         if artifact:
             return artifact.get_local_copy()
         self._task.get_logger().report_text("Can not find artifact {}".format(filename))
+
+    @idist.one_rank_only()
+    def remove(self, filename: str) -> None:
+        super(TrainsSaver, self).remove(filename)
+        for slots in self._checkpoint_slots.values():
+            try:
+                slots[slots.index(filename)] = None
+            except ValueError:
+                pass
+            else:
+                break
diff --git a/requirements-dev.txt b/requirements-dev.txt
index bc16fdfe..47b63d5f 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -17,7 +17,7 @@ mlflow
 neptune-client
 tensorboard
 pynvml; python_version > '3.5'
-trains
+trains>=0.15.1
 # Examples dependencies
 pandas
 gym


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_trains_logger.py b/tests/ignite/contrib/handlers/test_trains_logger.py
index 33ef12f8..321de1f1 100644
--- a/tests/ignite/contrib/handlers/test_trains_logger.py
+++ b/tests/ignite/contrib/handlers/test_trains_logger.py
@@ -1,10 +1,13 @@
 import math
 import os
+from collections import defaultdict
 from unittest.mock import ANY, MagicMock, Mock, call
 
 import pytest
 import torch
 import trains
+from trains.binding.frameworks import WeightsFileHandler
+from trains.model import Framework
 
 import ignite.distributed as idist
 from ignite.contrib.handlers.trains_logger import *
@@ -670,6 +673,88 @@ def test_trains_disk_saver_integration_no_logger():
         assert saved_files[0] == "model_1.pt"
 
 
+def test_trains_saver_callbacks():
+    mock_task = MagicMock(spec=trains.Task)
+    mock_task.name = "check-task"
+
+    mock_model = MagicMock(spec=trains.OutputModel)
+
+    model_info = WeightsFileHandler.ModelInfo(
+        model=mock_model,
+        upload_filename="test.pt",
+        local_model_path="",
+        local_model_id="",
+        framework=Framework.pytorch,
+        task=mock_task,
+    )
+
+    mock_model_info = MagicMock(spec_set=model_info)
+
+    # Simulate 4 calls to save model and 2 to remove (n_saved=2)
+    filenames = [
+        "best_model_5_val_acc=0.123.pt",
+        "best_model_6_val_acc=0.234.pt",
+        "best_model_7_val_acc=0.356.pt",
+        "best_model_8_val_acc=0.456.pt",
+    ]
+    metadata_list = [
+        {"basename": "best_model", "score_name": "val_acc", "priority": 0.123},
+        {"basename": "best_model", "score_name": "val_acc", "priority": 0.234},
+        {"basename": "best_model", "score_name": "val_acc", "priority": 0.345},
+        {"basename": "best_model", "score_name": "val_acc", "priority": 0.456},
+    ]
+    dirname = "/tmp/test"
+
+    _checkpoint_slots = defaultdict(list)
+
+    n_saved = 2
+
+    for i, (filename, metadata) in enumerate(zip(filenames, metadata_list)):
+
+        mock_model_info.upload_filename = filename
+
+        if i >= n_saved:
+            # Remove
+            filename_to_remove = filenames[i % n_saved]
+            for slots in _checkpoint_slots.values():
+                try:
+                    slots[slots.index(filename_to_remove)] = None
+                except ValueError:
+                    pass
+                else:
+                    i = i % n_saved
+                    break
+
+        basename = metadata["basename"]
+        checkpoint_key = (dirname, basename)
+
+        context = TrainsSaver._CallbacksContext(
+            callback_type=WeightsFileHandler.CallbackType,
+            slots=_checkpoint_slots[checkpoint_key],
+            checkpoint_key=str(checkpoint_key),
+            filename=filename,
+            basename=basename,
+            metadata=metadata,
+        )
+
+        output_model_info = context.pre_callback(str(WeightsFileHandler.CallbackType.save), mock_model_info)
+        assert (
+            hasattr(output_model_info, "upload_filename")
+            and "{}_{}.pt".format(basename, i) in output_model_info.upload_filename
+        )
+        assert hasattr(output_model_info, "local_model_id") and str(checkpoint_key) in output_model_info.local_model_id
+
+        output_model_info = context.post_callback(str(WeightsFileHandler.CallbackType.save), mock_model_info)
+        assert hasattr(output_model_info, "model") and hasattr(output_model_info.model, "name")
+        assert hasattr(output_model_info, "model") and hasattr(output_model_info.model, "comment")
+        assert isinstance(output_model_info.model.name, str) and filename in output_model_info.model.name
+        assert (
+            isinstance(output_model_info.model.comment, str)
+            and metadata["basename"] in output_model_info.model.comment
+            and metadata["score_name"] in output_model_info.model.comment
+        )
+
+
 class DummyModel(torch.nn.Module):
     def __init__(self):
         super(DummyModel, self).__init__()