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

mlrun/utils/notifications/notification_pusher.py

@@ -15,10 +15,17 @@
 import asyncio
 import datetime
 import os
+import re
 import traceback
 import typing
 from concurrent.futures import ThreadPoolExecutor
 
+import kfp
+import mlrun_pipelines.common.ops
+import mlrun_pipelines.models
+
+import mlrun.common.constants as mlrun_constants
+import mlrun.common.runtimes.constants
 import mlrun.common.schemas
 import mlrun.config
 import mlrun.db.base
@@ -233,25 +240,12 @@ class NotificationPusher(_NotificationPusherBase):
         resource = "Run"
         runs = [run.to_dict()]
 
-        if "workflow" in run.metadata.labels:
-            resource = "Workflow"
+        if mlrun_constants.MLRunInternalLabels.workflow in run.metadata.labels:
+            resource = mlrun_constants.MLRunInternalLabels.workflow
             custom_message = (
                 f" (workflow: {run.metadata.labels['workflow']}){custom_message}"
             )
-            db = mlrun.get_run_db()
-
-            workflow_id = run.status.results.get("workflow_id", None)
-            if workflow_id:
-                workflow_runs = db.list_runs(
-                    project=run.metadata.project,
-                    labels=f"workflow={workflow_id}",
-                )
-                logger.debug(
-                    "Found workflow runs, extending notification runs",
-                    workflow_id=workflow_id,
-                    workflow_runs_amount=len(workflow_runs),
-                )
-                runs.extend(workflow_runs)
+            runs.extend(self.get_workflow_steps(run))
 
         message = (
             self.messages.get(run.state(), "").format(resource=resource)
@@ -395,6 +389,137 @@ 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"] = (
+                    mlrun.common.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 = mlrun.common.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.ops.PipelineRunType.run: _add_run_step,
+            mlrun_pipelines.common.ops.PipelineRunType.build: _add_deploy_function_step,
+            mlrun_pipelines.common.ops.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_url = mlrun.mlconf.resolve_kfp_url(mlrun.mlconf.namespace)
+        if not kfp_url:
+            raise mlrun.errors.MLRunNotFoundError(
+                "KubeFlow Pipelines is not configured"
+            )
+
+        kfp_client = kfp.Client(host=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: list[str] = None):
@@ -413,6 +538,12 @@ class CustomNotificationPusher(_NotificationPusherBase):
             if notification.is_async
         }
 
+    @property
+    def notifications(self):
+        notifications = self._sync_notifications.copy()
+        notifications.update(self._async_notifications)
+        return notifications
+
     def push(
         self,
         message: str,

mlrun/utils/regex.py

@@ -92,3 +92,12 @@ artifact_key = [r"[^\/\\]+$"]
 # must be alphanumeric or _
 # max 256 length
 v3io_stream_consumer_group = [r"^(?!_)[a-zA-Z0-9_]{1,256}$"]
+
+# URI patterns
+run_uri_pattern = r"^(?P<project>.*)@(?P<uid>.*)\#(?P<iteration>.*?)(:(?P<tag>.*))?$"
+
+artifact_uri_pattern = r"^((?P<project>.*)/)?(?P<key>.*?)(\#(?P<iteration>.*?))?(:(?P<tag>.*?))?(@(?P<tree>.*))?$"
+
+artifact_producer_uri_pattern = (
+    r"^((?P<project>.*)/)?(?P<uid>.*?)(\-(?P<iteration>.*?))?$"
+)

mlrun/utils/v3io_clients.py

@@ -11,7 +11,6 @@
 # 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.
-#
 
 from v3io.dataplane import Client as V3IOClient
 from v3io_frames import Client as get_client

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "0967b2a28f5a75cc377483f5267159eb75925822",
-  "version": "1.7.0-rc14"
+  "git_commit": "0cbcbb9417af7ac58072ad981672602728fee57e",
+  "version": "1.7.0-rc22"
 }

{mlrun-1.7.0rc14.dist-info → mlrun-1.7.0rc22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlrun
-Version: 1.7.0rc14
+Version: 1.7.0rc22
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv
@@ -26,10 +26,9 @@ Requires-Dist: GitPython >=3.1.41,~=3.1
 Requires-Dist: aiohttp ~=3.9
 Requires-Dist: aiohttp-retry ~=2.8
 Requires-Dist: click ~=8.1
-Requires-Dist: kfp ~=1.8
 Requires-Dist: nest-asyncio ~=1.0
 Requires-Dist: ipython ~=8.10
-Requires-Dist: nuclio-jupyter ~=0.9.16
+Requires-Dist: nuclio-jupyter ~=0.9.17
 Requires-Dist: numpy <1.27.0,>=1.16.5
 Requires-Dist: pandas <2.2,>=1.2
 Requires-Dist: pyarrow <15,>=10.0
@@ -44,13 +43,15 @@ Requires-Dist: semver ~=3.0
 Requires-Dist: dependency-injector ~=4.41
 Requires-Dist: fsspec <2024.4,>=2023.9.2
 Requires-Dist: v3iofs ~=0.1.17
-Requires-Dist: storey ~=1.7.7
+Requires-Dist: storey ~=1.7.17
 Requires-Dist: inflection ~=0.5.0
 Requires-Dist: python-dotenv ~=0.17.0
 Requires-Dist: setuptools ~=69.1
 Requires-Dist: deprecated ~=1.2
 Requires-Dist: jinja2 >=3.1.3,~=3.1
 Requires-Dist: orjson <4,>=3.9.15
+Requires-Dist: mlrun-pipelines-kfp-common >0.1.0,~=0.1.1
+Requires-Dist: mlrun-pipelines-kfp-v1-8 >0.1.0,~=0.1.1
 Provides-Extra: alibaba-oss
 Requires-Dist: ossfs ==2023.12.0 ; extra == 'alibaba-oss'
 Requires-Dist: oss2 ==2.18.1 ; extra == 'alibaba-oss'
@@ -81,18 +82,20 @@ Requires-Dist: pyopenssl >=23 ; extra == 'all'
 Requires-Dist: redis ~=4.3 ; extra == 'all'
 Requires-Dist: s3fs <2024.4,>=2023.9.2 ; extra == 'all'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'all'
+Requires-Dist: taos-ws-py ~=0.3.2 ; extra == 'all'
 Provides-Extra: api
 Requires-Dist: uvicorn ~=0.27.1 ; extra == 'api'
 Requires-Dist: dask-kubernetes ~=0.11.0 ; extra == 'api'
 Requires-Dist: apscheduler <4,>=3.10.3 ; extra == 'api'
 Requires-Dist: objgraph ~=3.6 ; extra == 'api'
-Requires-Dist: igz-mgmt ~=0.1.3 ; extra == 'api'
+Requires-Dist: igz-mgmt ~=0.2.0 ; extra == 'api'
 Requires-Dist: humanfriendly ~=10.0 ; extra == 'api'
 Requires-Dist: fastapi ~=0.110.0 ; extra == 'api'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'api'
 Requires-Dist: pymysql ~=1.0 ; extra == 'api'
 Requires-Dist: alembic ~=1.9 ; extra == 'api'
 Requires-Dist: timelength ~=1.1 ; extra == 'api'
+Requires-Dist: memray ~=1.12 ; extra == 'api'
 Provides-Extra: azure-blob-storage
 Requires-Dist: msrest ~=0.6.21 ; extra == 'azure-blob-storage'
 Requires-Dist: azure-core ~=1.24 ; extra == 'azure-blob-storage'
@@ -128,6 +131,7 @@ Requires-Dist: pyopenssl >=23 ; extra == 'complete'
 Requires-Dist: redis ~=4.3 ; extra == 'complete'
 Requires-Dist: s3fs <2024.4,>=2023.9.2 ; extra == 'complete'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'complete'
+Requires-Dist: taos-ws-py ~=0.3.2 ; extra == 'complete'
 Provides-Extra: complete-api
 Requires-Dist: adlfs ==2023.9.0 ; extra == 'complete-api'
 Requires-Dist: aiobotocore <2.8,>=2.5.0 ; extra == 'complete-api'
@@ -147,8 +151,9 @@ Requires-Dist: gcsfs <2024.4,>=2023.9.2 ; extra == 'complete-api'
 Requires-Dist: google-cloud-bigquery[bqstorage,pandas] ==3.14.1 ; extra == 'complete-api'
 Requires-Dist: graphviz ~=0.20.0 ; extra == 'complete-api'
 Requires-Dist: humanfriendly ~=10.0 ; extra == 'complete-api'
-Requires-Dist: igz-mgmt ~=0.1.3 ; extra == 'complete-api'
+Requires-Dist: igz-mgmt ~=0.2.0 ; extra == 'complete-api'
 Requires-Dist: kafka-python ~=2.0 ; extra == 'complete-api'
+Requires-Dist: memray ~=1.12 ; extra == 'complete-api'
 Requires-Dist: mlflow ~=2.8 ; extra == 'complete-api'
 Requires-Dist: msrest ~=0.6.21 ; extra == 'complete-api'
 Requires-Dist: objgraph ~=3.6 ; extra == 'complete-api'
@@ -160,6 +165,7 @@ Requires-Dist: pyopenssl >=23 ; extra == 'complete-api'
 Requires-Dist: redis ~=4.3 ; extra == 'complete-api'
 Requires-Dist: s3fs <2024.4,>=2023.9.2 ; extra == 'complete-api'
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'complete-api'
+Requires-Dist: taos-ws-py ~=0.3.2 ; extra == 'complete-api'
 Requires-Dist: timelength ~=1.1 ; extra == 'complete-api'
 Requires-Dist: uvicorn ~=0.27.1 ; extra == 'complete-api'
 Provides-Extra: dask
@@ -192,6 +198,8 @@ Requires-Dist: aiobotocore <2.8,>=2.5.0 ; extra == 's3'
 Requires-Dist: s3fs <2024.4,>=2023.9.2 ; extra == 's3'
 Provides-Extra: sqlalchemy
 Requires-Dist: sqlalchemy ~=1.4 ; extra == 'sqlalchemy'
+Provides-Extra: tdengine
+Requires-Dist: taos-ws-py ~=0.3.2 ; extra == 'tdengine'
 
 <a id="top"></a>
 [![Build Status](https://github.com/mlrun/mlrun/actions/workflows/build.yaml/badge.svg?branch=development)](https://github.com/mlrun/mlrun/actions/workflows/build.yaml?query=branch%3Adevelopment)