zenml-nightly 0.70.0.dev20241201__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/models/v2/core/model_version.py

@@ -77,10 +77,6 @@ class ModelVersionRequest(WorkspaceScopedRequest):
         default=None,
     )
 
-    number: Optional[int] = Field(
-        description="The number of the model version",
-        default=None,
-    )
     model: UUID = Field(
         description="The ID of the model containing version",
     )
@@ -585,7 +581,6 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
-        "user",
         "run_metadata",
     ]
 
@@ -597,25 +592,11 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
         default=None,
         description="The number of the Model Version",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="The workspace of the Model Version",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="The user of the Model Version",
-        union_mode="left_to_right",
-    )
     stage: Optional[Union[str, ModelStages]] = Field(
         description="The model version stage",
         default=None,
         union_mode="left_to_right",
     )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the model version.",
-    )
     run_metadata: Optional[Dict[str, str]] = Field(
         default=None,
         description="The run_metadata to filter the model versions by.",
@@ -639,14 +620,17 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
         self._model_id = model_id
 
     def get_custom_filters(
-        self,
+        self, 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 sqlmodel import and_
 
@@ -654,20 +638,8 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
             ModelVersionSchema,
             RunMetadataResourceSchema,
             RunMetadataSchema,
-            UserSchema,
         )
 
-        if self.user:
-            user_filter = and_(
-                ModelVersionSchema.user_id == UserSchema.id,
-                self.generate_name_or_id_query_conditions(
-                    value=self.user,
-                    table=UserSchema,
-                    additional_columns=["full_name"],
-                ),
-            )
-            custom_filters.append(user_filter)
-
         if self.run_metadata is not None:
             from zenml.enums import MetadataResourceTypes
 

zenml/models/v2/core/model_version_artifact.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing the link between model versions and artifacts."""
 
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Type, TypeVar, Union
 from uuid import UUID
 
 from pydantic import ConfigDict, Field
@@ -32,6 +32,9 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 # ------------------ Request Model ------------------
@@ -164,13 +167,18 @@ class ModelVersionArtifactFilter(BaseFilter):
     #  careful we might overwrite some fields protected by pydantic.
     model_config = ConfigDict(protected_namespaces=())
 
-    def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List[Union["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 sqlmodel import and_, col
 

zenml/models/v2/core/model_version_pipeline_run.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing the link between model versions and pipeline runs."""
 
-from typing import List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Type, TypeVar, Union
 from uuid import UUID
 
 from pydantic import ConfigDict, Field
@@ -30,6 +30,12 @@ from zenml.models.v2.base.base import (
 from zenml.models.v2.base.filter import BaseFilter, StrFilter
 from zenml.models.v2.core.pipeline_run import PipelineRunResponse
 
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
+
 # ------------------ Request Model ------------------
 
 
@@ -147,13 +153,18 @@ class ModelVersionPipelineRunFilter(BaseFilter):
     #  careful we might overwrite some fields protected by pydantic.
     model_config = ConfigDict(protected_namespaces=())
 
-    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
+    def get_custom_filters(
+        self, 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 sqlmodel import and_
 

zenml/models/v2/core/pipeline.py

@@ -21,7 +21,6 @@ from typing import (
     Optional,
     Type,
     TypeVar,
-    Union,
 )
 from uuid import UUID
 
@@ -45,9 +44,7 @@ from zenml.models.v2.base.scoped import (
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import ColumnElement
-
-    from zenml.models.v2.core.pipeline_run import PipelineRunResponse
+    from zenml.models import PipelineRunResponse, UserResponse
     from zenml.zen_stores.schemas import BaseSchema
 
     AnySchema = TypeVar("AnySchema", bound=BaseSchema)
@@ -122,6 +119,10 @@ class PipelineResponseMetadata(WorkspaceScopedResponseMetadata):
 class PipelineResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the pipeline entity."""
 
+    latest_run_user: Optional["UserResponse"] = Field(
+        default=None,
+        title="The user that created the latest run of this pipeline.",
+    )
     tags: List[TagResponse] = Field(
         title="Tags associated with the pipeline.",
     )
@@ -258,10 +259,12 @@ class PipelineResponse(
 class PipelineFilter(WorkspaceScopedTaggableFilter):
     """Pipeline filter model."""
 
-    CUSTOM_SORTING_OPTIONS = [SORT_PIPELINES_BY_LATEST_RUN_KEY]
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *WorkspaceScopedTaggableFilter.CUSTOM_SORTING_OPTIONS,
+        SORT_PIPELINES_BY_LATEST_RUN_KEY,
+    ]
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
-        "user",
         "latest_run_status",
     ]
 
@@ -274,20 +277,6 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
         description="Filter by the status of the latest run of a pipeline. "
         "This will always be applied as an `AND` filter for now.",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the Pipeline",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the Pipeline",
-        union_mode="left_to_right",
-    )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the pipeline.",
-    )
 
     def apply_filter(
         self, query: AnyQuery, table: Type["AnySchema"]
@@ -343,36 +332,6 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
 
         return query
 
-    def get_custom_filters(
-        self,
-    ) -> List["ColumnElement[bool]"]:
-        """Get custom filters.
-
-        Returns:
-            A list of custom filters.
-        """
-        custom_filters = super().get_custom_filters()
-
-        from sqlmodel import and_
-
-        from zenml.zen_stores.schemas import (
-            PipelineSchema,
-            UserSchema,
-        )
-
-        if self.user:
-            user_filter = and_(
-                PipelineSchema.user_id == UserSchema.id,
-                self.generate_name_or_id_query_conditions(
-                    value=self.user,
-                    table=UserSchema,
-                    additional_columns=["full_name"],
-                ),
-            )
-            custom_filters.append(user_filter)
-
-        return custom_filters
-
     def apply_sorting(
         self,
         query: AnyQuery,
@@ -387,12 +346,45 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
         Returns:
             The query with sorting applied.
         """
-        column, _ = self.sorting_params
+        from sqlmodel import asc, case, col, desc, func, select
+
+        from zenml.enums import SorterOps
+        from zenml.zen_stores.schemas import PipelineRunSchema, PipelineSchema
+
+        sort_by, operand = self.sorting_params
+
+        if sort_by == SORT_PIPELINES_BY_LATEST_RUN_KEY:
+            # Subquery to find the latest run per pipeline
+            latest_run_subquery = (
+                select(
+                    PipelineRunSchema.pipeline_id,
+                    case(
+                        (
+                            func.max(PipelineRunSchema.created).is_(None),
+                            PipelineSchema.created,
+                        ),
+                        else_=func.max(PipelineRunSchema.created),
+                    ).label("latest_run"),
+                )
+                .group_by(col(PipelineRunSchema.pipeline_id))
+                .subquery()
+            )
+
+            # Join the subquery with the pipelines
+            query = query.outerjoin(
+                latest_run_subquery,
+                PipelineSchema.id == latest_run_subquery.c.pipeline_id,
+            )
+
+            if operand == SorterOps.ASCENDING:
+                query = query.order_by(
+                    asc(latest_run_subquery.c.latest_run)
+                ).order_by(col(PipelineSchema.id))
+            else:
+                query = query.order_by(
+                    desc(latest_run_subquery.c.latest_run)
+                ).order_by(col(PipelineSchema.id))
 
-        if column == SORT_PIPELINES_BY_LATEST_RUN_KEY:
-            # If sorting by the latest run, the sorting is already done in the
-            # base query in `SqlZenStore.list_pipelines(...)` and we don't need
-            # to to anything here
             return query
         else:
             return super().apply_sorting(query=query, table=table)

zenml/models/v2/core/pipeline_build.py

@@ -14,7 +14,17 @@
 """Models representing pipeline builds."""
 
 import json
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import Field
@@ -31,8 +41,13 @@ from zenml.models.v2.base.scoped import (
 from zenml.models.v2.misc.build_item import BuildItem
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.pipeline import PipelineResponse
     from zenml.models.v2.core.stack import StackResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 # ------------------ Request Model ------------------
@@ -446,23 +461,23 @@ class PipelineBuildResponse(
 class PipelineBuildFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of all pipeline builds."""
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        description="Workspace for this pipeline build.",
-        default=None,
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        description="User that produced this pipeline build.",
-        default=None,
-        union_mode="left_to_right",
-    )
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "container_registry_id",
+    ]
+
     pipeline_id: Optional[Union[UUID, str]] = Field(
         description="Pipeline associated with the pipeline build.",
         default=None,
         union_mode="left_to_right",
     )
     stack_id: Optional[Union[UUID, str]] = Field(
-        description="Stack used for the Pipeline Run",
+        description="Stack associated with the pipeline build.",
+        default=None,
+        union_mode="left_to_right",
+    )
+    container_registry_id: Optional[Union[UUID, str]] = Field(
+        description="Container registry associated with the pipeline build.",
         default=None,
         union_mode="left_to_right",
     )
@@ -484,3 +499,43 @@ class PipelineBuildFilter(WorkspaceScopedFilter):
     checksum: Optional[str] = Field(
         description="The build checksum.", default=None
     )
+    stack_checksum: Optional[str] = Field(
+        description="The stack checksum.", default=None
+    )
+
+    def get_custom_filters(
+        self,
+        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(table)
+
+        from sqlmodel import and_
+
+        from zenml.enums import StackComponentType
+        from zenml.zen_stores.schemas import (
+            PipelineBuildSchema,
+            StackComponentSchema,
+            StackCompositionSchema,
+            StackSchema,
+        )
+
+        if self.container_registry_id:
+            container_registry_filter = and_(
+                PipelineBuildSchema.stack_id == StackSchema.id,
+                StackSchema.id == StackCompositionSchema.stack_id,
+                StackCompositionSchema.component_id == StackComponentSchema.id,
+                StackComponentSchema.type
+                == StackComponentType.CONTAINER_REGISTRY.value,
+                StackComponentSchema.id == self.container_registry_id,
+            )
+            custom_filters.append(container_registry_filter)
+
+        return custom_filters

zenml/models/v2/core/pipeline_deployment.py

@@ -358,16 +358,6 @@ class PipelineDeploymentResponse(
 class PipelineDeploymentFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of all pipeline deployments."""
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace for this deployment.",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that created this deployment.",
-        union_mode="left_to_right",
-    )
     pipeline_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Pipeline associated with the deployment.",

zenml/models/v2/core/pipeline_run.py

@@ -16,10 +16,13 @@
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
     Dict,
     List,
     Optional,
+    Type,
+    TypeVar,
     Union,
     cast,
 )
@@ -55,6 +58,11 @@ if TYPE_CHECKING:
     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 ------------------
@@ -584,6 +592,15 @@ class PipelineRunResponse(
 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]] = [
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "unlisted",
@@ -592,7 +609,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "schedule_id",
         "stack_id",
         "template_id",
-        "user",
         "pipeline",
         "stack",
         "code_repository",
@@ -615,16 +631,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         description="Pipeline associated with the Pipeline Run",
         union_mode="left_to_right",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the Pipeline Run",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that created the Pipeline Run",
-        union_mode="left_to_right",
-    )
     stack_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Stack used for the Pipeline Run",
@@ -675,16 +681,12 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         union_mode="left_to_right",
     )
     unlisted: Optional[bool] = None
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the run.",
-    )
     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.
+    #   generic `pipeline` filter below.
     pipeline_name: Optional[str] = Field(
         default=None,
         description="Name of the pipeline associated with the run",
@@ -716,13 +718,17 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
 
     def get_custom_filters(
         self,
+        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 sqlmodel import and_, col, or_
 
@@ -741,7 +747,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             StackComponentSchema,
             StackCompositionSchema,
             StackSchema,
-            UserSchema,
         )
 
         if self.unlisted is not None:
@@ -792,17 +797,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(run_template_filter)
 
-        if self.user:
-            user_filter = and_(
-                PipelineRunSchema.user_id == UserSchema.id,
-                self.generate_name_or_id_query_conditions(
-                    value=self.user,
-                    table=UserSchema,
-                    additional_columns=["full_name"],
-                ),
-            )
-            custom_filters.append(user_filter)
-
         if self.pipeline:
             pipeline_filter = and_(
                 PipelineRunSchema.pipeline_id == PipelineSchema.id,
@@ -926,3 +920,71 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
                 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)
+
+        if operand == SorterOps.ASCENDING:
+            query = query.order_by(asc(column))
+        else:
+            query = query.order_by(desc(column))
+
+        return query