Repository: Lightning-AI/litdata

Issue:
Add support for GCS
## 🚀 Feature

<!-- A clear and concise description of the feature proposal -->

### Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
 
I am trying to write a dataset using the litdata optimize function. My plan was to write directly to a GCS bucket using the gs://<bucket_name> syntax (see code sample below).

```python
ld.optimize(
    fn=read_audio_files,  # the function applied to each input
    inputs=list(
        range(num_files)
    ),  # the inputs to the function (here it's a list of numbers)
    output_dir=os.path.join(
        f"gs://{DESTINATION_GCP_BUCKET}",
        storage_name.format(chunk_bytes=chunk_bytes, num_files=num_files),
    ),  # optimized data is stored here
    num_workers=num_workers,  # The number of workers on the same machine
    chunk_bytes=chunk_bytes,  # size of each chunk
)
```

The error is:

```python
    raise ValueError(f"The provided folder should start with s3://. Found {output_dir.url}.")
ValueError: The provided folder should start with s3://. Found gs://es-ml-bucket-eu-west4/fast_audio_512MB_1000.
```

Then, we tried using gcs fuse with a mounted bucket. This failed as well, with the following error below.
Basically, we could only upload the first shard of each process, and all other chunks couldn’t be written.

```python
  File "/home/mauro_luzzatto_epidemicsound_com/litdata-run/.venv/lib/python3.11/site-packages/litdata/processing/data_processor.py", line 824, in _done
    raise RuntimeError(f"All the chunks should have been deleted. Found {chunks} in cache: {cache_dir}")
RuntimeError: All the chunks should have been deleted. Found ['chunk-3-29.bin', 'chunk-1-28.bin', 'chunk-2-6.bin', 'chunk-25-19.bin', 'chunk-28-6.bin', 'chunk-20-1.bin', 'chunk-23-14.bin', 'chunk-6-29.bin', 'chunk-19-
...
'chunk-26-15.bin', 'chunk-20-23.bin', 'chunk-29-5.bin', 'chunk-17-1.bin', 'chunk-12-17.bin', 'chunk-4-0.bin', 'chunk-17-8.bin', 'chunk-15-19.bin', 'chunk-22-7.bin', 'chunk-10-19.bin', 'chunk-18-9.bin'] in cache: /var/tmp/chunks
```

### Pitch

<!-- A clear and concise description of what you want to happen. -->

### Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->


Discussion:

Patch:
diff --git a/src/litdata/constants.py b/src/litdata/constants.py
index 685e8a2..ada4a05 100644
--- a/src/litdata/constants.py
+++ b/src/litdata/constants.py
@@ -24,6 +24,7 @@ _DEFAULT_CHUNK_BYTES = 1 << 26  # 64M B
 _DEFAULT_FAST_DEV_RUN_ITEMS = 10
 _DEFAULT_CACHE_DIR = os.path.join(Path.home(), ".lightning", "chunks")
 _DEFAULT_LIGHTNING_CACHE_DIR = os.path.join("/cache", "chunks")
+_SUPPORTED_PROVIDERS = ("s3", "gs")  # cloud providers supported by litdata for uploading (optimize, map, merge, etc)
 
 # This is required for full pytree serialization / deserialization support
 _TORCH_GREATER_EQUAL_2_1_0 = RequirementCache("torch>=2.1.0")
diff --git a/src/litdata/processing/data_processor.py b/src/litdata/processing/data_processor.py
index 4268e7b..26965fb 100644
--- a/src/litdata/processing/data_processor.py
+++ b/src/litdata/processing/data_processor.py
@@ -32,8 +32,6 @@ from time import sleep, time
 from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union
 from urllib import parse
 
-import boto3
-import botocore
 import numpy as np
 import torch
 
@@ -42,14 +40,15 @@ from litdata.constants import (
     _ENABLE_STATUS,
     _INDEX_FILENAME,
     _IS_IN_STUDIO,
+    _SUPPORTED_PROVIDERS,
     _TQDM_AVAILABLE,
 )
 from litdata.processing.readers import BaseReader, StreamingDataLoaderReader
-from litdata.processing.utilities import _create_dataset, download_directory_from_S3, remove_uuid_from_filename
+from litdata.processing.utilities import _create_dataset, remove_uuid_from_filename
 from litdata.streaming import Cache
 from litdata.streaming.cache import Dir
-from litdata.streaming.client import S3Client
 from litdata.streaming.dataloader import StreamingDataLoader
+from litdata.streaming.fs_provider import _get_fs_provider, not_supported_provider
 from litdata.streaming.item_loader import BaseItemLoader
 from litdata.streaming.resolver import _resolve_dir
 from litdata.utilities._pytree import tree_flatten, tree_unflatten, treespec_loads
@@ -92,22 +91,18 @@ def _get_cache_data_dir(name: Optional[str] = None) -> str:
     """Returns the cache data directory used by the DataProcessor workers to download the files."""
     cache_dir = os.getenv("DATA_OPTIMIZER_DATA_CACHE_FOLDER", f"{_get_default_cache()}/data")
     if name is None:
-        #! TODO: Remove `join`. Why use join when only one path is provided?
-        return os.path.join(cache_dir)
+        return cache_dir
     return os.path.join(cache_dir, name.lstrip("/"))
 
 
-#! TODO: Move the checking logic to separate file/Class `StorageClient`.
-def _wait_for_file_to_exist(s3: S3Client, obj: parse.ParseResult, sleep_time: int = 2) -> Any:
-    """This function check."""
-    while True:
-        try:
-            return s3.client.head_object(Bucket=obj.netloc, Key=obj.path.lstrip("/"))
-        except botocore.exceptions.ClientError as e:
-            if "the HeadObject operation: Not Found" in str(e):
-                sleep(sleep_time)
-            else:
-                raise e
+def _wait_for_file_to_exist(remote_filepath: str, sleep_time: int = 2) -> Any:
+    """Wait until the file exists."""
+    file_exists = False
+    fs_provider = _get_fs_provider(remote_filepath)
+    while not file_exists:
+        file_exists = fs_provider.exists(remote_filepath)
+        if not file_exists:
+            sleep(sleep_time)
 
 
 def _wait_for_disk_usage_higher_than_threshold(input_dir: str, threshold_in_gb: int = 25, sleep_time: int = 3) -> None:
@@ -128,7 +123,7 @@ def _wait_for_disk_usage_higher_than_threshold(input_dir: str, threshold_in_gb:
 #
 def _download_data_target(input_dir: Dir, cache_dir: str, queue_in: Queue, queue_out: Queue) -> None:
     """Download data from a remote directory to a cache directory to optimise reading."""
-    s3 = S3Client()
+    fs_provider = None
 
     while True:
         # 2. Fetch from the queue
@@ -164,13 +159,12 @@ def _download_data_target(input_dir: Dir, cache_dir: str, queue_in: Queue, queue
 
                 obj = parse.urlparse(path)
 
-                if obj.scheme == "s3":
+                if obj.scheme in _SUPPORTED_PROVIDERS:
                     dirpath = os.path.dirname(local_path)
-
                     os.makedirs(dirpath, exist_ok=True)
-
-                    with open(local_path, "wb") as f:
-                        s3.client.download_fileobj(obj.netloc, obj.path.lstrip("/"), f)
+                    if fs_provider is None:
+                        fs_provider = _get_fs_provider(input_dir.url)
+                    fs_provider.download_file(path, local_path)
 
                 elif os.path.isfile(path):
                     if not path.startswith("/teamspace/studios/this_studio"):
@@ -231,8 +225,8 @@ def _upload_fn(upload_queue: Queue, remove_queue: Queue, cache_dir: str, output_
     """Upload optimised chunks from a local to remote dataset directory."""
     obj = parse.urlparse(output_dir.url if output_dir.url else output_dir.path)
 
-    if obj.scheme == "s3":
-        s3 = S3Client()
+    if obj.scheme in _SUPPORTED_PROVIDERS:
+        fs_provider = _get_fs_provider(output_dir.url)
 
     while True:
         data: Optional[Union[str, Tuple[str, str]]] = upload_queue.get()
@@ -252,9 +246,9 @@ def _upload_fn(upload_queue: Queue, remove_queue: Queue, cache_dir: str, output_
         if not local_filepath.startswith(cache_dir):
             local_filepath = os.path.join(cache_dir, local_filepath)
 
-        if obj.scheme == "s3":
+        if obj.scheme in _SUPPORTED_PROVIDERS:
             try:
-                output_filepath = str(obj.path).lstrip("/")
+                output_filepath = output_dir.url
 
                 if local_filepath.__contains__(".checkpoints"):
                     output_filepath = os.path.join(output_filepath, ".checkpoints")
@@ -265,9 +259,8 @@ def _upload_fn(upload_queue: Queue, remove_queue: Queue, cache_dir: str, output_
 
                 output_filepath = remove_uuid_from_filename(output_filepath)  # remove unique id from checkpoints
 
-                s3.client.upload_file(
+                fs_provider.upload_file(
                     local_filepath,
-                    obj.netloc,
                     output_filepath,
                 )
             except Exception as e:
@@ -944,10 +937,11 @@ class DataChunkRecipe(DataRecipe):
         else:
             local_filepath = os.path.join(cache_dir, _INDEX_FILENAME)
 
-        if obj.scheme == "s3":
-            s3 = S3Client()
-            s3.client.upload_file(
-                local_filepath, obj.netloc, os.path.join(str(obj.path).lstrip("/"), os.path.basename(local_filepath))
+        if obj.scheme in _SUPPORTED_PROVIDERS:
+            fs_provider = _get_fs_provider(output_dir.url)
+            fs_provider.upload_file(
+                local_filepath,
+                os.path.join(output_dir.url, os.path.basename(local_filepath)),
             )
         elif output_dir.path and os.path.isdir(output_dir.path):
             shutil.copyfile(local_filepath, os.path.join(output_dir.path, os.path.basename(local_filepath)))
@@ -965,11 +959,9 @@ class DataChunkRecipe(DataRecipe):
                 assert output_dir_path
                 remote_filepath = os.path.join(output_dir_path, f"{node_rank}-{_INDEX_FILENAME}")
                 node_index_filepath = os.path.join(cache_dir, os.path.basename(remote_filepath))
-                if obj.scheme == "s3":
-                    obj = parse.urlparse(remote_filepath)
-                    _wait_for_file_to_exist(s3, obj)
-                    with open(node_index_filepath, "wb") as f:
-                        s3.client.download_fileobj(obj.netloc, obj.path.lstrip("/"), f)
+                if obj.scheme in _SUPPORTED_PROVIDERS:
+                    fs_provider = _get_fs_provider(remote_filepath)
+                    fs_provider.download_file(remote_filepath, node_index_filepath)
                 elif output_dir.path and os.path.isdir(output_dir.path):
                     shutil.copyfile(remote_filepath, node_index_filepath)
 
@@ -1342,22 +1334,14 @@ class DataProcessor:
             return
 
         obj = parse.urlparse(self.output_dir.url)
+        if obj.scheme not in _SUPPORTED_PROVIDERS:
+            not_supported_provider(self.output_dir.url)
 
-        if obj.scheme != "s3":
-            raise ValueError(f"The provided folder should start with s3://. Found {self.output_dir.path}.")
-
-        s3 = boto3.client("s3")
-
-        prefix = obj.path.lstrip("/").rstrip("/") + "/"
-
-        # Delete all the files (including the index file in overwrite mode)
-        bucket_name = obj.netloc
-        s3 = boto3.resource("s3")
-
+        prefix = self.output_dir.url.rstrip("/") + "/"
         checkpoint_prefix = os.path.join(prefix, ".checkpoints")
 
-        for obj in s3.Bucket(bucket_name).objects.filter(Prefix=checkpoint_prefix):
-            s3.Object(bucket_name, obj.key).delete()
+        fs_provider = _get_fs_provider(self.output_dir.url)
+        fs_provider.delete_file_or_directory(checkpoint_prefix)
 
     def _save_current_config(self, workers_user_items: List[List[Any]]) -> None:
         if not self.use_checkpoint:
@@ -1383,23 +1367,20 @@ class DataProcessor:
 
             obj = parse.urlparse(self.output_dir.url)
 
-            if obj.scheme != "s3":
-                raise ValueError(f"The provided folder should start with s3://. Found {self.output_dir.path}.")
-
-            # TODO: Add support for all cloud providers
+            if obj.scheme not in _SUPPORTED_PROVIDERS:
+                not_supported_provider(self.output_dir.url)
 
-            s3 = S3Client()
+            fs_provider = _get_fs_provider(self.output_dir.url)
 
-            prefix = obj.path.lstrip("/").rstrip("/") + "/" + ".checkpoints/"
+            prefix = self.output_dir.url.rstrip("/") + "/" + ".checkpoints/"
 
-            # write config.json file to temp directory and upload it to s3
+            # write config.json file to temp directory and upload it to the cloud provider
             with tempfile.TemporaryDirectory() as temp_dir:
                 temp_file_name = os.path.join(temp_dir, "config.json")
                 with open(temp_file_name, "w") as f:
                     json.dump(config, f)
-                s3.client.upload_file(
+                fs_provider.upload_file(
                     temp_file_name,
-                    obj.netloc,
                     os.path.join(prefix, "config.json"),
                 )
         except Exception as e:
@@ -1451,19 +1432,17 @@ class DataProcessor:
 
         obj = parse.urlparse(self.output_dir.url)
 
-        if obj.scheme != "s3":
-            raise ValueError(f"The provided folder should start with s3://. Found {self.output_dir.path}.")
-
-        # TODO: Add support for all cloud providers
+        if obj.scheme not in _SUPPORTED_PROVIDERS:
+            not_supported_provider(self.output_dir.url)
 
-        prefix = obj.path.lstrip("/").rstrip("/") + "/" + ".checkpoints/"
+        prefix = self.output_dir.url.rstrip("/") + "/" + ".checkpoints/"
 
         # Delete all the files (including the index file in overwrite mode)
-        bucket_name = obj.netloc
 
         # download all the checkpoint files in tempdir and read them
         with tempfile.TemporaryDirectory() as temp_dir:
-            saved_file_dir = download_directory_from_S3(bucket_name, prefix, temp_dir)
+            fs_provider = _get_fs_provider(self.output_dir.url)
+            saved_file_dir = fs_provider.download_directory(prefix, temp_dir)
 
             if not os.path.exists(os.path.join(saved_file_dir, "config.json")):
                 # if the config.json file doesn't exist, we don't have any checkpoint saved
diff --git a/src/litdata/processing/functions.py b/src/litdata/processing/functions.py
index 9776519..f2f9dd1 100644
--- a/src/litdata/processing/functions.py
+++ b/src/litdata/processing/functions.py
@@ -28,7 +28,7 @@ from urllib import parse
 import torch
 
 from litdata import __version__
-from litdata.constants import _INDEX_FILENAME, _IS_IN_STUDIO
+from litdata.constants import _INDEX_FILENAME, _IS_IN_STUDIO, _SUPPORTED_PROVIDERS
 from litdata.helpers import _check_version_and_prompt_upgrade
 from litdata.processing.data_processor import DataChunkRecipe, DataProcessor, MapRecipe
 from litdata.processing.readers import BaseReader
@@ -38,8 +38,8 @@ from litdata.processing.utilities import (
     optimize_dns_context,
     read_index_file_content,
 )
-from litdata.streaming.client import S3Client
 from litdata.streaming.dataloader import StreamingDataLoader
+from litdata.streaming.fs_provider import _get_fs_provider
 from litdata.streaming.item_loader import BaseItemLoader
 from litdata.streaming.resolver import (
     Dir,
@@ -58,7 +58,7 @@ if TYPE_CHECKING:
 
 def _is_remote_file(path: str) -> bool:
     obj = parse.urlparse(path)
-    return obj.scheme in ["s3", "gcs"]
+    return obj.scheme in _SUPPORTED_PROVIDERS
 
 
 def _get_indexed_paths(data: Any) -> Dict[int, str]:
@@ -641,15 +641,11 @@ def _apply_copy(copy_info: CopyInfo, output_dir: Dir) -> None:
         shutil.copyfile(input_filepath, output_filepath)
 
     elif output_dir.url and copy_info.input_dir.url:
-        input_obj = parse.urlparse(os.path.join(copy_info.input_dir.url, copy_info.old_filename))
-        output_obj = parse.urlparse(os.path.join(output_dir.url, copy_info.new_filename))
-
-        s3 = S3Client()
-        s3.client.copy(
-            {"Bucket": input_obj.netloc, "Key": input_obj.path.lstrip("/")},
-            output_obj.netloc,
-            output_obj.path.lstrip("/"),
-        )
+        input_filepath = os.path.join(copy_info.input_dir.url, copy_info.old_filename)
+        output_filepath = os.path.join(output_dir.url, copy_info.new_filename)
+
+        fs_provider = _get_fs_provider(output_dir.url)
+        fs_provider.copy(input_filepath, output_filepath)
     else:
         raise NotImplementedError
 
@@ -665,11 +661,7 @@ def _save_index(index_json: Dict, output_dir: Dir) -> None:
 
             f.flush()
 
-            obj = parse.urlparse(os.path.join(output_dir.url, _INDEX_FILENAME))
+            remote_path = os.path.join(output_dir.url, _INDEX_FILENAME)
 
-            s3 = S3Client()
-            s3.client.upload_file(
-                f.name,
-                obj.netloc,
-                obj.path.lstrip("/"),
-            )
+            fs_provider = _get_fs_provider(output_dir.url)
+            fs_provider.upload_file(f.name, remote_path)
diff --git a/src/litdata/processing/utilities.py b/src/litdata/processing/utilities.py
index 866d460..66dc5e1 100644
--- a/src/litdata/processing/utilities.py
+++ b/src/litdata/processing/utilities.py
@@ -21,11 +21,9 @@ from subprocess import DEVNULL, Popen
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from urllib import parse
 
-import boto3
-import botocore
-
-from litdata.constants import _INDEX_FILENAME, _IS_IN_STUDIO
+from litdata.constants import _INDEX_FILENAME, _IS_IN_STUDIO, _SUPPORTED_PROVIDERS
 from litdata.streaming.cache import Dir
+from litdata.streaming.fs_provider import _get_fs_provider, not_supported_provider
 
 
 #! TODO: Not sure what this function is used for.
@@ -132,7 +130,6 @@ def make_request(
         return io.BytesIO(r.read())
 
 
-#! TODO: Why `enable` parameter is needed? If you're using this context, you will always enable it (True).
 @contextmanager
 def optimize_dns_context(enable: bool) -> Any:
     """Optimize the DNS resolution for the Lightning Studio machine.
@@ -203,7 +200,6 @@ def _get_work_dir() -> str:
     return f"s3://{bucket_name}/projects/{project_id}/lightningapps/{app_id}/artifacts/{work_id}/content/"
 
 
-#! TODO: Move this to a separate file/class (StorageClient) and add support for other storage providers
 def read_index_file_content(output_dir: Dir) -> Optional[Dict[str, Any]]:
     """Read the index file content."""
     if not isinstance(output_dir, Dir):
@@ -219,30 +215,29 @@ def read_index_file_content(output_dir: Dir) -> Optional[Dict[str, Any]]:
             return json.load(f)
 
     else:
-        # download the index file from s3, and read it
+        # download the index file from the cloud provider, and read it
         obj = parse.urlparse(output_dir.url)
 
-        if obj.scheme != "s3":
-            raise ValueError(f"The provided folder should start with s3://. Found {output_dir.path}.")
+        if obj.scheme not in _SUPPORTED_PROVIDERS:
+            not_supported_provider(output_dir.url)
 
-        # TODO: Add support for all cloud providers
-        s3 = boto3.client("s3")
+        fs_provider = _get_fs_provider(output_dir.url)
 
-        prefix = obj.path.lstrip("/").rstrip("/") + "/"
+        prefix = output_dir.url.rstrip("/") + "/"
 
         # Check the index file exists
         try:
             # Create a temporary file
             with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as temp_file:
                 temp_file_name = temp_file.name
-                s3.download_file(obj.netloc, os.path.join(prefix, _INDEX_FILENAME), temp_file_name)
+                fs_provider.download_file(os.path.join(prefix, _INDEX_FILENAME), temp_file_name)
             # Read data from the temporary file
             with open(temp_file_name) as temp_file:
                 data = json.load(temp_file)
             # Delete the temporary file
             os.remove(temp_file_name)
             return data
-        except botocore.exceptions.ClientError:
+        except Exception:
             return None
 
 
@@ -277,22 +272,3 @@ def remove_uuid_from_filename(filepath: str) -> str:
 
     # uuid is of 32 characters, '.json' is 5 characters and '-' is 1 character
     return filepath[:-38] + ".json"
-
-
-#! TODO: Move this to a separate file/class (StorageClient)
-def download_directory_from_S3(bucket_name: str, remote_directory_name: str, local_directory_name: str) -> str:
-    s3_resource = boto3.resource("s3")
-    bucket = s3_resource.Bucket(bucket_name)
-
-    saved_file_dir = "."
-
-    for obj in bucket.objects.filter(Prefix=remote_directory_name):
-        local_filename = os.path.join(local_directory_name, obj.key)
-
-        if not os.path.exists(os.path.dirname(local_filename)):
-            os.makedirs(os.path.dirname(local_filename))
-        with open(local_filename, "wb") as f:
-            s3_resource.meta.client.download_fileobj(bucket_name, obj.key, f)
-            saved_file_dir = os.path.dirname(local_filename)
-
-    return saved_file_dir
diff --git a/src/litdata/streaming/fs_provider.py b/src/litdata/streaming/fs_provider.py
new file mode 100644
index 0000000..d7dbe31
--- /dev/null
+++ b/src/litdata/streaming/fs_provider.py
@@ -0,0 +1,268 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+from urllib import parse
+
+from litdata.constants import _GOOGLE_STORAGE_AVAILABLE, _SUPPORTED_PROVIDERS
+from litdata.streaming.client import S3Client
+
+
+class FsProvider(ABC):
+    def __init__(self, storage_options: Optional[Dict[str, Any]] = {}):
+        self.storage_options = storage_options
+
+    @abstractmethod
+    def upload_file(self, local_path: str, remote_path: str) -> None:
+        raise NotImplementedError
+
+    def download_file(self, remote_path: str, local_path: str) -> None:
+        raise NotImplementedError
+
+    def download_directory(self, remote_path: str, local_directory_name: str) -> str:
+        raise NotImplementedError
+
+    def copy(self, remote_source: str, remote_destination: str) -> None:
+        raise NotImplementedError
+
+    def list_directory(self, path: str) -> List[str]:
+        raise NotImplementedError
+
+    def delete_file_or_directory(self, path: str) -> None:
+        raise NotImplementedError
+
+    def exists(self, path: str) -> bool:
+        raise NotImplementedError
+
+    def is_empty(self, path: str) -> bool:
+        raise NotImplementedError
+
+
+class GCPFsProvider(FsProvider):
+    def __init__(self, storage_options: Optional[Dict[str, Any]] = {}):
+        if not _GOOGLE_STORAGE_AVAILABLE:
+            raise ModuleNotFoundError(str(_GOOGLE_STORAGE_AVAILABLE))
+        from google.cloud import storage
+
+        super().__init__(storage_options=storage_options)
+        self.client = storage.Client(**self.storage_options)
+
+    def upload_file(self, local_path: str, remote_path: str) -> None:
+        bucket_name, blob_path = get_bucket_and_path(remote_path, "gs")
+        bucket = self.client.bucket(bucket_name)
+        blob = bucket.blob(blob_path)
+        blob.upload_from_filename(local_path)
+
+    def download_file(self, remote_path: str, local_path: str) -> None:
+        bucket_name, blob_path = get_bucket_and_path(remote_path, "gs")
+        bucket = self.client.bucket(bucket_name)
+        blob = bucket.blob(blob_path)
+        blob.download_to_filename(local_path)
+
+    def download_directory(self, remote_path: str, local_directory_name: str) -> str:
+        bucket_name, blob_path = get_bucket_and_path(remote_path, "gs")
+        bucket = self.client.get_bucket(bucket_name)
+        blobs = bucket.list_blobs(prefix=blob_path)  # Get list of files
+        local_directory_name = os.path.abspath(local_directory_name)
+        os.makedirs(local_directory_name, exist_ok=True)
+        saved_file_dir = "."
+
+        for blob in blobs:
+            if blob.name.endswith("/") or blob.name == blob_path:  # Skip directories
+                continue
+            file_split = blob.name.split("/")
+            local_filename = os.path.join(local_directory_name, *file_split[1:])
+            os.makedirs(os.path.dirname(local_filename), exist_ok=True)  # Create local directory
+            blob.download_to_filename(local_filename)  # Download to the correct local path
+            saved_file_dir = os.path.dirname(local_filename)
+
+        return saved_file_dir
+
+    def list_directory(self, path: str) -> List[str]:
+        raise NotImplementedError
+
+    def copy(self, remote_source: str, remote_destination: str) -> None:
+        # WARNING: you need to pass complete path (+file_name.ext), else it will fail silently.
+        source_bucket_name, source_blob_path = get_bucket_and_path(remote_source, "gs")
+        destination_bucket_name, destination_blob_path = get_bucket_and_path(remote_destination, "gs")
+
+        source_bucket = self.client.bucket(source_bucket_name)
+        destination_bucket = self.client.bucket(destination_bucket_name)
+
+        source_blob = source_bucket.blob(source_blob_path)
+
+        # Check if the source blob exists
+        if not source_blob.exists():
+            raise FileNotFoundError(f"Source blob {source_blob_path} not found.")
+
+        # Use the correct `copy_blob` method
+        source_bucket.copy_blob(source_blob, destination_bucket, destination_blob_path)
+
+    def delete_file_or_directory(self, path: str) -> None:
+        bucket_name, blob_path = get_bucket_and_path(path, "gs")
+        bucket = self.client.bucket(bucket_name)
+        # if it's a single file, only one will match the prefix
+        blobs = bucket.list_blobs(prefix=blob_path)
+        for blob in blobs:
+            blob.delete()
+
+    def exists(self, path: str) -> bool:
+        bucket_name, blob_path = get_bucket_and_path(path, "gs")
+        bucket = self.client.bucket(bucket_name)
+        blob = bucket.blob(blob_path)
+        return blob.exists()
+
+    def is_empty(self, path: str) -> bool:
+        bucket_name, blob_path = get_bucket_and_path(path, "gs")
+        bucket = self.client.bucket(bucket_name)
+        # List blobs with the given prefix
+        blobs = bucket.list_blobs(prefix=blob_path)
+        # If no blobs are found, it's considered empty
+        return not any(blobs)
+
+
+class S3FsProvider(FsProvider):
+    def __init__(self, storage_options: Optional[Dict[str, Any]] = {}):
+        super().__init__(storage_options=storage_options)
+        self.client = S3Client(storage_options=storage_options)
+
+    def upload_file(self, local_path: str, remote_path: str) -> None:
+        bucket_name, blob_path = get_bucket_and_path(remote_path, "s3")
+        self.client.client.upload_file(local_path, bucket_name, blob_path)
+
+    def download_file(self, remote_path: str, local_path: str) -> None:
+        bucket_name, blob_path = get_bucket_and_path(remote_path, "s3")
+        with open(local_path, "wb") as f:
+            self.client.client.download_fileobj(bucket_name, blob_path, f)
+
+    def download_directory(self, remote_path: str, local_directory_name: str) -> str:
+        """Download all objects under a given S3 prefix (directory) using the existing client."""
+        bucket_name, remote_directory_name = get_bucket_and_path(remote_path, "s3")
+
+        # Ensure local directory exists
+        local_directory_name = os.path.abspath(local_directory_name)
+        os.makedirs(local_directory_name, exist_ok=True)
+
+        saved_file_dir = "."
+
+        # List objects under the given prefix
+        objects = self.client.client.list_objects_v2(Bucket=bucket_name, Prefix=remote_directory_name)
+
+        # Check if objects exist
+        if "Contents" in objects:
+            for obj in objects["Contents"]:
+                local_filename = os.path.join(local_directory_name, obj["Key"])
+
+                # Ensure parent directories exist
+                os.makedirs(os.path.dirname(local_filename), exist_ok=True)
+
+                # Download each file
+                with open(local_filename, "wb") as f:
+                    self.client.client.download_fileobj(bucket_name, obj["Key"], f)
+                    saved_file_dir = os.path.dirname(local_filename)
+
+        return saved_file_dir
+
+    def copy(self, remote_source: str, remote_destination: str) -> None:
+        input_obj = parse.urlparse(remote_source)
+        output_obj = parse.urlparse(remote_destination)
+        self.client.client.copy(
+            {"Bucket": input_obj.netloc, "Key": input_obj.path.lstrip("/")},
+            output_obj.netloc,
+            output_obj.path.lstrip("/"),
+        )
+
+    def list_directory(self, path: str) -> List[str]:
+        raise NotImplementedError
+
+    def delete_file_or_directory(self, path: str) -> None:
+        """Delete the file or the directory."""
+        bucket_name, blob_path = get_bucket_and_path(path, "s3")
+
+        # List objects under the given path
+        objects = self.client.client.list_objects_v2(Bucket=bucket_name, Prefix=blob_path)
+
+        # Check if objects exist
+        if "Contents" in objects:
+            for obj in objects["Contents"]:
+                self.client.client.delete_object(Bucket=bucket_name, Key=obj["Key"])
+
+    def exists(self, path: str) -> bool:
+        import botocore
+
+        bucket_name, blob_path = get_bucket_and_path(path, "s3")
+        try:
+            _ = self.client.client.head_object(Bucket=bucket_name, Key=blob_path)
+            return True
+        except botocore.exceptions.ClientError as e:
+            if "the HeadObject operation: Not Found" in str(e):
+                return False
+            raise e
+        except Exception as e:
+            raise e
+
+    def is_empty(self, path: str) -> bool:
+        obj = parse.urlparse(path)
+
+        objects = self.client.list_objects_v2(
+            Bucket=obj.netloc,
+            Delimiter="/",
+            Prefix=obj.path.lstrip("/").rstrip("/") + "/",
+        )
+
+        return not objects["KeyCount"] > 0
+
+
+def get_bucket_and_path(remote_filepath: str, expected_scheme: str = "s3") -> Tuple[str, str]:
+    """Parse the remote filepath and return the bucket name and the blob path.
+
+    Args:
+        remote_filepath (str): The remote filepath to parse.
+        expected_scheme (str, optional): The expected scheme of the remote filepath. Defaults to "s3".
+
+    Raises:
+        ValueError: If the scheme of the remote filepath is not as expected.
+
+    Returns:
+        Tuple[str, str]: The bucket name and the blob_path.
+    """
+    obj = parse.urlparse(remote_filepath)
+
+    if obj.scheme != expected_scheme:
+        raise ValueError(
+            f"Expected obj.scheme to be `{expected_scheme}`, instead, got {obj.scheme} for remote={remote_filepath}."
+        )
+
+    bucket_name = obj.netloc
+    blob_path = obj.path
+    # Remove the leading "/":
+    if blob_path[0] == "/":
+        blob_path = blob_path[1:]
+
+    return bucket_name, blob_path
+
+
+def _get_fs_provider(remote_filepath: str, storage_options: Optional[Dict[str, Any]] = {}) -> FsProvider:
+    obj = parse.urlparse(remote_filepath)
+    if obj.scheme == "gs":
+        return GCPFsProvider(storage_options=storage_options)
+    if obj.scheme == "s3":
+        return S3FsProvider(storage_options=storage_options)
+    raise ValueError(f"Unsupported scheme: {obj.scheme}")
+
+
+def not_supported_provider(remote_filepath: str) -> bool:
+    raise ValueError(
+        f"URL should start with one of {[el + '://' for el in _SUPPORTED_PROVIDERS]}." f"Found {remote_filepath}."
+    )
diff --git a/src/litdata/streaming/resolver.py b/src/litdata/streaming/resolver.py
index 238cd8f..9bdf238 100644
--- a/src/litdata/streaming/resolver.py
+++ b/src/litdata/streaming/resolver.py
@@ -23,10 +23,8 @@ from time import sleep
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 from urllib import parse
 
-import boto3
-import botocore
-
-from litdata.constants import _LIGHTNING_SDK_AVAILABLE
+from litdata.constants import _LIGHTNING_SDK_AVAILABLE, _SUPPORTED_PROVIDERS
+from litdata.streaming.fs_provider import _get_fs_provider, not_supported_provider
 
 if TYPE_CHECKING:
     from lightning_sdk import Machine
@@ -240,20 +238,16 @@ def _assert_dir_is_empty(output_dir: Dir, append: bool = False, overwrite: bool
 
     obj = parse.urlparse(output_dir.url)
 
-    if obj.scheme != "s3":
-        raise ValueError(f"The provided folder should start with s3://. Found {output_dir.url}.")
+    if obj.scheme not in _SUPPORTED_PROVIDERS:
+        not_supported_provider(output_dir.url)
 
-    s3 = boto3.client("s3")
+    fs_provider = _get_fs_provider(output_dir.url)
 
-    objects = s3.list_objects_v2(
-        Bucket=obj.netloc,
-        Delimiter="/",
-        Prefix=obj.path.lstrip("/").rstrip("/") + "/",
-    )
+    is_empty = fs_provider.is_empty(output_dir.url)
 
-    # We aren't alloweing to add more data
+    # We aren't allowing to add more data
     # TODO: Add support for `append` and `overwrite`.
-    if objects["KeyCount"] > 0:
+    if not is_empty:
         raise RuntimeError(
             f"The provided output_dir `{output_dir.path}` already contains data and datasets are meant to be immutable."
             "\n HINT: Did you consider changing the `output_dir` with your own versioning as a suffix?"
@@ -305,28 +299,24 @@ def _assert_dir_has_index_file(
 
     obj = parse.urlparse(output_dir.url)
 
-    if obj.scheme != "s3":
-        raise ValueError(f"The provided folder should start with s3://. Found {output_dir.url}.")
+    if obj.scheme not in _SUPPORTED_PROVIDERS:
+        not_supported_provider(output_dir.url)
 
-    s3 = boto3.client("s3")
+    fs_provider = _get_fs_provider(output_dir.url)
 
-    prefix = obj.path.lstrip("/").rstrip("/") + "/"
+    prefix = output_dir.url.rstrip("/") + "/"
 
-    objects = s3.list_objects_v2(
-        Bucket=obj.netloc,
-        Delimiter="/",
-        Prefix=prefix,
-    )
+    is_empty = fs_provider.is_empty(output_dir.url)
 
     # No files are found in this folder
-    if objects["KeyCount"] == 0:
+    if is_empty:
         return
 
     # Check the index file exists
     try:
-        s3.head_object(Bucket=obj.netloc, Key=os.path.join(prefix, "index.json"))
+        fs_provider.exists(os.path.join(prefix, "index.json"))
         has_index_file = True
-    except botocore.exceptions.ClientError:
+    except Exception:
         has_index_file = False
 
     if has_index_file and mode is None:
@@ -336,13 +326,11 @@ def _assert_dir_has_index_file(
             "\n HINT: If you want to append/overwrite to the existing dataset, use `mode='append | overwrite'`."
         )
 
-    # Delete all the files (including the index file in overwrite mode)
-    bucket_name = obj.netloc
-    s3 = boto3.resource("s3")
+    # all the files (including the index file in overwrite mode)
+    fs_provider = _get_fs_provider(output_dir.url)
 
     if mode == "overwrite" or (mode is None and not use_checkpoint):
-        for obj in s3.Bucket(bucket_name).objects.filter(Prefix=prefix):
-            s3.Object(bucket_name, obj.key).delete()
+        fs_provider.delete_file_or_directory(output_dir.url)
 
 
 def _get_lightning_cloud_url() -> str:


Test patch:
diff --git a/tests/processing/test_data_processor.py b/tests/processing/test_data_processor.py
index 059e2f9..aaa0b36 100644
--- a/tests/processing/test_data_processor.py
+++ b/tests/processing/test_data_processor.py
@@ -110,7 +110,7 @@ def test_upload_s3_fn(tmpdir, monkeypatch):
 
     remove_queue = mock.MagicMock()
 
-    s3_client = mock.MagicMock()
+    fs_provider = mock.MagicMock()
 
     called = False
 
@@ -121,9 +121,9 @@ def test_upload_s3_fn(tmpdir, monkeypatch):
 
         copyfile(local_filepath, os.path.join(remote_output_dir, os.path.basename(local_filepath)))
 
-    s3_client.client.upload_file = copy_file
+    fs_provider.upload_file = copy_file
 
-    monkeypatch.setattr(data_processor_module, "S3Client", mock.MagicMock(return_value=s3_client))
+    monkeypatch.setattr(data_processor_module, "_get_fs_provider", mock.MagicMock(return_value=fs_provider))
 
     assert os.listdir(remote_output_dir) == []
 
@@ -218,32 +218,29 @@ def test_wait_for_disk_usage_higher_than_threshold():
 
 
 @pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows")
-def test_wait_for_file_to_exist():
-    import botocore
-
-    s3 = mock.MagicMock()
-    obj = mock.MagicMock()
-    raise_error = [True, True, False]
+def test_wait_for_file_to_exist(monkeypatch):
+    url = "s3://url"
+    fs_provider = mock.MagicMock()
+    raise_error = [False, False, True]
 
     def fn(*_, **__):
-        value = raise_error.pop(0)
-        if value:
-            raise botocore.exceptions.ClientError({"Error": {"Code": "404", "Message": "Not Found"}}, "HeadObject")
-        return
+        return raise_error.pop(0)
 
-    s3.client.head_object = fn
+    fs_provider.exists = fn
+    monkeypatch.setattr(data_processor_module, "_get_fs_provider", mock.MagicMock(return_value=fs_provider))
 
-    _wait_for_file_to_exist(s3, obj, sleep_time=0.01)
+    _wait_for_file_to_exist(url, sleep_time=0.01)
 
     assert len(raise_error) == 0
 
     def fn(*_, **__):
         raise ValueError("HERE")
 
-    s3.client.head_object = fn
+    fs_provider.exists = fn
+    monkeypatch.setattr(data_processor_module, "_get_fs_provider", mock.MagicMock(return_value=fs_provider))
 
     with pytest.raises(ValueError, match="HERE"):
-        _wait_for_file_to_exist(s3, obj, sleep_time=0.01)
+        _wait_for_file_to_exist(url, sleep_time=0.01)
 
 
 def test_cache_dir_cleanup(tmpdir, monkeypatch):
@@ -1030,11 +1027,9 @@ def test_data_processing_map_non_absolute_path(monkeypatch, tmpdir):
 
 @pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows")
 def test_map_error_when_not_empty(monkeypatch):
-    boto3 = mock.MagicMock()
-    client_s3_mock = mock.MagicMock()
-    client_s3_mock.list_objects_v2.return_value = {"KeyCount": 1, "Contents": []}
-    boto3.client.return_value = client_s3_mock
-    monkeypatch.setattr(resolver, "boto3", boto3)
+    fs_provider = mock.MagicMock()
+    fs_provider.is_empty = mock.MagicMock(return_value=False)
+    monkeypatch.setattr(resolver, "_get_fs_provider", mock.MagicMock(return_value=fs_provider))
 
     with pytest.raises(RuntimeError, match="data and datasets are meant to be immutable"):
         map(
diff --git a/tests/processing/test_utilities.py b/tests/processing/test_utilities.py
index 411779e..31c0d1c 100644
--- a/tests/processing/test_utilities.py
+++ b/tests/processing/test_utilities.py
@@ -72,7 +72,7 @@ def test_extract_rank_and_index_from_filename():
         assert index == rank_and_index[idx][1]
 
 
-def test_read_index_file_content(tmpdir):
+def test_read_index_file_content(tmpdir, monkeypatch):
     output_dir = tmpdir / "output_dir"
 
     assert read_index_file_content(_resolve_dir(str(output_dir))) is None
@@ -86,6 +86,16 @@ def test_read_index_file_content(tmpdir):
 
     assert read_index_file_content(_resolve_dir(str(output_dir))) == dummy_dict
 
+    def _fn(remote_path, local_path):
+        with open(local_path, "w") as f:
+            json.dump(dummy_dict, f)
+
+    fs_provider = MagicMock()
+    fs_provider.download_file = _fn
+
+    monkeypatch.setattr(utilities_module, "_get_fs_provider", MagicMock(return_value=fs_provider))
+    assert read_index_file_content(_resolve_dir("s3://bucket/path")) == dummy_dict
+
 
 def test_remove_uuid_from_filename():
     filepaths = [
diff --git a/tests/streaming/test_fs_provider.py b/tests/streaming/test_fs_provider.py
new file mode 100644
index 0000000..1cf924d
--- /dev/null
+++ b/tests/streaming/test_fs_provider.py
@@ -0,0 +1,46 @@
+from unittest.mock import Mock
+
+import pytest
+
+from litdata.streaming import fs_provider as fs_provider_module
+from litdata.streaming.fs_provider import (
+    GCPFsProvider,
+    S3FsProvider,
+    _get_fs_provider,
+    get_bucket_and_path,
+    not_supported_provider,
+)
+
+
+def test_get_bucket_and_path():
+    bucket, path = get_bucket_and_path("s3://bucket/path/to/file.txt")
+    assert bucket == "bucket"
+    assert path == "path/to/file.txt"
+
+    bucket, path = get_bucket_and_path("s3://bucket/path/to/file.txt", "s3")
+    assert bucket == "bucket"
+    assert path == "path/to/file.txt"
+
+    bucket, path = get_bucket_and_path("gs://bucket/path/to/file.txt", "gs")
+    assert bucket == "bucket"
+    assert path == "path/to/file.txt"
+
+
+def test_get_fs_provider(monkeypatch, google_mock):
+    google_mock.cloud.storage.Client = Mock()
+    monkeypatch.setattr(fs_provider_module, "_GOOGLE_STORAGE_AVAILABLE", True)
+    monkeypatch.setattr(fs_provider_module, "S3Client", Mock())
+
+    fs_provider = _get_fs_provider("s3://bucket/path/to/file.txt")
+    assert isinstance(fs_provider, S3FsProvider)
+
+    fs_provider = _get_fs_provider("gs://bucket/path/to/file.txt")
+    assert isinstance(fs_provider, GCPFsProvider)
+
+    with pytest.raises(ValueError, match="Unsupported scheme"):
+        _get_fs_provider("http://bucket/path/to/file.txt")
+
+
+def test_not_supported_provider():
+    with pytest.raises(ValueError, match="URL should start with one of"):
+        not_supported_provider("http://bucket/path/to/file.txt")
diff --git a/tests/streaming/test_resolver.py b/tests/streaming/test_resolver.py
index 3c487c9..a1e70d8 100644
--- a/tests/streaming/test_resolver.py
+++ b/tests/streaming/test_resolver.py
@@ -329,52 +329,46 @@ def test_execute(phase, monkeypatch, lightning_sdk_mock):
 
 
 def test_assert_dir_is_empty(monkeypatch):
-    boto3 = mock.MagicMock()
-    client_s3_mock = mock.MagicMock()
-    client_s3_mock.list_objects_v2.return_value = {"KeyCount": 1, "Contents": []}
-    boto3.client.return_value = client_s3_mock
-    resolver.boto3 = boto3
+    fs_provider = mock.MagicMock()
+    fs_provider.is_empty = mock.MagicMock(return_value=False)
+    monkeypatch.setattr(resolver, "_get_fs_provider", mock.MagicMock(return_value=fs_provider))
 
     with pytest.raises(RuntimeError, match="The provided output_dir"):
         resolver._assert_dir_is_empty(resolver.Dir(path="/teamspace/...", url="s3://"))
 
-    client_s3_mock.list_objects_v2.return_value = {"KeyCount": 0, "Contents": []}
-    boto3.client.return_value = client_s3_mock
-    resolver.boto3 = boto3
+    fs_provider.is_empty = mock.MagicMock(return_value=True)
 
     resolver._assert_dir_is_empty(resolver.Dir(path="/teamspace/...", url="s3://"))
 
 
 def test_assert_dir_has_index_file(monkeypatch):
-    boto3 = mock.MagicMock()
-    client_s3_mock = mock.MagicMock()
-    client_s3_mock.list_objects_v2.return_value = {"KeyCount": 1, "Contents": []}
-    boto3.client.return_value = client_s3_mock
-    resolver.boto3 = boto3
+    fs_provider = mock.MagicMock()
+    fs_provider.is_empty = mock.MagicMock(return_value=False)
+    monkeypatch.setattr(resolver, "_get_fs_provider", mock.MagicMock(return_value=fs_provider))
 
     with pytest.raises(RuntimeError, match="The provided output_dir"):
         resolver._assert_dir_has_index_file(resolver.Dir(path="/teamspace/...", url="s3://"))
 
-    client_s3_mock.list_objects_v2.return_value = {"KeyCount": 0, "Contents": []}
-    boto3.client.return_value = client_s3_mock
-    resolver.boto3 = boto3
+    fs_provider.is_empty = mock.MagicMock(return_value=True)
 
     resolver._assert_dir_has_index_file(resolver.Dir(path="/teamspace/...", url="s3://"))
 
-    client_s3_mock.list_objects_v2.return_value = {"KeyCount": 1, "Contents": []}
+    fs_provider.exists = mock.MagicMock(return_value=False)
 
-    def head_object(*args, **kwargs):
-        import botocore
+    fs_provider.is_empty = mock.MagicMock(return_value=True)
 
-        raise botocore.exceptions.ClientError({"Error": {"Code": "404", "Message": "Not Found"}}, "HeadObject")
+    resolver._assert_dir_has_index_file(resolver.Dir(path="/teamspace/...", url="s3://"))
 
-    client_s3_mock.head_object = head_object
-    boto3.client.return_value = client_s3_mock
-    resolver.boto3 = boto3
+    fs_provider.exists = mock.MagicMock(return_value=True)
 
-    resolver._assert_dir_has_index_file(resolver.Dir(path="/teamspace/...", url="s3://"))
+    fs_provider.is_empty = mock.MagicMock(return_value=False)
+    fs_provider.delete_file_or_directory = mock.MagicMock()
+
+    resolver._assert_dir_has_index_file(resolver.Dir(path="/teamspace/...", url="s3://"), mode="overwrite")
+
+    resolver._assert_dir_has_index_file(resolver.Dir(path="/teamspace/...", url="s3://"), mode="append")
 
-    boto3.resource.assert_called()
+    assert fs_provider.delete_file_or_directory.call_count == 1
 
 
 def test_resolve_dir_absolute(tmp_path, monkeypatch):