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

zenml/orchestrators/utils.py

@@ -13,31 +13,37 @@
 #  permissions and limitations under the License.
 """Utility functions for the orchestrator."""
 
+import os
 import random
-from datetime import datetime
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional, cast
 from uuid import UUID
 
 from zenml.client import Client
 from zenml.config.global_config import (
-    CONFIG_ENV_VAR_PREFIX,
     GlobalConfiguration,
 )
 from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
-    ENV_ZENML_BACKUP_SECRETS_STORE_PREFIX,
-    ENV_ZENML_SECRETS_STORE_PREFIX,
+    ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
+    ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION,
+    ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
-    PIPELINE_API_TOKEN_EXPIRES_MINUTES,
+    ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION,
 )
-from zenml.enums import StoreType
+from zenml.enums import APITokenType, AuthScheme, StackComponentType, StoreType
+from zenml.logger import get_logger
+from zenml.stack import StackComponent
+
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
 
 
-def get_orchestrator_run_name(pipeline_name: str) -> str:
+def get_orchestrator_run_name(
+    pipeline_name: str, max_length: Optional[int] = None
+) -> str:
     """Gets an orchestrator run name.
 
     This run name is not the same as the ZenML run name but can instead be
@@ -45,11 +51,31 @@ def get_orchestrator_run_name(pipeline_name: str) -> str:
 
     Args:
         pipeline_name: Name of the pipeline that will run.
+        max_length: Maximum length of the generated name.
+
+    Raises:
+        ValueError: If the max length is below 8 characters.
 
     Returns:
         The orchestrator run name.
     """
-    return f"{pipeline_name}_{random.Random().getrandbits(128):032x}"
+    suffix_length = 32
+    pipeline_name = f"{pipeline_name}_"
+
+    if max_length:
+        if max_length < 8:
+            raise ValueError(
+                "Maximum length for orchestrator run name must be 8 or above."
+            )
+
+        # Make sure we always have a certain suffix to guarantee no overlap
+        # with other runs
+        suffix_length = min(32, max(8, max_length - len(pipeline_name)))
+        pipeline_name = pipeline_name[: (max_length - suffix_length)]
+
+    suffix = "".join(random.choices("0123456789abcdef", k=suffix_length))
+
+    return f"{pipeline_name}{suffix}"
 
 
 def is_setting_enabled(
@@ -77,79 +103,113 @@ def is_setting_enabled(
 
 
 def get_config_environment_vars(
-    deployment: Optional["PipelineDeploymentResponse"] = None,
+    schedule_id: Optional[UUID] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
 ) -> Dict[str, str]:
     """Gets environment variables to set for mirroring the active config.
 
-    If a pipeline deployment is given, the environment variables will be set to
-    include a newly generated API token valid for the duration of the pipeline
-    run instead of the API token from the global config.
+    If a schedule ID, pipeline run ID or step run ID is given, and the current
+    client is not authenticated to a server with an API key, the environment
+    variables will be updated to include a newly generated workload API token
+    that will be valid for the duration of the schedule, pipeline run, or step
+    run instead of the current API token used to authenticate the client.
 
     Args:
-        deployment: Optional deployment to use for the environment variables.
+        schedule_id: Optional schedule ID to use to generate a new API token.
+        pipeline_run_id: Optional pipeline run ID to use to generate a new API
+            token.
+        step_run_id: Optional step run ID to use to generate a new API token.
 
     Returns:
         Environment variable dict.
     """
+    from zenml.login.credentials_store import get_credentials_store
     from zenml.zen_stores.rest_zen_store import RestZenStore
 
     global_config = GlobalConfiguration()
-    environment_vars = {}
-
-    for key in global_config.__fields__.keys():
-        if key == "store":
-            continue
-
-        value = getattr(global_config, key)
-        if value is not None:
-            environment_vars[CONFIG_ENV_VAR_PREFIX + key.upper()] = str(value)
-
-    if global_config.store:
-        store_dict = global_config.store.dict(exclude_none=True)
-        secrets_store_dict = store_dict.pop("secrets_store", None) or {}
-        backup_secrets_store_dict = (
-            store_dict.pop("backup_secrets_store", None) or {}
-        )
-
-        for key, value in store_dict.items():
-            if key in ["username", "password"]:
-                # Don't include the username and password as we use the token to
-                # authenticate with the server
-                continue
-
-            environment_vars[ENV_ZENML_STORE_PREFIX + key.upper()] = str(value)
-
-        for key, value in secrets_store_dict.items():
-            environment_vars[
-                ENV_ZENML_SECRETS_STORE_PREFIX + key.upper()
-            ] = str(value)
-
-        for key, value in backup_secrets_store_dict.items():
-            environment_vars[
-                ENV_ZENML_BACKUP_SECRETS_STORE_PREFIX + key.upper()
-            ] = str(value)
-
-        if deployment and global_config.store.type == StoreType.REST:
-            # When connected to a ZenML server, if a pipeline deployment is
-            # supplied, we need to fetch an API token that will be valid for the
-            # duration of the pipeline run.
+    environment_vars = global_config.get_config_environment_vars()
+
+    if (
+        global_config.store_configuration.type == StoreType.REST
+        and global_config.zen_store.get_store_info().auth_scheme
+        != AuthScheme.NO_AUTH
+    ):
+        credentials_store = get_credentials_store()
+        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:
             assert isinstance(global_config.zen_store, RestZenStore)
-            pipeline_id: Optional[UUID] = None
-            if deployment.pipeline:
-                pipeline_id = deployment.pipeline.id
-            schedule_id: Optional[UUID] = None
-            expires_minutes: Optional[int] = PIPELINE_API_TOKEN_EXPIRES_MINUTES
-            if deployment.schedule:
-                schedule_id = deployment.schedule.id
-                # If a schedule is given, this is a long running pipeline that
-                # should not have an API token that expires.
-                expires_minutes = None
-            api_token = global_config.zen_store.get_api_token(
-                pipeline_id=pipeline_id,
-                schedule_id=schedule_id,
-                expires_minutes=expires_minutes,
+
+            # 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(
+                    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,
+                )
+
+            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
             )
-            environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = api_token
+        elif api_token:
+            # For all other cases, the pipeline run environment is configured
+            # with the current access token.
+            environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
+                api_token.access_token
+            )
+
+    # Disable credentials caching to avoid storing sensitive information
+    # in the pipeline run environment
+    environment_vars[ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING] = "true"
 
     # Make sure to use the correct active stack/workspace which might come
     # from a .zen repository and not the global config
@@ -163,24 +223,99 @@ def get_config_environment_vars(
     return environment_vars
 
 
-def get_run_name(run_name_template: str) -> str:
-    """Fill out the run name template to get a complete run name.
+class register_artifact_store_filesystem:
+    """Context manager for the artifact_store/filesystem_registry dependency.
 
-    Args:
-        run_name_template: The run name template to fill out.
+    Even though it is rare, sometimes we bump into cases where we are trying to
+    load artifacts that belong to an artifact store which is different from
+    the active artifact store.
 
-    Raises:
-        ValueError: If the run name is empty.
+    In cases like this, we will try to instantiate the target artifact store
+    by creating the corresponding artifact store Python object, which ends up
+    registering the right filesystem in the filesystem registry.
 
-    Returns:
-        The run name derived from the template.
+    The problem is, the keys in the filesystem registry are schemes (such as
+    "s3://" or "gcs://"). If we have two artifact stores with the same set of
+    supported schemes, we might end up overwriting the filesystem that belongs
+    to the active artifact store (and its authentication). That's why we have
+    to re-instantiate the active artifact store again, so the correct filesystem
+    will be restored.
     """
-    date = datetime.utcnow().strftime("%Y_%m_%d")
-    time = datetime.utcnow().strftime("%H_%M_%S_%f")
-
-    run_name = run_name_template.format(date=date, time=time)
 
-    if run_name == "":
-        raise ValueError("Empty run names are not allowed.")
+    def __init__(self, target_artifact_store_id: Optional[UUID]) -> None:
+        """Initialization of the context manager.
+
+        Args:
+            target_artifact_store_id: the ID of the artifact store to load.
+        """
+        self.target_artifact_store_id = target_artifact_store_id
+
+    def __enter__(self) -> "BaseArtifactStore":
+        """Entering the context manager.
+
+        It creates an instance of the target artifact store to register the
+        correct filesystem in the registry.
+
+        Returns:
+            The target artifact store object.
+
+        Raises:
+            RuntimeError: If the target artifact store can not be fetched or
+                initiated due to missing dependencies.
+        """
+        try:
+            if self.target_artifact_store_id is not None:
+                if (
+                    Client().active_stack.artifact_store.id
+                    != self.target_artifact_store_id
+                ):
+                    get_logger(__name__).debug(
+                        f"Trying to use the artifact store with ID:"
+                        f"'{self.target_artifact_store_id}'"
+                        f"which is currently not the active artifact store."
+                    )
+
+                artifact_store_model_response = Client().get_stack_component(
+                    component_type=StackComponentType.ARTIFACT_STORE,
+                    name_id_or_prefix=self.target_artifact_store_id,
+                )
+                return cast(
+                    "BaseArtifactStore",
+                    StackComponent.from_model(artifact_store_model_response),
+                )
+            else:
+                return Client().active_stack.artifact_store
+
+        except KeyError:
+            raise RuntimeError(
+                "Unable to fetch the artifact store with id: "
+                f"'{self.target_artifact_store_id}'. Check whether the "
+                "artifact store still exists and you have the right "
+                "permissions to access it."
+            )
+        except ImportError:
+            raise RuntimeError(
+                "Unable to load the implementation of the artifact store with"
+                f"id: '{self.target_artifact_store_id}'. Please make sure that "
+                "the environment that you are loading this artifact from "
+                "has the right dependencies."
+            )
 
-    return run_name
+    def __exit__(
+        self,
+        exc_type: Optional[Any],
+        exc_value: Optional[Any],
+        traceback: Optional[Any],
+    ) -> None:
+        """Set it back to the original state.
+
+        Args:
+            exc_type: The class of the exception
+            exc_value: The instance of the exception
+            traceback: The traceback of the exception
+        """
+        if ENV_ZENML_SERVER not in os.environ:
+            # As we exit the handler, we have to re-register the filesystem
+            # that belongs to the active artifact store as it may have been
+            # overwritten.
+            Client().active_stack.artifact_store._register()

zenml/orchestrators/wheeled_orchestrator.py

@@ -0,0 +1,147 @@
+#  Copyright (c) ZenML GmbH 2021. 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.
+"""Wheeled orchestrator class."""
+
+import os
+import re
+import subprocess
+import tempfile
+from abc import ABC
+
+from zenml import __version__
+from zenml.io import fileio
+from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestrator
+from zenml.utils.io_utils import copy_dir
+from zenml.utils.source_utils import get_source_root
+
+logger = get_logger(__name__)
+
+DEFAULT_PACKAGE_NAME = "zenmlproject"
+
+
+class WheeledOrchestrator(BaseOrchestrator, ABC):
+    """Base class for wheeled orchestrators."""
+
+    package_name = DEFAULT_PACKAGE_NAME
+    package_version = __version__
+
+    def copy_repository_to_temp_dir_and_add_setup_py(self) -> str:
+        """Copy the repository to a temporary directory and add a setup.py file.
+
+        Returns:
+            Path to the temporary directory containing the copied repository.
+        """
+        repo_path = get_source_root()
+
+        self.package_name = f"{DEFAULT_PACKAGE_NAME}_{self.sanitize_name(os.path.basename(repo_path))}"
+
+        # Create a temporary folder
+        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
+
+        # Create a folder within the temporary directory
+        temp_repo_path = os.path.join(temp_dir, self.package_name)
+        fileio.mkdir(temp_repo_path)
+
+        # Copy the repository to the temporary directory
+        copy_dir(repo_path, temp_repo_path)
+
+        # Create init file in the copied directory
+        init_file_path = os.path.join(temp_repo_path, "__init__.py")
+        with fileio.open(init_file_path, "w") as f:
+            f.write("")
+
+        # Create a setup.py file
+        setup_py_content = f"""
+from setuptools import setup, find_packages
+
+setup(
+    name="{self.package_name}",
+    version="{self.package_version}",
+    packages=find_packages(),
+)
+"""
+        setup_py_path = os.path.join(temp_dir, "setup.py")
+        with fileio.open(setup_py_path, "w") as f:
+            f.write(setup_py_content)
+
+        return temp_dir
+
+    def create_wheel(self, temp_dir: str) -> str:
+        """Create a wheel for the package in the given temporary directory.
+
+        Args:
+            temp_dir (str): Path to the temporary directory containing the package.
+
+        Raises:
+            RuntimeError: If the wheel file could not be created.
+
+        Returns:
+            str: Path to the created wheel file.
+        """
+        # Change to the temporary directory
+        original_dir = os.getcwd()
+        os.chdir(temp_dir)
+
+        try:
+            # Run the `pip wheel` command to create the wheel
+            result = subprocess.run(
+                ["pip", "wheel", "."], check=True, capture_output=True
+            )
+            logger.debug(f"Wheel creation stdout: {result.stdout.decode()}")
+            logger.debug(f"Wheel creation stderr: {result.stderr.decode()}")
+
+            # Find the created wheel file
+            wheel_file = next(
+                (
+                    file
+                    for file in os.listdir(temp_dir)
+                    if file.endswith(".whl")
+                ),
+                None,
+            )
+
+            if wheel_file is None:
+                raise RuntimeError("Failed to create wheel file.")
+
+            wheel_path = os.path.join(temp_dir, wheel_file)
+
+            # Verify the wheel file is a valid zip file
+            import zipfile
+
+            if not zipfile.is_zipfile(wheel_path):
+                raise RuntimeError(
+                    f"The file {wheel_path} is not a valid zip file."
+                )
+
+            return wheel_path
+        finally:
+            # Change back to the original directory
+            os.chdir(original_dir)
+
+    def sanitize_name(self, name: str) -> str:
+        """Sanitize the value to be used in a cluster name.
+
+        Args:
+            name: Arbitrary input cluster name.
+
+        Returns:
+            Sanitized cluster name.
+        """
+        name = re.sub(
+            r"[^a-z0-9-]", "-", name.lower()
+        )  # replaces any character that is not a lowercase letter, digit, or hyphen with a hyphen
+        name = re.sub(r"^[-]+", "", name)  # trim leading hyphens
+        name = re.sub(r"[-]+$", "", name)  # trim trailing hyphens
+        return name

zenml/pipelines/__init__.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  Copyright (c) ZenML GmbH 2023. 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.
@@ -11,21 +11,8 @@
 #  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.
-"""A ZenML pipeline consists of tasks that execute in order and yield artifacts.
-
-The artifacts are automatically stored within the artifact store and metadata 
-is tracked by ZenML. Each individual task within a pipeline is known as a
-step. The standard pipelines within ZenML are designed to have easy interfaces
-to add pre-decided steps, with the order also pre-decided. Other sorts of
-pipelines can be created as well from scratch, building on the `BasePipeline` class.
-
-Pipelines can be written as simple functions. They are created by using decorators appropriate to the specific use case you have. The moment it is `run`, a pipeline is compiled and passed directly to the orchestrator.
-"""
-
-
-from zenml.config import DockerSettings
 from zenml.config.schedule import Schedule
-from zenml.pipelines.base_pipeline import BasePipeline
 from zenml.pipelines.pipeline_decorator import pipeline
+from zenml.pipelines.pipeline_context import PipelineContext, get_pipeline_context
 
-__all__ = ["BasePipeline", "DockerSettings", "pipeline", "Schedule"]
+__all__ = ["pipeline", "Schedule", "PipelineContext", "get_pipeline_context"]