mlrun 1.8.0rc10__py3-none-any.whl → 1.8.0rc13__py3-none-any.whl

mlrun/artifacts/document.py

@@ -89,12 +89,17 @@ 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.
+        artifact_key (str, optional): The key for the artifact to be logged. Special characters and symbols
+            not valid in artifact names will be encoded as their hexadecimal representation. The '%%' pattern
+            in the key will be replaced by the hex-encoded version of the source path. Defaults to "doc%%".
         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
+        producer (Optional[Union[MlrunProject, str, MLClientCtx]], optional): The producer of the document.
+                                If not specified, will try to get the current MLRun context or project.
+                                Defaults to None.
         upload (bool, optional): Flag indicating whether to upload the document.
+        labels (Optional[Dict[str, str]], optional): Key-value labels to attach to the artifact. Defaults to None.
+        tag (str, optional): Version tag for the artifact. Defaults to "".
 
     Returns:
         DynamicDocumentLoader: An instance of a dynamically defined subclass of BaseLoader.
@@ -146,6 +151,8 @@ class MLRunLoader:
         artifact_key="doc%%",
         producer: Optional[Union["MlrunProject", str, "MLClientCtx"]] = None,  # noqa: F821
         upload: bool = False,
+        tag: str = "",
+        labels: Optional[dict[str, str]] = None,
     ):
         # Dynamically import BaseLoader
         from langchain_community.document_loaders.base import BaseLoader
@@ -158,6 +165,8 @@ class MLRunLoader:
                 artifact_key,
                 producer,
                 upload,
+                tag,
+                labels,
             ):
                 self.producer = producer
                 self.artifact_key = (
@@ -168,6 +177,8 @@ class MLRunLoader:
                 self.loader_spec = loader_spec
                 self.local_path = local_path
                 self.upload = upload
+                self.tag = tag
+                self.labels = labels
 
                 # Resolve the producer
                 if not self.producer:
@@ -181,9 +192,11 @@ class MLRunLoader:
                     document_loader_spec=self.loader_spec,
                     local_path=self.local_path,
                     upload=self.upload,
+                    labels=self.labels,
+                    tag=self.tag,
                 )
                 res = artifact.to_langchain_documents()
-                yield res[0]
+                return res
 
         # Return an instance of the dynamically defined subclass
         instance = DynamicDocumentLoader(
@@ -192,6 +205,8 @@ class MLRunLoader:
             loader_spec=loader_spec,
             producer=producer,
             upload=upload,
+            tag=tag,
+            labels=labels,
         )
         return instance
 
@@ -257,6 +272,9 @@ class DocumentArtifact(Artifact):
     METADATA_CHUNK_KEY = "mlrun_chunk"
     METADATA_ARTIFACT_URI_KEY = "mlrun_object_uri"
     METADATA_ARTIFACT_TARGET_PATH_KEY = "mlrun_target_path"
+    METADATA_ARTIFACT_TAG = "mlrun_tag"
+    METADATA_ARTIFACT_KEY = "mlrun_key"
+    METADATA_ARTIFACT_PROJECT = "mlrun_project"
 
     def __init__(
         self,
@@ -331,6 +349,10 @@ class DocumentArtifact(Artifact):
             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
+            metadata[self.METADATA_ARTIFACT_TAG] = self.tag or "latest"
+            metadata[self.METADATA_ARTIFACT_KEY] = self.key
+            metadata[self.METADATA_ARTIFACT_PROJECT] = self.metadata.project
+
             if self.get_target_path():
                 metadata[self.METADATA_ARTIFACT_TARGET_PATH_KEY] = (
                     self.get_target_path()
@@ -346,7 +368,7 @@ class DocumentArtifact(Artifact):
                 idx = idx + 1
         return results
 
-    def collection_add(self, collection_id: str) -> None:
+    def collection_add(self, collection_id: str) -> bool:
         """
         Add a collection ID to the artifact's collection list.
 
@@ -361,8 +383,10 @@ class DocumentArtifact(Artifact):
         """
         if collection_id not in self.spec.collections:
             self.spec.collections[collection_id] = "1"
+            return True
+        return False
 
-    def collection_remove(self, collection_id: str) -> None:
+    def collection_remove(self, collection_id: str) -> bool:
         """
         Remove a collection ID from the artifact's collection list.
 
@@ -376,3 +400,5 @@ class DocumentArtifact(Artifact):
         """
         if collection_id in self.spec.collections:
             self.spec.collections.pop(collection_id)
+            return True
+        return False

mlrun/common/constants.py

@@ -25,6 +25,7 @@ MYSQL_MEDIUMBLOB_SIZE_BYTES = 16 * 1024 * 1024
 MLRUN_LABEL_PREFIX = "mlrun/"
 DASK_LABEL_PREFIX = "dask.org/"
 NUCLIO_LABEL_PREFIX = "nuclio.io/"
+RESERVED_TAG_NAME_LATEST = "latest"
 
 
 class MLRunInternalLabels:

mlrun/common/formatters/artifact.py

@@ -32,7 +32,7 @@ class ArtifactFormat(ObjectFormat, mlrun.common.types.StrEnum):
                 [
                     "kind",
                     "metadata",
-                    "status",
+                    "status.state",
                     "project",
                     "spec.producer",
                     "spec.db_key",

mlrun/common/schemas/__init__.py

@@ -146,8 +146,10 @@ from .model_monitoring import (
     GrafanaTable,
     GrafanaTimeSeriesTarget,
     ModelEndpoint,
+    ModelEndpointCreationStrategy,
     ModelEndpointList,
     ModelEndpointMetadata,
+    ModelEndpointSchema,
     ModelEndpointSpec,
     ModelEndpointStatus,
     ModelMonitoringMode,

mlrun/common/schemas/model_monitoring/__init__.py

@@ -26,6 +26,7 @@ from .constants import (
     FileTargetKind,
     FunctionURI,
     MetricData,
+    ModelEndpointCreationStrategy,
     ModelEndpointMonitoringMetricType,
     ModelEndpointSchema,
     ModelEndpointTarget,

mlrun/common/schemas/model_monitoring/constants.py

@@ -71,6 +71,12 @@ class ModelEndpointSchema(MonitoringStrEnum):
     DRIFT_MEASURES = "drift_measures"
 
 
+class ModelEndpointCreationStrategy(MonitoringStrEnum):
+    INPLACE = "inplace"
+    ARCHIVE = "archive"
+    OVERWRITE = "overwrite"
+
+
 class EventFieldType:
     FUNCTION_URI = "function_uri"
     FUNCTION = "function"

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -117,6 +117,10 @@ class ModelEndpointMetadata(ObjectMetadata, ModelEndpointParser):
     endpoint_type: EndpointType = EndpointType.NODE_EP
     uid: Optional[constr(regex=MODEL_ENDPOINT_ID_PATTERN)]
 
+    @classmethod
+    def mutable_fields(cls):
+        return ["labels"]
+
 
 class ModelEndpointSpec(ObjectSpec, ModelEndpointParser):
     model_uid: Optional[str] = ""
@@ -136,6 +140,21 @@ class ModelEndpointSpec(ObjectSpec, ModelEndpointParser):
     children_uids: Optional[list[str]] = []
     monitoring_feature_set_uri: Optional[str] = ""
 
+    @classmethod
+    def mutable_fields(cls):
+        return [
+            "model_uid",
+            "model_name",
+            "model_db_key",
+            "model_tag",
+            "model_class",
+            "function_uid",
+            "feature_names",
+            "label_names",
+            "children",
+            "children_uids",
+        ]
+
 
 class ModelEndpointStatus(ObjectStatus, ModelEndpointParser):
     state: Optional[str] = "unknown"  # will be updated according to the function state
@@ -152,6 +171,14 @@ class ModelEndpointStatus(ObjectStatus, ModelEndpointParser):
     drift_measures: Optional[dict] = {}
     drift_measures_timestamp: Optional[datetime] = None
 
+    @classmethod
+    def mutable_fields(cls):
+        return [
+            "monitoring_mode",
+            "first_request",
+            "last_request",
+        ]
+
 
 class ModelEndpoint(BaseModel):
     kind: ObjectKind = Field(ObjectKind.model_endpoint, const=True)
@@ -159,6 +186,14 @@ class ModelEndpoint(BaseModel):
     spec: ModelEndpointSpec
     status: ModelEndpointStatus
 
+    @classmethod
+    def mutable_fields(cls):
+        return (
+            ModelEndpointMetadata.mutable_fields()
+            + ModelEndpointSpec.mutable_fields()
+            + ModelEndpointStatus.mutable_fields()
+        )
+
     def flat_dict(self) -> dict[str, Any]:
         """Generate a flattened `ModelEndpoint` dictionary. The flattened dictionary result is important for storing
         the model endpoint object in the database.

mlrun/common/schemas/partition.py

@@ -46,24 +46,23 @@ class PartitionInterval(StrEnum):
             return timedelta(weeks=1)
 
     @classmethod
-    def from_function(cls, partition_function: str):
+    def from_expression(cls, partition_expression: str):
         """
-        Returns the corresponding PartitionInterval for a given partition function,
+        Returns the corresponding PartitionInterval for a given partition expression,
         or None if the function is not mapped.
 
-        :param partition_function: The partition function to map to an interval.
-        :return: PartitionInterval corresponding to the function, or None if no match is found.
+        :param partition_expression: The partition expression to map to an interval.
+        :return: PartitionInterval corresponding to the expression, or `month` if no match is found.
         """
-        partition_function_to_partitions_interval = {
-            "DAY": "DAY",
-            "DAYOFMONTH": "DAY",
-            "MONTH": "MONTH",
-            "YEARWEEK": "YEARWEEK",
-        }
-        interval = partition_function_to_partitions_interval.get(partition_function)
-        if interval and cls.is_valid(interval):
-            return cls[interval]
-        raise KeyError(f"Partition function: {partition_function} isn't supported")
+
+        # Match the provided function string to the correct interval
+        partition_expression = partition_expression.upper()
+        if "YEARWEEK" in partition_expression:
+            return cls.YEARWEEK
+        elif "DAYOFMONTH" in partition_expression:
+            return cls.DAY
+        else:
+            return cls.MONTH
 
     def get_partition_info(
         self,
@@ -120,11 +119,17 @@ class PartitionInterval(StrEnum):
             year, week, _ = current_datetime.isocalendar()
             return f"{year}{week:02d}"
 
-    def get_partition_expression(self):
+    def get_partition_expression(self, column_name: str):
         if self == PartitionInterval.YEARWEEK:
-            return "YEARWEEK(activation_time, 1)"
-        else:
-            return f"{self}(activation_time)"
+            return f"YEARWEEK({column_name}, 1)"
+        elif self == PartitionInterval.DAY:
+            # generates value in format %Y%m%d in mysql
+            # mysql query example: `select YEAR(NOW())*10000 + MONTH(NOW())*100 + DAY(NOW());`
+            return f"YEAR({column_name}) * 10000 + MONTH({column_name}) * 100 + DAY({column_name})"
+        elif self == PartitionInterval.MONTH:
+            # generates value in format %Y%m in mysql
+            # mysql query example: `select YEAR(NOW())*100 + MONTH(NOW());`
+            return f"YEAR({column_name}) * 100 + MONTH({column_name})"
 
     def get_number_of_partitions(self, days: int) -> int:
         # Calculate the number partitions based on given number of days

mlrun/datastore/vectorstore.py

@@ -19,19 +19,42 @@ from typing import Optional, Union
 from mlrun.artifacts import DocumentArtifact
 
 
-def _extract_collection_name(vectorstore: "VectorStore") -> str:  # noqa: F821
-    # List of possible attribute names for collection name
-    possible_attributes = ["collection_name", "_collection_name"]
+def find_existing_attribute(obj, base_name="name", parent_name="collection"):
+    # Define all possible patterns
+
+    return None
+
 
-    for attr in possible_attributes:
-        if hasattr(vectorstore, attr):
-            collection_name = getattr(vectorstore, attr)
-            if collection_name:
-                return collection_name
+def _extract_collection_name(vectorstore: "VectorStore") -> str:  # noqa: F821
+    patterns = [
+        "collection.name",
+        "collection._name",
+        "_collection.name",
+        "_collection._name",
+        "collection_name",
+        "_collection_name",
+    ]
+
+    def resolve_attribute(obj, pattern):
+        if "." in pattern:
+            parts = pattern.split(".")
+            current = vectorstore
+            for part in parts:
+                if hasattr(current, part):
+                    current = getattr(current, part)
+                else:
+                    return None
+            return current
+        else:
+            return getattr(obj, pattern, None)
 
-    store_class = vectorstore.__class__.__name__.lower()
-    if store_class == "mongodbatlasvectorsearch":
-        return vectorstore.collection.name
+    for pattern in patterns:
+        try:
+            value = resolve_attribute(vectorstore, pattern)
+            if value is not None:
+                return value
+        except (AttributeError, TypeError):
+            continue
 
     # If we get here, we couldn't find a valid collection name
     raise ValueError(
@@ -82,6 +105,19 @@ class VectorStoreCollection:
             # Forward the attribute setting to _collection_impl
             setattr(self._collection_impl, name, value)
 
+    def _get_mlrun_project_name(self):
+        import mlrun
+
+        if self._mlrun_context and isinstance(
+            self._mlrun_context, mlrun.projects.MlrunProject
+        ):
+            return self._mlrun_context.name
+        if self._mlrun_context and isinstance(
+            self._mlrun_context, mlrun.execution.MLClientCtx
+        ):
+            return self._mlrun_context.get_project_object().name
+        return None
+
     def delete(self, *args, **kwargs):
         self._collection_impl.delete(*args, **kwargs)
 
@@ -106,13 +142,22 @@ class VectorStoreCollection:
         """
         if self._mlrun_context:
             for document in documents:
-                mlrun_uri = document.metadata.get(
-                    DocumentArtifact.METADATA_ARTIFACT_URI_KEY
+                mlrun_key = document.metadata.get(
+                    DocumentArtifact.METADATA_ARTIFACT_KEY, None
+                )
+                mlrun_project = document.metadata.get(
+                    DocumentArtifact.METADATA_ARTIFACT_PROJECT, None
                 )
-                if mlrun_uri:
-                    artifact = self._mlrun_context.get_store_resource(mlrun_uri)
-                    artifact.collection_add(self.collection_name)
-                    self._mlrun_context.update_artifact(artifact)
+
+                if mlrun_key and mlrun_project == self._get_mlrun_project_name():
+                    mlrun_tag = document.metadata.get(
+                        DocumentArtifact.METADATA_ARTIFACT_TAG, None
+                    )
+                    artifact = self._mlrun_context.get_artifact(
+                        key=mlrun_key, tag=mlrun_tag
+                    )
+                    if artifact.collection_add(self.collection_name):
+                        self._mlrun_context.update_artifact(artifact)
 
         return self._collection_impl.add_documents(documents, **kwargs)
 
@@ -159,8 +204,7 @@ class VectorStoreCollection:
                 )
         for index, artifact in enumerate(artifacts):
             documents = artifact.to_langchain_documents(splitter)
-            artifact.collection_add(self.collection_name)
-            if self._mlrun_context:
+            if artifact.collection_add(self.collection_name) and self._mlrun_context:
                 self._mlrun_context.update_artifact(artifact)
             if user_ids:
                 num_of_documents = len(documents)
@@ -182,8 +226,8 @@ class VectorStoreCollection:
         Args:
             artifact (DocumentArtifact): The artifact from which the current object should be removed.
         """
-        artifact.collection_remove(self.collection_name)
-        if self._mlrun_context:
+
+        if artifact.collection_remove(self.collection_name) and self._mlrun_context:
             self._mlrun_context.update_artifact(artifact)
 
     def delete_artifacts(self, artifacts: list[DocumentArtifact]):
@@ -201,16 +245,15 @@ class VectorStoreCollection:
         """
         store_class = self._collection_impl.__class__.__name__.lower()
         for artifact in artifacts:
-            artifact.collection_remove(self.collection_name)
-            if self._mlrun_context:
+            if artifact.collection_remove(self.collection_name) and self._mlrun_context:
                 self._mlrun_context.update_artifact(artifact)
 
             if store_class == "milvus":
                 expr = f"{DocumentArtifact.METADATA_SOURCE_KEY} == '{artifact.get_source()}'"
-                return self._collection_impl.delete(expr=expr)
+                self._collection_impl.delete(expr=expr)
             elif store_class == "chroma":
                 where = {DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_source()}
-                return self._collection_impl.delete(where=where)
+                self._collection_impl.delete(where=where)
 
             elif (
                 hasattr(self._collection_impl, "delete")
@@ -222,7 +265,7 @@ class VectorStoreCollection:
                         DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_source()
                     }
                 }
-                return self._collection_impl.delete(filter=filter)
+                self._collection_impl.delete(filter=filter)
             else:
                 raise NotImplementedError(
                     f"delete_artifacts() operation not supported for {store_class}"

mlrun/db/base.py

@@ -23,6 +23,7 @@ import mlrun.common
 import mlrun.common.formatters
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.common.schemas.model_monitoring.model_endpoints as mm_endpoints
 import mlrun.model_monitoring
 
@@ -58,6 +59,15 @@ class RunDBInterface(ABC):
     def abort_run(self, uid, project="", iter=0, timeout=45, status_text=""):
         pass
 
+    @abstractmethod
+    def push_run_notifications(
+        self,
+        uid,
+        project="",
+        timeout=45,
+    ):
+        pass
+
     @abstractmethod
     def read_run(
         self,
@@ -666,6 +676,9 @@ class RunDBInterface(ABC):
     def create_model_endpoint(
         self,
         model_endpoint: mlrun.common.schemas.ModelEndpoint,
+        creation_strategy: Optional[
+            mm_constants.ModelEndpointCreationStrategy
+        ] = mm_constants.ModelEndpointCreationStrategy.INPLACE,
     ) -> mlrun.common.schemas.ModelEndpoint:
         pass
 
@@ -688,6 +701,7 @@ class RunDBInterface(ABC):
         function_name: Optional[str] = None,
         function_tag: Optional[str] = None,
         model_name: Optional[str] = None,
+        model_tag: Optional[str] = None,
         labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         start: Optional[datetime.datetime] = None,
         end: Optional[datetime.datetime] = None,

mlrun/db/httpdb.py

@@ -35,6 +35,7 @@ import mlrun.common.constants
 import mlrun.common.formatters
 import mlrun.common.runtimes
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.common.schemas.model_monitoring.model_endpoints as mm_endpoints
 import mlrun.common.types
 import mlrun.platforms
@@ -755,6 +756,34 @@ class HTTPRunDB(RunDBInterface):
             )
         return None
 
+    def push_run_notifications(
+        self,
+        uid,
+        project="",
+        timeout=45,
+    ):
+        """
+        Push notifications for a run.
+
+        :param uid: Unique ID of the run.
+        :param project: Project that the run belongs to.
+        :returns: :py:class:`~mlrun.common.schemas.BackgroundTask`.
+        """
+        project = project or config.default_project
+
+        response = self.api_call(
+            "POST",
+            path=f"projects/{project}/runs/{uid}/push_notifications",
+            error="Failed push notifications",
+            timeout=timeout,
+        )
+        if response.status_code == http.HTTPStatus.ACCEPTED:
+            background_task = mlrun.common.schemas.BackgroundTask(**response.json())
+            return self._wait_for_background_task_to_reach_terminal_state(
+                background_task.metadata.name, project=project
+            )
+        return None
+
     def read_run(
         self,
         uid,
@@ -3582,12 +3611,24 @@ class HTTPRunDB(RunDBInterface):
     def create_model_endpoint(
         self,
         model_endpoint: mlrun.common.schemas.ModelEndpoint,
+        creation_strategy: Optional[
+            mm_constants.ModelEndpointCreationStrategy
+        ] = mm_constants.ModelEndpointCreationStrategy.INPLACE,
     ) -> mlrun.common.schemas.ModelEndpoint:
         """
         Creates a DB record with the given model_endpoint record.
 
         :param model_endpoint: An object representing the model endpoint.
-
+        :param creation_strategy: Strategy for creating or updating the model endpoint:
+            * **overwrite**:
+            1. If model endpoints with the same name exist, delete the `latest` one.
+            2. Create a new model endpoint entry and set it as `latest`.
+            * **inplace** (default):
+            1. If model endpoints with the same name exist, update the `latest` entry.
+            2. Otherwise, create a new entry.
+            * **archive**:
+            1. If model endpoints with the same name exist, preserve them.
+            2. Create a new model endpoint with the same name and set it to `latest`.
         :return: The created model endpoint object.
         """
 
@@ -3596,6 +3637,9 @@ class HTTPRunDB(RunDBInterface):
             method=mlrun.common.types.HTTPMethod.POST,
             path=path,
             body=model_endpoint.json(),
+            params={
+                "creation_strategy": creation_strategy,
+            },
         )
         return mlrun.common.schemas.ModelEndpoint(**response.json())
 
@@ -3637,6 +3681,7 @@ class HTTPRunDB(RunDBInterface):
         function_name: Optional[str] = None,
         function_tag: Optional[str] = None,
         model_name: Optional[str] = None,
+        model_tag: Optional[str] = None,
         labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
@@ -3653,6 +3698,7 @@ class HTTPRunDB(RunDBInterface):
         :param function_name:   The name of the function
         :param function_tag:    The tag of the function
         :param model_name:      The name of the model
+        :param model_tag:       The tag of the model
         :param labels:          A list of labels to filter by. (see mlrun.common.schemas.LabelsModel)
         :param start:           The start time to filter by.Corresponding to the `created` field.
         :param end:             The end time to filter by. Corresponding to the `created` field.
@@ -3671,6 +3717,7 @@ class HTTPRunDB(RunDBInterface):
             params={
                 "name": name,
                 "model_name": model_name,
+                "model_tag": model_tag,
                 "function_name": function_name,
                 "function_tag": function_tag,
                 "label": labels,

mlrun/db/nopdb.py

@@ -20,6 +20,7 @@ import mlrun.alerts
 import mlrun.common.formatters
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.errors
 import mlrun.lists
 import mlrun.model_monitoring
@@ -75,6 +76,14 @@ class NopDB(RunDBInterface):
     def abort_run(self, uid, project="", iter=0, timeout=45, status_text=""):
         pass
 
+    def push_run_notifications(
+        self,
+        uid,
+        project="",
+        timeout=45,
+    ):
+        pass
+
     def list_runtime_resources(
         self,
         project: Optional[str] = None,
@@ -575,6 +584,9 @@ class NopDB(RunDBInterface):
     def create_model_endpoint(
         self,
         model_endpoint: mlrun.common.schemas.ModelEndpoint,
+        creation_strategy: Optional[
+            mm_constants.ModelEndpointCreationStrategy
+        ] = mm_constants.ModelEndpointCreationStrategy.INPLACE,
     ) -> mlrun.common.schemas.ModelEndpoint:
         pass
 
@@ -595,6 +607,7 @@ class NopDB(RunDBInterface):
         function_name: Optional[str] = None,
         function_tag: Optional[str] = None,
         model_name: Optional[str] = None,
+        model_tag: Optional[str] = None,
         labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
         start: Optional[datetime.datetime] = None,
         end: Optional[datetime.datetime] = None,