mlrun 1.7.2rc4__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/common/schemas/alert.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from collections import defaultdict
+from collections.abc import Iterator
 from datetime import datetime
-from typing import Annotated, Optional, Union
+from typing import Annotated, Any, Callable, Optional, Union
 
-import pydantic
+import pydantic.v1
 
-from mlrun.common.schemas.notification import Notification
+import mlrun.common.schemas.notification as notification_objects
 from mlrun.common.types import StrEnum
 
 
@@ -27,10 +29,10 @@ class EventEntityKind(StrEnum):
     JOB = "job"
 
 
-class EventEntities(pydantic.BaseModel):
+class EventEntities(pydantic.v1.BaseModel):
     kind: EventEntityKind
     project: str
-    ids: pydantic.conlist(str, min_items=1, max_items=1)
+    ids: pydantic.v1.conlist(str, min_items=1, max_items=1)
 
 
 class EventKind(StrEnum):
@@ -64,11 +66,11 @@ _event_kind_entity_map = {
 }
 
 
-class Event(pydantic.BaseModel):
+class Event(pydantic.v1.BaseModel):
     kind: EventKind
     timestamp: Union[str, datetime] = None  # occurrence time
     entity: EventEntities
-    value_dict: Optional[dict] = pydantic.Field(default_factory=dict)
+    value_dict: Optional[dict] = pydantic.v1.Field(default_factory=dict)
 
     def is_valid(self):
         return self.entity.kind in _event_kind_entity_map[self.kind]
@@ -86,7 +88,7 @@ class AlertSeverity(StrEnum):
 
 
 # what should trigger the alert. must be either event (at least 1), or prometheus query
-class AlertTrigger(pydantic.BaseModel):
+class AlertTrigger(pydantic.v1.BaseModel):
     events: list[EventKind] = []
     prometheus_alert: str = None
 
@@ -97,16 +99,16 @@ class AlertTrigger(pydantic.BaseModel):
         )
 
 
-class AlertCriteria(pydantic.BaseModel):
+class AlertCriteria(pydantic.v1.BaseModel):
     count: Annotated[
         int,
-        pydantic.Field(
+        pydantic.v1.Field(
             description="Number of events to wait until notification is sent"
         ),
     ] = 1
     period: Annotated[
         str,
-        pydantic.Field(
+        pydantic.v1.Field(
             description="Time period during which event occurred. e.g. 1d, 3h, 5m, 15s"
         ),
     ] = None
@@ -120,11 +122,11 @@ class ResetPolicy(StrEnum):
     AUTO = "auto"
 
 
-class AlertNotification(pydantic.BaseModel):
-    notification: Notification
+class AlertNotification(pydantic.v1.BaseModel):
+    notification: notification_objects.Notification
     cooldown_period: Annotated[
         str,
-        pydantic.Field(
+        pydantic.v1.Field(
             description="Period during which notifications "
             "will not be sent after initial send. The format of this would be in time."
             " e.g. 1d, 3h, 5m, 15s"
@@ -132,14 +134,14 @@ class AlertNotification(pydantic.BaseModel):
     ] = None
 
 
-class AlertConfig(pydantic.BaseModel):
+class AlertConfig(pydantic.v1.BaseModel):
     project: str
     id: int = None
     name: str
     description: Optional[str] = ""
     summary: Annotated[
         str,
-        pydantic.Field(
+        pydantic.v1.Field(
             description=(
                 "String to be sent in the notifications generated."
                 "e.g. 'Model {{project}}/{{entity}} is drifting.'"
@@ -153,11 +155,15 @@ class AlertConfig(pydantic.BaseModel):
     trigger: AlertTrigger
     criteria: Optional[AlertCriteria]
     reset_policy: ResetPolicy = ResetPolicy.AUTO
-    notifications: pydantic.conlist(AlertNotification, min_items=1)
+    notifications: pydantic.v1.conlist(AlertNotification, min_items=1)
     state: AlertActiveState = AlertActiveState.INACTIVE
     count: Optional[int] = 0
+    updated: datetime = None
 
-    def get_raw_notifications(self) -> list[Notification]:
+    class Config:
+        extra = pydantic.v1.Extra.allow
+
+    def get_raw_notifications(self) -> list[notification_objects.Notification]:
         return [
             alert_notification.notification for alert_notification in self.notifications
         ]
@@ -169,7 +175,7 @@ class AlertsModes(StrEnum):
 
 
 class AlertTemplate(
-    pydantic.BaseModel
+    pydantic.v1.BaseModel
 ):  # Template fields that are not shared with created configs
     template_id: int = None
     template_name: str
@@ -200,3 +206,92 @@ class AlertTemplate(
             or self.reset_policy != other.reset_policy
             or self.criteria != other.criteria
         )
+
+
+class AlertActivation(pydantic.v1.BaseModel):
+    id: int
+    name: str
+    project: str
+    severity: AlertSeverity
+    activation_time: datetime
+    entity_id: str
+    entity_kind: EventEntityKind
+    criteria: AlertCriteria
+    event_kind: EventKind
+    number_of_events: int
+    notifications: list[notification_objects.NotificationState]
+    reset_time: Optional[datetime] = None
+
+    def group_key(self, attributes: list[str]) -> Union[Any, tuple]:
+        """
+        Dynamically create a key for grouping based on the provided attributes.
+        - If there's only one attribute, return the value directly (not a single-element tuple).
+        - If there are multiple attributes, return them as a tuple for grouping.
+
+        This ensures grouping behaves intuitively without redundant tuple representations.
+        """
+        if len(attributes) == 1:
+            # Avoid single-element tuple like (high,) when only one grouping attribute is used
+            return getattr(self, attributes[0])
+        # Otherwise, return a tuple of all specified attributes
+        return tuple(getattr(self, attr) for attr in attributes)
+
+
+class AlertActivations(pydantic.v1.BaseModel):
+    activations: list[AlertActivation]
+    pagination: Optional[dict]
+
+    def __iter__(self) -> Iterator[AlertActivation]:
+        return iter(self.activations)
+
+    def __getitem__(self, index: int) -> AlertActivation:
+        return self.activations[index]
+
+    def __len__(self) -> int:
+        return len(self.activations)
+
+    def group_by(self, *attributes: str) -> dict:
+        """
+        Group alert activations by specified attributes.
+
+        Args:
+        :param attributes: Attributes to group by.
+
+        :returns: A dictionary where keys are tuples of attribute values and values are lists of
+            AlertActivation objects.
+
+        Example:
+            # Group by project and severity
+            grouped = activations.group_by("project", "severity")
+        """
+        grouped = defaultdict(list)
+        for activation in self.activations:
+            key = activation.group_key(attributes)
+            grouped[key].append(activation)
+        return dict(grouped)
+
+    def aggregate_by(
+        self,
+        group_by_attrs: list[str],
+        aggregation_function: Callable[[list[AlertActivation]], Any],
+    ) -> dict:
+        """
+        Aggregate alert activations by specified attributes using a given aggregation function.
+
+        Args:
+        :param group_by_attrs: Attributes to group by.
+        :param aggregation_function: Function to aggregate grouped activations.
+
+        :returns: A dictionary where keys are tuples of attribute values and values are the result
+            of the aggregation function.
+
+        Example:
+            # Aggregate by name and entity_id and count number of activations in each group
+            activations.aggregate_by(["name", "entity_id"], lambda activations: len(activations))
+        """
+        grouped = self.group_by(*group_by_attrs)
+        aggregated = {
+            key: aggregation_function(activations)
+            for key, activations in grouped.items()
+        }
+        return aggregated

mlrun/common/schemas/api_gateway.py

@@ -15,7 +15,7 @@
 import typing
 from typing import Optional
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.constants as mlrun_constants
 import mlrun.common.types
@@ -49,9 +49,9 @@ class APIGatewayState(mlrun.common.types.StrEnum):
     waiting_for_provisioning = "waitingForProvisioning"
 
 
-class _APIGatewayBaseModel(pydantic.BaseModel):
+class _APIGatewayBaseModel(pydantic.v1.BaseModel):
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class APIGatewayMetadata(_APIGatewayBaseModel):

mlrun/common/schemas/artifact.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 from deprecated import deprecated
 
 import mlrun.common.types
@@ -25,6 +25,7 @@ from .object import ObjectStatus
 class ArtifactCategories(mlrun.common.types.StrEnum):
     model = "model"
     dataset = "dataset"
+    document = "document"
     other = "other"
 
     # we define the link as a category to prevent import cycles, but it's not a real category
@@ -38,23 +39,35 @@ class ArtifactCategories(mlrun.common.types.StrEnum):
             return [ArtifactCategories.model.value, link_kind], False
         if self.value == ArtifactCategories.dataset.value:
             return [ArtifactCategories.dataset.value, link_kind], False
+        if self.value == ArtifactCategories.document.value:
+            return [ArtifactCategories.document.value, link_kind], False
         if self.value == ArtifactCategories.other.value:
             return (
                 [
                     ArtifactCategories.model.value,
                     ArtifactCategories.dataset.value,
+                    ArtifactCategories.document.value,
                 ],
                 True,
             )
 
     @classmethod
     def from_kind(cls, kind: str) -> "ArtifactCategories":
-        if kind in [cls.model.value, cls.dataset.value]:
+        if kind in [cls.model.value, cls.dataset.value, cls.document.value]:
             return cls(kind)
         return cls.other
 
+    @staticmethod
+    def all():
+        """Return all applicable artifact categories"""
+        return [
+            ArtifactCategories.model,
+            ArtifactCategories.dataset,
+            ArtifactCategories.document,
+        ]
 
-class ArtifactIdentifier(pydantic.BaseModel):
+
+class ArtifactIdentifier(pydantic.v1.BaseModel):
     # artifact kind
     kind: typing.Optional[str]
     key: typing.Optional[str]
@@ -67,7 +80,7 @@ class ArtifactIdentifier(pydantic.BaseModel):
 
 @deprecated(
     version="1.7.0",
-    reason="mlrun.common.schemas.ArtifactsFormat is deprecated and will be removed in 1.9.0. "
+    reason="mlrun.common.schemas.ArtifactsFormat is deprecated and will be removed in 1.10.0. "
     "Use mlrun.common.formatters.ArtifactFormat instead.",
     category=FutureWarning,
 )
@@ -75,7 +88,7 @@ class ArtifactsFormat(mlrun.common.types.StrEnum):
     full = "full"
 
 
-class ArtifactMetadata(pydantic.BaseModel):
+class ArtifactMetadata(pydantic.v1.BaseModel):
     key: str
     project: str
     iter: typing.Optional[int]
@@ -83,10 +96,10 @@ class ArtifactMetadata(pydantic.BaseModel):
     tag: typing.Optional[str]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class ArtifactSpec(pydantic.BaseModel):
+class ArtifactSpec(pydantic.v1.BaseModel):
     src_path: typing.Optional[str]
     target_path: typing.Optional[str]
     viewer: typing.Optional[str]
@@ -97,10 +110,10 @@ class ArtifactSpec(pydantic.BaseModel):
     unpackaging_instructions: typing.Optional[dict[str, typing.Any]]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class Artifact(pydantic.BaseModel):
+class Artifact(pydantic.v1.BaseModel):
     kind: str
     metadata: ArtifactMetadata
     spec: ArtifactSpec

mlrun/common/schemas/auth.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 from nuclio.auth import AuthInfo as NuclioAuthInfo
 from nuclio.auth import AuthKinds as NuclioAuthKinds
 
@@ -59,6 +59,7 @@ class AuthorizationResourceTypes(mlrun.common.types.StrEnum):
     hub_source = "hub-source"
     workflow = "workflow"
     alert = "alert"
+    alert_activations = "alert-activations"
     alert_templates = "alert-templates"
     event = "event"
     datastore_profile = "datastore-profile"
@@ -90,6 +91,7 @@ class AuthorizationResourceTypes(mlrun.common.types.StrEnum):
             AuthorizationResourceTypes.run: "/projects/{project_name}/runs/{resource_name}",
             AuthorizationResourceTypes.event: "/projects/{project_name}/events/{resource_name}",
             AuthorizationResourceTypes.alert: "/projects/{project_name}/alerts/{resource_name}",
+            AuthorizationResourceTypes.alert_activations: "/projects/{project_name}/alerts/{resource_name}/activations",
             AuthorizationResourceTypes.alert_templates: "/alert-templates/{resource_name}",
             # runtime resource doesn't have an identifier, we don't need any auth granularity behind project level
             AuthorizationResourceTypes.runtime_resource: "/projects/{project_name}/runtime-resources",
@@ -106,12 +108,14 @@ class AuthorizationResourceTypes(mlrun.common.types.StrEnum):
         }[self].format(project_name=project_name, resource_name=resource_name)
 
 
-class AuthorizationVerificationInput(pydantic.BaseModel):
+class AuthorizationVerificationInput(pydantic.v1.BaseModel):
     resource: str
     action: AuthorizationAction
 
 
-class AuthInfo(pydantic.BaseModel):
+class AuthInfo(pydantic.v1.BaseModel):
+    # Keep request headers for inter-service communication
+    request_headers: typing.Optional[dict[str, str]] = None
     # Basic + Iguazio auth
     username: typing.Optional[str] = None
     # Basic auth
@@ -145,5 +149,5 @@ class AuthInfo(pydantic.BaseModel):
         return self.data_session or self.session
 
 
-class Credentials(pydantic.BaseModel):
+class Credentials(pydantic.v1.BaseModel):
     access_key: typing.Optional[str]

mlrun/common/schemas/background_task.py

@@ -15,7 +15,7 @@
 import datetime
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
@@ -35,7 +35,7 @@ class BackgroundTaskState(mlrun.common.types.StrEnum):
         ]
 
 
-class BackgroundTaskMetadata(pydantic.BaseModel):
+class BackgroundTaskMetadata(pydantic.v1.BaseModel):
     name: str
     kind: typing.Optional[str]
     project: typing.Optional[str]
@@ -44,21 +44,21 @@ class BackgroundTaskMetadata(pydantic.BaseModel):
     timeout: typing.Optional[int]
 
 
-class BackgroundTaskSpec(pydantic.BaseModel):
+class BackgroundTaskSpec(pydantic.v1.BaseModel):
     pass
 
 
-class BackgroundTaskStatus(pydantic.BaseModel):
+class BackgroundTaskStatus(pydantic.v1.BaseModel):
     state: BackgroundTaskState
     error: typing.Optional[str]
 
 
-class BackgroundTask(pydantic.BaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.background_task, const=True)
+class BackgroundTask(pydantic.v1.BaseModel):
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.background_task, const=True)
     metadata: BackgroundTaskMetadata
     spec: BackgroundTaskSpec
     status: BackgroundTaskStatus
 
 
-class BackgroundTaskList(pydantic.BaseModel):
+class BackgroundTaskList(pydantic.v1.BaseModel):
     background_tasks: list[BackgroundTask]

mlrun/common/schemas/client_spec.py

@@ -14,13 +14,13 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 from .function import Function
 from .k8s import Resources
 
 
-class ClientSpec(pydantic.BaseModel):
+class ClientSpec(pydantic.v1.BaseModel):
     version: typing.Optional[str]
     namespace: typing.Optional[str]
     docker_registry: typing.Optional[str]
@@ -57,8 +57,6 @@ class ClientSpec(pydantic.BaseModel):
     redis_url: typing.Optional[str]
     redis_type: typing.Optional[str]
     sql_url: typing.Optional[str]
-    model_endpoint_monitoring_endpoint_store_connection: typing.Optional[str]
-    model_monitoring_tsdb_connection: typing.Optional[str]
     ce: typing.Optional[dict]
     # not passing them as one object as it possible client user would like to override only one of the params
     calculate_artifact_hash: typing.Optional[str]
@@ -67,3 +65,5 @@ class ClientSpec(pydantic.BaseModel):
     packagers: typing.Optional[dict]
     external_platform_tracking: typing.Optional[dict]
     alerts_mode: typing.Optional[str]
+    system_id: typing.Optional[str]
+    model_endpoint_monitoring_store_prefixes: typing.Optional[dict[str, str]]

mlrun/common/schemas/clusterization_spec.py

@@ -14,12 +14,12 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
 
-class ClusterizationSpec(pydantic.BaseModel):
+class ClusterizationSpec(pydantic.v1.BaseModel):
     chief_api_state: typing.Optional[str]
     chief_version: typing.Optional[str]
 

mlrun/common/schemas/common.py

@@ -14,10 +14,12 @@
 
 import typing
 
-import pydantic
+import pydantic.v1
 
+import mlrun.errors
 
-class ImageBuilder(pydantic.BaseModel):
+
+class ImageBuilder(pydantic.v1.BaseModel):
     functionSourceCode: typing.Optional[str] = None  # noqa: N815
     codeEntryType: typing.Optional[str] = None  # noqa: N815
     codeEntryAttributes: typing.Optional[str] = None  # noqa: N815
@@ -40,4 +42,52 @@ class ImageBuilder(pydantic.BaseModel):
     source_code_target_dir: typing.Optional[str] = None
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
+
+
+class LabelsModel(pydantic.v1.BaseModel):
+    """
+    This class accepts either a dictionary, a list, or a string for filtering by labels.
+
+    :param labels:
+        - If a dictionary is provided, it should be in the format {'label_name': 'value'}.
+          The values can also be `None`, which will result in the format 'label_name' (without a value).
+          This will be converted to a list of strings in the format 'label_name=value'.
+        - If a list is provided, all items must be strings. Each string can either
+          be a simple label name (e.g., 'label1') or a key-value pair in the format
+          'label=value'.
+        - If a string is provided, it should be a comma-separated list of labels
+          (e.g., 'label1,label2').
+        - If no labels are specified, the default is an empty list.
+    """
+
+    labels: typing.Optional[
+        typing.Union[str, dict[str, typing.Optional[str]], list[str]]
+    ]
+
+    @pydantic.v1.validator("labels")
+    @classmethod
+    def validate(cls, labels) -> list[str]:
+        if labels is None:
+            return []
+
+        # If labels is a string, split it by commas
+        if isinstance(labels, str):
+            return [label.strip() for label in labels.split(",") if label.strip()]
+
+        if isinstance(labels, list):
+            if not all(isinstance(item, str) for item in labels):
+                raise mlrun.errors.MLRunValueError(
+                    "All items in the list must be strings."
+                )
+            return labels
+
+        if isinstance(labels, dict):
+            return [
+                f"{key}={value}" if value is not None else key
+                for key, value in labels.items()
+            ]
+
+        raise mlrun.errors.MLRunValueError(
+            "Invalid labels format. Must be a string, dictionary of strings, or a list of strings."
+        )

mlrun/common/schemas/constants.py

@@ -133,6 +133,21 @@ class RunPartitionByField(mlrun.common.types.StrEnum):
             )
 
 
+class ArtifactPartitionByField(mlrun.common.types.StrEnum):
+    name = "name"  # Supported for artifacts objects
+    project_and_name = "project_and_name"  # Supported for artifacts objects
+
+    def to_partition_by_db_field(self, db_cls):
+        if self.value == ArtifactPartitionByField.name:
+            return db_cls.key
+        elif self.value == ArtifactPartitionByField.project_and_name:
+            return db_cls.project, db_cls.key
+        else:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Unknown group by field: {self.value}"
+            )
+
+
 class SortField(mlrun.common.types.StrEnum):
     created = "created"
     updated = "updated"

mlrun/common/schemas/datastore_profile.py

@@ -15,7 +15,7 @@
 
 import typing
 
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 
 
 class DatastoreProfile(BaseModel):

mlrun/common/schemas/feature_store.py

@@ -14,7 +14,7 @@
 #
 from typing import Optional
 
-import pydantic
+import pydantic.v1
 
 from .auth import AuthorizationResourceTypes, Credentials
 from .object import (
@@ -27,7 +27,7 @@ from .object import (
 )
 
 
-class FeatureStoreBaseModel(pydantic.BaseModel):
+class FeatureStoreBaseModel(pydantic.v1.BaseModel):
     """
     Intermediate base class, in order to override pydantic's configuration, as per
     https://docs.pydantic.dev/1.10/usage/model_config/#change-behaviour-globally
@@ -43,7 +43,7 @@ class Feature(FeatureStoreBaseModel):
     labels: Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class Entity(FeatureStoreBaseModel):
@@ -52,17 +52,17 @@ class Entity(FeatureStoreBaseModel):
     labels: Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class FeatureSetSpec(ObjectSpec):
     entities: list[Entity] = []
     features: list[Feature] = []
-    engine: Optional[str] = pydantic.Field(default="storey")
+    engine: Optional[str] = pydantic.v1.Field(default="storey")
 
 
 class FeatureSet(FeatureStoreBaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.feature_set, const=True)
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.feature_set, const=True)
     metadata: ObjectMetadata
     spec: FeatureSetSpec
     status: ObjectStatus
@@ -155,7 +155,7 @@ class EntitiesOutput(FeatureStoreBaseModel):
 
 
 class FeatureVector(FeatureStoreBaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.feature_vector, const=True)
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.feature_vector, const=True)
     metadata: ObjectMetadata
     spec: ObjectSpec
     status: ObjectStatus
@@ -183,7 +183,7 @@ class DataSource(FeatureStoreBaseModel):
     path: str
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class DataTarget(FeatureStoreBaseModel):
@@ -192,7 +192,7 @@ class DataTarget(FeatureStoreBaseModel):
     path: Optional[str]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class FeatureSetIngestInput(FeatureStoreBaseModel):

mlrun/common/schemas/frontend_spec.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
@@ -43,20 +43,20 @@ class NuclioStreamsFeatureFlag(mlrun.common.types.StrEnum):
     disabled = "disabled"
 
 
-class FeatureFlags(pydantic.BaseModel):
+class FeatureFlags(pydantic.v1.BaseModel):
     project_membership: ProjectMembershipFeatureFlag
     authentication: AuthenticationFeatureFlag
     nuclio_streams: NuclioStreamsFeatureFlag
     preemption_nodes: PreemptionNodesFeatureFlag
 
 
-class ArtifactLimits(pydantic.BaseModel):
+class ArtifactLimits(pydantic.v1.BaseModel):
     max_chunk_size: int
     max_preview_size: int
     max_download_size: int
 
 
-class FrontendSpec(pydantic.BaseModel):
+class FrontendSpec(pydantic.v1.BaseModel):
     jobs_dashboard_url: typing.Optional[str]
     model_monitoring_dashboard_url: typing.Optional[str]
     abortable_function_kinds: list[str] = []