mlrun 1.8.0rc18__py3-none-any.whl → 1.8.0rc20__py3-none-any.whl

mlrun/serving/v2_serving.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import random
 import threading
 import time
 import traceback
@@ -283,7 +284,6 @@ class V2ModelServer(StepToDict):
             }
             if self.version:
                 response["model_version"] = self.version
-
         elif op == "ready" and event.method == "GET":
             # get model health operation
             setattr(event, "terminated", True)
@@ -468,13 +468,9 @@ class _ModelLogPusher:
         self.hostname = context.stream.hostname
         self.function_uri = context.stream.function_uri
         self.stream_path = context.stream.stream_uri
-        self.stream_batch = int(context.get_param("log_stream_batch", 1))
-        self.stream_sample = int(context.get_param("log_stream_sample", 1))
+        self.sampling_percentage = float(context.get_param("sampling_percentage", 100))
         self.output_stream = output_stream or context.stream.output_stream
         self._worker = context.worker_id
-        self._sample_iter = 0
-        self._batch_iter = 0
-        self._batch = []
 
     def base_data(self):
         base_data = {
@@ -485,6 +481,7 @@ class _ModelLogPusher:
             "host": self.hostname,
             "function_uri": self.function_uri,
             "endpoint_id": self.model.model_endpoint_uid,
+            "sampling_percentage": self.sampling_percentage,
         }
         if getattr(self.model, "labels", None):
             base_data["labels"] = self.model.labels
@@ -504,37 +501,55 @@ class _ModelLogPusher:
             self.output_stream.push([data], partition_key=partition_key)
             return
 
-        self._sample_iter = (self._sample_iter + 1) % self.stream_sample
-        if self.output_stream and self._sample_iter == 0:
+        if self.output_stream:
+            # Ensure that the inputs are a list of lists
+            request["inputs"] = (
+                request["inputs"]
+                if not any(not isinstance(req, list) for req in request["inputs"])
+                else [request["inputs"]]
+            )
             microsec = (now_date() - start).microseconds
 
-            if self.stream_batch > 1:
-                if self._batch_iter == 0:
-                    self._batch = []
-                self._batch.append(
-                    [request, op, resp, str(start), microsec, self.model.metrics]
+            if self.sampling_percentage != 100:
+                # Randomly select a subset of the requests based on the percentage
+                num_of_inputs = len(request["inputs"])
+                sampled_requests_indices = self._pick_random_requests(
+                    num_of_inputs, self.sampling_percentage
                 )
-                self._batch_iter = (self._batch_iter + 1) % self.stream_batch
-
-                if self._batch_iter == 0:
-                    data = self.base_data()
-                    data["headers"] = [
-                        "request",
-                        "op",
-                        "resp",
-                        "when",
-                        "microsec",
-                        "metrics",
+                if not sampled_requests_indices:
+                    # No events were selected for sampling
+                    return
+
+                request["inputs"] = [
+                    request["inputs"][i] for i in sampled_requests_indices
+                ]
+
+                if resp and "outputs" in resp and isinstance(resp["outputs"], list):
+                    resp["outputs"] = [
+                        resp["outputs"][i] for i in sampled_requests_indices
                     ]
-                    data["values"] = self._batch
-                    self.output_stream.push([data], partition_key=partition_key)
-            else:
-                data = self.base_data()
-                data["request"] = request
-                data["op"] = op
-                data["resp"] = resp
-                data["when"] = start_str
-                data["microsec"] = microsec
-                if getattr(self.model, "metrics", None):
-                    data["metrics"] = self.model.metrics
-                self.output_stream.push([data], partition_key=partition_key)
+
+            data = self.base_data()
+            data["request"] = request
+            data["op"] = op
+            data["resp"] = resp
+            data["when"] = start_str
+            data["microsec"] = microsec
+            if getattr(self.model, "metrics", None):
+                data["metrics"] = self.model.metrics
+            data["effective_sample_count"] = len(request["inputs"])
+            self.output_stream.push([data], partition_key=partition_key)
+
+    @staticmethod
+    def _pick_random_requests(num_of_reqs: int, percentage: float) -> list[int]:
+        """
+        Randomly selects indices of requests to sample based on the given percentage
+
+        :param num_of_reqs: Number of requests to select from
+        :param percentage: Sample percentage for each request
+        :return: A list containing the indices of the selected requests
+        """
+
+        return [
+            req for req in range(num_of_reqs) if random.random() < (percentage / 100)
+        ]

mlrun/utils/helpers.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 import asyncio
+import base64
 import enum
 import functools
+import gzip
 import hashlib
 import inspect
 import itertools
@@ -23,6 +25,7 @@ import os
 import re
 import string
 import sys
+import traceback
 import typing
 import uuid
 import warnings
@@ -44,11 +47,16 @@ from pandas import Timedelta, Timestamp
 from yaml.representer import RepresenterError
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.helpers
+import mlrun.common.runtimes.constants as runtimes_constants
 import mlrun.common.schemas
 import mlrun.errors
 import mlrun.utils.regex
 import mlrun.utils.version.version
+import mlrun_pipelines.common.constants
+import mlrun_pipelines.models
+import mlrun_pipelines.utils
 from mlrun.common.constants import MYSQL_MEDIUMBLOB_SIZE_BYTES
 from mlrun.config import config
 from mlrun_pipelines.models import PipelineRun
@@ -1703,7 +1711,14 @@ def get_serving_spec():
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Failed to find serving spec in env var or config file"
             )
-    spec = json.loads(data)
+    # Attempt to decode and decompress, or use as-is for backward compatibility
+    try:
+        decoded_data = base64.b64decode(data)
+        decompressed_data = gzip.decompress(decoded_data)
+        spec = json.loads(decompressed_data.decode("utf-8"))
+    except (OSError, gzip.BadGzipFile, base64.binascii.Error, json.JSONDecodeError):
+        spec = json.loads(data)
+
     return spec
 
 
@@ -1904,3 +1919,150 @@ def join_urls(base_url: Optional[str], path: Optional[str]) -> str:
     if base_url is None:
         base_url = ""
     return f"{base_url.rstrip('/')}/{path.lstrip('/')}" if path else base_url
+
+
+class Workflow:
+    @staticmethod
+    def get_workflow_steps(workflow_id: str, project: str) -> list:
+        steps = []
+        db = mlrun.get_run_db()
+
+        def _add_run_step(_step: mlrun_pipelines.models.PipelineStep):
+            try:
+                _run = db.list_runs(
+                    project=project,
+                    labels=f"{mlrun_constants.MLRunInternalLabels.runner_pod}={_step.node_name}",
+                )[0]
+            except IndexError:
+                _run = {
+                    "metadata": {
+                        "name": _step.display_name,
+                        "project": project,
+                    },
+                }
+            _run["step_kind"] = _step.step_type
+            if _step.skipped:
+                _run.setdefault("status", {})["state"] = (
+                    runtimes_constants.RunStates.skipped
+                )
+            steps.append(_run)
+
+        def _add_deploy_function_step(_step: mlrun_pipelines.models.PipelineStep):
+            project, name, hash_key = Workflow._extract_function_uri(
+                _step.get_annotation("mlrun/function-uri")
+            )
+            if name:
+                try:
+                    function = db.get_function(
+                        project=project, name=name, hash_key=hash_key
+                    )
+                except mlrun.errors.MLRunNotFoundError:
+                    # If the function is not found (if build failed for example), we will create a dummy
+                    # function object for the notification to display the function name
+                    function = {
+                        "metadata": {
+                            "name": name,
+                            "project": project,
+                            "hash_key": hash_key,
+                        },
+                    }
+                pod_phase = _step.phase
+                if _step.skipped:
+                    state = mlrun.common.schemas.FunctionState.skipped
+                else:
+                    state = runtimes_constants.PodPhases.pod_phase_to_run_state(
+                        pod_phase
+                    )
+                function["status"] = {"state": state}
+                if isinstance(function["metadata"].get("updated"), datetime.datetime):
+                    function["metadata"]["updated"] = function["metadata"][
+                        "updated"
+                    ].isoformat()
+                function["step_kind"] = _step.step_type
+                steps.append(function)
+
+        step_methods = {
+            mlrun_pipelines.common.constants.PipelineRunType.run: _add_run_step,
+            mlrun_pipelines.common.constants.PipelineRunType.build: _add_deploy_function_step,
+            mlrun_pipelines.common.constants.PipelineRunType.deploy: _add_deploy_function_step,
+        }
+
+        if not workflow_id:
+            return steps
+
+        try:
+            workflow_manifest = Workflow._get_workflow_manifest(workflow_id)
+        except Exception:
+            logger.warning(
+                "Failed to extract workflow steps from workflow manifest, "
+                "returning all runs with the workflow id label",
+                workflow_id=workflow_id,
+                traceback=traceback.format_exc(),
+            )
+            return db.list_runs(
+                project=project,
+                labels=f"workflow={workflow_id}",
+            )
+
+        if not workflow_manifest:
+            return steps
+
+        try:
+            for step in workflow_manifest.get_steps():
+                step_method = step_methods.get(step.step_type)
+                if step_method:
+                    step_method(step)
+            return steps
+        except Exception:
+            # If we fail to read the pipeline steps, we will return the list of runs that have the same workflow id
+            logger.warning(
+                "Failed to extract workflow steps from workflow manifest, "
+                "returning all runs with the workflow id label",
+                workflow_id=workflow_id,
+                traceback=traceback.format_exc(),
+            )
+            return db.list_runs(
+                project=project,
+                labels=f"workflow={workflow_id}",
+            )
+
+    @staticmethod
+    def _extract_function_uri(function_uri: str) -> tuple[str, str, str]:
+        """
+        Extract the project, name, and hash key from a function uri.
+        Examples:
+            - "project/name@hash_key" returns project, name, hash_key
+            - "project/name returns" project, name, ""
+        """
+        project, name, hash_key = None, None, None
+        hashed_pattern = r"^(.+)/(.+)@(.+)$"
+        pattern = r"^(.+)/(.+)$"
+        match = re.match(hashed_pattern, function_uri)
+        if match:
+            project, name, hash_key = match.groups()
+        else:
+            match = re.match(pattern, function_uri)
+            if match:
+                project, name = match.groups()
+                hash_key = ""
+        return project, name, hash_key
+
+    @staticmethod
+    def _get_workflow_manifest(
+        workflow_id: str,
+    ) -> typing.Optional[mlrun_pipelines.models.PipelineManifest]:
+        kfp_client = mlrun_pipelines.utils.get_client(mlrun.mlconf.kfp_url)
+
+        # arbitrary timeout of 5 seconds, the workflow should be done by now
+        kfp_run = kfp_client.wait_for_run_completion(workflow_id, 5)
+        if not kfp_run:
+            return None
+
+        kfp_run = mlrun_pipelines.models.PipelineRun(kfp_run)
+        return kfp_run.workflow_manifest()
+
+
+def as_dict(data: typing.Union[dict, str]) -> dict:
+    if isinstance(data, str):
+        return json.loads(data)
+    return data

mlrun/utils/notifications/notification/webhook.py

@@ -118,6 +118,9 @@ class WebhookNotification(NotificationBase):
 
         if isinstance(override_body, dict):
             for key, value in override_body.items():
+                if not isinstance(value, str):
+                    # If the value is not a string, we don't want to parse it
+                    continue
                 if re.search(r"{{\s*runs\s*}}", value):
                     str_parsed_runs = parse_runs()
                     override_body[key] = re.sub(

mlrun/utils/notifications/notification_pusher.py

@@ -15,7 +15,6 @@
 import asyncio
 import datetime
 import os
-import re
 import traceback
 import typing
 from concurrent.futures import ThreadPoolExecutor
@@ -31,11 +30,7 @@ import mlrun.model
 import mlrun.utils.helpers
 import mlrun.utils.notifications.notification as notification_module
 import mlrun.utils.notifications.notification.base as base
-import mlrun_pipelines.common.constants
-import mlrun_pipelines.common.ops
-import mlrun_pipelines.models
-import mlrun_pipelines.utils
-from mlrun.utils import logger
+from mlrun.utils import Workflow, logger
 from mlrun.utils.condition_evaluator import evaluate_condition_in_separate_process
 
 
@@ -144,15 +139,25 @@ class NotificationPusher(_NotificationPusherBase):
                     error=mlrun.errors.err_to_str(exc),
                 )
 
-    def _process_notification(self, notification, run):
-        notification.status = run.status.notifications.get(notification.name, {}).get(
+    def _process_notification(self, notification_object, run):
+        notification_object.status = run.status.notifications.get(
+            notification_object.name, {}
+        ).get(
             "status",
             mlrun.common.schemas.NotificationStatus.PENDING,
         )
-        if self._should_notify(run, notification):
-            self._load_notification(run, notification)
+        if self._should_notify(run, notification_object):
+            notification = self._load_notification(notification_object)
+            if notification.is_async:
+                self._async_notifications.append(
+                    (notification, run, notification_object)
+                )
+            else:
+                self._sync_notifications.append(
+                    (notification, run, notification_object)
+                )
 
-    def push(self):
+    def push(self, sync_push_callback=None, async_push_callback=None):
         """
         Asynchronously push notifications for all runs in the initialized runs list (if they should be pushed).
         When running from a sync environment, the notifications will be pushed asynchronously however the function will
@@ -206,8 +211,9 @@ class NotificationPusher(_NotificationPusherBase):
             notifications_amount=len(self._sync_notifications)
             + len(self._async_notifications),
         )
-
-        self._push(sync_push, async_push)
+        sync_push_callback = sync_push_callback or sync_push
+        async_push_callback = async_push_callback or async_push
+        self._push(sync_push_callback, async_push_callback)
 
     @staticmethod
     def _should_notify(
@@ -246,24 +252,19 @@ class NotificationPusher(_NotificationPusherBase):
         return False
 
     def _load_notification(
-        self, run: mlrun.model.RunObject, notification_object: mlrun.model.Notification
+        self, notification_object: mlrun.model.Notification
     ) -> base.NotificationBase:
         name = notification_object.name
         notification_type = notification_module.NotificationTypes(
             notification_object.kind or notification_module.NotificationTypes.console
         )
         params = {}
-        params.update(notification_object.secret_params)
-        params.update(notification_object.params)
+        params.update(notification_object.secret_params or {})
+        params.update(notification_object.params or {})
         default_params = self._default_params.get(notification_type.value, {})
         notification = notification_type.get_notification()(
             name, params, default_params
         )
-        if notification.is_async:
-            self._async_notifications.append((notification, run, notification_object))
-        else:
-            self._sync_notifications.append((notification, run, notification_object))
-
         logger.debug(
             "Loaded notification", notification=name, type=notification_type.value
         )
@@ -283,7 +284,9 @@ class NotificationPusher(_NotificationPusherBase):
             custom_message = (
                 f" (workflow: {run.metadata.labels['workflow']}){custom_message}"
             )
-            runs.extend(self.get_workflow_steps(run))
+            project = run.metadata.project
+            workflow_id = run.status.results.get("workflow_id", None)
+            runs.extend(Workflow.get_workflow_steps(workflow_id, project))
 
         message = (
             self.messages.get(run.state(), "").format(resource=resource)
@@ -442,131 +445,6 @@ class NotificationPusher(_NotificationPusherBase):
             mask_params=False,
         )
 
-    def get_workflow_steps(self, run: mlrun.model.RunObject) -> list:
-        steps = []
-        db = mlrun.get_run_db()
-
-        def _add_run_step(_step: mlrun_pipelines.models.PipelineStep):
-            try:
-                _run = db.list_runs(
-                    project=run.metadata.project,
-                    labels=f"{mlrun_constants.MLRunInternalLabels.runner_pod}={_step.node_name}",
-                )[0]
-            except IndexError:
-                _run = {
-                    "metadata": {
-                        "name": _step.display_name,
-                        "project": run.metadata.project,
-                    },
-                }
-            _run["step_kind"] = _step.step_type
-            if _step.skipped:
-                _run.setdefault("status", {})["state"] = (
-                    runtimes_constants.RunStates.skipped
-                )
-            steps.append(_run)
-
-        def _add_deploy_function_step(_step: mlrun_pipelines.models.PipelineStep):
-            project, name, hash_key = self._extract_function_uri(
-                _step.get_annotation("mlrun/function-uri")
-            )
-            if name:
-                try:
-                    function = db.get_function(
-                        project=project, name=name, hash_key=hash_key
-                    )
-                except mlrun.errors.MLRunNotFoundError:
-                    # If the function is not found (if build failed for example), we will create a dummy
-                    # function object for the notification to display the function name
-                    function = {
-                        "metadata": {
-                            "name": name,
-                            "project": project,
-                            "hash_key": hash_key,
-                        },
-                    }
-                pod_phase = _step.phase
-                if _step.skipped:
-                    state = mlrun.common.schemas.FunctionState.skipped
-                else:
-                    state = runtimes_constants.PodPhases.pod_phase_to_run_state(
-                        pod_phase
-                    )
-                function["status"] = {"state": state}
-                if isinstance(function["metadata"].get("updated"), datetime.datetime):
-                    function["metadata"]["updated"] = function["metadata"][
-                        "updated"
-                    ].isoformat()
-                function["step_kind"] = _step.step_type
-                steps.append(function)
-
-        step_methods = {
-            mlrun_pipelines.common.constants.PipelineRunType.run: _add_run_step,
-            mlrun_pipelines.common.constants.PipelineRunType.build: _add_deploy_function_step,
-            mlrun_pipelines.common.constants.PipelineRunType.deploy: _add_deploy_function_step,
-        }
-
-        workflow_id = run.status.results.get("workflow_id", None)
-        if not workflow_id:
-            return steps
-
-        workflow_manifest = self._get_workflow_manifest(workflow_id)
-        if not workflow_manifest:
-            return steps
-
-        try:
-            for step in workflow_manifest.get_steps():
-                step_method = step_methods.get(step.step_type)
-                if step_method:
-                    step_method(step)
-            return steps
-        except Exception:
-            # If we fail to read the pipeline steps, we will return the list of runs that have the same workflow id
-            logger.warning(
-                "Failed to extract workflow steps from workflow manifest, "
-                "returning all runs with the workflow id label",
-                workflow_id=workflow_id,
-                traceback=traceback.format_exc(),
-            )
-            return db.list_runs(
-                project=run.metadata.project,
-                labels=f"workflow={workflow_id}",
-            )
-
-    @staticmethod
-    def _get_workflow_manifest(
-        workflow_id: str,
-    ) -> typing.Optional[mlrun_pipelines.models.PipelineManifest]:
-        kfp_client = mlrun_pipelines.utils.get_client(mlrun.mlconf.kfp_url)
-
-        # arbitrary timeout of 5 seconds, the workflow should be done by now
-        kfp_run = kfp_client.wait_for_run_completion(workflow_id, 5)
-        if not kfp_run:
-            return None
-
-        kfp_run = mlrun_pipelines.models.PipelineRun(kfp_run)
-        return kfp_run.workflow_manifest()
-
-    def _extract_function_uri(self, function_uri: str) -> tuple[str, str, str]:
-        """
-        Extract the project, name, and hash key from a function uri.
-        Examples:
-            - "project/name@hash_key" returns project, name, hash_key
-            - "project/name returns" project, name, ""
-        """
-        project, name, hash_key = None, None, None
-        hashed_pattern = r"^(.+)/(.+)@(.+)$"
-        pattern = r"^(.+)/(.+)$"
-        match = re.match(hashed_pattern, function_uri)
-        if match:
-            project, name, hash_key = match.groups()
-        else:
-            match = re.match(pattern, function_uri)
-            if match:
-                project, name = match.groups()
-                hash_key = ""
-        return project, name, hash_key
-
 
 class CustomNotificationPusher(_NotificationPusherBase):
     def __init__(self, notification_types: typing.Optional[list[str]] = None):
@@ -624,6 +502,14 @@ class CustomNotificationPusher(_NotificationPusherBase):
         notification_type: str,
         params: typing.Optional[dict[str, str]] = None,
     ):
+        if notification_type not in [
+            notification_module.NotificationTypes.console,
+            notification_module.NotificationTypes.ipython,
+        ]:
+            # We want that only the console and ipython notifications will be notified by the client.
+            # The rest of the notifications will be notified by the BE.
+            return
+
         if notification_type in self._async_notifications:
             self._async_notifications[notification_type].load_notification(params)
         elif notification_type in self._sync_notifications:
@@ -693,25 +579,9 @@ class CustomNotificationPusher(_NotificationPusherBase):
         pipeline_id: typing.Optional[str] = None,
         has_workflow_url: bool = False,
     ):
-        message = f"Workflow started in project {project}"
-        if pipeline_id:
-            message += f" id={pipeline_id}"
-        commit_id = (
-            commit_id or os.environ.get("GITHUB_SHA") or os.environ.get("CI_COMMIT_SHA")
+        html, message = self.generate_start_message(
+            commit_id, has_workflow_url, pipeline_id, project
         )
-        if commit_id:
-            message += f", commit={commit_id}"
-        if has_workflow_url:
-            url = mlrun.utils.helpers.get_workflow_url(project, pipeline_id)
-        else:
-            url = mlrun.utils.helpers.get_ui_url(project)
-        html = ""
-        if url:
-            html = (
-                message
-                + f'<div><a href="{url}" target="_blank">click here to view progress</a></div>'
-            )
-            message = message + f", check progress in {url}"
         self.push(message, "info", custom_html=html)
 
     def push_pipeline_run_results(
@@ -744,6 +614,30 @@ class CustomNotificationPusher(_NotificationPusherBase):
             text += f", state={state}"
         self.push(text, "info", runs=runs_list)
 
+    def generate_start_message(
+        self, commit_id=None, has_workflow_url=None, pipeline_id=None, project=None
+    ):
+        message = f"Workflow started in project {project}"
+        if pipeline_id:
+            message += f" id={pipeline_id}"
+        commit_id = (
+            commit_id or os.environ.get("GITHUB_SHA") or os.environ.get("CI_COMMIT_SHA")
+        )
+        if commit_id:
+            message += f", commit={commit_id}"
+        if has_workflow_url:
+            url = mlrun.utils.helpers.get_workflow_url(project, pipeline_id)
+        else:
+            url = mlrun.utils.helpers.get_ui_url(project)
+        html = ""
+        if url:
+            html = (
+                message
+                + f'<div><a href="{url}" target="_blank">click here to view progress</a></div>'
+            )
+            message = message + f", check progress in {url}"
+        return html, message
+
 
 def sanitize_notification(notification_dict: dict):
     notification_dict.pop("secret_params", None)

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "2d5cc0ee3dacf5d31a3b6a2158cd2e2b4a378b56",
-  "version": "1.8.0-rc18"
+  "git_commit": "91ab2fb641a35e93a5056d93c55b98b96d74ebc1",
+  "version": "1.8.0-rc20"
 }

{mlrun-1.8.0rc18.dist-info → mlrun-1.8.0rc20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlrun
-Version: 1.8.0rc18
+Version: 1.8.0rc20
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv