zenml-nightly 0.68.0.dev20241027__py3-none-any.whl → 0.68.1.dev20241101__py3-none-any.whl

zenml/metadata/metadata_types.py

@@ -13,13 +13,18 @@
 #  permissions and limitations under the License.
 """Custom types that can be used as metadata of ZenML artifacts."""
 
+import json
 from typing import Any, Dict, List, Set, Tuple, Union
 
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import CoreSchema, core_schema
 
+from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
+from zenml.logger import get_logger
 from zenml.utils.enum_utils import StrEnum
 
+logger = get_logger(__name__)
+
 
 class Uri(str):
     """Special string class to indicate a URI."""
@@ -203,3 +208,47 @@ def cast_to_metadata_type(
     metadata_type = metadata_enum_to_type_mapping[type_]
     typed_value = metadata_type(value)
     return typed_value  # type: ignore[no-any-return]
+
+
+def validate_metadata(
+    metadata: Dict[str, MetadataType],
+) -> Dict[str, MetadataType]:
+    """Validate metadata.
+
+    This function excludes and warns about metadata values that are too long
+    or of an unsupported type.
+
+    Args:
+        metadata: The metadata to validate.
+
+    Returns:
+        The validated metadata.
+    """
+    validated_metadata = {}
+
+    for key, value in metadata.items():
+        if len(key) > STR_FIELD_MAX_LENGTH:
+            logger.warning(
+                f"Metadata key '{key}' is too large to be "
+                "stored in the database. Skipping."
+            )
+            continue
+
+        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
+
+        try:
+            get_metadata_type(value)
+        except ValueError as e:
+            logger.warning(
+                f"Metadata value for key '{key}' is not of a supported "
+                f"type. Skipping. Full error: {e}"
+            )
+
+        validated_metadata[key] = value
+
+    return validated_metadata

zenml/model/model.py

@@ -85,7 +85,6 @@ class Model(BaseModel):
     # technical attributes
     model_version_id: Optional[UUID] = None
     suppress_class_validation_warnings: bool = False
-    was_created_in_this_run: bool = False
     _model_id: UUID = PrivateAttr(None)
     _number: Optional[int] = PrivateAttr(None)
     _created_model_version: bool = PrivateAttr(False)
@@ -691,42 +690,6 @@ class Model(BaseModel):
         )
         mv_request = ModelVersionRequest.model_validate(model_version_request)
         try:
-            if not self.version:
-                try:
-                    from zenml import get_step_context
-
-                    context = get_step_context()
-                except RuntimeError:
-                    pass
-                else:
-                    # if inside a step context we loop over all
-                    # model version configuration to find, if the
-                    # model version for current model was already
-                    # created in the current run, not to create
-                    # new model versions
-                    pipeline_mv = context.pipeline_run.config.model
-                    if (
-                        pipeline_mv
-                        and pipeline_mv.was_created_in_this_run
-                        and pipeline_mv.name == self.name
-                        and pipeline_mv.version is not None
-                    ):
-                        self.version = pipeline_mv.version
-                        self.model_version_id = pipeline_mv.model_version_id
-                    else:
-                        for step in context.pipeline_run.steps.values():
-                            step_mv = step.config.model
-                            if (
-                                step_mv
-                                and step_mv.was_created_in_this_run
-                                and step_mv.name == self.name
-                                and step_mv.version is not None
-                            ):
-                                self.version = step_mv.version
-                                self.model_version_id = (
-                                    step_mv.model_version_id
-                                )
-                                break
             if self.version or self.model_version_id:
                 model_version = self._get_model_version()
             else:
@@ -783,7 +746,6 @@ class Model(BaseModel):
                     time.sleep(sleep)
                     retries_made += 1
             self.version = model_version.name
-            self.was_created_in_this_run = True
             self._created_model_version = True
 
             logger.info(

zenml/models/__init__.py

@@ -132,6 +132,7 @@ from zenml.models.v2.core.component import (
     ComponentResponse,
     ComponentResponseBody,
     ComponentResponseMetadata,
+    ComponentResponseResources
 )
 from zenml.models.v2.core.event_source_flavor import (
     EventSourceFlavorResponse,
@@ -410,6 +411,7 @@ CodeRepositoryResponseBody.model_rebuild()
 CodeRepositoryResponseMetadata.model_rebuild()
 ComponentResponseBody.model_rebuild()
 ComponentResponseMetadata.model_rebuild()
+ComponentResponseResources.model_rebuild()
 EventSourceResponseBody.model_rebuild()
 EventSourceResponseMetadata.model_rebuild()
 EventSourceResponseResources.model_rebuild()
@@ -557,6 +559,7 @@ __all__ = [
     "ComponentResponse",
     "ComponentResponseBody",
     "ComponentResponseMetadata",
+    "ComponentResponseResources",
     "EventSourceFlavorResponse",
     "EventSourceFlavorResponseBody",
     "EventSourceFlavorResponseMetadata",

zenml/models/v2/base/base.py

@@ -215,6 +215,14 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
                         f"`{hydrated_value}`"
                     )
 
+    def hydrate(self) -> None:
+        """Hydrate the response."""
+        hydrated_version = self.get_hydrated_version()
+        self._validate_hydrated_version(hydrated_version)
+
+        self.resources = hydrated_version.resources
+        self.metadata = hydrated_version.metadata
+
     def get_hydrated_version(
         self,
     ) -> "BaseResponse[AnyBody, AnyMetadata, AnyResources]":
@@ -269,9 +277,7 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
             if len(metadata_type.model_fields):
                 # If the metadata class defines any fields, fetch the metadata
                 # through the hydrated version.
-                hydrated_version = self.get_hydrated_version()
-                self._validate_hydrated_version(hydrated_version)
-                self.metadata = hydrated_version.metadata
+                self.hydrate()
             else:
                 # Otherwise, use the metadata class to create an empty metadata
                 # object.
@@ -298,8 +304,8 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
                 "field should exist."
             )
 
-            # metadata is defined as:
-            #   metadata: Optional[....ResponseMetadata] = Field(default=None)
+            # resources is defined as:
+            #   resources: Optional[....ResponseResources] = Field(default=None)
             # We need to find the actual class inside the Optional annotation.
             from zenml.utils.typing_utils import get_args
 
@@ -309,9 +315,7 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
             if len(resources_type.model_fields):
                 # If the resources class defines any fields, fetch the resources
                 # through the hydrated version.
-                hydrated_version = self.get_hydrated_version()
-                self._validate_hydrated_version(hydrated_version)
-                self.resources = hydrated_version.resources
+                self.hydrate()
             else:
                 # Otherwise, use the resources class to create an empty
                 # resources object.

zenml/models/v2/base/filter.py

@@ -604,12 +604,15 @@ class BaseFilter(BaseModel):
         self,
         value: Union[UUID, str],
         table: Type["NamedSchema"],
+        additional_columns: Optional[List[str]] = None,
     ) -> "ColumnElement[bool]":
         """Generate filter conditions for name or id of a table.
 
         Args:
             value: The filter value.
             table: The table to filter.
+            additional_columns: Additional table columns that should also
+                filtered for the given value as part of the or condition.
 
         Returns:
             The query conditions.
@@ -637,6 +640,12 @@ class BaseFilter(BaseModel):
         )
         conditions.append(filter_.generate_query_conditions(table=table))
 
+        for column in additional_columns or []:
+            filter_ = FilterGenerator(table).define_filter(
+                column=column, value=value, operator=operator
+            )
+            conditions.append(filter_.generate_query_conditions(table=table))
+
         return or_(*conditions)
 
     def generate_custom_query_conditions_for_column(

zenml/models/v2/core/artifact_version.py

@@ -24,12 +24,19 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+)
 
 from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.filter import StrFilter
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
@@ -63,11 +70,16 @@ logger = get_logger(__name__)
 class ArtifactVersionRequest(WorkspaceScopedRequest):
     """Request model for artifact versions."""
 
-    artifact_id: UUID = Field(
+    artifact_id: Optional[UUID] = Field(
+        default=None,
         title="ID of the artifact to which this version belongs.",
     )
-    version: Union[str, int] = Field(
-        title="Version of the artifact.", union_mode="left_to_right"
+    artifact_name: Optional[str] = Field(
+        default=None,
+        title="Name of the artifact to which this version belongs.",
+    )
+    version: Optional[Union[int, str]] = Field(
+        default=None, title="Version of the artifact."
     )
     has_custom_name: bool = Field(
         title="Whether the name is custom (True) or auto-generated (False).",
@@ -95,6 +107,9 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
     visualizations: Optional[List["ArtifactVisualizationRequest"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
+    metadata: Optional[Dict[str, MetadataType]] = Field(
+        default=None, title="Metadata of the artifact version."
+    )
 
     @field_validator("version")
     @classmethod
@@ -117,6 +132,28 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
         )
         return value
 
+    @model_validator(mode="after")
+    def _validate_request(self) -> "ArtifactVersionRequest":
+        """Validate the request values.
+
+        Raises:
+            ValueError: If the request is invalid.
+
+        Returns:
+            The validated request.
+        """
+        if self.artifact_id and self.artifact_name:
+            raise ValueError(
+                "Only one of artifact_name and artifact_id can be set."
+            )
+
+        if not (self.artifact_id or self.artifact_name):
+            raise ValueError(
+                "Either artifact_name or artifact_id must be set."
+            )
+
+        return self
+
 
 # ------------------ Update Model ------------------
 
@@ -156,6 +193,10 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
         title="The ID of the pipeline run that generated this artifact version.",
         default=None,
     )
+    artifact_store_id: Optional[UUID] = Field(
+        title="ID of the artifact store in which this artifact is stored.",
+        default=None,
+    )
 
     @field_validator("version")
     @classmethod
@@ -182,10 +223,6 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
 class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for artifact versions."""
 
-    artifact_store_id: Optional[UUID] = Field(
-        title="ID of the artifact store in which this artifact is stored.",
-        default=None,
-    )
     producer_step_run_id: Optional[UUID] = Field(
         title="ID of the step run that produced this artifact.",
         default=None,
@@ -283,7 +320,7 @@ class ArtifactVersionResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().artifact_store_id
+        return self.get_body().artifact_store_id
 
     @property
     def producer_step_run_id(self) -> Optional[UUID]:
@@ -556,7 +593,9 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 ArtifactVersionSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/component.py

@@ -44,9 +44,7 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
     from sqlmodel import SQLModel
 
-    from zenml.models.v2.core.service_connector import (
-        ServiceConnectorResponse,
-    )
+    from zenml.models import FlavorResponse, ServiceConnectorResponse
 
 # ------------------ Base Model ------------------
 
@@ -139,22 +137,11 @@ class InternalComponentRequest(ComponentRequest):
 class ComponentUpdate(BaseUpdate):
     """Update model for stack components."""
 
-    ANALYTICS_FIELDS: ClassVar[List[str]] = ["type", "flavor"]
-
     name: Optional[str] = Field(
         title="The name of the stack component.",
         max_length=STR_FIELD_MAX_LENGTH,
         default=None,
     )
-    type: Optional[StackComponentType] = Field(
-        title="The type of the stack component.",
-        default=None,
-    )
-    flavor: Optional[str] = Field(
-        title="The flavor of the stack component.",
-        max_length=STR_FIELD_MAX_LENGTH,
-        default=None,
-    )
     configuration: Optional[Dict[str, Any]] = Field(
         title="The stack component configuration.",
         default=None,
@@ -187,7 +174,7 @@ class ComponentResponseBody(WorkspaceScopedResponseBody):
     type: StackComponentType = Field(
         title="The type of the stack component.",
     )
-    flavor: str = Field(
+    flavor_name: str = Field(
         title="The flavor of the stack component.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
@@ -232,6 +219,10 @@ class ComponentResponseMetadata(WorkspaceScopedResponseMetadata):
 class ComponentResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the component entity."""
 
+    flavor: "FlavorResponse" = Field(
+        title="The flavor of this stack component.",
+    )
+
 
 class ComponentResponse(
     WorkspaceScopedResponse[
@@ -242,7 +233,7 @@ class ComponentResponse(
 ):
     """Response model for components."""
 
-    ANALYTICS_FIELDS: ClassVar[List[str]] = ["type", "flavor"]
+    ANALYTICS_FIELDS: ClassVar[List[str]] = ["type"]
 
     name: str = Field(
         title="The name of the stack component.",
@@ -265,6 +256,8 @@ class ComponentResponse(
                     if label.startswith("zenml:")
                 }
             )
+        metadata["flavor"] = self.flavor_name
+
         return metadata
 
     def get_hydrated_version(self) -> "ComponentResponse":
@@ -288,13 +281,13 @@ class ComponentResponse(
         return self.get_body().type
 
     @property
-    def flavor(self) -> str:
-        """The `flavor` property.
+    def flavor_name(self) -> str:
+        """The `flavor_name` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_body().flavor
+        return self.get_body().flavor_name
 
     @property
     def integration(self) -> Optional[str]:
@@ -359,6 +352,15 @@ class ComponentResponse(
         """
         return self.get_metadata().connector
 
+    @property
+    def flavor(self) -> "FlavorResponse":
+        """The `flavor` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().flavor
+
 
 # ------------------ Filter Model ------------------
 
@@ -376,6 +378,7 @@ class ComponentFilter(WorkspaceScopedFilter):
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "scope_type",
         "stack_id",
+        "user",
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -418,6 +421,10 @@ class ComponentFilter(WorkspaceScopedFilter):
         description="Stack of the stack component",
         union_mode="left_to_right",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the component.",
+    )
 
     def set_scope_type(self, component_type: str) -> None:
         """Set the type of component on which to perform the filtering to scope the response.
@@ -464,3 +471,31 @@ class ComponentFilter(WorkspaceScopedFilter):
             base_filter = operator(base_filter, stack_filter)
 
         return base_filter
+
+    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Returns:
+            A list of custom filters.
+        """
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import (
+            StackComponentSchema,
+            UserSchema,
+        )
+
+        custom_filters = super().get_custom_filters()
+
+        if self.user:
+            user_filter = and_(
+                StackComponentSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        return custom_filters

zenml/models/v2/core/flavor.py

@@ -192,6 +192,10 @@ class FlavorResponseBody(UserScopedResponseBody):
         title="The name of the integration that the Flavor belongs to.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
+    source: str = Field(
+        title="The path to the module which contains this Flavor.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
     logo_url: Optional[str] = Field(
         default=None,
         title="Optionally, a url pointing to a png,"
@@ -225,10 +229,6 @@ class FlavorResponseMetadata(UserScopedResponseMetadata):
         "connector.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    source: str = Field(
-        title="The path to the module which contains this Flavor.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     docs_url: Optional[str] = Field(
         default=None,
         title="Optionally, a url pointing to docs, within docs.zenml.io.",
@@ -319,6 +319,15 @@ class FlavorResponse(
         """
         return self.get_body().integration
 
+    @property
+    def source(self) -> str:
+        """The `source` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().source
+
     @property
     def logo_url(self) -> Optional[str]:
         """The `logo_url` property.
@@ -373,15 +382,6 @@ class FlavorResponse(
         """
         return self.get_metadata().connector_resource_id_attr
 
-    @property
-    def source(self) -> str:
-        """The `source` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().source
-
     @property
     def docs_url(self) -> Optional[str]:
         """The `docs_url` property.

zenml/models/v2/core/model.py

@@ -368,7 +368,9 @@ class ModelFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 ModelSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/model_version.py

@@ -325,14 +325,11 @@ class ModelVersionResponse(
     # Helper functions
     def to_model_class(
         self,
-        was_created_in_this_run: bool = False,
         suppress_class_validation_warnings: bool = True,
     ) -> "Model":
         """Convert response model to Model object.
 
         Args:
-            was_created_in_this_run: Whether model version was created during
-                the current run.
             suppress_class_validation_warnings: internally used to suppress
                 repeated warnings.
 
@@ -352,7 +349,6 @@ class ModelVersionResponse(
             ethics=self.model.ethics,
             tags=[t.name for t in self.tags],
             version=self.name,
-            was_created_in_this_run=was_created_in_this_run,
             suppress_class_validation_warnings=suppress_class_validation_warnings,
             model_version_id=self.id,
         )
@@ -664,7 +660,9 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 ModelVersionSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/model_version_artifact.py

@@ -294,7 +294,9 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
                 == ArtifactVersionSchema.id,
                 ArtifactVersionSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/model_version_pipeline_run.py

@@ -218,7 +218,9 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
                 == PipelineRunSchema.id,
                 PipelineRunSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/pipeline.py

@@ -364,7 +364,9 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 PipelineSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)