wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

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,290 @@
+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
+
+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:
+        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)
+        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)
+        return rec

wandb/sdk/lib/sparkline.py

@@ -0,0 +1,44 @@
+#
+# 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
+        ]
+    )

wandb/sdk/lib/telemetry.py

@@ -0,0 +1,100 @@
+import re
+import sys
+from types import TracebackType
+from typing import TYPE_CHECKING, ContextManager, Dict, List, Optional, Set, Type
+
+import wandb
+from wandb.proto.wandb_telemetry_pb2 import Imports as TelemetryImports
+from wandb.proto.wandb_telemetry_pb2 import TelemetryRecord
+
+# avoid cycle, use string type reference
+
+if TYPE_CHECKING:
+    from .. import wandb_run
+
+
+_LABEL_TOKEN: str = "@wandbcode{"
+
+
+class _TelemetryObject:
+    _run: Optional["wandb_run.Run"]
+    _obj: TelemetryRecord
+
+    def __init__(
+        self,
+        run: Optional["wandb_run.Run"] = None,
+        obj: Optional[TelemetryRecord] = None,
+    ) -> None:
+        self._run = run or wandb.run
+        self._obj = obj or TelemetryRecord()
+
+    def __enter__(self) -> TelemetryRecord:
+        return self._obj
+
+    def __exit__(
+        self,
+        exctype: Optional[Type[BaseException]],
+        excinst: Optional[BaseException],
+        exctb: Optional[TracebackType],
+    ) -> None:
+        if not self._run:
+            return
+        self._run._telemetry_callback(self._obj)
+
+
+def context(
+    run: Optional["wandb_run.Run"] = None, obj: Optional[TelemetryRecord] = None
+) -> ContextManager[TelemetryRecord]:
+    return _TelemetryObject(run=run, obj=obj)
+
+
+MATCH_RE = re.compile(r"(?P<code>[a-zA-Z0-9_-]+)[,}](?P<rest>.*)")
+
+
+def _parse_label_lines(lines: List[str]) -> Dict[str, str]:
+    seen = False
+    ret = {}
+    for line in lines:
+        idx = line.find(_LABEL_TOKEN)
+        if idx < 0:
+            # Stop parsing on first non token line after match
+            if seen:
+                break
+            continue
+        seen = True
+        label_str = line[idx + len(_LABEL_TOKEN) :]
+
+        # match identifier (first token without key=value syntax (optional)
+        # Note: Parse is fairly permissive as it doesnt enforce strict syntax
+        r = MATCH_RE.match(label_str)
+        if r:
+            ret["code"] = r.group("code").replace("-", "_")
+            label_str = r.group("rest")
+
+        # match rest of tokens on one line
+        tokens = re.findall(
+            r'([a-zA-Z0-9_]+)\s*=\s*("[a-zA-Z0-9_-]*"|[a-zA-Z0-9_-]*)[,}]', label_str
+        )
+        for k, v in tokens:
+            ret[k] = v.strip('"').replace("-", "_")
+    return ret
+
+
+def list_telemetry_imports(only_imported: bool = False) -> Set[str]:
+    import_telemetry_set = {
+        desc.name
+        for desc in TelemetryImports.DESCRIPTOR.fields
+        if desc.type == desc.TYPE_BOOL
+    }
+    if only_imported:
+        imported_modules_set = set(sys.modules)
+        return imported_modules_set.intersection(import_telemetry_set)
+    return import_telemetry_set
+
+
+__all__ = [
+    "TelemetryImports",
+    "TelemetryRecord",
+    "context",
+    "list_telemetry_imports",
+]

wandb/sdk/lib/timed_input.py

@@ -0,0 +1,133 @@
+"""timed_input: add a timeout to standard input.
+
+Approach was inspired by: https://github.com/johejo/inputimeout
+"""
+
+import sys
+import threading
+
+import wandb
+
+SP = " "
+CR = "\r"
+LF = "\n"
+CRLF = CR + LF
+
+
+def _echo(prompt: str) -> None:
+    sys.stdout.write(prompt)
+    sys.stdout.flush()
+
+
+def _posix_timed_input(prompt: str, timeout: float) -> str:
+    _echo(prompt)
+    sel = selectors.DefaultSelector()
+    sel.register(sys.stdin, selectors.EVENT_READ, data=sys.stdin.readline)
+    events = sel.select(timeout=timeout)
+
+    for key, _ in events:
+        input_callback = key.data
+        input_data: str = input_callback()
+        if not input_data:  # end-of-file - treat as timeout
+            raise TimeoutError
+        return input_data.rstrip(LF)
+
+    _echo(LF)
+    termios.tcflush(sys.stdin, termios.TCIFLUSH)
+    raise TimeoutError
+
+
+def _windows_timed_input(prompt: str, timeout: float) -> str:
+    interval = 0.1
+
+    _echo(prompt)
+    begin = time.monotonic()
+    end = begin + timeout
+    line = ""
+
+    while time.monotonic() < end:
+        if msvcrt.kbhit():  # type: ignore[attr-defined]
+            c = msvcrt.getwche()  # type: ignore[attr-defined]
+            if c in (CR, LF):
+                _echo(CRLF)
+                return line
+            if c == "\003":
+                raise KeyboardInterrupt
+            if c == "\b":
+                line = line[:-1]
+                cover = SP * len(prompt + line + SP)
+                _echo("".join([CR, cover, CR, prompt, line]))
+            else:
+                line += c
+        time.sleep(interval)
+
+    _echo(CRLF)
+    raise TimeoutError
+
+
+def _jupyter_timed_input(prompt: str, timeout: float) -> str:
+    clear = True
+    try:
+        from IPython.core.display import clear_output  # type: ignore
+    except ImportError:
+        clear = False
+        wandb.termwarn(
+            "Unable to clear output, can't import clear_output from ipython.core"
+        )
+
+    _echo(prompt)
+
+    user_inp = None
+    event = threading.Event()
+
+    def get_input() -> None:
+        nonlocal user_inp
+        raw = input()
+        if event.is_set():
+            return
+        user_inp = raw
+
+    t = threading.Thread(target=get_input)
+    t.start()
+    t.join(timeout)
+    event.set()
+    if user_inp:
+        return user_inp
+    if clear:
+        clear_output()
+    raise TimeoutError
+
+
+def timed_input(
+    prompt: str, timeout: float, show_timeout: bool = True, jupyter: bool = False
+) -> str:
+    """Behaves like builtin `input()` but adds timeout.
+
+    Args:
+        prompt (str): Prompt to output to stdout.
+        timeout (float): Timeout to wait for input.
+        show_timeout (bool): Show timeout in prompt
+        jupyter (bool): If True, use jupyter specific code.
+
+    Raises:
+        TimeoutError: exception raised if timeout occurred.
+    """
+    if show_timeout:
+        prompt = f"{prompt}({timeout:.0f} second timeout) "
+    if jupyter:
+        return _jupyter_timed_input(prompt=prompt, timeout=timeout)
+
+    return _timed_input(prompt=prompt, timeout=timeout)
+
+
+try:
+    import msvcrt
+except ImportError:
+    import selectors
+    import termios
+
+    _timed_input = _posix_timed_input
+else:
+    import time
+
+    _timed_input = _windows_timed_input

wandb/sdk/lib/timer.py

@@ -0,0 +1,19 @@
+import time
+from typing import Any
+
+
+class Timer:
+    def __init__(self) -> None:
+        self.start_time: float = time.time()
+        self.start: float = time.perf_counter()
+        self.stop: float = self.start
+
+    def __enter__(self) -> "Timer":
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.stop = time.perf_counter()
+
+    @property
+    def elapsed(self) -> float:
+        return self.stop - self.start

wandb/sdk/service/_startup_debug.py

@@ -0,0 +1,22 @@
+"""_startup_debug.
+
+Temporary helper to debug issues with wandb service startup
+"""
+
+import os
+import time
+
+
+def is_enabled() -> bool:
+    # This is very temporary to help diagnose problems seen by some
+    # customers which we are having trouble reproducing. It should be
+    # replaced by something more permanent in the future when we have
+    # proper logging for wandb-service
+    if os.environ.get("_WANDB_STARTUP_DEBUG"):
+        return True
+    return False
+
+
+def print_message(message: str) -> None:
+    time_now = time.time()
+    print("WANDB_STARTUP_DEBUG", time_now, message)