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:
Thanks for the report @achigeor ! Let us reproduce the issue and have deeper look on where is the problem.

cc @jkhenning 
Checked the number of temporary saved checkpoints and only one remains in temp folder. Probably, the problem is that we need to explicitly remove them from Trains' storage: override 
remove method
https://github.com/pytorch/ignite/blob/0452e4199ae0da627274f563a901493d9654b80d/ignite/handlers/checkpoint.py#L370-L372

@jkhenning do you have an API to remove stored artifact ?
Hi @achigeor ,

When you say you see:
```
epoch_checkpoint_1.pt
epoch_checkpoint_2.pt
epoch_checkpoint_3.pt
...
```
Do you mean you see them in the Trains UI models page, or somewhere else? As @vfdev-5 said, I believe these files will not be present in the `dirname` folder (or the temp folder used if `dirname` is not provided)
Hey @jkhenning,

I mean I see them in the directory I pass as `output_uri`. E.g. if you pass `output_uri="checkpoints"`,  you will see these checkpoints in `checkpoints/trains_project_name/trains_task_name.task_id/models/`:
```
|   +-- <project-name>
|       +-- <task-name>.<task-Id>
|           +-- models
|               +-- <epoch_checkpoint_1.pt>
|               +-- <epoch_checkpoint_2.pt>
|               +-- <epoch_checkpoint_3.pt>
```

In the Trains UI, you would only see the _last_ model model saved in `artifacts/Output Model`, i.e `task_name-epoch_checkpoint_3.pt`.

EDIT: I also checked the temp folder. It indeed only keeps the last one. However, the problem is that trains **_uploads_** all the checkpoints in the `output_uri` location defined - in this case a local directory.

Following is a fully reproducible example, based on the mnist example:
```
from argparse import ArgumentParser

import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor

from ignite.contrib.handlers.trains_logger import *
from ignite.engine import Events, create_supervised_trainer
from ignite.handlers import Checkpoint, DiskSaver

LOG_INTERVAL = 10


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=-1)


def get_data_loaders(train_batch_size):
    data_transform = Compose([ToTensor()])

    train_loader = DataLoader(
        MNIST(download=True, root=".", transform=data_transform, train=True), batch_size=train_batch_size, shuffle=True
    )

    return train_loader


def run(train_batch_size, epochs, lr=0.01):
    train_loader = get_data_loaders(train_batch_size)
    model = Net()
    device = "cpu"

    if torch.cuda.is_available():
        device = "cuda"

    model.to(device)  # Move model before creating optimizer
    optimizer = SGD(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    trainer = create_supervised_trainer(model, optimizer, criterion, device=device)

    TrainsLogger(project_name="trains-ignite", task_name="test-checkpoints")

    handler = Checkpoint(
        {"model": model},
        TrainsSaver(output_uri="test_checkpoints"),
        # DiskSaver(dirname="test_checkpoints", require_empty=False),
        n_saved=1,
        filename_prefix="epoch",
        global_step_transform=global_step_from_engine(trainer),
    )

    trainer.add_event_handler(Events.EPOCH_COMPLETED, handler)
    trainer.run(train_loader, max_epochs=epochs)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=32, help="input batch size for training (default: 64)")
    parser.add_argument("--epochs", type=int, default=3, help="number of epochs to train (default: 10)")

    args = parser.parse_args()

    run(args.batch_size, args.epochs)
```
I just removed the unnecessary stuff, focusing on the checkpoints.

You can comment  `TrainsSaver` or `DiskSaver` out, here:
```
    handler = Checkpoint(
        {"model": model},
        TrainsSaver(output_uri="test_checkpoints"),
        # DiskSaver(dirname="test_checkpoints", require_empty=False),
        n_saved=1,
        filename_prefix="epoch",
        global_step_transform=global_step_from_engine(trainer),
    )
```

Let me know if that helps. I just run the above and reproduced the behaviour I explained on the first message.

Hey @achigeor 

I see what you mean, and you're quite correct about Trains "uploading" the model files to the `output_uri` which in this case is a folder.
The thing is, `output_uri` is also (most commonly) used to represent remote storage services, like an S3 bucket, in which cases supporting a `delete` action might add additional constraints (requiring the credentials to allow these permissions) or introduce a considerable amount of risk (automated processes with delete permission for buckets containing "production" grade data) - for us, the compromise there was to take into account the relatively low cost of cloud storage and forego support for `delete` in this case.

An approach we did take in other cases, which is much more conservative, is reusing resource names, so if for example you'd like no more than the 3 last checkpoints to be stored, we can cycle through uploaded checkpoint file names, and upload as follows:
```
epoch_checkpoint_1.pt
epoch_checkpoint_2.pt
epoch_checkpoint_3.pt
epoch_checkpoint_1.pt
epoch_checkpoint_2.pt
...
```

As you can see, checkpoint 4 will simply overwrite old checkpoint 1 (and checkpoint 5 will override checkpoint 2), and assuming you're listing the files in modified order (or viewing through the Trains UI) - you still get the last 3 checkpoints, in the correct order (luckily, write access also mean overwrite access :)

This can be easily implemented in this case in the `TrainsSaver` itself (with a little help from Trains).

What do you say?
> An approach we did take in other cases, which is much more conservative, is reusing resource names, so if for example you'd like no more than the 3 last checkpoints to be stored, we can cycle through uploaded checkpoint file names

@jkhenning maybe this approach is not suitable for `epoch_checkpoint_<N>.pt` or best model file `best_model_val_score=<0.XYZ>.pt` where `<N>` indicates training epoch and `<0.XYZ>` best validation score. We can not cycle on those values, otherwise they do not make sense...

About `output_uri` and `dirname` args of `TrainSaver` shouldn't we limit to only one provided possible (in case of when `output_uri` is a local storage) ?

If I understand correctly in case of @achigeor the checkpoints are stored twice: at temp `dirname` and  at `output_uri` local storage ?
@vfdev-5 if I may chime in.

We are using the filename as a way to add meta-data on the model weight itself.

Maybe the best usage for TrainsLogger would be to store that meta-data as a comment on the model, this way it is also searchable? (we could also use the model name in the Trains system and store the meta-data there)
Another option is using the model configuration entry (that can store anything, and add the meta-data there) [see example](https://demoapp.trains.allegro.ai/projects/*/models/13432e3f49dd456ca87f963a59ad501e/network?order=created&filter=framework:PyTorch)

This removes the need to preserve the exact "filename" in the system, as the way to look for it would be to go to the Trains UI and search for the Model with specific score or checkpoint. Obviously you can also get it from a previously executed Task, including meta-data.

What do think?

p.s.
We are thinking on adding a more specific key/value to the model object, so one could add more diverse information.
Hey @vfdev-5 

Well, I agree it makes little sense to store twice, however I'm not sure detecting when `output_uri` is local and  when not is a good approach (from experience, likely to break on edge-cases).

I would say we should instruct users to use `dirname` for local storage, and `output_uri` for remote storage, but this means users who'd like to test their code with `output_uri` before switching to a remote storage will still see the same side-effects @achigeor noticed...
@bmartinn thanks for the idea ! If I understand correctly your suggestions is to : 
- pick filename which is need to be stored, e.g. `epoch_checkpoint_1.pt` and setup some meta-data
- save checkpoint on Trains system with another name that we can rewrite in a cycle manner as proposed by @jkhenning 
=> on Trains UI we have a good number of stored artifacts and their meta-data contain the name we intendeded to store ?


@jkhenning, could you please detail why we need both arguments actually ? Sorry, I become a bit lost with how it work with Trains... 
Naively, can't we keep one argument, e.g. `output_uri`?
@vfdev-5 yes exactly!
a minor detail though:
>on Trains UI we have a good number of stored artifacts and their meta-data contain the name we intended to store ?

Trains separates artifacts from Models. Basically artifacts you store once, and are tied with an experiment, where as  Models you store multiple times (e.g. snapshots) and have their own entity. This means we can easily track models between experiments. Unfortunately Artifacts already received the meta-data feature, and Models are still missing it, but I'm hoping it is coming soon ;)

To the point, models have names that are unrelated to the experiment that created them, (kind of like Artifacts). We can basically use these names as the meta-data. Notice that a user can always choose to rename a model, even after the experiment is completed, with no effect on the actual stored file. This gives you the possibility of renaming your model (in the UI) from "`best_model_val_score=<0.XYZ>.pt`" to "`this is the chosen one! score 0.999.pt`"


Hi all,

What's not clear to me regarding `output_uri` is what @jkhenning mentioned about remote saving. We will indeed use S3 to save our models and artifacts - we're now finalizing integrating Trains in our workflow and were using a local dir as `output_uri` for testing.
But from my understanding, trains would behave the same regardless, i.e, what I'm experiencing locally will still happen if I store in S3. Is there a mechanism that I miss that makes this behaviour relevant only for local storing? 

To give some perspective, my goal is to store a checkpoint every `N` epochs keeping the last `K` (as kind of a "backup") as well as the best model(s). With `DiskSaver` this is straightforward to do.

One solution could be to name the models `checkpoint` and `best_model` as @jkhenning suggested, but this won't work as @vfdev-5 pointed out, since afaik ignite adds a relevant suffix under the hood and no option to disable it yourself (please @vfdev-5 correct me if I'm wrong on this one). If this can work, it's fine for my usecase.

Regarding trains and output models I like the idea of adding metadata to them, especially since it will allow to programmatically query based on them (instead of having to do `torch.load()` to read from the `checkpoint` dictionary).
> One solution could be to name the models checkpoint and best_model as @jkhenning suggested, but this won't work as @vfdev-5 pointed out, since afaik ignite adds a relevant suffix under the hood and no option to disable it yourself (please @vfdev-5 correct me if I'm wrong on this one). If this can work, it's fine for my usecase.

@achigeor you are right, there is no easy way to store checkpoints without suffix. Maybe we can create a feature request for that inside `Checkpoint`. However, meta-data (e.g. epoch index or score) may be missing in the `.pt` file.
Hey @achigeor 

>But from my understanding, trains would behave the same regardless, i.e, what I'm experiencing locally will still happen if I store in S3. Is there a mechanism that I miss that makes this behaviour relevant only for local storing?

You're absolutely right - that was my point, that Trains currently handles remote and local storage locations exactly the same, without using any delete-like feature.

@vfdev-5 

>@jkhenning, could you please detail why we need both arguments actually ? Sorry, I become a bit lost with how it work with Trains...
Naively, can't we keep one argument, e.g. output_uri?

Well, the original intent of keeping `dirname`was to keep part of the `DiskSaver` interface since we need a local storage where checkpoint are stored prior to being uploaded anyway (if you remember, we create a temp dir in case `dirname` was not provided). I agree this represents some kind of duplication since `output_uri` (which is essential since it allows the user to provide a remote storage for Trains to use) also supports local storage.
So I guess the _actual_ original intent was _"Use `dirname` for local storage and `output_uri` for remote storage. If you don't use `dirname`, we'll use a temp dir instead"_ - but maybe that was not reflected very well in the documentation, and it's possible it's too complicated and we should just decide we'll drop `dirname` altogether and always use a temp dir...
> Hey @achigeor
> 
> > But from my understanding, trains would behave the same regardless, i.e, what I'm experiencing locally will still happen if I store in S3. Is there a mechanism that I miss that makes this behaviour relevant only for local storing?
> 
> You're absolutely right - that was my point, that Trains currently handles remote and local storage locations exactly the same, without using any delete-like feature.

Thanks got it now! I have this followup question then:

Is there a specific reason that uploading happens while the task is running? What I see happening now is:
* `Checkpoint` with `TrainsSaver` writes to `/tmp` or `dirname`, respecting `n_saved` param
* `TrainsSaver`, through trains task's `output_uri` uploads the model from `/tmp` to `output_uri` everytime a model is saved in the previous step

Would it be possible to instead upload the contents of the `/tmp` directory _after_ the task is finished? This way you upload only once, instead of everytime you save a model while the training task is running.
> Would it be possible to instead upload the contents of the /tmp directory after the task is finished? 

@achigeor this can be possible, but we should keep in mind that we can loose everything if script crashes before the `task is finished` and we didn't catch that.
From what I understand, [Checkpoint](https://github.com/pytorch/ignite/blob/master/ignite/handlers/checkpoint.py#L68) is doing the filename creation, and it seems that "`_`" is a hard-coded separator.

In case `output_uri` is  passed, we could parse the filename, extract the filename_prefix/name and use that as the base filename for upload, or we can use a hard-coded fixed name.

This way if `output_uri` is not passed, the behavior is consistent with `DiskSaver` (as is already the case now), and if `output_uri` is passed, the `Checkpoint` filename (e.g. `{filename_prefix}_{name}_{global_step}_{score_name}={score}.{ext}`) will be used as the Trains Model name (i.e. in the UI), making it intuitive to look for.

What do you guys think?
> In case output_uri is passed, we could parse the filename, extract the filename_prefix/name and use that as the base filename for upload, or we can use a hard-coded fixed name.

Well, both approaches are not ideal, but between parsing and hard-coded fixed name, IMO, hard-coded fixed name probably better... especially if in UI full correct filename is used.


@vfdev-5 I agree, I think hard-coded is probably safer, but we should have a way to differentiate between two models being stored, for example GAN training. We brings me back to parsing...
How about we inherit Checkpoint, that way we know what's the prefix and can easily make those decisions. What do you think?
I was also thinking about if we can pass more info to `save_handler` in addition to `object_to_save` and `filename`:
https://github.com/pytorch/ignite/blob/c012166f93e56f8e9538741f5745a5010983ba38/ignite/handlers/checkpoint.py#L21

For example, we can opt to pass some meta-info about the checkpoint to save:
```python
class BaseSaveHandler(metaclass=ABCMeta):
    """Base class for save handlers"""

    @abstractmethod
    def __call__(self, checkpoint: Mapping, filename: str, metadata=None) -> None:
        pass
```
and in metadata we can pass `prefix`, `name` and all scores which compose the filename.

This certainly requires minor API change for `DiskSaver` and other savers. However, we recently introduced `BaseSaveHandler` as base class for savers, so we still can change thing now...
@vfdev-5 I like it! I think it makes a lot of sense. That means the DiskSaver can gain more insight into the model being stored (so for example, if the model score is lower then the previous, it can only store it locally and choose to skip the upload).

Yes, for a cloud disk savers or other more complex savers where we can set some metadata it can be helpful and avoid parsing the filename :)

> so for example, if the model score is lower then the previous, it can only store it locally and choose to skip the upload

Normally, it is the role of `Checkpoint` with score function etc.

So, let us create a FR for that.


@bmartinn PR with metadata is landed. So, you may try to incorporate it into `TrainsSaver` and give a feedback on that....
Thanks @vfdev-5 ! kudos for the quick PR!
We will add it into `TrainsServer` (probably needs to add a bit of support in `Trains` as well)
I'll update here once the PR is ready
@vfdev-5 Great idea, and 👍 for pushing things so quickly - once we'll add the required support in trains I'll open a PR 😄 
@bmartinn as 0.15.0 is out from your side, what can be done here to fix this issue ?
@jkhenning could you help with a PR fix?
@bmartinn, @vfdev-5 I'll get on it :)
Hi @vfdev-5,

I've pushed a draft into my fork, branch `fix_trains_checkpoint_n_saved` (see [here](https://github.com/jkhenning/ignite/tree/fix_trains_checkpoint_n_saved)) - I'd appreciate your input.

Basically, I changed the `TrainsSaver` behavior to use the new Trains pre/post registration/upload callbacks, which always uses the filename as passed to the saver (which removed the need to have a special case for `atomic` mode).

However, in order to preserve the `n_saved` behavior which keeps the `n` checkpoints with the highest priority, the logic inside the pre-save callback had to be aware of the `n_saved` value stored in the checkpoint, as well as the list of saved files and their priorities - this is required in order for the callback to know which previously registered/uploaded file must be overwritten. Even though I could have achieved that by inheriting the Checkpoint class, it seems a bit too much - I was wondering if you think we can simply expose these values in the Checkpoint's interface (for read, obviously, not write), or whether you have a better idea :)

@jkhenning thanks for working on that !
Could you please confirm me that overriding `DiskSaver.remove` method can not help for that ?
For example, `TrainsSaver` receives checkpoints with the filenames and metadata:
```
T1 : save: training_checkpoint_100.pt meta1
---
T2 : save: training_checkpoint_200.pt meta2
---
T3 : save: training_checkpoint_300.pt meta3
---
T4 : save: training_checkpoint_400.pt meta4
T4 : remove: training_checkpoint_100.pt
---
T5 : save: training_checkpoint_500.pt meta5
T5 : remove: training_checkpoint_200.pt
---
...
```
if you store somewhere a mapping between `training_checkpoint_X.pt` and the name you save `n` checkpoints to store on server, maybe it could be possible to override `remove: training_checkpoint_X.pt` and perform necessary action ?
What do you think ?

PS: I tried to reach out you on allegro slack to talk about this in more fluent way.
I think the reason `remove` might not be enough, is that it will basically be too late.
In other words if `n_saved` equals to 5, the first time the TrainsLogger will see a `remove` called will be after the 6th model was created, so essentially although you wanted the "last" 5 models, TrainsLogger will store the "last" 6.

@vfdev-5 makes sense ?

Resolving patch:
diff --git a/examples/contrib/mnist/mnist_with_trains_logger.py b/examples/contrib/mnist/mnist_with_trains_logger.py
--- a/examples/contrib/mnist/mnist_with_trains_logger.py
+++ b/examples/contrib/mnist/mnist_with_trains_logger.py
@@ -122,7 +122,7 @@ def compute_metrics(engine):
 
     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
--- 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 @@ def __setattr__(self, attr, val):
 
             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 @@ def __init__(self, logger: TrainsLogger = None, output_uri: str = None, dirname:
         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 @@ def _setup_check_trains(self, logger, output_uri):
         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 @@ def get_local_copy(self, filename: str) -> Optional[str]:
         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


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_trains_logger.py b/tests/ignite/contrib/handlers/test_trains_logger.py
--- 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__()