wandb 0.21.2__py3-none-macosx_12_0_arm64.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,334 @@
+"""apikey util."""
+
+from __future__ import annotations
+
+import dataclasses
+import os
+import platform
+import stat
+import sys
+import textwrap
+from functools import partial
+
+# import Literal
+from typing import TYPE_CHECKING, Callable, Literal
+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.sdk import wandb_setup
+from wandb.util import _is_databricks, isatty, prompt_choices
+
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
+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,
+]
+
+
+@dataclasses.dataclass(frozen=True)
+class _NetrcPermissions:
+    exists: bool
+    read_access: bool
+    write_access: bool
+
+
+class WriteNetrcError(Exception):
+    """Raised when we cannot write to the netrc file."""
+
+
+Mode = Literal["allow", "must", "never", "false", "true"]
+
+
+getpass = partial(click.prompt, hide_input=True, err=True)
+
+
+def _fixup_anon_mode(default: Mode | None) -> Mode | None:
+    # 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 _api_key_prompt_str(app_url: str, referrer: str | None = None) -> str:
+    """Generate a prompt string for API key authorization.
+
+    Creates a URL string that directs users to the authorization page where they
+    can find their API key.
+
+    Args:
+        app_url: The base URL of the W&B application.
+        referrer: Optional referrer parameter to include in the URL.
+
+    Returns:
+        A formatted string with instructions and the authorization URL.
+    """
+    ref = ""
+    if referrer:
+        ref = f"?ref={referrer}"
+    return f"You can find your API key in your browser here: {app_url}/authorize{ref}"
+
+
+def prompt_api_key(  # noqa: C901
+    settings: Settings,
+    api: InternalApi | None = None,
+    input_callback: Callable | None = None,
+    browser_callback: Callable | None = None,
+    no_offline: bool = False,
+    no_create: bool = False,
+    local: bool = False,
+    referrer: str | None = None,
+) -> 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:
+            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
+        )
+
+    key = None
+    api_ask = (
+        f"{log_string}: Paste an API key from your profile and hit enter"
+        if jupyter
+        else 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()
+    elif result == LOGIN_CHOICE_NEW:
+        key = browser_callback(signup=True) if browser_callback else None
+
+        if not key:
+            ref = f"&ref={referrer}" if referrer else ""
+            wandb.termlog(
+                f"Create an account here: {app_url}/authorize?signup=true{ref}"
+            )
+            key = input_callback(api_ask).strip()
+    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(_api_key_prompt_str(app_url, referrer))
+            key = input_callback(api_ask).strip()
+    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)
+        )
+
+    if not key:
+        raise ValueError("No API key specified.")
+    return key
+
+
+def check_netrc_access(
+    netrc_path: str,
+) -> _NetrcPermissions:
+    """Check if we can read and write to the netrc file."""
+    file_exists = False
+    write_access = False
+    read_access = False
+    try:
+        st = os.stat(netrc_path)
+        file_exists = True
+        write_access = bool(st.st_mode & stat.S_IWUSR)
+        read_access = bool(st.st_mode & stat.S_IRUSR)
+    except FileNotFoundError:
+        # If the netrc file doesn't exist, we will create it.
+        write_access = True
+        read_access = True
+    except OSError as e:
+        wandb.termerror(f"Unable to read permissions for {netrc_path}, {e}")
+
+    return _NetrcPermissions(
+        exists=file_exists,
+        write_access=write_access,
+        read_access=read_access,
+    )
+
+
+def write_netrc(host: str, entity: str, key: str):
+    """Add our host and key to .netrc."""
+    _, key_suffix = key.split("-", 1) if "-" in key else ("", key)
+    if len(key_suffix) != 40:
+        raise ValueError(
+            f"API-key must be exactly 40 characters long: {key_suffix} ({len(key_suffix)} chars)"
+        )
+
+    normalized_host = urlparse(host).netloc
+    netrc_path = get_netrc_file_path()
+    netrc_access = check_netrc_access(netrc_path)
+
+    if not netrc_access.write_access or not netrc_access.read_access:
+        raise WriteNetrcError(
+            f"Cannot access {netrc_path}. In order to persist your API key, "
+            "grant read and write permissions for your user to the file "
+            'or specify a different file with the environment variable "NETRC=<new_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 FileNotFoundError:
+        wandb.termlog("No netrc file found, creating one.")
+    except OSError as e:
+        raise WriteNetrcError(f"Unable to read {netrc_path}") from e
+
+    try:
+        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(f"{line}\n")
+
+            wandb.termlog(
+                f"Appending key for {normalized_host} to your netrc file: {netrc_path}"
+            )
+            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)
+    except OSError as e:
+        raise WriteNetrcError(f"Unable to write {netrc_path}") from e
+
+
+def write_key(
+    settings: Settings,
+    key: str | None,
+    api: InternalApi | None = None,
+) -> 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(f"API key must be 40 characters long, yours was {len(key)}")
+
+    write_netrc(settings.base_url, "user", key)
+
+
+def api_key(settings: Settings | None = None) -> str | None:
+    if settings is None:
+        settings = wandb_setup.singleton().settings
+    if settings.api_key:
+        return settings.api_key
+
+    netrc_access = check_netrc_access(get_netrc_file_path())
+    if netrc_access.exists and not netrc_access.read_access:
+        wandb.termwarn(f"Cannot access {get_netrc_file_path()}.")
+        return None
+
+    if netrc_access.exists:
+        auth = get_netrc_auth(settings.base_url)
+        if auth:
+            return auth[-1]
+
+    return None

wandb/sdk/lib/asyncio_compat.py

@@ -0,0 +1,213 @@
+"""Functions for compatibility with asyncio."""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent
+import concurrent.futures
+import contextlib
+import threading
+from typing import Any, AsyncIterator, Callable, Coroutine, Iterator, TypeVar
+
+_T = TypeVar("_T")
+
+
+def run(fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
+    """Run `fn` in an asyncio loop in a new thread.
+
+    This must always be used instead of `asyncio.run` which fails if there is
+    an active `asyncio` event loop in the current thread. Since `wandb` was not
+    originally designed with `asyncio` in mind, using `asyncio.run` would break
+    users who were calling `wandb` methods from an `asyncio` loop.
+
+    Note that due to starting a new thread, this is slightly slow.
+    """
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        runner = CancellableRunner()
+        future = executor.submit(runner.run, fn)
+
+        try:
+            return future.result()
+
+        finally:
+            runner.cancel()
+
+
+class RunnerCancelledError(Exception):
+    """The `CancellableRunner.run()` invocation was cancelled."""
+
+
+class CancellableRunner:
+    """Runs an asyncio event loop allowing cancellation.
+
+    The `run()` method is like `asyncio.run()`. The `cancel()` method may
+    be used in a different thread, for instance in a `finally` block, to cancel
+    all tasks, and it is a no-op if `run()` completed.
+
+    Without this, it is impossible to make `asyncio.run()` stop if it runs
+    in a non-main thread. In particular, a KeyboardInterrupt causes the
+    ThreadPoolExecutor above to block until the asyncio thread completes,
+    but there is no way to tell the asyncio thread to cancel its work.
+    A second KeyboardInterrupt makes ThreadPoolExecutor give up while the
+    asyncio thread still runs in the background, with terrible effects if it
+    prints to the user's terminal.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+
+        self._is_cancelled = False
+        self._started = False
+        self._done = False
+
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._cancel_event: asyncio.Event | None = None
+
+    def run(self, fn: Callable[[], Coroutine[Any, Any, _T]]) -> _T:
+        """Run a coroutine in asyncio, cancelling it on `cancel()`.
+
+        Returns:
+            The result of the coroutine returned by `fn`.
+
+        Raises:
+            RunnerCancelledError: If `cancel()` is called.
+        """
+        return asyncio.run(self._run_or_cancel(fn))
+
+    async def _run_or_cancel(
+        self,
+        fn: Callable[[], Coroutine[Any, Any, _T]],
+    ) -> _T:
+        with self._lock:
+            if self._is_cancelled:
+                raise RunnerCancelledError()
+
+            self._loop = asyncio.get_running_loop()
+            self._cancel_event = asyncio.Event()
+            self._started = True
+
+        cancellation_task = asyncio.create_task(self._cancel_event.wait())
+        fn_task = asyncio.create_task(fn())
+
+        try:
+            await asyncio.wait(
+                [cancellation_task, fn_task],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if fn_task.done():
+                return fn_task.result()
+            else:
+                raise RunnerCancelledError()
+
+        finally:
+            # NOTE: asyncio.run() cancels all tasks after the main task exits,
+            #   but this is not documented, so we cancel them explicitly here
+            #   as well. It also blocks until canceled tasks complete.
+            cancellation_task.cancel()
+            fn_task.cancel()
+
+            with self._lock:
+                self._done = True
+
+    def cancel(self) -> None:
+        """Cancel all asyncio work started by `run()`."""
+        with self._lock:
+            if self._is_cancelled:
+                return
+            self._is_cancelled = True
+
+            if self._done or not self._started:
+                # If the runner already finished, no need to cancel it.
+                #
+                # If the runner hasn't started the loop yet, then it will not
+                # as we already set _is_cancelled.
+                return
+
+            assert self._loop
+            assert self._cancel_event
+            self._loop.call_soon_threadsafe(self._cancel_event.set)
+
+
+class TaskGroup:
+    """Object that `open_task_group()` yields."""
+
+    def __init__(self) -> None:
+        self._tasks: list[asyncio.Task] = []
+
+    def start_soon(self, coro: Coroutine[Any, Any, Any]) -> None:
+        """Schedule a task in the group.
+
+        Args:
+            coro: The return value of the `async` function defining the task.
+        """
+        self._tasks.append(asyncio.create_task(coro))
+
+    async def _wait_all(self) -> None:
+        """Block until all tasks complete.
+
+        Raises:
+            Exception: If one or more tasks raises an exception, one of these
+                is raised arbitrarily.
+        """
+        done, _ = await asyncio.wait(
+            self._tasks,
+            # NOTE: Cancelling a task counts as a normal exit,
+            #   not an exception.
+            return_when=concurrent.futures.FIRST_EXCEPTION,
+        )
+
+        for task in done:
+            with contextlib.suppress(asyncio.CancelledError):
+                if exc := task.exception():
+                    raise exc
+
+    def _cancel_all(self) -> None:
+        """Cancel all tasks."""
+        for task in self._tasks:
+            # NOTE: It is safe to cancel tasks that have already completed.
+            task.cancel()
+
+
+@contextlib.asynccontextmanager
+async def open_task_group() -> AsyncIterator[TaskGroup]:
+    """Create a task group.
+
+    `asyncio` gained task groups in Python 3.11.
+
+    This is an async context manager, meant to be used with `async with`.
+    On exit, it blocks until all subtasks complete. If any subtask fails, or if
+    the current task is cancelled, it cancels all subtasks in the group and
+    raises the subtask's exception. If multiple subtasks fail simultaneously,
+    one of their exceptions is chosen arbitrarily.
+
+    NOTE: Subtask exceptions do not propagate until the context manager exits.
+    This means that the task group cannot cancel code running inside the
+    `async with` block .
+    """
+    task_group = TaskGroup()
+
+    try:
+        yield task_group
+        await task_group._wait_all()
+    finally:
+        task_group._cancel_all()
+
+
+@contextlib.contextmanager
+def cancel_on_exit(coro: Coroutine[Any, Any, Any]) -> Iterator[None]:
+    """Schedule a task, cancelling it when exiting the context manager.
+
+    If the context manager exits successfully but the given coroutine raises
+    an exception, that exception is reraised. The exception is suppressed
+    if the context manager raises an exception.
+    """
+    task = asyncio.create_task(coro)
+
+    try:
+        yield
+
+        if task.done() and (exception := task.exception()):
+            raise exception
+    finally:
+        task.cancel()