From b6f3d882445163b0c1d6937b516e5ccdec213ef3 Mon Sep 17 00:00:00 2001 From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com> Date: Fri, 2 Oct 2026 10:17:24 -0700 Subject: [PATCH] fix Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com> --- flytekit/core/data_persistence.py | 32 ++++++++- .../unit/core/test_data_persistence.py | 68 +++++++++++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index 329bfe3ff8..9f3fb5d157 100644 --- a/flytekit/core/data_persistence.py +++ b/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( diff --git a/tests/flytekit/unit/core/test_data_persistence.py b/tests/flytekit/unit/core/test_data_persistence.py index ab56f5d07d..7da44e0324 100644 --- a/tests/flytekit/unit/core/test_data_persistence.py +++ b/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("") + + # 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")