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

zenml/config/docker_settings.py

@@ -16,11 +16,12 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import Extra, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.logger import get_logger
 from zenml.utils import deprecation_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -47,13 +48,28 @@ class PythonEnvironmentExportMethod(Enum):
         }[self]
 
 
-class SourceFileMode(Enum):
-    """Different methods to handle source files in Docker images."""
+class PythonPackageInstaller(Enum):
+    """Different installers for python packages."""
 
-    INCLUDE = "include"
-    DOWNLOAD_OR_INCLUDE = "download_or_include"
-    DOWNLOAD = "download"
-    IGNORE = "ignore"
+    PIP = "pip"
+    UV = "uv"
+
+
+class DockerBuildConfig(BaseModel):
+    """Configuration for a Docker build.
+
+    Attributes:
+        build_options: Additional options that will be passed unmodified to the
+            Docker build call when building an image. You can use this to for
+            example specify build args or a target stage. See
+            https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
+            for a full list of available options.
+        dockerignore: Path to a dockerignore file to use when building the
+            Docker image.
+    """
+
+    build_options: Dict[str, Any] = {}
+    dockerignore: Optional[str] = None
 
 
 class DockerSettings(BaseSettings):
@@ -63,7 +79,7 @@ class DockerSettings(BaseSettings):
     --------------
     * No `dockerfile` specified: If any of the options regarding
     requirements, environment variables or copying files require us to build an
-    image, ZenML will build this image. Otherwise the `parent_image` will be
+    image, ZenML will build this image. Otherwise, the `parent_image` will be
     used to run the pipeline.
     * `dockerfile` specified: ZenML will first build an image based on the
     specified Dockerfile. If any of the options regarding
@@ -76,10 +92,10 @@ class DockerSettings(BaseSettings):
     Depending on the configuration of this object, requirements will be
     installed in the following order (each step optional):
     - The packages installed in your local python environment
+    - The packages required by the stack unless this is disabled by setting
+      `install_stack_requirements=False`.
+    - The packages specified via the `required_integrations`
     - The packages specified via the `requirements` attribute
-    - The packages specified via the `required_integrations` and potentially
-      stack requirements
-    - The packages specified via the `required_hub_plugins` attribute
 
     Attributes:
         parent_image: Full name of the Docker image that should be
@@ -107,18 +123,20 @@ class DockerSettings(BaseSettings):
         build_context_root: Build context root for the Docker build, only used
             when the `dockerfile` attribute is set. If this is left empty, the
             build context will only contain the Dockerfile.
-        build_options: Additional options that will be passed unmodified to the
-            Docker build call when building an image using the specified
-            `dockerfile`. You can use this to for example specify build
-            args or a target stage. See
-            https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
-            for a full list of available options.
+        parent_image_build_config: Configuration for the parent image build.
         skip_build: If set to `True`, the parent image will be used directly to
             run the steps of your pipeline.
+        prevent_build_reuse: Prevent the reuse of an existing build.
         target_repository: Name of the Docker repository to which the
             image should be pushed. This repository will be appended to the
             registry URI of the container registry of your stack and should
-            therefore **not** include any registry.
+            therefore **not** include any registry. If not specified, the
+            default repository name configured in the container registry
+            stack component settings will be used.
+        python_package_installer: The package installer to use for python
+            packages.
+        python_package_installer_args: Arguments to pass to the python package
+            installer.
         replicate_local_python_environment: If not `None`, ZenML will use the
             specified method to generate a requirements file that replicates
             the packages installed in the currently running python environment.
@@ -131,11 +149,7 @@ class DockerSettings(BaseSettings):
         required_integrations: List of ZenML integrations that should be
             installed. All requirements for the specified integrations will
             be installed inside the Docker image.
-        required_hub_plugins: List of ZenML Hub plugins to install.
-            Expected format: '(<author_username>/)<plugin_name>==<version>'.
-            If no version is specified, the latest version is taken. The
-            packages of required plugins and all their dependencies will be
-            installed inside the Docker image.
+        required_hub_plugins: DEPRECATED/UNUSED.
         install_stack_requirements: If `True`, ZenML will automatically detect
             if components of your active stack are part of a ZenML integration
             and install the corresponding requirements and apt packages.
@@ -145,89 +159,135 @@ class DockerSettings(BaseSettings):
         apt_packages: APT packages to install inside the Docker image.
         environment: Dictionary of environment variables to set inside the
             Docker image.
-        dockerignore: Path to a dockerignore file to use when building the
-            Docker image.
-        copy_files: DEPRECATED, use the `source_files` attribute instead.
-        copy_global_config: DEPRECATED/UNUSED.
+        build_config: Configuration for the main image build.
         user: If not `None`, will set the user, make it owner of the `/app`
             directory which contains all the user code and run the container
             entrypoint as this user.
-        source_files: Defines how the user source files will be handled when
-            building the Docker image.
-            * INCLUDE: The files will be included in the Docker image.
-            * DOWNLOAD: The files will be downloaded when running the image. If
-              this is specified, the files must be inside a registered code
-              repository and the repository must have no local changes,
-              otherwise the build will fail.
-            * DOWNLOAD_OR_INCLUDE: The files will be downloaded if they're
-              inside a registered code repository and the repository has no
-              local changes, otherwise they will be included in the image.
-            * IGNORE: The files will not be included or downloaded in the image.
-              If you use this option, you're responsible that all the files
-              to run your steps exist in the right place.
+        allow_including_files_in_images: If `True`, code can be included in the
+            Docker images if code download from a code repository or artifact
+            store is disabled or not possible.
+        allow_download_from_code_repository: If `True`, code can be downloaded
+            from a code repository if possible.
+        allow_download_from_artifact_store: If `True`, code can be downloaded
+            from the artifact store.
+        build_options: DEPRECATED, use parent_image_build_config.build_options
+            instead.
+        dockerignore: DEPRECATED, use build_config.dockerignore instead.
+        copy_files: DEPRECATED/UNUSED.
+        copy_global_config: DEPRECATED/UNUSED.
+        source_files: DEPRECATED. Use allow_including_files_in_images,
+            allow_download_from_code_repository and
+            allow_download_from_artifact_store instead.
     """
 
     parent_image: Optional[str] = None
     dockerfile: Optional[str] = None
     build_context_root: Optional[str] = None
-    build_options: Dict[str, Any] = {}
+    parent_image_build_config: Optional[DockerBuildConfig] = None
     skip_build: bool = False
-    target_repository: str = "zenml"
+    prevent_build_reuse: bool = False
+    target_repository: Optional[str] = None
+    python_package_installer: PythonPackageInstaller = (
+        PythonPackageInstaller.PIP
+    )
+    python_package_installer_args: Dict[str, Any] = {}
     replicate_local_python_environment: Optional[
         Union[List[str], PythonEnvironmentExportMethod]
-    ] = None
-    requirements: Union[None, str, List[str]] = None
+    ] = Field(default=None, union_mode="left_to_right")
+    requirements: Union[None, str, List[str]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     required_integrations: List[str] = []
-    required_hub_plugins: List[str] = []
     install_stack_requirements: bool = True
     apt_packages: List[str] = []
     environment: Dict[str, Any] = {}
+    user: Optional[str] = None
+    build_config: Optional[DockerBuildConfig] = None
+
+    allow_including_files_in_images: bool = True
+    allow_download_from_code_repository: bool = True
+    allow_download_from_artifact_store: bool = True
+
+    # Deprecated attributes
+    build_options: Dict[str, Any] = {}
     dockerignore: Optional[str] = None
     copy_files: bool = True
     copy_global_config: bool = True
-    user: Optional[str] = None
-
-    source_files: SourceFileMode = SourceFileMode.DOWNLOAD_OR_INCLUDE
+    source_files: Optional[str] = None
+    required_hub_plugins: List[str] = []
 
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
-        "copy_files", "copy_global_config"
+        "copy_files",
+        "copy_global_config",
+        "source_files",
+        "required_hub_plugins",
     )
 
-    @root_validator(pre=True)
-    def _migrate_copy_files(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        """Migrates the value from the old copy_files attribute.
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _migrate_source_files(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Migrate old source_files values.
 
         Args:
-            values: The settings values.
+            data: The model data.
+
+        Raises:
+            ValueError: If an invalid source file mode is specified.
 
         Returns:
-            The migrated settings values.
+            The migrated data.
         """
-        copy_files = values.get("copy_files", None)
+        source_files = data.get("source_files", None)
 
-        if copy_files is None:
-            return values
+        if source_files is None:
+            return data
 
-        if values.get("source_files", None):
-            # Ignore the copy files value in favor of the new source files
+        replacement_attributes = [
+            "allow_including_files_in_images",
+            "allow_download_from_code_repository",
+            "allow_download_from_artifact_store",
+        ]
+        if any(v in data for v in replacement_attributes):
             logger.warning(
-                "Both `copy_files` and `source_files` specified for the "
-                "DockerSettings, ignoring the `copy_files` value."
+                "Both `source_files` and one of %s specified for the "
+                "DockerSettings, ignoring the `source_files` value.",
+                replacement_attributes,
             )
-        elif copy_files is True:
-            values["source_files"] = SourceFileMode.INCLUDE
-        elif copy_files is False:
-            values["source_files"] = SourceFileMode.IGNORE
-
-        return values
-
-    @root_validator
-    def _validate_skip_build(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+            return data
+
+        allow_including_files_in_images = False
+        allow_download_from_code_repository = False
+        allow_download_from_artifact_store = False
+
+        if source_files == "download":
+            allow_download_from_code_repository = True
+        elif source_files == "include":
+            allow_including_files_in_images = True
+        elif source_files == "download_or_include":
+            allow_including_files_in_images = True
+            allow_download_from_code_repository = True
+        elif source_files == "ignore":
+            pass
+        else:
+            raise ValueError(f"Invalid source file mode `{source_files}`.")
+
+        data["allow_including_files_in_images"] = (
+            allow_including_files_in_images
+        )
+        data["allow_download_from_code_repository"] = (
+            allow_download_from_code_repository
+        )
+        data["allow_download_from_artifact_store"] = (
+            allow_download_from_artifact_store
+        )
+
+        return data
+
+    @model_validator(mode="after")
+    def _validate_skip_build(self) -> "DockerSettings":
         """Ensures that a parent image is passed when trying to skip the build.
 
-        Args:
-            values: The settings values.
-
         Returns:
             The validated settings values.
 
@@ -235,10 +295,7 @@ class DockerSettings(BaseSettings):
             ValueError: If the build should be skipped but no parent image
                 was specified.
         """
-        skip_build = values.get("skip_build", False)
-        parent_image = values.get("parent_image")
-
-        if skip_build and not parent_image:
+        if self.skip_build and not self.parent_image:
             raise ValueError(
                 "Docker settings that specify `skip_build=True` must always "
                 "contain a `parent_image`. This parent image will be used "
@@ -246,12 +303,11 @@ class DockerSettings(BaseSettings):
                 "Docker builds on top of it."
             )
 
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
+        return self
 
+    model_config = ConfigDict(
         # public attributes are immutable
-        allow_mutation = False
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )