wandb 0.22.1__py3-none-win_arm64.whl

wandb/automations/_utils.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+from typing import Any, Collection, Final, Optional, Protocol, TypedDict
+
+from pydantic import Field
+from typing_extensions import Annotated, Self, Unpack
+
+from wandb._pydantic import GQLBase, GQLId, computed_field, model_validator, to_json
+
+from ._filters import MongoLikeFilter
+from ._generated import (
+    CreateFilterTriggerInput,
+    QueueJobActionInput,
+    TriggeredActionConfig,
+    UpdateFilterTriggerInput,
+)
+from ._validators import to_input_action
+from .actions import (
+    ActionType,
+    DoNothing,
+    InputAction,
+    SavedAction,
+    SendNotification,
+    SendWebhook,
+)
+from .automations import Automation, NewAutomation
+from .events import EventType, InputEvent, RunMetricFilter, _WrappedSavedEventFilter
+from .scopes import AutomationScope, ScopeType
+
+EXCLUDED_INPUT_EVENTS: Final[Collection[EventType]] = frozenset(
+    {
+        EventType.UPDATE_ARTIFACT_ALIAS,
+    }
+)
+"""Event types that should not be assigned when creating/updating automations."""
+
+EXCLUDED_INPUT_ACTIONS: Final[Collection[ActionType]] = frozenset(
+    {
+        ActionType.QUEUE_JOB,
+    }
+)
+"""Action types that should not be assigned when creating/updating automations."""
+
+ALWAYS_SUPPORTED_EVENTS: Final[Collection[EventType]] = frozenset(
+    {
+        EventType.CREATE_ARTIFACT,
+        EventType.LINK_ARTIFACT,
+        EventType.ADD_ARTIFACT_ALIAS,
+    }
+)
+"""Event types that we can safely assume all contemporary server versions support."""
+
+ALWAYS_SUPPORTED_ACTIONS: Final[Collection[ActionType]] = frozenset(
+    {
+        ActionType.NOTIFICATION,
+        ActionType.GENERIC_WEBHOOK,
+    }
+)
+"""Action types that we can safely assume all contemporary server versions support."""
+
+
+class HasId(Protocol):
+    id: str
+
+
+def extract_id(obj: HasId | str) -> str:
+    return obj.id if hasattr(obj, "id") else obj
+
+
+# ---------------------------------------------------------------------------
+ACTION_CONFIG_KEYS: dict[ActionType, str] = {
+    ActionType.NOTIFICATION: "notification_action_input",
+    ActionType.GENERIC_WEBHOOK: "generic_webhook_action_input",
+    ActionType.NO_OP: "no_op_action_input",
+    ActionType.QUEUE_JOB: "queue_job_action_input",
+}
+
+
+class InputActionConfig(TriggeredActionConfig):
+    """A `TriggeredActionConfig` that prepares the action config for saving an automation."""
+
+    # NOTE: `QueueJobActionInput` for defining a Launch job is deprecated,
+    # so while it's allowed here to update EXISTING mutations, we don't
+    # currently expose it through the public API to create NEW automations.
+    queue_job_action_input: Optional[QueueJobActionInput] = None
+
+    notification_action_input: Optional[SendNotification] = None
+    generic_webhook_action_input: Optional[SendWebhook] = None
+    no_op_action_input: Optional[DoNothing] = None
+
+
+def prepare_action_config_input(obj: SavedAction | InputAction) -> dict[str, Any]:
+    """Prepare the `TriggeredActionConfig` input, nesting the action input inside the appropriate key.
+
+    This is necessary to conform to the schemas for:
+    - CreateFilterTriggerInput
+    - UpdateFilterTriggerInput
+    """
+    # Delegate to inner validators to convert SavedAction -> InputAction types, if needed.
+    obj = to_input_action(obj)
+    return InputActionConfig(**{ACTION_CONFIG_KEYS[obj.action_type]: obj}).model_dump()
+
+
+def prepare_event_filter_input(
+    obj: _WrappedSavedEventFilter | MongoLikeFilter | RunMetricFilter,
+) -> str:
+    """Prepare the `EventFilter` input, unnesting the filter if needed and serializing to JSON.
+
+    This is necessary to conform to the schemas for:
+    - CreateFilterTriggerInput
+    - UpdateFilterTriggerInput
+    """
+    # Input event filters are nested one level deeper than saved event filters.
+    # Note that this is NOT the case for run/run metric filters.
+    #
+    # Yes, this is confusing.  It's also necessary to conform to under-the-hood
+    # schemas and logic in the backend.
+    filter_to_serialize = (
+        obj.filter if isinstance(obj, _WrappedSavedEventFilter) else obj
+    )
+    return to_json(filter_to_serialize)
+
+
+class WriteAutomationsKwargs(TypedDict, total=False):
+    """Keyword arguments that can be passed to create or update an automation."""
+
+    name: str
+    description: str
+    enabled: bool
+    scope: AutomationScope
+    event: InputEvent
+    action: InputAction
+
+
+class ValidatedCreateInput(GQLBase, extra="forbid", frozen=True):
+    """Validated automation parameters, prepared for creating a new automation.
+
+    Note: Users should never need to instantiate this class directly.
+    """
+
+    name: str
+    description: Optional[str] = None
+    enabled: bool = True
+
+    # ------------------------------------------------------------------------------
+    # Set on instantiation, but used to derive other fields and deliberately
+    # EXCLUDED from the final GraphQL request vars
+    event: Annotated[InputEvent, Field(exclude=True)]
+    action: Annotated[InputAction, Field(exclude=True)]
+
+    # ------------------------------------------------------------------------------
+    # Derived fields to match the input schemas
+    @computed_field
+    def scope_type(self) -> ScopeType:
+        return self.event.scope.scope_type
+
+    @computed_field
+    def scope_id(self) -> GQLId:
+        return self.event.scope.id
+
+    @computed_field
+    def triggering_event_type(self) -> EventType:
+        return self.event.event_type
+
+    @computed_field
+    def event_filter(self) -> str:
+        return prepare_event_filter_input(self.event.filter)
+
+    @computed_field
+    def triggered_action_type(self) -> ActionType:
+        return self.action.action_type
+
+    @computed_field
+    def triggered_action_config(self) -> dict[str, Any]:
+        return prepare_action_config_input(self.action)
+
+    # ------------------------------------------------------------------------------
+    # Custom validation
+    @model_validator(mode="after")
+    def _forbid_legacy_event_types(self) -> Self:
+        if (type_ := self.event.event_type) in EXCLUDED_INPUT_EVENTS:
+            raise ValueError(f"{type_!r} events cannot be assigned to automations.")
+        return self
+
+    @model_validator(mode="after")
+    def _forbid_legacy_action_types(self) -> Self:
+        if (type_ := self.action.action_type) in EXCLUDED_INPUT_ACTIONS:
+            raise ValueError(f"{type_!r} actions cannot be assigned to automations.")
+        return self
+
+
+def prepare_to_create(
+    obj: NewAutomation | None = None,
+    /,
+    **kwargs: Unpack[WriteAutomationsKwargs],
+) -> CreateFilterTriggerInput:
+    """Prepares the payload to create an automation in a GraphQL request."""
+    # Validate all input variables, and prepare as expected by the GraphQL request.
+    # - if an object is provided, override its fields with any keyword args
+    # - otherwise, instantiate from the keyword args
+
+    # NOTE: `exclude_none=True` drops fields that are still `None`.
+    #
+    # This assumes that `None` is good enough for now as a sentinel
+    # "unset" value.  If this proves insufficient, revisit in the future,
+    # as it should be reasonably easy to implement a custom sentinel
+    # type later on.
+    obj_dict = {**obj.model_dump(exclude_none=True), **kwargs} if obj else kwargs
+    validated = ValidatedCreateInput(**obj_dict)
+    return CreateFilterTriggerInput.model_validate(validated)
+
+
+def prepare_to_update(
+    obj: Automation | None = None,
+    /,
+    **kwargs: Unpack[WriteAutomationsKwargs],
+) -> UpdateFilterTriggerInput:
+    """Prepares the payload to update an automation in a GraphQL request."""
+    # Validate all values:
+    # - if an object is provided, override its fields with any keyword args
+    # - otherwise, instantiate from the keyword args
+    v_obj = Automation(**{**dict(obj or {}), **kwargs})
+
+    return UpdateFilterTriggerInput(
+        id=v_obj.id,
+        name=v_obj.name,
+        description=v_obj.description,
+        enabled=v_obj.enabled,
+        scope_type=v_obj.scope.scope_type,
+        scope_id=v_obj.scope.id,
+        triggering_event_type=v_obj.event.event_type,
+        event_filter=prepare_event_filter_input(v_obj.event.filter),
+        triggered_action_type=v_obj.action.action_type,
+        triggered_action_config=prepare_action_config_input(v_obj.action),
+    )

wandb/automations/_validators.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from enum import Enum
+from functools import singledispatch
+from itertools import chain
+from typing import Any, TypeVar
+
+from pydantic import BeforeValidator, Json, PlainSerializer
+from pydantic_core import PydanticUseDefault
+from typing_extensions import Annotated
+
+from wandb._pydantic import to_json
+
+from ._filters import And, FilterExpr, In, Nor, Not, NotIn, Op, Or
+
+T = TypeVar("T")
+
+
+def ensure_json(v: Any) -> Any:
+    """In case the incoming value isn't serialized JSON, reserialize it.
+
+    This lets us use `Json[...]` fields with values that are already deserialized.
+    """
+    # NOTE: Assumes that the deserialized type is not itself a string.
+    # Revisit this if we need to support deserialized types that are str/bytes.
+    return v if isinstance(v, (str, bytes)) else to_json(v)
+
+
+# Allow lenient instantiation/validation: incoming data may already be deserialized.
+SerializedToJson = Annotated[
+    Json[T], BeforeValidator(ensure_json), PlainSerializer(to_json)
+]
+
+
+class LenientStrEnum(str, Enum):
+    """A string enum allowing for case-insensitive lookups by value.
+
+    May include other internal customizations if needed.
+
+    Note: This is a bespoke, internal implementation and NOT intended as a
+    backport of `enum.StrEnum` from Python 3.11+.
+    """
+
+    def __repr__(self) -> str:
+        return self.name
+
+    @classmethod
+    def _missing_(cls, value: object) -> Any:
+        # Accept case-insensitive enum values
+        if isinstance(value, str):
+            v = value.lower()
+            return next((e for e in cls if e.value.lower() == v), None)
+        return None
+
+
+def default_if_none(v: Any) -> Any:
+    """A before-validator validator that coerces `None` to the default field value instead."""
+    # https://docs.pydantic.dev/2.11/api/pydantic_core/#pydantic_core.PydanticUseDefault
+    if v is None:
+        raise PydanticUseDefault
+    return v
+
+
+def upper_if_str(v: Any) -> Any:
+    return v.strip().upper() if isinstance(v, str) else v
+
+
+# ----------------------------------------------------------------------------
+def to_scope(v: Any) -> Any:
+    """Convert eligible objects, including pre-existing `wandb` types, to an automation scope."""
+    from wandb.apis.public import ArtifactCollection, Project
+
+    from .scopes import ProjectScope, _ArtifactPortfolioScope, _ArtifactSequenceScope
+
+    if isinstance(v, Project):
+        return ProjectScope(id=v.id, name=v.name)
+    if isinstance(v, ArtifactCollection):
+        cls = _ArtifactSequenceScope if v.is_sequence() else _ArtifactPortfolioScope
+        return cls(id=v.id, name=v.name)
+    return v
+
+
+def to_saved_action(v: Any) -> Any:
+    """If necessary (and possible), convert the object to a saved action."""
+    from .actions import (
+        DoNothing,
+        SavedNoOpAction,
+        SavedNotificationAction,
+        SavedWebhookAction,
+        SendNotification,
+        SendWebhook,
+    )
+
+    if isinstance(v, SendNotification):
+        return SavedNotificationAction(
+            integration={"id": v.integration_id},
+            **v.model_dump(exclude={"integration_id"}),
+        )
+    if isinstance(v, SendWebhook):
+        return SavedWebhookAction(
+            integration={"id": v.integration_id},
+            **v.model_dump(exclude={"integration_id"}),
+        )
+    if isinstance(v, DoNothing):
+        return SavedNoOpAction.model_validate(v)
+
+    return v
+
+
+def to_input_action(v: Any) -> Any:
+    """If necessary (and possible), convert the object to an input action."""
+    from .actions import (
+        DoNothing,
+        SavedNoOpAction,
+        SavedNotificationAction,
+        SavedWebhookAction,
+        SendNotification,
+        SendWebhook,
+    )
+
+    if isinstance(v, SavedNotificationAction):
+        return SendNotification(
+            integration_id=v.integration.id,
+            **v.model_dump(exclude={"integration"}),
+        )
+    if isinstance(v, SavedWebhookAction):
+        return SendWebhook(
+            integration_id=v.integration.id,
+            **v.model_dump(exclude={"integration"}),
+        )
+    if isinstance(v, SavedNoOpAction):
+        return DoNothing.model_validate(v)
+
+    return v
+
+
+# ----------------------------------------------------------------------------
+@singledispatch
+def simplify_op(op: Op | FilterExpr) -> Op | FilterExpr:
+    """Simplify a MongoDB filter by removing and unnesting redundant operators."""
+    return op
+
+
+@simplify_op.register
+def _(op: And) -> Op:
+    # {"$and": []} -> {"$and": []}
+    if not (args := op.and_):
+        return op
+
+    # {"$and": [op]} -> op
+    if len(args) == 1:
+        return simplify_op(args[0])
+
+    # {"$and": [op, {"$and": [op2, ...]}]} -> {"$and": [op, op2, ...]}
+    flattened = chain.from_iterable(x.and_ if isinstance(x, And) else [x] for x in args)
+    return And(and_=map(simplify_op, flattened))
+
+
+@simplify_op.register
+def _(op: Or) -> Op:
+    # {"$or": []} -> {"$or": []}
+    if not (args := op.or_):
+        return op
+
+    # {"$or": [op]} -> op
+    if len(args) == 1:
+        return simplify_op(args[0])
+
+    # {"$or": [op, {"$or": [op2, ...]}]} -> {"$or": [op, op2, ...]}
+    flattened = chain.from_iterable(x.or_ if isinstance(x, Or) else [x] for x in args)
+    return Or(or_=map(simplify_op, flattened))
+
+
+@simplify_op.register
+def _(op: Not) -> Op:
+    inner = op.not_
+
+    # {"$not": {"$not": op}} -> op
+    # {"$not": {"$or": [op, ...]}} -> {"$nor": [op, ...]}
+    # {"$not": {"$nor": [op, ...]}} -> {"$or": [op, ...]}
+    # {"$not": {"$in": [op, ...]}} -> {"$nin": [op, ...]}
+    # {"$not": {"$nin": [op, ...]}} -> {"$in": [op, ...]}
+    if isinstance(inner, (Not, Or, Nor, In, NotIn)):
+        return simplify_op(~inner)
+    return Not(not_=simplify_op(inner))

wandb/automations/actions.py

@@ -0,0 +1,220 @@
+"""Actions that are triggered by W&B Automations."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, Optional, Union
+
+from pydantic import BeforeValidator, Field
+from typing_extensions import Annotated, Self, get_args
+
+from wandb._pydantic import GQLBase, GQLId, Typename
+from wandb._strutils import nameof
+
+from ._generated import (
+    AlertSeverity,
+    GenericWebhookActionFields,
+    GenericWebhookActionInput,
+    NoOpActionFields,
+    NoOpTriggeredActionInput,
+    NotificationActionFields,
+    NotificationActionInput,
+    QueueJobActionFields,
+)
+from ._validators import (
+    LenientStrEnum,
+    SerializedToJson,
+    default_if_none,
+    to_input_action,
+    to_saved_action,
+    upper_if_str,
+)
+from .integrations import SlackIntegration, WebhookIntegration
+
+
+# NOTE: Name shortened for readability and defined publicly for easier access
+class ActionType(LenientStrEnum):
+    """The type of action triggered by an automation."""
+
+    QUEUE_JOB = "QUEUE_JOB"  # NOTE: Deprecated for creation
+    NOTIFICATION = "NOTIFICATION"
+    GENERIC_WEBHOOK = "GENERIC_WEBHOOK"
+    NO_OP = "NO_OP"
+
+
+# ------------------------------------------------------------------------------
+# Saved types: for parsing response data from saved automations
+
+
+# NOTE: `QueueJobActionInput` for defining a Launch job is deprecated,
+# so while we allow parsing it from previously saved Automations, we deliberately
+# don't currently expose it in the API for creating automations.
+class SavedLaunchJobAction(QueueJobActionFields):
+    action_type: Literal[ActionType.QUEUE_JOB] = ActionType.QUEUE_JOB
+
+
+# FIXME: Find a better place to put these OR a better way to handle the
+#   conversion from `InputAction` -> `SavedAction`.
+#
+# Necessary placeholder class defs for converting:
+# - `SendNotification -> SavedNotificationAction`
+# - `SendWebhook -> SavedWebhookAction`
+#
+# The "input" types (`Send{Notification,Webhook}`) will only have an `integration_id`,
+# and we don't want/need to fetch the other `{Slack,Webhook}Integration` fields if
+# we can avoid it.
+class _SavedActionSlackIntegration(GQLBase, extra="allow"):
+    typename__: Typename[Literal["SlackIntegration"]] = "SlackIntegration"
+    id: GQLId
+
+
+class _SavedActionWebhookIntegration(GQLBase, extra="allow"):
+    typename__: Typename[Literal["GenericWebhookIntegration"]] = (
+        "GenericWebhookIntegration"
+    )
+    id: GQLId
+
+
+class SavedNotificationAction(NotificationActionFields):
+    action_type: Literal[ActionType.NOTIFICATION] = ActionType.NOTIFICATION
+    integration: _SavedActionSlackIntegration
+
+
+class SavedWebhookAction(GenericWebhookActionFields):
+    action_type: Literal[ActionType.GENERIC_WEBHOOK] = ActionType.GENERIC_WEBHOOK
+    integration: _SavedActionWebhookIntegration
+
+    # We override the type of the `requestPayload` field since the original GraphQL
+    # schema (and generated class) effectively defines it as a string, when we know
+    # and need to anticipate the expected structure of the JSON-serialized data.
+    request_payload: Annotated[
+        Optional[SerializedToJson[dict[str, Any]]],
+        Field(alias="requestPayload"),
+    ] = None  # type: ignore[assignment]
+
+
+class SavedNoOpAction(NoOpActionFields, frozen=True):
+    action_type: Literal[ActionType.NO_OP] = ActionType.NO_OP
+
+    no_op: Annotated[bool, BeforeValidator(default_if_none)] = True
+    """Placeholder field, only needed to conform to schema requirements.
+
+    There should never be a need to set this field explicitly, as its value is ignored.
+    """
+
+
+# for type annotations
+SavedAction = Annotated[
+    Union[
+        SavedLaunchJobAction,
+        SavedNotificationAction,
+        SavedWebhookAction,
+        SavedNoOpAction,
+    ],
+    BeforeValidator(to_saved_action),
+    Field(discriminator="typename__"),
+]
+# for runtime type checks
+SavedActionTypes: tuple[type, ...] = get_args(SavedAction.__origin__)  # type: ignore[attr-defined]
+
+
+# ------------------------------------------------------------------------------
+# Input types: for creating or updating automations
+class _BaseActionInput(GQLBase):
+    action_type: Annotated[ActionType, Field(frozen=True)]
+    """The kind of action to be triggered."""
+
+
+class SendNotification(_BaseActionInput, NotificationActionInput):
+    """Defines an automation action that sends a (Slack) notification."""
+
+    action_type: Literal[ActionType.NOTIFICATION] = ActionType.NOTIFICATION
+
+    integration_id: GQLId
+    """The ID of the Slack integration that will be used to send the notification."""
+
+    # Note: Validation aliases are meant to provide continuity with prior `wandb.alert()` API.
+    title: str = ""
+    """The title of the sent notification."""
+
+    message: Annotated[str, Field(validation_alias="text")] = ""
+    """The message body of the sent notification."""
+
+    severity: Annotated[
+        AlertSeverity,
+        BeforeValidator(upper_if_str),  # Be helpful by ensuring uppercase strings
+        Field(validation_alias="level"),
+    ] = AlertSeverity.INFO
+    """The severity (`INFO`, `WARN`, `ERROR`) of the sent notification."""
+
+    @classmethod
+    def from_integration(
+        cls,
+        integration: SlackIntegration,
+        *,
+        title: str = "",
+        text: str = "",
+        level: AlertSeverity = AlertSeverity.INFO,
+    ) -> Self:
+        """Define a notification action that sends to the given (Slack) integration."""
+        return cls(
+            integration_id=integration.id,
+            title=title,
+            message=text,
+            severity=level,
+        )
+
+
+class SendWebhook(_BaseActionInput, GenericWebhookActionInput):
+    """Defines an automation action that sends a webhook request."""
+
+    action_type: Literal[ActionType.GENERIC_WEBHOOK] = ActionType.GENERIC_WEBHOOK
+
+    integration_id: GQLId
+    """The ID of the webhook integration that will be used to send the request."""
+
+    # overrides the generated field type to parse/serialize JSON strings
+    request_payload: Optional[SerializedToJson[dict[str, Any]]] = Field(  # type: ignore[assignment]
+        default=None, alias="requestPayload"
+    )
+    """The payload, possibly with template variables, to send in the webhook request."""
+
+    @classmethod
+    def from_integration(
+        cls,
+        integration: WebhookIntegration,
+        *,
+        payload: Optional[SerializedToJson[dict[str, Any]]] = None,
+    ) -> Self:
+        """Define a webhook action that sends to the given (webhook) integration."""
+        return cls(integration_id=integration.id, request_payload=payload)
+
+
+class DoNothing(_BaseActionInput, NoOpTriggeredActionInput, frozen=True):
+    """Defines an automation action that intentionally does nothing."""
+
+    action_type: Literal[ActionType.NO_OP] = ActionType.NO_OP
+
+    no_op: Annotated[bool, BeforeValidator(default_if_none)] = True
+    """Placeholder field which exists only to satisfy backend schema requirements.
+
+    There should never be a need to set this field explicitly, as its value is ignored.
+    """
+
+
+# for type annotations
+InputAction = Annotated[
+    Union[
+        SendNotification,
+        SendWebhook,
+        DoNothing,
+    ],
+    BeforeValidator(to_input_action),
+    Field(discriminator="action_type"),
+]
+# for runtime type checks
+InputActionTypes: tuple[type, ...] = get_args(InputAction.__origin__)  # type: ignore[attr-defined]
+
+__all__ = [
+    "ActionType",
+    *(nameof(cls) for cls in InputActionTypes),
+]