wandb 0.22.1__py3-none-win_arm64.whl → 0.22.2__py3-none-win_arm64.whl

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -3,32 +3,35 @@
 from __future__ import annotations
 
 import concurrent.futures
-import functools
 import hashlib
 import logging
-import math
 import os
-import queue
 import shutil
-import threading
 from collections import deque
-from typing import IO, TYPE_CHECKING, Any, NamedTuple, Sequence
+from operator import itemgetter
+from typing import TYPE_CHECKING, Any
 from urllib.parse import quote
 
 import requests
 
-from wandb import env
 from wandb.errors.term import termwarn
 from wandb.proto.wandb_internal_pb2 import ServerFeature
 from wandb.sdk.artifacts.artifact_file_cache import (
     ArtifactFileCache,
-    Opener,
     get_artifact_file_cache,
 )
 from wandb.sdk.artifacts.staging import get_staging_dir
 from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler
 from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler
 from wandb.sdk.artifacts.storage_layout import StorageLayout
+from wandb.sdk.artifacts.storage_policies._multipart import (
+    MAX_MULTI_UPLOAD_SIZE,
+    MIN_MULTI_UPLOAD_SIZE,
+    KiB,
+    calc_part_size,
+    multipart_download,
+    scan_chunks,
+)
 from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY
 from wandb.sdk.artifacts.storage_policy import StoragePolicy
 from wandb.sdk.internal.internal_api import Api as InternalApi
@@ -44,34 +47,9 @@ if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
     from wandb.sdk.internal import progress
 
-
-# AWS S3 max upload parts without having to make additional requests for extra parts
-S3_MAX_PART_NUMBERS = 1000
-S3_MIN_MULTI_UPLOAD_SIZE = 2 * 1024**3
-S3_MAX_MULTI_UPLOAD_SIZE = 5 * 1024**4
-
-
-# Minimum size to switch to multipart download, same as upload, 2GB.
-_MULTIPART_DOWNLOAD_SIZE = S3_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
-_DOWNLOAD_PART_SIZE_BYTES = 100 * 1024 * 1024
-# Chunk size for reading http response and writing to disk. 1MB.
-_HTTP_RES_CHUNK_SIZE_BYTES = 1 * 1024 * 1024
-# Signal end of _ChunkQueue, consumer (file writer) should stop after getting this item.
-# NOTE: it should only be used for multithread executor, it does notwork for multiprocess executor.
-# multipart download is using the executor from artifact.download() which is a multithread executor.
-_CHUNK_QUEUE_SENTINEL = object()
-
 logger = logging.getLogger(__name__)
 
 
-class _ChunkContent(NamedTuple):
-    offset: int
-    data: bytes
-
-
 class WandbStoragePolicy(StoragePolicy):
     @classmethod
     def name(cls) -> str:
@@ -109,7 +87,7 @@ class WandbStoragePolicy(StoragePolicy):
         if not storage_region.strip():
             raise ValueError("storageRegion must be a non-empty string")
 
-    def config(self) -> dict:
+    def config(self) -> dict[str, Any]:
         return self._config
 
     def load_file(
@@ -117,8 +95,9 @@ class WandbStoragePolicy(StoragePolicy):
         artifact: Artifact,
         manifest_entry: ArtifactManifestEntry,
         dest_path: str | None = None,
+        # FIXME: We should avoid passing the executor into multiple inner functions,
+        # it leads to confusing code and opaque tracebacks/call stacks.
         executor: concurrent.futures.Executor | None = None,
-        multipart: bool | None = None,
     ) -> FilePathStr:
         """Use cache or download the file using signed url.
 
@@ -126,10 +105,8 @@ class WandbStoragePolicy(StoragePolicy):
             executor: Passed from caller, artifact has a thread pool for multi file download.
                 Reuse the thread pool for multi part download. The thread pool is closed when
                 artifact download is done.
-            multipart: If set to `None` (default), the artifact will be downloaded
-                in parallel using multipart download if individual file size is greater than
-                2GB. If set to `True` or `False`, the artifact will be downloaded in
-                parallel or serially regardless of the file size.
+
+                If this is None, download the file serially.
         """
         if dest_path is not None:
             self._cache._override_cache_path = dest_path
@@ -141,14 +118,10 @@ class WandbStoragePolicy(StoragePolicy):
         if hit:
             return path
 
-        if (url := manifest_entry._download_url) is not None:
+        if url := manifest_entry._download_url:
             # Use multipart parallel download for large file
-            if (
-                executor
-                and (size := manifest_entry.size)
-                and self._should_multipart_download(size, multipart)
-            ):
-                self._multipart_file_download(executor, url, size, cache_open)
+            if executor and (size := manifest_entry.size):
+                multipart_download(executor, self._session, url, size, cache_open)
                 return path
 
             # Serial download
@@ -171,142 +144,16 @@ class WandbStoragePolicy(StoragePolicy):
             else:
                 auth = ("api", self._api.api_key or "")
 
-            file_url = self._file_url(
-                self._api,
-                artifact.entity,
-                artifact.project,
-                artifact.name.split(":")[0],
-                manifest_entry,
-            )
+            file_url = self._file_url(self._api, artifact, manifest_entry)
             response = self._session.get(
                 file_url, auth=auth, cookies=cookies, headers=headers, stream=True
             )
 
         with cache_open(mode="wb") as file:
-            for data in response.iter_content(chunk_size=16 * 1024):
+            for data in response.iter_content(chunk_size=16 * KiB):
                 file.write(data)
         return path
 
-    def _should_multipart_download(
-        self,
-        file_size: int,
-        multipart: bool | None,
-    ) -> bool:
-        if multipart is not None:
-            return multipart
-        return file_size >= _MULTIPART_DOWNLOAD_SIZE
-
-    def _write_chunks_to_file(
-        self,
-        f: IO,
-        q: queue.Queue,
-        download_has_error: threading.Event,
-    ):
-        while not download_has_error.is_set():
-            item = q.get()
-            if item is _CHUNK_QUEUE_SENTINEL:
-                # Normal shutdown, all the chunks are written
-                return
-            elif isinstance(item, _ChunkContent):
-                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(item.offset)
-                    f.write(item.data)
-                except Exception as e:
-                    if env.is_debug():
-                        logger.debug(f"Error writing chunk to file: {e}")
-                    download_has_error.set()
-                    raise
-            else:
-                raise ValueError(f"Unknown queue item type: {type(item)}")
-
-    def _download_part(
-        self,
-        download_url: str,
-        headers: dict,
-        start: int,
-        q: queue.Queue,
-        download_has_error: threading.Event,
-    ):
-        # Other threads has error, no need to start
-        if download_has_error.is_set():
-            return
-        response = self._session.get(url=download_url, headers=headers, stream=True)
-
-        file_offset = start
-        for content in response.iter_content(chunk_size=_HTTP_RES_CHUNK_SIZE_BYTES):
-            if download_has_error.is_set():
-                return
-            q.put(_ChunkContent(offset=file_offset, data=content))
-            file_offset += len(content)
-
-    def _multipart_file_download(
-        self,
-        executor: concurrent.futures.Executor,
-        download_url: str,
-        file_size_bytes: int,
-        cache_open: Opener,
-    ):
-        """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.
-        """
-        q: queue.Queue[_ChunkContent | object] = queue.Queue(maxsize=500)
-        download_has_error = threading.Event()
-
-        # Put cache_open at top so we remove the tmp file when there is network error.
-        with cache_open("wb") as f:
-            # Start writer thread first.
-            write_handler = functools.partial(
-                self._write_chunks_to_file, f, q, download_has_error
-            )
-            write_future = executor.submit(write_handler)
-
-            # Start download threads for each part.
-            download_futures: deque[concurrent.futures.Future] = deque()
-            part_size = _DOWNLOAD_PART_SIZE_BYTES
-            num_parts = int(math.ceil(file_size_bytes / float(part_size)))
-            for i in range(num_parts):
-                # 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.
-                start = i * part_size
-                bytes_range = f"bytes={start}-"
-                if i != (num_parts - 1):
-                    # bytes=0-499
-                    bytes_range += f"{start + part_size - 1}"
-                headers = {"Range": bytes_range}
-                download_handler = functools.partial(
-                    self._download_part,
-                    download_url,
-                    headers,
-                    start,
-                    q,
-                    download_has_error,
-                )
-                download_futures.append(executor.submit(download_handler))
-
-            # Wait for download
-            done, not_done = concurrent.futures.wait(
-                download_futures, return_when=concurrent.futures.FIRST_EXCEPTION
-            )
-            try:
-                for fut in done:
-                    fut.result()
-            except Exception as e:
-                if env.is_debug():
-                    logger.debug(f"Error downloading file: {e}")
-                download_has_error.set()
-                raise
-            finally:
-                # Always signal the writer to stop
-                q.put(_CHUNK_QUEUE_SENTINEL)
-                write_future.result()
-
     def store_reference(
         self,
         artifact: Artifact,
@@ -314,7 +161,7 @@ class WandbStoragePolicy(StoragePolicy):
         name: str | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         return self._handler.store_path(
             artifact, path, name=name, checksum=checksum, max_objects=max_objects
         )
@@ -334,13 +181,16 @@ class WandbStoragePolicy(StoragePolicy):
     def _file_url(
         self,
         api: InternalApi,
-        entity_name: str,
-        project_name: str,
-        artifact_name: str,
+        artifact: Artifact,
         entry: ArtifactManifestEntry,
     ) -> str:
         layout = self._config.get("storageLayout", StorageLayout.V1)
         region = self._config.get("storageRegion", "default")
+
+        entity_name = artifact.entity
+        project_name = artifact.project
+        artifact_name = artifact.name.split(":")[0]
+
         md5_hex = b64_to_hex_id(entry.digest)
 
         base_url: str = api.settings("base_url")
@@ -367,30 +217,21 @@ class WandbStoragePolicy(StoragePolicy):
         multipart_urls: dict[int, str],
         extra_headers: dict[str, str],
     ) -> list[dict[str, Any]]:
-        etags = []
-        part_number = 1
-
-        with open(file_path, "rb") as f:
-            while True:
-                data = f.read(chunk_size)
-                if not data:
-                    break
-                md5_b64_str = str(hex_to_b64_id(hex_digests[part_number]))
-                upload_resp = self._api.upload_multipart_file_chunk_retry(
-                    multipart_urls[part_number],
-                    data,
-                    extra_headers={
-                        "content-md5": md5_b64_str,
-                        "content-length": str(len(data)),
-                        "content-type": extra_headers.get("Content-Type", ""),
-                    },
-                )
-                assert upload_resp is not None
-                etags.append(
-                    {"partNumber": part_number, "hexMD5": upload_resp.headers["ETag"]}
-                )
-                part_number += 1
-        return etags
+        etags: deque[dict[str, Any]] = deque()
+        file_chunks = scan_chunks(file_path, chunk_size)
+        for num, data in enumerate(file_chunks, start=1):
+            rsp = self._api.upload_multipart_file_chunk_retry(
+                multipart_urls[num],
+                data,
+                extra_headers={
+                    "content-md5": hex_to_b64_id(hex_digests[num]),
+                    "content-length": str(len(data)),
+                    "content-type": extra_headers.get("Content-Type") or "",
+                },
+            )
+            assert rsp is not None
+            etags.append({"partNumber": num, "hexMD5": rsp.headers["ETag"]})
+        return list(etags)
 
     def default_file_upload(
         self,
@@ -403,20 +244,9 @@ class WandbStoragePolicy(StoragePolicy):
         with open(file_path, "rb") as file:
             # This fails if we don't send the first byte before the signed URL expires.
             self._api.upload_file_retry(
-                upload_url,
-                file,
-                progress_callback,
-                extra_headers=extra_headers,
+                upload_url, file, progress_callback, extra_headers=extra_headers
             )
 
-    def calc_chunk_size(self, file_size: int) -> int:
-        # Default to chunk size of 100MiB. S3 has cap of 10,000 upload parts.
-        # If file size exceeds the default chunk size, recalculate chunk size.
-        default_chunk_size = 100 * 1024**2
-        if default_chunk_size * S3_MAX_PART_NUMBERS < file_size:
-            return math.ceil(file_size / S3_MAX_PART_NUMBERS)
-        return default_chunk_size
-
     def store_file(
         self,
         artifact_id: str,
@@ -432,28 +262,20 @@ class WandbStoragePolicy(StoragePolicy):
             False if it needed to be uploaded or was a reference (nothing to dedupe).
         """
         file_size = entry.size or 0
-        chunk_size = self.calc_chunk_size(file_size)
-        upload_parts = []
-        hex_digests = {}
-        file_path = entry.local_path if entry.local_path is not None else ""
+        chunk_size = calc_part_size(file_size)
+        file_path = entry.local_path or ""
         # Logic for AWS s3 multipart upload.
         # Only chunk files if larger than 2 GiB. Currently can only support up to 5TiB.
-        if (
-            file_size >= S3_MIN_MULTI_UPLOAD_SIZE
-            and file_size <= S3_MAX_MULTI_UPLOAD_SIZE
-        ):
-            part_number = 1
-            with open(file_path, "rb") as f:
-                while True:
-                    data = f.read(chunk_size)
-                    if not data:
-                        break
-                    hex_digest = hashlib.md5(data).hexdigest()
-                    upload_parts.append(
-                        {"hexMD5": hex_digest, "partNumber": part_number}
-                    )
-                    hex_digests[part_number] = hex_digest
-                    part_number += 1
+        if MIN_MULTI_UPLOAD_SIZE <= file_size <= MAX_MULTI_UPLOAD_SIZE:
+            file_chunks = scan_chunks(file_path, chunk_size)
+            upload_parts = [
+                {"partNumber": num, "hexMD5": hashlib.md5(data).hexdigest()}
+                for num, data in enumerate(file_chunks, start=1)
+            ]
+            hex_digests = dict(map(itemgetter("partNumber", "hexMD5"), upload_parts))
+        else:
+            upload_parts = []
+            hex_digests = {}
 
         resp = preparer.prepare(
             {
@@ -467,24 +289,21 @@ class WandbStoragePolicy(StoragePolicy):
 
         entry.birth_artifact_id = resp.birth_artifact_id
 
-        multipart_urls = resp.multipart_upload_urls
         if resp.upload_url is None:
             return True
         if entry.local_path is None:
             return False
-        extra_headers = {
-            header.split(":", 1)[0]: header.split(":", 1)[1]
-            for header in (resp.upload_headers or {})
-        }
+
+        extra_headers = dict(hdr.split(":", 1) for hdr in (resp.upload_headers or []))
 
         # This multipart upload isn't available, do a regular single url upload
-        if multipart_urls is None and resp.upload_url:
+        if (multipart_urls := resp.multipart_upload_urls) is None and resp.upload_url:
             self.default_file_upload(
                 resp.upload_url, file_path, extra_headers, progress_callback
             )
+        elif multipart_urls is None:
+            raise ValueError(f"No multipart urls to upload for file: {file_path}")
         else:
-            if multipart_urls is None:
-                raise ValueError(f"No multipart urls to upload for file: {file_path}")
             # Upload files using s3 multipart upload urls
             etags = self.s3_multipart_file_upload(
                 file_path,
@@ -513,7 +332,7 @@ class WandbStoragePolicy(StoragePolicy):
 
         staging_dir = get_staging_dir()
         try:
-            if not entry.skip_cache and not hit:
+            if not (entry.skip_cache or hit):
                 with cache_open("wb") as f, open(entry.local_path, "rb") as src:
                     shutil.copyfileobj(src, f)
             if entry.local_path.startswith(staging_dir):

wandb/sdk/artifacts/storage_policy.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 import concurrent.futures
-from typing import TYPE_CHECKING, Sequence
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
 
 from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.lib.paths import FilePathStr, URIStr
@@ -15,37 +16,47 @@ if TYPE_CHECKING:
     from wandb.sdk.internal.progress import ProgressFn
 
 
-class StoragePolicy:
+_POLICY_REGISTRY: dict[str, type[StoragePolicy]] = {}
+
+
+class StoragePolicy(ABC):
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        _POLICY_REGISTRY[cls.name()] = cls
+
     @classmethod
     def lookup_by_name(cls, name: str) -> type[StoragePolicy]:
-        import wandb.sdk.artifacts.storage_policies  # noqa: F401
-
-        for sub in cls.__subclasses__():
-            if sub.name() == name:
-                return sub
-        raise NotImplementedError(f"Failed to find storage policy '{name}'")
+        if policy := _POLICY_REGISTRY.get(name):
+            return policy
+        raise ValueError(f"Failed to find storage policy {name!r}")
 
     @classmethod
+    @abstractmethod
     def name(cls) -> str:
         raise NotImplementedError
 
     @classmethod
-    def from_config(cls, config: dict, api: InternalApi | None = None) -> StoragePolicy:
+    @abstractmethod
+    def from_config(
+        cls, config: dict[str, Any], api: InternalApi | None = None
+    ) -> StoragePolicy:
         raise NotImplementedError
 
-    def config(self) -> dict:
+    @abstractmethod
+    def config(self) -> dict[str, Any]:
         raise NotImplementedError
 
+    @abstractmethod
     def load_file(
         self,
         artifact: Artifact,
         manifest_entry: ArtifactManifestEntry,
         dest_path: str | None = None,
         executor: concurrent.futures.Executor | None = None,
-        multipart: bool | None = None,
     ) -> FilePathStr:
         raise NotImplementedError
 
+    @abstractmethod
     def store_file(
         self,
         artifact_id: str,
@@ -56,6 +67,7 @@ class StoragePolicy:
     ) -> bool:
         raise NotImplementedError
 
+    @abstractmethod
     def store_reference(
         self,
         artifact: Artifact,
@@ -63,9 +75,10 @@ class StoragePolicy:
         name: str | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         raise NotImplementedError
 
+    @abstractmethod
     def load_reference(
         self,
         manifest_entry: ArtifactManifestEntry,

wandb/sdk/data_types/object_3d.py

@@ -138,7 +138,7 @@ def box3d(
     label: "Optional[str]" = None,
     score: "Optional[numeric]" = None,
 ) -> "Box3D":
-    """Returns a Box3D.
+    """A 3D bounding box. The box is specified by its center, size and orientation.
 
     Args:
         center: The center point of the box as a length-3 ndarray.
@@ -149,7 +149,72 @@ def box3d(
             r + xi + yj + zk.
         color: The box's color as an (r, g, b) tuple with 0 <= r,g,b <= 1.
         label: An optional label for the box.
-        score: An optional score for the box.
+        score: An optional score for the box. Typically used to indicate
+            the confidence of a detection.
+
+    Returns:
+        A Box3D object.
+
+    Example:
+    The following example creates a point cloud with 60 boxes rotating
+    around the X, Y and Z axes.
+
+    ```python
+    import wandb
+
+    import math
+    import numpy as np
+    from scipy.spatial.transform import Rotation
+
+
+    with wandb.init() as run:
+        run.log(
+            {
+                "points": wandb.Object3D.from_point_cloud(
+                    points=np.random.uniform(-5, 5, size=(100, 3)),
+                    boxes=[
+                        wandb.box3d(
+                            center=(0.3 * t - 3, 0, 0),
+                            size=(0.1, 0.1, 0.1),
+                            orientation=Rotation.from_euler(
+                                "xyz", [t * math.pi / 10, 0, 0]
+                            ).as_quat(),
+                            color=(0.5 + t / 40, 0.5, 0.5),
+                            label=f"box {t}",
+                            score=0.9,
+                        )
+                        for t in range(20)
+                    ]
+                    + [
+                        wandb.box3d(
+                            center=(0, 0.3 * t - 3, 0.3),
+                            size=(0.1, 0.1, 0.1),
+                            orientation=Rotation.from_euler(
+                                "xyz", [0, t * math.pi / 10, 0]
+                            ).as_quat(),
+                            color=(0.5, 0.5 + t / 40, 0.5),
+                            label=f"box {t}",
+                            score=0.9,
+                        )
+                        for t in range(20)
+                    ]
+                    + [
+                        wandb.box3d(
+                            center=(0.3, 0.3, 0.3 * t - 3),
+                            size=(0.1, 0.1, 0.1),
+                            orientation=Rotation.from_euler(
+                                "xyz", [0, 0, t * math.pi / 10]
+                            ).as_quat(),
+                            color=(0.5, 0.5, 0.5 + t / 40),
+                            label=f"box {t}",
+                            score=0.9,
+                        )
+                        for t in range(20)
+                    ],
+                ),
+            }
+        )
+    ```
     """
     try:
         import numpy as np