wandb 0.18.0rc1__py3-none-win32.whl → 0.18.2__py3-none-win32.whl

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -1,10 +1,12 @@
 """WandB storage policy."""
 
+from __future__ import annotations
+
 import hashlib
 import math
 import os
 import shutil
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Sequence
 from urllib.parse import quote
 
 import requests
@@ -64,15 +66,15 @@ class WandbStoragePolicy(StoragePolicy):
 
     @classmethod
     def from_config(
-        cls, config: Dict, api: Optional[InternalApi] = None
-    ) -> "WandbStoragePolicy":
+        cls, config: dict, api: InternalApi | None = None
+    ) -> WandbStoragePolicy:
         return cls(config=config, api=api)
 
     def __init__(
         self,
-        config: Optional[Dict] = None,
-        cache: Optional[ArtifactFileCache] = None,
-        api: Optional[InternalApi] = None,
+        config: dict | None = None,
+        cache: ArtifactFileCache | None = None,
+        api: InternalApi | None = None,
     ) -> None:
         self._cache = cache or get_artifact_file_cache()
         self._config = config or {}
@@ -109,14 +111,14 @@ class WandbStoragePolicy(StoragePolicy):
             default_handler=TrackingHandler(),
         )
 
-    def config(self) -> Dict:
+    def config(self) -> dict:
         return self._config
 
     def load_file(
         self,
-        artifact: "Artifact",
-        manifest_entry: "ArtifactManifestEntry",
-        dest_path: Optional[str] = None,
+        artifact: Artifact,
+        manifest_entry: ArtifactManifestEntry,
+        dest_path: str | None = None,
     ) -> FilePathStr:
         if dest_path is not None:
             self._cache._override_cache_path = dest_path
@@ -159,22 +161,22 @@ class WandbStoragePolicy(StoragePolicy):
 
     def store_reference(
         self,
-        artifact: "Artifact",
-        path: Union[URIStr, FilePathStr],
-        name: Optional[str] = None,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: str | None = None,
         checksum: bool = True,
-        max_objects: Optional[int] = None,
-    ) -> Sequence["ArtifactManifestEntry"]:
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
         return self._handler.store_path(
             artifact, path, name=name, checksum=checksum, max_objects=max_objects
         )
 
     def load_reference(
         self,
-        manifest_entry: "ArtifactManifestEntry",
+        manifest_entry: ArtifactManifestEntry,
         local: bool = False,
-        dest_path: Optional[str] = None,
-    ) -> Union[FilePathStr, URIStr]:
+        dest_path: str | None = None,
+    ) -> FilePathStr | URIStr:
         assert manifest_entry.ref is not None
         used_handler = self._handler._get_handler(manifest_entry.ref)
         if hasattr(used_handler, "_cache") and (dest_path is not None):
@@ -185,7 +187,7 @@ class WandbStoragePolicy(StoragePolicy):
         self,
         api: InternalApi,
         entity_name: str,
-        manifest_entry: "ArtifactManifestEntry",
+        manifest_entry: ArtifactManifestEntry,
     ) -> str:
         storage_layout = self._config.get("storageLayout", StorageLayout.V1)
         storage_region = self._config.get("storageRegion", "default")
@@ -214,10 +216,10 @@ class WandbStoragePolicy(StoragePolicy):
         self,
         file_path: str,
         chunk_size: int,
-        hex_digests: Dict[int, str],
-        multipart_urls: Dict[int, str],
-        extra_headers: Dict[str, str],
-    ) -> List[Dict[str, Any]]:
+        hex_digests: dict[int, str],
+        multipart_urls: dict[int, str],
+        extra_headers: dict[str, str],
+    ) -> list[dict[str, Any]]:
         etags = []
         part_number = 1
 
@@ -247,8 +249,8 @@ class WandbStoragePolicy(StoragePolicy):
         self,
         upload_url: str,
         file_path: str,
-        extra_headers: Dict[str, Any],
-        progress_callback: Optional["progress.ProgressFn"] = None,
+        extra_headers: dict[str, Any],
+        progress_callback: progress.ProgressFn | None = None,
     ) -> None:
         """Upload a file to the artifact store and write to cache."""
         with open(file_path, "rb") as file:
@@ -272,9 +274,9 @@ class WandbStoragePolicy(StoragePolicy):
         self,
         artifact_id: str,
         artifact_manifest_id: str,
-        entry: "ArtifactManifestEntry",
-        preparer: "StepPrepare",
-        progress_callback: Optional["progress.ProgressFn"] = None,
+        entry: ArtifactManifestEntry,
+        preparer: StepPrepare,
+        progress_callback: progress.ProgressFn | None = None,
     ) -> bool:
         """Upload a file to the artifact store.
 
@@ -352,7 +354,7 @@ class WandbStoragePolicy(StoragePolicy):
 
         return False
 
-    def _write_cache(self, entry: "ArtifactManifestEntry") -> None:
+    def _write_cache(self, entry: ArtifactManifestEntry) -> None:
         if entry.local_path is None:
             return
 

wandb/sdk/artifacts/storage_policy.py

@@ -1,6 +1,8 @@
 """Storage policy."""
 
-from typing import TYPE_CHECKING, Dict, Optional, Sequence, Type, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
 
 from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.lib.paths import FilePathStr, URIStr
@@ -14,7 +16,7 @@ if TYPE_CHECKING:
 
 class StoragePolicy:
     @classmethod
-    def lookup_by_name(cls, name: str) -> Type["StoragePolicy"]:
+    def lookup_by_name(cls, name: str) -> type[StoragePolicy]:
         import wandb.sdk.artifacts.storage_policies  # noqa: F401
 
         for sub in cls.__subclasses__():
@@ -27,19 +29,17 @@ class StoragePolicy:
         raise NotImplementedError
 
     @classmethod
-    def from_config(
-        cls, config: Dict, api: Optional[InternalApi] = None
-    ) -> "StoragePolicy":
+    def from_config(cls, config: dict, api: InternalApi | None = None) -> StoragePolicy:
         raise NotImplementedError
 
-    def config(self) -> Dict:
+    def config(self) -> dict:
         raise NotImplementedError
 
     def load_file(
         self,
-        artifact: "Artifact",
-        manifest_entry: "ArtifactManifestEntry",
-        dest_path: Optional[str] = None,
+        artifact: Artifact,
+        manifest_entry: ArtifactManifestEntry,
+        dest_path: str | None = None,
     ) -> FilePathStr:
         raise NotImplementedError
 
@@ -47,26 +47,26 @@ class StoragePolicy:
         self,
         artifact_id: str,
         artifact_manifest_id: str,
-        entry: "ArtifactManifestEntry",
-        preparer: "StepPrepare",
-        progress_callback: Optional["ProgressFn"] = None,
+        entry: ArtifactManifestEntry,
+        preparer: StepPrepare,
+        progress_callback: ProgressFn | None = None,
     ) -> bool:
         raise NotImplementedError
 
     def store_reference(
         self,
-        artifact: "Artifact",
-        path: Union[URIStr, FilePathStr],
-        name: Optional[str] = None,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: str | None = None,
         checksum: bool = True,
-        max_objects: Optional[int] = None,
-    ) -> Sequence["ArtifactManifestEntry"]:
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
         raise NotImplementedError
 
     def load_reference(
         self,
-        manifest_entry: "ArtifactManifestEntry",
+        manifest_entry: ArtifactManifestEntry,
         local: bool = False,
-        dest_path: Optional[str] = None,
-    ) -> Union[FilePathStr, URIStr]:
+        dest_path: str | None = None,
+    ) -> FilePathStr | URIStr:
         raise NotImplementedError

wandb/sdk/backend/backend.py

@@ -11,7 +11,7 @@ import os
 import queue
 import sys
 import threading
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
 import wandb
 from wandb.sdk.interface.interface import InterfaceBase
@@ -19,17 +19,16 @@ from wandb.sdk.interface.interface_queue import InterfaceQueue
 from wandb.sdk.internal.internal import wandb_internal
 from wandb.sdk.internal.settings_static import SettingsStatic
 from wandb.sdk.lib.mailbox import Mailbox
-from wandb.sdk.wandb_manager import _Manager
 from wandb.sdk.wandb_settings import Settings
 
 if TYPE_CHECKING:
     from wandb.proto.wandb_internal_pb2 import Record, Result
+    from wandb.sdk.lib import service_connection
 
-    from ..service.service_sock import ServiceSockInterface
     from ..wandb_run import Run
 
-    RecordQueue = Union[queue.Queue[Record], multiprocessing.Queue[Record]]
-    ResultQueue = Union[queue.Queue[Result], multiprocessing.Queue[Result]]
+    RecordQueue = Union["queue.Queue[Record]", multiprocessing.Queue[Record]]
+    ResultQueue = Union["queue.Queue[Result]", multiprocessing.Queue[Result]]
 
 logger = logging.getLogger("wandb")
 
@@ -65,7 +64,7 @@ class Backend:
         mailbox: Mailbox,
         settings: Optional[Settings] = None,
         log_level: Optional[int] = None,
-        manager: Optional[_Manager] = None,
+        service: "Optional[service_connection.ServiceConnection]" = None,
     ) -> None:
         self._done = False
         self.record_q = None
@@ -75,7 +74,7 @@ class Backend:
         self._internal_pid = None
         self._settings = settings
         self._log_level = log_level
-        self._manager = manager
+        self._service = service
         self._mailbox = mailbox
 
         self._multiprocessing = multiprocessing  # type: ignore
@@ -139,27 +138,10 @@ class Backend:
         if self._save_mod_path:
             main_module.__file__ = self._save_mod_path
 
-    def _ensure_launched_manager(self) -> None:
-        assert self._manager
-        svc = self._manager._get_service()
-        assert svc
-        svc_iface = svc.service_interface
-
-        svc_transport = svc_iface.get_transport()
-        if svc_transport == "tcp":
-            from ..interface.interface_sock import InterfaceSock
-
-            svc_iface_sock = cast("ServiceSockInterface", svc_iface)
-            sock_client = svc_iface_sock._get_sock_client()
-            sock_interface = InterfaceSock(sock_client, mailbox=self._mailbox)
-            self.interface = sock_interface
-        else:
-            raise AssertionError(f"Unsupported service transport: {svc_transport}")
-
     def ensure_launched(self) -> None:
         """Launch backend worker if not running."""
-        if self._manager:
-            self._ensure_launched_manager()
+        if self._service:
+            self.interface = self._service.make_interface(self._mailbox)
             return
 
         assert self._settings

wandb/sdk/data_types/audio.py

@@ -0,0 +1,165 @@
+import hashlib
+import os
+from typing import Optional
+
+from wandb import util
+from wandb.sdk.lib import filesystem, runid
+
+from . import _dtypes
+from ._private import MEDIA_TMP
+from .base_types.media import BatchableMedia
+
+
+class Audio(BatchableMedia):
+    """Wandb class for audio clips.
+
+    Arguments:
+        data_or_path: (string or numpy array) A path to an audio file
+            or a numpy array of audio data.
+        sample_rate: (int) Sample rate, required when passing in raw
+            numpy array of audio data.
+        caption: (string) Caption to display with audio.
+    """
+
+    _log_type = "audio-file"
+
+    def __init__(self, data_or_path, sample_rate=None, caption=None):
+        """Accept a path to an audio file or a numpy array of audio data."""
+        super().__init__()
+        self._duration = None
+        self._sample_rate = sample_rate
+        self._caption = caption
+
+        if isinstance(data_or_path, str):
+            if self.path_is_reference(data_or_path):
+                self._path = data_or_path
+                self._sha256 = hashlib.sha256(data_or_path.encode("utf-8")).hexdigest()
+                self._is_tmp = False
+            else:
+                self._set_file(data_or_path, is_tmp=False)
+        else:
+            if sample_rate is None:
+                raise ValueError(
+                    'Argument "sample_rate" is required when instantiating wandb.Audio with raw data.'
+                )
+
+            soundfile = util.get_module(
+                "soundfile",
+                required='Raw audio requires the soundfile package. To get it, run "pip install soundfile"',
+            )
+
+            tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".wav")
+            soundfile.write(tmp_path, data_or_path, sample_rate)
+            self._duration = len(data_or_path) / float(sample_rate)
+
+            self._set_file(tmp_path, is_tmp=True)
+
+    @classmethod
+    def get_media_subdir(cls):
+        return os.path.join("media", "audio")
+
+    @classmethod
+    def from_json(cls, json_obj, source_artifact):
+        return cls(
+            source_artifact.get_entry(json_obj["path"]).download(),
+            caption=json_obj["caption"],
+        )
+
+    def bind_to_run(
+        self, run, key, step, id_=None, ignore_copy_err: Optional[bool] = None
+    ):
+        if self.path_is_reference(self._path):
+            raise ValueError(
+                "Audio media created by a reference to external storage cannot currently be added to a run"
+            )
+
+        return super().bind_to_run(run, key, step, id_, ignore_copy_err)
+
+    def to_json(self, run):
+        json_dict = super().to_json(run)
+        json_dict.update(
+            {
+                "_type": self._log_type,
+                "caption": self._caption,
+            }
+        )
+        return json_dict
+
+    @classmethod
+    def seq_to_json(cls, seq, run, key, step):
+        audio_list = list(seq)
+
+        util.get_module(
+            "soundfile",
+            required="wandb.Audio requires the soundfile package. To get it, run: pip install soundfile",
+        )
+        base_path = os.path.join(run.dir, "media", "audio")
+        filesystem.mkdir_exists_ok(base_path)
+        meta = {
+            "_type": "audio",
+            "count": len(audio_list),
+            "audio": [a.to_json(run) for a in audio_list],
+        }
+        sample_rates = cls.sample_rates(audio_list)
+        if sample_rates:
+            meta["sampleRates"] = sample_rates
+        durations = cls.durations(audio_list)
+        if durations:
+            meta["durations"] = durations
+        captions = cls.captions(audio_list)
+        if captions:
+            meta["captions"] = captions
+
+        return meta
+
+    @classmethod
+    def durations(cls, audio_list):
+        return [a._duration for a in audio_list]
+
+    @classmethod
+    def sample_rates(cls, audio_list):
+        return [a._sample_rate for a in audio_list]
+
+    @classmethod
+    def captions(cls, audio_list):
+        captions = [a._caption for a in audio_list]
+        if all(c is None for c in captions):
+            return False
+        else:
+            return ["" if c is None else c for c in captions]
+
+    def resolve_ref(self):
+        if self.path_is_reference(self._path):
+            # this object was already created using a ref:
+            return self._path
+        source_artifact = self._artifact_source.artifact
+
+        resolved_name = source_artifact._local_path_to_name(self._path)
+        if resolved_name is not None:
+            target_entry = source_artifact.manifest.get_entry_by_path(resolved_name)
+            if target_entry is not None:
+                return target_entry.ref
+
+        return None
+
+    def __eq__(self, other):
+        if self.path_is_reference(self._path) or self.path_is_reference(other._path):
+            # one or more of these objects is an unresolved reference -- we'll compare
+            # their reference paths instead of their SHAs:
+            return (
+                self.resolve_ref() == other.resolve_ref()
+                and self._caption == other._caption
+            )
+
+        return super().__eq__(other) and self._caption == other._caption
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+class _AudioFileType(_dtypes.Type):
+    name = "audio-file"
+    types = [Audio]
+
+
+_dtypes.TypeRegistry.add(_AudioFileType)

wandb/sdk/data_types/base_types/wb_value.py

@@ -88,9 +88,7 @@ class WBValue:
         raise NotImplementedError
 
     @classmethod
-    def from_json(
-        cls: Type["WBValue"], json_obj: dict, source_artifact: "Artifact"
-    ) -> "WBValue":
+    def from_json(cls, json_obj: dict, source_artifact: "Artifact") -> "WBValue":
         """Deserialize a `json_obj` into it's class representation.
 
         If additional resources were stored in the `run_or_artifact` artifact during the

wandb/sdk/data_types/bokeh.py

@@ -0,0 +1,70 @@
+import codecs
+import json
+import os
+
+from wandb import util
+from wandb.sdk.lib import runid
+
+from . import _dtypes
+from ._private import MEDIA_TMP
+from .base_types.media import Media
+
+
+class Bokeh(Media):
+    """Wandb class for Bokeh plots.
+
+    Arguments:
+        val: Bokeh plot
+    """
+
+    _log_type = "bokeh-file"
+
+    def __init__(self, data_or_path):
+        super().__init__()
+        bokeh = util.get_module("bokeh", required=True)
+        if isinstance(data_or_path, str) and os.path.exists(data_or_path):
+            with open(data_or_path) as file:
+                b_json = json.load(file)
+            self.b_obj = bokeh.document.Document.from_json(b_json)
+            self._set_file(data_or_path, is_tmp=False, extension=".bokeh.json")
+        elif isinstance(data_or_path, bokeh.model.Model):
+            _data = bokeh.document.Document()
+            _data.add_root(data_or_path)
+            # serialize/deserialize pairing followed by sorting attributes ensures
+            # that the file's sha's are equivalent in subsequent calls
+            self.b_obj = bokeh.document.Document.from_json(_data.to_json())
+            b_json = self.b_obj.to_json()
+            if "references" in b_json["roots"]:
+                b_json["roots"]["references"].sort(key=lambda x: x["id"])
+
+            tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".bokeh.json")
+            with codecs.open(tmp_path, "w", encoding="utf-8") as fp:
+                util.json_dump_safer(b_json, fp)
+            self._set_file(tmp_path, is_tmp=True, extension=".bokeh.json")
+        elif not isinstance(data_or_path, bokeh.document.Document):
+            raise TypeError(
+                "Bokeh constructor accepts Bokeh document/model or path to Bokeh json file"
+            )
+
+    def get_media_subdir(self):
+        return os.path.join("media", "bokeh")
+
+    def to_json(self, run):
+        # TODO: (tss) this is getting redundant for all the media objects. We can probably
+        # pull this into Media#to_json and remove this type override for all the media types.
+        # There are only a few cases where the type is different between artifacts and runs.
+        json_dict = super().to_json(run)
+        json_dict["_type"] = self._log_type
+        return json_dict
+
+    @classmethod
+    def from_json(cls, json_obj, source_artifact):
+        return cls(source_artifact.get_entry(json_obj["path"]).download())
+
+
+class _BokehFileType(_dtypes.Type):
+    name = "bokeh-file"
+    types = [Bokeh]
+
+
+_dtypes.TypeRegistry.add(_BokehFileType)