wandb 0.19.7__py3-none-any.whl → 0.19.8__py3-none-any.whl

wandb/sdk/artifacts/artifact.py

@@ -41,6 +41,7 @@ from wandb.apis.normalize import normalize_exceptions
 from wandb.apis.public import ArtifactCollection, ArtifactFiles, RetryingClient, Run
 from wandb.data_types import WBValue
 from wandb.errors.term import termerror, termlog, termwarn
+from wandb.proto import wandb_internal_pb2 as pb
 from wandb.sdk.artifacts._graphql_fragments import _gql_artifact_fragment
 from wandb.sdk.artifacts._validators import (
     ensure_logged,
@@ -73,7 +74,7 @@ from wandb.sdk.lib.deprecate import Deprecated, deprecate
 from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64
 from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
 from wandb.sdk.lib.runid import generate_id
-from wandb.sdk.mailbox import Mailbox, MailboxHandle
+from wandb.sdk.mailbox import MailboxHandle
 
 reset_path = util.vendor_setup()
 
@@ -117,6 +118,9 @@ class Artifact:
             description as markdown in the W&B App.
         metadata: Additional information about an artifact. Specify metadata as a
             dictionary of key-value pairs. You can specify no more than 100 total keys.
+        incremental: Use `Artifact.new_draft()` method instead to modify an
+            existing artifact.
+        use_as: W&B Launch specific parameter. Not recommended for general use.
 
     Returns:
         An `Artifact` object.
@@ -158,7 +162,7 @@ class Artifact:
         self._tmp_dir: tempfile.TemporaryDirectory | None = None
         self._added_objs: dict[int, tuple[WBValue, ArtifactManifestEntry]] = {}
         self._added_local_paths: dict[str, ArtifactManifestEntry] = {}
-        self._save_handle: MailboxHandle | None = None
+        self._save_handle: MailboxHandle[pb.Result] | None = None
         self._download_roots: set[str] = set()
         # Set by new_draft(), otherwise the latest artifact will be used as the base.
         self._base_id: str | None = None
@@ -421,7 +425,9 @@ class Artifact:
     def new_draft(self) -> Artifact:
         """Create a new draft artifact with the same content as this committed artifact.
 
-        The artifact returned can be extended or modified and logged as a new version.
+        Modifying an existing artifact creates a new artifact version known
+        as an "incremental artifact". The artifact returned can be extended or
+        modified and logged as a new version.
 
         Returns:
             An `Artifact` object.
@@ -956,7 +962,7 @@ class Artifact:
 
     def _set_save_handle(
         self,
-        save_handle: MailboxHandle,
+        save_handle: MailboxHandle[pb.Result],
         client: RetryingClient,
     ) -> None:
         self._save_handle = save_handle
@@ -1828,12 +1834,7 @@ class Artifact:
             service = wl.ensure_service()
             service.inform_init(settings=settings, run_id=stream_id)
 
-            mailbox = Mailbox()
-            backend = Backend(
-                settings=wl.settings,
-                service=service,
-                mailbox=mailbox,
-            )
+            backend = Backend(settings=wl.settings, service=service)
             backend.ensure_launched()
 
             assert backend.interface

wandb/sdk/backend/backend.py

@@ -16,6 +16,7 @@ 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.mailbox import Mailbox
@@ -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/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,
-            )