Repository: horovod/horovod
Pull request: Refactored FilesystemStore to use fsspec to support additional remote filesystems
This adds support for various cloud storage systems like S3, GCS, ADLS, etc. in the Spark Estimators.

Fixes #2905.
Fixes #1732.

Problem statement:
S3Store on hovorod.spark.common.store
I'm interested in using Hovorod on our Spark cluster, but I'd like to use it with S3 rather than HDFS. 

Supporting ADLS storage
Currently, access to ADLS (Azure Data Lake filesystem) is not there. Is there any plans of adding support to ADLS in near term ?

Hints from issue discussion:
Hey @nanne-aben, thanks for the feature request.  I think this should be doable given that PyArrow now natively supports S3.  I can try putting together a quick PR for you to try out.
Hi @tgaddair, that'd be great! Thanks! Let me know when it's there, would be more than happy to look at the PR.
Hey @nanne-aben, quick update on this.  We're making some changes to our intermediate data storage API to make it more generic, so it will work with S3 if your Spark application is configured to use S3.  I'll let you know when these features are available, should be soon.  Thanks.
Thanks, Travis!

On Tue, 10 Mar 2020 at 17:04, Travis Addair <notifications@github.com>
wrote:

> Hey @nanne-aben <https://github.com/nanne-aben>, quick update on this.
> We're making some changes to our intermediate data storage API to make it
> more generic, so it will work with S3 if your Spark application is
> configured to use S3. I'll let you know when these features are available,
> should be soon. Thanks.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/horovod/horovod/issues/1732?email_source=notifications&email_token=ALOBCX7BYY7WUIVCSXCPJO3RGZQKRA5CNFSM4KZFKK4KYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEOMBJUI#issuecomment-597169361>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ALOBCX7MTUVS6WKNKCAOHNLRGZQKRANCNFSM4KZFKK4A>
> .
>

Any update on this feature request?
Hey @umashankark, Horovod itself is agnostic to where the data is coming from. One way to use it with ADLS is to use a framework like [Petastorm](https://github.com/uber/petastorm) to read Parquet data. I recently added support for reading from ADLS in v0.10.  Check it out, and let me know if you run into any issues.
Thanks for your suggestions, @tgaddair.  Appreciate some inputs from you on these:
1. In Horovod on Spark  scenario, with Petastorm, Horovod will not be using any intermediate storage - is this understanding correct ?
~~2. Regarding ADLS in v0.10, I don't see Store classes supporting 'abfs' URI. Can you please point me to details ?~~
3.  In Horovod on Spark scenario, it should be advisable to  use a distributed data store - is that right ?
TorchEstimator() seems to invoke Petastorm->fsspec->adlfs for accessing ADLS, for reading train data.
Is there anyway to pass authentication parameters (like account_name, storage_options, sas_token) to fsspec ?
Hey @umashankark, yes you're right the stores have not been updated to use Petastorm's new features. I'm putting together #2927 to address this. Feel free to check it out and let me know how it goes.

Commit message: Refactored LocalStore into FilesystemStore

Signed-off-by: Travis Addair <tgaddair@gmail.com>

Patch to review:
diff --git a/horovod/spark/common/store.py b/horovod/spark/common/store.py
index e5b893722c..73e5044f5f 100644
--- a/horovod/spark/common/store.py
+++ b/horovod/spark/common/store.py
@@ -16,6 +16,7 @@
 import contextlib
 import errno
 import os
+import pathlib
 import re
 import shutil
 import tempfile
@@ -26,6 +27,9 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 
+import fsspec
+from fsspec.core import split_protocol
+
 from horovod.spark.common.util import is_databricks
 
 
@@ -155,10 +159,10 @@ def create(prefix_path, *args, **kwargs):
         elif is_databricks() and DBFSLocalStore.matches_dbfs(prefix_path):
             return DBFSLocalStore(prefix_path, *args, **kwargs)
         else:
-            return LocalStore(prefix_path, *args, **kwargs)
+            return FilesystemStore(prefix_path, *args, **kwargs)
 
 
-class FilesystemStore(Store):
+class AbstractFilesystemStore(Store):
     """Abstract class for stores that use a filesystem for underlying storage."""
 
     def __init__(self, prefix_path, train_path=None, val_path=None, test_path=None, runs_path=None, save_runs=True):
@@ -168,13 +172,13 @@ def __init__(self, prefix_path, train_path=None, val_path=None, test_path=None,
         self._test_path = self._get_full_path_or_default(test_path, 'intermediate_test_data')
         self._runs_path = self._get_full_path_or_default(runs_path, 'runs')
         self._save_runs = save_runs
-        super(FilesystemStore, self).__init__()
+        super().__init__()
 
     def exists(self, path):
-        return self.get_filesystem().exists(self.get_localized_path(path))
+        return self.fs.exists(self.get_localized_path(path))
 
     def read(self, path):
-        with self.get_filesystem().open(self.get_localized_path(path), 'rb') as f:
+        with self.fs.open(self.get_localized_path(path), 'rb') as f:
             return f.read()
 
     def read_serialized_keras_model(self, ckpt_path, model, custom_objects):
@@ -193,7 +197,7 @@ def read_serialized_keras_model(self, ckpt_path, model, custom_objects):
         return codec.dumps_base64(model_bytes)
 
     def write_text(self, path, text):
-        with self.get_filesystem().open(self.get_localized_path(path), 'w') as f:
+        with self.fs.open(self.get_localized_path(path), 'w') as f:
             f.write(text)
 
     def is_parquet_dataset(self, path):
@@ -204,7 +208,7 @@ def is_parquet_dataset(self, path):
             return False
 
     def get_parquet_dataset(self, path):
-        return pq.ParquetDataset(self.get_localized_path(path), filesystem=self.get_filesystem())
+        return pq.ParquetDataset(self.get_localized_path(path), filesystem=self.fs)
 
     def get_train_data_path(self, idx=None):
         return '{}.{}'.format(self._train_path, idx) if idx is not None else self._train_path
@@ -237,7 +241,7 @@ def get_checkpoint_path(self, run_id):
 
     def get_checkpoints(self, run_id, suffix='.ckpt'):
         checkpoint_dir = self.get_localized_path(self.get_checkpoint_path(run_id))
-        filenames = self.get_filesystem().ls(checkpoint_dir)
+        filenames = self.fs.ls(checkpoint_dir)
         return sorted([name for name in filenames if name.endswith(suffix)])
 
     def get_logs_path(self, run_id):
@@ -250,23 +254,6 @@ def get_checkpoint_filename(self):
     def get_logs_subdir(self):
         return 'logs'
 
-    def get_full_path(self, path):
-        if not self.matches(path):
-            return self.path_prefix() + path
-        return path
-
-    def get_localized_path(self, path):
-        if self.matches(path):
-            return path[len(self.path_prefix()):]
-        return path
-
-    def get_full_path_fn(self):
-        prefix = self.path_prefix()
-
-        def get_path(path):
-            return prefix + path
-        return get_path
-
     def _get_full_path_or_default(self, path, default_key):
         if path is not None:
             return self.get_full_path(path)
@@ -275,66 +262,83 @@ def _get_full_path_or_default(self, path, default_key):
     def _get_path(self, key):
         return os.path.join(self.prefix_path, key)
 
-    def path_prefix(self):
-        raise NotImplementedError()
+    def get_local_output_dir_fn(self, run_id):
+        @contextlib.contextmanager
+        def local_run_path():
+            with tempfile.TemporaryDirectory() as tmpdir:
+                yield tmpdir
+        return local_run_path
 
-    def get_filesystem(self):
+    def get_localized_path(self, path):
         raise NotImplementedError()
 
-    @classmethod
-    def matches(cls, path):
-        return path.startswith(cls.filesystem_prefix())
+    def get_full_path(self, path):
+        raise NotImplementedError()
 
-    @classmethod
-    def filesystem_prefix(cls):
+    def get_full_path_fn(self):
         raise NotImplementedError()
 
+    @property
+    def fs(self):
+        raise NotImplementedError()
 
-class LocalStore(FilesystemStore):
-    """Uses the local filesystem as a store of intermediate data and training artifacts."""
 
-    FS_PREFIX = 'file://'
+class FilesystemStore(AbstractFilesystemStore):
+    """Concrete filesystems store that delegates to `fsspec`."""
 
     def __init__(self, prefix_path, *args, **kwargs):
-        self._fs = pa.LocalFileSystem()
-        super(LocalStore, self).__init__(prefix_path, *args, **kwargs)
+        self._fs, self.protocol = FilesystemStore._get_fs_and_protocol(prefix_path)
+        super().__init__(prefix_path, *args, **kwargs)
 
-    def path_prefix(self):
-        return self.FS_PREFIX
+    def sync_fn(self, run_id):
+        run_path = self.get_run_path(run_id)
 
-    def get_filesystem(self):
-        return self._fs
+        def fn(local_run_path):
+            self.fs.put(local_run_path, run_path, recursive=True)
 
-    def get_local_output_dir_fn(self, run_id):
-        run_path = self.get_localized_path(self.get_run_path(run_id))
+        return fn
 
-        @contextlib.contextmanager
-        def local_run_path():
-            if not os.path.exists(run_path):
-                try:
-                    os.makedirs(run_path, mode=0o755)
-                except OSError as e:
-                    # Race condition from workers on the same host: ignore
-                    if e.errno != errno.EEXIST:
-                        raise
-            yield run_path
+    def get_localized_path(self, path):
+        _, lpath = split_protocol(path)
+        return lpath
 
-        return local_run_path
+    def get_full_path(self, path):
+        return self.get_full_path_fn()(path)
 
-    def sync_fn(self, run_id):
-        run_path = self.get_localized_path(self.get_run_path(run_id))
+    def get_full_path_fn(self):
+        def get_path(path):
+            protocol, _ = split_protocol(path)
+            if protocol is not None:
+                return path
+            return pathlib.Path(os.path.abspath(path)).as_uri()
+        return get_path
 
-        def fn(local_run_path):
-            # No-op for LocalStore since the `local_run_path` will be the same as the run path
-            assert run_path == local_run_path
-        return fn
+    @property
+    def fs(self):
+        return self._fs
+
+    @staticmethod
+    def _get_fs_and_protocol(url):
+        protocol, path = split_protocol(url)
+        fs = fsspec.filesystem(protocol)
+        return fs, protocol
 
     @classmethod
-    def filesystem_prefix(cls):
-        return cls.FS_PREFIX
+    def matches(cls, path):
+        return True
+
+
+class LocalStore(FilesystemStore):
+    """Uses the local filesystem as a store of intermediate data and training artifacts.
+
+    This class is deprecated and now just resolves to FilesystemStore.
+    """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
-class HDFSStore(FilesystemStore):
+
+class HDFSStore(AbstractFilesystemStore):
     """Uses HDFS as a store of intermediate data and training artifacts.
 
     Initialized from a `prefix_path` that can take one of the following forms:
@@ -358,9 +362,7 @@ class HDFSStore(FilesystemStore):
 
     def __init__(self, prefix_path,
                  host=None, port=None, user=None, kerb_ticket=None,
-                 driver='libhdfs', extra_conf=None, temp_dir=None, *args, **kwargs):
-        self._temp_dir = temp_dir
-
+                 driver='libhdfs', extra_conf=None, *args, **kwargs):
         prefix, url_host, url_port, path, path_offset = self.parse_url(prefix_path)
         self._check_url(prefix_path, prefix, path)
         self._url_prefix = prefix_path[:path_offset] if prefix else self.FS_PREFIX
@@ -391,24 +393,21 @@ def parse_url(self, url):
         path_offset = match.start(4)
         return prefix, host, port, path, path_offset
 
-    def path_prefix(self):
-        return self._url_prefix
-
-    def get_filesystem(self):
-        return self._hdfs
+    def get_full_path(self, path):
+        if not self.matches(path):
+            return self._url_prefix + path
+        return path
 
-    def get_local_output_dir_fn(self, run_id):
-        temp_dir = self._temp_dir
+    def get_full_path_fn(self):
+        prefix = self._url_prefix
 
-        @contextlib.contextmanager
-        def local_run_path():
-            dirpath = tempfile.mkdtemp(dir=temp_dir)
-            try:
-                yield dirpath
-            finally:
-                shutil.rmtree(dirpath)
+        def get_path(path):
+            return prefix + path
+        return get_path
 
-        return local_run_path
+    @property
+    def fs(self):
+        return self._hdfs
 
     def sync_fn(self, run_id):
         class SyncState(object):
@@ -466,11 +465,11 @@ def _check_url(self, url, prefix, path):
             raise ValueError('Failed to parse path from URL: {}'.format(url))
 
     @classmethod
-    def filesystem_prefix(cls):
-        return cls.FS_PREFIX
+    def matches(cls, path):
+        return path.startswith(cls.FS_PREFIX)
 
 
-class DBFSLocalStore(LocalStore):
+class DBFSLocalStore(FilesystemStore):
     """Uses Databricks File System (DBFS) local file APIs as a store of intermediate data and
     training artifacts.
 
diff --git a/setup.py b/setup.py
index a3368c787d..43373e7045 100644
--- a/setup.py
+++ b/setup.py
@@ -112,7 +112,7 @@ def build_extensions(self):
 pyspark_require_list = ['pyspark>=2.3.2;python_version<"3.8"',
                         'pyspark>=3.0.0;python_version>="3.8"']
 # Pin h5py: https://github.com/h5py/h5py/issues/1732
-spark_require_list = ['h5py<3', 'numpy', 'petastorm>=0.11.0', 'pyarrow>=0.15.0']
+spark_require_list = ['h5py<3', 'numpy', 'petastorm>=0.11.0', 'pyarrow>=0.15.0', 'fsspec']
 ray_require_list = ['ray']
 pytorch_spark_require_list = pytorch_require_list + \
                              spark_require_list + \


Code review comments:
diff_hunk: @@ -275,66 +262,83 @@ def _get_full_path_or_default(self, path, default_key):
     def _get_path(self, key):
         return os.path.join(self.prefix_path, key)
 
-    def path_prefix(self):
-        raise NotImplementedError()
+    def get_local_output_dir_fn(self, run_id):
+        @contextlib.contextmanager
+        def local_run_path():
+            with tempfile.TemporaryDirectory() as tmpdir:
+                yield tmpdir
+        return local_run_path
 
-    def get_filesystem(self):
+    def get_localized_path(self, path):
         raise NotImplementedError()
 
-    @classmethod
-    def matches(cls, path):
-        return path.startswith(cls.filesystem_prefix())
+    def get_full_path(self, path):
+        raise NotImplementedError()
 
-    @classmethod
-    def filesystem_prefix(cls):
+    def get_full_path_fn(self):
         raise NotImplementedError()
 
+    @property
+    def fs(self):
+        raise NotImplementedError()
 
-class LocalStore(FilesystemStore):
-    """Uses the local filesystem as a store of intermediate data and training artifacts."""
 
-    FS_PREFIX = 'file://'
+class FilesystemStore(AbstractFilesystemStore):
+    """Concrete filesystems store that delegates to `fsspec`."""
 
     def __init__(self, prefix_path, *args, **kwargs):
-        self._fs = pa.LocalFileSystem()
-        super(LocalStore, self).__init__(prefix_path, *args, **kwargs)
+        self._fs, self.protocol = FilesystemStore._get_fs_and_protocol(prefix_path)
+        super().__init__(prefix_path, *args, **kwargs)
 
-    def path_prefix(self):
-        return self.FS_PREFIX
+    def sync_fn(self, run_id):
+        run_path = self.get_run_path(run_id)
 
-    def get_filesystem(self):
-        return self._fs
+        def fn(local_run_path):
+            self.fs.put(local_run_path, run_path, recursive=True)
 
-    def get_local_output_dir_fn(self, run_id):
-        run_path = self.get_localized_path(self.get_run_path(run_id))
+        return fn
 
-        @contextlib.contextmanager
-        def local_run_path():
-            if not os.path.exists(run_path):
-                try:
-                    os.makedirs(run_path, mode=0o755)
-                except OSError as e:
-                    # Race condition from workers on the same host: ignore
-                    if e.errno != errno.EEXIST:
-                        raise
-            yield run_path
+    def get_localized_path(self, path):
+        _, lpath = split_protocol(path)
+        return lpath
 
-        return local_run_path
+    def get_full_path(self, path):
+        return self.get_full_path_fn()(path)
 
-    def sync_fn(self, run_id):
-        run_path = self.get_localized_path(self.get_run_path(run_id))
+    def get_full_path_fn(self):
+        def get_path(path):
+            protocol, _ = split_protocol(path)
+            if protocol is not None:
+                return path
+            return pathlib.Path(os.path.abspath(path)).as_uri()
+        return get_path
 
-        def fn(local_run_path):
-            # No-op for LocalStore since the `local_run_path` will be the same as the run path
-            assert run_path == local_run_path
-        return fn
+    @property
+    def fs(self):
+        return self._fs
+
+    @staticmethod
+    def _get_fs_and_protocol(url):
+        protocol, path = split_protocol(url)
+        fs = fsspec.filesystem(protocol)
path: horovod/spark/common/store.py
text: @user1:
Can storage_options be provided as second argument [fsspec.filesystem(protocol, **storage_options)] ? This will help provide authentication options to fsspec.

diff_hunk: @@ -275,66 +262,83 @@ def _get_full_path_or_default(self, path, default_key):
     def _get_path(self, key):
         return os.path.join(self.prefix_path, key)
 
-    def path_prefix(self):
-        raise NotImplementedError()
+    def get_local_output_dir_fn(self, run_id):
+        @contextlib.contextmanager
+        def local_run_path():
+            with tempfile.TemporaryDirectory() as tmpdir:
+                yield tmpdir
+        return local_run_path
 
-    def get_filesystem(self):
+    def get_localized_path(self, path):
         raise NotImplementedError()
 
-    @classmethod
-    def matches(cls, path):
-        return path.startswith(cls.filesystem_prefix())
+    def get_full_path(self, path):
+        raise NotImplementedError()
 
-    @classmethod
-    def filesystem_prefix(cls):
+    def get_full_path_fn(self):
         raise NotImplementedError()
 
+    @property
+    def fs(self):
+        raise NotImplementedError()
 
-class LocalStore(FilesystemStore):
-    """Uses the local filesystem as a store of intermediate data and training artifacts."""
 
-    FS_PREFIX = 'file://'
+class FilesystemStore(AbstractFilesystemStore):
+    """Concrete filesystems store that delegates to `fsspec`."""
 
     def __init__(self, prefix_path, *args, **kwargs):
-        self._fs = pa.LocalFileSystem()
-        super(LocalStore, self).__init__(prefix_path, *args, **kwargs)
+        self._fs, self.protocol = FilesystemStore._get_fs_and_protocol(prefix_path)
path: horovod/spark/common/store.py
text: @user1:
Can this be added ?:
self.storage_options = kwargs['storage_options'] if 'storage_options' in kwargs else dict()
self._fs, self.protocol = FilesystemStore._get_fs_and_protocol(prefix_path, self.storage_options)