zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__py3-none-any.whl

zenml/zen_stores/schemas/model_schemas.py

@@ -44,6 +44,7 @@ from zenml.models import (
 )
 from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
+from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
@@ -53,7 +54,7 @@ from zenml.zen_stores.schemas.utils import get_page_from_list
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
-    from zenml.zen_stores.schemas import ServiceSchema
+    from zenml.zen_stores.schemas import ServiceSchema, StepRunSchema
 
 
 class ModelSchema(NamedSchema, table=True):
@@ -180,7 +181,6 @@ class ModelSchema(NamedSchema, table=True):
 
         body = ModelResponseBody(
             user=self.user.to_model() if self.user else None,
-            workspace=self.workspace.to_model(),
             created=self.created,
             updated=self.updated,
             tags=tags,
@@ -218,7 +218,7 @@ class ModelSchema(NamedSchema, table=True):
 class ModelVersionSchema(NamedSchema, table=True):
     """SQL Model for model version."""
 
-    __tablename__ = "model_version"
+    __tablename__ = MODEL_VERSION_TABLENAME
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -286,6 +286,12 @@ class ModelVersionSchema(NamedSchema, table=True):
             overlaps="run_metadata",
         ),
     )
+    pipeline_runs: List["PipelineRunSchema"] = Relationship(
+        back_populates="model_version"
+    )
+    step_runs: List["StepRunSchema"] = Relationship(
+        back_populates="model_version"
+    )
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -18,6 +18,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
+from pydantic import ConfigDict
 from sqlalchemy import UniqueConstraint
 from sqlmodel import TEXT, Column, Field, Relationship
 
@@ -36,6 +37,7 @@ from zenml.models import (
 )
 from zenml.models.v2.core.pipeline_run import PipelineRunResponseResources
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
+from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
     PipelineDeploymentSchema,
@@ -52,6 +54,7 @@ if TYPE_CHECKING:
     from zenml.zen_stores.schemas.logs_schemas import LogsSchema
     from zenml.zen_stores.schemas.model_schemas import (
         ModelVersionPipelineRunSchema,
+        ModelVersionSchema,
     )
     from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
     from zenml.zen_stores.schemas.service_schemas import ServiceSchema
@@ -113,6 +116,14 @@ class PipelineRunSchema(NamedSchema, table=True):
         ondelete="SET NULL",
         nullable=True,
     )
+    model_version_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=MODEL_VERSION_TABLENAME,
+        source_column="model_version_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
 
     # Relationships
     deployment: Optional["PipelineDeploymentSchema"] = Relationship(
@@ -141,6 +152,9 @@ class PipelineRunSchema(NamedSchema, table=True):
     step_runs: List["StepRunSchema"] = Relationship(
         sa_relationship_kwargs={"cascade": "delete"},
     )
+    model_version: "ModelVersionSchema" = Relationship(
+        back_populates="pipeline_runs",
+    )
 
     # Temporary fields and foreign keys to be deprecated
     pipeline_configuration: Optional[str] = Field(
@@ -200,6 +214,8 @@ class PipelineRunSchema(NamedSchema, table=True):
         ),
     )
 
+    model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
+
     @classmethod
     def from_request(
         cls, request: "PipelineRunRequest"
@@ -225,6 +241,7 @@ class PipelineRunSchema(NamedSchema, table=True):
             pipeline_id=request.pipeline,
             deployment_id=request.deployment,
             trigger_execution_id=request.trigger_execution_id,
+            model_version_id=request.model_version_id,
         )
 
     def to_model(
@@ -307,6 +324,7 @@ class PipelineRunSchema(NamedSchema, table=True):
             created=self.created,
             updated=self.updated,
             deployment_id=self.deployment_id,
+            model_version_id=self.model_version_id,
         )
         metadata = None
         if include_metadata:
@@ -333,14 +351,8 @@ class PipelineRunSchema(NamedSchema, table=True):
         resources = None
         if include_resources:
             model_version = None
-            if config.model and config.model.model_version_id:
-                try:
-                    model_version = config.model._get_model_version(
-                        hydrate=False
-                    )
-                except KeyError:
-                    # Unable to find the model version, it was probably deleted
-                    pass
+            if self.model_version:
+                model_version = self.model_version.to_model()
 
             resources = PipelineRunResponseResources(
                 model_version=model_version,
@@ -367,6 +379,8 @@ class PipelineRunSchema(NamedSchema, table=True):
         if run_update.status:
             self.status = run_update.status.value
             self.end_time = run_update.end_time
+        if run_update.model_version_id and self.model_version_id is None:
+            self.model_version_id = run_update.model_version_id
 
         self.updated = datetime.utcnow()
         return self

zenml/zen_stores/schemas/step_run_schemas.py

@@ -18,6 +18,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
+from pydantic import ConfigDict
 from sqlalchemy import TEXT, Column, String
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlmodel import Field, Relationship, SQLModel
@@ -37,6 +38,7 @@ from zenml.models import (
 )
 from zenml.models.v2.core.step_run import StepRunResponseResources
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
+from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
     PipelineDeploymentSchema,
 )
@@ -48,6 +50,7 @@ from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
     from zenml.zen_stores.schemas.logs_schemas import LogsSchema
+    from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
     from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
 
 
@@ -116,6 +119,14 @@ class StepRunSchema(NamedSchema, table=True):
         ondelete="CASCADE",
         nullable=False,
     )
+    model_version_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=MODEL_VERSION_TABLENAME,
+        source_column="model_version_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
 
     # Relationships
     workspace: "WorkspaceSchema" = Relationship(back_populates="step_runs")
@@ -147,6 +158,11 @@ class StepRunSchema(NamedSchema, table=True):
             "primaryjoin": "StepRunParentsSchema.child_id == StepRunSchema.id",
         },
     )
+    model_version: "ModelVersionSchema" = Relationship(
+        back_populates="step_runs",
+    )
+
+    model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
     @classmethod
     def from_request(cls, request: StepRunRequest) -> "StepRunSchema":
@@ -172,6 +188,7 @@ class StepRunSchema(NamedSchema, table=True):
             cache_key=request.cache_key,
             code_hash=request.code_hash,
             source_code=request.source_code,
+            model_version_id=request.model_version_id,
         )
 
     def to_model(
@@ -249,6 +266,7 @@ class StepRunSchema(NamedSchema, table=True):
             outputs=output_artifacts,
             created=self.created,
             updated=self.updated,
+            model_version_id=self.model_version_id,
         )
         metadata = None
         if include_metadata:
@@ -273,18 +291,8 @@ class StepRunSchema(NamedSchema, table=True):
         resources = None
         if include_resources:
             model_version = None
-            if full_step_config.config.model:
-                # 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
+            if self.model_version:
+                model_version = self.model_version.to_model()
 
             resources = StepRunResponseResources(model_version=model_version)
 
@@ -312,6 +320,9 @@ class StepRunSchema(NamedSchema, table=True):
                 self.status = value.value
             if key == "end_time":
                 self.end_time = value
+            if key == "model_version_id":
+                if value and self.model_version_id is None:
+                    self.model_version_id = value
 
         self.updated = datetime.utcnow()