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

wandb/sdk/data_types/trace_tree.py

@@ -14,9 +14,9 @@ from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import wandb
-import wandb.data_types
 from wandb.sdk.data_types import _dtypes
 from wandb.sdk.data_types.base_types.media import Media
+from wandb.sdk.data_types.utils import _json_helper
 
 if TYPE_CHECKING:  # pragma: no cover
     from wandb.sdk.artifacts.artifact import Artifact
@@ -142,7 +142,7 @@ def _fallback_serialize(obj: Any) -> str:
 def _safe_serialize(obj: dict) -> str:
     try:
         return json.dumps(
-            wandb.data_types._json_helper(obj, None),
+            _json_helper(obj, None),
             skipkeys=True,
             default=_fallback_serialize,
         )

wandb/sdk/data_types/utils.py

@@ -1,6 +1,8 @@
+import datetime
 import logging
 import os
 import re
+from decimal import Decimal
 from typing import TYPE_CHECKING, Optional, Sequence, Union, cast
 
 import wandb
@@ -178,3 +180,50 @@ def _prune_max_seq(seq: Sequence["BatchableMedia"]) -> Sequence["BatchableMedia"
         )
         items = seq[: seq[0].MAX_ITEMS]
     return items
+
+
+def _json_helper(val, artifact):
+    if isinstance(val, WBValue):
+        return val.to_json(artifact)
+    elif val.__class__ is dict:
+        res = {}
+        for key in val:
+            res[key] = _json_helper(val[key], artifact)
+        return res
+
+    if hasattr(val, "tolist"):
+        py_val = val.tolist()
+        if val.__class__.__name__ == "datetime64" and isinstance(py_val, int):
+            # when numpy datetime64 .tolist() returns an int, it is nanoseconds.
+            # need to convert to milliseconds
+            return _json_helper(py_val / int(1e6), artifact)
+        return _json_helper(py_val, artifact)
+    elif hasattr(val, "item"):
+        return _json_helper(val.item(), artifact)
+
+    if isinstance(val, datetime.datetime):
+        if val.tzinfo is None:
+            val = datetime.datetime(
+                val.year,
+                val.month,
+                val.day,
+                val.hour,
+                val.minute,
+                val.second,
+                val.microsecond,
+                tzinfo=datetime.timezone.utc,
+            )
+        return int(val.timestamp() * 1000)
+    elif isinstance(val, datetime.date):
+        return int(
+            datetime.datetime(
+                val.year, val.month, val.day, tzinfo=datetime.timezone.utc
+            ).timestamp()
+            * 1000
+        )
+    elif isinstance(val, (list, tuple)):
+        return [_json_helper(i, artifact) for i in val]
+    elif isinstance(val, Decimal):
+        return float(val)
+    else:
+        return util.json_friendly(val)[0]

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",

wandb/sdk/lib/hashutil.py

@@ -1,19 +1,22 @@
+from __future__ import annotations
+
 import base64
 import hashlib
 import mmap
-import os
 import sys
-from pathlib import Path
-from typing import NewType, Union
+from typing import TYPE_CHECKING, NewType
 
 from wandb.sdk.lib.paths import StrPath
 
+if TYPE_CHECKING:
+    import _hashlib  # type: ignore[import-not-found]
+
 ETag = NewType("ETag", str)
 HexMD5 = NewType("HexMD5", str)
 B64MD5 = NewType("B64MD5", str)
 
 
-def _md5(data: bytes = b"") -> "hashlib._Hash":
+def _md5(data: bytes = b"") -> _hashlib.HASH:
     """Allow FIPS-compliant md5 hash when supported."""
     if sys.version_info >= (3, 9):
         return hashlib.md5(data, usedforsecurity=False)
@@ -25,7 +28,7 @@ def md5_string(string: str) -> B64MD5:
     return _b64_from_hasher(_md5(string.encode("utf-8")))
 
 
-def _b64_from_hasher(hasher: "hashlib._Hash") -> B64MD5:
+def _b64_from_hasher(hasher: _hashlib.HASH) -> B64MD5:
     return B64MD5(base64.b64encode(hasher.digest()).decode("ascii"))
 
 
@@ -33,7 +36,7 @@ def b64_to_hex_id(string: B64MD5) -> HexMD5:
     return HexMD5(base64.standard_b64decode(string).hex())
 
 
-def hex_to_b64_id(encoded_string: Union[str, bytes]) -> B64MD5:
+def hex_to_b64_id(encoded_string: str | bytes) -> B64MD5:
     if isinstance(encoded_string, bytes):
         encoded_string = encoded_string.decode("utf-8")
     as_str = bytes.fromhex(encoded_string)
@@ -48,15 +51,34 @@ def md5_file_hex(*paths: StrPath) -> HexMD5:
     return HexMD5(_md5_file_hasher(*paths).hexdigest())
 
 
-def _md5_file_hasher(*paths: StrPath) -> "hashlib._Hash":
+_KB: int = 1_024
+_CHUNKSIZE: int = 128 * _KB
+"""Chunk size (in bytes) for iteratively reading from file, if needed."""
+
+
+def _md5_file_hasher(*paths: StrPath) -> _hashlib.HASH:
     md5_hash = _md5()
 
-    for path in sorted(Path(p) for p in paths):
-        with path.open("rb") as f:
-            if os.stat(f.fileno()).st_size <= 1024 * 1024:
-                md5_hash.update(f.read())
-            else:
+    # Note: We use str paths (instead of pathlib.Path objs) for minor perf improvements.
+    for path in sorted(map(str, paths)):
+        with open(path, "rb") as f:
+            try:
                 with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as mview:
                     md5_hash.update(mview)
+            except OSError:
+                # This occurs if the mmap-ed file is on a different/mounted filesystem,
+                # so we'll fall back on a less performant implementation.
+
+                # Note: At the time of implementation, the walrus operator `:=`
+                # is avoided to maintain support for users on python 3.7.
+                # Consider revisiting once 3.7 support is no longer needed.
+                chunk = f.read(_CHUNKSIZE)
+                while chunk:
+                    md5_hash.update(chunk)
+                    chunk = f.read(_CHUNKSIZE)
+            except ValueError:
+                # This occurs when mmap-ing an empty file, which can be skipped.
+                # See: https://github.com/python/cpython/blob/986a4e1b6fcae7fe7a1d0a26aea446107dd58dd2/Modules/mmapmodule.c#L1589
+                pass
 
     return md5_hash

wandb/sdk/lib/service_connection.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import atexit
+import os
+from typing import Callable
+
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.proto import wandb_settings_pb2
+from wandb.sdk import wandb_settings
+from wandb.sdk.interface.interface import InterfaceBase
+from wandb.sdk.interface.interface_sock import InterfaceSock
+from wandb.sdk.lib import service_token
+from wandb.sdk.lib.exit_hooks import ExitHooks
+from wandb.sdk.lib.mailbox import Mailbox
+from wandb.sdk.lib.sock_client import SockClient, SockClientTimeoutError
+from wandb.sdk.service import service
+
+
+class WandbServiceNotOwnedError(Exception):
+    """Raised when the current process does not own the service process."""
+
+
+class WandbServiceConnectionError(Exception):
+    """Raised on failure to connect to the service process."""
+
+
+class WandbAttachFailedError(Exception):
+    """Raised if attaching to a run fails."""
+
+
+def connect_to_service(
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Connects to the service process, starting one up if necessary."""
+    conn = _try_connect_to_existing_service()
+    if conn:
+        return conn
+
+    return _start_and_connect_service(settings)
+
+
+def _try_connect_to_existing_service() -> ServiceConnection | None:
+    """Attemps to connect to an existing service process."""
+    token = service_token.get_service_token()
+    if not token:
+        return None
+
+    # Only localhost sockets are supported below.
+    assert token.host == "localhost"
+    client = SockClient()
+
+    try:
+        # TODO: This may block indefinitely if the service is unhealthy.
+        client.connect(token.port)
+
+    except Exception as e:
+        raise WandbServiceConnectionError(
+            "Failed to connect to internal service."
+        ) from e
+
+    return ServiceConnection(client=client, proc=None)
+
+
+def _start_and_connect_service(
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Starts a service process and returns a connection to it.
+
+    An atexit hook is registered to tear down the service process and wait for
+    it to complete. The hook does not run in processes started using the
+    multiprocessing module.
+    """
+    proc = service._Service(settings)
+    proc.start()
+
+    port = proc.sock_port
+    assert port
+    client = SockClient()
+    client.connect(port)
+
+    service_token.set_service_token(
+        parent_pid=os.getpid(),
+        transport="tcp",
+        host="localhost",
+        port=port,
+    )
+
+    hooks = ExitHooks()
+    hooks.hook()
+
+    def teardown_atexit():
+        conn.teardown(hooks.exit_code)
+
+    conn = ServiceConnection(
+        client=client,
+        proc=proc,
+        cleanup=lambda: atexit.unregister(teardown_atexit),
+    )
+
+    atexit.register(teardown_atexit)
+
+    return conn
+
+
+class ServiceConnection:
+    """A connection to the W&B internal service process."""
+
+    def __init__(
+        self,
+        client: SockClient,
+        proc: service._Service | None,
+        cleanup: Callable[[], None] | None = None,
+    ):
+        """Returns a new ServiceConnection.
+
+        Args:
+            client: A socket that's connected to the service.
+            proc: The service process if we own it, or None otherwise.
+            cleanup: A callback to run on teardown before doing anything.
+        """
+        self._client = client
+        self._proc = proc
+        self._torn_down = False
+        self._cleanup = cleanup
+
+    def make_interface(self, mailbox: Mailbox) -> InterfaceBase:
+        """Returns an interface for communicating with the service."""
+        return InterfaceSock(self._client, mailbox)
+
+    def send_record(self, record: pb.Record) -> None:
+        """Sends data to the service."""
+        self._client.send_record_publish(record)
+
+    def inform_init(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Sends an init request to the service."""
+        request = spb.ServerInformInitRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.send(inform_init=request)
+
+    def inform_finish(self, run_id: str) -> None:
+        """Sends an finish request to the service."""
+        request = spb.ServerInformFinishRequest()
+        request._info.stream_id = run_id
+        self._client.send(inform_finish=request)
+
+    def inform_attach(
+        self,
+        attach_id: str,
+    ) -> wandb_settings_pb2.Settings:
+        """Sends an attach request to the service.
+
+        Raises a WandbAttachFailedError if attaching is not possible.
+        """
+        request = spb.ServerInformAttachRequest()
+        request._info.stream_id = attach_id
+
+        try:
+            response = self._client.send_and_recv(inform_attach=request)
+            return response.inform_attach_response.settings
+        except SockClientTimeoutError:
+            raise WandbAttachFailedError(
+                "Could not attach because the run does not belong to"
+                " the current service process, or because the service"
+                " process is busy (unlikely)."
+            )
+
+    def inform_start(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Sends a start request to the service."""
+        request = spb.ServerInformStartRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.send(inform_start=request)
+
+    def teardown(self, exit_code: int) -> int:
+        """Shuts down the service process and returns its exit code.
+
+        This may only be called once.
+
+        Returns:
+            The exit code of the service process.
+
+        Raises:
+            WandbServiceNotOwnedError: If the current process did not start
+                the service process.
+        """
+        if not self._proc:
+            raise WandbServiceNotOwnedError(
+                "Cannot tear down service started by different process",
+            )
+
+        assert not self._torn_down
+        self._torn_down = True
+
+        if self._cleanup:
+            self._cleanup()
+
+        # Clear the service token to prevent new connections from being made.
+        service_token.clear_service_token()
+
+        self._client.send(
+            inform_teardown=spb.ServerInformTeardownRequest(
+                exit_code=exit_code,
+            )
+        )
+
+        return self._proc.join()