adding DB field dag to comp_runs

Author: matusdrobuliak66

packages/postgres-database/src/simcore_postgres_database/migration/versions/27b5ea128a2a_add_comp_run_snapshot_tasks_table_3.py

@@ -0,0 +1,30 @@
+"""add comp_run_snapshot_tasks table 3
+
+Revision ID: 27b5ea128a2a
+Revises: 47faea7f09c2
+Create Date: 2025-05-27 13:30:27.592568+00:00
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "27b5ea128a2a"
+down_revision = "47faea7f09c2"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "comp_runs", sa.Column("dag_adjacency_list", sa.JSON(), nullable=True)
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("comp_runs", "dag_adjacency_list")
+    # ### end Alembic commands ###

packages/postgres-database/src/simcore_postgres_database/models/comp_runs.py

@@ -98,6 +98,9 @@
         nullable=False,
         doc="the run uses on demand clusters",
     ),
+    sa.Column(
+        "dag_adjacency_list", sa.JSON, doc="Adjancey list for the pipeline's graph"
+    ),
     sa.UniqueConstraint("project_uuid", "user_id", "iteration"),
     sa.Index("ix_comp_runs_user_id", "user_id"),
 )

services/director-v2/src/simcore_service_director_v2/models/comp_runs.py

@@ -56,6 +56,7 @@ class CompRunsAtDB(BaseModel):
     use_on_demand_clusters: bool
     scheduled: datetime.datetime | None
     processed: datetime.datetime | None
+    dag_adjacency_list: dict[str, list[str]]
 
     @field_validator("result", mode="before")
     @classmethod
@@ -102,6 +103,7 @@ def convert_null_to_empty_metadata(cls, v):
                     "use_on_demand_clusters": False,
                     "scheduled": None,
                     "processed": None,
+                    "dag_adjacency_list": {},
                 },
                 {
                     "run_id": 432,
@@ -117,6 +119,7 @@ def convert_null_to_empty_metadata(cls, v):
                     "use_on_demand_clusters": False,
                     "scheduled": None,
                     "processed": None,
+                    "dag_adjacency_list": {},
                 },
                 {
                     "run_id": 43243,
@@ -139,6 +142,7 @@ def convert_null_to_empty_metadata(cls, v):
                     "use_on_demand_clusters": False,
                     "scheduled": None,
                     "processed": None,
+                    "dag_adjacency_list": {},
                 },
                 {
                     "run_id": 43243,
@@ -155,6 +159,7 @@ def convert_null_to_empty_metadata(cls, v):
                     "use_on_demand_clusters": False,
                     "scheduled": None,
                     "processed": None,
+                    "dag_adjacency_list": {},
                 },
             ]
         },

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py

@@ -1,7 +1,6 @@
 import logging
 from typing import Final
 
-import networkx as nx
 from fastapi import FastAPI
 from models_library.projects import ProjectID
 from models_library.users import UserID
@@ -11,6 +10,7 @@
 from servicelib.logging_utils import log_context
 from servicelib.redis import CouldNotAcquireLockError, exclusive
 from servicelib.utils import limited_gather
+from simcore_service_director_v2.models.comp_pipelines import CompPipelineAtDB
 from simcore_service_director_v2.modules.db.repositories.comp_runs_snapshot_tasks import (
     CompRunsSnapshotTasksRepository,
 )
@@ -46,7 +46,8 @@ async def run_new_pipeline(
     """Sets a new pipeline to be scheduled on the computational resources."""
     # ensure the pipeline exists and is populated with something
     db_engine = get_db_engine(app)
-    dag = await _get_pipeline_dag(project_id, db_engine)
+    comp_pipeline_at_db = await _get_pipeline_at_db(project_id, db_engine)
+    dag = comp_pipeline_at_db.get_graph()
     if not dag:
         _logger.warning(
             "project %s has no computational dag defined. not scheduled for a run.",
@@ -59,6 +60,7 @@ async def run_new_pipeline(
         project_id=project_id,
         metadata=run_metadata,
         use_on_demand_clusters=use_on_demand_clusters,
+        dag_adjacency_list=comp_pipeline_at_db.dag_adjacency_list,
     )
 
     db_create_snaphot_tasks = [
@@ -120,12 +122,12 @@ async def stop_pipeline(
         )
 
 
-async def _get_pipeline_dag(
+async def _get_pipeline_at_db(
     project_id: ProjectID, db_engine: AsyncEngine
-) -> nx.DiGraph:
+) -> CompPipelineAtDB:
     comp_pipeline_repo = CompPipelinesRepository.instance(db_engine)
     pipeline_at_db = await comp_pipeline_repo.get_pipeline(project_id)
-    return pipeline_at_db.get_graph()
+    return pipeline_at_db
 
 
 _LOST_TASKS_FACTOR: Final[int] = 10

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

@@ -358,6 +358,7 @@ async def create(
         iteration: PositiveInt | None = None,
         metadata: RunMetadataDict,
         use_on_demand_clusters: bool,
+        dag_adjacency_list: dict[str, list[str]],
     ) -> CompRunsAtDB:
         try:
             async with transaction_context(self.db_engine) as conn:
@@ -373,6 +374,7 @@ async def create(
                         result=RUNNING_STATE_TO_DB[RunningState.PUBLISHED],
                         metadata=jsonable_encoder(metadata),
                         use_on_demand_clusters=use_on_demand_clusters,
+                        dag_adjacency_list=dag_adjacency_list,
                     )
                     .returning(literal_column("*"))
                 )

services/director-v2/tests/unit/with_dbs/comp_scheduler/test_db_repositories_comp_runs.py

@@ -93,6 +93,7 @@ async def test_list(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     assert await CompRunsRepository(sqlalchemy_async_engine).list_() == [created]
 
@@ -104,6 +105,7 @@ async def test_list(
                 iteration=created.iteration + n + 1,
                 metadata=run_metadata,
                 use_on_demand_clusters=faker.pybool(),
+                dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
             )
             for n in range(50)
         )
@@ -268,6 +270,7 @@ async def test_create(
             iteration=None,
             metadata=run_metadata,
             use_on_demand_clusters=faker.pybool(),
+            dag_adjacency_list={},
         )
     published_project = await publish_project()
     with pytest.raises(UserNotFoundError):
@@ -277,6 +280,7 @@ async def test_create(
             iteration=None,
             metadata=run_metadata,
             use_on_demand_clusters=faker.pybool(),
+            dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
         )
 
     created = await CompRunsRepository(sqlalchemy_async_engine).create(
@@ -285,6 +289,7 @@ async def test_create(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     got = await CompRunsRepository(sqlalchemy_async_engine).get(
         user_id=published_project.user["id"],
@@ -299,6 +304,7 @@ async def test_create(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     assert created != got
     assert created.iteration == got.iteration + 1
@@ -332,6 +338,7 @@ async def test_update(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
 
     got = await CompRunsRepository(sqlalchemy_async_engine).get(
@@ -365,6 +372,7 @@ async def test_set_run_result(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     got = await CompRunsRepository(sqlalchemy_async_engine).get(
         user_id=published_project.user["id"],
@@ -412,6 +420,7 @@ async def test_mark_for_cancellation(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     got = await CompRunsRepository(sqlalchemy_async_engine).get(
         user_id=published_project.user["id"],
@@ -443,6 +452,7 @@ async def test_mark_for_scheduling(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     got = await CompRunsRepository(sqlalchemy_async_engine).get(
         user_id=published_project.user["id"],
@@ -476,6 +486,7 @@ async def test_mark_scheduling_done(
         iteration=None,
         metadata=run_metadata,
         use_on_demand_clusters=faker.pybool(),
+        dag_adjacency_list=published_project.pipeline.dag_adjacency_list,
     )
     got = await CompRunsRepository(sqlalchemy_async_engine).get(
         user_id=published_project.user["id"],

services/director-v2/tests/unit/with_dbs/conftest.py

@@ -196,6 +196,7 @@ async def _(
             "result": StateType.NOT_STARTED,
             "metadata": jsonable_encoder(run_metadata),
             "use_on_demand_clusters": False,
+            "dag_adjacency_list": {},
         }
         run_config.update(**run_kwargs)
         async with sqlalchemy_async_engine.begin() as conn: