From 3f696e0da5a8ab53e6ad4fa6560c6c6e3c048893 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:34:26 +0900 Subject: [PATCH 01/11] Size the asyncio filesystem's threads by max_workers and add open_async() AioS3FileSystem ran its blocking calls with asyncio.to_thread(), so the default executor of the event loop capped its parallel requests, and a write through a loop whose default executor was busy could deadlock. - Add S3AioThreadPool, a lazily started thread pool for the blocking calls of asyncio code that keeps the caller's context variables and creates new threads in a forked child. - AioS3FileSystem runs the blocking calls of its coroutines in a pool of max_workers threads of its own; a multipart copy runs its requests in a pool sized by the max_workers of the call, shut down after a detached cleanup ends. - S3AioExecutor dispatches through the event loop to a pool of its own, and shuts it down once the submitted functions are done. - Add AioS3FileSystem.open_async(), returning AioS3StreamedFile, whose reads and writes are coroutines: parallel range reads, awaited backpressure of at most max_workers parts in flight, and an abort on an exception, a failure, or a cancellation. - Extract the argument checks and lookups of S3File into S3FileTarget, shared by both file types. Refs #1159 Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 85 ++- docs/api/filesystem.rst | 9 + docs/filesystem.md | 3 +- pyathena/filesystem/s3.py | 284 +++++--- pyathena/filesystem/s3_async.py | 636 ++++++++++++++++-- pyathena/filesystem/s3_executor.py | 199 +++++- tests/pyathena/filesystem/test_s3_async.py | 383 ++++++++++- tests/pyathena/filesystem/test_s3_executor.py | 202 +++++- tests/pyathena/util.py | 29 + 9 files changed, 1602 insertions(+), 228 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 97599137e..a4f2a93c0 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -249,43 +249,84 @@ See each cursor's documentation page for detailed usage examples. ## AioS3FileSystem `AioS3FileSystem` is a native asyncio filesystem interface for Amazon S3, built on -fsspec's `AsyncFileSystem`. It provides the same functionality as `S3FileSystem` but -uses `asyncio.gather` with `asyncio.to_thread` for parallel operations instead of -`ThreadPoolExecutor`. - -### Why AioS3FileSystem? - -The synchronous `S3FileSystem` uses `ThreadPoolExecutor` for parallel S3 operations -(batch deletes, multipart uploads, range reads). When used from within an asyncio -application via `AioS3FSCursor`, this creates a thread-in-thread pattern: -the cursor wraps calls in `asyncio.to_thread()`, and inside that thread -`S3FileSystem` spawns additional threads via `ThreadPoolExecutor`. - -`AioS3FileSystem` eliminates this inefficiency by dispatching all parallel -operations through the asyncio event loop. +fsspec's `AsyncFileSystem`. It provides the same functionality as `S3FileSystem`, and +its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized by +`max_workers`, instead of the default executor of the event loop. | | S3FileSystem | AioS3FileSystem | |---|---|---| -| **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` + `asyncio.to_thread` | -| **File handles** | `S3File` with thread pool | `AioS3File` with `S3AioExecutor` | +| **Blocking calls** | Caller's thread | Thread pool of the filesystem | +| **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` on a thread pool | +| **File handles** | `S3File` with thread pool | `AioS3File` with `S3AioExecutor`; `AioS3StreamedFile` from `open_async()` | | **Bulk delete** | Thread pool per batch | `asyncio.gather` per batch | -| **Multipart copy** | Thread pool per part | `asyncio.gather` per part | +| **Multipart copy** | Thread pool per part | Asyncio task per part, on a thread pool of the copy | | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) | +### Threads + +The threads of `AioS3FileSystem` belong to `S3AioThreadPool` instances, not to the +default executor of the event loop, so `loop.set_default_executor()` neither limits nor +blocks them: + +- The filesystem has a pool of `max_workers` threads for the blocking calls of its + coroutines, such as `_info()`, `_cat_file()`, and the requests of `_rm()`. +- A multipart copy has a pool of its own, sized by the `max_workers` of the call. +- Each file opened with `open()` or `open_async()` has a pool of its own, sized by the + `max_workers` of the call. + +The pools create threads as the calls need them, and each call runs with the caller's +context variables, as with `asyncio.to_thread()`. +In a forked child process, a pool creates new threads. + ### Executor strategy `S3FileSystem` and `S3File` use a pluggable executor abstraction (`S3Executor`) for parallel operations. Two implementations are provided: - `S3ThreadPoolExecutor` — wraps `ThreadPoolExecutor` (default for sync usage) -- `S3AioExecutor` — dispatches work via `asyncio.run_coroutine_threadsafe` + `asyncio.to_thread` +- `S3AioExecutor` — dispatches work via `asyncio.run_coroutine_threadsafe` to an + `S3AioThreadPool` of its own -`AioS3FileSystem` automatically uses `S3AioExecutor` for file handles, so multipart -uploads and parallel range reads are dispatched through the event loop with -`asyncio.to_thread()` instead of a separate `ThreadPoolExecutor` per file. -At most `max_workers` of them run at once. +`AioS3FileSystem` automatically uses `S3AioExecutor` for the file handles of `open()`, +so multipart uploads and parallel range reads are dispatched through the event loop to +a pool of `max_workers` threads of each file. An instance created with `asynchronous=True` has no event loop of its own, so its file handles use `S3ThreadPoolExecutor`. +### Async file handles + +`open_async()` returns an `AioS3StreamedFile`, whose `read()`, `write()`, and `close()` +are coroutines. It supports the `rb`, `wb`, `ab`, and `xb` modes: + +```python +from pyathena.filesystem.s3_async import AioS3FileSystem + +fs = AioS3FileSystem(asynchronous=True) + +async with await fs.open_async("s3://my-bucket/data/file.bin", "wb") as f: + await f.write(data) + +async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: + data = await f.read() +``` + +- A read splits the requested range into ranges of the block size, and reads them in + parallel. The file has no read cache. +- A write uploads each full block as a part of a multipart upload, with at most + `max_workers` parts in flight. A write that would upload another part awaits one of + them, so the event loop keeps running. A part found failed by that wait stops the + write with its error, and the multipart upload is aborted. A part that fails + afterward raises its error when the file is closed. +- Closing the file commits the written data, as `S3File` does. +- Leaving an `async with` block with an exception, or cancelling a write, closes the file + without committing it, and aborts the multipart upload after the running parts finish. + The abort completes even if the task is cancelled again. +- A commit that has started completes even if the task that closes the file is + cancelled. +- A file opened for writing must be closed with `await f.close()` or an `async with` + block. A file that is garbage collected without being closed is not committed. +- `open_async()` raises `NotImplementedError` for the write modes inside a transaction, + because the file commits when it is closed. + ### Usage with AioS3FSCursor `AioS3FSCursor` automatically uses `AioS3FileSystem` internally. No additional diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index d9d2772ee..722efcd4e 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -22,6 +22,9 @@ S3 FileSystem .. autoclass:: pyathena.filesystem.s3.S3File :members: +.. autoclass:: pyathena.filesystem.s3.S3FileTarget + :members: + Async S3 FileSystem ------------------- @@ -31,6 +34,9 @@ Async S3 FileSystem .. autoclass:: pyathena.filesystem.s3_async.AioS3File :members: +.. autoclass:: pyathena.filesystem.s3_async.AioS3StreamedFile + :members: + Read-ahead Stream ----------------- @@ -49,6 +55,9 @@ S3 Executor .. autoclass:: pyathena.filesystem.s3_executor.S3AioExecutor :members: +.. autoclass:: pyathena.filesystem.s3_executor.S3AioThreadPool + :members: + S3 Paths -------- diff --git a/docs/filesystem.md b/docs/filesystem.md index acc96f81f..1855bb2f9 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -625,5 +625,6 @@ for source, destination in pairing.copy_pairs(sources): ## Async filesystem `AioS3FileSystem` provides the same functionality on top of fsspec's -`AsyncFileSystem`, dispatching parallel operations through the asyncio event loop. +`AsyncFileSystem`, running its blocking calls in thread pools sized by `max_workers`. +Its `open_async()` returns a file whose reads and writes are coroutines. See {ref}`aio-s3-filesystem` for details. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index ab3121f4c..2ad7d1b1a 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -11,6 +11,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, as_completed, wait from copy import deepcopy +from dataclasses import dataclass, field from datetime import datetime from io import BytesIO from multiprocessing import cpu_count @@ -2721,71 +2722,75 @@ def _call(self, method: str | Callable[..., Any], **kwargs) -> dict[str, Any]: return self._core.call(method, **kwargs) -class S3File(AbstractBufferedFile): - """A buffered file object for reading and writing an S3 object. +@dataclass +class S3FileTarget: + """The object that a file reads or writes, resolved when it is opened. - Instances are returned by ``S3FileSystem.open()``. + In read mode, the object is looked up, and the reads are made conditional + on its ETag (``IfMatch``). In append mode, an existing object is looked + up; one smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is also + read, so that the file rewrites it from its buffer, and a larger one is + copied by the file as the first parts of a multipart upload. In + exclusive-create mode, the object must not exist, and the upload is + committed with ``IfNoneMatch="*"``. + + Attributes: + path: The path of the file. A version that is read is carried in the + path as a ``?versionId=`` suffix, so that a reopened file reads + the same version. + bucket: The bucket of the object. + key: The key of the object. + version_id: The version that is read, if any. + request_kwargs: The parameters of the S3 requests of the file, with + the conditions above and, in append mode, the metadata of the + existing object, which the rewritten object keeps. + details: The object that is read or appended, or an empty dictionary. + size: The size of the object that is read, if known. + appended: The existing object in append mode, if any. + append_data: The data of an appended object that the file rewrites + from its buffer, if any. """ - fs: S3FileSystem - buffer: BytesIO | None + path: str + bucket: str + key: str + version_id: str | None + request_kwargs: dict[str, Any] + details: S3Object | dict[str, Any] = field(default_factory=dict) + size: int | None = None + appended: S3Object | None = None + append_data: bytes | None = None - def __init__( - self, + @classmethod + def resolve( + cls, fs: S3FileSystem, path: str, - mode: str = "rb", + mode: str, + *, + block_size: int, version_id: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, - executor: S3Executor | None = None, - block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, - cache_type: str = "bytes", - autocommit: bool = True, - cache_options: dict[Any, Any] | None = None, size: int | None = None, - s3_additional_kwargs: dict[str, Any] | None = None, - **kwargs, - ) -> None: - """Initialize the file for the path and mode. + request_kwargs: Mapping[str, Any] | None = None, + ) -> S3FileTarget: + """Validate the arguments of a file and look up its object. - In read mode, the object is looked up with ``info()`` and the reads - are made conditional on its ETag (``IfMatch``). In append mode, an - existing object smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is - read into the write buffer; a larger one is copied with - ``UploadPartCopy`` as the first parts of a multipart upload, whatever - the block size. In exclusive-create mode, the object must not exist - when the file is opened, and the upload is committed with - ``IfNoneMatch="*"`` so that it does not replace an object created in - the meantime. + Sends the lookup requests of the mode, which block. Args: fs: The filesystem that the file belongs to. path: S3 path (s3://bucket/key) of the file. mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + block_size: The block size of the file. version_id: The version ID to read. Must match the version ID in the path if both are given. A version cannot be given, in either form, for writing or appending. - max_workers: The number of parallel workers for range reads, part - copies, and part uploads. A write keeps at most - ``max_workers`` uploaded parts in flight, waiting for one of - them to finish before it submits another. - executor: The executor for parallel operations. If None, a new - ``S3ThreadPoolExecutor`` is created. - block_size: The block size for reads and writes. Must be between - ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and - ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless reading. - cache_type: The fsspec cache type for reads; see - ``S3FileSystem.__init__()``. - autocommit: Whether to commit the written data when the file is - closed. If False, :meth:`commit` must be called. - cache_options: Options for the fsspec cache. - size: The size of the object, if known. Passed to - ``fsspec.spec.AbstractBufferedFile``. - s3_additional_kwargs: Additional parameters for the S3 requests of - the file, such as ``ContentType`` or ``RequestPayer``. Each - request receives those that its operation accepts. - **kwargs: Additional parameters for the S3 requests of the file, - which take precedence over ``s3_additional_kwargs``. + size: The size of the object to read, if known. + request_kwargs: The parameters of the S3 requests of the file. A + copy is kept. + + Returns: + The resolved target. Raises: FileExistsError: If an object exists at the path in @@ -2797,36 +2802,26 @@ def __init__( between ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing. """ - self.max_workers = max_workers # A new dictionary, so that the caller's is not modified. - self.s3_additional_kwargs: dict[str, Any] = {**(s3_additional_kwargs or {}), **kwargs} - - # The arguments are validated, and the objects looked up, before the - # base class initializer: a file that fails here is never opened, so - # its garbage collection does not close (flush and commit) it. + kwargs: dict[str, Any] = dict(request_kwargs or {}) s3_path = S3Path.parse(path) path_version_id = s3_path.version_id - self.bucket = s3_path.bucket if not s3_path.key: raise ValueError("The path does not contain a key.") - self.key = s3_path.key if version_id and path_version_id: if version_id != path_version_id: raise ValueError( f"The version_id: {version_id} specified in the argument and " f"the version_id: {path_version_id} specified in the path do not match." ) - self.version_id: str | None = version_id elif path_version_id: - self.version_id = path_version_id - else: - self.version_id = version_id - if self.version_id and "r" not in mode: + version_id = path_version_id + if version_id and "r" not in mode: raise ValueError("Cannot write to the file with the version specified.") - if self.version_id and not path_version_id: + if version_id and not path_version_id: # Carry the version in the path, as with the ?versionId= suffix, # so that a reopened (e.g., unpickled) file reads the same version. - path = f"{path}?versionId={self.version_id}" + path = f"{path}?versionId={version_id}" if "r" not in mode and not ( fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE <= block_size @@ -2840,63 +2835,170 @@ def __init__( f"5 GiB ({fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), inclusive: {block_size}." ) - self._details: S3Object | dict[str, Any] = {} - append_info: S3Object | None = None + details: S3Object | dict[str, Any] = {} + appended: S3Object | None = None append_data: bytes | None = None # The lookups need the parameters of the file on which their # authorization depends, such as the customer-provided key. - lookup_kwargs = fs._get_lookup_kwargs(self.s3_additional_kwargs) + lookup_kwargs = fs._get_lookup_kwargs(kwargs) if "r" in mode: - # Looked up before the base class initializer, which would + # Looked up here; the base class initializer of a file would # otherwise take the size from the latest version of the object. - info = fs.info(path, version_id=self.version_id, **lookup_kwargs) + info = fs.info(path, version_id=version_id, **lookup_kwargs) if info.get("type") == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY: # A prefix has no object to read. raise FileNotFoundError(path) - if fs.version_aware and not self.version_id: + if fs.version_aware and not version_id: # Pin the version observed at open time so that reads are # consistent even if the object is overwritten. info() heads # the object when the cached entry carries no version. - self.version_id = info.get("version_id") - if self.version_id: + version_id = info.get("version_id") + if version_id: # Carried in the path as an explicit version is, so that # the methods of the file use it too. - path = f"{path}?versionId={self.version_id}" + path = f"{path}?versionId={version_id}" if etag := info.get("etag"): - self.s3_additional_kwargs.update({"IfMatch": etag}) - self._details = info + kwargs.update({"IfMatch": etag}) + details = info if size is None: size = info.get("size") elif "a" in mode: # The rewritten object keeps the metadata of the existing one, # which a cached listing entry lacks, so look up the object. with contextlib.suppress(FileNotFoundError): - append_info = fs.info(path, refresh=True, **lookup_kwargs) - if ( - append_info is not None - and append_info.get("size", 0) < fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE - ): - # Too small to be a part of a multipart upload: rewritten - # from the buffer. Only the lookup parameters are sent, so - # that the whole object is read. - append_data = fs.cat_file(path, **lookup_kwargs) + appended = fs.info(path, refresh=True, **lookup_kwargs) + if appended is not None: + if appended.get("size", 0) < fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE: + # Too small to be a part of a multipart upload: rewritten + # from the buffer. Only the lookup parameters are sent, so + # that the whole object is read. + append_data = fs.cat_file(path, **lookup_kwargs) + kwargs.update(appended.to_api_repr()) + details = appended elif "x" in mode: # Checked up front so that no data is uploaded for an existing # object, and on commit with IfNoneMatch for one created since. if fs.exists(path, **lookup_kwargs): raise FileExistsError(path) - self.s3_additional_kwargs.update({"IfNoneMatch": "*"}) + kwargs.update({"IfNoneMatch": "*"}) + return cls( + path=path, + bucket=s3_path.bucket, + key=s3_path.key, + version_id=version_id, + request_kwargs=kwargs, + details=details, + size=size, + appended=appended, + append_data=append_data, + ) + + +class S3File(AbstractBufferedFile): + """A buffered file object for reading and writing an S3 object. + + Instances are returned by ``S3FileSystem.open()``. + """ + + fs: S3FileSystem + buffer: BytesIO | None + + def __init__( + self, + fs: S3FileSystem, + path: str, + mode: str = "rb", + version_id: str | None = None, + max_workers: int = (cpu_count() or 1) * 5, + executor: S3Executor | None = None, + block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, + cache_type: str = "bytes", + autocommit: bool = True, + cache_options: dict[Any, Any] | None = None, + size: int | None = None, + s3_additional_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> None: + """Initialize the file for the path and mode. + + In read mode, the object is looked up with ``info()`` and the reads + are made conditional on its ETag (``IfMatch``). In append mode, an + existing object smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is + read into the write buffer; a larger one is copied with + ``UploadPartCopy`` as the first parts of a multipart upload, whatever + the block size. In exclusive-create mode, the object must not exist + when the file is opened, and the upload is committed with + ``IfNoneMatch="*"`` so that it does not replace an object created in + the meantime. + + Args: + fs: The filesystem that the file belongs to. + path: S3 path (s3://bucket/key) of the file. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + version_id: The version ID to read. Must match the version ID in + the path if both are given. A version cannot be given, in + either form, for writing or appending. + max_workers: The number of parallel workers for range reads, part + copies, and part uploads. A write keeps at most + ``max_workers`` uploaded parts in flight, waiting for one of + them to finish before it submits another. + executor: The executor for parallel operations. If None, a new + ``S3ThreadPoolExecutor`` is created. + block_size: The block size for reads and writes. Must be between + ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless reading. + cache_type: The fsspec cache type for reads; see + ``S3FileSystem.__init__()``. + autocommit: Whether to commit the written data when the file is + closed. If False, :meth:`commit` must be called. + cache_options: Options for the fsspec cache. + size: The size of the object, if known. Passed to + ``fsspec.spec.AbstractBufferedFile``. + s3_additional_kwargs: Additional parameters for the S3 requests of + the file, such as ``ContentType`` or ``RequestPayer``. Each + request receives those that its operation accepts. + **kwargs: Additional parameters for the S3 requests of the file, + which take precedence over ``s3_additional_kwargs``. + + Raises: + FileExistsError: If an object exists at the path in + exclusive-create mode. + FileNotFoundError: If no object exists at the path when reading, + including when the path is a prefix. + ValueError: If the path has no key, the version IDs do not match, + a version is given for writing, or the block size is not + between ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing. + """ + self.max_workers = max_workers + # The arguments are validated, and the objects looked up, before the + # base class initializer: a file that fails here is never opened, so + # its garbage collection does not close (flush and commit) it. + target = S3FileTarget.resolve( + fs, + path, + mode, + block_size=block_size, + version_id=version_id, + size=size, + request_kwargs={**(s3_additional_kwargs or {}), **kwargs}, + ) + self.bucket = target.bucket + self.key = target.key + self.version_id: str | None = target.version_id + self.s3_additional_kwargs: dict[str, Any] = target.request_kwargs + self._details: S3Object | dict[str, Any] = target.details self._executor: S3Executor = executor or S3ThreadPoolExecutor(max_workers=max_workers) super().__init__( fs=fs, - path=path, + path=target.path, mode=mode, block_size=block_size, autocommit=autocommit, cache_type=cache_type, cache_options=cache_options, - size=size, + size=target.size, ) self.append_block = False @@ -2907,15 +3009,13 @@ def __init__( self._uploads_in_flight: set[Future[S3MultipartUploadPart]] = set() # The part copies of an append that may not have finished. self._copies_in_flight: set[Future[S3MultipartUploadPart]] = set() - if append_info is not None: - if append_data is not None: - self.write(append_data) + if target.appended is not None: + if target.append_data is not None: + self.write(target.append_data) else: # Copied with UploadPartCopy as the leading part(s). self.append_block = True - self.loc = append_info.get("size", 0) - self.s3_additional_kwargs.update(append_info.to_api_repr()) - self._details = append_info + self.loc = target.appended.get("size", 0) def _get_request_kwargs(self, method: str) -> dict[str, Any]: """Select the parameters of the file that an S3 operation accepts. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 99a86907e..3be1515b9 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -13,18 +13,25 @@ import logging import mimetypes import os +from concurrent.futures import Future, wait +from io import BytesIO from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast -from fsspec.asyn import AsyncFileSystem, sync +from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK from fsspec.core import get_compression from fsspec.implementations.local import LocalFileSystem, make_path_posix from fsspec.utils import check_contained -from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem +from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import ( + S3AioExecutor, + S3AioThreadPool, + S3Executor, + S3ThreadPoolExecutor, +) from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -36,6 +43,7 @@ ) from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing +from pyathena.filesystem.s3_writer import S3MultipartWriter if TYPE_CHECKING: from datetime import datetime @@ -50,19 +58,22 @@ class AioS3FileSystem(AsyncFileSystem): """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem. - This class wraps ``S3FileSystem`` to provide native asyncio support. Instead of - using ``ThreadPoolExecutor`` for parallel operations, it uses ``asyncio.gather`` - with ``asyncio.to_thread`` for natural integration with the asyncio event loop. + This class wraps ``S3FileSystem`` to provide native asyncio support. Its + coroutines run the blocking boto3 calls in an ``S3AioThreadPool`` with + ``max_workers`` threads that the filesystem owns, instead of the default + executor of the event loop, and run parallel requests with + ``asyncio.gather``. The implementation uses composition: an internal ``S3FileSystem`` instance handles - all boto3 calls, while this class delegates to it via ``asyncio.to_thread()``. + all boto3 calls, while this class delegates to it in the threads of its pool. This avoids diamond inheritance issues and keeps all boto3 logic in one place. File handles created by ``_open`` use ``S3AioExecutor`` so that parallel operations (range reads, multipart uploads) are dispatched through the event - loop with ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file. - An instance created with ``asynchronous=True`` has no event loop of its own, - so its file handles use a ``ThreadPoolExecutor``. + loop to a pool of ``max_workers`` threads of each file. An instance created + with ``asynchronous=True`` has no event loop of its own, so its file handles + use a ``ThreadPoolExecutor``. :meth:`open_async` returns an + ``AioS3StreamedFile``, whose reads and writes are coroutines. Attributes: _sync_fs: The internal synchronous S3FileSystem instance. @@ -105,7 +116,10 @@ def __init__( connection: Passed to the internal ``S3FileSystem``. default_block_size: Passed to the internal ``S3FileSystem``. default_cache_type: Passed to the internal ``S3FileSystem``. - max_workers: Passed to the internal ``S3FileSystem``. + max_workers: The number of threads that run the blocking calls of + the filesystem, and the default number of parallel requests + of its files and copies. Passed to the internal + ``S3FileSystem``. s3_additional_kwargs: Passed to the internal ``S3FileSystem``. allow_bucket_creation: Passed to the internal ``S3FileSystem``. allow_bucket_deletion: Passed to the internal ``S3FileSystem``. @@ -137,6 +151,9 @@ def __init__( ) # Share dircache for cache coherence between async and sync instances self.dircache = self._sync_fs.dircache + # The blocking calls of the coroutines run in threads of this pool, + # not in the default executor of the event loop. + self._thread_pool = S3AioThreadPool(max_workers) @property def core(self) -> S3Core: @@ -144,21 +161,23 @@ def core(self) -> S3Core: return self._sync_fs.core async def _info(self, path: str, **kwargs) -> S3Object: - return await asyncio.to_thread(self._sync_fs.info, path, **kwargs) + return await self._thread_pool.run(self._sync_fs.info, path, **kwargs) async def _ls(self, path: str, detail: bool = False, **kwargs) -> list[S3Object] | list[str]: - return await asyncio.to_thread(self._sync_fs.ls, path, detail=detail, **kwargs) + return await self._thread_pool.run(self._sync_fs.ls, path, detail=detail, **kwargs) async def _cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs ) -> bytes: - return await asyncio.to_thread(self._sync_fs.cat_file, path, start=start, end=end, **kwargs) + return await self._thread_pool.run( + self._sync_fs.cat_file, path, start=start, end=end, **kwargs + ) async def _exists(self, path: str, **kwargs) -> bool: - return await asyncio.to_thread(self._sync_fs.exists, path, **kwargs) + return await self._thread_pool.run(self._sync_fs.exists, path, **kwargs) async def _rm_file(self, path: str, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.rm_file, path, **kwargs) + await self._thread_pool.run(self._sync_fs.rm_file, path, **kwargs) async def _pipe_file( self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs @@ -166,9 +185,9 @@ async def _pipe_file( if self._intrans: # The transaction belongs to this filesystem, not to the internal # S3FileSystem, so write through open() to defer the commit to it. - await asyncio.to_thread(self._pipe_file_in_transaction, path, value, mode, **kwargs) + await self._thread_pool.run(self._pipe_file_in_transaction, path, value, mode, **kwargs) return - await asyncio.to_thread(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) + await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) def _pipe_file_in_transaction( self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs @@ -211,11 +230,11 @@ async def _put_file( ) -> None: if self._intrans: # See _pipe_file. - await asyncio.to_thread( + await self._thread_pool.run( self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs ) return - await asyncio.to_thread( + await self._thread_pool.run( self._sync_fs.put_file, lpath, rpath, callback=callback, mode=mode, **kwargs ) @@ -274,13 +293,17 @@ def _put_file_in_transaction( self.invalidate_cache(rpath) async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs) + await self._thread_pool.run( + self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs + ) async def _mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs) + await self._thread_pool.run( + self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs + ) async def _makedirs(self, path: str, exist_ok: bool = False) -> None: - await asyncio.to_thread(self._sync_fs.makedirs, path, exist_ok=exist_ok) + await self._thread_pool.run(self._sync_fs.makedirs, path, exist_ok=exist_ok) async def _rm( self, @@ -293,7 +316,7 @@ async def _rm( """Delete objects with DeleteObjects requests. See :meth:`S3FileSystem.rm`. The requests run in parallel with - ``asyncio.gather`` and ``asyncio.to_thread``. + ``asyncio.gather`` in the threads of the filesystem. Args: path: S3 path (s3://bucket/key) or list of paths to delete. @@ -307,7 +330,7 @@ async def _rm( ValueError: If a path is a bucket. OSError: If S3 could not delete some of the objects. """ - paths = await asyncio.to_thread( + paths = await self._thread_pool.run( self._sync_fs._delete_paths, path, recursive=recursive, maxdepth=maxdepth ) await self._delete_objects(paths, **kwargs) @@ -328,7 +351,10 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: """ batches, params = self._sync_fs._delete_batches(paths, **kwargs) results = await asyncio.gather( - *[asyncio.to_thread(self._sync_fs._delete_batch, batch, **params) for batch in batches], + *[ + self._thread_pool.run(self._sync_fs._delete_batch, batch, **params) + for batch in batches + ], return_exceptions=True, ) self._sync_fs._raise_delete_errors(results) @@ -356,7 +382,7 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N """ if path1 == path2: return - pairs = await asyncio.to_thread( + pairs = await self._thread_pool.run( self._sync_fs._move_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth ) # Every copy finishes before a failure is raised, as in fsspec's @@ -418,7 +444,7 @@ async def _copy( """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = await asyncio.to_thread( + pairs = await self._thread_pool.run( self._sync_fs._copy_pairs, S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth), ) @@ -459,7 +485,7 @@ async def _get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) else rpath if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): root = make_path_posix(lpath) - pairs = await asyncio.to_thread( + pairs = await self._thread_pool.run( self._sync_fs._copy_pairs, S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), isdir=LocalFileSystem().isdir, @@ -521,7 +547,7 @@ async def _copy_file(self, path1: str, path2: str, **kwargs) -> bool: size1 = info1.get("size", 0) try: if size1 <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: - await asyncio.to_thread(self.core.copy_object, source, destination, **kwargs) + await self._thread_pool.run(self.core.copy_object, source, destination, **kwargs) else: await self._copy_object_with_multipart_upload( source, @@ -546,8 +572,8 @@ async def _copy_object_with_multipart_upload( """Copy an object with a multipart upload of its byte ranges. See :meth:`S3FileSystem._copy_object_with_multipart_upload`. The part - and annotation copies run in parallel as asyncio tasks with - ``asyncio.to_thread``. On a cancellation, the creation of the upload, + and annotation copies run in parallel as asyncio tasks, in a pool of + ``max_workers`` threads of the copy. On a cancellation, the creation of the upload, the running part copies and the completion are waited for, the upload is aborted if it was created and has not completed, and the cancellation is re-raised. A repeated cancellation returns without @@ -567,15 +593,20 @@ async def _copy_object_with_multipart_upload( size. """ max_workers = max_workers if max_workers else self._sync_fs.max_workers - plan = await asyncio.to_thread( + plan = await self._thread_pool.run( self.core.plan_multipart_copy, source, destination, block_size, **kwargs ) if plan.fits_single_request: # See S3FileSystem._copy_object_with_multipart_upload. - await asyncio.to_thread(self.core.copy_object, plan.source, plan.destination, **kwargs) + await self._thread_pool.run( + self.core.copy_object, plan.source, plan.destination, **kwargs + ) return multipart_upload: S3MultipartUpload + # The requests of the copy run in a pool of its own, so that + # max_workers of the call bounds them. + pool = S3AioThreadPool(max_workers) semaphore = asyncio.Semaphore(max_workers) failed = False @@ -586,7 +617,7 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart # The upload is being aborted; do not start more parts. return None try: - return await asyncio.to_thread( + return await pool.run( self.core.upload_part_copy, upload=multipart_upload, part_number=i + 1, @@ -600,7 +631,7 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart raise tasks: list[asyncio.Future[S3MultipartUploadPart | None]] = [] - completion: asyncio.Task[S3CompleteMultipartUpload] | None = None + completion: asyncio.Future[S3CompleteMultipartUpload] | None = None async def _abort() -> None: await asyncio.wait([creation]) @@ -616,7 +647,7 @@ async def _abort() -> None: # The upload completed despite the cancellation, so # there is nothing to abort. return - await asyncio.to_thread( + await pool.run( self._sync_fs._abort_multipart_upload, creation.result(), plan.abort_params, @@ -626,9 +657,7 @@ async def _abort() -> None: # after a cancellation; scheduled right before the try, so that # nothing can fail between the two. creation = asyncio.ensure_future( - asyncio.to_thread( - self.core.create_multipart_upload, plan.destination, **plan.create_params - ) + pool.run(self.core.create_multipart_upload, plan.destination, **plan.create_params) ) try: # shield keeps a cancellation from cancelling the creation, whose @@ -645,7 +674,7 @@ async def _abort() -> None: # The tasks are in part-number order. parts = [task.result() for task in tasks] completion = asyncio.ensure_future( - asyncio.to_thread( + pool.run( self.core.complete_multipart_upload, multipart_upload, cast(list[S3MultipartUploadPart], parts), @@ -664,6 +693,9 @@ async def _abort() -> None: # the cleanup; the event loop keeps only weak references to tasks. _cleanup_tasks.add(cleanup) cleanup.add_done_callback(_cleanup_tasks.discard) + # The cleanup sends its requests through the pool, so the pool is + # shut down when the cleanup ends, also after this task returns. + cleanup.add_done_callback(lambda _: pool.shutdown(wait=False)) await asyncio.shield(cleanup) raise @@ -677,7 +709,7 @@ async def _copy_annotation(name: str) -> None: # S3FileSystem does. return try: - await asyncio.to_thread( + await pool.run( self.core.copy_object_annotation, name, plan.source, @@ -690,9 +722,12 @@ async def _copy_annotation(name: str) -> None: failed = True raise - results = await asyncio.gather( - *[_copy_annotation(name) for name in plan.annotations], return_exceptions=True - ) + try: + results = await asyncio.gather( + *[_copy_annotation(name) for name in plan.annotations], return_exceptions=True + ) + finally: + pool.shutdown(wait=False) for result in results: if isinstance(result, BaseException): raise result @@ -705,7 +740,7 @@ async def _find( **kwargs, ) -> dict[str, S3Object] | list[str]: detail = kwargs.pop("detail", False) - files = await asyncio.to_thread( + files = await self._thread_pool.run( self._sync_fs._find, path, maxdepth=maxdepth, withdirs=withdirs, **kwargs ) if detail: @@ -807,8 +842,78 @@ def _open( **kwargs, ) + async def open_async( + self, + path: str, + mode: str = "rb", + block_size: int | None = None, + max_workers: int | None = None, + version_id: str | None = None, + s3_additional_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> AioS3StreamedFile: + """Open a file whose reads and writes are coroutines. + + The lookups of the mode, as in ``open()``, run in the threads of the + filesystem. See :class:`AioS3StreamedFile` for how the file reads, + writes, and commits. + + Args: + path: S3 path (s3://bucket/key) of the file. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + block_size: The block size for reads and writes. Defaults to the + default block size of the filesystem. + max_workers: The number of threads of the file, which bounds its + parallel requests. Defaults to the ``max_workers`` of the + filesystem. + version_id: The version ID to read. + s3_additional_kwargs: Additional parameters for the S3 requests of + the file, which take precedence over those of the filesystem. + **kwargs: Additional parameters for the S3 requests of the file, + which take precedence over ``s3_additional_kwargs``. + + Returns: + The opened file. + + Raises: + ValueError: If the mode is not binary, compression is requested, + or the arguments are invalid, as in ``open()``. + NotImplementedError: If the mode is not supported, or the file is + opened for writing in a transaction, whose commit the file + cannot defer. + FileExistsError: If an object exists at the path in + exclusive-create mode. + FileNotFoundError: If no object exists at the path when reading. + """ + if "b" not in mode or kwargs.get("compression"): + raise ValueError("open_async() supports only binary modes without compression.") + if mode not in {"rb", "wb", "ab", "xb"}: + raise NotImplementedError("File mode not supported") + if mode != "rb" and self._intrans: + raise NotImplementedError("open_async() cannot write in a transaction.") + path = self._strip_protocol(path) + if block_size is None: + block_size = self._sync_fs.default_block_size + if max_workers is None: + max_workers = self._sync_fs.max_workers + s3_additional_kwargs = { + **self._sync_fs.s3_additional_kwargs, + **(s3_additional_kwargs or {}), + } + return await self._thread_pool.run( + AioS3StreamedFile, + self, + path, + mode, + version_id=version_id, + max_workers=max_workers, + block_size=block_size, + s3_additional_kwargs=s3_additional_kwargs, + **kwargs, + ) + async def _rmdir(self, path: str) -> None: - await asyncio.to_thread(self._sync_fs.rmdir, path) + await self._thread_pool.run(self._sync_fs.rmdir, path) def rmdir(self, path: str) -> None: """Remove an S3 bucket, which must be empty. @@ -1008,7 +1113,7 @@ def invalidate_cache(self, path: str | None = None) -> None: self._sync_fs.invalidate_cache(path) async def _touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: - return await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs) + return await self._thread_pool.run(self._sync_fs.touch, path, truncate=truncate, **kwargs) def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: """Create an empty object with PutObject. @@ -1034,8 +1139,437 @@ class AioS3File(S3File): ``isinstance`` checks and to document the async execution model. All parallel operations (range reads, multipart uploads) are dispatched through the ``S3Executor`` interface — the ``S3AioExecutor`` - provided by ``AioS3FileSystem`` dispatches them through the event loop with - ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file. + provided by ``AioS3FileSystem`` dispatches them through the event loop to + a pool of ``max_workers`` threads of the file. For an ``AioS3FileSystem`` created with ``asynchronous=True``, it is an - ``S3ThreadPoolExecutor``. + ``S3ThreadPoolExecutor``. For reads and writes that are coroutines, see + ``AioS3StreamedFile``. + """ + + +class AioS3StreamedFile(AbstractAsyncStreamedFile): + """An S3 file whose reads and writes are coroutines. + + Returned by :meth:`AioS3FileSystem.open_async`. The requests of the file + run in an ``S3AioThreadPool`` of ``max_workers`` threads of its own, and + the coroutines wait for them without blocking the event loop: + + - A read splits the requested range into ranges of the block size, read + in parallel. Nothing is cached. As with ``S3File``, the reads are + conditional on the ETag of the object, and on its version when the + filesystem is version-aware. + - A write buffers the data, and uploads each full block as a part of a + multipart upload, with at most ``max_workers`` parts in flight. A write + that would submit another part waits for one of them to finish first, + and a part found failed by that wait stops the write with its error. + An append copies an existing object with ``UploadPartCopy``, or + rewrites one smaller than a part from the buffer, as ``S3File`` does. + - Closing the file commits the written data, as ``S3File.close()`` does: + with PutObject if no part was uploaded, otherwise by completing the + multipart upload, which is aborted if the completion fails. + - Leaving an ``async with`` block with an exception, a failed write, and + a cancelled write close the file without committing it: the requests + that have started are waited for and the multipart upload is aborted, + so that the existing object is left unchanged. + + The commit and the abort run to their end in the threads of the file, + also if the task that waits for them is cancelled. A file opened for + writing must be closed with ``await close()`` or an ``async with`` block; + garbage collection cannot await it, so the data of an unclosed file is + never committed. """ + + fs: AioS3FileSystem + buffer: BytesIO | None + closed: bool + forced: bool + offset: int | None + + def __init__( + self, + fs: AioS3FileSystem, + path: str, + mode: str = "rb", + version_id: str | None = None, + max_workers: int = (cpu_count() or 1) * 5, + block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, + size: int | None = None, + s3_additional_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> None: + """Initialize the file for the path and mode. + + Sends the lookup requests of the mode (see :class:`S3FileTarget`), + which block; :meth:`AioS3FileSystem.open_async` runs this in a thread. + + Args: + fs: The filesystem that the file belongs to. + path: S3 path (s3://bucket/key) of the file. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + version_id: The version ID to read. + max_workers: The number of threads of the file, which bounds its + parallel requests. + block_size: The block size for reads and writes. Must be between + ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless + reading. + size: The size of the object to read, if known. + s3_additional_kwargs: Additional parameters for the S3 requests of + the file. Each request receives those that its operation + accepts. + **kwargs: Additional parameters for the S3 requests of the file, + which take precedence over ``s3_additional_kwargs``. + + Raises: + FileExistsError: If an object exists at the path in + exclusive-create mode. + FileNotFoundError: If no object exists at the path when reading. + NotImplementedError: If the mode is not supported. + ValueError: If ``max_workers`` is not positive, or the arguments + are invalid as for ``S3FileTarget.resolve()``. + """ + if mode not in {"rb", "wb", "ab", "xb"}: + raise NotImplementedError("File mode not supported") + self.max_workers = max_workers + # Created first, so that an invalid max_workers raises before the lookups. + self._pool = S3AioThreadPool(max_workers) + target = S3FileTarget.resolve( + fs._sync_fs, + path, + mode, + block_size=block_size, + version_id=version_id, + size=size, + request_kwargs={**(s3_additional_kwargs or {}), **kwargs}, + ) + self.bucket = target.bucket + self.key = target.key + self.version_id = target.version_id + self.s3_additional_kwargs: dict[str, Any] = target.request_kwargs + self._details: S3Object | dict[str, Any] = target.details + super().__init__( + fs=fs, + path=target.path, + mode=mode, + block_size=block_size, + cache_type="none", + size=target.size, + ) + + self.append_block = False + self._writer: S3MultipartWriter | None = None + if mode != "rb": + self._writer = S3MultipartWriter( + fs.core, S3Path(self.bucket, self.key), block_size=block_size + ) + self._creation: Future[S3MultipartUpload] | None = None + # The futures of the part uploads and copies, in part-number order. + self._parts: list[Future[S3MultipartUploadPart]] = [] + # The parts uploaded from the buffer that may not have finished. The + # part copies of an append hold no data and are not counted. + self._uploads_in_flight: set[asyncio.Future[S3MultipartUploadPart]] = set() + if target.appended is not None: + if target.append_data is not None: + cast(BytesIO, self.buffer).write(target.append_data) + else: + # Copied with UploadPartCopy as the leading part(s). + self.append_block = True + self.loc = target.appended.get("size", 0) + + def __del__(self) -> None: + # close() is a coroutine, which garbage collection cannot await, so + # an unclosed file is not committed; only its threads are released. + pool = getattr(self, "_pool", None) + if pool is not None: + pool.shutdown(wait=False) + + def _get_request_kwargs(self, method: str) -> dict[str, Any]: + """Select the parameters of the file that an S3 operation accepts. + + Args: + method: The name of the client method, such as ``upload_part``. + + Returns: + The parameters in ``s3_additional_kwargs`` that the operation + accepts. + """ + return self.fs.core.operation_params(method, self.s3_additional_kwargs) + + async def write(self, data: bytes) -> int: + """Write data to the buffer, uploading each full block as a part. + + Args: + data: The bytes to write. + + Returns: + The number of bytes written. + + Raises: + ValueError: If the file is not open for writing, is closed, or + has been force-flushed. + BaseException: The error of a part that the write found failed + while it waited for a part in flight. + """ + if self.mode not in {"wb", "ab", "xb"}: + raise ValueError("File not in write mode") + if self.closed: + raise ValueError("I/O operation on closed file.") + if self.forced: + raise ValueError("This file has been force-flushed, can only close") + buffer = cast(BytesIO, self.buffer) + out = buffer.write(data) + self.loc += out + if buffer.tell() >= self.blocksize: + await self.flush() + return out + + async def flush(self, force: bool = False) -> None: + """Upload the buffered data if it fills a block, or all of it if forced. + + Args: + force: Whether this is the final flush of the file, after which + the written data is committed. + + Raises: + ValueError: If the file is closed, or force-flushed twice. + """ + if self.closed: + raise ValueError("Flush on closed file") + if force and self.forced: + raise ValueError("Force flush cannot be called more than once") + if force: + self.forced = True + if self.mode not in {"wb", "ab", "xb"}: + return + buffer = cast(BytesIO, self.buffer) + if not force and buffer.tell() < self.blocksize: + return + if self.offset is None: + self.offset = 0 + await self._initiate_upload() + if await self._upload_chunk(final=force) is not False: + self.offset += cast(BytesIO, self.buffer).seek(0, 2) + self.buffer = BytesIO() + + async def close(self) -> None: + """Close the file, committing the written data.""" + if self.closed: + return + try: + if self.mode != "rb" and not self.forced: + await self.flush(force=True) + finally: + self.closed = True + # A commit that is still running keeps its thread. + self._pool.shutdown(wait=False) + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + if exc_type is not None and self.mode != "rb": + await self._close_without_commit() + else: + await self.close() + + async def _close_without_commit(self) -> None: + """Close the file without uploading the written data. + + Drops the buffered data, and aborts the multipart upload, if any, + once its started requests have ended (see :meth:`_discard`). + """ + if self.closed and self.buffer is None: + return + self.buffer = None + self.closed = True + try: + await self._pool.run_shielded(self._discard) + finally: + self._pool.shutdown(wait=False) + + async def _initiate_upload(self) -> None: + if not self.append_block and self.tell() < self.blocksize: + return + writer = cast(S3MultipartWriter, self._writer) + try: + self._creation = self._pool.submit( + writer.initiate, **self._get_request_kwargs("create_multipart_upload") + ) + await asyncio.wrap_future(self._creation) + if self.append_block: + whole_source = self.tell() <= self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE + size = self.tell() + if not whole_source: + info = await self.fs._info( + self.path, **self.fs._sync_fs._get_lookup_kwargs(self.s3_additional_kwargs) + ) + size = info.get("size", 0) + for part_number, range_ in writer.iter_copy_parts(size): + self._parts.append( + self._pool.submit( + writer.upload_part_copy, + part_number=part_number, + source=S3Path(self.bucket, self.key), + range_=None if whole_source else range_, + **self._get_request_kwargs("upload_part_copy"), + ) + ) + except BaseException: + await self._close_without_commit() + raise + + async def _upload_chunk(self, final: bool = False) -> bool: + """Upload the buffered data as parts of the multipart upload. + + See :meth:`S3File._upload_chunk`; a write waits for a part in flight + without blocking the event loop. + + Args: + final: Whether this is the last flush of the file, after which + the upload is committed. + + Returns: + Whether the buffer is reset afterwards. + + Raises: + BaseException: The error of an uploaded part that the write found + failed while it waited for a part in flight. + RuntimeError: If no multipart upload is initialized. + ValueError: If the data needs more parts than a multipart upload + allows. + """ + if not self.append_block and self.tell() < self.blocksize: + # Smaller than a block: uploaded with PutObject by the commit. + if final: + await self._commit() + return not final + writer = cast(S3MultipartWriter, self._writer) + if writer.upload is None: + raise RuntimeError("Multipart upload is not initialized.") + buffer = cast(BytesIO, self.buffer) + buffer.seek(0) + try: + for part_number, data in writer.iter_parts( + buffer, first_part_number=len(self._parts) + 1 + ): + if len(self._uploads_in_flight) >= self.max_workers: + # Otherwise the parts waiting for a thread pile up in memory. + done, self._uploads_in_flight = await asyncio.wait( + self._uploads_in_flight, return_when=asyncio.FIRST_COMPLETED + ) + for upload in done: + # A failed part would otherwise keep its data until + # the commit while the rest of the data is uploaded. + upload.result() + part = self._pool.submit( + writer.upload_part, + part_number=part_number, + body=data, + **self._get_request_kwargs("upload_part"), + ) + self._parts.append(part) + waiter = asyncio.wrap_future(part) + # The commit reports a failure that no wait found; retrieved + # here so that asyncio does not log it as never retrieved. + waiter.add_done_callback(lambda w: w.cancelled() or w.exception()) + self._uploads_in_flight.add(waiter) + except BaseException: + await self._close_without_commit() + raise + if final: + await self._commit() + return not final + + async def _commit(self) -> None: + """Commit the written data in a thread of the file. + + See :meth:`_complete_upload`. + """ + if self.buffer is None: + return + await self._pool.run_shielded(self._complete_upload) + + def _complete_upload(self) -> None: + """Upload or complete the written data; blocks, and runs in a thread. + + Creates an empty object if nothing was written, uploads the buffered + data with PutObject if no part was uploaded, and otherwise completes + the multipart upload once its parts have ended, aborting it if a + part or the completion fails. Invalidates the cache of the path + afterwards. + + Raises: + FileExistsError: If an object was created at the path after the + file was opened in exclusive-create mode. + """ + buffer = cast(BytesIO, self.buffer) + writer = cast(S3MultipartWriter, self._writer) + if self.tell() == 0: + self.fs._sync_fs.touch(self.path, **self._get_request_kwargs("put_object")) + elif not self._parts: + buffer.seek(0) + self.fs.core.put_object( + S3Path(self.bucket, self.key), + buffer.read(), + **self._get_request_kwargs("put_object"), + ) + else: + try: + # The parts were submitted before this call, so they have + # all started. + parts = [part.result() for part in self._parts] + writer.complete(parts, **self._get_request_kwargs("complete_multipart_upload")) + except BaseException: + self._discard() + raise + self.fs.invalidate_cache(self.path) + + def _discard(self) -> None: + """Abort the multipart upload, if any, logging a failure; blocks. + + The failure is logged instead of raised, so that it does not mask + the error that the caller is handling, also when no task waits for + the abort any more. + """ + try: + self._abort_upload() + except Exception: + _logger.exception( + f"Failed to abort the multipart upload to s3://{self.bucket}/{self.key}." + ) + + def _abort_upload(self) -> None: + """Abort the multipart upload, if any; blocks, and runs in a thread. + + The requests that have not started are cancelled, and the running + ones are waited for first, since a part that is still uploading when + the upload is aborted may be stored after the abort. + """ + if self._creation is not None and not self._creation.cancel(): + wait([self._creation]) + if self._writer is not None and self._writer.upload is not None: + wait([part for part in self._parts if not part.cancel()]) + self._writer.abort(**self._get_request_kwargs("abort_multipart_upload")) + + async def _fetch_range(self, start: int, end: int) -> bytes: + """Read a byte range of the object with parallel range requests. + + The range is clamped to the size of the object, and split into ranges + of the block size, read in the threads of the file. + + Args: + start: The offset of the first byte to read. + end: The offset to stop reading at (exclusive). + + Returns: + The bytes read, empty if the clamped range is empty. + """ + end = min(end, self.size) + if start >= end: + return b"" + ranges = S3File._get_ranges( + start, end, max_workers=self.max_workers, worker_block_size=self.blocksize + ) + path = S3Path(self.bucket, self.key, self.version_id) + request_kwargs = self._get_request_kwargs("get_object") + results = await asyncio.gather( + *[self._pool.run(self.fs.core.get_object, path, r, **request_kwargs) for r in ranges] + ) + return b"".join(results) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index ae943b9a7..f513600ce 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -10,9 +10,12 @@ from __future__ import annotations import asyncio +import contextvars +import os import threading from abc import ABCMeta, abstractmethod from collections.abc import Callable +from concurrent import futures from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count @@ -127,18 +130,148 @@ def shutdown(self, wait: bool = True) -> None: self._executor.shutdown(wait=wait) +class S3AioThreadPool: + """A thread pool that runs the blocking calls of asyncio code. + + Unlike ``asyncio.to_thread()``, which runs a call in the default executor + of the event loop, this pool runs at most ``max_workers`` calls at once in + threads of its own, so that neither the size nor the load of the default + executor limits or blocks them. The threads are created as the calls need + them, and each call runs with a copy of the caller's context variables, + as with ``asyncio.to_thread()``. In a forked child process, whose copy + has no threads, the pool creates new ones. The pool can be shared by + several event loops and threads. + + Args: + max_workers: The maximum number of calls that run at once. + + Raises: + ValueError: If ``max_workers`` is not positive. + """ + + def __init__(self, max_workers: int) -> None: + """Initialize the pool without starting any thread. + + Args: + max_workers: The maximum number of calls that run at once. + + Raises: + ValueError: If ``max_workers`` is not positive. + """ + if max_workers <= 0: + raise ValueError("max_workers must be greater than 0") + self.max_workers = max_workers + self._lock = threading.Lock() + self._executor: ThreadPoolExecutor | None = None + self._pid: int | None = None + self._shutdown = False + + def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: + """Schedule a blocking call in a thread of the pool. + + As with ``ThreadPoolExecutor``, the returned future can be cancelled + only before the call starts. The calls start in the order they were + submitted, so that a call that waits for the calls submitted before it + does not wait for a thread that it holds. + + Args: + fn: The blocking callable. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The future of the call. + + Raises: + RuntimeError: If the pool has been shut down. + """ + context = contextvars.copy_context() + return self._get_executor().submit(context.run, fn, *args, **kwargs) + + async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """Run a blocking call in a thread of the pool and wait for it. + + Cancelling the waiting task cancels a call that has not started; a + running call keeps running, as with ``asyncio.to_thread()``. + + Args: + fn: The blocking callable. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The return value of the call. + + Raises: + RuntimeError: If the pool has been shut down. + """ + return await asyncio.wrap_future(self.submit(fn, *args, **kwargs)) + + async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """Run a blocking call that a cancellation of the waiting task does not stop. + + The call runs to its end, even if it has not started when the waiting + task is cancelled; its outcome is then dropped. + + Args: + fn: The blocking callable. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The return value of the call. + + Raises: + RuntimeError: If the pool has been shut down. + """ + future = asyncio.wrap_future(self.submit(fn, *args, **kwargs)) + # Retrieved, so that asyncio does not log the failure of a call whose + # waiting task was cancelled as never retrieved. + future.add_done_callback(lambda f: f.cancelled() or f.exception()) + return await asyncio.shield(future) + + def shutdown(self, wait: bool = True) -> None: + """Stop accepting calls and release the threads once they are idle. + + The calls that were submitted before still run. + + Args: + wait: Whether to wait for the submitted calls to finish. + """ + with self._lock: + self._shutdown = True + executor = self._executor if self._pid == os.getpid() else None + if executor is not None: + executor.shutdown(wait=wait) + + def _get_executor(self) -> ThreadPoolExecutor: + """Return the thread pool of this process, creating it if needed. + + Returns: + The thread pool. + + Raises: + RuntimeError: If the pool has been shut down. + """ + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + if self._executor is None or self._pid != os.getpid(): + self._executor = ThreadPoolExecutor(max_workers=self.max_workers) + self._pid = os.getpid() + return self._executor + + class S3AioExecutor(S3Executor): """Executor that schedules work on an asyncio event loop. - Uses ``asyncio.run_coroutine_threadsafe(asyncio.to_thread(fn), loop)`` to - dispatch blocking functions onto the event loop's thread pool, returning + Uses ``asyncio.run_coroutine_threadsafe()`` to dispatch blocking functions + through the event loop to an ``S3AioThreadPool`` of its own, returning ``concurrent.futures.Future`` objects that are compatible with ``as_completed()``, ``wait()`` and ``Future.cancel()``. As with ``ThreadPoolExecutor``, a future cannot be cancelled once its function has - started. At most ``max_workers`` of the submitted functions run at once. - - This avoids thread-in-thread nesting when ``S3File`` is used from within - ``asyncio.to_thread()`` calls (the pattern used by ``AioS3FileSystem``). + started. At most ``max_workers`` of the submitted functions run at once, + whatever the size of the default executor of the event loop. Args: loop: A running asyncio event loop. @@ -164,28 +297,19 @@ def __init__( Raises: ValueError: If ``max_workers`` is not positive. """ - if max_workers <= 0: - # As ThreadPoolExecutor does; a semaphore of 0 would never run anything. - raise ValueError("max_workers must be greater than 0") self._loop = loop - self._semaphore = asyncio.Semaphore(max_workers) - - async def _run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: - """Run the function in a thread once fewer than ``max_workers`` run. - - Args: - fn: The blocking function to run. - *args: Positional arguments passed to the function. - **kwargs: Keyword arguments passed to the function. - - Returns: - The return value of the function. - """ - async with self._semaphore: - return await asyncio.to_thread(fn, *args, **kwargs) + self._pool = S3AioThreadPool(max_workers) + self._lock = threading.Lock() + # The returned futures that are not done. A function reaches the pool + # only when the event loop runs its task, so the pool is shut down + # once they are all done. + self._pending: set[Future[Any]] = set() + self._shutdown = False @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") if self._loop is not None and self._loop.is_running(): # The future of run_coroutine_threadsafe can be cancelled while # the function keeps running in its thread, so the returned future @@ -225,7 +349,10 @@ def settle(task: Future[None]) -> None: elif future.set_running_or_notify_cancel(): future.set_exception(task.exception()) - task = asyncio.run_coroutine_threadsafe(self._run(run), self._loop) + with self._lock: + self._pending.add(future) + future.add_done_callback(self._release) + task = asyncio.run_coroutine_threadsafe(self._pool.run(run), self._loop) task.add_done_callback(settle) return future raise RuntimeError( @@ -233,7 +360,25 @@ def settle(task: Future[None]) -> None: "Use S3ThreadPoolExecutor for synchronous usage." ) + def _release(self, future: Future[Any]) -> None: + """Forget a done future, shutting down the pool after the last one. + + Args: + future: The done future. + """ + with self._lock: + self._pending.discard(future) + last = self._shutdown and not self._pending + if last: + self._pool.shutdown(wait=False) + @override def shutdown(self, wait: bool = True) -> None: - # No resources to release — work is dispatched to the event loop. - pass + with self._lock: + self._shutdown = True + pending = list(self._pending) + if wait: + # Also the functions that have not reached the pool yet. + futures.wait(pending) + if not pending or wait: + self._pool.shutdown(wait=wait) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 0440d35e0..87ab091f9 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -8,7 +8,7 @@ import urllib.parse import urllib.request import uuid -from concurrent.futures import FIRST_COMPLETED, wait +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain from pathlib import Path @@ -22,7 +22,7 @@ from fsspec import Callback from pyathena.filesystem.s3 import S3File, S3FileSystem -from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem +from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core from pyathena.filesystem.s3_executor import S3Executor from pyathena.filesystem.s3_object import ( @@ -39,6 +39,7 @@ from tests.pyathena.util import ( MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, + overlapping_calls, stub_multipart_copy, ) @@ -257,6 +258,63 @@ def upload_part_copy(**kw): assert events[2:] == ["end 1", "abort"] sync_fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_max_workers(self): + # GH-1159: max_workers of the copy bounds its parallel part copies, + # whatever the size of the default executor of the event loop, in + # which they used to run. + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + sync_fs = fs._sync_fs + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} + ) + ) + call, state = overlapping_calls(3) + + def upload_part_copy(**kw): + call() + return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock() + # The HeadObject of the source, with 4 parts of the default block size. + size = 4 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + + await fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + max_workers=3, + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + + assert sync_fs.core.upload_part_copy.call_count == 4 + assert state["peak"] == 3 + sync_fs.core.complete_multipart_upload.assert_called_once() + + @pytest.mark.asyncio + async def test_blocking_calls_bypass_default_executor(self): + # GH-1159: the blocking calls of the coroutines run in the threads of + # the filesystem, so that a busy default executor of the event loop + # does not hold them. They used to run in it. + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) + release = threading.Event() + blocker = loop.run_in_executor(None, release.wait, 5) + try: + info = await asyncio.wait_for(fs._info("s3://bucket/key"), 5) + finally: + release.set() + await blocker + + assert info == {"name": "bucket/key"} + @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_cancelled(self, cancellations): @@ -1525,6 +1583,43 @@ def test_append(self, fs, base, exp): assert len(actual) == len(data + extra) assert actual == data + extra + @pytest.mark.parametrize( + "size", + [ + 2**10, # < block size: PutObject + 10 * 2**20, # > block size (5 MiB): two parts of a multipart upload + ], + ) + @pytest.mark.asyncio + async def test_open_async_write_read(self, fs, size): + data = os.urandom(size) + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_async_open_async/{uuid.uuid4()}" + ) + async with await fs.open_async(path, "wb", block_size=5 * 2**20) as f: + assert isinstance(f, AioS3StreamedFile) + await f.write(data) + async with await fs.open_async(path, "rb", block_size=5 * 2**20) as f: + assert await f.read() == data + f.seek(size // 2) + assert await f.read(10) == data[size // 2 : size // 2 + 10] + + @pytest.mark.asyncio + async def test_open_async_append_exclusive(self, fs): + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_async_open_async_append/{uuid.uuid4()}" + ) + async with await fs.open_async(path, "xb") as f: + await f.write(b"hello") + with pytest.raises(FileExistsError): + await fs.open_async(path, "xb") + async with await fs.open_async(path, "ab") as f: + await f.write(b" world") + async with await fs.open_async(path, "rb") as f: + assert await f.read() == b"hello world" + @pytest.mark.asyncio async def test_ls_buckets(self, fs): fs.invalidate_cache() @@ -2539,3 +2634,287 @@ def test_get_ranges(self, start, end, max_workers, worker_block_size, ranges): ) == ranges ) + + +class TestAioS3StreamedFile: + @pytest.fixture + def fs(self): + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + core = fs._sync_fs.core + core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + ) + ) + core.upload_part = mock.MagicMock( + side_effect=lambda **kw: S3MultipartUploadPart(kw["part_number"], {"ETag": '"e"'}) + ) + core.complete_multipart_upload = mock.MagicMock() + core.abort_multipart_upload = mock.MagicMock() + core.put_object = mock.MagicMock() + # The mocked client has no operation model to select the parameters. + core.operation_params = mock.MagicMock(side_effect=lambda method, params: dict(params)) + return fs + + @pytest.mark.asyncio + async def test_write_parts_in_flight(self, fs): + # GH-1159: a write keeps at most max_workers parts in flight, and waits + # for one to finish without blocking the event loop, on which the + # parts are released here. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + lock = threading.Lock() + gates: dict[int, threading.Event] = {} + active: list[int] = [] + in_flight: list[int] = [] + + def upload_part(**kw): + part_number = kw["part_number"] + with lock: + gate = gates.setdefault(part_number, threading.Event()) + active.append(part_number) + in_flight.append(len(active)) + if not gate.wait(5): + raise TimeoutError(part_number) + with lock: + active.remove(part_number) + return S3MultipartUploadPart(part_number, {"ETag": '"e"'}) + + fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + written = asyncio.Event() + + async def release(): + while True: + with lock: + if written.is_set() or len(active) >= 2: + # Only the newest part finishes until the data is + # written, so a write that waits for the oldest one + # would time out. + for part_number in active if written.is_set() else active[-1:]: + gates.setdefault(part_number, threading.Event()).set() + if written.is_set() and not active: + return + await asyncio.sleep(0.01) + + releaser = asyncio.ensure_future(release()) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + await f.write(b"a" * block_size * 4) + written.set() + await f.close() + await releaser + + assert in_flight == [1, 2, 2, 2] + (_, parts), _ = fs.core.complete_multipart_upload.call_args + assert [part.part_number for part in parts] == [1, 2, 3, 4] + fs.core.abort_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_write_failed_part_while_waiting(self, fs): + # A part found failed when a write waits for a part stops the write + # with its error and aborts the upload. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + fs.core.upload_part = mock.MagicMock(side_effect=PermissionError("Access Denied")) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + + with pytest.raises(PermissionError, match="Access Denied"): + await f.write(b"a" * block_size * 3) + + assert f.closed + assert sorted(c.kwargs["part_number"] for c in fs.core.upload_part.call_args_list) == [ + 1, + 2, + ] + fs.core.abort_multipart_upload.assert_called_once() + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.parametrize("blocks", [0, 2]) + @pytest.mark.asyncio + async def test_exit_with_exception(self, fs, blocks): + # Leaving the block with an exception does not commit the written + # data, and aborts the multipart upload, if any. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + + async def write_and_fail(): + async with f: + await f.write(b"a" * (block_size * blocks + 1)) + raise RuntimeError("failed") + + with pytest.raises(RuntimeError, match="failed"): + await write_and_fail() + + assert f.closed + fs.core.put_object.assert_not_called() + fs.core.complete_multipart_upload.assert_not_called() + assert fs.core.abort_multipart_upload.call_count == (1 if blocks else 0) + + @pytest.mark.parametrize("cancellations", [1, 2]) + @pytest.mark.asyncio + async def test_write_cancelled(self, fs, cancellations): + # A cancelled write waits for the running parts in the threads of the + # file and aborts the upload, also after a repeated cancellation. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + events = [] + lock = threading.Lock() + started = threading.Semaphore(0) + release = threading.Event() + aborted = threading.Event() + + def upload_part(**kw): + part_number = kw["part_number"] + with lock: + events.append(f"start {part_number}") + started.release() + # The finally block of the test always releases it. + release.wait() + with lock: + events.append(f"end {part_number}") + return S3MultipartUploadPart(part_number, {"ETag": '"e"'}) + + def abort_multipart_upload(*args, **kwargs): + events.append("abort") + aborted.set() + + fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + fs.core.abort_multipart_upload = mock.MagicMock(side_effect=abort_multipart_upload) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + task = asyncio.ensure_future(f.write(b"a" * block_size * 3)) + try: + for _ in range(2): + assert await asyncio.to_thread(started.acquire, timeout=5) + for _ in range(cancellations): + task.cancel() + # Lets the write enter its cleanup. + await asyncio.sleep(0) + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert await asyncio.to_thread(aborted.wait, 5) + + assert f.closed + assert sorted(events[:2]) == ["start 1", "start 2"] + assert sorted(events[2:4]) == ["end 1", "end 2"] + assert events[4:] == ["abort"] + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.parametrize("data", [b"", b"abc"]) + @pytest.mark.asyncio + async def test_close_small(self, fs, data): + # Data smaller than a block is uploaded with PutObject, and no data + # creates an empty object. + fs._sync_fs.touch = mock.MagicMock() + async with await fs.open_async("s3://bucket/key", "wb", ContentType="text/plain") as f: + await f.write(data) + + if data: + fs.core.put_object.assert_called_once_with( + S3Path("bucket", "key"), data, ContentType="text/plain" + ) + fs._sync_fs.touch.assert_not_called() + else: + fs._sync_fs.touch.assert_called_once_with("bucket/key", ContentType="text/plain") + fs.core.put_object.assert_not_called() + fs.core.create_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_read_parallel(self, fs): + # A read runs at most max_workers range requests at once, conditional + # on the ETag of the object. + block_size = 2**10 + size = block_size * 4 + call, state = overlapping_calls(2) + + def get_object(path, range_, **kwargs): + call() + return b"a" * (range_[1] - range_[0]) + + fs.core.get_object = mock.MagicMock(side_effect=get_object) + info = S3Object( + init={"Key": "key", "ETag": '"etag"'}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = size + fs._sync_fs.info = mock.MagicMock(return_value=info) + + async with await fs.open_async( + "s3://bucket/key", "rb", block_size=block_size, max_workers=2 + ) as f: + assert await f.read() == b"a" * size + assert await f.read() == b"" + + assert fs.core.get_object.call_count == 4 + assert state["peak"] == 2 + for c in fs.core.get_object.call_args_list: + assert c.kwargs == {"IfMatch": '"etag"'} + + @pytest.mark.parametrize( + ("mode", "kwargs", "error"), + [ + ("r", {}, ValueError), + ("wb", {"compression": "gzip"}, ValueError), + ("r+b", {}, NotImplementedError), + ("wb", {"max_workers": 0}, ValueError), + ("wb", {"block_size": 2**10}, ValueError), + ("wb", {"version_id": "v"}, ValueError), + ], + ) + @pytest.mark.asyncio + async def test_open_async_invalid(self, fs, mode, kwargs, error): + with pytest.raises(error): + await fs.open_async("s3://bucket/key", mode, **kwargs) + + @pytest.mark.asyncio + async def test_open_async_in_transaction(self, fs): + # The file always commits when it is closed, so it cannot write in a + # transaction, which defers the commit. + with fs.transaction, pytest.raises(NotImplementedError, match="transaction"): + await fs.open_async("s3://bucket/key", "wb") + + @pytest.mark.parametrize("exists", [False, True]) + @pytest.mark.asyncio + async def test_exclusive_create(self, fs, exists): + fs._sync_fs.exists = mock.MagicMock(return_value=exists) + if exists: + with pytest.raises(FileExistsError): + await fs.open_async("s3://bucket/key", "xb") + return + async with await fs.open_async("s3://bucket/key", "xb") as f: + await f.write(b"abc") + + fs.core.put_object.assert_called_once_with(S3Path("bucket", "key"), b"abc", IfNoneMatch="*") + + @pytest.mark.parametrize("size", [3, S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE]) + @pytest.mark.asyncio + async def test_append(self, fs, size): + # An existing object smaller than a part is rewritten from the buffer, + # and a larger one is copied as the first part. + info = S3Object( + init={"Key": "key", "ContentType": "text/plain"}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = size + fs._sync_fs.info = mock.MagicMock(return_value=info) + fs._sync_fs.cat_file = mock.MagicMock(return_value=b"a" * size) + fs.core.upload_part_copy = mock.MagicMock( + side_effect=lambda **kw: S3MultipartUploadPart(kw["part_number"], {"ETag": '"c"'}) + ) + + async with await fs.open_async("s3://bucket/key", "ab") as f: + await f.write(b"b") + + if size < S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE: + fs.core.put_object.assert_called_once() + assert fs.core.put_object.call_args.args == (S3Path("bucket", "key"), b"aaab") + # The rewritten object keeps the metadata of the existing one. + assert fs.core.put_object.call_args.kwargs["ContentType"] == "text/plain" + fs.core.upload_part_copy.assert_not_called() + else: + fs.core.put_object.assert_not_called() + fs.core.upload_part_copy.assert_called_once() + assert fs.core.upload_part_copy.call_args.kwargs["range_"] is None + (_, parts), _ = fs.core.complete_multipart_upload.call_args + assert [part.part_number for part in parts] == [1, 2] diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index cb79a2a9e..71e6b5140 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -6,13 +6,111 @@ # SPDX-License-Identifier: MIT import asyncio +import contextvars +import os import threading import time from concurrent.futures import Future, ThreadPoolExecutor, wait +from unittest import mock import pytest -from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor +from tests.pyathena.util import overlapping_calls + +EVENT_TIMEOUT = 5 + + +class TestS3AioThreadPool: + def test_init(self): + with pytest.raises(ValueError, match="max_workers must be greater than 0"): + S3AioThreadPool(0) + + def test_run_beyond_default_executor(self): + # GH-1159: the calls run in threads of the pool, so that max_workers, + # not the default executor of the event loop, bounds them. + call, state = overlapping_calls(4) + pool = S3AioThreadPool(4) + + async def main(): + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + await asyncio.gather(*[pool.run(call) for _ in range(6)]) + + asyncio.run(main()) + pool.shutdown() + + assert state["peak"] == 4 + + def test_run_context(self): + # The call sees the context variables of the caller, as with + # asyncio.to_thread(). + variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") + pool = S3AioThreadPool(1) + + async def main(): + variable.set("caller") + return await pool.run(variable.get) + + assert asyncio.run(main()) == "caller" + pool.shutdown() + + @pytest.mark.parametrize("shielded", [False, True]) + def test_run_cancelled_before_start(self, shielded): + # Cancelling the waiting task drops a call that has not started, + # unless it runs shielded. + events = [] + started = threading.Event() + release = threading.Event() + pool = S3AioThreadPool(1) + + def work(): + started.set() + release.wait(EVENT_TIMEOUT) + + async def main(): + running = pool.submit(work) + run = pool.run_shielded if shielded else pool.run + pending = asyncio.ensure_future(run(events.append, "pending finished")) + await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # Lets the task submit the call, which waits for the thread. + while not pending.done() and not pool._get_executor()._work_queue.qsize(): + await asyncio.sleep(0.01) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + release.set() + await asyncio.wrap_future(running) + + asyncio.run(main()) + pool.shutdown() + + assert events == (["pending finished"] if shielded else []) + + def test_shutdown(self): + # The calls submitted before the shutdown still run. + events = [] + release = threading.Event() + pool = S3AioThreadPool(1) + pool.submit(release.wait, EVENT_TIMEOUT) + pending = pool.submit(events.append, "pending finished") + pool.shutdown(wait=False) + with pytest.raises(RuntimeError, match="after shutdown"): + pool.submit(sum, [1, 2]) + release.set() + pending.result(timeout=EVENT_TIMEOUT) + + assert events == ["pending finished"] + + def test_fork(self): + # A forked child, whose copy of the pool has no threads, creates new ones. + pool = S3AioThreadPool(1) + parent = pool._get_executor() + with mock.patch("pyathena.filesystem.s3_executor.os.getpid", return_value=os.getpid() + 1): + child = pool._get_executor() + assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 + pool.shutdown() + assert child is not parent + parent.shutdown() class TestS3AioExecutor: @@ -36,6 +134,57 @@ async def main(): with pytest.raises(ValueError, match="invalid literal"): failed.result() + def test_submit_beyond_default_executor(self): + # GH-1159: max_workers functions run at once, whatever the size of the + # default executor of the event loop. They used to run in it. + call, state = overlapping_calls(3) + + async def main(): + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + executor = S3AioExecutor(loop=loop, max_workers=3) + futures = [executor.submit(call) for _ in range(3)] + await asyncio.gather(*[asyncio.wrap_future(f) for f in futures]) + executor.shutdown() + + asyncio.run(main()) + + assert state["peak"] == 3 + + def test_submit_from_default_executor(self): + # GH-1159: a function submitted from a thread of the default executor + # of the event loop does not wait for that thread. With a default + # executor of one thread, this used to deadlock a multipart write + # through the loop's asyncio.to_thread(). + async def main(): + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + executor = S3AioExecutor(loop=loop) + # Waits in the only thread of the default executor, as a + # multipart write waits for the creation of its upload. + result = await asyncio.to_thread( + lambda: executor.submit(sum, [1, 2]).result(EVENT_TIMEOUT) + ) + executor.shutdown() + return result + + assert asyncio.run(main()) == 3 + + def test_shutdown_before_dispatch(self): + # GH-1159: a function submitted before the shutdown runs, although it + # reaches the pool of the executor only when the event loop runs its + # task. An AioS3File in a transaction shuts its executor down when it + # is closed, before its parts may have reached the pool. + async def main(): + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + future = executor.submit(sum, [1, 2]) + executor.shutdown(wait=False) + with pytest.raises(RuntimeError, match="after shutdown"): + executor.submit(sum, [3, 4]) + return await asyncio.wrap_future(future) + + assert asyncio.run(main()) == 3 + def test_submit_without_running_loop(self): with pytest.raises(RuntimeError, match="requires a running event loop"): S3AioExecutor().submit(sum, [1, 2]) @@ -66,10 +215,9 @@ def cancel(executor: S3AioExecutor) -> tuple[Future[None], Future[None]]: return running, pending async def main(): - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=2)) - executor = S3AioExecutor(loop=loop) - # The function runs in one of the two threads, and cancel() in the other. + # The function runs in the only thread of the executor, so that + # the pending one waits for it, and cancel() in another thread. + executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) return await asyncio.to_thread(cancel, executor) running, pending = asyncio.run(main()) @@ -79,14 +227,7 @@ async def main(): assert not running.cancelled() assert pending.cancelled() - @pytest.mark.parametrize( - ("threads", "max_workers"), - [ - (1, 5), # The pending function waits for a thread. - (2, 1), # The pending function waits for a permit. - ], - ) - def test_loop_shutdown(self, threads, max_workers): + def test_loop_shutdown(self): # A function that has not started when the event loop shuts down is # never run, and its future is cancelled and settled, so that wait() # returns, instead of left pending. @@ -102,9 +243,8 @@ def work(): events.append("finished") async def main(): - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=threads)) - executor = S3AioExecutor(loop=loop, max_workers=max_workers) + # The pending function waits for the only thread of the executor. + executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) running = executor.submit(work) pending = executor.submit(events.append, "pending finished") pending.add_done_callback(lambda _: settled.set()) @@ -123,11 +263,10 @@ async def main(): def test_executor_shut_down(self): # An error raised before the function starts is set on the future. async def main(): - loop = asyncio.get_running_loop() - default_executor = ThreadPoolExecutor(max_workers=1) - loop.set_default_executor(default_executor) - default_executor.shutdown() - future = S3AioExecutor(loop=loop).submit(sum, [1, 2]) + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + # Fails the dispatch of the function to the pool. + executor._pool.shutdown() + future = executor.submit(sum, [1, 2]) with pytest.raises(RuntimeError, match="after shutdown"): await asyncio.wrap_future(future) return future @@ -188,9 +327,7 @@ def work(): pending.add_done_callback(lambda _: settled.set()) async def main(): - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) - executor = S3AioExecutor(loop=loop) + executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) running = executor.submit(work) executor.submit_to(pending, events.append, "pending finished") while not started.is_set(): @@ -210,11 +347,10 @@ def test_submit_to_executor_shut_down(self): future: Future[int] = Future() async def main(): - loop = asyncio.get_running_loop() - default_executor = ThreadPoolExecutor(max_workers=1) - loop.set_default_executor(default_executor) - default_executor.shutdown() - S3AioExecutor(loop=loop).submit_to(future, sum, [1, 2]) + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + # Fails the dispatch of the callable to the pool. + executor._pool.shutdown() + executor.submit_to(future, sum, [1, 2]) with pytest.raises(RuntimeError, match="after shutdown"): await asyncio.wrap_future(future) @@ -229,10 +365,10 @@ def test_submit_to_cancelled_executor_shut_down(self): async def main(): loop = asyncio.get_running_loop() - default_executor = ThreadPoolExecutor(max_workers=1) - loop.set_default_executor(default_executor) - default_executor.shutdown() - S3AioExecutor(loop=loop).submit_to(future, sum, [1, 2]) + executor = S3AioExecutor(loop=loop) + # Fails the dispatch of the callable to the pool. + executor._pool.shutdown() + executor.submit_to(future, sum, [1, 2]) assert future.cancel() with ThreadPoolExecutor(max_workers=1) as waiter: return await loop.run_in_executor(waiter, lambda: wait([future], timeout=5)) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 50f6712e5..2b2784f6d 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import io +import threading import time from concurrent.futures import Future, wait from datetime import UTC, datetime, timedelta, timezone @@ -30,6 +31,34 @@ ) +def overlapping_calls(calls, timeout=5): + """Return a callable that records the peak number of its overlapping calls. + + Each call waits until ``calls`` calls overlap, or until the timeout. + + Args: + calls: The number of calls to wait for. + timeout: The longest wait of a call, in seconds. + + Returns: + The callable, which accepts and ignores any arguments, and a + dictionary whose ``peak`` item is the peak number. + """ + condition = threading.Condition() + state = {"active": 0, "peak": 0} + + def call(*args, **kwargs): + with condition: + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + condition.notify_all() + condition.wait_for(lambda: state["active"] >= calls, timeout=timeout) + with condition: + state["active"] -= 1 + + return call, state + + def read_query(name, **kwargs): template = _queries.get_template(name) return [q.strip() for q in template.render(**kwargs).split(";") if q and q.strip()] From 5bd1605a4f06bba1f145bee2db5c690e61425d4d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:40:35 +0900 Subject: [PATCH 02/11] Share the filesystem's threads among copies and keep an error after close - Run the requests of multipart copies in the threads of the filesystem. A pool per copy let a bulk copy of large objects, which gathers the copies, start max_workers threads per object. - Leaving an async with block of a closed AioS3StreamedFile with an exception does nothing. It used to submit an abort to the shut-down pool, and the RuntimeError replaced the caller's exception. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 8 +-- pyathena/filesystem/s3_async.py | 36 ++++++-------- tests/pyathena/filesystem/test_s3_async.py | 58 ++++++++++++++++++++++ 3 files changed, 79 insertions(+), 23 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index a4f2a93c0..4062dedf0 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -259,7 +259,7 @@ its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized | **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` on a thread pool | | **File handles** | `S3File` with thread pool | `AioS3File` with `S3AioExecutor`; `AioS3StreamedFile` from `open_async()` | | **Bulk delete** | Thread pool per batch | `asyncio.gather` per batch | -| **Multipart copy** | Thread pool per part | Asyncio task per part, on a thread pool of the copy | +| **Multipart copy** | Thread pool per part | Asyncio task per part | | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) | ### Threads @@ -269,8 +269,10 @@ default executor of the event loop, so `loop.set_default_executor()` neither lim blocks them: - The filesystem has a pool of `max_workers` threads for the blocking calls of its - coroutines, such as `_info()`, `_cat_file()`, and the requests of `_rm()`. -- A multipart copy has a pool of its own, sized by the `max_workers` of the call. + coroutines, such as `_info()`, `_cat_file()`, the requests of `_rm()`, and the part + copies of multipart copies. A multipart copy runs at most the `max_workers` of the + call of its part copies at once, and the copies together run at most the + `max_workers` of the filesystem. - Each file opened with `open()` or `open_async()` has a pool of its own, sized by the `max_workers` of the call. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 3be1515b9..4977b5b85 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -572,8 +572,9 @@ async def _copy_object_with_multipart_upload( """Copy an object with a multipart upload of its byte ranges. See :meth:`S3FileSystem._copy_object_with_multipart_upload`. The part - and annotation copies run in parallel as asyncio tasks, in a pool of - ``max_workers`` threads of the copy. On a cancellation, the creation of the upload, + and annotation copies run in parallel as asyncio tasks, at most + ``max_workers`` at once, in the threads of the filesystem, which bound + the requests of all its copies together. On a cancellation, the creation of the upload, the running part copies and the completion are waited for, the upload is aborted if it was created and has not completed, and the cancellation is re-raised. A repeated cancellation returns without @@ -604,9 +605,6 @@ async def _copy_object_with_multipart_upload( return multipart_upload: S3MultipartUpload - # The requests of the copy run in a pool of its own, so that - # max_workers of the call bounds them. - pool = S3AioThreadPool(max_workers) semaphore = asyncio.Semaphore(max_workers) failed = False @@ -617,7 +615,7 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart # The upload is being aborted; do not start more parts. return None try: - return await pool.run( + return await self._thread_pool.run( self.core.upload_part_copy, upload=multipart_upload, part_number=i + 1, @@ -647,7 +645,7 @@ async def _abort() -> None: # The upload completed despite the cancellation, so # there is nothing to abort. return - await pool.run( + await self._thread_pool.run( self._sync_fs._abort_multipart_upload, creation.result(), plan.abort_params, @@ -657,7 +655,9 @@ async def _abort() -> None: # after a cancellation; scheduled right before the try, so that # nothing can fail between the two. creation = asyncio.ensure_future( - pool.run(self.core.create_multipart_upload, plan.destination, **plan.create_params) + self._thread_pool.run( + self.core.create_multipart_upload, plan.destination, **plan.create_params + ) ) try: # shield keeps a cancellation from cancelling the creation, whose @@ -674,7 +674,7 @@ async def _abort() -> None: # The tasks are in part-number order. parts = [task.result() for task in tasks] completion = asyncio.ensure_future( - pool.run( + self._thread_pool.run( self.core.complete_multipart_upload, multipart_upload, cast(list[S3MultipartUploadPart], parts), @@ -693,9 +693,6 @@ async def _abort() -> None: # the cleanup; the event loop keeps only weak references to tasks. _cleanup_tasks.add(cleanup) cleanup.add_done_callback(_cleanup_tasks.discard) - # The cleanup sends its requests through the pool, so the pool is - # shut down when the cleanup ends, also after this task returns. - cleanup.add_done_callback(lambda _: pool.shutdown(wait=False)) await asyncio.shield(cleanup) raise @@ -709,7 +706,7 @@ async def _copy_annotation(name: str) -> None: # S3FileSystem does. return try: - await pool.run( + await self._thread_pool.run( self.core.copy_object_annotation, name, plan.source, @@ -722,12 +719,9 @@ async def _copy_annotation(name: str) -> None: failed = True raise - try: - results = await asyncio.gather( - *[_copy_annotation(name) for name in plan.annotations], return_exceptions=True - ) - finally: - pool.shutdown(wait=False) + results = await asyncio.gather( + *[_copy_annotation(name) for name in plan.annotations], return_exceptions=True + ) for result in results: if isinstance(result, BaseException): raise result @@ -1364,6 +1358,8 @@ async def close(self) -> None: self._pool.shutdown(wait=False) async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + if self.closed: + return if exc_type is not None and self.mode != "rb": await self._close_without_commit() else: @@ -1375,7 +1371,7 @@ async def _close_without_commit(self) -> None: Drops the buffered data, and aborts the multipart upload, if any, once its started requests have ended (see :meth:`_discard`). """ - if self.closed and self.buffer is None: + if self.closed: return self.buffer = None self.closed = True diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 87ab091f9..24e5f930a 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -296,6 +296,47 @@ def upload_part_copy(**kw): assert state["peak"] == 3 sync_fs.core.complete_multipart_upload.assert_called_once() + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_shares_threads(self): + # GH-1159: concurrent copies share the threads of the filesystem, so + # that a bulk copy of large objects does not start max_workers threads + # per object. + fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=2, skip_instance_cache=True) + sync_fs = fs._sync_fs + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} + ) + ) + # Waits briefly for a third overlapping call, which only copies + # with threads of their own run. + call, state = overlapping_calls(3, timeout=0.5) + + def upload_part_copy(**kw): + call() + return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock() + size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + + await asyncio.gather( + *[ + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", f"dst{i}"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + for i in range(2) + ] + ) + + assert sync_fs.core.upload_part_copy.call_count == 4 + assert state["peak"] == 2 + @pytest.mark.asyncio async def test_blocking_calls_bypass_default_executor(self): # GH-1159: the blocking calls of the coroutines run in the threads of @@ -2747,6 +2788,23 @@ async def write_and_fail(): fs.core.complete_multipart_upload.assert_not_called() assert fs.core.abort_multipart_upload.call_count == (1 if blocks else 0) + @pytest.mark.asyncio + async def test_exit_with_exception_after_close(self, fs): + # An exception raised after the file was closed propagates unchanged. + f = await fs.open_async("s3://bucket/key", "wb") + + async def close_and_fail(): + async with f: + await f.write(b"abc") + await f.close() + raise KeyError("failed") + + with pytest.raises(KeyError, match="failed"): + await close_and_fail() + + fs.core.put_object.assert_called_once() + fs.core.abort_multipart_upload.assert_not_called() + @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio async def test_write_cancelled(self, fs, cancellations): From 38281eac1e775654b27027e5dd654bbe2197481f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:43:01 +0900 Subject: [PATCH 03/11] Test a cancelled commit and describe the aio threads of the benchmark - Test that a commit of AioS3StreamedFile that has started completes when the task closing the file is cancelled, as documented. - The benchmark README said that the aio requests of the filesystem suite run in the default executors, which no longer holds. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 7 +++-- tests/pyathena/filesystem/test_s3_async.py | 32 ++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index e6a323fc3..00d3546b7 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -218,9 +218,10 @@ The `[filesystem]` section of `config.toml` sets the swept values, 8 and 40 work The suite's scales are the object sizes in bytes in `[filesystem] sizes`, not the row counts in `[scales]`: 64 MiB (small), 256 MiB (medium), 1 GiB (large), and 6 GiB (xlarge) by default. The `sync` API calls the filesystem on the trial's thread, and a file's parallel requests run in an `S3ThreadPoolExecutor`. -The `aio` API runs a read or write in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run through `S3AioExecutor` in `asyncio.to_thread()` on the filesystem's event loop, fsspec's IO thread. -Its copy is a coroutine awaited on the trial's event loop, whose part copies run in `asyncio.to_thread()`. -Both event loops keep Python's default executor, whose `min(32, os.cpu_count() + 4)` threads bound the parallel requests of the `aio` API. +The `aio` API runs a read or write in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run through `S3AioExecutor`, dispatched on the filesystem's event loop, fsspec's IO thread, to a pool of `max_workers` threads of the file. +Its copy is a coroutine awaited on the trial's event loop, whose part copies run in the filesystem's pool of `max_workers` threads. +Both event loops keep Python's default executor, which runs only the trial's `asyncio.to_thread()` call. +Before #1159, the parallel requests of the `aio` API ran in the default executors, whose `min(32, os.cpu_count() + 4)` threads bound them. The source objects, of zero bytes, are uploaded with boto3 below the run's scratch prefix and recorded under `objects` in the manifest. `prepare --filesystem-scale` uploads them with the manifest, so that every fleet host reads the same objects. diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 24e5f930a..a2e6b2af9 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2855,6 +2855,38 @@ def abort_multipart_upload(*args, **kwargs): assert events[4:] == ["abort"] fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_close_cancelled(self, fs): + # A commit that has started completes even if the task that closes + # the file is cancelled, and the upload is not aborted. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + started = threading.Event() + release = threading.Event() + completed = threading.Event() + + def complete_multipart_upload(*args, **kwargs): + started.set() + # The finally block of the test always releases it. + release.wait() + completed.set() + + fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + await f.write(b"a" * block_size * 2) + task = asyncio.ensure_future(f.close()) + try: + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + assert await asyncio.to_thread(completed.wait, 5) + + assert f.closed + fs.core.complete_multipart_upload.assert_called_once() + fs.core.abort_multipart_upload.assert_not_called() + @pytest.mark.parametrize("data", [b"", b"abc"]) @pytest.mark.asyncio async def test_close_small(self, fs, data): From 42f1453ee10b883a48528c2bb9282e6e6c29ce7e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 01:58:55 +0900 Subject: [PATCH 04/11] Harden the streamed file and the executors found by code review - An AioS3StreamedFile that is garbage collected unclosed warns with a ResourceWarning and aborts its multipart upload; a synchronous with block raises TypeError; open_async(autocommit=False) raises NotImplementedError instead of being ignored. - The creation and the parts of the streamed file are tracked before they are scheduled, with S3AioThreadPool.submit_to(), as S3File does; S3AioThreadPool is an S3Executor. - S3AioExecutor.shutdown(wait=True) does not wait on the thread of its event loop, which would keep the loop from dispatching the pending functions; submit() checks the shutdown under the lock. - S3AioThreadPool takes None as the default of ThreadPoolExecutor, so that AioS3FileSystem(max_workers=None) works again, and replaces a lock that a forked child copied. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 6 +- pyathena/filesystem/s3_async.py | 82 +++++++++++++------ pyathena/filesystem/s3_executor.py | 28 +++++-- tests/pyathena/filesystem/test_s3_async.py | 42 ++++++++++ tests/pyathena/filesystem/test_s3_executor.py | 19 ++++- 5 files changed, 145 insertions(+), 32 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 4062dedf0..64571ab6b 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -325,9 +325,11 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: - A commit that has started completes even if the task that closes the file is cancelled. - A file opened for writing must be closed with `await f.close()` or an `async with` - block. A file that is garbage collected without being closed is not committed. + block. A file that is garbage collected without being closed is not committed: it + emits a `ResourceWarning`, and its multipart upload is aborted. A synchronous `with` + block raises `TypeError`. - `open_async()` raises `NotImplementedError` for the write modes inside a transaction, - because the file commits when it is closed. + and for `autocommit=False`, because the file commits when it is closed. ### Usage with AioS3FSCursor diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 4977b5b85..0b36126b5 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -10,9 +10,11 @@ from __future__ import annotations import asyncio +import contextlib import logging import mimetypes import os +import warnings from concurrent.futures import Future, wait from io import BytesIO from multiprocessing import cpu_count @@ -873,14 +875,16 @@ async def open_async( ValueError: If the mode is not binary, compression is requested, or the arguments are invalid, as in ``open()``. NotImplementedError: If the mode is not supported, or the file is - opened for writing in a transaction, whose commit the file - cannot defer. + opened for writing in a transaction or with + ``autocommit=False``, as the file commits when it is closed. FileExistsError: If an object exists at the path in exclusive-create mode. FileNotFoundError: If no object exists at the path when reading. """ - if "b" not in mode or kwargs.get("compression"): + if "b" not in mode or kwargs.pop("compression", None): raise ValueError("open_async() supports only binary modes without compression.") + if not kwargs.pop("autocommit", True): + raise NotImplementedError("The file of open_async() commits when it is closed.") if mode not in {"rb", "wb", "ab", "xb"}: raise NotImplementedError("File mode not supported") if mode != "rb" and self._intrans: @@ -1170,7 +1174,8 @@ class AioS3StreamedFile(AbstractAsyncStreamedFile): also if the task that waits for them is cancelled. A file opened for writing must be closed with ``await close()`` or an ``async with`` block; garbage collection cannot await it, so the data of an unclosed file is - never committed. + never committed; its multipart upload, if any, is aborted, with a + ``ResourceWarning``. """ fs: AioS3FileSystem @@ -1224,9 +1229,9 @@ def __init__( """ if mode not in {"rb", "wb", "ab", "xb"}: raise NotImplementedError("File mode not supported") - self.max_workers = max_workers # Created first, so that an invalid max_workers raises before the lookups. self._pool = S3AioThreadPool(max_workers) + self.max_workers = self._pool.max_workers target = S3FileTarget.resolve( fs._sync_fs, path, @@ -1272,10 +1277,29 @@ def __init__( def __del__(self) -> None: # close() is a coroutine, which garbage collection cannot await, so - # an unclosed file is not committed; only its threads are released. + # an unclosed file is not committed: its multipart upload, if any, is + # aborted, and its threads are released. pool = getattr(self, "_pool", None) - if pool is not None: - pool.shutdown(wait=False) + if pool is None: + return + if not getattr(self, "closed", True) and self.mode != "rb": + warnings.warn( + f"s3://{self.bucket}/{self.key} was not closed, and the written data " + "was not committed.", + ResourceWarning, + stacklevel=2, + ) + with contextlib.suppress(RuntimeError): + pool.submit(self._discard) + pool.shutdown(wait=False) + + def __enter__(self) -> AioS3StreamedFile: + """Reject a synchronous ``with`` block, whose exit cannot close the file. + + Raises: + TypeError: Always; use ``async with``. + """ + raise TypeError("AioS3StreamedFile is asynchronous; use 'async with'.") def _get_request_kwargs(self, method: str) -> dict[str, Any]: """Select the parameters of the file that an S3 operation accepts. @@ -1384,9 +1408,14 @@ async def _initiate_upload(self) -> None: if not self.append_block and self.tell() < self.blocksize: return writer = cast(S3MultipartWriter, self._writer) + # Tracked before it is scheduled, so that an interrupt while it is + # being scheduled still lets the cleanup cancel it or wait for it. + self._creation = Future() try: - self._creation = self._pool.submit( - writer.initiate, **self._get_request_kwargs("create_multipart_upload") + self._pool.submit_to( + self._creation, + writer.initiate, + **self._get_request_kwargs("create_multipart_upload"), ) await asyncio.wrap_future(self._creation) if self.append_block: @@ -1398,14 +1427,16 @@ async def _initiate_upload(self) -> None: ) size = info.get("size", 0) for part_number, range_ in writer.iter_copy_parts(size): - self._parts.append( - self._pool.submit( - writer.upload_part_copy, - part_number=part_number, - source=S3Path(self.bucket, self.key), - range_=None if whole_source else range_, - **self._get_request_kwargs("upload_part_copy"), - ) + # Tracked before it is scheduled; see the creation. + copy: Future[S3MultipartUploadPart] = Future() + self._parts.append(copy) + self._pool.submit_to( + copy, + writer.upload_part_copy, + part_number=part_number, + source=S3Path(self.bucket, self.key), + range_=None if whole_source else range_, + **self._get_request_kwargs("upload_part_copy"), ) except BaseException: await self._close_without_commit() @@ -1454,18 +1485,21 @@ async def _upload_chunk(self, final: bool = False) -> bool: # A failed part would otherwise keep its data until # the commit while the rest of the data is uploaded. upload.result() - part = self._pool.submit( - writer.upload_part, - part_number=part_number, - body=data, - **self._get_request_kwargs("upload_part"), - ) + # Tracked before it is scheduled; see _initiate_upload(). + part: Future[S3MultipartUploadPart] = Future() self._parts.append(part) waiter = asyncio.wrap_future(part) # The commit reports a failure that no wait found; retrieved # here so that asyncio does not log it as never retrieved. waiter.add_done_callback(lambda w: w.cancelled() or w.exception()) self._uploads_in_flight.add(waiter) + self._pool.submit_to( + part, + writer.upload_part, + part_number=part_number, + body=data, + **self._get_request_kwargs("upload_part"), + ) except BaseException: await self._close_without_commit() raise diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index f513600ce..a56a08064 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -130,7 +130,7 @@ def shutdown(self, wait: bool = True) -> None: self._executor.shutdown(wait=wait) -class S3AioThreadPool: +class S3AioThreadPool(S3Executor): """A thread pool that runs the blocking calls of asyncio code. Unlike ``asyncio.to_thread()``, which runs a call in the default executor @@ -149,15 +149,19 @@ class S3AioThreadPool: ValueError: If ``max_workers`` is not positive. """ - def __init__(self, max_workers: int) -> None: + def __init__(self, max_workers: int | None = None) -> None: """Initialize the pool without starting any thread. Args: - max_workers: The maximum number of calls that run at once. + max_workers: The maximum number of calls that run at once. If + None, ``min(32, os.cpu_count() + 4)``, as with + ``ThreadPoolExecutor``. Raises: ValueError: If ``max_workers`` is not positive. """ + if max_workers is None: + max_workers = min(32, (os.cpu_count() or 1) + 4) if max_workers <= 0: raise ValueError("max_workers must be greater than 0") self.max_workers = max_workers @@ -166,6 +170,7 @@ def __init__(self, max_workers: int) -> None: self._pid: int | None = None self._shutdown = False + @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: """Schedule a blocking call in a thread of the pool. @@ -230,6 +235,7 @@ async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> future.add_done_callback(lambda f: f.cancelled() or f.exception()) return await asyncio.shield(future) + @override def shutdown(self, wait: bool = True) -> None: """Stop accepting calls and release the threads once they are idle. @@ -253,6 +259,10 @@ def _get_executor(self) -> ThreadPoolExecutor: Raises: RuntimeError: If the pool has been shut down. """ + if self._pid is not None and self._pid != os.getpid(): + # A forked child: its copy of the lock may be held by a thread + # of the parent, which does not exist here. + self._lock = threading.Lock() with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") @@ -308,8 +318,6 @@ def __init__( @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: - if self._shutdown: - raise RuntimeError("cannot schedule new futures after shutdown") if self._loop is not None and self._loop.is_running(): # The future of run_coroutine_threadsafe can be cancelled while # the function keeps running in its thread, so the returned future @@ -350,6 +358,8 @@ def settle(task: Future[None]) -> None: future.set_exception(task.exception()) with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") self._pending.add(future) future.add_done_callback(self._release) task = asyncio.run_coroutine_threadsafe(self._pool.run(run), self._loop) @@ -377,6 +387,14 @@ def shutdown(self, wait: bool = True) -> None: with self._lock: self._shutdown = True pending = list(self._pending) + try: + on_loop = asyncio.get_running_loop() is self._loop + except RuntimeError: + on_loop = False + # On the thread of the event loop, such as when a file is garbage + # collected there, waiting would keep the loop from dispatching the + # pending functions, so the pool is shut down after the last one. + wait = wait and not on_loop if wait: # Also the functions that have not reached the pool yet. futures.wait(pending) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index a2e6b2af9..bc81cb417 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -356,6 +356,17 @@ async def test_blocking_calls_bypass_default_executor(self): assert info == {"name": "bucket/key"} + @pytest.mark.asyncio + async def test_max_workers_none(self): + # None takes the default of ThreadPoolExecutor for the threads of the + # filesystem, as for those of S3FileSystem's files. + fs = AioS3FileSystem( + connection=mock.MagicMock(), max_workers=None, skip_instance_cache=True + ) + fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) + + assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} + @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_cancelled(self, cancellations): @@ -2805,6 +2816,36 @@ async def close_and_fail(): fs.core.put_object.assert_called_once() fs.core.abort_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_sync_with(self, fs): + # A synchronous with block cannot await close(), so it is rejected + # instead of leaving the data uncommitted. + f = await fs.open_async("s3://bucket/key", "wb") + with pytest.raises(TypeError, match="async with"), f: + pass + await f.close() + + @pytest.mark.asyncio + async def test_garbage_collected_unclosed(self, fs): + # An unclosed file that is garbage collected warns, and aborts its + # multipart upload instead of leaving it incomplete. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + aborted = threading.Event() + fs.core.abort_multipart_upload = mock.MagicMock(side_effect=lambda *a, **k: aborted.set()) + + async def write_unclosed(): + # Created directly: the futures of open_async() reference the + # file until the garbage collector breaks their cycle. + f = AioS3StreamedFile(fs, "bucket/key", "wb", block_size=block_size) + await f.write(b"a" * block_size * 2) + + # The file is released when the coroutine returns. + with pytest.warns(ResourceWarning, match="was not closed"): + await write_unclosed() + assert await asyncio.to_thread(aborted.wait, 5) + + fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio async def test_write_cancelled(self, fs, cancellations): @@ -2948,6 +2989,7 @@ def get_object(path, range_, **kwargs): ("wb", {"max_workers": 0}, ValueError), ("wb", {"block_size": 2**10}, ValueError), ("wb", {"version_id": "v"}, ValueError), + ("wb", {"autocommit": False}, NotImplementedError), ], ) @pytest.mark.asyncio diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 71e6b5140..3aeabf1d9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -23,6 +23,8 @@ class TestS3AioThreadPool: def test_init(self): + # None takes the default of ThreadPoolExecutor. + assert S3AioThreadPool(None).max_workers == min(32, (os.cpu_count() or 1) + 4) with pytest.raises(ValueError, match="max_workers must be greater than 0"): S3AioThreadPool(0) @@ -102,9 +104,11 @@ def test_shutdown(self): assert events == ["pending finished"] def test_fork(self): - # A forked child, whose copy of the pool has no threads, creates new ones. + # A forked child, whose copy of the pool has no threads, creates new + # ones, also if a thread of the parent held the lock of the pool. pool = S3AioThreadPool(1) parent = pool._get_executor() + pool._lock.acquire() with mock.patch("pyathena.filesystem.s3_executor.os.getpid", return_value=os.getpid() + 1): child = pool._get_executor() assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 @@ -185,6 +189,19 @@ async def main(): assert asyncio.run(main()) == 3 + def test_shutdown_on_loop_thread(self): + # Shutting down with wait=True on the thread of the event loop, such + # as when a file is garbage collected there, does not wait for the + # functions that only the loop dispatches; they still run. + async def main(): + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + future = executor.submit(sum, [1, 2]) + executor.shutdown(wait=True) + assert not future.done() + return await asyncio.wrap_future(future) + + assert asyncio.run(main()) == 3 + def test_submit_without_running_loop(self): with pytest.raises(RuntimeError, match="requires a running event loop"): S3AioExecutor().submit(sum, [1, 2]) From b2b7565571e88841711dcd79932e631e4f6a5d59 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 02:10:33 +0900 Subject: [PATCH 05/11] Simplify the pool shutdown, the fork reset, and the streamed file - S3AioExecutor shuts its pool down from a task of the event loop, which starts after the tasks that dispatch the functions submitted before, instead of tracking every pending future. - S3AioThreadPool resets its threads and lock in a forked child with os.register_at_fork(), as concurrent.futures does, instead of comparing the process ID on every call. - AioS3StreamedFile uses the size of the appended object found when it was opened, instead of a second HeadObject; computes the request parameters of its parts once per flush; and merges the abort into _discard(). open_async() leaves the mode check to the file. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 53 +++++------- pyathena/filesystem/s3_executor.py | 83 +++++++++++-------- tests/pyathena/filesystem/test_s3_executor.py | 14 ++-- 3 files changed, 77 insertions(+), 73 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 0b36126b5..9bdb1896a 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -885,8 +885,6 @@ async def open_async( raise ValueError("open_async() supports only binary modes without compression.") if not kwargs.pop("autocommit", True): raise NotImplementedError("The file of open_async() commits when it is closed.") - if mode not in {"rb", "wb", "ab", "xb"}: - raise NotImplementedError("File mode not supported") if mode != "rb" and self._intrans: raise NotImplementedError("open_async() cannot write in a transaction.") path = self._strip_protocol(path) @@ -1256,6 +1254,8 @@ def __init__( ) self.append_block = False + # The size of the appended object, as looked up when opening. + self._append_size = 0 self._writer: S3MultipartWriter | None = None if mode != "rb": self._writer = S3MultipartWriter( @@ -1273,7 +1273,7 @@ def __init__( else: # Copied with UploadPartCopy as the leading part(s). self.append_block = True - self.loc = target.appended.get("size", 0) + self._append_size = self.loc = target.appended.get("size", 0) def __del__(self) -> None: # close() is a coroutine, which garbage collection cannot await, so @@ -1419,13 +1419,10 @@ async def _initiate_upload(self) -> None: ) await asyncio.wrap_future(self._creation) if self.append_block: - whole_source = self.tell() <= self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE - size = self.tell() - if not whole_source: - info = await self.fs._info( - self.path, **self.fs._sync_fs._get_lookup_kwargs(self.s3_additional_kwargs) - ) - size = info.get("size", 0) + size = self._append_size + whole_source = size <= self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE + source = S3Path(self.bucket, self.key) + copy_kwargs = self._get_request_kwargs("upload_part_copy") for part_number, range_ in writer.iter_copy_parts(size): # Tracked before it is scheduled; see the creation. copy: Future[S3MultipartUploadPart] = Future() @@ -1434,9 +1431,9 @@ async def _initiate_upload(self) -> None: copy, writer.upload_part_copy, part_number=part_number, - source=S3Path(self.bucket, self.key), + source=source, range_=None if whole_source else range_, - **self._get_request_kwargs("upload_part_copy"), + **copy_kwargs, ) except BaseException: await self._close_without_commit() @@ -1472,6 +1469,7 @@ async def _upload_chunk(self, final: bool = False) -> bool: raise RuntimeError("Multipart upload is not initialized.") buffer = cast(BytesIO, self.buffer) buffer.seek(0) + upload_kwargs = self._get_request_kwargs("upload_part") try: for part_number, data in writer.iter_parts( buffer, first_part_number=len(self._parts) + 1 @@ -1498,7 +1496,7 @@ async def _upload_chunk(self, final: bool = False) -> bool: writer.upload_part, part_number=part_number, body=data, - **self._get_request_kwargs("upload_part"), + **upload_kwargs, ) except BaseException: await self._close_without_commit() @@ -1552,32 +1550,25 @@ def _complete_upload(self) -> None: self.fs.invalidate_cache(self.path) def _discard(self) -> None: - """Abort the multipart upload, if any, logging a failure; blocks. + """Abort the multipart upload, if any; blocks, and runs in a thread. - The failure is logged instead of raised, so that it does not mask - the error that the caller is handling, also when no task waits for - the abort any more. + The requests that have not started are cancelled, and the running + ones are waited for first, since a part that is still uploading when + the upload is aborted may be stored after the abort. A failure is + logged instead of raised, so that it does not mask the error that the + caller is handling, also when no task waits for the abort any more. """ try: - self._abort_upload() + if self._creation is not None and not self._creation.cancel(): + wait([self._creation]) + if self._writer is not None and self._writer.upload is not None: + wait([part for part in self._parts if not part.cancel()]) + self._writer.abort(**self._get_request_kwargs("abort_multipart_upload")) except Exception: _logger.exception( f"Failed to abort the multipart upload to s3://{self.bucket}/{self.key}." ) - def _abort_upload(self) -> None: - """Abort the multipart upload, if any; blocks, and runs in a thread. - - The requests that have not started are cancelled, and the running - ones are waited for first, since a part that is still uploading when - the upload is aborted may be stored after the abort. - """ - if self._creation is not None and not self._creation.cancel(): - wait([self._creation]) - if self._writer is not None and self._writer.upload is not None: - wait([part for part in self._parts if not part.cancel()]) - self._writer.abort(**self._get_request_kwargs("abort_multipart_upload")) - async def _fetch_range(self, start: int, end: int) -> bytes: """Read a byte range of the object with parallel range requests. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index a56a08064..f5e7fac7f 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -10,16 +10,18 @@ from __future__ import annotations import asyncio +import contextlib import contextvars import os import threading +import weakref from abc import ABCMeta, abstractmethod from collections.abc import Callable from concurrent import futures from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count -from typing import Any, TypeVar +from typing import Any, ClassVar, TypeVar from pyathena.util import override @@ -149,6 +151,9 @@ class S3AioThreadPool(S3Executor): ValueError: If ``max_workers`` is not positive. """ + # The pools of this process, which a forked child resets. + _instances: ClassVar[weakref.WeakSet[S3AioThreadPool]] = weakref.WeakSet() + def __init__(self, max_workers: int | None = None) -> None: """Initialize the pool without starting any thread. @@ -167,8 +172,8 @@ def __init__(self, max_workers: int | None = None) -> None: self.max_workers = max_workers self._lock = threading.Lock() self._executor: ThreadPoolExecutor | None = None - self._pid: int | None = None self._shutdown = False + S3AioThreadPool._instances.add(self) @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: @@ -246,10 +251,21 @@ def shutdown(self, wait: bool = True) -> None: """ with self._lock: self._shutdown = True - executor = self._executor if self._pid == os.getpid() else None + executor = self._executor if executor is not None: executor.shutdown(wait=wait) + @classmethod + def _reset_after_fork(cls) -> None: + """Drop the threads and locks that a forked child copied from its parent. + + The threads of the parent do not exist in the child, and a lock that + one of them held stays held, so each pool creates new ones. + """ + for pool in list(cls._instances): + pool._lock = threading.Lock() + pool._executor = None + def _get_executor(self) -> ThreadPoolExecutor: """Return the thread pool of this process, creating it if needed. @@ -259,19 +275,21 @@ def _get_executor(self) -> ThreadPoolExecutor: Raises: RuntimeError: If the pool has been shut down. """ - if self._pid is not None and self._pid != os.getpid(): - # A forked child: its copy of the lock may be held by a thread - # of the parent, which does not exist here. - self._lock = threading.Lock() + if (executor := self._executor) is not None: + # Its submit() raises after the shutdown. + return executor with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - if self._executor is None or self._pid != os.getpid(): + if self._executor is None: self._executor = ThreadPoolExecutor(max_workers=self.max_workers) - self._pid = os.getpid() return self._executor +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=S3AioThreadPool._reset_after_fork) + + class S3AioExecutor(S3Executor): """Executor that schedules work on an asyncio event loop. @@ -309,11 +327,8 @@ def __init__( """ self._loop = loop self._pool = S3AioThreadPool(max_workers) + # Orders the dispatches with the shutdown. self._lock = threading.Lock() - # The returned futures that are not done. A function reaches the pool - # only when the event loop runs its task, so the pool is shut down - # once they are all done. - self._pending: set[Future[Any]] = set() self._shutdown = False @override @@ -360,9 +375,7 @@ def settle(task: Future[None]) -> None: with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - self._pending.add(future) - future.add_done_callback(self._release) - task = asyncio.run_coroutine_threadsafe(self._pool.run(run), self._loop) + task = asyncio.run_coroutine_threadsafe(self._pool.run(run), self._loop) task.add_done_callback(settle) return future raise RuntimeError( @@ -370,33 +383,35 @@ def settle(task: Future[None]) -> None: "Use S3ThreadPoolExecutor for synchronous usage." ) - def _release(self, future: Future[Any]) -> None: - """Forget a done future, shutting down the pool after the last one. + async def _shutdown_pool(self) -> None: + """Shut down the pool once the functions submitted before have reached it. - Args: - future: The done future. + Run as a task of the event loop, which starts its tasks in the order + they were scheduled, so the tasks that dispatch the functions + submitted before have handed them to the pool. """ - with self._lock: - self._pending.discard(future) - last = self._shutdown and not self._pending - if last: - self._pool.shutdown(wait=False) + self._pool.shutdown(wait=False) @override def shutdown(self, wait: bool = True) -> None: with self._lock: self._shutdown = True - pending = list(self._pending) + closing = None + if self._loop is not None and self._loop.is_running(): + # A function reaches the pool only when the event loop runs + # its task, so the pool is shut down after the dispatches. + closing = asyncio.run_coroutine_threadsafe(self._shutdown_pool(), self._loop) + if closing is None: + self._pool.shutdown(wait=wait) + return try: on_loop = asyncio.get_running_loop() is self._loop except RuntimeError: on_loop = False # On the thread of the event loop, such as when a file is garbage - # collected there, waiting would keep the loop from dispatching the - # pending functions, so the pool is shut down after the last one. - wait = wait and not on_loop - if wait: - # Also the functions that have not reached the pool yet. - futures.wait(pending) - if not pending or wait: - self._pool.shutdown(wait=wait) + # collected there, waiting would keep the loop from dispatching. + if wait and not on_loop: + # Cancelled if the event loop shuts down before it runs. + with contextlib.suppress(futures.CancelledError): + closing.result() + self._pool.shutdown(wait=True) diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 3aeabf1d9..7029ea422 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -11,14 +11,11 @@ import threading import time from concurrent.futures import Future, ThreadPoolExecutor, wait -from unittest import mock import pytest from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor -from tests.pyathena.util import overlapping_calls - -EVENT_TIMEOUT = 5 +from tests.pyathena.util import EVENT_TIMEOUT, overlapping_calls class TestS3AioThreadPool: @@ -109,10 +106,11 @@ def test_fork(self): pool = S3AioThreadPool(1) parent = pool._get_executor() pool._lock.acquire() - with mock.patch("pyathena.filesystem.s3_executor.os.getpid", return_value=os.getpid() + 1): - child = pool._get_executor() - assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 - pool.shutdown() + # As os.register_at_fork() runs it in a forked child. + S3AioThreadPool._reset_after_fork() + child = pool._get_executor() + assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 + pool.shutdown() assert child is not parent parent.shutdown() From 5af66bad12f034613b58ac77b39cc883a76c53ff Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 02:29:43 +0900 Subject: [PATCH 06/11] Keep a running commit from an abort and a stopped loop from a shutdown - Once AioS3StreamedFile submits its commit, the commit completes the upload or aborts it on a failure, and closing the file without committing does not abort it too. Leaving an async with block with an exception after a cancelled final flush used to race the completion with an abort. - S3AioExecutor.shutdown(wait=True) stops waiting for the task that shuts the pool down when the event loop stops before running it. - The busy-executor test holds the default executor until it releases it, so that an implementation using it cannot pass on a timeout. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 5 ++- pyathena/filesystem/s3_async.py | 19 +++++++-- pyathena/filesystem/s3_executor.py | 19 +++++---- tests/pyathena/filesystem/test_s3_async.py | 42 ++++++++++++++++++- tests/pyathena/filesystem/test_s3_executor.py | 38 +++++++++++++++++ 5 files changed, 109 insertions(+), 14 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 64571ab6b..95965cccf 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -322,8 +322,9 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: - Leaving an `async with` block with an exception, or cancelling a write, closes the file without committing it, and aborts the multipart upload after the running parts finish. The abort completes even if the task is cancelled again. -- A commit that has started completes even if the task that closes the file is - cancelled. +- A commit that has started completes even if the task that closes or flushes the + file is cancelled, and leaving an `async with` block with an exception afterward does + not abort it. - A file opened for writing must be closed with `await f.close()` or an `async with` block. A file that is garbage collected without being closed is not committed: it emits a `ResourceWarning`, and its multipart upload is aborted. A synchronous `with` diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 9bdb1896a..b2928c516 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1169,7 +1169,8 @@ class AioS3StreamedFile(AbstractAsyncStreamedFile): so that the existing object is left unchanged. The commit and the abort run to their end in the threads of the file, - also if the task that waits for them is cancelled. A file opened for + also if the task that waits for them is cancelled; once the commit has + started, nothing else aborts the upload. A file opened for writing must be closed with ``await close()`` or an ``async with`` block; garbage collection cannot await it, so the data of an unclosed file is never committed; its multipart upload, if any, is aborted, with a @@ -1262,6 +1263,9 @@ def __init__( fs.core, S3Path(self.bucket, self.key), block_size=block_size ) self._creation: Future[S3MultipartUpload] | None = None + # Set when the commit is submitted, after which it completes the + # upload, or aborts it if it fails, and nothing else aborts it. + self._committing = False # The futures of the part uploads and copies, in part-number order. self._parts: list[Future[S3MultipartUploadPart]] = [] # The parts uploaded from the buffer that may not have finished. The @@ -1282,7 +1286,11 @@ def __del__(self) -> None: pool = getattr(self, "_pool", None) if pool is None: return - if not getattr(self, "closed", True) and self.mode != "rb": + if ( + not getattr(self, "closed", True) + and self.mode != "rb" + and not getattr(self, "_committing", False) + ): warnings.warn( f"s3://{self.bucket}/{self.key} was not closed, and the written data " "was not committed.", @@ -1397,8 +1405,12 @@ async def _close_without_commit(self) -> None: """ if self.closed: return - self.buffer = None self.closed = True + if self._committing: + # A cancelled commit runs on; an abort would race its completion. + self._pool.shutdown(wait=False) + return + self.buffer = None try: await self._pool.run_shielded(self._discard) finally: @@ -1512,6 +1524,7 @@ async def _commit(self) -> None: """ if self.buffer is None: return + self._committing = True await self._pool.run_shielded(self._complete_upload) def _complete_upload(self) -> None: diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index f5e7fac7f..1f83c39d9 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -10,7 +10,6 @@ from __future__ import annotations import asyncio -import contextlib import contextvars import os import threading @@ -394,24 +393,28 @@ async def _shutdown_pool(self) -> None: @override def shutdown(self, wait: bool = True) -> None: + loop = self._loop with self._lock: self._shutdown = True closing = None - if self._loop is not None and self._loop.is_running(): + if loop is not None and loop.is_running(): # A function reaches the pool only when the event loop runs # its task, so the pool is shut down after the dispatches. - closing = asyncio.run_coroutine_threadsafe(self._shutdown_pool(), self._loop) - if closing is None: + closing = asyncio.run_coroutine_threadsafe(self._shutdown_pool(), loop) + if loop is None or closing is None: self._pool.shutdown(wait=wait) return try: - on_loop = asyncio.get_running_loop() is self._loop + on_loop = asyncio.get_running_loop() is loop except RuntimeError: on_loop = False # On the thread of the event loop, such as when a file is garbage # collected there, waiting would keep the loop from dispatching. if wait and not on_loop: - # Cancelled if the event loop shuts down before it runs. - with contextlib.suppress(futures.CancelledError): - closing.result() + # The task is cancelled if the event loop shuts down before it + # runs, and never runs if the loop stops instead; the functions + # that it has not dispatched then do not reach the pool either. + while not futures.wait([closing], timeout=0.1).done: + if not loop.is_running(): + break self._pool.shutdown(wait=True) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index bc81cb417..0bfa1c2e1 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -347,7 +347,8 @@ async def test_blocking_calls_bypass_default_executor(self): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) release = threading.Event() - blocker = loop.run_in_executor(None, release.wait, 5) + # Held until it is released, after the lookup has had to finish. + blocker = loop.run_in_executor(None, release.wait, 60) try: info = await asyncio.wait_for(fs._info("s3://bucket/key"), 5) finally: @@ -2928,6 +2929,45 @@ def complete_multipart_upload(*args, **kwargs): fs.core.complete_multipart_upload.assert_called_once() fs.core.abort_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_flush_cancelled_then_exit_with_exception(self, fs): + # Leaving the block with an exception after a cancelled final flush + # does not abort the upload that the running commit completes. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + started = threading.Event() + release = threading.Event() + completed = threading.Event() + + def complete_multipart_upload(*args, **kwargs): + started.set() + # The finally block of the test always releases it. + release.wait() + completed.set() + + fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + + async def flush_and_exit(): + async with f: + await f.write(b"a" * block_size * 2) + flush = asyncio.ensure_future(f.flush(force=True)) + assert await asyncio.to_thread(started.wait, 5) + flush.cancel() + await flush + + try: + with pytest.raises(asyncio.CancelledError): + await flush_and_exit() + finally: + release.set() + assert await asyncio.to_thread(completed.wait, 5) + # Lets an abort that was submitted run. + await asyncio.sleep(0.1) + + assert f.closed + fs.core.complete_multipart_upload.assert_called_once() + fs.core.abort_multipart_upload.assert_not_called() + @pytest.mark.parametrize("data", [b"", b"abc"]) @pytest.mark.asyncio async def test_close_small(self, fs, data): diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 7029ea422..f1ca6ff0b 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -200,6 +200,44 @@ async def main(): assert asyncio.run(main()) == 3 + def test_shutdown_loop_stopped(self): + # Shutting down with wait=True returns when the event loop stops + # before it runs the task that shuts the pool down. + loop = asyncio.new_event_loop() + busy = threading.Event() + release = threading.Event() + + def block(): + busy.set() + release.wait(EVENT_TIMEOUT) + # The loop exits after this callback, without running the task + # that the shutdown schedules meanwhile. + loop.stop() + + thread = threading.Thread(target=loop.run_forever) + thread.start() + try: + loop.call_soon_threadsafe(block) + assert busy.wait(EVENT_TIMEOUT) + executor = S3AioExecutor(loop=loop) + shut_down = threading.Event() + waiter = threading.Thread( + target=lambda: (executor.shutdown(wait=True), shut_down.set()) + ) + waiter.start() + # The shutdown waits for the task while the loop is busy. + assert not shut_down.wait(0.2) + release.set() + assert shut_down.wait(EVENT_TIMEOUT) + waiter.join() + finally: + release.set() + thread.join() + # Runs the pending task, which would otherwise be dropped + # unawaited when the loop is closed. + loop.run_until_complete(asyncio.sleep(0)) + loop.close() + def test_submit_without_running_loop(self): with pytest.raises(RuntimeError, match="requires a running event loop"): S3AioExecutor().submit(sum, [1, 2]) From 683b5b370a6119f9989400a3a3962e222381d1a0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 02:40:15 +0900 Subject: [PATCH 07/11] Synchronize the stopped-loop test and qualify the collection guarantee - The stopped-loop shutdown test releases the loop only after the shutdown has scheduled its task, so that it cannot pass by the loop stopping first. - The garbage collection of an unclosed AioS3StreamedFile warns and aborts only when its commit has not started; the docs say so. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 6 +++--- pyathena/filesystem/s3_async.py | 9 +++++---- tests/pyathena/filesystem/test_s3_executor.py | 18 +++++++++++++++--- 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 95965cccf..3454808dc 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -326,9 +326,9 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: file is cancelled, and leaving an `async with` block with an exception afterward does not abort it. - A file opened for writing must be closed with `await f.close()` or an `async with` - block. A file that is garbage collected without being closed is not committed: it - emits a `ResourceWarning`, and its multipart upload is aborted. A synchronous `with` - block raises `TypeError`. + block. A file that is garbage collected before it is closed or its commit has started + is not committed: it emits a `ResourceWarning`, and its multipart upload is aborted. + A synchronous `with` block raises `TypeError`. - `open_async()` raises `NotImplementedError` for the write modes inside a transaction, and for `autocommit=False`, because the file commits when it is closed. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index b2928c516..6dacced45 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1172,8 +1172,9 @@ class AioS3StreamedFile(AbstractAsyncStreamedFile): also if the task that waits for them is cancelled; once the commit has started, nothing else aborts the upload. A file opened for writing must be closed with ``await close()`` or an ``async with`` block; - garbage collection cannot await it, so the data of an unclosed file is - never committed; its multipart upload, if any, is aborted, with a + garbage collection cannot await it, so the data of a file that is + collected before it is closed or its commit has started is never + committed; its multipart upload, if any, is aborted, with a ``ResourceWarning``. """ @@ -1281,8 +1282,8 @@ def __init__( def __del__(self) -> None: # close() is a coroutine, which garbage collection cannot await, so - # an unclosed file is not committed: its multipart upload, if any, is - # aborted, and its threads are released. + # an unclosed file whose commit has not started is not committed: its + # multipart upload, if any, is aborted, and its threads are released. pool = getattr(self, "_pool", None) if pool is None: return diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index f1ca6ff0b..f15613873 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -11,6 +11,7 @@ import threading import time from concurrent.futures import Future, ThreadPoolExecutor, wait +from unittest import mock import pytest @@ -214,6 +215,14 @@ def block(): # that the shutdown schedules meanwhile. loop.stop() + scheduled = threading.Event() + run_coroutine_threadsafe = asyncio.run_coroutine_threadsafe + + def schedule(coro, loop): + future = run_coroutine_threadsafe(coro, loop) + scheduled.set() + return future + thread = threading.Thread(target=loop.run_forever) thread.start() try: @@ -224,9 +233,12 @@ def block(): waiter = threading.Thread( target=lambda: (executor.shutdown(wait=True), shut_down.set()) ) - waiter.start() - # The shutdown waits for the task while the loop is busy. - assert not shut_down.wait(0.2) + with mock.patch("asyncio.run_coroutine_threadsafe", side_effect=schedule): + waiter.start() + # The loop stops only after the task has been scheduled, so + # that the shutdown waits for a task that never runs. + assert scheduled.wait(EVENT_TIMEOUT) + assert not shut_down.is_set() release.set() assert shut_down.wait(EVENT_TIMEOUT) waiter.join() From 1c50b5ee6877cdd87354932508fa2a3d9d809b27 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 07:55:40 +0900 Subject: [PATCH 08/11] Return from a shutdown at interpreter exit and normalize max_workers=None From the review of claude-fable-5-1: - S3AioExecutor.shutdown(wait=True) stops waiting when the interpreter is finalizing. fsspec's loop runs in a daemon thread, which stops at exit while the loop still reports running, so closing a file at exit used to wait forever. - AioS3FileSystem normalizes max_workers=None to the pool's default before it creates the internal S3FileSystem, so that the files and copies get a number too. - A forked child keeps the executors copied from its parent instead of letting them be collected, which would put to a queue whose lock a thread of the parent may have held. - The parts-in-flight test counts the parts submitted and not finished, which the pool size alone does not bound; the docs say that a commit that still waits for a thread completes too. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 8 ++--- pyathena/filesystem/s3_async.py | 20 ++++++----- pyathena/filesystem/s3_executor.py | 12 ++++++- tests/pyathena/filesystem/test_s3_async.py | 33 +++++++++++++++---- tests/pyathena/filesystem/test_s3_executor.py | 33 +++++++++++++++++++ 5 files changed, 86 insertions(+), 20 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 3454808dc..0493bac91 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -322,11 +322,11 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: - Leaving an `async with` block with an exception, or cancelling a write, closes the file without committing it, and aborts the multipart upload after the running parts finish. The abort completes even if the task is cancelled again. -- A commit that has started completes even if the task that closes or flushes the - file is cancelled, and leaving an `async with` block with an exception afterward does - not abort it. +- Once closing or the final flush has begun the commit, the commit completes even if + the task that closes or flushes the file is cancelled, and leaving an `async with` + block with an exception afterward does not abort it. - A file opened for writing must be closed with `await f.close()` or an `async with` - block. A file that is garbage collected before it is closed or its commit has started + block. A file that is garbage collected before it is closed or its commit has begun is not committed: it emits a `ResourceWarning`, and its multipart upload is aborted. A synchronous `with` block raises `TypeError`. - `open_async()` raises `NotImplementedError` for the write modes inside a transaction, diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 6dacced45..6bd1f6b62 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -120,7 +120,8 @@ def __init__( default_cache_type: Passed to the internal ``S3FileSystem``. max_workers: The number of threads that run the blocking calls of the filesystem, and the default number of parallel requests - of its files and copies. Passed to the internal + of its files and copies. If None, ``min(32, os.cpu_count() + + 4)``, as with ``ThreadPoolExecutor``. Passed to the internal ``S3FileSystem``. s3_additional_kwargs: Passed to the internal ``S3FileSystem``. allow_bucket_creation: Passed to the internal ``S3FileSystem``. @@ -138,6 +139,11 @@ def __init__( batch_size=batch_size, **kwargs, ) + # The blocking calls of the coroutines run in threads of this pool, + # not in the default executor of the event loop. Created first, so + # that None takes its default for the files and copies too. + self._thread_pool = S3AioThreadPool(max_workers) + max_workers = self._thread_pool.max_workers self._sync_fs = S3FileSystem( connection=connection, default_block_size=default_block_size, @@ -153,9 +159,6 @@ def __init__( ) # Share dircache for cache coherence between async and sync instances self.dircache = self._sync_fs.dircache - # The blocking calls of the coroutines run in threads of this pool, - # not in the default executor of the event loop. - self._thread_pool = S3AioThreadPool(max_workers) @property def core(self) -> S3Core: @@ -1169,11 +1172,12 @@ class AioS3StreamedFile(AbstractAsyncStreamedFile): so that the existing object is left unchanged. The commit and the abort run to their end in the threads of the file, - also if the task that waits for them is cancelled; once the commit has - started, nothing else aborts the upload. A file opened for + also if the task that waits for them is cancelled, including a commit + that still waits for a thread; once the commit has been submitted, + nothing else aborts the upload. A file opened for writing must be closed with ``await close()`` or an ``async with`` block; garbage collection cannot await it, so the data of a file that is - collected before it is closed or its commit has started is never + collected before it is closed or its commit has been submitted is never committed; its multipart upload, if any, is aborted, with a ``ResourceWarning``. """ @@ -1282,7 +1286,7 @@ def __init__( def __del__(self) -> None: # close() is a coroutine, which garbage collection cannot await, so - # an unclosed file whose commit has not started is not committed: its + # an unclosed file whose commit was not submitted is not committed: its # multipart upload, if any, is aborted, and its threads are released. pool = getattr(self, "_pool", None) if pool is None: diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 1f83c39d9..b3068da6a 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -12,6 +12,7 @@ import asyncio import contextvars import os +import sys import threading import weakref from abc import ABCMeta, abstractmethod @@ -152,6 +153,10 @@ class S3AioThreadPool(S3Executor): # The pools of this process, which a forked child resets. _instances: ClassVar[weakref.WeakSet[S3AioThreadPool]] = weakref.WeakSet() + # The executors that a forked child copied from its parent. They are + # kept, since collecting one would put to a queue whose lock a thread of + # the parent may have held. + _forked_executors: ClassVar[list[ThreadPoolExecutor]] = [] def __init__(self, max_workers: int | None = None) -> None: """Initialize the pool without starting any thread. @@ -262,6 +267,8 @@ def _reset_after_fork(cls) -> None: one of them held stays held, so each pool creates new ones. """ for pool in list(cls._instances): + if pool._executor is not None: + cls._forked_executors.append(pool._executor) pool._lock = threading.Lock() pool._executor = None @@ -415,6 +422,9 @@ def shutdown(self, wait: bool = True) -> None: # runs, and never runs if the loop stops instead; the functions # that it has not dispatched then do not reach the pool either. while not futures.wait([closing], timeout=0.1).done: - if not loop.is_running(): + # The thread of a loop that runs in a daemon thread, such as + # fsspec's, stops when the interpreter exits, while the loop + # still reports running. + if not loop.is_running() or sys.is_finalizing(): break self._pool.shutdown(wait=True) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 0bfa1c2e1..8da3e9cc0 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -8,7 +8,7 @@ import urllib.parse import urllib.request import uuid -from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain from pathlib import Path @@ -24,7 +24,7 @@ from pyathena.filesystem.s3 import S3File, S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3Executor +from pyathena.filesystem.s3_executor import S3AioThreadPool, S3Executor from pyathena.filesystem.s3_object import ( S3MultipartUpload, S3MultipartUploadPart, @@ -360,13 +360,20 @@ async def test_blocking_calls_bypass_default_executor(self): @pytest.mark.asyncio async def test_max_workers_none(self): # None takes the default of ThreadPoolExecutor for the threads of the - # filesystem, as for those of S3FileSystem's files. + # filesystem and as the default max_workers of its files and copies. + # The files and copies used to fail with it. fs = AioS3FileSystem( connection=mock.MagicMock(), max_workers=None, skip_instance_cache=True ) fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} + # The files and copies take the same default. + assert ( + fs._sync_fs.max_workers + == fs._thread_pool.max_workers + == min(32, (os.cpu_count() or 1) + 4) + ) @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio @@ -2718,14 +2725,23 @@ async def test_write_parts_in_flight(self, fs): lock = threading.Lock() gates: dict[int, threading.Event] = {} active: list[int] = [] + uploads: dict[int, Future[S3MultipartUploadPart]] = {} in_flight: list[int] = [] + submit_to = S3AioThreadPool.submit_to + + def record_upload(pool, future, fn, *args, **kwargs): + # The parts submitted and not finished, including those that + # wait for a thread with their data. + if "body" in kwargs: + uploads[kwargs["part_number"]] = future + in_flight.append(sum(not upload.done() for upload in uploads.values())) + submit_to(pool, future, fn, *args, **kwargs) def upload_part(**kw): part_number = kw["part_number"] with lock: gate = gates.setdefault(part_number, threading.Event()) active.append(part_number) - in_flight.append(len(active)) if not gate.wait(5): raise TimeoutError(part_number) with lock: @@ -2750,9 +2766,12 @@ async def release(): releaser = asyncio.ensure_future(release()) f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) - await f.write(b"a" * block_size * 4) - written.set() - await f.close() + with mock.patch.object( + S3AioThreadPool, "submit_to", autospec=True, side_effect=record_upload + ): + await f.write(b"a" * block_size * 4) + written.set() + await f.close() await releaser assert in_flight == [1, 2, 2, 2] diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index f15613873..07a51fcba 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -109,10 +109,14 @@ def test_fork(self): pool._lock.acquire() # As os.register_at_fork() runs it in a forked child. S3AioThreadPool._reset_after_fork() + # The copied executor is kept, so that its collection does not touch + # the locks that the parent's threads may have held. + assert parent in S3AioThreadPool._forked_executors child = pool._get_executor() assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 pool.shutdown() assert child is not parent + S3AioThreadPool._forked_executors.remove(parent) parent.shutdown() @@ -250,6 +254,35 @@ def schedule(coro, loop): loop.run_until_complete(asyncio.sleep(0)) loop.close() + def test_shutdown_interpreter_finalizing(self): + # Shutting down with wait=True while the interpreter exits returns, + # although the loop, whose daemon thread no longer runs, still + # reports running. + loop = asyncio.new_event_loop() + busy = threading.Event() + release = threading.Event() + + def block(): + busy.set() + release.wait(EVENT_TIMEOUT) + + thread = threading.Thread(target=loop.run_forever) + thread.start() + try: + loop.call_soon_threadsafe(block) + assert busy.wait(EVENT_TIMEOUT) + executor = S3AioExecutor(loop=loop) + with mock.patch("pyathena.filesystem.s3_executor.sys.is_finalizing", return_value=True): + started = time.monotonic() + executor.shutdown(wait=True) + assert time.monotonic() - started < EVENT_TIMEOUT + finally: + release.set() + loop.call_soon_threadsafe(loop.stop) + thread.join() + loop.run_until_complete(asyncio.sleep(0)) + loop.close() + def test_submit_without_running_loop(self): with pytest.raises(RuntimeError, match="requires a running event loop"): S3AioExecutor().submit(sum, [1, 2]) From aef5dc2377cee25a6a93d19afa94d0d2ffc7ce0d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 08:06:04 +0900 Subject: [PATCH 09/11] Hold the loop through the finalization test and count usable CPUs - The finalization shutdown test keeps the loop blocked until the shutdown has returned, so that it cannot pass by the loop running the task after a timeout. - The default max_workers counts the CPUs that the process may use on Python 3.13 and later, as ThreadPoolExecutor does. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 4 ++-- pyathena/filesystem/s3_executor.py | 9 ++++++--- tests/pyathena/filesystem/test_s3_async.py | 7 ++----- tests/pyathena/filesystem/test_s3_executor.py | 17 ++++++++++++----- 4 files changed, 22 insertions(+), 15 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 6bd1f6b62..c3137a4ed 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -120,8 +120,8 @@ def __init__( default_cache_type: Passed to the internal ``S3FileSystem``. max_workers: The number of threads that run the blocking calls of the filesystem, and the default number of parallel requests - of its files and copies. If None, ``min(32, os.cpu_count() + - 4)``, as with ``ThreadPoolExecutor``. Passed to the internal + of its files and copies. If None, the default of + ``ThreadPoolExecutor``: ``min(32, CPUs + 4)``. Passed to the internal ``S3FileSystem``. s3_additional_kwargs: Passed to the internal ``S3FileSystem``. allow_bucket_creation: Passed to the internal ``S3FileSystem``. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index b3068da6a..ce9b57112 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -163,14 +163,17 @@ def __init__(self, max_workers: int | None = None) -> None: Args: max_workers: The maximum number of calls that run at once. If - None, ``min(32, os.cpu_count() + 4)``, as with - ``ThreadPoolExecutor``. + None, the default of ``ThreadPoolExecutor``: ``min(32, CPUs + + 4)``. Raises: ValueError: If ``max_workers`` is not positive. """ if max_workers is None: - max_workers = min(32, (os.cpu_count() or 1) + 4) + # As ThreadPoolExecutor, which counts the CPUs that the process + # may use since Python 3.13. + cpus = getattr(os, "process_cpu_count", os.cpu_count)() + max_workers = min(32, (cpus or 1) + 4) if max_workers <= 0: raise ValueError("max_workers must be greater than 0") self.max_workers = max_workers diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 8da3e9cc0..aaea78a8d 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -369,11 +369,8 @@ async def test_max_workers_none(self): assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} # The files and copies take the same default. - assert ( - fs._sync_fs.max_workers - == fs._thread_pool.max_workers - == min(32, (os.cpu_count() or 1) + 4) - ) + cpus = getattr(os, "process_cpu_count", os.cpu_count)() + assert fs._sync_fs.max_workers == fs._thread_pool.max_workers == min(32, (cpus or 1) + 4) @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 07a51fcba..b4aa849f9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -22,7 +22,8 @@ class TestS3AioThreadPool: def test_init(self): # None takes the default of ThreadPoolExecutor. - assert S3AioThreadPool(None).max_workers == min(32, (os.cpu_count() or 1) + 4) + cpus = getattr(os, "process_cpu_count", os.cpu_count)() + assert S3AioThreadPool(None).max_workers == min(32, (cpus or 1) + 4) with pytest.raises(ValueError, match="max_workers must be greater than 0"): S3AioThreadPool(0) @@ -264,7 +265,9 @@ def test_shutdown_interpreter_finalizing(self): def block(): busy.set() - release.wait(EVENT_TIMEOUT) + # Held until the shutdown has returned, so that the loop cannot + # run the task that the shutdown waits for. + release.wait(6 * EVENT_TIMEOUT) thread = threading.Thread(target=loop.run_forever) thread.start() @@ -272,10 +275,14 @@ def block(): loop.call_soon_threadsafe(block) assert busy.wait(EVENT_TIMEOUT) executor = S3AioExecutor(loop=loop) + shut_down = threading.Event() + waiter = threading.Thread( + target=lambda: (executor.shutdown(wait=True), shut_down.set()) + ) with mock.patch("pyathena.filesystem.s3_executor.sys.is_finalizing", return_value=True): - started = time.monotonic() - executor.shutdown(wait=True) - assert time.monotonic() - started < EVENT_TIMEOUT + waiter.start() + assert shut_down.wait(EVENT_TIMEOUT) + waiter.join() finally: release.set() loop.call_soon_threadsafe(loop.stop) From 55b8c7616ebe30887a0a0b741a80f6392b60e98b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 08:21:13 +0900 Subject: [PATCH 10/11] Stop the streamed file at a failed part copy and a commit at a failed part As S3File does since #1183: - A write of an AioS3StreamedFile raises a part copy of an append found failed before it submits a part, instead of uploading the rest of the data. - The commit raises a failed part without waiting for the parts before it, so that the parts that have not started are cancelled. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 8 +-- pyathena/filesystem/s3_async.py | 34 ++++++++++--- tests/pyathena/filesystem/test_s3_async.py | 57 ++++++++++++++++++++++ 3 files changed, 89 insertions(+), 10 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 0493bac91..7b60ec05a 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -315,9 +315,11 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: parallel. The file has no read cache. - A write uploads each full block as a part of a multipart upload, with at most `max_workers` parts in flight. A write that would upload another part awaits one of - them, so the event loop keeps running. A part found failed by that wait stops the - write with its error, and the multipart upload is aborted. A part that fails - afterward raises its error when the file is closed. + them, so the event loop keeps running. A part found failed by that wait, or a part + copied from the existing object in an append that has failed before the write uploads + another part, stops the write with its error, and the multipart upload is aborted. + A part that fails later raises its error when the file is closed, without waiting for + the parts before it. - Closing the file commits the written data, as `S3File` does. - Leaving an `async with` block with an exception, or cancelling a write, closes the file without committing it, and aborts the multipart upload after the running parts finish. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index c3137a4ed..d56e966ad 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -15,7 +15,7 @@ import mimetypes import os import warnings -from concurrent.futures import Future, wait +from concurrent.futures import FIRST_EXCEPTION, Future, wait from io import BytesIO from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast @@ -1160,7 +1160,9 @@ class AioS3StreamedFile(AbstractAsyncStreamedFile): - A write buffers the data, and uploads each full block as a part of a multipart upload, with at most ``max_workers`` parts in flight. A write that would submit another part waits for one of them to finish first, - and a part found failed by that wait stops the write with its error. + and a part found failed by that wait, or a part copy of an append + found failed before a part is submitted, stops the write with its + error. An append copies an existing object with ``UploadPartCopy``, or rewrites one smaller than a part from the buffer, as ``S3File`` does. - Closing the file commits the written data, as ``S3File.close()`` does: @@ -1276,6 +1278,8 @@ def __init__( # The parts uploaded from the buffer that may not have finished. The # part copies of an append hold no data and are not counted. self._uploads_in_flight: set[asyncio.Future[S3MultipartUploadPart]] = set() + # The part copies of an append that may not have finished. + self._copies_in_flight: set[Future[S3MultipartUploadPart]] = set() if target.appended is not None: if target.append_data is not None: cast(BytesIO, self.buffer).write(target.append_data) @@ -1339,7 +1343,8 @@ async def write(self, data: bytes) -> int: ValueError: If the file is not open for writing, is closed, or has been force-flushed. BaseException: The error of a part that the write found failed - while it waited for a part in flight. + while it waited for a part in flight, or of a part copy of an + append found failed before a part is submitted. """ if self.mode not in {"wb", "ab", "xb"}: raise ValueError("File not in write mode") @@ -1444,6 +1449,7 @@ async def _initiate_upload(self) -> None: # Tracked before it is scheduled; see the creation. copy: Future[S3MultipartUploadPart] = Future() self._parts.append(copy) + self._copies_in_flight.add(copy) self._pool.submit_to( copy, writer.upload_part_copy, @@ -1471,7 +1477,8 @@ async def _upload_chunk(self, final: bool = False) -> bool: Raises: BaseException: The error of an uploaded part that the write found - failed while it waited for a part in flight. + failed while it waited for a part in flight, or of a part copy + of an append found failed before a part is submitted. RuntimeError: If no multipart upload is initialized. ValueError: If the data needs more parts than a multipart upload allows. @@ -1491,6 +1498,13 @@ async def _upload_chunk(self, final: bool = False) -> bool: for part_number, data in writer.iter_parts( buffer, first_part_number=len(self._parts) + 1 ): + if self._copies_in_flight: + copied = {copy for copy in self._copies_in_flight if copy.done()} + self._copies_in_flight -= copied + for copy in copied: + # Otherwise the rest of the data is uploaded after a + # failed part copy, as S3File does. + copy.result() if len(self._uploads_in_flight) >= self.max_workers: # Otherwise the parts waiting for a thread pile up in memory. done, self._uploads_in_flight = await asyncio.wait( @@ -1538,8 +1552,9 @@ def _complete_upload(self) -> None: Creates an empty object if nothing was written, uploads the buffered data with PutObject if no part was uploaded, and otherwise completes the multipart upload once its parts have ended, aborting it if a - part or the completion fails. Invalidates the cache of the path - afterwards. + part or the completion fails. A failed part is handled as soon as it + fails, without waiting for the parts before it. Invalidates the cache + of the path afterwards. Raises: FileExistsError: If an object was created at the path after the @@ -1559,7 +1574,12 @@ def _complete_upload(self) -> None: else: try: # The parts were submitted before this call, so they have - # all started. + # all started. A failed part is raised without waiting for + # the parts before it, as S3File does. + done, _ = wait(self._parts, return_when=FIRST_EXCEPTION) + for part in done: + if not part.cancelled() and part.exception() is not None: + part.result() parts = [part.result() for part in self._parts] writer.complete(parts, **self._get_request_kwargs("complete_multipart_upload")) except BaseException: diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index aaea78a8d..70a3f298b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2863,6 +2863,63 @@ async def write_unclosed(): fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_append_failed_part_copy(self, fs): + # A part copy of an append found failed before a part is submitted + # stops the write with its error, as with S3File (GH-1178). + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + info = S3Object( + init={"Key": "key"}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = block_size + fs._sync_fs.info = mock.MagicMock(return_value=info) + fs.core.upload_part_copy = mock.MagicMock(side_effect=PermissionError("Access Denied")) + f = await fs.open_async("s3://bucket/key", "ab", block_size=block_size) + await f.write(b"a" * block_size) + # The copy has failed before the next write submits a part. + await asyncio.to_thread(wait, f._parts[:1], 5) + + with pytest.raises(PermissionError, match="Access Denied"): + await f.write(b"a" * block_size) + + assert f.closed + assert fs.core.upload_part.call_count == 1 + fs.core.abort_multipart_upload.assert_called_once() + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_commit_failed_part_first(self, fs): + # The commit raises a failed part without waiting for the parts + # before it, and cancels those that have not started, as S3File does. + f = await fs.open_async("s3://bucket/key", "wb") + await f.write(b"a") + f._writer.upload = S3MultipartUpload( + {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + ) + failed: Future[S3MultipartUploadPart] = Future() + failed.set_exception(PermissionError("Access Denied")) + pending: list[Future[S3MultipartUploadPart]] = [Future(), Future()] + f._parts = [pending[0], failed, pending[1]] + + try: + with pytest.raises(PermissionError, match="Access Denied"): + await asyncio.wait_for(asyncio.to_thread(f._complete_upload), 5) + finally: + # Releases a commit that waits for the parts in order. + for part in pending: + if not part.done(): + part.set_result(S3MultipartUploadPart(1, {"ETag": '"e"'})) + + assert all(part.cancelled() for part in pending) + fs.core.abort_multipart_upload.assert_called_once() + fs.core.complete_multipart_upload.assert_not_called() + # Closed without committing again. + f.closed = True + f._pool.shutdown() + @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio async def test_write_cancelled(self, fs, cancellations): From 5f32a53b97a0161309ffed2024b27b51430a86de Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 08:27:01 +0900 Subject: [PATCH 11/11] Hold the failed part copy in the append test and describe the wait - The append test holds the failing part copy until the first write has returned, so that the first write cannot find the failure. - The docs say that the commit cancels the parts that have not started as soon as one fails, and waits for the running ones before it aborts. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 5 +++-- pyathena/filesystem/s3_async.py | 7 ++++--- tests/pyathena/filesystem/test_s3_async.py | 18 +++++++++++++++--- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 7b60ec05a..9841119b7 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -318,8 +318,9 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: them, so the event loop keeps running. A part found failed by that wait, or a part copied from the existing object in an append that has failed before the write uploads another part, stops the write with its error, and the multipart upload is aborted. - A part that fails later raises its error when the file is closed, without waiting for - the parts before it. + A part that fails later raises its error when the file is closed. The commit cancels + the parts that have not started as soon as one fails, and waits for the running ones + before it aborts the upload and raises the error. - Closing the file commits the written data, as `S3File` does. - Leaving an `async with` block with an exception, or cancelling a write, closes the file without committing it, and aborts the multipart upload after the running parts finish. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index d56e966ad..1bdee1556 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1552,9 +1552,10 @@ def _complete_upload(self) -> None: Creates an empty object if nothing was written, uploads the buffered data with PutObject if no part was uploaded, and otherwise completes the multipart upload once its parts have ended, aborting it if a - part or the completion fails. A failed part is handled as soon as it - fails, without waiting for the parts before it. Invalidates the cache - of the path afterwards. + part or the completion fails. A failed part is found as soon as it + fails, without waiting for the parts before it, and the parts that + have not started are cancelled; the running ones are waited for + before the abort. Invalidates the cache of the path afterwards. Raises: FileExistsError: If an object was created at the path after the diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 70a3f298b..80de1537d 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -37,6 +37,7 @@ from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( + EVENT_TIMEOUT, MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, overlapping_calls, @@ -2876,11 +2877,22 @@ async def test_append_failed_part_copy(self, fs): ) info.size = block_size fs._sync_fs.info = mock.MagicMock(return_value=info) - fs.core.upload_part_copy = mock.MagicMock(side_effect=PermissionError("Access Denied")) + release = threading.Event() + + def upload_part_copy(**kwargs): + # Held until the first write returns, so that it does not see + # the failure. + release.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + + fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) f = await fs.open_async("s3://bucket/key", "ab", block_size=block_size) - await f.write(b"a" * block_size) + try: + await f.write(b"a" * block_size) + finally: + release.set() # The copy has failed before the next write submits a part. - await asyncio.to_thread(wait, f._parts[:1], 5) + await asyncio.to_thread(wait, f._parts[:1], EVENT_TIMEOUT) with pytest.raises(PermissionError, match="Access Denied"): await f.write(b"a" * block_size)