wandb 0.19.8__py3-none-macosx_11_0_arm64.whl → 0.19.10__py3-none-macosx_11_0_arm64.whl

wandb/apis/public/automations.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from itertools import chain
+from typing import TYPE_CHECKING, Any, Iterable, Mapping
+
+from pydantic import ValidationError
+from typing_extensions import override
+from wandb_graphql.language.ast import Document
+
+from wandb.apis.paginator import Paginator, _Client
+
+if TYPE_CHECKING:
+    from wandb.automations import Automation
+    from wandb.automations._generated import ProjectConnectionFields
+
+
+class Automations(Paginator["Automation"]):
+    last_response: ProjectConnectionFields | None
+    _query: Document
+
+    def __init__(
+        self,
+        client: _Client,
+        variables: Mapping[str, Any],
+        per_page: int = 50,
+        _query: Document | None = None,
+    ):
+        super().__init__(client, variables, per_page=per_page)
+        if _query is None:
+            raise RuntimeError(f"Query required for {type(self).__qualname__}")
+        self._query = _query
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more items to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch the raw response data for the current page."""
+        from wandb.automations._generated import ProjectConnectionFields
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["searchScope"]["projects"]
+            self.last_response = ProjectConnectionFields.model_validate(page_data)
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[Automation]:
+        """Parse the page data into a list of objects."""
+        from wandb.automations import Automation
+
+        page = self.last_response
+        return [
+            Automation.model_validate_json(obj.model_dump_json())
+            for obj in chain.from_iterable(edge.node.triggers for edge in page.edges)
+        ]

wandb/apis/public/files.py

@@ -12,7 +12,7 @@ import wandb
 from wandb import util
 from wandb.apis.attrs import Attrs
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import SizedPaginator
 from wandb.apis.public import utils
 from wandb.apis.public.api import Api
 from wandb.apis.public.const import RETRY_TIMEDELTA
@@ -41,7 +41,7 @@ FILE_FRAGMENT = """fragment RunFilesFragment on Run {
 }"""
 
 
-class Files(Paginator):
+class Files(SizedPaginator["File"]):
     """An iterable collection of `File` objects."""
 
     QUERY = gql(

wandb/apis/public/integrations.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Iterable
+
+from pydantic import ValidationError
+from typing_extensions import override
+from wandb_gql import gql
+from wandb_graphql.language.ast import Document
+
+from wandb.apis.paginator import Paginator
+
+if TYPE_CHECKING:
+    from wandb.apis.paginator import _Client
+    from wandb.automations import Integration, SlackIntegration, WebhookIntegration
+    from wandb.automations._generated import (
+        GenericWebhookIntegrationConnectionFields,
+        IntegrationConnectionFields,
+        SlackIntegrationConnectionFields,
+    )
+
+
+class Integrations(Paginator["Integration"]):
+    last_response: IntegrationConnectionFields | None
+    _query: Document
+
+    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
+        from wandb.automations._generated import INTEGRATIONS_BY_ENTITY_GQL
+
+        super().__init__(client, variables, per_page=per_page)
+        # All integrations for entity
+        self._query = gql(INTEGRATIONS_BY_ENTITY_GQL)
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more Integrations to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch and parse the response data for the current page."""
+        from wandb.automations._generated import IntegrationConnectionFields
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["entity"]["integrations"]
+            self.last_response = IntegrationConnectionFields.model_validate(page_data)
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[Integration]:
+        """Parse the page data into a list of Integrations."""
+        from wandb.automations.integrations import _IntegrationEdge
+
+        page = self.last_response
+        return [_IntegrationEdge.model_validate(edge).node for edge in page.edges]
+
+
+class WebhookIntegrations(Paginator["WebhookIntegration"]):
+    last_response: GenericWebhookIntegrationConnectionFields | None
+    _query: Document
+
+    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
+        from wandb.automations._generated import (
+            GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL,
+        )
+
+        super().__init__(client, variables, per_page=per_page)
+        # Webhook integrations for entity
+        self._query = gql(GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL)
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more Integrations to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch and parse the response data for the current page."""
+        from wandb.automations._generated import (
+            GenericWebhookIntegrationConnectionFields,
+        )
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["entity"]["integrations"]
+            self.last_response = (
+                GenericWebhookIntegrationConnectionFields.model_validate(page_data)
+            )
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[WebhookIntegration]:
+        """Parse the page data into a list of Integrations."""
+        from wandb.automations import WebhookIntegration
+
+        page = self.last_response
+        return [WebhookIntegration.model_validate(edge.node) for edge in page.edges]
+
+
+class SlackIntegrations(Paginator["SlackIntegration"]):
+    last_response: SlackIntegrationConnectionFields | None
+    _query: Document
+
+    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
+        from wandb.automations._generated import SLACK_INTEGRATIONS_BY_ENTITY_GQL
+
+        super().__init__(client, variables, per_page=per_page)
+        # Slack integrations for entity
+        self._query = gql(SLACK_INTEGRATIONS_BY_ENTITY_GQL)
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more Integrations to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch and parse the response data for the current page."""
+        from wandb.automations._generated import SlackIntegrationConnectionFields
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["entity"]["integrations"]
+            self.last_response = SlackIntegrationConnectionFields.model_validate(
+                page_data
+            )
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> list[SlackIntegration]:
+        """Parse the page data into a list of Integrations."""
+        from wandb.automations import SlackIntegration
+
+        page = self.last_response
+        return [SlackIntegration.model_validate(edge.node) for edge in page.edges]

wandb/apis/public/projects.py

@@ -1,5 +1,8 @@
 """Public API: projects."""
 
+from contextlib import suppress
+
+from requests import HTTPError
 from wandb_gql import gql
 
 from wandb.apis import public
@@ -17,7 +20,7 @@ PROJECT_FRAGMENT = """fragment ProjectFragment on Project {
 }"""
 
 
-class Projects(Paginator):
+class Projects(Paginator["Project"]):
     """An iterable collection of `Project` objects."""
 
     QUERY = gql(
@@ -49,7 +52,8 @@ class Projects(Paginator):
         super().__init__(client, variables, per_page)
 
     @property
-    def length(self):
+    def length(self) -> None:
+        # For backwards compatibility, even though this isn't a SizedPaginator
         return None
 
     @property
@@ -152,3 +156,29 @@ class Project(Attrs):
             )
             for e in ret["project"]["sweeps"]["edges"]
         ]
+
+    _PROJECT_ID = gql(
+        """
+        query ProjectID($projectName: String!, $entityName: String!) {
+            project(name: $projectName, entityName: $entityName) {
+                id
+            }
+        }
+        """
+    )
+
+    @property
+    def id(self) -> str:
+        # This is a workaround to ensure that the project ID can be retrieved
+        # on demand, as it generally is not set or fetched on instantiation.
+        # This is necessary if using this project as the scope of a new Automation.
+        with suppress(LookupError):
+            return self._attrs["id"]
+
+        variable_values = {"projectName": self.name, "entityName": self.entity}
+        try:
+            data = self.client.execute(self._PROJECT_ID, variable_values)
+            self._attrs["id"] = data["project"]["id"]
+            return self._attrs["id"]
+        except (HTTPError, LookupError, TypeError) as e:
+            raise ValueError(f"Unable to fetch project ID: {variable_values!r}") from e

wandb/apis/public/reports.py

@@ -9,11 +9,11 @@ from wandb_gql import gql
 import wandb
 from wandb.apis import public
 from wandb.apis.attrs import Attrs
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import SizedPaginator
 from wandb.sdk.lib import ipython
 
 
-class Reports(Paginator):
+class Reports(SizedPaginator["BetaReport"]):
     """Reports is an iterable collection of `BetaReport` objects."""
 
     QUERY = gql(

wandb/apis/public/runs.py

@@ -24,7 +24,7 @@ from wandb.apis import public
 from wandb.apis.attrs import Attrs
 from wandb.apis.internal import Api as InternalApi
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.paginator import Paginator
+from wandb.apis.paginator import SizedPaginator
 from wandb.apis.public.const import RETRY_TIMEDELTA
 from wandb.sdk.lib import ipython, json_util, runid
 from wandb.sdk.lib.paths import LogicalPath
@@ -61,7 +61,7 @@ RUN_FRAGMENT = """fragment RunFragment on Run {
 }"""
 
 
-class Runs(Paginator):
+class Runs(SizedPaginator["Run"]):
     """An iterable collection of runs associated with a project and optional filter.
 
     This is generally used indirectly via the `Api`.runs method.
@@ -421,16 +421,15 @@ class Run(Attrs):
             """
         query Run($project: String!, $entity: String!, $name: String!) {{
             project(name: $project, entityName: $entity) {{
-                {}
                 run(name: $name) {{
+                    {}
                     ...RunFragment
                 }}
             }}
         }}
         {}
         """.format(
-                # Only query internalId if the server supports it
-                "internalId" if self._server_provides_internal_id_for_project() else "",
+                "projectId" if self._server_provides_internal_id_for_project() else "",
                 RUN_FRAGMENT,
             )
         )
@@ -444,7 +443,11 @@ class Run(Attrs):
                 raise ValueError("Could not find run {}".format(self))
             self._attrs = response["project"]["run"]
             self._state = self._attrs["state"]
-            self._project_internal_id = response["project"].get("internalId", None)
+
+            self._project_internal_id = (
+                int(self._attrs["projectId"]) if "projectId" in self._attrs else None
+            )
+
             if self._include_sweeps and self.sweep_name and not self.sweep:
                 # There may be a lot of runs. Don't bother pulling them all
                 # just for the sake of this one.
@@ -847,7 +850,12 @@ class Run(Attrs):
         api.set_current_run_id(self.id)
 
         if isinstance(artifact, wandb.Artifact) and not artifact.is_draft():
-            api.use_artifact(artifact.id, use_as=use_as or artifact.name)
+            api.use_artifact(
+                artifact.id,
+                use_as=use_as or artifact.name,
+                artifact_entity_name=artifact.entity,
+                artifact_project_name=artifact.project,
+            )
             return artifact
         elif isinstance(artifact, wandb.Artifact) and artifact.is_draft():
             raise ValueError(
@@ -911,8 +919,8 @@ class Run(Attrs):
         This check is done by utilizing GraphQL introspection in the available fields on the Project type.
         """
         query_string = """
-           query ProbeProjectInput {
-                ProjectType: __type(name:"Project") {
+           query ProbeRunInput {
+                RunType: __type(name:"Run") {
                     fields {
                         name
                     }
@@ -924,8 +932,8 @@ class Run(Attrs):
         if self.server_provides_internal_id_field is None:
             query = gql(query_string)
             res = self.client.execute(query)
-            self.server_provides_internal_id_field = "internalId" in [
-                x["name"] for x in (res.get("ProjectType", {}).get("fields", [{}]))
+            self.server_provides_internal_id_field = "projectId" in [
+                x["name"] for x in (res.get("RunType", {}).get("fields", [{}]))
             ]
 
         return self.server_provides_internal_id_field

wandb/apis/public/utils.py

@@ -1,8 +1,11 @@
 import re
 from enum import Enum
-from typing import Optional
+from typing import Any, Dict, Iterable, Mapping, Optional, Set
 from urllib.parse import urlparse
 
+from wandb_gql import gql
+from wandb_graphql.language import ast, visitor
+
 from wandb._iterutils import one
 from wandb.sdk.artifacts._validators import is_artifact_registry_project
 from wandb.sdk.internal.internal_api import Api as InternalApi
@@ -102,3 +105,106 @@ def fetch_org_from_settings_or_entity(
         )
         organization = entity_org.display_name
     return organization
+
+
+class _GQLCompatRewriter(visitor.Visitor):
+    """GraphQL AST visitor to rewrite queries/mutations to be compatible with older server versions."""
+
+    omit_variables: Set[str]
+    omit_fragments: Set[str]
+    omit_fields: Set[str]
+    rename_fields: Dict[str, str]
+
+    def __init__(
+        self,
+        omit_variables: Optional[Iterable[str]] = None,
+        omit_fragments: Optional[Iterable[str]] = None,
+        omit_fields: Optional[Iterable[str]] = None,
+        rename_fields: Optional[Mapping[str, str]] = None,
+    ):
+        self.omit_variables = set(omit_variables or ())
+        self.omit_fragments = set(omit_fragments or ())
+        self.omit_fields = set(omit_fields or ())
+        self.rename_fields = dict(rename_fields or {})
+
+    def enter_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any:  # noqa: N802
+        if node.variable.name.value in self.omit_variables:
+            return visitor.REMOVE
+        # return node
+
+    def enter_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any:  # noqa: N802
+        # For context, note that e.g.:
+        #
+        #   {description: $description
+        #   ...}
+        #
+        # Is parsed as:
+        #
+        #   ObjectValue(fields=[
+        #     ObjectField(name=Name(value='description'), value=Variable(name=Name(value='description'))),
+        #   ...])
+        if (
+            isinstance(var := node.value, ast.Variable)
+            and var.name.value in self.omit_variables
+        ):
+            return visitor.REMOVE
+
+    def enter_Argument(self, node: ast.Argument, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_variables:
+            return visitor.REMOVE
+
+    def enter_FragmentDefinition(self, node: ast.FragmentDefinition, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_fragments:
+            return visitor.REMOVE
+
+    def enter_FragmentSpread(self, node: ast.FragmentSpread, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_fragments:
+            return visitor.REMOVE
+
+    def enter_Field(self, node: ast.Field, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_fields:
+            return visitor.REMOVE
+        if new_name := self.rename_fields.get(node.name.value):
+            node.name.value = new_name
+            return node
+
+    def leave_Field(self, node: ast.Field, *_, **__) -> Any:  # noqa: N802
+        # If the field had a selection set, but now it's empty, remove the field entirely
+        if (node.selection_set is not None) and (not node.selection_set.selections):
+            return visitor.REMOVE
+
+
+def gql_compat(
+    request_string: str,
+    omit_variables: Optional[Iterable[str]] = None,
+    omit_fragments: Optional[Iterable[str]] = None,
+    omit_fields: Optional[Iterable[str]] = None,
+    rename_fields: Optional[Mapping[str, str]] = None,
+) -> ast.Document:
+    """Rewrite a GraphQL request string to ensure compatibility with older server versions.
+
+    Args:
+        request_string (str): The GraphQL request string to rewrite.
+        omit_variables (Iterable[str] | None): Names of variables to remove from the request string.
+        omit_fragments (Iterable[str] | None): Names of fragments to remove from the request string.
+        omit_fields (Iterable[str] | None): Names of fields to remove from the request string.
+        rename_fields (Mapping[str, str] | None):
+            A mapping of fields to rename in the request string, given as `{old_name -> new_name}`.
+
+    Returns:
+        str: Modified GraphQL request string with fragments on omitted types removed.
+    """
+    # Parse the request into a GraphQL AST
+    doc = gql(request_string)
+
+    if not (omit_variables or omit_fragments or omit_fields or rename_fields):
+        return doc
+
+    # Visit the AST with our visitor to filter out unwanted fragments
+    rewriter = _GQLCompatRewriter(
+        omit_variables=omit_variables,
+        omit_fragments=omit_fragments,
+        omit_fields=omit_fields,
+        rename_fields=rename_fields,
+    )
+    return visitor.visit(doc, rewriter)

wandb/automations/__init__.py

@@ -0,0 +1,81 @@
+from wandb._pydantic import IS_PYDANTIC_V2
+
+from . import _filters as filters
+from . import actions, automations, events, scopes
+from .actions import ActionType, DoNothing, DoNotification, DoWebhook
+from .automations import Automation, NewAutomation, PreparedAutomation
+from .events import (
+    ArtifactEvent,
+    EventType,
+    OnAddArtifactAlias,
+    OnCreateArtifact,
+    OnLinkArtifact,
+    OnRunMetric,
+    RunEvent,
+)
+from .integrations import Integration, SlackIntegration, WebhookIntegration
+from .scopes import ArtifactCollectionScope, ProjectScope, ScopeType
+
+# ----------------------------------------------------------------------------
+# WARNINGS on import
+if not IS_PYDANTIC_V2:
+    # Raises an error in Pydantic v1 environments, where the Automations API
+    # has not been tested and is unlikely to work as expected.
+    #
+    # Remove this when we either:
+    # - Drop support for Pydantic v1
+    # - Are able to implement (limited) Pydantic v1 support
+    raise ImportError(
+        "The W&B Automations API is not supported in Pydantic v1 at this time. "
+        "If at all possible, we currently recommend upgrading to Pydantic v2 to use this feature.",
+    )
+
+else:
+    # If Pydantic v2 is available, we can use the full Automations API
+    # but communicate to users that the API is still experimental and
+    # may change rapidly.
+    import warnings
+
+    warnings.warn(
+        "The W&B Automations API is currently experimental. Although we'll communicate "
+        "breaking changes in release notes and attempt to minimize them in general, "
+        "please know that such changes may occur between release versions without notice. "
+        "We strongly recommend pinning your `wandb` version when using the Automations API "
+        "to avoid unexpected breakages.",
+        FutureWarning,
+        stacklevel=1,
+    )
+# ----------------------------------------------------------------------------
+
+__all__ = [
+    "filters",
+    "scopes",
+    "events",
+    "actions",
+    "automations",
+    # Scopes
+    "ScopeType",
+    "ArtifactCollectionScope",
+    "ProjectScope",
+    # Events
+    "EventType",
+    "OnAddArtifactAlias",
+    "OnCreateArtifact",
+    "OnLinkArtifact",
+    "OnRunMetric",
+    "ArtifactEvent",
+    "RunEvent",
+    # Actions
+    "ActionType",
+    "DoNotification",
+    "DoWebhook",
+    "DoNothing",
+    # Automations
+    "Automation",
+    "NewAutomation",
+    "PreparedAutomation",
+    # Integrations
+    "Integration",
+    "SlackIntegration",
+    "WebhookIntegration",
+]

wandb/automations/_filters/__init__.py

@@ -0,0 +1,40 @@
+from .expressions import FilterExpr, MongoLikeFilter
+from .operators import (
+    And,
+    Contains,
+    Eq,
+    Exists,
+    Gt,
+    Gte,
+    In,
+    Lt,
+    Lte,
+    Ne,
+    Nor,
+    Not,
+    NotIn,
+    Op,
+    Or,
+    Regex,
+)
+
+__all__ = [
+    "And",
+    "Or",
+    "Nor",
+    "Not",
+    "Op",
+    "Gt",
+    "Lt",
+    "Gte",
+    "Lte",
+    "Eq",
+    "Ne",
+    "In",
+    "NotIn",
+    "Contains",
+    "Exists",
+    "Regex",
+    "FilterExpr",
+    "MongoLikeFilter",
+]