diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 49dfc71e..a825bef9 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -13,7 +13,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.11", "3.12", "3.13", "3.14"] + python-version: ["3.11", "3.12", "3.13", "3.14", "3.15", "3.15t"] steps: - uses: actions/checkout@v7 @@ -40,6 +40,12 @@ jobs: uv pip install --system cryptography - name: Install WebDAV packages for testing + # wsgidav needs bcrypt, which publishes no free-threaded wheel for + # this Python and cannot be built from source because the pyo3 its + # Rust extension uses refuses a Python newer than 3.13 and cannot + # fall back to the limited API on a free-threaded build. The WebDAV + # tests skip when wsgidav is missing. + if: matrix.python-version != '3.15t' run: | uv pip install --system cheroot wsgidav diff --git a/doc/changes/DM-56097.bugfix.rst b/doc/changes/DM-56097.bugfix.rst new file mode 100644 index 00000000..51ae4e59 --- /dev/null +++ b/doc/changes/DM-56097.bugfix.rst @@ -0,0 +1 @@ +``clean_test_environment_for_s3()`` now also clears ``$AWS_REQUEST_CHECKSUM_CALCULATION`` and ``$AWS_RESPONSE_CHECKSUM_VALIDATION``, so that a site which disables checksums does not change the object metadata seen by tests. diff --git a/doc/changes/DM-56097.perf.rst b/doc/changes/DM-56097.perf.rst new file mode 100644 index 00000000..97a06ee2 --- /dev/null +++ b/doc/changes/DM-56097.perf.rst @@ -0,0 +1,6 @@ +``ResourcePath.mremove()``, ``ResourcePath.mexists()`` and ``ResourcePath.mtransfer()`` now send a batch of URIs to each worker rather than submitting one task per URI. +The batch size is fixed per scheme rather than derived from how many URIs a call is given, so a batch never grows large enough that one worker drawing a run of slow URIs stalls the operation with no way to rebalance. +A scheme sets it from the cost of one of its operations: a local file check is cheap enough that a batch has to reach a thousand URIs before threading it beats a loop in the calling thread, while a scheme whose every operation is a network round trip is worth overlapping immediately. +Transfers batch separately from existence checks and removals, since a transfer costs far more and scales with a file size that is not known in advance. +The upper bound on the number of workers is now set per scheme, so a scheme that holds no connection pool can raise it, and ``ResourcePath.mtransfer()`` now uses the bound for the target scheme instead of the global default. +A batch of URIs that fits in a single chunk is handled in the calling thread rather than being given to a pool. diff --git a/doc/changes/DM-56097.removal.rst b/doc/changes/DM-56097.removal.rst new file mode 100644 index 00000000..6ea5037f --- /dev/null +++ b/doc/changes/DM-56097.removal.rst @@ -0,0 +1,3 @@ +Removed support for the ``$LSST_RESOURCES_EXECUTOR`` environment variable. +Bulk operations now always use a thread pool. +Process pools were measured to be no faster for remote URIs and they silently discarded the undo actions registered by ``ResourcePath.mtransfer()`` on a caller's transaction, since those were recorded in the worker process and never reached the parent. diff --git a/python/lsst/resources/_resourcePath.py b/python/lsst/resources/_resourcePath.py index 692d8142..abcf4026 100644 --- a/python/lsst/resources/_resourcePath.py +++ b/python/lsst/resources/_resourcePath.py @@ -29,7 +29,7 @@ from collections import defaultdict from pathlib import Path, PurePath, PurePosixPath from random import Random -from typing import TYPE_CHECKING, TypeAlias +from typing import TYPE_CHECKING try: import fsspec @@ -42,10 +42,12 @@ AbstractFileSystem = type from collections.abc import Generator, Iterable, Iterator -from typing import Any, Literal, NamedTuple, overload +from typing import Any, Literal, NamedTuple, TypeVar, overload + +from lsst.utils.iteration import chunk_iterable from ._resourceHandles._baseResourceHandle import ResourceHandleProtocol -from .utils import _get_num_workers, get_tempdir +from .utils import MAX_WORKERS, _get_num_workers, get_tempdir if TYPE_CHECKING: from .utils import TransactionProtocol @@ -59,6 +61,8 @@ # Precomputed escaped hash ESCAPED_HASH = urllib.parse.quote("#") +_T = TypeVar("_T") + class MBulkResult(NamedTuple): """Report on a bulk operation.""" @@ -67,75 +71,6 @@ class MBulkResult(NamedTuple): exception: Exception | None -_EXECUTOR_TYPE: TypeAlias = type[ - concurrent.futures.ThreadPoolExecutor | concurrent.futures.ProcessPoolExecutor -] - -# Cache value for executor class so as not to issue warning multiple -# times but still allow tests to override the value. -_POOL_EXECUTOR_CLASS: _EXECUTOR_TYPE | None = None - - -def _get_executor_class() -> _EXECUTOR_TYPE: - """Return the executor class used for parallelized execution. - - Returns - ------- - cls : `concurrent.futures.Executor` - The ``Executor`` class. Default is - `concurrent.futures.ThreadPoolExecutor`. Can be set explicitly by - setting the ``$LSST_RESOURCES_EXECUTOR`` environment variable to - "thread" or "process". Returns "thread" pool if the value of the - variable is not recognized. - """ - global _POOL_EXECUTOR_CLASS - - if _POOL_EXECUTOR_CLASS is not None: - return _POOL_EXECUTOR_CLASS - - pool_executor_classes = { - "threads": concurrent.futures.ThreadPoolExecutor, - "process": concurrent.futures.ProcessPoolExecutor, - } - default_executor = "threads" - external = os.getenv("LSST_RESOURCES_EXECUTOR", default_executor) - if not external: - external = default_executor - if external not in pool_executor_classes: - log.warning( - "Unrecognized value of '%s' for LSST_RESOURCES_EXECUTOR env var. Using '%s'", - external, - default_executor, - ) - external = default_executor - _POOL_EXECUTOR_CLASS = pool_executor_classes[external] - return _POOL_EXECUTOR_CLASS - - -@contextlib.contextmanager -def _patch_environ(new_values: dict[str, str]) -> Generator[None]: - """Patch os.environ temporarily using the supplied values. - - Parameters - ---------- - new_values : `dict` [ `str`, `str` ] - New values to be stored in the environment. - """ - old_values: dict[str, str] = {} - for k, v in new_values.items(): - if k in os.environ: - old_values[k] = os.environ[k] - os.environ[k] = v - - try: - yield - finally: - for k in new_values: - del os.environ[k] - if k in old_values: - os.environ[k] = old_values[k] - - @dataclasses.dataclass(frozen=True) class ResourceInfo: """Information about this resource.""" @@ -227,6 +162,33 @@ class ResourcePath: # numpydoc ignore=PR02 isLocal = False """If `True` this URI refers to a local file.""" + _max_workers: int = MAX_WORKERS + """Upper bound on workers for parallel operations on this scheme. + + Schemes backed by a connection pool keep this modest because the pool is + sized to match it; schemes with no pool can raise it. + """ + + _chunk_size: int = 1 + """Number of URIs given to each worker by `mexists` and `mremove`. + + A batch that fits in a single chunk is handled in the calling thread + rather than being sent to a pool. The default suits a scheme where every + operation is a network round trip, which is expensive enough that handing + over one URI at a time costs nothing and gives the pool the freedom to + balance itself. A scheme whose operations are cheap should raise it until + a chunk is worth handing over. + """ + + _transfer_chunk_size: int = 1 + """Number of files given to each worker by `mtransfer`. + + Kept separate from ``_chunk_size`` because a transfer costs orders of + magnitude more than an existence check on the same scheme, and its cost + scales with a file size the caller does not know in advance, so chunks + have to stay small enough for the pool queue to balance them. + """ + # This is not an ABC with abstract methods because the __new__ being # a factory confuses mypy such that it assumes that every constructor # returns a ResourcePath and then determines that all the abstract methods @@ -996,10 +958,9 @@ def mexists( uris : iterable of `ResourcePath` The URIs to test. num_workers : `int` or `None`, optional - The number of parallel workers to use when checking for existence - If `None`, the default value will be taken from the environment. - If this number is higher than the default and a thread pool is - used, there may not be enough cached connections available. + The number of parallel workers to use when checking for existence. + If `None`, the default value will be taken from the environment + and bounded by the limit for this scheme. Returns ------- @@ -1020,7 +981,6 @@ def _mexists( Implementation helper method for `mexists`. - Parameters ---------- uris : iterable of `ResourcePath` @@ -1034,55 +994,50 @@ def _mexists( existence : `dict` of [`ResourcePath`, `bool`] Mapping of original URI to boolean indicating existence. """ - pool_executor_class = _get_executor_class() - if issubclass(pool_executor_class, concurrent.futures.ProcessPoolExecutor): - # Patch the environment to make it think there is only one worker - # for each subprocess. - with _patch_environ({"LSST_RESOURCES_NUM_WORKERS": "1"}): - return cls._mexists_pool(pool_executor_class, uris) - else: - return cls._mexists_pool(pool_executor_class, uris, num_workers=num_workers) + uri_list = list(uris) + max_workers = num_workers if num_workers is not None else _get_num_workers(cls._max_workers) + chunks = cls._chunk_work(uri_list, cls._chunk_size) + if not chunks: + return {} + if len(chunks) == 1: + # Not enough work to be worth handing to another thread. + return cls._exists_chunk(chunks[0]) + + results: dict[ResourcePath, bool] = {} + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as exists_executor: + future_exists = {exists_executor.submit(cls._exists_chunk, chunk): chunk for chunk in chunks} + for future in concurrent.futures.as_completed(future_exists): + try: + results.update(future.result()) + except Exception: + # The chunk failed as a whole, for example because the + # pool could not start a thread. + for uri in future_exists[future]: + results[uri] = False + return results @classmethod - def _mexists_pool( - cls, - pool_executor_class: _EXECUTOR_TYPE, - uris: Iterable[ResourcePath], - *, - num_workers: int | None = None, - ) -> dict[ResourcePath, bool]: - """Check for existence of multiple URIs at once using specified pool - executor. - - Implementation helper method for `_mexists`. + def _exists_chunk(cls, uris: tuple[ResourcePath, ...]) -> dict[ResourcePath, bool]: + """Check a batch of URIs for existence. Parameters ---------- - pool_executor_class : `type` [ `concurrent.futures.Executor` ] - Type of executor pool to use. - uris : iterable of `ResourcePath` - The URIs to test. - num_workers : `int` or `None`, optional - The number of parallel workers to use when checking for existence - If `None`, the default value will be taken from the environment. + uris : `tuple` [ `ResourcePath`, ... ] + The URIs to check. Returns ------- - existence : `dict` of [`ResourcePath`, `bool`] - Mapping of original URI to boolean indicating existence. + results : `dict` [ `ResourcePath`, `bool` ] + An entry for every URI in ``uris``. A URI that cannot be checked + is reported as absent, and does not prevent the URIs after it in + the batch from being checked. """ - max_workers = num_workers if num_workers is not None else _get_num_workers() - with pool_executor_class(max_workers=max_workers) as exists_executor: - future_exists = {exists_executor.submit(uri.exists): uri for uri in uris} - - results: dict[ResourcePath, bool] = {} - for future in concurrent.futures.as_completed(future_exists): - uri = future_exists[future] - try: - exists = future.result() - except Exception: - exists = False - results[uri] = exists + results: dict[ResourcePath, bool] = {} + for uri in uris: + try: + results[uri] = uri.exists() + except Exception: + results[uri] = False return results @classmethod @@ -1123,47 +1078,44 @@ def mtransfer( whether the transfer succeeded for the target URI. If ``do_raise`` is `True`, this will only be returned if there are no errors. """ - pool_executor_class = _get_executor_class() - if issubclass(pool_executor_class, concurrent.futures.ProcessPoolExecutor): - # Patch the environment to make it think there is only one worker - # for each subprocess. - with _patch_environ({"LSST_RESOURCES_NUM_WORKERS": "1"}): - return cls._mtransfer( - pool_executor_class, - transfer, - from_to, - overwrite=overwrite, - transaction=transaction, - do_raise=do_raise, - ) - return cls._mtransfer( - pool_executor_class, - transfer, - from_to, - overwrite=overwrite, - transaction=transaction, - do_raise=do_raise, - ) + # A transfer is driven by the target, so group by the target scheme + # and let each scheme decide how many workers to use. + grouped: dict[type[ResourcePath], list[tuple[ResourcePath, ResourcePath]]] = defaultdict(list) + for from_uri, to_uri in from_to: + grouped[to_uri.__class__].append((from_uri, to_uri)) + + results: dict[ResourcePath, MBulkResult] = {} + for uri_class, group in grouped.items(): + results.update( + uri_class._mtransfer(transfer, group, overwrite=overwrite, transaction=transaction) + ) + + if do_raise and any(not res.success for res in results.values()): + raise ExceptionGroup( + f"Errors transferring {len(results)} artifacts", + tuple(res.exception for res in results.values() if res.exception is not None), + ) + + return results @classmethod def _mtransfer( cls, - pool_executor_class: _EXECUTOR_TYPE, transfer: str, from_to: Iterable[tuple[ResourcePath, ResourcePath]], + *, overwrite: bool = False, transaction: TransactionProtocol | None = None, - do_raise: bool = True, ) -> dict[ResourcePath, MBulkResult]: - """Transfer many files in bulk. + """Transfer many files in bulk to targets of this scheme. + + Implementation helper method for `mtransfer`. Parameters ---------- transfer : `str` - Mode to use for transferring the resource. Generically there are - many standard options: copy, link, symlink, hardlink, relsymlink. - Not all URIs support all modes. - from_to : `list` [ `tuple` [ `ResourcePath`, `ResourcePath` ] ] + Mode to use for transferring the resource. + from_to : iterable [ `tuple` [ `ResourcePath`, `ResourcePath` ] ] A sequence of the source URIs and the target URIs. overwrite : `bool`, optional Allow an existing file to be overwritten. Defaults to `False`. @@ -1171,10 +1123,6 @@ def _mtransfer( A transaction object that can (depending on implementation) rollback transfers on error. Not guaranteed to be implemented. The transaction object must be thread safe. - do_raise : `bool`, optional - If `True` an `ExceptionGroup` will be raised containing any - exceptions raised by the individual transfers. Else a dict - reporting the status of each `ResourcePath` will be returned. Returns ------- @@ -1182,37 +1130,73 @@ def _mtransfer( A dict of all the transfer attempts with a value indicating whether the transfer succeeded for the target URI. """ - with pool_executor_class(max_workers=_get_num_workers()) as transfer_executor: + pairs = list(from_to) + max_workers = _get_num_workers(cls._max_workers) + chunks = cls._chunk_work(pairs, cls._transfer_chunk_size) + if not chunks: + return {} + if len(chunks) == 1: + # Not enough work to be worth handing to another thread. + return cls._transfer_chunk(chunks[0], transfer, overwrite, transaction) + + results: dict[ResourcePath, MBulkResult] = {} + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as transfer_executor: future_transfers = { - transfer_executor.submit( - to_uri.transfer_from, - from_uri, - transfer=transfer, - overwrite=overwrite, - transaction=transaction, - multithreaded=False, - ): to_uri - for from_uri, to_uri in from_to + transfer_executor.submit(cls._transfer_chunk, chunk, transfer, overwrite, transaction): chunk + for chunk in chunks } - results: dict[ResourcePath, MBulkResult] = {} - failed = False for future in concurrent.futures.as_completed(future_transfers): - to_uri = future_transfers[future] try: - future.result() + results.update(future.result()) except Exception as e: - transferred = MBulkResult(False, e) - failed = True - else: - transferred = MBulkResult(True, None) - results[to_uri] = transferred + # The chunk failed as a whole, for example because the + # pool could not start a thread. + for _, to_uri in future_transfers[future]: + results[to_uri] = MBulkResult(False, e) + return results - if do_raise and failed: - raise ExceptionGroup( - f"Errors transferring {len(results)} artifacts", - tuple(res.exception for res in results.values() if res.exception is not None), - ) + @classmethod + def _transfer_chunk( + cls, + from_to: tuple[tuple[ResourcePath, ResourcePath], ...], + transfer: str, + overwrite: bool, + transaction: TransactionProtocol | None, + ) -> dict[ResourcePath, MBulkResult]: + """Transfer a batch of files, reporting each result independently. + Parameters + ---------- + from_to : `tuple` [ `tuple` [ `ResourcePath`, `ResourcePath` ], ... ] + The source and target URIs to transfer. + transfer : `str` + Mode to use for transferring the resource. + overwrite : `bool` + Allow an existing file to be overwritten. + transaction : `~lsst.resources.utils.TransactionProtocol` or `None` + A transaction object that can (depending on implementation) + rollback transfers on error. + + Returns + ------- + results : `dict` [ `ResourcePath`, `MBulkResult` ] + An entry for every target URI in ``from_to``. A transfer that + fails does not prevent the transfers after it in the batch. + """ + results: dict[ResourcePath, MBulkResult] = {} + for from_uri, to_uri in from_to: + try: + to_uri.transfer_from( + from_uri, + transfer=transfer, + overwrite=overwrite, + transaction=transaction, + multithreaded=False, + ) + except Exception as e: + results[to_uri] = MBulkResult(False, e) + else: + results[to_uri] = MBulkResult(True, None) return results def remove(self) -> None: @@ -1258,40 +1242,97 @@ def mremove( return results + @staticmethod + def _chunk_work(items: list[_T], chunk_size: int) -> list[tuple[_T, ...]]: + """Split work items into batches of a fixed size. + + Parameters + ---------- + items : `list` + The work items to split. + chunk_size : `int` + Number of items to put in each batch. + + Returns + ------- + chunks : `list` [ `tuple` ] + The batches. Empty if ``items`` is empty. A single batch means the + work is not worth spreading, and callers run it directly. + + Notes + ----- + The batch size does not depend on how many items there are. Sizing it + from the total would make a batch grow without bound as the total + grows, and a batch that draws a run of slow URIs then stalls a worker + for the rest of the operation with no way to rebalance. Asking for + more batches than there are workers is harmless, since a pool only + starts a thread when there is a batch waiting for it. + """ + if not items: + return [] + return list(chunk_iterable(items, chunk_size=chunk_size)) + + @classmethod + def _remove_chunk(cls, uris: tuple[ResourcePath, ...]) -> dict[ResourcePath, MBulkResult]: + """Remove a batch of URIs, reporting each result independently. + + Parameters + ---------- + uris : `tuple` [ `ResourcePath`, ... ] + The URIs to remove. + + Returns + ------- + results : `dict` [ `ResourcePath`, `MBulkResult` ] + An entry for every URI in ``uris``. A URI that cannot be removed + does not prevent the removal of the URIs after it. + """ + results: dict[ResourcePath, MBulkResult] = {} + for uri in uris: + try: + uri.remove() + except Exception as e: + results[uri] = MBulkResult(False, e) + else: + results[uri] = MBulkResult(True, None) + return results + @classmethod def _mremove(cls, uris: Iterable[ResourcePath]) -> dict[ResourcePath, MBulkResult]: - """Remove multiple URIs using futures.""" - pool_executor_class = _get_executor_class() - if issubclass(pool_executor_class, concurrent.futures.ProcessPoolExecutor): - # Patch the environment to make it think there is only one worker - # for each subprocess. - with _patch_environ({"LSST_RESOURCES_NUM_WORKERS": "1"}): - return cls._mremove_pool(pool_executor_class, uris) - else: - return cls._mremove_pool(pool_executor_class, uris) + """Remove multiple URIs using threads. + + Implementation helper method for `mremove`. + + Parameters + ---------- + uris : iterable of `ResourcePath` + The URIs to remove. + + Returns + ------- + removal : `dict` of [`ResourcePath`, `MBulkResult`] + Mapping of original URI to the result of removing it. + """ + uri_list = list(uris) + max_workers = _get_num_workers(cls._max_workers) + chunks = cls._chunk_work(uri_list, cls._chunk_size) + if not chunks: + return {} + if len(chunks) == 1: + # Not enough work to be worth handing to another thread. + return cls._remove_chunk(chunks[0]) - @classmethod - def _mremove_pool( - cls, - pool_executor_class: _EXECUTOR_TYPE, - uris: Iterable[ResourcePath], - *, - num_workers: int | None = None, - ) -> dict[ResourcePath, MBulkResult]: - """Remove URIs using a futures pool.""" - max_workers = num_workers if num_workers is not None else _get_num_workers() results: dict[ResourcePath, MBulkResult] = {} - with pool_executor_class(max_workers=max_workers) as remove_executor: - future_remove = {remove_executor.submit(uri.remove): uri for uri in uris} + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as remove_executor: + future_remove = {remove_executor.submit(cls._remove_chunk, chunk): chunk for chunk in chunks} for future in concurrent.futures.as_completed(future_remove): try: - future.result() + results.update(future.result()) except Exception as e: - removed = MBulkResult(False, e) - else: - removed = MBulkResult(True, None) - uri = future_remove[future] - results[uri] = removed + # The chunk failed as a whole, for example because the + # pool could not start a thread. + for uri in future_remove[future]: + results[uri] = MBulkResult(False, e) return results def isabs(self) -> bool: diff --git a/python/lsst/resources/file.py b/python/lsst/resources/file.py index 5bef1b73..7dd80f3a 100644 --- a/python/lsst/resources/file.py +++ b/python/lsst/resources/file.py @@ -84,6 +84,17 @@ class FileResourcePath(ResourcePath): # By definition refers to a local file isLocal = True + # A warm existence check or removal here costs a few microseconds, almost + # all of it holding the GIL, so a batch has to be this large before + # spreading it over threads beats a plain loop in the calling thread. + # Measurements above this size are flat, so it also matches the 1000 keys + # an S3 bulk delete takes. + _chunk_size = 1000 + + # A transfer costs orders of magnitude more than an existence check and + # scales with the file size, so batches stay small enough to balance. + _transfer_chunk_size = 25 + @property def ospath(self) -> str: """Path component of the URI localized to current OS. diff --git a/python/lsst/resources/s3.py b/python/lsst/resources/s3.py index 276d54c8..80366691 100644 --- a/python/lsst/resources/s3.py +++ b/python/lsst/resources/s3.py @@ -35,14 +35,7 @@ from ._resourceHandles._baseResourceHandle import ResourceHandleProtocol from ._resourceHandles._s3ResourceHandle import S3ResourceHandle -from ._resourcePath import ( - _EXECUTOR_TYPE, - MBulkResult, - ResourceInfo, - ResourcePath, - _get_executor_class, - _patch_environ, -) +from ._resourcePath import MBulkResult, ResourceInfo, ResourcePath from .s3utils import ( _get_s3_connection_parameters, _s3_disable_bucket_validation, @@ -295,28 +288,14 @@ def _mremove_select(cls, chunks: list[tuple[ResourcePath, ...]]) -> dict[Resourc if len(chunks) == 1: # Do the removal directly without futures. return cls._delete_objects_wrapper(chunks[0]) - pool_executor_class = _get_executor_class() - if issubclass(pool_executor_class, concurrent.futures.ProcessPoolExecutor): - # Patch the environment to make it think there is only one worker - # for each subprocess. - with _patch_environ({"LSST_RESOURCES_NUM_WORKERS": "1"}): - return cls._mremove_with_pool(pool_executor_class, chunks) - else: - return cls._mremove_with_pool(pool_executor_class, chunks) + return cls._mremove_with_pool(chunks) @classmethod - def _mremove_with_pool( - cls, - pool_executor_class: _EXECUTOR_TYPE, - chunks: list[tuple[ResourcePath, ...]], - *, - num_workers: int | None = None, - ) -> dict[ResourcePath, MBulkResult]: + def _mremove_with_pool(cls, chunks: list[tuple[ResourcePath, ...]]) -> dict[ResourcePath, MBulkResult]: # Different name because different API to base class. - # No need to make more workers than we have chunks. - max_workers = num_workers if num_workers is not None else min(len(chunks), _get_num_workers()) + max_workers = _get_num_workers(cls._max_workers) results: dict[ResourcePath, MBulkResult] = {} - with pool_executor_class(max_workers=max_workers) as remove_executor: + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as remove_executor: future_remove = { remove_executor.submit(cls._delete_objects_wrapper, chunk): i for i, chunk in enumerate(chunks) diff --git a/python/lsst/resources/s3utils.py b/python/lsst/resources/s3utils.py index 18aaab26..4f81aaa6 100644 --- a/python/lsst/resources/s3utils.py +++ b/python/lsst/resources/s3utils.py @@ -125,7 +125,8 @@ class _TooManyRequestsError(Exception): @contextmanager def clean_test_environment_for_s3() -> Generator[None]: """Reset S3 environment to ensure that unit tests with a mock S3 can't - accidentally reference real infrastructure. + accidentally reference real infrastructure, and that site configuration + cannot change how the client behaves. """ with patch.dict( os.environ, @@ -142,6 +143,10 @@ def clean_test_environment_for_s3() -> Generator[None]: "AWS_PROFILE", "AWS_SHARED_CREDENTIALS_FILE", "AWS_CONFIG_FILE", + # A site that turns checksums off changes what an object's + # metadata contains, so let the library defaults apply. + "AWS_REQUEST_CHECKSUM_CALCULATION", + "AWS_RESPONSE_CHECKSUM_VALIDATION", ): patched_environ.pop(var, None) # Clear the cached boto3 S3 client instances. diff --git a/python/lsst/resources/utils.py b/python/lsst/resources/utils.py index e29f687b..14f1e1f0 100644 --- a/python/lsst/resources/utils.py +++ b/python/lsst/resources/utils.py @@ -34,9 +34,10 @@ # posix means posix and only determine explicitly in the non-posix case. OS_ROOT_PATH = posixpath.sep if IS_POSIX else Path().resolve().root -# Maximum number of worker threads for parallelized operations. -# If greater than 10, be aware that this number has to be consistent -# with connection pool sizing (for example in urllib3). +# Default upper bound on the number of workers for parallelized operations. +# Subclasses of ResourcePath override this for schemes that have no connection +# pool to contend with. Backends that do have one size that pool from the +# worker count, so they need no separate coordination here. MAX_WORKERS = 10 log = logging.getLogger(__name__) @@ -244,28 +245,50 @@ def _get_int_env_var(env_var: str) -> int | None: @cache -def _get_num_workers() -> int: - f"""Calculate the number of workers to use. +def _get_configured_num_workers() -> int | None: + """Return the explicitly requested number of workers. + + Returns + ------- + num : `int` or `None` + Value of the ``LSST_RESOURCES_NUM_WORKERS`` environment variable, or + `None` if it is unset or unparsable. + """ + return _get_int_env_var("LSST_RESOURCES_NUM_WORKERS") + + +@cache +def _get_default_num_workers() -> int: + """Return the number of workers implied by the available CPUs. Returns ------- num : `int` - The number of workers to use. Will use the value of the - ``LSST_RESOURCES_NUM_WORKERS`` environment variable if set. Will fall - back to using the CPU count (plus 2) but capped at {MAX_WORKERS}. + The CPU count plus two. Uncapped. """ - num_workers: int | None = None - num_workers = _get_int_env_var("LSST_RESOURCES_NUM_WORKERS") + # CPU_LIMIT is used on nublado. + cpu_limit = _get_int_env_var("CPU_LIMIT") or multiprocessing.cpu_count() + return cpu_limit + 2 - # If someone is explicitly specifying a number, let them use that number. - if num_workers is not None: - return num_workers - if num_workers is None: - # CPU_LIMIT is used on nublado. - cpu_limit = _get_int_env_var("CPU_LIMIT") or multiprocessing.cpu_count() - if cpu_limit is not None: - num_workers = cpu_limit + 2 +def _get_num_workers(max_workers: int = MAX_WORKERS) -> int: + """Calculate the number of workers to use. - # But don't ever return more than the maximum allowed. - return min([num_workers, MAX_WORKERS]) + Parameters + ---------- + max_workers : `int`, optional + Upper bound to apply to the calculated default. Ignored when the + number of workers has been requested explicitly. + + Returns + ------- + num : `int` + The number of workers to use. The value of + ``$LSST_RESOURCES_NUM_WORKERS`` is used if set, and the CPU count plus + two bounded by ``max_workers`` if not. + """ + configured = _get_configured_num_workers() + if configured is not None: + # An explicit request is honored without capping. + return configured + return min(_get_default_num_workers(), max_workers) diff --git a/tests/test_file.py b/tests/test_file.py index 45a31e62..427f37e8 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -11,14 +11,20 @@ import contextlib import datetime +import functools import os import pathlib +import threading import unittest import unittest.mock import urllib.parse +from collections.abc import Callable, Iterator +from typing import Any from lsst.resources import ResourceInfo, ResourcePath, ResourcePathExpression +from lsst.resources.file import FileResourcePath from lsst.resources.tests import GenericReadWriteTestCase, GenericTestCase +from lsst.resources.utils import makeTestTempDir, removeTestTempDir TESTDIR = os.path.abspath(os.path.dirname(__file__)) @@ -248,23 +254,154 @@ def _test_with_restrictive_umask(self, callback): mode = os.stat(dir).st_mode self.assertEqual(mode & TEST_UMASK, 0o0300, f"Permissions incorrect for {dir}: {mode:o}") - @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) - @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "process"}) - def test_mexists_process(self) -> None: - """Test mexists with override executor pool. - Force test with process pool. - """ - super().test_mexists() +class BulkOperationTestCase(unittest.TestCase): + """Tests for batched bulk operations on local files.""" - @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) - @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "process"}) - def test_mtransfer_process(self) -> None: - """Test transfer with override executor pool. + def setUp(self) -> None: + self.tmpdir = ResourcePath(makeTestTempDir(TESTDIR), forceDirectory=True) + + def tearDown(self) -> None: + removeTestTempDir(self.tmpdir.ospath) - Force test with process pool. - """ - super().test_mtransfer() + def _make_files(self, prefix: str, count: int) -> list[ResourcePath]: + uris = [self.tmpdir.join(f"{prefix}{n}.txt") for n in range(count)] + for uri in uris: + uri.write(b"") + return uris + + def test_chunk_sizes(self) -> None: + items = list(range(10)) + + # The batch size does not depend on how many items there are. + self.assertEqual([len(c) for c in FileResourcePath._chunk_work(items, 4)], [4, 4, 2]) + self.assertEqual([len(c) for c in FileResourcePath._chunk_work(items, 1)], [1] * 10) + + # A batch that fits in one chunk is the signal to the caller to handle + # it without a pool. + self.assertEqual([len(c) for c in FileResourcePath._chunk_work(items, 100)], [10]) + + # An empty input yields no chunks at all. + self.assertEqual(FileResourcePath._chunk_work([], 4), []) + + def test_schemes_size_their_own_chunks(self) -> None: + # A local operation is cheap enough that a batch has to be large + # before threading it pays, while a scheme whose every operation is a + # round trip is worth overlapping immediately. + self.assertEqual(FileResourcePath._chunk_size, 1000) + self.assertEqual(ResourcePath._chunk_size, 1) + + # A transfer is far more expensive than an existence check on the same + # scheme, so it batches separately. + self.assertLess(FileResourcePath._transfer_chunk_size, FileResourcePath._chunk_size) + + def test_empty_bulk_operations_are_no_ops(self) -> None: + self.assertEqual(ResourcePath.mremove([]), {}) + self.assertEqual(ResourcePath.mexists([]), {}) + self.assertEqual(ResourcePath.mtransfer("copy", []), {}) + + @unittest.mock.patch.object(FileResourcePath, "_chunk_size", 1) + def test_removal_failure_does_not_abandon_the_rest(self) -> None: + uris = self._make_files("f", 20) + # Remove one out from under the batch so that its own removal raises. + uris[1].remove() + + results = ResourcePath.mremove(uris, do_raise=False) + + self.assertEqual(len(results), len(uris)) + self.assertFalse(results[uris[1]].success) + self.assertIsInstance(results[uris[1]].exception, FileNotFoundError) + for uri in uris[2:]: + self.assertTrue(results[uri].success, f"{uri} should have been removed") + self.assertFalse(uri.exists()) + + @unittest.mock.patch.object(FileResourcePath, "_chunk_size", 1) + def test_existence_check_reports_each_uri(self) -> None: + present = self._make_files("p", 20) + absent = self.tmpdir.join("gone.txt") + + results = ResourcePath.mexists([*present, absent]) + + self.assertEqual(len(results), len(present) + 1) + self.assertTrue(all(results[uri] for uri in present)) + self.assertFalse(results[absent]) + + @unittest.mock.patch.object(FileResourcePath, "_chunk_size", 1) + def test_existence_check_treats_an_error_as_missing(self) -> None: + uris = self._make_files("e", 20) + failing = uris[1].ospath + real_exists = FileResourcePath.exists + + def flaky(self: FileResourcePath) -> bool: + if self.ospath == failing: + raise PermissionError("cannot stat") + return real_exists(self) + + with unittest.mock.patch.object(FileResourcePath, "exists", flaky): + results = ResourcePath.mexists(uris) + + # The failure is reported as absent and the rest are still checked. + self.assertFalse(results[uris[1]]) + self.assertTrue(all(results[uri] for uri in uris if uri != uris[1])) + + def test_transfer_of_many_files(self) -> None: + sources = self._make_files("src", 50) + for i, src in enumerate(sources): + src.write(f"{i}".encode()) + destinations = [self.tmpdir.join(f"dest{n}.txt") for n in range(len(sources))] + + results = ResourcePath.mtransfer("copy", zip(sources, destinations, strict=True)) + + self.assertEqual(len(results), len(sources)) + self.assertTrue(all(res.success for res in results.values())) + for i, dest in enumerate(destinations): + self.assertEqual(dest.read().decode(), str(i)) + + def test_transfer_failure_does_not_abandon_the_rest(self) -> None: + sources = self._make_files("s", 20) + destinations = [self.tmpdir.join(f"d{n}.txt") for n in range(len(sources))] + # An existing target fails when overwriting is not allowed. + destinations[1].write(b"in the way") + + results = ResourcePath.mtransfer("copy", zip(sources, destinations, strict=True), do_raise=False) + + self.assertEqual(len(results), len(sources)) + self.assertFalse(results[destinations[1]].success) + for dest in destinations[2:]: + self.assertTrue(results[dest].success, f"{dest} should have been written") + self.assertTrue(dest.exists()) + + def test_transfer_registers_undo_actions(self) -> None: + sources = self._make_files("u", 20) + destinations = [self.tmpdir.join(f"undo{n}.txt") for n in range(len(sources))] + transaction = _RecordingTransaction() + + results = ResourcePath.mtransfer( + "copy", zip(sources, destinations, strict=True), overwrite=True, transaction=transaction + ) + + self.assertTrue(all(res.success for res in results.values())) + # Every transfer must be undoable by the caller that supplied the + # transaction, no matter which worker performed it. + self.assertEqual(len(transaction.undone), len(sources)) + + for undo in transaction.undone: + undo() + self.assertFalse(any(dest.exists() for dest in destinations)) + + +class _RecordingTransaction: + """Transaction that collects the undo actions registered against it.""" + + def __init__(self) -> None: + self.undone: list[Callable[[], Any]] = [] + self._lock = threading.Lock() + + @contextlib.contextmanager + def undoWith(self, name: str, undoFunc: Callable, *args: Any, **kwargs: Any) -> Iterator[None]: + yield None + with self._lock: + self.undone.append(functools.partial(undoFunc, *args, **kwargs)) @contextlib.contextmanager diff --git a/tests/test_s3.py b/tests/test_s3.py index 9791cf41..b8b9fb64 100644 --- a/tests/test_s3.py +++ b/tests/test_s3.py @@ -11,6 +11,7 @@ import datetime import os +import sys import time import unittest from inspect import signature @@ -22,6 +23,19 @@ from lsst.resources.s3utils import clean_test_environment_for_s3 from lsst.resources.tests import GenericReadWriteTestCase, GenericTestCase +# moto mocks S3 by patching botocore in this process, and its stubber reloads +# the module holding a backend's URL table on every request. Concurrent +# reloads of the same module race, and the loser raises ImportError, which a +# bulk operation reports as a URI that could not be reached. Only a +# free-threaded interpreter issues those requests at the same time, and which +# request loses varies from run to run, so every test that drives requests in +# parallel is affected rather than one in particular. +_MOTO_IS_THREAD_SAFE = getattr(sys, "_is_gil_enabled", lambda: True)() +skip_if_moto_races = unittest.skipIf( + not _MOTO_IS_THREAD_SAFE, + "moto reloads modules per request, which races without the GIL", +) + try: import boto3 import botocore @@ -274,31 +288,18 @@ def test_fsspec_constructor(self) -> None: def test_fsspec(self) -> None: raise unittest.SkipTest("fsspec s3fs incompatible with moto") - @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) - @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "threads"}) + # These drive many S3 requests in parallel, so under a free-threaded + # interpreter they are the tests moto's reload race shows up in. + @skip_if_moto_races def test_mexists(self) -> None: - """Test mexists with override executor pool. - - moto does not work with process pool. - """ super().test_mexists() - @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) - @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "threads"}) + @skip_if_moto_races def test_mtransfer(self) -> None: - """Test mtransfer with override executor pool. - - moto does not work with process pool. - """ super().test_mtransfer() - @unittest.mock.patch("lsst.resources._resourcePath._POOL_EXECUTOR_CLASS", None) - @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_EXECUTOR": "threads"}) + @skip_if_moto_races def test_local_mtransfer(self) -> None: - """Test local mtransfer with override executor pool. - - moto does not work with process pool. - """ super().test_local_mtransfer() diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..74988710 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,88 @@ +# This file is part of lsst-resources. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# Use of this source code is governed by a 3-clause BSD-style +# license that can be found in the LICENSE file. + +import os +import unittest +import unittest.mock + +from lsst.resources import ResourcePath +from lsst.resources.file import FileResourcePath +from lsst.resources.s3 import S3ResourcePath +from lsst.resources.utils import ( + MAX_WORKERS, + _get_configured_num_workers, + _get_default_num_workers, + _get_num_workers, +) + + +def _clear_worker_caches() -> None: + """Discard memoized worker-count lookups.""" + _get_configured_num_workers.cache_clear() + _get_default_num_workers.cache_clear() + + +class NumWorkersTestCase(unittest.TestCase): + """Tests for the worker-count calculation.""" + + def setUp(self) -> None: + _clear_worker_caches() + + def tearDown(self) -> None: + _clear_worker_caches() + + @unittest.mock.patch.dict(os.environ, {}, clear=False) + def test_default_is_capped(self) -> None: + os.environ.pop("LSST_RESOURCES_NUM_WORKERS", None) + _clear_worker_caches() + self.assertLessEqual(_get_num_workers(), MAX_WORKERS) + self.assertEqual(_get_num_workers(2), 2) + + @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_NUM_WORKERS": "99"}) + def test_explicit_request_bypasses_cap(self) -> None: + _clear_worker_caches() + self.assertEqual(_get_num_workers(), 99) + self.assertEqual(_get_num_workers(2), 99) + + def test_docstring_is_present(self) -> None: + # An f-string in the leading position is not a docstring. + self.assertIsNotNone(_get_num_workers.__doc__) + + +class WorkerCapTestCase(unittest.TestCase): + """Tests for per-scheme worker caps.""" + + def setUp(self) -> None: + _clear_worker_caches() + + def tearDown(self) -> None: + _clear_worker_caches() + + def test_schemes_share_the_default_cap(self) -> None: + self.assertEqual(ResourcePath._max_workers, MAX_WORKERS) + self.assertEqual(FileResourcePath._max_workers, MAX_WORKERS) + self.assertEqual(S3ResourcePath._max_workers, MAX_WORKERS) + + @unittest.mock.patch.dict(os.environ, {}, clear=False) + @unittest.mock.patch.object(FileResourcePath, "_max_workers", 2) + def test_an_overridden_cap_is_honored(self) -> None: + os.environ.pop("LSST_RESOURCES_NUM_WORKERS", None) + _clear_worker_caches() + self.assertEqual(_get_num_workers(FileResourcePath._max_workers), 2) + + @unittest.mock.patch.dict(os.environ, {"LSST_RESOURCES_NUM_WORKERS": "99"}) + def test_explicit_request_overrides_the_scheme_cap(self) -> None: + _clear_worker_caches() + self.assertEqual(_get_num_workers(S3ResourcePath._max_workers), 99) + + +if __name__ == "__main__": + unittest.main()