mlrun 1.10.0rc1__py3-none-any.whl → 1.10.0rc2__py3-none-any.whl

mlrun/model_monitoring/applications/evidently/base.py

@@ -14,19 +14,18 @@
 
 import json
 import posixpath
-import uuid
 import warnings
 from abc import ABC
+from tempfile import NamedTemporaryFile
+from typing import Optional
 
-import pandas as pd
 import semver
-from evidently.ui.storage.local.base import METADATA_PATH, FSLocation
 
 import mlrun.model_monitoring.applications.base as mm_base
 import mlrun.model_monitoring.applications.context as mm_context
-from mlrun.errors import MLRunIncompatibleVersionError
+from mlrun.errors import MLRunIncompatibleVersionError, MLRunValueError
 
-SUPPORTED_EVIDENTLY_VERSION = semver.Version.parse("0.6.0")
+SUPPORTED_EVIDENTLY_VERSION = semver.Version.parse("0.7.5")
 
 
 def _check_evidently_version(*, cur: semver.Version, ref: semver.Version) -> None:
@@ -60,36 +59,66 @@ except ModuleNotFoundError:
 
 
 if _HAS_EVIDENTLY:
-    from evidently.suite.base_suite import Display
-    from evidently.ui.type_aliases import STR_UUID
-    from evidently.ui.workspace import Workspace
-    from evidently.utils.dashboard import TemplateParams, file_html_template
+    from evidently.core.report import Snapshot
+    from evidently.legacy.ui.storage.local.base import METADATA_PATH, FSLocation
+    from evidently.ui.workspace import (
+        STR_UUID,
+        CloudWorkspace,
+        Project,
+        Workspace,
+        WorkspaceBase,
+    )
 
 
 class EvidentlyModelMonitoringApplicationBase(
     mm_base.ModelMonitoringApplicationBase, ABC
 ):
     def __init__(
-        self, evidently_workspace_path: str, evidently_project_id: "STR_UUID"
+        self,
+        evidently_project_id: "STR_UUID",
+        evidently_workspace_path: Optional[str] = None,
+        cloud_workspace: bool = False,
     ) -> None:
         """
-        A class for integrating Evidently for mlrun model monitoring within a monitoring application.
-        Note: evidently is not installed by default in the mlrun/mlrun image.
-        It must be installed separately to use this class.
+        A class for integrating Evidently for MLRun model monitoring within a monitoring application.
+
+        .. note::
+
+            The ``evidently`` package is not installed by default in the mlrun/mlrun image.
+            It must be installed separately to use this class.
 
-        :param evidently_workspace_path:    (str) The path to the Evidently workspace.
         :param evidently_project_id:        (str) The ID of the Evidently project.
+        :param evidently_workspace_path:    (str) The path to the Evidently workspace.
+        :param cloud_workspace:             (bool) Whether the workspace is an Evidently Cloud workspace.
         """
-
-        # TODO : more then one project (mep -> project)
         if not _HAS_EVIDENTLY:
             raise ModuleNotFoundError("Evidently is not installed - the app cannot run")
-        self._log_location(evidently_workspace_path)
-        self.evidently_workspace = Workspace.create(evidently_workspace_path)
+        self.evidently_workspace_path = evidently_workspace_path
+        if cloud_workspace:
+            self.get_workspace = self.get_cloud_workspace
+        self.evidently_workspace = self.get_workspace()
         self.evidently_project_id = evidently_project_id
-        self.evidently_project = self.evidently_workspace.get_project(
-            evidently_project_id
-        )
+        self.evidently_project = self.load_project()
+
+    def load_project(self) -> Project:
+        """Load the Evidently project."""
+        return self.evidently_workspace.get_project(self.evidently_project_id)
+
+    def get_workspace(self) -> WorkspaceBase:
+        """Get the Evidently workspace. Override this method for customize access to the workspace."""
+        if self.evidently_workspace_path:
+            self._log_location(self.evidently_workspace_path)
+            return Workspace.create(self.evidently_workspace_path)
+        else:
+            raise MLRunValueError(
+                "A local workspace could not be created as `evidently_workspace_path` is not set.\n"
+                "If you intend to use a cloud workspace, please use `cloud_workspace=True` and set the "
+                "`EVIDENTLY_API_KEY` environment variable. In other cases, override this method."
+            )
+
+    def get_cloud_workspace(self) -> CloudWorkspace:
+        """Load the Evidently cloud workspace according to the `EVIDENTLY_API_KEY` environment variable."""
+        return CloudWorkspace()
 
     @staticmethod
     def _log_location(evidently_workspace_path):
@@ -128,7 +157,7 @@ class EvidentlyModelMonitoringApplicationBase(
     @staticmethod
     def log_evidently_object(
         monitoring_context: mm_context.MonitoringApplicationContext,
-        evidently_object: "Display",
+        evidently_object: "Snapshot",
         artifact_name: str,
         unique_per_endpoint: bool = True,
     ) -> None:
@@ -141,56 +170,15 @@ class EvidentlyModelMonitoringApplicationBase(
             This method should be called on special occasions only.
 
         :param monitoring_context:  (MonitoringApplicationContext) The monitoring context to process.
-        :param evidently_object:    (Display) The Evidently display to log, e.g. a report or a test suite object.
-        :param artifact_name:       (str) The name for the logged artifact.
-        :param unique_per_endpoint: by default ``True``, we will log different artifact for each model endpoint,
-                                    set to ``False`` without changing item key will cause artifact override.
-        """
-        evidently_object_html = evidently_object.get_html()
-        monitoring_context.log_artifact(
-            artifact_name,
-            body=evidently_object_html.encode("utf-8"),
-            format="html",
-            unique_per_endpoint=unique_per_endpoint,
-        )
-
-    def log_project_dashboard(
-        self,
-        monitoring_context: mm_context.MonitoringApplicationContext,
-        timestamp_start: pd.Timestamp,
-        timestamp_end: pd.Timestamp,
-        artifact_name: str = "dashboard",
-        unique_per_endpoint: bool = True,
-    ) -> None:
-        """
-        Logs an Evidently project dashboard.
-
-        .. caution::
-
-            Logging Evidently dashboards in every model monitoring window may cause scale issues.
-            This method should be called on special occasions only.
-
-        :param monitoring_context:  (MonitoringApplicationContext) The monitoring context to process.
-        :param timestamp_start:     (pd.Timestamp) The start timestamp for the dashboard data.
-        :param timestamp_end:       (pd.Timestamp) The end timestamp for the dashboard data.
+        :param evidently_object:    (Snapshot) The Evidently run to log, e.g. a report run.
         :param artifact_name:       (str) The name for the logged artifact.
         :param unique_per_endpoint: by default ``True``, we will log different artifact for each model endpoint,
                                     set to ``False`` without changing item key will cause artifact override.
         """
-
-        dashboard_info = self.evidently_project.build_dashboard_info(
-            timestamp_start, timestamp_end
-        )
-        template_params = TemplateParams(
-            dashboard_id="pd_" + str(uuid.uuid4()).replace("-", ""),
-            dashboard_info=dashboard_info,
-            additional_graphs={},
-        )
-
-        dashboard_html = file_html_template(params=template_params)
-        monitoring_context.log_artifact(
-            artifact_name,
-            body=dashboard_html.encode("utf-8"),
-            format="html",
-            unique_per_endpoint=unique_per_endpoint,
-        )
+        with NamedTemporaryFile(suffix=".html") as file:
+            evidently_object.save_html(filename=file.name)
+            monitoring_context.log_artifact(
+                artifact_name,
+                local_path=file.name,
+                unique_per_endpoint=unique_per_endpoint,
+            )

mlrun/model_monitoring/controller.py

@@ -25,6 +25,7 @@ from types import TracebackType
 from typing import Any, NamedTuple, Optional, Union, cast
 
 import nuclio_sdk
+import pandas as pd
 
 import mlrun
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
@@ -673,9 +674,15 @@ class MonitoringApplicationController:
         """
         logger.info("Starting monitoring controller chief")
         applications_names = []
-        endpoints = self.project_obj.list_model_endpoints(
-            metric_list=["last_request"]
-        ).endpoints
+        endpoints = self.project_obj.list_model_endpoints(tsdb_metrics=False).endpoints
+        last_request_dict = self.tsdb_connector.get_last_request(
+            endpoint_ids=[mep.metadata.uid for mep in endpoints]
+        )
+        if isinstance(last_request_dict, pd.DataFrame):
+            last_request_dict = last_request_dict.set_index(
+                mm_constants.EventFieldType.ENDPOINT_ID
+            )[mm_constants.ModelEndpointSchema.LAST_REQUEST].to_dict()
+
         if not endpoints:
             logger.info("No model endpoints found", project=self.project)
             return
@@ -721,16 +728,22 @@ class MonitoringApplicationController:
             with schedules.ModelMonitoringSchedulesFileChief(
                 self.project
             ) as schedule_file:
-                futures = {
-                    pool.submit(
-                        self.endpoint_to_regular_event,
-                        endpoint,
-                        policy,
-                        set(applications_names),
-                        schedule_file,
-                    ): endpoint
-                    for endpoint in endpoints
-                }
+                for endpoint in endpoints:
+                    last_request = last_request_dict.get(endpoint.metadata.uid, None)
+                    if isinstance(last_request, float):
+                        last_request = pd.to_datetime(last_request, unit="s", utc=True)
+                    endpoint.status.last_request = (
+                        last_request or endpoint.status.last_request
+                    )
+                    futures = {
+                        pool.submit(
+                            self.endpoint_to_regular_event,
+                            endpoint,
+                            policy,
+                            set(applications_names),
+                            schedule_file,
+                        ): endpoint
+                    }
                 for future in concurrent.futures.as_completed(futures):
                     if future.exception():
                         exception = future.exception()

mlrun/model_monitoring/db/tsdb/v3io/v3io_connector.py

@@ -455,12 +455,20 @@ class V3IOTSDBConnector(TSDBConnector):
             # Delete all tables
             tables = mm_schemas.V3IOTSDBTables.list()
         for table_to_delete in tables:
-            try:
-                self.frames_client.delete(backend=_TSDB_BE, table=table_to_delete)
-            except v3io_frames.DeleteError as e:
+            if table_to_delete in self.tables:
+                try:
+                    self.frames_client.delete(
+                        backend=_TSDB_BE, table=self.tables[table_to_delete]
+                    )
+                except v3io_frames.DeleteError as e:
+                    logger.warning(
+                        f"Failed to delete TSDB table '{table_to_delete}'",
+                        err=mlrun.errors.err_to_str(e),
+                    )
+            else:
                 logger.warning(
-                    f"Failed to delete TSDB table '{table}'",
-                    err=mlrun.errors.err_to_str(e),
+                    f"Skipping deletion: table '{table_to_delete}' is not among the initialized tables.",
+                    initialized_tables=list(self.tables.keys()),
                 )
 
         # Final cleanup of tsdb path

mlrun/projects/project.py

@@ -30,6 +30,7 @@ from copy import deepcopy
 from os import environ, makedirs, path
 from typing import Callable, Optional, Union, cast
 
+import deprecated
 import dotenv
 import git
 import git.exc
@@ -2860,6 +2861,20 @@ class MlrunProject(ModelObj):
 
         self.spec.set_function(name, function_object, func)
 
+    # TODO: Remove this in 1.11.0
+    @deprecated.deprecated(
+        version="1.8.0",
+        reason="'remove_function' is deprecated and will be removed in 1.11.0. "
+        "Please use `delete_function` instead.",
+        category=FutureWarning,
+    )
+    def remove_function(self, name):
+        """remove the specified function from the project
+
+        :param name:    name of the function (under the project)
+        """
+        self.spec.remove_function(name)
+
     def delete_function(self, name, delete_from_db=False):
         """deletes the specified function from the project
 
@@ -3695,7 +3710,7 @@ class MlrunProject(ModelObj):
                 brokers=["<kafka-broker-ip-address>:9094"],
                 topics=[],  # Keep the topics list empty
                 ## SASL is supported
-                # sasl_user="user1",
+                # sasl_user="<kafka-sasl-user>",
                 # sasl_pass="<kafka-sasl-password>",
             )
             project.register_datastore_profile(stream_profile)
@@ -3728,6 +3743,29 @@ class MlrunProject(ModelObj):
 
         In the V3IO datastore, you must provide an explicit access key to the stream, but not to the TSDB.
 
+        An external Confluent Kafka stream is also supported. Here is an example:
+
+        .. code-block:: python
+
+            from mlrun.datastore.datastore_profile import DatastoreProfileKafkaSource
+
+            stream_profile = DatastoreProfileKafkaSource(
+                name="confluent-kafka",
+                brokers=["<server-domain-start>.confluent.cloud:9092"],
+                topics=[],
+                sasl_user="<API-key>",
+                sasl_pass="<API-secret>",
+                kwargs_public={
+                    "security_protocol": "SASL_SSL",
+                    "api_version_auto_timeout_ms": 15_000,  # 15 seconds
+                    "tls": {"enable": True},
+                    "new_topic": {"replication_factor": 3},
+                },
+            )
+
+        The replication factor and timeout configuration might need to be adjusted according to your Confluent cluster
+        type and settings.
+
         :param tsdb_profile_name:         The datastore profile name of the time-series database to be used in model
                                           monitoring. The supported profiles are:
 
@@ -3781,7 +3819,7 @@ class MlrunProject(ModelObj):
         top_level: bool = False,
         uids: Optional[list[str]] = None,
         latest_only: bool = False,
-        tsdb_metrics: bool = True,
+        tsdb_metrics: bool = False,
         metric_list: Optional[list[str]] = None,
     ) -> mlrun.common.schemas.ModelEndpointList:
         """
@@ -4003,8 +4041,10 @@ class MlrunProject(ModelObj):
             e.g. builder_env={"GIT_TOKEN": token}, does not work yet in KFP
         :param overwrite_build_params:  Overwrite existing build configuration (currently applies to
             requirements and commands)
+
             * False: The new params are merged with the existing
             * True: The existing params are replaced by the new ones
+
         :param extra_args:  A string containing additional builder arguments in the format of command-line options,
             e.g. extra_args="--skip-tls-verify --build-arg A=val"
         :param force_build:  force building the image, even when no changes were made
@@ -4055,8 +4095,10 @@ class MlrunProject(ModelObj):
         :param requirements_file: requirements file to install on the built image
         :param overwrite_build_params:  Overwrite existing build configuration (currently applies to
             requirements and commands)
+
             * False: The new params are merged with the existing
             * True: The existing params are replaced by the new ones
+
         :param builder_env: Kaniko builder pod env vars dict (for config/credentials)
             e.g. builder_env={"GIT_TOKEN": token}, does not work yet in KFP
         :param extra_args:  A string containing additional builder arguments in the format of command-line options,
@@ -4128,8 +4170,10 @@ class MlrunProject(ModelObj):
             e.g. builder_env={"GIT_TOKEN": token}, does not work yet in KFP
         :param overwrite_build_params:  Overwrite existing build configuration (currently applies to
             requirements and commands)
+
             * False: The new params are merged with the existing
             * True: The existing params are replaced by the new ones
+
         :param extra_args:  A string containing additional builder arguments in the format of command-line options,
             e.g. extra_args="--skip-tls-verify --build-arg A=val"
         :param target_dir: Path on the image where source code would be extracted (by default `/home/mlrun_code`)
@@ -4307,12 +4351,14 @@ class MlrunProject(ModelObj):
             ``my_Name_1`` or ``surname``.
         :param tag: Return artifacts assigned this tag.
         :param labels: Filter artifacts by label key-value pairs or key existence. This can be provided as:
-            - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
-            or `{"label": None}` to check for key existence.
-            - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
-            or just `"label"` for key existence.
-            - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
-            the specified key-value pairs or key existence.
+
+                       - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
+                         or `{"label": None}` to check for key existence.
+                       - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
+                         or just `"label"` for key existence.
+                       - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
+                         the specified key-value pairs or key existence.
+
         :param since: Not in use in :py:class:`HTTPRunDB`.
         :param until: Not in use in :py:class:`HTTPRunDB`.
         :param iter: Return artifacts from a specific iteration (where ``iter=0`` means the root iteration). If
@@ -4323,7 +4369,7 @@ class MlrunProject(ModelObj):
         :param kind: Return artifacts of the requested kind.
         :param category: Return artifacts of the requested category.
         :param tree: Return artifacts of the requested tree.
-        :param limit: Maximum number of artifacts to return.
+        :param limit: Deprecated - Maximum number of artifacts to return (will be removed in 1.11.0).
         :param format_: The format in which to return the artifacts. Default is 'full'.
         :param partition_by: Field to group results by. When `partition_by` is specified, the `partition_sort_by`
             parameter must be provided as well.
@@ -4334,6 +4380,14 @@ class MlrunProject(ModelObj):
         :param partition_order: Order of sorting within partitions - `asc` or `desc`. Default is `desc`.
         """
         db = mlrun.db.get_run_db(secrets=self._secrets)
+
+        if limit:
+            # TODO: Remove this in 1.11.0
+            warnings.warn(
+                "'limit' is deprecated and will be removed in 1.11.0. Use 'page' and 'page_size' instead.",
+                FutureWarning,
+            )
+
         return db.list_artifacts(
             name,
             self.metadata.name,
@@ -4446,12 +4500,14 @@ class MlrunProject(ModelObj):
             ``my_Name_1`` or ``surname``.
         :param tag: Return artifacts assigned this tag.
         :param labels: Filter model artifacts by label key-value pairs or key existence. This can be provided as:
-            - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
-            or `{"label": None}` to check for key existence.
-            - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
-            or just `"label"` for key existence.
-            - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
-            the specified key-value pairs or key existence.
+
+                       - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
+                         or `{"label": None}` to check for key existence.
+                       - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
+                         or just `"label"` for key existence.
+                       - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
+                         the specified key-value pairs or key existence.
+
         :param since: Not in use in :py:class:`HTTPRunDB`.
         :param until: Not in use in :py:class:`HTTPRunDB`.
         :param iter: Return artifacts from a specific iteration (where ``iter=0`` means the root iteration). If
@@ -4557,12 +4613,14 @@ class MlrunProject(ModelObj):
         :param name: Return only functions with a specific name.
         :param tag: Return function versions with specific tags. To return only tagged functions, set tag to ``"*"``.
         :param labels: Filter functions by label key-value pairs or key existence. This can be provided as:
-            - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
-            or `{"label": None}` to check for key existence.
-            - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
-            or just `"label"` for key existence.
-            - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
-            the specified key-value pairs or key existence.
+
+                       - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
+                         or `{"label": None}` to check for key existence.
+                       - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
+                         or just `"label"` for key existence.
+                       - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
+                         the specified key-value pairs or key existence.
+
         :param kind: Return functions of the specified kind. If not provided, all function kinds will be returned.
         :param format_: The format in which to return the functions. Default is 'full'.
         :returns: List of function objects.
@@ -4656,12 +4714,14 @@ class MlrunProject(ModelObj):
         :param name:    Return only functions with a specific name.
         :param tag:     Return function versions with specific tags.
         :param labels: Filter functions by label key-value pairs or key existence. This can be provided as:
-            - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
-            or `{"label": None}` to check for key existence.
-            - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
-            or just `"label"` for key existence.
-            - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
-            the specified key-value pairs or key existence.
+
+                       - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
+                         or `{"label": None}` to check for key existence.
+                       - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
+                         or just `"label"` for key existence.
+                       - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
+                         the specified key-value pairs or key existence.
+
         :returns: List of function objects.
         """
 
@@ -4717,12 +4777,14 @@ class MlrunProject(ModelObj):
         :param name: Name of the run to retrieve.
         :param uid: Unique ID of the run.
         :param labels: Filter runs by label key-value pairs or key existence. This can be provided as:
-            - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
-            or `{"label": None}` to check for key existence.
-            - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
-            or just `"label"` for key existence.
-            - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
-            the specified key-value pairs or key existence.
+
+                       - A dictionary in the format `{"label": "value"}` to match specific label key-value pairs,
+                         or `{"label": None}` to check for key existence.
+                       - A list of strings formatted as `"label=value"` to match specific label key-value pairs,
+                         or just `"label"` for key existence.
+                       - A comma-separated string formatted as `"label1=value1,label2"` to match entities with
+                         the specified key-value pairs or key existence.
+
         :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

mlrun/render.py

@@ -361,9 +361,6 @@ def get_tblframe(df, display, classes=None):
     return ipython_display(html, display)
 
 
-uid_template = '<div title="{}"><a href="{}/{}/{}/jobs/monitor/{}/overview" target="_blank" >...{}</a></div>'
-
-
 def runs_to_html(
     df: pd.DataFrame,
     display: bool = True,
@@ -379,15 +376,14 @@ def runs_to_html(
     df["results"] = df["results"].apply(dict_html)
     df["start"] = df["start"].apply(time_str)
     df["parameters"] = df["parameters"].apply(dict_html)
+    uid_template = '<div title="{}"><a href="{}" target="_blank" >...{}</a></div>'
+
     if config.resolve_ui_url():
         df["uid"] = df.apply(
             lambda x: uid_template.format(
-                x.uid,
-                config.resolve_ui_url(),
-                config.ui.projects_prefix,
-                x.project,
-                x.uid,
-                x.uid[-8:],
+                x["uid"],
+                mlrun.utils.get_run_url(x["project"], x["uid"], x["name"]),
+                x["uid"][-8:],
             ),
             axis=1,
         )

mlrun/runtimes/base.py

@@ -499,7 +499,7 @@ class BaseRuntime(ModelObj):
     def _store_function(self, runspec, meta, db):
         meta.labels["kind"] = self.kind
         mlrun.runtimes.utils.enrich_run_labels(
-            meta.labels, [mlrun.common.runtimes.constants.RunLabels.owner]
+            meta.labels, [mlrun_constants.MLRunInternalLabels.owner]
         )
         if runspec.spec.output_path:
             runspec.spec.output_path = runspec.spec.output_path.replace(

mlrun/runtimes/utils.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import enum
 import getpass
 import hashlib
 import json
@@ -28,7 +29,6 @@ import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas
 import mlrun.utils.regex
 from mlrun.artifacts import TableArtifact
-from mlrun.common.runtimes.constants import RunLabels
 from mlrun.config import config
 from mlrun.errors import err_to_str
 from mlrun.frameworks.parallel_coordinates import gen_pcp_plot
@@ -433,18 +433,35 @@ def enrich_function_from_dict(function, function_dict):
 
 def enrich_run_labels(
     labels: dict,
-    labels_to_enrich: Optional[list[RunLabels]] = None,
+    labels_to_enrich: Optional[list[mlrun_constants.MLRunInternalLabels]] = None,
 ):
+    """
+    Enrich the run labels with the internal labels and the labels enrichment extension
+    :param labels: The run labels dict
+    :param labels_to_enrich: The label keys to enrich from MLRunInternalLabels.default_run_labels_to_enrich
+    :return: The enriched labels dict
+    """
+    # Merge the labels with the labels enrichment extension
     labels_enrichment = {
-        RunLabels.owner: os.environ.get("V3IO_USERNAME") or getpass.getuser(),
-        # TODO: remove this in 1.9.0
-        RunLabels.v3io_user: os.environ.get("V3IO_USERNAME"),
+        mlrun_constants.MLRunInternalLabels.owner: os.environ.get("V3IO_USERNAME")
+        or getpass.getuser(),
+        # TODO: remove this in 1.10.0
+        mlrun_constants.MLRunInternalLabels.v3io_user: os.environ.get("V3IO_USERNAME"),
     }
-    labels_to_enrich = labels_to_enrich or RunLabels.all()
+
+    # Resolve which label keys to enrich
+    if labels_to_enrich is None:
+        labels_to_enrich = (
+            mlrun_constants.MLRunInternalLabels.default_run_labels_to_enrich()
+        )
+
+    # Enrich labels
     for label in labels_to_enrich:
+        if isinstance(label, enum.Enum):
+            label = label.value
         enrichment = labels_enrichment.get(label)
-        if label.value not in labels and enrichment:
-            labels[label.value] = enrichment
+        if label not in labels and enrichment:
+            labels[label] = enrichment
     return labels