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

zenml/models/v2/core/model.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing models."""
 
-from typing import TYPE_CHECKING, ClassVar, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional
 from uuid import UUID
 
 from pydantic import BaseModel, Field
@@ -30,8 +30,6 @@ from zenml.models.v2.base.scoped import (
 from zenml.utils.pagination_utils import depaginate
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import ColumnElement
-
     from zenml.model.model import Model
     from zenml.models.v2.core.tag import TagResponse
 
@@ -318,61 +316,7 @@ class ModelResponse(
 class ModelFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of all Workspaces."""
 
-    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedTaggableFilter.CLI_EXCLUDE_FIELDS,
-        "workspace_id",
-        "user_id",
-    ]
-    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
-        "user",
-    ]
-
     name: Optional[str] = Field(
         default=None,
         description="Name of the Model",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the Model",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the Model",
-        union_mode="left_to_right",
-    )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the model.",
-    )
-
-    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 (
-            ModelSchema,
-            UserSchema,
-        )
-
-        if self.user:
-            user_filter = and_(
-                ModelSchema.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

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,43 +620,37 @@ 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_
 
         from zenml.zen_stores.schemas import (
             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
 
             for key, value in self.run_metadata.items():
                 additional_filter = and_(
-                    RunMetadataSchema.resource_id == ModelVersionSchema.id,
-                    RunMetadataSchema.resource_type
+                    RunMetadataResourceSchema.resource_id
+                    == ModelVersionSchema.id,
+                    RunMetadataResourceSchema.resource_type
                     == MetadataResourceTypes.MODEL_VERSION,
-                    RunMetadataSchema.key == key,
+                    RunMetadataResourceSchema.run_metadata_id
+                    == RunMetadataSchema.id,
                     self.generate_custom_query_conditions_for_column(
                         value=value,
                         table=RunMetadataSchema,

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.",