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..6bbab094b 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,19 @@ 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` 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, 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/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..e79b2a663 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, @@ -1099,18 +1099,19 @@ 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 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 +2883,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 +2929,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..029054cab 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 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() @@ -836,28 +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. - - An instance created with ``asynchronous=True`` has no event loop of - its own, so its files run the operations in a thread pool. - - 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. - """ - if self._loop is None: - return S3ThreadPoolExecutor(max_workers=max_workers) - return S3AioExecutor(loop=self.loop, max_workers=max_workers) - def _open( self, path: str, @@ -888,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, @@ -1194,17 +1165,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 +1178,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 +1264,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..2da2b57e7 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 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 - :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 ``S3AioThreadPool`` 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,16 +248,17 @@ 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]: """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. @@ -483,148 +458,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..8abc70fb3 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 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) @@ -992,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 the shutdown of the event loop of - # an S3AioExecutor, 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.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..f1bf51820 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,36 @@ 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(): + calls = fs.core.get_object.call_count + with fs.open("s3://bucket/key", "rb", block_size=block_size, cache_type="none") as f: + 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 - 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 @pytest.mark.parametrize("size", [0, 10]) @pytest.mark.asyncio @@ -1080,7 +1095,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 +1122,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 +1157,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 +1177,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. @@ -2095,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 executor of the file dispatches through the event loop, which - # the blocking calls must not hold. - 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): @@ -2269,7 +2242,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 +2283,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 +2503,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 +2531,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 +2694,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 +2715,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 +2737,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 +2753,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 +2968,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..2402d6a4e 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -13,35 +13,34 @@ 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.util import RetryConfig, retry_api_call, stop_retries_on +from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPool +from pyathena.util import RetryConfig, override, retry_api_call, stop_retries_on from tests.pyathena.util import ( EVENT_TIMEOUT, ThrottledOperation, 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,39 @@ 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]] = [] + @override + 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 + @override + 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 +357,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 +365,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 +376,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 +392,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 +412,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 +423,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: