wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

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()

wandb/sdk/lib/service_token.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import dataclasses
+import os
+
+from wandb import env
+
+_CURRENT_VERSION = "2"
+_SUPPORTED_TRANSPORTS = "tcp"
+
+
+def get_service_token() -> ServiceToken | None:
+    """Reads the token from environment variables.
+
+    Returns:
+        The token if the correct environment variable is set, or None.
+
+    Raises:
+        ValueError: If the environment variable is set but cannot be
+            parsed.
+    """
+    token = os.environ.get(env.SERVICE)
+    if not token:
+        return None
+
+    parts = token.split("-")
+    if len(parts) != 5:
+        raise ValueError(f"Invalid token: {token}")
+
+    version, pid_str, transport, host, port_str = parts
+
+    if version != _CURRENT_VERSION:
+        raise ValueError(
+            f"Expected version {_CURRENT_VERSION},"
+            f" but got {version} (token={token})"
+        )
+    if transport not in _SUPPORTED_TRANSPORTS:
+        raise ValueError(
+            f"Unsupported transport: {transport} (token={token})",
+        )
+
+    try:
+        return ServiceToken(
+            version=version,
+            pid=int(pid_str),
+            transport=transport,
+            host=host,
+            port=int(port_str),
+        )
+    except ValueError as e:
+        raise ValueError(f"Invalid token: {token}") from e
+
+
+def set_service_token(parent_pid: int, transport: str, host: str, port: int) -> None:
+    """Stores a service token in an environment variable.
+
+    Args:
+        parent_pid: The process ID of the process that started the service.
+        transport: The transport used to communicate with the service.
+        host: The host part of the internet address on which the service
+            is listening (e.g. localhost).
+        port: The port the service is listening on.
+
+    Raises:
+        ValueError: If given an unsupported transport.
+    """
+    if transport not in _SUPPORTED_TRANSPORTS:
+        raise ValueError(f"Unsupported transport: {transport}")
+
+    os.environ[env.SERVICE] = "-".join(
+        (
+            _CURRENT_VERSION,
+            str(parent_pid),
+            transport,
+            host,
+            str(port),
+        )
+    )
+
+
+def clear_service_token() -> None:
+    """Clears the environment variable storing the service token."""
+    os.environ.pop(env.SERVICE, None)
+
+
+@dataclasses.dataclass(frozen=True)
+class ServiceToken:
+    """An identifier for a running service process."""
+
+    version: str
+    pid: int
+    transport: str
+    host: str
+    port: int

wandb/sdk/lib/sock_client.py

@@ -0,0 +1,295 @@
+import socket
+import struct
+import threading
+import time
+import uuid
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from wandb.proto import wandb_server_pb2 as spb
+
+from . import tracelog
+
+if TYPE_CHECKING:
+    from wandb.proto import wandb_internal_pb2 as pb
+
+
+class SockClientClosedError(Exception):
+    """Raised on operations on a closed socket."""
+
+
+class SockClientTimeoutError(Exception):
+    """Raised if the server didn't respond before the timeout."""
+
+
+class SockBuffer:
+    _buf_list: List[bytes]
+    _buf_lengths: List[int]
+    _buf_total: int
+
+    def __init__(self) -> None:
+        self._buf_list = []
+        self._buf_lengths = []
+        self._buf_total = 0
+
+    @property
+    def length(self) -> int:
+        return self._buf_total
+
+    def _get(self, start: int, end: int, peek: bool = False) -> bytes:
+        index: Optional[int] = None
+        buffers = []
+        need = end
+
+        # compute buffers needed
+        for i, (buf_len, buf_data) in enumerate(zip(self._buf_lengths, self._buf_list)):
+            buffers.append(buf_data[:need] if need < buf_len else buf_data)
+            if need <= buf_len:
+                index = i
+                break
+            need -= buf_len
+
+        # buffer not large enough, caller should have made sure there was enough data
+        if index is None:
+            raise IndexError("SockBuffer index out of range")
+
+        # advance buffer internals if we are not peeking into the data
+        if not peek:
+            self._buf_total -= end
+            if need < buf_len:
+                # update partially used buffer list
+                self._buf_list = self._buf_list[index:]
+                self._buf_lengths = self._buf_lengths[index:]
+                self._buf_list[0] = self._buf_list[0][need:]
+                self._buf_lengths[0] -= need
+            else:
+                # update fully used buffer list
+                self._buf_list = self._buf_list[index + 1 :]
+                self._buf_lengths = self._buf_lengths[index + 1 :]
+
+        return b"".join(buffers)[start:end]
+
+    def get(self, start: int, end: int) -> bytes:
+        return self._get(start, end)
+
+    def peek(self, start: int, end: int) -> bytes:
+        return self._get(start, end, peek=True)
+
+    def put(self, data: bytes, data_len: int) -> None:
+        self._buf_list.append(data)
+        self._buf_lengths.append(data_len)
+        self._buf_total += data_len
+
+
+class SockClient:
+    _sock: socket.socket
+    _sockid: str
+    _retry_delay: float
+    _lock: "threading.Lock"
+    _bufsize: int
+    _buffer: SockBuffer
+
+    # current header is magic byte "W" followed by 4 byte length of the message
+    HEADLEN = 1 + 4
+
+    def __init__(self) -> None:
+        # TODO: use safe uuid's (python3.7+) or emulate this
+        self._sockid = uuid.uuid4().hex
+        self._retry_delay = 0.1
+        self._lock = threading.Lock()
+        self._bufsize = 4096
+        self._buffer = SockBuffer()
+
+    def connect(self, port: int) -> None:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect(("localhost", port))
+        self._sock = s
+        self._detect_bufsize()
+
+    def _detect_bufsize(self) -> None:
+        sndbuf_size = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
+        rcvbuf_size = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)
+        self._bufsize = min(sndbuf_size, rcvbuf_size, 65536)
+
+    def close(self) -> None:
+        self._sock.close()
+
+    def shutdown(self, val: int) -> None:
+        self._sock.shutdown(val)
+
+    def set_socket(self, sock: socket.socket) -> None:
+        self._sock = sock
+        self._detect_bufsize()
+
+    def _sendall_with_error_handle(self, data: bytes) -> None:
+        # This is a helper function for sending data in a retry fashion.
+        # Similar to the sendall() function in the socket module, but with
+        # an error handling in case of timeout.
+        total_sent = 0
+        total_data = len(data)
+        while total_sent < total_data:
+            start_time = time.monotonic()
+            try:
+                sent = self._sock.send(data)
+                # sent equal to 0 indicates a closed socket
+                if sent == 0:
+                    raise SockClientClosedError("socket connection broken")
+                total_sent += sent
+                # truncate our data to save memory
+                data = data[sent:]
+            # we handle the timeout case for the cases when timeout is set
+            # on a system level by another application
+            except socket.timeout:
+                # adding sleep to avoid tight loop
+                delta_time = time.monotonic() - start_time
+                if delta_time < self._retry_delay:
+                    time.sleep(self._retry_delay - delta_time)
+
+    def _send_message(self, msg: Any) -> None:
+        tracelog.log_message_send(msg, self._sockid)
+        raw_size = msg.ByteSize()
+        data = msg.SerializeToString()
+        assert len(data) == raw_size, "invalid serialization"
+        header = struct.pack("<BI", ord("W"), raw_size)
+        with self._lock:
+            self._sendall_with_error_handle(header + data)
+
+    def send_server_request(self, msg: Any) -> None:
+        self._send_message(msg)
+
+    def send_server_response(self, msg: Any) -> None:
+        try:
+            self._send_message(msg)
+        except BrokenPipeError:
+            # TODO(jhr): user thread might no longer be around to receive responses to
+            #  things like network status poll loop, there might be a better way to quiesce
+            pass
+
+    def send_and_recv(
+        self,
+        *,
+        inform_init: Optional[spb.ServerInformInitRequest] = None,
+        inform_start: Optional[spb.ServerInformStartRequest] = None,
+        inform_attach: Optional[spb.ServerInformAttachRequest] = None,
+        inform_finish: Optional[spb.ServerInformFinishRequest] = None,
+        inform_teardown: Optional[spb.ServerInformTeardownRequest] = None,
+    ) -> spb.ServerResponse:
+        self.send(
+            inform_init=inform_init,
+            inform_start=inform_start,
+            inform_attach=inform_attach,
+            inform_finish=inform_finish,
+            inform_teardown=inform_teardown,
+        )
+        # TODO: this solution is fragile, but for checking attach
+        # it should be relatively stable.
+        # This pass would be solved as part of the fix in https://wandb.atlassian.net/browse/WB-8709
+        response = self.read_server_response(timeout=1)
+
+        if response is None:
+            raise SockClientTimeoutError("No response after 1 second.")
+
+        return response
+
+    def send(
+        self,
+        *,
+        inform_init: Optional[spb.ServerInformInitRequest] = None,
+        inform_start: Optional[spb.ServerInformStartRequest] = None,
+        inform_attach: Optional[spb.ServerInformAttachRequest] = None,
+        inform_finish: Optional[spb.ServerInformFinishRequest] = None,
+        inform_teardown: Optional[spb.ServerInformTeardownRequest] = None,
+    ) -> None:
+        server_req = spb.ServerRequest()
+        if inform_init:
+            server_req.inform_init.CopyFrom(inform_init)
+        elif inform_start:
+            server_req.inform_start.CopyFrom(inform_start)
+        elif inform_attach:
+            server_req.inform_attach.CopyFrom(inform_attach)
+        elif inform_finish:
+            server_req.inform_finish.CopyFrom(inform_finish)
+        elif inform_teardown:
+            server_req.inform_teardown.CopyFrom(inform_teardown)
+        else:
+            raise Exception("unmatched")
+        self.send_server_request(server_req)
+
+    def send_record_communicate(self, record: "pb.Record") -> None:
+        server_req = spb.ServerRequest()
+        server_req.record_communicate.CopyFrom(record)
+        self.send_server_request(server_req)
+
+    def send_record_publish(self, record: "pb.Record") -> None:
+        server_req = spb.ServerRequest()
+        server_req.record_publish.CopyFrom(record)
+        self.send_server_request(server_req)
+
+    def _extract_packet_bytes(self) -> Optional[bytes]:
+        # Do we have enough data to read the header?
+        start_offset = self.HEADLEN
+        if self._buffer.length >= start_offset:
+            header = self._buffer.peek(0, start_offset)
+            fields = struct.unpack("<BI", header)
+            magic, dlength = fields
+            assert magic == ord("W")
+            # Do we have enough data to read the full record?
+            end_offset = self.HEADLEN + dlength
+            if self._buffer.length >= end_offset:
+                rec_data = self._buffer.get(start_offset, end_offset)
+                return rec_data
+        return None
+
+    def _read_packet_bytes(self, timeout: Optional[int] = None) -> Optional[bytes]:
+        """Read full message from socket.
+
+        Args:
+            timeout: number of seconds to wait on socket data.
+
+        Raises:
+            SockClientClosedError: socket has been closed.
+        """
+        while True:
+            rec = self._extract_packet_bytes()
+            if rec:
+                return rec
+
+            if timeout:
+                self._sock.settimeout(timeout)
+            try:
+                data = self._sock.recv(self._bufsize)
+            except socket.timeout:
+                break
+            except ConnectionResetError:
+                raise SockClientClosedError
+            except OSError:
+                raise SockClientClosedError
+            finally:
+                if timeout:
+                    self._sock.settimeout(None)
+            data_len = len(data)
+            if data_len == 0:
+                # socket.recv() will return 0 bytes if socket was shutdown
+                # caller will handle this condition like other connection problems
+                raise SockClientClosedError
+            self._buffer.put(data, data_len)
+        return None
+
+    def read_server_request(self) -> Optional[spb.ServerRequest]:
+        data = self._read_packet_bytes()
+        if not data:
+            return None
+        rec = spb.ServerRequest()
+        rec.ParseFromString(data)
+        tracelog.log_message_recv(rec, self._sockid)
+        return rec
+
+    def read_server_response(
+        self, timeout: Optional[int] = None
+    ) -> Optional[spb.ServerResponse]:
+        data = self._read_packet_bytes(timeout=timeout)
+        if not data:
+            return None
+        rec = spb.ServerResponse()
+        rec.ParseFromString(data)
+        tracelog.log_message_recv(rec, self._sockid)
+        return rec

wandb/sdk/lib/sparkline.py

@@ -0,0 +1,45 @@
+#
+# From pysparklines (BSD License): https://pypi.python.org/pypi/pysparklines
+
+import math
+from typing import List, Union
+
+spark_chars = "▁▂▃▄▅▆▇█"
+
+
+# math.isfinite doesn't exist in python2, so provider our own
+def isfinite(f):
+    return not (math.isinf(f) or math.isnan(f))
+
+
+def sparkify(series: List[Union[float, int]]) -> str:
+    """Convert <series> to a sparkline string.
+
+    Example:
+    >>> sparkify([ 0.5, 1.2, 3.5, 7.3, 8.0, 12.5, 13.2, 15.0, 14.2, 11.8, 6.1,
+    ... 1.9 ])
+    u'▁▁▂▄▅▇▇██▆▄▂'
+
+    >>> sparkify([1, 1, -2, 3, -5, 8, -13])
+    u'▆▆▅▆▄█▁'
+
+    Raises ValueError if input data cannot be converted to float.
+    Raises TypeError if series is not an iterable.
+    """
+    series = [float(i) for i in series]
+    finite_series = [x for x in series if isfinite(x)]
+    if not finite_series:
+        return ""
+    minimum = min(finite_series)
+    maximum = max(finite_series)
+    data_range = maximum - minimum
+    if data_range == 0.0:
+        # Graph a baseline if every input value is equal.
+        return "".join([spark_chars[0] if isfinite(x) else " " for x in series])
+    coefficient = (len(spark_chars) - 1.0) / data_range
+    return "".join(
+        [
+            spark_chars[int(round((x - minimum) * coefficient))] if isfinite(x) else " "
+            for x in series
+        ]
+    )