wandb 0.21.2__py3-none-macosx_12_0_x86_64.whl

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