Allow filtering pipeline runs by stack component

zenml-io · Oct 24, 2024 · ee8c0b4 · ee8c0b4
1 parent 024d612
commit ee8c0b4
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/zenml/client.py b/src/zenml/client.py
@@ -3778,6 +3778,7 @@ def list_pipeline_runs(
         code_repository: Optional[Union[UUID, str]] = None,
         model: Optional[Union[UUID, str]] = None,
         stack: Optional[Union[UUID, str]] = None,
+        stack_component: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[PipelineRunResponse]:
         """List all pipeline runs.
@@ -3816,6 +3817,7 @@ def list_pipeline_runs(
             code_repository: Filter by code repository name/ID.
             model: Filter by model name/ID.
             stack: Filter by stack name/ID.
+            stack_component: Filter by stack component name/ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -3854,6 +3856,7 @@ def list_pipeline_runs(
             code_repository=code_repository,
             stack=stack,
             model=model,
+            stack_component=stack_component,
             templatable=templatable,
         )
         runs_filter_model.set_scope_workspace(self.active_workspace.id)

diff --git a/src/zenml/models/v2/core/pipeline_run.py b/src/zenml/models/v2/core/pipeline_run.py
@@ -586,6 +586,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "stack",
         "code_repository",
         "model",
+        "component",
         "pipeline_name",
         "templatable",
     ]
@@ -688,6 +689,10 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         default=None,
         description="Name/ID of the model associated with the run.",
     )
+    stack_component: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the stack component associated with the run.",
+    )
     templatable: Optional[bool] = Field(
         default=None, description="Whether the run is templatable."
     )
@@ -716,6 +721,8 @@ def get_custom_filters(
             PipelineRunSchema,
             PipelineSchema,
             ScheduleSchema,
+            StackComponentSchema,
+            StackCompositionSchema,
             StackSchema,
             UserSchema,
         )
@@ -821,6 +828,19 @@ def get_custom_filters(
             )
             custom_filters.append(model_filter)
 
+        if self.stack_component:
+            component_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+                StackSchema.id == StackCompositionSchema.stack_id,
+                StackCompositionSchema.component_id == StackComponentSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.stack_component,
+                    table=StackComponentSchema,
+                ),
+            )
+            custom_filters.append(component_filter)
+
         if self.pipeline_name:
             pipeline_name_filter = and_(
                 PipelineRunSchema.pipeline_id == PipelineSchema.id,