wandb 0.19.9__py3-none-win32.whl → 0.19.11__py3-none-win32.whl

wandb/sdk/artifacts/artifact_file_cache.py

@@ -11,7 +11,7 @@ import subprocess
 import sys
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import IO, TYPE_CHECKING, ContextManager, Iterator
+from typing import IO, ContextManager, Iterator, Protocol
 
 import wandb
 from wandb import env, util
@@ -19,12 +19,10 @@ from wandb.sdk.lib.filesystem import files_in
 from wandb.sdk.lib.hashutil import B64MD5, ETag, b64_to_hex_id
 from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 
-if TYPE_CHECKING:
-    from typing import Protocol
 
-    class Opener(Protocol):
-        def __call__(self, mode: str = ...) -> ContextManager[IO]:
-            pass
+class Opener(Protocol):
+    def __call__(self, mode: str = ...) -> ContextManager[IO]:
+        pass
 
 
 def _get_sys_umask_threadsafe() -> int:

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import concurrent.futures
 import json
 import logging
 import os
@@ -9,8 +10,9 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 from urllib.parse import urlparse
 
+from wandb.proto.wandb_deprecated import Deprecated
 from wandb.sdk.lib import filesystem
-from wandb.sdk.lib.deprecate import Deprecated, deprecate
+from wandb.sdk.lib.deprecate import deprecate
 from wandb.sdk.lib.hashutil import (
     B64MD5,
     ETag,
@@ -129,7 +131,11 @@ class ArtifactManifestEntry:
         return self._parent_artifact
 
     def download(
-        self, root: str | None = None, skip_cache: bool | None = None
+        self,
+        root: str | None = None,
+        skip_cache: bool | None = None,
+        executor: concurrent.futures.Executor | None = None,
+        multipart: bool | None = None,
     ) -> FilePathStr:
         """Download this artifact entry to the specified root path.
 
@@ -169,7 +175,11 @@ class ArtifactManifestEntry:
             )
         else:
             cache_path = self._parent_artifact.manifest.storage_policy.load_file(
-                self._parent_artifact, self, dest_path=override_cache_path
+                self._parent_artifact,
+                self,
+                dest_path=override_cache_path,
+                executor=executor,
+                multipart=multipart,
             )
 
         if skip_cache:

wandb/sdk/artifacts/storage_handlers/azure_handler.py

@@ -171,6 +171,7 @@ class AzureHandler(StorageHandler):
     def _get_credential(
         self, account_url: str
     ) -> azure.identity.DefaultAzureCredential | str:
+        # NOTE: Always returns default credential for reinit="create_new" runs.
         if (
             wandb.run
             and wandb.run.settings.azure_account_url_to_access_key is not None

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -2,20 +2,28 @@
 
 from __future__ import annotations
 
+import concurrent.futures
+import functools
 import hashlib
+import logging
 import math
 import os
+import queue
 import shutil
-from typing import TYPE_CHECKING, Any, Sequence
+import threading
+from collections import deque
+from typing import IO, TYPE_CHECKING, Any, NamedTuple, Sequence
 from urllib.parse import quote
 
 import requests
 import urllib3
 
+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
@@ -60,6 +68,27 @@ 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:
@@ -120,7 +149,20 @@ class WandbStoragePolicy(StoragePolicy):
         artifact: Artifact,
         manifest_entry: ArtifactManifestEntry,
         dest_path: str | None = None,
+        executor: concurrent.futures.Executor | None = None,
+        multipart: bool | None = None,
     ) -> FilePathStr:
+        """Use cache or download the file using signed url.
+
+        Args:
+            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 dest_path is not None:
             self._cache._override_cache_path = dest_path
 
@@ -132,6 +174,20 @@ class WandbStoragePolicy(StoragePolicy):
             return path
 
         if manifest_entry._download_url is not None:
+            # Use multipart parallel download for large file
+            if (
+                executor is not None
+                and manifest_entry.size is not None
+                and self._should_multipart_download(manifest_entry.size, multipart)
+            ):
+                self._multipart_file_download(
+                    executor,
+                    manifest_entry._download_url,
+                    manifest_entry.size,
+                    cache_open,
+                )
+                return path
+            # Serial download
             response = self._session.get(manifest_entry._download_url, stream=True)
             try:
                 response.raise_for_status()
@@ -165,6 +221,131 @@ class WandbStoragePolicy(StoragePolicy):
                 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 e
+            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,
+        )
+        response.raise_for_status()
+
+        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 e
+            finally:
+                # Always signal the writer to stop
+                q.put(_CHUNK_QUEUE_SENTINEL)
+                write_future.result()
+
     def store_reference(
         self,
         artifact: Artifact,

wandb/sdk/artifacts/storage_policy.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import concurrent.futures
 from typing import TYPE_CHECKING, Sequence
 
 from wandb.sdk.internal.internal_api import Api as InternalApi
@@ -40,6 +41,8 @@ class StoragePolicy:
         artifact: Artifact,
         manifest_entry: ArtifactManifestEntry,
         dest_path: str | None = None,
+        executor: concurrent.futures.Executor | None = None,
+        multipart: bool | None = None,
     ) -> FilePathStr:
         raise NotImplementedError
 

wandb/sdk/data_types/base_types/media.py

@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type, Union, ca
 
 import wandb
 from wandb import util
-from wandb._globals import _datatypes_callback
 from wandb.sdk.lib import filesystem
 from wandb.sdk.lib.paths import LogicalPath
 
@@ -192,7 +191,7 @@ class Media(WBValue):
             shutil.move(self._path, new_path)
             self._path = new_path
             self._is_tmp = False
-            _datatypes_callback(media_path)
+            run._publish_file(media_path)
         else:
             try:
                 shutil.copy(self._path, new_path)
@@ -200,7 +199,7 @@ class Media(WBValue):
                 if not ignore_copy_err:
                     raise e
             self._path = new_path
-            _datatypes_callback(media_path)
+            run._publish_file(media_path)
 
     def to_json(self, run: Union["LocalRun", "Artifact"]) -> dict:
         """Serialize the object into a JSON blob.

wandb/sdk/data_types/base_types/wb_value.py

@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Type, Union
 
-import wandb
 from wandb import util
+from wandb.sdk import wandb_setup
 
 if TYPE_CHECKING:  # pragma: no cover
     from wandb.sdk.artifacts.artifact import Artifact
@@ -11,6 +11,31 @@ if TYPE_CHECKING:  # pragma: no cover
     TypeMappingType = Dict[str, Type["WBValue"]]
 
 
+def _is_maybe_offline() -> bool:
+    """Guess whether wandb is configured to be offline.
+
+    This is an anti-pattern because there is no library-level "offline" mode:
+    only runs can be offline. Online and offline runs can exist in the same
+    process. This function is a heuristic that works only if there is at most
+    one run in the process, and could otherwise produce unexpected results.
+
+    Returns:
+        Whether the user likely configured wandb to be offline.
+    """
+    singleton = wandb_setup._setup(start_service=False)
+
+    # First check: if there's a run, check if it is offline.
+    #
+    # This covers uses like `wandb.init(mode="offline")` which don't modify
+    # the singleton's settings.
+    if run := singleton.most_recent_active_run:
+        return run.offline
+
+    # Second check: default to global defaults derived from environment
+    # variables or passed explicitly to `wandb.setup()`.
+    return singleton.settings._offline
+
+
 def _server_accepts_client_ids() -> bool:
     from wandb.util import parse_version
 
@@ -25,15 +50,13 @@ def _server_accepts_client_ids() -> bool:
     # AS OF NOW, 2024/11/06, we assume that all customer's server deployments accept
     # client IDs.
 
-    if util._is_offline():
-        # If there are any users with issues on an older backend, customers can disable the
-        # setting `allow_offline_artifacts` to revert the SDK's behavior back to not
-        # using client IDs in offline mode.
-        if wandb.run and not wandb.run.settings.allow_offline_artifacts:
-            return False
-        # Assume client IDs are accepted
+    if _is_maybe_offline():
+        singleton = wandb_setup._setup(start_service=False)
+
+        if run := singleton.most_recent_active_run:
+            return run._settings.allow_offline_artifacts
         else:
-            return True
+            return singleton.settings.allow_offline_artifacts
 
     # If the script is online, request the max_cli_version and ensure the server
     # is of a high enough version.
@@ -240,7 +263,7 @@ class WBValue:
             self._artifact_target
             and self._artifact_target.name
             and self._artifact_target.artifact._is_draft_save_started()
-            and not util._is_offline()
+            and not _is_maybe_offline()
             and not _server_accepts_client_ids()
         ):
             self._artifact_target.artifact.wait()
@@ -271,7 +294,7 @@ class WBValue:
             self._artifact_target
             and self._artifact_target.name
             and self._artifact_target.artifact._is_draft_save_started()
-            and not util._is_offline()
+            and not _is_maybe_offline()
             and not _server_accepts_client_ids()
         ):
             self._artifact_target.artifact.wait()

wandb/sdk/data_types/html.py

@@ -16,19 +16,46 @@ if TYPE_CHECKING:  # pragma: no cover
 
 
 class Html(BatchableMedia):
-    """Wandb class for arbitrary html.
-
-    Args:
-        data: (string or io object) HTML to display in wandb
-        inject: (boolean) Add a stylesheet to the HTML object.  If set
-            to False the HTML will pass through unchanged.
-    """
+    """A class for logging HTML content to W&B."""
 
     _log_type = "html-file"
 
-    def __init__(self, data: Union[str, "TextIO"], inject: bool = True) -> None:
+    def __init__(
+        self,
+        data: Union[str, "TextIO"],
+        inject: bool = True,
+        data_is_not_path: bool = False,
+    ) -> None:
+        """Creates a W&B HTML object.
+
+        It can be initialized by providing a path to a file:
+        ```
+        with wandb.init() as run:
+            run.log({"html": wandb.Html("./index.html")})
+        ```
+
+        Alternatively, it can be initialized by providing literal HTML,
+        in either a string or IO object:
+        ```
+        with wandb.init() as run:
+            run.log({"html": wandb.Html("<h1>Hello, world!</h1>")})
+        ```
+
+        Args:
+            data:
+                A string that is a path to a file with the extension ".html",
+                or a string or IO object containing literal HTML.
+            inject: Add a stylesheet to the HTML object. If set
+                to False the HTML will pass through unchanged.
+            data_is_not_path: If set to False, the data will be
+                treated as a path to a file.
+        """
         super().__init__()
-        data_is_path = isinstance(data, str) and os.path.exists(data)
+        data_is_path = (
+            isinstance(data, str)
+            and os.path.isfile(data)
+            and os.path.splitext(data)[1] == ".html"
+        ) and not data_is_not_path
         data_path = ""
         if data_is_path:
             assert isinstance(data, str)

wandb/sdk/data_types/image.py

@@ -34,8 +34,8 @@ if TYPE_CHECKING:  # pragma: no cover
     TorchTensorType = Union["torch.Tensor", "torch.Variable"]
 
 
-def _server_accepts_image_filenames() -> bool:
-    if util._is_offline():
+def _server_accepts_image_filenames(run: "LocalRun") -> bool:
+    if run.offline:
         return True
 
     # Newer versions of wandb accept large image filenames arrays
@@ -51,15 +51,15 @@ def _server_accepts_image_filenames() -> bool:
     return accepts_image_filenames
 
 
-def _server_accepts_artifact_path() -> bool:
-    from wandb.util import parse_version
+def _server_accepts_artifact_path(run: "LocalRun") -> bool:
+    if run.offline:
+        return False
+
+    max_cli_version = util._get_max_cli_version()
+    if max_cli_version is None:
+        return False
 
-    target_version = "0.12.14"
-    max_cli_version = util._get_max_cli_version() if not util._is_offline() else None
-    accepts_artifact_path: bool = max_cli_version is not None and parse_version(
-        target_version
-    ) <= parse_version(max_cli_version)
-    return accepts_artifact_path
+    return util.parse_version("0.12.14") <= util.parse_version(max_cli_version)
 
 
 class Image(BatchableMedia):
@@ -401,7 +401,7 @@ class Image(BatchableMedia):
             )
 
         if (
-            not _server_accepts_artifact_path()
+            not _server_accepts_artifact_path(run)
             or self._get_artifact_entry_ref_url() is None
         ):
             super().bind_to_run(run, key, step, id_, ignore_copy_err=ignore_copy_err)
@@ -575,7 +575,7 @@ class Image(BatchableMedia):
             "format": format,
             "count": num_images_to_log,
         }
-        if _server_accepts_image_filenames():
+        if _server_accepts_image_filenames(run):
             meta["filenames"] = [
                 obj.get("path", obj.get("artifact_path")) for obj in jsons
             ]

wandb/sdk/data_types/table.py

@@ -480,6 +480,11 @@ class Table(Media):
             max_rows = Table.MAX_ROWS
         n_rows = len(self.data)
         if n_rows > max_rows and warn:
+            # NOTE: Never raises for reinit="create_new" runs.
+            #   Since this is called by bind_to_run(), this can be fixed by
+            #   propagating the run. It cannot be fixed for to_json() calls
+            #   that are given an artifact, other than by deferring to singleton
+            #   settings.
             if wandb.run and (
                 wandb.run.settings.table_raise_on_max_row_limit_exceeded
                 or wandb.run.settings.strict

wandb/sdk/data_types/trace_tree.py

@@ -431,6 +431,8 @@ class Trace:
             name: The name of the trace to be logged
         """
         trace_tree = WBTraceTree(self._span, self._model_dict)
+        # NOTE: Does not work for reinit="create_new" runs.
+        #   This method should be deprecated and users should call run.log().
         assert (
             wandb.run is not None
         ), "You must call wandb.init() before logging a trace"

wandb/sdk/data_types/utils.py

@@ -101,7 +101,7 @@ def val_to_json(
 
             items = _prune_max_seq(val)
 
-            if _server_accepts_image_filenames():
+            if _server_accepts_image_filenames(run):
                 for item in items:
                     item.bind_to_run(
                         run=run,