[Projects] Make store_project async and refresh cached result (mlrun#…

…5434)
alonmr · Apr 24, 2024 · 2c92579 · 2c92579
1 parent 0e8b192
commit 2c92579
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 3 deletions.
diff --git a/server/api/api/endpoints/projects.py b/server/api/api/endpoints/projects.py
@@ -72,7 +72,7 @@ def create_project(
         http.HTTPStatus.ACCEPTED.value: {},
     },
 )
-def store_project(
+async def store_project(
     project: mlrun.common.schemas.Project,
     name: str,
     # TODO: we're in a http request context here, therefore it doesn't make sense that by default it will hold the
@@ -85,7 +85,8 @@ def store_project(
         server.api.api.deps.get_db_session
     ),
 ):
-    project, is_running_in_background = get_project_member().store_project(
+    project, is_running_in_background = await run_in_threadpool(
+        get_project_member().store_project,
         db_session,
         name,
         project,

diff --git a/server/api/db/sqldb/db.py b/server/api/db/sqldb/db.py
@@ -2377,6 +2377,9 @@ def _update_project_record_from_project(
         project_record.source = project.spec.source
         project_record.owner = project.spec.owner
         project_record.state = project.status.state
+        project_record.default_function_node_selector = (
+            project.spec.default_function_node_selector
+        )
         labels = project.metadata.labels or {}
         update_labels(project_record, labels)
         self._upsert(session, [project_record])
@@ -3904,6 +3907,7 @@ def _transform_project_record_to_schema(
                 spec=mlrun.common.schemas.ProjectSpec(
                     description=project_record.description,
                     source=project_record.source,
+                    default_function_node_selector=project_record.default_function_node_selector,
                 ),
                 status=mlrun.common.schemas.ObjectStatus(
                     state=project_record.state,

diff --git a/server/api/utils/clients/iguazio.py b/server/api/utils/clients/iguazio.py
@@ -755,6 +755,13 @@ def _transform_mlrun_project_to_iguazio_project(
             )
         if project.spec.owner:
             body["data"]["attributes"]["owner_username"] = project.spec.owner
+
+        if project.spec.default_function_node_selector is not None:
+            body["data"]["attributes"]["default_function_node_selector"] = (
+                Client._transform_mlrun_labels_to_iguazio_labels(
+                    project.spec.default_function_node_selector
+                )
+            )
         return body
 
     @staticmethod

diff --git a/server/api/utils/projects/follower.py b/server/api/utils/projects/follower.py
@@ -150,7 +150,9 @@ def store_project(
                 )
             else:
                 self._leader_client.update_project(leader_session, name, project)
-                return self.get_project(db_session, name, leader_session), False
+                return server.api.db.session.run_function_with_new_db_session(
+                    self.get_project, name, leader_session
+                ), False
 
     def patch_project(
         self,

diff --git a/tests/api/utils/clients/test_iguazio.py b/tests/api/utils/clients/test_iguazio.py
@@ -1088,11 +1088,14 @@ def _generate_project(
     annotations=None,
     created=None,
     owner="project-owner",
+    default_function_node_selector=None,
 ) -> mlrun.common.schemas.Project:
     if labels is None:
         labels = {
             "some-label": "some-label-value",
         }
+    if default_function_node_selector is None:
+        default_function_node_selector = {"zone": "us-west-1"}
     if annotations is None:
         annotations = {
             "some-annotation": "some-annotation-value",
@@ -1110,6 +1113,7 @@ def _generate_project(
             desired_state=mlrun.common.schemas.ProjectState.online,
             owner=owner,
             some_extra_field="some value",
+            default_function_node_selector=default_function_node_selector,
         ),
         status=mlrun.common.schemas.ProjectStatus(
             some_extra_field="some value",
@@ -1135,6 +1139,7 @@ def _build_project_response(
             "updated_at": datetime.datetime.utcnow().isoformat(),
             "admin_status": project.spec.desired_state
             or mlrun.common.schemas.ProjectState.online,
+            "default_function_node_selector": [],
         },
     }
     if with_mlrun_project:
@@ -1155,6 +1160,12 @@ def _build_project_response(
                 project.metadata.labels
             )
         )
+    if project.spec.default_function_node_selector:
+        body["attributes"]["default_function_node_selector"] = (
+            iguazio_client._transform_mlrun_labels_to_iguazio_labels(
+                project.spec.default_function_node_selector
+            )
+        )
     if project.metadata.annotations:
         body["attributes"]["annotations"] = (
             iguazio_client._transform_mlrun_labels_to_iguazio_labels(