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

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

zenml/zen_stores/schemas/run_template_schemas.py

@@ -0,0 +1,264 @@
+#  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.
+"""SQLModel implementation of run template tables."""
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, List, Optional
+from uuid import UUID
+
+from sqlalchemy import Column, String, UniqueConstraint
+from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlmodel import Field, Relationship
+
+from zenml.constants import MEDIUMTEXT_MAX_LENGTH
+from zenml.enums import TaggableResourceTypes
+from zenml.models import (
+    RunTemplateRequest,
+    RunTemplateResponse,
+    RunTemplateResponseBody,
+    RunTemplateResponseMetadata,
+    RunTemplateResponseResources,
+    RunTemplateUpdate,
+)
+from zenml.zen_stores.schemas.base_schemas import BaseSchema
+from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
+from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
+
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
+        PipelineDeploymentSchema,
+    )
+    from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
+    from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
+
+
+class RunTemplateSchema(BaseSchema, table=True):
+    """SQL Model for run templates."""
+
+    __tablename__ = "run_template"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_template_name_in_workspace",
+        ),
+    )
+
+    name: str = Field(nullable=False)
+    description: Optional[str] = Field(
+        sa_column=Column(
+            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
+                MEDIUMTEXT, "mysql"
+            ),
+            nullable=True,
+        )
+    )
+
+    user_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target=UserSchema.__tablename__,
+        source_column="user_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
+    workspace_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=WorkspaceSchema.__tablename__,
+        source_column="workspace_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )
+    source_deployment_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target="pipeline_deployment",
+        source_column="source_deployment_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
+
+    user: Optional["UserSchema"] = Relationship()
+    workspace: "WorkspaceSchema" = Relationship()
+    source_deployment: Optional["PipelineDeploymentSchema"] = Relationship(
+        sa_relationship_kwargs={
+            "foreign_keys": "RunTemplateSchema.source_deployment_id",
+        }
+    )
+
+    runs: List["PipelineRunSchema"] = Relationship(
+        sa_relationship_kwargs={
+            "primaryjoin": "RunTemplateSchema.id==PipelineDeploymentSchema.template_id",
+            "secondaryjoin": "PipelineDeploymentSchema.id==PipelineRunSchema.deployment_id",
+            "secondary": "pipeline_deployment",
+            "cascade": "delete",
+            "viewonly": True,
+            "order_by": "PipelineRunSchema.created",
+        }
+    )
+
+    tags: List["TagResourceSchema"] = Relationship(
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.RUN_TEMPLATE.value}', foreign(TagResourceSchema.resource_id)==RunTemplateSchema.id)",
+            cascade="delete",
+            overlaps="tags",
+        ),
+    )
+
+    @classmethod
+    def from_request(
+        cls,
+        request: RunTemplateRequest,
+    ) -> "RunTemplateSchema":
+        """Create a schema from a request.
+
+        Args:
+            request: The request to convert.
+
+
+        Returns:
+            The created schema.
+        """
+        return cls(
+            user_id=request.user,
+            workspace_id=request.workspace,
+            name=request.name,
+            description=request.description,
+            source_deployment_id=request.source_deployment_id,
+        )
+
+    def update(self, update: RunTemplateUpdate) -> "RunTemplateSchema":
+        """Update the schema.
+
+        Args:
+            update: The update model.
+
+        Returns:
+            The updated schema.
+        """
+        for field, value in update.model_dump(
+            exclude_unset=True, exclude_none=True
+        ).items():
+            setattr(self, field, value)
+
+        self.updated = datetime.utcnow()
+        return self
+
+    def to_model(
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> RunTemplateResponse:
+        """Convert the schema to a response model.
+
+        Args:
+            include_metadata: Whether the metadata will be filled.
+            include_resources: Whether the resources will be filled.
+            **kwargs: Keyword arguments to allow schema specific logic
+
+        Returns:
+            Model representing this schema.
+        """
+        runnable = False
+        if (
+            self.source_deployment
+            and self.source_deployment.build
+            and self.source_deployment.build.stack
+        ):
+            runnable = True
+
+        body = RunTemplateResponseBody(
+            user=self.user.to_model() if self.user else None,
+            created=self.created,
+            updated=self.updated,
+            runnable=runnable,
+            latest_run_id=self.runs[-1].id if self.runs else None,
+            latest_run_status=self.runs[-1].status if self.runs else None,
+        )
+
+        metadata = None
+        if include_metadata:
+            pipeline_spec = None
+            config_template = None
+            config_schema = None
+
+            if self.source_deployment:
+                from zenml.zen_stores import template_utils
+
+                pipeline_spec = self.source_deployment.to_model(
+                    include_metadata=True, include_resources=True
+                ).pipeline_spec
+
+                if (
+                    self.source_deployment.build
+                    and self.source_deployment.build.stack
+                ):
+                    config_template = template_utils.generate_config_template(
+                        deployment=self.source_deployment
+                    )
+                    config_schema = template_utils.generate_config_schema(
+                        deployment=self.source_deployment
+                    )
+
+            metadata = RunTemplateResponseMetadata(
+                workspace=self.workspace.to_model(),
+                description=self.description,
+                pipeline_spec=pipeline_spec,
+                config_template=config_template,
+                config_schema=config_schema,
+            )
+
+        resources = None
+        if include_resources:
+            if self.source_deployment:
+                pipeline = (
+                    self.source_deployment.pipeline.to_model()
+                    if self.source_deployment.pipeline
+                    else None
+                )
+                build = (
+                    self.source_deployment.build.to_model()
+                    if self.source_deployment.build
+                    else None
+                )
+                code_reference = (
+                    self.source_deployment.code_reference.to_model()
+                    if self.source_deployment.code_reference
+                    else None
+                )
+            else:
+                pipeline = None
+                build = None
+                code_reference = None
+
+            resources = RunTemplateResponseResources(
+                source_deployment=self.source_deployment.to_model()
+                if self.source_deployment
+                else None,
+                pipeline=pipeline,
+                build=build,
+                code_reference=code_reference,
+                tags=[t.tag.to_model() for t in self.tags],
+            )
+
+        return RunTemplateResponse(
+            id=self.id,
+            name=self.name,
+            body=body,
+            metadata=metadata,
+            resources=resources,
+        )

zenml/zen_stores/schemas/step_run_schemas.py

@@ -274,11 +274,18 @@ class StepRunSchema(NamedSchema, table=True):
         if include_resources:
             model_version = None
             if full_step_config.config.model:
-                model_version = (
-                    full_step_config.config.model._get_model_version(
-                        hydrate=False
+                # TODO: Why is there no ID check similar to
+                # PipelineRunSchema.to_model()?
+                try:
+                    model_version = (
+                        full_step_config.config.model._get_model_version(
+                            hydrate=False
+                        )
                     )
-                )
+                except KeyError:
+                    # Unable to find the model version, it was probably deleted
+                    pass
+
             resources = StepRunResponseResources(model_version=model_version)
 
         return StepRunResponse(