diff --git a/wsds/pupyarrow/file_reader.py b/wsds/pupyarrow/file_reader.py index 142b3ec..20ccbd0 100644 --- a/wsds/pupyarrow/file_reader.py +++ b/wsds/pupyarrow/file_reader.py @@ -12,6 +12,10 @@ BLOCK_SIZE = 8192 # 8kB minimum sync read size MIN_ASYNC_READ = 4096 # 4kB minimum async read size +ASYNC_CACHE_MAX_BYTES = int(os.environ.get("WSDS_ASYNC_CACHE_MAX_BYTES", 8 * 1024 * 1024)) +# Additional attempts after the initial presigned range request. Set to zero +# when tail latency matters more than recovering transient S3 failures. +S3_PRESIGNED_RETRIES = max(0, int(os.environ.get("WSDS_S3_PRESIGNED_RETRIES", 3))) VERBOSE = False PRESIGN_EXPIRES = 3600 # presigned URL lifetime (seconds) @@ -86,6 +90,16 @@ def __init__(self): self._planned: contextvars.ContextVar[bool] = contextvars.ContextVar("_planned", default=False) self._pending: list[tuple[int, int, int, asyncio.Future]] = [] # (offset, actual, length, future) self._cache: list[tuple[int, bytes]] = [] # (offset, data) ranges + self._cache_bytes = 0 + + def _cache_put(self, offset: int, data: bytes) -> None: + if not data or len(data) > ASYNC_CACHE_MAX_BYTES: + return + self._cache.append((offset, data)) + self._cache_bytes += len(data) + while self._cache_bytes > ASYNC_CACHE_MAX_BYTES: + _, evicted = self._cache.pop(0) + self._cache_bytes -= len(evicted) def read(self, offset: int, length: int) -> bytes: """Read length bytes at absolute offset, using forward cache.""" @@ -160,7 +174,7 @@ async def async_read(self, offset: int, length: int) -> bytes: if not self._planned.get(): # Eager mode: read directly data = await self._async_read_impl(offset, actual) - self._cache.append((offset, data)) + self._cache_put(offset, data) return data[:length] # Plan mode: submit and await future @@ -178,7 +192,7 @@ async def flush(self): data_map: dict[int, bytes] = {} for region, data in zip(regions, fetched): - self._cache.append((region.offset, data)) + self._cache_put(region.offset, data) for abs_offset, start, end in region.members: data_map[abs_offset] = data[start:end] @@ -189,6 +203,7 @@ async def flush(self): def clear_cache(self): """Clear the async range cache.""" self._cache.clear() + self._cache_bytes = 0 @property def has_pending(self) -> bool: @@ -309,7 +324,7 @@ async def _presigned_read(self, offset: int, length: int) -> bytes | None: session = await _get_http_session() headers = {"Range": f"bytes={offset}-{offset + length - 1}"} auth_retried = False - for attempt in range(4): + for attempt in range(S3_PRESIGNED_RETRIES + 1): if self._url is None or time.monotonic() > self._url_deadline: await self._presign() try: @@ -323,10 +338,10 @@ async def _presigned_read(self, offset: int, length: int) -> bytes | None: r.raise_for_status() return await r.read() except aiohttp.ClientResponseError as e: - if e.status < 500 or attempt == 3: + if e.status < 500 or attempt == S3_PRESIGNED_RETRIES: raise except (aiohttp.ClientError, asyncio.TimeoutError): - if attempt == 3: + if attempt == S3_PRESIGNED_RETRIES: raise await asyncio.sleep(0.1 * 2**attempt) diff --git a/wsds/ws_s3_shard.py b/wsds/ws_s3_shard.py index b804b88..7484233 100644 --- a/wsds/ws_s3_shard.py +++ b/wsds/ws_s3_shard.py @@ -1,6 +1,7 @@ import atexit import os import re +import threading import typing from typing import TYPE_CHECKING, Optional, Tuple from urllib.parse import urlparse @@ -15,6 +16,27 @@ from .ws_dataset import WSDataset +_s3_clients: dict[tuple[str | None, ...], tuple[typing.Any, typing.Any]] = {} +_s3_clients_pid: int | None = None +_s3_clients_atexit_pid: int | None = None +_s3_clients_lock = threading.Lock() + + +def _close_s3_clients(): + """Close this process's shared aiobotocore clients at interpreter exit.""" + if _s3_clients_pid != os.getpid(): + return + + from .pupyarrow.file_reader import _get_io_loop + + for _, ctx in _s3_clients.values(): + try: + _get_io_loop().run(ctx.__aexit__(None, None, None)) + except Exception: + pass + _s3_clients.clear() + + def create_s3_client(link=None): """Create a shared aiobotocore S3 client. @@ -31,33 +53,54 @@ def create_s3_client(link=None): from .pupyarrow.file_reader import _get_io_loop + global _s3_clients_pid, _s3_clients_atexit_pid + link = link or {} - kwargs = {"config": Config(max_pool_connections=50, signature_version="s3v4")} endpoint_url = link.get("endpoint_url") or os.environ.get("WSDS_S3_ENDPOINT_URL") - if endpoint_url: - kwargs["endpoint_url"] = endpoint_url - for k in ("aws_access_key_id", "aws_secret_access_key", "aws_session_token", "region_name"): - if link.get(k): - kwargs[k] = link[k] - if "region_name" not in kwargs and endpoint_url: - # SigV4 embeds the region in the credential scope, so it must match - # the endpoint. Derive it from "s3.." hostnames - # (e.g. s3.us-east-005.backblazeb2.com); real regions contain "-", - # which also excludes bare hosts like s3.amazonaws.com. - m = re.match(r"https?://s3\.([a-z0-9-]+)\.", endpoint_url) - if m and "-" in m.group(1): - kwargs["region_name"] = m.group(1) - ctx = AioSession().create_client("s3", **kwargs) - client = _get_io_loop().run(ctx.__aenter__()) - - def _cleanup(): - try: - _get_io_loop().run(ctx.__aexit__(None, None, None)) - except Exception: - pass - - atexit.register(_cleanup) - return client, ctx + client_key = ( + endpoint_url, + link.get("aws_access_key_id"), + link.get("aws_secret_access_key"), + link.get("aws_session_token"), + link.get("region_name"), + ) + + with _s3_clients_lock: + pid = os.getpid() + if _s3_clients_pid != pid: + # A DataLoader worker may be forked after its parent opened S3. + # Its inherited client is tied to the parent's event loop, so do + # not reuse or close it from the child process. + _s3_clients.clear() + _s3_clients_pid = pid + _s3_clients_atexit_pid = None + + existing = _s3_clients.get(client_key) + if existing is not None: + return existing + + if _s3_clients_atexit_pid != pid: + atexit.register(_close_s3_clients) + _s3_clients_atexit_pid = pid + + kwargs = {"config": Config(max_pool_connections=50, signature_version="s3v4")} + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + for k in ("aws_access_key_id", "aws_secret_access_key", "aws_session_token", "region_name"): + if link.get(k): + kwargs[k] = link[k] + if "region_name" not in kwargs and endpoint_url: + # SigV4 embeds the region in the credential scope, so it must match + # the endpoint. Derive it from "s3.." hostnames + # (e.g. s3.us-east-005.backblazeb2.com); real regions contain "-", + # which also excludes bare hosts like s3.amazonaws.com. + m = re.match(r"https?://s3\.([a-z0-9-]+)\.", endpoint_url) + if m and "-" in m.group(1): + kwargs["region_name"] = m.group(1) + ctx = AioSession().create_client("s3", **kwargs) + client = _get_io_loop().run(ctx.__aenter__()) + _s3_clients[client_key] = (client, ctx) + return client, ctx class WSS3Shard(WSShardInterface):