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

zenml/{new/pipelines → pipelines}/build_utils.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Pipeline build utilities."""
+
 import hashlib
 import platform
 from typing import (
@@ -29,15 +30,15 @@ from zenml.code_repositories import BaseCodeRepository
 from zenml.logger import get_logger
 from zenml.models import (
     BuildItem,
+    CodeReferenceRequest,
     PipelineBuildBase,
     PipelineBuildRequest,
     PipelineBuildResponse,
     PipelineDeploymentBase,
+    StackResponse,
 )
 from zenml.stack import Stack
-from zenml.utils import (
-    source_utils,
-)
+from zenml.utils import source_utils
 from zenml.utils.pipeline_docker_image_builder import (
     PipelineDockerImageBuilder,
 )
@@ -49,6 +50,106 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def build_required(deployment: "PipelineDeploymentBase") -> bool:
+    """Checks whether a build is required for the deployment and active stack.
+
+    Args:
+        deployment: The deployment for which to check.
+
+    Returns:
+        If a build is required.
+    """
+    stack = Client().active_stack
+    return bool(stack.get_docker_builds(deployment=deployment))
+
+
+def requires_included_code(
+    deployment: "PipelineDeploymentBase",
+    code_repository: Optional["BaseCodeRepository"] = None,
+) -> bool:
+    """Checks whether the deployment requires included code.
+
+    Args:
+        deployment: The deployment.
+        code_repository: If provided, this code repository can be used to
+            download the code inside the container images.
+
+    Returns:
+        If the deployment requires code included in the container images.
+    """
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if docker_settings.allow_download_from_artifact_store:
+            return False
+
+        if docker_settings.allow_download_from_code_repository:
+            if code_repository:
+                continue
+
+        if docker_settings.allow_including_files_in_images:
+            return True
+
+    return False
+
+
+def requires_download_from_code_repository(
+    deployment: "PipelineDeploymentBase",
+) -> bool:
+    """Checks whether the deployment needs to download code from a repository.
+
+    Args:
+        deployment: The deployment.
+
+    Returns:
+        If the deployment needs to download code from a code repository.
+    """
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if docker_settings.allow_download_from_artifact_store:
+            return False
+
+        if docker_settings.allow_including_files_in_images:
+            return False
+
+        if docker_settings.allow_download_from_code_repository:
+            # The other two options are false, which means download from a
+            # code repo is required.
+            return True
+
+    return False
+
+
+def code_download_possible(
+    deployment: "PipelineDeploymentBase",
+    code_repository: Optional["BaseCodeRepository"] = None,
+) -> bool:
+    """Checks whether code download is possible for the deployment.
+
+    Args:
+        deployment: The deployment.
+        code_repository: If provided, this code repository can be used to
+            download the code inside the container images.
+
+    Returns:
+        Whether code download is possible for the deployment.
+    """
+    for step in deployment.step_configurations.values():
+        if step.config.docker_settings.allow_download_from_artifact_store:
+            continue
+
+        if (
+            step.config.docker_settings.allow_download_from_code_repository
+            and code_repository
+        ):
+            continue
+
+        return False
+
+    return True
+
+
 def reuse_or_create_pipeline_build(
     deployment: "PipelineDeploymentBase",
     allow_build_reuse: bool,
@@ -67,8 +168,8 @@ def reuse_or_create_pipeline_build(
         build: Optional existing build. If given, the build will be fetched
             (or registered) in the database. If not given, a new build will
             be created.
-        code_repository: If provided, this code repository will be used to
-            download inside the build images.
+        code_repository: If provided, this code repository can be used to
+            download code inside the container images.
 
     Returns:
         The build response.
@@ -76,8 +177,11 @@ def reuse_or_create_pipeline_build(
     if not build:
         if (
             allow_build_reuse
-            and code_repository
-            and not deployment.requires_included_files
+            and not deployment.should_prevent_build_reuse
+            and not requires_included_code(
+                deployment=deployment, code_repository=code_repository
+            )
+            and build_required(deployment=deployment)
         ):
             existing_build = find_existing_build(
                 deployment=deployment, code_repository=code_repository
@@ -90,6 +194,16 @@ def reuse_or_create_pipeline_build(
                     Client().active_stack.name,
                 )
                 return existing_build
+            else:
+                logger.info(
+                    "Unable to find a build to reuse. A previous build can be "
+                    "reused when the following conditions are met:\n"
+                    "  * The existing build was created for the same stack, "
+                    "ZenML version and Python version\n"
+                    "  * The stack contains a container registry\n"
+                    "  * The Docker settings of the pipeline and all its steps "
+                    "are the same as for the existing build."
+                )
 
         return create_pipeline_build(
             deployment=deployment,
@@ -105,7 +219,7 @@ def reuse_or_create_pipeline_build(
             workspace=Client().active_workspace.id,
             stack=Client().active_stack_model.id,
             pipeline=pipeline_id,
-            **build.dict(),
+            **build.model_dump(),
         )
         build_model = Client().zen_store.create_build(build=build_request)
 
@@ -120,7 +234,7 @@ def reuse_or_create_pipeline_build(
 
 def find_existing_build(
     deployment: "PipelineDeploymentBase",
-    code_repository: "BaseCodeRepository",
+    code_repository: Optional["BaseCodeRepository"] = None,
 ) -> Optional["PipelineBuildResponse"]:
     """Find an existing build for a deployment.
 
@@ -135,6 +249,11 @@ def find_existing_build(
     client = Client()
     stack = client.active_stack
 
+    if not stack.container_registry:
+        # There can be no non-local builds that we can reuse if there is no
+        # container registry in the stack.
+        return None
+
     python_version_prefix = ".".join(platform.python_version_tuple()[:2])
     required_builds = stack.get_docker_builds(deployment=deployment)
 
@@ -149,6 +268,13 @@ def find_existing_build(
         sort_by="desc:created",
         size=1,
         stack_id=stack.id,
+        # Until we implement stack versioning, users can still update their
+        # stack to update/remove the container registry. In that case, we might
+        # try to pull an image from a container registry that we don't have
+        # access to. This is why we add an additional check for the container
+        # registry ID here. (This is still not perfect as users can update the
+        # container registry URI or config, but the best we can do)
+        container_registry_id=stack.container_registry.id,
         # The build is local and it's not clear whether the images
         # exist on the current machine or if they've been overwritten.
         # TODO: Should we support this by storing the unique Docker ID for
@@ -190,6 +316,7 @@ def create_pipeline_build(
             settings were specified.
     """
     client = Client()
+    stack_model = Client().active_stack_model
     stack = client.active_stack
     required_builds = stack.get_docker_builds(deployment=deployment)
 
@@ -249,6 +376,11 @@ def create_pipeline_build(
             download_files = build_config.should_download_files(
                 code_repository=code_repository,
             )
+            pass_code_repo = (
+                build_config.should_download_files_from_code_repository(
+                    code_repository=code_repository
+                )
+            )
 
             (
                 image_name_or_digest,
@@ -262,7 +394,7 @@ def create_pipeline_build(
                 download_files=download_files,
                 entrypoint=build_config.entrypoint,
                 extra_files=build_config.extra_files,
-                code_repository=code_repository,
+                code_repository=code_repository if pass_code_repo else None,
             )
             contains_code = include_files
 
@@ -283,11 +415,11 @@ def create_pipeline_build(
     build_checksum = compute_build_checksum(
         required_builds, stack=stack, code_repository=code_repository
     )
-
+    stack_checksum = compute_stack_checksum(stack=stack_model)
     build_request = PipelineBuildRequest(
         user=client.active_user.id,
         workspace=client.active_workspace.id,
-        stack=client.active_stack_model.id,
+        stack=stack_model.id,
         pipeline=pipeline_id,
         is_local=is_local,
         contains_code=contains_code,
@@ -295,6 +427,7 @@ def create_pipeline_build(
         zenml_version=zenml.__version__,
         python_version=platform.python_version(),
         checksum=build_checksum,
+        stack_checksum=stack_checksum,
     )
     return client.zen_store.create_build(build_request)
 
@@ -356,36 +489,48 @@ def verify_local_repository_context(
         The code repository from which to download files for the runs of the
         deployment, or None if code download is not possible.
     """
-    if deployment.requires_code_download:
-        if not local_repo_context:
-            raise RuntimeError(
-                "The `DockerSettings` of the pipeline or one of its "
-                "steps specify that code should be included in the "
-                "Docker image (`source_files='download'`), but there is no "
-                "code repository active at your current source root "
-                f"`{source_utils.get_source_root()}`."
-            )
-        elif local_repo_context.is_dirty:
-            raise RuntimeError(
-                "The `DockerSettings` of the pipeline or one of its "
-                "steps specify that code should be included in the "
-                "Docker image (`source_files='download'`), but the code "
-                "repository active at your current source root "
-                f"`{source_utils.get_source_root()}` has uncommitted "
-                "changes."
-            )
-        elif local_repo_context.has_local_changes:
-            raise RuntimeError(
-                "The `DockerSettings` of the pipeline or one of its "
-                "steps specify that code should be included in the "
-                "Docker image (`source_files='download'`), but the code "
-                "repository active at your current source root "
-                f"`{source_utils.get_source_root()}` has unpushed "
-                "changes."
-            )
+    if build_required(deployment=deployment):
+        if requires_download_from_code_repository(deployment=deployment):
+            if not local_repo_context:
+                raise RuntimeError(
+                    "The `DockerSettings` of the pipeline or one of its "
+                    "steps specify that code should be downloaded from a "
+                    "code repository, but "
+                    "there is no code repository active at your current source "
+                    f"root `{source_utils.get_source_root()}`."
+                )
+            elif local_repo_context.is_dirty:
+                raise RuntimeError(
+                    "The `DockerSettings` of the pipeline or one of its "
+                    "steps specify that code should be downloaded from a "
+                    "code repository, but "
+                    "the code repository active at your current source root "
+                    f"`{source_utils.get_source_root()}` has uncommitted "
+                    "changes."
+                )
+            elif local_repo_context.has_local_changes:
+                raise RuntimeError(
+                    "The `DockerSettings` of the pipeline or one of its "
+                    "steps specify that code should be downloaded from a "
+                    "code repository, but "
+                    "the code repository active at your current source root "
+                    f"`{source_utils.get_source_root()}` has unpushed "
+                    "changes."
+                )
 
-    code_repository = None
+        if local_repo_context:
+            if local_repo_context.is_dirty:
+                logger.warning(
+                    "Unable to use code repository to download code for this "
+                    "run as there are uncommitted changes."
+                )
+            elif local_repo_context.has_local_changes:
+                logger.warning(
+                    "Unable to use code repository to download code for this "
+                    "run as there are unpushed changes."
+                )
 
+    code_repository = None
     if local_repo_context and not local_repo_context.has_local_changes:
         model = Client().get_code_repository(
             local_repo_context.code_repository_id
@@ -431,13 +576,41 @@ def verify_custom_build(
             "might differ from the local code in your client environment."
         )
 
-    if build.requires_code_download and not code_repository:
-        raise RuntimeError(
-            "The build you specified does not include code but code download "
-            "not possible. This might be because you don't have a code "
-            "repository registered or the code repository contains local "
-            "changes."
-        )
+    if build.requires_code_download:
+        if requires_included_code(
+            deployment=deployment, code_repository=code_repository
+        ):
+            raise RuntimeError(
+                "The `DockerSettings` of the pipeline or one of its "
+                "steps specify that code should be included in the Docker "
+                "image, but the build you "
+                "specified requires code download. Either update your "
+                "`DockerSettings` or specify a different build and try "
+                "again."
+            )
+
+        if (
+            requires_download_from_code_repository(deployment=deployment)
+            and not code_repository
+        ):
+            raise RuntimeError(
+                "The `DockerSettings` of the pipeline or one of its "
+                "steps specify that code should be downloaded from a "
+                "code repository but "
+                "there is no code repository active at your current source "
+                f"root `{source_utils.get_source_root()}`."
+            )
+
+        if not code_download_possible(
+            deployment=deployment, code_repository=code_repository
+        ):
+            raise RuntimeError(
+                "The `DockerSettings` of the pipeline or one of its "
+                "steps specify that code can not be downloaded from the "
+                "artifact store, but the build you specified requires code "
+                "download. Either update your `DockerSettings` or specify a "
+                "different build and try again."
+            )
 
     if build.checksum:
         build_checksum = compute_build_checksum(
@@ -488,3 +661,70 @@ def verify_custom_build(
             "your local machine or the image tags have been "
             "overwritten since the original build happened."
         )
+
+
+def compute_stack_checksum(stack: StackResponse) -> str:
+    """Compute a stack checksum.
+
+    Args:
+        stack: The stack for which to compute the checksum.
+
+    Returns:
+        The checksum.
+    """
+    hash_ = hashlib.md5()  # nosec
+
+    # This checksum is used to see if the stack has been updated since a build
+    # was created for it. We create this checksum not with specific requirements
+    # as these might change with new ZenML releases, but they don't actually
+    # invalidate those Docker images.
+    required_integrations = sorted(
+        {
+            component.integration
+            for components in stack.components.values()
+            for component in components
+            if component.integration and component.integration != "built-in"
+        }
+    )
+    for integration in required_integrations:
+        hash_.update(integration.encode())
+
+    return hash_.hexdigest()
+
+
+def should_upload_code(
+    deployment: PipelineDeploymentBase,
+    build: Optional[PipelineBuildResponse],
+    code_reference: Optional[CodeReferenceRequest],
+) -> bool:
+    """Checks whether the current code should be uploaded for the deployment.
+
+    Args:
+        deployment: The deployment.
+        build: The build for the deployment.
+        code_reference: The code reference for the deployment.
+
+    Returns:
+        Whether the current code should be uploaded for the deployment.
+    """
+    if not build:
+        # No build means we don't need to download code into a Docker container
+        # for step execution. In other remote orchestrators that don't use
+        # Docker containers but instead use e.g. Wheels to run, the code should
+        # already be included.
+        return False
+
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if (
+            code_reference
+            and docker_settings.allow_download_from_code_repository
+        ):
+            # No upload needed for this step
+            continue
+
+        if docker_settings.allow_download_from_artifact_store:
+            return True
+
+    return False

zenml/{new/pipelines → pipelines}/pipeline_context.py

@@ -15,9 +15,13 @@
 
 from typing import TYPE_CHECKING
 
+from zenml.logger import get_logger
+
 if TYPE_CHECKING:
     from zenml.config.pipeline_configurations import PipelineConfiguration
 
+logger = get_logger(__name__)
+
 
 def get_pipeline_context() -> "PipelineContext":
     """Get the context of the current pipeline.
@@ -29,11 +33,11 @@ def get_pipeline_context() -> "PipelineContext":
         RuntimeError: If no active pipeline is found.
         RuntimeError: If inside a running step.
     """
-    from zenml.new.pipelines.pipeline import Pipeline
+    from zenml.pipelines.pipeline_definition import Pipeline
 
     if Pipeline.ACTIVE_PIPELINE is None:
         try:
-            from zenml.new.steps.step_context import get_step_context
+            from zenml.steps.step_context import get_step_context
 
             get_step_context()
         except RuntimeError:

zenml/pipelines/pipeline_decorator.py

@@ -11,51 +11,35 @@
 #  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.
-"""Legacy ZenML pipeline decorator definition."""
-from types import FunctionType
+"""ZenML pipeline decorator definition."""
+
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
     Dict,
+    List,
     Optional,
-    Type,
     TypeVar,
     Union,
     overload,
 )
 
 from zenml.logger import get_logger
-from zenml.model.model import Model
-from zenml.pipelines.base_pipeline import (
-    CLASS_CONFIGURATION,
-    PARAM_ENABLE_ARTIFACT_METADATA,
-    PARAM_ENABLE_ARTIFACT_VISUALIZATION,
-    PARAM_ENABLE_CACHE,
-    PARAM_ENABLE_STEP_LOGS,
-    PARAM_EXTRA_OPTIONS,
-    PARAM_MODEL,
-    PARAM_ON_FAILURE,
-    PARAM_ON_SUCCESS,
-    PARAM_PIPELINE_NAME,
-    PARAM_SETTINGS,
-    PIPELINE_INNER_FUNC_NAME,
-    BasePipeline,
-)
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
+    from zenml.model.model import Model
+    from zenml.pipelines.pipeline_definition import Pipeline
+    from zenml.types import HookSpecification
 
-    HookSpecification = Union[str, FunctionType]
+    F = TypeVar("F", bound=Callable[..., None])
 
 logger = get_logger(__name__)
 
-F = TypeVar("F", bound=Callable[..., None])
-
 
 @overload
-def pipeline(_func: F) -> Type[BasePipeline]:
-    ...
+def pipeline(_func: "F") -> "Pipeline": ...
 
 
 @overload
@@ -64,30 +48,33 @@ def pipeline(
     name: Optional[str] = None,
     enable_cache: Optional[bool] = None,
     enable_artifact_metadata: Optional[bool] = None,
-    enable_artifact_visualization: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
     settings: Optional[Dict[str, "SettingsOrDict"]] = None,
+    tags: Optional[List[str]] = None,
     extra: Optional[Dict[str, Any]] = None,
+    on_failure: Optional["HookSpecification"] = None,
+    on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
-) -> Callable[[F], Type[BasePipeline]]:
-    ...
+    substitutions: Optional[Dict[str, str]] = None,
+) -> Callable[["F"], "Pipeline"]: ...
 
 
 def pipeline(
-    _func: Optional[F] = None,
+    _func: Optional["F"] = None,
     *,
     name: Optional[str] = None,
     enable_cache: Optional[bool] = None,
     enable_artifact_metadata: Optional[bool] = None,
-    enable_artifact_visualization: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
     settings: Optional[Dict[str, "SettingsOrDict"]] = None,
+    tags: Optional[List[str]] = None,
     extra: Optional[Dict[str, Any]] = None,
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
-) -> Union[Type[BasePipeline], Callable[[F], Type[BasePipeline]]]:
-    """Outer decorator function for the creation of a ZenML pipeline.
+    substitutions: Optional[Dict[str, str]] = None,
+) -> Union["Pipeline", Callable[["F"], "Pipeline"]]:
+    """Decorator to create a pipeline.
 
     Args:
         _func: The decorated function.
@@ -95,9 +82,9 @@ def pipeline(
             decorated function will be used as a fallback.
         enable_cache: Whether to use caching or not.
         enable_artifact_metadata: Whether to enable artifact metadata or not.
-        enable_artifact_visualization: Whether to enable artifact visualization.
-        enable_step_logs: Whether to enable step logs.
+        enable_step_logs: If step logs should be enabled for this pipeline.
         settings: Settings for this pipeline.
+        tags: Tags to apply to runs of the pipeline.
         extra: Extra configurations for this pipeline.
         on_failure: Callback function in event of failure of the step. Can be a
             function with a single argument of type `BaseException`, or a source
@@ -106,42 +93,31 @@ def pipeline(
             function with no arguments, or a source path to such a function
             (e.g. `module.my_function`).
         model: configuration of the model in the Model Control Plane.
+        substitutions: Extra placeholders to use in the name templates.
 
     Returns:
-        the inner decorator which creates the pipeline class based on the
-        ZenML BasePipeline
+        A pipeline instance.
     """
 
-    def inner_decorator(func: F) -> Type[BasePipeline]:
-        pipeline_name = name or func.__name__
-        logger.warning(
-            "The `@pipeline` decorator that you used to define your "
-            f"{pipeline_name} pipeline is deprecated. Check out the 0.40.0 "
-            "migration guide for more information on how to migrate your "
-            "pipelines to the new syntax: "
-            "https://docs.zenml.io/reference/migration-guide/migration-zero-forty.html"
+    def inner_decorator(func: "F") -> "Pipeline":
+        from zenml.pipelines.pipeline_definition import Pipeline
+
+        p = Pipeline(
+            name=name or func.__name__,
+            enable_cache=enable_cache,
+            enable_artifact_metadata=enable_artifact_metadata,
+            enable_step_logs=enable_step_logs,
+            settings=settings,
+            tags=tags,
+            extra=extra,
+            on_failure=on_failure,
+            on_success=on_success,
+            model=model,
+            entrypoint=func,
+            substitutions=substitutions,
         )
 
-        return type(
-            name or func.__name__,
-            (BasePipeline,),
-            {
-                PIPELINE_INNER_FUNC_NAME: staticmethod(func),
-                CLASS_CONFIGURATION: {
-                    PARAM_PIPELINE_NAME: name,
-                    PARAM_ENABLE_CACHE: enable_cache,
-                    PARAM_ENABLE_ARTIFACT_METADATA: enable_artifact_metadata,
-                    PARAM_ENABLE_ARTIFACT_VISUALIZATION: enable_artifact_visualization,
-                    PARAM_ENABLE_STEP_LOGS: enable_step_logs,
-                    PARAM_SETTINGS: settings,
-                    PARAM_EXTRA_OPTIONS: extra,
-                    PARAM_ON_FAILURE: on_failure,
-                    PARAM_ON_SUCCESS: on_success,
-                    PARAM_MODEL: model,
-                },
-                "__module__": func.__module__,
-                "__doc__": func.__doc__,
-            },
-        )
+        p.__doc__ = func.__doc__
+        return p
 
     return inner_decorator if _func is None else inner_decorator(_func)