Make s3fs dependency explicit (#2683)

* Make s3fs dependency explicit

* Define StorageProtocol type

* Use StorageProtocol type in assets configs

* Use StorageProtocol type as param in StorageClient

* Validate storage_protocol from environ

* Define STORAGE_PROTOCOL_VALUES variable

libs/libcommon/src/libcommon/config.py

@@ -4,19 +4,23 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Literal, Optional
 
 from environs import Env
+from marshmallow.validate import OneOf
+
+STORAGE_PROTOCOL_VALUES: list[str] = ["file", "s3"]
+StorageProtocol = Literal["file", "s3"]
 
 ASSETS_BASE_URL = "http://localhost/assets"
-ASSETS_STORAGE_PROTOCOL = "file"
+ASSETS_STORAGE_PROTOCOL: StorageProtocol = "file"
 ASSETS_STORAGE_ROOT = "/storage/assets"
 
 
 @dataclass(frozen=True)
 class AssetsConfig:
     base_url: str = ASSETS_BASE_URL
-    storage_protocol: str = ASSETS_STORAGE_PROTOCOL
+    storage_protocol: StorageProtocol = ASSETS_STORAGE_PROTOCOL
     storage_root: str = ASSETS_STORAGE_ROOT
 
     @classmethod
@@ -25,7 +29,11 @@ def from_env(cls) -> "AssetsConfig":
         with env.prefixed("ASSETS_"):
             return cls(
                 base_url=env.str(name="BASE_URL", default=ASSETS_BASE_URL),
-                storage_protocol=env.str(name="STORAGE_PROTOCOL", default=ASSETS_STORAGE_PROTOCOL),
+                storage_protocol=env.str(
+                    name="STORAGE_PROTOCOL",
+                    default=ASSETS_STORAGE_PROTOCOL,
+                    validate=OneOf(STORAGE_PROTOCOL_VALUES, error="ASSETS_STORAGE_PROTOCOL must be one of: {choices}"),
+                ),
                 storage_root=env.str(name="STORAGE_ROOT", default=ASSETS_STORAGE_ROOT),
             )
 
@@ -53,14 +61,14 @@ def from_env(cls) -> "S3Config":
 
 
 CACHED_ASSETS_BASE_URL = "http://localhost/cached-assets"
-CACHED_ASSETS_STORAGE_PROTOCOL = "file"
+CACHED_ASSETS_STORAGE_PROTOCOL: StorageProtocol = "file"
 CACHED_ASSETS_STORAGE_ROOT = "/storage/cached-assets"
 
 
 @dataclass(frozen=True)
 class CachedAssetsConfig:
     base_url: str = CACHED_ASSETS_BASE_URL
-    storage_protocol: str = CACHED_ASSETS_STORAGE_PROTOCOL
+    storage_protocol: StorageProtocol = CACHED_ASSETS_STORAGE_PROTOCOL
     storage_root: str = CACHED_ASSETS_STORAGE_ROOT
 
     @classmethod
@@ -69,7 +77,13 @@ def from_env(cls) -> "CachedAssetsConfig":
         with env.prefixed("CACHED_ASSETS_"):
             return cls(
                 base_url=env.str(name="BASE_URL", default=CACHED_ASSETS_BASE_URL),
-                storage_protocol=env.str(name="STORAGE_PROTOCOL", default=CACHED_ASSETS_STORAGE_PROTOCOL),
+                storage_protocol=env.str(
+                    name="STORAGE_PROTOCOL",
+                    default=CACHED_ASSETS_STORAGE_PROTOCOL,
+                    validate=OneOf(
+                        STORAGE_PROTOCOL_VALUES, error="CACHED_ASSETS_STORAGE_PROTOCOL must be one of: {choices}"
+                    ),
+                ),
                 storage_root=env.str(name="STORAGE_ROOT", default=CACHED_ASSETS_STORAGE_ROOT),
             )
 

libs/libcommon/src/libcommon/storage_client.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2023 The HuggingFace Authors.
 import logging
-from typing import Any, Optional
+from typing import Optional, Union
 
 import fsspec
+from fsspec.implementations.local import LocalFileSystem
+from s3fs import S3FileSystem  # type: ignore
 
-from libcommon.config import S3Config
+from libcommon.config import S3Config, StorageProtocol
 from libcommon.url_signer import URLSigner
 
 
@@ -26,16 +28,16 @@ class StorageClient:
         url_signer (`URLSigner`, *optional*): The url signer to use for signing urls
     """
 
-    _fs: Any
-    protocol: str
+    _fs: Union[LocalFileSystem, S3FileSystem]
+    protocol: StorageProtocol
     storage_root: str
     base_url: str
     overwrite: bool
     url_signer: Optional[URLSigner] = None
 
     def __init__(
         self,
-        protocol: str,
+        protocol: StorageProtocol,
         storage_root: str,
         base_url: str,
         overwrite: bool = False,