mlrun 1.7.0rc16__py3-none-any.whl → 1.7.0rc18__py3-none-any.whl

mlrun/model_monitoring/helpers.py

@@ -111,6 +111,24 @@ def get_connection_string(secret_provider: typing.Callable = None) -> str:
     )
 
 
+def get_tsdb_connection_string(
+    secret_provider: typing.Optional[typing.Callable] = None,
+) -> str:
+    """Get TSDB connection string from the project secret. If wasn't set, take it from the system
+    configurations.
+    :param secret_provider: An optional secret provider to get the connection string secret.
+    :return:                Valid TSDB connection string.
+    """
+
+    return (
+        mlrun.get_secret_or_env(
+            key=mlrun.common.schemas.model_monitoring.ProjectSecretKeys.TSDB_CONNECTION,
+            secret_provider=secret_provider,
+        )
+        or mlrun.mlconf.model_endpoint_monitoring.tsdb_connection
+    )
+
+
 def batch_dict2timedelta(batch_dict: _BatchDict) -> datetime.timedelta:
     """
     Convert a batch dictionary to timedelta.
@@ -260,3 +278,17 @@ def get_endpoint_record(project: str, endpoint_id: str):
         project=project,
     )
     return model_endpoint_store.get_model_endpoint(endpoint_id=endpoint_id)
+
+
+def get_result_instance_fqn(
+    model_endpoint_id: str, app_name: str, result_name: str
+) -> str:
+    return f"{model_endpoint_id}.{app_name}.result.{result_name}"
+
+
+def get_default_result_instance_fqn(model_endpoint_id: str) -> str:
+    return get_result_instance_fqn(
+        model_endpoint_id,
+        mm_constants.HistogramDataDriftApplicationConstants.NAME,
+        mm_constants.HistogramDataDriftApplicationConstants.GENERAL_RESULT_NAME,
+    )

mlrun/model_monitoring/stream_processing.py

@@ -136,7 +136,11 @@ class EventStreamProcessor:
         self.tsdb_batching_max_events = tsdb_batching_max_events
         self.tsdb_batching_timeout_secs = tsdb_batching_timeout_secs
 
-    def apply_monitoring_serving_graph(self, fn: mlrun.runtimes.ServingRuntime) -> None:
+    def apply_monitoring_serving_graph(
+        self,
+        fn: mlrun.runtimes.ServingRuntime,
+        tsdb_service_provider: typing.Optional[typing.Callable] = None,
+    ) -> None:
         """
         Apply monitoring serving graph to a given serving function. The following serving graph includes about 4 main
         parts that each one them includes several steps of different operations that are executed on the events from
@@ -163,6 +167,7 @@ class EventStreamProcessor:
            using CE, the parquet target path is based on the defined MLRun artifact path.
 
         :param fn: A serving function.
+        :param tsdb_service_provider: An optional callable function that provides the TSDB connection string.
         """
 
         graph = typing.cast(
@@ -322,15 +327,13 @@ class EventStreamProcessor:
 
         # TSDB branch (skip to Prometheus if in CE env)
         if not mlrun.mlconf.is_ce_mode():
-            # TSDB branch
             tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
-                project=self.project,
+                project=self.project, secret_provider=tsdb_service_provider
             )
             tsdb_connector.apply_monitoring_stream_steps(graph=graph)
 
         else:
             # Prometheus
-
             # Increase the prediction counter by 1 and update the latency value
             graph.add_step(
                 "IncCounter",

mlrun/model_monitoring/writer.py

@@ -17,7 +17,7 @@ from typing import Any, NewType
 
 import mlrun.common.model_monitoring
 import mlrun.common.schemas
-import mlrun.common.schemas.alert as alert_constants
+import mlrun.common.schemas.alert as alert_objects
 import mlrun.model_monitoring
 from mlrun.common.schemas.model_monitoring.constants import (
     EventFieldType,
@@ -29,7 +29,7 @@ from mlrun.common.schemas.model_monitoring.constants import (
     WriterEventKind,
 )
 from mlrun.common.schemas.notification import NotificationKind, NotificationSeverity
-from mlrun.model_monitoring.helpers import get_endpoint_record
+from mlrun.model_monitoring.helpers import get_endpoint_record, get_result_instance_fqn
 from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
 from mlrun.utils.notifications.notification_pusher import CustomNotificationPusher
@@ -101,7 +101,7 @@ class ModelMonitoringWriter(StepToDict):
 
     kind = "monitoring_application_stream_pusher"
 
-    def __init__(self, project: str) -> None:
+    def __init__(self, project: str, tsdb_secret_provider=None) -> None:
         self.project = project
         self.name = project  # required for the deployment process
 
@@ -113,24 +113,24 @@ class ModelMonitoringWriter(StepToDict):
             project=self.project
         )
         self._tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
-            project=self.project,
+            project=self.project, secret_provider=tsdb_secret_provider
         )
         self._endpoints_records = {}
 
     @staticmethod
     def _generate_event_on_drift(
-        model_endpoint: str, drift_status: str, event_value: dict, project_name: str
+        entity_id: str, drift_status: str, event_value: dict, project_name: str
     ) -> None:
-        logger.info("Sending an alert")
+        logger.info("Sending an event")
         entity = mlrun.common.schemas.alert.EventEntities(
-            kind=alert_constants.EventEntityKind.MODEL,
+            kind=alert_objects.EventEntityKind.MODEL_ENDPOINT_RESULT,
             project=project_name,
-            ids=[model_endpoint],
+            ids=[entity_id],
         )
         event_kind = (
-            alert_constants.EventKind.DRIFT_DETECTED
+            alert_objects.EventKind.DATA_DRIFT_DETECTED
             if drift_status == ResultStatusApp.detected.value
-            else alert_constants.EventKind.DRIFT_SUSPECTED
+            else alert_objects.EventKind.DATA_DRIFT_SUSPECTED
         )
         event_data = mlrun.common.schemas.Event(
             kind=event_kind, entity=entity, value_dict=event_value
@@ -138,7 +138,7 @@ class ModelMonitoringWriter(StepToDict):
         mlrun.get_run_db().generate_event(event_kind, event_data)
 
     @staticmethod
-    def _reconstruct_event(event: _RawEvent) -> tuple[_AppResultEvent, str]:
+    def _reconstruct_event(event: _RawEvent) -> tuple[_AppResultEvent, WriterEventKind]:
         """
         Modify the raw event into the expected monitoring application event
         schema as defined in `mlrun.common.schemas.model_monitoring.constants.WriterEvent`
@@ -179,12 +179,13 @@ class ModelMonitoringWriter(StepToDict):
     def do(self, event: _RawEvent) -> None:
         event, kind = self._reconstruct_event(event)
         logger.info("Starting to write event", event=event)
-
         self._tsdb_connector.write_application_event(event=event.copy(), kind=kind)
         self._app_result_store.write_application_event(event=event.copy(), kind=kind)
+
         logger.info("Completed event DB writes")
 
-        _Notifier(event=event, notification_pusher=self._custom_notifier).notify()
+        if kind == WriterEventKind.RESULT:
+            _Notifier(event=event, notification_pusher=self._custom_notifier).notify()
 
         if (
             mlrun.mlconf.alerts.mode == mlrun.common.schemas.alert.AlertsModes.enabled
@@ -208,7 +209,11 @@ class ModelMonitoringWriter(StepToDict):
                 "result_value": event[ResultData.RESULT_VALUE],
             }
             self._generate_event_on_drift(
-                event[WriterEvent.ENDPOINT_ID],
+                get_result_instance_fqn(
+                    event[WriterEvent.ENDPOINT_ID],
+                    event[WriterEvent.APPLICATION_NAME],
+                    event[ResultData.RESULT_NAME],
+                ),
                 event[ResultData.RESULT_STATUS],
                 event_value,
                 self.project,

mlrun/package/utils/_formatter.py

@@ -142,11 +142,11 @@ class _YAMLFormatter(_Formatter):
 
         :param obj:         The object to write.
         :param file_path:   The file path to write to.
-        :param dump_kwargs: Additional keyword arguments to pass to the `yaml.dump` method of the formatter in use.
+        :param dump_kwargs: Additional keyword arguments to pass to the `yaml.safe_dump` method of the formatter in use.
         """
         dump_kwargs = dump_kwargs or cls.DEFAULT_DUMP_KWARGS
         with open(file_path, "w") as file:
-            yaml.dump(obj, file, **dump_kwargs)
+            yaml.safe_dump(obj, file, **dump_kwargs)
 
     @classmethod
     def read(cls, file_path: str) -> Union[list, dict]:

mlrun/projects/project.py

@@ -39,6 +39,7 @@ import yaml
 from mlrun_pipelines.models import PipelineNodeWrapper
 
 import mlrun.common.helpers
+import mlrun.common.runtimes.constants
 import mlrun.common.schemas.artifact
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.db
@@ -2962,8 +2963,12 @@ class MlrunProject(ModelObj):
             engine = "remote"
         # The default engine is kfp if not given:
         workflow_engine = get_workflow_engine(engine or workflow_spec.engine, local)
-        if not inner_engine and engine == "remote":
-            inner_engine = get_workflow_engine(workflow_spec.engine, local).engine
+        if not inner_engine and workflow_engine.engine == "remote":
+            # if inner engine is set to remote, assume kfp as the default inner engine with remote as the runner
+            engine_kind = (
+                workflow_spec.engine if workflow_spec.engine != "remote" else "kfp"
+            )
+            inner_engine = get_workflow_engine(engine_kind, local).engine
         workflow_spec.engine = inner_engine or workflow_engine.engine
 
         run = workflow_engine.run(
@@ -2991,7 +2996,7 @@ class MlrunProject(ModelObj):
             # run's engine gets replaced with inner engine if engine is remote,
             # so in that case we need to get the status from the remote engine manually
             # TODO: support watch for remote:local
-            if engine == "remote" and status_engine.engine != "local":
+            if workflow_engine.engine == "remote" and status_engine.engine != "local":
                 status_engine = _RemoteRunner
 
             status_engine.get_run_status(project=self, run=run, timeout=timeout)
@@ -3094,17 +3099,18 @@ class MlrunProject(ModelObj):
 
     def set_model_monitoring_credentials(
         self,
-        access_key: str = None,
-        endpoint_store_connection: str = None,
-        stream_path: str = None,
+        access_key: Optional[str] = None,
+        endpoint_store_connection: Optional[str] = None,
+        stream_path: Optional[str] = None,
+        tsdb_connection: Optional[str] = None,
     ):
         """Set the credentials that will be used by the project's model monitoring
         infrastructure functions.
 
-        :param access_key:                Model Monitoring access key for managing user permissions
         :param access_key:                Model Monitoring access key for managing user permissions
         :param endpoint_store_connection: Endpoint store connection string
         :param stream_path:               Path to the model monitoring stream
+        :param tsdb_connection:           Connection string to the time series database
         """
 
         secrets_dict = {}
@@ -3127,6 +3133,16 @@ class MlrunProject(ModelObj):
                 mlrun.common.schemas.model_monitoring.ProjectSecretKeys.STREAM_PATH
             ] = stream_path
 
+        if tsdb_connection:
+            if not tsdb_connection.startswith("taosws://"):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Currently only TDEngine websocket connection is supported for non-v3io TSDB,"
+                    "please provide a full URL (e.g. taosws://user:password@host:port)"
+                )
+            secrets_dict[
+                mlrun.common.schemas.model_monitoring.ProjectSecretKeys.TSDB_CONNECTION
+            ] = tsdb_connection
+
         self.set_secrets(
             secrets=secrets_dict,
             provider=mlrun.common.schemas.SecretProviderName.kubernetes,
@@ -3685,7 +3701,10 @@ class MlrunProject(ModelObj):
         name: Optional[str] = None,
         uid: Optional[Union[str, list[str]]] = None,
         labels: Optional[Union[str, list[str]]] = None,
-        state: Optional[str] = None,
+        state: Optional[
+            mlrun.common.runtimes.constants.RunStates
+        ] = None,  # Backward compatibility
+        states: typing.Optional[list[mlrun.common.runtimes.constants.RunStates]] = None,
         sort: bool = True,
         last: int = 0,
         iter: bool = False,
@@ -3719,10 +3738,11 @@ class MlrunProject(ModelObj):
         :param labels:  A list of labels to filter by. Label filters work by either filtering a specific value
                 of a label (i.e. list("key=value")) or by looking for the existence of a given
                 key (i.e. "key").
-        :param state: List only runs whose state is specified.
+        :param state: Deprecated - List only runs whose state is specified.
+        :param states: List only runs whose state is one of the provided states.
         :param sort: Whether to sort the result according to their start time. Otherwise, results will be
             returned by their internal order in the DB (order will not be guaranteed).
-        :param last: Deprecated - currently not used (will be removed in 1.8.0).
+        :param last: Deprecated - currently not used (will be removed in 1.9.0).
         :param iter: If ``True`` return runs from all iterations. Otherwise, return only runs whose ``iter`` is 0.
         :param start_time_from: Filter by run start time in ``[start_time_from, start_time_to]``.
         :param start_time_to: Filter by run start time in ``[start_time_from, start_time_to]``.
@@ -3730,13 +3750,22 @@ class MlrunProject(ModelObj):
             last_update_time_to)``.
         :param last_update_time_to: Filter by run last update time in ``(last_update_time_from, last_update_time_to)``.
         """
+        if state:
+            # TODO: Remove this in 1.9.0
+            warnings.warn(
+                "'state' is deprecated and will be removed in 1.9.0. Use 'states' instead.",
+                FutureWarning,
+            )
+
         db = mlrun.db.get_run_db(secrets=self._secrets)
         return db.list_runs(
             name,
             uid,
             self.metadata.name,
             labels=labels,
-            state=state,
+            states=mlrun.utils.helpers.as_list(state)
+            if state is not None
+            else states or None,
             sort=sort,
             last=last,
             iter=iter,

mlrun/render.py

@@ -126,7 +126,7 @@ def artifacts_html(
 
         if not attribute_value:
             mlrun.utils.logger.warning(
-                "Artifact is incomplete, omitting from output (most likely due to a failed artifact logging)",
+                f"Artifact required attribute {attribute_name} is missing, omitting from output",
                 artifact_key=key,
             )
             continue
@@ -400,14 +400,17 @@ def runs_to_html(
     else:
         df["labels"] = df["labels"].apply(dict_html)
         df["inputs"] = df["inputs"].apply(inputs_html)
-        if df["artifact_uris"][0]:
-            df["artifact_uris"] = df["artifact_uris"].apply(dict_html)
-            df.drop("artifacts", axis=1, inplace=True)
-        else:
+        if df["artifacts"][0]:
             df["artifacts"] = df["artifacts"].apply(
                 lambda artifacts: artifacts_html(artifacts, "target_path"),
             )
             df.drop("artifact_uris", axis=1, inplace=True)
+        elif df["artifact_uris"][0]:
+            df["artifact_uris"] = df["artifact_uris"].apply(dict_html)
+            df.drop("artifacts", axis=1, inplace=True)
+        else:
+            df.drop("artifacts", axis=1, inplace=True)
+            df.drop("artifact_uris", axis=1, inplace=True)
 
     def expand_error(x):
         if x["state"] == "error":

mlrun/runtimes/__init__.py

@@ -43,6 +43,7 @@ from .nuclio import (
     new_v2_model_server,
     nuclio_init_hook,
 )
+from .nuclio.api_gateway import APIGateway
 from .nuclio.application import ApplicationRuntime
 from .nuclio.serving import serving_subkind
 from .remotesparkjob import RemoteSparkRuntime

mlrun/runtimes/databricks_job/databricks_wrapper.py

@@ -99,7 +99,7 @@ def save_credentials(
         credentials["DATABRICKS_CLUSTER_ID"] = cluster_id
 
     with open(credentials_path, "w") as yaml_file:
-        yaml.dump(credentials, yaml_file, default_flow_style=False)
+        yaml.safe_dump(credentials, yaml_file, default_flow_style=False)
 
 
 def run_mlrun_databricks_job(