wandb 0.21.4__py3-none-macosx_12_0_arm64.whl → 0.22.1__py3-none-macosx_12_0_arm64.whl

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -16,7 +16,6 @@ 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
@@ -27,40 +26,24 @@ from wandb.sdk.artifacts.artifact_file_cache import (
     get_artifact_file_cache,
 )
 from wandb.sdk.artifacts.staging import get_staging_dir
-from wandb.sdk.artifacts.storage_handlers.azure_handler import AzureHandler
-from wandb.sdk.artifacts.storage_handlers.gcs_handler import GCSHandler
-from wandb.sdk.artifacts.storage_handlers.http_handler import HTTPHandler
-from wandb.sdk.artifacts.storage_handlers.local_file_handler import LocalFileHandler
 from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler
-from wandb.sdk.artifacts.storage_handlers.s3_handler import S3Handler
 from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler
-from wandb.sdk.artifacts.storage_handlers.wb_artifact_handler import WBArtifactHandler
-from wandb.sdk.artifacts.storage_handlers.wb_local_artifact_handler import (
-    WBLocalArtifactHandler,
-)
 from wandb.sdk.artifacts.storage_layout import StorageLayout
 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
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
-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, URIStr
 
+from ._factories import make_http_session, make_storage_handlers
+
 if TYPE_CHECKING:
     from wandb.filesync.step_prepare import StepPrepare
     from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
     from wandb.sdk.internal import progress
 
-# Sleep length: 0, 2, 4, 8, 16, 32, 64, 120, 120, 120, 120, 120, 120, 120, 120, 120
-# seconds, i.e. a total of 20min 6s.
-_REQUEST_RETRY_STRATEGY = urllib3.util.retry.Retry(
-    backoff_factor=1,
-    total=16,
-    status_forcelist=(308, 408, 409, 429, 500, 502, 503, 504),
-)
-_REQUEST_POOL_CONNECTIONS = 64
-_REQUEST_POOL_MAXSIZE = 64
 
 # AWS S3 max upload parts without having to make additional requests for extra parts
 S3_MAX_PART_NUMBERS = 1000
@@ -96,51 +79,36 @@ class WandbStoragePolicy(StoragePolicy):
 
     @classmethod
     def from_config(
-        cls, config: dict, api: InternalApi | None = None
+        cls, config: dict[str, Any], api: InternalApi | None = None
     ) -> WandbStoragePolicy:
         return cls(config=config, api=api)
 
     def __init__(
         self,
-        config: dict | None = None,
+        config: dict[str, Any] | None = None,
         cache: ArtifactFileCache | None = None,
         api: InternalApi | None = None,
+        session: requests.Session | None = None,
     ) -> None:
-        self._cache = cache or get_artifact_file_cache()
         self._config = config or {}
-        self._session = requests.Session()
-        adapter = requests.adapters.HTTPAdapter(
-            max_retries=_REQUEST_RETRY_STRATEGY,
-            pool_connections=_REQUEST_POOL_CONNECTIONS,
-            pool_maxsize=_REQUEST_POOL_MAXSIZE,
-        )
-        self._session.mount("http://", adapter)
-        self._session.mount("https://", adapter)
-
-        s3 = S3Handler()
-        gcs = GCSHandler()
-        azure = AzureHandler()
-        http = HTTPHandler(self._session)
-        https = HTTPHandler(self._session, scheme="https")
-        artifact = WBArtifactHandler()
-        local_artifact = WBLocalArtifactHandler()
-        file_handler = LocalFileHandler()
-
+        if (storage_region := self._config.get("storageRegion")) is not None:
+            self._validate_storage_region(storage_region)
+        self._cache = cache or get_artifact_file_cache()
+        self._session = session or make_http_session()
         self._api = api or InternalApi()
         self._handler = MultiHandler(
-            handlers=[
-                s3,
-                gcs,
-                azure,
-                http,
-                https,
-                artifact,
-                local_artifact,
-                file_handler,
-            ],
+            handlers=make_storage_handlers(self._session),
             default_handler=TrackingHandler(),
         )
 
+    def _validate_storage_region(self, storage_region: Any) -> None:
+        if not isinstance(storage_region, str):
+            raise TypeError(
+                f"storageRegion must be a string, got {type(storage_region).__name__}: {storage_region!r}"
+            )
+        if not storage_region.strip():
+            raise ValueError("storageRegion must be a non-empty string")
+
     def config(self) -> dict:
         return self._config
 
@@ -167,54 +135,52 @@ class WandbStoragePolicy(StoragePolicy):
             self._cache._override_cache_path = dest_path
 
         path, hit, cache_open = self._cache.check_md5_obj_path(
-            B64MD5(manifest_entry.digest),
-            manifest_entry.size if manifest_entry.size is not None else 0,
+            manifest_entry.digest,
+            size=manifest_entry.size or 0,
         )
         if hit:
             return path
 
-        if manifest_entry._download_url is not None:
+        if (url := 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)
+                executor
+                and (size := manifest_entry.size)
+                and self._should_multipart_download(size, multipart)
             ):
-                self._multipart_file_download(
-                    executor,
-                    manifest_entry._download_url,
-                    manifest_entry.size,
-                    cache_open,
-                )
+                self._multipart_file_download(executor, url, size, cache_open)
                 return path
+
             # Serial download
-            response = self._session.get(manifest_entry._download_url, stream=True)
             try:
-                response.raise_for_status()
-            except Exception:
+                response = self._session.get(url, stream=True)
+            except requests.HTTPError:
                 # Signed URL might have expired, fall back to fetching it one by one.
                 manifest_entry._download_url = None
+
         if manifest_entry._download_url is None:
             auth = None
-            http_headers = _thread_local_api_settings.headers or {}
-            if self._api.access_token is not None:
-                http_headers["Authorization"] = f"Bearer {self._api.access_token}"
-            elif _thread_local_api_settings.cookies is None:
+            headers = _thread_local_api_settings.headers
+            cookies = _thread_local_api_settings.cookies
+
+            # For auth, prefer using (in order): auth header, cookies, HTTP Basic Auth
+            if token := self._api.access_token:
+                headers = {**(headers or {}), "Authorization": f"Bearer {token}"}
+            elif cookies is not None:
+                pass
+            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,
+            )
             response = self._session.get(
-                self._file_url(
-                    self._api,
-                    artifact.entity,
-                    artifact.project,
-                    artifact.name.split(":")[0],
-                    manifest_entry,
-                ),
-                auth=auth,
-                cookies=_thread_local_api_settings.cookies,
-                headers=http_headers,
-                stream=True,
+                file_url, auth=auth, cookies=cookies, headers=headers, stream=True
             )
-            response.raise_for_status()
 
         with cache_open(mode="wb") as file:
             for data in response.iter_content(chunk_size=16 * 1024):
@@ -269,12 +235,7 @@ class WandbStoragePolicy(StoragePolicy):
         # 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()
+        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):
@@ -376,43 +337,27 @@ class WandbStoragePolicy(StoragePolicy):
         entity_name: str,
         project_name: str,
         artifact_name: str,
-        manifest_entry: ArtifactManifestEntry,
+        entry: ArtifactManifestEntry,
     ) -> str:
-        storage_layout = self._config.get("storageLayout", StorageLayout.V1)
-        storage_region = self._config.get("storageRegion", "default")
-        md5_hex = b64_to_hex_id(B64MD5(manifest_entry.digest))
+        layout = self._config.get("storageLayout", StorageLayout.V1)
+        region = self._config.get("storageRegion", "default")
+        md5_hex = b64_to_hex_id(entry.digest)
 
-        if storage_layout == StorageLayout.V1:
-            return "{}/artifacts/{}/{}".format(
-                api.settings("base_url"), entity_name, md5_hex
-            )
-        elif storage_layout == StorageLayout.V2:
+        base_url: str = api.settings("base_url")
+
+        if layout == StorageLayout.V1:
+            return f"{base_url}/artifacts/{entity_name}/{md5_hex}"
+
+        if layout == StorageLayout.V2:
+            birth_artifact_id = entry.birth_artifact_id or ""
             if api._server_supports(
-                ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER  # type: ignore
+                ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER
             ):
-                return "{}/artifactsV2/{}/{}/{}/{}/{}/{}/{}".format(
-                    api.settings("base_url"),
-                    storage_region,
-                    quote(entity_name),
-                    quote(project_name),
-                    quote(artifact_name),
-                    quote(manifest_entry.birth_artifact_id or ""),
-                    md5_hex,
-                    manifest_entry.path.name,
-                )
-            return "{}/artifactsV2/{}/{}/{}/{}".format(
-                api.settings("base_url"),
-                storage_region,
-                entity_name,
-                quote(
-                    manifest_entry.birth_artifact_id
-                    if manifest_entry.birth_artifact_id is not None
-                    else ""
-                ),
-                md5_hex,
-            )
-        else:
-            raise Exception(f"unrecognized storage layout: {storage_layout}")
+                return f"{base_url}/artifactsV2/{region}/{quote(entity_name)}/{quote(project_name)}/{quote(artifact_name)}/{quote(birth_artifact_id)}/{md5_hex}/{entry.path.name}"
+
+            return f"{base_url}/artifactsV2/{region}/{entity_name}/{quote(birth_artifact_id)}/{md5_hex}"
+
+        raise ValueError(f"unrecognized storage layout: {layout!r}")
 
     def s3_multipart_file_upload(
         self,
@@ -486,7 +431,7 @@ class WandbStoragePolicy(StoragePolicy):
             True if the file was a duplicate (did not need to be uploaded),
             False if it needed to be uploaded or was a reference (nothing to dedupe).
         """
-        file_size = entry.size if entry.size is not None else 0
+        file_size = entry.size or 0
         chunk_size = self.calc_chunk_size(file_size)
         upload_parts = []
         hex_digests = {}
@@ -562,8 +507,8 @@ class WandbStoragePolicy(StoragePolicy):
 
         # Cache upon successful upload.
         _, hit, cache_open = self._cache.check_md5_obj_path(
-            B64MD5(entry.digest),
-            entry.size if entry.size is not None else 0,
+            entry.digest,
+            size=entry.size or 0,
         )
 
         staging_dir = get_staging_dir()

wandb/sdk/data_types/bokeh.py

@@ -5,6 +5,7 @@ import pathlib
 from typing import TYPE_CHECKING, Union
 
 from wandb import util
+from wandb._strutils import nameof
 from wandb.sdk.lib import runid
 
 from . import _dtypes
@@ -34,7 +35,10 @@ class Bokeh(Media):
         ],
     ):
         super().__init__()
-        bokeh = util.get_module("bokeh", required=True)
+        bokeh = util.get_module(
+            "bokeh",
+            required=f"{nameof(Bokeh)!r} requires the bokeh package.  Please install it with `pip install bokeh`.",
+        )
         if isinstance(data_or_path, (str, pathlib.Path)) and os.path.exists(
             data_or_path
         ):

wandb/sdk/data_types/image.py

@@ -161,6 +161,17 @@ class Image(BatchableMedia):
     ) -> None:
         """Initialize a `wandb.Image` object.
 
+        This class handles various image data formats and automatically normalizes
+        pixel values to the range [0, 255] when needed, ensuring compatibility
+        with the W&B backend.
+
+        * Data in range [0, 1] is multiplied by 255 and converted to uint8
+        * Data in range [-1, 1] is rescaled from [-1, 1] to [0, 255] by mapping
+            -1 to 0 and 1 to 255, then converted to uint8
+        * Data outside [-1, 1] but not in [0, 255] is clipped to [0, 255] and
+            converted to uint8 (with a warning if values fall outside [0, 255])
+        * Data already in [0, 255] is converted to uint8 without modification
+
         Args:
             data_or_path: Accepts NumPy array/pytorch tensor of image data,
                 a PIL image object, or a path to an image file. If a NumPy
@@ -168,7 +179,7 @@ class Image(BatchableMedia):
                 the image data will be saved to the given file type.
                 If the values are not in the range [0, 255] or all values are in the range [0, 1],
                 the image pixel values will be normalized to the range [0, 255]
-                unless `normalize` is set to False.
+                unless `normalize` is set to `False`.
             - pytorch tensor should be in the format (channel, height, width)
             - NumPy array should be in the format (height, width, channel)
             mode: The PIL mode for an image. Most common are "L", "RGB",
@@ -178,13 +189,13 @@ class Image(BatchableMedia):
             classes: A list of class information for the image,
                 used for labeling bounding boxes, and image masks.
             boxes: A dictionary containing bounding box information for the image.
-                see: https://docs.wandb.ai/ref/python/data-types/boundingboxes2d/
+                see https://docs.wandb.ai/ref/python/data-types/boundingboxes2d/
             masks: A dictionary containing mask information for the image.
-                see: https://docs.wandb.ai/ref/python/data-types/imagemask/
+                see https://docs.wandb.ai/ref/python/data-types/imagemask/
             file_type: The file type to save the image as.
-                This parameter has no effect if data_or_path is a path to an image file.
-            normalize: If True, normalize the image pixel values to fall within the range of [0, 255].
-                Normalize is only applied if data_or_path is a numpy array or pytorch tensor.
+                This parameter has no effect if `data_or_path` is a path to an image file.
+            normalize: If `True`, normalize the image pixel values to fall within the range of [0, 255].
+                Normalize is only applied if `data_or_path` is a numpy array or pytorch tensor.
 
         Examples:
         Create a wandb.Image from a numpy array

wandb/sdk/interface/interface.py

@@ -87,11 +87,38 @@ def file_enum_to_policy(enum: "pb.FilesItem.PolicyType.V") -> "PolicyName":
 
 
 class InterfaceBase:
+    """Methods for sending different types of Records to the service.
+
+    None of the methods may be called from an asyncio context other than
+    deliver_async().
+    """
+
     _drop: bool
 
     def __init__(self) -> None:
         self._drop = False
 
+    @abstractmethod
+    async def deliver_async(
+        self,
+        record: pb.Record,
+    ) -> MailboxHandle[pb.Result]:
+        """Send a record and create a handle to wait for the response.
+
+        The synchronous publish and deliver methods on this class cannot be
+        called in the asyncio thread because they block. Instead of having
+        an async copy of every method, this is a general method for sending
+        any kind of record in the asyncio thread.
+
+        Args:
+            record: The record to send. This method takes ownership of the
+                record and it must not be used afterward.
+
+        Returns:
+            A handle to wait for a response to the record.
+        """
+        raise NotImplementedError
+
     def publish_header(self) -> None:
         header = pb.HeaderRecord()
         self._publish_header(header)
@@ -392,9 +419,13 @@ class InterfaceBase:
             proto_manifest.manifest_file_path = path
             return proto_manifest
 
+        # Set storage policy on storageLayout (always V2) and storageRegion, only allow coreweave-us on wandb.ai for now.
+        # NOTE: the decode logic is NewManifestFromProto in core/pkg/artifacts/manifest.go
+        # The creation logic is in artifacts/_factories.py make_storage_policy
         for k, v in artifact_manifest.storage_policy.config().items() or {}.items():
             cfg = proto_manifest.storage_policy_config.add()
             cfg.key = k
+            # TODO: Why json.dumps when existing values are plain string? We want to send complex structure without defining the proto?
             cfg.value_json = json.dumps(v)
 
         for entry in sorted(artifact_manifest.entries.values(), key=lambda k: k.path):
@@ -883,6 +914,16 @@ class InterfaceBase:
     ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
+    def publish_probe_system_info(self) -> None:
+        probe_system_info = pb.ProbeSystemInfoRequest()
+        return self._publish_probe_system_info(probe_system_info)
+
+    @abstractmethod
+    def _publish_probe_system_info(
+        self, probe_system_info: pb.ProbeSystemInfoRequest
+    ) -> None:
+        raise NotImplementedError
+
     def join(self) -> None:
         # Drop indicates that the internal process has already been shutdown
         if self._drop:
@@ -1010,10 +1051,6 @@ class InterfaceBase:
     ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    @abstractmethod
-    def deliver_operation_stats(self) -> MailboxHandle[pb.Result]:
-        raise NotImplementedError
-
     def deliver_poll_exit(self) -> MailboxHandle[pb.Result]:
         poll_exit = pb.PollExitRequest()
         return self._deliver_poll_exit(poll_exit)

wandb/sdk/interface/interface_queue.py

@@ -8,12 +8,15 @@ import logging
 from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from .interface_shared import InterfaceShared
 
 if TYPE_CHECKING:
     from queue import Queue
 
     from wandb.proto import wandb_internal_pb2 as pb
+    from wandb.sdk.mailbox.mailbox_handle import MailboxHandle
 
 
 logger = logging.getLogger("wandb")
@@ -31,6 +34,13 @@ class InterfaceQueue(InterfaceShared):
         self._process = process
         super().__init__()
 
+    @override
+    async def deliver_async(
+        self,
+        record: "pb.Record",
+    ) -> "MailboxHandle[pb.Result]":
+        raise NotImplementedError
+
     def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
         if self._process and not self._process.is_alive():
             raise Exception("The wandb backend process has shutdown")

wandb/sdk/interface/interface_shared.py

@@ -87,7 +87,6 @@ class InterfaceShared(InterfaceBase):
         stop_status: Optional[pb.StopStatusRequest] = None,
         internal_messages: Optional[pb.InternalMessagesRequest] = None,
         network_status: Optional[pb.NetworkStatusRequest] = None,
-        operation_stats: Optional[pb.OperationStatsRequest] = None,
         poll_exit: Optional[pb.PollExitRequest] = None,
         partial_history: Optional[pb.PartialHistoryRequest] = None,
         sampled_history: Optional[pb.SampledHistoryRequest] = None,
@@ -112,6 +111,7 @@ class InterfaceShared(InterfaceBase):
         python_packages: Optional[pb.PythonPackagesRequest] = None,
         job_input: Optional[pb.JobInputRequest] = None,
         run_finish_without_exit: Optional[pb.RunFinishWithoutExitRequest] = None,
+        probe_system_info: Optional[pb.ProbeSystemInfoRequest] = None,
     ) -> pb.Record:
         request = pb.Request()
         if get_summary:
@@ -128,8 +128,6 @@ class InterfaceShared(InterfaceBase):
             request.internal_messages.CopyFrom(internal_messages)
         elif network_status:
             request.network_status.CopyFrom(network_status)
-        elif operation_stats:
-            request.operations.CopyFrom(operation_stats)
         elif poll_exit:
             request.poll_exit.CopyFrom(poll_exit)
         elif partial_history:
@@ -178,6 +176,8 @@ class InterfaceShared(InterfaceBase):
             request.job_input.CopyFrom(job_input)
         elif run_finish_without_exit:
             request.run_finish_without_exit.CopyFrom(run_finish_without_exit)
+        elif probe_system_info:
+            request.probe_system_info.CopyFrom(probe_system_info)
         else:
             raise Exception("Invalid request")
         record = self._make_record(request=request)
@@ -330,6 +330,12 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(use_artifact=use_artifact)
         self._publish(rec)
 
+    def _publish_probe_system_info(
+        self, probe_system_info: pb.ProbeSystemInfoRequest
+    ) -> None:
+        record = self._make_request(probe_system_info=probe_system_info)
+        self._publish(record)
+
     def _deliver_artifact(
         self,
         log_artifact: pb.LogArtifactRequest,
@@ -415,10 +421,6 @@ class InterfaceShared(InterfaceBase):
         record = self._make_record(exit=exit_data)
         return self._deliver(record)
 
-    def deliver_operation_stats(self):
-        record = self._make_request(operation_stats=pb.OperationStatsRequest())
-        return self._deliver(record)
-
     def _deliver_poll_exit(
         self,
         poll_exit: pb.PollExitRequest,

wandb/sdk/interface/interface_sock.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any
 from typing_extensions import override
 
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .interface_shared import InterfaceShared
 
@@ -21,10 +22,12 @@ logger = logging.getLogger("wandb")
 class InterfaceSock(InterfaceShared):
     def __init__(
         self,
+        asyncer: asyncio_manager.AsyncioManager,
         client: ServiceClient,
         stream_id: str,
     ) -> None:
         super().__init__()
+        self._asyncer = asyncer
         self._client = client
         self._stream_id = stream_id
 
@@ -37,13 +40,16 @@ class InterfaceSock(InterfaceShared):
         self._assign(record)
         request = spb.ServerRequest()
         request.record_publish.CopyFrom(record)
-        self._client.publish(request)
+        self._asyncer.run(lambda: self._client.publish(request))
 
-    @override
     def _deliver(self, record: pb.Record) -> MailboxHandle[pb.Result]:
+        return self._asyncer.run(lambda: self.deliver_async(record))
+
+    @override
+    async def deliver_async(self, record: pb.Record) -> MailboxHandle[pb.Result]:
         self._assign(record)
         request = spb.ServerRequest()
         request.record_publish.CopyFrom(record)
 
-        handle = self._client.deliver(request)
+        handle = await self._client.deliver(request)
         return handle.map(lambda response: response.result_communicate)

wandb/sdk/internal/_generated/__init__.py

@@ -1,15 +1,5 @@
 # Generated by ariadne-codegen
 
+__all__ = ["SERVER_FEATURES_QUERY_GQL", "ServerFeaturesQuery"]
 from .operations import SERVER_FEATURES_QUERY_GQL
-from .server_features_query import (
-    ServerFeaturesQuery,
-    ServerFeaturesQueryServerInfo,
-    ServerFeaturesQueryServerInfoFeatures,
-)
-
-__all__ = [
-    "SERVER_FEATURES_QUERY_GQL",
-    "ServerFeaturesQuery",
-    "ServerFeaturesQueryServerInfo",
-    "ServerFeaturesQueryServerInfoFeatures",
-]
+from .server_features_query import ServerFeaturesQuery

wandb/sdk/internal/sender.py

@@ -343,7 +343,7 @@ class SendManager:
         publish_interface = InterfaceQueue(record_q=record_q)
         context_keeper = context.ContextKeeper()
         return SendManager(
-            settings=SettingsStatic(settings.to_proto()),
+            settings=SettingsStatic(dict(settings)),
             record_q=record_q,
             result_q=result_q,
             interface=publish_interface,