diff --git a/src/litdata/streaming/client.py b/src/litdata/streaming/client.py index 18d6182e..5cdf4cd5 100644 --- a/src/litdata/streaming/client.py +++ b/src/litdata/streaming/client.py @@ -15,7 +15,9 @@ import logging import os import random +import re import threading +from datetime import datetime, timezone from time import sleep, time from typing import Any @@ -41,13 +43,17 @@ # Default timeout for each HTTP request in seconds _DEFAULT_REQUEST_TIMEOUT = 30 # seconds -# The control plane mints credentials with a 1 hour TTL for S3 project-role data connections, and -# the response carries no expiry for us to read, so refresh well inside it. The remaining time is -# the window in which a failed refresh can be retried while the credentials in hand still work. +# Fallback for a control plane that reports no expiry: assume the 1 hour TTL S3 project-role +# connections have always had, and refresh well inside it. Also an upper bound on how long any +# credentials are held, so a longer reported TTL does not stretch the window between refreshes. _DEFAULT_REFETCH_INTERVAL = 2700 # seconds +# Fraction of a reported lifetime to hold credentials for. The remainder is the window in which +# a failed refresh can be retried while the credentials in hand still work. 0.75 of the 1 hour +# TTL is the 2700s above, so a control plane that reports its expiry changes nothing. +_REFETCH_FRACTION = 0.75 # How long past the refetch interval we keep serving existing credentials while refreshes fail. -# Sized against the TTL, not comfort: 2700 + 600 leaves ~5 minutes before the 1 hour S3 expiry, -# so we stop before reads start failing as unexplained S3 403s. +# Sized against the TTL, not comfort: 2700 + 600 leaves ~5 minutes before a 1 hour expiry, so we +# stop before reads start failing as unexplained S3 403s. A reported expiry bounds this directly. _REFRESH_GRACE_PERIOD = 600 # seconds # How long to wait for the control plane when there are no credentials yet. No TTL constrains # this one — nothing is being served — so it can be more patient than the refresh grace. @@ -60,6 +66,48 @@ _REFETCH_JITTER_RATIO = 0.1 +def _parse_reported_expiry(value: Any) -> float | None: + """Parse the control plane's RFC 3339 ``expiresAt`` into a unix timestamp. + + Anything unreadable is reported as absent rather than raised: a deadline we cannot + parse should fall back to the assumed TTL, not fail the read. + """ + if not isinstance(value, str) or not value.strip(): + return None + text = value.strip() + # The proto3 JSON mapping emits Z-normalized RFC 3339 with 0, 3, 6 or 9 fractional digits. + # `fromisoformat` rejects the Z before 3.11, and the 9-digit form on every version we support. + if text[-1] in "Zz": + text = text[:-1] + "+00:00" + text = re.sub(r"(\.\d{6})\d+", r"\1", text) + try: + parsed = datetime.fromisoformat(text) + except ValueError: + logger.warning("Ignoring unparsable credential expiry %r from the control plane", value) + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.timestamp() + + +def _credentials_expiry(creds: dict[str, Any] | None) -> float | None: + """Unix timestamp at which ``creds`` stop working, or ``None`` when none was reported.""" + if not creds: + return None + return _parse_reported_expiry(creds.get("expiresAt")) + + +def _refetch_interval_for(expires_at: float | None, fetched_at: float, upper_bound: float) -> float: + """How long after ``fetched_at`` credentials expiring at ``expires_at`` should be replaced. + + ``upper_bound`` still applies once an expiry is known, because it is what the caller + asked for: a 12 hour R2 lifetime should not stretch the gap between refreshes to match. + """ + if expires_at is None: + return upper_bound + return max(0.0, min(upper_bound, (expires_at - fetched_at) * _REFETCH_FRACTION)) + + class _CredentialsError(RuntimeError): """Raised when credentials could not be obtained from the control plane or IMDS. @@ -161,7 +209,8 @@ def _cached_temp_bucket_credentials( cached = _temp_creds_cache.get(data_connection_id) if cached is not None: fetched_at, creds = cached - if time() - fetched_at < _DEFAULT_REFETCH_INTERVAL: + reuse_for = _refetch_interval_for(_credentials_expiry(creds), fetched_at, _DEFAULT_REFETCH_INTERVAL) + if time() - fetched_at < reuse_for: return fetched_at, dict(creds) creds = _fetch_temp_bucket_credentials(data_connection_id) fetched_at = time() @@ -175,7 +224,8 @@ def _login_and_get_temp_bucket_credentials(data_connection_id: str, *, force_ref Shared by R2 (lightning storage) connections and by S3 connections marked ``available_in_non_aws_providers``: both bypass the FUSE mount and need short-lived creds from the same control-plane ``temp-bucket-credentials`` endpoint. Returns the raw response - (accessKeyId/secretAccessKey/sessionToken, plus accountId for R2). + (accessKeyId/secretAccessKey/sessionToken, plus accountId for R2, and expiresAt/region/ + endpoint from control planes new enough to report them). Results are cached per process (see ``_cached_temp_bucket_credentials``). """ @@ -278,15 +328,29 @@ def _reset_process_state(self) -> None: # Guards lazy create + credential refresh (range GETs hit .client from many threads). self._client_lock = threading.Lock() self._owner_pid = os.getpid() + # Set before the deadline below, which is derived from them once credentials exist. + self._creds_fetched_at: float | None = None + self._creds_expires_at: float | None = None self._refetch_deadline = self._jittered_refetch_interval() self._refresh_retry_time: float | None = None self._force_refresh_credentials = False - self._creds_fetched_at: float | None = None def _jittered_refetch_interval(self) -> float: # Only ever early, never late: callers set the interval as an upper bound on how long # a set of credentials is held, and the grace period below is measured against it. - return self._refetch_interval * (1.0 - random.uniform(0.0, _REFETCH_JITTER_RATIO)) # noqa: S311 + interval = _refetch_interval_for( + self._creds_expires_at, + self._creds_fetched_at if self._creds_fetched_at is not None else time(), + self._refetch_interval, + ) + return interval * (1.0 - random.uniform(0.0, _REFETCH_JITTER_RATIO)) # noqa: S311 + + def _record_credentials_window(self, data_connection_id: str, creds: dict[str, Any]) -> None: + """Note when these credentials were minted and when the control plane says they die.""" + with _temp_creds_lock_for_pid(): + cached = _temp_creds_cache.get(data_connection_id) + self._creds_fetched_at = cached[0] if cached is not None else None + self._creds_expires_at = _credentials_expiry(creds) def __getstate__(self) -> dict[str, Any]: state = self.__dict__.copy() @@ -346,9 +410,7 @@ def _create_client_from_temp_credentials(self, data_connection_id: str) -> None: temp_credentials = _login_and_get_temp_bucket_credentials( data_connection_id, force_refresh=self._force_refresh_credentials ) - with _temp_creds_lock_for_pid(): - cached = _temp_creds_cache.get(data_connection_id) - self._creds_fetched_at = cached[0] if cached is not None else None + self._record_credentials_window(data_connection_id, temp_credentials) # data_connection_id is our own metadata; drop it before handing options to boto3. storage_options = {k: v for k, v in self._storage_options.items() if k != "data_connection_id"} @@ -378,13 +440,26 @@ def _create_client_from_temp_credentials(self, data_connection_id: str) -> None: ) def _mark_refreshed(self) -> None: - # Prefer the mint time from the process cache so a new client that reused - # credentials still refreshes before the 1 hour S3/R2 TTL, not 2700s from now. + # Prefer the mint time from the process cache so a new client that reused credentials + # still refreshes before they expire, rather than a full interval from now. fetched_at = getattr(self, "_creds_fetched_at", None) self._last_time = fetched_at if fetched_at is not None else time() self._refetch_deadline = self._jittered_refetch_interval() self._refresh_retry_time = None + def next_refresh_time(self) -> datetime: + """When :attr:`client` will next mint credentials, as an aware UTC datetime. + + Callers that cache credentials of their own need this rather than the expiry, so they + come back while the credentials they hold are still the ones this client is serving. + Bounded by the reported expiry so it can never name a time the credentials are dead by. + """ + last = self._last_time if self._last_time is not None else time() + deadline = last + self._refetch_deadline + if self._creds_expires_at is not None: + deadline = min(deadline, self._creds_expires_at) + return datetime.fromtimestamp(deadline, tz=timezone.utc) + def _create_initial_client(self) -> None: """Create the first client, waiting out a control plane that is briefly unavailable. @@ -440,10 +515,12 @@ def _refresh_client(self) -> None: # misbehaving mid-deploy — and if it is a real revocation, the deadline still catches it. except _CredentialsError as e: held_for = 0.0 if self._last_time is None else now - self._last_time - if held_for > self._refetch_deadline + _REFRESH_GRACE_PERIOD: - raise RuntimeError( - f"Failed to refresh credentials for {held_for:.0f}s, so they are assumed expired: {e}" - ) from e + # Once the control plane has told us when these die there is nothing left to serve + # past that point, and continuing only turns the failure into an opaque S3 403. + expired = self._creds_expires_at is not None and now >= self._creds_expires_at + if expired or held_for > self._refetch_deadline + _REFRESH_GRACE_PERIOD: + reason = "they have expired" if expired else "they are assumed expired" + raise RuntimeError(f"Failed to refresh credentials for {held_for:.0f}s, so {reason}: {e}") from e self._refresh_retry_time = now + _REFRESH_RETRY_INTERVAL logger.warning( "Could not refresh credentials (%.0fs since the last successful refresh); reusing the current " @@ -497,9 +574,7 @@ def get_r2_bucket_credentials(self, data_connection_id: str, *, force_refresh: b """Fetch temporary R2 credentials for the current lightning storage connection.""" try: temp_credentials = _login_and_get_temp_bucket_credentials(data_connection_id, force_refresh=force_refresh) - with _temp_creds_lock_for_pid(): - cached = _temp_creds_cache.get(data_connection_id) - self._creds_fetched_at = cached[0] if cached is not None else None + self._record_credentials_window(data_connection_id, temp_credentials) endpoint_url = f"https://{temp_credentials['accountId']}.r2.cloudflarestorage.com" @@ -528,9 +603,15 @@ def _create_client(self) -> None: cached_client = _temp_creds_boto_clients.get(data_connection_id) if cached_client is not None: fetched_at, boto_client = cached_client - if time() - fetched_at < _DEFAULT_REFETCH_INTERVAL: + # The credentials behind this client are cached separately; their reported + # deadline is what says whether it is still safe to hand back. + with _temp_creds_lock_for_pid(): + cached_creds = _temp_creds_cache.get(data_connection_id) + expires_at = _credentials_expiry(cached_creds[1] if cached_creds is not None else None) + if time() - fetched_at < _refetch_interval_for(expires_at, fetched_at, _DEFAULT_REFETCH_INTERVAL): self._client = boto_client self._creds_fetched_at = fetched_at + self._creds_expires_at = expires_at return # Get R2 credentials (process cache on first use; mint on scheduled refresh). diff --git a/src/litdata/streaming/dataloader.py b/src/litdata/streaming/dataloader.py index 2fbf7388..99ed64d8 100644 --- a/src/litdata/streaming/dataloader.py +++ b/src/litdata/streaming/dataloader.py @@ -275,7 +275,9 @@ def __init__(self, loader: DataLoader) -> None: # Patch PyTorch worker loop to call the `cache.done()` method. from torch.utils.data._utils import worker - worker._worker_loop = WorkerLoop(loader._global_rank, loader._profile) + # WorkerLoop takes *args/**kwargs so it survives torch changing the private + # signature, which mypy checks against the exact arity torch declares today. + worker._worker_loop = WorkerLoop(loader._global_rank, loader._profile) # type: ignore[assignment] super().__init__(loader) def _shutdown_workers(self) -> None: @@ -568,12 +570,14 @@ def __init__(self, loader: DataLoader) -> None: if distributed_env.global_rank == 0: if self._loader._profile_batches and _VIZ_TRACKER_AVAILABLE: original_worker_loop = worker._worker_loop - worker._worker_loop = _ProfileWorkerLoop( + worker._worker_loop = _ProfileWorkerLoop( # type: ignore[assignment] self._loader._profile_batches, self._loader._profile_skip_batches, self._loader._profile_dir ) elif profile_cprofile: original_worker_loop = worker._worker_loop - worker._worker_loop = _CProfileWorkerLoop(_cprofile_output_dir(self._loader._profile_dir)) + worker._worker_loop = _CProfileWorkerLoop( # type: ignore[assignment] + _cprofile_output_dir(self._loader._profile_dir) + ) # Workers fork/spawn here. Enable the parent profiler only after that so # the child does not inherit an active cProfile (one profiler per process). diff --git a/src/litdata/streaming/downloader.py b/src/litdata/streaming/downloader.py index 38f74b80..9de6ccb3 100644 --- a/src/litdata/streaming/downloader.py +++ b/src/litdata/streaming/downloader.py @@ -266,8 +266,6 @@ def _obstore_credential_provider(s3_client: S3Client) -> Any: """ def _provider() -> dict[str, Any]: - from datetime import datetime, timedelta, timezone - boto_client = s3_client.client frozen = boto_client._get_credentials().get_frozen_credentials() if frozen.access_key is None or frozen.secret_key is None: @@ -276,7 +274,10 @@ def _provider() -> dict[str, Any]: "access_key_id": frozen.access_key, "secret_access_key": frozen.secret_key, "token": frozen.token, - "expires_at": datetime.now(timezone.utc) + timedelta(minutes=30), + # obstore holds these until this moment and then asks again, so it has to be when + # the client rolls over. A fixed guess from now outlives credentials that were + # already most of the way through their life when we read them off a warm client. + "expires_at": s3_client.next_refresh_time(), } return _provider diff --git a/tests/streaming/test_client.py b/tests/streaming/test_client.py index 7299f01f..fecf9e05 100644 --- a/tests/streaming/test_client.py +++ b/tests/streaming/test_client.py @@ -1,6 +1,7 @@ import logging import sys import threading +from datetime import datetime, timezone from time import sleep, time from unittest import mock @@ -708,20 +709,28 @@ def test_r2_client_api_call_format(monkeypatch): ) -def _successful_login_session(monkeypatch): - """Wire requests.Session so a full credential fetch succeeds, and hand back the mock.""" +def _successful_login_session(monkeypatch, expires_at=None): + """Wire requests.Session so a full credential fetch succeeds, and hand back the mock. + + ``expires_at`` is left out by default, which is what a control plane predating the field + returns — so every caller that does not ask for one covers the fallback path. + """ login_response = mock.MagicMock() login_response.status_code = 200 login_response.json.return_value = {"token": "test-token"} - credentials_response = mock.MagicMock() - credentials_response.status_code = 200 - credentials_response.json.return_value = { + credentials = { "accessKeyId": "test-access-key", "secretAccessKey": "test-secret-key", "sessionToken": "test-session-token", "accountId": "test-account-id", } + if expires_at is not None: + credentials["expiresAt"] = expires_at + + credentials_response = mock.MagicMock() + credentials_response.status_code = 200 + credentials_response.json.return_value = credentials requests_mock = mock.MagicMock() requests_mock.post = mock.MagicMock(return_value=login_response) @@ -1100,3 +1109,137 @@ def worker() -> None: assert not errors assert requests_mock.post.call_count == 1 + + +@pytest.mark.parametrize( + "reported", + [ + "2026-09-02T17:30:00Z", + "2026-09-02T17:30:00.000Z", + "2026-09-02T17:30:00.000000Z", + "2026-09-02T17:30:00.000000000Z", + "2026-09-02T17:30:00+00:00", + ], +) +def test_reported_expiry_parses_every_form_the_control_plane_emits(reported): + """The proto3 JSON mapping emits whichever of these is shortest for the value. + + An STS expiry lands on a whole second and arrives bare; R2's is derived from a wall + clock and arrives with nanoseconds, which `fromisoformat` will not take. + """ + assert client._parse_reported_expiry(reported) == datetime(2026, 9, 2, 17, 30, tzinfo=timezone.utc).timestamp() + + +@pytest.mark.parametrize("reported", [None, "", " ", "whenever", 1788370200]) +def test_unreadable_expiry_is_reported_as_absent(reported): + """A deadline we cannot read falls back to the assumed TTL rather than failing the read.""" + assert client._parse_reported_expiry(reported) is None + + +def test_reported_expiry_shortens_the_refetch_interval(): + """A credential that dies sooner than assumed has to be replaced sooner.""" + minted = 1000.0 + dies_in_20_minutes = minted + 1200 + + interval = client._refetch_interval_for(dies_in_20_minutes, minted, client._DEFAULT_REFETCH_INTERVAL) + + assert interval == 900 # 0.75 of the 20 minutes it actually has + assert interval < client._DEFAULT_REFETCH_INTERVAL + + +def test_reported_expiry_does_not_stretch_the_refetch_interval(): + """R2 lifetimes run to 12 hours; holding one set of credentials that long is not the fix here.""" + minted = 1000.0 + dies_in_12_hours = minted + 12 * 3600 + + interval = client._refetch_interval_for(dies_in_12_hours, minted, client._DEFAULT_REFETCH_INTERVAL) + + assert interval == client._DEFAULT_REFETCH_INTERVAL + + +def test_missing_expiry_falls_back_to_the_assumed_ttl(): + """Control planes predating the field must behave exactly as before.""" + assert client._refetch_interval_for(None, 1000.0, client._DEFAULT_REFETCH_INTERVAL) == ( + client._DEFAULT_REFETCH_INTERVAL + ) + + +def test_already_expired_credentials_are_never_reused(): + assert client._refetch_interval_for(500.0, 1000.0, client._DEFAULT_REFETCH_INTERVAL) == 0.0 + + +def _rfc3339(timestamp): + """Format a unix timestamp the way the control plane reports ``expiresAt``.""" + return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def test_cached_credentials_are_refetched_on_the_reported_expiry(monkeypatch): + """The process cache holds a short-lived credential for its life, not the assumed 45 minutes.""" + _mock_login_env(monkeypatch) + now = {"t": 1000.0} + monkeypatch.setattr(client, "time", lambda: now["t"]) + # 20 minutes, so the cache should let go at 900s rather than the 2700s default. + requests_mock = _successful_login_session(monkeypatch, expires_at=_rfc3339(now["t"] + 1200)) + + client._login_and_get_temp_bucket_credentials("conn-short") + now["t"] = 1899.0 + client._login_and_get_temp_bucket_credentials("conn-short") + assert requests_mock.post.call_count == 1 + + now["t"] = 1900.0 + client._login_and_get_temp_bucket_credentials("conn-short") + assert requests_mock.post.call_count == 2 + + +def test_next_refresh_time_never_outlives_credentials_read_off_a_warm_client(monkeypatch): + """The regression this exists for. + + A client built late in a cached credential's life used to promise obstore a flat 30 more + minutes, which ran past the point the credential stopped working — and the read then failed + as an unexplained InvalidAccessKeyId rather than triggering a refresh. + """ + _mock_login_env(monkeypatch) + monkeypatch.setattr(client, "_REFETCH_JITTER_RATIO", 0.0) + monkeypatch.setattr(client, "boto3", mock.MagicMock()) + monkeypatch.setattr(client, "botocore", mock.MagicMock()) + + minted = 1000.0 + expires_at = minted + 3600 + now = {"t": minted} + monkeypatch.setattr(client, "time", lambda: now["t"]) + _successful_login_session(monkeypatch, expires_at=_rfc3339(expires_at)) + + first = client.R2Client(storage_options={"data_connection_id": "conn-warm"}) + assert first.client is not None + + # 40 minutes in: inside the cache window, so a new client inherits these credentials with + # only 20 minutes left on them. + now["t"] = minted + 2400 + second = client.R2Client(storage_options={"data_connection_id": "conn-warm"}) + assert second.client is not None + + deadline = second.next_refresh_time().timestamp() + assert deadline <= expires_at + assert deadline < now["t"] + 30 * 60 # what the old fixed guess would have promised + + +def test_next_refresh_time_falls_back_to_the_client_schedule(monkeypatch): + """Without a reported expiry there is still an honest answer: when this client rolls over.""" + monkeypatch.setattr(client, "_REFETCH_JITTER_RATIO", 0.0) + monkeypatch.setattr(client, "boto3", mock.MagicMock()) + monkeypatch.setattr(client, "botocore", mock.MagicMock()) + + s3 = client.S3Client(refetch_interval=600, storage_options={"region_name": "us-east-1"}) + assert s3.client is not None + + assert s3._creds_expires_at is None + assert s3.next_refresh_time().timestamp() == pytest.approx(s3._last_time + 600, abs=1) + + +def test_refresh_gives_up_once_the_reported_expiry_passes(monkeypatch): + """Past a known expiry there is nothing left to serve, so say so instead of retrying.""" + s3, _, _ = _client_with_failing_refresh(monkeypatch) + s3._creds_expires_at = time() - 1 + + with pytest.raises(RuntimeError, match="they have expired"): + _ = s3.client diff --git a/tests/streaming/test_downloader.py b/tests/streaming/test_downloader.py index 37173295..31ad1bf0 100644 --- a/tests/streaming/test_downloader.py +++ b/tests/streaming/test_downloader.py @@ -1,6 +1,7 @@ import contextlib import io import os +from datetime import datetime, timedelta, timezone from unittest import mock from unittest.mock import MagicMock @@ -16,6 +17,7 @@ R2Downloader, S3Downloader, _indexed_object_bytes, + _obstore_credential_provider, _range_parts, get_downloader, register_downloader, @@ -832,3 +834,35 @@ def test_r2_tiny_chunk_uses_get_object(r2_client_mock, monkeypatch, tmpdir): client.client.get_object.assert_called_once_with(Bucket="bucket", Key="chunk-0-0.zstd.bin") client.client.download_file.assert_not_called() get_store.assert_not_called() + + +def test_obstore_credential_provider_reports_the_client_refresh_time(): + """Obstore must be told when the client rolls over, not a flat guess from now. + + It caches what the provider hands back until ``expires_at``, so a guess that outlives + credentials read off a warm client leaves it signing with ones already dead. + """ + rolls_over_at = datetime.now(timezone.utc) + timedelta(minutes=4) + + s3_client = MagicMock() + s3_client.next_refresh_time.return_value = rolls_over_at + frozen = s3_client.client._get_credentials.return_value.get_frozen_credentials.return_value + frozen.access_key = "AKIATEST" + frozen.secret_key = "secret" + frozen.token = "token" + + credentials = _obstore_credential_provider(s3_client)() + + assert credentials["access_key_id"] == "AKIATEST" + assert credentials["expires_at"] == rolls_over_at + s3_client.next_refresh_time.assert_called_once_with() + + +def test_obstore_credential_provider_rejects_incomplete_credentials(): + s3_client = MagicMock() + frozen = s3_client.client._get_credentials.return_value.get_frozen_credentials.return_value + frozen.access_key = None + frozen.secret_key = "secret" + + with pytest.raises(ValueError, match="incomplete credentials"): + _obstore_credential_provider(s3_client)()