fix(s3): prevent per-shard client and cache growth

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)
 

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.<region>.<provider>" 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.<region>.<provider>" 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):