zenml-nightly 0.62.0.dev20240728__py3-none-any.whl → 0.62.0.dev20240730__py3-none-any.whl

zenml/models/v2/core/pipeline_deployment.py

@@ -20,6 +20,7 @@ from pydantic import Field
 
 from zenml.config.docker_settings import SourceFileMode
 from zenml.config.pipeline_configurations import PipelineConfiguration
+from zenml.config.pipeline_spec import PipelineSpec
 from zenml.config.step_configurations import Step
 from zenml.models.v2.base.base import BaseZenModel
 from zenml.models.v2.base.page import Page
@@ -71,6 +72,14 @@ class PipelineDeploymentBase(BaseZenModel):
         default=None,
         title="The version of the ZenML installation on the server side.",
     )
+    pipeline_version_hash: Optional[str] = Field(
+        default=None,
+        title="The pipeline version hash of the deployment.",
+    )
+    pipeline_spec: Optional[PipelineSpec] = Field(
+        default=None,
+        title="The pipeline spec of the deployment.",
+    )
 
     @property
     def requires_included_files(self) -> bool:
@@ -116,6 +125,10 @@ class PipelineDeploymentRequest(
         default=None,
         title="The code reference associated with the deployment.",
     )
+    template: Optional[UUID] = Field(
+        default=None,
+        description="Template used for the deployment.",
+    )
 
 
 # ------------------ Update Model ------------------
@@ -150,6 +163,13 @@ class PipelineDeploymentResponseMetadata(WorkspaceScopedResponseMetadata):
     server_version: Optional[str] = Field(
         title="The version of the ZenML installation on the server side."
     )
+    pipeline_version_hash: Optional[str] = Field(
+        default=None, title="The pipeline version hash of the deployment."
+    )
+    pipeline_spec: Optional[PipelineSpec] = Field(
+        default=None, title="The pipeline spec of the deployment."
+    )
+
     pipeline: Optional[PipelineResponse] = Field(
         default=None, title="The pipeline associated with the deployment."
     )
@@ -167,6 +187,10 @@ class PipelineDeploymentResponseMetadata(WorkspaceScopedResponseMetadata):
         default=None,
         title="The code reference associated with the deployment.",
     )
+    template_id: Optional[UUID] = Field(
+        default=None,
+        description="Template used for the pipeline run.",
+    )
 
 
 class PipelineDeploymentResponseResources(WorkspaceScopedResponseResources):
@@ -251,6 +275,24 @@ class PipelineDeploymentResponse(
         """
         return self.get_metadata().server_version
 
+    @property
+    def pipeline_version_hash(self) -> Optional[str]:
+        """The `pipeline_version_hash` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().pipeline_version_hash
+
+    @property
+    def pipeline_spec(self) -> Optional[PipelineSpec]:
+        """The `pipeline_spec` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().pipeline_spec
+
     @property
     def pipeline(self) -> Optional[PipelineResponse]:
         """The `pipeline` property.
@@ -296,6 +338,15 @@ class PipelineDeploymentResponse(
         """
         return self.get_metadata().code_reference
 
+    @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 requires_code_download(self) -> bool:
         """Whether the deployment requires downloading some code files.
@@ -345,3 +396,8 @@ class PipelineDeploymentFilter(WorkspaceScopedFilter):
         description="Schedule associated with the deployment.",
         union_mode="left_to_right",
     )
+    template_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Template used as base for the deployment.",
+        union_mode="left_to_right",
+    )

zenml/models/v2/core/pipeline_run.py

@@ -37,8 +37,10 @@ from zenml.models.v2.base.scoped import (
     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 ColumnElement
@@ -109,6 +111,10 @@ class PipelineRunRequest(WorkspaceScopedRequest):
         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.",
+    )
 
 
 # ------------------ Update Model ------------------
@@ -119,6 +125,13 @@ class PipelineRunUpdate(BaseModel):
 
     status: Optional[ExecutionStatus] = None
     end_time: Optional[datetime] = None
+    # TODO: we should maybe have a different update model here, the upper two 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."
+    )
 
 
 # ------------------ Response Model ------------------
@@ -193,12 +206,19 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=STR_FIELD_MAX_LENGTH,
         default=None,
     )
+    template_id: Optional[UUID] = Field(
+        default=None,
+        description="Template used for the pipeline run.",
+    )
 
 
 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.
@@ -405,6 +425,15 @@ class PipelineRunResponse(
         """
         return self.get_metadata().orchestrator_run_id
 
+    @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 model_version(self) -> Optional[ModelVersionResponse]:
         """The `model_version` property.
@@ -414,11 +443,20 @@ class PipelineRunResponse(
         """
         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."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
@@ -428,6 +466,7 @@ class PipelineRunFilter(WorkspaceScopedFilter):
         "build_id",
         "schedule_id",
         "stack_id",
+        "template_id",
         "pipeline_name",
     ]
     name: Optional[str] = Field(
@@ -482,6 +521,11 @@ class PipelineRunFilter(WorkspaceScopedFilter):
         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",
+    )
     status: Optional[str] = Field(
         default=None,
         description="Name of the Pipeline Run",
@@ -574,4 +618,11 @@ class PipelineRunFilter(WorkspaceScopedFilter):
             )
             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)
+
         return custom_filters

zenml/models/v2/core/run_template.py

@@ -0,0 +1,393 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Models representing pipeline templates."""
+
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
+from uuid import UUID
+
+from pydantic import Field
+
+from zenml.config.pipeline_spec import PipelineSpec
+from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
+from zenml.enums import ExecutionStatus
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.scoped import (
+    WorkspaceScopedRequest,
+    WorkspaceScopedResponse,
+    WorkspaceScopedResponseBody,
+    WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
+    WorkspaceScopedTaggableFilter,
+)
+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.pipeline_deployment import (
+    PipelineDeploymentResponse,
+)
+from zenml.models.v2.core.tag import TagResponse
+
+if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
+# ------------------ Request Model ------------------
+
+
+class RunTemplateRequest(WorkspaceScopedRequest):
+    """Request model for run templates."""
+
+    name: str = Field(
+        title="The name of the run template.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    description: Optional[str] = Field(
+        default=None,
+        title="The description of the run template.",
+        max_length=TEXT_FIELD_MAX_LENGTH,
+    )
+    source_deployment_id: UUID = Field(
+        title="The deployment that should be the base of the created template."
+    )
+    tags: Optional[List[str]] = Field(
+        default=None,
+        title="Tags of the run template.",
+    )
+
+
+# ------------------ Update Model ------------------
+
+
+class RunTemplateUpdate(BaseUpdate):
+    """Run template update model."""
+
+    name: Optional[str] = Field(
+        default=None,
+        title="The name of the run template.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    description: Optional[str] = Field(
+        default=None,
+        title="The description of the run template.",
+        max_length=TEXT_FIELD_MAX_LENGTH,
+    )
+    add_tags: Optional[List[str]] = Field(
+        default=None, title="New tags to add to the run template."
+    )
+    remove_tags: Optional[List[str]] = Field(
+        default=None, title="Tags to remove from the run template."
+    )
+
+
+# ------------------ Response Model ------------------
+
+
+class RunTemplateResponseBody(WorkspaceScopedResponseBody):
+    """Response body for run templates."""
+
+    runnable: bool = Field(
+        title="If a run can be started from the template.",
+    )
+    latest_run_id: Optional[UUID] = Field(
+        default=None,
+        title="The ID of the latest run of the run template.",
+    )
+    latest_run_status: Optional[ExecutionStatus] = Field(
+        default=None,
+        title="The status of the latest run of the run template.",
+    )
+
+
+class RunTemplateResponseMetadata(WorkspaceScopedResponseMetadata):
+    """Response metadata for run templates."""
+
+    description: Optional[str] = Field(
+        default=None,
+        title="The description of the run template.",
+    )
+    pipeline_spec: Optional[PipelineSpec] = Field(
+        default=None, title="The spec of the pipeline."
+    )
+    config_template: Optional[Dict[str, Any]] = Field(
+        default=None, title="Run configuration template."
+    )
+    config_schema: Optional[Dict[str, Any]] = Field(
+        default=None, title="Run configuration schema."
+    )
+
+
+class RunTemplateResponseResources(WorkspaceScopedResponseResources):
+    """All resource models associated with the run template."""
+
+    source_deployment: Optional[PipelineDeploymentResponse] = Field(
+        default=None,
+        title="The deployment that is the source of the template.",
+    )
+    pipeline: Optional[PipelineResponse] = Field(
+        default=None, title="The pipeline associated with the template."
+    )
+    build: Optional[PipelineBuildResponse] = Field(
+        default=None,
+        title="The pipeline build associated with the template.",
+    )
+    code_reference: Optional[CodeReferenceResponse] = Field(
+        default=None,
+        title="The code reference associated with the template.",
+    )
+    tags: List[TagResponse] = Field(
+        title="Tags associated with the run template.",
+    )
+
+
+class RunTemplateResponse(
+    WorkspaceScopedResponse[
+        RunTemplateResponseBody,
+        RunTemplateResponseMetadata,
+        RunTemplateResponseResources,
+    ]
+):
+    """Response model for run templates."""
+
+    name: str = Field(
+        title="The name of the run template.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+
+    def get_hydrated_version(self) -> "RunTemplateResponse":
+        """Return the hydrated version of this run template.
+
+        Returns:
+            The hydrated run template.
+        """
+        from zenml.client import Client
+
+        return Client().zen_store.get_run_template(
+            template_id=self.id, hydrate=True
+        )
+
+    # Body and metadata properties
+    @property
+    def runnable(self) -> bool:
+        """The `runnable` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().runnable
+
+    @property
+    def latest_run_id(self) -> Optional[UUID]:
+        """The `latest_run_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().latest_run_id
+
+    @property
+    def latest_run_status(self) -> Optional[ExecutionStatus]:
+        """The `latest_run_status` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().latest_run_status
+
+    @property
+    def description(self) -> Optional[str]:
+        """The `description` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().description
+
+    @property
+    def pipeline_spec(self) -> Optional[PipelineSpec]:
+        """The `pipeline_spec` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().pipeline_spec
+
+    @property
+    def config_template(self) -> Optional[Dict[str, Any]]:
+        """The `config_template` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().config_template
+
+    @property
+    def config_schema(self) -> Optional[Dict[str, Any]]:
+        """The `config_schema` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().config_schema
+
+    @property
+    def source_deployment(self) -> Optional[PipelineDeploymentResponse]:
+        """The `source_deployment` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().source_deployment
+
+    @property
+    def pipeline(self) -> Optional[PipelineResponse]:
+        """The `pipeline` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().pipeline
+
+    @property
+    def build(self) -> Optional[PipelineBuildResponse]:
+        """The `build` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().build
+
+    @property
+    def code_reference(self) -> Optional[CodeReferenceResponse]:
+        """The `code_reference` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().code_reference
+
+    @property
+    def tags(self) -> List[TagResponse]:
+        """The `tags` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().tags
+
+
+# ------------------ Filter Model ------------------
+
+
+class RunTemplateFilter(WorkspaceScopedTaggableFilter):
+    """Model for filtering of run templates."""
+
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
+        "code_repository_id",
+        "stack_id",
+        "build_id" "pipeline_id",
+    ]
+
+    name: Optional[str] = Field(
+        default=None,
+        description="Name of the run template.",
+    )
+    workspace_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Workspace associated with the template.",
+        union_mode="left_to_right",
+    )
+    user_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="User that created the template.",
+        union_mode="left_to_right",
+    )
+    pipeline_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Pipeline associated with the template.",
+        union_mode="left_to_right",
+    )
+    build_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Build associated with the template.",
+        union_mode="left_to_right",
+    )
+    stack_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Stack associated with the template.",
+        union_mode="left_to_right",
+    )
+    code_repository_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Code repository associated with the template.",
+        union_mode="left_to_right",
+    )
+
+    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 (
+            CodeReferenceSchema,
+            PipelineDeploymentSchema,
+            RunTemplateSchema,
+        )
+
+        if self.code_repository_id:
+            code_repo_filter = and_(
+                RunTemplateSchema.source_deployment_id
+                == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.code_reference_id
+                == CodeReferenceSchema.id,
+                CodeReferenceSchema.code_repository_id
+                == self.code_repository_id,
+            )
+            custom_filters.append(code_repo_filter)
+
+        if self.stack_id:
+            stack_filter = and_(
+                RunTemplateSchema.source_deployment_id
+                == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == self.stack_id,
+            )
+            custom_filters.append(stack_filter)
+
+        if self.build_id:
+            build_filter = and_(
+                RunTemplateSchema.source_deployment_id
+                == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.build_id == self.build_id,
+            )
+            custom_filters.append(build_filter)
+
+        if self.pipeline_id:
+            pipeline_filter = and_(
+                RunTemplateSchema.source_deployment_id
+                == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.pipeline_id == self.pipeline_id,
+            )
+            custom_filters.append(pipeline_filter)
+
+        return custom_filters

zenml/models/v2/misc/stack_deployment.py

@@ -72,9 +72,14 @@ class StackDeploymentConfig(BaseModel):
         title="A textual description for the cloud provider console URL.",
     )
     configuration: Optional[str] = Field(
+        default=None,
         title="Configuration for the stack deployment that the user must "
         "manually configure into the cloud provider console.",
     )
+    instructions: Optional[str] = Field(
+        default=None,
+        title="Instructions for deploying the stack.",
+    )
 
 
 class DeployedStack(BaseModel):