From 79129d7a3147df241b5da9cd481bdcf977f7375d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 4 Nov 2026 01:05:36 +0900 Subject: [PATCH 1/5] Merge the filesystem thread pools into one S3ThreadPool type (GH-1332) S3AioThreadPool becomes S3ThreadPool, a neutral name, and S3ThreadPoolExecutor and S3AioThreadPool remain as aliases of it. S3AioExecutor remains as a subclass for compatibility: it no longer dispatches the calls through an asyncio event loop, and ignores its loop argument with a DeprecationWarning. The files of AioS3FileSystem run their range reads and parts in an S3ThreadPool of their own, as those of S3FileSystem do, so they need no running loop and also work in a forked child process. The tests of the loop dispatch are removed; the contract of submit_to() for a callable that an executor drops or fails is tested with an executor that holds the calls. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 4 +- benchmarks/pyathena_bench/filesystem.py | 8 +- docs/aio.md | 34 +- docs/api/filesystem.rst | 9 +- pyathena/filesystem/s3.py | 8 +- pyathena/filesystem/s3_async.py | 60 +-- pyathena/filesystem/s3_executor.py | 214 ++------- pyathena/filesystem/s3_writer.py | 7 +- tests/pyathena/filesystem/test_s3.py | 58 +-- tests/pyathena/filesystem/test_s3_async.py | 73 +-- tests/pyathena/filesystem/test_s3_core.py | 4 +- tests/pyathena/filesystem/test_s3_executor.py | 422 +++--------------- tests/pyathena/filesystem/test_s3_writer.py | 18 +- tests/pyathena/util.py | 2 +- 14 files changed, 232 insertions(+), 689 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 2c0135431..d0c67b7c9 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -238,8 +238,8 @@ 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. A `get` also writes the object to the trial's temporary directory, which the parent process measures as the peak temp dir and removes; on the Amazon Linux 2023 hosts, that directory is on the tmpfs `/tmp`, so the object stays in memory outside the process's RSS. A `put` reads a sparse file, which no disk bounds. -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, get, write, or put 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 `get_file()` and `put_file()` run the synchronous filesystem's in the filesystem's thread pool. +The `sync` API calls the filesystem on the trial's thread, and a file's parallel requests run in an `S3ThreadPool`. +The `aio` API runs a read, get, write, or put in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run in a pool of `max_workers` threads of the file, as in the `sync` API; its `get_file()` and `put_file()` run the synchronous filesystem's in the filesystem's thread pool. 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. diff --git a/benchmarks/pyathena_bench/filesystem.py b/benchmarks/pyathena_bench/filesystem.py index 8492c55a6..d10d7a1bc 100644 --- a/benchmarks/pyathena_bench/filesystem.py +++ b/benchmarks/pyathena_bench/filesystem.py @@ -135,10 +135,10 @@ async def transfer( The synchronous filesystem and s3fs run on the calling thread; s3fs waits for each operation, a coroutine on fsspec's IO loop. As ``AioS3FSCursor`` reads result files, the asyncio filesystem's file and local-file transfers - run in ``asyncio.to_thread()``: its files dispatch their range requests - and part uploads through ``S3AioExecutor`` to the filesystem's event loop, - and its ``get_file()`` and ``put_file()`` run the synchronous filesystem's - in the filesystem's thread pool. Its copy is a coroutine awaited on the + run in ``asyncio.to_thread()``: its files run their range requests and + part uploads in a pool of ``max_workers`` threads of the file, and its + ``get_file()`` and ``put_file()`` run the synchronous filesystem's in the + filesystem's thread pool. Its copy is a coroutine awaited on the trial's event loop. Args: diff --git a/docs/aio.md b/docs/aio.md index ab83aea8a..375e244a7 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -282,14 +282,14 @@ its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized |---|---|---| | **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()` | +| **File handles** | `S3File` with thread pool | `AioS3File` with thread pool; `AioS3StreamedFile` from `open_async()` | | **Bulk delete** | Thread pool per batch | `asyncio.gather` per batch | | **Multipart copy** | Thread pool per copy | Thread pool of the filesystem | | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) | ### Threads -The threads of `AioS3FileSystem` belong to `S3AioThreadPool` instances, not to the +The threads of `AioS3FileSystem` belong to `S3ThreadPool` instances, not to the default executor of the event loop, so `loop.set_default_executor()` neither limits nor blocks them: @@ -332,22 +332,20 @@ multipart uploads. ### Executor strategy `S3FileSystem` and `S3File` use a pluggable executor abstraction (`S3Executor`) for -parallel operations. Two implementations are provided: - -- `S3ThreadPoolExecutor` — runs the calls in a pool of `max_workers` threads (default - for sync usage); in a forked child process, the pool creates new threads -- `S3AioExecutor` — dispatches work via `asyncio.run_coroutine_threadsafe` to an - `S3AioThreadPool` of its own - -`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`. - -In a forked child process, the event loop of the parent does not run. -An `AioS3FileSystem` created in the parent without `asynchronous=True` raises -`RuntimeError` from `open()` there, as fsspec's synchronous methods do, and so does -`S3AioExecutor.submit()`; a filesystem created in the child works. +parallel operations. `S3ThreadPool` runs the calls in a pool of `max_workers` threads, +in the order they were submitted; in a forked child process, the pool creates new +threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3ThreadPool`. +`S3AioExecutor` is kept for compatibility: it is an `S3ThreadPool` that ignores its +`loop` argument, with a `DeprecationWarning` if one is given. + +The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel +range reads in a pool of `max_workers` threads of each file, as `S3File` does, without +the event loop. + +In a forked child process, the event loop of the parent does not run. The coroutines of +an `AioS3FileSystem` created in the parent, and fsspec's synchronous methods that run +them, raise `RuntimeError` there, while its `open()` works. A filesystem created in the +child works. ### Async file handles diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index 1cff7949a..d87d36e16 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -49,13 +49,14 @@ S3 Executor .. autoclass:: pyathena.filesystem.s3_executor.S3Executor :members: -.. autoclass:: pyathena.filesystem.s3_executor.S3ThreadPoolExecutor +.. autoclass:: pyathena.filesystem.s3_executor.S3ThreadPool :members: -.. autoclass:: pyathena.filesystem.s3_executor.S3AioExecutor - :members: +.. autodata:: pyathena.filesystem.s3_executor.S3ThreadPoolExecutor -.. autoclass:: pyathena.filesystem.s3_executor.S3AioThreadPool +.. autodata:: pyathena.filesystem.s3_executor.S3AioThreadPool + +.. autoclass:: pyathena.filesystem.s3_executor.S3AioExecutor :members: S3 Paths diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 1d2456db0..9b5ce5e88 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -36,7 +36,7 @@ from pyathena.connection import Connection from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool from pyathena.filesystem.s3_object import ( S3Metadata, S3MultipartUpload, @@ -1110,7 +1110,7 @@ def _create_executor(self, max_workers: int) -> S3Executor: Returns: An S3Executor instance. """ - return S3ThreadPoolExecutor(max_workers=max_workers) + return S3ThreadPool(max_workers=max_workers) def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None: """Delete objects with DeleteObjects requests grouped by bucket. @@ -2882,7 +2882,7 @@ def __init__( 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. + ``S3ThreadPool`` 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. @@ -2928,7 +2928,7 @@ def __init__( 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) + self._executor: S3Executor = executor or S3ThreadPool(max_workers=max_workers) # The requests and the state of the write; None when reading. self._session: _S3WriteSession | None = None if mode != "rb": diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 57203a7d0..f899ebaf9 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -32,12 +32,7 @@ from pyathena.aio.util import create_in_thread from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import ( - S3AioExecutor, - S3AioThreadPool, - S3Executor, - S3ThreadPoolExecutor, -) +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -63,7 +58,7 @@ class AioS3FileSystem(AsyncFileSystem): """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem. This class wraps ``S3FileSystem`` to provide native asyncio support. Its - coroutines run the blocking boto3 calls in an ``S3AioThreadPool`` with + coroutines run the blocking boto3 calls in an ``S3ThreadPool`` with ``max_workers`` threads that the filesystem owns, instead of the default executor of the event loop, and run parallel requests with ``asyncio.gather``. @@ -72,11 +67,9 @@ class AioS3FileSystem(AsyncFileSystem): 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 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 + File handles created by ``_open`` run their parallel operations (range + reads, multipart uploads) in an ``S3ThreadPool`` of ``max_workers`` + threads of each file, as ``S3File`` does. :meth:`open_async` returns an ``AioS3StreamedFile``, whose reads and writes are coroutines. :meth:`close`, or the end of a ``with`` or ``async with`` block, releases the threads of the filesystem, which stays usable. @@ -146,7 +139,7 @@ def __init__( # 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) + self._thread_pool = S3ThreadPool(max_workers) max_workers = self._thread_pool.max_workers self._sync_fs = S3FileSystem( connection=connection, @@ -172,8 +165,8 @@ def core(self) -> S3Core: def close(self) -> None: """Release the threads of the filesystem, which stays usable, once they are idle. - See :meth:`S3AioThreadPool.release() - `. Files opened + See :meth:`S3ThreadPool.release() + `. Files opened before have threads of their own and are not affected. """ self._thread_pool.release() @@ -839,24 +832,16 @@ async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> def _create_executor(self, max_workers: int) -> S3Executor: """Create the executor for the parallel operations of a file. - An instance created with ``asynchronous=True`` has no event loop of - its own, so its files run the operations in a thread pool. + The operations run in threads, as those of ``S3File`` do, without the + event loop. Args: max_workers: The maximum number of operations that run at once. Returns: - An ``S3AioExecutor`` on the event loop of this filesystem, or an - ``S3ThreadPoolExecutor`` if it has none. - - Raises: - RuntimeError: If this instance has an event loop and the call is - made in a forked child process of the process that created - it, where the loop does not run. + The executor of the wrapped ``S3FileSystem``. """ - if self._loop is None: - return S3ThreadPoolExecutor(max_workers=max_workers) - return S3AioExecutor(loop=self.loop, max_workers=max_workers) + return self._sync_fs._create_executor(max_workers=max_workers) def _open( self, @@ -1194,17 +1179,12 @@ def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: class AioS3File(S3File): - """Async-aware S3 file handle using ``S3AioExecutor``. - - Functionally identical to ``S3File``; exists as a distinct type for - ``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 to - a pool of ``max_workers`` threads of the file. - For an ``AioS3FileSystem`` created with ``asynchronous=True``, it is an - ``S3ThreadPoolExecutor``. For reads and writes that are coroutines, see - ``AioS3StreamedFile``. + """The file handle that ``AioS3FileSystem.open()`` returns. + + Functionally identical to ``S3File``, whose parallel operations (range + reads, multipart uploads) it runs in a pool of ``max_workers`` threads of + the file; it exists as a distinct type for ``isinstance`` checks. For + reads and writes that are coroutines, see ``AioS3StreamedFile``. """ @@ -1212,7 +1192,7 @@ 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 + run in an ``S3ThreadPool`` 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 @@ -1298,7 +1278,7 @@ def __init__( if mode not in {"rb", "wb", "ab", "xb"}: raise NotImplementedError("File mode not supported") # Created first, so that an invalid max_workers raises before the lookups. - self._pool = S3AioThreadPool(max_workers) + self._pool = S3ThreadPool(max_workers) self.max_workers = self._pool.max_workers target = S3FileTarget.resolve( fs._sync_fs, diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 1e9b3f2e6..a3c126940 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -12,13 +12,12 @@ import asyncio import contextvars import os -import sys import threading import traceback +import warnings 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 @@ -84,8 +83,8 @@ def run(self) -> None: def settle(self, submitted: Future[Any]) -> None: """Resolve the future if the executor dropped the callable before it ran. - This happens, for example, when the event loop of an - :class:`S3AioExecutor` shuts down. + This happens, for example, when an executor cancels the calls that + have not started as it shuts down. Args: submitted: The finished future of the submitted :meth:`run`. @@ -150,9 +149,9 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg scheduling midway, before :meth:`submit` returns. Cancelling the future keeps a callable that has not started from running; a started callable sets the future's result or exception. If the executor drops - the callable before it starts, as :class:`S3AioExecutor` does when its - event loop shuts down, the future is cancelled or gets the error that - the executor reported. + the callable before it starts, such as by cancelling it at its + shutdown, the future is cancelled or gets the error that the + executor reported. The arguments of a failed callable, such as the body of a part, are freed when it fails: its error keeps the frames that it was raised @@ -177,47 +176,22 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.shutdown(wait=True) -class S3ThreadPoolExecutor(S3Executor): - """Executor that runs the submitted calls in a pool of threads. +class S3ThreadPool(S3Executor): + """A thread pool for the requests of the S3 filesystems. - This is the default executor used by ``S3File`` and ``S3FileSystem`` - for synchronous parallel operations. The calls run in an - :class:`S3AioThreadPool`, so that a forked child process, whose copy of - the executor has no threads, creates new ones. - """ - - def __init__(self, max_workers: int) -> None: - """Initialize the executor without starting any thread. - - Args: - max_workers: The maximum number of threads of the thread pool. - - Raises: - ValueError: If ``max_workers`` is not positive. - """ - self._pool = S3AioThreadPool(max_workers) - - @override - def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: - return self._pool.submit(fn, *args, **kwargs) - - @override - def shutdown(self, wait: bool = True) -> None: - self._pool.shutdown(wait=wait) - - -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 - 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 + It runs at most ``max_workers`` calls at once in threads of its own, so + that, unlike ``asyncio.to_thread()``, which runs a call in the default + executor of the event loop, neither the size nor the load of that + executor limits or blocks the calls of asyncio code. The calls start in + the order they were submitted. 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, as it does after - :meth:`release`. The pool can be shared by several event loops and - threads. + :meth:`release`. The pool can be shared by several threads and event + loops. + + ``S3ThreadPoolExecutor`` and ``S3ThreadPool`` are aliases of this + class. Args: max_workers: The maximum number of calls that run at once. @@ -227,7 +201,7 @@ class S3AioThreadPool(S3Executor): """ # The pools of this process, which a forked child resets. - _instances: ClassVar[weakref.WeakSet[S3AioThreadPool]] = weakref.WeakSet() + _instances: ClassVar[weakref.WeakSet[S3ThreadPool]] = 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. @@ -274,7 +248,7 @@ def end_call(future: Future[Any]) -> None: live._end_call(future) self._end_call_callback = end_call - S3AioThreadPool._instances.add(self) + S3ThreadPool._instances.add(self) @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: @@ -483,148 +457,50 @@ def _get_executor(self) -> ThreadPoolExecutor: if hasattr(os, "register_at_fork"): - os.register_at_fork(after_in_child=S3AioThreadPool._reset_after_fork) + os.register_at_fork(after_in_child=S3ThreadPool._reset_after_fork) -class S3AioExecutor(S3Executor): - """Executor that schedules work on an asyncio event loop. +class S3AioExecutor(S3ThreadPool): + """A thread pool kept for compatibility, an :class:`S3ThreadPool`. - 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, - whatever the size of the default executor of the event loop. + It used to dispatch the calls through an asyncio event loop. The calls + now run directly in the pool, so they need no running loop, and they run + to their end when the loop stops. Args: - loop: A running asyncio event loop. - max_workers: The maximum number of submitted functions that run at once. + loop: Ignored, with a ``DeprecationWarning`` if given. + max_workers: The maximum number of calls that run at once. Raises: - RuntimeError: If the event loop is not running when ``submit`` is - called, or ``submit`` is called in a forked child process, where - the event loop of the parent does not run. + ValueError: If ``max_workers`` is not positive. """ - # The executors of this process, which a forked child detaches from the - # event loop of its parent. - _instances: ClassVar[weakref.WeakSet[S3AioExecutor]] = weakref.WeakSet() - def __init__( self, loop: asyncio.AbstractEventLoop | None = None, max_workers: int = (cpu_count() or 1) * 5, ) -> None: - """Initialize the executor with the event loop to schedule work on. + """Initialize the pool without starting any thread. Args: - loop: The asyncio event loop. ``submit`` raises ``RuntimeError`` - if it is None or not running. - max_workers: The maximum number of submitted functions that run - at once. + loop: Ignored, with a ``DeprecationWarning`` if given. + max_workers: The maximum number of calls that run at once. Raises: ValueError: If ``max_workers`` is not positive. """ - self._loop = loop - self._pool = S3AioThreadPool(max_workers) - # Orders the dispatches with the shutdown. - self._lock = threading.Lock() - self._shutdown = False - self._forked = False - S3AioExecutor._instances.add(self) - - @override - def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: - if self._forked: - raise RuntimeError("S3AioExecutor cannot be used in a forked child process.") - 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 - # is started and resolved by the function's thread instead. - future: Future[T] = Future() - item = _S3WorkItem(future, fn, args, kwargs) - with self._lock: - if self._shutdown: - raise RuntimeError("cannot schedule new futures after shutdown") - # The task runs in a copy of the submitter's context, which - # call_soon_threadsafe() schedules it with, and so does the - # function in the thread of the pool. - task = asyncio.run_coroutine_threadsafe(self._dispatch(item.run), self._loop) - task.add_done_callback(item.settle) - return future - raise RuntimeError( - "S3AioExecutor requires a running event loop. " - "Use S3ThreadPoolExecutor for synchronous usage." - ) - - async def _dispatch(self, run: Callable[[], None]) -> None: - """Hand a submitted function to the pool and wait for it. - - Unlike :meth:`S3AioThreadPool.run`, a cancellation of this task, such - as by the shutdown of the event loop, does not stop the retries of a - function that has started, which the caller still waits for. - - Args: - run: The function of the work item. - """ - await asyncio.wrap_future(self._pool.submit(run)) - - async def _shutdown_pool(self) -> None: - """Shut down the pool once the functions submitted before have reached it. - - 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. - """ - self._pool.shutdown(wait=False) - - @override - def shutdown(self, wait: bool = True) -> None: - loop = self._loop - with self._lock: - self._shutdown = True - closing = None - 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(), loop) - if loop is None or closing is None: - self._pool.shutdown(wait=wait) - return - try: - 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: - # 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: - # 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) - - @classmethod - def _reset_after_fork(cls) -> None: - """Detach the executors that a forked child copied from the event loop of its parent. - - The copy of the loop reports running, but nothing runs it, and a lock - that a thread of the parent held stays held. Each executor then - raises from :meth:`submit`, and :meth:`shutdown` only shuts down its - pool. - """ - for executor in list(cls._instances): - executor._loop = None - executor._lock = threading.Lock() - executor._forked = True + if loop is not None: + warnings.warn( + "The loop argument of S3AioExecutor is ignored: the calls run in the pool " + "without the event loop. S3AioExecutor is kept for compatibility; use " + "S3ThreadPool instead.", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(max_workers) -if hasattr(os, "register_at_fork"): - os.register_at_fork(after_in_child=S3AioExecutor._reset_after_fork) +#: An alias of :class:`S3ThreadPool`, kept for compatibility. +S3ThreadPoolExecutor = S3ThreadPool +#: An alias of :class:`S3ThreadPool`, kept for compatibility. +S3AioThreadPool = S3ThreadPool diff --git a/pyathena/filesystem/s3_writer.py b/pyathena/filesystem/s3_writer.py index be5defa14..af6b8b3e5 100644 --- a/pyathena/filesystem/s3_writer.py +++ b/pyathena/filesystem/s3_writer.py @@ -969,7 +969,8 @@ def _consume_creation(self, creation: Future[S3MultipartUpload]) -> None: return if self._state is _S3WriteState.CREATING: if creation.cancelled(): - # Such as by the shutdown of the event loop of an S3AioExecutor. + # Such as by an executor that cancels its queued calls as it + # shuts down. self._fail(CancelledError("The creation of the multipart upload was cancelled.")) elif (error := creation.exception()) is not None: self._fail(error) @@ -992,8 +993,8 @@ def _consume(self, request: Future[S3MultipartUploadPart]) -> None: if request not in self._requests: return if request.cancelled(): - # Not by abort(), but such as by the shutdown of the event loop of - # an S3AioExecutor, the upload lacks the part. + # Not by abort(), but such as by an executor that cancels its + # queued calls as it shuts down, the upload lacks the part. error: BaseException | None = CancelledError( "A part request was cancelled before it ran." ) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 2bdef9302..50038d249 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -42,7 +42,7 @@ from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3ThreadPool from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession @@ -3706,7 +3706,7 @@ def upload_part_copy(**kwargs): fs.core.abort_multipart_upload = mock.MagicMock( side_effect=lambda *_, **__: events.append("abort") ) - executor = S3ThreadPoolExecutor(max_workers=2) + executor = S3ThreadPool(max_workers=2) submit = executor.submit submitted = [] @@ -4267,7 +4267,7 @@ def put_object_acl(path, acl, **kwargs): return throttled() fs.core.put_object_acl = put_object_acl - executor = S3ThreadPoolExecutor(max_workers=1) + executor = S3ThreadPool(max_workers=1) submit = executor.submit submitted = [] @@ -5987,7 +5987,7 @@ def create(*args, **kwargs): return upload fs.core.create_multipart_upload.side_effect = create - executor = S3ThreadPoolExecutor(max_workers=1) + executor = S3ThreadPool(max_workers=1) ready = _S3WriteSession.ready def interrupt_ready(session): @@ -6087,7 +6087,7 @@ def test_part_scheduling_interrupted(self, operation): fs = self._make_append_fs(b"a" * 6 if operation == "append" else b"") events = [] fs._call.side_effect = lambda *args, **kwargs: events.append("abort") - executor = S3ThreadPoolExecutor(max_workers=2) + executor = S3ThreadPool(max_workers=2) submit = executor.submit submitted = [] @@ -6196,7 +6196,7 @@ def create(*args, **kwargs): "s3://bucket/key.txt", mode="wb", block_size=4, - executor=S3AioExecutor(asyncio.get_running_loop(), max_workers=1), + executor=S3ThreadPool(max_workers=1), ) def write(): @@ -6278,7 +6278,7 @@ def _make_multipart_write_file(data: bytes, autocommit: bool, executor=None): data, autocommit=autocommit, block_size=4, - executor=executor or S3ThreadPoolExecutor(max_workers=1), + executor=executor or S3ThreadPool(max_workers=1), ) file.fs.core.create_multipart_upload.return_value = S3MultipartUpload( {"Bucket": "bucket", "Key": "key.txt", "UploadId": "uploadid"} @@ -6499,7 +6499,7 @@ def test_write_max_parts(self, existing, mode, writes): def test_write_exceeding_max_parts(self, existing, mode, writes, closing, autocommit): fs = self._make_append_fs(existing) fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 - executor = S3ThreadPoolExecutor(max_workers=1) + executor = S3ThreadPool(max_workers=1) f = S3File( fs, "s3://bucket/key.txt", @@ -6564,7 +6564,7 @@ def upload_part(**kwargs): fs.core.upload_part.side_effect = upload_part # More threads than parts, so that only the write limits the parts # in flight. - executor = S3ThreadPoolExecutor(max_workers=8) + executor = S3ThreadPool(max_workers=8) submit_to = executor.submit_to uploads: dict[int, Future] = {} in_flight = [] @@ -6625,7 +6625,7 @@ def upload_part(**kwargs): block_size=4, autocommit=autocommit, max_workers=max_workers, - executor=S3ThreadPoolExecutor(max_workers=max_workers), + executor=S3ThreadPool(max_workers=max_workers), ) try: for _ in range(max_workers): @@ -6715,7 +6715,7 @@ def upload_part_copy(**kwargs): raise PermissionError("Access Denied") fs.core.upload_part_copy.side_effect = upload_part_copy - executor = S3ThreadPoolExecutor(max_workers=max_workers) + executor = S3ThreadPool(max_workers=max_workers) f = S3File( fs, "s3://bucket/key.txt", @@ -6972,7 +6972,7 @@ def test_commit_stops_at_failed_part(self): # copies of a large append, are cancelled. The parts used to be # collected in part-number order. file = self._make_multipart_write_file( - b"x" * 16, autocommit=False, executor=S3ThreadPoolExecutor(max_workers=2) + b"x" * 16, autocommit=False, executor=S3ThreadPool(max_workers=2) ) events = [] file.fs._call.side_effect = lambda *args, **kwargs: events.append("abort") @@ -7018,7 +7018,7 @@ def test_commit_failure_and_interrupted_abort(self): # interrupted, the interrupt propagates, the executor is shut down, # and the upload is kept so that discard() retries the abort; a later # commit() raises instead of completing it. - executor = S3ThreadPoolExecutor(max_workers=1) + executor = S3ThreadPool(max_workers=1) file = self._make_multipart_write_file(b"x" * 16, autocommit=False, executor=executor) file._upload_chunk(final=True) file.fs.core.complete_multipart_upload.side_effect = RuntimeError("complete failed") @@ -7079,7 +7079,7 @@ def test_interrupted_creation_stops_retries(self): fs.core.create_multipart_upload = retried( operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) ) - executor = S3ThreadPoolExecutor(max_workers=1) + executor = S3ThreadPool(max_workers=1) ready = _S3WriteSession.ready def interrupt_ready(session): @@ -7106,36 +7106,6 @@ def interrupted(timeout=None): assert operation.calls == 1 assert file.closed - def test_discard_on_event_loop_thread(self): - # GH-976: the parts that have not started are cancelled and not - # waited for, so a rollback on the thread of the event loop that - # would run them does not block. - files = [] - - async def rollback(): - executor = S3AioExecutor(loop=asyncio.get_running_loop()) - file = self._make_write_file(b"", autocommit=False, block_size=4, executor=executor) - files.append(file) - file.fs.core.create_multipart_upload.return_value = S3MultipartUpload( - {"Bucket": "bucket", "Key": "key.txt", "UploadId": "uploadid"} - ) - file._session.start() - await asyncio.wrap_future(file._session.ready()) - # Dispatched by the event loop, which this rollback blocks. - for part_number in (1, 2): - file._session.submit_part(part_number, b"x" * 4) - file.discard() - - thread = threading.Thread(target=asyncio.run, args=(rollback(),), daemon=True) - thread.start() - thread.join(5) - - assert not thread.is_alive() - (file,) = files - file.fs.core.upload_part.assert_not_called() - file.fs._call.assert_called_once() - assert file.fs._call.call_args.args[0] == S3_CLIENT.abort_multipart_upload - @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): # A file larger than the block size is uploaded in parts and completed diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index ff98fd09b..15d8c6b75 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -23,7 +23,7 @@ from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan -from pyathena.filesystem.s3_executor import S3AioThreadPool +from pyathena.filesystem.s3_executor import S3ThreadPool from pyathena.filesystem.s3_object import ( S3Metadata, S3MultipartUpload, @@ -130,21 +130,34 @@ async def _multipart_copy(release, fs=None, **kwargs): stubber.assert_no_pending_responses() @forking - def test_open_in_forked_child(self): - # GH-1336: in a forked child, a filesystem created in the parent - # raises on open(), as fsspec's synchronous methods do there. It used - # to return a file whose reads waited forever for the event loop of - # the parent. - fs = AioS3FileSystem( - key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + def test_read_in_forked_child(self): + # GH-1332: in a forked child, a file that a filesystem created in the + # parent opens reads its ranges in parallel. Its open() raised there + # since GH-1336, and before that, its reads waited forever for the + # event loop of the parent. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + block_size = 2**10 + size = block_size * 4 + fs.core.get_object = mock.MagicMock( + side_effect=lambda path, range_, **kwargs: bytes(range_[1] - range_[0]) + ) + 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) - def child(): - with pytest.raises(RuntimeError, match="not fork-safe"): - fs.open("s3://bucket/key", "rb") - return True + def read(): + with fs.open("s3://bucket/key", "rb", block_size=block_size, cache_type="none") as f: + return f.read() == bytes(size) - assert run_in_forked_child(child) == 0 + # Read in the parent too, whose file threads the child does not have. + assert read() + assert run_in_forked_child(read) == 0 + assert fs.core.get_object.call_count > 1 @pytest.mark.parametrize("size", [0, 10]) @pytest.mark.asyncio @@ -1080,7 +1093,7 @@ def test_copy_object_with_multipart_upload_interrupted_while_scheduling_completi # The HeadObject of the source, with 2 parts of the default block size. size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) - submit_to = S3AioThreadPool.submit_to + submit_to = S3ThreadPool.submit_to async def copy(): async with asyncio.timeout(EVENT_TIMEOUT): @@ -1107,7 +1120,7 @@ def interrupt_completion(pool, future, fn, *args, **kwargs): with ( mock.patch.object( - S3AioThreadPool, "submit_to", autospec=True, side_effect=interrupt_completion + S3ThreadPool, "submit_to", autospec=True, side_effect=interrupt_completion ), pytest.raises(KeyboardInterrupt), ): @@ -1142,7 +1155,7 @@ async def test_copy_object_with_multipart_upload_failed_part_with_threads_taken( operation = ThrottledOperation(operation_name="UploadPartCopy") throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) submitted = threading.Event() - submit_to = S3AioThreadPool.submit_to + submit_to = S3ThreadPool.submit_to def record_part(pool, future, fn, *args, **kwargs): submit_to(pool, future, fn, *args, **kwargs) @@ -1162,7 +1175,7 @@ def upload_part_copy(part_number, **kwargs): sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) with ( - mock.patch.object(S3AioThreadPool, "submit_to", autospec=True, side_effect=record_part), + mock.patch.object(S3ThreadPool, "submit_to", autospec=True, side_effect=record_part), pytest.raises(PermissionError, match="Access Denied"), ): # The retries stop, or part 2 is cancelled, before the bound. @@ -2124,8 +2137,8 @@ def get_object(path, range_, **kwargs): ) info.size = block_size * 4 fs._sync_fs.info = mock.MagicMock(return_value=info) - # The executor of the file dispatches through the event loop, which - # the blocking calls must not hold. + # The blocking calls of the file run in threads, as AioS3FSCursor + # reads result files, not on the event loop. f = await asyncio.to_thread( fs.open, "s3://bucket/key", "rb", block_size=block_size, cache_type="none" ) @@ -2269,7 +2282,7 @@ async def test_write_parts_in_flight(self, fs): active: list[int] = [] uploads: dict[int, Future[S3MultipartUploadPart]] = {} in_flight: list[int] = [] - submit_to = S3AioThreadPool.submit_to + submit_to = S3ThreadPool.submit_to def record_upload(pool, future, fn, *args, **kwargs): # The parts submitted and not finished, including those that @@ -2310,9 +2323,7 @@ async def release(): releaser = asyncio.ensure_future(release()) f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) - with mock.patch.object( - S3AioThreadPool, "submit_to", autospec=True, side_effect=record_upload - ): + with mock.patch.object(S3ThreadPool, "submit_to", autospec=True, side_effect=record_upload): await f.write(b"a" * block_size * 4) written.set() await f.close() @@ -2532,7 +2543,7 @@ async def test_commit_failed_part_first(self, fs): } queued = threading.Event() - submit = S3AioThreadPool.submit + submit = S3ThreadPool.submit def record_commit(pool, fn, *args, **kwargs): if getattr(fn, "__name__", None) == "_complete_upload": @@ -2560,7 +2571,7 @@ def upload_part(part_number, **kwargs): await f.write(b"a" * block_size * 3) started = time.monotonic() with ( - mock.patch.object(S3AioThreadPool, "submit", autospec=True, side_effect=record_commit), + mock.patch.object(S3ThreadPool, "submit", autospec=True, side_effect=record_commit), pytest.raises(PermissionError, match="Access Denied"), ): await f.close() @@ -2723,7 +2734,7 @@ def upload_part(**kwargs): fs.core.upload_part = mock.MagicMock(side_effect=upload_part) queued = threading.Event() - submit = S3AioThreadPool.submit + submit = S3ThreadPool.submit def record_commit(pool, fn, *args, **kwargs): if getattr(fn, "__name__", None) == "_complete_upload": @@ -2744,9 +2755,7 @@ async def flush_and_exit(): try: with ( - mock.patch.object( - S3AioThreadPool, "submit", autospec=True, side_effect=record_commit - ), + mock.patch.object(S3ThreadPool, "submit", autospec=True, side_effect=record_commit), pytest.raises(asyncio.CancelledError), ): # Bounded, so that a regression whose exit waits for the @@ -2768,7 +2777,7 @@ def test_close_interrupted_while_queuing_commit(self, fs): # committing it and aborts the upload. The upload used to be left # incomplete with a commit that nothing sends. block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - submit = S3AioThreadPool.submit + submit = S3ThreadPool.submit def interrupt_commit(pool, fn, *args, **kwargs): if getattr(fn, "__name__", None) == "_complete_upload": @@ -2784,7 +2793,7 @@ async def write_and_close_file(): await f.write(b"a" * block_size) try: with mock.patch.object( - S3AioThreadPool, "submit", autospec=True, side_effect=interrupt_commit + S3ThreadPool, "submit", autospec=True, side_effect=interrupt_commit ): await f.close() finally: @@ -2999,7 +3008,7 @@ def running_reads(): return [ t for t in asyncio.all_tasks() - if t.get_coro().__qualname__ == "S3AioThreadPool.run" and not t.done() + if t.get_coro().__qualname__ == "S3ThreadPool.run" and not t.done() ] # At least the held read is found. diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index 3c4e2a3f0..52d8041c3 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -39,7 +39,7 @@ S3MultipartCopyPlan, S3ObjectSummary, ) -from pyathena.filesystem.s3_executor import S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3ThreadPool from pyathena.filesystem.s3_object import S3Metadata, S3MultipartUpload, S3MultipartUploadPart from pyathena.filesystem.s3_path import S3Path from pyathena.util import RetryConfig, stop_retries_on @@ -201,7 +201,7 @@ def send(**params): future: Future[dict[str, Any]] = Future() gc.disable() try: - with S3ThreadPoolExecutor(max_workers=1) as executor: + with S3ThreadPool(max_workers=1) as executor: executor.submit_to(future, core.call, send, Body=body) del body assert isinstance(future.exception(), PermissionError) diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index d6a7bdc51..b86db428e 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -13,13 +13,13 @@ import traceback import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait -from typing import cast +from typing import Any, cast from unittest import mock import pytest from botocore.exceptions import ClientError -from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPool from pyathena.util import RetryConfig, retry_api_call, stop_retries_on from tests.pyathena.util import ( EVENT_TIMEOUT, @@ -27,21 +27,20 @@ forking, overlapping_calls, run_in_forked_child, - shutdown_bounded, ) -class TestS3AioThreadPool: +class TestS3ThreadPool: def test_default_max_workers(self): # None takes the default of ThreadPoolExecutor, as documented. with ThreadPoolExecutor() as executor: - assert S3AioThreadPool().max_workers == executor._max_workers + assert S3ThreadPool().max_workers == executor._max_workers 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) + pool = S3ThreadPool(4) async def main(): asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) @@ -56,7 +55,7 @@ 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) + pool = S3ThreadPool(1) async def main(): variable.set("caller") @@ -72,7 +71,7 @@ def test_run_cancelled_while_retrying(self, shielded): # would take at least EVENT_TIMEOUT seconds. operation = ThrottledOperation(block=True, failures=None if not shielded else 2) finished = threading.Event() - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) def work(): try: @@ -108,7 +107,7 @@ def test_run_cancelled_after_return(self): operation = ThrottledOperation(block=True, failures=1) left = [] go = threading.Event() - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) submit = pool.submit def submit_and_cancel_after_return(fn, *args, **kwargs): @@ -160,7 +159,7 @@ def test_run_cancelled_before_start(self, shielded): events = [] started = threading.Event() release = threading.Event() - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) def work(): started.set() @@ -198,7 +197,7 @@ def test_shutdown(self): # The calls submitted before the shutdown still run. events = [] release = threading.Event() - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) pool.submit(release.wait, EVENT_TIMEOUT) pending = pool.submit(events.append, "pending finished") pool.shutdown(wait=False) @@ -214,7 +213,7 @@ def test_shutdown(self): def test_release_after_calls(self): # GH-1209: release() keeps the threads, and the max_workers bound, # until no call runs or waits, and then ends them. - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) threads = [] started, release = threading.Event(), threading.Event() @@ -239,7 +238,7 @@ def block(): def test_kept_future_keeps_no_pool(self): # GH-1209: a future that the caller keeps after its call has ended does # not keep the pool, whose threads end once the pool is collected. - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) threads = [] release = threading.Event() @@ -265,7 +264,7 @@ def call(): def test_release_after_interrupted_submit(self): # GH-1209: a call whose submit() an interrupt stopped after scheduling # it does not keep release() from ending the threads once it has ended. - pool = S3AioThreadPool(1) + pool = S3ThreadPool(1) threads = [] with ( mock.patch.object(Future, "add_done_callback", side_effect=KeyboardInterrupt), @@ -284,7 +283,7 @@ def test_release_after_interrupted_submit(self): def test_fork(self): # 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) + pool = S3ThreadPool(1) assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 # No public call holds the lock across a fork; a thread of the parent # holds it while the main thread forks. @@ -312,293 +311,37 @@ def hold_lock(): class TestS3AioExecutor: - def test_submit_context(self): - # GH-1203: a function runs with the context of its submitter, as the - # task that dispatches it does, so that the stop scope of the - # submitter reaches its retries. - variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") - stop = threading.Event() - stop.set() - operation = ThrottledOperation() - - def work(): - with pytest.raises(ClientError): - retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) - return variable.get() - - async def main(): - executor = S3AioExecutor(loop=asyncio.get_running_loop()) - - def submit(): - variable.set("submitter") - with stop_retries_on(stop): - return executor.submit(work).result(EVENT_TIMEOUT) - - variable.set("loop") - result = await asyncio.to_thread(submit) - executor.shutdown() - return result - - assert asyncio.run(main()) == "submitter" - assert operation.calls == 1 - - 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_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_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() - - 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: - 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("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() - 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_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() - # 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() - 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()) - ) - with mock.patch("pyathena.filesystem.s3_executor.sys.is_finalizing", return_value=True): - waiter.start() - assert shut_down.wait(EVENT_TIMEOUT) - waiter.join() - finally: - release.set() - loop.call_soon_threadsafe(loop.stop) - thread.join() - loop.run_until_complete(asyncio.sleep(0)) - loop.close() - - @forking - def test_fork(self): - # GH-1336: in a forked child, where the event loop of the parent does - # not run, submit() raises and shutdown() returns, also if a thread - # of the parent held the lock of the executor. Both used to wait - # forever for the copy of the loop, which reports running. + def test_loop_ignored(self): + # GH-1332: the loop argument is ignored with a warning, and the calls + # run in the pool without a running loop. They used to be dispatched + # through the loop, which had to run. loop = asyncio.new_event_loop() - thread = threading.Thread(target=loop.run_forever, daemon=True) - thread.start() - executor = S3AioExecutor(loop=loop) - held, release = threading.Event(), threading.Event() - - def hold_lock(): - with executor._lock: - held.set() - release.wait() - - def child(): - with pytest.raises(RuntimeError, match="forked child"): - executor.submit(sum, [3, 4]) - executor.shutdown(wait=True) - return True - - holder = threading.Thread(target=hold_lock, daemon=True) - # The executor submits only to a loop that runs. - running = threading.Event() - loop.call_soon_threadsafe(running.set) try: - assert running.wait(EVENT_TIMEOUT) - assert executor.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 - holder.start() - assert held.wait(EVENT_TIMEOUT) - code = run_in_forked_child(child) + with pytest.warns(DeprecationWarning, match="loop argument of S3AioExecutor"): + executor = S3AioExecutor(loop=loop, max_workers=1) + with executor: + assert executor.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 finally: - release.set() - if holder.ident is not None: - holder.join(EVENT_TIMEOUT) - shutdown_bounded(executor) - loop.call_soon_threadsafe(loop.stop) - thread.join(EVENT_TIMEOUT) - assert not thread.is_alive() - # Runs the dispatch tasks that the stop left pending. - loop.run_until_complete(asyncio.sleep(0)) loop.close() - assert code == 0 - - def test_submit_without_running_loop(self): - with pytest.raises(RuntimeError, match="requires a running event loop"): - S3AioExecutor().submit(sum, [1, 2]) - - def test_cancel(self): - # GH-976: as with ThreadPoolExecutor, a future can be cancelled only - # before its function starts, so that wait() waits for a running one. - events = [] - started = threading.Event() - cancelled = threading.Event() - - def work(): - started.set() - cancelled.wait(5) - time.sleep(0.05) - events.append("finished") - - def cancel(executor: S3AioExecutor) -> tuple[Future[None], Future[None]]: - running = executor.submit(work) - pending = executor.submit(events.append, "pending finished") - started.wait(5) - assert not running.cancel() - assert pending.cancel() - cancelled.set() - _, not_done = wait([running, pending], timeout=5) - assert not not_done - events.append("waited") - return running, pending - async def main(): - # 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()) +class _DeferringExecutor(S3Executor): + """Holds the submitted calls, which a test drops or fails before they run.""" - assert events == ["finished", "waited"] - assert running.done() - assert not running.cancelled() - assert pending.cancelled() + def __init__(self) -> None: + self.submitted: list[Future[Any]] = [] + def submit(self, fn, *args, **kwargs): + future: Future[Any] = Future() + self.submitted.append(future) + return future -class TestS3ThreadPoolExecutor: - @forking - def test_fork(self): - # GH-1336: a forked child, whose copy of the executor has no threads, - # runs the calls submitted there. After the executor had run a call - # in the parent, they used to wait forever. - executor = S3ThreadPoolExecutor(max_workers=1) - try: - assert executor.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 - code = run_in_forked_child( - lambda: executor.submit(sum, [3, 4]).result(timeout=EVENT_TIMEOUT) == 7 - ) - finally: - executor.shutdown() - assert code == 0 + def shutdown(self, wait=True): + pass class TestS3Executor: - @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) - def test_submit_context(self, executor_class): + def test_submit_context(self): # GH-1203: a callable runs with a copy of the context of its submitter, # so that the stop scope of the submitter reaches its retries. variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") @@ -612,7 +355,7 @@ def work(): return variable.get() variable.set("submitter") - with executor_class(max_workers=1) as executor, stop_retries_on(stop): + with S3ThreadPool(max_workers=1) as executor, stop_retries_on(stop): future: Future[str] = Future() executor.submit_to(future, work) assert future.result(EVENT_TIMEOUT) == "submitter" @@ -620,7 +363,7 @@ def work(): def test_submit_to(self): # The callable's outcome resolves the future of the caller. - with S3ThreadPoolExecutor(max_workers=1) as executor: + with S3ThreadPool(max_workers=1) as executor: succeeded: Future[int] = Future() failed: Future[int] = Future() executor.submit_to(succeeded, sum, [1, 2]) @@ -631,8 +374,7 @@ def test_submit_to(self): with pytest.raises(ValueError, match="invalid literal"): failed.result() - @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) - def test_submit_to_failure_frees_arguments(self, executor_class): + def test_submit_to_failure_frees_arguments(self): # GH-1186: the arguments of a failed callable, such as the body of a # part, are freed when it fails, while the error keeps the frames of # its traceback. The traceback used to keep them, in a reference @@ -648,7 +390,7 @@ def upload(body): future: Future[None] = Future() gc.disable() try: - with executor_class(max_workers=1) as executor: + with S3ThreadPool(max_workers=1) as executor: executor.submit_to(future, upload, body) del body error = future.exception() @@ -668,7 +410,7 @@ def work(): started.set() release.wait(5) - with S3ThreadPoolExecutor(max_workers=1) as executor: + with S3ThreadPool(max_workers=1) as executor: executor.submit(work) pending: Future[None] = Future() executor.submit_to(pending, events.append, "pending finished") @@ -679,80 +421,46 @@ def work(): assert events == [] assert pending.cancelled() - def test_submit_to_loop_shutdown(self): - # A callable that S3AioExecutor drops when its event loop shuts down - # is never run, and its future, whether submit() returned it or the - # caller passed it to submit_to(), is cancelled and settled, so that + def test_submit_to_dropped(self): + # A callable that the executor drops before it starts is never run, + # and the future of the caller is cancelled and settled, so that # wait() returns, instead of left pending. events = [] - started = threading.Event() - settled = threading.Semaphore(0) - - def work(): - started.set() - # Running until both pending futures are settled, so that their - # callables cannot start before the shutdown. - for _ in range(2): - assert settled.acquire(timeout=EVENT_TIMEOUT) - events.append("finished") - + executor = _DeferringExecutor() pending: Future[None] = Future() - pending.add_done_callback(lambda _: settled.release()) + executor.submit_to(pending, events.append, "pending finished") + (submitted,) = executor.submitted + assert submitted.cancel() - async def main(): - # The pending callables wait for the only thread of the executor. - executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) - running = executor.submit(work) - submitted = executor.submit(events.append, "submitted finished") - submitted.add_done_callback(lambda _: settled.release()) - executor.submit_to(pending, events.append, "pending finished") - assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - return running, submitted - - running, submitted = asyncio.run(main()) - - _, not_done = wait([running, submitted, pending], timeout=EVENT_TIMEOUT) + _, not_done = wait([pending], timeout=EVENT_TIMEOUT) assert not not_done - assert events == ["finished"] - assert not running.cancelled() - assert submitted.cancelled() assert pending.cancelled() + assert events == [] - def test_submit_to_executor_shut_down(self): - # An error that S3AioExecutor reports before the callable starts is + def test_submit_to_executor_failure(self): + # An error that the executor reports before the callable starts is # set on the future of the caller. + executor = _DeferringExecutor() future: Future[int] = Future() + executor.submit_to(future, sum, [1, 2]) + (submitted,) = executor.submitted + submitted.set_running_or_notify_cancel() + submitted.set_exception(RuntimeError("cannot schedule new futures after shutdown")) - async def main(): - 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.wait_for(asyncio.wrap_future(future), EVENT_TIMEOUT) - - asyncio.run(main()) + with pytest.raises(RuntimeError, match="after shutdown"): + future.result(timeout=EVENT_TIMEOUT) - assert not future.cancelled() - - def test_submit_to_cancelled_executor_shut_down(self): - # A future that the caller cancels before S3AioExecutor reports the + def test_submit_to_cancelled_executor_failure(self): + # A future that the caller cancels before the executor reports the # error is marked as notified, so that wait() counts it as done. + executor = _DeferringExecutor() future: Future[int] = Future() + executor.submit_to(future, sum, [1, 2]) + assert future.cancel() + (submitted,) = executor.submitted + submitted.set_running_or_notify_cancel() + submitted.set_exception(RuntimeError("cannot schedule new futures after shutdown")) - async def main(): - loop = asyncio.get_running_loop() - 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=EVENT_TIMEOUT) - ) - - _, not_done = asyncio.run(main()) - + _, not_done = wait([future], timeout=EVENT_TIMEOUT) assert not not_done assert future.cancelled() diff --git a/tests/pyathena/filesystem/test_s3_writer.py b/tests/pyathena/filesystem/test_s3_writer.py index 94104036b..5a38092a9 100644 --- a/tests/pyathena/filesystem/test_s3_writer.py +++ b/tests/pyathena/filesystem/test_s3_writer.py @@ -18,7 +18,7 @@ from botocore.stub import Stubber from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3ThreadPool from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3MultipartUpload, @@ -231,8 +231,8 @@ def test_writer_requires_upload(self): writer.abort() -class _DroppingExecutor(S3ThreadPoolExecutor): - """Drops requests unrun, as S3AioExecutor does when its event loop shuts down.""" +class _DroppingExecutor(S3ThreadPool): + """Drops requests unrun, as an executor that cancels its queued calls at shutdown does.""" def __init__(self, drop) -> None: super().__init__(max_workers=1) @@ -316,7 +316,7 @@ class AccessDeniedError(PermissionError): error_ref = weakref.ref(error) gc.disable() try: - with S3ThreadPoolExecutor(max_workers=1) as executor: + with S3ThreadPool(max_workers=1) as executor: session, requests = _make_session(executor, part_errors={1: error}) del error _start(session) @@ -348,9 +348,9 @@ class AccessDeniedError(PermissionError): @pytest.mark.parametrize("dropped", ["creation", "part"]) def test_dropped_request_fails_the_write(self, dropped): - # A request that the executor drops without running it, as - # S3AioExecutor does when its event loop shuts down, fails the write, - # so that the upload is never completed without it. + # A request that the executor drops without running it, as an + # executor that cancels its queued calls at shutdown does, fails the + # write, so that the upload is never completed without it. executor = _DroppingExecutor( lambda kwargs: dropped == "creation" or "part_number" in kwargs ) @@ -382,7 +382,7 @@ def defer(future): deferred.append(future) called.release() - with S3ThreadPoolExecutor(max_workers=2) as executor: + with S3ThreadPool(max_workers=2) as executor: session, requests = _make_session( executor, max_workers=2 if operation == "abort" else 1, @@ -444,7 +444,7 @@ def test_cancelled_waiter(self, dropped_by): # whether the notification or a later ready() drops it, so that # concurrent.futures.wait() returns for it. gate = threading.Event() - with S3ThreadPoolExecutor(max_workers=1) as executor: + with S3ThreadPool(max_workers=1) as executor: session, requests = _make_session(executor, gate=gate) _start(session) session.submit_part(1, b"abcd") diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 277b020fe..d967fe793 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -346,7 +346,7 @@ def shutdown_bounded(pool, timeout=EVENT_TIMEOUT): from exiting through this wait. Args: - pool: The executor, such as an ``S3AioThreadPool``. + pool: The executor, such as an ``S3ThreadPool``. timeout: The longest wait, in seconds. Raises: From 5d56657e7c3be033b241fdc75e48a7b427a0bf2b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 4 Nov 2026 01:11:03 +0900 Subject: [PATCH 2/5] Correct the fork docs and drop the AioS3FileSystem executor override (GH-1332) In a forked child, only fsspec's synchronous methods of an AioS3FileSystem created in the parent raise; its coroutines and open() work. The files of AioS3FileSystem.open() take the executor that the wrapped S3FileSystem creates, without an override that only forwarded to it. The fork read test counts the requests of the child's read, and the docstring of S3ThreadPool names both aliases. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 8 ++++---- pyathena/filesystem/s3_async.py | 18 ++---------------- pyathena/filesystem/s3_executor.py | 2 +- pyathena/filesystem/s3_writer.py | 9 +++++---- tests/pyathena/filesystem/test_s3_async.py | 6 ++++-- tests/pyathena/filesystem/test_s3_executor.py | 4 +++- 6 files changed, 19 insertions(+), 28 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 375e244a7..ce5b2a496 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -342,10 +342,10 @@ The file handles of `AioS3FileSystem.open()` run their multipart uploads and par range reads in a pool of `max_workers` threads of each file, as `S3File` does, without the event loop. -In a forked child process, the event loop of the parent does not run. The coroutines of -an `AioS3FileSystem` created in the parent, and fsspec's synchronous methods that run -them, raise `RuntimeError` there, while its `open()` works. A filesystem created in the -child works. +In a forked child process, the event loop of the parent does not run, so fsspec's +synchronous methods of an `AioS3FileSystem` created in the parent, such as `ls()`, raise +`RuntimeError` there. Its `open()`, and its coroutines run on an event loop of the child, +work, and so does a filesystem created in the child. ### Async file handles diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index f899ebaf9..029054cab 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -32,7 +32,7 @@ from pyathena.aio.util import create_in_thread from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool +from pyathena.filesystem.s3_executor import S3ThreadPool from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -829,20 +829,6 @@ async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> raise FileNotFoundError(path) return sorted(out) - def _create_executor(self, max_workers: int) -> S3Executor: - """Create the executor for the parallel operations of a file. - - The operations run in threads, as those of ``S3File`` do, without the - event loop. - - Args: - max_workers: The maximum number of operations that run at once. - - Returns: - The executor of the wrapped ``S3FileSystem``. - """ - return self._sync_fs._create_executor(max_workers=max_workers) - def _open( self, path: str, @@ -873,7 +859,7 @@ def _open( path, mode, max_workers=max_workers, - executor=self._create_executor(max_workers=max_workers), + executor=self._sync_fs._create_executor(max_workers=max_workers), block_size=block_size, cache_type=cache_type, autocommit=autocommit, diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index a3c126940..ae77d0f74 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -190,7 +190,7 @@ class S3ThreadPool(S3Executor): :meth:`release`. The pool can be shared by several threads and event loops. - ``S3ThreadPoolExecutor`` and ``S3ThreadPool`` are aliases of this + ``S3ThreadPoolExecutor`` and ``S3AioThreadPool`` are aliases of this class. Args: diff --git a/pyathena/filesystem/s3_writer.py b/pyathena/filesystem/s3_writer.py index af6b8b3e5..8abc70fb3 100644 --- a/pyathena/filesystem/s3_writer.py +++ b/pyathena/filesystem/s3_writer.py @@ -969,8 +969,8 @@ def _consume_creation(self, creation: Future[S3MultipartUpload]) -> None: return if self._state is _S3WriteState.CREATING: if creation.cancelled(): - # Such as by an executor that cancels its queued calls as it - # shuts down. + # Such as by an executor passed to the file that cancels its + # queued calls as it shuts down. self._fail(CancelledError("The creation of the multipart upload was cancelled.")) elif (error := creation.exception()) is not None: self._fail(error) @@ -993,8 +993,9 @@ def _consume(self, request: Future[S3MultipartUploadPart]) -> None: if request not in self._requests: return if request.cancelled(): - # Not by abort(), but such as by an executor that cancels its - # queued calls as it shuts down, the upload lacks the part. + # Not by abort(), but such as by an executor passed to the file + # that cancels its queued calls as it shuts down, the upload + # lacks the part. error: BaseException | None = CancelledError( "A part request was cancelled before it ran." ) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 15d8c6b75..47dd4a5b4 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -151,13 +151,15 @@ def test_read_in_forked_child(self): fs._sync_fs.info = mock.MagicMock(return_value=info) def read(): + calls = fs.core.get_object.call_count with fs.open("s3://bucket/key", "rb", block_size=block_size, cache_type="none") as f: - return f.read() == bytes(size) + data = f.read() + # The ranges of the read, each a request of its own. + return data == bytes(size) and fs.core.get_object.call_count - calls > 1 # Read in the parent too, whose file threads the child does not have. assert read() assert run_in_forked_child(read) == 0 - assert fs.core.get_object.call_count > 1 @pytest.mark.parametrize("size", [0, 10]) @pytest.mark.asyncio diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index b86db428e..2402d6a4e 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -20,7 +20,7 @@ from botocore.exceptions import ClientError from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPool -from pyathena.util import RetryConfig, retry_api_call, stop_retries_on +from pyathena.util import RetryConfig, override, retry_api_call, stop_retries_on from tests.pyathena.util import ( EVENT_TIMEOUT, ThrottledOperation, @@ -331,11 +331,13 @@ class _DeferringExecutor(S3Executor): def __init__(self) -> None: self.submitted: list[Future[Any]] = [] + @override def submit(self, fn, *args, **kwargs): future: Future[Any] = Future() self.submitted.append(future) return future + @override def shutdown(self, wait=True): pass From f855e7fe11c7e12368a345b6b692b82f791b00ec Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 4 Nov 2026 01:12:38 +0900 Subject: [PATCH 3/5] Describe _create_executor as the executor source of files and operations (GH-1332) Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 3 +-- pyathena/filesystem/s3.py | 9 +++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index ce5b2a496..70062e10a 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -339,8 +339,7 @@ threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3ThreadPo `loop` argument, with a `DeprecationWarning` if one is given. The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel -range reads in a pool of `max_workers` threads of each file, as `S3File` does, without -the event loop. +range reads in a pool of `max_workers` threads of each file, as `S3File` does. In a forked child process, the event loop of the parent does not run, so fsspec's synchronous methods of an `AioS3FileSystem` created in the parent, such as `ls()`, raise diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 9b5ce5e88..e79b2a663 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1099,16 +1099,17 @@ def _delete_paths( ) def _create_executor(self, max_workers: int) -> S3Executor: - """Create an executor strategy for parallel operations. + """Create the executor of the parallel requests of a file or an operation. - Subclasses can override to provide alternative execution strategies - (e.g., asyncio-based execution). + The files of ``S3FileSystem`` and ``AioS3FileSystem`` and the + parallel operations of ``S3FileSystem`` take their executors from + this method. Args: max_workers: Maximum number of parallel workers. Returns: - An S3Executor instance. + A new ``S3ThreadPool``. """ return S3ThreadPool(max_workers=max_workers) From dbb262c1efa0031c469940fd0bf461ced76ecc0a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 4 Nov 2026 01:26:40 +0900 Subject: [PATCH 4/5] Narrow the fork and ordering docs and drop a duplicate test (GH-1332) A filesystem created in the parent keeps the parent's boto3 client, so the docs say to create a new filesystem in a forked child instead of promising that the parent's works there. S3ThreadPool queues the calls in submission order; with several threads it does not guarantee the order in which they start. The asyncio variant of the failed-read memory test ran the same S3File code as the synchronous one. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 14 ++++---- pyathena/filesystem/s3_executor.py | 4 +-- tests/pyathena/filesystem/test_s3_async.py | 42 ---------------------- 3 files changed, 9 insertions(+), 51 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 70062e10a..6bbab094b 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -332,19 +332,19 @@ multipart uploads. ### Executor strategy `S3FileSystem` and `S3File` use a pluggable executor abstraction (`S3Executor`) for -parallel operations. `S3ThreadPool` runs the calls in a pool of `max_workers` threads, -in the order they were submitted; in a forked child process, the pool creates new -threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3ThreadPool`. +parallel operations. `S3ThreadPool` queues the calls in the order they were submitted +and runs them in a pool of `max_workers` threads; in a forked child process, the pool +creates new threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3ThreadPool`. `S3AioExecutor` is kept for compatibility: it is an `S3ThreadPool` that ignores its `loop` argument, with a `DeprecationWarning` if one is given. The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel range reads in a pool of `max_workers` threads of each file, as `S3File` does. -In a forked child process, the event loop of the parent does not run, so fsspec's -synchronous methods of an `AioS3FileSystem` created in the parent, such as `ls()`, raise -`RuntimeError` there. Its `open()`, and its coroutines run on an event loop of the child, -work, and so does a filesystem created in the child. +In a forked child process, create a new filesystem. A filesystem created in the parent +keeps the parent's boto3 client, which is not meant to be shared across processes, and +the event loop of the parent does not run in the child, so fsspec's synchronous methods +of an `AioS3FileSystem` created in the parent, such as `ls()`, raise `RuntimeError` there. ### Async file handles diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index ae77d0f74..965a04443 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -182,8 +182,8 @@ class S3ThreadPool(S3Executor): It runs at most ``max_workers`` calls at once in threads of its own, so that, unlike ``asyncio.to_thread()``, which runs a call in the default executor of the event loop, neither the size nor the load of that - executor limits or blocks the calls of asyncio code. The calls start in - the order they were submitted. The threads are created as the calls need + executor limits or blocks the calls of asyncio code. The calls are + queued in the order they were submitted. 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, as it does after diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 47dd4a5b4..f1bf51820 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2110,48 +2110,6 @@ def test_open_max_workers(self, max_workers): assert isinstance(f, AioS3File) assert f.max_workers == (max_workers or fs._sync_fs.max_workers) - @pytest.mark.asyncio - async def test_failed_read_frees_other_ranges(self): - # GH-1186: a parallel read that fails frees the data of the other - # ranges, also while its error is held, as with S3File. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - block_size = 2**10 - read = [] - others_read = threading.Semaphore(0) - - def get_object(path, range_, **kwargs): - if range_[0] == 0: - for _ in range(3): - assert others_read.acquire(timeout=5) - raise PermissionError("Access Denied") - # Unlike bytes, an array can be weakly referenced. - data = array.array("B", bytes(range_[1] - range_[0])) - read.append(weakref.ref(data)) - others_read.release() - return data - - 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 = block_size * 4 - fs._sync_fs.info = mock.MagicMock(return_value=info) - # The blocking calls of the file run in threads, as AioS3FSCursor - # reads result files, not on the event loop. - f = await asyncio.to_thread( - fs.open, "s3://bucket/key", "rb", block_size=block_size, cache_type="none" - ) - with pytest.raises(PermissionError, match="Access Denied") as info_: - await asyncio.to_thread(f.read) - # Waits for the reads, whose threads may still set the results. - await asyncio.to_thread(f.close) - assert len(read) == 3 - assert [ref() for ref in read] == [None] * 3 - assert info_.value.__traceback__ is not None - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_open_parallel_requests(self, asynchronous): From 06ff6632333f6961ffa029c6b418e491eba97276 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 4 Nov 2026 01:30:32 +0900 Subject: [PATCH 5/5] State the queue order in the S3ThreadPool.submit() docstring (GH-1332) Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 965a04443..2da2b57e7 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -255,9 +255,10 @@ 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. + only before the call starts. The calls are queued in the order they + were submitted, so that the calls submitted before a call that waits + for them are taken by threads before it, and it does not wait for a + thread that it holds. Args: fn: The blocking callable.