wandb 0.19.7__py3-none-macosx_11_0_arm64.whl → 0.19.9__py3-none-macosx_11_0_arm64.whl

wandb/sdk/data_types/base_types/media.py

@@ -3,7 +3,7 @@ import os
 import platform
 import re
 import shutil
-from typing import TYPE_CHECKING, Optional, Sequence, Type, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type, Union, cast
 
 import wandb
 from wandb import util
@@ -222,7 +222,10 @@ class Media(WBValue):
         from wandb.data_types import Audio
         from wandb.sdk.wandb_run import Run
 
-        json_obj = {}
+        json_obj: Dict[str, Any] = {}
+
+        if self._caption is not None:
+            json_obj["caption"] = self._caption
 
         if isinstance(run, Run):
             json_obj.update(
@@ -232,6 +235,7 @@ class Media(WBValue):
                     "size": self._size,
                 }
             )
+
             artifact_entry_url = self._get_artifact_entry_ref_url()
             if artifact_entry_url is not None:
                 json_obj["artifact_path"] = artifact_entry_url
@@ -337,8 +341,11 @@ class BatchableMedia(Media):
     organize files by name in the media directory.
     """
 
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(
+        self,
+        caption: Optional[str] = None,
+    ) -> None:
+        super().__init__(caption=caption)
 
     @classmethod
     def seq_to_json(

wandb/sdk/data_types/image.py

@@ -152,12 +152,11 @@ class Image(BatchableMedia):
         masks: Optional[Union[Dict[str, "ImageMask"], Dict[str, dict]]] = None,
         file_type: Optional[str] = None,
     ) -> None:
-        super().__init__()
+        super().__init__(caption=caption)
         # TODO: We should remove grouping, it's a terrible name and I don't
         # think anyone uses it.
 
         self._grouping = None
-        self._caption = None
         self._width = None
         self._height = None
         self._image = None
@@ -193,9 +192,6 @@ class Image(BatchableMedia):
         if grouping is not None:
             self._grouping = grouping
 
-        if caption is not None:
-            self._caption = caption
-
         total_classes = {}
 
         if boxes:
@@ -297,10 +293,19 @@ class Image(BatchableMedia):
             "PIL.Image",
             required='wandb.Image needs the PIL package. To get it, run "pip install pillow".',
         )
+
+        accepted_formats = ["png", "jpg", "jpeg", "bmp"]
+        self.format = file_type or "png"
+
+        if self.format not in accepted_formats:
+            raise ValueError(f"file_type must be one of {accepted_formats}")
+
+        tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + "." + self.format)
+
         if util.is_matplotlib_typename(util.get_full_typename(data)):
             buf = BytesIO()
-            util.ensure_matplotlib_figure(data).savefig(buf, format="png")
-            self._image = pil_image.open(buf, formats=["PNG"])
+            util.ensure_matplotlib_figure(data).savefig(buf, format=self.format)
+            self._image = pil_image.open(buf)
         elif isinstance(data, pil_image.Image):
             self._image = data
         elif util.is_pytorch_tensor_typename(util.get_full_typename(data)):
@@ -312,26 +317,23 @@ class Image(BatchableMedia):
             if hasattr(data, "dtype") and str(data.dtype) == "torch.uint8":
                 data = data.to(float)
             data = vis_util.make_grid(data, normalize=True)
+            mode = mode or self.guess_mode(data, file_type)
             self._image = pil_image.fromarray(
-                data.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
+                data.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy(),
+                mode=mode,
             )
         else:
             if hasattr(data, "numpy"):  # TF data eager tensors
                 data = data.numpy()
             if data.ndim > 2:
                 data = data.squeeze()  # get rid of trivial dimensions as a convenience
+
+            mode = mode or self.guess_mode(data, file_type)
             self._image = pil_image.fromarray(
-                self.to_uint8(data), mode=mode or self.guess_mode(data)
+                self.to_uint8(data),
+                mode=mode,
             )
-        accepted_formats = ["png", "jpg", "jpeg", "bmp"]
-        if file_type is None:
-            self.format = "png"
-        else:
-            self.format = file_type
-        assert (
-            self.format in accepted_formats
-        ), f"file_type must be one of {accepted_formats}"
-        tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + "." + self.format)
+
         assert self._image is not None
         self._image.save(tmp_path, transparency=None)
         self._set_file(tmp_path, is_tmp=True)
@@ -430,8 +432,6 @@ class Image(BatchableMedia):
             json_dict["height"] = self._height
         if self._grouping:
             json_dict["grouping"] = self._grouping
-        if self._caption:
-            json_dict["caption"] = self._caption
 
         if isinstance(run_or_artifact, wandb.Artifact):
             artifact = run_or_artifact
@@ -471,15 +471,34 @@ class Image(BatchableMedia):
             }
         return json_dict
 
-    def guess_mode(self, data: "np.ndarray") -> str:
+    def guess_mode(
+        self,
+        data: Union["np.ndarray", "torch.Tensor"],
+        file_type: Optional[str] = None,
+    ) -> str:
         """Guess what type of image the np.array is representing."""
         # TODO: do we want to support dimensions being at the beginning of the array?
-        if data.ndim == 2:
+        ndims = data.ndim
+        if util.is_pytorch_tensor_typename(util.get_full_typename(data)):
+            # Torch tenors typically have the channels dimension first
+            num_channels = data.shape[0]
+        else:
+            num_channels = data.shape[-1]
+
+        if ndims == 2:
             return "L"
-        elif data.shape[-1] == 3:
+        elif num_channels == 3:
             return "RGB"
-        elif data.shape[-1] == 4:
-            return "RGBA"
+        elif num_channels == 4:
+            if file_type in ["jpg", "jpeg"]:
+                wandb.termwarn(
+                    "JPEG format does not support transparency. "
+                    "Ignoring alpha channel.",
+                    repeat=False,
+                )
+                return "RGB"
+            else:
+                return "RGBA"
         else:
             raise ValueError(
                 "Un-supported shape for image conversion {}".format(list(data.shape))

wandb/sdk/data_types/molecule.py

@@ -53,9 +53,7 @@ class Molecule(BatchableMedia):
         caption: Optional[str] = None,
         **kwargs: str,
     ) -> None:
-        super().__init__()
-
-        self._caption = caption
+        super().__init__(caption=caption)
 
         if hasattr(data_or_path, "name"):
             # if the file has a path, we just detect the type and copy it from there
@@ -208,8 +206,6 @@ class Molecule(BatchableMedia):
     def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self._log_type
-        if self._caption:
-            json_dict["caption"] = self._caption
         return json_dict
 
     @classmethod

wandb/sdk/data_types/object_3d.py

@@ -215,9 +215,10 @@ class Object3D(BatchableMedia):
     def __init__(
         self,
         data_or_path: Union["np.ndarray", str, "TextIO", dict],
+        caption: Optional[str] = None,
         **kwargs: Optional[Union[str, "FileFormat3D"]],
     ) -> None:
-        super().__init__()
+        super().__init__(caption=caption)
 
         if hasattr(data_or_path, "name"):
             # if the file has a path, we just detect the type and copy it from there

wandb/sdk/data_types/saved_model.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import os
 import shutil
 import sys
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
 
 import wandb
@@ -15,9 +16,6 @@ from ._private import MEDIA_TMP
 from .base_types.wb_value import WBValue
 
 if TYPE_CHECKING:
-    from types import ModuleType
-
-    import cloudpickle  # type: ignore
     import sklearn  # type: ignore
     import tensorflow  # type: ignore
     import torch  # type: ignore
@@ -264,9 +262,9 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         self._serialize(self._model_obj, target_path)
 
 
-def _get_cloudpickle() -> "cloudpickle":
+def _get_cloudpickle() -> ModuleType:
     return cast(
-        "cloudpickle",
+        ModuleType,
         util.get_module("cloudpickle", "ModelAdapter requires `cloudpickle`"),
     )
 
@@ -338,9 +336,9 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
         return json_obj
 
 
-def _get_torch() -> "torch":
+def _get_torch() -> ModuleType:
     return cast(
-        "torch",
+        ModuleType,
         util.get_module("torch", "ModelAdapter requires `torch`"),
     )
 
@@ -366,9 +364,9 @@ class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]):
         )
 
 
-def _get_sklearn() -> "sklearn":
+def _get_sklearn() -> ModuleType:
     return cast(
-        "sklearn",
+        ModuleType,
         util.get_module("sklearn", "ModelAdapter requires `sklearn`"),
     )
 

wandb/sdk/data_types/video.py

@@ -90,13 +90,12 @@ class Video(BatchableMedia):
         fps: Optional[int] = None,
         format: Optional[str] = None,
     ):
-        super().__init__()
+        super().__init__(caption=caption)
 
         self._format = format or "gif"
         self._width = None
         self._height = None
         self._channels = None
-        self._caption = caption
         if self._format not in Video.EXTS:
             raise ValueError(
                 "wandb.Video accepts {} formats".format(", ".join(Video.EXTS))
@@ -190,8 +189,6 @@ class Video(BatchableMedia):
             json_dict["width"] = self._width
         if self._height is not None:
             json_dict["height"] = self._height
-        if self._caption:
-            json_dict["caption"] = self._caption
 
         return json_dict
 

wandb/sdk/interface/interface.py

@@ -4,8 +4,6 @@ InterfaceBase: The abstract class
 InterfaceShared: Common routines for socket and queue based implementations
 InterfaceQueue: Use multiprocessing queues to send and receive messages
 InterfaceSock: Use socket to send and receive messages
-InterfaceRelay: Responses are routed to a relay queue (not matching uuids)
-
 """
 
 import gzip
@@ -102,14 +100,14 @@ class InterfaceBase:
     def _publish_header(self, header: pb.HeaderRecord) -> None:
         raise NotImplementedError
 
-    def deliver_status(self) -> MailboxHandle:
+    def deliver_status(self) -> MailboxHandle[pb.Result]:
         return self._deliver_status(pb.StatusRequest())
 
     @abstractmethod
     def _deliver_status(
         self,
         status: pb.StatusRequest,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     def _make_config(
@@ -435,7 +433,7 @@ class InterfaceBase:
         entity: Optional[str] = None,
         project: Optional[str] = None,
         organization: Optional[str] = None,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         link_artifact = pb.LinkArtifactRequest()
         if artifact.is_draft():
             link_artifact.client_id = artifact._client_id
@@ -452,7 +450,7 @@ class InterfaceBase:
     @abstractmethod
     def _deliver_link_artifact(
         self, link_artifact: pb.LinkArtifactRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     @staticmethod
@@ -570,7 +568,7 @@ class InterfaceBase:
         is_user_created: bool = False,
         use_after_commit: bool = False,
         finalize: bool = True,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         proto_run = self._make_run(run)
         proto_artifact = self._make_artifact(artifact)
         proto_artifact.run_id = proto_run.run_id
@@ -595,7 +593,7 @@ class InterfaceBase:
     def _deliver_artifact(
         self,
         log_artifact: pb.LogArtifactRequest,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     def deliver_download_artifact(
@@ -605,7 +603,7 @@ class InterfaceBase:
         allow_missing_references: bool,
         skip_cache: bool,
         path_prefix: Optional[str],
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         download_artifact = pb.DownloadArtifactRequest()
         download_artifact.artifact_id = artifact_id
         download_artifact.download_root = download_root
@@ -618,7 +616,7 @@ class InterfaceBase:
     @abstractmethod
     def _deliver_download_artifact(
         self, download_artifact: pb.DownloadArtifactRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     def publish_artifact(
@@ -870,7 +868,9 @@ class InterfaceBase:
         return self._publish_job_input(request)
 
     @abstractmethod
-    def _publish_job_input(self, request: pb.JobInputRequest) -> MailboxHandle:
+    def _publish_job_input(
+        self, request: pb.JobInputRequest
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     def join(self) -> None:
@@ -892,143 +892,165 @@ class InterfaceBase:
             logger.warning("handle abandoned while communicating shutdown")
 
     @abstractmethod
-    def _deliver_shutdown(self) -> MailboxHandle:
+    def _deliver_shutdown(self) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_run(self, run: "Run") -> MailboxHandle:
+    def deliver_run(self, run: "Run") -> MailboxHandle[pb.Result]:
         run_record = self._make_run(run)
         return self._deliver_run(run_record)
 
     def deliver_finish_sync(
         self,
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         sync = pb.SyncFinishRequest()
         return self._deliver_finish_sync(sync)
 
     @abstractmethod
-    def _deliver_finish_sync(self, sync: pb.SyncFinishRequest) -> MailboxHandle:
+    def _deliver_finish_sync(
+        self, sync: pb.SyncFinishRequest
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     @abstractmethod
-    def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle:
+    def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_run_start(self, run: "Run") -> MailboxHandle:
+    def deliver_run_start(self, run: "Run") -> MailboxHandle[pb.Result]:
         run_start = pb.RunStartRequest(run=self._make_run(run))
         return self._deliver_run_start(run_start)
 
     @abstractmethod
-    def _deliver_run_start(self, run_start: pb.RunStartRequest) -> MailboxHandle:
+    def _deliver_run_start(
+        self, run_start: pb.RunStartRequest
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_attach(self, attach_id: str) -> MailboxHandle:
+    def deliver_attach(self, attach_id: str) -> MailboxHandle[pb.Result]:
         attach = pb.AttachRequest(attach_id=attach_id)
         return self._deliver_attach(attach)
 
     @abstractmethod
-    def _deliver_attach(self, status: pb.AttachRequest) -> MailboxHandle:
+    def _deliver_attach(
+        self,
+        status: pb.AttachRequest,
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_stop_status(self) -> MailboxHandle:
+    def deliver_stop_status(self) -> MailboxHandle[pb.Result]:
         status = pb.StopStatusRequest()
         return self._deliver_stop_status(status)
 
     @abstractmethod
-    def _deliver_stop_status(self, status: pb.StopStatusRequest) -> MailboxHandle:
+    def _deliver_stop_status(
+        self,
+        status: pb.StopStatusRequest,
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_network_status(self) -> MailboxHandle:
+    def deliver_network_status(self) -> MailboxHandle[pb.Result]:
         status = pb.NetworkStatusRequest()
         return self._deliver_network_status(status)
 
     @abstractmethod
-    def _deliver_network_status(self, status: pb.NetworkStatusRequest) -> MailboxHandle:
+    def _deliver_network_status(
+        self,
+        status: pb.NetworkStatusRequest,
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_internal_messages(self) -> MailboxHandle:
+    def deliver_internal_messages(self) -> MailboxHandle[pb.Result]:
         internal_message = pb.InternalMessagesRequest()
         return self._deliver_internal_messages(internal_message)
 
     @abstractmethod
     def _deliver_internal_messages(
         self, internal_message: pb.InternalMessagesRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_get_summary(self) -> MailboxHandle:
+    def deliver_get_summary(self) -> MailboxHandle[pb.Result]:
         get_summary = pb.GetSummaryRequest()
         return self._deliver_get_summary(get_summary)
 
     @abstractmethod
-    def _deliver_get_summary(self, get_summary: pb.GetSummaryRequest) -> MailboxHandle:
+    def _deliver_get_summary(
+        self,
+        get_summary: pb.GetSummaryRequest,
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_get_system_metrics(self) -> MailboxHandle:
+    def deliver_get_system_metrics(self) -> MailboxHandle[pb.Result]:
         get_system_metrics = pb.GetSystemMetricsRequest()
         return self._deliver_get_system_metrics(get_system_metrics)
 
     @abstractmethod
     def _deliver_get_system_metrics(
         self, get_summary: pb.GetSystemMetricsRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_get_system_metadata(self) -> MailboxHandle:
+    def deliver_get_system_metadata(self) -> MailboxHandle[pb.Result]:
         get_system_metadata = pb.GetSystemMetadataRequest()
         return self._deliver_get_system_metadata(get_system_metadata)
 
     @abstractmethod
     def _deliver_get_system_metadata(
         self, get_system_metadata: pb.GetSystemMetadataRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_exit(self, exit_code: Optional[int]) -> MailboxHandle:
+    def deliver_exit(self, exit_code: Optional[int]) -> MailboxHandle[pb.Result]:
         exit_data = self._make_exit(exit_code)
         return self._deliver_exit(exit_data)
 
     @abstractmethod
-    def _deliver_exit(self, exit_data: pb.RunExitRecord) -> MailboxHandle:
+    def _deliver_exit(
+        self,
+        exit_data: pb.RunExitRecord,
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
     @abstractmethod
-    def deliver_operation_stats(self) -> MailboxHandle:
+    def deliver_operation_stats(self) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_poll_exit(self) -> MailboxHandle:
+    def deliver_poll_exit(self) -> MailboxHandle[pb.Result]:
         poll_exit = pb.PollExitRequest()
         return self._deliver_poll_exit(poll_exit)
 
     @abstractmethod
-    def _deliver_poll_exit(self, poll_exit: pb.PollExitRequest) -> MailboxHandle:
+    def _deliver_poll_exit(
+        self,
+        poll_exit: pb.PollExitRequest,
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_finish_without_exit(self) -> MailboxHandle:
+    def deliver_finish_without_exit(self) -> MailboxHandle[pb.Result]:
         run_finish_without_exit = pb.RunFinishWithoutExitRequest()
         return self._deliver_finish_without_exit(run_finish_without_exit)
 
     @abstractmethod
     def _deliver_finish_without_exit(
         self, run_finish_without_exit: pb.RunFinishWithoutExitRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_request_sampled_history(self) -> MailboxHandle:
+    def deliver_request_sampled_history(self) -> MailboxHandle[pb.Result]:
         sampled_history = pb.SampledHistoryRequest()
         return self._deliver_request_sampled_history(sampled_history)
 
     @abstractmethod
     def _deliver_request_sampled_history(
         self, sampled_history: pb.SampledHistoryRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    def deliver_request_run_status(self) -> MailboxHandle:
+    def deliver_request_run_status(self) -> MailboxHandle[pb.Result]:
         run_status = pb.RunStatusRequest()
         return self._deliver_request_run_status(run_status)
 
     @abstractmethod
     def _deliver_request_run_status(
         self, run_status: pb.RunStatusRequest
-    ) -> MailboxHandle:
+    ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError

wandb/sdk/interface/interface_queue.py

@@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Optional
 from wandb.sdk.mailbox import Mailbox
 
 from .interface_shared import InterfaceShared
-from .router_queue import MessageQueueRouter
 
 if TYPE_CHECKING:
     from queue import Queue
@@ -35,12 +34,6 @@ class InterfaceQueue(InterfaceShared):
         self._process = process
         super().__init__(mailbox=mailbox)
 
-    def _init_router(self) -> None:
-        if self.record_q and self.result_q:
-            self._router = MessageQueueRouter(
-                self.record_q, self.result_q, mailbox=self._mailbox
-            )
-
     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_relay.py

@@ -1,17 +1,17 @@
 """InterfaceRelay - Derived from InterfaceQueue using RelayRouter to preserve uuid req/resp.
 
 See interface.py for how interface classes relate to each other.
-
 """
 
+from __future__ import annotations
+
 import logging
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.sdk.mailbox import Mailbox
 
 from .interface_queue import InterfaceQueue
-from .router_relay import MessageRelayRouter
 
 if TYPE_CHECKING:
     from queue import Queue
@@ -22,14 +22,13 @@ logger = logging.getLogger("wandb")
 
 class InterfaceRelay(InterfaceQueue):
     _mailbox: Mailbox
-    relay_q: Optional["Queue[pb.Result]"]
 
     def __init__(
         self,
         mailbox: Mailbox,
-        record_q: Optional["Queue[pb.Record]"] = None,
-        result_q: Optional["Queue[pb.Result]"] = None,
-        relay_q: Optional["Queue[pb.Result]"] = None,
+        record_q: Queue[pb.Record],
+        result_q: Queue[pb.Result],
+        relay_q: Queue[pb.Result],
     ) -> None:
         self.relay_q = relay_q
         super().__init__(
@@ -37,12 +36,3 @@ class InterfaceRelay(InterfaceQueue):
             result_q=result_q,
             mailbox=mailbox,
         )
-
-    def _init_router(self) -> None:
-        if self.record_q and self.result_q and self.relay_q:
-            self._router = MessageRelayRouter(
-                request_queue=self.record_q,
-                response_queue=self.result_q,
-                relay_queue=self.relay_q,
-                mailbox=self._mailbox,
-            )