fix tests

Author: giancarloromeo

packages/models-library/src/models_library/projects_nodes.py

@@ -165,7 +165,7 @@ class Node(BaseModel):
     ] = DEFAULT_FACTORY
 
     inputs_required: Annotated[
-        list[InputID],
+        list[InputID] | None,
         Field(
             default_factory=list,
             description="Defines inputs that are required in order to run the service",

services/director-v2/src/simcore_service_director_v2/api/routes/computations.py

@@ -27,7 +27,7 @@
     ComputationGet,
     ComputationStop,
 )
-from models_library.projects import ProjectAtDB, ProjectID
+from models_library.projects import NodesDict, ProjectAtDB, ProjectID
 from models_library.projects_nodes_io import NodeID
 from models_library.projects_state import RunningState
 from models_library.services import ServiceKeyVersion
@@ -38,6 +38,9 @@
 from servicelib.logging_utils import log_decorator
 from servicelib.rabbitmq import RabbitMQRPCClient
 from simcore_postgres_database.utils_projects_metadata import DBProjectNotFoundError
+from simcore_service_director_v2.modules.db.repositories.projects_nodes import (
+    ProjectsNodesRepository,
+)
 from starlette import status
 from starlette.requests import Request
 from tenacity import retry
@@ -196,6 +199,7 @@ async def _try_start_pipeline(
     complete_dag: nx.DiGraph,
     minimal_dag: nx.DiGraph,
     project: ProjectAtDB,
+    workbench: NodesDict,
     users_repo: UsersRepository,
     projects_metadata_repo: ProjectsMetadataRepository,
 ) -> None:
@@ -226,7 +230,7 @@ async def _try_start_pipeline(
         run_metadata=RunMetadataDict(
             node_id_names_map={
                 NodeID(node_idstr): node_data.label
-                for node_idstr, node_data in project.workbench.items()
+                for node_idstr, node_data in workbench.items()
             },
             product_name=computation.product_name,
             project_name=project.name,
@@ -273,6 +277,9 @@ async def create_computation(  # noqa: PLR0913 # pylint: disable=too-many-positi
     project_repo: Annotated[
         ProjectsRepository, Depends(get_repository(ProjectsRepository))
     ],
+    project_nodes_repo: Annotated[
+        ProjectsNodesRepository, Depends(get_repository(ProjectsNodesRepository))
+    ],
     comp_pipelines_repo: Annotated[
         CompPipelinesRepository, Depends(get_repository(CompPipelinesRepository))
     ],
@@ -302,8 +309,12 @@ async def create_computation(  # noqa: PLR0913 # pylint: disable=too-many-positi
         # check if current state allow to modify the computation
         await _check_pipeline_not_running_or_raise_409(comp_tasks_repo, computation)
 
+        workbench: NodesDict = await project_nodes_repo.get_nodes(
+            computation.project_id
+        )
+
         # create the complete DAG graph
-        complete_dag = create_complete_dag(project.workbench)
+        complete_dag = create_complete_dag(workbench)
         # find the minimal viable graph to be run
         minimal_computational_dag: nx.DiGraph = (
             await create_minimal_computational_graph_based_on_selection(
@@ -330,6 +341,7 @@ async def create_computation(  # noqa: PLR0913 # pylint: disable=too-many-positi
         ]
         comp_tasks = await comp_tasks_repo.upsert_tasks_from_project(
             project=project,
+            workbench=workbench,
             catalog_client=catalog_client,
             published_nodes=min_computation_nodes if computation.start_pipeline else [],
             user_id=computation.user_id,
@@ -347,6 +359,7 @@ async def create_computation(  # noqa: PLR0913 # pylint: disable=too-many-positi
                 complete_dag=complete_dag,
                 minimal_dag=minimal_computational_dag,
                 project=project,
+                workbench=workbench,
                 users_repo=users_repo,
                 projects_metadata_repo=projects_metadata_repo,
             )

services/director-v2/src/simcore_service_director_v2/modules/db/repositories/comp_tasks/_core.py

@@ -26,6 +26,7 @@
 from ....catalog import CatalogClient
 from ...tables import NodeClass, StateType, comp_tasks
 from .._base import BaseRepository
+from ..projects_nodes import NodesDict
 from . import _utils
 
 _logger = logging.getLogger(__name__)
@@ -89,6 +90,7 @@ async def upsert_tasks_from_project(
         self,
         *,
         project: ProjectAtDB,
+        workbench: NodesDict,
         catalog_client: CatalogClient,
         published_nodes: list[NodeID],
         user_id: UserID,
@@ -103,6 +105,7 @@ async def upsert_tasks_from_project(
                 CompTaskAtDB
             ] = await _utils.generate_tasks_list_from_project(
                 project=project,
+                workbench=workbench,
                 catalog_client=catalog_client,
                 published_nodes=published_nodes,
                 user_id=user_id,
@@ -121,7 +124,7 @@ async def upsert_tasks_from_project(
             # remove the tasks that were removed from project workbench
             if all_nodes := await result.fetchall():
                 node_ids_to_delete = [
-                    t.node_id for t in all_nodes if t.node_id not in project.workbench
+                    t.node_id for t in all_nodes if t.node_id not in workbench
                 ]
                 for node_id in node_ids_to_delete:
                     await conn.execute(

services/director-v2/src/simcore_service_director_v2/modules/db/repositories/comp_tasks/_utils.py

@@ -60,6 +60,7 @@
 from ....catalog import CatalogClient
 from ....comp_scheduler._utils import COMPLETED_STATES
 from ...tables import NodeClass
+from ..projects_nodes import NodesDict
 
 _logger = logging.getLogger(__name__)
 
@@ -331,6 +332,7 @@ def _by_type_name(ec2: EC2InstanceTypeGet) -> bool:
 async def generate_tasks_list_from_project(
     *,
     project: ProjectAtDB,
+    workbench: NodesDict,
     catalog_client: CatalogClient,
     published_nodes: list[NodeID],
     user_id: UserID,
@@ -346,7 +348,7 @@ async def generate_tasks_list_from_project(
         ServiceKeyVersion(
             key=node.key, version=node.version
         )  # the service key version is frozen
-        for node in project.workbench.values()
+        for node in workbench.values()
     }
 
     key_version_to_node_infos = {
@@ -359,8 +361,8 @@ async def generate_tasks_list_from_project(
         for key_version in unique_service_key_versions
     }
 
-    for internal_id, node_id in enumerate(project.workbench, 1):
-        node: Node = project.workbench[node_id]
+    for internal_id, node_id in enumerate(workbench, start=1):
+        node: Node = workbench[node_id]
         node_key_version = ServiceKeyVersion(key=node.key, version=node.version)
         node_details, node_extras, node_labels = key_version_to_node_infos.get(
             node_key_version,

services/director-v2/src/simcore_service_director_v2/modules/db/repositories/projects_nodes.py

@@ -0,0 +1,55 @@
+import logging
+
+import sqlalchemy as sa
+from models_library.projects import NodesDict, ProjectID
+from models_library.projects_nodes import Node
+from simcore_postgres_database.utils_projects_nodes import ProjectNode
+
+from ..tables import projects_nodes
+from ._base import BaseRepository
+
+logger = logging.getLogger(__name__)
+
+
+class ProjectsNodesRepository(BaseRepository):
+    async def get_nodes(self, project_uuid: ProjectID) -> NodesDict:
+        nodes_dict = {}
+        async with self.db_engine.acquire() as conn:
+            rows = await (
+                await conn.execute(
+                    sa.select(projects_nodes).where(
+                        projects_nodes.c.project_uuid == f"{project_uuid}"
+                    )
+                )
+            ).fetchall()
+
+            for row in rows:
+                nodes_dict[f"{row.node_id}"] = Node.model_validate(
+                    ProjectNode.model_validate(row, from_attributes=True).model_dump(
+                        exclude={
+                            "node_id",
+                            "required_resources",
+                            "created",
+                            "modified",
+                        },
+                        exclude_none=True,
+                        exclude_unset=True,
+                    )
+                )
+
+        return nodes_dict
+
+    # async def is_node_present_in_workbench(
+    #     self, project_id: ProjectID, node_uuid: NodeID
+    # ) -> bool:
+    #     try:
+    #         project = await self.get_project(project_id)
+    #         return f"{node_uuid}" in project.workbench
+    #     except ProjectNotFoundError:
+    #         return False
+
+    # async def get_project_id_from_node(self, node_id: NodeID) -> ProjectID:
+    #     async with self.db_engine.acquire() as conn:
+    #         return await ProjectNodesRepo.get_project_id_from_node_id(
+    #             conn, node_id=node_id
+    #         )

services/director-v2/src/simcore_service_director_v2/modules/db/tables.py

@@ -7,6 +7,7 @@
 )
 from simcore_postgres_database.models.projects import ProjectType, projects
 from simcore_postgres_database.models.projects_networks import projects_networks
+from simcore_postgres_database.models.projects_nodes import projects_nodes
 
 __all__ = [
     "comp_pipeline",
@@ -15,6 +16,7 @@
     "groups_extra_properties",
     "NodeClass",
     "projects_networks",
+    "projects_nodes",
     "projects",
     "ProjectType",
     "StateType",

services/director-v2/src/simcore_service_director_v2/utils/dags.py

@@ -42,8 +42,7 @@ def create_complete_dag(workbench: NodesDict) -> nx.DiGraph:
         )
         if node.input_nodes:
             for input_node_id in node.input_nodes:
-                predecessor_node = workbench.get(NodeIDStr(input_node_id))
-                if predecessor_node:
+                if f"{input_node_id}" in workbench:  # predecessor node
                     dag_graph.add_edge(str(input_node_id), node_id)
 
     return dag_graph
@@ -188,7 +187,7 @@ def compute_pipeline_started_timestamp(
     if not pipeline_dag.nodes:
         return None
     node_id_to_comp_task: dict[NodeIDStr, CompTaskAtDB] = {
-        NodeIDStr(f"{task.node_id}"): task for task in comp_tasks
+        f"{task.node_id}": task for task in comp_tasks
     }
     TOMORROW = arrow.utcnow().shift(days=1).datetime
     pipeline_started_at: datetime.datetime | None = min(
@@ -206,7 +205,7 @@ def compute_pipeline_stopped_timestamp(
     if not pipeline_dag.nodes:
         return None
     node_id_to_comp_task: dict[NodeIDStr, CompTaskAtDB] = {
-        NodeIDStr(f"{task.node_id}"): task for task in comp_tasks
+        f"{task.node_id}": task for task in comp_tasks
     }
     TOMORROW = arrow.utcnow().shift(days=1).datetime
     pipeline_stopped_at: datetime.datetime | None = max(
@@ -227,7 +226,7 @@ async def compute_pipeline_details(
 
     # NOTE: the latest progress is available in comp_tasks only
     node_id_to_comp_task: dict[NodeIDStr, CompTaskAtDB] = {
-        NodeIDStr(f"{task.node_id}"): task for task in comp_tasks
+        f"{task.node_id}": task for task in comp_tasks
     }
     pipeline_progress = None
     if len(pipeline_dag.nodes) > 0: