wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/lib/fsm.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python
+"""Finite state machine.
+
+Simple FSM implementation.
+
+Usage:
+    ```python
+    class A:
+        def on_output(self, inputs) -> None:
+            pass
+
+
+    class B:
+        def on_output(self, inputs) -> None:
+            pass
+
+
+    def to_b(inputs) -> bool:
+        return True
+
+
+    def to_a(inputs) -> bool:
+        return True
+
+
+    f = Fsm(states=[A(), B()], table={A: [(to_b, B)], B: [(to_a, A)]})
+    f.run({"input1": 1, "input2": 2})
+    ```
+"""
+
+import sys
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import (
+    Callable,
+    Dict,
+    Generic,
+    Optional,
+    Protocol,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+    runtime_checkable,
+)
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+T_FsmInputs = TypeVar("T_FsmInputs", contravariant=True)
+T_FsmContext = TypeVar("T_FsmContext")
+T_FsmContext_cov = TypeVar("T_FsmContext_cov", covariant=True)
+T_FsmContext_contra = TypeVar("T_FsmContext_contra", contravariant=True)
+
+
+@runtime_checkable
+class FsmStateCheck(Protocol[T_FsmInputs]):
+    @abstractmethod
+    def on_check(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
+
+
+@runtime_checkable
+class FsmStateOutput(Protocol[T_FsmInputs]):
+    @abstractmethod
+    def on_state(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
+
+
+@runtime_checkable
+class FsmStateEnter(Protocol[T_FsmInputs]):
+    @abstractmethod
+    def on_enter(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
+
+
+@runtime_checkable
+class FsmStateEnterWithContext(Protocol[T_FsmInputs, T_FsmContext_contra]):
+    @abstractmethod
+    def on_enter(
+        self, inputs: T_FsmInputs, context: T_FsmContext_contra
+    ) -> None: ...  # pragma: no cover
+
+
+@runtime_checkable
+class FsmStateStay(Protocol[T_FsmInputs]):
+    @abstractmethod
+    def on_stay(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
+
+
+@runtime_checkable
+class FsmStateExit(Protocol[T_FsmInputs, T_FsmContext_cov]):
+    @abstractmethod
+    def on_exit(self, inputs: T_FsmInputs) -> T_FsmContext_cov: ...  # pragma: no cover
+
+
+# It would be nice if python provided optional protocol members, but it doesnt as described here:
+#   https://peps.python.org/pep-0544/#support-optional-protocol-members
+# Until then, we can only enforce that a state at least supports one protocol interface.  This
+# unfortunately will not check the signature of other potential protocols.
+FsmState: TypeAlias = Union[
+    FsmStateCheck[T_FsmInputs],
+    FsmStateOutput[T_FsmInputs],
+    FsmStateEnter[T_FsmInputs],
+    FsmStateEnterWithContext[T_FsmInputs, T_FsmContext],
+    FsmStateStay[T_FsmInputs],
+    FsmStateExit[T_FsmInputs, T_FsmContext],
+]
+
+
+@dataclass
+class FsmEntry(Generic[T_FsmInputs, T_FsmContext]):
+    condition: Callable[[T_FsmInputs], bool]
+    target_state: Type[FsmState[T_FsmInputs, T_FsmContext]]
+    action: Optional[Callable[[T_FsmInputs], None]] = None
+
+
+FsmTableWithContext: TypeAlias = Dict[
+    Type[FsmState[T_FsmInputs, T_FsmContext]],
+    Sequence[FsmEntry[T_FsmInputs, T_FsmContext]],
+]
+
+
+FsmTable: TypeAlias = FsmTableWithContext[T_FsmInputs, None]
+
+
+class FsmWithContext(Generic[T_FsmInputs, T_FsmContext]):
+    _state_dict: Dict[Type[FsmState], FsmState]
+    _table: FsmTableWithContext[T_FsmInputs, T_FsmContext]
+    _state: FsmState[T_FsmInputs, T_FsmContext]
+    _states: Sequence[FsmState]
+
+    def __init__(
+        self,
+        states: Sequence[FsmState],
+        table: FsmTableWithContext[T_FsmInputs, T_FsmContext],
+    ) -> None:
+        self._states = states
+        self._table = table
+        self._state_dict = {type(s): s for s in states}
+        self._state = self._state_dict[type(states[0])]
+
+    def _transition(
+        self,
+        inputs: T_FsmInputs,
+        new_state: Type[FsmState[T_FsmInputs, T_FsmContext]],
+        action: Optional[Callable[[T_FsmInputs], None]],
+    ) -> None:
+        if action:
+            action(inputs)
+
+        context = None
+        if isinstance(self._state, FsmStateExit):
+            context = self._state.on_exit(inputs)
+
+        prev_state = type(self._state)
+        if prev_state == new_state:
+            if isinstance(self._state, FsmStateStay):
+                self._state.on_stay(inputs)
+        else:
+            self._state = self._state_dict[new_state]
+            if context and isinstance(self._state, FsmStateEnterWithContext):
+                self._state.on_enter(inputs, context=context)
+            elif isinstance(self._state, FsmStateEnter):
+                self._state.on_enter(inputs)
+
+    def _check_transitions(self, inputs: T_FsmInputs) -> None:
+        for entry in self._table[type(self._state)]:
+            if entry.condition(inputs):
+                self._transition(inputs, entry.target_state, entry.action)
+                return
+
+    def input(self, inputs: T_FsmInputs) -> None:
+        if isinstance(self._state, FsmStateCheck):
+            self._state.on_check(inputs)
+        self._check_transitions(inputs)
+        if isinstance(self._state, FsmStateOutput):
+            self._state.on_state(inputs)
+
+
+Fsm: TypeAlias = FsmWithContext[T_FsmInputs, None]

wandb/sdk/lib/gitlib.py

@@ -0,0 +1,239 @@
+import configparser
+import logging
+import os
+from typing import TYPE_CHECKING, Any, Optional
+from urllib.parse import urlparse, urlunparse
+
+import wandb
+
+try:
+    from git import (  # type: ignore
+        GitCommandError,
+        InvalidGitRepositoryError,
+        NoSuchPathError,
+        Repo,
+    )
+except ImportError:
+    Repo = None  # type: ignore
+
+if TYPE_CHECKING:
+    from git import Repo
+
+
+logger = logging.getLogger(__name__)
+
+
+class GitRepo:
+    def __init__(
+        self,
+        root: Optional[str] = None,
+        remote: str = "origin",
+        lazy: bool = True,
+        remote_url: Optional[str] = None,
+        commit: Optional[str] = None,
+    ) -> None:
+        self.remote_name = remote if remote_url is None else None
+        self._root = root
+        self._remote_url = remote_url
+        self._commit = commit
+        self._repo = None
+        self._repo_initialized = False
+        if not lazy:
+            self._repo = self._init_repo()
+
+    def _init_repo(self) -> Optional[Repo]:
+        self._repo_initialized = True
+        if Repo is None:
+            return None
+        if self.remote_name is None:
+            return None
+        try:
+            return Repo(self._root or os.getcwd(), search_parent_directories=True)
+        except FileNotFoundError:
+            wandb.termwarn("current working directory has been invalidated")
+            logger.warn("current working directory has been invalidated")
+        except InvalidGitRepositoryError:
+            logger.debug("git repository is invalid")
+        except NoSuchPathError:
+            wandb.termwarn(f"git root {self._root} does not exist")
+            logger.warn(f"git root {self._root} does not exist")
+        return None
+
+    @property
+    def repo(self) -> Optional[Repo]:
+        if not self._repo_initialized:
+            self._repo = self._init_repo()
+        return self._repo
+
+    @property
+    def auto(self) -> bool:
+        return self._remote_url is None
+
+    def is_untracked(self, file_name: str) -> Optional[bool]:
+        if not self.repo:
+            return True
+        try:
+            return file_name in self.repo.untracked_files
+        except GitCommandError:
+            return None
+
+    @property
+    def enabled(self) -> bool:
+        return bool(self.repo)
+
+    @property
+    def root(self) -> Any:
+        if not self.repo:
+            return None
+        try:
+            return self.repo.git.rev_parse("--show-toplevel")
+        except GitCommandError as e:
+            # todo: collect telemetry on this
+            logger.error(f"git root error: {e}")
+            return None
+
+    @property
+    def dirty(self) -> Any:
+        if not self.repo:
+            return False
+        try:
+            return self.repo.is_dirty()
+        except GitCommandError:
+            return False
+
+    @property
+    def email(self) -> Optional[str]:
+        if not self.repo:
+            return None
+        try:
+            return self.repo.config_reader().get_value("user", "email")  # type: ignore
+        except configparser.Error:
+            return None
+
+    @property
+    def last_commit(self) -> Any:
+        if self._commit:
+            return self._commit
+        if not self.repo:
+            return None
+        if not self.repo.head or not self.repo.head.is_valid():
+            return None
+        # TODO: Saw a user getting a Unicode decode error when parsing refs,
+        # more details on implementing a real fix in [WB-4064]
+        try:
+            if len(self.repo.refs) > 0:  # type: ignore[arg-type]
+                return self.repo.head.commit.hexsha
+            else:
+                return self.repo.git.show_ref("--head").split(" ")[0]
+        except Exception:
+            logger.exception("Unable to find most recent commit in git")
+            return None
+
+    @property
+    def branch(self) -> Any:
+        if not self.repo:
+            return None
+        return self.repo.head.ref.name
+
+    @property
+    def remote(self) -> Any:
+        if not self.repo:
+            return None
+        try:
+            return self.repo.remotes[self.remote_name]  # type: ignore[index]
+        except IndexError:
+            return None
+
+    # the --submodule=diff option doesn't exist in pre-2.11 versions of git (november 2016)
+    # https://stackoverflow.com/questions/10757091/git-list-of-all-changed-files-including-those-in-submodules
+    @property
+    def has_submodule_diff(self) -> bool:
+        if not self.repo:
+            return False
+        return bool(self.repo.git.version_info >= (2, 11, 0))
+
+    @property
+    def remote_url(self) -> Any:
+        if self._remote_url:
+            return self._remote_url
+        if not self.remote:
+            return None
+        parsed = urlparse(self.remote.url)
+        hostname = parsed.hostname
+        if parsed.port is not None:
+            hostname = f"{hostname}:{parsed.port}"
+        if parsed.password is not None:
+            return urlunparse(parsed._replace(netloc=f"{parsed.username}:@{hostname}"))
+        return urlunparse(parsed._replace(netloc=hostname))
+
+    @property
+    def root_dir(self) -> Any:
+        if not self.repo:
+            return None
+        try:
+            return self.repo.git.rev_parse("--show-toplevel")
+        except GitCommandError:
+            return None
+
+    def get_upstream_fork_point(self) -> Any:
+        """Get the most recent ancestor of HEAD that occurs on an upstream branch.
+
+        First looks at the current branch's tracking branch, if applicable. If
+        that doesn't work, looks at every other branch to find the most recent
+        ancestor of HEAD that occurs on a tracking branch.
+
+        Returns:
+            git.Commit object or None
+        """
+        possible_relatives = []
+        try:
+            if not self.repo:
+                return None
+            try:
+                active_branch = self.repo.active_branch
+            except (TypeError, ValueError):
+                logger.debug("git is in a detached head state")
+                return None  # detached head
+            else:
+                tracking_branch = active_branch.tracking_branch()
+                if tracking_branch:
+                    possible_relatives.append(tracking_branch.commit)
+
+            if not possible_relatives:
+                for branch in self.repo.branches:  # type: ignore[attr-defined]
+                    tracking_branch = branch.tracking_branch()
+                    if tracking_branch is not None:
+                        possible_relatives.append(tracking_branch.commit)
+
+            head = self.repo.head
+            most_recent_ancestor = None
+            for possible_relative in possible_relatives:
+                # at most one:
+                for ancestor in self.repo.merge_base(head, possible_relative):
+                    if most_recent_ancestor is None:
+                        most_recent_ancestor = ancestor
+                    elif self.repo.is_ancestor(most_recent_ancestor, ancestor):  # type: ignore
+                        most_recent_ancestor = ancestor
+            return most_recent_ancestor
+        except GitCommandError as e:
+            logger.debug("git remote upstream fork point could not be found")
+            logger.debug(str(e))
+            return None
+
+    def tag(self, name: str, message: Optional[str]) -> Any:
+        if not self.repo:
+            return None
+        try:
+            return self.repo.create_tag(f"wandb/{name}", message=message, force=True)
+        except GitCommandError:
+            print("Failed to tag repository.")
+            return None
+
+    def push(self, name: str) -> Any:
+        if not self.remote:
+            return None
+        try:
+            return self.remote.push(f"wandb/{name}", force=True)
+        except GitCommandError:
+            logger.debug("failed to push git")
+            return None

wandb/sdk/lib/gql_request.py

@@ -0,0 +1,65 @@
+"""A simple GraphQL client for sending queries and mutations.
+
+Note: This was originally wandb/vendor/gql-0.2.0/wandb_gql/transport/requests.py
+The only substantial change is to re-use a requests.Session object.
+"""
+
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+import requests
+from wandb_gql.transport.http import HTTPTransport
+from wandb_graphql.execution import ExecutionResult
+from wandb_graphql.language import ast
+from wandb_graphql.language.printer import print_ast
+
+
+class GraphQLSession(HTTPTransport):
+    def __init__(
+        self,
+        url: str,
+        auth: Optional[Union[Tuple[str, str], Callable]] = None,
+        use_json: bool = False,
+        timeout: Optional[Union[int, float]] = None,
+        proxies: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Setup a session for sending GraphQL queries and mutations.
+
+        Args:
+            url (str): The GraphQL URL
+            auth (tuple or callable): Auth tuple or callable for Basic/Digest/Custom HTTP Auth
+            use_json (bool): Send request body as JSON instead of form-urlencoded
+            timeout (int, float): Specifies a default timeout for requests (Default: None)
+        """
+        super().__init__(url, **kwargs)
+        self.session = requests.Session()
+        if proxies:
+            self.session.proxies.update(proxies)
+        self.session.auth = auth
+        self.default_timeout = timeout
+        self.use_json = use_json
+
+    def execute(
+        self,
+        document: ast.Node,
+        variable_values: Optional[Dict] = None,
+        timeout: Optional[Union[int, float]] = None,
+    ) -> ExecutionResult:
+        query_str = print_ast(document)
+        payload = {"query": query_str, "variables": variable_values or {}}
+
+        data_key = "json" if self.use_json else "data"
+        post_args = {
+            "headers": self.headers,
+            "cookies": self.cookies,
+            "timeout": timeout or self.default_timeout,
+            data_key: payload,
+        }
+        request = self.session.post(self.url, **post_args)
+        request.raise_for_status()
+
+        result = request.json()
+        data, errors = result.get("data"), result.get("errors")
+        if data is None and errors is None:
+            raise RuntimeError(f"Received non-compatible response: {result}")
+        return ExecutionResult(data=data, errors=errors)

wandb/sdk/lib/handler_util.py

@@ -0,0 +1,21 @@
+import wandb.data_types as data_types
+
+
+def get_types():
+    classes = map(data_types.__dict__.get, data_types.__all__)
+    types = []
+    for cls in classes:
+        if hasattr(cls, "_log_type") and cls._log_type is not None:
+            types.append(cls._log_type)
+    # add table-file type because this is a special case
+    # that does not have a matching _log_type for artifacts
+    # and files
+    types.append("table-file")
+    return types
+
+
+WANDB_TYPES = get_types()
+
+
+def metric_is_wandb_dict(metric):
+    return "_type" in list(metric.keys()) and metric["_type"] in WANDB_TYPES

wandb/sdk/lib/hashutil.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import mmap
+import sys
+from typing import TYPE_CHECKING, NewType
+
+from wandb.sdk.lib.paths import StrPath
+
+if TYPE_CHECKING:
+    import _hashlib  # type: ignore[import-not-found]
+
+ETag = NewType("ETag", str)
+HexMD5 = NewType("HexMD5", str)
+B64MD5 = NewType("B64MD5", str)
+
+
+def _md5(data: bytes = b"") -> _hashlib.HASH:
+    """Allow FIPS-compliant md5 hash when supported."""
+    if sys.version_info >= (3, 9):
+        return hashlib.md5(data, usedforsecurity=False)
+    else:
+        return hashlib.md5(data)
+
+
+def md5_string(string: str) -> B64MD5:
+    return _b64_from_hasher(_md5(string.encode("utf-8")))
+
+
+def _b64_from_hasher(hasher: _hashlib.HASH) -> B64MD5:
+    return B64MD5(base64.b64encode(hasher.digest()).decode("ascii"))
+
+
+def b64_to_hex_id(string: B64MD5) -> HexMD5:
+    return HexMD5(base64.standard_b64decode(string).hex())
+
+
+def hex_to_b64_id(encoded_string: str | bytes) -> B64MD5:
+    if isinstance(encoded_string, bytes):
+        encoded_string = encoded_string.decode("utf-8")
+    as_str = bytes.fromhex(encoded_string)
+    return B64MD5(base64.standard_b64encode(as_str).decode("utf-8"))
+
+
+def md5_file_b64(*paths: StrPath) -> B64MD5:
+    return _b64_from_hasher(_md5_file_hasher(*paths))
+
+
+def md5_file_hex(*paths: StrPath) -> HexMD5:
+    return HexMD5(_md5_file_hasher(*paths).hexdigest())
+
+
+_KB: int = 1_024
+_CHUNKSIZE: int = 128 * _KB
+"""Chunk size (in bytes) for iteratively reading from file, if needed."""
+
+
+def _md5_file_hasher(*paths: StrPath) -> _hashlib.HASH:
+    md5_hash = _md5()
+
+    # Note: We use str paths (instead of pathlib.Path objs) for minor perf improvements.
+    for path in sorted(map(str, paths)):
+        with open(path, "rb") as f:
+            try:
+                with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as mview:
+                    md5_hash.update(mview)
+            except OSError:
+                # This occurs if the mmap-ed file is on a different/mounted filesystem,
+                # so we'll fall back on a less performant implementation.
+
+                # Note: At the time of implementation, the walrus operator `:=`
+                # is avoided to maintain support for users on python 3.7.
+                # Consider revisiting once 3.7 support is no longer needed.
+                chunk = f.read(_CHUNKSIZE)
+                while chunk:
+                    md5_hash.update(chunk)
+                    chunk = f.read(_CHUNKSIZE)
+            except ValueError:
+                # This occurs when mmap-ing an empty file, which can be skipped.
+                # See: https://github.com/python/cpython/blob/986a4e1b6fcae7fe7a1d0a26aea446107dd58dd2/Modules/mmapmodule.c#L1589
+                pass
+
+    return md5_hash