mlrun 1.7.0rc57__py3-none-any.whl → 1.7.1__py3-none-any.whl

mlrun/common/formatters/artifact.py

@@ -21,6 +21,7 @@ from .base import ObjectFormat
 
 
 class ArtifactFormat(ObjectFormat, mlrun.common.types.StrEnum):
+    full = "full"
     minimal = "minimal"
 
     @staticmethod

mlrun/common/formatters/feature_set.py

@@ -28,6 +28,17 @@ class FeatureSetFormat(ObjectFormat, mlrun.common.types.StrEnum):
         return {
             FeatureSetFormat.full: None,
             FeatureSetFormat.minimal: FeatureSetFormat.filter_obj_method(
-                ["kind", "metadata", "spec", "status.state"]
+                [
+                    "metadata.name",
+                    "metadata.project",
+                    "metadata.tag",
+                    "metadata.uid",
+                    "metadata.labels",
+                    "spec.entities",
+                    "spec.description",
+                    "spec.targets",
+                    "spec.engine",  # It's not needed by the UI, but we override it anyway to storey if empty
+                    "status.state",
+                ]
             ),
         }[_format]

mlrun/config.py

@@ -49,6 +49,7 @@ _load_lock = Lock()
 _none_type = type(None)
 default_env_file = os.getenv("MLRUN_DEFAULT_ENV_FILE", "~/.mlrun.env")
 
+
 default_config = {
     "namespace": "",  # default kubernetes namespace
     "kubernetes": {
@@ -532,8 +533,55 @@ default_config = {
         },
     },
     "model_endpoint_monitoring": {
-        "serving_stream_args": {"shard_count": 1, "retention_period_hours": 24},
-        "application_stream_args": {"shard_count": 1, "retention_period_hours": 24},
+        "serving_stream": {
+            "v3io": {
+                "shard_count": 2,
+                "retention_period_hours": 24,
+                "num_workers": 1,
+                "min_replicas": 2,
+                "max_replicas": 2,
+            },
+            "kafka": {
+                "partition_count": 8,
+                "replication_factor": 1,
+                "num_workers": 2,
+                "min_replicas": 1,
+                "max_replicas": 4,
+            },
+        },
+        "application_stream_args": {
+            "v3io": {
+                "shard_count": 1,
+                "retention_period_hours": 24,
+                "num_workers": 1,
+                "min_replicas": 1,
+                "max_replicas": 1,
+            },
+            "kafka": {
+                "partition_count": 1,
+                "replication_factor": 1,
+                "num_workers": 1,
+                "min_replicas": 1,
+                "max_replicas": 1,
+            },
+        },
+        "writer_stream_args": {
+            "v3io": {
+                "shard_count": 1,
+                "retention_period_hours": 24,
+                "num_workers": 1,
+                "min_replicas": 1,
+                "max_replicas": 1,
+            },
+            "kafka": {
+                "partition_count": 1,
+                # TODO: add retention period configuration
+                "replication_factor": 1,
+                "num_workers": 1,
+                "min_replicas": 1,
+                "max_replicas": 1,
+            },
+        },
         # Store prefixes are used to handle model monitoring storing policies based on project and kind, such as events,
         # stream, and endpoints.
         "store_prefixes": {
@@ -556,6 +604,10 @@ default_config = {
         "tsdb_connection": "",
         # See mlrun.common.schemas.model_monitoring.constants.StreamKind for available options
         "stream_connection": "",
+        "tdengine": {
+            "timeout": 10,
+            "retries": 1,
+        },
     },
     "secret_stores": {
         # Use only in testing scenarios (such as integration tests) to avoid using k8s for secrets (will use in-memory
@@ -746,7 +798,6 @@ default_config = {
         "request_timeout": 5,
     },
 }
-
 _is_running_as_api = None
 
 

mlrun/datastore/__init__.py

@@ -131,9 +131,9 @@ class _DummyStream:
     def __init__(self, event_list=None, **kwargs):
         self.event_list = event_list or []
 
-    def push(self, data):
+    def push(self, data, **kwargs):
         if not isinstance(data, list):
             data = [data]
         for item in data:
-            logger.info(f"dummy stream got event: {item}")
+            logger.info(f"dummy stream got event: {item}, kwargs={kwargs}")
             self.event_list.append(item)

mlrun/db/httpdb.py

@@ -1075,7 +1075,9 @@ class HTTPRunDB(RunDBInterface):
         category: Union[str, mlrun.common.schemas.ArtifactCategories] = None,
         tree: str = None,
         producer_uri: str = None,
-        format_: mlrun.common.formatters.ArtifactFormat = mlrun.common.formatters.ArtifactFormat.full,
+        format_: Optional[
+            mlrun.common.formatters.ArtifactFormat
+        ] = mlrun.common.formatters.ArtifactFormat.full,
         limit: int = None,
     ) -> ArtifactList:
         """List artifacts filtered by various parameters.

mlrun/features.py

@@ -100,7 +100,8 @@ class Feature(ModelObj):
         :param name:        name of the feature
         :param validator:   feature validation policy
         :param default:     default value
-        :param labels:      a set of key/value labels (tags)
+        :param labels:      a set of key/value labels (tags). Labels can be used to filter featues, for example,
+                            in the UI Feature store page.
         """
         self.name = name or ""
         if isinstance(value_type, ValueType):

mlrun/model_monitoring/applications/_application_steps.py

@@ -162,10 +162,17 @@ class _ApplicationErrorHandler(StepToDict):
         :param event: Application event.
         """
 
-        exception_with_trace = "".join(
-            traceback.format_exception(None, event.error, event.error.__traceback__)
-        )
-        logger.error(f"Error in application step: {exception_with_trace}")
+        error_data = {
+            "Endpoint ID": event.body.endpoint_id,
+            "Application Class": event.body.application_name,
+            "Error": "".join(
+                traceback.format_exception(None, event.error, event.error.__traceback__)
+            ),
+            "Timestamp": event.timestamp,
+        }
+        logger.error("Error in application step", **error_data)
+
+        error_data["Error"] = event.error
 
         event_data = alert_objects.Event(
             kind=alert_objects.EventKind.MM_APP_FAILED,
@@ -174,12 +181,7 @@ class _ApplicationErrorHandler(StepToDict):
                 project=self.project,
                 ids=[f"{self.project}_{event.body.application_name}"],
             ),
-            value_dict={
-                "Error": event.error,
-                "Timestamp": event.timestamp,
-                "Application Class": event.body.application_name,
-                "Endpoint ID": event.body.endpoint_id,
-            },
+            value_dict=error_data,
         )
 
         mlrun.get_run_db().generate_event(

mlrun/model_monitoring/applications/evidently_base.py

@@ -23,7 +23,7 @@ import mlrun.model_monitoring.applications.base as mm_base
 import mlrun.model_monitoring.applications.context as mm_context
 from mlrun.errors import MLRunIncompatibleVersionError
 
-SUPPORTED_EVIDENTLY_VERSION = semver.Version.parse("0.4.32")
+SUPPORTED_EVIDENTLY_VERSION = semver.Version.parse("0.4.39")
 
 
 def _check_evidently_version(*, cur: semver.Version, ref: semver.Version) -> None:

mlrun/model_monitoring/db/tsdb/tdengine/schemas.py

@@ -82,9 +82,10 @@ class TDEngineSchema:
         super_table: str,
         columns: dict[str, _TDEngineColumn],
         tags: dict[str, str],
+        project: str,
         database: Optional[str] = None,
     ):
-        self.super_table = super_table
+        self.super_table = f"{super_table}_{project.replace('-', '_')}"
         self.columns = columns
         self.tags = tags
         self.database = database or _MODEL_MONITORING_DATABASE
@@ -148,6 +149,9 @@ class TDEngineSchema:
     ) -> str:
         return f"DROP TABLE if EXISTS {self.database}.{subtable};"
 
+    def drop_supertable_query(self) -> str:
+        return f"DROP STABLE if EXISTS {self.database}.{self.super_table};"
+
     def _get_subtables_query(
         self,
         values: dict[str, Union[str, int, float, datetime.datetime]],
@@ -159,7 +163,7 @@ class TDEngineSchema:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"values must contain at least one tag: {self.tags.keys()}"
             )
-        return f"SELECT tbname FROM {self.database}.{self.super_table} WHERE {values};"
+        return f"SELECT DISTINCT tbname FROM {self.database}.{self.super_table} WHERE {values};"
 
     @staticmethod
     def _get_records_query(
@@ -227,7 +231,7 @@ class TDEngineSchema:
 
 @dataclass
 class AppResultTable(TDEngineSchema):
-    def __init__(self, database: Optional[str] = None):
+    def __init__(self, project: str, database: Optional[str] = None):
         super_table = mm_schemas.TDEngineSuperTables.APP_RESULTS
         columns = {
             mm_schemas.WriterEvent.END_INFER_TIME: _TDEngineColumn.TIMESTAMP,
@@ -236,18 +240,23 @@ class AppResultTable(TDEngineSchema):
             mm_schemas.ResultData.RESULT_STATUS: _TDEngineColumn.INT,
         }
         tags = {
-            mm_schemas.EventFieldType.PROJECT: _TDEngineColumn.BINARY_64,
             mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
             mm_schemas.WriterEvent.APPLICATION_NAME: _TDEngineColumn.BINARY_64,
             mm_schemas.ResultData.RESULT_NAME: _TDEngineColumn.BINARY_64,
             mm_schemas.ResultData.RESULT_KIND: _TDEngineColumn.INT,
         }
-        super().__init__(super_table, columns, tags, database)
+        super().__init__(
+            super_table=super_table,
+            columns=columns,
+            tags=tags,
+            database=database,
+            project=project,
+        )
 
 
 @dataclass
 class Metrics(TDEngineSchema):
-    def __init__(self, database: Optional[str] = None):
+    def __init__(self, project: str, database: Optional[str] = None):
         super_table = mm_schemas.TDEngineSuperTables.METRICS
         columns = {
             mm_schemas.WriterEvent.END_INFER_TIME: _TDEngineColumn.TIMESTAMP,
@@ -255,17 +264,22 @@ class Metrics(TDEngineSchema):
             mm_schemas.MetricData.METRIC_VALUE: _TDEngineColumn.FLOAT,
         }
         tags = {
-            mm_schemas.EventFieldType.PROJECT: _TDEngineColumn.BINARY_64,
             mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
             mm_schemas.WriterEvent.APPLICATION_NAME: _TDEngineColumn.BINARY_64,
             mm_schemas.MetricData.METRIC_NAME: _TDEngineColumn.BINARY_64,
         }
-        super().__init__(super_table, columns, tags, database)
+        super().__init__(
+            super_table=super_table,
+            columns=columns,
+            tags=tags,
+            database=database,
+            project=project,
+        )
 
 
 @dataclass
 class Predictions(TDEngineSchema):
-    def __init__(self, database: Optional[str] = None):
+    def __init__(self, project: str, database: Optional[str] = None):
         super_table = mm_schemas.TDEngineSuperTables.PREDICTIONS
         columns = {
             mm_schemas.EventFieldType.TIME: _TDEngineColumn.TIMESTAMP,
@@ -273,7 +287,12 @@ class Predictions(TDEngineSchema):
             mm_schemas.EventKeyMetrics.CUSTOM_METRICS: _TDEngineColumn.BINARY_10000,
         }
         tags = {
-            mm_schemas.EventFieldType.PROJECT: _TDEngineColumn.BINARY_64,
             mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
         }
-        super().__init__(super_table, columns, tags, database)
+        super().__init__(
+            super_table=super_table,
+            columns=columns,
+            tags=tags,
+            database=database,
+            project=project,
+        )

mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py

@@ -56,6 +56,9 @@ class TDEngineConnector(TSDBConnector):
         self._connection = None
         self._init_super_tables()
 
+        self._timeout = mlrun.mlconf.model_endpoint_monitoring.tdengine.timeout
+        self._retries = mlrun.mlconf.model_endpoint_monitoring.tdengine.retries
+
     @property
     def connection(self) -> TDEngineConnection:
         if not self._connection:
@@ -66,7 +69,11 @@ class TDEngineConnector(TSDBConnector):
         """Establish a connection to the TSDB server."""
         logger.debug("Creating a new connection to TDEngine", project=self.project)
         conn = TDEngineConnection(self._tdengine_connection_string)
-        conn.run(statements=f"CREATE DATABASE IF NOT EXISTS {self.database}")
+        conn.run(
+            statements=f"CREATE DATABASE IF NOT EXISTS {self.database}",
+            timeout=self._timeout,
+            retries=self._retries,
+        )
         conn.prefix_statements = [f"USE {self.database}"]
         logger.debug("Connected to TDEngine", project=self.project)
         return conn
@@ -75,13 +82,13 @@ class TDEngineConnector(TSDBConnector):
         """Initialize the super tables for the TSDB."""
         self.tables = {
             mm_schemas.TDEngineSuperTables.APP_RESULTS: tdengine_schemas.AppResultTable(
-                self.database
+                project=self.project, database=self.database
             ),
             mm_schemas.TDEngineSuperTables.METRICS: tdengine_schemas.Metrics(
-                self.database
+                project=self.project, database=self.database
             ),
             mm_schemas.TDEngineSuperTables.PREDICTIONS: tdengine_schemas.Predictions(
-                self.database
+                project=self.project, database=self.database
             ),
         }
 
@@ -89,7 +96,11 @@ class TDEngineConnector(TSDBConnector):
         """Create TDEngine supertables."""
         for table in self.tables:
             create_table_query = self.tables[table]._create_super_table_query()
-            self.connection.run(statements=create_table_query)
+            self.connection.run(
+                statements=create_table_query,
+                timeout=self._timeout,
+                retries=self._retries,
+            )
 
     def write_application_event(
         self,
@@ -101,11 +112,9 @@ class TDEngineConnector(TSDBConnector):
         """
 
         table_name = (
-            f"{self.project}_"
             f"{event[mm_schemas.WriterEvent.ENDPOINT_ID]}_"
-            f"{event[mm_schemas.WriterEvent.APPLICATION_NAME]}_"
+            f"{event[mm_schemas.WriterEvent.APPLICATION_NAME]}"
         )
-        event[mm_schemas.EventFieldType.PROJECT] = self.project
 
         if kind == mm_schemas.WriterEventKind.RESULT:
             # Write a new result
@@ -145,7 +154,9 @@ class TDEngineConnector(TSDBConnector):
             statements=[
                 create_table_sql,
                 insert_statement,
-            ]
+            ],
+            timeout=self._timeout,
+            retries=self._retries,
         )
 
     @staticmethod
@@ -174,7 +185,9 @@ class TDEngineConnector(TSDBConnector):
                 name=name,
                 after=after,
                 url=self._tdengine_connection_string,
-                supertable=mm_schemas.TDEngineSuperTables.PREDICTIONS,
+                supertable=self.tables[
+                    mm_schemas.TDEngineSuperTables.PREDICTIONS
+                ].super_table,
                 table_col=mm_schemas.EventFieldType.TABLE_COLUMN,
                 time_col=mm_schemas.EventFieldType.TIME,
                 database=self.database,
@@ -207,17 +220,24 @@ class TDEngineConnector(TSDBConnector):
             "Deleting all project resources using the TDEngine connector",
             project=self.project,
         )
+        drop_statements = []
         for table in self.tables:
-            get_subtable_names_query = self.tables[table]._get_subtables_query(
-                values={mm_schemas.EventFieldType.PROJECT: self.project}
+            drop_statements.append(self.tables[table].drop_supertable_query())
+
+        try:
+            self.connection.run(
+                statements=drop_statements,
+                timeout=self._timeout,
+                retries=self._retries,
+            )
+        except Exception as e:
+            logger.warning(
+                "Failed to drop TDEngine tables. You may need to drop them manually. "
+                "These can be found under the following supertables: app_results, "
+                "metrics, and predictions.",
+                project=self.project,
+                error=mlrun.errors.err_to_str(e),
             )
-            subtables = self.connection.run(query=get_subtable_names_query).data
-            drop_statements = []
-            for subtable in subtables:
-                drop_statements.append(
-                    self.tables[table]._drop_subtable_query(subtable=subtable[0])
-                )
-            self.connection.run(statements=drop_statements)
         logger.debug(
             "Deleted all project resources using the TDEngine connector",
             project=self.project,
@@ -269,13 +289,6 @@ class TDEngineConnector(TSDBConnector):
         :raise:  MLRunInvalidArgumentError if query the provided table failed.
         """
 
-        project_condition = f"project = '{self.project}'"
-        filter_query = (
-            f"({filter_query}) AND ({project_condition})"
-            if filter_query
-            else project_condition
-        )
-
         full_query = tdengine_schemas.TDEngineSchema._get_records_query(
             table=table,
             start=start,
@@ -291,7 +304,9 @@ class TDEngineConnector(TSDBConnector):
         )
         logger.debug("Querying TDEngine", query=full_query)
         try:
-            query_result = self.connection.run(query=full_query)
+            query_result = self.connection.run(
+                query=full_query, timeout=self._timeout, retries=self._retries
+            )
         except taosws.QueryError as e:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"Failed to query table {table} in database {self.database}, {str(e)}"
@@ -325,12 +340,12 @@ class TDEngineConnector(TSDBConnector):
         timestamp_column = mm_schemas.WriterEvent.END_INFER_TIME
         columns = [timestamp_column, mm_schemas.WriterEvent.APPLICATION_NAME]
         if type == "metrics":
-            table = mm_schemas.TDEngineSuperTables.METRICS
+            table = self.tables[mm_schemas.TDEngineSuperTables.METRICS].super_table
             name = mm_schemas.MetricData.METRIC_NAME
             columns += [name, mm_schemas.MetricData.METRIC_VALUE]
             df_handler = self.df_to_metrics_values
         elif type == "results":
-            table = mm_schemas.TDEngineSuperTables.APP_RESULTS
+            table = self.tables[mm_schemas.TDEngineSuperTables.APP_RESULTS].super_table
             name = mm_schemas.ResultData.RESULT_NAME
             columns += [
                 name,
@@ -396,7 +411,7 @@ class TDEngineConnector(TSDBConnector):
                 "both or neither of `aggregation_window` and `agg_funcs` must be provided"
             )
         df = self._get_records(
-            table=mm_schemas.TDEngineSuperTables.PREDICTIONS,
+            table=self.tables[mm_schemas.TDEngineSuperTables.PREDICTIONS].super_table,
             start=start,
             end=end,
             columns=[mm_schemas.EventFieldType.LATENCY],

mlrun/platforms/iguazio.py

@@ -97,34 +97,43 @@ class OutputStream:
 
         self._v3io_client = v3io.dataplane.Client(**v3io_client_kwargs)
         self._container, self._stream_path = split_path(stream_path)
+        self._shards = shards
+        self._retention_in_hours = retention_in_hours
+        self._create = create
+        self._endpoint = endpoint
         self._mock = mock
         self._mock_queue = []
 
-        if create and not mock:
-            # this import creates an import loop via the utils module, so putting it in execution path
-            from mlrun.utils.helpers import logger
+    def create_stream(self):
+        # this import creates an import loop via the utils module, so putting it in execution path
+        from mlrun.utils.helpers import logger
+
+        logger.debug(
+            "Creating output stream",
+            endpoint=self._endpoint,
+            container=self._container,
+            stream_path=self._stream_path,
+            shards=self._shards,
+            retention_in_hours=self._retention_in_hours,
+        )
+        response = self._v3io_client.stream.create(
+            container=self._container,
+            stream_path=self._stream_path,
+            shard_count=self._shards or 1,
+            retention_period_hours=self._retention_in_hours or 24,
+            raise_for_status=v3io.dataplane.RaiseForStatus.never,
+        )
+        if not (response.status_code == 400 and "ResourceInUse" in str(response.body)):
+            response.raise_for_status([409, 204])
 
-            logger.debug(
-                "Creating output stream",
-                endpoint=endpoint,
-                container=self._container,
-                stream_path=self._stream_path,
-                shards=shards,
-                retention_in_hours=retention_in_hours,
-            )
-            response = self._v3io_client.stream.create(
-                container=self._container,
-                stream_path=self._stream_path,
-                shard_count=shards or 1,
-                retention_period_hours=retention_in_hours or 24,
-                raise_for_status=v3io.dataplane.RaiseForStatus.never,
-            )
-            if not (
-                response.status_code == 400 and "ResourceInUse" in str(response.body)
-            ):
-                response.raise_for_status([409, 204])
+    def _lazy_init(self):
+        if self._create and not self._mock:
+            self._create = False
+            self.create_stream()
+
+    def push(self, data, partition_key=None):
+        self._lazy_init()
 
-    def push(self, data):
         def dump_record(rec):
             if not isinstance(rec, (str, bytes)):
                 return dict_to_json(rec)
@@ -132,7 +141,14 @@ class OutputStream:
 
         if not isinstance(data, list):
             data = [data]
-        records = [{"data": dump_record(rec)} for rec in data]
+
+        records = []
+        for rec in data:
+            record = {"data": dump_record(rec)}
+            if partition_key is not None:
+                record["partition_key"] = partition_key
+            records.append(record)
+
         if self._mock:
             # for mock testing
             self._mock_queue.extend(records)
@@ -205,7 +221,7 @@ class KafkaOutputStream:
 
         self._initialized = True
 
-    def push(self, data):
+    def push(self, data, partition_key=None):
         self._lazy_init()
 
         def dump_record(rec):
@@ -226,7 +242,11 @@ class KafkaOutputStream:
         else:
             for record in data:
                 serialized_record = dump_record(record)
-                self._kafka_producer.send(self._topic, serialized_record)
+                if isinstance(partition_key, str):
+                    partition_key = partition_key.encode("UTF-8")
+                self._kafka_producer.send(
+                    self._topic, serialized_record, key=partition_key
+                )
 
 
 class V3ioStreamClient:

mlrun/projects/__init__.py

@@ -27,7 +27,12 @@ __all__ = [
 ]
 
 from .operations import build_function, deploy_function, run_function  # noqa
-from .pipelines import load_and_run, pipeline_context  # noqa
+from .pipelines import (
+    import_remote_project,
+    load_and_run_workflow,
+    load_and_run,
+    pipeline_context,
+)  # noqa
 from .project import (
     MlrunProject,
     ProjectMetadata,