[Data] Fix silent credential drop for fsspec-S3 in download expression

[xyuzh]

Summary

When a user passes an fsspec S3FileSystem with Okta/STS/profile-based auth (wrapped as PyFileSystem(FSSpecHandler(s3fs_fs))), _extract_credentials_from_filesystem read only static attrs (key/secret/token, storage_options). Those attrs are None for session-backed s3fs — credentials live on fs.session.get_credentials() and may rotate. The function returned {}, obstore's from_url got no keys, and obstore silently fell back to its own credential chain (IMDS on EC2). Under concurrent S3 workloads this manifested as intermittent NoCredentialsError — users never saw a warning, and their Okta credentials had been silently dropped.

This is the first of two independent PRs addressing related bugs in the download expression pipeline. The matching threaded-path pre-resolve fix (which avoids an IMDS thundering herd on the PyArrow fallback path, and is triggered whenever this PR routes an unextractable filesystem to threaded) is in a separate PR. Either can land first.

Changes

• _extract_credentials_from_filesystem now returns Optional[Dict]. None signals "route to threaded" so the user's filesystem stays authoritative. Unrecognized non-None filesystems also return None per the new contract.
• _frozen_s3fs_credentials (new helper) snapshots credentials via session.get_credentials().get_frozen_credentials() for both sync (botocore) and async (aiobotocore) sessions.
• _plan_obstore_routing(fs) -> (use_obstore, fs_kwargs) (new helper) centralizes dispatch at plan time with a one-shot WARNING per filesystem when credentials can't be extracted.
• plan_download_op wires the routing helper so both the partition actor choice (AsyncPartitionActor vs PartitionActor) and the download function choice (download_bytes_async vs download_bytes_threaded) use the same gate.
• download_bytes_async extracts credentials once in sync context and threads the kwargs through to _download_uris_with_obstore. Fixes a latent correctness issue where aiobotocore's async get_credentials can't run from inside an already-running event loop.
• AsyncPartitionActor and _download_uris_with_obstore both fail closed with a clear RuntimeError when constructed / called with a filesystem whose credentials can't be statically extracted. Dispatch should route such filesystems to the PyArrow path; reaching these code paths anyway indicates a dispatch bug and the guard prevents regressions back to the silent-drop behavior.

Test plan

• pre-commit run passes (ruff, pydoclint, black, docstyle, semgrep, import order, mock-method / logger checks).
• New unit tests in test_obstore_download.py:
  • TestSessionBackedFsspecCredentials — sync botocore session, aiobotocore async session, _session fallback for older s3fs, unresolvable-session → None, no-access-key → None, anon=True skips the session, static attrs win over session.
  • TestPlanObstoreRouting — None filesystem → obstore, non-S3 fsspec → threaded, unextractable fsspec-S3 → threaded with warning, warning dedup (same FS twice = one warning), extractable fsspec-S3 → obstore, AsyncPartitionActor raises on unextractable creds, _download_uris_with_obstore raises on unextractable creds.
• End-to-end against moto/minio with Okta-style fsspec S3FileSystem over ~100 URIs (no NoCredentialsError, moto sees signed requests with the session's current keys).

[gemini-code-assist]

Code Review

This pull request enhances the obstore download path in Ray Data by implementing a robust credential extraction and routing mechanism. It introduces _plan_obstore_routing to decide between obstore and threaded PyArrow downloads, specifically handling session-backed s3fs configurations (e.g., Okta, STS) where credentials are resolved lazily. The changes ensure that if credentials cannot be extracted, the system falls back to the authoritative user filesystem rather than silently using default credentials. Review feedback suggests using isawaitable instead of iscoroutine to better handle various asynchronous return types from fsspec implementations.

[xyuzh]

@goutamvenkat-anyscale

[goutamvenkat-anyscale]

General Comment:

Might be better to use https://developmentseed.org/obstore/latest/authentication/ credential_provider so that the authentication mechanism can refresh the creds on expiry instead of just using a single cached value.

from botocore.credentials import Credentials, ReadOnlyCredentials
from cachetools import TLRUCache

class _S3FSSession(Protocol):
    def get_credentials(
        self,
    ) -> Optional[Union["Credentials", Awaitable["Credentials"]]]: ...


class _S3FSSessionCredentialProvider:
    """obstore credential provider backed by an s3fs-style session."""

    _CACHE_KEY = "credentials"
    _DEFAULT_TTL = timedelta(minutes=30)

    def __init__(
        self, session: _S3FSSession, ttl: Optional[timedelta] = None
    ) -> None:
        self._session = session
        self._ttl = self._DEFAULT_TTL if ttl is None else ttl
        self._cache = TLRUCache(
            maxsize=1,
            ttu=lambda _key, value, _now: value["expires_at"],
            timer=lambda: datetime.now(timezone.utc),
        )
        self._cache_lock = RLock()

    @classmethod
    def from_fsspec_fs(
        cls, fsspec_fs: Any, *, validate: bool = False
    ) -> Optional["_S3FSSessionCredentialProvider"]:
        session = getattr(fsspec_fs, "session", None) or getattr(
            fsspec_fs, "_session", None
        )
        if session is None or not hasattr(session, "get_credentials"):
            return None

        provider = cls(cast(_S3FSSession, session))
        if validate and not provider.can_fetch_credentials():
            return None
        return provider

    async def __call__(self) -> Dict[str, Any]:
        with self._cache_lock:
            cached = self._cache.get(self._CACHE_KEY)
            if cached is not None:
                return dict(cached)

        session_credentials = await self._get_session_credentials()
        frozen_credentials = await self._get_frozen_credentials(session_credentials)
        credentials = self._to_obstore_credentials(
            session_credentials, frozen_credentials
        )

        with self._cache_lock:
            self._cache[self._CACHE_KEY] = credentials
        return dict(credentials)

    async def _get_session_credentials(self) -> "Credentials":
        credentials = self._session.get_credentials()
        if inspect.isawaitable(credentials):
            credentials = await credentials
        if credentials is None:
            raise RuntimeError("fsspec S3 session returned no credentials")
        return credentials

    async def _get_frozen_credentials(
        self, session_credentials: "Credentials"
    ) -> "ReadOnlyCredentials":
        frozen_credentials = session_credentials.get_frozen_credentials()
        if inspect.isawaitable(frozen_credentials):
            frozen_credentials = await frozen_credentials
        if not frozen_credentials.access_key:
            raise RuntimeError("fsspec S3 session returned no access key")
        return cast("ReadOnlyCredentials", frozen_credentials)

    def _to_obstore_credentials(
        self,
        session_credentials: "Credentials",
        frozen_credentials: "ReadOnlyCredentials",
    ) -> Dict[str, Any]:
        credentials: Dict[str, Any] = {
            "access_key_id": frozen_credentials.access_key,
            "secret_access_key": frozen_credentials.secret_key or "",
            "expires_at": self._expires_at(session_credentials),
        }
        if frozen_credentials.token:
            credentials["token"] = frozen_credentials.token
        return credentials

    def _expires_at(self, session_credentials: "Credentials") -> datetime:
        for attr in ("expiry_time", "_expiry_time"):
            expires_at = getattr(session_credentials, attr, None)
            if isinstance(expires_at, datetime):
                if expires_at.tzinfo is None:
                    return expires_at.replace(tzinfo=timezone.utc)
                return expires_at
        return datetime.now(timezone.utc) + self._ttl

    def can_fetch_credentials(self) -> bool:
        """Return True if credentials can be fetched from sync code now."""
        try:
            asyncio.run(self())
            return True
        except Exception as e:
            logger.debug("Could not fetch fsspec session credentials: %r", e)
            return False

let me know what you think

[goutamvenkat-anyscale]

Can you expand more on why this function is needed?

[xyuzh]

_await was only needed to drive Tasks/Futures through asyncio.run from sync code; the new
_S3FSSessionCredentialProvider.__call__ runs inside obstore's event loop and uses plain await via _await_maybe, so the wrapper is no longer required.

[goutamvenkat-anyscale]

Can we strongly type this param?

[xyuzh]

The helper was replaced by _S3FSSessionCredentialProvider.from_fsspec_fs(fsspec_fs: "s3fs.S3FileSystem").

[goutamvenkat-anyscale]

We can replace the getattr lookups. Does this not suffice?

def _is_fsspec_s3_filesystem(filesystem: "pyarrow.fs.FileSystem") -> bool:
    """True if *filesystem* is a PyFileSystem wrapping an fsspec s3/s3a handler."""
    if isinstance(filesystem, RetryingPyFileSystem):
        filesystem = filesystem.unwrap()

    if not isinstance(filesystem, pyarrow.fs.PyFileSystem):
        return False
    if not isinstance(filesystem.handler, pyarrow.fs.FSSpecHandler):
        return False

    try:
        from s3fs import S3FileSystem
    except ImportError:
        return False
    return isinstance(filesystem.handler.fs, S3FileSystem)

[xyuzh]

Removed getattr lookup

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 87d1442. Configure here.}

[cursor]

Redundant credential validation with I/O discarded at plan time

Low Severity

_plan_obstore_routing(filesystem) in plan_download_op creates a _S3FSSessionCredentialProvider, calls can_fetch_credentials() (which performs asyncio.run → real session get_credentials() potentially involving network I/O to STS/Okta), and then discards the result with _. The identical validation and provider creation is repeated on every worker in download_bytes_async. For session-backed filesystems requiring network calls, this adds unnecessary latency to the planning phase with no reuse of the fetched credentials.

Additional Locations (1)

• python/ray/data/_internal/planner/_obstore_download.py#L613-L618

 

^{Reviewed by Cursor Bugbot for commit 87d1442. Configure here.}

[cursor]

Race condition in async credential provider cache refresh

Low Severity

_S3FSSessionCredentialProvider.__call__ releases the lock after detecting a cache miss (line 162) before performing the credential fetch (line 164). Multiple concurrent callers (from obstore's async runtime) can all observe the cache miss and simultaneously call self._session.get_credentials(). For botocore sessions that are not guaranteed thread-safe, this thundering-herd pattern could cause undefined behavior under concurrent refresh. A double-checked lock after fetching, or holding the lock across the full refresh (using an asyncio lock instead of a threading lock), would be more robust.

 

^{Reviewed by Cursor Bugbot for commit 87d1442. Configure here.}

[goutamvenkat-anyscale]

Looks good. Just running the release test to see if there's any regression