mlrun 1.5.0rc1__py3-none-any.whl → 1.5.0rc2__py3-none-any.whl

mlrun/api/launcher.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 from typing import Dict, List, Optional, Union
 
+from dependency_injector import containers, providers
+
 import mlrun.api.crud
 import mlrun.common.db.sql_session
 import mlrun.common.schemas.schedule
@@ -26,20 +28,25 @@ import mlrun.runtimes.utils
 import mlrun.utils
 import mlrun.utils.regex
 
-
-def initialize_launcher():
-    """Set the factory custom launcher to the server side launcher"""
-    mlrun.launcher.factory.LauncherFactory().set_launcher(ServerSideLauncher)
+# 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, **kwargs):
+    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
+
     def launch(
         self,
         runtime: mlrun.runtimes.BaseRuntime,
@@ -158,15 +165,21 @@ class ServerSideLauncher(launcher.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
@@ -207,3 +220,9 @@ class ServerSideLauncher(launcher.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

@@ -36,8 +36,8 @@ import mlrun.utils.version
 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.launcher import initialize_launcher
 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,
@@ -53,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"
@@ -141,7 +141,6 @@ async def startup_event():
 
     initialize_logs_dir()
     initialize_db()
-    initialize_launcher()
 
     if (
         config.httpdb.clusterization.worker.sync_with_chief.mode
@@ -618,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/{mlutils → api/rundb}/__init__.py

@@ -1,17 +1,13 @@
-# Copyright 2023 Iguazio
+# 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
+#     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.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-# for backwards compatibility
-from ..utils.helpers import create_class, create_function

mlrun/{db → api/rundb}/sqldb.py

@@ -15,24 +15,27 @@
 import datetime
 from typing import List, Optional, Union
 
+from dependency_injector import containers, providers
+
+import mlrun.api.crud
 import mlrun.common.schemas
+import mlrun.db.factory
 import mlrun.model_monitoring.model_endpoint
 from mlrun.api.db.base import DBError
-from mlrun.api.db.sqldb.db import SQLDB as SQLAPIDB
+from mlrun.api.db.sqldb.db import SQLDB
 from mlrun.common.db.sql_session import create_session
+from mlrun.db import RunDBInterface
 
 # This class is a proxy for the real implementation that sits under mlrun.api.db.sqldb
 # The runtime objects (which manages the resources that do the real logic, like Nuclio functions, Dask jobs, etc...)
-# require a RunDB to manage their state, when a user run them locally this db will either be the
-# local filedb or the remote httpdb (we decided that we don't want to support SQLDB as an optional RunDB).
+# require a RunDB to manage their state, when a user run them locally this db will be a remote httpdb.
 # When the user submits something to run (task, function etc...) this runtime managers actually runs inside the api
 # service, in order to prevent the api from calling itself several times for each submission request (since the runDB
 # will be httpdb to that same api service) we have this class which is kind of a proxy between the RunDB interface to
 # the api service's DB interface
-from .base import RunDBError, RunDBInterface
 
 
-class SQLDB(RunDBInterface):
+class SQLRunDB(RunDBInterface):
     def __init__(
         self,
         dsn,
@@ -45,12 +48,10 @@ class SQLDB(RunDBInterface):
     def connect(self, secrets=None):
         if not self.session:
             self.session = create_session()
-        self.db = SQLAPIDB(self.dsn)
+        self.db = SQLDB(self.dsn)
         return self
 
     def store_log(self, uid, project="", body=b"", append=False):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Logs().store_log,
             body,
@@ -71,8 +72,6 @@ class SQLDB(RunDBInterface):
         )
 
     def store_run(self, struct, uid, project="", iter=0):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Runs().store_run,
             self.session,
@@ -83,8 +82,6 @@ class SQLDB(RunDBInterface):
         )
 
     def update_run(self, updates: dict, uid, project="", iter=0):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Runs().update_run,
             self.session,
@@ -98,8 +95,6 @@ class SQLDB(RunDBInterface):
         raise NotImplementedError()
 
     def read_run(self, uid, project=None, iter=None):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Runs().get_run,
             self.session,
@@ -131,8 +126,6 @@ class SQLDB(RunDBInterface):
         max_partitions: int = 0,
         with_notifications: bool = False,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Runs().list_runs,
             self.session,
@@ -157,8 +150,6 @@ class SQLDB(RunDBInterface):
         )
 
     def del_run(self, uid, project=None, iter=None):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Runs().delete_run,
             self.session,
@@ -168,8 +159,6 @@ class SQLDB(RunDBInterface):
         )
 
     def del_runs(self, name=None, project=None, labels=None, state=None, days_ago=0):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Runs().delete_runs,
             self.session,
@@ -181,8 +170,6 @@ class SQLDB(RunDBInterface):
         )
 
     def store_artifact(self, key, artifact, uid, iter=None, tag="", project=""):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Artifacts().store_artifact,
             self.session,
@@ -195,8 +182,6 @@ class SQLDB(RunDBInterface):
         )
 
     def read_artifact(self, key, tag="", iter=None, project=""):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Artifacts().get_artifact,
             self.session,
@@ -219,8 +204,6 @@ class SQLDB(RunDBInterface):
         kind: str = None,
         category: Union[str, mlrun.common.schemas.ArtifactCategories] = None,
     ):
-        import mlrun.api.crud
-
         if category and isinstance(category, str):
             category = mlrun.common.schemas.ArtifactCategories(category)
 
@@ -240,8 +223,6 @@ class SQLDB(RunDBInterface):
         )
 
     def del_artifact(self, key, tag="", project=""):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Artifacts().delete_artifact,
             self.session,
@@ -251,8 +232,6 @@ class SQLDB(RunDBInterface):
         )
 
     def del_artifacts(self, name="", project="", tag="", labels=None):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Artifacts().delete_artifacts,
             self.session,
@@ -263,8 +242,6 @@ class SQLDB(RunDBInterface):
         )
 
     def store_function(self, function, name, project="", tag="", versioned=False):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Functions().store_function,
             self.session,
@@ -276,8 +253,6 @@ class SQLDB(RunDBInterface):
         )
 
     def get_function(self, name, project="", tag="", hash_key=""):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Functions().get_function,
             self.session,
@@ -288,8 +263,6 @@ class SQLDB(RunDBInterface):
         )
 
     def delete_function(self, name: str, project: str = ""):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Functions().delete_function,
             self.session,
@@ -298,8 +271,6 @@ class SQLDB(RunDBInterface):
         )
 
     def list_functions(self, name=None, project=None, tag=None, labels=None):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Functions().list_functions,
             db_session=self.session,
@@ -325,8 +296,6 @@ class SQLDB(RunDBInterface):
         tag_objects: mlrun.common.schemas.TagObjects,
         replace: bool = False,
     ):
-        import mlrun.api.crud
-
         if replace:
             return self._transform_db_error(
                 mlrun.api.crud.Tags().overwrite_object_tags_with_tag,
@@ -350,8 +319,6 @@ class SQLDB(RunDBInterface):
         tag_name: str,
         tag_objects: mlrun.common.schemas.TagObjects,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Tags().delete_tag_from_objects,
             self.session,
@@ -399,8 +366,6 @@ class SQLDB(RunDBInterface):
         name: str,
         project: mlrun.common.schemas.Project,
     ) -> mlrun.common.schemas.Project:
-        import mlrun.api.crud
-
         if isinstance(project, dict):
             project = mlrun.common.schemas.Project(**project)
 
@@ -417,8 +382,6 @@ class SQLDB(RunDBInterface):
         project: dict,
         patch_mode: mlrun.common.schemas.PatchMode = mlrun.common.schemas.PatchMode.replace,
     ) -> mlrun.common.schemas.Project:
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Projects().patch_project,
             self.session,
@@ -431,8 +394,6 @@ class SQLDB(RunDBInterface):
         self,
         project: mlrun.common.schemas.Project,
     ) -> mlrun.common.schemas.Project:
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Projects().create_project,
             self.session,
@@ -444,8 +405,6 @@ class SQLDB(RunDBInterface):
         name: str,
         deletion_strategy: mlrun.common.schemas.DeletionStrategy = mlrun.common.schemas.DeletionStrategy.default(),
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Projects().delete_project,
             self.session,
@@ -456,8 +415,6 @@ class SQLDB(RunDBInterface):
     def get_project(
         self, name: str = None, project_id: int = None
     ) -> mlrun.common.schemas.Project:
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Projects().get_project,
             self.session,
@@ -471,8 +428,6 @@ class SQLDB(RunDBInterface):
         labels: List[str] = None,
         state: mlrun.common.schemas.ProjectState = None,
     ) -> mlrun.common.schemas.ProjectsOutput:
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.Projects().list_projects,
             self.session,
@@ -487,11 +442,9 @@ class SQLDB(RunDBInterface):
         try:
             return func(*args, **kwargs)
         except DBError as exc:
-            raise RunDBError(exc.args)
+            raise mlrun.db.RunDBError(exc.args)
 
     def create_feature_set(self, feature_set, project="", versioned=True):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().create_feature_set,
             self.session,
@@ -503,8 +456,6 @@ class SQLDB(RunDBInterface):
     def get_feature_set(
         self, name: str, project: str = "", tag: str = None, uid: str = None
     ):
-        import mlrun.api.crud
-
         feature_set = self._transform_db_error(
             mlrun.api.crud.FeatureStore().get_feature_set,
             self.session,
@@ -523,8 +474,6 @@ class SQLDB(RunDBInterface):
         entities: List[str] = None,
         labels: List[str] = None,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_features,
             self.session,
@@ -542,8 +491,6 @@ class SQLDB(RunDBInterface):
         tag: str = None,
         labels: List[str] = None,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_entities,
             self.session,
@@ -567,8 +514,6 @@ class SQLDB(RunDBInterface):
         partition_sort_by: mlrun.common.schemas.SortField = None,
         partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_feature_sets,
             self.session,
@@ -594,8 +539,6 @@ class SQLDB(RunDBInterface):
         uid=None,
         versioned=True,
     ):
-        import mlrun.api.crud
-
         if isinstance(feature_set, dict):
             feature_set = mlrun.common.schemas.FeatureSet(**feature_set)
 
@@ -615,8 +558,6 @@ class SQLDB(RunDBInterface):
     def patch_feature_set(
         self, name, feature_set, project="", tag=None, uid=None, patch_mode="replace"
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().patch_feature_set,
             self.session,
@@ -629,8 +570,6 @@ class SQLDB(RunDBInterface):
         )
 
     def delete_feature_set(self, name, project="", tag=None, uid=None):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().delete_feature_set,
             self.session,
@@ -641,8 +580,6 @@ class SQLDB(RunDBInterface):
         )
 
     def create_feature_vector(self, feature_vector, project="", versioned=True):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().create_feature_vector,
             self.session,
@@ -654,8 +591,6 @@ class SQLDB(RunDBInterface):
     def get_feature_vector(
         self, name: str, project: str = "", tag: str = None, uid: str = None
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().get_feature_vector,
             self.session,
@@ -677,8 +612,6 @@ class SQLDB(RunDBInterface):
         partition_sort_by: mlrun.common.schemas.SortField = None,
         partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_feature_vectors,
             self.session,
@@ -702,8 +635,6 @@ class SQLDB(RunDBInterface):
         uid=None,
         versioned=True,
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().store_feature_vector,
             self.session,
@@ -724,8 +655,6 @@ class SQLDB(RunDBInterface):
         uid=None,
         patch_mode="replace",
     ):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().patch_feature_vector,
             self.session,
@@ -738,8 +667,6 @@ class SQLDB(RunDBInterface):
         )
 
     def delete_feature_vector(self, name, project="", tag=None, uid=None):
-        import mlrun.api.crud
-
         return self._transform_db_error(
             mlrun.api.crud.FeatureStore().delete_feature_vector,
             self.session,
@@ -749,6 +676,22 @@ class SQLDB(RunDBInterface):
             uid,
         )
 
+    def store_run_notifications(
+        self,
+        notification_objects: List[mlrun.model.Notification],
+        run_uid: str,
+        project: str = None,
+        mask_params: bool = True,
+    ):
+        return self._transform_db_error(
+            mlrun.api.crud.Notifications().store_run_notifications,
+            self.session,
+            notification_objects,
+            run_uid,
+            project,
+            mask_params,
+        )
+
     def list_pipelines(
         self,
         project: str,
@@ -909,3 +852,12 @@ class SQLDB(RunDBInterface):
         # on server side authorization is done in endpoint anyway, so for server side we can "pass" on check
         # done from ingest()
         pass
+
+    def watch_log(self, uid, project="", watch=True, offset=0):
+        raise NotImplementedError("Watching logs is not supported on the server")
+
+
+# Once this file is imported it will override the default RunDB implementation (RunDBContainer)
+@containers.override(mlrun.db.factory.RunDBContainer)
+class SQLRunDBContainer(containers.DeclarativeContainer):
+    run_db = providers.Factory(SQLRunDB)

mlrun/api/runtime_handlers/__init__.py

@@ -0,0 +1,56 @@
+# 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.
+#
+from mlrun.api.runtime_handlers.base import BaseRuntimeHandler
+from mlrun.api.runtime_handlers.daskjob import DaskRuntimeHandler
+from mlrun.api.runtime_handlers.kubejob import KubeRuntimeHandler
+from mlrun.api.runtime_handlers.mpijob import (
+    MpiV1Alpha1RuntimeHandler,
+    MpiV1RuntimeHandler,
+)
+from mlrun.api.runtime_handlers.remotesparkjob import RemoteSparkRuntimeHandler
+from mlrun.api.runtime_handlers.sparkjob import SparkRuntimeHandler
+from mlrun.runtimes import MPIJobCRDVersions, RuntimeKinds, resolve_mpijob_crd_version
+
+runtime_handler_instances_cache = {}
+
+
+def get_runtime_handler(kind: str) -> BaseRuntimeHandler:
+    global runtime_handler_instances_cache
+    if kind == RuntimeKinds.mpijob:
+        # TODO: split resolve_mpijob_crd_version to client and server side
+        mpijob_crd_version = resolve_mpijob_crd_version()
+        crd_version_to_runtime_handler_class = {
+            MPIJobCRDVersions.v1alpha1: MpiV1Alpha1RuntimeHandler,
+            MPIJobCRDVersions.v1: MpiV1RuntimeHandler,
+        }
+        runtime_handler_class = crd_version_to_runtime_handler_class[mpijob_crd_version]
+        if not runtime_handler_instances_cache.setdefault(RuntimeKinds.mpijob, {}).get(
+            mpijob_crd_version
+        ):
+            runtime_handler_instances_cache[RuntimeKinds.mpijob][
+                mpijob_crd_version
+            ] = runtime_handler_class()
+        return runtime_handler_instances_cache[RuntimeKinds.mpijob][mpijob_crd_version]
+
+    kind_runtime_handler_map = {
+        RuntimeKinds.dask: DaskRuntimeHandler,
+        RuntimeKinds.spark: SparkRuntimeHandler,
+        RuntimeKinds.remotespark: RemoteSparkRuntimeHandler,
+        RuntimeKinds.job: KubeRuntimeHandler,
+    }
+    runtime_handler_class = kind_runtime_handler_map[kind]
+    if not runtime_handler_instances_cache.get(kind):
+        runtime_handler_instances_cache[kind] = runtime_handler_class()
+    return runtime_handler_instances_cache[kind]