zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.58.2.dev20240614
+0.58.2.dev20240622

zenml/_hub/client.py

@@ -76,7 +76,8 @@ class HubClient:
         if not isinstance(response, list):
             return []
         return [
-            HubPluginResponseModel.parse_obj(plugin) for plugin in response
+            HubPluginResponseModel.model_validate(plugin)
+            for plugin in response
         ]
 
     def get_plugin(
@@ -113,7 +114,7 @@ class HubClient:
             return None
         if not isinstance(response, list) or len(response) == 0:
             return None
-        return HubPluginResponseModel.parse_obj(response[0])
+        return HubPluginResponseModel.model_validate(response[0])
 
     def create_plugin(
         self, plugin_request: HubPluginRequestModel
@@ -127,8 +128,10 @@ class HubClient:
             The plugin response model.
         """
         route = "/plugins"
-        response = self._request("POST", route, data=plugin_request.json())
-        return HubPluginResponseModel.parse_obj(response)
+        response = self._request(
+            "POST", route, data=plugin_request.model_dump_json()
+        )
+        return HubPluginResponseModel.model_validate(response)
 
     # TODO: Potentially reenable this later if hub adds logs streaming endpoint
     # def stream_plugin_build_logs(
@@ -224,7 +227,7 @@ class HubClient:
         """
         try:
             response = self._request("GET", "/users/me")
-            return HubUserResponseModel.parse_obj(response)
+            return HubUserResponseModel.model_validate(response)
         except HubAPIError:
             return None
 

zenml/actions/base_action.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Base implementation of actions."""
 
-import json
 from abc import ABC, abstractmethod
 from typing import Any, ClassVar, Dict, Optional, Type
 
@@ -66,10 +65,7 @@ class BaseActionFlavor(BasePluginFlavor, ABC):
         Returns:
             The config schema.
         """
-        config_schema: Dict[str, Any] = json.loads(
-            cls.ACTION_CONFIG_CLASS.schema_json()
-        )
-        return config_schema
+        return cls.ACTION_CONFIG_CLASS.model_json_schema()
 
     @classmethod
     def get_flavor_response_model(cls, hydrate: bool) -> ActionFlavorResponse:
@@ -260,7 +256,7 @@ class BaseActionHandler(BasePlugin, ABC):
         # before it is sent to the database
         self._validate_trigger_request(trigger=trigger, config=config)
         # Serialize the configuration back into the request
-        trigger.action = config.dict(exclude_none=True)
+        trigger.action = config.model_dump(exclude_none=True)
         # Create the trigger in the database
         trigger_response = self.zen_store.create_trigger(trigger=trigger)
         try:
@@ -290,7 +286,7 @@ class BaseActionHandler(BasePlugin, ABC):
             raise
 
         # Serialize the configuration back into the response
-        trigger_response.set_action(config.dict(exclude_none=True))
+        trigger_response.set_action(config.model_dump(exclude_none=True))
         # Return the response to the user
         return trigger_response
 
@@ -330,7 +326,7 @@ class BaseActionHandler(BasePlugin, ABC):
             config_update=config_update,
         )
         # Serialize the configuration update back into the update request
-        trigger_update.action = config_update.dict(exclude_none=True)
+        trigger_update.action = config_update.model_dump(exclude_none=True)
 
         # Update the trigger in the database
         trigger_response = self.zen_store.update_trigger(
@@ -373,7 +369,9 @@ class BaseActionHandler(BasePlugin, ABC):
             raise
 
         # Serialize the configuration back into the response
-        trigger_response.set_action(response_config.dict(exclude_none=True))
+        trigger_response.set_action(
+            response_config.model_dump(exclude_none=True)
+        )
         # Return the response to the user
         return trigger_response
 
@@ -436,7 +434,7 @@ class BaseActionHandler(BasePlugin, ABC):
             # Call the implementation specific method to process the response
             self._process_trigger_response(trigger=trigger, config=config)
             # Serialize the configuration back into the response
-            trigger.set_action(config.dict(exclude_none=True))
+            trigger.set_action(config.model_dump(exclude_none=True))
             # Add any implementation specific related resources to the trigger
             # response
             self._populate_trigger_response_resources(

zenml/artifact_stores/base_artifact_store.py

@@ -32,7 +32,7 @@ from typing import (
     cast,
 )
 
-from pydantic import root_validator
+from pydantic import model_validator
 
 from zenml.enums import StackComponentType
 from zenml.exceptions import ArtifactStoreInterfaceError
@@ -40,6 +40,7 @@ from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent, StackComponentConfig
 from zenml.utils import io_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -173,15 +174,17 @@ class BaseArtifactStoreConfig(StackComponentConfig):
     SUPPORTED_SCHEMES: ClassVar[Set[str]]
     IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = False
 
-    @root_validator(skip_on_failure=True)
-    def _ensure_artifact_store(cls, values: Dict[str, Any]) -> Any:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _ensure_artifact_store(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validator function for the Artifact Stores.
 
         Checks whether supported schemes are defined and the given path is
         supported.
 
         Args:
-            values: The values to validate.
+            data: the input data to construct the artifact store.
 
         Returns:
             The validated values.
@@ -211,18 +214,20 @@ class BaseArtifactStoreConfig(StackComponentConfig):
                     """
                 )
             )
-        values["path"] = values["path"].strip("'\"`")
-        if not any(
-            values["path"].startswith(i) for i in cls.SUPPORTED_SCHEMES
-        ):
-            raise ArtifactStoreInterfaceError(
-                f"The path: '{values['path']}' you defined for your "
-                f"artifact store is not supported by the implementation of "
-                f"{cls.schema()['title']}, because it does not start with "
-                f"one of its supported schemes: {cls.SUPPORTED_SCHEMES}."
-            )
 
-        return values
+        if "path" in data:
+            data["path"] = data["path"].strip("'\"`")
+            if not any(
+                data["path"].startswith(i) for i in cls.SUPPORTED_SCHEMES
+            ):
+                raise ArtifactStoreInterfaceError(
+                    f"The path: '{data['path']}' you defined for your "
+                    f"artifact store is not supported by the implementation of "
+                    f"{cls.schema()['title']}, because it does not start with "
+                    f"one of its supported schemes: {cls.SUPPORTED_SCHEMES}."
+                )
+
+        return data
 
 
 class BaseArtifactStore(StackComponent):

zenml/artifact_stores/local_artifact_store.py

@@ -21,7 +21,7 @@ artifacts get stored.
 import os
 from typing import TYPE_CHECKING, ClassVar, Optional, Set, Type, Union
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.artifact_stores import (
     BaseArtifactStore,
@@ -50,7 +50,8 @@ class LocalArtifactStoreConfig(BaseArtifactStoreConfig):
 
     path: str = ""
 
-    @validator("path")
+    @field_validator("path")
+    @classmethod
     def ensure_path_local(cls, path: str) -> str:
         """Pydantic validator which ensures that the given path is a local path.
 

zenml/artifacts/artifact_config.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Artifact Config classes to support Model Control Plane feature."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from zenml.enums import ModelStages
 from zenml.exceptions import StepContextError
@@ -63,32 +63,47 @@ class ArtifactConfig(BaseModel):
     """
 
     name: Optional[str] = None
-    version: Optional[Union[str, int]] = None
+    version: Optional[Union[str, int]] = Field(
+        default=None, union_mode="smart"
+    )
     tags: Optional[List[str]] = None
     run_metadata: Optional[Dict[str, MetadataType]] = None
 
     model_name: Optional[str] = None
-    model_version: Optional[Union[ModelStages, str, int]] = None
+    model_version: Optional[Union[ModelStages, str, int]] = Field(
+        default=None, union_mode="smart"
+    )
     is_model_artifact: bool = False
     is_deployment_artifact: bool = False
 
-    @root_validator
-    def _root_validator(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        model_name = values.get("model_name", None)
-        if model_name and values.get("model_version", None) is None:
-            raise ValueError(
-                f"Creation of new model version from `{cls}` is not allowed. "
-                "Please either keep `model_name` and `model_version` both "
-                "`None` to get the model version from the step context or "
-                "specify both at the same time. You can use `ModelStages.LATEST` "
-                "as `model_version` when latest model version is desired."
-            )
-        return values
+    # 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
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
-    class Config:
-        """Config class for ArtifactConfig."""
+    @model_validator(mode="after")
+    def artifact_config_validator(self) -> "ArtifactConfig":
+        """Model validator for the artifact config.
 
-        smart_union = True
+        Raises:
+            ValueError: If both model_name and model_version is set incorrectly.
+
+        Returns:
+            the validated instance.
+        """
+        if self.model_name is not None and self.model_version is None:
+            raise ValueError(
+                f"Creation of new model version from {self.__class__.__name__} "
+                "is not allowed. Please either keep `model_name` and "
+                "`model_version` both `None` to get the model version from the "
+                "step context or specify both at the same time. You can use "
+                "`ModelStages.LATEST` as `model_version` when latest model "
+                "version is desired."
+            )
+        return self
 
     @property
     def _model(self) -> Optional["Model"]:

zenml/artifacts/external_artifact.py

@@ -14,10 +14,10 @@
 """External artifact definition."""
 
 import os
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Optional, Type, Union
 from uuid import UUID, uuid4
 
-from pydantic import root_validator
+from pydantic import Field, model_validator
 
 from zenml.artifacts.external_artifact_config import (
     ExternalArtifactConfiguration,
@@ -79,12 +79,22 @@ class ExternalArtifact(ExternalArtifactConfiguration):
     """
 
     value: Optional[Any] = None
-    materializer: Optional[MaterializerClassOrSource] = None
+    materializer: Optional[MaterializerClassOrSource] = Field(
+        default=None, union_mode="left_to_right"
+    )
     store_artifact_metadata: bool = True
     store_artifact_visualizations: bool = True
 
-    @root_validator
-    def _validate_all(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def external_artifact_validator(self) -> "ExternalArtifact":
+        """Model validator for the external artifact.
+
+        Raises:
+            ValueError: if the value, id and name fields are set incorrectly.
+
+        Returns:
+            the validated instance.
+        """
         deprecation_msg = (
             "Parameter `{param}` or `ExternalArtifact` will be deprecated "
             "in upcoming releases. Please use `{substitute}` instead."
@@ -104,7 +114,7 @@ class ExternalArtifact(ExternalArtifactConfiguration):
                 "Client().get_model_version(<model_name>,<model_version>).get_artifact(name)",
             ],
         ]:
-            if _ := values.get(param, None):
+            if getattr(self, param, None):
                 logger.warning(
                     deprecation_msg.format(
                         param=param,
@@ -112,7 +122,7 @@ class ExternalArtifact(ExternalArtifactConfiguration):
                     )
                 )
         options = [
-            values.get(field, None) is not None
+            getattr(self, field, None) is not None
             for field in ["value", "id", "name"]
         ]
         if sum(options) > 1:
@@ -125,7 +135,7 @@ class ExternalArtifact(ExternalArtifactConfiguration):
                 "Either `value`, `id`, or `name` must be provided when "
                 "creating an external artifact."
             )
-        return values
+        return self
 
     def upload_by_value(self) -> UUID:
         """Uploads the artifact by value.

zenml/artifacts/external_artifact_config.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """External artifact definition."""
 
-from typing import Any, Dict, Optional
+from typing import Optional
 from uuid import UUID
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, model_validator
 
 from zenml.logger import get_logger
 from zenml.model.model import Model
@@ -36,14 +36,22 @@ class ExternalArtifactConfiguration(BaseModel):
     version: Optional[str] = None
     model: Optional[Model] = None
 
-    @root_validator
-    def _validate_all_eac(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("version", None) and values.get("model", None):
+    @model_validator(mode="after")
+    def external_artifact_validator(self) -> "ExternalArtifactConfiguration":
+        """Model validator for the external artifact configuration.
+
+        Raises:
+            ValueError: if both version and model fields are set.
+
+        Returns:
+            the validated instance.
+        """
+        if self.version and self.model:
             raise ValueError(
                 "Cannot provide both `version` and `model` when "
                 "creating an external artifact."
             )
-        return values
+        return self
 
     def get_artifact_version_id(self) -> UUID:
         """Get the artifact.

zenml/artifacts/unmaterialized_artifact.py

@@ -13,12 +13,7 @@
 #  permissions and limitations under the License.
 """Unmaterialized artifact class."""
 
-from zenml.models import (
-    ArtifactVersionResponse,
-    RunMetadataResponse,
-    UserResponse,
-    WorkspaceResponse,
-)
+from zenml.models import ArtifactVersionResponse
 
 
 class UnmaterializedArtifact(ArtifactVersionResponse):
@@ -41,8 +36,4 @@ class UnmaterializedArtifact(ArtifactVersionResponse):
     """
 
 
-UnmaterializedArtifact.update_forward_refs(
-    UserResponseModel=UserResponse,
-    WorkspaceResponseModel=WorkspaceResponse,
-    RunMetadataResponseModel=RunMetadataResponse,
-)
+UnmaterializedArtifact.model_rebuild()

zenml/cli/__init__.py

@@ -1746,6 +1746,12 @@ pipeline runs, use:
 zenml artifact prune
 ```
 
+You might find that some artifacts throw errors when you try to prune them,
+likely because they were stored locally and no longer exist. If you wish to
+continue pruning and to ignore these errors, please add the `--ignore-errors`
+flag. Warning messages will still be output to the terminal during this
+process.
+
 Each pipeline run that requires Docker images also stores a build which
 contains the image names used for this run. To list all builds, use:
 

zenml/cli/artifact.py

@@ -229,8 +229,17 @@ def update_artifact_version(
     is_flag=True,
     help="Don't ask for confirmation.",
 )
+@click.option(
+    "--ignore-errors",
+    "-i",
+    is_flag=True,
+    help="Ignore errors and continue with the next artifact version.",
+)
 def prune_artifacts(
-    only_artifact: bool = False, only_metadata: bool = False, yes: bool = False
+    only_artifact: bool = False,
+    only_metadata: bool = False,
+    yes: bool = False,
+    ignore_errors: bool = False,
 ) -> None:
     """Delete all unused artifacts and artifact versions.
 
@@ -244,6 +253,8 @@ def prune_artifacts(
         only_metadata: If set, only delete metadata and not the actual artifact
             objects stored in the artifact store.
         yes: If set, don't ask for confirmation.
+        ignore_errors: If set, ignore errors and continue with the next
+            artifact version.
     """
     client = Client()
     unused_artifact_versions = depaginate(
@@ -275,7 +286,14 @@ def prune_artifacts(
                 Client().delete_artifact(unused_artifact.id)
 
         except Exception as e:
-            cli_utils.error(str(e))
+            if ignore_errors:
+                cli_utils.warning(
+                    f"Failed to delete artifact version {unused_artifact_version.id}: {str(e)}"
+                )
+            else:
+                cli_utils.error(
+                    f"Failed to delete artifact version {unused_artifact_version.id}: {str(e)}"
+                )
     cli_utils.declare("All unused artifacts and artifact versions deleted.")
 
 

zenml/cli/base.py

@@ -87,11 +87,11 @@ ZENML_PROJECT_TEMPLATES = dict(
     ),
     nlp=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-nlp",
-        github_tag="2024.06.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.06.14",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     llm_finetuning=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-llm-finetuning",
-        github_tag="2024.06.06",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.06.20",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
 )
 

zenml/cli/served_model.py

@@ -29,7 +29,6 @@ from zenml.cli.utils import (
 )
 from zenml.console import console
 from zenml.enums import StackComponentType
-from zenml.model_deployers import BaseModelDeployer
 
 if TYPE_CHECKING:
     from zenml.model_deployers import BaseModelDeployer

zenml/cli/server.py

@@ -392,7 +392,7 @@ def deploy(
 
         from zenml.zen_server.deploy.deployment import ServerDeploymentConfig
 
-        server_config = ServerDeploymentConfig.parse_obj(config_dict)
+        server_config = ServerDeploymentConfig.model_validate(config_dict)
 
         from zenml.zen_server.deploy.deployer import ServerDeployer
 
@@ -685,7 +685,7 @@ def connect(
             )
 
         if raw_config:
-            store_config = StoreConfiguration.parse_obj(store_dict)
+            store_config = StoreConfiguration.model_validate(store_dict)
             GlobalConfiguration().set_store(store_config)
             return
 
@@ -756,7 +756,7 @@ def connect(
     store_config_class = BaseZenStore.get_store_config_class(store_type)
     assert store_config_class is not None
 
-    store_config = store_config_class.parse_obj(store_dict)
+    store_config = store_config_class.model_validate(store_dict)
     try:
         GlobalConfiguration().set_store(store_config)
     except IllegalOperationError: