mlrun 1.8.0rc19__py3-none-any.whl → 1.8.0rc26__py3-none-any.whl

mlrun/datastore/sources.py

@@ -1089,9 +1089,10 @@ class KafkaSource(OnlineSource):
             attributes["partitions"] = partitions
         sasl = attributes.pop("sasl", {})
         if sasl_user and sasl_pass:
-            sasl["enabled"] = True
+            sasl["enable"] = True
             sasl["user"] = sasl_user
             sasl["password"] = sasl_pass
+            sasl["mechanism"] = "PLAIN"
         if sasl:
             attributes["sasl"] = sasl
         super().__init__(attributes=attributes, **kwargs)
@@ -1127,8 +1128,13 @@ class KafkaSource(OnlineSource):
             extra_attributes["workerAllocationMode"] = extra_attributes.get(
                 "worker_allocation_mode", "static"
             )
+        else:
+            extra_attributes["workerAllocationMode"] = extra_attributes.get(
+                "worker_allocation_mode", "pool"
+            )
 
         trigger_kwargs = {}
+
         if "max_workers" in extra_attributes:
             trigger_kwargs = {"max_workers": extra_attributes.pop("max_workers")}
 

mlrun/datastore/vectorstore.py

@@ -48,6 +48,18 @@ def _extract_collection_name(vectorstore: "VectorStore") -> str:  # noqa: F821
         else:
             return getattr(obj, pattern, None)
 
+    if type(vectorstore).__name__ == "PineconeVectorStore":
+        try:
+            url = (
+                vectorstore._index.config.host
+                if hasattr(vectorstore._index, "config")
+                else vectorstore._index._config.host
+            )
+            index_name = url.split("//")[1].split("-")[0]
+            return index_name
+        except Exception:
+            pass
+
     for pattern in patterns:
         try:
             value = resolve_attribute(vectorstore, pattern)
@@ -254,7 +266,14 @@ class VectorStoreCollection:
             elif store_class == "chroma":
                 where = {DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_source()}
                 self._collection_impl.delete(where=where)
-
+            elif store_class == "pineconevectorstore":
+                filter = {
+                    DocumentArtifact.METADATA_SOURCE_KEY: {"$eq": artifact.get_source()}
+                }
+                self._collection_impl.delete(filter=filter)
+            elif store_class == "mongodbatlasvectorsearch":
+                filter = {DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_source()}
+                self._collection_impl.collection.delete_many(filter=filter)
             elif (
                 hasattr(self._collection_impl, "delete")
                 and "filter"

mlrun/db/base.py

@@ -68,6 +68,18 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    def refresh_smtp_configuration(self):
+        pass
+
+    def push_pipeline_notifications(
+        self,
+        pipeline_id,
+        project="",
+        notifications=None,
+        timeout=45,
+    ):
+        pass
+
     @abstractmethod
     def read_run(
         self,
@@ -917,6 +929,14 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    @abstractmethod
+    def get_alert_activation(
+        self,
+        project,
+        activation_id,
+    ) -> mlrun.common.schemas.AlertActivation:
+        pass
+
     def update_alert_activation(
         self,
         activation_id: int,

mlrun/db/httpdb.py

@@ -559,14 +559,6 @@ class HTTPRunDB(RunDBInterface):
                 server_cfg.get("external_platform_tracking")
                 or config.external_platform_tracking
             )
-            config.model_endpoint_monitoring.tsdb_connection = (
-                server_cfg.get("model_monitoring_tsdb_connection")
-                or config.model_endpoint_monitoring.tsdb_connection
-            )
-            config.model_endpoint_monitoring.stream_connection = (
-                server_cfg.get("stream_connection")
-                or config.model_endpoint_monitoring.stream_connection
-            )
             config.packagers = server_cfg.get("packagers") or config.packagers
             server_data_prefixes = server_cfg.get("feature_store_data_prefixes") or {}
             for prefix in ["default", "nosql", "redisnosql"]:
@@ -771,7 +763,6 @@ class HTTPRunDB(RunDBInterface):
         :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",
@@ -780,9 +771,84 @@ class HTTPRunDB(RunDBInterface):
         )
         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 = self._wait_for_background_task_to_reach_terminal_state(
                 background_task.metadata.name, project=project
             )
+            if (
+                background_task.status.state
+                == mlrun.common.schemas.BackgroundTaskState.succeeded
+            ):
+                logger.info(
+                    "Notifications for the run have been pushed",
+                    project=project,
+                    run_id=uid,
+                )
+            elif (
+                background_task.status.state
+                == mlrun.common.schemas.BackgroundTaskState.failed
+            ):
+                logger.error(
+                    "Failed to push run notifications",
+                    project=project,
+                    run_id=uid,
+                    error=background_task.status.error,
+                )
+        return None
+
+    def push_pipeline_notifications(
+        self,
+        pipeline_id,
+        project="",
+        notifications=None,
+        timeout=45,
+    ):
+        """
+        Push notifications for a pipeline.
+
+        :param pipeline_id: Unique ID of the pipeline(KFP).
+        :param project: Project that the run belongs to.
+        :param notifications: List of notifications to push.
+        :returns: :py:class:`~mlrun.common.schemas.BackgroundTask`.
+        """
+        if notifications is None or type(notifications) is not list:
+            raise MLRunInvalidArgumentError(
+                "The 'notifications' parameter must be a list."
+            )
+
+        project = project or config.default_project
+
+        response = self.api_call(
+            "POST",
+            path=f"projects/{project}/pipelines/{pipeline_id}/push-notifications",
+            error="Failed push notifications",
+            body=_as_json([notification.to_dict() for notification in notifications]),
+            timeout=timeout,
+        )
+        if response.status_code == http.HTTPStatus.ACCEPTED:
+            background_task = mlrun.common.schemas.BackgroundTask(**response.json())
+            background_task = self._wait_for_background_task_to_reach_terminal_state(
+                background_task.metadata.name, project=project
+            )
+            if (
+                background_task.status.state
+                == mlrun.common.schemas.BackgroundTaskState.succeeded
+            ):
+                logger.info(
+                    "Pipeline notifications have been pushed",
+                    project=project,
+                    pipeline_id=pipeline_id,
+                )
+            elif (
+                background_task.status.state
+                == mlrun.common.schemas.BackgroundTaskState.failed
+            ):
+                logger.error(
+                    "Failed to push pipeline notifications",
+                    project=project,
+                    pipeline_id=pipeline_id,
+                    error=background_task.status.error,
+                )
+
         return None
 
     def read_run(
@@ -4955,6 +5021,27 @@ class HTTPRunDB(RunDBInterface):
             **kwargs,
         )
 
+    def get_alert_activation(
+        self,
+        project,
+        activation_id,
+    ) -> mlrun.common.schemas.AlertActivation:
+        """
+        Retrieve the alert activation by id
+
+        :param project: Project name for which the summary belongs.
+        :param activation_id: alert activation id.
+        :returns: alert activation object.
+        """
+        project = project or config.default_project
+
+        error = "get alert activation"
+        path = f"projects/{project}/alert-activations/{activation_id}"
+
+        response = self.api_call("GET", path, error)
+
+        return mlrun.common.schemas.AlertActivation(**response.json())
+
     def get_project_summary(
         self, project: Optional[str] = None
     ) -> mlrun.common.schemas.ProjectSummary:

mlrun/db/nopdb.py

@@ -84,6 +84,18 @@ class NopDB(RunDBInterface):
     ):
         pass
 
+    def refresh_smtp_configuration(self):
+        pass
+
+    def push_pipeline_notifications(
+        self,
+        pipeline_id,
+        project="",
+        notifications=None,
+        timeout=45,
+    ):
+        pass
+
     def list_runtime_resources(
         self,
         project: Optional[str] = None,
@@ -936,5 +948,12 @@ class NopDB(RunDBInterface):
     ):
         pass
 
+    def get_alert_activation(
+        self,
+        project,
+        activation_id,
+    ) -> mlrun.common.schemas.AlertActivation:
+        pass
+
     def get_project_summary(self, project: str):
         pass

mlrun/errors.py

@@ -174,6 +174,10 @@ class MLRunInvalidArgumentError(MLRunHTTPStatusError, ValueError):
     error_status_code = HTTPStatus.BAD_REQUEST.value
 
 
+class MLRunModelLimitExceededError(MLRunHTTPStatusError, ValueError):
+    error_status_code = HTTPStatus.BAD_REQUEST.value
+
+
 class MLRunInvalidArgumentTypeError(MLRunHTTPStatusError, TypeError):
     error_status_code = HTTPStatus.BAD_REQUEST.value
 

mlrun/execution.py

@@ -876,7 +876,7 @@ class MLClientCtx:
 
     def log_document(
         self,
-        key: str,
+        key: str = "",
         tag: str = "",
         local_path: str = "",
         artifact_path: Optional[str] = None,
@@ -890,7 +890,8 @@ class MLClientCtx:
         """
         Log a document as an artifact.
 
-        :param key: Artifact key
+        :param key: Optional artifact key. If not provided, will be derived from local_path
+                or target_path using DocumentArtifact.key_from_source()
         :param tag: Version tag
         :param local_path: path to the local file we upload, will also be use
                         as the destination subpath (under "artifact_path")
@@ -923,7 +924,6 @@ class MLClientCtx:
         Example:
             >>> # Log a PDF document with custom loader
             >>> project.log_document(
-            ...     key="my_doc",
             ...     local_path="path/to/doc.pdf",
             ...     document_loader_spec=DocumentLoaderSpec(
             ...         loader_class_name="langchain_community.document_loaders.PDFLoader",
@@ -932,10 +932,19 @@ class MLClientCtx:
             ...     ),
             ... )
         """
+
+        if not key and not local_path and not target_path:
+            raise ValueError(
+                "Must provide either 'key' parameter or 'local_path'/'target_path' to derive the key from"
+            )
+        if not key:
+            key = DocumentArtifact.key_from_source(local_path or target_path)
+
         doc_artifact = DocumentArtifact(
             key=key,
             original_source=local_path or target_path,
             document_loader_spec=document_loader_spec,
+            collections=kwargs.pop("collections", None),
             **kwargs,
         )
 
@@ -964,12 +973,12 @@ class MLClientCtx:
     def get_artifact(
         self, key, tag=None, iter=None, tree=None, uid=None
     ) -> Optional[Artifact]:
-        if tag or iter or tree or uid:
+        cached_artifact_uri = self._artifacts_manager.artifact_uris.get(key, None)
+        if tag or iter or tree or uid or (not cached_artifact_uri):
             project = self.get_project_object()
             return project.get_artifact(key=key, tag=tag, iter=iter, tree=tree, uid=uid)
         else:
-            artifact_uri = self._artifacts_manager.artifact_uris[key]
-            return self.get_store_resource(artifact_uri)
+            return self.get_store_resource(cached_artifact_uri)
 
     def update_artifact(self, artifact_object: Artifact):
         """Update an artifact object in the DB and the cached uri"""

mlrun/frameworks/_common/model_handler.py

@@ -976,7 +976,6 @@ class ModelHandler(ABC, Generic[CommonTypes.ModelType, CommonTypes.IOSampleType]
             custom_objects_map_json,
             local_path=custom_objects_map_json,
             artifact_path=self._context.artifact_path,
-            db_key=False,
         )
 
         # Zip the custom objects directory:
@@ -997,7 +996,6 @@ class ModelHandler(ABC, Generic[CommonTypes.ModelType, CommonTypes.IOSampleType]
             custom_objects_zip,
             local_path=custom_objects_zip,
             artifact_path=self._context.artifact_path,
-            db_key=False,
         )
 
         return artifacts

mlrun/launcher/client.py

@@ -134,7 +134,7 @@ class ClientBaseLauncher(launcher.BaseLauncher, abc.ABC):
         if mlrun.utils.is_jupyter and mlrun.mlconf.ipython_widget:
             results_tbl.show()
             print()
-            ui_url = mlrun.utils.get_ui_url(project, uid)
+            ui_url = mlrun.utils.get_run_url(project, uid=uid, name=run.metadata.name)
             if ui_url:
                 ui_url = f' or <a href="{ui_url}" target="_blank">click here</a> to open in UI'
             IPython.display.display(
@@ -150,6 +150,6 @@ class ClientBaseLauncher(launcher.BaseLauncher, abc.ABC):
             mlrun.utils.logger.info(
                 "To track results use the CLI", info_cmd=info_cmd, logs_cmd=logs_cmd
             )
-            ui_url = mlrun.utils.get_ui_url(project, uid)
+            ui_url = mlrun.utils.get_run_url(project, uid=uid, name=run.metadata.name)
             if ui_url:
                 mlrun.utils.logger.info("Or click for UI", ui_url=ui_url)

mlrun/launcher/local.py

@@ -281,5 +281,9 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
         # once the run is completed, and we can just push the notifications.
         # Only push from jupyter, not from the CLI.
         # "handler" and "dask" kinds are special cases of local runs which don't set local=True
-        if self._is_run_local or runtime.kind in ["handler", "dask"]:
+        if self._is_run_local or runtime.kind in ["handler"]:
             mlrun.utils.notifications.NotificationPusher([runobj]).push()
+        elif runtime.kind in ["dask"]:
+            runtime._get_db().push_run_notifications(
+                uid=runobj.metadata.uid, project=runobj.metadata.project
+            )

mlrun/model_monitoring/applications/_application_steps.py

@@ -166,7 +166,9 @@ class _ApplicationErrorHandler(StepToDict):
             "Endpoint ID": event.body.endpoint_id,
             "Application Class": event.body.application_name,
             "Error": "".join(
-                traceback.format_exception(None, event.error, event.error.__traceback__)
+                traceback.format_exception(
+                    None, value=event.error, tb=event.error.__traceback__
+                )
             ),
             "Timestamp": event.timestamp,
         }