wandb 0.19.10__py3-none-any.whl → 0.19.11__py3-none-any.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
@@ -130,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.
 
@@ -170,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_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/video.py

@@ -2,7 +2,7 @@ import functools
 import logging
 import os
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Optional, Sequence, Type, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Sequence, Type, Union
 
 import wandb
 from wandb import util
@@ -48,36 +48,7 @@ def write_gif_with_image_io(
 
 
 class Video(BatchableMedia):
-    """Format a video for logging to W&B.
-
-    Args:
-        data_or_path: (numpy array, string, io)
-            Video can be initialized with a path to a file or an io object.
-            The format must be "gif", "mp4", "webm" or "ogg".
-            The format must be specified with the format argument.
-            Video can be initialized with a numpy tensor.
-            The numpy tensor must be either 4 dimensional or 5 dimensional.
-            Channels should be (time, channel, height, width) or
-            (batch, time, channel, height width)
-        caption: (string) caption associated with the video for display
-        fps: (int)
-            The frame rate to use when encoding raw video frames. Default value is 4.
-            This parameter has no effect when data_or_path is a string, or bytes.
-        format: (string) format of video, necessary if initializing with path or io object.
-
-    Examples:
-        ### Log a numpy array as a video
-        <!--yeadoc-test:log-video-numpy-->
-        ```python
-        import numpy as np
-        import wandb
-
-        run = wandb.init()
-        # axes are (time, channel, height, width)
-        frames = np.random.randint(low=0, high=256, size=(10, 3, 100, 100), dtype=np.uint8)
-        run.log({"video": wandb.Video(frames, fps=4)})
-        ```
-    """
+    """A class for logging videos to W&B."""
 
     _log_type = "video-file"
     EXTS = ("gif", "mp4", "webm", "ogg")
@@ -89,10 +60,53 @@ class Video(BatchableMedia):
         data_or_path: Union["np.ndarray", str, "TextIO", "BytesIO"],
         caption: Optional[str] = None,
         fps: Optional[int] = None,
-        format: Optional[str] = None,
+        format: Optional[Literal["gif", "mp4", "webm", "ogg"]] = None,
     ):
+        """Initialize a W&B Video object.
+
+        Args:
+            data_or_path:
+                Video can be initialized with a path to a file or an io object.
+                Video can be initialized with a numpy tensor.
+                The numpy tensor must be either 4 dimensional or 5 dimensional.
+                The dimensions should be (number of frames, channel, height, width) or
+                (batch, number of frames, channel, height, width)
+                The format parameter must be specified with the format argument
+                when initializing with a numpy array
+                or io object.
+            caption: Caption associated with the video for display.
+            fps:
+                The frame rate to use when encoding raw video frames.
+                Default value is 4.
+                This parameter has no effect when data_or_path is a string, or bytes.
+            format:
+                Format of video, necessary if initializing with a numpy array
+                or io object. This parameter will be used to determine the format
+                to use when encoding the video data. Accepted values are "gif",
+                "mp4", "webm", or "ogg".
+
+        Examples:
+            ### Log a numpy array as a video
+            ```python
+            import numpy as np
+            import wandb
+
+            with wandb.init() as run:
+                # axes are (number of frames, channel, height, width)
+                frames = np.random.randint(
+                    low=0, high=256, size=(10, 3, 100, 100), dtype=np.uint8
+                )
+                run.log({"video": wandb.Video(frames, format="mp4", fps=4)})
+            ```
+        """
         super().__init__(caption=caption)
 
+        if format is None:
+            wandb.termwarn(
+                "`format` argument was not provided, defaulting to `gif`. "
+                "This parameter will be required in v0.20.0, "
+                "please specify the format explicitly."
+            )
         self._format = format or "gif"
         self._width = None
         self._height = None

wandb/sdk/interface/interface.py

@@ -428,7 +428,6 @@ class InterfaceBase:
 
     def deliver_link_artifact(
         self,
-        run: "Run",
         artifact: "Artifact",
         portfolio_name: str,
         aliases: Iterable[str],
@@ -442,9 +441,9 @@ class InterfaceBase:
         else:
             link_artifact.server_id = artifact.id if artifact.id else ""
         link_artifact.portfolio_name = portfolio_name
-        link_artifact.portfolio_entity = entity or run.entity
+        link_artifact.portfolio_entity = entity or ""
         link_artifact.portfolio_organization = organization or ""
-        link_artifact.portfolio_project = project or run.project
+        link_artifact.portfolio_project = project or ""
         link_artifact.portfolio_aliases.extend(aliases)
 
         return self._deliver_link_artifact(link_artifact)

wandb/sdk/internal/internal_api.py

@@ -12,6 +12,7 @@ import sys
 import threading
 from copy import deepcopy
 from pathlib import Path
+from types import MappingProxyType
 from typing import (
     IO,
     TYPE_CHECKING,
@@ -189,11 +190,6 @@ def _match_org_with_fetched_org_entities(
     """
     for org_names in orgs:
         if organization in org_names:
-            wandb.termwarn(
-                "Registries can be linked/fetched using a shorthand form without specifying the organization name. "
-                "Try using shorthand path format: <my_registry_name>/<artifact_name> or "
-                "just <my_registry_name> if fetching just the project."
-            )
             return org_names.entity_name
 
     if len(orgs) == 1:
@@ -873,30 +869,29 @@ class Api:
         _, _, mutations = self.server_info_introspection()
         return "updateRunQueueItemWarning" in mutations
 
-    def _check_server_feature(self, feature_value: ServerFeature) -> bool:
-        """Check if a server feature is enabled.
-
-        Args:
-            feature_value (ServerFeature): The enum value of the feature to check.
-
-        Returns:
-            bool: True if the feature is enabled, False otherwise.
-
-        Raises:
-            Exception: If server doesn't support feature queries or other errors occur
-        """
+    def _server_features(self) -> Mapping[str, bool]:
+        """Returns a cached, read-only lookup of current server feature flags."""
         if self._server_features_cache is None:
             query = gql(SERVER_FEATURES_QUERY_GQL)
-            response = self.gql(query)
-            server_info = ServerFeaturesQuery.model_validate(response).server_info
-            if server_info and (features := server_info.features):
-                self._server_features_cache = {
-                    f.name: f.is_enabled for f in features if f
-                }
+
+            try:
+                response = self.gql(query)
+            except Exception as e:
+                # Unfortunately we currently have to match on the text of the error message
+                if 'Cannot query field "features" on type "ServerInfo".' in str(e):
+                    self._server_features_cache = {}
+                else:
+                    raise
             else:
-                self._server_features_cache = {}
+                info = ServerFeaturesQuery.model_validate(response).server_info
+                if info and (feats := info.features):
+                    self._server_features_cache = {
+                        f.name: f.is_enabled for f in feats if f
+                    }
+                else:
+                    self._server_features_cache = {}
 
-        return self._server_features_cache.get(ServerFeature.Name(feature_value), False)
+        return MappingProxyType(self._server_features_cache)
 
     def _check_server_feature_with_fallback(self, feature_value: ServerFeature) -> bool:
         """Wrapper around check_server_feature that warns and returns False for older unsupported servers.
@@ -912,12 +907,7 @@ class Api:
         Exceptions:
             Exception: If an error other than the server not supporting feature queries occurs.
         """
-        try:
-            return self._check_server_feature(feature_value)
-        except Exception as e:
-            if 'Cannot query field "features" on type "ServerInfo".' in str(e):
-                return False
-            raise e
+        return self._server_features().get(ServerFeature.Name(feature_value), False)
 
     @normalize_exceptions
     def update_run_queue_item_warning(

wandb/sdk/internal/sender.py

@@ -1444,7 +1444,7 @@ class SendManager:
         )
         if (client_id or server_id) and portfolio_name and entity and project:
             try:
-                self._api.link_artifact(
+                response = self._api.link_artifact(
                     client_id,
                     server_id,
                     portfolio_name,
@@ -1453,9 +1453,12 @@ class SendManager:
                     aliases,
                     organization,
                 )
+                result.response.link_artifact_response.version_index = response[
+                    "versionIndex"
+                ]
             except Exception as e:
                 org_or_entity = organization or entity
-                result.response.log_artifact_response.error_message = (
+                result.response.link_artifact_response.error_message = (
                     f"error linking artifact to "
                     f'"{org_or_entity}/{project}/{portfolio_name}"; error: {e}'
                 )

wandb/sdk/launch/sweeps/utils.py

@@ -223,6 +223,10 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
     flags_dict: Dict[str, Any] = {}
     # (5) flags without equals (e.g. --foo bar)
     args_no_equals: List[str] = []
+    # (6) flags for hydra append config value (e.g. +foo=bar)
+    flags_append_hydra: List[str] = []
+    # (7) flags for hydra override config value (e.g. ++foo=bar)
+    flags_override_hydra: List[str] = []
     for param, config in command["args"].items():
         # allow 'None' as a valid value, but error if no value is found
         try:
@@ -234,6 +238,8 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
         flags.append("--" + _flag)
         flags_no_hyphens.append(_flag)
         args_no_equals += [f"--{param}", str(_value)]
+        flags_append_hydra.append("+" + _flag)
+        flags_override_hydra.append("++" + _flag)
         if isinstance(_value, bool):
             # omit flags if they are boolean and false
             if _value:
@@ -248,6 +254,8 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
         "args_no_boolean_flags": flags_no_booleans,
         "args_json": [json.dumps(flags_dict)],
         "args_dict": flags_dict,
+        "args_append_hydra": flags_append_hydra,
+        "args_override_hydra": flags_override_hydra,
     }
 
 

wandb/sdk/projects/_generated/__init__.py

@@ -0,0 +1,47 @@
+# Generated by ariadne-codegen
+
+from .delete_project import DeleteProject, DeleteProjectDeleteModel
+from .fetch_registry import FetchRegistry, FetchRegistryEntity
+from .fragments import (
+    RegistryFragment,
+    RegistryFragmentArtifactTypes,
+    RegistryFragmentArtifactTypesEdges,
+    RegistryFragmentArtifactTypesEdgesNode,
+)
+from .input_types import ArtifactTypeInput
+from .operations import (
+    DELETE_PROJECT_GQL,
+    FETCH_REGISTRY_GQL,
+    RENAME_PROJECT_GQL,
+    UPSERT_REGISTRY_PROJECT_GQL,
+)
+from .rename_project import (
+    RenameProject,
+    RenameProjectRenameProject,
+    RenameProjectRenameProjectProject,
+)
+from .upsert_registry_project import (
+    UpsertRegistryProject,
+    UpsertRegistryProjectUpsertModel,
+)
+
+__all__ = [
+    "DELETE_PROJECT_GQL",
+    "FETCH_REGISTRY_GQL",
+    "RENAME_PROJECT_GQL",
+    "UPSERT_REGISTRY_PROJECT_GQL",
+    "FetchRegistry",
+    "FetchRegistryEntity",
+    "RenameProject",
+    "RenameProjectRenameProject",
+    "RenameProjectRenameProjectProject",
+    "UpsertRegistryProject",
+    "UpsertRegistryProjectUpsertModel",
+    "DeleteProject",
+    "DeleteProjectDeleteModel",
+    "ArtifactTypeInput",
+    "RegistryFragment",
+    "RegistryFragmentArtifactTypes",
+    "RegistryFragmentArtifactTypesEdges",
+    "RegistryFragmentArtifactTypesEdgesNode",
+]

wandb/sdk/projects/_generated/delete_project.py

@@ -0,0 +1,22 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import Literal, Optional
+
+from pydantic import Field
+
+from wandb._pydantic import GQLBase, Typename
+
+
+class DeleteProject(GQLBase):
+    delete_model: Optional[DeleteProjectDeleteModel] = Field(alias="deleteModel")
+
+
+class DeleteProjectDeleteModel(GQLBase):
+    success: Optional[bool]
+    typename__: Typename[Literal["DeleteModelPayload"]]
+
+
+DeleteProject.model_rebuild()

wandb/sdk/projects/_generated/enums.py

@@ -0,0 +1,4 @@
+# Generated by ariadne-codegen
+# Source: core/api/graphql/schemas/schema-latest.graphql
+
+from __future__ import annotations

wandb/sdk/projects/_generated/fetch_registry.py

@@ -0,0 +1,22 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import Optional
+
+from wandb._pydantic import GQLBase
+
+from .fragments import RegistryFragment
+
+
+class FetchRegistry(GQLBase):
+    entity: Optional[FetchRegistryEntity]
+
+
+class FetchRegistryEntity(GQLBase):
+    project: Optional[RegistryFragment]
+
+
+FetchRegistry.model_rebuild()
+FetchRegistryEntity.model_rebuild()