zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__py3-none-any.whl

zenml/steps/base_step.py

@@ -22,6 +22,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
+    List,
     Mapping,
     Optional,
     Sequence,
@@ -38,7 +39,11 @@ from pydantic import BaseModel, ConfigDict, ValidationError
 from zenml.client_lazy_loader import ClientLazyLoader
 from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import Source
-from zenml.constants import STEP_SOURCE_PARAMETER_NAME
+from zenml.constants import (
+    ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK,
+    STEP_SOURCE_PARAMETER_NAME,
+    handle_bool_env_var,
+)
 from zenml.exceptions import MissingStepParameterError, StepInterfaceError
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
@@ -51,15 +56,16 @@ from zenml.steps.entrypoint_function_utils import (
 )
 from zenml.steps.utils import (
     resolve_type_annotation,
+    run_as_single_step_pipeline,
 )
 from zenml.utils import (
     dict_utils,
+    materializer_utils,
     notebook_utils,
     pydantic_utils,
     settings_utils,
     source_code_utils,
     source_utils,
-    typing_utils,
 )
 
 if TYPE_CHECKING:
@@ -586,9 +592,16 @@ class BaseStep(metaclass=BaseStepMeta):
         from zenml.new.pipelines.pipeline import Pipeline
 
         if not Pipeline.ACTIVE_PIPELINE:
-            # The step is being called outside the context of a pipeline,
-            # we simply call the entrypoint
-            return self.call_entrypoint(*args, **kwargs)
+            # The step is being called outside the context of a pipeline, either
+            # run the step function or run it as a single step pipeline on the
+            # active stack
+            run_without_stack = handle_bool_env_var(
+                ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK, default=False
+            )
+            if run_without_stack:
+                return self.call_entrypoint(*args, **kwargs)
+            else:
+                return run_as_single_step_pipeline(self, *args, **kwargs)
 
         (
             input_artifacts,
@@ -1098,6 +1111,11 @@ To avoid this consider setting step parameters only in one place (config or code
                 this step.
             client_lazy_loaders: The client lazy loaders of this step.
 
+        Raises:
+            StepInterfaceError: If explicit materializers were specified for an
+                output but they do not work for the data type(s) defined by
+                the type annotation.
+
         Returns:
             The finalized step configuration.
         """
@@ -1119,9 +1137,45 @@ To avoid this consider setting step parameters only in one place (config or code
                 output_name, PartialArtifactConfiguration()
             )
 
-            from zenml.steps.utils import get_args
+            if output.materializer_source:
+                # The materializer source was configured by the user. We
+                # validate that their configured materializer supports the
+                # output type. If the output annotation is a Union, we check
+                # that at least one of the specified materializers works with at
+                # least one of the types in the Union. If that's not the case,
+                # it would be a guaranteed failure at runtime and we fail early
+                # here.
+                if output_annotation.resolved_annotation is Any:
+                    continue
 
-            if not output.materializer_source:
+                materializer_classes: List[Type["BaseMaterializer"]] = [
+                    source_utils.load(materializer_source)
+                    for materializer_source in output.materializer_source
+                ]
+
+                for data_type in output_annotation.get_output_types():
+                    try:
+                        materializer_utils.select_materializer(
+                            data_type=data_type,
+                            materializer_classes=materializer_classes,
+                        )
+                        break
+                    except RuntimeError:
+                        pass
+                else:
+                    materializer_strings = [
+                        materializer_source.import_path
+                        for materializer_source in output.materializer_source
+                    ]
+                    raise StepInterfaceError(
+                        "Invalid materializers specified for output "
+                        f"{output_name} of step {self.name}. None of the "
+                        f"materializers ({materializer_strings}) are "
+                        "able to save or load data of the type that is defined "
+                        "for the output "
+                        f"({output_annotation.resolved_annotation})."
+                    )
+            else:
                 if output_annotation.resolved_annotation is Any:
                     outputs[output_name]["materializer_source"] = ()
                     outputs[output_name]["default_materializer_source"] = (
@@ -1131,26 +1185,9 @@ To avoid this consider setting step parameters only in one place (config or code
                     )
                     continue
 
-                if typing_utils.is_union(
-                    typing_utils.get_origin(
-                        output_annotation.resolved_annotation
-                    )
-                    or output_annotation.resolved_annotation
-                ):
-                    output_types = tuple(
-                        type(None)
-                        if typing_utils.is_none_type(output_type)
-                        else output_type
-                        for output_type in get_args(
-                            output_annotation.resolved_annotation
-                        )
-                    )
-                else:
-                    output_types = (output_annotation.resolved_annotation,)
-
                 materializer_sources = []
 
-                for output_type in output_types:
+                for output_type in output_annotation.get_output_types():
                     materializer_class = materializer_registry[output_type]
                     materializer_sources.append(
                         source_utils.resolve(materializer_class)

zenml/steps/utils.py

@@ -18,7 +18,7 @@ import ast
 import contextlib
 import inspect
 import textwrap
-from typing import Any, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
 from uuid import UUID
 
 from pydantic import BaseModel
@@ -26,12 +26,17 @@ from typing_extensions import Annotated
 
 from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.client import Client
-from zenml.enums import MetadataResourceTypes
+from zenml.enums import ExecutionStatus, MetadataResourceTypes
+from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.new.steps.step_context import get_step_context
 from zenml.steps.step_output import Output
-from zenml.utils import source_code_utils, typing_utils
+from zenml.utils import settings_utils, source_code_utils, typing_utils
+
+if TYPE_CHECKING:
+    from zenml.steps import BaseStep
+
 
 logger = get_logger(__name__)
 
@@ -45,6 +50,28 @@ class OutputSignature(BaseModel):
     artifact_config: Optional[ArtifactConfig] = None
     has_custom_name: bool = False
 
+    def get_output_types(self) -> Tuple[Any, ...]:
+        """Get all output types that match the type annotation.
+
+        Returns:
+            All output types that match the type annotation.
+        """
+        if self.resolved_annotation is Any:
+            return ()
+
+        if typing_utils.is_union(
+            typing_utils.get_origin(self.resolved_annotation)
+            or self.resolved_annotation
+        ):
+            return tuple(
+                type(None)
+                if typing_utils.is_none_type(output_type)
+                else output_type
+                for output_type in get_args(self.resolved_annotation)
+            )
+        else:
+            return (self.resolved_annotation,)
+
 
 def get_args(obj: Any) -> Tuple[Any, ...]:
     """Get arguments of a type annotation.
@@ -464,3 +491,88 @@ def log_step_metadata(
         resource_id=step_run_id,
         resource_type=MetadataResourceTypes.STEP_RUN,
     )
+
+
+def run_as_single_step_pipeline(
+    __step: "BaseStep", *args: Any, **kwargs: Any
+) -> Any:
+    """Runs the step as a single step pipeline.
+
+    - All inputs that are not JSON serializable will be uploaded to the
+    artifact store before the pipeline is being executed.
+    - All output artifacts of the step will be loaded using the materializer
+    that was used to store them.
+
+    Args:
+        *args: Entrypoint function arguments.
+        **kwargs: Entrypoint function keyword arguments.
+
+    Raises:
+        RuntimeError: If the step execution failed.
+        StepInterfaceError: If the arguments to the entrypoint function are
+            invalid.
+
+    Returns:
+        The output of the step entrypoint function.
+    """
+    from zenml import ExternalArtifact, pipeline
+    from zenml.config.base_settings import BaseSettings
+    from zenml.new.pipelines.run_utils import (
+        wait_for_pipeline_run_to_finish,
+    )
+
+    logger.info(
+        "Running single step pipeline to execute step `%s`", __step.name
+    )
+
+    try:
+        validated_arguments = (
+            inspect.signature(__step.entrypoint)
+            .bind(*args, **kwargs)
+            .arguments
+        )
+    except TypeError as e:
+        raise StepInterfaceError(
+            "Invalid step function entrypoint arguments. Check out the "
+            "error above for more details."
+        ) from e
+
+    inputs: Dict[str, Any] = {}
+    for key, value in validated_arguments.items():
+        try:
+            __step.entrypoint_definition.validate_input(key=key, value=value)
+            inputs[key] = value
+        except Exception:
+            inputs[key] = ExternalArtifact(value=value)
+
+    orchestrator = Client().active_stack.orchestrator
+
+    pipeline_settings: Any = {}
+    if "synchronous" in orchestrator.config.model_fields:
+        # Make sure the orchestrator runs sync so we stream the logs
+        key = settings_utils.get_stack_component_setting_key(orchestrator)
+        pipeline_settings[key] = BaseSettings(synchronous=True)
+
+    @pipeline(name=__step.name, enable_cache=False, settings=pipeline_settings)
+    def single_step_pipeline() -> None:
+        __step(**inputs)
+
+    run = single_step_pipeline.with_options(unlisted=True)()
+    run = wait_for_pipeline_run_to_finish(run.id)
+
+    if run.status != ExecutionStatus.COMPLETED:
+        raise RuntimeError("Failed to execute step %s.", __step.name)
+
+    # 4. Load output artifacts
+    step_run = next(iter(run.steps.values()))
+    outputs = [
+        step_run.outputs[output_name].load()
+        for output_name in step_run.config.outputs.keys()
+    ]
+
+    if len(outputs) == 0:
+        return None
+    elif len(outputs) == 1:
+        return outputs[0]
+    else:
+        return tuple(outputs)

zenml/utils/cloud_utils.py

@@ -20,21 +20,21 @@ from zenml.utils.dashboard_utils import get_model_version_url
 logger = get_logger(__name__)
 
 
-def is_cloud_model_version(model_version: ModelVersionResponse) -> bool:
-    """Check if a model version is from a ZenML Pro server.
+def try_get_model_version_url(model_version: ModelVersionResponse) -> str:
+    """Check if a model version is from a ZenML Pro server and return its' URL.
 
     Args:
         model_version: The model version to check.
 
     Returns:
-        True if the model version is from a ZenML Pro server, else False.
+        URL if the model version is from a ZenML Pro server, else empty string.
     """
     model_version_url = get_model_version_url(model_version.id)
     if model_version_url:
-        logger.info(
-            f"Dashboard URL for Model Version with name {model_version.name} "
-            f": {model_version_url}"
+        return (
+            "Dashboard URL for created Model Version "
+            f"`{model_version.model.name}::{model_version.name}`:\n"
+            + model_version_url
         )
-        return True
     else:
-        return False
+        return ""

zenml/utils/code_utils.py

@@ -16,6 +16,7 @@
 import hashlib
 import os
 import shutil
+import sys
 import tempfile
 from pathlib import Path
 from typing import IO, TYPE_CHECKING, Dict, Optional
@@ -23,12 +24,14 @@ from typing import IO, TYPE_CHECKING, Dict, Optional
 from zenml.client import Client
 from zenml.io import fileio
 from zenml.logger import get_logger
-from zenml.utils import string_utils
+from zenml.utils import source_utils, string_utils
 from zenml.utils.archivable import Archivable
 
 if TYPE_CHECKING:
     from git.repo.base import Repo
 
+    from zenml.artifact_stores import BaseArtifactStore
+
 
 logger = get_logger(__name__)
 
@@ -168,6 +171,30 @@ class CodeArchive(Archivable):
             )
 
 
+def compute_file_hash(file: IO[bytes]) -> str:
+    """Compute a hash of the content of a file.
+
+    This function will not seek the file before or after the hash computation.
+    This means that the content will be computed based on the current cursor
+    until the end of the file.
+
+    Args:
+        file: The file for which to compute the hash.
+
+    Returns:
+        A hash of the file content.
+    """
+    hash_ = hashlib.sha1()  # nosec
+
+    while True:
+        data = file.read(64 * 1024)
+        if not data:
+            break
+        hash_.update(data)
+
+    return hash_.hexdigest()
+
+
 def upload_code_if_necessary(code_archive: CodeArchive) -> str:
     """Upload code to the artifact store if necessary.
 
@@ -179,7 +206,7 @@ def upload_code_if_necessary(code_archive: CodeArchive) -> str:
         code_archive: The code archive to upload.
 
     Returns:
-        The path where to archived code is uploaded.
+        The path where the archived code is uploaded.
     """
     artifact_store = Client().active_stack.artifact_store
 
@@ -187,36 +214,27 @@ def upload_code_if_necessary(code_archive: CodeArchive) -> str:
         mode="w+b", delete=False, suffix=".tar.gz"
     ) as f:
         code_archive.write_archive(f)
-
-        hash_ = hashlib.sha1()  # nosec
-
-        while True:
-            data = f.read(64 * 1024)
-            if not data:
-                break
-            hash_.update(data)
-
-        filename = f"{hash_.hexdigest()}.tar.gz"
-        upload_dir = os.path.join(artifact_store.path, "code_uploads")
-        fileio.makedirs(upload_dir)
-        upload_path = os.path.join(upload_dir, filename)
-
-        if not fileio.exists(upload_path):
-            archive_size = string_utils.get_human_readable_filesize(
-                os.path.getsize(f.name)
-            )
-            logger.info(
-                "Uploading code to `%s` (Size: %s).", upload_path, archive_size
-            )
-            fileio.copy(f.name, upload_path)
-            logger.info("Code upload finished.")
-        else:
-            logger.info(
-                "Code already exists in artifact store, skipping upload."
-            )
-
-    if os.path.exists(f.name):
-        os.remove(f.name)
+        archive_path = f.name
+        archive_hash = compute_file_hash(f)
+
+    upload_dir = os.path.join(artifact_store.path, "code_uploads")
+    fileio.makedirs(upload_dir)
+    upload_path = os.path.join(upload_dir, f"{archive_hash}.tar.gz")
+
+    if not fileio.exists(upload_path):
+        archive_size = string_utils.get_human_readable_filesize(
+            os.path.getsize(archive_path)
+        )
+        logger.info(
+            "Uploading code to `%s` (Size: %s).", upload_path, archive_size
+        )
+        fileio.copy(archive_path, upload_path)
+        logger.info("Code upload finished.")
+    else:
+        logger.info("Code already exists in artifact store, skipping upload.")
+
+    if os.path.exists(archive_path):
+        os.remove(archive_path)
 
     return upload_path
 
@@ -242,3 +260,83 @@ def download_and_extract_code(code_path: str, extract_dir: str) -> None:
 
     shutil.unpack_archive(filename=download_path, extract_dir=extract_dir)
     os.remove(download_path)
+
+
+def download_code_from_artifact_store(code_path: str) -> None:
+    """Download code from the artifact store.
+
+    Args:
+        code_path: Path where the code is stored.
+    """
+    logger.info("Downloading code from artifact store path `%s`.", code_path)
+
+    # Do not remove this line, we need to instantiate the artifact store to
+    # register the filesystem needed for the file download
+    _ = Client().active_stack.artifact_store
+
+    extract_dir = os.path.abspath("code")
+    os.makedirs(extract_dir)
+
+    download_and_extract_code(code_path=code_path, extract_dir=extract_dir)
+
+    source_utils.set_custom_source_root(extract_dir)
+    sys.path.insert(0, extract_dir)
+    os.chdir(extract_dir)
+
+
+def _get_notebook_upload_dir(artifact_store: "BaseArtifactStore") -> str:
+    """Get the upload directory for code extracted from notebook cells.
+
+    Args:
+        artifact_store: The artifact store in which the directory should be.
+
+    Returns:
+        The upload directory for code extracted from notebook cells.
+    """
+    return os.path.join(artifact_store.path, "notebook_code")
+
+
+def upload_notebook_code(
+    artifact_store: "BaseArtifactStore", cell_code: str, file_name: str
+) -> None:
+    """Upload code extracted from a notebook cell.
+
+    Args:
+        artifact_store: The artifact store in which to upload the code.
+        cell_code: The notebook cell code.
+        file_name: The filename to use for storing the cell code.
+    """
+    upload_dir = _get_notebook_upload_dir(artifact_store=artifact_store)
+    fileio.makedirs(upload_dir)
+    upload_path = os.path.join(upload_dir, file_name)
+
+    if not fileio.exists(upload_path):
+        with fileio.open(upload_path, "wb") as f:
+            f.write(cell_code.encode())
+
+        logger.info("Uploaded notebook cell code to %s.", upload_path)
+
+
+def download_notebook_code(
+    artifact_store: "BaseArtifactStore", file_name: str, download_path: str
+) -> None:
+    """Download code extracted from a notebook cell.
+
+    Args:
+        artifact_store: The artifact store from which to download the code.
+        file_name: The name of the code file.
+        download_path: The local path where the file should be downloaded to.
+
+    Raises:
+        FileNotFoundError: If no file with the given filename exists in this
+            artifact store.
+    """
+    code_dir = _get_notebook_upload_dir(artifact_store=artifact_store)
+    code_path = os.path.join(code_dir, file_name)
+
+    if not fileio.exists(code_path):
+        raise FileNotFoundError(
+            f"Notebook code at path {code_path} not found."
+        )
+
+    fileio.copy(code_path, download_path)

zenml/utils/function_utils.py

@@ -34,12 +34,12 @@ from zenml.utils.function_utils import _cli_wrapped_function
 import sys
 sys.path.append(r"{func_path}")
 
-from {func_module} import {func_name} as func_to_wrap
+from {func_module} import {func_name} as step_function
 
-if entrypoint:=getattr(func_to_wrap, "entrypoint", None):
-    func = _cli_wrapped_function(entrypoint)
+if unwrapped_entrypoint:=getattr(step_function, "unwrapped_entrypoint", None):
+    func = _cli_wrapped_function(unwrapped_entrypoint)
 else:
-    func = _cli_wrapped_function(func_to_wrap)
+    func = _cli_wrapped_function(step_function.entrypoint)
 """
 _CLI_WRAPPED_MAINS = {
     "accelerate": """
@@ -204,13 +204,13 @@ def _cli_wrapped_function(func: F) -> F:
 
 @contextmanager
 def create_cli_wrapped_script(
-    func: F, flavour: str = "accelerate"
+    func: F, flavor: str = "accelerate"
 ) -> Iterator[Tuple[Path, Path]]:
     """Create a script with the CLI-wrapped function.
 
     Args:
         func: The function to use.
-        flavour: The flavour to use.
+        flavor: The flavor to use.
 
     Yields:
         The paths of the script and the output.
@@ -241,7 +241,7 @@ def create_cli_wrapped_script(
                     func_module=clean_module_name,
                     func_name=func.__name__,
                 )
-                script += _CLI_WRAPPED_MAINS[flavour].format(
+                script += _CLI_WRAPPED_MAINS[flavor].format(
                     output_file=str(output_path.absolute())
                 )
                 f.write(script)

zenml/utils/notebook_utils.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Notebook utilities."""
 
+import hashlib
 from typing import Any, Callable, Optional, TypeVar, Union
 
 from zenml.environment import Environment
@@ -120,3 +121,16 @@ def warn_about_notebook_cell_magic_commands(cell_code: str) -> None:
             "of these lines contain Jupyter notebook magic commands, "
             "remove them and try again."
         )
+
+
+def compute_cell_replacement_module_name(cell_code: str) -> str:
+    """Compute the replacement module name for a given cell code.
+
+    Args:
+        cell_code: The code of the notebook cell.
+
+    Returns:
+        The replacement module name.
+    """
+    code_hash = hashlib.sha1(cell_code.encode()).hexdigest()  # nosec
+    return f"extracted_notebook_code_{code_hash}"

zenml/utils/pipeline_docker_image_builder.py

@@ -38,7 +38,6 @@ from zenml.constants import (
     ENV_ZENML_CONFIG_PATH,
     ENV_ZENML_ENABLE_REPO_INIT_WARNINGS,
     ENV_ZENML_LOGGING_COLORS_DISABLED,
-    ENV_ZENML_REQUIRES_CODE_DOWNLOAD,
     handle_bool_env_var,
 )
 from zenml.enums import OperatingSystemType
@@ -55,10 +54,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 DOCKER_IMAGE_WORKDIR = "/app"
-DOCKER_IMAGE_ZENML_CONFIG_DIR = ".zenconfig"
-DOCKER_IMAGE_ZENML_CONFIG_PATH = (
-    f"{DOCKER_IMAGE_WORKDIR}/{DOCKER_IMAGE_ZENML_CONFIG_DIR}"
-)
+DOCKER_IMAGE_ZENML_CONFIG_PATH = f"{DOCKER_IMAGE_WORKDIR}/.zenconfig"
 
 DEFAULT_DOCKER_PARENT_IMAGE = (
     f"zenmldocker/zenml:{zenml.__version__}-"
@@ -334,7 +330,6 @@ class PipelineDockerImageBuilder:
             dockerfile = self._generate_zenml_pipeline_dockerfile(
                 parent_image=parent_image,
                 docker_settings=docker_settings,
-                download_files=download_files,
                 requirements_files=requirements_files,
                 apt_packages=apt_packages,
                 entrypoint=entrypoint,
@@ -563,7 +558,6 @@ class PipelineDockerImageBuilder:
     def _generate_zenml_pipeline_dockerfile(
         parent_image: str,
         docker_settings: DockerSettings,
-        download_files: bool,
         requirements_files: Sequence[Tuple[str, str, List[str]]] = (),
         apt_packages: Sequence[str] = (),
         entrypoint: Optional[str] = None,
@@ -573,7 +567,6 @@ class PipelineDockerImageBuilder:
         Args:
             parent_image: The image to use as parent for the Dockerfile.
             docker_settings: Docker settings for this image build.
-            download_files: Whether to download files in the build context.
             requirements_files: List of tuples that contain three items:
                 - the name of a requirements file,
                 - the content of that file,
@@ -641,9 +634,6 @@ class PipelineDockerImageBuilder:
             )
 
         lines.append(f"ENV {ENV_ZENML_ENABLE_REPO_INIT_WARNINGS}=False")
-        if download_files:
-            lines.append(f"ENV {ENV_ZENML_REQUIRES_CODE_DOWNLOAD}=True")
-
         lines.append(
             f"ENV {ENV_ZENML_CONFIG_PATH}={DOCKER_IMAGE_ZENML_CONFIG_PATH}"
         )

zenml/utils/pydantic_utils.py

@@ -297,7 +297,7 @@ def model_validator_data_handler(
     this can create conflicts after the migration and this function will be
     used as a helper function to handle different types of raw input data.
 
-    A code snippet to showcase how the behaviour changes. The "before" validator
+    A code snippet to showcase how the behavior changes. The "before" validator
     prints the type of the input:
 
         class Base(BaseModel):
@@ -336,7 +336,7 @@ def model_validator_data_handler(
         # raw data such as MyClass(...) or MyClass.model_validate()
 
         if isinstance(raw_data, dict):
-            # In most cases, this is the behaviour as the raw input is a dict
+            # In most cases, this is the behavior as the raw input is a dict
             return raw_data
 
         elif isinstance(raw_data, base_class):
@@ -354,7 +354,7 @@ def model_validator_data_handler(
                 f"During the validation of a `{base_class}` object, an instance"
                 f"of `{raw_data.__class__}` (super class of `{base_class}`) "
                 f"has been passed as raw input. This might lead to unexpected "
-                f"behaviour in case `{base_class}` have features which can not"
+                f"behavior in case `{base_class}` have features which can not"
                 f"be extracted from an instance of a `{raw_data.__class__}`."
             )
             return dict(raw_data)