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

zenml/client_lazy_loader.py

@@ -0,0 +1,224 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Lazy loading functionality for Client methods."""
+
+import contextlib
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type
+
+from pydantic import BaseModel, Field
+
+from zenml.logger import get_logger
+
+if TYPE_CHECKING:
+    from zenml.client import Client
+
+logger = get_logger(__name__)
+
+
+class _CallStep(BaseModel):
+    attribute_name: Optional[str] = None
+    is_call: bool = False
+    call_args: List[Any] = Field(default_factory=list)
+    call_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    selector: Optional[Any] = None
+
+
+class ClientLazyLoader(BaseModel):
+    """Lazy loader for Client methods."""
+
+    method_name: str
+    call_chain: List[_CallStep] = []
+    exclude_next_call: bool = False
+
+    def __getattr__(self, name: str) -> "ClientLazyLoader":
+        """Get attribute not defined in ClientLazyLoader.
+
+        Args:
+            name: Name of the attribute to get.
+
+        Returns:
+            self
+        """
+        self_ = ClientLazyLoader(
+            method_name=self.method_name, call_chain=self.call_chain.copy()
+        )
+        # workaround to protect from infinitely looping over in deepcopy called in invocations
+        if name != "__deepcopy__":
+            self_.call_chain.append(_CallStep(attribute_name=name))
+        else:
+            self_.exclude_next_call = True
+        return self_
+
+    def __call__(self, *args: Any, **kwargs: Any) -> "ClientLazyLoader":
+        """Call mocked attribute.
+
+        Args:
+            args: Positional arguments.
+            kwargs: Keyword arguments.
+
+        Returns:
+            self
+        """
+        # workaround to protect from infinitely looping over in deepcopy called in invocations
+        if not self.exclude_next_call:
+            self.call_chain.append(
+                _CallStep(is_call=True, call_args=args, call_kwargs=kwargs)
+            )
+        self.exclude_next_call = False
+        return self
+
+    def __getitem__(self, item: Any) -> "ClientLazyLoader":
+        """Get item from mocked attribute.
+
+        Args:
+            item: Item to get.
+
+        Returns:
+            self
+        """
+        self.call_chain.append(_CallStep(selector=item))
+        return self
+
+    def evaluate(self) -> Any:
+        """Evaluate lazy loaded Client method.
+
+        Returns:
+            Evaluated lazy loader chain of calls.
+        """
+        from zenml.client import Client
+
+        def _iterate_over_lazy_chain(
+            self: "ClientLazyLoader", self_: Any, call_chain_: List[_CallStep]
+        ) -> Any:
+            next_step = call_chain_.pop(0)
+            try:
+                if next_step.is_call:
+                    self_ = self_(
+                        *next_step.call_args, **next_step.call_kwargs
+                    )
+                elif next_step.selector:
+                    self_ = self_[next_step.selector]
+                elif next_step.attribute_name:
+                    self_ = getattr(self_, next_step.attribute_name)
+                else:
+                    raise ValueError(
+                        "Invalid call chain. Reach out to the ZenML team."
+                    )
+            except Exception as e:
+                logger.debug(
+                    f"Failed to evaluate lazy load chain `{self.method_name}` "
+                    f"+ `{next_step}` + `{self.call_chain}`."
+                )
+                msg = f"`{self.method_name}("
+                if next_step:
+                    for arg in next_step.call_args:
+                        msg += f"'{arg}',"
+                    for k, v in next_step.call_kwargs.items():
+                        msg += f"{k}='{v}',"
+                    msg = msg[:-1]
+                msg += f")` failed during lazy load with error: {e}"
+                logger.error(msg)
+                raise RuntimeError(msg)
+            return self_
+
+        self_ = getattr(Client(), self.method_name)
+        call_chain_ = self.call_chain.copy()
+        while call_chain_:
+            self_ = _iterate_over_lazy_chain(self, self_, call_chain_)
+        return self_
+
+
+def client_lazy_loader(
+    method_name: str, *args: Any, **kwargs: Any
+) -> Optional[ClientLazyLoader]:
+    """Lazy loader for Client methods helper.
+
+    Usage:
+    ```
+    def get_something(self, arg1: Any)->SomeResponse:
+        if cll:=client_lazy_loader("get_something", arg1):
+            return cll # type: ignore[return-value]
+        return SomeResponse()
+    ```
+
+    Args:
+        method_name: The name of the method to be called.
+        *args: The arguments to be passed to the method.
+        **kwargs: The keyword arguments to be passed to the method.
+
+    Returns:
+        The result of the method call.
+    """
+    from zenml import get_pipeline_context
+
+    try:
+        get_pipeline_context()
+        cll = ClientLazyLoader(
+            method_name=method_name,
+        )
+        return cll(*args, **kwargs)
+    except RuntimeError:
+        return None
+
+
+def evaluate_all_lazy_load_args_in_client_methods(
+    cls: Type["Client"],
+) -> Type["Client"]:
+    """Class wrapper to evaluate lazy loader arguments of all methods.
+
+    Args:
+        cls: The class to wrap.
+
+    Returns:
+        Wrapped class.
+    """
+    import inspect
+
+    def _evaluate_args(
+        func: Callable[..., Any], is_instance_method: bool
+    ) -> Any:
+        def _inner(*args: Any, **kwargs: Any) -> Any:
+            args_ = list(args)
+            if not is_instance_method:
+                from zenml.client import Client
+
+                if args and isinstance(args[0], Client):
+                    args_ = list(args[1:])
+
+            for i in range(len(args_)):
+                if isinstance(args_[i], dict):
+                    with contextlib.suppress(ValueError):
+                        args_[i] = ClientLazyLoader(**args_[i]).evaluate()
+                elif isinstance(args_[i], ClientLazyLoader):
+                    args_[i] = args_[i].evaluate()
+
+            for k, v in kwargs.items():
+                if isinstance(v, dict):
+                    with contextlib.suppress(ValueError):
+                        kwargs[k] = ClientLazyLoader(**v).evaluate()
+
+            return func(*args_, **kwargs)
+
+        return _inner
+
+    def _decorate() -> Type["Client"]:
+        for name, fn in inspect.getmembers(cls, inspect.isfunction):
+            setattr(
+                cls,
+                name,
+                _evaluate_args(fn, "self" in inspect.getfullargspec(fn).args),
+            )
+        return cls
+
+    return _decorate()

zenml/code_repositories/base_code_repository.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Base class for code repositories."""
+
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Type
 from uuid import UUID
@@ -74,10 +75,10 @@ class BaseCodeRepository(ABC):
         Returns:
             The loaded code repository object.
         """
-        class_: Type[
-            BaseCodeRepository
-        ] = source_utils.load_and_validate_class(
-            source=model.source, expected_class=BaseCodeRepository
+        class_: Type[BaseCodeRepository] = (
+            source_utils.load_and_validate_class(
+                source=model.source, expected_class=BaseCodeRepository
+            )
         )
         return class_(id=model.id, config=model.config)
 

zenml/code_repositories/git/local_git_repository_context.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Implementation of the Local git repository context."""
+
 from typing import TYPE_CHECKING, Callable, Optional, cast
 from uuid import UUID
 

zenml/code_repositories/local_repository_context.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Base class for local code repository contexts."""
+
 from abc import ABC, abstractmethod
 from uuid import UUID
 

zenml/config/__init__.py

@@ -25,8 +25,10 @@ order to persist the configuration across sessions.
 """
 from zenml.config.docker_settings import DockerSettings
 from zenml.config.resource_settings import ResourceSettings
+from zenml.config.retry_config import StepRetryConfig
 
 __all__ = [
     "DockerSettings",
     "ResourceSettings",
+    "StepRetryConfig",
 ]

zenml/config/base_settings.py

@@ -12,10 +12,11 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Base class for all ZenML settings."""
+
 from enum import IntFlag, auto
 from typing import Any, ClassVar, Dict, Union
 
-from pydantic import Extra
+from pydantic import ConfigDict
 
 from zenml.config.secret_reference_mixin import SecretReferenceMixin
 
@@ -45,11 +46,10 @@ class BaseSettings(SecretReferenceMixin):
         ConfigurationLevel.PIPELINE | ConfigurationLevel.STEP
     )
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # public attributes are immutable
-        allow_mutation = False
+        frozen=True,
         # allow extra attributes so this class can be used to parse dicts
         # of arbitrary subclasses
-        extra = Extra.allow
+        extra="allow",
+    )

zenml/config/build_configuration.py

@@ -14,11 +14,13 @@
 """Build configuration class."""
 
 import hashlib
+import json
 from typing import TYPE_CHECKING, Dict, Optional
 
 from pydantic import BaseModel
 
-from zenml.config.docker_settings import DockerSettings, SourceFileMode
+from zenml.config.docker_settings import DockerSettings
+from zenml.utils import json_utils
 
 if TYPE_CHECKING:
     from zenml.code_repositories import BaseCodeRepository
@@ -60,7 +62,14 @@ class BuildConfiguration(BaseModel):
             The checksum.
         """
         hash_ = hashlib.md5()  # nosec
-        hash_.update(self.settings.json().encode())
+        settings_json = json.dumps(
+            self.settings.model_dump(
+                mode="json", exclude={"prevent_build_reuse"}
+            ),
+            sort_keys=True,
+            default=json_utils.pydantic_encoder,
+        )
+        hash_.update(settings_json.encode())
         if self.entrypoint:
             hash_.update(self.entrypoint.encode())
 
@@ -72,7 +81,7 @@ class BuildConfiguration(BaseModel):
             PipelineDockerImageBuilder,
         )
 
-        pass_code_repo = self.should_download_files(
+        pass_code_repo = self.should_download_files_from_code_repository(
             code_repository=code_repository
         )
         requirements_files = (
@@ -101,34 +110,51 @@ class BuildConfiguration(BaseModel):
         Returns:
             Whether files should be included in the image.
         """
-        if self.settings.source_files == SourceFileMode.INCLUDE:
-            return True
+        if self.should_download_files(code_repository=code_repository):
+            return False
 
-        if (
-            self.settings.source_files == SourceFileMode.DOWNLOAD_OR_INCLUDE
-            and not code_repository
+        return self.settings.allow_including_files_in_images
+
+    def should_download_files(
+        self,
+        code_repository: Optional["BaseCodeRepository"],
+    ) -> bool:
+        """Whether files should be downloaded in the image.
+
+        Args:
+            code_repository: Code repository that can be used to download files
+                inside the image.
+
+        Returns:
+            Whether files should be downloaded in the image.
+        """
+        if self.should_download_files_from_code_repository(
+            code_repository=code_repository
         ):
             return True
 
+        if self.settings.allow_download_from_artifact_store:
+            return True
+
         return False
 
-    def should_download_files(
+    def should_download_files_from_code_repository(
         self,
         code_repository: Optional["BaseCodeRepository"],
     ) -> bool:
-        """Whether files should be downloaded in the image.
+        """Whether files should be downloaded from the code repository.
 
         Args:
             code_repository: Code repository that can be used to download files
                 inside the image.
 
         Returns:
-            Whether files should be downloaded in the image.
+            Whether files should be downloaded from the code repository.
         """
-        if not code_repository:
-            return False
+        if (
+            code_repository
+            and self.settings.allow_download_from_code_repository
+        ):
+            return True
 
-        return self.settings.source_files in {
-            SourceFileMode.DOWNLOAD,
-            SourceFileMode.DOWNLOAD_OR_INCLUDE,
-        }
+        return False

zenml/config/compiler.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Class for compiling ZenML pipelines into a serializable format."""
+
 import copy
 import string
 from typing import (
@@ -38,11 +39,12 @@ from zenml.config.step_configurations import (
 from zenml.environment import get_run_environment_dict
 from zenml.exceptions import StackValidationError
 from zenml.models import PipelineDeploymentBase
+from zenml.pipelines.run_utils import get_default_run_name
 from zenml.utils import pydantic_utils, settings_utils
 
 if TYPE_CHECKING:
     from zenml.config.source import Source
-    from zenml.new.pipelines.pipeline import Pipeline
+    from zenml.pipelines.pipeline_definition import Pipeline
     from zenml.stack import Stack, StackComponent
     from zenml.steps.step_invocation import StepInvocation
 
@@ -73,7 +75,7 @@ class Compiler:
         pipeline: "Pipeline",
         stack: "Stack",
         run_configuration: PipelineRunConfiguration,
-    ) -> Tuple[PipelineDeploymentBase, PipelineSpec]:
+    ) -> PipelineDeploymentBase:
         """Compiles a ZenML pipeline to a serializable representation.
 
         Args:
@@ -82,7 +84,7 @@ class Compiler:
             run_configuration: The run configuration for this pipeline.
 
         Returns:
-            The compiled pipeline deployment and spec
+            The compiled pipeline deployment.
         """
         logger.debug("Compiling pipeline `%s`.", pipeline.name)
         # Copy the pipeline before we apply any run-level configurations, so
@@ -91,9 +93,16 @@ class Compiler:
         self._apply_run_configuration(
             pipeline=pipeline, config=run_configuration
         )
+        convert_component_shortcut_settings_keys(
+            pipeline.configuration.settings, stack=stack
+        )
+
         self._apply_stack_default_settings(pipeline=pipeline, stack=stack)
         if run_configuration.run_name:
-            self._verify_run_name(run_configuration.run_name)
+            self._verify_run_name(
+                run_configuration.run_name,
+                pipeline.configuration.substitutions,
+            )
 
         pipeline_settings = self._filter_and_validate_settings(
             settings=pipeline.configuration.settings,
@@ -126,12 +135,17 @@ class Compiler:
 
         self._ensure_required_stack_components_exist(stack=stack, steps=steps)
 
-        run_name = run_configuration.run_name or self._get_default_run_name(
+        run_name = run_configuration.run_name or get_default_run_name(
             pipeline_name=pipeline.name
         )
 
         client_version, server_version = get_zenml_versions()
 
+        step_specs = [step.spec for step in steps.values()]
+        pipeline_spec = self._compute_pipeline_spec(
+            pipeline=pipeline, step_specs=step_specs
+        )
+
         deployment = PipelineDeploymentBase(
             run_name_template=run_name,
             pipeline_configuration=pipeline.configuration,
@@ -139,17 +153,15 @@ class Compiler:
             client_environment=get_run_environment_dict(),
             client_version=client_version,
             server_version=server_version,
-        )
-
-        step_specs = [step.spec for step in steps.values()]
-        pipeline_spec = self._compute_pipeline_spec(
-            pipeline=pipeline, step_specs=step_specs
+            pipeline_version_hash=pipeline._compute_unique_identifier(
+                pipeline_spec=pipeline_spec
+            ),
+            pipeline_spec=pipeline_spec,
         )
 
         logger.debug("Compiled pipeline deployment: %s", deployment)
-        logger.debug("Compiled pipeline spec: %s", pipeline_spec)
 
-        return deployment, pipeline_spec
+        return deployment
 
     def compile_spec(self, pipeline: "Pipeline") -> PipelineSpec:
         """Compiles a ZenML pipeline to a pipeline spec.
@@ -191,10 +203,6 @@ class Compiler:
         Args:
             pipeline: The pipeline to configure.
             config: The run configurations.
-
-        Raises:
-            KeyError: If the run configuration contains options for a
-                non-existent step.
         """
         with pipeline.__suppress_configure_warnings__():
             pipeline.configure(
@@ -203,16 +211,22 @@ class Compiler:
                 enable_artifact_visualization=config.enable_artifact_visualization,
                 enable_step_logs=config.enable_step_logs,
                 settings=config.settings,
+                tags=config.tags,
                 extra=config.extra,
                 model=config.model,
                 parameters=config.parameters,
             )
 
-        for invocation_id in config.steps:
-            if invocation_id not in pipeline.invocations:
-                raise KeyError(
-                    f"Configuration for step {invocation_id} cannot be applied to any pipeline step. Make sure that all configured steps are present in your pipeline."
-                )
+        invalid_step_configs = set(config.steps) - set(pipeline.invocations)
+        if invalid_step_configs:
+            logger.warning(
+                f"Configuration for step invocations {invalid_step_configs} "
+                "cannot be applied to any pipeline step invocations, "
+                "ignoring..."
+            )
+
+        for key in invalid_step_configs:
+            config.steps.pop(key)
 
         # Override `enable_cache` of all steps if set at run level
         if config.enable_cache is not None:
@@ -285,25 +299,31 @@ class Compiler:
         """
         assert stack_component.settings_class
         # Exclude additional config attributes that aren't part of the settings
-        field_names = set(stack_component.settings_class.__fields__)
-        default_settings = stack_component.settings_class.parse_obj(
-            stack_component.config.dict(
+        field_names = set(stack_component.settings_class.model_fields.keys())
+        default_settings = stack_component.settings_class.model_validate(
+            stack_component.config.model_dump(
                 include=field_names, exclude_unset=True, exclude_defaults=True
             )
         )
         return default_settings
 
     @staticmethod
-    def _verify_run_name(run_name: str) -> None:
+    def _verify_run_name(
+        run_name: str,
+        substitutions: Dict[str, str],
+    ) -> None:
         """Verifies that the run name contains only valid placeholders.
 
         Args:
             run_name: The run name to verify.
+            substitutions: The substitutions to be used in the run name.
 
         Raises:
             ValueError: If the run name contains invalid placeholders.
         """
-        valid_placeholder_names = {"date", "time"}
+        valid_placeholder_names = {"date", "time"}.union(
+            set(substitutions.keys())
+        )
         placeholders = {
             v[1] for v in string.Formatter().parse(run_name) if v[1]
         }
@@ -439,6 +459,9 @@ class Compiler:
                 step_config, runtime_parameters=invocation.parameters
             )
 
+        convert_component_shortcut_settings_keys(
+            step.configuration.settings, stack=stack
+        )
         step_spec = self._get_step_spec(invocation=invocation)
         step_settings = self._filter_and_validate_settings(
             settings=step.configuration.settings,
@@ -472,18 +495,6 @@ class Compiler:
         )
         return Step(spec=step_spec, config=complete_step_configuration)
 
-    @staticmethod
-    def _get_default_run_name(pipeline_name: str) -> str:
-        """Gets the default name for a pipeline run.
-
-        Args:
-            pipeline_name: Name of the pipeline which will be run.
-
-        Returns:
-            Run name.
-        """
-        return f"{pipeline_name}-{{date}}-{{time}}"
-
     def _get_sorted_invocations(
         self,
         pipeline: "Pipeline",
@@ -548,7 +559,7 @@ class Compiler:
             step_operator = step.config.step_operator
             if step_operator and step_operator not in available_step_operators:
                 raise StackValidationError(
-                    f"Step '{name}' requires step operator "
+                    f"Step `{name}` requires step operator "
                     f"'{step_operator}' which is not configured in "
                     f"the stack '{stack.name}'. Available step operators: "
                     f"{available_step_operators}."
@@ -560,7 +571,7 @@ class Compiler:
                 and experiment_tracker not in available_experiment_trackers
             ):
                 raise StackValidationError(
-                    f"Step '{name}' requires experiment tracker "
+                    f"Step `{name}` requires experiment tracker "
                     f"'{experiment_tracker}' which is not "
                     f"configured in the stack '{stack.name}'. Available "
                     f"experiment trackers: {available_experiment_trackers}."
@@ -582,8 +593,6 @@ class Compiler:
         Raises:
             ValueError: If the pipeline has no steps.
         """
-        from zenml.pipelines import BasePipeline
-
         if not step_specs:
             raise ValueError(
                 f"Pipeline '{pipeline.name}' cannot be compiled because it has "
@@ -593,12 +602,36 @@ class Compiler:
                 "https://docs.zenml.io/user-guide/starter-guide."
             )
 
-        additional_spec_args: Dict[str, Any] = {}
-        if isinstance(pipeline, BasePipeline):
-            # use older spec version for legacy pipelines
-            additional_spec_args["version"] = "0.3"
-        else:
-            additional_spec_args["source"] = pipeline.resolve()
-            additional_spec_args["parameters"] = pipeline._parameters
+        additional_spec_args: Dict[str, Any] = {
+            "source": pipeline.resolve(),
+            "parameters": pipeline._parameters,
+        }
 
         return PipelineSpec(steps=step_specs, **additional_spec_args)
+
+
+def convert_component_shortcut_settings_keys(
+    settings: Dict[str, "BaseSettings"], stack: "Stack"
+) -> None:
+    """Convert component shortcut settings keys.
+
+    Args:
+        settings: Dictionary of settings.
+        stack: The stack that the pipeline will run on.
+
+    Raises:
+        ValueError: If stack component settings were defined both using the
+            full and the shortcut key.
+    """
+    for component in stack.components.values():
+        shortcut_key = str(component.type)
+        if component_settings := settings.pop(shortcut_key, None):
+            key = settings_utils.get_stack_component_setting_key(component)
+            if key in settings:
+                raise ValueError(
+                    f"Duplicate settings provided for your {shortcut_key} "
+                    f"using the keys {shortcut_key} and {key}. Remove settings "
+                    "for one of them to fix this error."
+                )
+
+            settings[key] = component_settings