zenml-nightly 0.70.0.dev20241120__tar.gz → 0.70.0.dev20241121__tar.gz

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/PKG-INFO

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

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.70.0.dev20241120"
+version = "0.70.0.dev20241121"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.70.0.dev20241121/src/zenml/VERSION

@@ -0,0 +1 @@
+0.70.0.dev20241121

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/artifacts/artifact_config.py

@@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field, model_validator
 
+from zenml.enums import ArtifactType
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.utils.pydantic_utils import before_validator_handler
@@ -36,6 +37,7 @@ class ArtifactConfig(BaseModel):
         int, ArtifactConfig(
             name="my_artifact",  # override the default artifact name
             version=42,  # set a custom version
+            artifact_type=ArtifactType.MODEL,  # Specify the artifact type
             tags=["tag1", "tag2"],  # set custom tags
         )
     ]:
@@ -47,8 +49,9 @@ class ArtifactConfig(BaseModel):
         version: The version of the artifact.
         tags: The tags of the artifact.
         run_metadata: Metadata to add to the artifact.
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
+        artifact_type: Optional type of the artifact. If not given, the type
+            specified by the materializer that is used to save this artifact
+            is used.
     """
 
     name: Optional[str] = None
@@ -58,8 +61,7 @@ class ArtifactConfig(BaseModel):
     tags: Optional[List[str]] = None
     run_metadata: Optional[Dict[str, MetadataType]] = None
 
-    is_model_artifact: bool = False
-    is_deployment_artifact: bool = False
+    artifact_type: Optional[ArtifactType] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -70,6 +72,10 @@ class ArtifactConfig(BaseModel):
         Args:
             data: The model data.
 
+        Raises:
+            ValueError: If the artifact is configured to be
+                both a model and a deployment artifact.
+
         Returns:
             Model data without the removed attributes.
         """
@@ -82,4 +88,26 @@ class ArtifactConfig(BaseModel):
                 "artifact is not supported anymore."
             )
 
+        is_model_artifact = data.pop("is_model_artifact", None)
+        is_deployment_artifact = data.pop("is_deployment_artifact", None)
+
+        if is_model_artifact and is_deployment_artifact:
+            raise ValueError(
+                "An artifact can only be a model artifact or deployment "
+                "artifact."
+            )
+        elif is_model_artifact:
+            logger.warning(
+                "`ArtifactConfig.is_model_artifact` is deprecated and will be "
+                "removed soon. Use `ArtifactConfig.artifact_type` instead."
+            )
+            data.setdefault("artifact_type", ArtifactType.MODEL)
+        elif is_deployment_artifact:
+            logger.warning(
+                "`ArtifactConfig.is_deployment_artifact` is deprecated and "
+                "will be removed soon. Use `ArtifactConfig.artifact_type` "
+                "instead."
+            )
+            data.setdefault("artifact_type", ArtifactType.SERVICE)
+
         return data

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/artifacts/utils.py

@@ -30,7 +30,6 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.artifacts.preexisting_data_materializer import (
     PreexistingDataMaterializer,
 )
@@ -118,6 +117,7 @@ def _store_artifact_data_and_prepare_request(
     materializer_class: Type["BaseMaterializer"],
     save_type: ArtifactSaveType,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     store_metadata: bool = True,
     store_visualizations: bool = True,
@@ -134,6 +134,8 @@ def _store_artifact_data_and_prepare_request(
             artifact data.
         save_type: Save type of the artifact version.
         version: The artifact version.
+        artifact_type: The artifact type. If not given, the type will be defined
+            by the materializer that is used to save the artifact.
         tags: Tags for the artifact version.
         store_metadata: Whether to store metadata for the artifact version.
         store_visualizations: Whether to store visualizations for the artifact
@@ -176,7 +178,7 @@ def _store_artifact_data_and_prepare_request(
         artifact_name=name,
         version=version,
         tags=tags,
-        type=materializer.ASSOCIATED_ARTIFACT_TYPE,
+        type=artifact_type or materializer.ASSOCIATED_ARTIFACT_TYPE,
         uri=materializer.uri,
         materializer=source_utils.resolve(materializer.__class__),
         data_type=source_utils.resolve(data_type),
@@ -198,14 +200,13 @@ def save_artifact(
     data: Any,
     name: str,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     extract_metadata: bool = True,
     include_visualizations: bool = True,
     user_metadata: Optional[Dict[str, "MetadataType"]] = None,
     materializer: Optional["MaterializerClassOrSource"] = None,
     uri: Optional[str] = None,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
     # TODO: remove these once external artifact does not use this function anymore
     save_type: ArtifactSaveType = ArtifactSaveType.MANUAL,
     has_custom_name: bool = True,
@@ -218,6 +219,8 @@ def save_artifact(
         version: The version of the artifact. If not provided, a new
             auto-incremented version will be used.
         tags: Tags to associate with the artifact.
+        artifact_type: The artifact type. If not given, the type will be defined
+            by the materializer that is used to save the artifact.
         extract_metadata: If artifact metadata should be extracted and returned.
         include_visualizations: If artifact visualizations should be generated.
         user_metadata: User-provided metadata to store with the artifact.
@@ -226,8 +229,6 @@ def save_artifact(
         uri: The URI within the artifact store to upload the artifact
             to. If not provided, the artifact will be uploaded to
             `custom_artifacts/{name}/{version}`.
-        is_model_artifact: If the artifact is a model artifact.
-        is_deployment_artifact: If the artifact is a deployment artifact.
         save_type: The type of save operation that created the artifact version.
         has_custom_name: If the artifact name is custom and should be listed in
             the dashboard "Artifacts" tab.
@@ -273,6 +274,7 @@ def save_artifact(
         materializer_class=materializer_class,
         save_type=save_type,
         version=version,
+        artifact_type=artifact_type,
         tags=tags,
         store_metadata=extract_metadata,
         store_visualizations=include_visualizations,
@@ -286,8 +288,6 @@ def save_artifact(
     if save_type == ArtifactSaveType.MANUAL:
         _link_artifact_version_to_the_step_and_model(
             artifact_version=artifact_version,
-            is_model_artifact=is_model_artifact,
-            is_deployment_artifact=is_deployment_artifact,
         )
 
     return artifact_version
@@ -297,10 +297,9 @@ def register_artifact(
     folder_or_file_uri: str,
     name: str,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     has_custom_name: bool = True,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
     artifact_metadata: Dict[str, "MetadataType"] = {},
 ) -> "ArtifactVersionResponse":
     """Register existing data stored in the artifact store as a ZenML Artifact.
@@ -311,11 +310,11 @@ def register_artifact(
         name: The name of the artifact.
         version: The version of the artifact. If not provided, a new
             auto-incremented version will be used.
+        artifact_type: The artifact type. If not given, the type will default
+            to `data`.
         tags: Tags to associate with the artifact.
         has_custom_name: If the artifact name is custom and should be listed in
             the dashboard "Artifacts" tab.
-        is_model_artifact: If the artifact is a model artifact.
-        is_deployment_artifact: If the artifact is a deployment artifact.
         artifact_metadata: Metadata dictionary to attach to the artifact version.
 
     Returns:
@@ -346,7 +345,7 @@ def register_artifact(
         artifact_name=name,
         version=version,
         tags=tags,
-        type=ArtifactType.DATA,
+        type=artifact_type or ArtifactType.DATA,
         save_type=ArtifactSaveType.PREEXISTING,
         uri=folder_or_file_uri,
         materializer=source_utils.resolve(PreexistingDataMaterializer),
@@ -365,8 +364,6 @@ def register_artifact(
 
     _link_artifact_version_to_the_step_and_model(
         artifact_version=artifact_version,
-        is_model_artifact=is_model_artifact,
-        is_deployment_artifact=is_deployment_artifact,
     )
 
     return artifact_version
@@ -674,8 +671,6 @@ def _check_if_artifact_with_given_uri_already_registered(
 
 def _link_artifact_version_to_the_step_and_model(
     artifact_version: ArtifactVersionResponse,
-    is_model_artifact: bool,
-    is_deployment_artifact: bool,
 ) -> None:
     """Link an artifact version to the step and its' context model.
 
@@ -685,8 +680,6 @@ def _link_artifact_version_to_the_step_and_model(
 
     Args:
         artifact_version: The artifact version to link.
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
     """
     client = Client()
     try:
@@ -706,14 +699,9 @@ def _link_artifact_version_to_the_step_and_model(
                 link_artifact_version_to_model_version,
             )
 
-            artifact_config = ArtifactConfig(
-                is_model_artifact=is_model_artifact,
-                is_deployment_artifact=is_deployment_artifact,
-            )
             link_artifact_version_to_model_version(
                 artifact_version=artifact_version,
                 model_version=step_context.model_version,
-                artifact_config=artifact_config,
             )
     except (RuntimeError, StepContextError):
         logger.debug(f"Unable to link saved artifact to {error_message}.")

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/cli/base.py

@@ -79,7 +79,7 @@ class ZenMLProjectTemplateLocation(BaseModel):
 ZENML_PROJECT_TEMPLATES = dict(
     e2e_batch=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-e2e-batch",
-        github_tag="2024.10.30",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.11.20",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/client.py

@@ -4220,6 +4220,7 @@ class Client(metaclass=ClientMetaClass):
         materializer: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        model_version_id: Optional[Union[str, UUID]] = None,
         only_unused: Optional[bool] = False,
         has_custom_name: Optional[bool] = None,
         user: Optional[Union[UUID, str]] = None,
@@ -4249,7 +4250,8 @@ class Client(metaclass=ClientMetaClass):
             uri: The uri of the artifact to filter by.
             materializer: The materializer of the artifact to filter by.
             workspace_id: The id of the workspace to filter by.
-            user_id: The  id of the user to filter by.
+            user_id: The id of the user to filter by.
+            model_version_id: Filter by model version ID.
             only_unused: Only return artifact versions that are not used in
                 any pipeline runs.
             has_custom_name: Filter artifacts with/without custom names.
@@ -4283,6 +4285,7 @@ class Client(metaclass=ClientMetaClass):
             materializer=materializer,
             workspace_id=workspace_id,
             user_id=user_id,
+            model_version_id=model_version_id,
             only_unused=only_unused,
             has_custom_name=has_custom_name,
             tag=tag,
@@ -6435,9 +6438,6 @@ class Client(metaclass=ClientMetaClass):
         logical_operator: LogicalOperators = LogicalOperators.AND,
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
-        workspace_id: Optional[Union[UUID, str]] = None,
-        user_id: Optional[Union[UUID, str]] = None,
-        model_id: Optional[Union[UUID, str]] = None,
         model_version_id: Optional[Union[UUID, str]] = None,
         artifact_version_id: Optional[Union[UUID, str]] = None,
         artifact_name: Optional[str] = None,
@@ -6457,9 +6457,6 @@ class Client(metaclass=ClientMetaClass):
             logical_operator: Which logical operator to use [and, or]
             created: Use to filter by time of creation
             updated: Use the last updated date for filtering
-            workspace_id: Use the workspace id for filtering
-            user_id: Use the user id for filtering
-            model_id: Use the model id for filtering
             model_version_id: Use the model version id for filtering
             artifact_version_id: Use the artifact id for filtering
             artifact_name: Use the artifact name for filtering
@@ -6482,9 +6479,6 @@ class Client(metaclass=ClientMetaClass):
                 size=size,
                 created=created,
                 updated=updated,
-                workspace_id=workspace_id,
-                user_id=user_id,
-                model_id=model_id,
                 model_version_id=model_version_id,
                 artifact_version_id=artifact_version_id,
                 artifact_name=artifact_name,
@@ -6556,9 +6550,6 @@ class Client(metaclass=ClientMetaClass):
         logical_operator: LogicalOperators = LogicalOperators.AND,
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
-        workspace_id: Optional[Union[UUID, str]] = None,
-        user_id: Optional[Union[UUID, str]] = None,
-        model_id: Optional[Union[UUID, str]] = None,
         model_version_id: Optional[Union[UUID, str]] = None,
         pipeline_run_id: Optional[Union[UUID, str]] = None,
         pipeline_run_name: Optional[str] = None,
@@ -6574,9 +6565,6 @@ class Client(metaclass=ClientMetaClass):
             logical_operator: Which logical operator to use [and, or]
             created: Use to filter by time of creation
             updated: Use the last updated date for filtering
-            workspace_id: Use the workspace id for filtering
-            user_id: Use the user id for filtering
-            model_id: Use the model id for filtering
             model_version_id: Use the model version id for filtering
             pipeline_run_id: Use the pipeline run id for filtering
             pipeline_run_name: Use the pipeline run name for filtering
@@ -6595,9 +6583,6 @@ class Client(metaclass=ClientMetaClass):
                 size=size,
                 created=created,
                 updated=updated,
-                workspace_id=workspace_id,
-                user_id=user_id,
-                model_id=model_id,
                 model_version_id=model_version_id,
                 pipeline_run_id=pipeline_run_id,
                 pipeline_run_name=pipeline_run_name,

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/model/utils.py

@@ -16,7 +16,6 @@
 from typing import Dict, Optional, Union
 from uuid import UUID
 
-from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.client import Client
 from zenml.enums import ModelStages
 from zenml.exceptions import StepContextError
@@ -81,32 +80,18 @@ def log_model_metadata(
 def link_artifact_version_to_model_version(
     artifact_version: ArtifactVersionResponse,
     model_version: ModelVersionResponse,
-    artifact_config: Optional[ArtifactConfig] = None,
 ) -> None:
     """Link an artifact version to a model version.
 
     Args:
         artifact_version: The artifact version to link.
         model_version: The model version to link.
-        artifact_config: Output artifact configuration.
     """
-    if artifact_config:
-        is_model_artifact = artifact_config.is_model_artifact
-        is_deployment_artifact = artifact_config.is_deployment_artifact
-    else:
-        is_model_artifact = False
-        is_deployment_artifact = False
-
     client = Client()
     client.zen_store.create_model_version_artifact_link(
         ModelVersionArtifactRequest(
-            user=client.active_user.id,
-            workspace=client.active_workspace.id,
             artifact_version=artifact_version.id,
-            model=model_version.model.id,
             model_version=model_version.id,
-            is_model_artifact=is_model_artifact,
-            is_deployment_artifact=is_deployment_artifact,
         )
     )
 
@@ -114,16 +99,12 @@ def link_artifact_version_to_model_version(
 def link_artifact_to_model(
     artifact_version: ArtifactVersionResponse,
     model: Optional["Model"] = None,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
 ) -> None:
     """Link the artifact to the model.
 
     Args:
         artifact_version: The artifact version to link.
         model: The model to link to.
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
 
     Raises:
         RuntimeError: If called outside of a step.
@@ -145,14 +126,9 @@ def link_artifact_to_model(
             )
 
     model_version = model._get_or_create_model_version()
-    artifact_config = ArtifactConfig(
-        is_model_artifact=is_model_artifact,
-        is_deployment_artifact=is_deployment_artifact,
-    )
     link_artifact_version_to_model_version(
         artifact_version=artifact_version,
         model_version=model_version,
-        artifact_config=artifact_config,
     )
 
 

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/models/v2/core/artifact_version.py

@@ -37,7 +37,7 @@ from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ArtifactSaveType, 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.filter import FilterGenerator, StrFilter
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
@@ -474,6 +474,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         "user",
         "model",
         "pipeline_run",
+        "model_version_id",
         "run_metadata",
     ]
     artifact_id: Optional[Union[UUID, str]] = Field(
@@ -525,6 +526,11 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         description="User that produced this artifact",
         union_mode="left_to_right",
     )
+    model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="ID of the model version that is associated with this artifact version.",
+        union_mode="left_to_right",
+    )
     only_unused: Optional[bool] = Field(
         default=False, description="Filter only for unused artifacts"
     )
@@ -568,6 +574,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             ArtifactVersionSchema,
             ModelSchema,
             ModelVersionArtifactSchema,
+            ModelVersionSchema,
             PipelineRunSchema,
             RunMetadataSchema,
             StepRunInputArtifactSchema,
@@ -600,6 +607,20 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(unused_filter)
 
+        if self.model_version_id:
+            value, operator = self._resolve_operator(self.model_version_id)
+
+            model_version_filter = and_(
+                ArtifactVersionSchema.id
+                == ModelVersionArtifactSchema.artifact_version_id,
+                ModelVersionArtifactSchema.model_version_id
+                == ModelVersionSchema.id,
+                FilterGenerator(ModelVersionSchema)
+                .define_filter(column="id", value=value, operator=operator)
+                .generate_query_conditions(ModelVersionSchema),
+            )
+            custom_filters.append(model_version_filter)
+
         if self.has_custom_name is not None:
             custom_name_filter = and_(
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
@@ -622,7 +643,9 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             model_filter = and_(
                 ArtifactVersionSchema.id
                 == ModelVersionArtifactSchema.artifact_version_id,
-                ModelVersionArtifactSchema.model_id == ModelSchema.id,
+                ModelVersionArtifactSchema.model_version_id
+                == ModelVersionSchema.id,
+                ModelVersionSchema.model_id == ModelSchema.id,
                 self.generate_name_or_id_query_conditions(
                     value=self.model, table=ModelSchema
                 ),

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241121}/src/zenml/models/v2/core/model_version.py

@@ -576,10 +576,6 @@ class ModelVersionResponse(
             force=force,
         )
 
-    # TODO in https://zenml.atlassian.net/browse/OSS-2433
-    # def generate_model_card(self, template_name: str) -> str:
-    #     """Return HTML/PDF based on input template"""
-
 
 # ------------------ Filter Model ------------------