wandb 0.18.1__py3-none-macosx_11_0_arm64.whl → 0.18.3__py3-none-macosx_11_0_arm64.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/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/video.py

@@ -34,7 +34,7 @@ def write_gif_with_image_io(
 ) -> None:
     imageio = util.get_module(
         "imageio",
-        required='wandb.Video requires imageio when passing raw data. Install with "pip install imageio"',
+        required='wandb.Video requires imageio when passing raw data. Install with "pip install wandb[media]"',
     )
 
     writer = imageio.save(filename, fps=clip.fps, quantizer=0, palettesize=256, loop=0)
@@ -130,7 +130,7 @@ class Video(BatchableMedia):
     def encode(self) -> None:
         mpy = util.get_module(
             "moviepy.editor",
-            required='wandb.Video requires moviepy and imageio when passing raw data.  Install with "pip install moviepy imageio"',
+            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

wandb/sdk/interface/interface.py

@@ -891,20 +891,6 @@ class InterfaceBase:
     def _deliver_attach(self, status: pb.AttachRequest) -> MailboxHandle:
         raise NotImplementedError
 
-    def deliver_check_version(
-        self, current_version: Optional[str] = None
-    ) -> MailboxHandle:
-        check_version = pb.CheckVersionRequest()
-        if current_version:
-            check_version.current_version = current_version
-        return self._deliver_check_version(check_version)
-
-    @abstractmethod
-    def _deliver_check_version(
-        self, check_version: pb.CheckVersionRequest
-    ) -> MailboxHandle:
-        raise NotImplementedError
-
     def deliver_stop_status(self) -> MailboxHandle:
         status = pb.StopStatusRequest()
         return self._deliver_stop_status(status)
@@ -965,16 +951,6 @@ class InterfaceBase:
     def _deliver_poll_exit(self, poll_exit: pb.PollExitRequest) -> MailboxHandle:
         raise NotImplementedError
 
-    def deliver_request_server_info(self) -> MailboxHandle:
-        server_info = pb.ServerInfoRequest()
-        return self._deliver_request_server_info(server_info)
-
-    @abstractmethod
-    def _deliver_request_server_info(
-        self, server_info: pb.ServerInfoRequest
-    ) -> MailboxHandle:
-        raise NotImplementedError
-
     def deliver_request_sampled_history(self) -> MailboxHandle:
         sampled_history = pb.SampledHistoryRequest()
         return self._deliver_request_sampled_history(sampled_history)

wandb/sdk/interface/interface_shared.py

@@ -490,12 +490,6 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(attach=attach)
         return self._deliver_record(record)
 
-    def _deliver_check_version(
-        self, check_version: pb.CheckVersionRequest
-    ) -> MailboxHandle:
-        record = self._make_request(check_version=check_version)
-        return self._deliver_record(record)
-
     def _deliver_network_status(
         self, network_status: pb.NetworkStatusRequest
     ) -> MailboxHandle:
@@ -508,12 +502,6 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(internal_messages=internal_message)
         return self._deliver_record(record)
 
-    def _deliver_request_server_info(
-        self, server_info: pb.ServerInfoRequest
-    ) -> MailboxHandle:
-        record = self._make_request(server_info=server_info)
-        return self._deliver_record(record)
-
     def _deliver_request_sampled_history(
         self, sampled_history: pb.SampledHistoryRequest
     ) -> MailboxHandle:

wandb/sdk/internal/handler.py

@@ -659,13 +659,6 @@ class HandleManager:
     def handle_footer(self, record: Record) -> None:
         self._dispatch_record(record)
 
-    def handle_request_check_version(self, record: Record) -> None:
-        if self._settings._offline:
-            result = proto_util._result_from_record(record)
-            self._respond_result(result)
-        else:
-            self._dispatch_record(record)
-
     def handle_request_attach(self, record: Record) -> None:
         result = proto_util._result_from_record(record)
         attach_id = record.request.attach.attach_id
@@ -862,9 +855,6 @@ class HandleManager:
             result.response.sampled_history_response.item.append(item)
         self._respond_result(result)
 
-    def handle_request_server_info(self, record: Record) -> None:
-        self._dispatch_record(record, always_send=True)
-
     def handle_request_keepalive(self, record: Record) -> None:
         """Handle a keepalive request.
 

wandb/sdk/internal/internal_api.py

@@ -53,6 +53,8 @@ from .progress import Progress
 
 logger = logging.getLogger(__name__)
 
+LAUNCH_DEFAULT_PROJECT = "model-registry"
+
 if TYPE_CHECKING:
     if sys.version_info >= (3, 8):
         from typing import Literal, TypedDict
@@ -674,6 +676,11 @@ class Api:
             self.server_create_run_queue_supports_priority,
         )
 
+    @normalize_exceptions
+    def upsert_run_queue_introspection(self) -> bool:
+        _, _, mutations = self.server_info_introspection()
+        return "upsertRunQueue" in mutations
+
     @normalize_exceptions
     def push_to_run_queue_introspection(self) -> Tuple[bool, bool]:
         query_string = """
@@ -1580,6 +1587,70 @@ class Api:
         ]
         return result
 
+    @normalize_exceptions
+    def upsert_run_queue(
+        self,
+        queue_name: str,
+        entity: str,
+        resource_type: str,
+        resource_config: dict,
+        project: str = LAUNCH_DEFAULT_PROJECT,
+        prioritization_mode: Optional[str] = None,
+        template_variables: Optional[dict] = None,
+        external_links: Optional[dict] = None,
+    ) -> Optional[Dict[str, Any]]:
+        if not self.upsert_run_queue_introspection():
+            raise UnsupportedError(
+                "upserting run queues is not supported by this version of "
+                "wandb server. Consider updating to the latest version."
+            )
+        query = gql(
+            """
+            mutation upsertRunQueue(
+                $entityName: String!
+                $projectName: String!
+                $queueName: String!
+                $resourceType: String!
+                $resourceConfig: JSONString!
+                $templateVariables: JSONString
+                $prioritizationMode: RunQueuePrioritizationMode
+                $externalLinks: JSONString
+                $clientMutationId: String
+            ) {
+                upsertRunQueue(
+                    input: {
+                        entityName: $entityName
+                        projectName: $projectName
+                        queueName: $queueName
+                        resourceType: $resourceType
+                        resourceConfig: $resourceConfig
+                        templateVariables: $templateVariables
+                        prioritizationMode: $prioritizationMode
+                        externalLinks: $externalLinks
+                        clientMutationId: $clientMutationId
+                    }
+                ) {
+                    success
+                    configSchemaValidationErrors
+                }
+            }
+            """
+        )
+        variable_values = {
+            "entityName": entity,
+            "projectName": project,
+            "queueName": queue_name,
+            "resourceType": resource_type,
+            "resourceConfig": json.dumps(resource_config),
+            "templateVariables": (
+                json.dumps(template_variables) if template_variables else None
+            ),
+            "prioritizationMode": prioritization_mode,
+            "externalLinks": json.dumps(external_links) if external_links else None,
+        }
+        result: Dict[str, Any] = self.gql(query, variable_values)
+        return result["upsertRunQueue"]
+
     @normalize_exceptions
     def push_to_run_queue_by_name(
         self,

wandb/sdk/internal/sender.py

@@ -42,7 +42,6 @@ from wandb.sdk.internal import (
     file_stream,
     internal_api,
     sender_config,
-    update,
 )
 from wandb.sdk.internal.file_pusher import FilePusher
 from wandb.sdk.internal.job_builder import JobBuilder
@@ -51,7 +50,6 @@ from wandb.sdk.lib import (
     config_util,
     filenames,
     filesystem,
-    printer,
     proto_util,
     redirect,
     telemetry,
@@ -483,25 +481,6 @@ class SendManager:
         # make sure that we always update writer for every sended read request
         self._maybe_report_status(always=True)
 
-    def send_request_check_version(self, record: "Record") -> None:
-        assert record.control.req_resp or record.control.mailbox_slot
-        result = proto_util._result_from_record(record)
-        current_version = (
-            record.request.check_version.current_version or wandb.__version__
-        )
-        messages = update.check_available(current_version)
-        if messages:
-            upgrade_message = messages.get("upgrade_message")
-            if upgrade_message:
-                result.response.check_version_response.upgrade_message = upgrade_message
-            yank_message = messages.get("yank_message")
-            if yank_message:
-                result.response.check_version_response.yank_message = yank_message
-            delete_message = messages.get("delete_message")
-            if delete_message:
-                result.response.check_version_response.delete_message = delete_message
-        self._respond_result(result)
-
     def send_request_stop_status(self, record: "Record") -> None:
         result = proto_util._result_from_record(record)
         status_resp = result.response.stop_status_response
@@ -724,28 +703,6 @@ class SendManager:
 
         self._respond_result(result)
 
-    def send_request_server_info(self, record: "Record") -> None:
-        assert record.control.req_resp or record.control.mailbox_slot
-        result = proto_util._result_from_record(record)
-
-        result.response.server_info_response.local_info.CopyFrom(self.get_local_info())
-        for message in self._server_messages:
-            # guard against the case the message level returns malformed from server
-            message_level = str(message.get("messageLevel"))
-            message_level_sanitized = int(
-                printer.INFO if not message_level.isdigit() else message_level
-            )
-            result.response.server_info_response.server_messages.item.append(
-                wandb_internal_pb2.ServerMessage(
-                    utf_text=message.get("utfText", ""),
-                    plain_text=message.get("plainText", ""),
-                    html_text=message.get("htmlText", ""),
-                    type=message.get("messageType", ""),
-                    level=message_level_sanitized,
-                )
-            )
-        self._respond_result(result)
-
     def _setup_resume(
         self, run: "RunRecord"
     ) -> Optional["wandb_internal_pb2.ErrorInfo"]:

wandb/sdk/internal/tb_watcher.py

@@ -12,9 +12,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import wandb
 from wandb import util
+from wandb.plot.viz import CustomChart
 from wandb.sdk.interface.interface import GlobStr
 from wandb.sdk.lib import filesystem
-from wandb.sdk.lib.viz import CustomChart
 
 from . import run as internal_run
 

wandb/sdk/lib/_settings_toposort_generated.py

@@ -27,6 +27,7 @@ _Setting = Literal[
     "_executable",
     "_extra_http_headers",
     "_file_stream_max_bytes",
+    "_file_stream_transmit_interval",
     "_file_stream_retry_max",
     "_file_stream_retry_wait_min_seconds",
     "_file_stream_retry_wait_max_seconds",