wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/launch/wandb_reference.py

@@ -0,0 +1,138 @@
+"""Support for parsing W&B URLs (which might be user provided) into constituent parts."""
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Optional
+from urllib.parse import urlparse
+
+PREFIX_HTTP = "http://"
+PREFIX_HTTPS = "https://"
+
+
+class ReferenceType(IntEnum):
+    RUN = 1
+    JOB = 2
+
+
+# Ideally we would not overload the URL paths as we do.
+# TODO: Not sure these are exhaustive, and even if so more special paths might get added.
+#       Would be good to have restrictions that we could check.
+RESERVED_NON_ENTITIES = (
+    "create-team",
+    "fully-connected",
+    "registry",
+    "settings",
+    "subscriptions",
+)
+RESERVED_NON_PROJECTS = (
+    "likes",
+    "projects",
+)
+RESERVED_JOB_PATHS = ("_view",)
+
+
+@dataclass
+class WandbReference:
+    # TODO: This will include port, should we separate that out?
+    host: Optional[str] = None
+
+    entity: Optional[str] = None
+    project: Optional[str] = None
+
+    # Set when we don't know how to parse yet
+    path: Optional[str] = None
+
+    # Reference type will determine what other fields are set
+    ref_type: Optional[ReferenceType] = None
+
+    run_id: Optional[str] = None
+
+    job_name: Optional[str] = None
+    job_alias: str = "latest"  # In addition to an alias can be a version specifier
+
+    def is_bare(self) -> bool:
+        return self.host is None
+
+    def is_job(self) -> bool:
+        return self.ref_type == ReferenceType.JOB
+
+    def is_run(self) -> bool:
+        return self.ref_type == ReferenceType.RUN
+
+    def is_job_or_run(self) -> bool:
+        return self.is_job() or self.is_run()
+
+    def job_reference(self) -> str:
+        assert self.is_job()
+        return f"{self.job_name}:{self.job_alias}"
+
+    def job_reference_scoped(self) -> str:
+        assert self.entity
+        assert self.project
+        unscoped = self.job_reference()
+        return f"{self.entity}/{self.project}/{unscoped}"
+
+    def url_host(self) -> str:
+        return f"{PREFIX_HTTPS}{self.host}" if self.host else ""
+
+    def url_entity(self) -> str:
+        assert self.entity
+        return f"{self.url_host()}/{self.entity}"
+
+    def url_project(self) -> str:
+        assert self.project
+        return f"{self.url_entity()}/{self.project}"
+
+    @staticmethod
+    def parse(uri: str) -> Optional["WandbReference"]:
+        """Attempt to parse a string as a W&B URL."""
+        # TODO: Error if HTTP and host is not localhost?
+        if (
+            not uri.startswith("/")
+            and not uri.startswith(PREFIX_HTTP)
+            and not uri.startswith(PREFIX_HTTPS)
+        ):
+            return None
+
+        ref = WandbReference()
+
+        # This takes care of things like query and fragment
+        parsed = urlparse(uri)
+        if parsed.netloc:
+            ref.host = parsed.netloc
+
+        if not parsed.path.startswith("/"):
+            return ref
+
+        ref.path = parsed.path[1:]
+        parts = ref.path.split("/")
+        if len(parts) > 0:
+            if parts[0] not in RESERVED_NON_ENTITIES:
+                ref.path = None
+                ref.entity = parts[0]
+                if len(parts) > 1:
+                    if parts[1] not in RESERVED_NON_PROJECTS:
+                        ref.project = parts[1]
+                        if len(parts) > 3 and parts[2] == "runs":
+                            ref.ref_type = ReferenceType.RUN
+                            ref.run_id = parts[3]
+                        elif (
+                            len(parts) > 4
+                            and parts[2] == "artifacts"
+                            and parts[3] == "job"
+                        ):
+                            ref.ref_type = ReferenceType.JOB
+                            ref.job_name = parts[4]
+                            if len(parts) > 5 and parts[5] not in RESERVED_JOB_PATHS:
+                                ref.job_alias = parts[5]
+                        # TODO: Right now we are not tracking selection as part of URL state in the Jobs tab.
+                        #       If that changes we'll want to update this.
+
+        return ref
+
+    @staticmethod
+    def is_uri_job_or_run(uri: str) -> bool:
+        ref = WandbReference.parse(uri)
+        if ref and ref.is_job_or_run():
+            return True
+        return False

wandb/sdk/lib/__init__.py

@@ -0,0 +1,5 @@
+from . import lazyloader
+from .disabled import RunDisabled, SummaryDisabled
+from .run_moment import RunMoment
+
+__all__ = ("lazyloader", "RunDisabled", "SummaryDisabled", "RunMoment")

wandb/sdk/lib/apikey.py

@@ -0,0 +1,269 @@
+"""apikey util."""
+
+import os
+import platform
+import stat
+import sys
+import textwrap
+from functools import partial
+
+# import Literal
+from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional, Union
+from urllib.parse import urlparse
+
+import click
+from requests.utils import NETRC_FILES, get_netrc_auth
+
+import wandb
+from wandb.apis import InternalApi
+from wandb.errors import term
+from wandb.errors.links import url_registry
+from wandb.util import _is_databricks, isatty, prompt_choices
+
+LOGIN_CHOICE_ANON = "Private W&B dashboard, no account required"
+LOGIN_CHOICE_NEW = "Create a W&B account"
+LOGIN_CHOICE_EXISTS = "Use an existing W&B account"
+LOGIN_CHOICE_DRYRUN = "Don't visualize my results"
+LOGIN_CHOICE_NOTTY = "Unconfigured"
+LOGIN_CHOICES = [
+    LOGIN_CHOICE_ANON,
+    LOGIN_CHOICE_NEW,
+    LOGIN_CHOICE_EXISTS,
+    LOGIN_CHOICE_DRYRUN,
+]
+
+Mode = Literal["allow", "must", "never", "false", "true"]
+
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
+
+getpass = partial(click.prompt, hide_input=True, err=True)
+
+
+def _fixup_anon_mode(default: Optional[Mode]) -> Optional[Mode]:
+    # Convert weird anonymode values from legacy settings files
+    # into one of our expected values.
+    anon_mode = default or "never"
+    mapping: Dict[Mode, Mode] = {"true": "allow", "false": "never"}
+    return mapping.get(anon_mode, anon_mode)
+
+
+def get_netrc_file_path() -> str:
+    """Return the path to the netrc file."""
+    # if the NETRC environment variable is set, use that
+    netrc_file = os.environ.get("NETRC")
+    if netrc_file:
+        return os.path.expanduser(netrc_file)
+
+    # if either .netrc or _netrc exists in the home directory, use that
+    for netrc_file in NETRC_FILES:
+        home_dir = os.path.expanduser("~")
+        if os.path.exists(os.path.join(home_dir, netrc_file)):
+            return os.path.join(home_dir, netrc_file)
+
+    # otherwise, use .netrc on non-Windows platforms and _netrc on Windows
+    netrc_file = ".netrc" if platform.system() != "Windows" else "_netrc"
+
+    return os.path.join(os.path.expanduser("~"), netrc_file)
+
+
+def prompt_api_key(  # noqa: C901
+    settings: "Settings",
+    api: Optional[InternalApi] = None,
+    input_callback: Optional[Callable] = None,
+    browser_callback: Optional[Callable] = None,
+    no_offline: bool = False,
+    no_create: bool = False,
+    local: bool = False,
+) -> Union[str, bool, None]:
+    """Prompt for api key.
+
+    Returns:
+        str - if key is configured
+        None - if dryrun is selected
+        False - if unconfigured (notty)
+    """
+    input_callback = input_callback or getpass
+    log_string = term.LOG_STRING
+    api = api or InternalApi(settings)
+    anon_mode = _fixup_anon_mode(settings.anonymous)  # type: ignore
+    jupyter = settings._jupyter or False
+    app_url = api.app_url
+
+    choices = [choice for choice in LOGIN_CHOICES]
+    if anon_mode == "never":
+        # Omit LOGIN_CHOICE_ANON as a choice if the env var is set to never
+        choices.remove(LOGIN_CHOICE_ANON)
+    if (jupyter and not settings.login_timeout) or no_offline:
+        choices.remove(LOGIN_CHOICE_DRYRUN)
+    if (jupyter and not settings.login_timeout) or no_create:
+        choices.remove(LOGIN_CHOICE_NEW)
+
+    if jupyter and "google.colab" in sys.modules:
+        log_string = term.LOG_STRING_NOCOLOR
+        key = wandb.jupyter.attempt_colab_login(app_url)  # type: ignore
+        if key is not None:
+            write_key(settings, key, api=api)
+            return key  # type: ignore
+
+    if anon_mode == "must":
+        result = LOGIN_CHOICE_ANON
+    # If we're not in an interactive environment, default to dry-run.
+    elif (
+        not jupyter and (not isatty(sys.stdout) or not isatty(sys.stdin))
+    ) or _is_databricks():
+        result = LOGIN_CHOICE_NOTTY
+    elif local:
+        result = LOGIN_CHOICE_EXISTS
+    elif len(choices) == 1:
+        result = choices[0]
+    else:
+        result = prompt_choices(
+            choices, input_timeout=settings.login_timeout, jupyter=jupyter
+        )
+
+    api_ask = (
+        f"{log_string}: Paste an API key from your profile and hit enter, "
+        "or press ctrl+c to quit"
+    )
+    if result == LOGIN_CHOICE_ANON:
+        key = api.create_anonymous_api_key()
+
+        write_key(settings, key, api=api, anonymous=True)
+        return key  # type: ignore
+    elif result == LOGIN_CHOICE_NEW:
+        key = browser_callback(signup=True) if browser_callback else None
+
+        if not key:
+            wandb.termlog(f"Create an account here: {app_url}/authorize?signup=true")
+            key = input_callback(api_ask).strip()
+
+        write_key(settings, key, api=api)
+        return key  # type: ignore
+    elif result == LOGIN_CHOICE_EXISTS:
+        key = browser_callback() if browser_callback else None
+
+        if not key:
+            if not (settings.is_local or local):
+                host = app_url
+                for prefix in ("http://", "https://"):
+                    if app_url.startswith(prefix):
+                        host = app_url[len(prefix) :]
+                wandb.termlog(
+                    f"Logging into {host}. (Learn how to deploy a W&B server "
+                    f"locally: {url_registry.url('wandb-server')})"
+                )
+            wandb.termlog(
+                f"You can find your API key in your browser here: {app_url}/authorize"
+            )
+            key = input_callback(api_ask).strip()
+        write_key(settings, key, api=api)
+        return key  # type: ignore
+    elif result == LOGIN_CHOICE_NOTTY:
+        # TODO: Needs refactor as this needs to be handled by caller
+        return False
+    elif result == LOGIN_CHOICE_DRYRUN:
+        return None
+    else:
+        # Jupyter environments don't have a tty, but we can still try logging in using
+        # the browser callback if one is supplied.
+        key, anonymous = (
+            browser_callback() if jupyter and browser_callback else (None, False)
+        )
+
+        write_key(settings, key, api=api)
+        return key  # type: ignore
+
+
+def write_netrc(host: str, entity: str, key: str) -> Optional[bool]:
+    """Add our host and key to .netrc."""
+    _, key_suffix = key.split("-", 1) if "-" in key else ("", key)
+    if len(key_suffix) != 40:
+        wandb.termerror(
+            "API-key must be exactly 40 characters long: {} ({} chars)".format(
+                key_suffix, len(key_suffix)
+            )
+        )
+        return None
+    try:
+        normalized_host = urlparse(host).netloc.split(":")[0]
+        netrc_path = get_netrc_file_path()
+        wandb.termlog(
+            f"Appending key for {normalized_host} to your netrc file: {netrc_path}"
+        )
+        machine_line = f"machine {normalized_host}"
+        orig_lines = None
+        try:
+            with open(netrc_path) as f:
+                orig_lines = f.read().strip().split("\n")
+        except OSError:
+            pass
+        with open(netrc_path, "w") as f:
+            if orig_lines:
+                # delete this machine from the file if it's already there.
+                skip = 0
+                for line in orig_lines:
+                    # we fix invalid netrc files with an empty host that we wrote before
+                    # verifying host...
+                    if line == "machine " or machine_line in line:
+                        skip = 2
+                    elif skip:
+                        skip -= 1
+                    else:
+                        f.write("{}\n".format(line))
+            f.write(
+                textwrap.dedent(
+                    """\
+            machine {host}
+              login {entity}
+              password {key}
+            """
+                ).format(host=normalized_host, entity=entity, key=key)
+            )
+        os.chmod(netrc_path, stat.S_IRUSR | stat.S_IWUSR)
+        return True
+    except OSError:
+        wandb.termerror(f"Unable to read {netrc_path}")
+        return None
+
+
+def write_key(
+    settings: "Settings",
+    key: Optional[str],
+    api: Optional["InternalApi"] = None,
+    anonymous: bool = False,
+) -> None:
+    if not key:
+        raise ValueError("No API key specified.")
+
+    # TODO(jhr): api shouldn't be optional or it shouldn't be passed, clean up callers
+    api = api or InternalApi()
+
+    # Normal API keys are 40-character hex strings. On-prem API keys have a
+    # variable-length prefix, a dash, then the 40-char string.
+    _, suffix = key.split("-", 1) if "-" in key else ("", key)
+
+    if len(suffix) != 40:
+        raise ValueError(
+            "API key must be 40 characters long, yours was {}".format(len(key))
+        )
+
+    if anonymous:
+        api.set_setting("anonymous", "true", globally=True, persist=True)
+    else:
+        api.clear_setting("anonymous", globally=True, persist=True)
+
+    write_netrc(settings.base_url, "user", key)
+
+
+def api_key(settings: Optional["Settings"] = None) -> Optional[str]:
+    if settings is None:
+        settings = wandb.setup().settings  # type: ignore
+        assert settings is not None
+    if settings.api_key:
+        return settings.api_key
+    auth = get_netrc_auth(settings.base_url)
+    if auth:
+        return auth[-1]
+    return None

wandb/sdk/lib/capped_dict.py

@@ -0,0 +1,26 @@
+import collections
+from typing import Any, Optional
+
+
+class CappedDict(collections.OrderedDict):
+    default_max_size = 50
+
+    def __init__(self, max_size: Optional[int] = None) -> None:
+        self.max_size = max_size or self.default_max_size
+        super().__init__()
+
+    def __setitem__(self, key: str, val: Any) -> None:
+        if key not in self:
+            max_size = self.max_size - 1
+            self._prune_dict(max_size)
+        super().__setitem__(key, val)
+
+    def update(self, **kwargs: Any) -> None:  # type: ignore[override]
+        super().update(**kwargs)
+        self._prune_dict(self.max_size)
+
+    def _prune_dict(self, max_size: int) -> None:
+        if len(self) >= max_size:
+            diff = len(self) - max_size
+            for k in list(self.keys())[:diff]:
+                del self[k]

wandb/sdk/lib/config_util.py

@@ -0,0 +1,101 @@
+import json
+import logging
+import os
+from typing import Any, Dict, Optional
+
+import yaml
+
+import wandb
+from wandb.errors import Error
+from wandb.util import load_yaml
+
+from . import filesystem
+
+logger = logging.getLogger("wandb")
+
+
+class ConfigError(Error):
+    pass
+
+
+def dict_from_proto_list(obj_list):
+    d = dict()
+    for item in obj_list:
+        d[item.key] = dict(desc=None, value=json.loads(item.value_json))
+    return d
+
+
+def dict_strip_value_dict(config_dict):
+    d = dict()
+    for k, v in config_dict.items():
+        d[k] = v["value"]
+    return d
+
+
+def dict_no_value_from_proto_list(obj_list):
+    d = dict()
+    for item in obj_list:
+        possible_dict = json.loads(item.value_json)
+        if not isinstance(possible_dict, dict) or "value" not in possible_dict:
+            continue
+        d[item.key] = possible_dict["value"]
+
+    return d
+
+
+# TODO(jhr): these functions should go away once we merge jobspec PR
+def save_config_file_from_dict(config_filename, config_dict):
+    s = b"wandb_version: 1"
+    if config_dict:  # adding an empty dictionary here causes a parse error
+        s += b"\n\n" + yaml.dump(
+            config_dict,
+            Dumper=yaml.SafeDumper,
+            default_flow_style=False,
+            allow_unicode=True,
+            encoding="utf-8",
+            sort_keys=False,
+        )
+    data = s.decode("utf-8")
+    filesystem.mkdir_exists_ok(os.path.dirname(config_filename))
+    with open(config_filename, "w") as conf_file:
+        conf_file.write(data)
+
+
+def dict_from_config_file(
+    filename: str, must_exist: bool = False
+) -> Optional[Dict[str, Any]]:
+    if not os.path.exists(filename):
+        if must_exist:
+            raise ConfigError("config file {} doesn't exist".format(filename))
+        logger.debug("no default config file found in {}".format(filename))
+        return None
+    try:
+        conf_file = open(filename)
+    except OSError:
+        raise ConfigError("Couldn't read config file: {}".format(filename))
+    try:
+        loaded = load_yaml(conf_file)
+    except yaml.parser.ParserError:
+        raise ConfigError("Invalid YAML in config yaml")
+    if loaded is None:
+        wandb.termwarn(
+            "Found an empty default config file (config-defaults.yaml). Proceeding with no defaults."
+        )
+        return None
+    config_version = loaded.pop("wandb_version", None)
+    if config_version is not None and config_version != 1:
+        raise ConfigError("Unknown config version")
+    data = dict()
+    for k, v in loaded.items():
+        data[k] = v["value"]
+    return data
+
+
+def merge_dicts(dest: dict, src: dict) -> dict:
+    """Recursively merge two dictionaries. Similar to Lodash's _.merge()."""
+    for key, value in src.items():
+        if isinstance(value, dict) and key in dest and isinstance(dest[key], dict):
+            merge_dicts(dest[key], value)
+        else:
+            dest[key] = value
+    return dest

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