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

wandb/automations/_filters/expressions.py

@@ -0,0 +1,179 @@
+"""Pydantic-compatible representations of MongoDB expressions."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any, Union
+
+from pydantic import ConfigDict, model_serializer
+from typing_extensions import Self, TypeAlias
+
+from wandb._pydantic import CompatBaseModel, model_validator
+
+from .operators import (
+    Contains,
+    Eq,
+    Exists,
+    Gt,
+    Gte,
+    In,
+    Lt,
+    Lte,
+    Ne,
+    NotIn,
+    Op,
+    Regex,
+    RichReprResult,
+    Scalar,
+    ScalarTypes,
+    SupportsLogicalOpSyntax,
+)
+
+
+class FilterableField:
+    """A descriptor that can be used to define a "filterable" field on a class.
+
+    Internal helper to support syntactic sugar for defining event filters.
+    """
+
+    _python_name: str  #: The name of the field this descriptor was assigned to in the Python class.
+    _server_name: str | None  #: If set, the actual server-side field name to filter on.
+
+    def __init__(self, server_name: str | None = None):
+        self._server_name = server_name
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        self._python_name = name
+
+    def __get__(self, obj: Any, objtype: type) -> Self:
+        # By default, if we didn't explicitly provide a backend name for
+        # filtering, assume the field has the same name in the backend as
+        # the python attribute.
+        return self
+
+    @property
+    def _name(self) -> str:
+        return self._server_name or self._python_name
+
+    def __str__(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self._name!r})"
+
+    # Methods to define filter expressions through chaining
+    def matches_regex(self, pattern: str) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Regex(regex_=pattern))
+
+    def contains(self, text: str) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Contains(contains_=text))
+
+    def exists(self, exists: bool = True) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Exists(exists_=exists))
+
+    def lt(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Lt(lt_=value))
+
+    def gt(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Gt(gt_=value))
+
+    def lte(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Lte(lte_=value))
+
+    def gte(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Gte(gte_=value))
+
+    def ne(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Ne(ne_=value))
+
+    def eq(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Eq(eq_=value))
+
+    def in_(self, values: Iterable[Scalar]) -> FilterExpr:
+        return FilterExpr(field=self._name, op=In(in_=values))
+
+    def not_in(self, values: Iterable[Scalar]) -> FilterExpr:
+        return FilterExpr(field=self._name, op=NotIn(nin_=values))
+
+    # Override the default behavior of comparison operators: <, >=, ==, etc
+    def __lt__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.lt(other)
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __gt__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.gt(other)
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __le__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.lte(other)
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __ge__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.gte(other)
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    # Operator behavior is intentionally overridden to allow defining
+    # filter expressions like `field == "value"`.  See similar overrides
+    # of built-in dunder methods in sqlalchemy, polars, pandas, numpy, etc.
+    #
+    # sqlalchemy example for illustrative purposes:
+    # https://github.com/sqlalchemy/sqlalchemy/blob/f21ae633486380a26dc0b67b70ae1c0efc6b4dc4/lib/sqlalchemy/orm/descriptor_props.py#L808-L812
+    def __eq__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.eq(other)
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __ne__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.ne(other)
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+
+# ------------------------------------------------------------------------------
+class FilterExpr(CompatBaseModel, SupportsLogicalOpSyntax):
+    """A MongoDB filter expression on a specific field."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+    field: str
+    op: Op
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.field!s}={self.op!r})"
+
+    def __rich_repr__(self) -> RichReprResult:  # type: ignore[override]
+        # https://rich.readthedocs.io/en/stable/pretty.html
+        yield self.field, self.op
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate(cls, data: Any) -> Any:
+        """Parse a MongoDB dict representation of the filter expression."""
+        if (
+            isinstance(data, dict)
+            and len(data) == 1
+            and not any(key.startswith("$") for key in data)
+        ):
+            # This looks like a MongoDB filter dict.  E.g.:
+            # - in:  `{"display_name": {"$contains": "my-run"}}`
+            # - out: `FilterExpr(field="display_name", op=Contains(contains_="my-run"))`
+            ((field, op),) = data.items()
+            return {"field": field, "op": op}
+        return data
+
+    @model_serializer(mode="plain")
+    def _serialize(self) -> dict[str, Any]:
+        """Return a MongoDB dict representation of the expression."""
+        from pydantic_core import to_jsonable_python  # Only valid in pydantic v2
+
+        op_dict = to_jsonable_python(self.op, by_alias=True, round_trip=True)
+        return {self.field: op_dict}
+
+
+MongoLikeFilter: TypeAlias = Union[Op, FilterExpr]

wandb/automations/_filters/operators.py

@@ -0,0 +1,267 @@
+"""Types that represent operators in MongoDB filter expressions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Tuple, TypeVar, Union, overload
+
+from pydantic import ConfigDict, Field, StrictBool, StrictFloat, StrictInt, StrictStr
+from typing_extensions import TypeAlias, get_args
+
+from wandb._pydantic import Base
+
+if TYPE_CHECKING:
+    from wandb.automations._filters.run_metrics import MetricThresholdFilter
+    from wandb.automations.events import RunMetricFilter
+
+
+# for type annotations
+Scalar = Union[StrictStr, StrictInt, StrictFloat, StrictBool]
+# for runtime type checks
+ScalarTypes = get_args(Scalar)
+
+# See: https://rich.readthedocs.io/en/stable/pretty.html#rich-repr-protocol
+RichReprResult: TypeAlias = Iterable[
+    Union[
+        Any,
+        Tuple[Any],
+        Tuple[str, Any],
+        Tuple[str, Any, Any],
+    ]
+]
+
+T = TypeVar("T")
+TupleOf: TypeAlias = Tuple[T, ...]
+
+
+# NOTE: Wherever class descriptions that are not docstrings, this is deliberate.
+# This is done to ensure the descriptions are omitted from generated API docs.
+
+
+# Mixin to support syntactic sugar for MongoDB expressions with (bitwise) logical operators,
+# e.g. `a | b` -> `{"$or": [a, b]}` or `~a` -> `{"$not": a}`.
+class SupportsLogicalOpSyntax:
+    def __or__(self, other: Any) -> Or:
+        """Syntactic sugar for: `a | b` -> `Or(a, b)`."""
+        return Or(or_=[self, other])
+
+    @overload
+    def __and__(self, other: MetricThresholdFilter) -> RunMetricFilter: ...
+    @overload
+    def __and__(self, other: Any) -> And: ...
+    def __and__(self, other: Any) -> Any:
+        """Syntactic sugar for: `a & b` -> `And(a, b)`."""
+        from wandb.automations._filters.run_metrics import MetricThresholdFilter
+
+        # Special handling `run_filter & metric_filter`
+        if isinstance(other, MetricThresholdFilter):
+            return other.__and__(self)
+        return And(and_=[self, other])
+
+    def __invert__(self) -> Not:
+        """Syntactic sugar for: `~a` -> `Not(a)`."""
+        return Not(not_=self)
+
+
+# Base class for parsed MongoDB filter/query operators, e.g. `{"$and": [...]}`.
+class BaseOp(Base, SupportsLogicalOpSyntax):
+    model_config = ConfigDict(
+        frozen=True,  # Make pseudo-immutable for easier comparison and hashing
+    )
+
+    def __repr__(self) -> str:
+        # Display operand as a positional arg
+        values_repr = ", ".join(map(repr, self.model_dump().values()))
+        return f"{type(self).__name__}({values_repr})"
+
+    def __rich_repr__(self) -> RichReprResult:  # type: ignore[override]
+        # Display field values as positional args:
+        # https://rich.readthedocs.io/en/stable/pretty.html
+        yield from ((None, v) for v in self.model_dump().values())
+
+
+# Logical operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/and/
+# https://www.mongodb.com/docs/manual/reference/operator/query/or/
+# https://www.mongodb.com/docs/manual/reference/operator/query/nor/
+# https://www.mongodb.com/docs/manual/reference/operator/query/not/
+class And(BaseOp):
+    and_: TupleOf[Any] = Field(default=(), alias="$and")
+
+
+class Or(BaseOp):
+    or_: TupleOf[Any] = Field(default=(), alias="$or")
+
+    def __invert__(self) -> Nor:
+        """Syntactic sugar for: `~Or(a, b)` -> `Nor(a, b)`."""
+        return Nor(nor_=self.or_)
+
+
+class Nor(BaseOp):
+    nor_: TupleOf[Any] = Field(default=(), alias="$nor")
+
+    def __invert__(self) -> Or:
+        """Syntactic sugar for: `~Nor(a, b)` -> `Or(a, b)`."""
+        return Or(or_=self.nor_)
+
+
+class Not(BaseOp):
+    not_: Any = Field(alias="$not")
+
+    def __invert__(self) -> Any:
+        """Syntactic sugar for: `~Not(a)` -> `a`."""
+        return self.not_
+
+
+# Comparison operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/lt/
+# https://www.mongodb.com/docs/manual/reference/operator/query/gt/
+# https://www.mongodb.com/docs/manual/reference/operator/query/lte/
+# https://www.mongodb.com/docs/manual/reference/operator/query/gte/
+# https://www.mongodb.com/docs/manual/reference/operator/query/eq/
+# https://www.mongodb.com/docs/manual/reference/operator/query/ne/
+# https://www.mongodb.com/docs/manual/reference/operator/query/in/
+# https://www.mongodb.com/docs/manual/reference/operator/query/nin/
+class Lt(BaseOp):
+    lt_: Scalar = Field(alias="$lt")
+
+    def __invert__(self) -> Gte:
+        """Syntactic sugar for: `~Lt(a)` -> `Gte(a)`."""
+        return Gte(gte_=self.lt_)
+
+
+class Gt(BaseOp):
+    gt_: Scalar = Field(alias="$gt")
+
+    def __invert__(self) -> Lte:
+        """Syntactic sugar for: `~Gt(a)` -> `Lte(a)`."""
+        return Lte(lte_=self.gt_)
+
+
+class Lte(BaseOp):
+    lte_: Scalar = Field(alias="$lte")
+
+    def __invert__(self) -> Gt:
+        """Syntactic sugar for: `~Lte(a)` -> `Gt(a)`."""
+        return Gt(gt_=self.lte_)
+
+
+class Gte(BaseOp):
+    gte_: Scalar = Field(alias="$gte")
+
+    def __invert__(self) -> Lt:
+        """Syntactic sugar for: `~Gte(a)` -> `Lt(a)`."""
+        return Lt(lt_=self.gte_)
+
+
+class Eq(BaseOp):
+    eq_: Scalar = Field(alias="$eq")
+
+    def __invert__(self) -> Ne:
+        """Syntactic sugar for: `~Eq(a)` -> `Ne(a)`."""
+        return Ne(ne_=self.eq_)
+
+
+class Ne(BaseOp):
+    ne_: Scalar = Field(alias="$ne")
+
+    def __invert__(self) -> Eq:
+        """Syntactic sugar for: `~Ne(a)` -> `Eq(a)`."""
+        return Eq(eq_=self.ne_)
+
+
+class In(BaseOp):
+    in_: TupleOf[Scalar] = Field(default=(), alias="$in")
+
+    def __invert__(self) -> NotIn:
+        """Syntactic sugar for: `~In(a)` -> `NotIn(a)`."""
+        return NotIn(nin_=self.in_)
+
+
+class NotIn(BaseOp):
+    nin_: TupleOf[Scalar] = Field(default=(), alias="$nin")
+
+    def __invert__(self) -> In:
+        """Syntactic sugar for: `~NotIn(a)` -> `In(a)`."""
+        return In(in_=self.nin_)
+
+
+# Element operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/exists/
+class Exists(BaseOp):
+    exists_: bool = Field(alias="$exists")
+
+
+# Evaluation operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/regex/
+#
+# Note: "$contains" is NOT a formal MongoDB operator, but the backend recognizes and
+# executes it as a substring-match filter.
+class Regex(BaseOp):
+    regex_: str = Field(alias="$regex")  #: The regex expression to match against.
+
+
+class Contains(BaseOp):
+    contains_: str = Field(alias="$contains")  #: The substring to match against.
+
+
+And.model_rebuild()
+Or.model_rebuild()
+Not.model_rebuild()
+Lt.model_rebuild()
+Gt.model_rebuild()
+Lte.model_rebuild()
+Gte.model_rebuild()
+Eq.model_rebuild()
+Ne.model_rebuild()
+In.model_rebuild()
+NotIn.model_rebuild()
+Exists.model_rebuild()
+Regex.model_rebuild()
+Contains.model_rebuild()
+
+
+# ------------------------------------------------------------------------------
+# Convenience helpers, constants, and utils for supported MongoDB operators
+# ------------------------------------------------------------------------------
+KEY_TO_OP: dict[str, type[BaseOp]] = {
+    "$and": And,
+    "$or": Or,
+    "$nor": Nor,
+    "$not": Not,
+    "$lt": Lt,
+    "$gt": Gt,
+    "$lte": Lte,
+    "$gte": Gte,
+    "$eq": Eq,
+    "$ne": Ne,
+    "$in": In,
+    "$nin": NotIn,
+    "$exists": Exists,
+    "$regex": Regex,
+    "$contains": Contains,
+}
+
+
+KnownOp = Union[
+    And,
+    Or,
+    Nor,
+    Not,
+    Lt,
+    Gt,
+    Lte,
+    Gte,
+    Eq,
+    Ne,
+    In,
+    NotIn,
+    Exists,
+    Regex,
+    Contains,
+]
+UnknownOp = Dict[str, Any]
+
+# for type annotations
+Op = Union[KnownOp, UnknownOp]
+# for runtime type checks
+OpTypes: tuple[type, ...] = (*get_args(KnownOp), dict)

wandb/automations/_filters/run_metrics.py

@@ -0,0 +1,183 @@
+# ruff: noqa: UP007
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, overload
+
+from pydantic import Field, PositiveInt, StrictFloat, StrictInt, field_validator
+from typing_extensions import Self, override
+
+from wandb._pydantic.base import Base, GQLBase
+
+from .expressions import FilterExpr
+from .operators import BaseOp, RichReprResult
+
+if TYPE_CHECKING:
+    from wandb.automations.events import RunMetricFilter
+
+# Maps MongoDB comparison operators -> Python literal (str) representations
+MONGO2PY_OPS: Final[dict[str, str]] = {
+    "$eq": "==",
+    "$ne": "!=",
+    "$gt": ">",
+    "$lt": "<",
+    "$gte": ">=",
+    "$lte": "<=",
+}
+# Reverse mapping from Python literal (str) -> MongoDB operator key
+PY2MONGO_OPS: Final[dict[str, str]] = {v: k for k, v in MONGO2PY_OPS.items()}
+
+
+class Agg(str, Enum):  # from `Aggregation`
+    """Supported run metric aggregation operations."""
+
+    MAX = "MAX"
+    MIN = "MIN"
+    AVERAGE = "AVERAGE"
+
+
+class ChangeType(str, Enum):  # from `RunMetricChangeType`
+    """Describes the metric change as absolute (arithmetic difference) or relative (decimal percentage)."""
+
+    ABSOLUTE = "ABSOLUTE"
+    RELATIVE = "RELATIVE"
+
+
+class ChangeDirection(str, Enum):  # from `RunMetricChangeDirection`
+    """Describes the direction of the metric change."""
+
+    INCREASE = "INCREASE"
+    DECREASE = "DECREASE"
+    ANY = "ANY"
+
+
+class _BaseMetricFilter(GQLBase):
+    name: str
+    """Name of the observed metric."""
+
+    agg: Optional[Agg]
+    """Aggregation operation, if any, to apply over the window size."""
+
+    window: PositiveInt
+    """Size of the window over which the metric is aggregated."""
+
+    # ------------------------------------------------------------------------------
+
+    threshold: Union[StrictInt, StrictFloat]
+    """Threshold value to compare against."""
+
+    @field_validator("agg", mode="before")
+    @classmethod
+    def _validate_agg(cls, v: Any) -> Any:
+        # Be helpful: e.g. "min" -> "MIN"
+        return v.strip().upper() if isinstance(v, str) else v
+
+    @overload
+    def __and__(self, other: BaseOp | FilterExpr) -> RunMetricFilter: ...
+    @overload
+    def __and__(self, other: Any) -> Any: ...
+    def __and__(self, other: BaseOp | FilterExpr | Any) -> RunMetricFilter | Any:
+        """Supports syntactic sugar for defining a triggering RunMetricEvent from `run_metric_filter & run_filter`."""
+        from wandb.automations.events import RunMetricFilter, _InnerRunMetricFilter
+
+        if isinstance(run_filter := other, (BaseOp, FilterExpr)):
+            # Assume `other` is a run filter, and we are building a RunMetricEvent.
+            # For the metric filter, delegate to the inner validator(s) to further wrap/nest as appropriate.
+            metric_filter = _InnerRunMetricFilter.model_validate(self)
+            return RunMetricFilter(
+                run_metric_filter=metric_filter, run_filter=run_filter
+            )
+        return other.__and__(self)  # Try switching the order of operands
+
+
+class MetricThresholdFilter(_BaseMetricFilter):  # from `RunMetricThresholdFilter`
+    """For run events, defines a metric filter comparing a metric against a user-defined threshold value."""
+
+    name: str
+    agg: Optional[Agg] = Field(default=None, alias="agg_op")
+    window: PositiveInt = Field(default=1, alias="window_size")
+
+    cmp: Literal["$gte", "$gt", "$lt", "$lte"] = Field(alias="cmp_op")
+    """Comparison operator used to compare the metric value (left) vs. the threshold value (right)."""
+
+    threshold: Union[StrictInt, StrictFloat]
+
+    @field_validator("cmp", mode="before")
+    @classmethod
+    def _validate_cmp(cls, v: Any) -> Any:
+        # Be helpful: e.g. ">" -> "$gt"
+        return PY2MONGO_OPS.get(v.strip(), v) if isinstance(v, str) else v
+
+    def __repr__(self) -> str:
+        metric = f"{self.agg.value}({self.name})" if self.agg else self.name
+        op = MONGO2PY_OPS.get(self.cmp, self.cmp)
+        expr = rf"{metric} {op} {self.threshold}"
+        return repr(expr)
+
+    @override
+    def __rich_repr__(self) -> RichReprResult:  # type: ignore[override]
+        yield None, repr(self)
+
+
+class MetricChangeFilter(_BaseMetricFilter):  # from `RunMetricChangeFilter`
+    # FIXME:
+    # - `prior_window` should be optional and default to `window` if not provided.
+    # - implement declarative syntax for `MetricChangeFilter` similar to `MetricThresholdFilter`.
+    # - split this into tagged union of relative/absolute change filters.
+
+    name: str
+    agg: Optional[Agg] = Field(default=None, alias="agg_op")
+
+    # FIXME: Set the `prior_window` to `window` if it's not provided, for convenience.
+    window: PositiveInt = Field(alias="current_window_size")
+    prior_window: PositiveInt = Field(alias="prior_window_size")
+    """Size of the preceding window over which the metric is aggregated."""
+
+    # NOTE: `cmp_op` isn't a field here.  In the backend, it's effectively `cmp_op` = "$gte"
+
+    change_type: ChangeType = Field(alias="change_type")
+    change_direction: ChangeDirection = Field(alias="change_dir")
+
+    threshold: Union[StrictInt, StrictFloat] = Field(alias="change_amount")
+
+
+class MetricOperand(Base):
+    name: str
+    agg: Optional[Agg] = Field(default=None, alias="agg_op")
+    window: PositiveInt = Field(default=1, alias="window_size")
+
+    def _agg(self, op: Agg, window: int) -> Self:
+        if self.agg is None:  # Prevent overwriting an existing aggregation operator
+            return self.model_copy(update={"agg": op, "window": window})
+        raise ValueError(f"Aggregation operator already set as: {self.agg!r}")
+
+    def max(self, window: int) -> Self:
+        return self._agg(Agg.MAX, window)
+
+    def min(self, window: int) -> Self:
+        return self._agg(Agg.MIN, window)
+
+    def average(self, window: int) -> Self:
+        return self._agg(Agg.AVERAGE, window)
+
+    # Aliased method for users familiar with e.g. torch/tf/numpy/pandas/polars/etc.
+    def mean(self, window: int) -> Self:
+        return self.average(window=window)
+
+    def gt(self, other: int | float) -> MetricThresholdFilter:
+        return MetricThresholdFilter(**dict(self), cmp="$gt", threshold=other)
+
+    def lt(self, other: int | float) -> MetricThresholdFilter:
+        return MetricThresholdFilter(**dict(self), cmp="$lt", threshold=other)
+
+    def gte(self, other: int | float) -> MetricThresholdFilter:
+        return MetricThresholdFilter(**dict(self), cmp="$gte", threshold=other)
+
+    def lte(self, other: int | float) -> MetricThresholdFilter:
+        return MetricThresholdFilter(**dict(self), cmp="$lte", threshold=other)
+
+    __gt__ = gt
+    __lt__ = lt
+    __ge__ = gte
+    __le__ = lte