fix: keep flytekit on s3fs/gcsfs/adlfs when another library registers a different fsspec implementation

flytekit/core/data_persistence.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import importlib
 import io
 import os
 import pathlib
@@ -17,6 +18,8 @@
 import fsspec
 from decorator import decorator
 from fsspec.asyn import AsyncFileSystem
+from fsspec.registry import known_implementations as fsspec_known_implementations
+from fsspec.registry import registry as fsspec_registry
 from fsspec.utils import get_protocol
 from typing_extensions import Unpack
 
@@ -101,6 +104,29 @@ def get_fsspec_storage_options(
     return {}
 
 
+# Protocols whose fsspec implementation flytekit passes implementation-specific kwargs to (s3fs, gcsfs, adlfs).
+_CANONICAL_FS_PROTOCOLS = ("s3", "gs", "abfs", "abfss")
+
+
+def _canonical_fs_class(protocol: str) -> type:
+    """
+    Returns the fsspec implementation to use for the protocol. flytekit relies on the constructor and call kwargs
+    of fsspec's default implementations (e.g. s3fs.S3FileSystem for s3, and ContentType on put), so a registered
+    implementation is only used if it extends the default one, as flytekitplugins-async-fsspec does. Otherwise,
+    e.g. when another library has registered an unrelated implementation with clobber=True, the default is used.
+    """
+    impl = fsspec_known_implementations[protocol]
+    module_name, class_name = impl["class"].rsplit(".", 1)
+    try:
+        default_cls = getattr(importlib.import_module(module_name), class_name)
+    except ImportError as e:
+        raise ImportError(impl.get("err")) from e
+    registered_cls = fsspec_registry.get(protocol)
+    if registered_cls is not None and issubclass(registered_cls, default_cls):
+        return registered_cls
+    return default_cls
+
+
 def get_additional_fsspec_call_kwargs(protocol: typing.Union[str, tuple], method_name: str) -> Dict[str, Any]:
     """
     These are different from the setup args functions defined above. Those kwargs are applied when asking fsspec
@@ -213,11 +239,11 @@ def get_filesystem(
         elif protocol == "s3":
             s3kwargs = s3_setup_args(self._data_config.s3, anonymous=anonymous)
             s3kwargs.update(kwargs)
-            return fsspec.filesystem(protocol, **s3kwargs)  # type: ignore
+            return _canonical_fs_class(protocol)(**s3kwargs)  # type: ignore
         elif protocol == "gs":
             if anonymous:
                 kwargs["token"] = _ANON
-            return fsspec.filesystem(protocol, **kwargs)  # type: ignore
+            return _canonical_fs_class(protocol)(**kwargs)  # type: ignore
         elif protocol == "ftp":
             kwargs.update(fsspec.implementations.ftp.FTPFileSystem._get_kwargs_from_urls(path))
             return fsspec.filesystem(protocol, **kwargs)
@@ -227,6 +253,8 @@ def get_filesystem(
         )
         kwargs.update(storage_options)
 
+        if protocol in _CANONICAL_FS_PROTOCOLS:
+            return _canonical_fs_class(protocol)(**kwargs)
         return fsspec.filesystem(protocol, **kwargs)
 
     async def get_async_filesystem_for_path(

tests/flytekit/unit/core/test_data_persistence.py

@@ -1,3 +1,4 @@
+import importlib
 import io
 import os
 import pathlib
@@ -211,6 +212,73 @@ def __init__(self, *args, **kwargs):
     fp.get_filesystem("testgetfs", test_arg="test_arg")
 
 
+@pytest.fixture
+def restore_fsspec_registry():
+    from fsspec.registry import _registry
+
+    saved = dict(_registry)
+    yield
+    _registry.clear()
+    _registry.update(saved)
+
+
+class UnrelatedFileSystem(fsspec.AbstractFileSystem):
+    pass
+
+
+@pytest.mark.parametrize(
+    "protocol, kwargs",
+    [("s3", {}), ("gs", {"token": "anon"})],
+)
+def test_get_filesystem_ignores_unrelated_registered_implementation(restore_fsspec_registry, protocol, kwargs):
+    from fsspec.registry import known_implementations
+
+    module_name, class_name = known_implementations[protocol]["class"].rsplit(".", 1)
+    default_cls = getattr(importlib.import_module(module_name), class_name)
+    fsspec.register_implementation(protocol, UnrelatedFileSystem, clobber=True)
+
+    fp = FileAccessProvider("/tmp", "s3://my-bucket")
+    assert type(fp.get_filesystem(protocol, **kwargs)) is default_cls
+
+
+def test_get_filesystem_keeps_registered_subclass_of_default(restore_fsspec_registry):
+    from s3fs import S3FileSystem
+
+    class CustomS3FileSystem(S3FileSystem):
+        pass
+
+    fsspec.register_implementation("s3", CustomS3FileSystem, clobber=True)
+
+    fp = FileAccessProvider("/tmp", "s3://my-bucket")
+    assert type(fp.get_filesystem("s3")) is CustomS3FileSystem
+
+
+@pytest.mark.parametrize(
+    "protocol, put_target, content_type_kwarg",
+    [
+        ("s3", "s3fs.S3FileSystem._put", "ContentType"),
+        ("gs", "gcsfs.GCSFileSystem._put", "content_type"),
+    ],
+)
+def test_put_data_passes_content_type_with_unrelated_implementation(
+    restore_fsspec_registry, tmp_path, protocol, put_target, content_type_kwarg
+):
+    fsspec.register_implementation(protocol, UnrelatedFileSystem, clobber=True)
+    local_path = tmp_path / "deck.html"
+    local_path.write_text("<html></html>")
+
+    # The upload itself is mocked, so skip gcsfs's credential lookup, which is slow without GCP credentials.
+    with (
+        mock.patch("gcsfs.credentials.GoogleCredentials.connect"),
+        mock.patch(put_target, new_callable=AsyncMock) as mock_put,
+    ):
+        fp = FileAccessProvider("/tmp", f"{protocol}://my-bucket")
+        fp.put_data(str(local_path), f"{protocol}://my-bucket/deck.html", **{content_type_kwarg: "text/html"})
+
+    mock_put.assert_called_once()
+    assert mock_put.call_args.kwargs[content_type_kwarg] == "text/html"
+
+
 def test_get_additional_fsspec_call_kwargs():
     with mock.patch("flytekit.core.data_persistence._WRITE_SIZE_CHUNK_BYTES", 12345):
         kwargs = get_additional_fsspec_call_kwargs(("s3", "s3a"), "put")