zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/stack/flavor.py

@@ -12,10 +12,11 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Base ZenML Flavor implementation."""
-import json
+
 from abc import abstractmethod
 from typing import Any, Dict, Optional, Type, cast
 
+from zenml.client import Client
 from zenml.enums import StackComponentType
 from zenml.models import (
     FlavorRequest,
@@ -25,6 +26,7 @@ from zenml.models import (
 from zenml.models.v2.core.flavor import InternalFlavorRequest
 from zenml.stack.stack_component import StackComponent, StackComponentConfig
 from zenml.utils import source_utils
+from zenml.utils.package_utils import is_latest_zenml_version
 
 
 class Flavor:
@@ -100,10 +102,7 @@ class Flavor:
         Returns:
             The config schema.
         """
-        config_schema: Dict[str, Any] = json.loads(
-            self.config_class.schema_json()
-        )
-        return config_schema
+        return self.config_class.model_json_schema()
 
     @property
     def service_connector_requirements(
@@ -148,9 +147,6 @@ class Flavor:
         Returns:
             The model.
         """
-        from zenml.client import Client
-
-        client = Client()
         connector_requirements = self.service_connector_requirements
         connector_type = (
             connector_requirements.connector_type
@@ -167,10 +163,16 @@ class Flavor:
             if connector_requirements
             else None
         )
+        user = None
+        workspace = None
+        if is_custom:
+            user = Client().active_user.id
+            workspace = Client().active_workspace.id
+
         model_class = FlavorRequest if is_custom else InternalFlavorRequest
         model = model_class(
-            user=client.active_user.id if is_custom else None,
-            workspace=client.active_workspace.id if is_custom else None,
+            user=user,
+            workspace=workspace,
             name=self.name,
             type=self.type,
             source=source_utils.resolve(self.__class__).import_path,
@@ -186,16 +188,12 @@ class Flavor:
         )
         return model
 
-    def generate_default_docs_url(self, component_name: str = "") -> str:
+    def generate_default_docs_url(self) -> str:
         """Generate the doc urls for all inbuilt and integration flavors.
 
         Note that this method is not going to be useful for custom flavors,
         which do not have any docs in the main zenml docs.
 
-        Args:
-            component_name: The name of the component for docs generation. Used
-                for legacy documentation before ZenML v0.34.0.
-
         Returns:
             The complete url to the zenml documentation
         """
@@ -203,9 +201,18 @@ class Flavor:
 
         component_type = self.type.plural.replace("_", "-")
         name = self.name.replace("_", "-")
-        docs_component_name = component_name or name
-        base = f"https://docs.zenml.io/v/{__version__}"
-        return f"{base}/user-guide/component-guide/{component_type}/{docs_component_name}"
+
+        try:
+            is_latest = is_latest_zenml_version()
+        except RuntimeError:
+            # We assume in error cases that we are on the latest version
+            is_latest = True
+
+        if is_latest:
+            base = "https://docs.zenml.io"
+        else:
+            base = f"https://zenml-io.gitbook.io/zenml-legacy-documentation/v/{__version__}"
+        return f"{base}/stack-components/{component_type}/{name}"
 
     def generate_default_sdk_docs_url(self) -> str:
         """Generate SDK docs url for a flavor.

zenml/stack/flavor_registry.py

@@ -122,8 +122,8 @@ class FlavorRegistry:
                 if len(existing_flavor) == 0:
                     store.create_flavor(flavor_request_model)
                 else:
-                    flavor_update_model = FlavorUpdate.parse_obj(
-                        flavor_request_model
+                    flavor_update_model = FlavorUpdate.model_validate(
+                        dict(flavor_request_model)
                     )
                     store.update_flavor(
                         existing_flavor[0].id, flavor_update_model
@@ -154,8 +154,8 @@ class FlavorRegistry:
                         if len(existing_flavor) == 0:
                             store.create_flavor(flavor_request_model)
                         else:
-                            flavor_update_model = FlavorUpdate.parse_obj(
-                                flavor_request_model
+                            flavor_update_model = FlavorUpdate.model_validate(
+                                dict(flavor_request_model)
                             )
                             store.update_flavor(
                                 existing_flavor[0].id, flavor_update_model

zenml/stack/stack.py

@@ -13,8 +13,8 @@
 #  permissions and limitations under the License.
 """Implementation of the ZenML Stack class."""
 
-import functools
 import itertools
+import json
 import os
 from datetime import datetime
 from typing import (
@@ -40,7 +40,7 @@ from zenml.constants import (
     handle_bool_env_var,
 )
 from zenml.enums import SecretValidationLevel, StackComponentType
-from zenml.exceptions import ProvisioningError, StackValidationError
+from zenml.exceptions import StackValidationError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models import StackResponse
@@ -62,7 +62,11 @@ if TYPE_CHECKING:
     from zenml.image_builders import BaseImageBuilder
     from zenml.model_deployers import BaseModelDeployer
     from zenml.model_registries import BaseModelRegistry
-    from zenml.models import PipelineDeploymentBase, PipelineDeploymentResponse
+    from zenml.models import (
+        PipelineDeploymentBase,
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+    )
     from zenml.orchestrators import BaseOrchestrator
     from zenml.stack import StackComponent
     from zenml.step_operators import BaseStepOperator
@@ -154,11 +158,9 @@ class Stack:
 
         # Run a hydrated list call once to avoid one request per component
         component_models = pagination_utils.depaginate(
-            list_method=functools.partial(
-                Client().list_stack_components,
-                stack_id=stack_model.id,
-                hydrate=True,
-            )
+            Client().list_stack_components,
+            stack_id=stack_model.id,
+            hydrate=True,
         )
 
         stack_components = {
@@ -481,7 +483,9 @@ class Stack:
             A dictionary containing the stack components.
         """
         component_dict = {
-            component_type.value: component.config.json(sort_keys=True)
+            component_type.value: json.dumps(
+                component.config.model_dump(mode="json"), sort_keys=True
+            )
             for component_type, component in self.components.items()
         }
         component_dict.update({"name": self.name})
@@ -751,21 +755,6 @@ class Stack:
                 updated=datetime.utcnow(),
             )
 
-            logger.warning(
-                "The stack `%s` contains components that require building "
-                "Docker images. Older versions of ZenML always built these "
-                "images locally, but since version 0.32.0 this behavior can be "
-                "configured using the `image_builder` stack component. This "
-                "stack will temporarily default to a local image builder that "
-                "mirrors the previous behavior, but this will be removed in "
-                "future versions of ZenML. Please add an image builder to this "
-                "stack:\n"
-                "`zenml image-builder register <NAME> ...\n"
-                "zenml stack update %s -i <NAME>`",
-                self.name,
-                self.id,
-            )
-
             self._image_builder = image_builder
 
     def prepare_pipeline_deployment(
@@ -779,39 +768,18 @@ class Stack:
             deployment: The pipeline deployment
 
         Raises:
-            StackValidationError: If the stack component is not running.
             RuntimeError: If trying to deploy a pipeline that requires a remote
                 ZenML server with a local one.
         """
         self.validate(fail_if_secrets_missing=True)
 
-        for component in self.components.values():
-            if not component.is_running:
-                raise StackValidationError(
-                    f"The '{component.name}' {component.type} stack component "
-                    f"is not currently running. Please run the following "
-                    f"command to provision and start the component:\n\n"
-                    f"    `zenml stack up`\n"
-                    f"It is worth noting that the provision command will "
-                    f" be deprecated in the future. ZenML will no longer "
-                    f"be responsible for provisioning infrastructure, "
-                    f"or port-forwarding directly. Instead of managing "
-                    f"the state of the components, ZenML will be utilizing "
-                    f"the already running stack or stack components directly. "
-                    f"Additionally, we are also providing a variety of "
-                    f" deployment recipes for popular Kubernetes-based "
-                    f"integrations such as Kubeflow, Tekton, and Seldon etc."
-                    f"Check out https://docs.zenml.io/stacks-and-components/stack-deployment/deploy-a-stack-using-stack-recipes"
-                    f"for more information."
-                )
-
         if self.requires_remote_server and Client().zen_store.is_local_store():
             raise RuntimeError(
                 "Stacks with remote components such as remote orchestrators "
                 "and step operators require a remote "
                 "ZenML server. To run a pipeline with this stack you need to "
                 "connect to a remote ZenML server first. Check out "
-                "https://docs.zenml.io/user-guide/starter-guide/switch-to-production "
+                "https://docs.zenml.io/getting-started/deploying-zenml "
                 "for more information on how to deploy ZenML."
             )
 
@@ -841,16 +809,17 @@ class Stack:
     def deploy_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
-    ) -> Any:
+        placeholder_run: Optional["PipelineRunResponse"] = None,
+    ) -> None:
         """Deploys a pipeline on this stack.
 
         Args:
             deployment: The pipeline deployment.
-
-        Returns:
-            The return value of the call to `orchestrator.run_pipeline(...)`.
+            placeholder_run: An optional placeholder run for the deployment.
         """
-        return self.orchestrator.run(deployment=deployment, stack=self)
+        self.orchestrator.run(
+            deployment=deployment, stack=self, placeholder_run=placeholder_run
+        )
 
     def _get_active_components_for_step(
         self, step_config: "StepConfiguration"
@@ -962,83 +931,3 @@ class Stack:
             info.config
         ).values():
             component.cleanup_step_run(info=info, step_failed=step_failed)
-
-    @property
-    def is_provisioned(self) -> bool:
-        """If the stack provisioned resources to run locally.
-
-        Returns:
-            True if the stack provisioned resources to run locally.
-        """
-        return all(
-            component.is_provisioned for component in self.components.values()
-        )
-
-    @property
-    def is_running(self) -> bool:
-        """If the stack is running locally.
-
-        Returns:
-            True if the stack is running locally, False otherwise.
-        """
-        return all(
-            component.is_running for component in self.components.values()
-        )
-
-    def provision(self) -> None:
-        """Provisions resources to run the stack locally."""
-        self.validate(fail_if_secrets_missing=True)
-        logger.info("Provisioning resources for stack '%s'.", self.name)
-        for component in self.components.values():
-            if not component.is_provisioned:
-                component.provision()
-                logger.info("Provisioned resources for %s.", component)
-
-    def deprovision(self) -> None:
-        """Deprovisions all local resources of the stack."""
-        logger.info("Deprovisioning resources for stack '%s'.", self.name)
-        for component in self.components.values():
-            if component.is_provisioned:
-                try:
-                    component.deprovision()
-                    logger.info("Deprovisioned resources for %s.", component)
-                except NotImplementedError as e:
-                    logger.warning(e)
-
-    def resume(self) -> None:
-        """Resumes the provisioned local resources of the stack.
-
-        Raises:
-            ProvisioningError: If any stack component is missing provisioned
-                resources.
-        """
-        logger.info("Resuming provisioned resources for stack %s.", self.name)
-        for component in self.components.values():
-            if component.is_running:
-                # the component is already running, no need to resume anything
-                pass
-            elif component.is_provisioned:
-                component.resume()
-                logger.info("Resumed resources for %s.", component)
-            else:
-                raise ProvisioningError(
-                    f"Unable to resume resources for {component}: No "
-                    f"resources have been provisioned for this component."
-                )
-
-    def suspend(self) -> None:
-        """Suspends the provisioned local resources of the stack."""
-        logger.info(
-            "Suspending provisioned resources for stack '%s'.", self.name
-        )
-        for component in self.components.values():
-            if not component.is_suspended:
-                try:
-                    component.suspend()
-                    logger.info("Suspended resources for %s.", component)
-                except NotImplementedError:
-                    logger.warning(
-                        "Suspending provisioned resources not implemented "
-                        "for %s. Continuing without suspending resources...",
-                        component,
-                    )

zenml/stack/stack_component.py

@@ -12,12 +12,16 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Implementation of the ZenML Stack Component class."""
+
+import json
 from abc import ABC
+from collections.abc import Mapping, Sequence
 from datetime import datetime
+from inspect import isclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from zenml.config.build_configuration import BuildConfiguration
 from zenml.config.step_configurations import Step
@@ -26,7 +30,12 @@ from zenml.enums import StackComponentType
 from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import ServiceConnectorRequirements, StepRunResponse
-from zenml.utils import secret_utils, settings_utils
+from zenml.utils import (
+    pydantic_utils,
+    secret_utils,
+    settings_utils,
+    typing_utils,
+)
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
@@ -60,7 +69,8 @@ class StackComponentConfig(BaseModel, ABC):
         custom pydantic validation are set as secret references.
 
         Args:
-            warn_about_plain_text_secrets: If true, then warns about using plain-text secrets.
+            warn_about_plain_text_secrets: If true, then warns about using
+                plain-text secrets.
             **kwargs: Arguments to initialize this stack component.
 
         Raises:
@@ -70,7 +80,7 @@ class StackComponentConfig(BaseModel, ABC):
         """
         for key, value in kwargs.items():
             try:
-                field = self.__class__.__fields__[key]
+                field = self.__class__.model_fields[key]
             except KeyError:
                 # Value for a private attribute or non-existing field, this
                 # will fail during the upcoming pydantic validation
@@ -92,12 +102,13 @@ class StackComponentConfig(BaseModel, ABC):
                         "in sensitive information as secrets. Check out the "
                         "documentation on how to configure your stack "
                         "components with secrets here: "
-                        "https://docs.zenml.io/user-guide/advanced-guide/secret-management"
+                        "https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                     )
                 continue
 
-            requires_validation = field.pre_validators or field.post_validators
-            if requires_validation:
+            if pydantic_utils.has_validators(
+                pydantic_class=self.__class__, field_name=key
+            ):
                 raise ValueError(
                     f"Passing the stack component attribute `{key}` as a "
                     "secret reference is not allowed as additional validation "
@@ -115,7 +126,7 @@ class StackComponentConfig(BaseModel, ABC):
         """
         return {
             secret_utils.parse_secret_reference(v)
-            for v in self.dict().values()
+            for v in self.model_dump().values()
             if secret_utils.is_secret_reference(v)
         }
 
@@ -241,16 +252,70 @@ class StackComponentConfig(BaseModel, ABC):
         # (see https://github.com/python/mypy/issues/13319).
         __getattribute__ = __custom_getattribute__
 
-    class Config:
-        """Pydantic configuration class."""
+    @model_validator(mode="before")
+    @classmethod
+    @pydantic_utils.before_validator_handler
+    def _convert_json_strings(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Converts potential JSON strings.
+
+        Args:
+            data: The model data.
+
+        Returns:
+            The potentially converted data.
 
+        Raises:
+            ValueError: If any of the values is an invalid JSON string.
+        """
+        for key, field in cls.model_fields.items():
+            if not field.annotation:
+                continue
+
+            value = data.get(key, None)
+
+            if isinstance(value, str):
+                if typing_utils.is_optional(field.annotation):
+                    args = list(typing_utils.get_args(field.annotation))
+                    if str in args:
+                        # Don't do any type coercion in case str is in the
+                        # possible types of the field
+                        continue
+
+                    # Remove `NoneType` from the arguments
+                    NoneType = type(None)
+                    if NoneType in args:
+                        args.remove(NoneType)
+
+                    # We just choose the first arg and match against this
+                    annotation = args[0]
+                else:
+                    annotation = field.annotation
+
+                if typing_utils.get_origin(annotation) in {
+                    dict,
+                    list,
+                    Mapping,
+                    Sequence,
+                }:
+                    try:
+                        data[key] = json.loads(value)
+                    except json.JSONDecodeError as e:
+                        raise ValueError(
+                            f"Invalid json string '{value}'"
+                        ) from e
+                elif isclass(annotation) and issubclass(annotation, BaseModel):
+                    data[key] = annotation.model_validate_json(
+                        value
+                    ).model_dump()
+
+        return data
+
+    model_config = ConfigDict(
         # public attributes are immutable
-        allow_mutation = False
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )
 
 
 class StackComponent:
@@ -333,12 +398,7 @@ class StackComponent:
         Raises:
             ImportError: If the flavor can't be imported.
         """
-        from zenml.client import Client
-
-        flavor_model = Client().get_flavor_by_name_and_type(
-            name=component_model.flavor,
-            component_type=component_model.type,
-        )
+        flavor_model = component_model.flavor
 
         try:
             from zenml.stack import Flavor
@@ -356,23 +416,56 @@ class StackComponent:
         else:
             user_id = None
 
-        return flavor.implementation_class(
-            user=user_id,
-            workspace=component_model.workspace.id,
-            name=component_model.name,
-            id=component_model.id,
-            config=configuration,
-            labels=component_model.labels,
-            flavor=component_model.flavor,
-            type=component_model.type,
-            created=component_model.created,
-            updated=component_model.updated,
-            connector_requirements=flavor.service_connector_requirements,
-            connector=component_model.connector.id
-            if component_model.connector
-            else None,
-            connector_resource_id=component_model.connector_resource_id,
-        )
+        try:
+            return flavor.implementation_class(
+                user=user_id,
+                workspace=component_model.workspace.id,
+                name=component_model.name,
+                id=component_model.id,
+                config=configuration,
+                labels=component_model.labels,
+                flavor=component_model.flavor_name,
+                type=component_model.type,
+                created=component_model.created,
+                updated=component_model.updated,
+                connector_requirements=flavor.service_connector_requirements,
+                connector=component_model.connector.id
+                if component_model.connector
+                else None,
+                connector_resource_id=component_model.connector_resource_id,
+            )
+        except ImportError as e:
+            from zenml.integrations.registry import integration_registry
+
+            integration_requirements = " ".join(
+                integration_registry.select_integration_requirements(
+                    flavor_model.integration
+                )
+            )
+
+            if integration_registry.is_installed(flavor_model.integration):
+                raise ImportError(
+                    f"{e}\n\n"
+                    f"Something went wrong while trying to import from the "
+                    f"`{flavor_model.integration}` integration. Please make "
+                    "sure that all its requirements are installed properly by "
+                    "reinstalling the integration either through our CLI: "
+                    f"`zenml integration install {flavor_model.integration} "
+                    "-y` or by manually installing its requirements: "
+                    f"`pip install {integration_requirements}`. If the error "
+                    "persists, please contact the ZenML team."
+                ) from e
+            else:
+                raise ImportError(
+                    f"{e}\n\n"
+                    f"The `{flavor_model.integration}` integration that you "
+                    "are trying to use is not installed in your current "
+                    "environment. Please make sure that it is installed by "
+                    "either using our CLI: `zenml integration install "
+                    f"{flavor_model.integration}` or by manually installing "
+                    f"its requirements: `pip install "
+                    f"{integration_requirements}`"
+                ) from e
 
     @property
     def config(self) -> StackComponentConfig:
@@ -439,7 +532,7 @@ class StackComponent:
         )
 
         if key in all_settings:
-            return self.settings_class.parse_obj(all_settings[key])
+            return self.settings_class.model_validate(dict(all_settings[key]))
         else:
             return self.settings_class()
 
@@ -691,76 +784,9 @@ class StackComponent:
         """
         return None
 
-    @property
-    def is_provisioned(self) -> bool:
-        """If the component provisioned resources to run.
-
-        Returns:
-            True if the component provisioned resources to run.
-        """
-        return True
-
-    @property
-    def is_running(self) -> bool:
-        """If the component is running.
-
-        Returns:
-            True if the component is running.
-        """
-        return True
-
-    @property
-    def is_suspended(self) -> bool:
-        """If the component is suspended.
-
-        Returns:
-            True if the component is suspended.
-        """
-        return not self.is_running
-
-    def provision(self) -> None:
-        """Provisions resources to run the component.
-
-        Raises:
-            NotImplementedError: If the component does not implement this
-                method.
-        """
-        raise NotImplementedError(
-            f"Provisioning resources not implemented for {self}."
-        )
-
-    def deprovision(self) -> None:
-        """Deprovisions all resources of the component.
-
-        Raises:
-            NotImplementedError: If the component does not implement this
-                method.
-        """
-        raise NotImplementedError(
-            f"Deprovisioning resource not implemented for {self}."
-        )
-
-    def resume(self) -> None:
-        """Resumes the provisioned resources of the component.
-
-        Raises:
-            NotImplementedError: If the component does not implement this
-                method.
-        """
-        raise NotImplementedError(
-            f"Resuming provisioned resources not implemented for {self}."
-        )
-
-    def suspend(self) -> None:
-        """Suspends the provisioned resources of the component.
-
-        Raises:
-            NotImplementedError: If the component does not implement this
-                method.
-        """
-        raise NotImplementedError(
-            f"Suspending provisioned resources not implemented for {self}."
-        )
+    def cleanup(self) -> None:
+        """Cleans up the component after it has been used."""
+        pass
 
     def __repr__(self) -> str:
         """String representation of the stack component.
@@ -769,7 +795,7 @@ class StackComponent:
             A string representation of the stack component.
         """
         attribute_representation = ", ".join(
-            f"{key}={value}" for key, value in self.config.dict().items()
+            f"{key}={value}" for key, value in self.config.model_dump().items()
         )
         return (
             f"{self.__class__.__qualname__}(type={self.type}, "