Modify ModelCheckpoint class to support additional options for cloud storages
🚀  Feature
Modify pytorch_lightning.callbacks.ModelCheckpoint class to make it possible to provide additional configuration options for cloud storage.
Motivation
We at @blue-yonder started building a prototype that uses PyTorch-Lighting as a high-level training loop library. It works great, but we stumbled upon an issue with saving a model's checkpoint into a cloud storage. It looks like currently, the ModelCheckpoint class doesn't work properly with cloud file systems. If you look at the source code, you'll find a function called get_filesystem that doesn't delegate any storage_options to the underlying fsspec.open function. Therefore, it is not possible to anyhow provide credentials for a blob store.
def get_filesystem(path: Union[str, Path]):
    path = str(path)
    if "://" in path:
        # use the fileystem from the protocol specified

        # we need to delegate storage_options to this call like
        # fsspec.filesystem(..., **storage_options)
        return fsspec.filesystem(path.split(":", 1)[0])
    else:
        # use local filesystem
        return fsspec.filesystem("file")
In our use-case, we want to store a checkpoint on Azure Blob Store. We need to provide parameters like account_name, account_key and sas_token in order to access the blob. The same is true when resuming training from a checkpoint. It is not possible to directly use the trainer's interface, like:
pl.Trainer(resume_from_checkpoint="az://checkpoints/checkpoint.ckpt")
Because there is no way to provide credentials. In general, the checkpoint callback doesn't work well with cloud storage that requires some additional configuration options in addition to a file path.
Pitch
We created an ad-hoc implementation that inherits from ModelCheckpoint and makes it capable of working with Azure Blob Storage.
class AzureModelCheckpoint(ModelCheckpoint):
    """An alternative implementation of ``pytorch_lightning.callbacks.ModelCheckpoint`` that
    saves checkpoints into Azure Blob Store.
    """

    def __init__(
        self, account_name: str, account_key: str, sas_token: str, **checkpoint_options
    ):
        self.account_name = account_name
        self.account_key = account_key
        self.sas_token = sas_token
        self._parent_method = getattr(self, "_ModelCheckpoint__init_ckpt_dir")
        setattr(self, "_ModelCheckpoint__init_ckpt_dir", self.__init_ckpt_dir)
        super().__init__(**checkpoint_options)

    def __init_ckpt_dir(self, filepath, dirpath, filename, save_top_k):
        if filename is not None and dirpath is None:
            raise NotImplementedError(
                "setting filename is not supported; use dirpath instead"
            )

        # passing "local" path to trick the original implementation and
        # initialize the local filesystem instead of azure blob store in
        # order to properly override the initialized object afterwards
        fake_path = dirpath.replace("az://", "/")

        self._parent_method(None, fake_path, filename, save_top_k)

        self._fs = get_filesystem(
            str(dirpath) if dirpath else "",
            account_name=self.account_name,
            account_key=self.account_key,
            sas_token=self.sas_token,
        )

        self.dirpath = dirpath or None
        self.filename = filename or None

    def on_pretrain_routine_start(self, trainer, pl_module):
        super().on_pretrain_routine_start(trainer, pl_module)

        def save_function(filepath: str, weights_only: bool = False):
            checkpoint = trainer.checkpoint_connector.dump_checkpoint(weights_only)

            if trainer.is_global_zero:
                try:
                    atomic_save(
                        checkpoint,
                        filepath,
                        account_name=self.account_name,
                        account_key=self.account_key,
                        sas_token=self.sas_token,
                    )
                except AttributeError as err:
                    if LightningModule.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
                        del checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY]
                    rank_zero_warn(
                        "Warning, `hyper_parameters` dropped from checkpoint."
                        f"An attribute is not picklable {err}"
                    )
                    atomic_save(
                        checkpoint,
                        filepath,
                        account_name=self.account_name,
                        account_key=self.account_key,
                        sas_token=self.sas_token,
                    )

        self.save_function = save_function
Sure enough, it is a rather hackish side way and should be easily achievable by modifying the original checkpoint class instead, so it properly forwards storage options to the fsspec calls. Also, maybe changing Trainer or CheckpointConnector a bit to forward store's credentials when restoring from a checkpoint as well.
Alternatives

An alternative solution is to provide a custom implementation of checkpoint callback to save checkpoints into a cloud storage that requires authentication. To resume training from a checkpoint stored in such storage, download it locally into some temporary location and then use restore_from_checkpoint option. (This approach we use currently.) While it does what we need, this alternative isn't too convenient and requires extra effort.
Additional context
We use pytorch-lightning==1.0.8 and Python 3.6.6..

Could you please inform us if this functionality is somewhere on the roadmap? Or was this issue already raised before? There is a similar one about generic file system support, but it doesn't seem to address the credentials question. Also, we're ready to open a PR with a basic implementation if it helps.
Thank you!