zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.64.0.dev20240809__py3-none-any.whl

zenml/zen_stores/rest_zen_store.py

@@ -80,6 +80,7 @@ from zenml.constants import (
     PIPELINE_DEPLOYMENTS,
     PIPELINES,
     RUN_METADATA,
+    RUN_TEMPLATES,
     RUNS,
     SCHEDULES,
     SECRETS,
@@ -192,6 +193,10 @@ from zenml.models import (
     RunMetadataFilter,
     RunMetadataRequest,
     RunMetadataResponse,
+    RunTemplateFilter,
+    RunTemplateRequest,
+    RunTemplateResponse,
+    RunTemplateUpdate,
     ScheduleFilter,
     ScheduleRequest,
     ScheduleResponse,
@@ -1524,35 +1529,6 @@ class RestZenStore(BaseZenStore):
             route=PIPELINE_BUILDS,
         )
 
-    def run_build(
-        self,
-        build_id: UUID,
-        run_configuration: Optional[PipelineRunConfiguration] = None,
-    ) -> PipelineRunResponse:
-        """Run a pipeline from a build.
-
-        Args:
-            build_id: The ID of the build to run.
-            run_configuration: Configuration for the run.
-
-        Raises:
-            RuntimeError: If the server does not support running a build.
-
-        Returns:
-            Model of the pipeline run.
-        """
-        run_configuration = run_configuration or PipelineRunConfiguration()
-        try:
-            response_body = self.post(
-                f"{PIPELINE_BUILDS}/{build_id}/runs", body=run_configuration
-            )
-        except MethodNotAllowedError as e:
-            raise RuntimeError(
-                "Running a build is not supported for this server."
-            ) from e
-
-        return PipelineRunResponse.model_validate(response_body)
-
     # -------------------------- Pipeline Deployments --------------------------
 
     def create_deployment(
@@ -1627,19 +1603,114 @@ class RestZenStore(BaseZenStore):
             route=PIPELINE_DEPLOYMENTS,
         )
 
-    def run_deployment(
+    # -------------------- Run templates --------------------
+
+    def create_run_template(
         self,
-        deployment_id: UUID,
+        template: RunTemplateRequest,
+    ) -> RunTemplateResponse:
+        """Create a new run template.
+
+        Args:
+            template: The template to create.
+
+        Returns:
+            The newly created template.
+        """
+        return self._create_workspace_scoped_resource(
+            resource=template,
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+        )
+
+    def get_run_template(
+        self, template_id: UUID, hydrate: bool = True
+    ) -> RunTemplateResponse:
+        """Get a run template with a given ID.
+
+        Args:
+            template_id: ID of the template.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The template.
+        """
+        return self._get_resource(
+            resource_id=template_id,
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_run_templates(
+        self,
+        template_filter_model: RunTemplateFilter,
+        hydrate: bool = False,
+    ) -> Page[RunTemplateResponse]:
+        """List all run templates matching the given filter criteria.
+
+        Args:
+            template_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all templates matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+            filter_model=template_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_run_template(
+        self,
+        template_id: UUID,
+        template_update: RunTemplateUpdate,
+    ) -> RunTemplateResponse:
+        """Updates a run template.
+
+        Args:
+            template_id: The ID of the template to update.
+            template_update: The update to apply.
+
+        Returns:
+            The updated template.
+        """
+        return self._update_resource(
+            resource_id=template_id,
+            resource_update=template_update,
+            route=RUN_TEMPLATES,
+            response_model=RunTemplateResponse,
+        )
+
+    def delete_run_template(self, template_id: UUID) -> None:
+        """Delete a run template.
+
+        Args:
+            template_id: The ID of the template to delete.
+        """
+        self._delete_resource(
+            resource_id=template_id,
+            route=RUN_TEMPLATES,
+        )
+
+    def run_template(
+        self,
+        template_id: UUID,
         run_configuration: Optional[PipelineRunConfiguration] = None,
     ) -> PipelineRunResponse:
-        """Run a pipeline from a deployment.
+        """Run a template.
 
         Args:
-            deployment_id: The ID of the deployment to run.
+            template_id: The ID of the template to run.
             run_configuration: Configuration for the run.
 
         Raises:
-            RuntimeError: If the server does not support running a deployment.
+            RuntimeError: If the server does not support running a template.
 
         Returns:
             Model of the pipeline run.
@@ -1648,12 +1719,12 @@ class RestZenStore(BaseZenStore):
 
         try:
             response_body = self.post(
-                f"{PIPELINE_DEPLOYMENTS}/{deployment_id}/runs",
+                f"{RUN_TEMPLATES}/{template_id}/runs",
                 body=run_configuration,
             )
         except MethodNotAllowedError as e:
             raise RuntimeError(
-                "Running a deployment is not supported for this server."
+                "Running a template is not supported for this server."
             ) from e
 
         return PipelineRunResponse.model_validate(response_body)
@@ -4224,7 +4295,7 @@ class RestZenStore(BaseZenStore):
         return self._request(
             "POST",
             self.url + API + VERSION_1 + path,
-            data=body.model_dump_json(),
+            json=body.model_dump(mode="json"),
             params=params,
             timeout=timeout,
             **kwargs,
@@ -4251,11 +4322,13 @@ class RestZenStore(BaseZenStore):
             The response body.
         """
         logger.debug(f"Sending PUT request to {path}...")
-        data = body.model_dump_json(exclude_unset=True) if body else None
+        json = (
+            body.model_dump(mode="json", exclude_unset=True) if body else None
+        )
         return self._request(
             "PUT",
             self.url + API + VERSION_1 + path,
-            data=data,
+            json=json,
             params=params,
             timeout=timeout,
             **kwargs,

zenml/zen_stores/schemas/__init__.py

@@ -69,6 +69,7 @@ from zenml.zen_stores.schemas.model_schemas import (
     ModelVersionArtifactSchema,
     ModelVersionPipelineRunSchema,
 )
+from zenml.zen_stores.schemas.run_template_schemas import RunTemplateSchema
 from zenml.zen_stores.schemas.server_settings_schemas import ServerSettingsSchema
 
 __all__ = [
@@ -102,6 +103,7 @@ __all__ = [
     "StepRunOutputArtifactSchema",
     "StepRunParentsSchema",
     "StepRunSchema",
+    "RunTemplateSchema",
     "TagSchema",
     "TagResourceSchema",
     "TriggerSchema",

zenml/zen_stores/schemas/pipeline_build_schemas.py

@@ -84,7 +84,6 @@ class PipelineBuildSchema(BaseSchema, table=True):
         back_populates="builds"
     )
 
-    template_deployment_id: Optional[UUID] = None
     images: str = Field(
         sa_column=Column(
             String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
@@ -100,6 +99,7 @@ class PipelineBuildSchema(BaseSchema, table=True):
     zenml_version: Optional[str]
     python_version: Optional[str]
     checksum: Optional[str]
+    stack_checksum: Optional[str]
 
     @classmethod
     def from_request(
@@ -124,7 +124,7 @@ class PipelineBuildSchema(BaseSchema, table=True):
             zenml_version=request.zenml_version,
             python_version=request.python_version,
             checksum=request.checksum,
-            template_deployment_id=request.template_deployment_id,
+            stack_checksum=request.stack_checksum,
         )
 
     def to_model(
@@ -159,9 +159,9 @@ class PipelineBuildSchema(BaseSchema, table=True):
                 zenml_version=self.zenml_version,
                 python_version=self.python_version,
                 checksum=self.checksum,
+                stack_checksum=self.stack_checksum,
                 is_local=self.is_local,
                 contains_code=self.contains_code,
-                template_deployment_id=self.template_deployment_id,
             )
         return PipelineBuildResponse(
             id=self.id,

zenml/zen_stores/schemas/pipeline_deployment_schemas.py

@@ -22,6 +22,7 @@ from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlmodel import Field, Relationship
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
+from zenml.config.pipeline_spec import PipelineSpec
 from zenml.config.step_configurations import Step
 from zenml.constants import MEDIUMTEXT_MAX_LENGTH
 from zenml.models import (
@@ -74,6 +75,16 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
     run_name_template: str = Field(nullable=False)
     client_version: str = Field(nullable=True)
     server_version: str = Field(nullable=True)
+    pipeline_version_hash: Optional[str] = Field(nullable=True, default=None)
+    pipeline_spec: Optional[str] = Field(
+        sa_column=Column(
+            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
+                MEDIUMTEXT, "mysql"
+            ),
+            nullable=True,
+        )
+    )
+    code_path: Optional[str] = Field(nullable=True)
 
     # Foreign keys
     user_id: Optional[UUID] = build_foreign_key_field(
@@ -132,12 +143,15 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         ondelete="SET NULL",
         nullable=True,
     )
+    # This is not a foreign key to remove a cycle which messes with our DB
+    # backup process
+    template_id: Optional[UUID] = None
 
     # SQLModel Relationships
     user: Optional["UserSchema"] = Relationship()
     workspace: "WorkspaceSchema" = Relationship()
-    stack: "StackSchema" = Relationship()
-    pipeline: "PipelineSchema" = Relationship()
+    stack: Optional["StackSchema"] = Relationship()
+    pipeline: Optional["PipelineSchema"] = Relationship()
     schedule: Optional["ScheduleSchema"] = Relationship()
     build: Optional["PipelineBuildSchema"] = Relationship(
         sa_relationship_kwargs={
@@ -176,6 +190,7 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             build_id=request.build,
             user_id=request.user,
             schedule_id=request.schedule,
+            template_id=request.template,
             code_reference_id=code_reference_id,
             run_name_template=request.run_name_template,
             pipeline_configuration=request.pipeline_configuration.model_dump_json(),
@@ -187,6 +202,13 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
             client_environment=json.dumps(request.client_environment),
             client_version=request.client_version,
             server_version=request.server_version,
+            pipeline_version_hash=request.pipeline_version_hash,
+            pipeline_spec=json.dumps(
+                request.pipeline_spec.model_dump(mode="json"), sort_keys=True
+            )
+            if request.pipeline_spec
+            else None,
+            code_path=request.code_path,
         )
 
     def to_model(
@@ -235,6 +257,14 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
                 code_reference=self.code_reference.to_model()
                 if self.code_reference
                 else None,
+                pipeline_version_hash=self.pipeline_version_hash,
+                pipeline_spec=PipelineSpec.model_validate_json(
+                    self.pipeline_spec
+                )
+                if self.pipeline_spec
+                else None,
+                code_path=self.code_path,
+                template_id=self.template_id,
             )
         return PipelineDeploymentResponse(
             id=self.id,

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -22,7 +22,11 @@ from sqlalchemy import UniqueConstraint
 from sqlmodel import TEXT, Column, Field, Relationship
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
-from zenml.enums import ExecutionStatus, MetadataResourceTypes
+from zenml.enums import (
+    ExecutionStatus,
+    MetadataResourceTypes,
+    TaggableResourceTypes,
+)
 from zenml.models import (
     PipelineRunRequest,
     PipelineRunResponse,
@@ -52,6 +56,7 @@ if TYPE_CHECKING:
     from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
     from zenml.zen_stores.schemas.service_schemas import ServiceSchema
     from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
+    from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 
 
 class PipelineRunSchema(NamedSchema, table=True):
@@ -187,6 +192,13 @@ class PipelineRunSchema(NamedSchema, table=True):
     services: List["ServiceSchema"] = Relationship(
         back_populates="pipeline_run",
     )
+    tags: List["TagResourceSchema"] = Relationship(
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.PIPELINE_RUN.value}', foreign(TagResourceSchema.resource_id)==PipelineRunSchema.id)",
+            cascade="delete",
+            overlaps="tags",
+        ),
+    )
 
     @classmethod
     def from_request(
@@ -310,15 +322,29 @@ class PipelineRunSchema(NamedSchema, table=True):
                 client_environment=client_environment,
                 orchestrator_environment=orchestrator_environment,
                 orchestrator_run_id=self.orchestrator_run_id,
+                code_path=self.deployment.code_path
+                if self.deployment
+                else None,
+                template_id=self.deployment.template_id
+                if self.deployment
+                else None,
             )
 
         resources = None
         if include_resources:
             model_version = None
             if config.model and config.model.model_version_id:
-                model_version = config.model._get_model_version(hydrate=False)
+                try:
+                    model_version = config.model._get_model_version(
+                        hydrate=False
+                    )
+                except KeyError:
+                    # Unable to find the model version, it was probably deleted
+                    pass
+
             resources = PipelineRunResponseResources(
-                model_version=model_version
+                model_version=model_version,
+                tags=[t.tag.to_model() for t in self.tags],
             )
 
         return PipelineRunResponse(

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -13,22 +13,20 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Pipelines and Pipeline Runs."""
 
-import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column, String
-from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlalchemy import TEXT, Column
 from sqlmodel import Field, Relationship
 
-from zenml.config.pipeline_spec import PipelineSpec
-from zenml.constants import MEDIUMTEXT_MAX_LENGTH
+from zenml.enums import TaggableResourceTypes
 from zenml.models import (
     PipelineRequest,
     PipelineResponse,
     PipelineResponseBody,
     PipelineResponseMetadata,
+    PipelineResponseResources,
     PipelineUpdate,
 )
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
@@ -45,6 +43,7 @@ if TYPE_CHECKING:
     )
     from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
     from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
+    from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 
 
 class PipelineSchema(NamedSchema, table=True):
@@ -53,17 +52,7 @@ class PipelineSchema(NamedSchema, table=True):
     __tablename__ = "pipeline"
 
     # Fields
-    version: str
-    version_hash: str
-    docstring: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
-    spec: str = Field(
-        sa_column=Column(
-            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
-                MEDIUMTEXT, "mysql"
-            ),
-            nullable=False,
-        )
-    )
+    description: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
 
     # Foreign keys
     workspace_id: UUID = build_foreign_key_field(
@@ -90,13 +79,23 @@ class PipelineSchema(NamedSchema, table=True):
     schedules: List["ScheduleSchema"] = Relationship(
         back_populates="pipeline",
     )
-    runs: List["PipelineRunSchema"] = Relationship(back_populates="pipeline")
+    runs: List["PipelineRunSchema"] = Relationship(
+        back_populates="pipeline",
+        sa_relationship_kwargs={"order_by": "PipelineRunSchema.created"},
+    )
     builds: List["PipelineBuildSchema"] = Relationship(
         back_populates="pipeline"
     )
     deployments: List["PipelineDeploymentSchema"] = Relationship(
         back_populates="pipeline",
     )
+    tags: List["TagResourceSchema"] = Relationship(
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.PIPELINE.value}', foreign(TagResourceSchema.resource_id)==PipelineSchema.id)",
+            cascade="delete",
+            overlaps="tags",
+        ),
+    )
 
     @classmethod
     def from_request(
@@ -113,21 +112,15 @@ class PipelineSchema(NamedSchema, table=True):
         """
         return cls(
             name=pipeline_request.name,
-            version=pipeline_request.version,
-            version_hash=pipeline_request.version_hash,
+            description=pipeline_request.description,
             workspace_id=pipeline_request.workspace,
             user_id=pipeline_request.user,
-            docstring=pipeline_request.docstring,
-            spec=json.dumps(
-                pipeline_request.spec.model_dump(mode="json"), sort_keys=True
-            ),
         )
 
     def to_model(
         self,
         include_metadata: bool = False,
         include_resources: bool = False,
-        last_x_runs: int = 3,
         **kwargs: Any,
     ) -> "PipelineResponse":
         """Convert a `PipelineSchema` to a `PipelineResponse`.
@@ -136,25 +129,29 @@ class PipelineSchema(NamedSchema, table=True):
             include_metadata: Whether the metadata will be filled.
             include_resources: Whether the resources will be filled.
             **kwargs: Keyword arguments to allow schema specific logic
-            last_x_runs: How many runs to use for the execution status
 
         Returns:
             The created PipelineResponse.
         """
         body = PipelineResponseBody(
             user=self.user.to_model() if self.user else None,
-            status=[run.status for run in self.runs[:last_x_runs]],
+            latest_run_id=self.runs[-1].id if self.runs else None,
+            latest_run_status=self.runs[-1].status if self.runs else None,
             created=self.created,
             updated=self.updated,
-            version=self.version,
         )
+
         metadata = None
         if include_metadata:
             metadata = PipelineResponseMetadata(
                 workspace=self.workspace.to_model(),
-                version_hash=self.version_hash,
-                spec=PipelineSpec.model_validate_json(self.spec),
-                docstring=self.docstring,
+                description=self.description,
+            )
+
+        resources = None
+        if include_resources:
+            resources = PipelineResponseResources(
+                tags=[t.tag.to_model() for t in self.tags],
             )
 
         return PipelineResponse(
@@ -162,6 +159,7 @@ class PipelineSchema(NamedSchema, table=True):
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )
 
     def update(self, pipeline_update: "PipelineUpdate") -> "PipelineSchema":
@@ -173,5 +171,6 @@ class PipelineSchema(NamedSchema, table=True):
         Returns:
             The updated `PipelineSchema`.
         """
+        self.description = pipeline_update.description
         self.updated = datetime.utcnow()
         return self