wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

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 does not 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/lib/wb_logging.py

@@ -0,0 +1,161 @@
+"""Logging configuration for the "wandb" logger.
+
+Most log statements in wandb are made in the context of a run and should be
+redirected to that run's log file (usually named 'debug.log'). This module
+provides a context manager to temporarily set the current run ID and registers
+a global handler for the 'wandb' logger that sends log statements to the right
+place.
+
+All functions in this module are threadsafe.
+
+NOTE: The pytest caplog fixture will fail to capture logs from the wandb logger
+because they are not propagated to the root logger.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import contextvars
+import logging
+import pathlib
+from typing import Iterator
+
+
+class _NotRunSpecific:
+    """Sentinel for `not_run_specific()`."""
+
+
+_NOT_RUN_SPECIFIC = _NotRunSpecific()
+
+
+_run_id: contextvars.ContextVar[str | _NotRunSpecific | None] = contextvars.ContextVar(
+    "_run_id",
+    default=None,
+)
+
+_logger = logging.getLogger("wandb")
+
+
+def configure_wandb_logger() -> None:
+    """Configures the global 'wandb' logger.
+
+    The wandb logger is not intended to be customized by users. Instead, it is
+    used as a mechanism to redirect log messages into wandb run-specific log
+    files.
+
+    This function is idempotent: calling it multiple times has the same effect.
+    """
+    # Send all DEBUG and above messages to registered handlers.
+    #
+    # Per-run handlers can set different levels.
+    _logger.setLevel(logging.DEBUG)
+
+    # Do not propagate wandb logs to the root logger, which the user may have
+    # configured to point elsewhere. All wandb log messages should go to a run's
+    # log file.
+    _logger.propagate = False
+
+    # If no handlers are configured for the 'wandb' logger, don't activate the
+    # "lastResort" handler which sends messages to stderr with a level of
+    # WARNING by default.
+    #
+    # This occurs in wandb code that runs outside the context of a Run and
+    # not as part of the CLI.
+    #
+    # Most such code uses the `termlog` / `termwarn` / `termerror` methods
+    # to communicate with the user. When that code executes while a run is
+    # active, its logger messages go to that run's log file.
+    if not _logger.handlers:
+        _logger.addHandler(logging.NullHandler())
+
+
+@contextlib.contextmanager
+def log_to_run(run_id: str | None) -> Iterator[None]:
+    """Direct all wandb log messages to the given run.
+
+    Args:
+        id: The current run ID, or None if actions in the context manager are
+            not associated to a specific run. In the latter case, log messages
+            will go to all runs.
+
+    Usage:
+
+        with wb_logging.run_id(...):
+            ... # Log messages here go to the specified run's logger.
+    """
+    token = _run_id.set(run_id)
+    try:
+        yield
+    finally:
+        _run_id.reset(token)
+
+
+@contextlib.contextmanager
+def log_to_all_runs() -> Iterator[None]:
+    """Direct wandb log messages to all runs.
+
+    Unlike `log_to_run(None)`, this indicates an intentional choice.
+    This is often convenient to use as a decorator:
+
+        @wb_logging.log_to_all_runs()
+        def my_func():
+            ... # Log messages here go to the specified run's logger.
+    """
+    token = _run_id.set(_NOT_RUN_SPECIFIC)
+    try:
+        yield
+    finally:
+        _run_id.reset(token)
+
+
+def add_file_handler(run_id: str, filepath: pathlib.Path) -> logging.Handler:
+    """Direct log messages for a run to a file.
+
+    Args:
+        run_id: The run for which to create a log file.
+        filepath: The file to write log messages to.
+
+    Returns:
+        The added handler which can then be configured further or removed
+        from the 'wandb' logger directly.
+
+        The default logging level is INFO.
+    """
+    handler = logging.FileHandler(filepath)
+    handler.setLevel(logging.INFO)
+    handler.addFilter(_RunIDFilter(run_id))
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d"
+            " [%(filename)s:%(funcName)s():%(lineno)s]%(run_id_tag)s"
+            " %(message)s"
+        )
+    )
+
+    _logger.addHandler(handler)
+    return handler
+
+
+class _RunIDFilter(logging.Filter):
+    """Filters out messages logged for a different run."""
+
+    def __init__(self, run_id: str) -> None:
+        """Create a _RunIDFilter.
+
+        Args:
+            run_id: Allows messages when the run ID is this or None.
+        """
+        self._run_id = run_id
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        run_id = _run_id.get()
+
+        if run_id is None:
+            record.run_id_tag = " [no run ID]"
+            return True
+        elif isinstance(run_id, _NotRunSpecific):
+            record.run_id_tag = " [all runs]"
+            return True
+        else:
+            record.run_id_tag = ""
+            return run_id == self._run_id

wandb/sdk/mailbox/__init__.py

@@ -0,0 +1,23 @@
+"""A message protocol for the internal service process.
+
+The core of W&B is implemented by a side process that asynchronously uploads
+data. The client process (such as this Python code) sends requests to the
+service, and for some requests, the service eventually sends a response.
+
+The client can send multiple requests before the service provides a response.
+The Mailbox handles matching responses to requests. An internal thread
+continuously reads data from the service and passes it to the mailbox.
+"""
+
+from .mailbox import Mailbox, MailboxClosedError
+from .mailbox_handle import HandleAbandonedError, MailboxHandle
+from .wait_with_progress import wait_all_with_progress, wait_with_progress
+
+__all__ = [
+    "Mailbox",
+    "MailboxClosedError",
+    "HandleAbandonedError",
+    "MailboxHandle",
+    "wait_all_with_progress",
+    "wait_with_progress",
+]

wandb/sdk/mailbox/mailbox.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import logging
+import secrets
+import string
+import threading
+
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
+
+from .mailbox_handle import MailboxHandle
+from .response_handle import MailboxResponseHandle
+
+_logger = logging.getLogger(__name__)
+
+
+class MailboxClosedError(Exception):
+    """The mailbox has been closed and cannot be used."""
+
+
+class Mailbox:
+    """Matches service responses to requests.
+
+    The mailbox can set an address on a server request and create a handle for
+    waiting for a response to that record. Responses are delivered by calling
+    `deliver()`. The `close()` method abandons all handles in case the
+    service process becomes unreachable.
+    """
+
+    def __init__(self, asyncer: asyncio_manager.AsyncioManager) -> None:
+        self._asyncer = asyncer
+        self._handles: dict[str, MailboxResponseHandle] = {}
+        self._handles_lock = threading.Lock()
+        self._closed = False
+
+    def require_response(
+        self,
+        request: spb.ServerRequest | pb.Record,
+    ) -> MailboxHandle[spb.ServerResponse]:
+        """Set a response address on a request.
+
+        Args:
+            request: The request on which to set a request ID or mailbox slot.
+                This is mutated. An address must not already be set.
+
+        Returns:
+            A handle for waiting for the response to the request.
+
+        Raises:
+            MailboxClosedError: If the mailbox has been closed, in which case
+                no new responses are expected to be delivered and new handles
+                cannot be created.
+        """
+        if isinstance(request, spb.ServerRequest):
+            if (address := request.request_id) or (
+                address := request.record_publish.control.mailbox_slot
+            ):
+                raise ValueError(f"Request already has an address ({address})")
+
+            address = self._new_address()
+            request.request_id = address
+            if request.HasField("record_publish"):
+                request.record_publish.control.mailbox_slot = address
+            if request.HasField("record_communicate"):
+                request.record_communicate.control.mailbox_slot = address
+        else:
+            if address := request.control.mailbox_slot:
+                raise ValueError(f"Request already has an address ({address})")
+
+            address = self._new_address()
+            request.control.mailbox_slot = address
+
+        with self._handles_lock:
+            if self._closed:
+                raise MailboxClosedError()
+
+            handle = MailboxResponseHandle(address, asyncer=self._asyncer)
+            self._handles[address] = handle
+
+        return handle
+
+    def _new_address(self) -> str:
+        """Returns an unused address for a request.
+
+        Assumes `_handles_lock` is held.
+        """
+
+        def generate():
+            return "".join(
+                secrets.choice(string.ascii_lowercase + string.digits)
+                for i in range(12)
+            )
+
+        address = generate()
+
+        # Being extra cautious. This loop will almost never be entered.
+        while address in self._handles:
+            address = generate()
+
+        return address
+
+    async def deliver(self, response: spb.ServerResponse) -> None:
+        """Deliver a response from the service.
+
+        If the response address is invalid, this does nothing.
+        It is a no-op if the mailbox has been closed.
+        """
+        address = response.request_id
+        if not address:
+            kind: str | None = response.WhichOneof("server_response_type")
+            if kind == "result_communicate":
+                result_type = response.result_communicate.WhichOneof("result_type")
+                kind = f"result_communicate.{result_type}"
+
+            _logger.error(f"Received response with no mailbox slot: {kind}")
+            return
+
+        with self._handles_lock:
+            # NOTE: If the mailbox is closed, this returns None because
+            # we clear the dict.
+            handle = self._handles.pop(address, None)
+
+        # It is not an error if there is no handle for the address:
+        # handles can be abandoned if the result is no longer needed.
+        if handle:
+            await handle.deliver(response)
+
+    def close(self) -> None:
+        """Indicate no further responses will be delivered.
+
+        Abandons all handles.
+        """
+        with self._handles_lock:
+            self._closed = True
+
+            _logger.info(
+                f"Closing mailbox, abandoning {len(self._handles)} handles.",
+            )
+
+            for handle in self._handles.values():
+                handle.abandon()
+            self._handles.clear()