lsst-resources 29.0.0rc7__py3-none-any.whl → 29.2025.4600__py3-none-any.whl

lsst/resources/s3.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 
 __all__ = ("S3ResourcePath",)
 
+import concurrent.futures
 import contextlib
 import io
 import logging
@@ -20,17 +21,19 @@ import os
 import re
 import sys
 import threading
+from collections import defaultdict
 from collections.abc import Iterable, Iterator
 from functools import cache, cached_property
 from typing import IO, TYPE_CHECKING, cast
 
 from botocore.exceptions import ClientError
 
+from lsst.utils.iteration import chunk_iterable
 from lsst.utils.timer import time_this
 
 from ._resourceHandles._baseResourceHandle import ResourceHandleProtocol
 from ._resourceHandles._s3ResourceHandle import S3ResourceHandle
-from ._resourcePath import ResourcePath
+from ._resourcePath import _EXECUTOR_TYPE, MBulkResult, ResourcePath, _get_executor_class, _patch_environ
 from .s3utils import (
     _get_s3_connection_parameters,
     _s3_disable_bucket_validation,
@@ -44,6 +47,7 @@ from .s3utils import (
     s3CheckFileExists,
     translate_client_error,
 )
+from .utils import _get_num_workers
 
 try:
     from boto3.s3.transfer import TransferConfig  # type: ignore
@@ -164,6 +168,14 @@ class S3ResourcePath(ResourcePath):
 
         return use_threads
 
+    @contextlib.contextmanager
+    def _use_threads_temp_override(self, multithreaded: bool) -> Iterator:
+        """Temporarily override the value of use_threads."""
+        original = self.use_threads
+        self.use_threads = multithreaded
+        yield
+        self.use_threads = original
+
     @property
     def _transfer_config(self) -> TransferConfig:
         if self.use_threads is None:
@@ -212,7 +224,9 @@ class S3ResourcePath(ResourcePath):
         return bucket
 
     @classmethod
-    def _mexists(cls, uris: Iterable[ResourcePath]) -> dict[ResourcePath, bool]:
+    def _mexists(
+        cls, uris: Iterable[ResourcePath], *, num_workers: int | None = None
+    ) -> dict[ResourcePath, bool]:
         # Force client to be created for each profile before creating threads.
         profiles = set[str | None]()
         for path in uris:
@@ -222,7 +236,125 @@ class S3ResourcePath(ResourcePath):
         for profile in profiles:
             getS3Client(profile)
 
-        return super()._mexists(uris)
+        return super()._mexists(uris, num_workers=num_workers)
+
+    @classmethod
+    def _mremove(cls, uris: Iterable[ResourcePath]) -> dict[ResourcePath, MBulkResult]:
+        # Delete multiple objects in one API call.
+        # Must group by profile and bucket.
+        grouped_uris: dict[tuple[str | None, str], list[S3ResourcePath]] = defaultdict(list)
+        for uri in uris:
+            uri = cast(S3ResourcePath, uri)
+            grouped_uris[uri._profile, uri._bucket].append(uri)
+
+        results: dict[ResourcePath, MBulkResult] = {}
+        for related_uris in grouped_uris.values():
+            # API requires no more than 1000 per call.
+            chunk_num = 0
+            chunks: list[tuple[ResourcePath, ...]] = []
+            key_to_uri: dict[str, ResourcePath] = {}
+            for chunk in chunk_iterable(related_uris, chunk_size=1_000):
+                for uri in chunk:
+                    key = uri.relativeToPathRoot
+                    key_to_uri[key] = uri
+                    # Default to assuming everything worked.
+                    results[uri] = MBulkResult(True, None)
+                chunk_num += 1
+                chunks.append(chunk)
+
+            # Bulk remove.
+            with time_this(
+                log,
+                msg="Bulk delete; %d chunk%s; totalling %d dataset%s",
+                args=(
+                    len(chunks),
+                    "s" if len(chunks) != 1 else "",
+                    len(related_uris),
+                    "s" if len(related_uris) != 1 else "",
+                ),
+            ):
+                errored = cls._mremove_select(chunks)
+
+            # Update with error information.
+            results.update(errored)
+
+        return results
+
+    @classmethod
+    def _mremove_select(cls, chunks: list[tuple[ResourcePath, ...]]) -> dict[ResourcePath, MBulkResult]:
+        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)
+
+    @classmethod
+    def _mremove_with_pool(
+        cls,
+        pool_executor_class: _EXECUTOR_TYPE,
+        chunks: list[tuple[ResourcePath, ...]],
+        *,
+        num_workers: int | None = None,
+    ) -> 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())
+        results: dict[ResourcePath, MBulkResult] = {}
+        with pool_executor_class(max_workers=max_workers) as remove_executor:
+            future_remove = {
+                remove_executor.submit(cls._delete_objects_wrapper, chunk): i
+                for i, chunk in enumerate(chunks)
+            }
+            for future in concurrent.futures.as_completed(future_remove):
+                try:
+                    results.update(future.result())
+                except Exception as e:
+                    # The chunk utterly failed.
+                    chunk = chunks[future_remove[future]]
+                    for uri in chunk:
+                        results[uri] = MBulkResult(False, e)
+        return results
+
+    @classmethod
+    def _delete_objects_wrapper(cls, uris: tuple[ResourcePath, ...]) -> dict[ResourcePath, MBulkResult]:
+        """Convert URIs to keys and call low-level API."""
+        if not uris:
+            return {}
+        keys: list[dict[str, str]] = []
+        key_to_uri: dict[str, ResourcePath] = {}
+        for uri in uris:
+            key = uri.relativeToPathRoot
+            key_to_uri[key] = uri
+            keys.append({"Key": key})
+
+        first_uri = cast(S3ResourcePath, uris[0])
+        results = cls._delete_related_objects(first_uri.client, first_uri._bucket, keys)
+
+        # Remap error object keys to uris.
+        return {key_to_uri[key]: result for key, result in results.items()}
+
+    @classmethod
+    @backoff.on_exception(backoff.expo, retryable_io_errors, max_time=max_retry_time)
+    def _delete_related_objects(
+        cls, client: boto3.client, bucket: str, keys: list[dict[str, str]]
+    ) -> dict[str, MBulkResult]:
+        # Delete multiple objects from the same bucket, allowing for backoff
+        # retry.
+        response = client.delete_objects(Bucket=bucket, Delete={"Objects": keys, "Quiet": True})
+        # Use Quiet mode so we assume everything worked unless told otherwise.
+        # Only returning errors -- indexed by Key name.
+        errors: dict[str, MBulkResult] = {}
+        for errored_key in response.get("Errors", []):
+            errors[errored_key["Key"]] = MBulkResult(
+                False, ClientError({"Error": errored_key}, f"delete_objects: {errored_key['Key']}")
+            )
+        return errors
 
     @backoff.on_exception(backoff.expo, retryable_io_errors, max_time=max_retry_time)
     def exists(self) -> bool:
@@ -345,18 +477,33 @@ class S3ResourcePath(ResourcePath):
 
         return s3, f"{self._bucket}/{self.relativeToPathRoot}"
 
-    def _as_local(self) -> tuple[str, bool]:
+    @contextlib.contextmanager
+    def _as_local(
+        self, multithreaded: bool = True, tmpdir: ResourcePath | None = None
+    ) -> Iterator[ResourcePath]:
         """Download object from S3 and place in temporary directory.
 
+        Parameters
+        ----------
+        multithreaded : `bool`, optional
+            If `True` the transfer will be allowed to attempt to improve
+            throughput by using parallel download streams. This may of no
+            effect if the URI scheme does not support parallel streams or
+            if a global override has been applied. If `False` parallel
+            streams will be disabled.
+        tmpdir : `ResourcePath` or `None`, optional
+            Explicit override of the temporary directory to use for remote
+            downloads.
+
         Returns
         -------
-        path : `str`
-            Path to local temporary file.
-        temporary : `bool`
-            Always returns `True`. This is always a temporary file.
+        local_uri : `ResourcePath`
+            A URI to a local POSIX file corresponding to a local temporary
+            downloaded copy of the resource.
         """
         with (
-            ResourcePath.temporary_uri(suffix=self.getExtension(), delete=False) as tmp_uri,
+            ResourcePath.temporary_uri(prefix=tmpdir, suffix=self.getExtension(), delete=True) as tmp_uri,
+            self._use_threads_temp_override(multithreaded),
             time_this(log, msg="Downloading %s to local file", args=(self,)),
         ):
             progress = (
@@ -366,7 +513,7 @@ class S3ResourcePath(ResourcePath):
             )
             with tmp_uri.open("wb") as tmpFile:
                 self._download_file(tmpFile, progress)
-        return tmp_uri.ospath, True
+            yield tmp_uri
 
     @backoff.on_exception(backoff.expo, all_retryable_errors, max_time=max_retry_time)
     def _upload_file(self, local_file: ResourcePath, progress: ProgressPercentage | None) -> None:
@@ -397,7 +544,7 @@ class S3ResourcePath(ResourcePath):
         try:
             self.client.copy_object(CopySource=copy_source, Bucket=self._bucket, Key=self.relativeToPathRoot)
         except (self.client.exceptions.NoSuchKey, self.client.exceptions.NoSuchBucket) as err:
-            raise FileNotFoundError("No such resource to transfer: {self}") from err
+            raise FileNotFoundError(f"No such resource to transfer: {src} -> {self}") from err
         except ClientError as err:
             translate_client_error(err, self)
             raise
@@ -408,6 +555,7 @@ class S3ResourcePath(ResourcePath):
         transfer: str = "copy",
         overwrite: bool = False,
         transaction: TransactionProtocol | None = None,
+        multithreaded: bool = True,
     ) -> None:
         """Transfer the current resource to an S3 bucket.
 
@@ -422,6 +570,12 @@ class S3ResourcePath(ResourcePath):
             Allow an existing file to be overwritten. Defaults to `False`.
         transaction : `~lsst.resources.utils.TransactionProtocol`, optional
             Currently unused.
+        multithreaded : `bool`, optional
+            If `True` the transfer will be allowed to attempt to improve
+            throughput by using parallel download streams. This may of no
+            effect if the URI scheme does not support parallel streams or
+            if a global override has been applied. If `False` parallel
+            streams will be disabled.
         """
         # Fail early to prevent delays if remote resources are requested
         if transfer not in self.transferModes:
@@ -457,22 +611,25 @@ class S3ResourcePath(ResourcePath):
         timer_msg = "Transfer from %s to %s"
         timer_args = (src, self)
 
-        if isinstance(src, type(self)):
-            # Looks like an S3 remote uri so we can use direct copy
-            # note that boto3.resource.meta.copy is cleverer than the low
-            # level copy_object
+        if isinstance(src, type(self)) and self.client == src.client:
+            # Looks like an S3 remote uri so we can use direct copy.
+            # This only works if the source and destination are using the same
+            # S3 endpoint and profile.
             with time_this(log, msg=timer_msg, args=timer_args):
                 self._copy_from(src)
 
         else:
             # Use local file and upload it
-            with src.as_local() as local_uri:
+            with src.as_local(multithreaded=multithreaded) as local_uri:
                 progress = (
                     ProgressPercentage(local_uri, file_for_msg=src, msg="Uploading:")
                     if log.isEnabledFor(ProgressPercentage.log_level)
                     else None
                 )
-                with time_this(log, msg=timer_msg, args=timer_args):
+                with (
+                    time_this(log, msg=timer_msg, args=timer_args),
+                    self._use_threads_temp_override(multithreaded),
+                ):
                     self._upload_file(local_uri, progress)
 
         # This was an explicit move requested from a remote resource

lsst/resources/s3utils.py

@@ -53,6 +53,7 @@ except ImportError:
 
 from ._resourcePath import ResourcePath
 from .location import Location
+from .utils import _get_num_workers
 
 # https://pypi.org/project/backoff/
 try:
@@ -246,7 +247,13 @@ def _s3_disable_bucket_validation(client: boto3.client) -> None:
 @functools.lru_cache
 def _get_s3_client(endpoint_config: _EndpointConfig, skip_validation: bool) -> boto3.client:
     # Helper function to cache the client for this endpoint
-    config = botocore.config.Config(read_timeout=180, retries={"mode": "adaptive", "max_attempts": 10})
+    # boto seems to assume it will always have at least 10 available.
+    max_pool_size = max(_get_num_workers(), 10)
+    config = botocore.config.Config(
+        read_timeout=180,
+        max_pool_connections=max_pool_size,
+        retries={"mode": "adaptive", "max_attempts": 10},
+    )
 
     session = boto3.Session(profile_name=endpoint_config.profile)
 

lsst/resources/schemeless.py

@@ -105,13 +105,16 @@ class SchemelessResourcePath(FileResourcePath):
             return stat.S_ISDIR(status.st_mode)
         return self.dirLike
 
-    def relative_to(self, other: ResourcePath) -> str | None:
+    def relative_to(self, other: ResourcePath, walk_up: bool = False) -> str | None:
         """Return the relative path from this URI to the other URI.
 
         Parameters
         ----------
         other : `ResourcePath`
             URI to use to calculate the relative path.
+        walk_up : `bool`, optional
+            Control whether "``..``" can be used to resolve a relative path.
+            Default is `False`. Can not be `True` on Python version 3.11.
 
         Returns
         -------
@@ -146,8 +149,8 @@ class SchemelessResourcePath(FileResourcePath):
             raise RuntimeError(f"Unexpected combination of {child}.relative_to({other}).")
 
         if child is None:
-            return super().relative_to(other)
-        return child.relative_to(other)
+            return super().relative_to(other, walk_up=walk_up)
+        return child.relative_to(other, walk_up=walk_up)
 
     @classmethod
     def _fixupPathUri(

lsst/resources/tests.py

@@ -17,6 +17,8 @@ import os
 import pathlib
 import random
 import string
+import sys
+import tempfile
 import unittest
 import urllib.parse
 import uuid
@@ -60,18 +62,18 @@ def _check_open(
     """
     text_content = "abcdefghijklmnopqrstuvwxyz🙂"
     bytes_content = uuid.uuid4().bytes
-    content_by_mode_suffix = {
+    content_by_mode_suffix: dict[str, str | bytes] = {
         "": text_content,
         "t": text_content,
         "b": bytes_content,
     }
-    empty_content_by_mode_suffix = {
+    empty_content_by_mode_suffix: dict[str, str | bytes] = {
         "": "",
         "t": "",
         "b": b"",
     }
     # To appease mypy
-    double_content_by_mode_suffix = {
+    double_content_by_mode_suffix: dict[str, str | bytes] = {
         "": text_content + text_content,
         "t": text_content + text_content,
         "b": bytes_content + bytes_content,
@@ -142,6 +144,16 @@ def _check_open(
             content_read = read_buffer.read()
             test_case.assertEqual(len(content_read), 0, f"Read: {content_read!r}, expected empty.")
 
+        # Write multiple chunks with flushing to ensure that any handles that
+        # cache without flushing work properly.
+        n = 3
+        with uri.open("w" + mode_suffix, **kwargs) as write_buffer:
+            for _ in range(n):
+                write_buffer.write(content)
+                write_buffer.flush()
+        with uri.open("r" + mode_suffix, **kwargs) as read_buffer:
+            test_case.assertEqual(read_buffer.read(), content * n)
+
         # Write two copies of the content, overwriting the single copy there.
         with uri.open("w" + mode_suffix, **kwargs) as write_buffer:
             write_buffer.write(double_content)
@@ -364,6 +376,19 @@ class GenericTestCase(_GenericTestCase):
         parent = ResourcePath("d/e.txt", forceAbsolute=False)
         self.assertIsNone(child.relative_to(parent), f"{child}.relative_to({parent})")
 
+        # Allow .. in response.
+        child = ResourcePath(self._make_uri("a/b/c/d.txt"), forceAbsolute=False)
+        parent = ResourcePath(self._make_uri("a/b/d/e/"), forceAbsolute=False)
+        self.assertIsNone(child.relative_to(parent), f"{child}.relative_to({parent})")
+
+        if sys.version_info >= (3, 12, 0):
+            # Fails on python 3.11.
+            self.assertEqual(
+                child.relative_to(parent, walk_up=True),
+                "../../c/d.txt",
+                f"{child}.relative_to({parent}, walk_up=True)",
+            )
+
     def test_parents(self) -> None:
         """Test of splitting and parent walking."""
         parent = ResourcePath(self._make_uri("somedir"), forceDirectory=True)
@@ -381,6 +406,14 @@ class GenericTestCase(_GenericTestCase):
         self.assertEqual(child_file.parent().parent(), parent)
         self.assertEqual(child_subdir.dirname(), child_subdir)
 
+        # Make sure that the parent doesn't retain any fragment from the
+        # child.
+        child_fragment = child_subdir.join("a.txt#fragment")
+        self.assertEqual(child_fragment.fragment, "fragment")
+        fragment_parent = child_fragment.parent()
+        self.assertEqual(fragment_parent.fragment, "")
+        self.assertTrue(str(fragment_parent).endswith("/"))
+
     def test_escapes(self) -> None:
         """Special characters in file paths."""
         src = self.root_uri.join("bbb/???/test.txt")
@@ -485,6 +518,13 @@ class GenericTestCase(_GenericTestCase):
         self.assertEqual(fnew3.fragment, "fragment")
         self.assertEqual(fnew3.basename(), "b.txt", msg=f"Got: {fnew3._uri}")
 
+        # Check that fragment on the directory is dropped on join.
+        frag_dir = add_dir.join("subdir/#dir_fragment")
+        self.assertEqual(frag_dir.fragment, "dir_fragment")
+        fnew4 = frag_dir.join("a.txt")
+        self.assertEqual(fnew4.fragment, "")
+        self.assertTrue(str(fnew4).endswith("/a.txt"))
+
         # Join a resource path.
         subpath = ResourcePath("a/b.txt#fragment2", forceAbsolute=False, forceDirectory=False)
         fnew3 = root.join(subpath)
@@ -556,6 +596,10 @@ class GenericReadWriteTestCase(_GenericTestCase):
 
     transfer_modes: tuple[str, ...] = ("copy", "move")
     testdir: str | None = None
+    # Number of files to use for mremove() testing to ensure difference code
+    # paths are hit. Do not want to generically use many files for schemes
+    # where it makes no difference.
+    n_mremove_files: int = 15
 
     def setUp(self) -> None:
         if self.scheme is None:
@@ -702,6 +746,37 @@ class GenericReadWriteTestCase(_GenericTestCase):
         with self.assertRaises(FileNotFoundError):
             dest.transfer_from(src, "auto")
 
+    def test_mtransfer(self) -> None:
+        n_files = 10
+        sources = [self.tmpdir.join(f"test{n}.txt") for n in range(n_files)]
+        destinations = [self.tmpdir.join(f"dest_test{n}.txt") for n in range(n_files)]
+
+        for i, src in enumerate(sources):
+            content = f"{i}\nContent is some content\nwith something to say\n\n"
+            src.write(content.encode())
+
+        results = ResourcePath.mtransfer("copy", zip(sources, destinations, strict=True))
+        self.assertTrue(all(res.success for res in results.values()))
+        self.assertTrue(all(dest.exists() for dest in results))
+
+        for i, dest in enumerate(destinations):
+            new_content = dest.read().decode()
+            self.assertTrue(new_content.startswith(f"{i}\n"))
+
+        # Overwrite should work.
+        results = ResourcePath.mtransfer("copy", zip(sources, destinations, strict=True), overwrite=True)
+
+        # Overwrite failure.
+        results = ResourcePath.mtransfer(
+            "copy", zip(sources, destinations, strict=True), overwrite=False, do_raise=False
+        )
+        self.assertFalse(all(res.success for res in results.values()))
+
+        with self.assertRaises(ExceptionGroup):
+            results = ResourcePath.mtransfer(
+                "copy", zip(sources, destinations, strict=True), overwrite=False, do_raise=True
+            )
+
     def test_local_transfer(self) -> None:
         """Test we can transfer to and from local file."""
         remote_src = self.tmpdir.join("src.json")
@@ -762,6 +837,48 @@ class GenericReadWriteTestCase(_GenericTestCase):
             with self.root_uri.as_local() as local_uri:
                 pass
 
+        if not src.isLocal:
+            # as_local tmpdir can not be a remote resource.
+            with self.assertRaises(ValueError):
+                with src.as_local(tmpdir=self.root_uri) as local_uri:
+                    pass
+
+            # tmpdir is ignored for local file.
+            with tempfile.TemporaryDirectory() as tmpdir:
+                temp_dir = ResourcePath(tmpdir, forceDirectory=True)
+                with src.as_local(tmpdir=temp_dir) as local_uri:
+                    self.assertEqual(local_uri.dirname(), temp_dir)
+                    self.assertTrue(local_uri.exists())
+
+    def test_local_mtransfer(self) -> None:
+        """Check that bulk transfer to/from local works."""
+        # Create remote resources
+        n_files = 10
+        sources = [self.tmpdir.join(f"test{n}.txt") for n in range(n_files)]
+
+        for i, src in enumerate(sources):
+            content = f"{i}\nContent is some content\nwith something to say\n\n"
+            src.write(content.encode())
+
+        # Potentially remote to local.
+        with tempfile.TemporaryDirectory() as tmpdir:
+            temp_dir = ResourcePath(tmpdir, forceDirectory=True)
+            destinations = [temp_dir.join(f"dest_test{n}.txt") for n in range(n_files)]
+
+            results = ResourcePath.mtransfer("copy", zip(sources, destinations, strict=True))
+            self.assertTrue(all(res.success for res in results.values()))
+            self.assertTrue(all(dest.exists() for dest in results))
+
+            # Overwrite should work.
+            results = ResourcePath.mtransfer("copy", zip(sources, destinations, strict=True), overwrite=True)
+
+            # Now reverse so local to potentially remote.
+            for src in sources:
+                src.remove()
+            results = ResourcePath.mtransfer("copy", zip(destinations, sources, strict=True), overwrite=False)
+            self.assertTrue(all(res.success for res in results.values()))
+            self.assertTrue(all(dest.exists() for dest in results))
+
     def test_walk(self) -> None:
         """Walk a directory hierarchy."""
         root = self.tmpdir.join("walk/")
@@ -949,22 +1066,33 @@ class GenericReadWriteTestCase(_GenericTestCase):
         # A file that is not there.
         file = root.join("config/basic/butler.yaml")
 
-        # Create some files.
-        expected_files = {
-            "dir1/a.yaml",
-            "dir1/b.yaml",
-            "dir2/e.yaml",
-        }
-        expected_uris = {root.join(f) for f in expected_files}
+        # Create some files. Most schemes the code paths do not change for 10
+        # vs 1000 files but in some schemes it does.
+        expected_files = [f"dir1/f{n}.yaml" for n in range(self.n_mremove_files)]
+        expected_uris = [root.join(f) for f in expected_files]
         for uri in expected_uris:
             uri.write(b"")
             self.assertTrue(uri.exists())
-        expected_uris.add(file)
+        expected_uris.append(file)
 
-        multi = ResourcePath.mexists(expected_uris)
+        # Force to run with fewer workers than there are files.
+        multi = ResourcePath.mexists(expected_uris, num_workers=3)
 
         for uri, is_there in multi.items():
             if uri == file:
                 self.assertFalse(is_there)
             else:
                 self.assertTrue(is_there)
+
+        # Clean up. Unfortunately POSIX raises a FileNotFoundError but
+        # S3 boto does not complain if there is no key.
+        ResourcePath.mremove(expected_uris, do_raise=False)
+
+        # Check they were really removed.
+        multi = ResourcePath.mexists(expected_uris, num_workers=3)
+        for uri, is_there in multi.items():
+            self.assertFalse(is_there)
+
+        # Clean up a subset of files that are already gone, but this can
+        # trigger a different code path.
+        ResourcePath.mremove(expected_uris[:5], do_raise=False)