Improved pipeline/step run fetching (#3776)

* Pass project ID when fetching steps of a run

* Improved loading of step runs

* Improved loading of pipeline runs

* Linting

* Viewonly relationship

* Optimize DAG DB queries

* Linting

Author: schustmi

src/zenml/models/v2/core/pipeline_run.py

@@ -480,6 +480,7 @@ def steps(self) -> Dict[str, "StepRunResponse"]:
             for step in pagination_utils.depaginate(
                 Client().list_run_steps,
                 pipeline_run_id=self.id,
+                project=self.project_id,
             )
         }
 

src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -20,7 +20,7 @@
 
 from pydantic import ConfigDict
 from sqlalchemy import UniqueConstraint
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import TEXT, Column, Field, Relationship
 
@@ -259,19 +259,19 @@ def get_query_options(
         from zenml.zen_stores.schemas import ModelVersionSchema
 
         options = [
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.pipeline)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.stack)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.build)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.schedule)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.code_reference)
             ),
         ]
@@ -286,14 +286,14 @@ def get_query_options(
         if include_resources:
             options.extend(
                 [
-                    joinedload(
+                    selectinload(
                         jl_arg(PipelineRunSchema.model_version)
                     ).joinedload(
                         jl_arg(ModelVersionSchema.model), innerjoin=True
                     ),
-                    joinedload(jl_arg(PipelineRunSchema.logs)),
-                    joinedload(jl_arg(PipelineRunSchema.user)),
-                    # joinedload(jl_arg(PipelineRunSchema.tags)),
+                    selectinload(jl_arg(PipelineRunSchema.logs)),
+                    selectinload(jl_arg(PipelineRunSchema.user)),
+                    selectinload(jl_arg(PipelineRunSchema.tags)),
                 ]
             )
 

src/zenml/zen_stores/schemas/step_run_schemas.py

@@ -21,7 +21,7 @@
 from pydantic import ConfigDict
 from sqlalchemy import TEXT, Column, String, UniqueConstraint
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship, SQLModel
 
@@ -50,6 +50,7 @@
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
     PipelineDeploymentSchema,
+    StepConfigurationSchema,
 )
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 from zenml.zen_stores.schemas.project_schemas import ProjectSchema
@@ -187,6 +188,14 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
     original_step_run: Optional["StepRunSchema"] = Relationship(
         sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
     )
+    step_configuration_schema: Optional["StepConfigurationSchema"] = (
+        Relationship(
+            sa_relationship_kwargs=dict(
+                viewonly=True,
+                primaryjoin="and_(foreign(StepConfigurationSchema.name) == StepRunSchema.name, foreign(StepConfigurationSchema.deployment_id) == StepRunSchema.deployment_id)",
+            ),
+        )
+    )
 
     model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
@@ -209,17 +218,25 @@ def get_query_options(
         Returns:
             A list of query options.
         """
-        from zenml.zen_stores.schemas import ModelVersionSchema
+        from zenml.zen_stores.schemas import (
+            ArtifactVersionSchema,
+            ModelVersionSchema,
+        )
 
         options = [
-            joinedload(jl_arg(StepRunSchema.deployment)),
-            joinedload(jl_arg(StepRunSchema.pipeline_run)),
+            selectinload(jl_arg(StepRunSchema.deployment)).load_only(
+                jl_arg(PipelineDeploymentSchema.pipeline_configuration)
+            ),
+            selectinload(jl_arg(StepRunSchema.pipeline_run)).load_only(
+                jl_arg(PipelineRunSchema.start_time)
+            ),
+            joinedload(jl_arg(StepRunSchema.step_configuration_schema)),
         ]
 
         if include_metadata:
             options.extend(
                 [
-                    joinedload(jl_arg(StepRunSchema.logs)),
+                    selectinload(jl_arg(StepRunSchema.logs)),
                     # joinedload(jl_arg(StepRunSchema.parents)),
                     # joinedload(jl_arg(StepRunSchema.run_metadata)),
                 ]
@@ -228,12 +245,28 @@ def get_query_options(
         if include_resources:
             options.extend(
                 [
-                    joinedload(jl_arg(StepRunSchema.model_version)).joinedload(
+                    selectinload(
+                        jl_arg(StepRunSchema.model_version)
+                    ).joinedload(
                         jl_arg(ModelVersionSchema.model), innerjoin=True
                     ),
-                    joinedload(jl_arg(StepRunSchema.user)),
-                    # joinedload(jl_arg(StepRunSchema.input_artifacts)),
-                    # joinedload(jl_arg(StepRunSchema.output_artifacts)),
+                    selectinload(jl_arg(StepRunSchema.user)),
+                    selectinload(jl_arg(StepRunSchema.input_artifacts))
+                    .joinedload(
+                        jl_arg(StepRunInputArtifactSchema.artifact_version),
+                        innerjoin=True,
+                    )
+                    .joinedload(
+                        jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
+                    ),
+                    selectinload(jl_arg(StepRunSchema.output_artifacts))
+                    .joinedload(
+                        jl_arg(StepRunOutputArtifactSchema.artifact_version),
+                        innerjoin=True,
+                    )
+                    .joinedload(
+                        jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
+                    ),
                 ]
             )
 
@@ -290,10 +323,7 @@ def to_model(
         """
         step = None
         if self.deployment is not None:
-            step_configurations = self.deployment.get_step_configurations(
-                include=[self.name]
-            )
-            if step_configurations:
+            if self.step_configuration_schema:
                 pipeline_configuration = (
                     PipelineConfiguration.model_validate_json(
                         self.deployment.pipeline_configuration
@@ -304,7 +334,7 @@ def to_model(
                     inplace=True,
                 )
                 step = Step.from_dict(
-                    json.loads(step_configurations[0].config),
+                    json.loads(self.step_configuration_schema.config),
                     pipeline_configuration=pipeline_configuration,
                 )
         if not step and self.step_configuration:

src/zenml/zen_stores/sql_zen_store.py

@@ -64,7 +64,7 @@
     ArgumentError,
     IntegrityError,
 )
-from sqlalchemy.orm import Mapped, joinedload, noload
+from sqlalchemy.orm import Mapped, noload, selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlalchemy.util import immutabledict
 from sqlmodel import Session as SqlModelSession
@@ -5332,13 +5332,17 @@ def get_pipeline_run_dag(self, pipeline_run_id: UUID) -> PipelineRunDAG:
                 schema_class=PipelineRunSchema,
                 session=session,
                 query_options=[
-                    joinedload(jl_arg(PipelineRunSchema.deployment)),
-                    # joinedload(jl_arg(PipelineRunSchema.step_runs)).sele(
-                    #     jl_arg(StepRunSchema.input_artifacts)
-                    # ),
-                    # joinedload(jl_arg(PipelineRunSchema.step_runs)).joinedload(
-                    #     jl_arg(StepRunSchema.output_artifacts)
-                    # ),
+                    selectinload(
+                        jl_arg(PipelineRunSchema.deployment)
+                    ).load_only(
+                        jl_arg(PipelineDeploymentSchema.pipeline_configuration)
+                    ),
+                    selectinload(
+                        jl_arg(PipelineRunSchema.step_runs)
+                    ).selectinload(jl_arg(StepRunSchema.input_artifacts)),
+                    selectinload(
+                        jl_arg(PipelineRunSchema.step_runs)
+                    ).selectinload(jl_arg(StepRunSchema.output_artifacts)),
                 ],
             )
             assert run.deployment is not None