mlrun 1.3.1rc5__py3-none-any.whl → 1.4.0rc2__py3-none-any.whl

mlrun/api/db/base.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import datetime
+import typing
 import warnings
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import mlrun.model
 from mlrun.api import schemas
 
 
@@ -103,6 +105,7 @@ class DBInterface(ABC):
         max_partitions: int = 0,
         requested_logs: bool = None,
         return_as_run_structs: bool = True,
+        with_notifications: bool = False,
     ):
         pass
 
@@ -575,3 +578,32 @@ class DBInterface(ABC):
         self, session, name: str, project: str
     ) -> schemas.BackgroundTask:
         pass
+
+    @abstractmethod
+    def store_run_notifications(
+        self,
+        session,
+        notification_objects: typing.List[mlrun.model.Notification],
+        run_uid: str,
+        project: str,
+    ):
+        pass
+
+    @abstractmethod
+    def list_run_notifications(
+        self,
+        session,
+        run_uid: str,
+        project: str,
+    ) -> typing.List[mlrun.model.Notification]:
+        pass
+
+    def delete_run_notifications(
+        self,
+        session,
+        name: str = None,
+        run_uid: str = None,
+        project: str = None,
+        commit: bool = True,
+    ):
+        pass

mlrun/api/db/session.py

@@ -15,22 +15,14 @@
 from sqlalchemy.orm import Session
 
 from mlrun.api.db.sqldb.session import create_session as sqldb_create_session
-from mlrun.config import config
 
 
-def create_session(db_type=None) -> Session:
-    db_type = db_type or config.httpdb.db_type
-    if db_type == "filedb":
-        return None
-    else:
-        return sqldb_create_session()
+def create_session() -> Session:
+    return sqldb_create_session()
 
 
 def close_session(db_session):
-
-    # will be None when it's filedb session
-    if db_session is not None:
-        db_session.close()
+    db_session.close()
 
 
 def run_function_with_new_db_session(func):

mlrun/api/db/sqldb/db.py

@@ -31,7 +31,9 @@ from sqlalchemy.orm import Session, aliased
 import mlrun
 import mlrun.api.db.session
 import mlrun.api.utils.projects.remotes.follower
+import mlrun.api.utils.singletons.k8s
 import mlrun.errors
+import mlrun.model
 from mlrun.api import schemas
 from mlrun.api.db.base import DBInterface
 from mlrun.api.db.sqldb.helpers import (
@@ -53,6 +55,7 @@ from mlrun.api.db.sqldb.models import (
     Function,
     Log,
     MarketplaceSource,
+    Notification,
     Project,
     Run,
     Schedule,
@@ -325,6 +328,7 @@ class SQLDB(DBInterface):
         max_partitions: int = 0,
         requested_logs: bool = None,
         return_as_run_structs: bool = True,
+        with_notifications: bool = False,
     ):
         project = project or config.default_project
         query = self._find_runs(session, uid, project, labels)
@@ -369,9 +373,28 @@ class SQLDB(DBInterface):
         if not return_as_run_structs:
             return query.all()
 
+        # Purposefully not using outer join to avoid returning runs without notifications
+        if with_notifications:
+            query = query.join(Notification, Run.id == Notification.run)
+
         runs = RunList()
         for run in query:
-            runs.append(run.struct)
+            run_struct = run.struct
+            if with_notifications:
+                run_struct.setdefault("spec", {}).setdefault("notifications", [])
+                run_struct.setdefault("status", {}).setdefault("notifications", {})
+                for notification in run.notifications:
+                    (
+                        notification_spec,
+                        notification_status,
+                    ) = self._transform_notification_record_to_spec_and_status(
+                        notification
+                    )
+                    run_struct["spec"]["notifications"].append(notification_spec)
+                    run_struct["status"]["notifications"][
+                        notification.name
+                    ] = notification_status
+            runs.append(run_struct)
 
         return runs
 
@@ -1689,6 +1712,10 @@ class SQLDB(DBInterface):
         self._verify_empty_list_of_project_related_resources(name, logs, "logs")
         runs = self._find_runs(session, None, name, []).all()
         self._verify_empty_list_of_project_related_resources(name, runs, "runs")
+        notifications = self._get_db_notifications(session, project=name)
+        self._verify_empty_list_of_project_related_resources(
+            name, notifications, "notifications"
+        )
         schedules = self.list_schedules(session, project=name)
         self._verify_empty_list_of_project_related_resources(
             name, schedules, "schedules"
@@ -1709,6 +1736,7 @@ class SQLDB(DBInterface):
     def delete_project_related_resources(self, session: Session, name: str):
         self.del_artifacts(session, project=name)
         self._delete_logs(session, name)
+        self.delete_run_notifications(session, project=name)
         self.del_runs(session, project=name)
         self.delete_schedules(session, name)
         self._delete_functions(session, name)
@@ -2832,6 +2860,13 @@ class SQLDB(DBInterface):
             query = query.filter(Run.uid.in_(uid))
         return self._add_labels_filter(session, query, Run, labels)
 
+    def _get_db_notifications(
+        self, session, name: str = None, run_id: int = None, project: str = None
+    ):
+        return self._query(
+            session, Notification, name=name, run=run_id, project=project
+        ).all()
+
     def _latest_uid_filter(self, session, query):
         # Create a sub query of latest uid (by updated) per (project,key)
         subq = (
@@ -3168,6 +3203,35 @@ class SQLDB(DBInterface):
         # TODO: handle transforming the functions/workflows/artifacts references to real objects
         return schemas.Project(**project_record.full_object)
 
+    def _transform_notification_record_to_spec_and_status(
+        self,
+        notification_record: Notification,
+    ) -> typing.Tuple[dict, dict]:
+        notification_spec = self._transform_notification_record_to_schema(
+            notification_record
+        ).to_dict()
+        notification_status = {
+            "status": notification_spec.pop("status", None),
+            "sent_time": notification_spec.pop("sent_time", None),
+        }
+        return notification_spec, notification_status
+
+    @staticmethod
+    def _transform_notification_record_to_schema(
+        notification_record: Notification,
+    ) -> mlrun.model.Notification:
+        return mlrun.model.Notification(
+            kind=notification_record.kind,
+            name=notification_record.name,
+            message=notification_record.message,
+            severity=notification_record.severity,
+            when=notification_record.when.split(","),
+            condition=notification_record.condition,
+            params=notification_record.params,
+            status=notification_record.status,
+            sent_time=notification_record.sent_time,
+        )
+
     def _move_and_reorder_table_items(
         self, session, moved_object, move_to=None, move_from=None
     ):
@@ -3543,3 +3607,100 @@ class SQLDB(DBInterface):
         ):
             return True
         return False
+
+    def store_run_notifications(
+        self,
+        session,
+        notification_objects: typing.List[mlrun.model.Notification],
+        run_uid: str,
+        project: str,
+    ):
+        # iteration is 0, as we don't support multiple notifications per hyper param run, only for the whole run
+        run = self._get_run(session, run_uid, project, 0)
+        if not run:
+            raise mlrun.errors.MLRunNotFoundError(
+                f"Run not found: uid={run_uid}, project={project}"
+            )
+
+        run_notifications = {
+            notification.name: notification
+            for notification in self._get_db_notifications(session, run_id=run.id)
+        }
+        notifications = []
+        for notification_model in notification_objects:
+            new_notification = False
+            notification = run_notifications.get(notification_model.name, None)
+            if not notification:
+                new_notification = True
+                notification = Notification(
+                    name=notification_model.name, run=run.id, project=project
+                )
+
+            notification.kind = notification_model.kind
+            notification.message = notification_model.message
+            notification.severity = notification_model.severity
+            notification.when = ",".join(notification_model.when)
+            notification.condition = notification_model.condition
+            notification.params = notification_model.params
+            notification.status = (
+                notification_model.status
+                or mlrun.api.schemas.NotificationStatus.PENDING
+            )
+            notification.sent_time = notification_model.sent_time
+
+            logger.debug(
+                f"Storing {'new' if new_notification else 'existing'} notification",
+                notification_name=notification.name,
+                run_uid=run_uid,
+                project=project,
+            )
+            notifications.append(notification)
+
+        self._upsert(session, notifications)
+
+    def list_run_notifications(
+        self,
+        session,
+        run_uid: str,
+        project: str = "",
+    ) -> typing.List[mlrun.model.Notification]:
+
+        # iteration is 0, as we don't support multiple notifications per hyper param run, only for the whole run
+        run = self._get_run(session, run_uid, project, 0)
+        if not run:
+            return []
+
+        return [
+            self._transform_notification_record_to_schema(notification)
+            for notification in self._query(session, Notification, run=run.id).all()
+        ]
+
+    def delete_run_notifications(
+        self,
+        session,
+        name: str = None,
+        run_uid: str = None,
+        project: str = None,
+        commit: bool = True,
+    ):
+        run_id = None
+        if run_uid:
+
+            # iteration is 0, as we don't support multiple notifications per hyper param run, only for the whole run
+            run = self._get_run(session, run_uid, project, 0)
+            if not run:
+                raise mlrun.errors.MLRunNotFoundError(
+                    f"Run not found: uid={run_uid}, project={project}"
+                )
+            run_id = run.id
+
+        project = project or config.default_project
+        if project == "*":
+            project = None
+
+        query = self._get_db_notifications(session, name, run_id, project)
+        for notification in query:
+            session.delete(notification)
+
+        if commit:
+            session.commit()

mlrun/api/db/sqldb/models/models_mysql.py

@@ -139,6 +139,46 @@ with warnings.catch_warnings():
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.name}/{self.uid}"
 
+    class Notification(Base, mlrun.utils.db.BaseModel):
+        __tablename__ = "notifications"
+        __table_args__ = (UniqueConstraint("name", "run", name="_notifications_uc"),)
+
+        id = Column(Integer, primary_key=True)
+        project = Column(String(255, collation=SQLCollationUtil.collation()))
+        name = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        kind = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        message = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        severity = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        when = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        condition = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        params = Column("params", JSON)
+        run = Column(Integer, ForeignKey("runs.id"))
+
+        # TODO: Separate table for notification state.
+        #   Currently, we are only supporting one notification being sent per DB row (either on completion or on error).
+        #   In the future, we might want to support multiple notifications per DB row, and we might want to support on
+        #   start, therefore we need to separate the state from the notification itself (e.g. this table can be  table
+        #   with notification_id, state, when, last_sent, etc.). This will require some refactoring in the code.
+        sent_time = Column(
+            sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3),
+            nullable=True,
+        )
+        status = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+
     class Log(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "logs"
 
@@ -182,6 +222,7 @@ with warnings.catch_warnings():
 
         labels = relationship(Label, cascade="all, delete-orphan")
         tags = relationship(Tag, cascade="all, delete-orphan")
+        notifications = relationship(Notification, cascade="all, delete-orphan")
 
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.uid}/{self.iteration}"

mlrun/api/db/sqldb/models/models_sqlite.py

@@ -151,6 +151,40 @@ with warnings.catch_warnings():
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.uid}"
 
+    class Notification(Base, mlrun.utils.db.BaseModel):
+        __tablename__ = "notifications"
+        __table_args__ = (UniqueConstraint("name", "run", name="_notifications_uc"),)
+
+        id = Column(Integer, primary_key=True)
+        project = Column(String(255, collation=SQLCollationUtil.collation()))
+        name = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        kind = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        message = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        severity = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        when = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        condition = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        params = Column("params", JSON)
+        run = Column(Integer, ForeignKey("runs.id"))
+        sent_time = Column(
+            TIMESTAMP(),
+            nullable=True,
+        )
+        status = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+
     class Run(Base, mlrun.utils.db.HasStruct):
         __tablename__ = "runs"
         __table_args__ = (
@@ -178,6 +212,7 @@ with warnings.catch_warnings():
         requested_logs = Column(BOOLEAN)
         updated = Column(TIMESTAMP, default=datetime.utcnow)
         labels = relationship(Label)
+        notifications = relationship(Notification, cascade="all, delete-orphan")
 
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.uid}/{self.iteration}"

mlrun/api/main.py

@@ -24,12 +24,14 @@ import sqlalchemy.orm
 import uvicorn
 from fastapi.exception_handlers import http_exception_handler
 
+import mlrun.api.db.base
 import mlrun.api.schemas
 import mlrun.api.utils.clients.chief
 import mlrun.api.utils.clients.log_collector
 import mlrun.errors
 import mlrun.lists
 import mlrun.utils
+import mlrun.utils.notifications
 import mlrun.utils.version
 from mlrun.api.api.api import api_router
 from mlrun.api.db.session import close_session, create_session
@@ -56,6 +58,14 @@ from mlrun.utils import logger
 API_PREFIX = "/api"
 BASE_VERSIONED_API_PREFIX = f"{API_PREFIX}/v1"
 
+# When pushing notifications, push notifications only for runs that entered a terminal state
+# since the last time we pushed notifications.
+# On the first time we push notifications, we'll push notifications for all runs that are in a terminal state
+# and their notifications haven't been sent yet.
+# TODO: find better solution than a global variable for chunking the list of runs
+#      for which to push notifications
+_last_notification_push_time: datetime.datetime = None
+
 
 app = fastapi.FastAPI(
     title="MLRun",
@@ -512,18 +522,20 @@ async def _align_worker_state_with_chief_state(
 
 
 def _monitor_runs():
+    db = get_db()
     db_session = create_session()
     try:
         for kind in RuntimeKinds.runtime_with_handlers():
             try:
                 runtime_handler = get_runtime_handler(kind)
-                runtime_handler.monitor_runs(get_db(), db_session)
+                runtime_handler.monitor_runs(db, db_session)
             except Exception as exc:
                 logger.warning(
                     "Failed monitoring runs. Ignoring",
                     exc=err_to_str(exc),
                     kind=kind,
                 )
+        _push_terminal_run_notifications(db, db_session)
     finally:
         close_session(db_session)
 
@@ -545,6 +557,47 @@ def _cleanup_runtimes():
         close_session(db_session)
 
 
+def _push_terminal_run_notifications(db: mlrun.api.db.base.DBInterface, db_session):
+    """
+    Get all runs with notification configs which became terminal since the last call to the function
+    and push their notifications if they haven't been pushed yet.
+    """
+
+    # Import here to avoid circular import
+    import mlrun.api.api.utils
+
+    # When pushing notifications, push notifications only for runs that entered a terminal state
+    # since the last time we pushed notifications.
+    # On the first time we push notifications, we'll push notifications for all runs that are in a terminal state
+    # and their notifications haven't been sent yet.
+    global _last_notification_push_time
+
+    runs = db.list_runs(
+        db_session,
+        project="*",
+        states=mlrun.runtimes.constants.RunStates.terminal_states(),
+        last_update_time_from=_last_notification_push_time,
+        with_notifications=True,
+    )
+
+    if not len(runs):
+        return
+
+    # Unmasking the run parameters from secrets before handing them over to the notification handler
+    # as importing the `Secrets` crud in the notification handler will cause a circular import
+    unmasked_runs = [
+        mlrun.api.api.utils.unmask_notification_params_secret_on_task(run)
+        for run in runs
+    ]
+
+    logger.debug(
+        "Got terminal runs with configured notifications", runs_amount=len(runs)
+    )
+    mlrun.utils.notifications.NotificationPusher(unmasked_runs).push(db)
+
+    _last_notification_push_time = datetime.datetime.now(datetime.timezone.utc)
+
+
 async def _stop_logs():
     """
     Stop logs for runs that are in terminal state and last updated in the previous interval

mlrun/api/migrations_mysql/versions/c905d15bd91d_notifications.py

@@ -0,0 +1,70 @@
+# Copyright 2018 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""notifications
+
+Revision ID: c905d15bd91d
+Revises: 88e656800d6a
+Create Date: 2022-09-20 10:44:41.727488
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = "c905d15bd91d"
+down_revision = "88e656800d6a"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "notifications",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("project", sa.String(length=255, collation="utf8_bin")),
+        sa.Column("name", sa.String(length=255, collation="utf8_bin"), nullable=False),
+        sa.Column("kind", sa.String(length=255, collation="utf8_bin"), nullable=False),
+        sa.Column(
+            "message", sa.String(length=255, collation="utf8_bin"), nullable=False
+        ),
+        sa.Column(
+            "severity", sa.String(length=255, collation="utf8_bin"), nullable=False
+        ),
+        sa.Column("when", sa.String(length=255, collation="utf8_bin"), nullable=False),
+        sa.Column(
+            "condition", sa.String(length=255, collation="utf8_bin"), nullable=False
+        ),
+        sa.Column("params", sa.JSON(), nullable=True),
+        sa.Column("run", sa.Integer(), nullable=True),
+        sa.Column("sent_time", mysql.TIMESTAMP(fsp=3), nullable=True),
+        sa.Column(
+            "status", sa.String(length=255, collation="utf8_bin"), nullable=False
+        ),
+        sa.ForeignKeyConstraint(
+            ["run"],
+            ["runs.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("name", "run", name="_notifications_uc"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("notifications")
+    # ### end Alembic commands ###

mlrun/api/migrations_sqlite/versions/959ae00528ad_notifications.py

@@ -0,0 +1,61 @@
+# Copyright 2018 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""notifications
+
+Revision ID: 959ae00528ad
+Revises: 803438ecd005
+Create Date: 2022-09-20 10:40:41.354209
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "959ae00528ad"
+down_revision = "803438ecd005"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "notifications",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("project", sa.String(length=255)),
+        sa.Column("name", sa.String(length=255), nullable=False),
+        sa.Column("kind", sa.String(length=255), nullable=False),
+        sa.Column("message", sa.String(length=255), nullable=False),
+        sa.Column("severity", sa.String(length=255), nullable=False),
+        sa.Column("when", sa.String(length=255), nullable=False),
+        sa.Column("condition", sa.String(length=255), nullable=False),
+        sa.Column("params", sa.JSON(), nullable=True),
+        sa.Column("run", sa.Integer(), nullable=True),
+        sa.Column("sent_time", sa.TIMESTAMP(), nullable=True),
+        sa.Column("status", sa.String(length=255), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["run"],
+            ["runs.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("name", "run", name="_notifications_uc"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("notifications")
+    # ### end Alembic commands ###

mlrun/api/schemas/__init__.py

@@ -107,6 +107,7 @@ from .model_endpoints import (
     ModelEndpointStatus,
     ModelMonitoringStoreKinds,
 )
+from .notification import NotificationSeverity, NotificationStatus
 from .object import ObjectKind, ObjectMetadata, ObjectSpec, ObjectStatus
 from .pipeline import PipelinesFormat, PipelinesOutput, PipelinesPagination
 from .project import (