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

mlrun/projects/pipelines.py

@@ -20,9 +20,10 @@ import tempfile
 import typing
 import uuid
 
-import kfp.compiler
-from kfp import dsl
+import mlrun_pipelines.common.models
+import mlrun_pipelines.patcher
 from kfp.compiler import compiler
+from mlrun_pipelines.helpers import new_pipe_metadata
 
 import mlrun
 import mlrun.common.schemas
@@ -31,7 +32,6 @@ from mlrun.errors import err_to_str
 from mlrun.utils import (
     get_ui_url,
     logger,
-    new_pipe_metadata,
     normalize_workflow_name,
     retry_until_successful,
 )
@@ -301,72 +301,6 @@ def _enrich_kfp_pod_security_context(kfp_pod_template, function):
     }
 
 
-# When we run pipelines, the kfp.compile.Compile.compile() method takes the decorated function with @dsl.pipeline and
-# converts it to a k8s object. As part of the flow in the Compile.compile() method,
-# we call _create_and_write_workflow, which builds a dictionary from the workflow and then writes it to a file.
-# Unfortunately, the kfp sdk does not provide an API for configuring priority_class_name and other attributes.
-# I ran across the following problem when seeking for a method to set the priority_class_name:
-# https://github.com/kubeflow/pipelines/issues/3594
-# When we patch the _create_and_write_workflow, we can eventually obtain the dictionary right before we write it
-# to a file and enrich it with argo compatible fields, make sure you looking for the same argo version we use
-# https://github.com/argoproj/argo-workflows/blob/release-2.7/pkg/apis/workflow/v1alpha1/workflow_types.go
-def _create_enriched_mlrun_workflow(
-    self,
-    pipeline_func: typing.Callable,
-    pipeline_name: typing.Optional[str] = None,
-    pipeline_description: typing.Optional[str] = None,
-    params_list: typing.Optional[list[dsl.PipelineParam]] = None,
-    pipeline_conf: typing.Optional[dsl.PipelineConf] = None,
-):
-    """Call internal implementation of create_workflow and enrich with mlrun functions attributes"""
-    workflow = self._original_create_workflow(
-        pipeline_func, pipeline_name, pipeline_description, params_list, pipeline_conf
-    )
-    # We don't want to interrupt the original flow and don't know all the scenarios the function could be called.
-    # that's why we have try/except on all the code of the enrichment and also specific try/except for errors that
-    # we know can be raised.
-    try:
-        functions = []
-        if pipeline_context.functions:
-            try:
-                functions = pipeline_context.functions.values()
-            except Exception as err:
-                logger.debug(
-                    "Unable to retrieve project functions, not enriching workflow with mlrun",
-                    error=err_to_str(err),
-                )
-                return workflow
-
-        # enrich each pipeline step with your desire k8s attribute
-        for kfp_step_template in workflow["spec"]["templates"]:
-            if kfp_step_template.get("container"):
-                for function_obj in functions:
-                    # we condition within each function since the comparison between the function and
-                    # the kfp pod may change depending on the attribute type.
-                    _set_function_attribute_on_kfp_pod(
-                        kfp_step_template,
-                        function_obj,
-                        "PriorityClassName",
-                        "priority_class_name",
-                    )
-                    _enrich_kfp_pod_security_context(
-                        kfp_step_template,
-                        function_obj,
-                    )
-    except mlrun.errors.MLRunInvalidArgumentError:
-        raise
-    except Exception as err:
-        logger.debug(
-            "Something in the enrichment of kfp pods failed", error=err_to_str(err)
-        )
-    return workflow
-
-
-# patching function as class method
-kfp.compiler.Compiler._original_create_workflow = kfp.compiler.Compiler._create_workflow
-kfp.compiler.Compiler._create_workflow = _create_enriched_mlrun_workflow
-
-
 def get_db_function(project, key) -> mlrun.runtimes.BaseRuntime:
     project_instance, name, tag, hash_key = parse_versioned_object_uri(
         key, project.metadata.name
@@ -457,7 +391,10 @@ class _PipelineRunStatus:
 
     @property
     def state(self):
-        if self._state not in mlrun.run.RunStatuses.stable_statuses():
+        if (
+            self._state
+            not in mlrun_pipelines.common.models.RunStatuses.stable_statuses()
+        ):
             self._state = self._engine.get_state(self.run_id, self.project)
         return self._state
 
@@ -754,7 +691,7 @@ class _LocalRunner(_PipelineRunner):
         err = None
         try:
             workflow_handler(**workflow_spec.args)
-            state = mlrun.run.RunStatuses.succeeded
+            state = mlrun_pipelines.common.models.RunStatuses.succeeded
         except Exception as exc:
             err = exc
             logger.exception("Workflow run failed")
@@ -762,7 +699,7 @@ class _LocalRunner(_PipelineRunner):
                 f":x: Workflow {workflow_id} run failed!, error: {err_to_str(exc)}",
                 mlrun.common.schemas.NotificationSeverity.ERROR,
             )
-            state = mlrun.run.RunStatuses.failed
+            state = mlrun_pipelines.common.models.RunStatuses.failed
         mlrun.run.wait_for_runs_completion(pipeline_context.runs_map.values())
         project.notifiers.push_pipeline_run_results(
             pipeline_context.runs_map.values(), state=state
@@ -921,9 +858,9 @@ class _RemoteRunner(_PipelineRunner):
                 f":x: Workflow {workflow_name} run failed!, error: {err_to_str(exc)}",
                 mlrun.common.schemas.NotificationSeverity.ERROR,
             )
-            state = mlrun.run.RunStatuses.failed
+            state = mlrun_pipelines.common.models.RunStatuses.failed
         else:
-            state = mlrun.run.RunStatuses.succeeded
+            state = mlrun_pipelines.common.models.RunStatuses.succeeded
             project.notifiers.push_pipeline_start_message(
                 project.metadata.name,
             )
@@ -1116,7 +1053,7 @@ def load_and_run(
     context.log_result(key="workflow_id", value=run.run_id)
     context.log_result(key="engine", value=run._engine.engine, commit=True)
 
-    if run.state == mlrun.run.RunStatuses.failed:
+    if run.state == mlrun_pipelines.common.models.RunStatuses.failed:
         raise RuntimeError(f"Workflow {workflow_log_message} failed") from run.exc
 
     if wait_for_completion:
@@ -1131,7 +1068,7 @@ def load_and_run(
 
         pipeline_state, _, _ = project.get_run_status(run)
         context.log_result(key="workflow_state", value=pipeline_state, commit=True)
-        if pipeline_state != mlrun.run.RunStatuses.succeeded:
+        if pipeline_state != mlrun_pipelines.common.models.RunStatuses.succeeded:
             raise RuntimeError(
                 f"Workflow {workflow_log_message} failed, state={pipeline_state}"
             )

mlrun/projects/project.py

@@ -31,23 +31,28 @@ from typing import Callable, Optional, Union
 import dotenv
 import git
 import git.exc
-import kfp
+import mlrun_pipelines.common.models
+import mlrun_pipelines.mounts
 import nuclio.utils
 import requests
 import yaml
+from mlrun_pipelines.models import PipelineNodeWrapper
 
 import mlrun.common.helpers
+import mlrun.common.schemas.artifact
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.db
 import mlrun.errors
 import mlrun.k8s_utils
+import mlrun.model_monitoring.applications as mm_app
 import mlrun.runtimes
 import mlrun.runtimes.nuclio.api_gateway
 import mlrun.runtimes.pod
 import mlrun.runtimes.utils
 import mlrun.serving
 import mlrun.utils.regex
-from mlrun.common.schemas import AlertConfig
+from mlrun.alerts.alert import AlertConfig
+from mlrun.common.schemas.alert import AlertTemplate
 from mlrun.datastore.datastore_profile import DatastoreProfile, DatastoreProfile2Json
 from mlrun.runtimes.nuclio.function import RemoteRuntime
 
@@ -56,14 +61,10 @@ from ..artifacts.manager import ArtifactManager, dict_to_artifact, extend_artifa
 from ..datastore import store_manager
 from ..features import Feature
 from ..model import EntrypointParam, ImageBuilder, ModelObj
-from ..model_monitoring.application import (
-    ModelMonitoringApplicationBase,
-)
 from ..run import code_to_function, get_object, import_function, new_function
 from ..secrets import SecretsStore
 from ..utils import (
     is_ipython,
-    is_legacy_artifact,
     is_relative_path,
     is_yaml_path,
     logger,
@@ -991,13 +992,9 @@ class ProjectSpec(ModelObj):
             if not isinstance(artifact, dict) and not hasattr(artifact, "to_dict"):
                 raise ValueError("artifacts must be a dict or class")
             if isinstance(artifact, dict):
-                # Support legacy artifacts
-                if is_legacy_artifact(artifact) or _is_imported_artifact(artifact):
-                    key = artifact.get("key")
-                else:
-                    key = artifact.get("metadata").get("key", "")
+                key = artifact.get("metadata", {}).get("key", "")
                 if not key:
-                    raise ValueError('artifacts "key" must be specified')
+                    raise ValueError('artifacts "metadata.key" must be specified')
             else:
                 key = artifact.key
                 artifact = artifact.to_dict()
@@ -1593,6 +1590,23 @@ class MlrunProject(ModelObj):
         )
         return item
 
+    def delete_artifact(
+        self,
+        item: Artifact,
+        deletion_strategy: mlrun.common.schemas.artifact.ArtifactsDeletionStrategies = (
+            mlrun.common.schemas.artifact.ArtifactsDeletionStrategies.metadata_only
+        ),
+        secrets: dict = None,
+    ):
+        """Delete an artifact object in the DB and optionally delete the artifact data
+
+        :param item: Artifact object (can be any type, such as dataset, model, feature store).
+        :param deletion_strategy: The artifact deletion strategy types.
+        :param secrets: Credentials needed to access the artifact data.
+        """
+        am = self._get_artifact_manager()
+        am.delete_artifact(item, deletion_strategy, secrets)
+
     def log_dataset(
         self,
         key,
@@ -1885,7 +1899,11 @@ class MlrunProject(ModelObj):
     def set_model_monitoring_function(
         self,
         func: typing.Union[str, mlrun.runtimes.BaseRuntime, None] = None,
-        application_class: typing.Union[str, ModelMonitoringApplicationBase] = None,
+        application_class: typing.Union[
+            str,
+            mm_app.ModelMonitoringApplicationBase,
+            mm_app.ModelMonitoringApplicationBaseV2,
+        ] = None,
         name: str = None,
         image: str = None,
         handler=None,
@@ -1923,11 +1941,6 @@ class MlrunProject(ModelObj):
                                         monitoring application's constructor.
         """
 
-        if name in mm_constants.MonitoringFunctionNames.list():
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"An application cannot have the following names: "
-                f"{mm_constants.MonitoringFunctionNames.list()}"
-            )
         function_object: RemoteRuntime = None
         (
             resolved_function_name,
@@ -1953,7 +1966,11 @@ class MlrunProject(ModelObj):
     def create_model_monitoring_function(
         self,
         func: str = None,
-        application_class: typing.Union[str, ModelMonitoringApplicationBase] = None,
+        application_class: typing.Union[
+            str,
+            mm_app.ModelMonitoringApplicationBase,
+            mm_app.ModelMonitoringApplicationBaseV2,
+        ] = None,
         name: str = None,
         image: str = None,
         handler: str = None,
@@ -2006,7 +2023,10 @@ class MlrunProject(ModelObj):
         self,
         func: typing.Union[str, mlrun.runtimes.BaseRuntime, None] = None,
         application_class: typing.Union[
-            str, ModelMonitoringApplicationBase, None
+            str,
+            mm_app.ModelMonitoringApplicationBase,
+            mm_app.ModelMonitoringApplicationBaseV2,
+            None,
         ] = None,
         name: typing.Optional[str] = None,
         image: typing.Optional[str] = None,
@@ -2869,7 +2889,7 @@ class MlrunProject(ModelObj):
                           (which will be converted to the class using its `from_crontab` constructor),
                           see this link for help:
                           https://apscheduler.readthedocs.io/en/3.x/modules/triggers/cron.html#module-apscheduler.triggers.cron
-                          for using the pre-defined workflow's schedule, set `schedule=True`
+                          For using the pre-defined workflow's schedule, set `schedule=True`
         :param timeout:   Timeout in seconds to wait for pipeline completion (watch will be activated)
         :param source:    Source to use instead of the actual `project.spec.source` (used when engine is remote).
                           Can be a one of:
@@ -2880,10 +2900,11 @@ class MlrunProject(ModelObj):
                           For other engines the source is used to validate that the code is up-to-date.
         :param cleanup_ttl:
                           Pipeline cleanup ttl in secs (time to wait after workflow completion, at which point the
-                          Workflow and all its resources are deleted)
+                          workflow and all its resources are deleted)
         :param notifications:
                           List of notifications to send for workflow completion
-        :returns: Run id
+
+        :returns: ~py:class:`~mlrun.projects.pipelines._PipelineRunStatus` instance
         """
 
         arguments = arguments or {}
@@ -2957,7 +2978,7 @@ class MlrunProject(ModelObj):
             notifications=notifications,
         )
         # run is None when scheduling
-        if run and run.state == mlrun.run.RunStatuses.failed:
+        if run and run.state == mlrun_pipelines.common.models.RunStatuses.failed:
             return run
         if not workflow_spec.schedule:
             # Failure and schedule messages already logged
@@ -3134,7 +3155,7 @@ class MlrunProject(ModelObj):
         notifications: list[mlrun.model.Notification] = None,
         returns: Optional[list[Union[str, dict[str, str]]]] = None,
         builder_env: Optional[dict] = None,
-    ) -> typing.Union[mlrun.model.RunObject, kfp.dsl.ContainerOp]:
+    ) -> typing.Union[mlrun.model.RunObject, PipelineNodeWrapper]:
         """Run a local or remote task as part of a local/kubeflow pipeline
 
         example (use with project)::
@@ -3190,7 +3211,7 @@ class MlrunProject(ModelObj):
                                   artifact type can be given there. The artifact key must appear in the dictionary as
                                   "key": "the_key".
         :param builder_env: env vars dict for source archive config/credentials e.g. builder_env={"GIT_TOKEN": token}
-        :return: MLRun RunObject or KubeFlow containerOp
+        :return: MLRun RunObject or PipelineNodeWrapper
         """
         return run_function(
             function,
@@ -3233,7 +3254,7 @@ class MlrunProject(ModelObj):
         requirements_file: str = None,
         extra_args: str = None,
         force_build: bool = False,
-    ) -> typing.Union[BuildStatus, kfp.dsl.ContainerOp]:
+    ) -> typing.Union[BuildStatus, PipelineNodeWrapper]:
         """deploy ML function, build container with its dependencies
 
         :param function:            name of the function (in the project) or function object
@@ -3344,7 +3365,6 @@ class MlrunProject(ModelObj):
         image: str = None,
         set_as_default: bool = True,
         with_mlrun: bool = None,
-        skip_deployed: bool = False,
         base_image: str = None,
         commands: list = None,
         secret_name: str = None,
@@ -3355,7 +3375,7 @@ class MlrunProject(ModelObj):
         requirements_file: str = None,
         extra_args: str = None,
         target_dir: str = None,
-    ) -> typing.Union[BuildStatus, kfp.dsl.ContainerOp]:
+    ) -> typing.Union[BuildStatus, PipelineNodeWrapper]:
         """Builder docker image for the project, based on the project's build config. Parameters allow to override
         the build config.
         If the project has a source configured and pull_at_runtime is not configured, this source will be cloned to the
@@ -3365,7 +3385,6 @@ class MlrunProject(ModelObj):
                         used. If not set, the `mlconf.default_project_image_name` value will be used
         :param set_as_default: set `image` to be the project's default image (default False)
         :param with_mlrun:      add the current mlrun package to the container build
-        :param skip_deployed:   *Deprecated* parameter is ignored
         :param base_image:      base image name/path (commands and source code will be added to it) defaults to
                                 mlrun.mlconf.default_base_image
         :param commands:        list of docker build (RUN) commands e.g. ['pip install pandas']
@@ -3390,14 +3409,6 @@ class MlrunProject(ModelObj):
                 base_image=base_image,
             )
 
-        if skip_deployed:
-            warnings.warn(
-                "The 'skip_deployed' parameter is deprecated and will be removed in 1.7.0. "
-                "This parameter is ignored.",
-                # TODO: remove in 1.7.0
-                FutureWarning,
-            )
-
         if not overwrite_build_params:
             # TODO: change overwrite_build_params default to True in 1.8.0
             warnings.warn(
@@ -3475,7 +3486,7 @@ class MlrunProject(ModelObj):
         verbose: bool = None,
         builder_env: dict = None,
         mock: bool = None,
-    ) -> typing.Union[DeployStatus, kfp.dsl.ContainerOp]:
+    ) -> typing.Union[DeployStatus, PipelineNodeWrapper]:
         """deploy real-time (nuclio based) functions
 
         :param function:    name of the function (in the project) or function object
@@ -3660,9 +3671,7 @@ class MlrunProject(ModelObj):
         :returns: List of function objects.
         """
 
-        model_monitoring_labels_list = [
-            f"{mm_constants.ModelMonitoringAppLabel.KEY}={mm_constants.ModelMonitoringAppLabel.VAL}"
-        ]
+        model_monitoring_labels_list = [str(mm_constants.ModelMonitoringAppLabel())]
         if labels:
             model_monitoring_labels_list += labels
         return self.list_functions(
@@ -3889,7 +3898,9 @@ class MlrunProject(ModelObj):
 
         mlrun.db.get_run_db().delete_api_gateway(name=name, project=self.name)
 
-    def store_alert_config(self, alert_data: AlertConfig, alert_name=None):
+    def store_alert_config(
+        self, alert_data: AlertConfig, alert_name=None
+    ) -> AlertConfig:
         """
         Create/modify an alert.
         :param alert_data: The data of the alert.
@@ -3899,7 +3910,7 @@ class MlrunProject(ModelObj):
         db = mlrun.db.get_run_db(secrets=self._secrets)
         if alert_name is None:
             alert_name = alert_data.name
-        return db.store_alert_config(alert_name, alert_data.dict(), self.metadata.name)
+        return db.store_alert_config(alert_name, alert_data, project=self.metadata.name)
 
     def get_alert_config(self, alert_name: str) -> AlertConfig:
         """
@@ -3910,7 +3921,7 @@ class MlrunProject(ModelObj):
         db = mlrun.db.get_run_db(secrets=self._secrets)
         return db.get_alert_config(alert_name, self.metadata.name)
 
-    def list_alerts_configs(self):
+    def list_alerts_configs(self) -> list[AlertConfig]:
         """
         Retrieve list of alerts of a project.
         :return: All the alerts objects of the project.
@@ -3956,6 +3967,23 @@ class MlrunProject(ModelObj):
             alert_name = alert_data.name
         db.reset_alert_config(alert_name, self.metadata.name)
 
+    def get_alert_template(self, template_name: str) -> AlertTemplate:
+        """
+        Retrieve a specific alert template.
+        :param template_name: The name of the template to retrieve.
+        :return: The template object.
+        """
+        db = mlrun.db.get_run_db(secrets=self._secrets)
+        return db.get_alert_template(template_name)
+
+    def list_alert_templates(self) -> list[AlertTemplate]:
+        """
+        Retrieve list of all alert templates.
+        :return: All the alert template objects in the database.
+        """
+        db = mlrun.db.get_run_db(secrets=self._secrets)
+        return db.list_alert_templates()
+
     def _run_authenticated_git_action(
         self,
         action: Callable,

mlrun/render.py

@@ -121,16 +121,8 @@ def artifacts_html(
     html = ""
 
     for artifact in artifacts:
-        # TODO: remove this in 1.7.0 once we no longer support legacy format
-        if mlrun.utils.is_legacy_artifact(artifact):
-            attribute_value = artifact.get(attribute_name)
-        else:
-            attribute_value = artifact["spec"].get(attribute_name)
-
-        if mlrun.utils.is_legacy_artifact(artifact):
-            key = artifact["key"]
-        else:
-            key = artifact["metadata"]["key"]
+        attribute_value = artifact["spec"].get(attribute_name)
+        key = artifact["metadata"]["key"]
 
         if not attribute_value:
             mlrun.utils.logger.warning(
@@ -404,12 +396,18 @@ def runs_to_html(
         df.drop("labels", axis=1, inplace=True)
         df.drop("inputs", axis=1, inplace=True)
         df.drop("artifacts", axis=1, inplace=True)
+        df.drop("artifact_uris", axis=1, inplace=True)
     else:
         df["labels"] = df["labels"].apply(dict_html)
         df["inputs"] = df["inputs"].apply(inputs_html)
-        df["artifacts"] = df["artifacts"].apply(
-            lambda artifacts: artifacts_html(artifacts, "target_path"),
-        )
+        if df["artifact_uris"][0]:
+            df["artifact_uris"] = df["artifact_uris"].apply(dict_html)
+            df.drop("artifacts", axis=1, inplace=True)
+        else:
+            df["artifacts"] = df["artifacts"].apply(
+                lambda artifacts: artifacts_html(artifacts, "target_path"),
+            )
+            df.drop("artifact_uris", axis=1, inplace=True)
 
     def expand_error(x):
         if x["state"] == "error":

mlrun/run.py

@@ -29,11 +29,13 @@ from typing import Optional, Union
 import nuclio
 import yaml
 from kfp import Client
+from mlrun_pipelines.common.models import RunStatuses
+from mlrun_pipelines.common.ops import format_summary_from_kfp_run, show_kfp_run
+from mlrun_pipelines.models import PipelineRun
 
 import mlrun.common.schemas
 import mlrun.errors
 import mlrun.utils.helpers
-from mlrun.kfpops import format_summary_from_kfp_run, show_kfp_run
 
 from .common.helpers import parse_versioned_object_uri
 from .config import config as mlconf
@@ -47,7 +49,6 @@ from .runtimes import (
     KubejobRuntime,
     LocalRuntime,
     MpiRuntimeV1,
-    MpiRuntimeV1Alpha1,
     RemoteRuntime,
     RemoteSparkRuntime,
     RuntimeKinds,
@@ -69,41 +70,6 @@ from .utils import (
 )
 
 
-class RunStatuses:
-    succeeded = "Succeeded"
-    failed = "Failed"
-    skipped = "Skipped"
-    error = "Error"
-    running = "Running"
-
-    @staticmethod
-    def all():
-        return [
-            RunStatuses.succeeded,
-            RunStatuses.failed,
-            RunStatuses.skipped,
-            RunStatuses.error,
-            RunStatuses.running,
-        ]
-
-    @staticmethod
-    def stable_statuses():
-        return [
-            RunStatuses.succeeded,
-            RunStatuses.failed,
-            RunStatuses.skipped,
-            RunStatuses.error,
-        ]
-
-    @staticmethod
-    def transient_statuses():
-        return [
-            status
-            for status in RunStatuses.all()
-            if status not in RunStatuses.stable_statuses()
-        ]
-
-
 def function_to_module(code="", workdir=None, secrets=None, silent=False):
     """Load code, notebook or mlrun function as .py module
     this function can import a local/remote py file or notebook
@@ -606,7 +572,6 @@ def code_to_function(
     ignored_tags: Optional[str] = None,
     requirements_file: Optional[str] = "",
 ) -> Union[
-    MpiRuntimeV1Alpha1,
     MpiRuntimeV1,
     RemoteRuntime,
     ServingRuntime,
@@ -1023,7 +988,7 @@ def get_pipeline(
     :param project:    the project of the pipeline run
     :param remote:     read kfp data from mlrun service (default=True)
 
-    :return: kfp run dict
+    :return: kfp run
     """
     namespace = namespace or mlconf.namespace
     if remote:
@@ -1047,7 +1012,7 @@ def get_pipeline(
                 not format_
                 or format_ == mlrun.common.schemas.PipelinesFormat.summary.value
             ):
-                resp = format_summary_from_kfp_run(resp)
+                resp = format_summary_from_kfp_run(PipelineRun(resp))
 
     show_kfp_run(resp)
     return resp
@@ -1150,7 +1115,7 @@ def wait_for_runs_completion(
         running = []
         for run in runs:
             state = run.state()
-            if state in mlrun.runtimes.constants.RunStates.terminal_states():
+            if state in mlrun.common.runtimes.constants.RunStates.terminal_states():
                 completed.append(run)
             else:
                 running.append(run)

mlrun/runtimes/__init__.py

@@ -30,13 +30,13 @@ __all__ = [
 
 from mlrun.runtimes.utils import resolve_spark_operator_version
 
+from ..common.runtimes.constants import MPIJobCRDVersions
 from .base import BaseRuntime, RunError, RuntimeClassMode  # noqa
-from .constants import MPIJobCRDVersions
 from .daskjob import DaskCluster  # noqa
 from .databricks_job.databricks_runtime import DatabricksRuntime
 from .kubejob import KubejobRuntime, KubeResource  # noqa
 from .local import HandlerRuntime, LocalRuntime  # noqa
-from .mpijob import MpiRuntimeContainer, MpiRuntimeV1, MpiRuntimeV1Alpha1  # noqa
+from .mpijob import MpiRuntimeV1  # noqa
 from .nuclio import (
     RemoteRuntime,
     ServingRuntime,
@@ -264,7 +264,7 @@ class RuntimeKinds:
 
 def get_runtime_class(kind: str):
     if kind == RuntimeKinds.mpijob:
-        return MpiRuntimeContainer.selector()
+        return MpiRuntimeV1
 
     if kind == RuntimeKinds.spark:
         return Spark3Runtime

mlrun/runtimes/base.py

@@ -21,6 +21,7 @@ from os import environ
 from typing import Callable, Optional, Union
 
 import requests.exceptions
+from mlrun_pipelines.common.ops import mlrun_op
 from nuclio.build import mlrun_footer
 
 import mlrun.common.constants
@@ -37,7 +38,6 @@ from mlrun.utils.helpers import generate_object_uri, verify_field_regex
 from ..config import config
 from ..datastore import store_manager
 from ..errors import err_to_str
-from ..kfpops import mlrun_op
 from ..lists import RunList
 from ..model import BaseMetadata, HyperParamOptions, ImageBuilder, ModelObj, RunObject
 from ..utils import (
@@ -469,7 +469,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.runtimes.constants.RunLabels.owner]
+            meta.labels, [mlrun.common.runtimes.constants.RunLabels.owner]
         )
         if runspec.spec.output_path:
             runspec.spec.output_path = runspec.spec.output_path.replace(
@@ -580,9 +580,9 @@ class BaseRuntime(ModelObj):
 
         elif (
             not was_none
-            and last_state != mlrun.runtimes.constants.RunStates.completed
+            and last_state != mlrun.common.runtimes.constants.RunStates.completed
             and last_state
-            not in mlrun.runtimes.constants.RunStates.error_and_abortion_states()
+            not in mlrun.common.runtimes.constants.RunStates.error_and_abortion_states()
         ):
             try:
                 runtime_cls = mlrun.runtimes.get_runtime_class(kind)
@@ -707,11 +707,11 @@ class BaseRuntime(ModelObj):
                                   "key": "the_key".
         :param auto_build:      when set to True and the function require build it will be built on the first
                                 function run, use only if you dont plan on changing the build config between runs
-        :return: KubeFlow containerOp
+        :return: mlrun_pipelines.models.PipelineNodeWrapper
         """
 
         # if the function contain KFP PipelineParams (futures) pass the full spec to the
-        # ContainerOp this way KFP will substitute the params with previous step outputs
+        # PipelineNodeWrapper this way KFP will substitute the params with previous step outputs
         if use_db and not self._has_pipeline_param():
             # if the same function is built as part of the pipeline we do not use the versioned function
             # rather the latest function w the same tag so we can pick up the updated image/status