zenml-nightly 0.71.0.dev20241212__py3-none-any.whl → 0.71.0.dev20241214__py3-none-any.whl

zenml/zen_stores/schemas/model_schemas.py

@@ -15,10 +15,16 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from pydantic import ConfigDict
-from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column, UniqueConstraint
+from sqlalchemy import (
+    BOOLEAN,
+    INTEGER,
+    TEXT,
+    Column,
+    UniqueConstraint,
+)
 from sqlmodel import Field, Relationship
 
 from zenml.enums import (
@@ -228,11 +234,13 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
 
     __tablename__ = MODEL_VERSION_TABLENAME
     __table_args__ = (
-        # We need two unique constraints here:
+        # We need three unique constraints here:
         # - The first to ensure that each model version for a
         #   model has a unique version number
         # - The second one to ensure that explicit names given by
         #   users are unique
+        # - The third one to ensure that a pipeline run only produces a single
+        #   auto-incremented version per model
         UniqueConstraint(
             "number",
             "model_id",
@@ -243,6 +251,11 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
             "model_id",
             name="unique_version_for_model_id",
         ),
+        UniqueConstraint(
+            "model_id",
+            "producer_run_id_if_numeric",
+            name="unique_numeric_version_for_pipeline_run",
+        ),
     )
 
     workspace_id: UUID = build_foreign_key_field(
@@ -312,12 +325,23 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
         ),
     )
     pipeline_runs: List["PipelineRunSchema"] = Relationship(
-        back_populates="model_version"
+        back_populates="model_version",
     )
     step_runs: List["StepRunSchema"] = Relationship(
         back_populates="model_version"
     )
 
+    # We want to make sure each pipeline run only creates a single numeric
+    # version for each model. To solve this, we need to add a unique constraint.
+    # If a value of a unique constraint is NULL it is ignored and the
+    # remaining values in the unique constraint have to be unique. In
+    # our case however, we only want the unique constraint applied in
+    # case there is a producer run and only for numeric versions. To solve this,
+    # we fall back to the model version ID (which is the primary key and
+    # therefore unique) in case there is no producer run or the version is not
+    # numeric.
+    producer_run_id_if_numeric: UUID
+
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
     #  It is possible to suppress this warning message with the following
@@ -328,24 +352,36 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
 
     @classmethod
     def from_request(
-        cls, model_version_request: ModelVersionRequest
+        cls,
+        model_version_request: ModelVersionRequest,
+        model_version_number: int,
+        producer_run_id: Optional[UUID] = None,
     ) -> "ModelVersionSchema":
         """Convert an `ModelVersionRequest` to an `ModelVersionSchema`.
 
         Args:
             model_version_request: The request model version to convert.
+            model_version_number: The model version number.
+            producer_run_id: The ID of the producer run.
 
         Returns:
             The converted schema.
         """
+        id_ = uuid4()
+        is_numeric = str(model_version_number) == model_version_request.name
+
         return cls(
+            id=id_,
             workspace_id=model_version_request.workspace,
             user_id=model_version_request.user,
             model_id=model_version_request.model,
             name=model_version_request.name,
-            number=model_version_request.number,
+            number=model_version_number,
             description=model_version_request.description,
             stage=model_version_request.stage,
+            producer_run_id_if_numeric=producer_run_id
+            if (producer_run_id and is_numeric)
+            else id_,
         )
 
     def to_model(

zenml/zen_stores/schemas/pipeline_deployment_schemas.py

@@ -228,13 +228,6 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         Returns:
             The created `PipelineDeploymentResponse`.
         """
-        pipeline_configuration = PipelineConfiguration.model_validate_json(
-            self.pipeline_configuration
-        )
-        step_configurations = json.loads(self.step_configurations)
-        for s, c in step_configurations.items():
-            step_configurations[s] = Step.model_validate(c)
-
         body = PipelineDeploymentResponseBody(
             user=self.user.to_model() if self.user else None,
             created=self.created,
@@ -242,6 +235,13 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         )
         metadata = None
         if include_metadata:
+            pipeline_configuration = PipelineConfiguration.model_validate_json(
+                self.pipeline_configuration
+            )
+            step_configurations = json.loads(self.step_configurations)
+            for s, c in step_configurations.items():
+                step_configurations[s] = Step.model_validate(c)
+
             metadata = PipelineDeploymentResponseMetadata(
                 workspace=self.workspace.to_model(),
                 run_name_template=self.run_name_template,

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -156,7 +156,12 @@ class PipelineSchema(NamedSchema, table=True):
 
         resources = None
         if include_resources:
+            latest_run_user = self.runs[-1].user if self.runs else None
+
             resources = PipelineResponseResources(
+                latest_run_user=latest_run_user.to_model()
+                if latest_run_user
+                else None,
                 tags=[t.tag.to_model() for t in self.tags],
             )
 

zenml/zen_stores/schemas/step_run_schemas.py

@@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
 from pydantic import ConfigDict
-from sqlalchemy import TEXT, Column, String
+from sqlalchemy import TEXT, Column, String, UniqueConstraint
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlmodel import Field, Relationship, SQLModel
 
@@ -67,6 +67,13 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for steps of pipeline runs."""
 
     __tablename__ = "step_run"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "pipeline_run_id",
+            name="unique_step_name_for_pipeline_run",
+        ),
+    )
 
     # Fields
     start_time: Optional[datetime] = Field(nullable=True)