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

zenml/integrations/skypilot_kubernetes/orchestrators/skypilot_kubernetes_vm_orchestrator.py

@@ -0,0 +1,74 @@
+#  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.
+#  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.
+"""Implementation of the a Skypilot based Kubernetes VM orchestrator."""
+
+from typing import TYPE_CHECKING, Optional, Type, cast
+
+import sky
+
+from zenml.integrations.skypilot.orchestrators.skypilot_base_vm_orchestrator import (
+    SkypilotBaseOrchestrator,
+)
+from zenml.integrations.skypilot_kubernetes.flavors.skypilot_orchestrator_kubernetes_vm_flavor import (
+    SkypilotKubernetesOrchestratorConfig,
+    SkypilotKubernetesOrchestratorSettings,
+)
+from zenml.logger import get_logger
+
+if TYPE_CHECKING:
+    from zenml.config.base_settings import BaseSettings
+
+logger = get_logger(__name__)
+
+
+class SkypilotKubernetesOrchestrator(SkypilotBaseOrchestrator):
+    """Orchestrator responsible for running pipelines remotely in a VM on Kubernetes.
+
+    This orchestrator does not support running on a schedule.
+    """
+
+    @property
+    def cloud(self) -> sky.clouds.Cloud:
+        """The type of sky cloud to use.
+
+        Returns:
+            A `sky.clouds.Cloud` instance.
+        """
+        return sky.clouds.Kubernetes()
+
+    @property
+    def config(self) -> SkypilotKubernetesOrchestratorConfig:
+        """Returns the `SkypilotKubernetesOrchestratorConfig` config.
+
+        Returns:
+            The configuration.
+        """
+        return cast(SkypilotKubernetesOrchestratorConfig, self._config)
+
+    @property
+    def settings_class(self) -> Optional[Type["BaseSettings"]]:
+        """Settings class for the Skypilot orchestrator.
+
+        Returns:
+            The settings class.
+        """
+        return SkypilotKubernetesOrchestratorSettings
+
+    def prepare_environment_variable(self, set: bool = True) -> None:
+        """Set up Environment variables that are required for the orchestrator.
+
+        Args:
+            set: Whether to set the environment variables or not.
+        """
+        pass

zenml/integrations/skypilot_lambda/__init__.py

@@ -0,0 +1,50 @@
+#  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.
+"""Initialization of the Skypilot Lambda integration for ZenML.
+
+The Skypilot integration sub-module powers an alternative to the local
+orchestrator for a remote orchestration of ZenML pipelines on VMs.
+"""
+from typing import List, Type
+
+from zenml.integrations.constants import (
+    SKYPILOT_LAMBDA,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+SKYPILOT_LAMBDA_ORCHESTRATOR_FLAVOR = "vm_lambda"
+
+
+class SkypilotLambdaIntegration(Integration):
+    """Definition of Skypilot Lambda Integration for ZenML."""
+
+    NAME = SKYPILOT_LAMBDA
+    REQUIREMENTS = ["skypilot[lambda]~=0.6.0"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Skypilot Lambda integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.skypilot_lambda.flavors import (
+            SkypilotLambdaOrchestratorFlavor,
+        )
+
+        return [SkypilotLambdaOrchestratorFlavor]
+
+
+SkypilotLambdaIntegration.check_installation()

zenml/integrations/{kserve/secret_schemas → skypilot_lambda/flavors}/__init__.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
+#  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.
@@ -11,19 +11,16 @@
 #  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.
-"""Initialization of Kserve Secret Schemas.
+"""Skypilot integration flavor for Skypilot Lambda orchestrator."""
 
-These are secret schemas that can be used to authenticate Kserve to the
-Artifact Store used to store served ML models.
-"""
-from zenml.integrations.kserve.secret_schemas.secret_schemas import (
-    KServeAzureSecretSchema,
-    KServeGSSecretSchema,
-    KServeS3SecretSchema,
+from zenml.integrations.skypilot_lambda.flavors.skypilot_orchestrator_lambda_vm_flavor import (
+    SkypilotLambdaOrchestratorConfig,
+    SkypilotLambdaOrchestratorFlavor,
+    SkypilotLambdaOrchestratorSettings,
 )
 
 __all__ = [
-    "KServeAzureSecretSchema",
-    "KServeGSSecretSchema",
-    "KServeS3SecretSchema",
+    "SkypilotLambdaOrchestratorConfig",
+    "SkypilotLambdaOrchestratorFlavor",
+    "SkypilotLambdaOrchestratorSettings",
 ]

zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

@@ -0,0 +1,130 @@
+#  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.
+"""Skypilot orchestrator Lambda flavor."""
+
+from typing import TYPE_CHECKING, Any, Optional, Type
+
+from zenml.integrations.skypilot.flavors.skypilot_orchestrator_base_vm_config import (
+    SkypilotBaseOrchestratorConfig,
+    SkypilotBaseOrchestratorSettings,
+)
+from zenml.integrations.skypilot_lambda import (
+    SKYPILOT_LAMBDA_ORCHESTRATOR_FLAVOR,
+)
+from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.skypilot_lambda.orchestrators import (
+        SkypilotLambdaOrchestrator,
+    )
+
+
+logger = get_logger(__name__)
+
+
+class SkypilotLambdaOrchestratorSettings(SkypilotBaseOrchestratorSettings):
+    """Skypilot orchestrator settings."""
+
+    _UNSUPPORTED_FEATURES = {
+        "use_spot": "Spot instances not supported for Lambda orchestrator.",
+        "job_recovery": "Job recovery not supported for Lambda orchestrator.",
+        "image_id": "Custom image IDs not supported for Lambda orchestrator.",
+        # Add other unsupported features as needed
+    }
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Set attribute.
+
+        Args:
+            name: Name of the attribute.
+            value: Value of the attribute.
+
+        Raises:
+            AttributeError: If the attribute is not supported.
+        """
+        if name in self._UNSUPPORTED_FEATURES:
+            raise AttributeError(f"{name} is not supported on Lambda.")
+        super().__setattr__(name, value)
+
+
+class SkypilotLambdaOrchestratorConfig(
+    SkypilotBaseOrchestratorConfig, SkypilotLambdaOrchestratorSettings
+):
+    """Skypilot orchestrator config."""
+
+    api_key: Optional[str] = SecretField(default=None)
+
+
+class SkypilotLambdaOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Flavor for the Skypilot Lambda orchestrator."""
+
+    @property
+    def name(self) -> str:
+        """Name of the orchestrator flavor.
+
+        Returns:
+            Name of the orchestrator flavor.
+        """
+        return SKYPILOT_LAMBDA_ORCHESTRATOR_FLAVOR
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/lambda.png"
+
+    @property
+    def config_class(self) -> Type[BaseOrchestratorConfig]:
+        """Config class for the base orchestrator flavor.
+
+        Returns:
+            The config class.
+        """
+        return SkypilotLambdaOrchestratorConfig
+
+    @property
+    def implementation_class(self) -> Type["SkypilotLambdaOrchestrator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            Implementation class for this flavor.
+        """
+        from zenml.integrations.skypilot_lambda.orchestrators import (
+            SkypilotLambdaOrchestrator,
+        )
+
+        return SkypilotLambdaOrchestrator

zenml/{lineage_graph/node → integrations/skypilot_lambda/orchestrators}/__init__.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
+#  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.
@@ -11,22 +11,15 @@
 #  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.
-"""Initialization of lineage nodes."""
+"""Initialization of the Skypilot Lambda ZenML orchestrator."""
 
-from zenml.lineage_graph.node.artifact_node import (  # noqa
-    ArtifactNode,
-    ArtifactNodeDetails,
+from zenml.integrations.skypilot.orchestrators.skypilot_base_vm_orchestrator import (  # noqa
+    SkypilotBaseOrchestrator,
 )
-from zenml.lineage_graph.node.base_node import BaseNode  # noqa
-from zenml.lineage_graph.node.step_node import (  # noqa
-    StepNode,
-    StepNodeDetails,
+from zenml.integrations.skypilot_lambda.orchestrators.skypilot_lambda_vm_orchestrator import (  # noqa
+    SkypilotLambdaOrchestrator,
 )
-
 __all__ = [
-    "BaseNode",
-    "ArtifactNode",
-    "StepNode",
-    "StepNodeDetails",
-    "ArtifactNodeDetails",
+    "SkypilotBaseOrchestrator",
+    "SkypilotLambdaOrchestrator",
 ]

zenml/integrations/skypilot_lambda/orchestrators/skypilot_lambda_vm_orchestrator.py

@@ -0,0 +1,92 @@
+#  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.
+"""Implementation of the a Skypilot based Lambda VM orchestrator."""
+
+import os
+from typing import TYPE_CHECKING, Optional, Type, cast
+
+import sky
+
+from zenml.integrations.skypilot.orchestrators.skypilot_base_vm_orchestrator import (
+    SkypilotBaseOrchestrator,
+)
+from zenml.integrations.skypilot_lambda.flavors.skypilot_orchestrator_lambda_vm_flavor import (
+    SkypilotLambdaOrchestratorConfig,
+    SkypilotLambdaOrchestratorSettings,
+)
+from zenml.io import fileio
+from zenml.logger import get_logger
+
+if TYPE_CHECKING:
+    from zenml.config.base_settings import BaseSettings
+
+logger = get_logger(__name__)
+
+
+class SkypilotLambdaOrchestrator(SkypilotBaseOrchestrator):
+    """Orchestrator responsible for running pipelines remotely in a VM on Lambda.
+
+    This orchestrator does not support running on a schedule.
+    """
+
+    DEFAULT_INSTANCE_TYPE: str = "gpu_1x_a10"
+
+    @property
+    def cloud(self) -> sky.clouds.Cloud:
+        """The type of sky cloud to use.
+
+        Returns:
+            A `sky.clouds.Cloud` instance.
+        """
+        return sky.clouds.Lambda()
+
+    @property
+    def config(self) -> SkypilotLambdaOrchestratorConfig:
+        """Returns the `SkypilotLambdaOrchestratorConfig` config.
+
+        Returns:
+            The configuration.
+        """
+        return cast(SkypilotLambdaOrchestratorConfig, self._config)
+
+    @property
+    def settings_class(self) -> Optional[Type["BaseSettings"]]:
+        """Settings class for the Skypilot orchestrator.
+
+        Returns:
+            The settings class.
+        """
+        return SkypilotLambdaOrchestratorSettings
+
+    def prepare_environment_variable(self, set: bool = True) -> None:
+        """Set up Environment variables that are required for the orchestrator.
+
+        Args:
+            set: Whether to set the environment variables or not.
+        """
+        pass
+
+    def setup_credentials(self) -> None:
+        """Set up credentials for the orchestrator."""
+        # Define the directory and file paths
+        directory = os.path.expanduser("~/.lambda_cloud")
+        file_path = os.path.join(directory, "lambda_keys")
+
+        # Check if the directory exists, and create it if it doesn't
+        if not os.path.exists(directory):
+            os.makedirs(directory)
+
+        # Write the API key to the file, creating or overwriting it
+        with fileio.open(file_path, "w") as file:
+            file.write(f"api_key = {self.config.api_key}")

zenml/integrations/slack/__init__.py

@@ -31,6 +31,7 @@ class SlackIntegration(Integration):
 
     NAME = SLACK
     REQUIREMENTS = ["slack-sdk>=3.16.1", "aiohttp>=3.8.1"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["aiohttp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/slack/alerters/slack_alerter.py

@@ -14,16 +14,18 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, List, Optional, Type, cast
 
 from pydantic import BaseModel
 from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
 from slack_sdk.rtm import RTMClient
 
+from zenml import get_step_context
 from zenml.alerter.base_alerter import BaseAlerter, BaseAlerterStepParameters
 from zenml.integrations.slack.flavors.slack_alerter_flavor import (
     SlackAlerterConfig,
+    SlackAlerterSettings,
 )
 from zenml.logger import get_logger
 
@@ -72,6 +74,15 @@ class SlackAlerter(BaseAlerter):
         """
         return cast(SlackAlerterConfig, self._config)
 
+    @property
+    def settings_class(self) -> Type[SlackAlerterSettings]:
+        """Settings class for the Slack alerter.
+
+        Returns:
+            The settings class.
+        """
+        return SlackAlerterSettings
+
     def _get_channel_id(
         self, params: Optional[BaseAlerterStepParameters] = None
     ) -> str:
@@ -99,10 +110,19 @@ class SlackAlerter(BaseAlerter):
             and params.slack_channel_id is not None
         ):
             return params.slack_channel_id
+
+        settings = cast(
+            SlackAlerterSettings,
+            self.get_settings(get_step_context().step_run),
+        )
+        if settings.slack_channel_id is not None:
+            return settings.slack_channel_id
+
         if self.config.default_slack_channel_id is not None:
             return self.config.default_slack_channel_id
+
         raise ValueError(
-            "Neither the `SlackAlerterConfig.slack_channel_id` in the runtime "
+            "Neither the `slack_channel_id` in the runtime "
             "configuration, nor the `default_slack_channel_id` in the alerter "
             "stack component is specified. Please specify at least one."
         )

zenml/integrations/slack/flavors/__init__.py

@@ -16,9 +16,11 @@
 from zenml.integrations.slack.flavors.slack_alerter_flavor import (
     SlackAlerterConfig,
     SlackAlerterFlavor,
+    SlackAlerterSettings,
 )
 
 __all__ = [
     "SlackAlerterFlavor",
     "SlackAlerterConfig",
+    "SlackAlerterSettings",
 ]

zenml/integrations/slack/flavors/slack_alerter_flavor.py

@@ -16,6 +16,7 @@
 from typing import TYPE_CHECKING, Optional, Type
 
 from zenml.alerter.base_alerter import BaseAlerterConfig, BaseAlerterFlavor
+from zenml.config.base_settings import BaseSettings
 from zenml.integrations.slack import SLACK_ALERTER_FLAVOR
 from zenml.logger import get_logger
 from zenml.utils.secret_utils import SecretField
@@ -26,18 +27,26 @@ if TYPE_CHECKING:
     from zenml.integrations.slack.alerters import SlackAlerter
 
 
+class SlackAlerterSettings(BaseSettings):
+    """Settings for the Slack alerter.
+
+    Attributes:
+        slack_channel_id: The ID of the Slack channel to use for communication.
+    """
+
+    slack_channel_id: Optional[str] = None
+
+
 class SlackAlerterConfig(BaseAlerterConfig):
     """Slack alerter config.
 
     Attributes:
         slack_token: The Slack token tied to the Slack account to be used.
-        default_slack_channel_id: The ID of the Slack channel to use for
-            communication if no channel ID is provided in the step config.
-
+        default_slack_channel_id: The ID of the default Slack channel to use for communication.
     """
 
     slack_token: str = SecretField()
-    default_slack_channel_id: Optional[str] = None  # TODO: Potential setting
+    default_slack_channel_id: Optional[str] = None
 
     @property
     def is_valid(self) -> bool:

zenml/integrations/slack/steps/slack_alerter_ask_step.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Step that allows you to send messages to Slack and wait for a response."""
+
 from typing import Optional
 
 from zenml import get_step_context, step

zenml/integrations/slack/steps/slack_alerter_post_step.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Step that allows you to post messages to Slack."""
+
 from typing import Optional
 
 from zenml import get_step_context, step

zenml/integrations/spark/flavors/spark_step_operator_flavor.py

@@ -13,10 +13,7 @@
 #  permissions and limitations under the License.
 """Spark step operator flavor."""
 
-import json
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
-
-from pydantic import validator
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
 from zenml.config.base_settings import BaseSettings
 from zenml.step_operators.base_step_operator import (
@@ -45,42 +42,8 @@ class SparkStepOperatorSettings(BaseSettings):
     deploy_mode: str = "cluster"
     submit_kwargs: Optional[Dict[str, Any]] = None
 
-    @validator("submit_kwargs", pre=True)
-    def _convert_json_string(
-        cls, value: Union[None, str, Dict[str, Any]]
-    ) -> Optional[Dict[str, Any]]:
-        """Converts potential JSON strings passed via the CLI to dictionaries.
-
-        Args:
-            value: The value to convert.
-
-        Returns:
-            The converted value.
-
-        Raises:
-            TypeError: If the value is not a `str`, `Dict` or `None`.
-            ValueError: If the value is an invalid json string or a json string
-                that does not decode into a dictionary.
-        """
-        if isinstance(value, str):
-            try:
-                dict_ = json.loads(value)
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Invalid json string '{value}'") from e
-
-            if not isinstance(dict_, Dict):
-                raise ValueError(
-                    f"Json string '{value}' did not decode into a dictionary."
-                )
-
-            return dict_
-        elif isinstance(value, Dict) or value is None:
-            return value
-        else:
-            raise TypeError(f"{value} is not a json string or a dictionary.")
-
 
-class SparkStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SparkStepOperatorConfig(
     BaseStepOperatorConfig, SparkStepOperatorSettings
 ):
     """Spark step operator config.

zenml/integrations/spark/step_operators/kubernetes_step_operator.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Implementation of the Kubernetes Spark Step Operator."""
+
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
 
 from pyspark.conf import SparkConf

zenml/integrations/spark/step_operators/spark_entrypoint_configuration.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Spark step operator entrypoint configuration."""
+
 from zenml.step_operators.step_operator_entrypoint_configuration import (
     StepOperatorEntrypointConfiguration,
 )

zenml/integrations/spark/step_operators/spark_step_operator.py

@@ -303,6 +303,8 @@ class SparkStepOperator(BaseStepOperator):
         # Add any additional configuration given by the user.
         self._additional_configuration(spark_config=conf, settings=settings)
 
+        info.force_write_logs()
+
         # Generate a spark-submit command given the configuration
         self._launch_spark_job(
             spark_config=conf,

zenml/integrations/tekton/__init__.py

@@ -31,7 +31,8 @@ class TektonIntegration(Integration):
     """Definition of Tekton Integration for ZenML."""
 
     NAME = TEKTON
-    REQUIREMENTS = ["kfp-tekton==1.7.1"]
+    REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kfp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]: