zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.63.0.dev20240731__py3-none-any.whl

zenml/zen_stores/migrations/versions/b59aa68fdb1f_simplify_pipelines.py

@@ -0,0 +1,139 @@
+"""Simplify pipelines [b59aa68fdb1f].
+
+Revision ID: b59aa68fdb1f
+Revises: 0.62.0
+Create Date: 2024-07-04 14:00:32.830722
+
+"""
+
+from typing import Dict, Optional
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = "b59aa68fdb1f"
+down_revision = "0.62.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "pipeline_version_hash",
+                sqlmodel.sql.sqltypes.AutoString(),
+                nullable=True,
+            )
+        )
+        batch_op.add_column(
+            sa.Column(
+                "pipeline_spec",
+                sa.String(length=16777215).with_variant(
+                    mysql.MEDIUMTEXT, "mysql"
+                ),
+                nullable=True,
+            )
+        )
+
+    connection = op.get_bind()
+    meta = sa.MetaData()
+    meta.reflect(
+        bind=connection,
+        only=(
+            "pipeline",
+            "pipeline_run",
+            "pipeline_deployment",
+            "pipeline_build",
+            "schedule",
+        ),
+    )
+    pipeline_table = sa.Table("pipeline", meta)
+    pipeline_run_table = sa.Table("pipeline_run", meta)
+    pipeline_deployment_table = sa.Table("pipeline_deployment", meta)
+    pipeline_build_table = sa.Table("pipeline_build", meta)
+    schedule_table = sa.Table("schedule", meta)
+
+    def _migrate_pipeline_columns(
+        pipeline_id: str,
+        version_hash: Optional[str],
+        pipeline_spec: Optional[str],
+    ) -> None:
+        connection.execute(
+            sa.update(pipeline_deployment_table)
+            .where(pipeline_deployment_table.c.pipeline_id == pipeline_id)
+            .values(
+                pipeline_version_hash=version_hash, pipeline_spec=pipeline_spec
+            )
+        )
+
+    def _update_pipeline_fks(pipeline_id: str, replacement_id: str) -> None:
+        for table in [
+            pipeline_run_table,
+            pipeline_deployment_table,
+            pipeline_build_table,
+            schedule_table,
+        ]:
+            connection.execute(
+                sa.update(table)
+                .where(table.c.pipeline_id == pipeline_id)
+                .values(pipeline_id=replacement_id)
+            )
+
+    all_pipelines = connection.execute(sa.select(pipeline_table)).fetchall()
+    replacement_mapping: Dict[str, str] = {}
+
+    for pipeline in all_pipelines:
+        _migrate_pipeline_columns(
+            pipeline_id=pipeline.id,
+            version_hash=pipeline.version_hash,
+            pipeline_spec=pipeline.spec,
+        )
+
+        if replacement_id := replacement_mapping.get(pipeline.name):
+            _update_pipeline_fks(
+                pipeline_id=pipeline.id, replacement_id=replacement_id
+            )
+            connection.execute(
+                sa.delete(pipeline_table).where(
+                    pipeline_table.c.id == pipeline.id
+                )
+            )
+        else:
+            replacement_mapping[pipeline.name] = pipeline.id
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("description", sa.TEXT(), nullable=True))
+        batch_op.drop_column("spec")
+        batch_op.drop_column("docstring")
+        batch_op.drop_column("version_hash")
+        batch_op.drop_column("version")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
+        batch_op.drop_column("pipeline_spec")
+        batch_op.drop_column("pipeline_version_hash")
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("version", sa.VARCHAR(), nullable=False))
+        batch_op.add_column(
+            sa.Column("version_hash", sa.VARCHAR(), nullable=False)
+        )
+        batch_op.add_column(sa.Column("docstring", sa.TEXT(), nullable=True))
+        batch_op.add_column(
+            sa.Column("spec", sa.VARCHAR(length=16777215), nullable=False)
+        )
+        batch_op.drop_column("description")
+
+    # ### end Alembic commands ###

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,
+        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,
-        deployment_id: UUID,
+        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)

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,15 @@ 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,
+        )
+    )
 
     # Foreign keys
     user_id: Optional[UUID] = build_foreign_key_field(
@@ -132,12 +142,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 +189,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 +201,12 @@ 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,
         )
 
     def to_model(
@@ -235,6 +255,13 @@ 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,
+                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,26 @@ class PipelineRunSchema(NamedSchema, table=True):
                 client_environment=client_environment,
                 orchestrator_environment=orchestrator_environment,
                 orchestrator_run_id=self.orchestrator_run_id,
+                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