From 422143dfdabc84a5786387037397aa689cbd9896 Mon Sep 17 00:00:00 2001 From: "Masaki.Nakano" Date: Fri, 5 Jun 2026 16:21:47 +0900 Subject: [PATCH] WIP: fsspec support --- docs/source/reference.rst | 40 ++++++ pfio/fsspec.py | 263 ++++++++++++++++++++++++++++++++++ pyproject.toml | 10 +- tests/v2_tests/test_fsspec.py | 254 ++++++++++++++++++++++++++++++++ 4 files changed, 566 insertions(+), 1 deletion(-) create mode 100644 pfio/fsspec.py create mode 100644 tests/v2_tests/test_fsspec.py diff --git a/docs/source/reference.rst b/docs/source/reference.rst index d93c52cd..257e978a 100644 --- a/docs/source/reference.rst +++ b/docs/source/reference.rst @@ -78,6 +78,46 @@ but several methods are not yet implemented. .. autoclass:: pfio.v2.pathlib.Path :members: +fsspec integration +------------------ + +PFIO backends can be used through the `fsspec +`_ ``AbstractFileSystem`` +interface, which lets fsspec-aware libraries (pandas, pyarrow, dask, +zarr, ...) read and write data via PFIO. Install the optional +dependency with ``pip install pfio[fsspec]``. + +Each PFIO backend is registered under its own protocol so that it does +not clobber fsspec's built-in implementations (such as ``s3`` provided +by ``s3fs``): + +* ``pfio-file://`` -> :class:`pfio.v2.Local` +* ``pfio-s3://`` -> :class:`pfio.v2.S3` +* ``pfio-hdfs://`` -> :class:`pfio.v2.Hdfs` + +These protocols are registered automatically through entry points, so +they are available as soon as both ``pfio`` and ``fsspec`` are +installed:: + + import fsspec + + with fsspec.open("pfio-s3://your-bucket/foo/bar.txt", "rb") as fp: + data = fp.read() + + fs = fsspec.filesystem("pfio-s3", endpoint="https://s3.example.com") + fs.ls("your-bucket/foo") + +Connection parameters (e.g. ``endpoint``, ``aws_access_key_id``) are +passed as keyword arguments to :func:`fsspec.filesystem`; the bucket is +taken from the path. To make PFIO handle a standard protocol such as +``s3://`` instead of ``pfio-s3://``, call :func:`pfio.fsspec.register`:: + + import pfio.fsspec + + pfio.fsspec.register(s3=True) # now "s3://..." is served by PFIO + +.. autofunction:: pfio.fsspec.register + Sparse File Cache ----------------- diff --git a/pfio/fsspec.py b/pfio/fsspec.py new file mode 100644 index 00000000..5bfa3fd4 --- /dev/null +++ b/pfio/fsspec.py @@ -0,0 +1,263 @@ +'''fsspec integration for PFIO. + +This module exposes :mod:`pfio.v2` backends (Local / S3 / HDFS) through the +`fsspec `_ ``AbstractFileSystem`` +interface so that libraries built on top of fsspec (pandas, pyarrow, dask, +zarr, ...) can use PFIO as their storage backend. + +Each PFIO backend is registered under its own protocol so that it does not +clobber fsspec's built-in implementations (e.g. ``s3`` provided by s3fs): + +* ``pfio-file://`` -> :class:`pfio.v2.Local` +* ``pfio-s3://`` -> :class:`pfio.v2.S3` +* ``pfio-hdfs://`` -> :class:`pfio.v2.Hdfs` + +These protocols are registered automatically through the ``fsspec.specs`` +entry points declared in ``pyproject.toml``. To make PFIO handle a standard +protocol such as ``s3://`` instead, call :func:`register`. + +``fsspec`` is an optional dependency. Importing :mod:`pfio` itself never +requires fsspec; this module is only imported lazily by fsspec's entry point +machinery (or explicitly by the user), at which point fsspec is guaranteed to +be installed. +''' +import threading +from datetime import datetime +from typing import Optional + +try: + from fsspec.spec import AbstractFileSystem + _HAS_FSSPEC = True +except ImportError: + _HAS_FSSPEC = False + + +if _HAS_FSSPEC: + class _PfioFileSystem(AbstractFileSystem): + '''Common base wrapping a :class:`pfio.v2.FS` behind fsspec. + + Subclasses bind a single PFIO backend via :attr:`_pfio_scheme` and + implement :meth:`_get_fs_and_path`, which resolves an fsspec path to a + concrete PFIO ``FS`` instance and the path relative to it. + + The wrapped PFIO ``FS`` instances are created lazily and reused. PFIO + backends detect ``fork()`` and reconnect at each operation boundary, so + sharing a single instance (as fsspec's instance cache does) is safe. + The PFIO instances are intentionally *not* part of the pickled state; + ``AbstractFileSystem.__reduce__`` reconstructs the filesystem from its + ``storage_options``, after which they are lazily recreated. + ''' + + _pfio_scheme: Optional[str] = None + + def __init__(self, **storage_options): + super().__init__(**storage_options) + self._fs = None + self._fs_lock = threading.Lock() + + def _get_fs_and_path(self, path): + '''Return ``(pfio_fs, path_relative_to_fs)`` for ``path``.''' + raise NotImplementedError + + @staticmethod + def _stat_to_info(name, st): + '''Convert a PFIO ``FileStat`` into an fsspec info dict.''' + if st.isdir(): + # S3PrefixStat reports size == -1; normalize directories to 0. + return {"name": name, "size": 0, "type": "directory"} + info = { + "name": name, + "size": st.size if st.size is not None else 0, + "type": "file", + } + mtime = getattr(st, "last_modified", None) + if mtime: + info["mtime"] = mtime + return info + + def _open(self, path, mode="rb", block_size=None, autocommit=True, + cache_options=None, **kwargs): + fs, rel = self._get_fs_and_path(path) + return fs.open(rel, mode) + + def ls(self, path, detail=True, **kwargs): + fs, rel = self._get_fs_and_path(path) + # fsspec contract: ls() of a file returns a single-element list + # with that file's info, not a directory listing. + if fs.exists(rel) and not fs.isdir(rel): + info = self.info(path) + return [info] if detail else [info["name"]] + base = self._strip_protocol(path).rstrip("/") + out = [] + for st in fs.list(rel, recursive=False, detail=True): + # PFIO returns names relative to the listed path (basename for + # Local, key-relative for S3); rejoin to a full fsspec path. + name = base + "/" + st.filename.rstrip("/") if base \ + else st.filename.rstrip("/") + out.append(self._stat_to_info(name, st)) + if detail: + return out + return sorted(e["name"] for e in out) + + def info(self, path, **kwargs): + fs, rel = self._get_fs_and_path(path) + name = self._strip_protocol(path) + st = fs.stat(rel) + return self._stat_to_info(name, st) + + def exists(self, path, **kwargs): + fs, rel = self._get_fs_and_path(path) + return fs.exists(rel) + + def isdir(self, path): + fs, rel = self._get_fs_and_path(path) + return fs.isdir(rel) + + def isfile(self, path): + fs, rel = self._get_fs_and_path(path) + return fs.exists(rel) and not fs.isdir(rel) + + def mkdir(self, path, create_parents=True, **kwargs): + fs, rel = self._get_fs_and_path(path) + if create_parents: + fs.makedirs(rel, exist_ok=True) + else: + fs.mkdir(rel) + + def makedirs(self, path, exist_ok=False): + fs, rel = self._get_fs_and_path(path) + fs.makedirs(rel, exist_ok=exist_ok) + + def rmdir(self, path): + fs, rel = self._get_fs_and_path(path) + fs.remove(rel, recursive=False) + + def _rm(self, path): + fs, rel = self._get_fs_and_path(path) + fs.remove(rel, recursive=False) + + def modified(self, path): + fs, rel = self._get_fs_and_path(path) + st = fs.stat(rel) + return datetime.fromtimestamp(st.last_modified) + + def created(self, path): + fs, rel = self._get_fs_and_path(path) + st = fs.stat(rel) + created = getattr(st, "created", None) + if created is None: + raise NotImplementedError( + "{} does not expose creation time".format( + type(self).__name__)) + return datetime.fromtimestamp(created) + + class PfioFileFileSystem(_PfioFileSystem): + '''fsspec filesystem backed by :class:`pfio.v2.Local`.''' + + protocol = "pfio-file" + root_marker = "/" + _pfio_scheme = "file" + + def _get_fs_and_path(self, path): + rel = self._strip_protocol(path) + with self._fs_lock: + if self._fs is None: + from pfio.v2 import Local + + # cwd="" falls back to os.getcwd(), but every operation + # receives an absolute path which os.path.join honors. + self._fs = Local(cwd="", scheme="file") + return self._fs, rel + + class PfioS3FileSystem(_PfioFileSystem): + '''fsspec filesystem backed by :class:`pfio.v2.S3`. + + Paths are ``pfio-s3:///``. Because PFIO's ``S3`` takes + the bucket as a constructor argument while fsspec uses a single + instance for all paths, one ``S3`` instance is created (lazily) per + bucket and reused. + ''' + + protocol = "pfio-s3" + root_marker = "" + _pfio_scheme = "s3" + + # Only connection-defining options; the bucket comes from the path. + _S3_CONN_KEYS = ( + "endpoint", "create_bucket", "aws_access_key_id", + "aws_secret_access_key", "mpu_chunksize", "buffering", + "connect_timeout", "read_timeout", + ) + + def __init__(self, **storage_options): + super().__init__(**storage_options) + self._fs_cache = {} + self._s3_kwargs = { + k: storage_options[k] + for k in self._S3_CONN_KEYS + if k in storage_options + } + + def _get_fs_and_path(self, path): + p = self._strip_protocol(path).lstrip("/") + bucket, _, key = p.partition("/") + if not bucket: + raise ValueError( + "S3 path must contain a bucket: {!r}".format(path)) + with self._fs_lock: + fs = self._fs_cache.get(bucket) + if fs is None: + from pfio.v2 import S3 + fs = S3(bucket=bucket, prefix="", scheme="s3", + **self._s3_kwargs) + self._fs_cache[bucket] = fs + return fs, key + + class PfioHdfsFileSystem(_PfioFileSystem): + '''fsspec filesystem backed by :class:`pfio.v2.Hdfs`.''' + + protocol = "pfio-hdfs" + root_marker = "/" + _pfio_scheme = "hdfs" + + @classmethod + def _strip_protocol(cls, path): + # PFIO resolves the HDFS nameservice from hdfs-site.xml and ignores + # any netloc, so drop it and keep only the path component. + if isinstance(path, str) and "://" in path: + from urllib.parse import urlparse + path = urlparse(path).path + return super()._strip_protocol(path) + + def _get_fs_and_path(self, path): + rel = self._strip_protocol(path) + with self._fs_lock: + if self._fs is None: + from pfio.v2 import Hdfs + self._fs = Hdfs(cwd="/", scheme="hdfs") + return self._fs, rel + + def register(*, s3=False, hdfs=False, file=False, clobber=True): + '''Register PFIO backends under fsspec's standard protocols. + + The ``pfio-file`` / ``pfio-s3`` / ``pfio-hdfs`` protocols are always + available via entry points. This helper additionally binds PFIO to the + standard ``file`` / ``s3`` / ``hdfs`` protocols on an opt-in basis, + overriding any implementation already registered for them. + + Args: + s3 (bool): Bind :class:`PfioS3FileSystem` to ``s3``. + hdfs (bool): Bind :class:`PfioHdfsFileSystem` to ``hdfs``. + file (bool): Bind :class:`PfioFileFileSystem` to ``file``. + clobber (bool): Overwrite an existing registration. + ''' + import fsspec + if s3: + fsspec.register_implementation( + "s3", PfioS3FileSystem, clobber=clobber) + if hdfs: + fsspec.register_implementation( + "hdfs", PfioHdfsFileSystem, clobber=clobber) + if file: + fsspec.register_implementation( + "file", PfioFileFileSystem, clobber=clobber) diff --git a/pyproject.toml b/pyproject.toml index e2c00185..b6396025 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,16 +33,22 @@ dependencies = [ ] [project.optional-dependencies] -test = ["tox>=4", "pytest", "flake8", "autopep8", "parameterized", "isort", "moto[server]>=5.0.0", "numpy", "mypy"] +test = ["tox>=4", "pytest", "flake8", "autopep8", "parameterized", "isort", "moto[server]>=5.0.0", "numpy", "mypy", "fsspec"] # When updating doc deps, docs/requirements.txt should be updated too doc = ["sphinx", "sphinx_rtd_theme"] bench = ["numpy>=1.19.5", "torch>=1.9.0", "Pillow<=8.2.0"] hdfs = ["pyarrow>=6.0.0"] trace = ["pytorch-pfn-extras"] +fsspec = ["fsspec"] [project.urls] Homepage = "https://github.com/pfnet/pfio" +[project.entry-points."fsspec.specs"] +pfio-file = "pfio.fsspec:PfioFileFileSystem" +pfio-s3 = "pfio.fsspec:PfioS3FileSystem" +pfio-hdfs = "pfio.fsspec:PfioHdfsFileSystem" + [dependency-groups] dev = ["tox", "tox-uv"] @@ -73,6 +79,8 @@ module = [ "pyarrow", "pyarrow.*", "pytorch_pfn_extras.*", + "fsspec", + "fsspec.*", ] ignore_missing_imports = true diff --git a/tests/v2_tests/test_fsspec.py b/tests/v2_tests/test_fsspec.py new file mode 100644 index 00000000..fd8f1e32 --- /dev/null +++ b/tests/v2_tests/test_fsspec.py @@ -0,0 +1,254 @@ +import pickle +from importlib.metadata import entry_points + +import boto3 +import fsspec +import pytest +from moto import mock_aws + +from pfio.fsspec import (PfioFileFileSystem, PfioHdfsFileSystem, + PfioS3FileSystem, register) + + +@pytest.fixture(autouse=True) +def _clear_fsspec_cache(): + # fsspec caches filesystem instances by class + storage_options. Clear it + # around every test so a cached S3 client cannot leak across moto mocks. + for cls in (PfioFileFileSystem, PfioS3FileSystem, PfioHdfsFileSystem): + cls.clear_instance_cache() + yield + for cls in (PfioFileFileSystem, PfioS3FileSystem, PfioHdfsFileSystem): + cls.clear_instance_cache() + + +# --------------------------------------------------------------------------- +# pfio-file +# --------------------------------------------------------------------------- + +def test_file_roundtrip(tmp_path): + fs = fsspec.filesystem("pfio-file") + assert isinstance(fs, PfioFileFileSystem) + + target = str(tmp_path / "a.txt") + with fs.open(target, "wb") as f: + f.write(b"hello") + + assert fs.cat(target) == b"hello" + assert fs.exists(target) + assert fs.isfile(target) + assert not fs.isdir(target) + + +def test_file_open_url(tmp_path): + target = tmp_path / "b.txt" + with fsspec.open("pfio-file://" + str(target), "wb") as f: + f.write(b"world") + + with fsspec.open("pfio-file://" + str(target), "rb") as f: + assert f.read() == b"world" + + +def test_file_info_and_ls(tmp_path): + fs = fsspec.filesystem("pfio-file") + (tmp_path / "sub").mkdir() + (tmp_path / "a.txt").write_bytes(b"12345") + + info = fs.info(str(tmp_path / "a.txt")) + assert info["type"] == "file" + assert info["size"] == 5 + assert info["name"] == str(tmp_path / "a.txt") + + dinfo = fs.info(str(tmp_path / "sub")) + assert dinfo["type"] == "directory" + assert dinfo["size"] == 0 + + listing = fs.ls(str(tmp_path), detail=True) + names = {e["name"]: e["type"] for e in listing} + assert names[str(tmp_path / "a.txt")] == "file" + assert names[str(tmp_path / "sub")] == "directory" + + # detail=False returns full paths + names = fs.ls(str(tmp_path), detail=False) + assert str(tmp_path / "a.txt") in names + assert str(tmp_path / "sub") in names + + +def test_file_ls_on_file(tmp_path): + fs = fsspec.filesystem("pfio-file") + target = str(tmp_path / "a.txt") + (tmp_path / "a.txt").write_bytes(b"12345") + + # fsspec contract: ls() of a file yields a single-element listing. + listing = fs.ls(target, detail=True) + assert len(listing) == 1 + assert listing[0]["name"] == target + assert listing[0]["type"] == "file" + assert fs.ls(target, detail=False) == [target] + + +def test_file_makedirs_and_rm(tmp_path): + fs = fsspec.filesystem("pfio-file") + d = str(tmp_path / "x" / "y") + fs.makedirs(d, exist_ok=True) + assert fs.isdir(d) + + target = str(tmp_path / "x" / "y" / "f.txt") + with fs.open(target, "wb") as f: + f.write(b"z") + fs.rm(target) + assert not fs.exists(target) + + +# --------------------------------------------------------------------------- +# pfio-s3 +# --------------------------------------------------------------------------- + +@pytest.fixture +def s3_mock(): + with mock_aws(): + client = boto3.client("s3") + client.create_bucket(Bucket="test-bucket") + client.create_bucket(Bucket="other-bucket") + yield + + +def test_s3_roundtrip(s3_mock): + fs = fsspec.filesystem("pfio-s3") + assert isinstance(fs, PfioS3FileSystem) + + with fs.open("pfio-s3://test-bucket/dir/a.txt", "wb") as f: + f.write(b"data") + + assert fs.cat("test-bucket/dir/a.txt") == b"data" + assert fs.exists("test-bucket/dir/a.txt") + assert fs.isfile("test-bucket/dir/a.txt") + + +def test_s3_ls_full_path(s3_mock): + fs = fsspec.filesystem("pfio-s3") + with fs.open("pfio-s3://test-bucket/dir/a.txt", "wb") as f: + f.write(b"data") + with fs.open("pfio-s3://test-bucket/dir/sub/b.txt", "wb") as f: + f.write(b"data") + + listing = fs.ls("test-bucket/dir", detail=True) + by_name = {e["name"]: e for e in listing} + assert "test-bucket/dir/a.txt" in by_name + assert by_name["test-bucket/dir/a.txt"]["type"] == "file" + # the common prefix shows up as a directory + assert "test-bucket/dir/sub" in by_name + assert by_name["test-bucket/dir/sub"]["type"] == "directory" + + +def test_s3_ls_on_file(s3_mock): + fs = fsspec.filesystem("pfio-s3") + with fs.open("pfio-s3://test-bucket/dir/a.txt", "wb") as f: + f.write(b"data") + + listing = fs.ls("test-bucket/dir/a.txt", detail=True) + assert len(listing) == 1 + assert listing[0]["name"] == "test-bucket/dir/a.txt" + assert listing[0]["type"] == "file" + assert fs.ls("test-bucket/dir/a.txt", detail=False) \ + == ["test-bucket/dir/a.txt"] + + +def test_s3_isdir_and_info_directory(s3_mock): + fs = fsspec.filesystem("pfio-s3") + with fs.open("pfio-s3://test-bucket/dir/a.txt", "wb") as f: + f.write(b"data") + + assert fs.isdir("test-bucket/dir") + # S3PrefixStat has size == -1; it must be normalized to a directory info. + info = fs.info("test-bucket/dir") + assert info["type"] == "directory" + assert info["size"] == 0 + + +def test_s3_multiple_buckets_cached(s3_mock): + fs = fsspec.filesystem("pfio-s3") + with fs.open("pfio-s3://test-bucket/a.txt", "wb") as f: + f.write(b"a") + with fs.open("pfio-s3://other-bucket/b.txt", "wb") as f: + f.write(b"b") + + assert fs.cat("test-bucket/a.txt") == b"a" + assert fs.cat("other-bucket/b.txt") == b"b" + # one pfio S3 instance lazily created per bucket + assert set(fs._fs_cache.keys()) == {"test-bucket", "other-bucket"} + + +def test_s3_pickle(s3_mock): + fs = fsspec.filesystem("pfio-s3") + with fs.open("pfio-s3://test-bucket/a.txt", "wb") as f: + f.write(b"hello") + + data = pickle.dumps(fs) + # Drop the cached instance so unpickling rebuilds from storage_options + # (as it would in a separate process), rather than hitting fsspec's + # in-process instance cache and returning the very same object. + PfioS3FileSystem.clear_instance_cache() + fs2 = pickle.loads(data) + assert fs2 is not fs + # internal pfio instances are not part of the pickled state + assert fs2._fs_cache == {} + assert fs2.cat("test-bucket/a.txt") == b"hello" + + +# --------------------------------------------------------------------------- +# Instance cache / class separation +# --------------------------------------------------------------------------- + +def test_instance_cache_same_options(): + a = fsspec.filesystem("pfio-file") + b = fsspec.filesystem("pfio-file") + assert a is b + + +def test_protocols_do_not_collide(): + f = fsspec.filesystem("pfio-file") + s = fsspec.filesystem("pfio-s3") + assert type(f) is PfioFileFileSystem + assert type(s) is PfioS3FileSystem + assert f is not s + + +# --------------------------------------------------------------------------- +# Entry points / registration +# --------------------------------------------------------------------------- + +def test_entry_points_registered(): + eps = entry_points(group="fsspec.specs") + names = {e.name for e in eps} + if not {"pfio-file", "pfio-s3", "pfio-hdfs"} <= names: + pytest.skip("entry points not visible (package not installed)") + assert fsspec.get_filesystem_class("pfio-file") is PfioFileFileSystem + assert fsspec.get_filesystem_class("pfio-s3") is PfioS3FileSystem + assert fsspec.get_filesystem_class("pfio-hdfs") is PfioHdfsFileSystem + + +def test_register_helper(): + from fsspec.registry import _registry as reg + original = reg.get("s3") + try: + register(s3=True) + assert fsspec.get_filesystem_class("s3") is PfioS3FileSystem + finally: + # restore the registry to avoid leaking into other tests + if original is None: + reg.pop("s3", None) + else: + reg["s3"] = original + + +# --------------------------------------------------------------------------- +# pfio-hdfs (path handling only; no live HDFS) +# --------------------------------------------------------------------------- + +def test_hdfs_strip_protocol_drops_netloc(): + assert PfioHdfsFileSystem._strip_protocol( + "pfio-hdfs://nameservice/a/b") == "/a/b" + assert PfioHdfsFileSystem._strip_protocol( + "pfio-hdfs:///a/b") == "/a/b" + assert PfioHdfsFileSystem._strip_protocol( + "pfio-hdfs://nameservice/a/b/") == "/a/b"