zenml-nightly 0.70.0.dev20241203__py3-none-any.whl → 0.70.0.dev20241205__py3-none-any.whl

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -134,6 +134,17 @@ class KubernetesOrchestratorConfig(
         """
         return True
 
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Kubernetes orchestrator starts step pods from a pipeline pod.
+        # This is currently not supported when using client-side caching.
+        return False
+
 
 class KubernetesOrchestratorFlavor(BaseOrchestratorFlavor):
     """Kubernetes orchestrator flavor."""

zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py

@@ -94,6 +94,17 @@ class LightningOrchestratorConfig(
         """
         return False
 
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Lightning orchestrator starts step studios from a pipeline studio.
+        # This is currently not supported when using client-side caching.
+        return False
+
 
 class LightningOrchestratorFlavor(BaseOrchestratorFlavor):
     """Lightning orchestrator flavor."""

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -144,3 +144,15 @@ class SkypilotBaseOrchestratorConfig(
             True if this config is for a local component, False otherwise.
         """
         return False
+
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Skypilot orchestrator runs the entire pipeline in a single VM, or
+        # starts additional VMs from the root VM. Both of those cases are
+        # currently not supported when using client-side caching.
+        return False

zenml/materializers/built_in_materializer.py

@@ -429,7 +429,7 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         # doesn't work for non-serializable types as they
         # are saved as list of lists in different files
         if _is_serializable(data):
-            return {self.data_path: VisualizationType.JSON}
+            return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
         return {}
 
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:

zenml/orchestrators/base_orchestrator.py

@@ -84,6 +84,15 @@ class BaseOrchestratorConfig(StackComponentConfig):
         """
         return False
 
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        return True
+
 
 class BaseOrchestrator(StackComponent, ABC):
     """Base class for all orchestrators.
@@ -205,6 +214,7 @@ class BaseOrchestrator(StackComponent, ABC):
 
         if (
             placeholder_run
+            and self.config.supports_client_side_caching
             and not deployment.schedule
             and not prevent_client_side_caching
         ):
@@ -232,6 +242,8 @@ class BaseOrchestrator(StackComponent, ABC):
                 self._cleanup_run()
                 logger.info("All steps of the pipeline run were cached.")
                 return
+        else:
+            logger.debug("Skipping client-side caching.")
 
         try:
             if metadata_iterator := self.prepare_or_run_pipeline(

zenml/orchestrators/output_utils.py

@@ -19,6 +19,7 @@ from uuid import uuid4
 
 from zenml.client import Client
 from zenml.logger import get_logger
+from zenml.utils import string_utils
 
 if TYPE_CHECKING:
     from zenml.artifact_stores import BaseArtifactStore
@@ -75,10 +76,13 @@ def prepare_output_artifact_uris(
     artifact_store = stack.artifact_store
     output_artifact_uris: Dict[str, str] = {}
     for output_name in step.config.outputs.keys():
+        substituted_output_name = string_utils.format_name_template(
+            output_name, substitutions=step_run.config.substitutions
+        )
         artifact_uri = generate_artifact_uri(
             artifact_store=stack.artifact_store,
             step_run=step_run,
-            output_name=output_name,
+            output_name=substituted_output_name,
         )
         if artifact_store.exists(artifact_uri):
             raise RuntimeError("Artifact already exists")

zenml/orchestrators/utils.py

@@ -26,10 +26,12 @@ from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
     ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
+    ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION,
     ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
+    ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION,
 )
-from zenml.enums import AuthScheme, StackComponentType, StoreType
+from zenml.enums import APITokenType, AuthScheme, StackComponentType, StoreType
 from zenml.logger import get_logger
 from zenml.stack import StackComponent
 
@@ -137,37 +139,63 @@ def get_config_environment_vars(
         url = global_config.store_configuration.url
         api_token = credentials_store.get_token(url, allow_expired=False)
         if schedule_id or pipeline_run_id or step_run_id:
-            # When connected to an authenticated ZenML server, if a schedule ID,
-            # pipeline run ID or step run ID is supplied, we need to fetch a new
-            # workload API token scoped to the schedule, pipeline run or step
-            # run.
             assert isinstance(global_config.zen_store, RestZenStore)
 
-            # If only a schedule is given, the pipeline run credentials will
-            # be valid for the entire duration of the schedule.
-            api_key = credentials_store.get_api_key(url)
-            if not api_key and not pipeline_run_id and not step_run_id:
+            # The user has the option to manually set an expiration for the API
+            # token generated for a pipeline run. In this case, we generate a new
+            # generic API token that will be valid for the indicated duration.
+            if (
+                pipeline_run_id
+                and ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION != 0
+            ):
                 logger.warning(
-                    "An API token without an expiration time will be generated "
-                    "and used to run this pipeline on a schedule. This is very "
-                    "insecure because the API token will be valid for the "
-                    "entire lifetime of the schedule and can be used to access "
-                    "your user account if accidentally leaked. When deploying "
-                    "a pipeline on a schedule, it is strongly advised to use a "
-                    "service account API key to authenticate to the ZenML "
-                    "server instead of your regular user account. For more "
-                    "information, see "
-                    "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
+                    f"An unscoped API token will be generated for this pipeline "
+                    f"run that will expire after "
+                    f"{ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION} "
+                    f"seconds instead of being scoped to the pipeline run "
+                    f"and not having an expiration time. This is more insecure "
+                    f"because the API token will remain valid even after the "
+                    f"pipeline run completes its execution. This option has "
+                    "been explicitly enabled by setting the "
+                    f"{ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION} environment "
+                    f"variable"
+                )
+                new_api_token = global_config.zen_store.get_api_token(
+                    token_type=APITokenType.GENERIC,
+                    expires_in=ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION,
                 )
 
-            # The schedule, pipeline run or step run credentials are scoped to
-            # the schedule, pipeline run or step run and will only be valid for
-            # the duration of the schedule/pipeline run/step run.
-            new_api_token = global_config.zen_store.get_api_token(
-                schedule_id=schedule_id,
-                pipeline_run_id=pipeline_run_id,
-                step_run_id=step_run_id,
-            )
+            else:
+                # If a schedule ID, pipeline run ID or step run ID is supplied,
+                # we need to fetch a new workload API token scoped to the
+                # schedule, pipeline run or step run.
+
+                # If only a schedule is given, the pipeline run credentials will
+                # be valid for the entire duration of the schedule.
+                api_key = credentials_store.get_api_key(url)
+                if not api_key and not pipeline_run_id and not step_run_id:
+                    logger.warning(
+                        "An API token without an expiration time will be generated "
+                        "and used to run this pipeline on a schedule. This is very "
+                        "insecure because the API token will be valid for the "
+                        "entire lifetime of the schedule and can be used to access "
+                        "your user account if accidentally leaked. When deploying "
+                        "a pipeline on a schedule, it is strongly advised to use a "
+                        "service account API key to authenticate to the ZenML "
+                        "server instead of your regular user account. For more "
+                        "information, see "
+                        "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
+                    )
+
+                # The schedule, pipeline run or step run credentials are scoped to
+                # the schedule, pipeline run or step run and will only be valid for
+                # the duration of the schedule/pipeline run/step run.
+                new_api_token = global_config.zen_store.get_api_token(
+                    token_type=APITokenType.WORKLOAD,
+                    schedule_id=schedule_id,
+                    pipeline_run_id=pipeline_run_id,
+                    step_run_id=step_run_id,
+                )
 
             environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
                 new_api_token

zenml/service_connectors/service_connector_utils.py

@@ -60,15 +60,9 @@ def _raise_specific_cloud_exception_if_needed(
     orchestrators: List[ResourcesInfo],
     container_registries: List[ResourcesInfo],
 ) -> None:
-    AWS_DOCS = (
-        "https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/aws-service-connector"
-    )
-    GCP_DOCS = (
-        "https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/gcp-service-connector"
-    )
-    AZURE_DOCS = (
-        "https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/azure-service-connector"
-    )
+    AWS_DOCS = "https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/aws-service-connector"
+    GCP_DOCS = "https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/gcp-service-connector"
+    AZURE_DOCS = "https://docs.zenml.io/how-to/infrastructure-deployment/auth-management/azure-service-connector"
 
     if not artifact_stores:
         error_msg = (

zenml/stack_deployments/aws_stack_deployment.py

@@ -73,6 +73,7 @@ of any potential costs:
 - An ECR repository registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/aws).
 - Sagemaker registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/sagemaker)
 as well as a [ZenML step operator](https://docs.zenml.io/stack-components/step-operators/sagemaker).
+- A CodeBuild project registered as a [ZenML image builder](https://docs.zenml.io/stack-components/image-builder/aws).
 - An IAM user and IAM role with the minimum necessary permissions to access the
 above resources.
 - An AWS access key used to give access to ZenML to connect to the above
@@ -158,6 +159,26 @@ console.
                 "ecr:PutImage",
                 "ecr:GetAuthorizationToken",
             ],
+            "CloudBuild (Client)": [
+                "codebuild:CreateProject",
+                "codebuild:BatchGetBuilds",
+            ],
+            "CloudBuild (Service)": [
+                "s3:GetObject",
+                "s3:GetObjectVersion",
+                "logs:CreateLogGroup",
+                "logs:CreateLogStream",
+                "logs:PutLogEvents",
+                "ecr:BatchGetImage",
+                "ecr:DescribeImages",
+                "ecr:BatchCheckLayerAvailability",
+                "ecr:GetDownloadUrlForLayer",
+                "ecr:InitiateLayerUpload",
+                "ecr:UploadLayerPart",
+                "ecr:CompleteLayerUpload",
+                "ecr:PutImage",
+                "ecr:GetAuthorizationToken",
+            ],
             "SageMaker (Client)": [
                 "sagemaker:CreatePipeline",
                 "sagemaker:StartPipelineExecution",
@@ -243,6 +264,7 @@ console.
             param_ResourceName=f"zenml-{random_str(6).lower()}",
             param_ZenMLServerURL=self.zenml_server_url,
             param_ZenMLServerAPIToken=self.zenml_server_api_token,
+            param_CodeBuild="true",
         )
         # Encode the parameters as URL query parameters
         query_params = "&".join([f"{k}={v}" for k, v in params.items()])

zenml/utils/archivable.py

@@ -15,11 +15,21 @@
 
 import io
 import tarfile
+import zipfile
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import IO, Any, Dict
+from typing import IO, Any, Dict, Optional
 
 from zenml.io import fileio
+from zenml.utils.enum_utils import StrEnum
+
+
+class ArchiveType(StrEnum):
+    """Archive types supported by the ZenML build context."""
+
+    TAR = "tar"
+    TAR_GZ = "tar.gz"
+    ZIP = "zip"
 
 
 class Archivable(ABC):
@@ -81,52 +91,71 @@ class Archivable(ABC):
                     self._extra_files[file_destination.as_posix()] = f.read()
 
     def write_archive(
-        self, output_file: IO[bytes], use_gzip: bool = True
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> None:
         """Writes an archive of the build context to the given file.
 
         Args:
             output_file: The file to write the archive to.
-            use_gzip: Whether to use `gzip` to compress the file.
+            archive_type: The type of archive to create.
         """
         files = self.get_files()
         extra_files = self.get_extra_files()
+        close_fileobj: Optional[Any] = None
+        fileobj: Any = output_file
 
-        if use_gzip:
-            from gzip import GzipFile
-
-            # We don't use the builtin gzip functionality of the `tarfile`
-            # library as that one includes the tar filename and creation
-            # timestamp in the archive which causes the hash of the resulting
-            # file to be different each time. We use this hash to avoid
-            # duplicate uploads, which is why we pass empty values for filename
-            # and mtime here.
-            fileobj: Any = GzipFile(
-                filename="", mode="wb", fileobj=output_file, mtime=0.0
-            )
+        if archive_type == ArchiveType.ZIP:
+            fileobj = zipfile.ZipFile(output_file, "w", zipfile.ZIP_DEFLATED)
         else:
-            fileobj = output_file
-
-        with tarfile.open(mode="w", fileobj=fileobj) as tf:
-            for archive_path, file_path in files.items():
-                if archive_path in extra_files:
-                    continue
-
-                if info := tf.gettarinfo(file_path, arcname=archive_path):
-                    if info.isfile():
-                        with open(file_path, "rb") as f:
-                            tf.addfile(info, f)
+            if archive_type == ArchiveType.TAR_GZ:
+                from gzip import GzipFile
+
+                # We don't use the builtin gzip functionality of the `tarfile`
+                # library as that one includes the tar filename and creation
+                # timestamp in the archive which causes the hash of the resulting
+                # file to be different each time. We use this hash to avoid
+                # duplicate uploads, which is why we pass empty values for filename
+                # and mtime here.
+                close_fileobj = fileobj = GzipFile(
+                    filename="", mode="wb", fileobj=output_file, mtime=0.0
+                )
+            fileobj = tarfile.open(mode="w", fileobj=fileobj)
+
+        try:
+            with fileobj as af:
+                for archive_path, file_path in files.items():
+                    if archive_path in extra_files:
+                        continue
+                    if archive_type == ArchiveType.ZIP:
+                        assert isinstance(af, zipfile.ZipFile)
+                        af.write(file_path, arcname=archive_path)
                     else:
-                        tf.addfile(info, None)
-
-            for archive_path, contents in extra_files.items():
-                info = tarfile.TarInfo(archive_path)
-                contents_encoded = contents.encode("utf-8")
-                info.size = len(contents_encoded)
-                tf.addfile(info, io.BytesIO(contents_encoded))
-
-        if use_gzip:
-            fileobj.close()
+                        assert isinstance(af, tarfile.TarFile)
+                        if info := af.gettarinfo(
+                            file_path, arcname=archive_path
+                        ):
+                            if info.isfile():
+                                with open(file_path, "rb") as f:
+                                    af.addfile(info, f)
+                            else:
+                                af.addfile(info, None)
+
+                for archive_path, contents in extra_files.items():
+                    contents_encoded = contents.encode("utf-8")
+
+                    if archive_type == ArchiveType.ZIP:
+                        assert isinstance(af, zipfile.ZipFile)
+                        af.writestr(archive_path, contents_encoded)
+                    else:
+                        assert isinstance(af, tarfile.TarFile)
+                        info = tarfile.TarInfo(archive_path)
+                        info.size = len(contents_encoded)
+                        af.addfile(info, io.BytesIO(contents_encoded))
+        finally:
+            if close_fileobj:
+                close_fileobj.close()
 
         output_file.seek(0)
 

zenml/utils/code_utils.py

@@ -25,7 +25,7 @@ from zenml.client import Client
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils import source_utils, string_utils
-from zenml.utils.archivable import Archivable
+from zenml.utils.archivable import Archivable, ArchiveType
 
 if TYPE_CHECKING:
     from git.repo.base import Repo
@@ -152,15 +152,19 @@ class CodeArchive(Archivable):
         return all_files
 
     def write_archive(
-        self, output_file: IO[bytes], use_gzip: bool = True
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> None:
         """Writes an archive of the build context to the given file.
 
         Args:
             output_file: The file to write the archive to.
-            use_gzip: Whether to use `gzip` to compress the file.
+            archive_type: The type of archive to create.
         """
-        super().write_archive(output_file=output_file, use_gzip=use_gzip)
+        super().write_archive(
+            output_file=output_file, archive_type=archive_type
+        )
         archive_size = os.path.getsize(output_file.name)
         if archive_size > 20 * 1024 * 1024:
             logger.warning(

zenml/utils/docker_utils.py

@@ -266,6 +266,14 @@ def push_image(
     logger.info("Finished pushing Docker image.")
 
     image_name_without_tag, _ = image_name.rsplit(":", maxsplit=1)
+
+    image = docker_client.images.get(image_name)
+    repo_digests: List[str] = image.attrs["RepoDigests"]
+
+    for digest in repo_digests:
+        if digest.startswith(f"{image_name_without_tag}@"):
+            return digest
+
     for info in reversed(aux_info):
         try:
             repo_digest = info["Digest"]
@@ -304,6 +312,7 @@ def get_image_digest(image_name: str) -> Optional[str]:
 
     image = docker_client.images.get(image_name)
     repo_digests = image.attrs["RepoDigests"]
+
     if len(repo_digests) == 1:
         return cast(str, repo_digests[0])
     else:

zenml/zen_server/auth.py

@@ -821,7 +821,10 @@ def generate_access_token(
         response: The FastAPI response object.
         device: The device used for authentication.
         api_key: The service account API key used for authentication.
-        expires_in: The number of seconds until the token expires.
+        expires_in: The number of seconds until the token expires. If not set,
+            the default value is determined automatically based on the server
+            configuration and type of token. If set to 0, the token will not
+            expire.
         schedule_id: The ID of the schedule to scope the token to.
         pipeline_run_id: The ID of the pipeline run to scope the token to.
         step_run_id: The ID of the step run to scope the token to.
@@ -835,7 +838,9 @@ def generate_access_token(
     # according to the values configured in the server config. Device tokens are
     # handled separately from regular user tokens.
     expires: Optional[datetime] = None
-    if expires_in:
+    if expires_in == 0:
+        expires_in = None
+    elif expires_in is not None:
         expires = datetime.utcnow() + timedelta(seconds=expires_in)
     elif device:
         # If a device was used for authentication, the token will expire

zenml/zen_server/dashboard/assets/{404-NVXKFp-x.js → 404-Cqu3EDCm.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-DeK6qiuw.js";import{f as s,r as t}from"./index-CCOPpudF.js";import{E as r}from"./EmptyState-BzdlCwp3.js";import{S as a}from"./help-Cc9bBIJH.js";import{L as o}from"./@react-router-B3Z5rLr2.js";import"./@tanstack-DT5WLu9C.js";import"./@reactflow-CK0KJUen.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-DeK6qiuw.js";import{f as s,r as t}from"./index-FO-p0GU7.js";import{E as r}from"./EmptyState-BzdlCwp3.js";import{S as a}from"./help-Cc9bBIJH.js";import{L as o}from"./@react-router-B3Z5rLr2.js";import"./@tanstack-DT5WLu9C.js";import"./@reactflow-D2Y7BWwz.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};