wandb 0.22.1__py3-none-win_arm64.whl

wandb/sdk/lib/console_capture.py

@@ -0,0 +1,219 @@
+"""Module for intercepting stdout/stderr.
+
+This patches the `write()` method of `stdout` and `stderr` on import.
+Once patched, it is not possible to unpatch or repatch, though individual
+callbacks can be removed.
+
+We assume that all other writing methods on the object delegate to `write()`,
+like `writelines()`. This is not guaranteed to be true, but it is true for
+common implementations. In particular, CPython's implementation of IOBase's
+`writelines()` delegates to `write()`.
+
+It is important to note that this technique interacts poorly with other
+code that performs similar patching if it also allows unpatching as this
+discards our modification. This is why we patch on import and do not support
+unpatching:
+
+    with contextlib.redirect_stderr(...):
+        from ... import console_capture
+        # Here, everything works fine.
+    # Here, callbacks are never called again.
+
+In particular, it does not work with some combinations of pytest's
+`capfd` / `capsys` fixtures and pytest's `--capture` option.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+import threading
+from typing import IO, AnyStr, Callable, Protocol
+
+from . import wb_logging
+
+_logger = logging.getLogger(__name__)
+
+
+class CannotCaptureConsoleError(Exception):
+    """The module failed to patch stdout or stderr."""
+
+
+class _WriteCallback(Protocol):
+    """A callback that receives intercepted bytes or string data.
+
+    This may be called from any thread, but is only called from one thread
+    at a time.
+
+    Note on errors: Any error raised during the callback will clear all
+    callbacks. This means that if a user presses Ctrl-C at an unlucky time
+    during a run, we will stop uploading console output---but it's not
+    likely to be a problem unless something catches the KeyboardInterrupt.
+
+    Regular Exceptions are caught and logged instead of bubbling up to the
+    user's print() statements; other exceptions like KeyboardInterrupt are
+    re-raised.
+
+    Callbacks should handle all exceptions---a callback that raises any
+    Exception is considered buggy.
+    """
+
+    def __call__(
+        self,
+        data: bytes | str,
+        written: int,
+        /,
+    ) -> None:
+        """Intercept data passed to `write()`.
+
+        See the protocol docstring for information about exceptions.
+
+        Args:
+            data: The object passed to stderr's or stdout's `write()`.
+            written: The number of bytes or characters written.
+                This is the return value of `write()`.
+        """
+
+
+# A reentrant lock is used to catch callbacks that write to stderr/stdout.
+_module_rlock = threading.RLock()
+_is_writing = False
+
+_patch_exception: CannotCaptureConsoleError | None = None
+
+_next_callback_id: int = 1
+
+_stdout_callbacks: dict[int, _WriteCallback] = {}
+_stderr_callbacks: dict[int, _WriteCallback] = {}
+
+
+def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
+    """Install a callback that runs after every write to sys.stdout.
+
+    Args:
+        callback: A callback to invoke after running `sys.stdout.write`.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_rlock:
+        if _patch_exception:
+            raise _patch_exception
+
+        return _insert_disposably(
+            _stdout_callbacks,
+            callback,
+        )
+
+
+def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
+    """Install a callback that runs after every write to sys.sdterr.
+
+    Args:
+        callback: A callback to invoke after running `sys.stderr.write`.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_rlock:
+        if _patch_exception:
+            raise _patch_exception
+
+        return _insert_disposably(
+            _stderr_callbacks,
+            callback,
+        )
+
+
+def _insert_disposably(
+    callback_dict: dict[int, _WriteCallback],
+    callback: _WriteCallback,
+) -> Callable[[], None]:
+    global _next_callback_id
+    id = _next_callback_id
+    _next_callback_id += 1
+
+    disposed = False
+
+    def dispose() -> None:
+        nonlocal disposed
+
+        with _module_rlock:
+            if disposed:
+                return
+
+            callback_dict.pop(id, None)
+
+            disposed = True
+
+    callback_dict[id] = callback
+    return dispose
+
+
+def _patch(
+    stdout_or_stderr: IO[AnyStr],
+    callbacks: dict[int, _WriteCallback],
+) -> None:
+    orig_write: Callable[[AnyStr], int]
+
+    @wb_logging.log_to_all_runs()
+    def write_with_callbacks(s: AnyStr, /) -> int:
+        global _is_writing
+        n = orig_write(s)
+
+        # NOTE: Since _module_rlock is reentrant, this is safe. It will not
+        # deadlock if a callback invokes write() again.
+        with _module_rlock:
+            if _is_writing:
+                return n
+
+            _is_writing = True
+            try:
+                for cb in callbacks.values():
+                    cb(s, n)
+
+            except BaseException as e:
+                # Clear all callbacks on any exception to avoid infinite loops:
+                #
+                # * If we re-raise, an exception handler is likely to print
+                #   the exception to the console and trigger callbacks again
+                # * If we log, we can't guarantee that this doesn't print
+                #   to console.
+                #
+                # This is especially important for KeyboardInterrupt.
+                _stderr_callbacks.clear()
+                _stdout_callbacks.clear()
+
+                if isinstance(e, Exception):
+                    # We suppress Exceptions so that bugs in W&B code don't
+                    # cause the user's print() statements to raise errors.
+                    _logger.exception("Error in console callback, clearing all!")
+                else:
+                    # Re-raise errors like KeyboardInterrupt.
+                    raise
+
+            finally:
+                _is_writing = False
+
+        return n
+
+    orig_write = stdout_or_stderr.write
+
+    # mypy==1.14.1 fails to type-check this:
+    #   Incompatible types in assignment (expression has type
+    #   "Callable[[bytes], int]", variable has type overloaded function)
+    stdout_or_stderr.write = write_with_callbacks  # type: ignore
+
+
+try:
+    _patch(sys.stdout, _stdout_callbacks)
+    _patch(sys.stderr, _stderr_callbacks)
+except Exception as _patch_exception_cause:
+    _patch_exception = CannotCaptureConsoleError()
+    _patch_exception.__cause__ = _patch_exception_cause

wandb/sdk/lib/credentials.py

@@ -0,0 +1,141 @@
+import json
+import os
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import requests.utils
+
+from wandb.errors import AuthenticationError
+
+DEFAULT_WANDB_CREDENTIALS_FILE = Path(
+    os.path.expanduser("~/.config/wandb/credentials.json")
+)
+
+_expires_at_fmt = "%Y-%m-%d %H:%M:%S"
+
+
+def access_token(base_url: str, token_file: Path, credentials_file: Path) -> str:
+    """Retrieve an access token from the credentials file.
+
+    If no access token exists, create a new one by exchanging the identity
+    token from the token file, and save it to the credentials file.
+
+    Args:
+        base_url (str): The base URL of the server
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+        credentials_file (pathlib.Path): The path to file used to save
+        temporary access tokens
+
+    Returns:
+        str: The access token
+    """
+    if not credentials_file.exists():
+        _write_credentials_file(base_url, token_file, credentials_file)
+
+    data = _fetch_credentials(base_url, token_file, credentials_file)
+    return data["access_token"]
+
+
+def _write_credentials_file(base_url: str, token_file: Path, credentials_file: Path):
+    """Obtain an access token from the server and write it to the credentials file.
+
+    Args:
+        base_url (str): The base URL of the server
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+        credentials_file (pathlib.Path): The path to file used to save
+        temporary access tokens
+    """
+    credentials = _create_access_token(base_url, token_file)
+    data = {"credentials": {base_url: credentials}}
+    with open(credentials_file, "w") as file:
+        json.dump(data, file, indent=4)
+
+        # Set file permissions to be read/write by the owner only
+        os.chmod(credentials_file, 0o600)
+
+
+def _fetch_credentials(base_url: str, token_file: Path, credentials_file: Path) -> dict:
+    """Fetch the access token from the credentials file.
+
+    If the access token has expired, fetch a new one from the server and save it
+    to the credentials file.
+
+    Args:
+        base_url (str): The base URL of the server
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+        credentials_file (pathlib.Path): The path to file used to save
+        temporary access tokens
+
+    Returns:
+        dict: The credentials including the access token.
+    """
+    creds = {}
+    with open(credentials_file) as file:
+        data = json.load(file)
+        if "credentials" not in data:
+            data["credentials"] = {}
+        if base_url in data["credentials"]:
+            creds = data["credentials"][base_url]
+
+    expires_at = datetime.utcnow()
+    if "expires_at" in creds:
+        expires_at = datetime.strptime(creds["expires_at"], _expires_at_fmt)
+
+    if expires_at <= datetime.utcnow():
+        creds = _create_access_token(base_url, token_file)
+        with open(credentials_file, "w") as file:
+            data["credentials"][base_url] = creds
+            json.dump(data, file, indent=4)
+
+    return creds
+
+
+def _create_access_token(base_url: str, token_file: Path) -> dict:
+    """Exchange an identity token for an access token from the server.
+
+    Args:
+        base_url (str): The base URL of the server.
+        token_file (pathlib.Path): The path to the file containing the
+        identity token
+
+    Returns:
+        dict: The access token and its expiration.
+
+    Raises:
+        FileNotFoundError: If the token file is not found.
+        OSError: If there is an issue reading the token file.
+        AuthenticationError: If the server fails to provide an access token.
+    """
+    try:
+        with open(token_file) as file:
+            token = file.read().strip()
+    except FileNotFoundError as e:
+        raise FileNotFoundError(f"Identity token file not found: {token_file}") from e
+    except OSError as e:
+        raise OSError(
+            f"Failed to read the identity token from file: {token_file}"
+        ) from e
+
+    url = f"{base_url}/oidc/token"
+    data = {
+        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+        "assertion": token,
+    }
+    headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
+    response = requests.post(url, data=data, headers=headers)
+
+    if response.status_code != 200:
+        raise AuthenticationError(
+            f"Failed to retrieve access token: {response.status_code}, {response.text}"
+        )
+
+    resp_json = response.json()
+    expires_at = datetime.utcnow() + timedelta(seconds=float(resp_json["expires_in"]))
+    resp_json["expires_at"] = expires_at.strftime(_expires_at_fmt)
+    del resp_json["expires_in"]
+
+    return resp_json

wandb/sdk/lib/deprecate.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import wandb
+from wandb.proto.wandb_deprecated import DEPRECATED_FEATURES
+from wandb.sdk.lib import telemetry
+
+
+def deprecate(
+    field_name: DEPRECATED_FEATURES,
+    warning_message: str,
+    run: wandb.Run | None = None,
+) -> None:
+    """Warn the user that a feature has been deprecated.
+
+    If a run is provided, the given field on its telemetry is updated.
+    Otherwise, the global run is used.
+
+    Args:
+        field_name: The field on the Deprecated proto for this deprecation.
+        warning_message: The message to display to the user.
+        run: The run whose telemetry to update.
+    """
+    _run = run or wandb.run
+    with telemetry.context(run=_run) as tel:
+        setattr(tel.deprecated, field_name, True)
+
+    wandb.termwarn(warning_message, repeat=False)

wandb/sdk/lib/disabled.py

@@ -0,0 +1,30 @@
+from typing import Any
+
+from wandb.proto.wandb_deprecated import Deprecated
+from wandb.sdk.lib import deprecate
+
+
+class SummaryDisabled(dict):
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+    def __getattr__(self, key):
+        return self[key]
+
+    def __getitem__(self, key):
+        val = dict.__getitem__(self, key)
+        if isinstance(val, dict) and not isinstance(val, SummaryDisabled):
+            val = SummaryDisabled(val)
+            self[key] = val
+        return val
+
+
+class RunDisabled:
+    """Compatibility class for integrations that explicitly check for wandb.RunDisabled."""
+
+    def __getattr__(self, name: str) -> Any:
+        deprecate.deprecate(
+            field_name=Deprecated.run_disabled,
+            warning_message="RunDisabled is deprecated and is a no-op. "
+            '`wandb.init(mode="disabled")` now returns an instance of `wandb.Run`.',
+        )

wandb/sdk/lib/exit_hooks.py

@@ -0,0 +1,54 @@
+import sys
+import traceback
+from types import TracebackType
+from typing import TYPE_CHECKING, Optional, Type
+
+import wandb
+from wandb.errors import Error
+
+if TYPE_CHECKING:
+    from typing import NoReturn
+
+
+class ExitHooks:
+    exception: Optional[BaseException] = None
+
+    def __init__(self) -> None:
+        self.exit_code = 0
+        self.exception = None
+
+    def hook(self) -> None:
+        self._orig_exit = sys.exit
+        sys.exit = self.exit
+        self._orig_excepthook = (
+            sys.excepthook
+            if sys.excepthook
+            != sys.__excepthook__  # respect hooks by other libraries like pdb
+            else None
+        )
+        sys.excepthook = self.exc_handler  # type: ignore
+
+    def exit(self, code: object = 0) -> "NoReturn":
+        orig_code = code
+        code = code if code is not None else 0
+        code = code if isinstance(code, int) else 1
+        self.exit_code = code
+        self._orig_exit(orig_code)  # type: ignore
+
+    def was_ctrl_c(self) -> bool:
+        return isinstance(self.exception, KeyboardInterrupt)
+
+    def exc_handler(
+        self, exc_type: Type[BaseException], exc: BaseException, tb: TracebackType
+    ) -> None:
+        self.exit_code = 1
+        self.exception = exc
+        if issubclass(exc_type, Error):
+            wandb.termerror(str(exc), repeat=False)
+
+        if self.was_ctrl_c():
+            self.exit_code = 255
+
+        traceback.print_exception(exc_type, exc, tb)
+        if self._orig_excepthook:
+            self._orig_excepthook(exc_type, exc, tb)

wandb/sdk/lib/file_stream_utils.py

@@ -0,0 +1,118 @@
+#
+from typing import Any, Dict, Iterable
+
+
+def split_files(
+    files: Dict[str, Any], max_bytes: int = 10 * 1024 * 1024
+) -> Iterable[Dict[str, Dict]]:
+    """Split a file's dict (see `files` arg) into smaller dicts.
+
+    Each smaller dict will have at most `MAX_BYTES` size.
+
+    This method is used in `FileStreamAPI._send()` to limit the size of post requests
+    sent to wandb server.
+
+    Args:
+    files (dict): `dict` of form {file_name: {'content': ".....", 'offset': 0}}
+        The key `file_name` can also be mapped to a List [{"offset": int, "content": str}]
+    `max_bytes`: max size for chunk in bytes
+    """
+    current_volume: Dict[str, Dict] = {}
+    current_size = 0
+
+    def _str_size(x):
+        return len(x) if isinstance(x, bytes) else len(x.encode("utf-8"))
+
+    def _file_size(file):
+        size = file.get("_size")
+        if size is None:
+            size = sum(map(_str_size, file["content"]))
+            file["_size"] = size
+        return size
+
+    def _split_file(file, num_lines):
+        offset = file["offset"]
+        content = file["content"]
+        name = file["name"]
+        f1 = {"offset": offset, "content": content[:num_lines], "name": name}
+        f2 = {
+            "offset": offset + num_lines,
+            "content": content[num_lines:],
+            "name": name,
+        }
+        return f1, f2
+
+    def _num_lines_from_num_bytes(file, num_bytes):
+        size = 0
+        num_lines = 0
+        content = file["content"]
+        while num_lines < len(content):
+            size += _str_size(content[num_lines])
+            if size > num_bytes:
+                break
+            num_lines += 1
+        return num_lines
+
+    files_stack = []
+    for k, v in files.items():
+        if isinstance(v, list):
+            for item in v:
+                files_stack.append(
+                    {"name": k, "offset": item["offset"], "content": item["content"]}
+                )
+        else:
+            files_stack.append(
+                {"name": k, "offset": v["offset"], "content": v["content"]}
+            )
+
+    while files_stack:
+        f = files_stack.pop()
+        if f["name"] in current_volume:
+            files_stack.append(f)
+            yield current_volume
+            current_volume = {}
+            current_size = 0
+            continue
+        # For each file, we have to do 1 of 4 things:
+        # - Add the file as such to the current volume if possible.
+        # - Split the file and add the first part to the current volume and push the second part back onto the stack.
+        # - If that's not possible, check if current volume is empty:
+        # - If empty, add first line of file to current volume and push rest onto stack (This volume will exceed MAX_MB).
+        # - If not, push file back to stack and yield current volume.
+        fsize = _file_size(f)
+        rem = max_bytes - current_size
+        if fsize <= rem:
+            current_volume[f["name"]] = {
+                "offset": f["offset"],
+                "content": f["content"],
+            }
+            current_size += fsize
+        else:
+            num_lines = _num_lines_from_num_bytes(f, rem)
+            if not num_lines and not current_volume:
+                num_lines = 1
+            if num_lines:
+                f1, f2 = _split_file(f, num_lines)
+                current_volume[f1["name"]] = {
+                    "offset": f1["offset"],
+                    "content": f1["content"],
+                }
+                files_stack.append(f2)
+                yield current_volume
+                current_volume = {}
+                current_size = 0
+                continue
+            else:
+                files_stack.append(f)
+                yield current_volume
+                current_volume = {}
+                current_size = 0
+                continue
+        if current_size >= max_bytes:
+            yield current_volume
+            current_volume = {}
+            current_size = 0
+            continue
+
+    if current_volume:
+        yield current_volume

wandb/sdk/lib/filenames.py

@@ -0,0 +1,64 @@
+import os
+from typing import Callable, Generator, Union
+
+WANDB_DIRS = ("wandb", ".wandb")
+
+CONFIG_FNAME = "config.yaml"
+OUTPUT_FNAME = "output.log"
+DIFF_FNAME = "diff.patch"
+SUMMARY_FNAME = "wandb-summary.json"
+METADATA_FNAME = "wandb-metadata.json"
+REQUIREMENTS_FNAME = "requirements.txt"
+HISTORY_FNAME = "wandb-history.jsonl"
+EVENTS_FNAME = "wandb-events.jsonl"
+JOBSPEC_FNAME = "wandb-jobspec.json"
+CONDA_ENVIRONMENTS_FNAME = "conda-environment.yaml"
+
+
+def is_wandb_file(name: str) -> bool:
+    return (
+        name.startswith("wandb")
+        or name == METADATA_FNAME
+        or name == CONFIG_FNAME
+        or name == REQUIREMENTS_FNAME
+        or name == OUTPUT_FNAME
+        or name == DIFF_FNAME
+        or name == CONDA_ENVIRONMENTS_FNAME
+    )
+
+
+def filtered_dir(
+    root: str,
+    include_fn: Union[Callable[[str, str], bool], Callable[[str], bool]],
+    exclude_fn: Union[Callable[[str, str], bool], Callable[[str], bool]],
+) -> Generator[str, None, None]:
+    """Simple generator to walk a directory."""
+    import inspect
+
+    # compatibility with old API, which didn't pass root
+    def _include_fn(path: str, root: str) -> bool:
+        return (
+            include_fn(path, root)  # type: ignore
+            if len(inspect.signature(include_fn).parameters) == 2
+            else include_fn(path)  # type: ignore
+        )
+
+    def _exclude_fn(path: str, root: str) -> bool:
+        return (
+            exclude_fn(path, root)  # type: ignore
+            if len(inspect.signature(exclude_fn).parameters) == 2
+            else exclude_fn(path)  # type: ignore
+        )
+
+    for dirpath, _, files in os.walk(root):
+        for fname in files:
+            file_path = os.path.join(dirpath, fname)
+            if _include_fn(file_path, root) and not _exclude_fn(file_path, root):
+                yield file_path
+
+
+def exclude_wandb_fn(path: str, root: str) -> bool:
+    return any(
+        os.path.relpath(path, root).startswith(wandb_dir + os.sep)
+        for wandb_dir in WANDB_DIRS
+    )