mlrun 1.7.0rc17__py3-none-any.whl → 1.7.0rc19__py3-none-any.whl

mlrun/execution.py

@@ -22,6 +22,7 @@ import yaml
 from dateutil import parser
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 from mlrun.artifacts import ModelArtifact
 from mlrun.datastore.store_resources import get_store_resource
 from mlrun.errors import MLRunInvalidArgumentError
@@ -129,7 +130,9 @@ class MLClientCtx:
     @property
     def tag(self):
         """Run tag (uid or workflow id if exists)"""
-        return self._labels.get("workflow") or self._uid
+        return (
+            self._labels.get(mlrun_constants.MLRunInternalLabels.workflow) or self._uid
+        )
 
     @property
     def state(self):
@@ -329,8 +332,10 @@ class MLClientCtx:
             "uri": uri,
             "owner": get_in(self._labels, "owner"),
         }
-        if "workflow" in self._labels:
-            resp["workflow"] = self._labels["workflow"]
+        if mlrun_constants.MLRunInternalLabels.workflow in self._labels:
+            resp[mlrun_constants.MLRunInternalLabels.workflow] = self._labels[
+                mlrun_constants.MLRunInternalLabels.workflow
+            ]
         return resp
 
     @classmethod
@@ -396,7 +401,7 @@ class MLClientCtx:
                         self._set_input(k, v)
 
         if host and not is_api:
-            self.set_label("host", host)
+            self.set_label(mlrun_constants.MLRunInternalLabels.host, host)
 
         start = get_in(attrs, "status.start_time")
         if start:
@@ -990,10 +995,15 @@ class MLClientCtx:
         # If it's a OpenMPI job, get the global rank and compare to the logging rank (worker) set in MLRun's
         # configuration:
         labels = self.labels
-        if "host" in labels and labels.get("kind", "job") == "mpijob":
+        if (
+            mlrun_constants.MLRunInternalLabels.host in labels
+            and labels.get(mlrun_constants.MLRunInternalLabels.kind, "job") == "mpijob"
+        ):
             # The host (pod name) of each worker is created by k8s, and by default it uses the rank number as the id in
             # the following template: ...-worker-<rank>
-            rank = int(labels["host"].rsplit("-", 1)[1])
+            rank = int(
+                labels[mlrun_constants.MLRunInternalLabels.host].rsplit("-", 1)[1]
+            )
             return rank == mlrun.mlconf.packagers.logging_worker
 
         # Single worker is always the logging worker:

mlrun/feature_store/ingestion.py

@@ -17,6 +17,7 @@ import uuid
 import pandas as pd
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 from mlrun.datastore.sources import get_source_from_dict, get_source_step
 from mlrun.datastore.targets import (
     add_target_steps,
@@ -263,13 +264,13 @@ def run_ingestion_job(name, featureset, run_config, schedule=None, spark_service
         out_path=featureset.spec.output_path,
     )
     task.spec.secret_sources = run_config.secret_sources
-    task.set_label("job-type", "feature-ingest").set_label(
-        "feature-set", featureset.uri
-    )
+    task.set_label(
+        mlrun_constants.MLRunInternalLabels.job_type, "feature-ingest"
+    ).set_label("feature-set", featureset.uri)
     if run_config.owner:
-        task.set_label("owner", run_config.owner).set_label(
-            "v3io_user", run_config.owner
-        )
+        task.set_label(
+            mlrun_constants.MLRunInternalLabels.owner, run_config.owner
+        ).set_label(mlrun_constants.MLRunInternalLabels.v3io_user, run_config.owner)
 
     # set run UID and save in the feature set status (linking the features et to the job)
     task.metadata.uid = uuid.uuid4().hex

mlrun/feature_store/retrieval/conversion.py

@@ -168,10 +168,10 @@ class PandasConversionMixin:
         column_counter = Counter(self.columns)
 
         dtype = [None] * len(self.schema)
-        for fieldIdx, field in enumerate(self.schema):
+        for field_idx, field in enumerate(self.schema):
             # For duplicate column name, we use `iloc` to access it.
             if column_counter[field.name] > 1:
-                pandas_col = pdf.iloc[:, fieldIdx]
+                pandas_col = pdf.iloc[:, field_idx]
             else:
                 pandas_col = pdf[field.name]
 
@@ -187,12 +187,12 @@ class PandasConversionMixin:
                 and field.nullable
                 and pandas_col.isnull().any()
             ):
-                dtype[fieldIdx] = pandas_type
+                dtype[field_idx] = pandas_type
             # Ensure we fall back to nullable numpy types, even when whole column is null:
             if isinstance(field.dataType, IntegralType) and pandas_col.isnull().any():
-                dtype[fieldIdx] = np.float64
+                dtype[field_idx] = np.float64
             if isinstance(field.dataType, BooleanType) and pandas_col.isnull().any():
-                dtype[fieldIdx] = object
+                dtype[field_idx] = object
 
         df = pd.DataFrame()
         for index, t in enumerate(dtype):

mlrun/feature_store/retrieval/job.py

@@ -15,6 +15,7 @@
 import uuid
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 from mlrun.config import config as mlconf
 from mlrun.model import DataTargetBase, new_task
 from mlrun.runtimes.function_reference import FunctionReference
@@ -122,7 +123,9 @@ def run_merge_job(
         inputs={"entity_rows": entity_rows} if entity_rows is not None else {},
     )
     task.spec.secret_sources = run_config.secret_sources
-    task.set_label("job-type", "feature-merge").set_label("feature-vector", vector.uri)
+    task.set_label(
+        mlrun_constants.MLRunInternalLabels.job_type, "feature-merge"
+    ).set_label(mlrun_constants.MLRunInternalLabels.feature_vector, vector.uri)
     task.metadata.uid = uuid.uuid4().hex
     vector.status.run_uri = task.metadata.uid
     vector.save()
@@ -198,7 +201,8 @@ import mlrun.feature_store.retrieval
 from mlrun.datastore.targets import get_target_driver
 def merge_handler(context, vector_uri, target, entity_rows=None, 
                   entity_timestamp_column=None, drop_columns=None, with_indexes=None, query=None,
-                  engine_args=None, order_by=None, start_time=None, end_time=None, timestamp_for_filtering=None):
+                  engine_args=None, order_by=None, start_time=None, end_time=None, timestamp_for_filtering=None,
+                  additional_filters=None):
     vector = context.get_store_resource(vector_uri)
     store_target = get_target_driver(target, vector)
     if entity_rows:
@@ -208,7 +212,7 @@ def merge_handler(context, vector_uri, target, entity_rows=None,
     merger = mlrun.feature_store.retrieval.{{{engine}}}(vector, **(engine_args or {}))
     merger.start(entity_rows, entity_timestamp_column, store_target, drop_columns, with_indexes=with_indexes, 
                  query=query, order_by=order_by, start_time=start_time, end_time=end_time,
-                 timestamp_for_filtering=timestamp_for_filtering)
+                 timestamp_for_filtering=timestamp_for_filtering, additional_filters=additional_filters)
 
     target = vector.status.targets[store_target.name].to_dict()
     context.log_result('feature_vector', vector.uri)

mlrun/feature_store/retrieval/spark_merger.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
 import pandas as pd
 import semver
 
@@ -252,13 +253,13 @@ class SparkFeatureMerger(BaseMerger):
         # handling case where there are multiple feature sets and user creates vector where
         # entity_timestamp_column is from a specific feature set (can't be entity timestamp)
         source_driver = mlrun.datastore.sources.source_kind_to_driver[source_kind]
-
         source = source_driver(
             name=self.vector.metadata.name,
             path=source_path,
             time_field=time_column,
             start_time=start_time,
             end_time=end_time,
+            additional_filters=additional_filters,
             **source_kwargs,
         )
 

mlrun/frameworks/_dl_common/loggers/tensorboard_logger.py

@@ -648,13 +648,13 @@ class TensorboardLogger(Logger, Generic[DLTypes.WeightType]):
         if isinstance(value, list):
             if len(value) == 0:
                 return ""
-            text = "\n" + yaml.dump(value)
+            text = "\n" + yaml.safe_dump(value)
             text = "  \n".join(["  " * tabs + line for line in text.splitlines()])
             return text
         if isinstance(value, dict):
             if len(value) == 0:
                 return ""
-            text = yaml.dump(value)
+            text = yaml.safe_dump(value)
             text = "  \n".join(
                 ["  " * tabs + "- " + line for line in text.splitlines()]
             )

mlrun/frameworks/parallel_coordinates.py

@@ -295,7 +295,7 @@ def compare_db_runs(
     iter=False,
     start_time_from: datetime = None,
     hide_identical: bool = True,
-    exclude: list = [],
+    exclude: list = None,
     show=None,
     colorscale: str = "Blues",
     filename=None,
@@ -332,6 +332,7 @@ def compare_db_runs(
         **query_args,
     )
 
+    exclude = exclude or []
     runs_df = _runs_list_to_df(runs_list)
     plot_as_html = gen_pcp_plot(
         runs_df,

mlrun/frameworks/tf_keras/__init__.py

@@ -18,6 +18,7 @@ from typing import Any, Union
 from tensorflow import keras
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 
 from .callbacks import MLRunLoggingCallback, TensorboardLoggingCallback
 from .mlrun_interface import TFKerasMLRunInterface
@@ -126,7 +127,9 @@ def apply_mlrun(
     # # Use horovod:
     if use_horovod is None:
         use_horovod = (
-            context.labels.get("kind", "") == "mpijob" if context is not None else False
+            context.labels.get(mlrun_constants.MLRunInternalLabels.kind, "") == "mpijob"
+            if context is not None
+            else False
         )
 
     # Create a model handler:

mlrun/launcher/client.py

@@ -16,6 +16,7 @@ from typing import Optional
 
 import IPython
 
+import mlrun.common.constants as mlrun_constants
 import mlrun.errors
 import mlrun.launcher.base as launcher
 import mlrun.lists
@@ -69,13 +70,14 @@ class ClientBaseLauncher(launcher.BaseLauncher, abc.ABC):
     def _store_function(
         runtime: "mlrun.runtimes.BaseRuntime", run: "mlrun.run.RunObject"
     ):
-        run.metadata.labels["kind"] = runtime.kind
+        run.metadata.labels[mlrun_constants.MLRunInternalLabels.kind] = runtime.kind
         mlrun.runtimes.utils.enrich_run_labels(
             run.metadata.labels, [mlrun.common.runtimes.constants.RunLabels.owner]
         )
         if run.spec.output_path:
             run.spec.output_path = run.spec.output_path.replace(
-                "{{run.user}}", run.metadata.labels["owner"]
+                "{{run.user}}",
+                run.metadata.labels[mlrun_constants.MLRunInternalLabels.owner],
             )
         db = runtime._get_db()
         if db and runtime.kind != "handler":

mlrun/launcher/local.py

@@ -15,6 +15,7 @@ import os
 import pathlib
 from typing import Callable, Optional, Union
 
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas.schedule
 import mlrun.errors
 import mlrun.launcher.client as launcher
@@ -132,8 +133,13 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
         runtime: "mlrun.runtimes.BaseRuntime",
         run: Optional[Union["mlrun.run.RunTemplate", "mlrun.run.RunObject"]] = None,
     ):
-        if "V3IO_USERNAME" in os.environ and "v3io_user" not in run.metadata.labels:
-            run.metadata.labels["v3io_user"] = os.environ.get("V3IO_USERNAME")
+        if (
+            "V3IO_USERNAME" in os.environ
+            and mlrun_constants.MLRunInternalLabels.v3io_user not in run.metadata.labels
+        ):
+            run.metadata.labels[mlrun_constants.MLRunInternalLabels.v3io_user] = (
+                os.environ.get("V3IO_USERNAME")
+            )
 
         # store function object in db unless running from within a run pod
         if not runtime.is_child:

mlrun/launcher/remote.py

@@ -17,6 +17,7 @@ from typing import Optional, Union
 import pandas as pd
 import requests
 
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas.schedule
 import mlrun.db
 import mlrun.errors
@@ -100,8 +101,13 @@ class ClientRemoteLauncher(launcher.ClientBaseLauncher):
         if runtime.verbose:
             logger.info(f"runspec:\n{run.to_yaml()}")
 
-        if "V3IO_USERNAME" in os.environ and "v3io_user" not in run.metadata.labels:
-            run.metadata.labels["v3io_user"] = os.environ.get("V3IO_USERNAME")
+        if (
+            "V3IO_USERNAME" in os.environ
+            and mlrun_constants.MLRunInternalLabels.v3io_user not in run.metadata.labels
+        ):
+            run.metadata.labels[mlrun_constants.MLRunInternalLabels.v3io_user] = (
+                os.environ.get("V3IO_USERNAME")
+            )
 
         logger.info(
             "Storing function",

mlrun/model.py

@@ -27,6 +27,7 @@ from typing import Any, Optional, Union
 import pydantic.error_wrappers
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas.notification
 
 from .utils import (
@@ -770,7 +771,10 @@ class RunMetadata(ModelObj):
     def is_workflow_runner(self):
         if not self.labels:
             return False
-        return self.labels.get("job-type", "") == "workflow-runner"
+        return (
+            self.labels.get(mlrun_constants.MLRunInternalLabels.job_type, "")
+            == "workflow-runner"
+        )
 
 
 class HyperParamStrategies:

mlrun/model_monitoring/db/stores/__init__.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-
 import enum
 import typing
 import warnings

mlrun/model_monitoring/db/stores/base/store.py

@@ -15,7 +15,7 @@
 import typing
 from abc import ABC, abstractmethod
 
-import mlrun.common.schemas.model_monitoring.constants as mm_constants
+import mlrun.common.schemas.model_monitoring as mm_schemas
 
 
 class StoreBase(ABC):
@@ -115,8 +115,8 @@ class StoreBase(ABC):
     def write_application_event(
         self,
         event: dict[str, typing.Any],
-        kind: mm_constants.WriterEventKind = mm_constants.WriterEventKind.RESULT,
-    ):
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ) -> None:
         """
         Write a new event in the target table.
 
@@ -125,7 +125,6 @@ class StoreBase(ABC):
                       object.
         :param kind: The type of the event, can be either "result" or "metric".
         """
-        pass
 
     @abstractmethod
     def get_last_analyzed(self, endpoint_id: str, application_name: str) -> int:
@@ -157,3 +156,16 @@ class StoreBase(ABC):
 
         """
         pass
+
+    @abstractmethod
+    def get_model_endpoint_metrics(
+        self, endpoint_id: str, type: mm_schemas.ModelEndpointMonitoringMetricType
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
+        """
+        Get the model monitoring results and metrics of the requested model endpoint.
+
+        :param: endpoint_id: The model endpoint identifier.
+        :param: type:        The type of the requested metrics ("result" or "metric").
+
+        :return:             A list of the available metrics.
+        """

mlrun/model_monitoring/db/stores/sqldb/models/__init__.py

@@ -12,38 +12,60 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Union
+from functools import partial
+from typing import Optional, TypeVar, Union
 
+from .mysql import ApplicationMetricsTable as MySQLApplicationMetricsTable
 from .mysql import ApplicationResultTable as MySQLApplicationResultTable
 from .mysql import ModelEndpointsTable as MySQLModelEndpointsTable
 from .mysql import MonitoringSchedulesTable as MySQLMonitoringSchedulesTable
+from .sqlite import ApplicationMetricsTable as SQLiteApplicationMetricsTable
 from .sqlite import ApplicationResultTable as SQLiteApplicationResultTable
 from .sqlite import ModelEndpointsTable as SQLiteModelEndpointsTable
 from .sqlite import MonitoringSchedulesTable as SQLiteMonitoringSchedulesTable
 
+MySQLTableType = TypeVar("MySQLTableType")
+SQLiteTableType = TypeVar("SQLiteTableType")
 
-def _get_model_endpoints_table(
-    connection_string: Optional[str] = None,
-) -> Union[type[MySQLModelEndpointsTable], type[SQLiteModelEndpointsTable]]:
-    """Return ModelEndpointsTable based on the provided connection string"""
-    if connection_string and "mysql:" in connection_string:
-        return MySQLModelEndpointsTable
-    return SQLiteModelEndpointsTable
+_MYSQL_SCHEME = "mysql:"
 
 
-def _get_application_result_table(
+def _get_sql_table(
+    *,
+    mysql_table: MySQLTableType,
+    sqlite_table: SQLiteTableType,
     connection_string: Optional[str] = None,
-) -> Union[type[MySQLApplicationResultTable], type[SQLiteApplicationResultTable]]:
-    """Return ModelEndpointsTable based on the provided connection string"""
-    if connection_string and "mysql:" in connection_string:
-        return MySQLApplicationResultTable
-    return SQLiteApplicationResultTable
+) -> Union[MySQLTableType, SQLiteTableType]:
+    """
+    Return a SQLAlchemy table for MySQL or SQLite according to the connection string.
+    Note: this function should not be directly used in other modules.
+    """
+    if connection_string and _MYSQL_SCHEME in connection_string:
+        return mysql_table
+    return sqlite_table
 
 
-def _get_monitoring_schedules_table(
-    connection_string: Optional[str] = None,
-) -> Union[type[MySQLMonitoringSchedulesTable], type[SQLiteMonitoringSchedulesTable]]:
-    """Return ModelEndpointsTable based on the provided connection string"""
-    if connection_string and "mysql:" in connection_string:
-        return MySQLMonitoringSchedulesTable
-    return SQLiteMonitoringSchedulesTable
+_get_model_endpoints_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLModelEndpointsTable,
+    sqlite_table=SQLiteModelEndpointsTable,
+)
+
+
+_get_application_result_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLApplicationResultTable,
+    sqlite_table=SQLiteApplicationResultTable,
+)
+
+_get_application_metrics_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLApplicationMetricsTable,
+    sqlite_table=SQLiteApplicationMetricsTable,
+)
+
+_get_monitoring_schedules_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLMonitoringSchedulesTable,
+    sqlite_table=SQLiteMonitoringSchedulesTable,
+)

mlrun/model_monitoring/db/stores/sqldb/models/base.py

@@ -24,6 +24,7 @@ from sqlalchemy import (
 from mlrun.common.schemas.model_monitoring import (
     EventFieldType,
     FileTargetKind,
+    MetricData,
     ResultData,
     SchedulingKeys,
     WriterEvent,
@@ -89,11 +90,11 @@ class ModelEndpointsBaseTable(BaseModel):
     metrics = Column(EventFieldType.METRICS, Text)
     first_request = Column(
         EventFieldType.FIRST_REQUEST,
-        TIMESTAMP,
+        TIMESTAMP(timezone=True),
     )
     last_request = Column(
         EventFieldType.LAST_REQUEST,
-        TIMESTAMP,
+        TIMESTAMP(timezone=True),
     )
 
 
@@ -135,6 +136,35 @@ class ApplicationResultBaseTable(BaseModel):
     current_stats = Column(ResultData.CURRENT_STATS, Text)
 
 
+class ApplicationMetricsBaseTable(BaseModel):
+    __tablename__ = FileTargetKind.APP_METRICS
+
+    uid = Column(EventFieldType.UID, String(120), primary_key=True)
+    application_name = Column(
+        WriterEvent.APPLICATION_NAME,
+        String(40),
+        nullable=True,
+    )
+    endpoint_id = Column(
+        WriterEvent.ENDPOINT_ID,
+        String(40),
+        nullable=True,
+    )
+    start_infer_time = Column(
+        WriterEvent.START_INFER_TIME,
+        TIMESTAMP(timezone=True),
+    )
+    end_infer_time = Column(
+        WriterEvent.END_INFER_TIME,
+        TIMESTAMP(timezone=True),
+    )
+    metric_name = Column(
+        MetricData.METRIC_NAME,
+        String(40),
+    )
+    metric_value = Column(MetricData.METRIC_VALUE, Float)
+
+
 class MonitoringSchedulesBaseTable(BaseModel):
     __tablename__ = FileTargetKind.MONITORING_SCHEDULES
 

mlrun/model_monitoring/db/stores/sqldb/models/mysql.py

@@ -22,6 +22,7 @@ from mlrun.common.schemas.model_monitoring import (
 )
 
 from .base import (
+    ApplicationMetricsBaseTable,
     ApplicationResultBaseTable,
     ModelEndpointsBaseTable,
     MonitoringSchedulesBaseTable,
@@ -33,22 +34,29 @@ Base = declarative_base()
 class ModelEndpointsTable(Base, ModelEndpointsBaseTable):
     first_request = Column(
         EventFieldType.FIRST_REQUEST,
-        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3),
+        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3, timezone=True),
     )
     last_request = Column(
         EventFieldType.LAST_REQUEST,
-        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3),
+        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3, timezone=True),
     )
 
 
-class ApplicationResultTable(Base, ApplicationResultBaseTable):
+class _ApplicationResultOrMetric:
+    """
+    This class sets common columns of `ApplicationResultTable` and `ApplicationMetricsTable`
+    to the correct values in MySQL.
+    Note: This class must come before the base tables in the inheritance order to override
+    the relevant columns.
+    """
+
     start_infer_time = Column(
         WriterEvent.START_INFER_TIME,
-        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3),
+        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3, timezone=True),
     )
     end_infer_time = Column(
         WriterEvent.END_INFER_TIME,
-        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3),
+        sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3, timezone=True),
     )
 
     @declared_attr
@@ -59,6 +67,18 @@ class ApplicationResultTable(Base, ApplicationResultBaseTable):
         )
 
 
+class ApplicationResultTable(
+    Base, _ApplicationResultOrMetric, ApplicationResultBaseTable
+):
+    pass
+
+
+class ApplicationMetricsTable(
+    Base, _ApplicationResultOrMetric, ApplicationMetricsBaseTable
+):
+    pass
+
+
 class MonitoringSchedulesTable(Base, MonitoringSchedulesBaseTable):
     @declared_attr
     def endpoint_id(cls):

mlrun/model_monitoring/db/stores/sqldb/models/sqlite.py

@@ -15,6 +15,7 @@
 from sqlalchemy.ext.declarative import declarative_base
 
 from .base import (
+    ApplicationMetricsBaseTable,
     ApplicationResultBaseTable,
     ModelEndpointsBaseTable,
     MonitoringSchedulesBaseTable,
@@ -31,5 +32,9 @@ class ApplicationResultTable(Base, ApplicationResultBaseTable):
     pass
 
 
+class ApplicationMetricsTable(Base, ApplicationMetricsBaseTable):
+    pass
+
+
 class MonitoringSchedulesTable(Base, MonitoringSchedulesBaseTable):
     pass