zenml-nightly 0.66.0.dev20240919__py3-none-any.whl → 0.66.0.dev20240927__py3-none-any.whl

zenml/integrations/skypilot_kubernetes/__init__.py

@@ -0,0 +1,52 @@
+#  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 Kubernetes 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_KUBERNETES,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+SKYPILOT_KUBERNETES_ORCHESTRATOR_FLAVOR = "vm_kubernetes"
+
+
+class SkypilotKubernetesIntegration(Integration):
+    """Definition of Skypilot Kubernetes Integration for ZenML."""
+
+    NAME = SKYPILOT_KUBERNETES
+    # all 0.6.x versions of skypilot[kubernetes] are compatible
+    REQUIREMENTS = ["skypilot[kubernetes]~=0.6.1"]
+    APT_PACKAGES = ["openssh-client", "rsync"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Skypilot Kubernetes integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.skypilot_kubernetes.flavors import (
+            SkypilotKubernetesOrchestratorFlavor,
+        )
+
+        return [SkypilotKubernetesOrchestratorFlavor]
+
+
+SkypilotKubernetesIntegration.check_installation()

zenml/integrations/skypilot_kubernetes/flavors/__init__.py

@@ -0,0 +1,26 @@
+#  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 integration flavor for Skypilot Kubernetes orchestrator."""
+
+from zenml.integrations.skypilot_kubernetes.flavors.skypilot_orchestrator_kubernetes_vm_flavor import (
+    SkypilotKubernetesOrchestratorConfig,
+    SkypilotKubernetesOrchestratorFlavor,
+    SkypilotKubernetesOrchestratorSettings,
+)
+
+__all__ = [
+    "SkypilotKubernetesOrchestratorConfig",
+    "SkypilotKubernetesOrchestratorFlavor",
+    "SkypilotKubernetesOrchestratorSettings",
+]

zenml/integrations/skypilot_kubernetes/flavors/skypilot_orchestrator_kubernetes_vm_flavor.py

@@ -0,0 +1,125 @@
+#  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.
+"""Skypilot orchestrator Kubernetes flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
+from zenml.integrations.skypilot.flavors.skypilot_orchestrator_base_vm_config import (
+    SkypilotBaseOrchestratorConfig,
+    SkypilotBaseOrchestratorSettings,
+)
+from zenml.integrations.skypilot_kubernetes import (
+    SKYPILOT_KUBERNETES_ORCHESTRATOR_FLAVOR,
+)
+from zenml.logger import get_logger
+from zenml.models import ServiceConnectorRequirements
+from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
+
+if TYPE_CHECKING:
+    from zenml.integrations.skypilot_kubernetes.orchestrators import (
+        SkypilotKubernetesOrchestrator,
+    )
+
+
+logger = get_logger(__name__)
+
+
+class SkypilotKubernetesOrchestratorSettings(SkypilotBaseOrchestratorSettings):
+    """Skypilot orchestrator settings."""
+
+
+class SkypilotKubernetesOrchestratorConfig(
+    SkypilotBaseOrchestratorConfig, SkypilotKubernetesOrchestratorSettings
+):
+    """Skypilot orchestrator config."""
+
+
+class SkypilotKubernetesOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Flavor for the Skypilot Kubernetes orchestrator."""
+
+    @property
+    def name(self) -> str:
+        """Name of the orchestrator flavor.
+
+        Returns:
+            Name of the orchestrator flavor.
+        """
+        return SKYPILOT_KUBERNETES_ORCHESTRATOR_FLAVOR
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(
+            resource_type=KUBERNETES_CLUSTER_RESOURCE_TYPE,
+        )
+
+    @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/kubernetes-skypilot.png"
+
+    @property
+    def config_class(self) -> Type[BaseOrchestratorConfig]:
+        """Config class for the base orchestrator flavor.
+
+        Returns:
+            The config class.
+        """
+        return SkypilotKubernetesOrchestratorConfig
+
+    @property
+    def implementation_class(self) -> Type["SkypilotKubernetesOrchestrator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            Implementation class for this flavor.
+        """
+        from zenml.integrations.skypilot_kubernetes.orchestrators import (
+            SkypilotKubernetesOrchestrator,
+        )
+
+        return SkypilotKubernetesOrchestrator

zenml/integrations/skypilot_kubernetes/orchestrators/__init__.py

@@ -0,0 +1,25 @@
+#  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 Kubernetes ZenML orchestrator."""
+
+from zenml.integrations.skypilot.orchestrators.skypilot_base_vm_orchestrator import (  # noqa
+    SkypilotBaseOrchestrator,
+)
+from zenml.integrations.skypilot_kubernetes.orchestrators.skypilot_kubernetes_vm_orchestrator import (  # noqa
+    SkypilotKubernetesOrchestrator,
+)
+__all__ = [
+    "SkypilotBaseOrchestrator",
+    "SkypilotKubernetesOrchestrator",
+]

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/tensorboard/visualizers/tensorboard_visualizer.py

@@ -194,7 +194,7 @@ def get_step(pipeline_name: str, step_name: str) -> "StepRunResponse":
     Raises:
         RuntimeError: If the step is not found.
     """
-    runs = Client().list_pipeline_runs(pipeline_name=pipeline_name)
+    runs = Client().list_pipeline_runs(pipeline=pipeline_name)
     if runs.total == 0:
         raise RuntimeError(
             f"No pipeline runs for pipeline `{pipeline_name}` were found"