mlrun 1.5.0rc4__py3-none-any.whl → 1.5.0rc6__py3-none-any.whl

mlrun/api/initial_data.py

@@ -192,7 +192,7 @@ def _perform_schema_migrations(alembic_util: mlrun.api.utils.db.alembic.AlembicU
 
 def _is_latest_data_version():
     db_session = create_session()
-    db = mlrun.api.db.sqldb.db.SQLDB("")
+    db = mlrun.api.db.sqldb.db.SQLDB()
 
     try:
         current_data_version = _resolve_current_data_version(db, db_session)
@@ -214,8 +214,7 @@ def _perform_database_migration(
 
 def _perform_data_migrations(db_session: sqlalchemy.orm.Session):
     if config.httpdb.db.data_migrations_mode == "enabled":
-        # FileDB is not really a thing anymore, so using SQLDB directly
-        db = mlrun.api.db.sqldb.db.SQLDB("")
+        db = mlrun.api.db.sqldb.db.SQLDB()
         current_data_version = int(db.get_current_data_version(db_session))
         if current_data_version != latest_data_version:
             logger.info(
@@ -235,8 +234,7 @@ def _perform_data_migrations(db_session: sqlalchemy.orm.Session):
 
 
 def _add_initial_data(db_session: sqlalchemy.orm.Session):
-    # FileDB is not really a thing anymore, so using SQLDB directly
-    db = mlrun.api.db.sqldb.db.SQLDB("")
+    db = mlrun.api.db.sqldb.db.SQLDB()
     _add_default_hub_source_if_needed(db, db_session)
     _add_data_version(db, db_session)
 

mlrun/api/launcher.py

@@ -203,8 +203,9 @@ class ServerSideLauncher(launcher.BaseLauncher):
         # in normal use cases if no project is found we will get an error
         if project:
             project = mlrun.projects.project.MlrunProject.from_dict(project.dict())
+            # there is no need to auto mount here as it was already done in the full spec enrichment with the auth info
             mlrun.projects.pipelines.enrich_function_object(
-                project, runtime, copy_function=False
+                project, runtime, copy_function=False, try_auto_mount=False
             )
 
     def _enrich_full_spec(

mlrun/api/migrations_mysql/versions/026c947c4487_altering_table_datastore_profiles_2.py

@@ -0,0 +1,46 @@
+# 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.
+#
+
+"""altering table datastore_profiles 2
+
+Revision ID: 026c947c4487
+Revises: b1d1e7ab5dec
+Create Date: 2023-08-10 14:15:30.523729
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "026c947c4487"
+down_revision = "b1d1e7ab5dec"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("datastore_profiles", sa.Column("object", sa.JSON(), nullable=True))
+    op.drop_column("datastore_profiles", "body")
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "datastore_profiles", sa.Column("body", sa.String(length=1024), nullable=True)
+    )
+    op.drop_column("datastore_profiles", "object")
+    # ### end Alembic commands ###

mlrun/api/migrations_sqlite/versions/026c947c4487_altering_table_datastore_profiles_2.py

@@ -0,0 +1,46 @@
+# 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.
+#
+
+"""altering table datastore_profiles 2
+
+Revision ID: 026c947c4487
+Revises: 6e0c9531edc7
+Create Date: 2023-08-10 14:15:30.523729
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "026c947c4487"
+down_revision = "6e0c9531edc7"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("datastore_profiles") as batch_op:
+        batch_op.add_column(sa.Column("object", sa.JSON(), nullable=True))
+        batch_op.drop_column("body")
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("datastore_profiles") as batch_op:
+        batch_op.add_column(sa.Column("body", sa.String(length=1024), nullable=True))
+        batch_op.drop_column("object")
+    # ### end Alembic commands ###

mlrun/api/rundb/sqldb.py

@@ -16,8 +16,10 @@ import datetime
 from typing import List, Optional, Union
 
 from dependency_injector import containers, providers
+from sqlalchemy.exc import SQLAlchemyError
 
 import mlrun.api.crud
+import mlrun.api.db.session
 import mlrun.common.schemas
 import mlrun.db.factory
 import mlrun.model_monitoring.model_endpoint
@@ -39,11 +41,15 @@ class SQLRunDB(RunDBInterface):
     def __init__(
         self,
         dsn,
+        session=None,
     ):
+        self.session = session
         self.dsn = dsn
         self.db = None
 
     def connect(self, secrets=None):
+        if not self.session:
+            self.session = create_session()
         self.db = SQLDB(self.dsn)
         return self
 
@@ -68,8 +74,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def store_run(self, struct, uid, project="", iter=0):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Runs().store_run,
+            self.session,
             struct,
             uid,
             iter,
@@ -77,8 +84,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def update_run(self, updates: dict, uid, project="", iter=0):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Runs().update_run,
+            self.session,
             project,
             uid,
             iter,
@@ -89,8 +97,9 @@ class SQLRunDB(RunDBInterface):
         raise NotImplementedError()
 
     def read_run(self, uid, project=None, iter=None):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Runs().get_run,
+            self.session,
             uid,
             iter,
             project,
@@ -119,8 +128,9 @@ class SQLRunDB(RunDBInterface):
         max_partitions: int = 0,
         with_notifications: bool = False,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Runs().list_runs,
+            db_session=self.session,
             name=name,
             uid=uid,
             project=project,
@@ -142,16 +152,18 @@ class SQLRunDB(RunDBInterface):
         )
 
     def del_run(self, uid, project=None, iter=None):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Runs().delete_run,
+            self.session,
             uid,
             iter,
             project,
         )
 
     def del_runs(self, name=None, project=None, labels=None, state=None, days_ago=0):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Runs().delete_runs,
+            self.session,
             name,
             project,
             labels,
@@ -160,8 +172,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def store_artifact(self, key, artifact, uid, iter=None, tag="", project=""):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Artifacts().store_artifact,
+            self.session,
             key,
             artifact,
             uid,
@@ -171,8 +184,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def read_artifact(self, key, tag="", iter=None, project=""):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Artifacts().get_artifact,
+            self.session,
             key,
             tag,
             iter,
@@ -195,8 +209,9 @@ class SQLRunDB(RunDBInterface):
         if category and isinstance(category, str):
             category = mlrun.common.schemas.ArtifactCategories(category)
 
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Artifacts().list_artifacts,
+            self.session,
             project,
             name,
             tag,
@@ -210,16 +225,18 @@ class SQLRunDB(RunDBInterface):
         )
 
     def del_artifact(self, key, tag="", project=""):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Artifacts().delete_artifact,
+            self.session,
             key,
             tag,
             project,
         )
 
     def del_artifacts(self, name="", project="", tag="", labels=None):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Artifacts().delete_artifacts,
+            self.session,
             project,
             name,
             tag,
@@ -227,8 +244,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def store_function(self, function, name, project="", tag="", versioned=False):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Functions().store_function,
+            self.session,
             function,
             name,
             project,
@@ -237,8 +255,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def get_function(self, name, project="", tag="", hash_key=""):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Functions().get_function,
+            self.session,
             name,
             project,
             tag,
@@ -246,15 +265,17 @@ class SQLRunDB(RunDBInterface):
         )
 
     def delete_function(self, name: str, project: str = ""):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Functions().delete_function,
+            self.session,
             project,
             name,
         )
 
     def list_functions(self, name=None, project=None, tag=None, labels=None):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Functions().list_functions,
+            db_session=self.session,
             project=project,
             name=name,
             tag=tag,
@@ -266,8 +287,8 @@ class SQLRunDB(RunDBInterface):
         project=None,
         category: Union[str, mlrun.common.schemas.ArtifactCategories] = None,
     ):
-        return self._create_session_and_transform_db_error(
-            self.db.list_artifact_tags, project
+        return self._transform_db_error(
+            self.db.list_artifact_tags, self.session, project
         )
 
     def tag_objects(
@@ -278,15 +299,17 @@ class SQLRunDB(RunDBInterface):
         replace: bool = False,
     ):
         if replace:
-            return self._create_session_and_transform_db_error(
+            return self._transform_db_error(
                 mlrun.api.crud.Tags().overwrite_object_tags_with_tag,
+                self.session,
                 project,
                 tag_name,
                 tag_objects,
             )
 
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Tags().append_tag_to_objects,
+            self.session,
             project,
             tag_name,
             tag_objects,
@@ -298,8 +321,9 @@ class SQLRunDB(RunDBInterface):
         tag_name: str,
         tag_objects: mlrun.common.schemas.TagObjects,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Tags().delete_tag_from_objects,
+            self.session,
             project,
             tag_name,
             tag_objects,
@@ -334,10 +358,10 @@ class SQLRunDB(RunDBInterface):
         )
 
     def store_schedule(self, data):
-        return self._create_session_and_transform_db_error(self.db.store_schedule, data)
+        return self._transform_db_error(self.db.store_schedule, self.session, data)
 
     def list_schedules(self):
-        return self._create_session_and_transform_db_error(self.db.list_schedules)
+        return self._transform_db_error(self.db.list_schedules, self.session)
 
     def store_project(
         self,
@@ -347,8 +371,9 @@ class SQLRunDB(RunDBInterface):
         if isinstance(project, dict):
             project = mlrun.common.schemas.Project(**project)
 
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Projects().store_project,
+            self.session,
             name=name,
             project=project,
         )
@@ -359,8 +384,9 @@ class SQLRunDB(RunDBInterface):
         project: dict,
         patch_mode: mlrun.common.schemas.PatchMode = mlrun.common.schemas.PatchMode.replace,
     ) -> mlrun.common.schemas.Project:
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Projects().patch_project,
+            self.session,
             name=name,
             project=project,
             patch_mode=patch_mode,
@@ -370,8 +396,9 @@ class SQLRunDB(RunDBInterface):
         self,
         project: mlrun.common.schemas.Project,
     ) -> mlrun.common.schemas.Project:
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Projects().create_project,
+            self.session,
             project=project,
         )
 
@@ -380,8 +407,9 @@ class SQLRunDB(RunDBInterface):
         name: str,
         deletion_strategy: mlrun.common.schemas.DeletionStrategy = mlrun.common.schemas.DeletionStrategy.default(),
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Projects().delete_project,
+            self.session,
             name=name,
             deletion_strategy=deletion_strategy,
         )
@@ -389,8 +417,9 @@ class SQLRunDB(RunDBInterface):
     def get_project(
         self, name: str = None, project_id: int = None
     ) -> mlrun.common.schemas.Project:
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Projects().get_project,
+            self.session,
             name=name,
         )
 
@@ -401,8 +430,9 @@ class SQLRunDB(RunDBInterface):
         labels: List[str] = None,
         state: mlrun.common.schemas.ProjectState = None,
     ) -> mlrun.common.schemas.ProjectsOutput:
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.Projects().list_projects,
+            self.session,
             owner=owner,
             format_=format_,
             labels=labels,
@@ -410,8 +440,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def create_feature_set(self, feature_set, project="", versioned=True):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().create_feature_set,
+            self.session,
             project,
             feature_set,
             versioned,
@@ -420,8 +451,9 @@ class SQLRunDB(RunDBInterface):
     def get_feature_set(
         self, name: str, project: str = "", tag: str = None, uid: str = None
     ):
-        feature_set = self._create_session_and_transform_db_error(
+        feature_set = self._transform_db_error(
             mlrun.api.crud.FeatureStore().get_feature_set,
+            self.session,
             project,
             name,
             tag,
@@ -437,8 +469,9 @@ class SQLRunDB(RunDBInterface):
         entities: List[str] = None,
         labels: List[str] = None,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_features,
+            self.session,
             project,
             name,
             tag,
@@ -453,8 +486,9 @@ class SQLRunDB(RunDBInterface):
         tag: str = None,
         labels: List[str] = None,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_entities,
+            self.session,
             project,
             name,
             tag,
@@ -475,8 +509,9 @@ class SQLRunDB(RunDBInterface):
         partition_sort_by: mlrun.common.schemas.SortField = None,
         partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_feature_sets,
+            self.session,
             project,
             name,
             tag,
@@ -504,8 +539,9 @@ class SQLRunDB(RunDBInterface):
 
         name = name or feature_set.metadata.name
         project = project or feature_set.metadata.project
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().store_feature_set,
+            self.session,
             project,
             name,
             feature_set,
@@ -517,8 +553,9 @@ class SQLRunDB(RunDBInterface):
     def patch_feature_set(
         self, name, feature_set, project="", tag=None, uid=None, patch_mode="replace"
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().patch_feature_set,
+            self.session,
             project,
             name,
             feature_set,
@@ -528,8 +565,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def delete_feature_set(self, name, project="", tag=None, uid=None):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().delete_feature_set,
+            self.session,
             project,
             name,
             tag,
@@ -537,8 +575,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def create_feature_vector(self, feature_vector, project="", versioned=True):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().create_feature_vector,
+            self.session,
             project,
             feature_vector,
             versioned,
@@ -547,8 +586,9 @@ class SQLRunDB(RunDBInterface):
     def get_feature_vector(
         self, name: str, project: str = "", tag: str = None, uid: str = None
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().get_feature_vector,
+            self.session,
             project,
             name,
             tag,
@@ -567,8 +607,9 @@ class SQLRunDB(RunDBInterface):
         partition_sort_by: mlrun.common.schemas.SortField = None,
         partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().list_feature_vectors,
+            self.session,
             project,
             name,
             tag,
@@ -589,8 +630,9 @@ class SQLRunDB(RunDBInterface):
         uid=None,
         versioned=True,
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().store_feature_vector,
+            self.session,
             project,
             name,
             feature_vector,
@@ -608,8 +650,9 @@ class SQLRunDB(RunDBInterface):
         uid=None,
         patch_mode="replace",
     ):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().patch_feature_vector,
+            self.session,
             project,
             name,
             feature_vector_update,
@@ -619,8 +662,9 @@ class SQLRunDB(RunDBInterface):
         )
 
     def delete_feature_vector(self, name, project="", tag=None, uid=None):
-        return self._create_session_and_transform_db_error(
+        return self._transform_db_error(
             mlrun.api.crud.FeatureStore().delete_feature_vector,
+            self.session,
             project,
             name,
             tag,
@@ -634,7 +678,10 @@ class SQLRunDB(RunDBInterface):
         project: str = None,
         mask_params: bool = True,
     ):
-        return self._create_session_and_transform_db_error(
+        # We run this function with a new session because it may run concurrently.
+        # Older sessions will not be able to see the changes made by this function until they are committed.
+        return self._transform_db_error(
+            mlrun.api.db.session.run_function_with_new_db_session,
             mlrun.api.crud.Notifications().store_run_notifications,
             notification_objects,
             run_uid,
@@ -670,6 +717,20 @@ class SQLRunDB(RunDBInterface):
             function,
         )
 
+    def list_hub_sources(
+        self,
+        item_name: Optional[str] = None,
+        tag: Optional[str] = None,
+        version: Optional[str] = None,
+    ):
+        return self._transform_db_error(
+            mlrun.api.crud.Hub().list_hub_sources,
+            self.session,
+            item_name,
+            tag,
+            version,
+        )
+
     def list_pipelines(
         self,
         project: str,
@@ -795,14 +856,6 @@ class SQLRunDB(RunDBInterface):
     ):
         raise NotImplementedError()
 
-    def list_hub_sources(
-        self,
-        item_name: Optional[str] = None,
-        tag: Optional[str] = None,
-        version: Optional[str] = None,
-    ):
-        raise NotImplementedError()
-
     def get_hub_source(self, source_name: str):
         raise NotImplementedError()
 
@@ -847,7 +900,7 @@ class SQLRunDB(RunDBInterface):
     def delete_datastore_profile(self, name: str, project: str):
         raise NotImplementedError()
 
-    def list_datastore_profile(
+    def list_datastore_profiles(
         self, project: str
     ) -> List[mlrun.common.schemas.DatastoreProfile]:
         raise NotImplementedError()
@@ -857,17 +910,16 @@ class SQLRunDB(RunDBInterface):
     ):
         raise NotImplementedError()
 
-    def _create_session_and_transform_db_error(self, func, *args, **kwargs):
-        session = create_session()
-        try:
-            return self._transform_db_error(func, session, *args, **kwargs)
-        finally:
-            session.close()
-
-    @staticmethod
-    def _transform_db_error(func, *args, **kwargs):
+    def _transform_db_error(self, func, *args, **kwargs):
         try:
             return func(*args, **kwargs)
+
+        except SQLAlchemyError as exc:
+            # If we got a SQLAlchemyError, it means the error was not handled by the SQLDB and we need to rollback
+            # to make the session usable again
+            self.session.rollback()
+            raise mlrun.db.RunDBError(exc.args) from exc
+
         except DBError as exc:
             raise mlrun.db.RunDBError(exc.args) from exc
 

mlrun/api/utils/db/sqlite_migration.py

@@ -57,6 +57,7 @@ class SQLiteMigrationUtil(object):
         "projects_labels",
         "data_versions",
         "background_tasks",
+        "datastore_profiles",
     ]
 
     def __init__(self):

mlrun/common/model_monitoring/helpers.py

@@ -65,5 +65,7 @@ def parse_monitoring_stream_path(stream_uri: str, project: str):
 
     elif stream_uri.startswith("v3io://") and mlrun.mlconf.is_ce_mode():
         # V3IO is not supported in CE mode, generating a default http stream path
-        stream_uri = mlrun.mlconf.model_endpoint_monitoring.default_http_sink
+        stream_uri = mlrun.mlconf.model_endpoint_monitoring.default_http_sink.format(
+            project=project
+        )
     return stream_uri

mlrun/common/schemas/client_spec.py

@@ -57,6 +57,7 @@ class ClientSpec(pydantic.BaseModel):
     redis_type: typing.Optional[str]
     sql_url: typing.Optional[str]
     model_endpoint_monitoring_store_type: typing.Optional[str]
+    model_endpoint_monitoring_endpoint_store_connection: typing.Optional[str]
     # ce_mode is deprecated, we will use the full ce config instead and ce_mode will be removed in 1.6.0
     ce_mode: typing.Optional[str]
     ce: typing.Optional[dict]