mosaicml · dakinggg · Oct 11, 2022 · Oct 8, 2022 · Oct 8, 2022 · Oct 8, 2022
@@ -36,7 +36,7 @@ class RemoteUploaderDownloaderHparams(hp.Hparams):
 
     Args:
         object_store_hparams (ObjectStoreHparams): The object store provider hparams.
-        object_name (str, optional): See :class:`.RemoteUploaderDownloader`.
+        path_format_string (str, optional): See :class:`.RemoteUploaderDownloader`.
         num_concurrent_uploads (int, optional): See :class:`.RemoteUploaderDownloader`.
         upload_staging_folder (str, optional): See :class:`.RemoteUploaderDownloader`.
         use_procs (bool, optional): See :class:`.RemoteUploaderDownloader`.
@@ -46,18 +46,19 @@ class RemoteUploaderDownloaderHparams(hp.Hparams):
         'object_store_hparams': object_store_registry,
     }
 
-    object_store_hparams: ObjectStoreHparams = hp.required('Object store provider hparams.')
-    object_name: str = hp.auto(RemoteUploaderDownloader, 'object_name')
+    bucket_uri: str = hp.required('Remote bucket uri')
+    object_store_hparams: Optional[ObjectStoreHparams] = hp.optional('Object store provider hparams.', default=None)
+    path_format_string: str = hp.auto(RemoteUploaderDownloader, 'path_format_string')
     num_concurrent_uploads: int = hp.auto(RemoteUploaderDownloader, 'num_concurrent_uploads')
     use_procs: bool = hp.auto(RemoteUploaderDownloader, 'use_procs')
     upload_staging_folder: Optional[str] = hp.auto(RemoteUploaderDownloader, 'upload_staging_folder')
     num_attempts: int = hp.auto(RemoteUploaderDownloader, 'num_attempts')
 
     def initialize_object(self) -> RemoteUploaderDownloader:
         return RemoteUploaderDownloader(
-            object_store_cls=self.object_store_hparams.get_object_store_cls(),
-            object_store_kwargs=self.object_store_hparams.get_kwargs(),
-            object_name=self.object_name,
+            bucket_uri=self.bucket_uri,
+            backend_kwargs=self.object_store_hparams.get_kwargs() if self.object_store_hparams is not None else {},
+            path_format_string=self.path_format_string,
             num_concurrent_uploads=self.num_concurrent_uploads,
             upload_staging_folder=self.upload_staging_folder,
             use_procs=self.use_procs,

diff --git a/docs/source/doctest_fixtures.py b/docs/source/doctest_fixtures.py
@@ -212,8 +212,8 @@ def _new_RemoteUploaderDownloader_init(self, fake_ellipses: None = None, **kwarg
     os.makedirs('./object_store', exist_ok=True)
     kwargs.update(use_procs=False,
                   num_concurrent_uploads=1,
-                  object_store_cls=LibcloudObjectStore,
-                  object_store_kwargs={
+                  bucket_uri='libcloud://.',
+                  backend_kwargs={
                       'provider': 'local',
                       'container': '.',
                       'provider_kwargs': {

diff --git a/docs/source/notes/resumption.rst b/docs/source/notes/resumption.rst
@@ -85,10 +85,7 @@ A typical use case is saving checkpoints to object store (e.g. S3) when there is
 
     # this assumes credentials are already configured via boto3
     remote_uploader_downloader = RemoteUploaderDownloader(
-        object_store_cls=S3ObjectStore,
-        object_store_kwargs={
-            "bucket": "checkpoint-debugging",
-            },
+        bucket_uri=f"s3://checkpoint-debugging",
     )
 
     trainer = Trainer(
@@ -118,10 +115,7 @@ To run fine-tuning on a spot instance, ``load_path`` would be set to the origina
     from composer.utils.object_store import S3ObjectStore
 
     remote_uploader_downloader = RemoteUploaderDownloader(
-        object_store_cls=S3ObjectStore,
-        object_store_kwargs={
-            "bucket": "checkpoint-debugging_2",
-        },
+        bucket_uri=f"s3://checkpoint-debugging_2",
     )
 
     # Train to generate and save the "pretrained_weights/model.pt",

diff --git a/docs/source/trainer/checkpointing.rst b/docs/source/trainer/checkpointing.rst
@@ -304,8 +304,8 @@ Behind the scenes, the :class:`.RemoteUploaderDownloader` uses :doc:`Apache Libc
     from composer.utils import LibcloudObjectStore
 
     remote_uploader_downloader = RemoteUploaderDownloader(
-        object_store_cls=LibcloudObjectStore,
-        object_store_kwargs={
+        bucket_uri="libcloud://my_bucket",
+        backend_kwargs={
             "provider": "s3",  # The Apache Libcloud provider name
             "container": "my_bucket",  # The name of the cloud container (i.e. bucket) to use.
             "provider_kwargs": {  # The Apache Libcloud provider driver initialization arguments
@@ -339,8 +339,8 @@ Once you've configured your object store logger per above, all that's left is to
     from composer.loggers import RemoteUploaderDownloader
 
     remote_uploader_downloader = RemoteUploaderDownloader(
-        object_store_cls=LibcloudObjectStore,
-        object_store_kwargs={
+        bucket_uri="libcloud://checkpoint-debugging",
+        backend_kwargs={
             "provider": "s3",  # The Apache Libcloud provider name
             "container": "checkpoint-debugging",  # The name of the cloud container (i.e. bucket) to use.
             "provider_kwargs": {  # The Apache Libcloud provider driver initialization arguments

diff --git a/docs/source/trainer/file_uploading.rst b/docs/source/trainer/file_uploading.rst
@@ -176,12 +176,7 @@ with the :class:`~composer.utils.object_store.s3_object_store.S3ObjectStore` bac
 
     # Configure the logger
     logger = RemoteUploaderDownloader(
-        object_store_cls=S3ObjectStore,
-        object_store_kwargs={
-            # Keyword arguments for the S3ObjectStore constructor.
-            # See the API reference for all available arguments
-            'bucket': 'my-bucket-name',
-        },
+        bucket_uri="s3://my-bucket-name",
     )
 
     # Define the trainer
@@ -212,12 +207,7 @@ Similar to the S3 Example above, we can upload files to a remote SFTP filesystem
 
     # Configure the logger
     logger = RemoteUploaderDownloader(
-        object_store_cls=SFTPObjectStore,
-        object_store_kwargs={
-            # Keyword arguments for the SFTPObjectStore constructor.
-            # See the API reference for all available arguments
-            'host': 'sftp_server.example.com',
-        },
+        bucket_uri="sftp://sftp_server.example.com",
     )
 
     # Define the trainer

diff --git a/examples/training_without_local_storage.ipynb b/examples/training_without_local_storage.ipynb
@@ -259,7 +259,6 @@
     "[wandb]: https://docs.mosaicml.com/en/stable/api_reference/generated/composer.loggers.WandBLogger.html\n",
     "[cometml]: https://docs.mosaicml.com/en/stable/api_reference/generated/composer.loggers.CometMLLogger.html\n",
     "[RemoteUploaderDownloader]: https://docs.mosaicml.com/en/stable/api_reference/generated/composer.loggers.RemoteUploaderDownloader.html#RemoteUploaderDownloader\n",
-    "[FileLogger]: https://docs.mosaicml.com/en/stable/api_reference/generated/composer.loggers.FileLogger.html#filelogger\n",
     "\n",
     "* Metrics: A metric is a scalar, such as accuracy, that can be logged. Usually you would want to plot metrics over time (e.g., to see how accuracy improves over batches).\n",
     "\n",
@@ -298,12 +297,9 @@
     "\n",
     "def get_remote_uploader_downloader():\n",
     "    return RemoteUploaderDownloader(\n",
-    "        object_store_cls=S3ObjectStore,\n",
-    "        # Keyword arguments passed to the S3ObjectStore constructor\n",
-    "        object_store_kwargs={\n",
-    "            'bucket': s3_bucket_name,\n",
-    "            'prefix': bucket_prefix,\n",
-    "        },\n",
+    "        bucket_uri=f\"s3://{s3_bucket_name}\",\n",
+    "        # This creates a format string for where to store the checkpoints within the S3 bucket\n",
+    "        path_format_string=bucket_prefix + \"/{remote_file_name}\"\n",
     "        # In Jupyter, we set use_procs to False, since subprocess do not work\n",
     "        # well within notebooks. Outside of Jupyter, it is recommended to let\n",
     "        # use_procs default to True for performance\n",
@@ -942,12 +938,12 @@
     ")\n",
     "\n",
     "cloud_logger = RemoteUploaderDownloader(\n",
-    "    object_store_cls=S3ObjectStore,\n",
-    "    # Keyword arguments passed to the S3ObjectStore constructor\n",
-    "    object_store_kwargs={\n",
-    "        'bucket': s3_bucket_name,\n",
-    "        'prefix': bucket_prefix,\n",
-    "    },\n",
+    "    bucket_uri=f\"s3://{s3_bucket_name}\",\n",
+    "    # This creates a format string for where to store the checkpoints within the S3 bucket\n",
+    "    path_format_string=bucket_prefix + \"/{remote_file_name}\"\n",
+    "    # In Jupyter, we set use_procs to False, since subprocess do not work\n",
+    "    # well within notebooks. Outside of Jupyter, it is recommended to let\n",
+    "    # use_procs default to True for performance\n",
     "    use_procs=False,\n",
     ")\n",
     "\n",

diff --git a/tests/callbacks/callback_settings.py b/tests/callbacks/callback_settings.py
@@ -19,7 +19,6 @@
 from composer.loggers.logger_destination import LoggerDestination
 from composer.loggers.logger_hparams_registry import RemoteUploaderDownloaderHparams, logger_registry
 from composer.loggers.progress_bar_logger import ProgressBarLogger
-from composer.utils.object_store.libcloud_object_store import LibcloudObjectStore
 from tests.common import get_module_subclasses
 
 try:
@@ -63,8 +62,8 @@
 
 _callback_kwargs: Dict[Union[Type[Callback], Type[hp.Hparams]], Dict[str, Any],] = {
     RemoteUploaderDownloader: {
-        'object_store_cls': LibcloudObjectStore,
-        'object_store_kwargs': {
+        'bucket_uri': 'libcloud://.',
+        'backend_kwargs': {
             'provider': 'local',
             'container': '.',
             'provider_kwargs': {
@@ -95,6 +94,7 @@
         'window_size': 1,
     },
     RemoteUploaderDownloaderHparams: {
+        'bucket_uri': 'libcloud://.',
         'object_store_hparams': {
             'libcloud': {
                 'provider': 'local',

diff --git a/tests/callbacks/test_loggers_across_callbacks.py b/tests/callbacks/test_loggers_across_callbacks.py
@@ -21,8 +21,8 @@ def test_loggers_on_callbacks(logger_cls: Type[LoggerDestination], callback_cls:
     logger_kwargs = get_cb_kwargs(logger_cls)
     if issubclass(logger_cls, RemoteUploaderDownloader):
         # Ensure that the remote directory does not conflict with any directory used by callbacks
-        logger_kwargs['object_store_kwargs']['provider_kwargs']['key'] = './remote'
-        os.makedirs(logger_kwargs['object_store_kwargs']['provider_kwargs']['key'], exist_ok=True)
+        logger_kwargs['backend_kwargs']['provider_kwargs']['key'] = './remote'
+        os.makedirs(logger_kwargs['backend_kwargs']['provider_kwargs']['key'], exist_ok=True)
     logger = logger_cls(**logger_kwargs)
     callback_kwargs = get_cb_kwargs(callback_cls)
     callback = callback_cls(**callback_kwargs)