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

zenml/config/server_config.py

@@ -381,6 +381,36 @@ class ServerConfiguration(BaseModel):
                     # Revert to the default value if the header is enabled
                     del data[k]
 
+        # Handle merging of user-defined secure_headers_csp value with the default value
+        if "secure_headers_csp" in data:
+            user_defined_csp = data["secure_headers_csp"]
+            if isinstance(user_defined_csp, str):
+                # Parse the user-defined CSP string into a dictionary
+                user_defined_csp_dict = {}
+                for directive in user_defined_csp.split(";"):
+                    directive = directive.strip()
+                    if directive:
+                        key, value = directive.split(" ", 1)
+                        user_defined_csp_dict[key] = value.strip("'\"")
+
+                # Merge the user-defined CSP dictionary with the default CSP dictionary
+                default_csp_dict = {}
+                for directive in DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP.split(
+                    ";"
+                ):
+                    directive = directive.strip()
+                    if directive:
+                        key, value = directive.split(" ", 1)
+                        default_csp_dict[key] = value.strip("'\"")
+
+                merged_csp_dict = {**default_csp_dict, **user_defined_csp_dict}
+
+                # Convert the merged CSP dictionary back to a string
+                merged_csp_str = "; ".join(
+                    f"{key} {value}" for key, value in merged_csp_dict.items()
+                )
+                data["secure_headers_csp"] = merged_csp_str
+
         return data
 
     @property

zenml/config/source.py

@@ -234,20 +234,16 @@ class NotebookSource(Source):
     """Source representing an object defined in a notebook.
 
     Attributes:
-        code_path: Path where the notebook cell code for this source is
-            uploaded.
         replacement_module: Name of the module from which this source should
             be loaded in case the code is not running in a notebook.
+        artifact_store_id: ID of the artifact store in which the replacement
+            module code is stored.
     """
 
-    code_path: Optional[str] = None
     replacement_module: Optional[str] = None
+    artifact_store_id: Optional[UUID] = None
     type: SourceType = SourceType.NOTEBOOK
 
-    # Private attribute that is used to store the code but should not be
-    # serialized
-    _cell_code: Optional[str] = None
-
     @field_validator("type")
     @classmethod
     def _validate_type(cls, value: SourceType) -> SourceType:

zenml/constants.py

@@ -162,7 +162,6 @@ ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING = (
 )
 ENV_ZENML_DISABLE_WORKSPACE_WARNINGS = "ZENML_DISABLE_WORKSPACE_WARNINGS"
 ENV_ZENML_SKIP_IMAGE_BUILDER_DEFAULT = "ZENML_SKIP_IMAGE_BUILDER_DEFAULT"
-ENV_ZENML_REQUIRES_CODE_DOWNLOAD = "ZENML_REQUIRES_CODE_DOWNLOAD"
 ENV_ZENML_SERVER = "ZENML_SERVER"
 ENV_ZENML_LOCAL_SERVER = "ZENML_LOCAL_SERVER"
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
@@ -173,6 +172,7 @@ ENV_ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES = (
 )
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
+ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
@@ -185,6 +185,9 @@ ENV_ZENML_SERVER_USE_LEGACY_DASHBOARD = (
     f"{ENV_ZENML_SERVER_PREFIX}USE_LEGACY_DASHBOARD"
 )
 ENV_ZENML_SERVER_AUTO_ACTIVATE = f"{ENV_ZENML_SERVER_PREFIX}AUTO_ACTIVATE"
+ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK = (
+    "ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK"
+)
 
 # Logging variables
 IS_DEBUG_ENV: bool = handle_bool_env_var(ENV_ZENML_DEBUG, default=False)
@@ -348,7 +351,6 @@ EMAIL_ANALYTICS = "/email-opt-in"
 EVENT_FLAVORS = "/event-flavors"
 EVENT_SOURCES = "/event-sources"
 FLAVORS = "/flavors"
-FULL_STACK = "/full-stack"
 GET_OR_CREATE = "/get-or-create"
 GRAPH = "/graph"
 HEALTH = "/health"
@@ -486,7 +488,7 @@ MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION = (
 )
 
 
-FINISHED_ONBOARDING_SURVEY_KEY = "awareness_channels"
+FINISHED_ONBOARDING_SURVEY_KEY = "finished_onboarding_survey"
 
 # Name validation
 BANNED_NAME_CHARACTERS = "\t\n\r\v\f"

zenml/entrypoints/base_entrypoint_configuration.py

@@ -17,15 +17,11 @@ import argparse
 import os
 import sys
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Set
+from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional, Set
 from uuid import UUID
 
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
-from zenml.constants import (
-    ENV_ZENML_REQUIRES_CODE_DOWNLOAD,
-    handle_bool_env_var,
-)
 from zenml.logger import get_logger
 from zenml.utils import (
     code_repository_utils,
@@ -194,22 +190,28 @@ class BaseEntrypointConfiguration(ABC):
         return Client().zen_store.get_deployment(deployment_id=deployment_id)
 
     def download_code_if_necessary(
-        self, deployment: "PipelineDeploymentResponse"
+        self,
+        deployment: "PipelineDeploymentResponse",
+        step_name: Optional[str] = None,
     ) -> None:
         """Downloads user code if necessary.
 
         Args:
             deployment: The deployment for which to download the code.
+            step_name: Name of the step to be run. This will be used to
+                determine whether code download is necessary. If not given,
+                the DockerSettings of the pipeline will be used to make that
+                decision instead.
 
         Raises:
             RuntimeError: If the current environment requires code download
                 but the deployment does not have a reference to any code.
         """
-        requires_code_download = handle_bool_env_var(
-            ENV_ZENML_REQUIRES_CODE_DOWNLOAD
+        should_download_code = self._should_download_code(
+            deployment=deployment, step_name=step_name
         )
 
-        if not requires_code_download:
+        if not should_download_code:
             return
 
         if code_reference := deployment.code_reference:
@@ -217,7 +219,7 @@ class BaseEntrypointConfiguration(ABC):
                 code_reference=code_reference
             )
         elif code_path := deployment.code_path:
-            self.download_code_from_artifact_store(code_path=code_path)
+            code_utils.download_code_from_artifact_store(code_path=code_path)
         else:
             raise RuntimeError(
                 "Code download required but no code reference or path provided."
@@ -259,30 +261,42 @@ class BaseEntrypointConfiguration(ABC):
         sys.path.insert(0, download_dir)
         os.chdir(download_dir)
 
-    def download_code_from_artifact_store(self, code_path: str) -> None:
-        """Download code from the artifact store.
+    def _should_download_code(
+        self,
+        deployment: "PipelineDeploymentResponse",
+        step_name: Optional[str] = None,
+    ) -> bool:
+        """Checks whether code should be downloaded.
 
         Args:
-            code_path: Path where the code is stored.
+            deployment: The deployment to check.
+            step_name: Name of the step to be run. This will be used to
+                determine whether code download is necessary. If not given,
+                the DockerSettings of the pipeline will be used to make that
+                decision instead.
+
+        Returns:
+            Whether code should be downloaded.
         """
-        logger.info(
-            "Downloading code from artifact store path `%s`.", code_path
+        docker_settings = (
+            deployment.step_configurations[step_name].config.docker_settings
+            if step_name
+            else deployment.pipeline_configuration.docker_settings
         )
 
-        # 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)
+        if (
+            deployment.code_reference
+            and docker_settings.allow_download_from_code_repository
+        ):
+            return True
 
-        code_utils.download_and_extract_code(
-            code_path=code_path, extract_dir=extract_dir
-        )
+        if (
+            deployment.code_path
+            and docker_settings.allow_download_from_artifact_store
+        ):
+            return True
 
-        source_utils.set_custom_source_root(extract_dir)
-        sys.path.insert(0, extract_dir)
-        os.chdir(extract_dir)
+        return False
 
     @abstractmethod
     def run(self) -> None:

zenml/entrypoints/step_entrypoint_configuration.py

@@ -151,9 +151,12 @@ class StepEntrypointConfiguration(BaseEntrypointConfiguration):
         # and stack component flavors are registered.
         integration_registry.activate_integrations()
 
-        self.download_code_if_necessary(deployment=deployment)
-
         step_name = self.entrypoint_args[STEP_NAME_OPTION]
+
+        self.download_code_if_necessary(
+            deployment=deployment, step_name=step_name
+        )
+
         pipeline_name = deployment.pipeline_configuration.name
 
         step = deployment.step_configurations[step_name]

zenml/enums.py

@@ -303,6 +303,8 @@ class EnvironmentType(StrEnum):
     PAPERSPACE = "paperspace"
     WSL = "wsl"
     LIGHTNING_AI_STUDIO = "lightning_ai_studio"
+    GITHUB_CODESPACES = "github_codespaces"
+    VSCODE_REMOTE_CONTAINER = "vscode_remote_container"
 
 
 class ModelStages(StrEnum):

zenml/environment.py

@@ -69,6 +69,10 @@ def get_environment() -> str:
         return EnvironmentType.BITBUCKET_CI
     elif Environment.in_ci():
         return EnvironmentType.GENERIC_CI
+    elif Environment.in_github_codespaces():
+        return EnvironmentType.GITHUB_CODESPACES
+    elif Environment.in_vscode_remote_container():
+        return EnvironmentType.VSCODE_REMOTE_CONTAINER
     elif Environment.in_lightning_ai_studio():
         return EnvironmentType.LIGHTNING_AI_STUDIO
     elif Environment.in_docker():
@@ -271,6 +275,33 @@ class Environment(metaclass=SingletonMetaClass):
             return True
         return False
 
+    @staticmethod
+    def in_github_codespaces() -> bool:
+        """If the current Python process is running in GitHub Codespaces.
+
+        Returns:
+            `True` if the current Python process is running in GitHub Codespaces,
+            `False` otherwise.
+        """
+        return (
+            "CODESPACES" in os.environ
+            or "GITHUB_CODESPACE_TOKEN" in os.environ
+            or "GITHUB_CODESPACES_PORT_FORWARDING_DOMAIN" in os.environ
+        )
+
+    @staticmethod
+    def in_vscode_remote_container() -> bool:
+        """If the current Python process is running in a VS Code Remote Container.
+
+        Returns:
+            `True` if the current Python process is running in a VS Code Remote Container,
+            `False` otherwise.
+        """
+        return (
+            "REMOTE_CONTAINERS" in os.environ
+            or "VSCODE_REMOTE_CONTAINERS_SESSION" in os.environ
+        )
+
     @staticmethod
     def in_paperspace_gradient() -> bool:
         """If the current Python process is running in Paperspace Gradient.

zenml/feature_stores/base_feature_store.py

@@ -14,9 +14,7 @@
 """The base class for feature stores."""
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Type, Union, cast
-
-import pandas as pd
+from typing import Any, Dict, List, Type, cast
 
 from zenml.enums import StackComponentType
 from zenml.stack import Flavor, StackComponent
@@ -42,10 +40,10 @@ class BaseFeatureStore(StackComponent, ABC):
     @abstractmethod
     def get_historical_features(
         self,
-        entity_df: Union[pd.DataFrame, str],
+        entity_df: Any,
         features: List[str],
         full_feature_names: bool = False,
-    ) -> pd.DataFrame:
+    ) -> Any:
         """Returns the historical features for training or batch scoring.
 
         Args:
@@ -54,7 +52,7 @@ class BaseFeatureStore(StackComponent, ABC):
             full_feature_names: Whether to return the full feature names.
 
         Returns:
-            The historical features as a Pandas DataFrame.
+            The historical features.
         """
 
     @abstractmethod

zenml/integrations/__init__.py

@@ -36,6 +36,7 @@ from zenml.integrations.gitlab import GitLabIntegration  # noqa
 from zenml.integrations.great_expectations import (  # noqa
     GreatExpectationsIntegration,
 )
+from zenml.integrations.lightning import LightningIntegration  # noqa
 from zenml.integrations.huggingface import HuggingfaceIntegration  # noqa
 from zenml.integrations.hyperai import HyperAIIntegration  # noqa
 from zenml.integrations.kaniko import KanikoIntegration  # noqa
@@ -48,7 +49,9 @@ from zenml.integrations.lightgbm import LightGBMIntegration  # noqa
 from zenml.integrations.mlflow import MlflowIntegration  # noqa
 from zenml.integrations.neptune import NeptuneIntegration  # noqa
 from zenml.integrations.neural_prophet import NeuralProphetIntegration  # noqa
+from zenml.integrations.numpy import NumpyIntegration  # noqa
 from zenml.integrations.openai import OpenAIIntegration  # noqa
+from zenml.integrations.pandas import PandasIntegration  # noqa
 from zenml.integrations.pigeon import PigeonIntegration # noqa
 from zenml.integrations.pillow import PillowIntegration  # noqa
 from zenml.integrations.polars import PolarsIntegration  # noqa

zenml/integrations/airflow/flavors/airflow_orchestrator_flavor.py

@@ -119,6 +119,15 @@ class AirflowOrchestratorConfig(
 
     local: bool = True
 
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return True
+
 
 class AirflowOrchestratorFlavor(BaseOrchestratorFlavor):
     """Flavor for the Airflow orchestrator."""

zenml/integrations/aws/__init__.py

@@ -45,8 +45,8 @@ class AWSIntegration(Integration):
     ]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
-    @staticmethod
-    def activate() -> None:
+    @classmethod
+    def activate(cls) -> None:
         """Activate the AWS integration."""
         from zenml.integrations.aws import service_connectors  # noqa
 

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -372,7 +372,8 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         # mainly for testing purposes, we wait for the pipeline to finish
         if self.config.synchronous:
             logger.info(
-                "Executing synchronously. Waiting for pipeline to finish..."
+                "Executing synchronously. Waiting for pipeline to finish... \n"
+                "At this point you can `Ctrl-C` out without cancelling the execution."
             )
             try:
                 pipeline_execution.wait(

zenml/integrations/azure/__init__.py

@@ -52,8 +52,8 @@ class AzureIntegration(Integration):
     ]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
-    @staticmethod
-    def activate() -> None:
+    @classmethod
+    def activate(cls) -> None:
         """Activate the Azure integration."""
         from zenml.integrations.azure import service_connectors  # noqa
 

zenml/integrations/azure/azureml_utils.py

@@ -0,0 +1,201 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""AzureML definitions."""
+
+from typing import Optional
+
+from azure.ai.ml import MLClient
+from azure.ai.ml.entities import Compute
+from azure.core.exceptions import (
+    ResourceNotFoundError,
+)
+
+from zenml.integrations.azure.flavors.azureml import (
+    AzureMLComputeSettings,
+    AzureMLComputeTypes,
+)
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def check_settings_and_compute_configuration(
+    parameter: str,
+    settings: AzureMLComputeSettings,
+    compute: Compute,
+) -> None:
+    """Utility function comparing a parameter between settings and compute.
+
+    Args:
+        parameter: the name of the parameter.
+        settings: The AzureML orchestrator settings.
+        compute: The compute instance or cluster from AzureML.
+    """
+    # Check the compute size
+    compute_value = getattr(compute, parameter)
+    settings_value = getattr(settings, parameter)
+
+    if settings_value is not None and settings_value != compute_value:
+        logger.warning(
+            f"The '{parameter}' defined in the settings '{settings_value}' "
+            "does not match the actual parameter of the instance: "
+            f"'{compute_value}'. Will ignore this setting for now."
+        )
+
+
+def create_or_get_compute(
+    client: MLClient,
+    settings: AzureMLComputeSettings,
+    default_compute_name: str,
+) -> Optional[str]:
+    """Creates or fetches the compute target if defined in the settings.
+
+    Args:
+        client: the AzureML client.
+        settings: the settings for the orchestrator.
+        default_compute_name: the default name for the compute target, if one
+            is not provided in the settings.
+
+    Returns:
+        None, if the orchestrator is using serverless compute or
+        str, the name of the compute target (instance or cluster).
+
+    Raises:
+        RuntimeError: if the fetched compute target is unsupported or the
+            mode defined in the setting does not match the type of the
+            compute target.
+    """
+    # If the mode is serverless, we can not fetch anything anyhow
+    if settings.mode == AzureMLComputeTypes.SERVERLESS:
+        return None
+
+    # If a name is not provided, generate one based on the orchestrator id
+    compute_name = settings.compute_name or default_compute_name
+
+    # Try to fetch the compute target
+    try:
+        compute = client.compute.get(compute_name)
+
+        logger.info(f"Using existing compute target: '{compute_name}'.")
+
+        # Check if compute size matches with the settings
+        check_settings_and_compute_configuration(
+            parameter="size", settings=settings, compute=compute
+        )
+
+        compute_type = compute.type
+
+        # Check the type and matches the settings
+        if compute_type == "computeinstance":  # Compute Instance
+            if settings.mode != AzureMLComputeTypes.COMPUTE_INSTANCE:
+                raise RuntimeError(
+                    "The mode of operation for the compute target defined"
+                    f"in the settings '{settings.mode}' does not match "
+                    f"the type of the compute target: `{compute_name}` "
+                    "which is a 'compute-instance'. Please make sure that "
+                    "the settings are adjusted properly."
+                )
+
+            if compute.state != "Running":
+                raise RuntimeError(
+                    f"The compute instance `{compute_name}` is not in a "
+                    "running state at the moment. Please make sure that "
+                    "the compute target is running, before executing the "
+                    "pipeline."
+                )
+
+            # Idle time before shutdown
+            check_settings_and_compute_configuration(
+                parameter="idle_time_before_shutdown_minutes",
+                settings=settings,
+                compute=compute,
+            )
+
+        elif compute_type == "amIcompute":  # Compute Cluster
+            if settings.mode != AzureMLComputeTypes.COMPUTE_CLUSTER:
+                raise RuntimeError(
+                    "The mode of operation for the compute target defined "
+                    f"in the settings '{settings.mode}' does not match "
+                    f"the type of the compute target: `{compute_name}` "
+                    "which is a 'compute-cluster'. Please make sure that "
+                    "the settings are adjusted properly."
+                )
+
+            if compute.provisioning_state != "Succeeded":
+                raise RuntimeError(
+                    f"The provisioning state '{compute.provisioning_state}'"
+                    f"of the compute cluster `{compute_name}` is not "
+                    "successful. Please make sure that the compute cluster "
+                    "is provisioned properly, before executing the "
+                    "pipeline."
+                )
+
+            for parameter in [
+                "idle_time_before_scale_down",
+                "max_instances",
+                "min_instances",
+                "tier",
+                "location",
+            ]:
+                # Check all possible configurations
+                check_settings_and_compute_configuration(
+                    parameter=parameter, settings=settings, compute=compute
+                )
+        else:
+            raise RuntimeError(f"Unsupported compute type: {compute_type}")
+        return compute_name
+
+    # If the compute target does not exist create it
+    except ResourceNotFoundError:
+        logger.info(
+            "Can not find the compute target with name: " f"'{compute_name}':"
+        )
+
+        if settings.mode == AzureMLComputeTypes.COMPUTE_INSTANCE:
+            logger.info(
+                "Creating a new compute instance. This might take a "
+                "few minutes."
+            )
+
+            from azure.ai.ml.entities import ComputeInstance
+
+            compute_instance = ComputeInstance(
+                name=compute_name,
+                size=settings.size,
+                idle_time_before_shutdown_minutes=settings.idle_time_before_shutdown_minutes,
+            )
+            client.begin_create_or_update(compute_instance).result()
+            return compute_name
+
+        elif settings.mode == AzureMLComputeTypes.COMPUTE_CLUSTER:
+            logger.info(
+                "Creating a new compute cluster. This might take a "
+                "few minutes."
+            )
+
+            from azure.ai.ml.entities import AmlCompute
+
+            compute_cluster = AmlCompute(
+                name=compute_name,
+                size=settings.size,
+                location=settings.location,
+                min_instances=settings.min_instances,
+                max_instances=settings.max_instances,
+                idle_time_before_scale_down=settings.idle_time_before_scaledown_down,
+                tier=settings.tier,
+            )
+            client.begin_create_or_update(compute_cluster).result()
+            return compute_name
+
+    return None