Add filter option for runs linked to model version (#4003)

* Add filter option for runs linked to model version

* Tests and filter on client method

* Typo

Author: schustmi

src/zenml/client.py

@@ -4168,6 +4168,7 @@ def list_pipeline_runs(
         template_id: Optional[Union[str, UUID]] = None,
         source_snapshot_id: Optional[Union[str, UUID]] = None,
         model_version_id: Optional[Union[str, UUID]] = None,
+        linked_to_model_version_id: Optional[Union[str, UUID]] = None,
         orchestrator_run_id: Optional[str] = None,
         status: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
@@ -4210,6 +4211,11 @@ def list_pipeline_runs(
             template_id: The ID of the template to filter by.
             source_snapshot_id: The ID of the source snapshot to filter by.
             model_version_id: The ID of the model version to filter by.
+            linked_to_model_version_id: Filter by model version linked to the
+                pipeline run. The difference to `model_version_id` is that this
+                filter will not only include pipeline runs which are directly
+                linked to the model version, but also if any step run is linked
+                to the model version.
             orchestrator_run_id: The run id of the orchestrator to filter by.
             name: The name of the run to filter by.
             status: The status of the pipeline run
@@ -4256,6 +4262,7 @@ def list_pipeline_runs(
             template_id=template_id,
             source_snapshot_id=source_snapshot_id,
             model_version_id=model_version_id,
+            linked_to_model_version_id=linked_to_model_version_id,
             orchestrator_run_id=orchestrator_run_id,
             stack_id=stack_id,
             status=status,

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

@@ -636,6 +636,7 @@ class PipelineRunFilter(
         "pipeline_name",
         "templatable",
         "triggered_by_step_run_id",
+        "linked_to_model_version_id",
     ]
     CLI_EXCLUDE_FIELDS = [
         *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -702,6 +703,14 @@ class PipelineRunFilter(
         description="Model version associated with the pipeline run.",
         union_mode="left_to_right",
     )
+    linked_to_model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Filter by model version linked to the pipeline run. "
+        "The difference to `model_version_id` is that this filter will "
+        "not only include pipeline runs which are directly linked to the model "
+        "version, but also if any step run is linked to the model version.",
+        union_mode="left_to_right",
+    )
     status: Optional[str] = Field(
         default=None,
         description="Name of the Pipeline Run",
@@ -777,6 +786,7 @@ def get_custom_filters(
             CodeReferenceSchema,
             CodeRepositorySchema,
             ModelSchema,
+            ModelVersionPipelineRunSchema,
             ModelVersionSchema,
             PipelineBuildSchema,
             PipelineRunSchema,
@@ -960,6 +970,20 @@ def get_custom_filters(
             )
             custom_filters.append(trigger_filter)
 
+        if self.linked_to_model_version_id:
+            linked_to_model_version_filter = and_(
+                PipelineRunSchema.id
+                == ModelVersionPipelineRunSchema.pipeline_run_id,
+                ModelVersionPipelineRunSchema.model_version_id
+                == ModelVersionSchema.id,
+                self.generate_custom_query_conditions_for_column(
+                    value=self.linked_to_model_version_id,
+                    table=ModelVersionSchema,
+                    column="id",
+                ),
+            )
+            custom_filters.append(linked_to_model_version_filter)
+
         return custom_filters
 
     def apply_sorting(

tests/integration/functional/zen_stores/test_zen_store.py

@@ -5075,6 +5075,37 @@ def test_link_create_duplicated(self):
             )
             assert link_1.id == link_2.id
 
+    def test_linked_pipeline_run_fetching(self):
+        """Test that model version pipeline run links can be fetched using
+        the list_pipeline_runs method.
+        """
+        with ModelContext(True, create_prs=2) as (
+            model_version,
+            prs,
+        ):
+            zs = Client().zen_store
+            zs.create_model_version_pipeline_run_link(
+                ModelVersionPipelineRunRequest(
+                    model_version=model_version.id,
+                    pipeline_run=prs[0].id,
+                )
+            )
+            zs.create_model_version_pipeline_run_link(
+                ModelVersionPipelineRunRequest(
+                    model_version=model_version.id,
+                    pipeline_run=prs[1].id,
+                )
+            )
+
+            assert (
+                Client()
+                .list_pipeline_runs(
+                    linked_to_model_version_id=model_version.id
+                )
+                .total
+                == 2
+            )
+
     def test_link_delete_found(self):
         with ModelContext(True, create_prs=1) as (
             model_version,