Implement a FSspec based ArtifactLoader

pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
     "Typing :: Typed",
 ]
 
-dependencies = ["tabulate"]
+dependencies = ["tabulate", "fsspec"]
 
 dynamic = ["version"]
 

src/nnbench/types.py

@@ -6,10 +6,17 @@
 import copy
 import inspect
 import os
+import shutil
+import weakref
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
+from pathlib import Path
+from tempfile import mkdtemp
 from typing import Any, Callable, Generic, Iterable, Iterator, Literal, TypeVar
 
+from fsspec import core, filesystem, utils
+from fsspec.implementations.local import LocalFileSystem
+
 from nnbench.context import Context
 
 T = TypeVar("T")
@@ -110,22 +117,72 @@ class LocalArtifactLoader(ArtifactLoader):
     Parameters
     ----------
     path : str | os.PathLike[str]
-        The file system pathto the artifact.
+        The file system path to the artifact.
     """
 
-    def __init__(self, path: str | os.PathLike[str]):
+    def __init__(self, path: str | os.PathLike[str]) -> None:
         self._path = path
 
-    def load(self):
+    def load(self) -> Path:
         """
         Returns the path to the artifact on the local file system.
         """
-        return self._path
+        return Path(self._path).resolve()
 
 
-class S3ArtifactLoader(ArtifactLoader):
-    # TODO: Implement this and other common ArtifactLoders here or in a util
-    pass
+class FilePathArtifactLoader(ArtifactLoader):
+    """
+    ArtifactLoader for loading artifacts using fsspec-supported file systems.
+
+    This allows for loading from various file systems like local, S3, GCS, etc.,
+    by using a unified API provided by fsspec.
+
+    Parameters
+    ----------
+    path : str | os.PathLike[str]
+        The path to the artifact, which can include a protocol specifier (like 's3://') for remote access.
+    to_localdir : str | os.PathLike[str] | None
+        The local directory to which remote artifacts will be downloaded. If provided, the model data will be persisted. Otherwise, local artifacts are cleaned.
+    storage_options : dict[str, Any] | None
+        Storage options for remote storage.
+    """
+
+    def __init__(
+        self,
+        path: str | os.PathLike[str],
+        to_localdir: str | os.PathLike[str] | None = None,
+        storage_options: dict[str, Any] | None = None,
+    ) -> None:
+        self.source_path = str(path)
+        if to_localdir:
+            self.target_path = Path(to_localdir)
+            delete = False
+        else:
+            self.target_path = Path(mkdtemp())
+            delete = True
+        self._finalizer = weakref.finalize(self, self._cleanup, delete=delete)
+        self.storage_options = storage_options or {}
+
+    def load(self) -> Path:
+        """
+        Loads the artifact, downloading it if necessary, and returns the local path.
+
+        Returns
+        -------
+        Path
+            The path to the artifact on the local filesystem.
+        """
+        fs, _, _ = core.get_fs_token_paths(self.source_path)
+        if isinstance(fs, LocalFileSystem):
+            return Path(self.source_path).resolve()
+        else:
+            fs = filesystem(utils.get_protocol(self.source_path), **self.storage_options)
+            fs.get(self.source_path, self.target_path, recursive=True)
+            return Path(self.target_path).resolve()
+
+    def _cleanup(self, delete: bool) -> None:
+        if delete:
+            shutil.rmtree(self.target_path)
 
 
 class Artifact(Generic[T], metaclass=ABCMeta):

tests/conftest.py

@@ -13,3 +13,10 @@
 def testfolder() -> str:
     """A test directory for benchmark collection."""
     return str(HERE / "benchmarks")
+
+
+@pytest.fixture
+def local_file(tmp_path: Path) -> Path:
+    file_path = tmp_path / "test_file.txt"
+    file_path.write_text("Test content")
+    return file_path

tests/test_artifacts.py

@@ -0,0 +1,11 @@
+from pathlib import Path
+
+from nnbench.types import FilePathArtifactLoader
+
+
+def test_load_local_file(local_file: Path, tmp_path: Path) -> None:
+    test_dir = tmp_path / "test_load_dir"
+    loader = FilePathArtifactLoader(str(local_file), str(test_dir))
+    loaded_path: Path = loader.load()
+    assert loaded_path.exists()
+    assert loaded_path.read_text() == "Test content"