mlrun 1.8.0rc5__py3-none-any.whl → 1.8.0rc9__py3-none-any.whl

mlrun/__init__.py

@@ -60,6 +60,7 @@ from .run import (
     get_pipeline,
     import_function,
     new_function,
+    retry_pipeline,
     wait_for_pipeline_completion,
 )
 from .runtimes import mounts, new_model_server

mlrun/artifacts/__init__.py

@@ -23,7 +23,7 @@ from .base import (
     get_artifact_meta,
 )
 from .dataset import DatasetArtifact, TableArtifact, update_dataset_meta
-from .document import DocumentArtifact, DocumentLoader, DocumentLoaderSpec
+from .document import DocumentArtifact, DocumentLoaderSpec, MLRunLoader
 from .manager import (
     ArtifactManager,
     ArtifactProducer,

mlrun/artifacts/base.py

@@ -36,7 +36,16 @@ from ..utils import (
 
 
 class ArtifactMetadata(ModelObj):
-    _dict_fields = ["key", "project", "iter", "tree", "description", "hash", "tag"]
+    _dict_fields = [
+        "key",
+        "project",
+        "iter",
+        "tree",
+        "description",
+        "hash",
+        "tag",
+        "uid",
+    ]
     _extra_fields = ["updated", "labels"]
 
     def __init__(
@@ -48,6 +57,7 @@ class ArtifactMetadata(ModelObj):
         description=None,
         hash=None,
         tag=None,
+        uid=None,
     ):
         self.key = key
         self.project = project
@@ -58,6 +68,7 @@ class ArtifactMetadata(ModelObj):
         self.labels = {}
         self.updated = None
         self.tag = tag  # temp store of the tag
+        self.uid = uid
 
     def base_dict(self):
         return super().to_dict()
@@ -368,6 +379,7 @@ class Artifact(ModelObj):
             iter=self.metadata.iter,
             tree=tree,
             tag=tag,
+            uid=self.uid,
         )
         return mlrun.datastore.get_store_uri(self._store_prefix, uri)
 
@@ -642,6 +654,14 @@ class Artifact(ModelObj):
     def hash(self, hash):
         self.metadata.hash = hash
 
+    @property
+    def uid(self):
+        return self.metadata.uid
+
+    @uid.setter
+    def uid(self, uid):
+        self.metadata.uid = uid
+
     def generate_target_path(self, artifact_path, producer):
         return generate_target_path(self, artifact_path, producer)
 

mlrun/artifacts/document.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import ast
 import re
 import tempfile
 from collections.abc import Iterator
@@ -74,14 +73,14 @@ class DocumentLoaderSpec(ModelObj):
         return loader
 
 
-class DocumentLoader:
+class MLRunLoader:
     """
     A factory class for creating instances of a dynamically defined document loader.
 
     Args:
         artifact_key (str): The key for the artifact to be logged.It can include '%%' which will be replaced
         by a hex-encoded version of the source path.
-        source_path (str): The source path of the document to be loaded.
+        local_path (str): The source path of the document to be loaded.
         loader_spec (DocumentLoaderSpec): Specification for the document loader.
         producer (Optional[Union[MlrunProject, str, MLClientCtx]], optional): The producer of the document
         upload (bool, optional): Flag indicating whether to upload the document.
@@ -104,7 +103,7 @@ class DocumentLoader:
         class DynamicDocumentLoader(BaseLoader):
             def __init__(
                 self,
-                source_path,
+                local_path,
                 loader_spec,
                 artifact_key,
                 producer,
@@ -112,12 +111,12 @@ class DocumentLoader:
             ):
                 self.producer = producer
                 self.artifact_key = (
-                    DocumentLoader.artifact_key_instance(artifact_key, source_path)
+                    MLRunLoader.artifact_key_instance(artifact_key, local_path)
                     if "%%" in artifact_key
                     else artifact_key
                 )
                 self.loader_spec = loader_spec
-                self.source_path = source_path
+                self.local_path = local_path
                 self.upload = upload
 
                 # Resolve the producer
@@ -129,16 +128,17 @@ class DocumentLoader:
             def lazy_load(self) -> Iterator["Document"]:  # noqa: F821
                 artifact = self.producer.log_document(
                     key=self.artifact_key,
-                    document_loader=self.loader_spec,
-                    src_path=self.source_path,
+                    document_loader_spec=self.loader_spec,
+                    local_path=self.local_path,
                     upload=self.upload,
                 )
-                yield artifact.to_langchain_documents()
+                res = artifact.to_langchain_documents()
+                yield res[0]
 
         # Return an instance of the dynamically defined subclass
         instance = DynamicDocumentLoader(
             artifact_key=artifact_key,
-            source_path=source_path,
+            local_path=source_path,
             loader_spec=loader_spec,
             producer=producer,
             upload=upload,
@@ -195,12 +195,15 @@ class DocumentArtifact(Artifact):
         def __init__(
             self,
             *args,
+            document_loader: Optional[DocumentLoaderSpec] = None,
+            collections: Optional[dict] = None,
+            original_source: Optional[str] = None,
             **kwargs,
         ):
             super().__init__(*args, **kwargs)
-            self.document_loader = None
-            self.collections = set()
-            self.original_source = None
+            self.document_loader = document_loader
+            self.collections = collections if collections is not None else {}
+            self.original_source = original_source
 
     """
     A specific artifact class inheriting from generic artifact, used to maintain Document meta-data.
@@ -216,14 +219,17 @@ class DocumentArtifact(Artifact):
 
     def __init__(
         self,
-        key=None,
-        document_loader: DocumentLoaderSpec = DocumentLoaderSpec(),
+        original_source: Optional[str] = None,
+        document_loader_spec: Optional[DocumentLoaderSpec] = None,
         **kwargs,
     ):
-        super().__init__(key, **kwargs)
-        self.spec.document_loader = document_loader.to_str()
-        if "src_path" in kwargs:
-            self.spec.original_source = kwargs["src_path"]
+        super().__init__(**kwargs)
+        self.spec.document_loader = (
+            document_loader_spec.to_dict()
+            if document_loader_spec
+            else self.spec.document_loader
+        )
+        self.spec.original_source = original_source or self.spec.original_source
 
     @property
     def spec(self) -> DocumentArtifactSpec:
@@ -234,17 +240,8 @@ class DocumentArtifact(Artifact):
         self._spec = self._verify_dict(
             spec, "spec", DocumentArtifact.DocumentArtifactSpec
         )
-        # _verify_dict doesn't handle set, so we need to convert it back
-        if isinstance(self._spec.collections, str):
-            self._spec.collections = ast.literal_eval(self._spec.collections)
-
-    @property
-    def inputs(self):
-        # To keep the interface consistent with the project.update_artifact() when we update the artifact
-        return None
 
-    @property
-    def source(self):
+    def get_source(self):
         return generate_artifact_uri(self.metadata.project, self.spec.db_key)
 
     def to_langchain_documents(
@@ -262,9 +259,8 @@ class DocumentArtifact(Artifact):
         Returns:
             list[Document]: A list of LangChain Document objects.
         """
-        dictionary = ast.literal_eval(self.spec.document_loader)
-        loader_spec = DocumentLoaderSpec.from_dict(dictionary)
 
+        loader_spec = DocumentLoaderSpec.from_dict(self.spec.document_loader)
         if self.get_target_path():
             with tempfile.NamedTemporaryFile() as tmp_file:
                 mlrun.datastore.store_manager.object(
@@ -272,8 +268,8 @@ class DocumentArtifact(Artifact):
                 ).download(tmp_file.name)
                 loader = loader_spec.make_loader(tmp_file.name)
                 documents = loader.load()
-        elif self.src_path:
-            loader = loader_spec.make_loader(self.src_path)
+        elif self.spec.original_source:
+            loader = loader_spec.make_loader(self.spec.original_source)
             documents = loader.load()
         else:
             raise ValueError(
@@ -281,6 +277,7 @@ class DocumentArtifact(Artifact):
             )
 
         results = []
+        idx = 0
         for document in documents:
             if splitter:
                 texts = splitter.split_text(document.page_content)
@@ -289,25 +286,51 @@ class DocumentArtifact(Artifact):
 
             metadata = document.metadata
 
-            metadata[self.METADATA_ORIGINAL_SOURCE_KEY] = self.src_path
-            metadata[self.METADATA_SOURCE_KEY] = self.source
+            metadata[self.METADATA_ORIGINAL_SOURCE_KEY] = self.spec.original_source
+            metadata[self.METADATA_SOURCE_KEY] = self.get_source()
             metadata[self.METADATA_ARTIFACT_URI_KEY] = self.uri
             if self.get_target_path():
                 metadata[self.METADATA_ARTIFACT_TARGET_PATH_KEY] = (
                     self.get_target_path()
                 )
 
-            for idx, text in enumerate(texts):
+            for text in texts:
                 metadata[self.METADATA_CHUNK_KEY] = str(idx)
                 doc = Document(
                     page_content=text,
-                    metadata=metadata,
+                    metadata=metadata.copy(),
                 )
                 results.append(doc)
+                idx = idx + 1
         return results
 
     def collection_add(self, collection_id: str) -> None:
-        self.spec.collections.add(collection_id)
+        """
+        Add a collection ID to the artifact's collection list.
+
+        Adds the specified collection ID to the artifact's collection mapping if it
+        doesn't already exist.
+        This method only modifies the client-side artifact object and does not persist
+        the changes to the MLRun DB. To save the changes permanently, you must call
+        project.update_artifact() after this method.
+
+        Args:
+            collection_id (str): The ID of the collection to add
+        """
+        if collection_id not in self.spec.collections:
+            self.spec.collections[collection_id] = "1"
 
     def collection_remove(self, collection_id: str) -> None:
-        return self.spec.collections.discard(collection_id)
+        """
+        Remove a collection ID from the artifact's collection list.
+
+        Removes the specified collection ID from the artifact's local collection mapping.
+        This method only modifies the client-side artifact object and does not persist
+        the changes to the MLRun DB. To save the changes permanently, you must call
+        project.update_artifact() or context.update_artifact() after this method.
+
+        Args:
+            collection_id (str): The ID of the collection to remove
+        """
+        if collection_id in self.spec.collections:
+            self.spec.collections.pop(collection_id)

mlrun/artifacts/manager.py

@@ -306,7 +306,6 @@ class ArtifactManager:
         item.target_path = target_path
 
         item.before_log()
-        self.artifact_uris[key] = item.uri
 
         if ((upload is None and item.kind != "dir") or upload) and not item.is_inline():
             # before uploading the item, we want to ensure that its tags are valid,
@@ -315,7 +314,12 @@ class ArtifactManager:
             item.upload(artifact_path=artifact_path)
 
         if db_key:
-            self._log_to_db(db_key, project, producer.inputs, item)
+            artifact_uid = self._log_to_db(db_key, project, producer.inputs, item)
+            if artifact_uid is not None:
+                item.uid = artifact_uid
+        # Generate the artifact URI after logging to the database and retrieving the artifact UID, if available.
+        self.artifact_uris[key] = item.uri
+
         size = str(item.size) or "?"
         db_str = "Y" if (self.artifact_db and db_key) else "N"
         logger.debug(
@@ -327,20 +331,21 @@ class ArtifactManager:
         self.artifact_uris[item.key] = item.uri
         self._log_to_db(item.db_key, producer.project, producer.inputs, item)
 
-    def _log_to_db(self, key, project, sources, item, tag=None):
+    def _log_to_db(self, key, project, sources, item, tag=None) -> typing.Optional[str]:
         """
         log artifact to db
         :param key: Identifying key of the artifact.
         :param project: Project that the artifact belongs to.
-        :param sources: List of artifact sources ( Mainly passed from the producer.items ).
+        :param sources: List of artifact sources ( Mainly passed from the `producer.items` ).
         :param item: The actual artifact to store.
         :param tag: The name of the Tag of the artifact.
+        :return: The logged artifact uid.
         """
         if self.artifact_db:
             item.updated = None
             if sources:
                 item.sources = [{"name": k, "path": str(v)} for k, v in sources.items()]
-            self.artifact_db.store_artifact(
+            artifact_item = self.artifact_db.store_artifact(
                 key,
                 item.to_dict(),
                 iter=item.iter,
@@ -348,6 +353,8 @@ class ArtifactManager:
                 project=project,
                 tree=item.tree,
             )
+            if artifact_item:
+                return artifact_item.get("metadata", {}).get("uid")
 
     def link_artifact(
         self,

mlrun/common/constants.py

@@ -69,6 +69,7 @@ class MLRunInternalLabels:
     producer_type = f"{MLRUN_LABEL_PREFIX}producer-type"
     app_name = f"{MLRUN_LABEL_PREFIX}app-name"
     endpoint_id = f"{MLRUN_LABEL_PREFIX}endpoint-id"
+    endpoint_name = f"{MLRUN_LABEL_PREFIX}endpoint-name"
     host = "host"
     job_type = "job-type"
     kind = "kind"

mlrun/common/model_monitoring/__init__.py

@@ -11,5 +11,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from .helpers import create_model_endpoint_uid

mlrun/common/model_monitoring/helpers.py

@@ -17,11 +17,6 @@ import typing
 
 import mlrun.common
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
-from mlrun.common.schemas.model_monitoring import (
-    EndpointUID,
-    FunctionURI,
-    VersionedModel,
-)
 
 FeatureStats = typing.NewType("FeatureStats", dict[str, dict[str, typing.Any]])
 Histogram = typing.NewType("Histogram", list[list])
@@ -31,29 +26,6 @@ BinEdges = typing.NewType("BinEdges", list[float])
 _MAX_FLOAT = sys.float_info.max
 
 
-def create_model_endpoint_uid(function_uri: str, versioned_model: str):
-    function_uri = FunctionURI.from_string(function_uri)
-    versioned_model = VersionedModel.from_string(versioned_model)
-
-    if (
-        not function_uri.project
-        or not function_uri.function
-        or not versioned_model.model
-    ):
-        raise ValueError("Both function_uri and versioned_model have to be initialized")
-
-    uid = EndpointUID(
-        function_uri.project,
-        function_uri.function,
-        function_uri.tag,
-        function_uri.hash_key,
-        versioned_model.model,
-        versioned_model.version,
-    )
-
-    return uid
-
-
 def parse_model_endpoint_project_prefix(path: str, project_name: str):
     return path.split(project_name, 1)[0] + project_name
 

mlrun/common/schemas/__init__.py

@@ -14,6 +14,7 @@
 
 from .alert import (
     AlertActivation,
+    AlertActivations,
     AlertActiveState,
     AlertConfig,
     AlertNotification,
@@ -149,10 +150,6 @@ from .model_monitoring import (
     ModelEndpointMetadata,
     ModelEndpointSpec,
     ModelEndpointStatus,
-    ModelEndpointV2,
-    ModelEndpointV2Metadata,
-    ModelEndpointV2Spec,
-    ModelEndpointV2Status,
     ModelMonitoringMode,
     ModelMonitoringStoreKinds,
     MonitoringFunctionNames,
@@ -165,6 +162,7 @@ from .notification import (
     NotificationSeverity,
     NotificationState,
     NotificationStatus,
+    NotificationSummary,
     SetNotificationRequest,
 )
 from .object import ObjectKind, ObjectMetadata, ObjectSpec, ObjectStatus

mlrun/common/schemas/alert.py

@@ -12,8 +12,10 @@
 # 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.v1
 
@@ -216,3 +218,77 @@ class AlertActivation(pydantic.v1.BaseModel):
     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/client_spec.py

@@ -57,7 +57,6 @@ class ClientSpec(pydantic.v1.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

mlrun/common/schemas/model_monitoring/__init__.py

@@ -55,12 +55,6 @@ from .grafana import (
     GrafanaTable,
     GrafanaTimeSeriesTarget,
 )
-from .model_endpoint_v2 import (
-    ModelEndpointV2,
-    ModelEndpointV2Metadata,
-    ModelEndpointV2Spec,
-    ModelEndpointV2Status,
-)
 from .model_endpoints import (
     Features,
     FeatureValues,

mlrun/common/schemas/model_monitoring/constants.py

@@ -41,6 +41,7 @@ class ModelEndpointSchema(MonitoringStrEnum):
 
     # spec
     FUNCTION_NAME = "function_name"
+    FUNCTION_TAG = "function_tag"
     FUNCTION_UID = "function_uid"
     MODEL_NAME = "model_name"
     MODEL_TAG = "model_tag"
@@ -48,23 +49,23 @@ class ModelEndpointSchema(MonitoringStrEnum):
     MODEL_UID = "model_uid"
     FEATURE_NAMES = "feature_names"
     LABEL_NAMES = "label_names"
-
-    # status
-    STATE = "state"
-    MONITORING_MODE = "monitoring_mode"
+    FEATURE_STATS = "feature_stats"
     MONITORING_FEATURE_SET_URI = "monitoring_feature_set_uri"
     CHILDREN = "children"
     CHILDREN_UIDS = "children_uids"
-    FIRST_REQUEST = "first_request"
     FUNCTION_URI = "function_uri"
     MODEL_URI = "model_uri"
 
+    # status
+    STATE = "state"
+    MONITORING_MODE = "monitoring_mode"
+    FIRST_REQUEST = "first_request"
+
     # status - operative
     LAST_REQUEST = "last_request"
-    DRIFT_STATUS = "drift_status"
+    RESULT_STATUS = "result_status"
     AVG_LATENCY = "avg_latency"
     ERROR_COUNT = "error_count"
-    FEATURE_STATS = "feature_stats"
     CURRENT_STATS = "current_stats"
     DRIFT_MEASURES = "drift_measures"
 
@@ -80,6 +81,7 @@ class EventFieldType:
     TIMESTAMP = "timestamp"
     # `endpoint_id` is deprecated as a field in the model endpoint schema since 1.3.1, replaced by `uid`.
     ENDPOINT_ID = "endpoint_id"
+    ENDPOINT_NAME = "endpoint_name"
     UID = "uid"
     ENDPOINT_TYPE = "endpoint_type"
     REQUEST_ID = "request_id"
@@ -148,10 +150,12 @@ class ApplicationEvent:
     START_INFER_TIME = "start_infer_time"
     END_INFER_TIME = "end_infer_time"
     ENDPOINT_ID = "endpoint_id"
+    ENDPOINT_NAME = "endpoint_name"
     OUTPUT_STREAM_URI = "output_stream_uri"
 
 
 class WriterEvent(MonitoringStrEnum):
+    ENDPOINT_NAME = "endpoint_name"
     APPLICATION_NAME = "application_name"
     ENDPOINT_ID = "endpoint_id"
     START_INFER_TIME = "start_infer_time"
@@ -222,7 +226,6 @@ class TSDBTarget(MonitoringStrEnum):
 
 
 class ProjectSecretKeys:
-    ENDPOINT_STORE_CONNECTION = "MODEL_MONITORING_ENDPOINT_STORE_CONNECTION"
     ACCESS_KEY = "MODEL_MONITORING_ACCESS_KEY"
     STREAM_PATH = "STREAM_PATH"
     TSDB_CONNECTION = "TSDB_CONNECTION"
@@ -230,7 +233,6 @@ class ProjectSecretKeys:
     @classmethod
     def mandatory_secrets(cls):
         return [
-            cls.ENDPOINT_STORE_CONNECTION,
             cls.STREAM_PATH,
             cls.TSDB_CONNECTION,
         ]