wandb 0.19.6rc4__py3-none-any.whl → 0.19.8__py3-none-any.whl

wandb/sdk/interface/router_sock.py

@@ -4,14 +4,14 @@ Router to manage responses from a socket client.
 
 """
 
-from typing import TYPE_CHECKING, Optional
+from __future__ import annotations
 
-from ..lib.mailbox import Mailbox
-from ..lib.sock_client import SockClient, SockClientClosedError
-from .router import MessageRouter, MessageRouterClosedError
+from wandb.proto import wandb_internal_pb2 as pb
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError
+from wandb.sdk.mailbox import Mailbox
 
-if TYPE_CHECKING:
-    from wandb.proto import wandb_internal_pb2 as pb
+from .router import MessageRouter, MessageRouterClosedError
 
 
 class MessageSockRouter(MessageRouter):
@@ -22,15 +22,11 @@ class MessageSockRouter(MessageRouter):
         self._sock_client = sock_client
         super().__init__(mailbox=mailbox)
 
-    def _read_message(self) -> Optional["pb.Result"]:
+    def _read_message(self) -> spb.ServerResponse | None:
         try:
-            resp = self._sock_client.read_server_response(timeout=1)
-        except SockClientClosedError:
-            raise MessageRouterClosedError
-        if not resp:
-            return None
-        msg = resp.result_communicate
-        return msg
-
-    def _send_message(self, record: "pb.Record") -> None:
+            return self._sock_client.read_server_response(timeout=1)
+        except SockClientClosedError as e:
+            raise MessageRouterClosedError from e
+
+    def _send_message(self, record: pb.Record) -> None:
         self._sock_client.send_record_communicate(record)

wandb/sdk/internal/handler.py

@@ -205,9 +205,6 @@ class HandleManager:
         # defer is used to drive the sender finish state machine
         self._dispatch_record(record, always_send=True)
 
-    def handle_request_login(self, record: Record) -> None:
-        self._dispatch_record(record)
-
     def handle_request_python_packages(self, record: Record) -> None:
         self._dispatch_record(record)
 
@@ -892,6 +889,10 @@ class HandleManager:
         self._respond_result(result)
         self._stopped.set()
 
+    def handle_request_operations(self, record: Record) -> None:
+        """No-op. Not implemented for the legacy-service."""
+        self._respond_result(proto_util._result_from_record(record))
+
     def finish(self) -> None:
         logger.info("shutting down handler")
         if self._system_monitor is not None:

wandb/sdk/internal/internal_api.py

@@ -115,6 +115,7 @@ if TYPE_CHECKING:
         root_dir: Optional[str]
         api_key: Optional[str]
         entity: Optional[str]
+        organization: Optional[str]
         project: Optional[str]
         _extra_http_headers: Optional[Mapping[str, str]]
         _proxies: Optional[Mapping[str, str]]
@@ -256,6 +257,7 @@ class Api:
             "root_dir": None,
             "api_key": None,
             "entity": None,
+            "organization": None,
             "project": None,
             "_extra_http_headers": None,
             "_proxies": None,
@@ -489,7 +491,8 @@ class Api:
                 {
                     "entity": "models",
                     "base_url": "https://api.wandb.ai",
-                    "project": None
+                    "project": None,
+                    "organization": "my-org",
                 }
         """
         result = self.default_settings.copy()
@@ -504,6 +507,14 @@ class Api:
                     ),
                     env=self._environ,
                 ),
+                "organization": env.get_organization(
+                    self._settings.get(
+                        Settings.DEFAULT_SECTION,
+                        "organization",
+                        fallback=result.get("organization"),
+                    ),
+                    env=self._environ,
+                ),
                 "project": env.get_project(
                     self._settings.get(
                         Settings.DEFAULT_SECTION,

wandb/sdk/internal/sender.py

@@ -1,6 +1,7 @@
 """sender."""
 
 import contextlib
+import glob
 import gzip
 import json
 import logging
@@ -544,23 +545,6 @@ class SendManager:
                 logger.warning(f"Error emptying retry queue: {e}")
         self._respond_result(result)
 
-    def send_request_login(self, record: "Record") -> None:
-        # TODO: do something with api_key or anonymous?
-        # TODO: return an error if we aren't logged in?
-        self._api.reauth()
-        viewer = self.get_viewer_info()
-        server_info = self.get_server_info()
-        # self._login_flags = json.loads(viewer.get("flags", "{}"))
-        # self._login_entity = viewer.get("entity")
-        if server_info:
-            logger.info(f"Login server info: {server_info}")
-        self._entity = viewer.get("entity")
-        if record.control.req_resp:
-            result = proto_util._result_from_record(record)
-            if self._entity:
-                result.response.login_response.active_entity = self._entity
-            self._respond_result(result)
-
     def send_exit(self, record: "Record") -> None:
         # track where the exit came from
         self._record_exit = record
@@ -1425,7 +1409,8 @@ class SendManager:
         for k in files.files:
             # TODO(jhr): fix paths with directories
             self._save_file(
-                interface.GlobStr(k.path), interface.file_enum_to_policy(k.policy)
+                interface.GlobStr(glob.escape(k.path)),
+                interface.file_enum_to_policy(k.policy),
             )
 
     def send_header(self, record: "Record") -> None:
@@ -1491,7 +1476,6 @@ class SendManager:
             self._job_builder.set_partial_source_id(use.id)
 
     def send_request_log_artifact(self, record: "Record") -> None:
-        assert record.control.req_resp
         result = proto_util._result_from_record(record)
         artifact = record.request.log_artifact.artifact
         history_step = record.request.log_artifact.history_step

wandb/sdk/lib/apikey.py

@@ -1,5 +1,8 @@
 """apikey util."""
 
+from __future__ import annotations
+
+import dataclasses
 import os
 import platform
 import stat
@@ -32,8 +35,21 @@ LOGIN_CHOICES = [
     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"]
 
+
 if TYPE_CHECKING:
     from wandb.sdk.wandb_settings import Settings
 
@@ -170,29 +186,64 @@ def prompt_api_key(  # noqa: C901
     return key
 
 
-def write_netrc(host: str, entity: str, key: str) -> Optional[bool]:
+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:
-        wandb.termerror(
+        raise ValueError(
             "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}"
+
+    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 OSError:
-            pass
+
+    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.
@@ -206,20 +257,22 @@ def write_netrc(host: str, entity: str, key: str) -> Optional[bool]:
                         skip -= 1
                     else:
                         f.write("{}\n".format(line))
+
+            wandb.termlog(
+                f"Appending key for {normalized_host} to your netrc file: {netrc_path}"
+            )
             f.write(
                 textwrap.dedent(
                     """\
-            machine {host}
-              login {entity}
-              password {key}
-            """
+                    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
+    except OSError as e:
+        raise WriteNetrcError(f"Unable to write {netrc_path}") from e
 
 
 def write_key(
@@ -250,7 +303,15 @@ def api_key(settings: Optional["Settings"] = None) -> Optional[str]:
         settings = wandb.setup().settings
     if settings.api_key:
         return settings.api_key
-    auth = get_netrc_auth(settings.base_url)
-    if auth:
-        return auth[-1]
+
+    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,210 @@
+"""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 = _Runner()
+        future = executor.submit(runner.run, fn)
+
+        try:
+            return future.result()
+
+        finally:
+            runner.cancel()
+
+
+class _RunnerCancelledError(Exception):
+    """The `_Runner.run()` invocation was cancelled."""
+
+
+class _Runner:
+    """Runs an asyncio event loop allowing cancellation.
+
+    This is like `asyncio.run()`, except it provides a `cancel()` method
+    meant to be called in a `finally` block.
+
+    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.Condition()
+
+        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:
+            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:
+            try:
+                if exc := task.exception():
+                    raise exc
+            except asyncio.CancelledError:
+                pass
+
+    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 given coroutine raises an exception, that exception is raised
+    when exiting the context manager.
+    """
+    task = asyncio.create_task(coro)
+
+    try:
+        yield
+    finally:
+        if task.done() and (exception := task.exception()):
+            raise exception
+
+        task.cancel()

wandb/sdk/lib/console_capture.py

@@ -0,0 +1,172 @@
+"""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 sys
+import threading
+from typing import IO, AnyStr, Callable, Protocol
+
+
+class CannotCaptureConsoleError(Exception):
+    """The module failed to patch stdout or stderr."""
+
+
+class _WriteCallback(Protocol):
+    """A callback that receives intercepted bytes or string data."""
+
+    def __call__(
+        self,
+        data: bytes | str,
+        written: int,
+        /,
+    ) -> None:
+        """Intercept data passed to `write()`.
+
+        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()`.
+        """
+
+
+_module_lock = threading.Lock()
+
+_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`.
+            This may be called from any thread, so it must be thread-safe.
+            Exceptions are propagated to the caller of `write`.
+            See `_WriteCallback` for the exact protocol.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_lock:
+        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`.
+            This may be called from any thread, so it must be thread-safe.
+            Exceptions are propagated to the caller of `write`.
+            See `_WriteCallback` for the exact protocol.
+
+    Returns:
+        A function to uninstall the callback.
+
+    Raises:
+        CannotCaptureConsoleError: If patching failed on import.
+    """
+    with _module_lock:
+        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_lock:
+            if disposed:
+                return
+
+            del callback_dict[id]
+
+            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]
+
+    def write_with_callbacks(s: AnyStr, /) -> int:
+        n = orig_write(s)
+
+        # We make a copy here because callbacks could, in theory, modify
+        # the list of callbacks.
+        with _module_lock:
+            callbacks_copy = list(callbacks.values())
+
+        for cb in callbacks_copy:
+            cb(s, n)
+
+        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