zenml-nightly 0.70.0.dev20241128__py3-none-any.whl → 0.70.0.dev20241130__py3-none-any.whl

zenml/zen_stores/schemas/step_run_schemas.py

@@ -23,6 +23,7 @@ from sqlalchemy import TEXT, Column, String
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlmodel import Field, Relationship, SQLModel
 
+from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.step_configurations import Step
 from zenml.constants import MEDIUMTEXT_MAX_LENGTH
 from zenml.enums import (
@@ -50,16 +51,19 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 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.utils import RunMetadataInterface
 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
+    from zenml.zen_stores.schemas.run_metadata_schemas import (
+        RunMetadataResourceSchema,
+    )
 
 
-class StepRunSchema(NamedSchema, table=True):
+class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for steps of pipeline runs."""
 
     __tablename__ = "step_run"
@@ -139,12 +143,12 @@ class StepRunSchema(NamedSchema, table=True):
     deployment: Optional["PipelineDeploymentSchema"] = Relationship(
         back_populates="step_runs"
     )
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="step_run",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="step_runs",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataSchema.resource_id)==StepRunSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     input_artifacts: List["StepRunInputArtifactSchema"] = Relationship(
@@ -163,9 +167,15 @@ class StepRunSchema(NamedSchema, table=True):
             "primaryjoin": "StepRunParentsSchema.child_id == StepRunSchema.id",
         },
     )
+    pipeline_run: "PipelineRunSchema" = Relationship(
+        back_populates="step_runs"
+    )
     model_version: "ModelVersionSchema" = Relationship(
         back_populates="step_runs",
     )
+    original_step_run: Optional["StepRunSchema"] = Relationship(
+        sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
+    )
 
     model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
@@ -218,11 +228,6 @@ class StepRunSchema(NamedSchema, table=True):
             RuntimeError: If the step run schema does not have a deployment_id
                 or a step_configuration.
         """
-        run_metadata = {
-            metadata_schema.key: json.loads(metadata_schema.value)
-            for metadata_schema in self.run_metadata
-        }
-
         input_artifacts = {
             artifact.name: StepRunInputResponse(
                 input_type=StepRunInputArtifactType(artifact.type),
@@ -248,6 +253,21 @@ class StepRunSchema(NamedSchema, table=True):
                 full_step_config = Step.model_validate(
                     step_configuration[self.name]
                 )
+                new_substitutions = (
+                    full_step_config.config._get_full_substitutions(
+                        PipelineConfiguration.model_validate_json(
+                            self.deployment.pipeline_configuration
+                        ),
+                        self.pipeline_run.start_time,
+                    )
+                )
+                full_step_config = full_step_config.model_copy(
+                    update={
+                        "config": full_step_config.config.model_copy(
+                            update={"substitutions": new_substitutions}
+                        )
+                    }
+                )
             elif not self.step_configuration:
                 raise ValueError(
                     f"Unable to load the configuration for step `{self.name}` from the"
@@ -294,7 +314,7 @@ class StepRunSchema(NamedSchema, table=True):
                 pipeline_run_id=self.pipeline_run_id,
                 original_step_run_id=self.original_step_run_id,
                 parent_step_ids=[p.parent_id for p in self.parents],
-                run_metadata=run_metadata,
+                run_metadata=self.fetch_metadata(),
             )
 
         resources = None

zenml/zen_stores/schemas/utils.py

@@ -13,11 +13,14 @@
 #  permissions and limitations under the License.
 """Utils for schemas."""
 
+import json
 import math
-from typing import List, Type, TypeVar
+from typing import Dict, List, Type, TypeVar
 
-from zenml.models.v2.base.base import BaseResponse
-from zenml.models.v2.base.page import Page
+from sqlmodel import Relationship
+
+from zenml.metadata.metadata_types import MetadataType
+from zenml.models import BaseResponse, Page, RunMetadataEntry
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 
 S = TypeVar("S", bound=BaseSchema)
@@ -67,3 +70,44 @@ def get_page_from_list(
         total=total,
         items=page_items,
     )
+
+
+class RunMetadataInterface:
+    """The interface for entities with run metadata."""
+
+    run_metadata_resources = Relationship()
+
+    def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
+        """Fetches all the metadata entries related to the artifact version.
+
+        Returns:
+            a dictionary, where the key is the key of the metadata entry
+                and the values represent the list of entries with this key.
+        """
+        metadata_collection: Dict[str, List[RunMetadataEntry]] = {}
+
+        # Fetch the metadata related to this step
+        for rm in self.run_metadata_resources:
+            if rm.run_metadata.key not in metadata_collection:
+                metadata_collection[rm.run_metadata.key] = []
+            metadata_collection[rm.run_metadata.key].append(
+                RunMetadataEntry(
+                    value=json.loads(rm.run_metadata.value),
+                    created=rm.run_metadata.created,
+                )
+            )
+
+        return metadata_collection
+
+    def fetch_metadata(self) -> Dict[str, MetadataType]:
+        """Fetches the latest metadata entry related to the artifact version.
+
+        Returns:
+            a dictionary, where the key is the key of the metadata entry
+                and the values represent the latest entry with this key.
+        """
+        metadata_collection = self.fetch_metadata_collection()
+        return {
+            k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
+            for k, v in metadata_collection.items()
+        }

zenml/zen_stores/sql_zen_store.py

@@ -26,6 +26,7 @@ from datetime import datetime, timezone
 from functools import lru_cache
 from pathlib import Path
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     ClassVar,
@@ -219,6 +220,7 @@ from zenml.models import (
     PipelineRunUpdate,
     PipelineUpdate,
     RunMetadataRequest,
+    RunMetadataResource,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -325,6 +327,7 @@ from zenml.zen_stores.schemas import (
     PipelineDeploymentSchema,
     PipelineRunSchema,
     PipelineSchema,
+    RunMetadataResourceSchema,
     RunMetadataSchema,
     RunTemplateSchema,
     ScheduleSchema,
@@ -354,6 +357,9 @@ from zenml.zen_stores.secrets_stores.sql_secrets_store import (
     SqlSecretsStoreConfiguration,
 )
 
+if TYPE_CHECKING:
+    from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
+
 AnyNamedSchema = TypeVar("AnyNamedSchema", bound=NamedSchema)
 AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
@@ -2726,7 +2732,9 @@ class SqlZenStore(BaseZenStore):
     # -------------------- Artifact Versions --------------------
 
     def _get_or_create_artifact_for_name(
-        self, name: str, has_custom_name: bool
+        self,
+        name: str,
+        has_custom_name: bool,
     ) -> ArtifactSchema:
         """Get or create an artifact with a specific name.
 
@@ -2747,7 +2755,8 @@ class SqlZenStore(BaseZenStore):
                 try:
                     with session.begin_nested():
                         artifact_request = ArtifactRequest(
-                            name=name, has_custom_name=has_custom_name
+                            name=name,
+                            has_custom_name=has_custom_name,
                         )
                         artifact = ArtifactSchema.from_request(
                             artifact_request
@@ -2915,17 +2924,41 @@ class SqlZenStore(BaseZenStore):
 
             # Save metadata of the artifact
             if artifact_version.metadata:
+                values: Dict[str, "MetadataType"] = {}
+                types: Dict[str, "MetadataTypeEnum"] = {}
                 for key, value in artifact_version.metadata.items():
-                    run_metadata_schema = RunMetadataSchema(
-                        workspace_id=artifact_version.workspace,
-                        user_id=artifact_version.user,
-                        resource_id=artifact_version_id,
-                        resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
-                        key=key,
-                        value=json.dumps(value),
-                        type=get_metadata_type(value),
+                    # Skip metadata that is too large to be stored in the DB.
+                    if len(json.dumps(value)) > TEXT_FIELD_MAX_LENGTH:
+                        logger.warning(
+                            f"Metadata value for key '{key}' is too large to be "
+                            "stored in the database. Skipping."
+                        )
+                        continue
+                    # Skip metadata that is not of a supported type.
+                    try:
+                        metadata_type = get_metadata_type(value)
+                    except ValueError as e:
+                        logger.warning(
+                            f"Metadata value for key '{key}' is not of a "
+                            f"supported type. Skipping. Full error: {e}"
+                        )
+                        continue
+                    values[key] = value
+                    types[key] = metadata_type
+                self.create_run_metadata(
+                    RunMetadataRequest(
+                        workspace=artifact_version.workspace,
+                        user=artifact_version.user,
+                        resources=[
+                            RunMetadataResource(
+                                id=artifact_version_id,
+                                type=MetadataResourceTypes.ARTIFACT_VERSION,
+                            )
+                        ],
+                        values=values,
+                        types=types,
                     )
-                    session.add(run_metadata_schema)
+                )
 
             session.commit()
             artifact_version_schema = session.exec(
@@ -5529,20 +5562,29 @@ class SqlZenStore(BaseZenStore):
             The created run metadata.
         """
         with Session(self.engine) as session:
-            for key, value in run_metadata.values.items():
-                type_ = run_metadata.types[key]
-                run_metadata_schema = RunMetadataSchema(
-                    workspace_id=run_metadata.workspace,
-                    user_id=run_metadata.user,
-                    resource_id=run_metadata.resource_id,
-                    resource_type=run_metadata.resource_type.value,
-                    stack_component_id=run_metadata.stack_component_id,
-                    key=key,
-                    value=json.dumps(value),
-                    type=type_,
-                )
-                session.add(run_metadata_schema)
-                session.commit()
+            if run_metadata.resources:
+                for key, value in run_metadata.values.items():
+                    type_ = run_metadata.types[key]
+                    run_metadata_schema = RunMetadataSchema(
+                        workspace_id=run_metadata.workspace,
+                        user_id=run_metadata.user,
+                        stack_component_id=run_metadata.stack_component_id,
+                        key=key,
+                        value=json.dumps(value),
+                        type=type_,
+                        publisher_step_id=run_metadata.publisher_step_id,
+                    )
+                    session.add(run_metadata_schema)
+                    session.commit()
+
+                    for resource in run_metadata.resources:
+                        rm_resource_link = RunMetadataResourceSchema(
+                            resource_id=resource.id,
+                            resource_type=resource.type.value,
+                            run_metadata_id=run_metadata_schema.id,
+                        )
+                        session.add(rm_resource_link)
+                        session.commit()
         return None
 
     # ----------------------------- Schedules -----------------------------
@@ -8153,6 +8195,46 @@ class SqlZenStore(BaseZenStore):
                 )
                 session.add(log_entry)
 
+            # If cached, attach metadata of the original step
+            if (
+                step_run.status == ExecutionStatus.CACHED
+                and step_run.original_step_run_id is not None
+            ):
+                original_metadata_links = session.exec(
+                    select(RunMetadataResourceSchema)
+                    .where(
+                        RunMetadataResourceSchema.run_metadata_id
+                        == RunMetadataSchema.id
+                    )
+                    .where(
+                        RunMetadataResourceSchema.resource_id
+                        == step_run.original_step_run_id
+                    )
+                    .where(
+                        RunMetadataResourceSchema.resource_type
+                        == MetadataResourceTypes.STEP_RUN
+                    )
+                    .where(
+                        RunMetadataSchema.publisher_step_id
+                        == step_run.original_step_run_id
+                    )
+                ).all()
+
+                # Create new links in a batch
+                new_links = [
+                    RunMetadataResourceSchema(
+                        resource_id=step_schema.id,
+                        resource_type=link.resource_type,
+                        run_metadata_id=link.run_metadata_id,
+                    )
+                    for link in original_metadata_links
+                ]
+                # Add all new links in a single operation
+                session.add_all(new_links)
+                # Commit the changes
+                session.commit()
+                session.refresh(step_schema)
+
             # Save parent step IDs into the database.
             for parent_step_id in step_run.parent_step_ids:
                 self._set_run_step_parent_step(
@@ -8182,12 +8264,12 @@ class SqlZenStore(BaseZenStore):
                 )
 
             # Save output artifact IDs into the database.
-            for output_name, artifact_version_ids in step_run.outputs.items():
+            for name, artifact_version_ids in step_run.outputs.items():
                 for artifact_version_id in artifact_version_ids:
                     self._set_run_step_output_artifact(
                         step_run_id=step_schema.id,
                         artifact_version_id=artifact_version_id,
-                        name=output_name,
+                        name=name,
                         session=session,
                     )
 
@@ -8291,13 +8373,14 @@ class SqlZenStore(BaseZenStore):
             session.add(existing_step_run)
 
             # Update the artifacts.
-            for name, artifact_version_id in step_run_update.outputs.items():
-                self._set_run_step_output_artifact(
-                    step_run_id=step_run_id,
-                    artifact_version_id=artifact_version_id,
-                    name=name,
-                    session=session,
-                )
+            for name, artifact_version_ids in step_run_update.outputs.items():
+                for artifact_version_id in artifact_version_ids:
+                    self._set_run_step_output_artifact(
+                        step_run_id=step_run_id,
+                        artifact_version_id=artifact_version_id,
+                        name=name,
+                        session=session,
+                    )
 
             # Update loaded artifacts.
             for (

{zenml_nightly-0.70.0.dev20241128.dist-info → zenml_nightly-0.70.0.dev20241130.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.70.0.dev20241128
+Version: 0.70.0.dev20241130
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0