zenml-nightly 0.58.0.dev20240529__py3-none-any.whl → 0.58.1.dev20240608__py3-none-any.whl

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -44,7 +44,9 @@ class KubernetesOrchestratorSettings(BaseSettings):
         step_pod_service_account_name: Name of the service account to use for the
             step pods. If not provided, the default service account will be used.
         privileged: If the container should be run in privileged mode.
-        pod_settings: Pod settings to apply.
+        pod_settings: Pod settings to apply to pods executing the steps.
+        orchestrator_pod_settings: Pod settings to apply to the pod which is
+            launching the actual steps.
     """
 
     synchronous: bool = True
@@ -53,6 +55,7 @@ class KubernetesOrchestratorSettings(BaseSettings):
     step_pod_service_account_name: Optional[str] = None
     privileged: bool = False
     pod_settings: Optional[KubernetesPodSettings] = None
+    orchestrator_pod_settings: Optional[KubernetesPodSettings] = None
 
 
 class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -391,20 +391,6 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         # Authorize pod to run Kubernetes commands inside the cluster.
         service_account_name = self._get_service_account_name(settings)
 
-        if settings.pod_settings:
-            # Remove all settings that specify on which pod to run for the
-            # orchestrator pod. These settings should only be used
-            # for the pods executing the actual steps.
-            pod_settings = settings.pod_settings.copy(
-                update={
-                    "resources": {},
-                    "node_selectors": {},
-                    "affinity": {},
-                    "tolerations": [],
-                }
-            )
-            settings = settings.copy(update={"pod_settings": pod_settings})
-
         # Schedule as CRON job if CRON schedule is given.
         if deployment.schedule:
             if not deployment.schedule.cron_expression:
@@ -423,7 +409,8 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 command=command,
                 args=args,
                 service_account_name=service_account_name,
-                settings=settings,
+                privileged=False,
+                pod_settings=settings.orchestrator_pod_settings,
                 env=environment,
                 mount_local_stores=self.config.is_local,
             )
@@ -446,8 +433,9 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             image_name=image,
             command=command,
             args=args,
+            privileged=False,
+            pod_settings=settings.orchestrator_pod_settings,
             service_account_name=service_account_name,
-            settings=settings,
             env=environment,
             mount_local_stores=self.config.is_local,
         )

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -116,7 +116,8 @@ def main() -> None:
             command=step_command,
             args=step_args,
             env=env,
-            settings=settings,
+            privileged=settings.privileged,
+            pod_settings=settings.pod_settings,
             service_account_name=settings.step_pod_service_account_name
             or settings.service_account_name,
             mount_local_stores=mount_local_stores,

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -25,9 +25,6 @@ from zenml.constants import ENV_ZENML_ENABLE_REPO_INIT_WARNINGS
 from zenml.integrations.airflow.orchestrators.dag_generator import (
     ENV_ZENML_LOCAL_STORES_PATH,
 )
-from zenml.integrations.kubernetes.flavors import (
-    KubernetesOrchestratorSettings,
-)
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 
 
@@ -100,7 +97,8 @@ def build_pod_manifest(
     image_name: str,
     command: List[str],
     args: List[str],
-    settings: KubernetesOrchestratorSettings,
+    privileged: bool,
+    pod_settings: Optional[KubernetesPodSettings] = None,
     service_account_name: Optional[str] = None,
     env: Optional[Dict[str, str]] = None,
     mount_local_stores: bool = False,
@@ -114,7 +112,8 @@ def build_pod_manifest(
         image_name: Name of the Docker image.
         command: Command to execute the entrypoint in the pod.
         args: Arguments provided to the entrypoint command.
-        settings: `KubernetesOrchestratorSettings` object
+        privileged: Whether to run the container in privileged mode.
+        pod_settings: Optional settings for the pod.
         service_account_name: Optional name of a service account.
             Can be used to assign certain roles to a pod, e.g., to allow it to
             run Kubernetes commands from within the cluster.
@@ -128,9 +127,7 @@ def build_pod_manifest(
     env = env.copy() if env else {}
     env.setdefault(ENV_ZENML_ENABLE_REPO_INIT_WARNINGS, "False")
 
-    security_context = k8s_client.V1SecurityContext(
-        privileged=settings.privileged
-    )
+    security_context = k8s_client.V1SecurityContext(privileged=privileged)
     container_spec = k8s_client.V1Container(
         name="main",
         image=image_name,
@@ -151,8 +148,8 @@ def build_pod_manifest(
     if service_account_name is not None:
         pod_spec.service_account_name = service_account_name
 
-    if settings.pod_settings:
-        add_pod_settings(pod_spec, settings.pod_settings)
+    if pod_settings:
+        add_pod_settings(pod_spec, pod_settings)
 
     pod_metadata = k8s_client.V1ObjectMeta(
         name=pod_name,
@@ -162,8 +159,8 @@ def build_pod_manifest(
         },
     )
 
-    if settings.pod_settings and settings.pod_settings.annotations:
-        pod_metadata.annotations = settings.pod_settings.annotations
+    if pod_settings and pod_settings.annotations:
+        pod_metadata.annotations = pod_settings.annotations
 
     pod_manifest = k8s_client.V1Pod(
         kind="Pod",
@@ -224,7 +221,8 @@ def build_cron_job_manifest(
     image_name: str,
     command: List[str],
     args: List[str],
-    settings: KubernetesOrchestratorSettings,
+    privileged: bool,
+    pod_settings: Optional[KubernetesPodSettings] = None,
     service_account_name: Optional[str] = None,
     env: Optional[Dict[str, str]] = None,
     mount_local_stores: bool = False,
@@ -239,7 +237,8 @@ def build_cron_job_manifest(
         image_name: Name of the Docker image.
         command: Command to execute the entrypoint in the pod.
         args: Arguments provided to the entrypoint command.
-        settings: `KubernetesOrchestratorSettings` object
+        privileged: Whether to run the container in privileged mode.
+        pod_settings: Optional settings for the pod.
         service_account_name: Optional name of a service account.
             Can be used to assign certain roles to a pod, e.g., to allow it to
             run Kubernetes commands from within the cluster.
@@ -257,7 +256,8 @@ def build_cron_job_manifest(
         image_name=image_name,
         command=command,
         args=args,
-        settings=settings,
+        privileged=privileged,
+        pod_settings=pod_settings,
         service_account_name=service_account_name,
         env=env,
         mount_local_stores=mount_local_stores,

zenml/integrations/s3/flavors/s3_artifact_store_flavor.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Amazon S3 artifact store flavor."""
 
-import json
 import re
 from typing import (
     TYPE_CHECKING,
@@ -23,7 +22,6 @@ from typing import (
     Optional,
     Set,
     Type,
-    Union,
 )
 
 from pydantic import validator
@@ -72,42 +70,6 @@ class S3ArtifactStoreConfig(
     config_kwargs: Optional[Dict[str, Any]] = None
     s3_additional_kwargs: Optional[Dict[str, Any]] = None
 
-    @validator(
-        "client_kwargs", "config_kwargs", "s3_additional_kwargs", pre=True
-    )
-    def _convert_json_string(
-        cls, value: Union[None, str, Dict[str, Any]]
-    ) -> Optional[Dict[str, Any]]:
-        """Converts potential JSON strings passed via the CLI to dictionaries.
-
-        Args:
-            value: The value to convert.
-
-        Returns:
-            The converted value.
-
-        Raises:
-            TypeError: If the value is not a `str`, `Dict` or `None`.
-            ValueError: If the value is an invalid json string or a json string
-                that does not decode into a dictionary.
-        """
-        if isinstance(value, str):
-            try:
-                dict_ = json.loads(value)
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Invalid json string '{value}'") from e
-
-            if not isinstance(dict_, Dict):
-                raise ValueError(
-                    f"Json string '{value}' did not decode into a dictionary."
-                )
-
-            return dict_
-        elif isinstance(value, Dict) or value is None:
-            return value
-        else:
-            raise TypeError(f"{value} is not a json string or a dictionary.")
-
     @validator("client_kwargs")
     def _validate_client_kwargs(
         cls, value: Optional[Dict[str, Any]]

zenml/integrations/spark/flavors/spark_step_operator_flavor.py

@@ -13,10 +13,7 @@
 #  permissions and limitations under the License.
 """Spark step operator flavor."""
 
-import json
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
-
-from pydantic import validator
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
 from zenml.config.base_settings import BaseSettings
 from zenml.step_operators.base_step_operator import (
@@ -45,40 +42,6 @@ class SparkStepOperatorSettings(BaseSettings):
     deploy_mode: str = "cluster"
     submit_kwargs: Optional[Dict[str, Any]] = None
 
-    @validator("submit_kwargs", pre=True)
-    def _convert_json_string(
-        cls, value: Union[None, str, Dict[str, Any]]
-    ) -> Optional[Dict[str, Any]]:
-        """Converts potential JSON strings passed via the CLI to dictionaries.
-
-        Args:
-            value: The value to convert.
-
-        Returns:
-            The converted value.
-
-        Raises:
-            TypeError: If the value is not a `str`, `Dict` or `None`.
-            ValueError: If the value is an invalid json string or a json string
-                that does not decode into a dictionary.
-        """
-        if isinstance(value, str):
-            try:
-                dict_ = json.loads(value)
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Invalid json string '{value}'") from e
-
-            if not isinstance(dict_, Dict):
-                raise ValueError(
-                    f"Json string '{value}' did not decode into a dictionary."
-                )
-
-            return dict_
-        elif isinstance(value, Dict) or value is None:
-            return value
-        else:
-            raise TypeError(f"{value} is not a json string or a dictionary.")
-
 
 class SparkStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
     BaseStepOperatorConfig, SparkStepOperatorSettings

zenml/materializers/base_materializer.py

@@ -65,7 +65,7 @@ class BaseMaterializerMeta(type):
                     f"Invalid materializer class '{name}'. When creating a "
                     f"custom materializer, make sure to specify at least one "
                     f"type in its ASSOCIATED_TYPES class variable.",
-                    url="https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types",
+                    url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
                 )
 
             # Validate associated artifact type.
@@ -80,7 +80,7 @@ class BaseMaterializerMeta(type):
                         f"custom materializer, make sure to specify a valid "
                         f"artifact type in its ASSOCIATED_ARTIFACT_TYPE class "
                         f"variable.",
-                        url="https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types",
+                        url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
                     )
 
             # Validate associated data types.
@@ -89,7 +89,7 @@ class BaseMaterializerMeta(type):
                     raise MaterializerInterfaceError(
                         f"Associated type {associated_type} for materializer "
                         f"{name} is not a class.",
-                        url="https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types",
+                        url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
                     )
 
             # Register the materializer.

zenml/materializers/cloudpickle_materializer.py

@@ -102,7 +102,7 @@ class CloudpickleMaterializer(BaseMaterializer):
                 "the artifacts cannot be loaded when running with a different "
                 "Python version. Please consider implementing a custom "
                 f"materializer for type `{type(data)}` according to the "
-                "instructions at https://docs.zenml.io/user-guide/advanced-guide/artifact-management/handle-custom-data-types"
+                "instructions at https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types"
             )
 
         # save python version for validation on loading

zenml/orchestrators/base_orchestrator.py

@@ -51,7 +51,7 @@ class BaseOrchestratorConfig(StackComponentConfig):
                     "The 'custom_docker_base_image_name' field has been "
                     "deprecated. To use a custom base container image with your "
                     "orchestrators, please use the DockerSettings in your "
-                    "pipeline (see https://docs.zenml.io/user-guide/advanced-guide/environment-management/containerize-your-pipeline)."
+                    "pipeline (see https://docs.zenml.io/how-to/customize-docker-builds)."
                 )
 
         return values

zenml/orchestrators/local_docker/local_docker_orchestrator.py

@@ -14,15 +14,13 @@
 """Implementation of the ZenML local Docker orchestrator."""
 
 import copy
-import json
 import os
 import sys
 import time
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
 from uuid import uuid4
 
 from docker.errors import ContainerError
-from pydantic import validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.config.global_config import GlobalConfiguration
@@ -207,40 +205,6 @@ class LocalDockerOrchestratorSettings(BaseSettings):
 
     run_args: Dict[str, Any] = {}
 
-    @validator("run_args", pre=True)
-    def _convert_json_string(
-        cls, value: Union[None, str, Dict[str, Any]]
-    ) -> Optional[Dict[str, Any]]:
-        """Converts potential JSON strings passed via the CLI to dictionaries.
-
-        Args:
-            value: The value to convert.
-
-        Returns:
-            The converted value.
-
-        Raises:
-            TypeError: If the value is not a `str`, `Dict` or `None`.
-            ValueError: If the value is an invalid json string or a json string
-                that does not decode into a dictionary.
-        """
-        if isinstance(value, str):
-            try:
-                dict_ = json.loads(value)
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Invalid json string '{value}'") from e
-
-            if not isinstance(dict_, Dict):
-                raise ValueError(
-                    f"Json string '{value}' did not decode into a dictionary."
-                )
-
-            return dict_
-        elif isinstance(value, Dict) or value is None:
-            return value
-        else:
-            raise TypeError(f"{value} is not a json string or a dictionary.")
-
 
 class LocalDockerOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
     BaseOrchestratorConfig, LocalDockerOrchestratorSettings

zenml/orchestrators/step_runner.py

@@ -348,7 +348,7 @@ class StepRunner:
                     "is deprecated and will be removed in a future release. "
                     f"Please adjust your '{step_name}' step to instead import "
                     "the `StepContext` inside your step, as shown here: "
-                    "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/fetch-metadata-within-steps"
+                    "https://docs.zenml.io/how-to/track-metrics-metadata/fetch-metadata-within-steps"
                 )
                 function_params[arg] = get_step_context()
             elif arg in input_artifacts:
@@ -409,7 +409,7 @@ class StepRunner:
                     "is deprecated and will be removed in a future release. "
                     f"Please adjust your '{step_name}' hook to instead import "
                     "the `StepContext` inside your hook, as shown here: "
-                    "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/fetch-metadata-within-steps"
+                    "https://docs.zenml.io/how-to/track-metrics-metadata/fetch-metadata-within-steps"
                 )
                 function_params[arg] = get_step_context()
 

zenml/stack/flavor.py

@@ -206,7 +206,9 @@ class Flavor:
         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}"
+        return (
+            f"{base}/stack-components/{component_type}/{docs_component_name}"
+        )
 
     def generate_default_sdk_docs_url(self) -> str:
         """Generate SDK docs url for a flavor.

zenml/stack/stack.py

@@ -801,7 +801,7 @@ class Stack:
                     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"Check out https://docs.zenml.io/how-to/stack-deployment/deploy-a-stack-using-mlstacks"
                     f"for more information."
                 )
 
@@ -811,7 +811,7 @@ class Stack:
                 "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."
             )
 

zenml/stack/stack_component.py

@@ -13,12 +13,16 @@
 #  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, Extra, root_validator
+from pydantic.typing import get_origin
 
 from zenml.config.build_configuration import BuildConfiguration
 from zenml.config.step_configurations import Step
@@ -93,7 +97,7 @@ 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/manage-the-deployed-services/secret-management"
                     )
                 continue
 
@@ -242,6 +246,42 @@ class StackComponentConfig(BaseModel, ABC):
         # (see https://github.com/python/mypy/issues/13319).
         __getattribute__ = __custom_getattribute__
 
+    @root_validator(pre=True)
+    def _convert_json_strings(cls, values: Dict[str, Any]) -> Any:
+        """Converts potential JSON strings.
+
+        Args:
+            values: The model values.
+
+        Returns:
+            The potentially converted values.
+
+        Raises:
+            ValueError: If any of the values is an invalid JSON string.
+        """
+        for key, field in cls.__fields__.items():
+            value = values.get(key, None)
+
+            if isinstance(value, str):
+                if get_origin(field.outer_type_) in {
+                    dict,
+                    list,
+                    Mapping,
+                    Sequence,
+                }:
+                    try:
+                        values[key] = json.loads(value)
+                    except json.JSONDecodeError as e:
+                        raise ValueError(
+                            f"Invalid json string '{value}'"
+                        ) from e
+                elif isclass(field.type_) and issubclass(
+                    field.type_, BaseModel
+                ):
+                    values[key] = field.type_.parse_raw(value).dict()
+
+        return values
+
     class Config:
         """Pydantic configuration class."""
 

zenml/steps/base_step.py

@@ -1233,7 +1233,7 @@ To avoid this consider setting step parameters only in one place (config or code
         logger.warning(
             "The `BaseParameters` class to define step parameters is "
             "deprecated. Check out our docs "
-            "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/configure-steps-pipelines "
+            "https://docs.zenml.io/how-to/use-configuration-files/how-to-use-config "
             "for information on how to parameterize your steps. As a quick "
             "fix to get rid of this warning, make sure your parameter class "
             "inherits from `pydantic.BaseModel` instead of the "

zenml/steps/entrypoint_function_utils.py

@@ -115,7 +115,7 @@ class StepArtifact:
             "you're trying to unpack the return value of your step but the "
             "step only returns a single artifact. For more information on how "
             "to add type annotations to your step to indicate multiple "
-            "artifacts visit https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/managing-steps#type-annotations."
+            "artifacts visit https://docs.zenml.io/how-to/build-pipelines/step-output-typing-and-annotation#type-annotations."
         )
 
 
@@ -214,7 +214,7 @@ class EntrypointFunctionDefinition(NamedTuple):
                 f"'{key}' is not JSON serializable and can not be passed as "
                 "a parameter. This input can either be provided by the "
                 "output of another step or as an external artifact: "
-                "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/configure-steps-pipelines#pass-any-kind-of-data-to-your-steps"
+                "https://docs.zenml.io/user-guide/starter-guide/manage-artifacts#managing-artifacts-not-produced-by-zenml-pipelines"
             )
 
         try:

zenml/steps/utils.py

@@ -110,7 +110,7 @@ def parse_return_type_annotations(
             "Using the `Output` class to define the outputs of your steps is "
             "deprecated. You should instead use the standard Python way of "
             "type annotating your functions. Check out our documentation "
-            "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/configure-steps-pipelines#step-output-names "
+            "https://docs.zenml.io/how-to/build-pipelines/step-output-typing-and-annotation "
             "for more information on how to assign custom names to your step "
             "outputs."
         )

zenml/utils/dashboard_utils.py

@@ -42,7 +42,7 @@ def is_cloud_server(server_info: ServerModel) -> bool:
     """
     return (
         "organization_id" in server_info.metadata
-        and "cloud.zenml.io" in server_info.server_url
+        and "cloud.zenml.io" in server_info.dashboard_url
     )
 
 

zenml/utils/pipeline_docker_image_builder.py

@@ -20,6 +20,7 @@ import sys
 from collections import defaultdict
 from typing import (
     TYPE_CHECKING,
+    Any,
     DefaultDict,
     Dict,
     List,
@@ -65,6 +66,15 @@ DEFAULT_DOCKER_PARENT_IMAGE = (
     f"py{sys.version_info.major}.{sys.version_info.minor}"
 )
 
+PIP_DEFAULT_ARGS = {
+    "no-cache-dir": None,
+    "default-timeout": 60,
+}
+UV_DEFAULT_ARGS = {
+    "no-cache-dir": None,
+    "system": None,
+}
+
 
 class PipelineDockerImageBuilder:
     """Builds Docker images to run a ZenML pipeline."""
@@ -636,22 +646,32 @@ class PipelineDockerImageBuilder:
             docker_settings.python_package_installer
             == PythonPackageInstaller.PIP
         ):
-            install_command = "pip install --default-timeout=60"
+            install_command = "pip install"
+            default_installer_args: Dict[str, Any] = PIP_DEFAULT_ARGS
         elif (
             docker_settings.python_package_installer
             == PythonPackageInstaller.UV
         ):
             lines.append("RUN pip install uv")
-            install_command = "uv pip install --system"
+            install_command = "uv pip install"
+            default_installer_args = UV_DEFAULT_ARGS
         else:
             raise ValueError("Unsupported python package installer.")
 
+        installer_args = {
+            **default_installer_args,
+            **docker_settings.python_package_installer_args,
+        }
+        installer_args_string = " ".join(
+            f"--{key}" if value is None else f"--{key}={value}"
+            for key, value in installer_args.items()
+        )
         for file, _, options in requirements_files:
             lines.append(f"COPY {file} .")
             option_string = " ".join(options)
 
             lines.append(
-                f"RUN {install_command} --no-cache-dir "
+                f"RUN {install_command} {installer_args_string}"
                 f"{option_string} -r {file}"
             )