From 8fbc628b609739b694ee1a481c6e65abe06c93df Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Fri, 14 Aug 2026 10:32:22 -0400 Subject: [PATCH] fix(storage): don't clobber the global fsspec registry on import Importing `flyte` (v2) inside a process that also runs flytekit (v1) broke all of flytekit's object-store I/O: `flyte.storage._storage` unconditionally ran `obstore.fsspec.register(["s3", "gs", "abfs", "abfss"], ...)` at import time, replacing s3fs/gcsfs/adlfs in the global fsspec registry. flytekit then called `fsspec.filesystem("s3", cache_regions=True, ...)` and obstore rejected the s3fs-only kwargs ("Configuration key: 'cache_regions' is not valid for store 'S3'"), failing every output upload and masking the real task error. Two-part fix: 1. flyte's own I/O is now registry-independent: for obstore-supported protocols, `get_underlying_filesystem` instantiates a flyte-owned per-protocol subclass of `obstore.fsspec.FsspecStore` directly (mirroring what `obstore.fsspec.register` builds, including instance caching via fsspec's `_Cached` metaclass) instead of resolving through the registry. 2. Global registration is now non-destructive gap-filling: obstore is only registered for a protocol when nothing else is registered and no other implementation (e.g. s3fs) is importable, so pure-v2 images keep working `s3://` paths for pandas/pyarrow while hybrid v1/v2 processes keep flytekit's filesystems intact regardless of import order. Co-Authored-By: Claude Fable 5 --- src/flyte/storage/_storage.py | 55 ++++++- .../internal/storage/test_fsspec_registry.py | 136 ++++++++++++++++++ 2 files changed, 189 insertions(+), 2 deletions(-) create mode 100644 tests/internal/storage/test_fsspec_registry.py diff --git a/src/flyte/storage/_storage.py b/src/flyte/storage/_storage.py index 509385c5b..9b4c6547d 100644 --- a/src/flyte/storage/_storage.py +++ b/src/flyte/storage/_storage.py @@ -1,5 +1,6 @@ from __future__ import annotations +import functools import os import pathlib import random @@ -13,7 +14,7 @@ from fsspec.asyn import AsyncFileSystem from fsspec.utils import get_protocol from obstore.exceptions import GenericError -from obstore.fsspec import register +from obstore.fsspec import FsspecStore, register from obstore.store import ObjectStore from flyte._initialize import get_storage @@ -159,6 +160,24 @@ def get_configured_fsspec_kwargs( return {} +@functools.lru_cache(maxsize=None) +def _obstore_filesystem_class(protocol: str) -> type[FsspecStore]: + """ + Build (and cache) a per-protocol subclass of obstore's ``FsspecStore``, mirroring what + ``obstore.fsspec.register(protocol, asynchronous=True)`` would register -- but without + touching the global fsspec registry. + + Subclasses of fsspec's ``AbstractFileSystem`` keep fsspec's instance caching (the + ``_Cached`` metaclass) even when instantiated directly, so constructing these classes + behaves the same as resolving them through ``fsspec.filesystem(...)``. + """ + return type( + f"FlyteFsspecStore_{protocol}", + (FsspecStore,), + {"protocol": protocol, "asynchronous": True}, + ) + + def get_underlying_filesystem( protocol: typing.Optional[str] = None, anonymous: bool = False, @@ -172,6 +191,13 @@ def get_underlying_filesystem( configured_kwargs = get_configured_fsspec_kwargs(protocol, anonymous=anonymous) configured_kwargs.update(kwargs) + if _is_obstore_supported_protocol(protocol): + # Instantiate flyte's own obstore-backed filesystem directly instead of resolving the + # protocol through the global fsspec registry: the configured kwargs are obstore-shaped + # (retry_config, client_options, config), and another SDK sharing the process (e.g. + # flyte v1 / flytekit) may have registered s3fs/gcsfs/adlfs for these protocols. + return _obstore_filesystem_class(protocol)(**configured_kwargs) + return fsspec.filesystem(protocol, **configured_kwargs) @@ -598,5 +624,30 @@ def get_credentials_error(uri: str, protocol: str) -> str: raise ValueError(f"Unsupported protocol: {protocol}") -register(_OBSTORE_SUPPORTED_PROTOCOLS, asynchronous=True) +def _register_obstore_for_missing_protocols() -> None: + """ + Register obstore with fsspec only for protocols that have no other usable implementation. + + flyte v1 (flytekit) may share the process (e.g. a v1 task importing ``flyte`` to launch v2 + runs) and owns these protocols via s3fs/gcsfs/adlfs; clobbering the global registry breaks + its I/O because obstore rejects s3fs/gcsfs/adlfs-only kwargs (e.g. ``cache_regions``). + flyte's own I/O never resolves these protocols through the registry (see + ``get_underlying_filesystem``), so this global registration only exists so that third-party + libraries (pandas, pyarrow, ...) can still resolve e.g. ``s3://`` paths in images that + don't ship s3fs/gcsfs/adlfs. + """ + for protocol in _OBSTORE_SUPPORTED_PROTOCOLS: + if protocol in fsspec.registry: + # Someone (e.g. flytekit v1) already registered an implementation; leave it alone. + continue + try: + # fsspec resolves known_implementations lazily, so an empty registry does not mean + # no implementation exists. This probes whether one (e.g. s3fs) is importable, and + # registers it in fsspec.registry as a side effect. + fsspec.get_filesystem_class(protocol) + except (ImportError, ValueError): + register(protocol, asynchronous=True) + + +_register_obstore_for_missing_protocols() fsspec.register_implementation("flyte", FlyteFS, clobber=True) diff --git a/tests/internal/storage/test_fsspec_registry.py b/tests/internal/storage/test_fsspec_registry.py new file mode 100644 index 000000000..4b4fac7ca --- /dev/null +++ b/tests/internal/storage/test_fsspec_registry.py @@ -0,0 +1,136 @@ +""" +Tests for how flyte interacts with the global fsspec registry. + +flyte v1 (flytekit) may share the process (e.g. a v1 task importing ``flyte`` to launch v2 +runs) and registers s3fs/gcsfs/adlfs for the object-store protocols. Importing ``flyte`` must +never clobber those registrations, and flyte's own I/O must not depend on what is registered. +""" + +import fsspec +import pytest +from fsspec.registry import _registry +from obstore.fsspec import FsspecStore + +from flyte.storage._storage import ( + _OBSTORE_SUPPORTED_PROTOCOLS, + _obstore_filesystem_class, + _register_obstore_for_missing_protocols, + get_underlying_filesystem, +) + + +@pytest.fixture +def restore_fsspec_registry(): + """Snapshot the global fsspec registry and restore it after the test.""" + saved = dict(_registry) + yield + _registry.clear() + _registry.update(saved) + + +class DummyS3FileSystem(fsspec.AbstractFileSystem): + """Stand-in for s3fs.S3FileSystem: accepts s3fs-only kwargs like ``cache_regions``.""" + + protocol = "s3" + cachable = False + + def __init__(self, *args, cache_regions: bool = False, **kwargs): + super().__init__(*args, **kwargs) + self.cache_regions = cache_regions + + +def test_pre_registered_implementation_survives(restore_fsspec_registry): + """A registration made before flyte's helper runs (e.g. by flytekit) must be preserved.""" + fsspec.register_implementation("s3", DummyS3FileSystem, clobber=True) + + _register_obstore_for_missing_protocols() + + assert fsspec.registry["s3"] is DummyS3FileSystem + assert isinstance(fsspec.filesystem("s3"), DummyS3FileSystem) + + +def test_get_underlying_filesystem_is_registry_independent(restore_fsspec_registry): + """flyte's own I/O gets an obstore-backed filesystem even when s3fs-style owns the registry.""" + fsspec.register_implementation("s3", DummyS3FileSystem, clobber=True) + + fs = get_underlying_filesystem("s3") + + assert isinstance(fs, FsspecStore) + assert fs.protocol == "s3" + # The obstore bypasses in _storage.py duck-type on these attributes. + assert hasattr(fs, "_split_path") + assert hasattr(fs, "_construct_store") + + +def test_helper_registers_obstore_when_nothing_importable(restore_fsspec_registry, monkeypatch): + """When no implementation is registered or importable, obstore is registered as a fallback.""" + + def raise_import_error(protocol): + raise ImportError(f"no implementation for {protocol}") + + monkeypatch.setattr(fsspec, "get_filesystem_class", raise_import_error) + for protocol in _OBSTORE_SUPPORTED_PROTOCOLS: + _registry.pop(protocol, None) + + _register_obstore_for_missing_protocols() + + for protocol in _OBSTORE_SUPPORTED_PROTOCOLS: + assert protocol in fsspec.registry + assert issubclass(fsspec.registry[protocol], FsspecStore) + + +def test_helper_leaves_importable_implementation_alone(restore_fsspec_registry, monkeypatch): + """ + fsspec resolves known_implementations lazily, so the registry may be empty at import time + even though e.g. s3fs is installed. The helper must probe importability, not just the + registry, and leave the importable implementation in place. + """ + _registry.pop("s3", None) + + def fake_get_filesystem_class(protocol): + if protocol == "s3": + # Mimic fsspec's lazy resolution side effect of registering the class. + fsspec.register_implementation("s3", DummyS3FileSystem, clobber=True) + return DummyS3FileSystem + raise ImportError(f"no implementation for {protocol}") + + monkeypatch.setattr(fsspec, "get_filesystem_class", fake_get_filesystem_class) + + _register_obstore_for_missing_protocols() + + assert fsspec.registry["s3"] is DummyS3FileSystem + + +def test_flytekit_style_kwargs_reach_registered_class(restore_fsspec_registry): + """ + Simulate the flytekit collision: flytekit calls fsspec.filesystem("s3", cache_regions=True). + With s3fs-style registered, those kwargs must reach that class instead of obstore (which + rejects them with "Configuration key: 'cache_regions' is not valid for store 'S3'"). + """ + fsspec.register_implementation("s3", DummyS3FileSystem, clobber=True) + _register_obstore_for_missing_protocols() + + fs = fsspec.filesystem("s3", cache_regions=True) + + assert isinstance(fs, DummyS3FileSystem) + assert fs.cache_regions is True + + +def test_obstore_filesystem_class_mirrors_obstore_register(): + """The flyte-owned class matches what obstore.fsspec.register(asynchronous=True) creates.""" + cls = _obstore_filesystem_class("s3") + + assert issubclass(cls, FsspecStore) + assert cls.protocol == "s3" + assert cls.asynchronous is True + # The per-protocol class is cached. + assert _obstore_filesystem_class("s3") is cls + # fsspec's _Cached metaclass instance caching still applies to direct instantiation. + assert cls() is cls() + + +def test_flyte_protocol_registration_intact(): + """The `flyte` protocol registration is unaffected by the gap-filling helper.""" + from flyte.storage._remote_fs import FlyteFS + + assert fsspec.registry["flyte"] is FlyteFS