zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.65.0.dev20240906__py3-none-any.whl

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/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

zenml/integrations/azure/flavors/azureml.py

@@ -0,0 +1,139 @@
+#  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 pydantic import model_validator
+
+from zenml.config.base_settings import BaseSettings
+from zenml.logger import get_logger
+from zenml.utils.enum_utils import StrEnum
+
+logger = get_logger(__name__)
+
+
+class AzureMLComputeTypes(StrEnum):
+    """Enum for different types of compute on AzureML."""
+
+    SERVERLESS = "serverless"
+    COMPUTE_INSTANCE = "compute-instance"
+    COMPUTE_CLUSTER = "compute-cluster"
+
+
+class AzureMLComputeSettings(BaseSettings):
+    """Settings for the AzureML compute.
+
+    These settings adjust the compute resources that will be used by the
+    pipeline execution.
+
+    There are three possible use cases for this implementation:
+
+        1. Serverless compute (default behavior):
+            - The `mode` is set to `serverless` (default behavior).
+            - All the other parameters become irrelevant and will throw a
+            warning if set.
+
+        2. Compute instance:
+            - The `mode` is set to `compute-instance`.
+            - In this case, users have to provide a `compute-name`.
+                - If a compute instance exists with this name, this instance
+                will be used and all the other parameters become irrelevant
+                and will throw a warning if set.
+                - If a compute instance does not already exist, ZenML will
+                create it. You can use the parameters `compute_size` and
+                `idle_type_before_shutdown_minutes` for this operation.
+
+        3. Compute cluster:
+            - The `mode` is set to `compute-cluster`.
+            - In this case, users have to provide a `compute-name`.
+                - If a compute cluster exists with this name, this instance
+                will be used and all the other parameters become irrelevant
+                and will throw a warning if set.
+                - If a compute cluster does not already exist, ZenML will
+                create it. You can set the additional parameters for this
+                operation.
+    """
+
+    # Mode for compute
+    mode: AzureMLComputeTypes = AzureMLComputeTypes.SERVERLESS
+
+    # Common Configuration for Compute Instances and Clusters
+    compute_name: Optional[str] = None
+    size: Optional[str] = None
+
+    # Additional configuration for a Compute Instance
+    idle_time_before_shutdown_minutes: Optional[int] = None
+
+    # Additional configuration for a Compute Cluster
+    idle_time_before_scaledown_down: Optional[int] = None
+    location: Optional[str] = None
+    min_instances: Optional[int] = None
+    max_instances: Optional[int] = None
+    tier: Optional[str] = None
+
+    @model_validator(mode="after")
+    def azureml_settings_validator(self) -> "AzureMLComputeSettings":
+        """Checks whether the right configuration is set based on mode.
+
+        Returns:
+            the instance itself.
+
+        Raises:
+            AssertionError: if a name is not provided when working with
+                instances and clusters.
+        """
+        viable_configuration_fields = {
+            AzureMLComputeTypes.SERVERLESS: {"mode"},
+            AzureMLComputeTypes.COMPUTE_INSTANCE: {
+                "mode",
+                "compute_name",
+                "size",
+                "idle_time_before_shutdown_minutes",
+            },
+            AzureMLComputeTypes.COMPUTE_CLUSTER: {
+                "mode",
+                "compute_name",
+                "size",
+                "idle_time_before_scaledown_down",
+                "location",
+                "min_instances",
+                "max_instances",
+                "tier",
+            },
+        }
+        viable_fields = viable_configuration_fields[self.mode]
+
+        for field in self.model_fields_set:
+            if (
+                field not in viable_fields
+                and field in AzureMLComputeSettings.model_fields
+            ):
+                logger.warning(
+                    f"In the {self.__class__.__name__} settings, the mode of "
+                    f"operation is set to {self.mode}. In this mode, you can "
+                    f"not configure the parameter '{field}'. This "
+                    "configuration will be ignored."
+                )
+
+        if (
+            self.mode == AzureMLComputeTypes.COMPUTE_INSTANCE
+            or self.mode == AzureMLComputeTypes.COMPUTE_CLUSTER
+        ):
+            assert self.compute_name is not None, (
+                "When you are working with compute instances and clusters, "
+                "please define a name for the compute target."
+            )
+
+        return self