wandb 0.19.6rc4__py3-none-macosx_11_0_arm64.whl → 0.19.8__py3-none-macosx_11_0_arm64.whl

wandb/sdk/backend/backend.py

@@ -16,9 +16,10 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 import wandb
 from wandb.sdk.interface.interface import InterfaceBase
 from wandb.sdk.interface.interface_queue import InterfaceQueue
+from wandb.sdk.interface.router_queue import MessageQueueRouter
 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.mailbox import Mailbox
 from wandb.sdk.wandb_settings import Settings
 
 if TYPE_CHECKING:
@@ -49,17 +50,18 @@ class BackendThread(threading.Thread):
 class Backend:
     # multiprocessing context or module
     _multiprocessing: multiprocessing.context.BaseContext
+
     interface: Optional[InterfaceBase]
+    _router: Optional[MessageQueueRouter]
+
     _internal_pid: Optional[int]
     wandb_process: Optional[multiprocessing.process.BaseProcess]
     _settings: Settings
     record_q: Optional["RecordQueue"]
     result_q: Optional["ResultQueue"]
-    _mailbox: Mailbox
 
     def __init__(
         self,
-        mailbox: Mailbox,
         settings: Settings,
         log_level: Optional[int] = None,
         service: "Optional[service_connection.ServiceConnection]" = None,
@@ -68,12 +70,14 @@ class Backend:
         self.record_q = None
         self.result_q = None
         self.wandb_process = None
+
         self.interface = None
+        self._router = None
+
         self._internal_pid = None
         self._settings = settings
         self._log_level = log_level
         self._service = service
-        self._mailbox = mailbox
 
         self._multiprocessing = multiprocessing  # type: ignore
         self._multiprocessing_setup()
@@ -136,7 +140,6 @@ class Backend:
         if self._service:
             assert self._settings.run_id
             self.interface = self._service.make_interface(
-                self._mailbox,
                 stream_id=self._settings.run_id,
             )
             return
@@ -190,11 +193,17 @@ class Backend:
 
         self._module_main_uninstall()
 
+        mailbox = Mailbox()
         self.interface = InterfaceQueue(
             process=self.wandb_process,
             record_q=self.record_q,  # type: ignore
             result_q=self.result_q,  # type: ignore
-            mailbox=self._mailbox,
+            mailbox=mailbox,
+        )
+        self._router = MessageQueueRouter(
+            request_queue=self.record_q,  # type: ignore
+            response_queue=self.result_q,  # type: ignore
+            mailbox=mailbox,
         )
 
     def server_status(self) -> None:
@@ -207,6 +216,8 @@ class Backend:
         self._done = True
         if self.interface:
             self.interface.join()
+        if self._router:
+            self._router.join()
         if self.wandb_process:
             self.wandb_process.join()
 

wandb/sdk/data_types/helper_types/bounding_boxes_2d.py

@@ -53,7 +53,7 @@ class BoundingBoxes2D(JSONMetadata):
         import numpy as np
         import wandb
 
-        wandb.init()
+        run = wandb.init()
         image = np.random.randint(low=0, high=256, size=(200, 300, 3))
 
         class_labels = {0: "person", 1: "car", 2: "road", 3: "building"}
@@ -77,7 +77,11 @@ class BoundingBoxes2D(JSONMetadata):
                         },
                         {
                             # another box expressed in the pixel domain
-                            "position": {"middle": [150, 20], "width": 68, "height": 112},
+                            "position": {
+                                "middle": [150, 20],
+                                "width": 68,
+                                "height": 112,
+                            },
                             "domain": "pixel",
                             "class_id": 3,
                             "box_caption": "a building",
@@ -90,7 +94,7 @@ class BoundingBoxes2D(JSONMetadata):
             },
         )
 
-        wandb.log({"driving_scene": img})
+        run.log({"driving_scene": img})
         ```
 
         ### Log a bounding box overlay to a Table
@@ -99,7 +103,7 @@ class BoundingBoxes2D(JSONMetadata):
         import numpy as np
         import wandb
 
-        wandb.init()
+        run = wandb.init()
         image = np.random.randint(low=0, high=256, size=(200, 300, 3))
 
         class_labels = {0: "person", 1: "car", 2: "road", 3: "building"}
@@ -132,7 +136,11 @@ class BoundingBoxes2D(JSONMetadata):
                         },
                         {
                             # another box expressed in the pixel domain
-                            "position": {"middle": [150, 20], "width": 68, "height": 112},
+                            "position": {
+                                "middle": [150, 20],
+                                "width": 68,
+                                "height": 112,
+                            },
                             "domain": "pixel",
                             "class_id": 3,
                             "box_caption": "a building",
@@ -148,7 +156,7 @@ class BoundingBoxes2D(JSONMetadata):
 
         table = wandb.Table(columns=["image"])
         table.add_data(img)
-        wandb.log({"driving_scene": table})
+        run.log({"driving_scene": table})
         ```
     """
 

wandb/sdk/data_types/helper_types/image_mask.py

@@ -38,7 +38,7 @@ class ImageMask(Media):
         import numpy as np
         import wandb
 
-        wandb.init()
+        run = wandb.init()
         image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
         predicted_mask = np.empty((100, 100), dtype=np.uint8)
         ground_truth_mask = np.empty((100, 100), dtype=np.uint8)
@@ -58,14 +58,17 @@ class ImageMask(Media):
         masked_image = wandb.Image(
             image,
             masks={
-                "predictions": {"mask_data": predicted_mask, "class_labels": class_labels},
+                "predictions": {
+                    "mask_data": predicted_mask,
+                    "class_labels": class_labels,
+                },
                 "ground_truth": {
                     "mask_data": ground_truth_mask,
                     "class_labels": class_labels,
                 },
             },
         )
-        wandb.log({"img_with_masks": masked_image})
+        run.log({"img_with_masks": masked_image})
         ```
 
         ### Log a masked image inside a Table
@@ -74,7 +77,7 @@ class ImageMask(Media):
         import numpy as np
         import wandb
 
-        wandb.init()
+        run = wandb.init()
         image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
         predicted_mask = np.empty((100, 100), dtype=np.uint8)
         ground_truth_mask = np.empty((100, 100), dtype=np.uint8)
@@ -103,7 +106,10 @@ class ImageMask(Media):
         masked_image = wandb.Image(
             image,
             masks={
-                "predictions": {"mask_data": predicted_mask, "class_labels": class_labels},
+                "predictions": {
+                    "mask_data": predicted_mask,
+                    "class_labels": class_labels,
+                },
                 "ground_truth": {
                     "mask_data": ground_truth_mask,
                     "class_labels": class_labels,
@@ -114,7 +120,7 @@ class ImageMask(Media):
 
         table = wandb.Table(columns=["image"])
         table.add_data(masked_image)
-        wandb.log({"random_field": table})
+        run.log({"random_field": table})
         ```
     """
 

wandb/sdk/data_types/saved_model.py

@@ -1,18 +1,9 @@
+from __future__ import annotations
+
 import os
 import shutil
 import sys
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    ClassVar,
-    Generic,
-    List,
-    Optional,
-    Type,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
 
 import wandb
 from wandb import util
@@ -23,22 +14,24 @@ from wandb.sdk.lib.paths import LogicalPath
 from ._private import MEDIA_TMP
 from .base_types.wb_value import WBValue
 
-if TYPE_CHECKING:  # pragma: no cover
+if TYPE_CHECKING:
+    from types import ModuleType
+
     import cloudpickle  # type: ignore
     import sklearn  # type: ignore
     import tensorflow  # type: ignore
     import torch  # type: ignore
+    from typing_extensions import Self
 
     from wandb.sdk.artifacts.artifact import Artifact
-
-    from ..wandb_run import Run as LocalRun
+    from wandb.sdk.wandb_run import Run as LocalRun
 
 
 DEBUG_MODE = False
 
 
 def _add_deterministic_dir_to_artifact(
-    artifact: "Artifact", dir_name: str, target_dir_root: str
+    artifact: Artifact, dir_name: str, target_dir_root: str
 ) -> str:
     file_paths = []
     for dirpath, _, filenames in os.walk(dir_name, topdown=True):
@@ -50,7 +43,7 @@ def _add_deterministic_dir_to_artifact(
     return target_path
 
 
-def _load_dir_from_artifact(source_artifact: "Artifact", path: str) -> str:
+def _load_dir_from_artifact(source_artifact: Artifact, path: str) -> str:
     dl_path = None
 
     # Look through the entire manifest to find all of the files in the directory.
@@ -79,14 +72,12 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
     _log_type: ClassVar[str]
     _path_extension: ClassVar[str]
 
-    _model_obj: Optional["SavedModelObjType"]
-    _path: Optional[str]
-    _input_obj_or_path: Union[SavedModelObjType, str]
+    _model_obj: SavedModelObjType | None
+    _path: str | None
+    _input_obj_or_path: SavedModelObjType | str
 
     # Public Methods
-    def __init__(
-        self, obj_or_path: Union[SavedModelObjType, str], **kwargs: Any
-    ) -> None:
+    def __init__(self, obj_or_path: SavedModelObjType | str, **kwargs: Any) -> None:
         super().__init__()
         if self.__class__ == _SavedModel:
             raise TypeError(
@@ -115,7 +106,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
             self._unset_obj()
 
     @staticmethod
-    def init(obj_or_path: Any, **kwargs: Any) -> "_SavedModel":
+    def init(obj_or_path: Any, **kwargs: Any) -> _SavedModel:
         maybe_instance = _SavedModel._maybe_init(obj_or_path, **kwargs)
         if maybe_instance is None:
             raise ValueError(
@@ -125,8 +116,8 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
 
     @classmethod
     def from_json(
-        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "Artifact"
-    ) -> "_SavedModel":
+        cls: type[_SavedModel], json_obj: dict, source_artifact: Artifact
+    ) -> _SavedModel:
         path = json_obj["path"]
 
         # First, if the entry is a file, the download it.
@@ -143,7 +134,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         # and specified adapter.
         return cls(dl_path)
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+    def to_json(self, run_or_artifact: LocalRun | Artifact) -> dict:
         # Unlike other data types, we do not allow adding to a Run directly. There is a
         # bit of tech debt in the other data types which requires the input to `to_json`
         # to accept a Run or Artifact. However, Run additions should be deprecated in the future.
@@ -218,8 +209,8 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
     # Private Class Methods
     @classmethod
     def _maybe_init(
-        cls: Type["_SavedModel"], obj_or_path: Any, **kwargs: Any
-    ) -> Optional["_SavedModel"]:
+        cls: type[_SavedModel], obj_or_path: Any, **kwargs: Any
+    ) -> _SavedModel | None:
         # _maybe_init is an exception-safe method that will return an instance of this class
         # (or any subclass of this class - recursively) OR None if no subclass constructor is found.
         # We first try the current class, then recursively call this method on children classes. This pattern
@@ -241,7 +232,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         return None
 
     @classmethod
-    def _tmp_path(cls: Type["_SavedModel"]) -> str:
+    def _tmp_path(cls: type[_SavedModel]) -> str:
         # Generates a tmp path under our MEDIA_TMP directory which confirms to the file
         # or folder preferences of the class.
         assert isinstance(cls._path_extension, str), "_path_extension must be a string"
@@ -286,13 +277,13 @@ PicklingSavedModelObjType = TypeVar("PicklingSavedModelObjType")
 
 
 class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
-    _dep_py_files: Optional[List[str]] = None
-    _dep_py_files_path: Optional[str] = None
+    _dep_py_files: list[str] | None = None
+    _dep_py_files_path: str | None = None
 
     def __init__(
         self,
-        obj_or_path: Union[SavedModelObjType, str],
-        dep_py_files: Optional[List[str]] = None,
+        obj_or_path: SavedModelObjType | str,
+        dep_py_files: list[str] | None = None,
     ):
         super().__init__(obj_or_path)
         if self.__class__ == _PicklingSavedModel:
@@ -319,9 +310,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
                     raise ValueError(f"Invalid dependency file: {extra_file}")
 
     @classmethod
-    def from_json(
-        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "Artifact"
-    ) -> "_PicklingSavedModel":
+    def from_json(cls, json_obj: dict, source_artifact: Artifact) -> Self:
         backup_path = [p for p in sys.path]
         if (
             "dep_py_files_path" in json_obj
@@ -337,7 +326,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
 
         return inst  # type: ignore
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+    def to_json(self, run_or_artifact: LocalRun | Artifact) -> dict:
         json_obj = super().to_json(run_or_artifact)
         assert isinstance(run_or_artifact, wandb.Artifact)
         if self._dep_py_files_path is not None:
@@ -361,7 +350,7 @@ class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]):
     _path_extension = "pt"
 
     @staticmethod
-    def _deserialize(dir_or_file_path: str) -> "torch.nn.Module":
+    def _deserialize(dir_or_file_path: str) -> torch.nn.Module:
         return _get_torch().load(dir_or_file_path, weights_only=False)
 
     @staticmethod
@@ -369,7 +358,7 @@ class _PytorchSavedModel(_PicklingSavedModel["torch.nn.Module"]):
         return isinstance(obj, _get_torch().nn.Module)
 
     @staticmethod
-    def _serialize(model_obj: "torch.nn.Module", dir_or_file_path: str) -> None:
+    def _serialize(model_obj: torch.nn.Module, dir_or_file_path: str) -> None:
         _get_torch().save(
             model_obj,
             dir_or_file_path,
@@ -391,7 +380,7 @@ class _SklearnSavedModel(_PicklingSavedModel["sklearn.base.BaseEstimator"]):
     @staticmethod
     def _deserialize(
         dir_or_file_path: str,
-    ) -> "sklearn.base.BaseEstimator":
+    ) -> sklearn.base.BaseEstimator:
         with open(dir_or_file_path, "rb") as file:
             model = _get_cloudpickle().load(file)
         return model
@@ -410,16 +399,16 @@ class _SklearnSavedModel(_PicklingSavedModel["sklearn.base.BaseEstimator"]):
 
     @staticmethod
     def _serialize(
-        model_obj: "sklearn.base.BaseEstimator", dir_or_file_path: str
+        model_obj: sklearn.base.BaseEstimator, dir_or_file_path: str
     ) -> None:
         dynamic_cloudpickle = _get_cloudpickle()
         with open(dir_or_file_path, "wb") as file:
             dynamic_cloudpickle.dump(model_obj, file)
 
 
-def _get_tf_keras() -> "tensorflow.keras":
+def _get_tf_keras() -> ModuleType:
     return cast(
-        "tensorflow",
+        ModuleType,
         util.get_module("tensorflow", "ModelAdapter requires `tensorflow`"),
     ).keras
 
@@ -431,7 +420,7 @@ class _TensorflowKerasSavedModel(_SavedModel["tensorflow.keras.Model"]):
     @staticmethod
     def _deserialize(
         dir_or_file_path: str,
-    ) -> "tensorflow.keras.Model":
+    ) -> tensorflow.keras.Model:
         return _get_tf_keras().models.load_model(dir_or_file_path)
 
     @staticmethod
@@ -439,7 +428,7 @@ class _TensorflowKerasSavedModel(_SavedModel["tensorflow.keras.Model"]):
         return isinstance(obj, _get_tf_keras().models.Model)
 
     @staticmethod
-    def _serialize(model_obj: "tensorflow.keras.Model", dir_or_file_path: str) -> None:
+    def _serialize(model_obj: tensorflow.keras.Model, dir_or_file_path: str) -> None:
         _get_tf_keras().models.save_model(
             model_obj, dir_or_file_path, include_optimizer=True
         )

wandb/sdk/data_types/video.py

@@ -71,10 +71,10 @@ class Video(BatchableMedia):
         import numpy as np
         import wandb
 
-        wandb.init()
+        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)
-        wandb.log({"video": wandb.Video(frames, fps=4)})
+        run.log({"video": wandb.Video(frames, fps=4)})
         ```
     """
 
@@ -138,20 +138,11 @@ class Video(BatchableMedia):
             self.encode(fps=fps)
 
     def encode(self, fps: int = 4) -> None:
-        # Try to import ImageSequenceClip from the appropriate MoviePy module
-        mpy = None
-        try:
-            # Attempt to load moviepy.editor for MoviePy < 2.0
-            mpy = util.get_module(
-                "moviepy.editor",
-                required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
-            )
-        except wandb.Error:
-            # Fallback to moviepy for MoviePy >= 2.0
-            mpy = util.get_module(
-                "moviepy",
-                required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
-            )
+        # import ImageSequenceClip from the appropriate MoviePy module
+        mpy = util.get_module(
+            "moviepy.video.io.ImageSequenceClip",
+            required='wandb.Video requires moviepy when passing raw data. Install with "pip install wandb[media]"',
+        )
 
         tensor = self._prepare_video(self.data)
         _, self._height, self._width, self._channels = tensor.shape  # type: ignore