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

wandb/sdk/interface/interface_sock.py

@@ -7,10 +7,10 @@ See interface.py for how interface classes relate to each other.
 import logging
 from typing import TYPE_CHECKING, Any, Optional
 
-from ..lib.mailbox import Mailbox
+from wandb.sdk.mailbox import Mailbox
+
 from ..lib.sock_client import SockClient
 from .interface_shared import InterfaceShared
-from .message_future import MessageFuture
 from .router_sock import MessageSockRouter
 
 if TYPE_CHECKING:
@@ -32,7 +32,6 @@ class InterfaceSock(InterfaceShared):
         # _sock_client is used when abstract method _init_router() is called by constructor
         self._sock_client = sock_client
         super().__init__(mailbox=mailbox)
-        self._process_check = False
         self._stream_id = stream_id
 
     def _init_router(self) -> None:
@@ -45,13 +44,3 @@ class InterfaceSock(InterfaceShared):
     def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
         self._assign(record)
         self._sock_client.send_record_publish(record)
-
-    def _communicate_async(
-        self, rec: "pb.Record", local: Optional[bool] = None
-    ) -> MessageFuture:
-        self._assign(rec)
-        assert self._router
-        if self._process_check and self._process and not self._process.is_alive():
-            raise Exception("The wandb backend process has shutdown")
-        future = self._router.send_and_receive(rec, local=local)
-        return future

wandb/sdk/interface/router.py

@@ -10,7 +10,8 @@ import uuid
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Dict, Optional
 
-from ..lib import mailbox
+from wandb.sdk import mailbox
+
 from .message_future import MessageFuture
 
 if TYPE_CHECKING:
@@ -63,20 +64,25 @@ class MessageRouter:
         raise NotImplementedError
 
     def message_loop(self) -> None:
-        while not self._join_event.is_set():
-            try:
-                msg = self._read_message()
-            except EOFError:
-                # On abnormal shutdown the queue will be destroyed underneath
-                # resulting in EOFError.  message_loop needs to exit..
-                logger.warning("EOFError seen in message_loop")
-                break
-            except MessageRouterClosedError:
-                logger.warning("message_loop has been closed")
-                break
-            if not msg:
-                continue
-            self._handle_msg_rcv(msg)
+        try:
+            while not self._join_event.is_set():
+                try:
+                    msg = self._read_message()
+                except EOFError:
+                    # On abnormal shutdown the queue will be destroyed underneath
+                    # resulting in EOFError.  message_loop needs to exit..
+                    logger.warning("EOFError seen in message_loop")
+                    break
+                except MessageRouterClosedError as e:
+                    logger.warning("message_loop has been closed", exc_info=e)
+                    break
+                if not msg:
+                    continue
+                self._handle_msg_rcv(msg)
+
+        finally:
+            if self._mailbox:
+                self._mailbox.close()
 
     def send_and_receive(
         self, rec: "pb.Record", local: Optional[bool] = None

wandb/sdk/interface/router_queue.py

@@ -7,7 +7,8 @@ Router to manage responses from a queue.
 import queue
 from typing import TYPE_CHECKING, Optional
 
-from ..lib.mailbox import Mailbox
+from wandb.sdk.mailbox import Mailbox
+
 from .router import MessageRouter
 
 if TYPE_CHECKING:

wandb/sdk/interface/router_relay.py

@@ -6,7 +6,8 @@ Router to manage responses from a queue with relay.
 
 from typing import TYPE_CHECKING
 
-from ..lib.mailbox import Mailbox
+from wandb.sdk.mailbox import Mailbox
+
 from .router_queue import MessageQueueRouter
 
 if TYPE_CHECKING:

wandb/sdk/interface/router_sock.py

@@ -6,8 +6,9 @@ Router to manage responses from a socket client.
 
 from typing import TYPE_CHECKING, Optional
 
-from ..lib.mailbox import Mailbox
-from ..lib.sock_client import SockClient, SockClientClosedError
+from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError
+from wandb.sdk.mailbox import Mailbox
+
 from .router import MessageRouter, MessageRouterClosedError
 
 if TYPE_CHECKING:
@@ -25,8 +26,8 @@ class MessageSockRouter(MessageRouter):
     def _read_message(self) -> Optional["pb.Result"]:
         try:
             resp = self._sock_client.read_server_response(timeout=1)
-        except SockClientClosedError:
-            raise MessageRouterClosedError
+        except SockClientClosedError as e:
+            raise MessageRouterClosedError from e
         if not resp:
             return None
         msg = resp.result_communicate

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

@@ -544,23 +544,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
@@ -1491,7 +1474,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()