[Feature] Upload checkpoints and logs to ceph

mmcv/fileio/file_client.py

@@ -11,6 +11,7 @@
 from typing import Iterable, Iterator, Optional, Tuple, Union
 from urllib.request import urlopen
 
+import mmcv
 from mmcv.utils.misc import has_method
 from mmcv.utils.path import is_filepath
 
@@ -23,6 +24,17 @@ class BaseStorageBackend(metaclass=ABCMeta):
     as texts.
     """
 
+    # a flag to indicate whether the backend can create a symlink for a file
+    _allow_symlink = False
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    @property
+    def allow_symlink(self):
+        return self._allow_symlink
+
     @abstractmethod
     def get(self, filepath):
         pass
@@ -41,8 +53,8 @@ class CephBackend(BaseStorageBackend):
             will be replaced by ``dst``. Default: None.
 
     .. warning::
-        :class:`CephBackend` will be deprecated, please use
-        :class:`PetrelBackend` instead
+        :class:`mmcv.fileio.file_client.CephBackend` will be deprecated,
+        please use :class:`mmcv.fileio.file_client.PetrelBackend` instead.
     """
 
     def __init__(self, path_mapping=None):
@@ -266,8 +278,8 @@ def isfile(self, filepath: Union[str, Path]) -> bool:
         filepath = self._format_path(filepath)
         return self._client.contains(filepath)
 
-    def concat_paths(self, filepath: Union[str, Path],
-                     *filepaths: Union[str, Path]) -> str:
+    def join_path(self, filepath: Union[str, Path],
+                  *filepaths: Union[str, Path]) -> str:
         """Concatenate all file paths.
 
         Args:
@@ -377,7 +389,7 @@ def _list_dir_or_file(dir_path, list_dir, list_file, suffix,
                 # is a directory, because `self.isdir` relies on
                 # `self._client.list`
                 if path.endswith('/'):  # a directory path
-                    next_dir_path = self.concat_paths(dir_path, path)
+                    next_dir_path = self.join_path(dir_path, path)
                     if list_dir:
                         # get the relative path and exclude the last
                         # character '/'
@@ -388,7 +400,7 @@ def _list_dir_or_file(dir_path, list_dir, list_file, suffix,
                                                      list_file, suffix,
                                                      recursive)
                 else:  # a file path
-                    absolute_path = self.concat_paths(dir_path, path)
+                    absolute_path = self.join_path(dir_path, path)
                     rel_path = absolute_path[len(root):]
                     if (suffix is None
                             or rel_path.endswith(suffix)) and list_file:
@@ -491,6 +503,8 @@ def get_text(self, filepath, encoding=None):
 class HardDiskBackend(BaseStorageBackend):
     """Raw hard disks storage backend."""
 
+    _allow_symlink = True
+
     def get(self, filepath: Union[str, Path]) -> bytes:
         """Read data from a given ``filepath`` with 'rb' mode.
 
@@ -524,10 +538,15 @@ def get_text(self,
     def put(self, obj: bytes, filepath: Union[str, Path]) -> None:
         """Write data to a given ``filepath`` with 'wb' mode.
 
+        Note:
+            ``put`` will create a directory if the directory of ``filepath``
+            does not exist.
+
         Args:
             obj (bytes): Data to be written.
             filepath (str or Path): Path to write data.
         """
+        mmcv.mkdir_or_exist(osp.dirname(filepath))
         with open(filepath, 'wb') as f:
             f.write(obj)
 
@@ -537,12 +556,17 @@ def put_text(self,
                  encoding: str = 'utf-8') -> None:
         """Write data to a given ``filepath`` with 'w' mode.
 
+        Note:
+            ``put_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
         Args:
             obj (str): Data to be written.
             filepath (str or Path): Path to write data.
             encoding (str): The encoding format used to open the ``filepath``.
                 Default: 'utf-8'.
         """
+        mmcv.mkdir_or_exist(osp.dirname(filepath))
         with open(filepath, 'w', encoding=encoding) as f:
             f.write(obj)
 
@@ -579,7 +603,7 @@ def isdir(self, filepath: Union[str, Path]) -> bool:
         return osp.isdir(filepath)
 
     def isfile(self, filepath: Union[str, Path]) -> bool:
-        """Check a ``filepath`` whether it is a file.
+        """Check whether a file path is a file.
 
         Args:
             filepath (str or Path): Path to be checked whether it is a file.
@@ -590,8 +614,8 @@ def isfile(self, filepath: Union[str, Path]) -> bool:
         """
         return osp.isfile(filepath)
 
-    def concat_paths(self, filepath: Union[str, Path],
-                     *filepaths: Union[str, Path]) -> str:
+    def join_path(self, filepath: Union[str, Path],
+                  *filepaths: Union[str, Path]) -> str:
         """Concatenate all file paths.
 
         Join one or more filepath components intelligently. The return value
@@ -714,7 +738,7 @@ class FileClient:
     Note that It can also register other backend accessor with a given name,
     prefixes, and backend class. In addition, We use the singleton pattern to
     avoid repeated object creation. If the arguments are the same, the same
-    object is returned.
+    object will be returned.
 
     Args:
         backend (str, optional): The storage backend type. Options are "disk",
@@ -788,18 +812,21 @@ def __new__(cls, backend=None, prefix=None, **kwargs):
             _instance = super().__new__(cls)
             if backend is not None:
                 _instance.client = cls._backends[backend](**kwargs)
-                _instance.backend_name = backend
             else:
                 _instance.client = cls._prefix_to_backends[prefix](**kwargs)
-                # infer the backend name according to the prefix
-                for backend_name, backend_cls in cls._backends.items():
-                    if isinstance(_instance.client, backend_cls):
-                        _instance.backend_name = backend_name
-                        break
+
             cls._instances[arg_key] = _instance
 
         return _instance
 
+    @property
+    def name(self):
+        return self.client.name
+
+    @property
+    def allow_symlink(self):
+        return self.client.allow_symlink
+
     @staticmethod
     def parse_uri_prefix(uri: Union[str, Path]) -> Optional[str]:
         """Parse the prefix of a uri.
@@ -980,6 +1007,10 @@ def get_text(self, filepath: Union[str, Path], encoding='utf-8') -> str:
     def put(self, obj: bytes, filepath: Union[str, Path]) -> None:
         """Write data to a given ``filepath`` with 'wb' mode.
 
+        Note:
+            ``put`` should create a directory if the directory of ``filepath``
+            does not exist.
+
         Args:
             obj (bytes): Data to be written.
             filepath (str or Path): Path to write data.
@@ -989,6 +1020,10 @@ def put(self, obj: bytes, filepath: Union[str, Path]) -> None:
     def put_text(self, obj: str, filepath: Union[str, Path]) -> None:
         """Write data to a given ``filepath`` with 'w' mode.
 
+        Note:
+            ``put_text`` should create a directory if the directory of
+            ``filepath`` does not exist.
+
         Args:
             obj (str): Data to be written.
             filepath (str or Path): Path to write data.
@@ -1041,8 +1076,8 @@ def isfile(self, filepath: Union[str, Path]) -> bool:
         """
         return self.client.isfile(filepath)
 
-    def concat_paths(self, filepath: Union[str, Path],
-                     *filepaths: Union[str, Path]) -> str:
+    def join_path(self, filepath: Union[str, Path],
+                  *filepaths: Union[str, Path]) -> str:
         """Concatenate all file paths.
 
         Join one or more filepath components intelligently. The return value
@@ -1054,7 +1089,7 @@ def concat_paths(self, filepath: Union[str, Path],
         Returns:
             str: The result of concatenation.
         """
-        return self.client.concat_paths(filepath, *filepaths)
+        return self.client.join_path(filepath, *filepaths)
 
     @contextmanager
     def get_local_path(self, filepath: Union[str, Path]) -> Iterable[str]:

mmcv/fileio/handlers/base.py

@@ -7,7 +7,6 @@ class BaseFileHandler(metaclass=ABCMeta):
     # str-like object or bytes-like object. Pickle only processes bytes-like
     # objects but json only processes str-like object. If it is str-like
     # object, `StringIO` will be used to process the buffer.
-
     str_like = True
 
     @abstractmethod

mmcv/runner/checkpoint.py

@@ -323,28 +323,43 @@ def load_from_pavi(filename, map_location=None):
 
 
 @CheckpointLoader.register_scheme(prefixes='s3://')
-def load_from_ceph(filename, map_location=None, backend='ceph'):
+def load_from_ceph(filename, map_location=None, backend='petrel'):
     """load checkpoint through the file path prefixed with s3.  In distributed
     setting, this function download ckpt at all ranks to different temporary
     directories.
 
     Args:
         filename (str): checkpoint file path with s3 prefix
         map_location (str, optional): Same as :func:`torch.load`.
-        backend (str): The storage backend type. Options are "disk", "ceph",
-            "memcached" and "lmdb". Default: 'ceph'
+        backend (str, optional): The storage backend type. Options are 'ceph',
+            'petrel'. Default: 'petrel'.
+
+    .. warning::
+        :class:`mmcv.fileio.file_client.CephBackend` will be deprecated,
+        please use :class:`mmcv.fileio.file_client.PetrelBackend` instead.
 
     Returns:
         dict or OrderedDict: The loaded checkpoint.
     """
-
-    allowed_backends = ['ceph']
+    allowed_backends = ['ceph', 'petrel']
     if backend not in allowed_backends:
         raise ValueError(f'Load from Backend {backend} is not supported.')
 
-    fileclient = FileClient(backend=backend)
-    buffer = io.BytesIO(fileclient.get(filename))
-    checkpoint = torch.load(buffer, map_location=map_location)
+    if backend == 'ceph':
+        warnings.warn(
+            'CephBackend will be deprecated, please use PetrelBackend instead')
+
+    # CephClient and PetrelBackend have the same prefix 's3://' and the latter
+    # will be chosen as default. If PetrelBackend can not be instantiated
+    # successfully, the CephClient will be chosen.
+    try:
+        file_client = FileClient(backend=backend)
+    except ImportError:
+        allowed_backends.remove(backend)
+        file_client = FileClient(backend=allowed_backends[0])
+
+    with io.BytesIO(file_client.get(filename)) as buffer:
+        checkpoint = torch.load(buffer, map_location=map_location)
     return checkpoint
 
 
@@ -506,7 +521,6 @@ def load_checkpoint(model,
             pair of the regular expression operations. Default: strip
             the prefix 'module.' by [(r'^module\\.', '')].
 
-
     Returns:
         dict or OrderedDict: The loaded checkpoint.
     """
@@ -616,7 +630,11 @@ def get_state_dict(module, destination=None, prefix='', keep_vars=False):
     return destination
 
 
-def save_checkpoint(model, filename, optimizer=None, meta=None):
+def save_checkpoint(model,
+                    filename,
+                    optimizer=None,
+                    meta=None,
+                    file_client_args=None):
     """Save checkpoint to file.
 
     The checkpoint will have 3 fields: ``meta``, ``state_dict`` and
@@ -627,6 +645,10 @@ def save_checkpoint(model, filename, optimizer=None, meta=None):
         filename (str): Checkpoint filename.
         optimizer (:obj:`Optimizer`, optional): Optimizer to be saved.
         meta (dict, optional): Metadata to be saved in checkpoint.
+        file_client_args (dict, optional): Arguments to instantiate a
+            FileClient. See :class:`mmcv.fileio.FileClient` for details.
+            Default: None.
+            `New in version 1.3.16.`
     """
     if meta is None:
         meta = {}
@@ -654,6 +676,10 @@ def save_checkpoint(model, filename, optimizer=None, meta=None):
             checkpoint['optimizer'][name] = optim.state_dict()
 
     if filename.startswith('pavi://'):
+        if file_client_args is not None:
+            raise ValueError(
+                'file_client_args should be "None" if filename starts with'
+                f'"pavi://", but got {file_client_args}')
         try:
             from pavi import modelcloud
             from pavi import exception
@@ -674,8 +700,7 @@ def save_checkpoint(model, filename, optimizer=None, meta=None):
                 f.flush()
             model.create_file(checkpoint_file, name=model_name)
     else:
-        mmcv.mkdir_or_exist(osp.dirname(filename))
-        # immediately flush buffer
-        with open(filename, 'wb') as f:
+        file_client = FileClient.infer_client(file_client_args, filename)
+        with io.BytesIO() as f:
             torch.save(checkpoint, f)
-            f.flush()
+            file_client.put(f.getvalue(), filename)