mlrun 1.4.0rc25__py3-none-any.whl → 1.5.0rc2__py3-none-any.whl

mlrun/api/db/sqldb/db.py

@@ -82,7 +82,6 @@ from mlrun.utils import (
 )
 
 NULL = None  # Avoid flake8 issuing warnings when comparing in filter
-run_time_fmt = "%Y-%m-%dT%H:%M:%S.%fZ"
 unversioned_tagged_object_uid_prefix = "unversioned-"
 
 conflict_messages = [
@@ -1175,6 +1174,63 @@ class SQLDB(DBInterface):
 
         return results
 
+    def store_schedule(
+        self,
+        session: Session,
+        project: str,
+        name: str,
+        kind: mlrun.common.schemas.ScheduleKinds = None,
+        scheduled_object: Any = None,
+        cron_trigger: mlrun.common.schemas.ScheduleCronTrigger = None,
+        labels: Dict = None,
+        last_run_uri: str = None,
+        concurrency_limit: int = None,
+        next_run_time: datetime = None,
+    ) -> typing.Optional[mlrun.common.schemas.ScheduleRecord]:
+        schedule = self.get_schedule(
+            session=session, project=project, name=name, raise_on_not_found=False
+        )
+        schedule_exist = schedule is not None
+
+        if not schedule_exist:
+            schedule = Schedule(
+                project=project,
+                name=name,
+                kind=kind.value,
+                creation_time=datetime.now(timezone.utc),
+                concurrency_limit=concurrency_limit,
+                next_run_time=next_run_time,
+                scheduled_object=scheduled_object,
+                cron_trigger=cron_trigger,
+            )
+            labels = labels or {}
+
+        self._update_schedule_body(
+            schedule=schedule,
+            scheduled_object=scheduled_object,
+            cron_trigger=cron_trigger,
+            labels=labels,
+            last_run_uri=last_run_uri,
+            concurrency_limit=concurrency_limit,
+            next_run_time=next_run_time,
+        )
+
+        logger.debug(
+            "Storing schedule to db",
+            project=project,
+            name=name,
+            kind=kind,
+            cron_trigger=cron_trigger,
+            labels=labels,
+            concurrency_limit=concurrency_limit,
+            scheduled_object=scheduled_object,
+        )
+
+        self._upsert(session, [schedule])
+
+        if schedule_exist:
+            return schedule
+
     def create_schedule(
         self,
         session: Session,
@@ -1233,6 +1289,37 @@ class SQLDB(DBInterface):
     ):
         schedule = self._get_schedule_record(session, project, name)
 
+        self._update_schedule_body(
+            schedule=schedule,
+            scheduled_object=scheduled_object,
+            cron_trigger=cron_trigger,
+            labels=labels,
+            last_run_uri=last_run_uri,
+            concurrency_limit=concurrency_limit,
+            next_run_time=next_run_time,
+        )
+
+        logger.debug(
+            "Updating schedule in db",
+            project=project,
+            name=name,
+            cron_trigger=cron_trigger,
+            labels=labels,
+            concurrency_limit=concurrency_limit,
+            next_run_time=next_run_time,
+        )
+        self._upsert(session, [schedule])
+
+    @staticmethod
+    def _update_schedule_body(
+        schedule: mlrun.common.schemas.ScheduleRecord,
+        scheduled_object: Any = None,
+        cron_trigger: mlrun.common.schemas.ScheduleCronTrigger = None,
+        labels: Dict = None,
+        last_run_uri: str = None,
+        concurrency_limit: int = None,
+        next_run_time: datetime = None,
+    ):
         # explicitly ensure the updated fields are not None, as they can be empty strings/dictionaries etc.
         if scheduled_object is not None:
             schedule.scheduled_object = scheduled_object
@@ -1254,17 +1341,6 @@ class SQLDB(DBInterface):
             # saved in the DB in UTC timezone, therefore we transform next_run_time to UTC as well.
             schedule.next_run_time = next_run_time.astimezone(pytz.utc)
 
-        logger.debug(
-            "Updating schedule in db",
-            project=project,
-            name=name,
-            cron_trigger=cron_trigger,
-            labels=labels,
-            concurrency_limit=concurrency_limit,
-            next_run_time=next_run_time,
-        )
-        self._upsert(session, [schedule])
-
     def list_schedules(
         self,
         session: Session,
@@ -1287,19 +1363,23 @@ class SQLDB(DBInterface):
         return schedules
 
     def get_schedule(
-        self, session: Session, project: str, name: str
-    ) -> mlrun.common.schemas.ScheduleRecord:
+        self, session: Session, project: str, name: str, raise_on_not_found: bool = True
+    ) -> typing.Optional[mlrun.common.schemas.ScheduleRecord]:
         logger.debug("Getting schedule from db", project=project, name=name)
-        schedule_record = self._get_schedule_record(session, project, name)
+        schedule_record = self._get_schedule_record(
+            session, project, name, raise_on_not_found
+        )
+        if not schedule_record:
+            return
         schedule = self._transform_schedule_record_to_scheme(schedule_record)
         return schedule
 
     def _get_schedule_record(
-        self, session: Session, project: str, name: str
+        self, session: Session, project: str, name: str, raise_on_not_found: bool = True
     ) -> mlrun.common.schemas.ScheduleRecord:
         query = self._query(session, Schedule, project=project, name=name)
         schedule_record = query.one_or_none()
-        if not schedule_record:
+        if not schedule_record and raise_on_not_found:
             raise mlrun.errors.MLRunNotFoundError(
                 f"Schedule not found: project={project}, name={name}"
             )
@@ -2866,32 +2946,50 @@ class SQLDB(DBInterface):
         def _try_commit_obj():
             try:
                 session.commit()
-            except SQLAlchemyError as err:
+            except SQLAlchemyError as sql_err:
                 session.rollback()
-                cls = objects[0].__class__.__name__
-                if "database is locked" in str(err):
+                classes = list(set([object_.__class__.__name__ for object_ in objects]))
+
+                # if the database is locked, we raise a retryable error
+                if "database is locked" in str(sql_err):
                     logger.warning(
-                        "Database is locked. Retrying", cls=cls, err=str(err)
+                        "Database is locked. Retrying",
+                        classes_to_commit=classes,
+                        err=str(sql_err),
                     )
                     raise mlrun.errors.MLRunRuntimeError(
                         "Failed committing changes, database is locked"
-                    ) from err
+                    ) from sql_err
+
+                # the error is not retryable, so we try to identify weather there was a conflict or not
+                # either way - we wrap the error with a fatal error so the retry mechanism will stop
                 logger.warning(
-                    "Failed committing changes to DB", cls=cls, err=err_to_str(err)
+                    "Failed committing changes to DB",
+                    classes=classes,
+                    err=err_to_str(sql_err),
                 )
                 if not ignore:
+                    # get the identifiers of the objects that failed to commit, for logging purposes
                     identifiers = ",".join(
                         object_.get_identifier_string() for object_ in objects
                     )
-                    # We want to retry only when database is locked so for any other scenario escalate to fatal failure
+
+                    mlrun_error = mlrun.errors.MLRunRuntimeError(
+                        f"Failed committing changes to DB. classes={classes} objects={identifiers}"
+                    )
+
+                    # check if the error is a conflict error
+                    if any([message in str(sql_err) for message in conflict_messages]):
+                        mlrun_error = mlrun.errors.MLRunConflictError(
+                            f"Conflict - at least one of the objects already exists: {identifiers}"
+                        )
+
+                    # we want to keep the exception stack trace, but we also want the retry mechanism to stop
+                    # so, we raise a new indicative exception from the original sql exception (this keeps
+                    # the stack trace intact), and then wrap it with a fatal error (which stops the retry mechanism).
+                    # Note - this way, the exception is raised from this code section, and not from the retry function.
                     try:
-                        if any([message in str(err) for message in conflict_messages]):
-                            raise mlrun.errors.MLRunConflictError(
-                                f"Conflict - {cls} already exists: {identifiers}"
-                            ) from err
-                        raise mlrun.errors.MLRunRuntimeError(
-                            f"Failed committing changes to DB. class={cls} objects={identifiers}"
-                        ) from err
+                        raise mlrun_error from sql_err
                     except (
                         mlrun.errors.MLRunRuntimeError,
                         mlrun.errors.MLRunConflictError,

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

@@ -38,7 +38,6 @@ from mlrun.api.utils.db.sql_collation import SQLCollationUtil
 
 Base = declarative_base()
 NULL = None  # Avoid flake8 issuing warnings when comparing in filter
-run_time_fmt = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 
 def make_label(table):
@@ -53,6 +52,9 @@ def make_label(table):
         value = Column(String(255, collation=SQLCollationUtil.collation()))
         parent = Column(Integer, ForeignKey(f"{table}.id"))
 
+        def get_identifier_string(self) -> str:
+            return f"{self.parent}/{self.name}/{self.value}"
+
     return Label
 
 
@@ -86,6 +88,9 @@ def make_tag_v2(table):
         obj_id = Column(Integer, ForeignKey(f"{table}.id"))
         obj_name = Column(String(255, collation=SQLCollationUtil.collation()))
 
+        def get_identifier_string(self) -> str:
+            return f"{self.project}/{self.name}"
+
     return Tag
 
 
@@ -258,6 +263,9 @@ with warnings.catch_warnings():
         state = Column(String(255, collation=SQLCollationUtil.collation()))
         timeout = Column(Integer)
 
+        def get_identifier_string(self) -> str:
+            return f"{self.project}/{self.name}"
+
     class Schedule(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "schedules_v2"
         __table_args__ = (UniqueConstraint("project", "name", name="_schedules_v2_uc"),)
@@ -317,6 +325,9 @@ with warnings.catch_warnings():
         id = Column(Integer, primary_key=True)
         name = Column(String(255, collation=SQLCollationUtil.collation()))
 
+        def get_identifier_string(self) -> str:
+            return f"{self.name}"
+
     class Project(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "projects"
         # For now since we use project name a lot
@@ -508,6 +519,9 @@ with warnings.catch_warnings():
             default=datetime.now(timezone.utc),
         )
 
+        def get_identifier_string(self) -> str:
+            return f"{self.version}"
+
 
 # Must be after all table definitions
 _tagged = [cls for cls in Base.__subclasses__() if hasattr(cls, "Tag")]

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

@@ -39,7 +39,6 @@ from mlrun.api.utils.db.sql_collation import SQLCollationUtil
 
 Base = declarative_base()
 NULL = None  # Avoid flake8 issuing warnings when comparing in filter
-run_time_fmt = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 
 def make_label(table):
@@ -54,6 +53,9 @@ def make_label(table):
         value = Column(String(255, collation=SQLCollationUtil.collation()))
         parent = Column(Integer, ForeignKey(f"{table}.id"))
 
+        def get_identifier_string(self) -> str:
+            return f"{self.parent}/{self.name}/{self.value}"
+
     return Label
 
 
@@ -90,6 +92,9 @@ def make_tag_v2(table):
             ForeignKey(f"{table}.name"),
         )
 
+        def get_identifier_string(self) -> str:
+            return f"{self.project}/{self.name}"
+
     return Tag
 
 
@@ -242,6 +247,9 @@ with warnings.catch_warnings():
         state = Column(String(255, collation=SQLCollationUtil.collation()))
         timeout = Column(Integer)
 
+        def get_identifier_string(self) -> str:
+            return f"{self.project}/{self.name}"
+
     class Schedule(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "schedules_v2"
         __table_args__ = (UniqueConstraint("project", "name", name="_schedules_v2_uc"),)
@@ -301,6 +309,9 @@ with warnings.catch_warnings():
         id = Column(Integer, primary_key=True)
         name = Column(String(255, collation=SQLCollationUtil.collation()))
 
+        def get_identifier_string(self) -> str:
+            return f"{self.name}"
+
     class Project(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "projects"
         # For now since we use project name a lot
@@ -347,7 +358,7 @@ with warnings.catch_warnings():
         labels = relationship(Label, cascade="all, delete-orphan")
 
         def get_identifier_string(self) -> str:
-            return f"{self.project}/{self.name}"
+            return f"{self.feature_set_id}/{self.name}"
 
     class Entity(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "entities"
@@ -464,6 +475,9 @@ with warnings.catch_warnings():
         version = Column(String(255, collation=SQLCollationUtil.collation()))
         created = Column(TIMESTAMP, default=datetime.now(timezone.utc))
 
+        def get_identifier_string(self) -> str:
+            return f"{self.version}"
+
 
 # Must be after all table definitions
 _tagged = [cls for cls in Base.__subclasses__() if hasattr(cls, "Tag")]

mlrun/api/launcher.py

@@ -13,20 +13,40 @@
 # limitations under the License.
 from typing import Dict, List, Optional, Union
 
+from dependency_injector import containers, providers
+
 import mlrun.api.crud
-import mlrun.api.db.sqldb.session
+import mlrun.common.db.sql_session
 import mlrun.common.schemas.schedule
 import mlrun.config
 import mlrun.execution
-import mlrun.launcher.base
+import mlrun.launcher.base as launcher
+import mlrun.launcher.factory
 import mlrun.runtimes
 import mlrun.runtimes.generators
 import mlrun.runtimes.utils
 import mlrun.utils
 import mlrun.utils.regex
 
+# must be at the bottom to avoid circular import conflicts and can't use 'from' notation because unit tests mock this
+import mlrun.api.api.utils  # isort:skip
+
+
+class ServerSideLauncher(launcher.BaseLauncher):
+    def __init__(
+        self,
+        local: bool = False,
+        auth_info: Optional[mlrun.common.schemas.AuthInfo] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if local:
+            raise mlrun.errors.MLRunInternalServerError(
+                "Launch of local run inside the server is not allowed"
+            )
+
+        self._auth_info = auth_info
 
-class ServerSideLauncher(mlrun.launcher.base.BaseLauncher):
     def launch(
         self,
         runtime: mlrun.runtimes.BaseRuntime,
@@ -145,15 +165,21 @@ class ServerSideLauncher(mlrun.launcher.base.BaseLauncher):
 
         return self._wrap_run_result(runtime, result, run, err=last_err)
 
-    @staticmethod
     def enrich_runtime(
-        runtime: "mlrun.runtimes.base.BaseRuntime", project_name: Optional[str] = ""
+        self,
+        runtime: "mlrun.runtimes.base.BaseRuntime",
+        project_name: Optional[str] = "",
     ):
         """
         Enrich the runtime object with the project spec and metadata.
         This is done only on the server side, since it's the source of truth for the project, and we want to keep the
         client side enrichment as minimal as possible.
         """
+        if self._auth_info:
+            mlrun.api.api.utils.apply_enrichment_and_validation_on_function(
+                runtime, self._auth_info
+            )
+
         # ensure the runtime has a project before we enrich it with the project's spec
         runtime.metadata.project = (
             project_name
@@ -175,7 +201,7 @@ class ServerSideLauncher(mlrun.launcher.base.BaseLauncher):
 
         # If in the api server, we can assume that watch=False, so we save notification
         # configs to the DB, for the run monitor to later pick up and push.
-        session = mlrun.api.db.sqldb.session.create_session()
+        session = mlrun.common.db.sql_session.create_session()
         mlrun.api.crud.Notifications().store_run_notifications(
             session,
             runobj.spec.notifications,
@@ -194,3 +220,9 @@ class ServerSideLauncher(mlrun.launcher.base.BaseLauncher):
                 struct, runtime.metadata.name, runtime.metadata.project, versioned=True
             )
             run.spec.function = runtime._function_uri(hash_key=hash_key)
+
+
+# Once this file is imported it will set the container server side launcher
+@containers.override(mlrun.launcher.factory.LauncherContainer)
+class ServerSideLauncherContainer(containers.DeclarativeContainer):
+    server_side_launcher = providers.Factory(ServerSideLauncher)

mlrun/api/main.py

@@ -37,6 +37,7 @@ from mlrun.api.api.api import api_router
 from mlrun.api.db.session import close_session, create_session
 from mlrun.api.initial_data import init_data
 from mlrun.api.middlewares import init_middlewares
+from mlrun.api.runtime_handlers import get_runtime_handler
 from mlrun.api.utils.periodic import (
     cancel_all_periodic_functions,
     cancel_periodic_function,
@@ -52,7 +53,7 @@ from mlrun.api.utils.singletons.project_member import (
 from mlrun.api.utils.singletons.scheduler import get_scheduler, initialize_scheduler
 from mlrun.config import config
 from mlrun.errors import err_to_str
-from mlrun.runtimes import RuntimeClassMode, RuntimeKinds, get_runtime_handler
+from mlrun.runtimes import RuntimeClassMode, RuntimeKinds
 from mlrun.utils import logger
 
 API_PREFIX = "/api"
@@ -616,7 +617,7 @@ def _push_terminal_run_notifications(db: mlrun.api.db.base.DBInterface, db_sessi
     logger.debug(
         "Got terminal runs with configured notifications", runs_amount=len(runs)
     )
-    mlrun.utils.notifications.NotificationPusher(unmasked_runs).push(db)
+    mlrun.utils.notifications.NotificationPusher(unmasked_runs).push()
 
     _last_notification_push_time = now
 

mlrun/api/rundb/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 MLRun Authors
+#
+# 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.