Support remote cache_dir (#4347)

* Support cache_dir when remote

* Cast cache_dir to str

* Force CI re-run

* Fix cast cache_dir to str

* Refactor data dirs preparation

* Make explicit when only local cache dir is considered

* Fix casting of cache_dir_root

* Fix casting of cache_downloaded_dir

* Use xjoin

* Revert "Use xjoin"

This reverts commit 43714db.

* Add suggested comments

Author: albertvillanova

src/datasets/builder.py

@@ -20,6 +20,7 @@
 import copy
 import inspect
 import os
+import posixpath
 import shutil
 import textwrap
 import urllib
@@ -286,25 +287,37 @@ def __init__(
         if features is not None:
             self.info.features = features
 
-        # prepare data dirs
-        self._cache_dir_root = os.path.expanduser(cache_dir or config.HF_DATASETS_CACHE)
+        # Prepare data dirs:
+        # cache_dir can be a remote bucket on GCS or S3 (when using BeamBasedBuilder for distributed data processing)
+        self._cache_dir_root = str(cache_dir or config.HF_DATASETS_CACHE)
+        self._cache_dir_root = (
+            self._cache_dir_root if is_remote_url(self._cache_dir_root) else os.path.expanduser(self._cache_dir_root)
+        )
+        path_join = posixpath.join if is_remote_url(self._cache_dir_root) else os.path.join
+        self._cache_downloaded_dir = (
+            path_join(self._cache_dir_root, config.DOWNLOADED_DATASETS_DIR)
+            if cache_dir
+            else str(config.DOWNLOADED_DATASETS_PATH)
+        )
         self._cache_downloaded_dir = (
-            os.path.join(cache_dir, config.DOWNLOADED_DATASETS_DIR) if cache_dir else config.DOWNLOADED_DATASETS_PATH
+            self._cache_downloaded_dir
+            if is_remote_url(self._cache_downloaded_dir)
+            else os.path.expanduser(self._cache_downloaded_dir)
         )
         self._cache_dir = self._build_cache_dir()
         if not is_remote_url(self._cache_dir_root):
             os.makedirs(self._cache_dir_root, exist_ok=True)
-        lock_path = os.path.join(self._cache_dir_root, self._cache_dir.replace(os.sep, "_") + ".lock")
-        with FileLock(lock_path):
-            if os.path.exists(self._cache_dir):  # check if data exist
-                if len(os.listdir(self._cache_dir)) > 0:
-                    logger.info("Overwrite dataset info from restored data version.")
-                    self.info = DatasetInfo.from_directory(self._cache_dir)
-                else:  # dir exists but no data, remove the empty dir as data aren't available anymore
-                    logger.warning(
-                        f"Old caching folder {self._cache_dir} for dataset {self.name} exists but not data were found. Removing it. "
-                    )
-                    os.rmdir(self._cache_dir)
+            lock_path = os.path.join(self._cache_dir_root, self._cache_dir.replace(os.sep, "_") + ".lock")
+            with FileLock(lock_path):
+                if os.path.exists(self._cache_dir):  # check if data exist
+                    if len(os.listdir(self._cache_dir)) > 0:
+                        logger.info("Overwrite dataset info from restored data version.")
+                        self.info = DatasetInfo.from_directory(self._cache_dir)
+                    else:  # dir exists but no data, remove the empty dir as data aren't available anymore
+                        logger.warning(
+                            f"Old caching folder {self._cache_dir} for dataset {self.name} exists but not data were found. Removing it. "
+                        )
+                        os.rmdir(self._cache_dir)
 
         # Set download manager
         self.dl_manager = None
@@ -439,7 +452,7 @@ def builder_configs(cls):
     def cache_dir(self):
         return self._cache_dir
 
-    def _relative_data_dir(self, with_version=True, with_hash=True) -> str:
+    def _relative_data_dir(self, with_version=True, with_hash=True, is_local=True) -> str:
         """Relative path of this dataset in cache_dir:
         Will be:
             self.name/self.config.version/self.hash/
@@ -451,19 +464,26 @@ def _relative_data_dir(self, with_version=True, with_hash=True) -> str:
         builder_data_dir = self.name if namespace is None else f"{namespace}___{self.name}"
         builder_config = self.config
         hash = self.hash
+        path_join = os.path.join if is_local else posixpath.join
         if builder_config:
             # use the enriched name instead of the name to make it unique
-            builder_data_dir = os.path.join(builder_data_dir, self.config_id)
+            builder_data_dir = path_join(builder_data_dir, self.config_id)
         if with_version:
-            builder_data_dir = os.path.join(builder_data_dir, str(self.config.version))
+            builder_data_dir = path_join(builder_data_dir, str(self.config.version))
         if with_hash and hash and isinstance(hash, str):
-            builder_data_dir = os.path.join(builder_data_dir, hash)
+            builder_data_dir = path_join(builder_data_dir, hash)
         return builder_data_dir
 
     def _build_cache_dir(self):
         """Return the data directory for the current version."""
-        builder_data_dir = os.path.join(self._cache_dir_root, self._relative_data_dir(with_version=False))
-        version_data_dir = os.path.join(self._cache_dir_root, self._relative_data_dir(with_version=True))
+        is_local = not is_remote_url(self._cache_dir_root)
+        path_join = os.path.join if is_local else posixpath.join
+        builder_data_dir = path_join(
+            self._cache_dir_root, self._relative_data_dir(with_version=False, is_local=is_local)
+        )
+        version_data_dir = path_join(
+            self._cache_dir_root, self._relative_data_dir(with_version=True, is_local=is_local)
+        )
 
         def _other_versions_on_disk():
             """Returns previous versions on disk."""
@@ -480,16 +500,17 @@ def _other_versions_on_disk():
             return version_dirnames
 
         # Check and warn if other versions exist on disk
-        version_dirs = _other_versions_on_disk()
-        if version_dirs:
-            other_version = version_dirs[0][0]
-            if other_version != self.config.version:
-                warn_msg = (
-                    f"Found a different version {str(other_version)} of dataset {self.name} in "
-                    f"cache_dir {self._cache_dir_root}. Using currently defined version "
-                    f"{str(self.config.version)}."
-                )
-                logger.warning(warn_msg)
+        if not is_remote_url(builder_data_dir):
+            version_dirs = _other_versions_on_disk()
+            if version_dirs:
+                other_version = version_dirs[0][0]
+                if other_version != self.config.version:
+                    warn_msg = (
+                        f"Found a different version {str(other_version)} of dataset {self.name} in "
+                        f"cache_dir {self._cache_dir_root}. Using currently defined version "
+                        f"{str(self.config.version)}."
+                    )
+                    logger.warning(warn_msg)
 
         return version_data_dir
 
@@ -571,18 +592,22 @@ def download_and_prepare(
         self.dl_manager = dl_manager
 
         # Prevent parallel disk operations
-        lock_path = os.path.join(self._cache_dir_root, self._cache_dir.replace(os.sep, "_") + ".lock")
-        with FileLock(lock_path):
-            data_exists = os.path.exists(self._cache_dir)
-            if data_exists and download_mode == DownloadMode.REUSE_DATASET_IF_EXISTS:
-                logger.warning(f"Reusing dataset {self.name} ({self._cache_dir})")
-                # We need to update the info in case some splits were added in the meantime
-                # for example when calling load_dataset from multiple workers.
-                self.info = self._load_info()
-                self.download_post_processing_resources(dl_manager)
-                return
+        is_local = not is_remote_url(self._cache_dir_root)
+        if is_local:
+            lock_path = os.path.join(self._cache_dir_root, self._cache_dir.replace(os.sep, "_") + ".lock")
+        # File locking only with local paths; no file locking on GCS or S3
+        with FileLock(lock_path) if is_local else contextlib.nullcontext():
+            if is_local:
+                data_exists = os.path.exists(self._cache_dir)
+                if data_exists and download_mode == DownloadMode.REUSE_DATASET_IF_EXISTS:
+                    logger.warning(f"Reusing dataset {self.name} ({self._cache_dir})")
+                    # We need to update the info in case some splits were added in the meantime
+                    # for example when calling load_dataset from multiple workers.
+                    self.info = self._load_info()
+                    self.download_post_processing_resources(dl_manager)
+                    return
             logger.info(f"Generating dataset {self.name} ({self._cache_dir})")
-            if not is_remote_url(self._cache_dir_root):  # if cache dir is local, check for available space
+            if is_local:  # if cache dir is local, check for available space
                 if not has_sufficient_disk_space(self.info.size_in_bytes or 0, directory=self._cache_dir_root):
                     raise OSError(
                         f"Not enough disk space. Needed: {size_str(self.info.size_in_bytes or 0)} (download: {size_str(self.info.download_size or 0)}, generated: {size_str(self.info.dataset_size or 0)}, post-processed: {size_str(self.info.post_processing_size or 0)})"