wandb 0.22.1__py3-none-win_arm64.whl

wandb/sdk/mailbox/mailbox_handle.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING, Callable, Generic
+
+from typing_extensions import TypeVar, override
+
+from wandb.sdk.lib import asyncio_manager
+
+# Necessary to break an import loop.
+if TYPE_CHECKING:
+    from wandb.sdk.interface import interface
+
+
+_T = TypeVar("_T")
+_S = TypeVar("_S")
+
+
+class HandleAbandonedError(Exception):
+    """The handle has no response and has been abandoned."""
+
+
+class MailboxHandle(abc.ABC, Generic[_T]):
+    """A handle for waiting on a response to a request."""
+
+    def __init__(self, asyncer: asyncio_manager.AsyncioManager) -> None:
+        self._asyncer = asyncer
+
+    @property
+    def asyncer(self) -> asyncio_manager.AsyncioManager:
+        """The asyncio thread to which the handle belongs.
+
+        The handle's async methods must be run using this object.
+        """
+        return self._asyncer
+
+    def map(self, fn: Callable[[_T], _S]) -> MailboxHandle[_S]:
+        """Returns a transformed handle.
+
+        Methods on the returned handle call methods on this handle, but the
+        response type is derived using the given function.
+
+        Args:
+            fn: A function to apply to this handle's result to get the new
+                handle's result. The function should be pure and fast.
+        """
+        return _MailboxMappedHandle(self, fn)
+
+    @abc.abstractmethod
+    def abandon(self) -> None:
+        """Abandon the handle, indicating it will not receive a response.
+
+        This may not happen immediately: it is possible for an existing
+        call to `wait_async` to complete successfully after this method returns.
+        """
+
+    @abc.abstractmethod
+    def cancel(self, iface: interface.InterfaceBase) -> None:
+        """Cancel the handle, requesting any associated work to not complete.
+
+        It is an error to call this from an async function.
+
+        This automatically abandons the handle, as a response is no longer
+        guaranteed.
+
+        Args:
+            iface: The interface on which to publish the cancel request.
+        """
+
+    @abc.abstractmethod
+    def wait_or(self, *, timeout: float | None) -> _T:
+        """Wait for a response or a timeout.
+
+        It is an error to call this from an async function.
+
+        Args:
+            timeout: A finite number of seconds or None to never time out.
+                If less than or equal to zero, times out immediately unless
+                the response is available.
+
+        Returns:
+            The response if it arrives before the timeout or has already arrived.
+
+        Raises:
+            TimeoutError: If the timeout is reached.
+            HandleAbandonedError: If the handle becomes abandoned.
+        """
+
+    @abc.abstractmethod
+    async def wait_async(self, *, timeout: float | None) -> _T:
+        """Wait for a response or timeout.
+
+        This must run in an `asyncio` event loop.
+
+        Args:
+            timeout: A finite number of seconds or None to never time out.
+
+        Returns:
+            The response if it arrives before the timeout or has already arrived.
+
+        Raises:
+            TimeoutError: If the timeout is reached.
+            HandleAbandonedError: If the handle becomes abandoned.
+        """
+
+
+class _MailboxMappedHandle(Generic[_S], MailboxHandle[_S]):
+    """A mailbox handle whose result is derived from another handle."""
+
+    def __init__(
+        self,
+        handle: MailboxHandle[_T],
+        fn: Callable[[_T], _S],
+    ) -> None:
+        super().__init__(handle.asyncer)
+        self._handle = handle
+        self._fn = fn
+
+    @override
+    def abandon(self) -> None:
+        self._handle.abandon()
+
+    @override
+    def cancel(self, iface: interface.InterfaceBase) -> None:
+        self._handle.cancel(iface)
+
+    @override
+    def wait_or(self, *, timeout: float | None) -> _S:
+        return self._fn(self._handle.wait_or(timeout=timeout))
+
+    @override
+    async def wait_async(self, *, timeout: float | None) -> _S:
+        response = await self._handle.wait_async(timeout=timeout)
+        return self._fn(response)

wandb/sdk/mailbox/response_handle.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import asyncio
+import math
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
+
+from .mailbox_handle import HandleAbandonedError, MailboxHandle
+
+# Necessary to break an import loop.
+if TYPE_CHECKING:
+    from wandb.sdk.interface import interface
+
+
+class MailboxResponseHandle(MailboxHandle[spb.ServerResponse]):
+    """A general handle for any ServerResponse."""
+
+    def __init__(
+        self,
+        address: str,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> None:
+        super().__init__(asyncer)
+
+        self._address = address
+
+        self._abandoned = False
+        self._response: spb.ServerResponse | None = None
+
+        # Initialized on first use in the asyncio thread.
+        self._done_event: asyncio.Event | None = None
+
+    async def deliver(self, response: spb.ServerResponse) -> None:
+        if self._abandoned:
+            return
+
+        if self._response:
+            raise ValueError(
+                f"A response has already been delivered to {self._address}."
+            )
+
+        self._response = response
+
+        if not self._done_event:
+            self._done_event = asyncio.Event()
+        self._done_event.set()
+
+    @override
+    def cancel(self, iface: interface.InterfaceBase) -> None:
+        iface.publish_cancel(self._address)
+        self.abandon()
+
+    @override
+    def abandon(self) -> None:
+        async def impl() -> None:
+            self._abandoned = True
+
+            if not self._done_event:
+                self._done_event = asyncio.Event()
+            self._done_event.set()
+
+        self.asyncer.run_soon(impl)
+
+    @override
+    def wait_or(self, *, timeout: float | None) -> spb.ServerResponse:
+        return self.asyncer.run(lambda: self.wait_async(timeout=timeout))
+
+    @override
+    async def wait_async(self, *, timeout: float | None) -> spb.ServerResponse:
+        if timeout is not None and not math.isfinite(timeout):
+            raise ValueError("Timeout must be finite or None.")
+
+        if not self._done_event:
+            self._done_event = asyncio.Event()
+
+        try:
+            await asyncio.wait_for(self._done_event.wait(), timeout=timeout)
+
+        except (asyncio.TimeoutError, TimeoutError) as e:
+            if self._response:
+                return self._response
+            elif self._abandoned:
+                raise HandleAbandonedError()
+            else:
+                raise TimeoutError(
+                    f"Timed out waiting for response on {self._address}"
+                ) from e
+
+        else:
+            if self._response:
+                return self._response
+
+            assert self._abandoned
+            raise HandleAbandonedError()

wandb/sdk/mailbox/wait_with_progress.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Callable, Coroutine, List, TypeVar, cast
+
+from wandb.sdk.lib import asyncio_compat
+
+from .mailbox_handle import MailboxHandle
+
+_T = TypeVar("_T")
+
+
+def wait_with_progress(
+    handle: MailboxHandle[_T],
+    *,
+    timeout: float | None,
+    display_progress: Callable[[], Coroutine[Any, Any, None]],
+) -> _T:
+    """Wait for a handle, possibly displaying progress to the user.
+
+    Equivalent to passing a single handle to `wait_all_with_progress`.
+    """
+    return wait_all_with_progress(
+        [handle],
+        timeout=timeout,
+        display_progress=display_progress,
+    )[0]
+
+
+def wait_all_with_progress(
+    handle_list: list[MailboxHandle[_T]],
+    *,
+    timeout: float | None,
+    display_progress: Callable[[], Coroutine[Any, Any, None]],
+) -> list[_T]:
+    """Wait for multiple handles, possibly displaying progress to the user.
+
+    Args:
+        handle_list: The handles to wait for.
+        timeout: A number of seconds after which to raise a TimeoutError,
+            or None if this should never timeout.
+        display_progress: An asyncio function that displays progress to
+            the user. This function runs using the handles' AsyncioManager.
+
+    Returns:
+        A list where the Nth item is the Nth handle's result.
+
+    Raises:
+        ValueError: If the handles live in different asyncio threads.
+        TimeoutError: If the overall timeout expires.
+        HandleAbandonedError: If any handle becomes abandoned.
+        Exception: Any exception from the display function is propagated.
+    """
+    if not handle_list:
+        return []
+
+    asyncer = handle_list[0].asyncer
+    for handle in handle_list:
+        if handle.asyncer is not asyncer:
+            raise ValueError("Handles have different AsyncioManagers.")
+
+    start_time = time.monotonic()
+
+    async def progress_loop_with_timeout() -> list[_T]:
+        async with asyncio_compat.cancel_on_exit(display_progress()):
+            if timeout is not None:
+                elapsed_time = time.monotonic() - start_time
+                remaining_timeout = timeout - elapsed_time
+            else:
+                remaining_timeout = None
+
+            return await _wait_handles_async(
+                handle_list,
+                timeout=remaining_timeout,
+            )
+
+    return asyncer.run(progress_loop_with_timeout)
+
+
+async def _wait_handles_async(
+    handle_list: list[MailboxHandle[_T]],
+    *,
+    timeout: float | None,
+) -> list[_T]:
+    """Asynchronously wait for multiple mailbox handles.
+
+    Just like _wait_handles.
+    """
+    results: list[_T | None] = [None for _ in handle_list]
+
+    async def wait_single(index: int) -> None:
+        handle = handle_list[index]
+        results[index] = await handle.wait_async(timeout=timeout)
+
+    async with asyncio_compat.open_task_group() as task_group:
+        for index in range(len(handle_list)):
+            task_group.start_soon(wait_single(index))
+
+    # NOTE: `list` is not subscriptable until Python 3.10, so we use List.
+    return cast(List[_T], results)

wandb/sdk/projects/_generated/__init__.py

@@ -0,0 +1,26 @@
+# Generated by ariadne-codegen
+
+__all__ = [
+    "DELETE_PROJECT_GQL",
+    "FETCH_REGISTRY_GQL",
+    "RENAME_PROJECT_GQL",
+    "UPSERT_REGISTRY_PROJECT_GQL",
+    "FetchRegistry",
+    "RenameProject",
+    "UpsertRegistryProject",
+    "DeleteProject",
+    "ArtifactTypeInput",
+    "RegistryFragment",
+]
+from .delete_project import DeleteProject
+from .fetch_registry import FetchRegistry
+from .fragments import RegistryFragment
+from .input_types import ArtifactTypeInput
+from .operations import (
+    DELETE_PROJECT_GQL,
+    FETCH_REGISTRY_GQL,
+    RENAME_PROJECT_GQL,
+    UPSERT_REGISTRY_PROJECT_GQL,
+)
+from .rename_project import RenameProject
+from .upsert_registry_project import UpsertRegistryProject

wandb/sdk/projects/_generated/delete_project.py

@@ -0,0 +1,22 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import Literal, Optional
+
+from pydantic import Field
+
+from wandb._pydantic import GQLBase, Typename
+
+
+class DeleteProject(GQLBase):
+    delete_model: Optional[DeleteProjectDeleteModel] = Field(alias="deleteModel")
+
+
+class DeleteProjectDeleteModel(GQLBase):
+    success: Optional[bool]
+    typename__: Typename[Literal["DeleteModelPayload"]]
+
+
+DeleteProject.model_rebuild()

wandb/sdk/projects/_generated/enums.py

@@ -0,0 +1,4 @@
+# Generated by ariadne-codegen
+# Source: core/api/graphql/schemas/schema-latest.graphql
+
+from __future__ import annotations

wandb/sdk/projects/_generated/fetch_registry.py

@@ -0,0 +1,22 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import Optional
+
+from wandb._pydantic import GQLBase
+
+from .fragments import RegistryFragment
+
+
+class FetchRegistry(GQLBase):
+    entity: Optional[FetchRegistryEntity]
+
+
+class FetchRegistryEntity(GQLBase):
+    project: Optional[RegistryFragment]
+
+
+FetchRegistry.model_rebuild()
+FetchRegistryEntity.model_rebuild()

wandb/sdk/projects/_generated/fragments.py

@@ -0,0 +1,41 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import Field
+
+from wandb._pydantic import GQLBase, GQLId
+
+
+class RegistryFragment(GQLBase):
+    id: GQLId
+    allow_all_artifact_types_in_registry: bool = Field(
+        alias="allowAllArtifactTypesInRegistry"
+    )
+    artifact_types: RegistryFragmentArtifactTypes = Field(alias="artifactTypes")
+    name: str
+    description: Optional[str]
+    created_at: str = Field(alias="createdAt")
+    updated_at: Optional[str] = Field(alias="updatedAt")
+    access: Optional[str]
+
+
+class RegistryFragmentArtifactTypes(GQLBase):
+    edges: List[RegistryFragmentArtifactTypesEdges]
+
+
+class RegistryFragmentArtifactTypesEdges(GQLBase):
+    node: Optional[RegistryFragmentArtifactTypesEdgesNode]
+
+
+class RegistryFragmentArtifactTypesEdgesNode(GQLBase):
+    name: str
+
+
+RegistryFragment.model_rebuild()
+RegistryFragmentArtifactTypes.model_rebuild()
+RegistryFragmentArtifactTypesEdges.model_rebuild()
+RegistryFragmentArtifactTypesEdgesNode.model_rebuild()

wandb/sdk/projects/_generated/input_types.py

@@ -0,0 +1,13 @@
+# Generated by ariadne-codegen
+# Source: core/api/graphql/schemas/schema-latest.graphql
+
+from __future__ import annotations
+
+from typing import Optional
+
+from wandb._pydantic import GQLBase
+
+
+class ArtifactTypeInput(GQLBase):
+    name: str
+    description: Optional[str] = None

wandb/sdk/projects/_generated/operations.py

@@ -0,0 +1,88 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+__all__ = [
+    "DELETE_PROJECT_GQL",
+    "FETCH_REGISTRY_GQL",
+    "RENAME_PROJECT_GQL",
+    "UPSERT_REGISTRY_PROJECT_GQL",
+]
+
+FETCH_REGISTRY_GQL = """
+query FetchRegistry($name: String, $entityName: String) {
+  entity(name: $entityName) {
+    project(name: $name) {
+      ...RegistryFragment
+    }
+  }
+}
+
+fragment RegistryFragment on Project {
+  id
+  allowAllArtifactTypesInRegistry
+  artifactTypes(includeAll: true) {
+    edges {
+      node {
+        name
+      }
+    }
+  }
+  name
+  description
+  createdAt
+  updatedAt
+  access
+}
+"""
+
+RENAME_PROJECT_GQL = """
+mutation renameProject($entityName: String!, $oldProjectName: String!, $newProjectName: String!) {
+  renameProject(
+    input: {entityName: $entityName, oldProjectName: $oldProjectName, newProjectName: $newProjectName}
+  ) {
+    project {
+      name
+    }
+    inserted
+  }
+}
+"""
+
+UPSERT_REGISTRY_PROJECT_GQL = """
+mutation UpsertRegistryProject($description: String, $entityName: String, $name: String, $access: String, $allowAllArtifactTypesInRegistry: Boolean, $artifactTypes: [ArtifactTypeInput!]) {
+  upsertModel(
+    input: {description: $description, entityName: $entityName, name: $name, access: $access, allowAllArtifactTypesInRegistry: $allowAllArtifactTypesInRegistry, artifactTypes: $artifactTypes}
+  ) {
+    project {
+      ...RegistryFragment
+    }
+    inserted
+  }
+}
+
+fragment RegistryFragment on Project {
+  id
+  allowAllArtifactTypesInRegistry
+  artifactTypes(includeAll: true) {
+    edges {
+      node {
+        name
+      }
+    }
+  }
+  name
+  description
+  createdAt
+  updatedAt
+  access
+}
+"""
+
+DELETE_PROJECT_GQL = """
+mutation deleteProject($id: String!) {
+  deleteModel(input: {id: $id}) {
+    success
+    __typename
+  }
+}
+"""

wandb/sdk/projects/_generated/rename_project.py

@@ -0,0 +1,27 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import Field
+
+from wandb._pydantic import GQLBase
+
+
+class RenameProject(GQLBase):
+    rename_project: Optional[RenameProjectRenameProject] = Field(alias="renameProject")
+
+
+class RenameProjectRenameProject(GQLBase):
+    project: Optional[RenameProjectRenameProjectProject]
+    inserted: Optional[bool]
+
+
+class RenameProjectRenameProjectProject(GQLBase):
+    name: str
+
+
+RenameProject.model_rebuild()
+RenameProjectRenameProject.model_rebuild()

wandb/sdk/projects/_generated/upsert_registry_project.py

@@ -0,0 +1,27 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/projects/
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import Field
+
+from wandb._pydantic import GQLBase
+
+from .fragments import RegistryFragment
+
+
+class UpsertRegistryProject(GQLBase):
+    upsert_model: Optional[UpsertRegistryProjectUpsertModel] = Field(
+        alias="upsertModel"
+    )
+
+
+class UpsertRegistryProjectUpsertModel(GQLBase):
+    project: Optional[RegistryFragment]
+    inserted: Optional[bool]
+
+
+UpsertRegistryProject.model_rebuild()
+UpsertRegistryProjectUpsertModel.model_rebuild()