zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/models/v2/core/pipeline_run.py

@@ -21,38 +21,49 @@ from typing import (
     Dict,
     List,
     Optional,
+    Type,
+    TypeVar,
     Union,
+    cast,
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
+    WorkspaceScopedTaggableFilter,
 )
+from zenml.models.v2.core.model_version import ModelVersionResponse
+from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+    from sqlalchemy.sql.elements import ColumnElement
 
+    from zenml.models import TriggerExecutionResponse
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.code_reference import CodeReferenceResponse
     from zenml.models.v2.core.pipeline import PipelineResponse
     from zenml.models.v2.core.pipeline_build import (
         PipelineBuildResponse,
     )
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.models.v2.core.schedule import ScheduleResponse
     from zenml.models.v2.core.stack import StackResponse
     from zenml.models.v2.core.step_run import StepRunResponse
+    from zenml.zen_stores.schemas.base_schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
+AnyQuery = TypeVar("AnyQuery", bound=Any)
+
 
 # ------------------ Request Model ------------------
 
@@ -68,7 +79,8 @@ class PipelineRunRequest(WorkspaceScopedRequest):
         title="The deployment associated with the pipeline run."
     )
     pipeline: Optional[UUID] = Field(
-        title="The pipeline associated with the pipeline run."
+        title="The pipeline associated with the pipeline run.",
+        default=None,
     )
     orchestrator_run_id: Optional[str] = Field(
         title="The orchestrator run ID.",
@@ -100,6 +112,21 @@ class PipelineRunRequest(WorkspaceScopedRequest):
             "(OS, Python version, etc.)."
         ),
     )
+    trigger_execution_id: Optional[UUID] = Field(
+        default=None,
+        title="ID of the trigger execution that triggered this run.",
+    )
+    tags: Optional[List[str]] = Field(
+        default=None,
+        title="Tags of the pipeline run.",
+    )
+    model_version_id: Optional[UUID] = Field(
+        title="The ID of the model version that was "
+        "configured by this pipeline run explicitly.",
+        default=None,
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 # ------------------ Update Model ------------------
@@ -110,6 +137,21 @@ class PipelineRunUpdate(BaseModel):
 
     status: Optional[ExecutionStatus] = None
     end_time: Optional[datetime] = None
+    model_version_id: Optional[UUID] = Field(
+        title="The ID of the model version that was "
+        "configured by this pipeline run explicitly.",
+        default=None,
+    )
+    # TODO: we should maybe have a different update model here, the upper
+    #  three attributes should only be for internal use
+    add_tags: Optional[List[str]] = Field(
+        default=None, title="New tags to add to the pipeline run."
+    )
+    remove_tags: Optional[List[str]] = Field(
+        default=None, title="Tags to remove from the pipeline run."
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 # ------------------ Response Model ------------------
@@ -136,12 +178,25 @@ class PipelineRunResponseBody(WorkspaceScopedResponseBody):
     code_reference: Optional["CodeReferenceResponse"] = Field(
         default=None, title="The code reference that was used for this run."
     )
+    deployment_id: Optional[UUID] = Field(
+        default=None, title="The deployment that was used for this run."
+    )
+    trigger_execution: Optional["TriggerExecutionResponse"] = Field(
+        default=None, title="The trigger execution that triggered this run."
+    )
+    model_version_id: Optional[UUID] = Field(
+        title="The ID of the model version that was "
+        "configured by this pipeline run explicitly.",
+        default=None,
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for pipeline runs."""
 
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         default={},
         title="Metadata associated with this pipeline run.",
     )
@@ -178,11 +233,46 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=STR_FIELD_MAX_LENGTH,
         default=None,
     )
+    code_path: Optional[str] = Field(
+        default=None,
+        title="Optional path where the code is stored in the artifact store.",
+    )
+    template_id: Optional[UUID] = Field(
+        default=None,
+        description="Template used for the pipeline run.",
+    )
+    is_templatable: bool = Field(
+        default=False,
+        description="Whether a template can be created from this run.",
+    )
+    step_substitutions: Dict[str, Dict[str, str]] = Field(
+        title="Substitutions used in the step runs of this pipeline run.",
+        default_factory=dict,
+    )
+
+
+class PipelineRunResponseResources(WorkspaceScopedResponseResources):
+    """Class for all resource models associated with the pipeline run entity."""
+
+    model_version: Optional[ModelVersionResponse] = None
+    tags: List[TagResponse] = Field(
+        title="Tags associated with the pipeline run.",
+    )
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class PipelineRunResponse(
     WorkspaceScopedResponse[
-        PipelineRunResponseBody, PipelineRunResponseMetadata
+        PipelineRunResponseBody,
+        PipelineRunResponseMetadata,
+        PipelineRunResponseResources,
     ]
 ):
     """Response model for pipeline runs."""
@@ -229,6 +319,64 @@ class PipelineRunResponse(
 
         return get_artifacts_versions_of_pipeline_run(self, only_produced=True)
 
+    def refresh_run_status(self) -> "PipelineRunResponse":
+        """Method to refresh the status of a run if it is initializing/running.
+
+        Returns:
+            The updated pipeline.
+
+        Raises:
+            ValueError: If the stack of the run response is None.
+        """
+        if self.status in [
+            ExecutionStatus.INITIALIZING,
+            ExecutionStatus.RUNNING,
+        ]:
+            # Check if the stack still accessible
+            if self.stack is None:
+                raise ValueError(
+                    "The stack that this pipeline run response was executed on"
+                    "has been deleted."
+                )
+
+            # Create the orchestrator instance
+            from zenml.enums import StackComponentType
+            from zenml.orchestrators.base_orchestrator import BaseOrchestrator
+            from zenml.stack.stack_component import StackComponent
+
+            # Check if the stack still accessible
+            orchestrator_list = self.stack.components.get(
+                StackComponentType.ORCHESTRATOR, []
+            )
+            if len(orchestrator_list) == 0:
+                raise ValueError(
+                    "The orchestrator that this pipeline run response was "
+                    "executed with has been deleted."
+                )
+
+            orchestrator = cast(
+                BaseOrchestrator,
+                StackComponent.from_model(
+                    component_model=orchestrator_list[0]
+                ),
+            )
+
+            # Fetch the status
+            status = orchestrator.fetch_status(run=self)
+
+            # If it is different from the current status, update it
+            if status != self.status:
+                from zenml.client import Client
+                from zenml.models import PipelineRunUpdate
+
+                client = Client()
+                return client.zen_store.update_run(
+                    run_id=self.id,
+                    run_update=PipelineRunUpdate(status=status),
+                )
+
+        return self
+
     # Body and metadata properties
     @property
     def status(self) -> ExecutionStatus:
@@ -275,6 +423,15 @@ class PipelineRunResponse(
         """
         return self.get_body().schedule
 
+    @property
+    def trigger_execution(self) -> Optional["TriggerExecutionResponse"]:
+        """The `trigger_execution` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().trigger_execution
+
     @property
     def code_reference(self) -> Optional["CodeReferenceResponse"]:
         """The `schedule` property.
@@ -285,7 +442,25 @@ class PipelineRunResponse(
         return self.get_body().code_reference
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def deployment_id(self) -> Optional["UUID"]:
+        """The `deployment_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().deployment_id
+
+    @property
+    def model_version_id(self) -> Optional[UUID]:
+        """The `model_version_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().model_version_id
+
+    @property
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `run_metadata` property.
 
         Returns:
@@ -356,20 +531,92 @@ class PipelineRunResponse(
         """
         return self.get_metadata().orchestrator_run_id
 
+    @property
+    def code_path(self) -> Optional[str]:
+        """The `code_path` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().code_path
+
+    @property
+    def template_id(self) -> Optional[UUID]:
+        """The `template_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().template_id
+
+    @property
+    def is_templatable(self) -> bool:
+        """The `is_templatable` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().is_templatable
+
+    @property
+    def step_substitutions(self) -> Dict[str, Dict[str, str]]:
+        """The `step_substitutions` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().step_substitutions
+
+    @property
+    def model_version(self) -> Optional[ModelVersionResponse]:
+        """The `model_version` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().model_version
+
+    @property
+    def tags(self) -> List[TagResponse]:
+        """The `tags` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().tags
+
 
 # ------------------ Filter Model ------------------
 
 
-class PipelineRunFilter(WorkspaceScopedFilter):
+class PipelineRunFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of all Workspaces."""
 
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *WorkspaceScopedTaggableFilter.CUSTOM_SORTING_OPTIONS,
+        "tag",
+        "stack",
+        "pipeline",
+        "model",
+        "model_version",
+    ]
+
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "unlisted",
         "code_repository_id",
         "build_id",
         "schedule_id",
         "stack_id",
+        "template_id",
+        "pipeline",
+        "stack",
+        "code_repository",
+        "model",
+        "stack_component",
+        "pipeline_name",
+        "templatable",
+        "run_metadata",
     ]
     name: Optional[str] = Field(
         default=None,
@@ -380,59 +627,125 @@ class PipelineRunFilter(WorkspaceScopedFilter):
         description="Name of the Pipeline Run within the orchestrator",
     )
     pipeline_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Pipeline associated with the Pipeline Run"
-    )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the Pipeline Run"
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User that created the Pipeline Run"
+        default=None,
+        description="Pipeline associated with the Pipeline Run",
+        union_mode="left_to_right",
     )
     stack_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Stack used for the Pipeline Run"
+        default=None,
+        description="Stack used for the Pipeline Run",
+        union_mode="left_to_right",
     )
     schedule_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Schedule that triggered the Pipeline Run"
+        default=None,
+        description="Schedule that triggered the Pipeline Run",
+        union_mode="left_to_right",
     )
     build_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Build used for the Pipeline Run"
+        default=None,
+        description="Build used for the Pipeline Run",
+        union_mode="left_to_right",
     )
     deployment_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Deployment used for the Pipeline Run"
+        default=None,
+        description="Deployment used for the Pipeline Run",
+        union_mode="left_to_right",
     )
     code_repository_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Code repository used for the Pipeline Run"
+        default=None,
+        description="Code repository used for the Pipeline Run",
+        union_mode="left_to_right",
+    )
+    template_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Template used for the pipeline run.",
+        union_mode="left_to_right",
+    )
+    model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Model version associated with the pipeline run.",
+        union_mode="left_to_right",
     )
     status: Optional[str] = Field(
         default=None,
         description="Name of the Pipeline Run",
     )
     start_time: Optional[Union[datetime, str]] = Field(
-        default=None, description="Start time for this run"
+        default=None,
+        description="Start time for this run",
+        union_mode="left_to_right",
     )
     end_time: Optional[Union[datetime, str]] = Field(
-        default=None, description="End time for this run"
+        default=None,
+        description="End time for this run",
+        union_mode="left_to_right",
     )
     unlisted: Optional[bool] = None
+    run_metadata: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="The run_metadata to filter the pipeline runs by.",
+    )
+    # TODO: Remove once frontend is ready for it. This is replaced by the more
+    #   generic `pipeline` filter below.
+    pipeline_name: Optional[str] = Field(
+        default=None,
+        description="Name of the pipeline associated with the run",
+    )
+    pipeline: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the pipeline associated with the run.",
+    )
+    stack: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the stack associated with the run.",
+    )
+    code_repository: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the code repository associated with the run.",
+    )
+    model: Optional[Union[UUID, str]] = Field(
+        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."
+    )
+    model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
         self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+        table: Type["AnySchema"],
+    ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
-        from sqlalchemy import and_
+        from sqlmodel import and_, col, or_
 
         from zenml.zen_stores.schemas import (
             CodeReferenceSchema,
+            CodeRepositorySchema,
+            ModelSchema,
+            ModelVersionSchema,
             PipelineBuildSchema,
             PipelineDeploymentSchema,
             PipelineRunSchema,
+            PipelineSchema,
+            RunMetadataResourceSchema,
+            RunMetadataSchema,
             ScheduleSchema,
+            StackComponentSchema,
+            StackCompositionSchema,
             StackSchema,
         )
 
@@ -444,7 +757,7 @@ class PipelineRunFilter(WorkspaceScopedFilter):
             custom_filters.append(unlisted_filter)
 
         if self.code_repository_id:
-            code_repo_filter = and_(  # type: ignore[type-var]
+            code_repo_filter = and_(
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
                 PipelineDeploymentSchema.code_reference_id
                 == CodeReferenceSchema.id,
@@ -454,7 +767,7 @@ class PipelineRunFilter(WorkspaceScopedFilter):
             custom_filters.append(code_repo_filter)
 
         if self.stack_id:
-            stack_filter = and_(  # type: ignore[type-var]
+            stack_filter = and_(
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
                 PipelineDeploymentSchema.stack_id == StackSchema.id,
                 StackSchema.id == self.stack_id,
@@ -462,7 +775,7 @@ class PipelineRunFilter(WorkspaceScopedFilter):
             custom_filters.append(stack_filter)
 
         if self.schedule_id:
-            schedule_filter = and_(  # type: ignore[type-var]
+            schedule_filter = and_(
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
                 PipelineDeploymentSchema.schedule_id == ScheduleSchema.id,
                 ScheduleSchema.id == self.schedule_id,
@@ -470,11 +783,210 @@ class PipelineRunFilter(WorkspaceScopedFilter):
             custom_filters.append(schedule_filter)
 
         if self.build_id:
-            pipeline_build_filter = and_(  # type: ignore[type-var]
+            pipeline_build_filter = and_(
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
                 PipelineDeploymentSchema.build_id == PipelineBuildSchema.id,
                 PipelineBuildSchema.id == self.build_id,
             )
             custom_filters.append(pipeline_build_filter)
 
+        if self.template_id:
+            run_template_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.template_id == self.template_id,
+            )
+            custom_filters.append(run_template_filter)
+
+        if self.pipeline:
+            pipeline_filter = and_(
+                PipelineRunSchema.pipeline_id == PipelineSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.pipeline, table=PipelineSchema
+                ),
+            )
+            custom_filters.append(pipeline_filter)
+
+        if self.stack:
+            stack_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.stack,
+                    table=StackSchema,
+                ),
+            )
+            custom_filters.append(stack_filter)
+
+        if self.code_repository:
+            code_repo_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.code_reference_id
+                == CodeReferenceSchema.id,
+                CodeReferenceSchema.code_repository_id
+                == CodeRepositorySchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.code_repository,
+                    table=CodeRepositorySchema,
+                ),
+            )
+            custom_filters.append(code_repo_filter)
+
+        if self.model:
+            model_filter = and_(
+                PipelineRunSchema.model_version_id == ModelVersionSchema.id,
+                ModelVersionSchema.model_id == ModelSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.model, table=ModelSchema
+                ),
+            )
+            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,
+                self.generate_custom_query_conditions_for_column(
+                    value=self.pipeline_name,
+                    table=PipelineSchema,
+                    column="name",
+                ),
+            )
+            custom_filters.append(pipeline_name_filter)
+
+        if self.templatable is not None:
+            if self.templatable is True:
+                templatable_filter = and_(
+                    # The following condition is not perfect as it does not
+                    # consider stacks with custom flavor components or local
+                    # components, but the best we can do currently with our
+                    # table columns.
+                    PipelineRunSchema.deployment_id
+                    == PipelineDeploymentSchema.id,
+                    PipelineDeploymentSchema.build_id
+                    == PipelineBuildSchema.id,
+                    col(PipelineBuildSchema.is_local).is_(False),
+                    col(PipelineBuildSchema.stack_id).is_not(None),
+                )
+            else:
+                templatable_filter = or_(
+                    col(PipelineRunSchema.deployment_id).is_(None),
+                    and_(
+                        PipelineRunSchema.deployment_id
+                        == PipelineDeploymentSchema.id,
+                        col(PipelineDeploymentSchema.build_id).is_(None),
+                    ),
+                    and_(
+                        PipelineRunSchema.deployment_id
+                        == PipelineDeploymentSchema.id,
+                        PipelineDeploymentSchema.build_id
+                        == PipelineBuildSchema.id,
+                        or_(
+                            col(PipelineBuildSchema.is_local).is_(True),
+                            col(PipelineBuildSchema.stack_id).is_(None),
+                        ),
+                    ),
+                )
+
+            custom_filters.append(templatable_filter)
+        if self.run_metadata is not None:
+            from zenml.enums import MetadataResourceTypes
+
+            for key, value in self.run_metadata.items():
+                additional_filter = and_(
+                    RunMetadataResourceSchema.resource_id
+                    == PipelineRunSchema.id,
+                    RunMetadataResourceSchema.resource_type
+                    == MetadataResourceTypes.PIPELINE_RUN,
+                    RunMetadataResourceSchema.run_metadata_id
+                    == RunMetadataSchema.id,
+                    self.generate_custom_query_conditions_for_column(
+                        value=value,
+                        table=RunMetadataSchema,
+                        column="value",
+                    ),
+                )
+                custom_filters.append(additional_filter)
+
         return custom_filters
+
+    def apply_sorting(
+        self,
+        query: AnyQuery,
+        table: Type["AnySchema"],
+    ) -> AnyQuery:
+        """Apply sorting to the query.
+
+        Args:
+            query: The query to which to apply the sorting.
+            table: The query table.
+
+        Returns:
+            The query with sorting applied.
+        """
+        from sqlmodel import asc, desc
+
+        from zenml.enums import SorterOps
+        from zenml.zen_stores.schemas import (
+            ModelSchema,
+            ModelVersionSchema,
+            PipelineDeploymentSchema,
+            PipelineRunSchema,
+            PipelineSchema,
+            StackSchema,
+        )
+
+        sort_by, operand = self.sorting_params
+
+        if sort_by == "pipeline":
+            query = query.join(
+                PipelineSchema,
+                PipelineRunSchema.pipeline_id == PipelineSchema.id,
+            )
+            column = PipelineSchema.name
+        elif sort_by == "stack":
+            query = query.join(
+                PipelineDeploymentSchema,
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+            ).join(
+                StackSchema,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+            )
+            column = StackSchema.name
+        elif sort_by == "model":
+            query = query.join(
+                ModelVersionSchema,
+                PipelineRunSchema.model_version_id == ModelVersionSchema.id,
+            ).join(
+                ModelSchema,
+                ModelVersionSchema.model_id == ModelSchema.id,
+            )
+            column = ModelSchema.name
+        elif sort_by == "model_version":
+            query = query.join(
+                ModelVersionSchema,
+                PipelineRunSchema.model_version_id == ModelVersionSchema.id,
+            )
+            column = ModelVersionSchema.name
+        else:
+            return super().apply_sorting(query=query, table=table)
+
+        query = query.add_columns(column)
+
+        if operand == SorterOps.ASCENDING:
+            query = query.order_by(asc(column))
+        else:
+            query = query.order_by(desc(column))
+
+        return query