mlrun 1.3.3rc1__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/api/crud/notifications.py

@@ -0,0 +1,149 @@
+# Copyright 2023 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.
+#
+import typing
+
+import sqlalchemy.orm
+
+import mlrun.api.api.utils
+import mlrun.api.db.sqldb.db
+import mlrun.api.utils.scheduler
+import mlrun.api.utils.singletons.db
+import mlrun.api.utils.singletons.scheduler
+import mlrun.common.schemas
+import mlrun.utils.singleton
+
+
+class Notifications(
+    metaclass=mlrun.utils.singleton.Singleton,
+):
+    def store_run_notifications(
+        self,
+        session: sqlalchemy.orm.Session,
+        notification_objects: typing.List[mlrun.model.Notification],
+        run_uid: str,
+        project: str = None,
+    ):
+        project = project or mlrun.mlconf.default_project
+        notification_objects_to_store = (
+            mlrun.api.api.utils.validate_and_mask_notification_list(
+                notification_objects, run_uid, project
+            )
+        )
+
+        mlrun.api.utils.singletons.db.get_db().store_run_notifications(
+            session, notification_objects_to_store, run_uid, project
+        )
+
+    def list_run_notifications(
+        self,
+        session: sqlalchemy.orm.Session,
+        run_uid: str,
+        project: str = "",
+    ) -> typing.List[mlrun.model.Notification]:
+        project = project or mlrun.mlconf.default_project
+        return mlrun.api.utils.singletons.db.get_db().list_run_notifications(
+            session, run_uid, project
+        )
+
+    def delete_run_notifications(
+        self,
+        session: sqlalchemy.orm.Session,
+        name: str = None,
+        run_uid: str = None,
+        project: str = None,
+    ):
+        project = project or mlrun.mlconf.default_project
+
+        # Delete notification param project secret
+        notifications = [
+            notification
+            for notification in self.list_run_notifications(session, run_uid, project)
+            if notification.name == name
+        ]
+        if notifications:
+            # unique constraint on name, run_uid, project, so the list will contain one item at most
+            notification = notifications[0]
+            mlrun.api.api.utils.delete_notification_params_secret(project, notification)
+
+        mlrun.api.utils.singletons.db.get_db().delete_run_notifications(
+            session, name, run_uid, project
+        )
+
+    @staticmethod
+    def set_object_notifications(
+        db_session: sqlalchemy.orm.Session,
+        auth_info: mlrun.common.schemas.AuthInfo,
+        project: str,
+        notifications: typing.List[mlrun.common.schemas.Notification],
+        notification_parent: typing.Union[
+            mlrun.common.schemas.RunIdentifier, mlrun.common.schemas.ScheduleIdentifier
+        ],
+    ):
+        """
+        Sets notifications on given object (run or schedule, might be extended in the future).
+        This will replace any existing notifications.
+        :param db_session: DB session
+        :param auth_info: Authorization info
+        :param project: Project name
+        :param notifications: List of notifications to set
+        :param notification_parent: Identifier of the object on which to set the notifications
+        """
+        set_notification_methods = {
+            "run": {
+                "factory": mlrun.api.utils.singletons.db.get_db,
+                "method_name": mlrun.api.db.sqldb.db.SQLDB.set_run_notifications.__name__,
+                "identifier_key": "uid",
+            },
+            "schedule": {
+                "factory": mlrun.api.utils.singletons.scheduler.get_scheduler,
+                "method_name": mlrun.api.utils.scheduler.Scheduler.set_schedule_notifications.__name__,
+                "identifier_key": "name",
+            },
+        }
+
+        set_notification_method = set_notification_methods.get(
+            notification_parent.kind, {}
+        )
+        factory = set_notification_method.get("factory")
+        if not factory:
+            raise mlrun.errors.MLRunNotFoundError(
+                f"couldn't find factory for object kind: {notification_parent.kind}"
+            )
+        set_func = set_notification_method.get("method_name")
+        if not set_func:
+            raise mlrun.errors.MLRunNotFoundError(
+                f"couldn't find set notification function for object kind: {notification_parent.kind}"
+            )
+        identifier_key = set_notification_method.get("identifier_key")
+        if not identifier_key:
+            raise mlrun.errors.MLRunNotFoundError(
+                f"couldn't find identifier key for object kind: {notification_parent.kind}"
+            )
+
+        notification_objects_to_set = (
+            mlrun.api.api.utils.validate_and_mask_notification_list(
+                notifications,
+                getattr(notification_parent, identifier_key),
+                project,
+            )
+        )
+
+        getattr(factory(), set_func)(
+            session=db_session,
+            project=project,
+            notifications=notification_objects_to_set,
+            identifier=notification_parent,
+            auth_info=auth_info,
+        )

mlrun/api/crud/pipelines.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,11 +20,12 @@ import traceback
 import typing
 
 import kfp
+import kfp_server_api
 import sqlalchemy.orm
 
 import mlrun
 import mlrun.api.api.utils
-import mlrun.api.schemas
+import mlrun.common.schemas
 import mlrun.errors
 import mlrun.kfpops
 import mlrun.utils.helpers
@@ -44,26 +45,22 @@ class Pipelines(
         sort_by: str = "",
         page_token: str = "",
         filter_: str = "",
-        format_: mlrun.api.schemas.PipelinesFormat = mlrun.api.schemas.PipelinesFormat.metadata_only,
+        format_: mlrun.common.schemas.PipelinesFormat = mlrun.common.schemas.PipelinesFormat.metadata_only,
         page_size: typing.Optional[int] = None,
     ) -> typing.Tuple[int, typing.Optional[int], typing.List[dict]]:
-        if project != "*" and (page_token or page_size or sort_by):
+        if project != "*" and (page_token or page_size):
             raise mlrun.errors.MLRunInvalidArgumentError(
-                "Filtering by project can not be used together with pagination, or sorting"
+                "Filtering by project can not be used together with pagination"
             )
-        if format_ == mlrun.api.schemas.PipelinesFormat.summary:
+        if format_ == mlrun.common.schemas.PipelinesFormat.summary:
             # we don't support summary format in list pipelines since the returned runs doesn't include the workflow
             # manifest status that includes the nodes section we use to generate the DAG.
             # (There is a workflow manifest under the run's pipeline_spec field, but it doesn't include the status)
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Summary format is not supported for list pipelines, use get instead"
             )
-        kfp_url = mlrun.mlconf.resolve_kfp_url(namespace)
-        if not kfp_url:
-            raise mlrun.errors.MLRunNotFoundError(
-                "KubeFlow Pipelines is not configured"
-            )
-        kfp_client = kfp.Client(host=kfp_url)
+
+        kfp_client = self.initialize_kfp_client(namespace)
         if project != "*":
             run_dicts = []
             while page_token is not None:
@@ -72,7 +69,8 @@ class Pipelines(
                 # the filter that was used to create it)
                 response = kfp_client._run_api.list_runs(
                     page_token=page_token,
-                    page_size=mlrun.api.schemas.PipelinesPagination.max_page_size,
+                    page_size=mlrun.common.schemas.PipelinesPagination.max_page_size,
+                    sort_by=sort_by,
                     filter=filter_ if page_token == "" else "",
                 )
                 run_dicts.extend([run.to_dict() for run in response.runs or []])
@@ -89,7 +87,7 @@ class Pipelines(
             response = kfp_client._run_api.list_runs(
                 page_token=page_token,
                 page_size=page_size
-                or mlrun.api.schemas.PipelinesPagination.default_page_size,
+                or mlrun.common.schemas.PipelinesPagination.default_page_size,
                 sort_by=sort_by,
                 filter=filter_,
             )
@@ -100,20 +98,45 @@ class Pipelines(
 
         return total_size, next_page_token, runs
 
+    def delete_pipelines_runs(self, db_session: sqlalchemy.orm.Session, project: str):
+        _, _, project_pipeline_runs = self.list_pipelines(
+            db_session=db_session,
+            project=project,
+            format_=mlrun.common.schemas.PipelinesFormat.metadata_only,
+        )
+        kfp_client = self.initialize_kfp_client()
+
+        if project_pipeline_runs:
+            logger.debug(
+                "Detected pipeline runs for project, deleting them",
+                project=project,
+                pipeline_run_ids=[run["id"] for run in project_pipeline_runs],
+            )
+
+        for pipeline_run in project_pipeline_runs:
+            try:
+                # delete pipeline run also terminates it if it is in progress
+                kfp_client._run_api.delete_run(pipeline_run["id"])
+            except Exception as exc:
+                # we don't want to fail the entire delete operation if we failed to delete a single pipeline run
+                # so it won't fail the delete project operation. we will log the error and continue
+                logger.warning(
+                    "Failed to delete pipeline run",
+                    project=project,
+                    pipeline_run_id=pipeline_run["id"],
+                    exc_info=exc,
+                )
+        logger.debug("Finished deleting pipeline runs", project=project)
+
     def get_pipeline(
         self,
         db_session: sqlalchemy.orm.Session,
         run_id: str,
         project: typing.Optional[str] = None,
         namespace: typing.Optional[str] = None,
-        format_: mlrun.api.schemas.PipelinesFormat = mlrun.api.schemas.PipelinesFormat.summary,
+        format_: mlrun.common.schemas.PipelinesFormat = mlrun.common.schemas.PipelinesFormat.summary,
     ):
-        kfp_url = mlrun.mlconf.resolve_kfp_url(namespace)
-        if not kfp_url:
-            raise mlrun.errors.MLRunBadRequestError(
-                "KubeFlow Pipelines is not configured"
-            )
-        kfp_client = kfp.Client(host=kfp_url)
+        kfp_client = self.initialize_kfp_client(namespace)
         run = None
         try:
             api_run_detail = kfp_client.get_run(run_id)
@@ -122,13 +145,16 @@ class Pipelines(
                 if project and project != "*":
                     run_project = self.resolve_project_from_pipeline(run)
                     if run_project != project:
-                        raise mlrun.errors.MLRunInvalidArgumentError(
+                        raise mlrun.errors.MLRunNotFoundError(
                             f"Pipeline run with id {run_id} is not of project {project}"
                         )
                 run = self._format_run(
                     db_session, run, format_, api_run_detail.to_dict()
                 )
-
+        except kfp_server_api.ApiException as exc:
+            mlrun.errors.raise_for_status_code(int(exc.status), err_to_str(exc))
+        except mlrun.errors.MLRunHTTPStatusError:
+            raise
         except Exception as exc:
             raise mlrun.errors.MLRunRuntimeError(
                 f"Failed getting kfp run: {err_to_str(exc)}"
@@ -158,7 +184,6 @@ class Pipelines(
             )
 
         logger.debug("Writing pipeline to temp file", content_type=content_type)
-        print(str(data))
 
         pipeline_file = tempfile.NamedTemporaryFile(suffix=content_type)
         with open(pipeline_file.name, "wb") as fp:
@@ -172,12 +197,7 @@ class Pipelines(
         )
 
         try:
-            kfp_url = mlrun.mlconf.resolve_kfp_url(namespace)
-            if not kfp_url:
-                raise mlrun.errors.MLRunBadRequestError(
-                    "KubeFlow Pipelines is not configured"
-                )
-            kfp_client = kfp.Client(host=kfp_url)
+            kfp_client = self.initialize_kfp_client(namespace)
             experiment = kfp_client.create_experiment(name=experiment_name)
             run = kfp_client.run_pipeline(
                 experiment.id, run_name, pipeline_file.name, params=arguments
@@ -196,11 +216,20 @@ class Pipelines(
 
         return run
 
+    @staticmethod
+    def initialize_kfp_client(namespace: typing.Optional[str] = None) -> kfp.Client:
+        kfp_url = mlrun.mlconf.resolve_kfp_url(namespace)
+        if not kfp_url:
+            raise mlrun.errors.MLRunNotFoundError(
+                "KubeFlow Pipelines is not configured"
+            )
+        return kfp.Client(host=kfp_url)
+
     def _format_runs(
         self,
         db_session: sqlalchemy.orm.Session,
         runs: typing.List[dict],
-        format_: mlrun.api.schemas.PipelinesFormat = mlrun.api.schemas.PipelinesFormat.metadata_only,
+        format_: mlrun.common.schemas.PipelinesFormat = mlrun.common.schemas.PipelinesFormat.metadata_only,
     ) -> typing.List[dict]:
         formatted_runs = []
         for run in runs:
@@ -211,32 +240,18 @@ class Pipelines(
         self,
         db_session: sqlalchemy.orm.Session,
         run: dict,
-        format_: mlrun.api.schemas.PipelinesFormat = mlrun.api.schemas.PipelinesFormat.metadata_only,
+        format_: mlrun.common.schemas.PipelinesFormat = mlrun.common.schemas.PipelinesFormat.metadata_only,
         api_run_detail: typing.Optional[dict] = None,
     ) -> dict:
         run["project"] = self.resolve_project_from_pipeline(run)
-        if format_ == mlrun.api.schemas.PipelinesFormat.full:
+        if format_ == mlrun.common.schemas.PipelinesFormat.full:
             return run
-        elif format_ == mlrun.api.schemas.PipelinesFormat.metadata_only:
-            return {
-                k: str(v)
-                for k, v in run.items()
-                if k
-                in [
-                    "id",
-                    "name",
-                    "project",
-                    "status",
-                    "error",
-                    "created_at",
-                    "scheduled_at",
-                    "finished_at",
-                    "description",
-                ]
-            }
-        elif format_ == mlrun.api.schemas.PipelinesFormat.name_only:
+        elif format_ == mlrun.common.schemas.PipelinesFormat.metadata_only:
+            return mlrun.utils.helpers.format_run(run, with_project=True)
+
+        elif format_ == mlrun.common.schemas.PipelinesFormat.name_only:
             return run.get("name")
-        elif format_ == mlrun.api.schemas.PipelinesFormat.summary:
+        elif format_ == mlrun.common.schemas.PipelinesFormat.summary:
             if not api_run_detail:
                 raise mlrun.errors.MLRunRuntimeError(
                     "The full kfp api_run_detail object is needed to generate the summary format"

mlrun/api/crud/projects.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,11 +23,12 @@ import sqlalchemy.orm
 
 import mlrun.api.crud
 import mlrun.api.db.session
-import mlrun.api.schemas
+import mlrun.api.utils.events.events_factory as events_factory
 import mlrun.api.utils.projects.remotes.follower
 import mlrun.api.utils.singletons.db
 import mlrun.api.utils.singletons.k8s
 import mlrun.api.utils.singletons.scheduler
+import mlrun.common.schemas
 import mlrun.errors
 import mlrun.utils.singleton
 from mlrun.utils import logger
@@ -44,7 +45,7 @@ class Projects(
         }
 
     def create_project(
-        self, session: sqlalchemy.orm.Session, project: mlrun.api.schemas.Project
+        self, session: sqlalchemy.orm.Session, project: mlrun.common.schemas.Project
     ):
         logger.debug(
             "Creating project",
@@ -63,7 +64,7 @@ class Projects(
         self,
         session: sqlalchemy.orm.Session,
         name: str,
-        project: mlrun.api.schemas.Project,
+        project: mlrun.common.schemas.Project,
     ):
         logger.debug(
             "Storing project",
@@ -83,7 +84,7 @@ class Projects(
         session: sqlalchemy.orm.Session,
         name: str,
         project: dict,
-        patch_mode: mlrun.api.schemas.PatchMode = mlrun.api.schemas.PatchMode.replace,
+        patch_mode: mlrun.common.schemas.PatchMode = mlrun.common.schemas.PatchMode.replace,
     ):
         logger.debug(
             "Patching project", name=name, project=project, patch_mode=patch_mode
@@ -96,12 +97,12 @@ class Projects(
         self,
         session: sqlalchemy.orm.Session,
         name: str,
-        deletion_strategy: mlrun.api.schemas.DeletionStrategy = mlrun.api.schemas.DeletionStrategy.default(),
+        deletion_strategy: mlrun.common.schemas.DeletionStrategy = mlrun.common.schemas.DeletionStrategy.default(),
     ):
         logger.debug("Deleting project", name=name, deletion_strategy=deletion_strategy)
         if (
             deletion_strategy.is_restricted()
-            or deletion_strategy == mlrun.api.schemas.DeletionStrategy.check
+            or deletion_strategy == mlrun.common.schemas.DeletionStrategy.check
         ):
             if not mlrun.api.utils.singletons.db.get_db().is_project_exists(
                 session, name
@@ -111,7 +112,7 @@ class Projects(
                 session, name
             )
             self._verify_project_has_no_external_resources(name)
-            if deletion_strategy == mlrun.api.schemas.DeletionStrategy.check:
+            if deletion_strategy == mlrun.common.schemas.DeletionStrategy.check:
                 return
         elif deletion_strategy.is_cascading():
             self.delete_project_resources(session, name)
@@ -134,7 +135,7 @@ class Projects(
         # Therefore, this check should remain at the end of the verification flow.
         if (
             mlrun.mlconf.is_api_running_on_k8s()
-            and mlrun.api.utils.singletons.k8s.get_k8s().get_project_secret_keys(
+            and mlrun.api.utils.singletons.k8s.get_k8s_helper().get_project_secret_keys(
                 project
             )
         ):
@@ -158,11 +159,16 @@ class Projects(
             label_selector=f"mlrun/project={name}",
             force=True,
         )
+        if mlrun.mlconf.resolve_kfp_url():
+            logger.debug("Removing KFP pipelines runs for project", project=name)
+            mlrun.api.crud.pipelines.Pipelines().delete_pipelines_runs(
+                db_session=session, project=name
+            )
 
         # log collector service will delete the logs, so we don't need to do it here
         if (
             mlrun.mlconf.log_collector.mode
-            == mlrun.api.schemas.LogsCollectorMode.legacy
+            == mlrun.common.schemas.LogsCollectorMode.legacy
         ):
             mlrun.api.crud.Logs().delete_logs(name)
 
@@ -176,22 +182,44 @@ class Projects(
 
         # delete project secrets - passing None will delete all secrets
         if mlrun.mlconf.is_api_running_on_k8s():
-            mlrun.api.utils.singletons.k8s.get_k8s().delete_project_secrets(name, None)
+            secrets = None
+            (
+                secret_name,
+                action,
+            ) = mlrun.api.utils.singletons.k8s.get_k8s_helper().delete_project_secrets(
+                name, secrets
+            )
+            if action:
+                events_client = events_factory.EventsFactory().get_events_client()
+                events_client.emit(
+                    events_client.generate_project_secret_event(
+                        name,
+                        secret_name,
+                        action=action,
+                    )
+                )
+
+            else:
+                logger.debug(
+                    "No project secrets to delete",
+                    action=action,
+                    secret_name=secret_name,
+                )
 
     def get_project(
         self, session: sqlalchemy.orm.Session, name: str
-    ) -> mlrun.api.schemas.Project:
+    ) -> mlrun.common.schemas.Project:
         return mlrun.api.utils.singletons.db.get_db().get_project(session, name)
 
     def list_projects(
         self,
         session: sqlalchemy.orm.Session,
         owner: str = None,
-        format_: mlrun.api.schemas.ProjectsFormat = mlrun.api.schemas.ProjectsFormat.full,
+        format_: mlrun.common.schemas.ProjectsFormat = mlrun.common.schemas.ProjectsFormat.full,
         labels: typing.List[str] = None,
-        state: mlrun.api.schemas.ProjectState = None,
+        state: mlrun.common.schemas.ProjectState = None,
         names: typing.Optional[typing.List[str]] = None,
-    ) -> mlrun.api.schemas.ProjectsOutput:
+    ) -> mlrun.common.schemas.ProjectsOutput:
         return mlrun.api.utils.singletons.db.get_db().list_projects(
             session, owner, format_, labels, state, names
         )
@@ -201,14 +229,14 @@ class Projects(
         session: sqlalchemy.orm.Session,
         owner: str = None,
         labels: typing.List[str] = None,
-        state: mlrun.api.schemas.ProjectState = None,
+        state: mlrun.common.schemas.ProjectState = None,
         names: typing.Optional[typing.List[str]] = None,
-    ) -> mlrun.api.schemas.ProjectSummariesOutput:
+    ) -> mlrun.common.schemas.ProjectSummariesOutput:
         projects_output = await fastapi.concurrency.run_in_threadpool(
             self.list_projects,
             session,
             owner,
-            mlrun.api.schemas.ProjectsFormat.name_only,
+            mlrun.common.schemas.ProjectsFormat.name_only,
             labels,
             state,
             names,
@@ -216,13 +244,13 @@ class Projects(
         project_summaries = await self.generate_projects_summaries(
             projects_output.projects
         )
-        return mlrun.api.schemas.ProjectSummariesOutput(
+        return mlrun.common.schemas.ProjectSummariesOutput(
             project_summaries=project_summaries
         )
 
     async def get_project_summary(
         self, session: sqlalchemy.orm.Session, name: str
-    ) -> mlrun.api.schemas.ProjectSummary:
+    ) -> mlrun.common.schemas.ProjectSummary:
         # Call get project so we'll explode if project doesn't exists
         await fastapi.concurrency.run_in_threadpool(self.get_project, session, name)
         project_summaries = await self.generate_projects_summaries([name])
@@ -230,7 +258,7 @@ class Projects(
 
     async def generate_projects_summaries(
         self, projects: typing.List[str]
-    ) -> typing.List[mlrun.api.schemas.ProjectSummary]:
+    ) -> typing.List[mlrun.common.schemas.ProjectSummary]:
         (
             project_to_files_count,
             project_to_schedule_count,
@@ -243,7 +271,7 @@ class Projects(
         project_summaries = []
         for project in projects:
             project_summaries.append(
-                mlrun.api.schemas.ProjectSummary(
+                mlrun.common.schemas.ProjectSummary(
                     name=project,
                     files_count=project_to_files_count.get(project, 0),
                     schedules_count=project_to_schedule_count.get(project, 0),
@@ -314,7 +342,7 @@ class Projects(
     @staticmethod
     def _list_pipelines(
         session,
-        format_: mlrun.api.schemas.PipelinesFormat = mlrun.api.schemas.PipelinesFormat.metadata_only,
+        format_: mlrun.common.schemas.PipelinesFormat = mlrun.common.schemas.PipelinesFormat.metadata_only,
     ):
         return mlrun.api.crud.Pipelines().list_pipelines(session, "*", format_=format_)
 

mlrun/api/crud/runs.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,10 +16,10 @@ import typing
 
 import sqlalchemy.orm
 
-import mlrun.api.schemas
 import mlrun.api.utils.projects.remotes.follower
 import mlrun.api.utils.singletons.db
 import mlrun.api.utils.singletons.project_member
+import mlrun.common.schemas
 import mlrun.config
 import mlrun.errors
 import mlrun.lists
@@ -119,13 +119,14 @@ class Runs(
         start_time_to=None,
         last_update_time_from=None,
         last_update_time_to=None,
-        partition_by: mlrun.api.schemas.RunPartitionByField = None,
+        partition_by: mlrun.common.schemas.RunPartitionByField = None,
         rows_per_partition: int = 1,
-        partition_sort_by: mlrun.api.schemas.SortField = None,
-        partition_order: mlrun.api.schemas.OrderType = mlrun.api.schemas.OrderType.desc,
+        partition_sort_by: mlrun.common.schemas.SortField = None,
+        partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
         max_partitions: int = 0,
         requested_logs: bool = None,
         return_as_run_structs: bool = True,
+        with_notifications: bool = False,
     ):
         project = project or mlrun.mlconf.default_project
         return mlrun.api.utils.singletons.db.get_db().list_runs(
@@ -149,6 +150,7 @@ class Runs(
             max_partitions,
             requested_logs,
             return_as_run_structs,
+            with_notifications,
         )
 
     def delete_run(