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

mlrun/model_monitoring/writer.py

@@ -12,25 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
 import json
 from typing import Any, NewType
 
-import pandas as pd
-from v3io.dataplane import Client as V3IOClient
-from v3io_frames.client import ClientBase as V3IOFramesClient
-from v3io_frames.errors import Error as V3IOFramesError
-from v3io_frames.frames_pb2 import IGNORE
-
 import mlrun.common.model_monitoring
+import mlrun.common.schemas
 import mlrun.common.schemas.alert as alert_constants
 import mlrun.model_monitoring
-import mlrun.model_monitoring.db.stores
-import mlrun.utils.v3io_clients
 from mlrun.common.schemas.model_monitoring.constants import (
     EventFieldType,
+    HistogramDataDriftApplicationConstants,
+    MetricData,
+    ResultData,
     ResultStatusApp,
     WriterEvent,
+    WriterEventKind,
 )
 from mlrun.common.schemas.notification import NotificationKind, NotificationSeverity
 from mlrun.model_monitoring.helpers import get_endpoint_record
@@ -38,9 +34,6 @@ from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
 from mlrun.utils.notifications.notification_pusher import CustomNotificationPusher
 
-_TSDB_BE = "tsdb"
-_TSDB_RATE = "1/s"
-_TSDB_TABLE = "app-results"
 _RawEvent = dict[str, Any]
 _AppResultEvent = NewType("_AppResultEvent", _RawEvent)
 
@@ -75,20 +68,20 @@ class _Notifier:
         self._severity = severity
 
     def _should_send_event(self) -> bool:
-        return self._event[WriterEvent.RESULT_STATUS] >= ResultStatusApp.detected
+        return self._event[ResultData.RESULT_STATUS] >= ResultStatusApp.detected.value
 
     def _generate_message(self) -> str:
         return f"""\
 The monitoring app `{self._event[WriterEvent.APPLICATION_NAME]}` \
-of kind `{self._event[WriterEvent.RESULT_KIND]}` \
+of kind `{self._event[ResultData.RESULT_KIND]}` \
 detected a problem in model endpoint ID `{self._event[WriterEvent.ENDPOINT_ID]}` \
 at time `{self._event[WriterEvent.START_INFER_TIME]}`.
 
 Result data:
-Name: `{self._event[WriterEvent.RESULT_NAME]}`
-Value: `{self._event[WriterEvent.RESULT_VALUE]}`
-Status: `{self._event[WriterEvent.RESULT_STATUS]}`
-Extra data: `{self._event[WriterEvent.RESULT_EXTRA_DATA]}`\
+Name: `{self._event[ResultData.RESULT_NAME]}`
+Value: `{self._event[ResultData.RESULT_VALUE]}`
+Status: `{self._event[ResultData.RESULT_STATUS]}`
+Extra data: `{self._event[ResultData.RESULT_EXTRA_DATA]}`\
 """
 
     def notify(self) -> None:
@@ -103,7 +96,7 @@ Extra data: `{self._event[WriterEvent.RESULT_EXTRA_DATA]}`\
 
 class ModelMonitoringWriter(StepToDict):
     """
-    Write monitoring app events to V3IO KV storage
+    Write monitoring application results to the target databases
     """
 
     kind = "monitoring_application_stream_pusher"
@@ -111,128 +104,97 @@ class ModelMonitoringWriter(StepToDict):
     def __init__(self, project: str) -> None:
         self.project = project
         self.name = project  # required for the deployment process
-        self._v3io_container = self.get_v3io_container(self.name)
-        self._tsdb_client = self._get_v3io_frames_client(self._v3io_container)
+
         self._custom_notifier = CustomNotificationPusher(
             notification_types=[NotificationKind.slack]
         )
-        self._create_tsdb_table()
-        self._endpoints_records = {}
-
-    @staticmethod
-    def get_v3io_container(project_name: str) -> str:
-        return f"users/pipelines/{project_name}/monitoring-apps"
-
-    @staticmethod
-    def _get_v3io_client() -> V3IOClient:
-        return mlrun.utils.v3io_clients.get_v3io_client(
-            endpoint=mlrun.mlconf.v3io_api,
-        )
-
-    @staticmethod
-    def _get_v3io_frames_client(v3io_container: str) -> V3IOFramesClient:
-        return mlrun.utils.v3io_clients.get_frames_client(
-            address=mlrun.mlconf.v3io_framesd,
-            container=v3io_container,
-        )
-
-    def _create_tsdb_table(self) -> None:
-        self._tsdb_client.create(
-            backend=_TSDB_BE,
-            table=_TSDB_TABLE,
-            if_exists=IGNORE,
-            rate=_TSDB_RATE,
-        )
 
-    def _update_kv_db(self, event: _AppResultEvent) -> None:
-        event = _AppResultEvent(event.copy())
-        application_result_store = mlrun.model_monitoring.get_store_object(
+        self._app_result_store = mlrun.model_monitoring.get_store_object(
             project=self.project
         )
-        application_result_store.write_application_result(event=event)
-
-    def _update_tsdb(self, event: _AppResultEvent) -> None:
-        event = _AppResultEvent(event.copy())
-        event[WriterEvent.END_INFER_TIME] = datetime.datetime.fromisoformat(
-            event[WriterEvent.END_INFER_TIME]
+        self._tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
+            project=self.project,
         )
-        del event[WriterEvent.RESULT_EXTRA_DATA]
-        try:
-            self._tsdb_client.write(
-                backend=_TSDB_BE,
-                table=_TSDB_TABLE,
-                dfs=pd.DataFrame.from_records([event]),
-                index_cols=[
-                    WriterEvent.END_INFER_TIME,
-                    WriterEvent.ENDPOINT_ID,
-                    WriterEvent.APPLICATION_NAME,
-                    WriterEvent.RESULT_NAME,
-                ],
-            )
-            logger.info("Updated V3IO TSDB successfully", table=_TSDB_TABLE)
-        except V3IOFramesError as err:
-            logger.warn(
-                "Could not write drift measures to TSDB",
-                err=err,
-                table=_TSDB_TABLE,
-                event=event,
-            )
+        self._endpoints_records = {}
 
     @staticmethod
     def _generate_event_on_drift(
-        uid: str, drift_status: str, event_value: dict, project_name: str
-    ):
-        if (
-            drift_status == ResultStatusApp.detected
-            or drift_status == ResultStatusApp.potential_detection
-        ):
-            entity = {
-                "kind": alert_constants.EventEntityKind.MODEL,
-                "project": project_name,
-                "id": uid,
-            }
-            event_kind = (
-                alert_constants.EventKind.DRIFT_DETECTED
-                if drift_status == ResultStatusApp.detected
-                else alert_constants.EventKind.DRIFT_SUSPECTED
-            )
-            event_data = mlrun.common.schemas.Event(
-                kind=event_kind, entity=entity, value_dict=event_value
-            )
-            mlrun.get_run_db().generate_event(event_kind, event_data)
+        model_endpoint: str, drift_status: str, event_value: dict, project_name: str
+    ) -> None:
+        logger.info("Sending an alert")
+        entity = mlrun.common.schemas.alert.EventEntities(
+            kind=alert_constants.EventEntityKind.MODEL,
+            project=project_name,
+            ids=[model_endpoint],
+        )
+        event_kind = (
+            alert_constants.EventKind.DRIFT_DETECTED
+            if drift_status == ResultStatusApp.detected.value
+            else alert_constants.EventKind.DRIFT_SUSPECTED
+        )
+        event_data = mlrun.common.schemas.Event(
+            kind=event_kind, entity=entity, value_dict=event_value
+        )
+        mlrun.get_run_db().generate_event(event_kind, event_data)
 
     @staticmethod
-    def _reconstruct_event(event: _RawEvent) -> _AppResultEvent:
+    def _reconstruct_event(event: _RawEvent) -> tuple[_AppResultEvent, str]:
         """
         Modify the raw event into the expected monitoring application event
         schema as defined in `mlrun.common.schemas.model_monitoring.constants.WriterEvent`
         """
-        try:
-            result_event = _AppResultEvent(
-                {key: event[key] for key in WriterEvent.list()}
+        if not isinstance(event, dict):
+            raise _WriterEventTypeError(
+                f"The event is of type: {type(event)}, expected a dictionary"
             )
-            result_event[WriterEvent.CURRENT_STATS] = json.loads(
-                event[WriterEvent.CURRENT_STATS]
+        kind = event.pop(WriterEvent.EVENT_KIND, WriterEventKind.RESULT)
+        result_event = _AppResultEvent(json.loads(event.pop(WriterEvent.DATA, "{}")))
+        if not result_event:  # BC for < 1.7.0, can be removed in 1.9.0
+            result_event = _AppResultEvent(event)
+        else:
+            result_event.update(_AppResultEvent(event))
+
+        expected_keys = list(
+            set(WriterEvent.list()).difference(
+                [WriterEvent.EVENT_KIND, WriterEvent.DATA]
             )
-            return result_event
-        except KeyError as err:
+        )
+        if kind == WriterEventKind.METRIC:
+            expected_keys.extend(MetricData.list())
+        elif kind == WriterEventKind.RESULT:
+            expected_keys.extend(ResultData.list())
+        else:
             raise _WriterEventValueError(
-                "The received event misses some keys compared to the expected "
-                "monitoring application event schema"
-            ) from err
-        except TypeError as err:
-            raise _WriterEventTypeError(
-                f"The event is of type: {type(event)}, expected a dictionary"
-            ) from err
+                f"Unknown event kind: {kind}, expected one of: {WriterEventKind.list()}"
+            )
+        missing_keys = [key for key in expected_keys if key not in result_event]
+        if missing_keys:
+            raise _WriterEventValueError(
+                f"The received event misses some keys compared to the expected "
+                f"monitoring application event schema: {missing_keys}"
+            )
+
+        return result_event, kind
 
     def do(self, event: _RawEvent) -> None:
-        event = self._reconstruct_event(event)
+        event, kind = self._reconstruct_event(event)
         logger.info("Starting to write event", event=event)
-        self._update_tsdb(event)
-        self._update_kv_db(event)
+
+        self._tsdb_connector.write_application_event(event=event.copy(), kind=kind)
+        self._app_result_store.write_application_event(event=event.copy(), kind=kind)
+        logger.info("Completed event DB writes")
+
         _Notifier(event=event, notification_pusher=self._custom_notifier).notify()
 
-        if mlrun.mlconf.alerts.mode == mlrun.common.schemas.alert.AlertsModes.enabled:
+        if (
+            mlrun.mlconf.alerts.mode == mlrun.common.schemas.alert.AlertsModes.enabled
+            and kind == WriterEventKind.RESULT
+            and (
+                event[ResultData.RESULT_STATUS] == ResultStatusApp.detected.value
+                or event[ResultData.RESULT_STATUS]
+                == ResultStatusApp.potential_detection.value
+            )
+        ):
             endpoint_id = event[WriterEvent.ENDPOINT_ID]
             endpoint_record = self._endpoints_records.setdefault(
                 endpoint_id,
@@ -242,13 +204,31 @@ class ModelMonitoringWriter(StepToDict):
                 "app_name": event[WriterEvent.APPLICATION_NAME],
                 "model": endpoint_record.get(EventFieldType.MODEL),
                 "model_endpoint_id": event[WriterEvent.ENDPOINT_ID],
-                "result_name": event[WriterEvent.RESULT_NAME],
-                "result_value": event[WriterEvent.RESULT_VALUE],
+                "result_name": event[ResultData.RESULT_NAME],
+                "result_value": event[ResultData.RESULT_VALUE],
             }
             self._generate_event_on_drift(
                 event[WriterEvent.ENDPOINT_ID],
-                event[WriterEvent.RESULT_STATUS],
+                event[ResultData.RESULT_STATUS],
                 event_value,
                 self.project,
             )
-        logger.info("Completed event DB writes")
+
+        if (
+            kind == WriterEventKind.RESULT
+            and event[WriterEvent.APPLICATION_NAME]
+            == HistogramDataDriftApplicationConstants.NAME
+            and event[ResultData.RESULT_NAME]
+            == HistogramDataDriftApplicationConstants.GENERAL_RESULT_NAME
+        ):
+            endpoint_id = event[WriterEvent.ENDPOINT_ID]
+            logger.info(
+                "Updating the model endpoint with metadata specific to the histogram "
+                "data drift app",
+                endpoint_id=endpoint_id,
+            )
+            store = mlrun.model_monitoring.get_store_object(project=self.project)
+            store.update_model_endpoint(
+                endpoint_id=endpoint_id,
+                attributes=json.loads(event[ResultData.RESULT_EXTRA_DATA]),
+            )

mlrun/platforms/__init__.py

@@ -17,22 +17,23 @@ import json
 from pprint import pprint
 from time import sleep
 
-from .iguazio import (
-    V3ioStreamClient,
-    VolumeMount,
-    add_or_refresh_credentials,
-    is_iguazio_session_cookie,
-    mount_v3io,
-    v3io_cred,
-)
-from .other import (
+from mlrun_pipelines.common.mounts import VolumeMount
+from mlrun_pipelines.mounts import (
     auto_mount,
     mount_configmap,
     mount_hostpath,
     mount_pvc,
     mount_s3,
     mount_secret,
+    mount_v3io,
     set_env_variables,
+    v3io_cred,
+)
+
+from .iguazio import (
+    V3ioStreamClient,
+    add_or_refresh_credentials,
+    is_iguazio_session_cookie,
 )
 
 

mlrun/platforms/iguazio.py

@@ -15,12 +15,9 @@
 import json
 import os
 import urllib
-from collections import namedtuple
 from urllib.parse import urlparse
 
-import kfp.dsl
 import requests
-import semver
 import v3io
 
 import mlrun.errors
@@ -29,203 +26,6 @@ from mlrun.utils import dict_to_json
 
 _cached_control_session = None
 
-VolumeMount = namedtuple("Mount", ["path", "sub_path"])
-
-
-def mount_v3io(
-    name="v3io",
-    remote="",
-    access_key="",
-    user="",
-    secret=None,
-    volume_mounts=None,
-):
-    """Modifier function to apply to a Container Op to volume mount a v3io path
-
-    :param name:            the volume name
-    :param remote:          the v3io path to use for the volume. ~/ prefix will be replaced with /users/<username>/
-    :param access_key:      the access key used to auth against v3io. if not given V3IO_ACCESS_KEY env var will be used
-    :param user:            the username used to auth against v3io. if not given V3IO_USERNAME env var will be used
-    :param secret:          k8s secret name which would be used to get the username and access key to auth against v3io.
-    :param volume_mounts:   list of VolumeMount. empty volume mounts & remote will default to mount /v3io & /User.
-    """
-    volume_mounts, user = _enrich_and_validate_v3io_mounts(
-        remote=remote,
-        volume_mounts=volume_mounts,
-        user=user,
-    )
-
-    def _attach_volume_mounts_and_creds(container_op: kfp.dsl.ContainerOp):
-        from kubernetes import client as k8s_client
-
-        vol = v3io_to_vol(name, remote, access_key, user, secret=secret)
-        container_op.add_volume(vol)
-        for volume_mount in volume_mounts:
-            container_op.container.add_volume_mount(
-                k8s_client.V1VolumeMount(
-                    mount_path=volume_mount.path,
-                    sub_path=volume_mount.sub_path,
-                    name=name,
-                )
-            )
-
-        if not secret:
-            container_op = v3io_cred(access_key=access_key, user=user)(container_op)
-        return container_op
-
-    return _attach_volume_mounts_and_creds
-
-
-def _enrich_and_validate_v3io_mounts(remote="", volume_mounts=None, user=""):
-    if remote and not volume_mounts:
-        raise mlrun.errors.MLRunInvalidArgumentError(
-            "volume_mounts must be specified when remote is given"
-        )
-
-    # Empty remote & volume_mounts defaults are volume mounts of /v3io and /User
-    if not remote and not volume_mounts:
-        user = _resolve_mount_user(user)
-        if not user:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "user name/env must be specified when using empty remote and volume_mounts"
-            )
-        volume_mounts = [
-            VolumeMount(path="/v3io", sub_path=""),
-            VolumeMount(path="/User", sub_path="users/" + user),
-        ]
-
-    if not isinstance(volume_mounts, list) and any(
-        [not isinstance(x, VolumeMount) for x in volume_mounts]
-    ):
-        raise TypeError("mounts should be a list of Mount")
-
-    return volume_mounts, user
-
-
-def _resolve_mount_user(user=None):
-    return user or os.environ.get("V3IO_USERNAME")
-
-
-def mount_spark_conf():
-    def _mount_spark(container_op: kfp.dsl.ContainerOp):
-        from kubernetes import client as k8s_client
-
-        container_op.container.add_volume_mount(
-            k8s_client.V1VolumeMount(
-                name="spark-master-config", mount_path="/etc/config/spark"
-            )
-        )
-        return container_op
-
-    return _mount_spark
-
-
-def mount_v3iod(namespace, v3io_config_configmap):
-    def _mount_v3iod(container_op: kfp.dsl.ContainerOp):
-        from kubernetes import client as k8s_client
-
-        def add_vol(name, mount_path, host_path):
-            vol = k8s_client.V1Volume(
-                name=name,
-                host_path=k8s_client.V1HostPathVolumeSource(path=host_path, type=""),
-            )
-            container_op.add_volume(vol)
-            container_op.container.add_volume_mount(
-                k8s_client.V1VolumeMount(mount_path=mount_path, name=name)
-            )
-
-        # this is a legacy path for the daemon shared memory
-        host_path = "/dev/shm/"
-
-        # path to shared memory for daemon was changed in Iguazio 3.2.3-b1
-        igz_version = mlrun.mlconf.get_parsed_igz_version()
-        if igz_version and igz_version >= semver.VersionInfo.parse("3.2.3-b1"):
-            host_path = "/var/run/iguazio/dayman-shm/"
-        add_vol(name="shm", mount_path="/dev/shm", host_path=host_path + namespace)
-
-        add_vol(
-            name="v3iod-comm",
-            mount_path="/var/run/iguazio/dayman",
-            host_path="/var/run/iguazio/dayman/" + namespace,
-        )
-
-        vol = k8s_client.V1Volume(
-            name="daemon-health", empty_dir=k8s_client.V1EmptyDirVolumeSource()
-        )
-        container_op.add_volume(vol)
-        container_op.container.add_volume_mount(
-            k8s_client.V1VolumeMount(
-                mount_path="/var/run/iguazio/daemon_health", name="daemon-health"
-            )
-        )
-
-        vol = k8s_client.V1Volume(
-            name="v3io-config",
-            config_map=k8s_client.V1ConfigMapVolumeSource(
-                name=v3io_config_configmap, default_mode=420
-            ),
-        )
-        container_op.add_volume(vol)
-        container_op.container.add_volume_mount(
-            k8s_client.V1VolumeMount(mount_path="/etc/config/v3io", name="v3io-config")
-        )
-
-        container_op.container.add_env_variable(
-            k8s_client.V1EnvVar(
-                name="CURRENT_NODE_IP",
-                value_from=k8s_client.V1EnvVarSource(
-                    field_ref=k8s_client.V1ObjectFieldSelector(
-                        api_version="v1", field_path="status.hostIP"
-                    )
-                ),
-            )
-        )
-        container_op.container.add_env_variable(
-            k8s_client.V1EnvVar(
-                name="IGZ_DATA_CONFIG_FILE", value="/igz/java/conf/v3io.conf"
-            )
-        )
-
-        return container_op
-
-    return _mount_v3iod
-
-
-def v3io_cred(api="", user="", access_key=""):
-    """
-    Modifier function to copy local v3io env vars to container
-
-    Usage::
-
-        train = train_op(...)
-        train.apply(use_v3io_cred())
-    """
-
-    def _use_v3io_cred(container_op: kfp.dsl.ContainerOp):
-        from os import environ
-
-        from kubernetes import client as k8s_client
-
-        web_api = api or environ.get("V3IO_API") or mlconf.v3io_api
-        _user = user or environ.get("V3IO_USERNAME")
-        _access_key = access_key or environ.get("V3IO_ACCESS_KEY")
-        v3io_framesd = mlconf.v3io_framesd or environ.get("V3IO_FRAMESD")
-
-        return (
-            container_op.container.add_env_variable(
-                k8s_client.V1EnvVar(name="V3IO_API", value=web_api)
-            )
-            .add_env_variable(k8s_client.V1EnvVar(name="V3IO_USERNAME", value=_user))
-            .add_env_variable(
-                k8s_client.V1EnvVar(name="V3IO_ACCESS_KEY", value=_access_key)
-            )
-            .add_env_variable(
-                k8s_client.V1EnvVar(name="V3IO_FRAMESD", value=v3io_framesd)
-            )
-        )
-
-    return _use_v3io_cred
-
 
 def split_path(mntpath=""):
     if mntpath[0] == "/":
@@ -525,8 +325,8 @@ def add_or_refresh_credentials(
     # different access keys for the 2 usages
     token = (
         token
-        # can't use mlrun.runtimes.constants.FunctionEnvironmentVariables.auth_session cause this is running in the
-        # import execution path (when we're initializing the run db) and therefore we can't import mlrun.runtimes
+        # can't use mlrun.common.runtimes.constants.FunctionEnvironmentVariables.auth_session cause this is running
+        # in the import execution path (when we're initializing the run db) and therefore we can't import mlrun.runtimes
         or os.environ.get("MLRUN_AUTH_SESSION")
         or os.environ.get("V3IO_ACCESS_KEY")
     )
@@ -582,3 +382,22 @@ def sanitize_username(username: str):
     So simply replace it with dash
     """
     return username.replace("_", "-")
+
+
+def min_iguazio_versions(*versions):
+    def decorator(function):
+        def wrapper(*args, **kwargs):
+            if mlrun.utils.helpers.validate_component_version_compatibility(
+                "iguazio", *versions
+            ):
+                return function(*args, **kwargs)
+
+            message = (
+                f"{function.__name__} is supported since Iguazio {' or '.join(versions)}, currently using "
+                f"Iguazio {mlconf.igz_version}."
+            )
+            raise mlrun.errors.MLRunIncompatibleVersionError(message)
+
+        return wrapper
+
+    return decorator

mlrun/projects/operations.py

@@ -15,7 +15,7 @@
 import warnings
 from typing import Optional, Union
 
-import kfp
+from mlrun_pipelines.models import PipelineNodeWrapper
 
 import mlrun
 from mlrun.utils import hub_prefix
@@ -76,7 +76,7 @@ def run_function(
     notifications: list[mlrun.model.Notification] = None,
     returns: Optional[list[Union[str, dict[str, str]]]] = None,
     builder_env: Optional[list] = None,
-) -> Union[mlrun.model.RunObject, kfp.dsl.ContainerOp]:
+) -> Union[mlrun.model.RunObject, PipelineNodeWrapper]:
     """Run a local or remote task as part of a local/kubeflow pipeline
 
     run_function() allow you to execute a function locally, on a remote cluster, or as part of an automated workflow
@@ -86,7 +86,7 @@ def run_function(
     when functions run as part of a workflow/pipeline (project.run()) some attributes can be set at the run level,
     e.g. local=True will run all the functions locally, setting artifact_path will direct all outputs to the same path.
     project runs provide additional notifications/reporting and exception handling.
-    inside a Kubeflow pipeline (KFP) run_function() generates KFP "ContainerOps" which are used to form a DAG
+    inside a Kubeflow pipeline (KFP) run_function() generates KFP node (see PipelineNodeWrapper) which forms a DAG
     some behavior may differ between regular runs and deferred KFP runs.
 
     example (use with function object)::
@@ -166,7 +166,7 @@ def run_function(
                               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
     """
     engine, function = _get_engine_and_function(function, project_object)
     task = mlrun.new_task(
@@ -254,7 +254,7 @@ def build_function(
     overwrite_build_params: bool = False,
     extra_args: str = None,
     force_build: bool = False,
-) -> Union[BuildStatus, kfp.dsl.ContainerOp]:
+) -> Union[BuildStatus, PipelineNodeWrapper]:
     """deploy ML function, build container with its dependencies
 
     :param function:        Name of the function (in the project) or function object
@@ -294,7 +294,11 @@ def build_function(
         if overwrite_build_params:
             function.spec.build.commands = None
         if requirements or requirements_file:
-            function.with_requirements(requirements, requirements_file, overwrite=True)
+            function.with_requirements(
+                requirements=requirements,
+                requirements_file=requirements_file,
+                overwrite=True,
+            )
         if commands:
             function.with_commands(commands)
         return function.deploy_step(
@@ -358,7 +362,7 @@ def deploy_function(
     builder_env: dict = None,
     project_object=None,
     mock: bool = None,
-) -> Union[DeployStatus, kfp.dsl.ContainerOp]:
+) -> Union[DeployStatus, PipelineNodeWrapper]:
     """deploy real-time (nuclio based) functions
 
     :param function:   name of the function (in the project) or function object