Merge pull request #49 from luxonis/luxonis_filesystem_kwargs

Cached LuxonisFileSystem

.github/workflows/unittests.yaml

@@ -13,22 +13,15 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
       with:
         ref: ${{ github.head_ref }}
 
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: '3.8'
-
-    - name: Cache pip dependencies
-      uses: actions/cache@v2
-      with:
-        path: ~/.cache/pip
-        key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
+        cache: pip
 
     - name: Install dependencies
       run: |

src/luxonis_ml/utils/filesystem.py

@@ -6,6 +6,9 @@
 from io import BytesIO
 from concurrent.futures import ThreadPoolExecutor
 from .environ import environ
+from logging import getLogger
+
+logger = getLogger(__name__)
 
 
 class LuxonisFileSystem:
@@ -14,15 +17,24 @@ def __init__(
         path: str,
         allow_active_mlflow_run: Optional[bool] = False,
         allow_local: Optional[bool] = True,
+        cache_storage: Optional[str] = None,
     ):
         """Helper class which abstracts uploading and downloading files from remote and local sources.
         Supports S3, MLflow and local file systems.
 
         Args:
             path (str): Input path consisting of protocol and actual path or just path for local files
             allow_active_mlflow_run (Optional[bool], optional): Flag if operations are allowed on active MLFlow run. Defaults to False.
-            allow_local (Optional[bool], optional): Flag if operations are allowed on local file system. Defaults to True.
+            allow_local (Optional[bool], optional): Flag if operations are
+                allowed on local file system. Defaults to True.
+            cache_storage (Optional[str], optional): Path to cache storage. No cache
+                is used if set to None. Defaults to None.
         """
+        if path is None:
+            raise ValueError("No path provided to LuxonisFileSystem.")
+
+        self.cache_storage = cache_storage
+
         self.protocol, self.path = _get_protocol_and_path(path)
         supported_protocols = ["s3", "gcs", "file", "mlflow"]
         if self.protocol not in supported_protocols:
@@ -69,22 +81,30 @@ def full_path(self) -> str:
         return f"{self.protocol}://{self.path}"
 
     def init_fsspec_filesystem(self) -> Any:
-        """Returns fsspec filesystem based on protocol"""
+        """Returns fsspec filesystem based on protocol."""
         if self.protocol == "s3":
             # NOTE: In theory boto3 should look in environment variables automatically but it doesn't seem to work
-            return fsspec.filesystem(
+            fs = fsspec.filesystem(
                 self.protocol,
                 key=environ.AWS_ACCESS_KEY_ID,
                 secret=environ.AWS_SECRET_ACCESS_KEY,
                 endpoint_url=environ.AWS_S3_ENDPOINT_URL,
             )
         elif self.protocol == "gcs":
             # NOTE: This should automatically read from GOOGLE_APPLICATION_CREDENTIALS
-            return fsspec.filesystem(self.protocol)
+            fs = fsspec.filesystem(self.protocol)
         elif self.protocol == "file":
-            return fsspec.filesystem(self.protocol)
+            fs = fsspec.filesystem(self.protocol)
         else:
             raise NotImplementedError
+        if self.cache_storage is None:
+            return fs
+
+        if self.protocol == "file":
+            logger.warning("Ignoring cache storage for local filesystem.")
+            return fs
+
+        return fsspec.filesystem("filecache", fs=fs, cache_storage=self.cache_storage)
 
     def put_file(
         self,