wandb 0.22.0__py3-none-win32.whl → 0.22.2__py3-none-win32.whl

wandb/sdk/artifacts/storage_handlers/azure_handler.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from pathlib import PurePosixPath
 from types import ModuleType
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import ParseResult, parse_qsl, urlparse
 
 import wandb
@@ -20,17 +20,22 @@ if TYPE_CHECKING:
     import azure.storage.blob  # type: ignore
 
     from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_file_cache import ArtifactFileCache
 
 
 class AzureHandler(StorageHandler):
+    _scheme: str
+    _cache: ArtifactFileCache
+
+    def __init__(self, scheme: str = "https") -> None:
+        self._scheme = scheme
+        self._cache = get_artifact_file_cache()
+
     def can_handle(self, parsed_url: ParseResult) -> bool:
-        return parsed_url.scheme == "https" and parsed_url.netloc.endswith(
+        return parsed_url.scheme == self._scheme and parsed_url.netloc.endswith(
             ".blob.core.windows.net"
         )
 
-    def __init__(self, scheme: str | None = None) -> None:
-        self._cache = get_artifact_file_cache()
-
     def load_path(
         self,
         manifest_entry: ArtifactManifestEntry,
@@ -101,7 +106,7 @@ class AzureHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         account_url, container_name, blob_name, query = self._parse_uri(path)
         path = URIStr(f"{account_url}/{container_name}/{blob_name}")
 

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import time
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import ParseResult, urlparse
 
 from wandb import util
@@ -19,6 +19,7 @@ if TYPE_CHECKING:
     import google.cloud.storage as gcs_module  # type: ignore
 
     from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_file_cache import ArtifactFileCache
 
 
 class _GCSIsADirectoryError(Exception):
@@ -26,10 +27,12 @@ class _GCSIsADirectoryError(Exception):
 
 
 class GCSHandler(StorageHandler):
+    _scheme: str
     _client: gcs_module.client.Client | None
+    _cache: ArtifactFileCache
 
-    def __init__(self, scheme: str | None = None) -> None:
-        self._scheme = scheme or "gs"
+    def __init__(self, scheme: str = "gs") -> None:
+        self._scheme = scheme
         self._client = None
         self._cache = get_artifact_file_cache()
 
@@ -111,7 +114,7 @@ class GCSHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         self.init_gcs()
         assert self._client is not None  # mypy: unwraps optionality
 
@@ -131,7 +134,7 @@ class GCSHandler(StorageHandler):
             raise ValueError(f"Object does not exist: {path}#{version}")
         multi = obj is None
         if multi:
-            start_time = time.time()
+            start_time = time.monotonic()
             termlog(
                 f'Generating checksum for up to {max_objects} objects with prefix "{key}"... ',
                 newline=False,
@@ -148,7 +151,7 @@ class GCSHandler(StorageHandler):
             if not obj.name.endswith("/")
         ]
         if start_time is not None:
-            termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+            termlog("Done. %.1fs" % (time.monotonic() - start_time), prefix=False)
         if len(entries) > max_objects:
             raise ValueError(
                 f"Exceeded {max_objects} objects tracked, pass max_objects to add_reference"

wandb/sdk/artifacts/storage_handlers/http_handler.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import ParseResult
 
 from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache
@@ -18,11 +18,16 @@ if TYPE_CHECKING:
     from requests.structures import CaseInsensitiveDict
 
     from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_file_cache import ArtifactFileCache
 
 
 class HTTPHandler(StorageHandler):
-    def __init__(self, session: requests.Session, scheme: str | None = None) -> None:
-        self._scheme = scheme or "http"
+    _scheme: str
+    _cache: ArtifactFileCache
+    _session: requests.Session
+
+    def __init__(self, session: requests.Session, scheme: str = "http") -> None:
+        self._scheme = scheme
         self._cache = get_artifact_file_cache()
         self._session = session
 
@@ -75,7 +80,7 @@ class HTTPHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         name = name or os.path.basename(path)
         if not checksum:
             return [ArtifactManifestEntry(path=name, ref=path, digest=path)]

wandb/sdk/artifacts/storage_handlers/local_file_handler.py

@@ -6,7 +6,7 @@ import os
 import shutil
 import time
 from pathlib import Path
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import ParseResult
 
 from wandb import util
@@ -20,17 +20,21 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 
 if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_file_cache import ArtifactFileCache
 
 
 class LocalFileHandler(StorageHandler):
     """Handles file:// references."""
 
-    def __init__(self, scheme: str | None = None) -> None:
+    _scheme: str
+    _cache: ArtifactFileCache
+
+    def __init__(self, scheme: str = "file") -> None:
         """Track files or directories on a local filesystem.
 
         Expand directories to create an entry for each file contained.
         """
-        self._scheme = scheme or "file"
+        self._scheme = scheme
         self._cache = get_artifact_file_cache()
 
     def can_handle(self, parsed_url: ParseResult) -> bool:
@@ -75,7 +79,7 @@ class LocalFileHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         local_path = util.local_file_uri_to_path(path)
         max_objects = max_objects or DEFAULT_MAX_OBJECTS
         # We have a single file or directory
@@ -95,7 +99,7 @@ class LocalFileHandler(StorageHandler):
 
         if os.path.isdir(local_path):
             i = 0
-            start_time = time.time()
+            start_time = time.monotonic()
             if checksum:
                 termlog(
                     f'Generating checksum for up to {max_objects} files in "{local_path}"... ',
@@ -126,7 +130,7 @@ class LocalFileHandler(StorageHandler):
                     )
                     entries.append(entry)
             if checksum:
-                termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+                termlog("Done. %.1fs" % (time.monotonic() - start_time), prefix=False)
         elif os.path.isfile(local_path):
             name = name or os.path.basename(local_path)
             entry = ArtifactManifestEntry(

wandb/sdk/artifacts/storage_handlers/multi_handler.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import urlparse
 
-from wandb.sdk.artifacts.storage_handler import StorageHandler
+from wandb.sdk.artifacts.storage_handler import StorageHandler, _BaseStorageHandler
 from wandb.sdk.lib.paths import FilePathStr, URIStr
 
 if TYPE_CHECKING:
@@ -13,8 +13,9 @@ if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
 
 
-class MultiHandler(StorageHandler):
+class MultiHandler(_BaseStorageHandler):
     _handlers: list[StorageHandler]
+    _default_handler: StorageHandler | None
 
     def __init__(
         self,
@@ -49,7 +50,7 @@ class MultiHandler(StorageHandler):
         name: str | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         handler = self._get_handler(path)
         return handler.store_path(
             artifact, path, name=name, checksum=checksum, max_objects=max_objects

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -6,7 +6,7 @@ import os
 import re
 import time
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import parse_qsl, urlparse
 
 from wandb import util
@@ -32,16 +32,18 @@ if TYPE_CHECKING:
     import boto3.session  # type: ignore
 
     from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_file_cache import ArtifactFileCache
 
 
 class S3Handler(StorageHandler):
-    _s3: boto3.resources.base.ServiceResource | None
     _scheme: str
+    _cache: ArtifactFileCache
+    _s3: boto3.resources.base.ServiceResource | None
 
-    def __init__(self, scheme: str | None = None) -> None:
-        self._scheme = scheme or "s3"
-        self._s3 = None
+    def __init__(self, scheme: str = "s3") -> None:
+        self._scheme = scheme
         self._cache = get_artifact_file_cache()
+        self._s3 = None
 
     def can_handle(self, parsed_url: ParseResult) -> bool:
         return parsed_url.scheme == self._scheme
@@ -160,7 +162,7 @@ class S3Handler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         self.init_boto()
         assert self._s3 is not None  # mypy: unwraps optionality
 
@@ -206,7 +208,7 @@ class S3Handler(StorageHandler):
             multi = True
 
         if multi:
-            start_time = time.time()
+            start_time = time.monotonic()
             termlog(
                 f'Generating checksum for up to {max_objects} objects in "{bucket}/{key}"... ',
                 newline=False,
@@ -227,7 +229,7 @@ class S3Handler(StorageHandler):
             if size(obj) > 0
         ]
         if start_time is not None:
-            termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+            termlog("Done. %.1fs" % (time.monotonic() - start_time), prefix=False)
         if len(entries) > max_objects:
             raise ValueError(
                 f"Exceeded {max_objects} objects tracked, pass max_objects to add_reference"

wandb/sdk/artifacts/storage_handlers/tracking_handler.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import urlparse
 
 from wandb.errors.term import termwarn
@@ -17,7 +17,9 @@ if TYPE_CHECKING:
 
 
 class TrackingHandler(StorageHandler):
-    def __init__(self, scheme: str | None = None) -> None:
+    _scheme: str
+
+    def __init__(self, scheme: str = "") -> None:
         """Track paths with no modification or special processing.
 
         Useful when paths being tracked are on file systems mounted at a standardized
@@ -26,7 +28,7 @@ class TrackingHandler(StorageHandler):
         For example, if the data to track is located on an NFS share mounted on
         `/data`, then it is sufficient to just track the paths.
         """
-        self._scheme = scheme or ""
+        self._scheme = scheme
 
     def can_handle(self, parsed_url: ParseResult) -> bool:
         return parsed_url.scheme == self._scheme
@@ -55,7 +57,7 @@ class TrackingHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         url = urlparse(path)
         if name is None:
             raise ValueError(

wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py

@@ -3,27 +3,29 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Literal
 from urllib.parse import urlparse
 
-import wandb
-from wandb import util
+from wandb._strutils import removeprefix
 from wandb.apis import PublicApi
 from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache
 from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
 from wandb.sdk.artifacts.storage_handler import StorageHandler
-from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id
+from wandb.sdk.lib.hashutil import b64_to_hex_id, hex_to_b64_id
 from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 
 if TYPE_CHECKING:
     from urllib.parse import ParseResult
 
     from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_file_cache import ArtifactFileCache
 
 
 class WBArtifactHandler(StorageHandler):
     """Handles loading and storing Artifact reference-type files."""
 
+    _scheme: Literal["wandb-artifact"]
+    _cache: ArtifactFileCache
     _client: PublicApi | None
 
     def __init__(self) -> None:
@@ -55,6 +57,8 @@ class WBArtifactHandler(StorageHandler):
         Returns:
             (os.PathLike): A path to the file represented by `index_entry`
         """
+        from wandb.sdk.artifacts.artifact import Artifact  # avoids circular import
+
         # We don't check for cache hits here. Since we have 0 for size (since this
         # is a cross-artifact reference which and we've made the choice to store 0
         # in the size field), we can't confirm if the file is complete. So we just
@@ -62,19 +66,17 @@ class WBArtifactHandler(StorageHandler):
         # check.
 
         # Parse the reference path and download the artifact if needed
-        artifact_id = util.host_from_path(manifest_entry.ref)
-        artifact_file_path = util.uri_from_path(manifest_entry.ref)
+        parsed = urlparse(manifest_entry.ref)
+        artifact_id = hex_to_b64_id(parsed.netloc)
+        artifact_file_path = removeprefix(str(parsed.path), "/")
 
-        dep_artifact = wandb.Artifact._from_id(
-            hex_to_b64_id(artifact_id), self.client.client
-        )
+        dep_artifact = Artifact._from_id(artifact_id, self.client.client)
         assert dep_artifact is not None
         link_target_path: URIStr | FilePathStr
         if local:
             link_target_path = dep_artifact.get_entry(artifact_file_path).download()
         else:
             link_target_path = dep_artifact.get_entry(artifact_file_path).ref_target()
-
         return link_target_path
 
     def store_path(
@@ -84,7 +86,7 @@ class WBArtifactHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         """Store the file or directory at the given path into the specified artifact.
 
         Recursively resolves the reference until the result is a concrete asset.
@@ -97,26 +99,27 @@ class WBArtifactHandler(StorageHandler):
             (list[ArtifactManifestEntry]): A list of manifest entries to store within
             the artifact
         """
+        from wandb.sdk.artifacts.artifact import Artifact  # avoids circular import
+
         # Recursively resolve the reference until a concrete asset is found
         # TODO: Consider resolving server-side for performance improvements.
-        iter_path: URIStr | FilePathStr | None = path
-        while iter_path is not None and urlparse(iter_path).scheme == self._scheme:
-            artifact_id = util.host_from_path(iter_path)
-            artifact_file_path = util.uri_from_path(iter_path)
-            target_artifact = wandb.Artifact._from_id(
-                hex_to_b64_id(artifact_id), self.client.client
-            )
+        curr_path: URIStr | FilePathStr | None = path
+        while curr_path and (parsed := urlparse(curr_path)).scheme == self._scheme:
+            artifact_id = hex_to_b64_id(parsed.netloc)
+            artifact_file_path = removeprefix(parsed.path, "/")
+
+            target_artifact = Artifact._from_id(artifact_id, self.client.client)
             assert target_artifact is not None
 
             entry = target_artifact.manifest.get_entry_by_path(artifact_file_path)
             assert entry is not None
-            iter_path = entry.ref
+            curr_path = entry.ref
 
         # Create the path reference
         assert target_artifact is not None
         assert target_artifact.id is not None
-        path = URIStr(
-            f"{self._scheme}://{b64_to_hex_id(B64MD5(target_artifact.id))}/{artifact_file_path}"
+        path = (
+            f"{self._scheme}://{b64_to_hex_id(target_artifact.id)}/{artifact_file_path}"
         )
 
         # Return the new entry

wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Literal
 
 import wandb
 from wandb import util
@@ -21,6 +21,8 @@ if TYPE_CHECKING:
 class WBLocalArtifactHandler(StorageHandler):
     """Handles loading and storing Artifact reference-type files."""
 
+    _scheme: Literal["wandb-client-artifact"]
+
     def __init__(self) -> None:
         self._scheme = "wandb-client-artifact"
 
@@ -43,7 +45,7 @@ class WBLocalArtifactHandler(StorageHandler):
         name: StrPath | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         """Store the file or directory at the given path within the specified artifact.
 
         Args:

wandb/sdk/artifacts/storage_policies/_multipart.py

@@ -0,0 +1,187 @@
+"""Helpers and constants for multipart upload and download."""
+
+from __future__ import annotations
+
+import logging
+import math
+import threading
+from concurrent.futures import FIRST_EXCEPTION, Executor, wait
+from dataclasses import dataclass, field
+from queue import Queue
+from typing import Any, Final, Iterator, Union
+
+from requests import Session
+from typing_extensions import TypeAlias, TypeIs, final
+
+from wandb import env
+from wandb.sdk.artifacts.artifact_file_cache import Opener
+
+logger = logging.getLogger(__name__)
+
+KiB: Final[int] = 1024
+MiB: Final[int] = 1024**2
+GiB: Final[int] = 1024**3
+TiB: Final[int] = 1024**4
+
+# AWS S3 max upload parts without having to make additional requests for extra parts
+MAX_PARTS = 1_000
+MIN_MULTI_UPLOAD_SIZE = 2 * GiB
+MAX_MULTI_UPLOAD_SIZE = 5 * TiB
+
+# Minimum size to switch to multipart download, same threshold as upload.
+MIN_MULTI_DOWNLOAD_SIZE = MIN_MULTI_UPLOAD_SIZE
+
+# Multipart download part size is same as multpart upload size, which is hard coded to 100MB.
+# https://github.com/wandb/wandb/blob/7b2a13cb8efcd553317167b823c8e52d8c3f7c4e/core/pkg/artifacts/saver.go#L496
+# https://docs.aws.amazon.com/AmazonS3/latest/userguide/optimizing-performance-guidelines.html#optimizing-performance-guidelines-get-range
+MULTI_DEFAULT_PART_SIZE = 100 * MiB
+
+# Chunk size for reading http response and writing to disk.
+RSP_CHUNK_SIZE = 1 * MiB
+
+
+@final
+class _ChunkSentinel:
+    """Signals the end of the multipart chunk queue.
+
+    Queue consumer(s) (file writer) should terminate on receiving an item of this type from the queue.
+    Do not instantiate this class directly, use the `END_CHUNK` constant as a pseudo-singleton instead.
+
+    NOTE: As implemented, this should only be used in multi-threaded (not multi-process) contexts, as
+    it's not currently guaranteed to be process-safe.
+    """
+
+    def __repr__(self) -> str:
+        return "ChunkSentinel"
+
+
+END_CHUNK: Final[_ChunkSentinel] = _ChunkSentinel()
+
+
+def is_end_chunk(obj: Any) -> TypeIs[_ChunkSentinel]:
+    """Returns True if the object is the terminal queue item for multipart downloads."""
+    # Needed for type checking, since _ChunkSentinel isn't formally a singleton.
+    return obj is END_CHUNK
+
+
+@dataclass(frozen=True)
+class ChunkContent:
+    __slots__ = ("offset", "data")  # slots=True only introduced in Python 3.10
+    offset: int
+    data: bytes
+
+
+QueuedChunk: TypeAlias = Union[ChunkContent, _ChunkSentinel]
+
+
+def should_multipart_download(size: int | None, override: bool | None = None) -> bool:
+    return ((size or 0) >= MIN_MULTI_DOWNLOAD_SIZE) if (override is None) else override
+
+
+def calc_part_size(file_size: int, min_part_size: int = MULTI_DEFAULT_PART_SIZE) -> int:
+    # Default to a chunk size of 100MiB. S3 has a cap of 10,000 upload parts.
+    return max(math.ceil(file_size / MAX_PARTS), min_part_size)
+
+
+def scan_chunks(path: str, chunk_size: int) -> Iterator[bytes]:
+    with open(path, "rb") as f:
+        while data := f.read(chunk_size):
+            yield data
+
+
+@dataclass
+class MultipartDownloadContext:
+    q: Queue[QueuedChunk]
+    cancel: threading.Event = field(default_factory=threading.Event)
+
+
+def multipart_download(
+    executor: Executor,
+    session: Session,
+    url: str,
+    size: int,
+    cached_open: Opener,
+    part_size: int = MULTI_DEFAULT_PART_SIZE,
+):
+    """Download file as multiple parts in parallel.
+
+    Only one thread for writing to file. Each part run one http request in one thread.
+    HTTP response chunk of a file part is sent to the writer thread via a queue.
+    """
+    # ------------------------------------------------------------------------------
+    # Shared between threads
+    ctx = MultipartDownloadContext(q=Queue(maxsize=500))
+
+    # Put cache_open at top so we remove the tmp file when there is network error.
+    with cached_open("wb") as f:
+
+        def download_chunk(start: int, end: int | None = None) -> None:
+            # Error from another thread, no need to start
+            if ctx.cancel.is_set():
+                return
+
+            # https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Range
+            # Start and end are both inclusive, empty end means use the actual end of the file.
+            # e.g. "bytes=0-499"
+            bytes_range = f"{start}-" if (end is None) else f"{start}-{end}"
+            headers = {"Range": f"bytes={bytes_range}"}
+            with session.get(url=url, headers=headers, stream=True) as rsp:
+                offset = start
+                for chunk in rsp.iter_content(chunk_size=RSP_CHUNK_SIZE):
+                    if ctx.cancel.is_set():
+                        return
+                    ctx.q.put(ChunkContent(offset=offset, data=chunk))
+                    offset += len(chunk)
+
+        def write_chunks() -> None:
+            # If all chunks are written or there's an error in another thread, shutdown
+            while not (ctx.cancel.is_set() or is_end_chunk(chunk := ctx.q.get())):
+                try:
+                    # NOTE: Seek works without pre allocating the file on disk.
+                    # It automatically creates a sparse file, e.g. ls -hl would show
+                    # a bigger size compared to du -sh * because downloading different
+                    # chunks is not a sequential write.
+                    # See https://man7.org/linux/man-pages/man2/lseek.2.html
+                    f.seek(chunk.offset)
+                    f.write(chunk.data)
+
+                except Exception as e:
+                    if env.is_debug():
+                        logger.debug(f"Error writing chunk to file: {e}")
+                    ctx.cancel.set()
+                    raise
+
+        # Start writer thread first.
+        write_future = executor.submit(write_chunks)
+
+        # Start download threads for each chunk.
+        download_futures = set()
+        for start in range(0, size, part_size):
+            # https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Range
+            # Start and end are both inclusive, empty end means use the actual end of the file.
+            # e.g. bytes=0-499
+            end = end if (end := (start + part_size - 1)) < size else None
+            download_futures.add(executor.submit(download_chunk, start=start, end=end))
+
+        # Wait for download
+        done, not_done = wait(download_futures, return_when=FIRST_EXCEPTION)
+        try:
+            for fut in done:
+                fut.result()
+        except Exception as e:
+            if env.is_debug():
+                logger.debug(f"Error downloading file: {e}")
+            ctx.cancel.set()
+
+            # Cancel any pending futures.  Note:
+            # - `Future.cancel()` does NOT stop the future if it's running, which is why
+            #   there's a separate `threading.Event` to ensure cooperative cancellation.
+            # - Once Python 3.8 support is dropped, replace these `fut.cancel()`
+            #   calls with `Executor.shutdown(cancel_futures=True)`.
+            for fut in not_done:
+                fut.cancel()
+            raise
+        finally:
+            # Always signal the writer to stop
+            ctx.q.put(END_CHUNK)
+            write_future.result()