zenml-nightly 0.66.0.dev20240924__py3-none-any.whl → 0.66.0.dev20240926__py3-none-any.whl

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/models/v2/core/pipeline_run.py

@@ -21,6 +21,7 @@ from typing import (
     List,
     Optional,
     Union,
+    cast,
 )
 from uuid import UUID
 
@@ -136,7 +137,8 @@ class PipelineRunUpdate(BaseModel):
         "configured by this pipeline run explicitly.",
         default=None,
     )
-    # TODO: we should maybe have a different update model here, the upper three attributes should only be for internal use
+    # TODO: we should maybe have a different update model here, the upper
+    #  three attributes should only be for internal use
     add_tags: Optional[List[str]] = Field(
         default=None, title="New tags to add to the pipeline run."
     )
@@ -235,6 +237,7 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         description="Template used for the pipeline run.",
     )
     is_templatable: bool = Field(
+        default=False,
         description="Whether a template can be created from this run.",
     )
 
@@ -307,6 +310,64 @@ class PipelineRunResponse(
 
         return get_artifacts_versions_of_pipeline_run(self, only_produced=True)
 
+    def refresh_run_status(self) -> "PipelineRunResponse":
+        """Method to refresh the status of a run if it is initializing/running.
+
+        Returns:
+            The updated pipeline.
+
+        Raises:
+            ValueError: If the stack of the run response is None.
+        """
+        if self.status in [
+            ExecutionStatus.INITIALIZING,
+            ExecutionStatus.RUNNING,
+        ]:
+            # Check if the stack still accessible
+            if self.stack is None:
+                raise ValueError(
+                    "The stack that this pipeline run response was executed on"
+                    "has been deleted."
+                )
+
+            # Create the orchestrator instance
+            from zenml.enums import StackComponentType
+            from zenml.orchestrators.base_orchestrator import BaseOrchestrator
+            from zenml.stack.stack_component import StackComponent
+
+            # Check if the stack still accessible
+            orchestrator_list = self.stack.components.get(
+                StackComponentType.ORCHESTRATOR, []
+            )
+            if len(orchestrator_list) == 0:
+                raise ValueError(
+                    "The orchestrator that this pipeline run response was "
+                    "executed with has been deleted."
+                )
+
+            orchestrator = cast(
+                BaseOrchestrator,
+                StackComponent.from_model(
+                    component_model=orchestrator_list[0]
+                ),
+            )
+
+            # Fetch the status
+            status = orchestrator.fetch_status(run=self)
+
+            # If it is different from the current status, update it
+            if status != self.status:
+                from zenml.client import Client
+                from zenml.models import PipelineRunUpdate
+
+                client = Client()
+                return client.zen_store.update_run(
+                    run_id=self.id,
+                    run_update=PipelineRunUpdate(status=status),
+                )
+
+        return self
+
     # Body and metadata properties
     @property
     def status(self) -> ExecutionStatus:

zenml/new/pipelines/run_utils.py

@@ -138,7 +138,10 @@ def deploy_pipeline(
     previous_value = constants.SHOULD_PREVENT_PIPELINE_EXECUTION
     constants.SHOULD_PREVENT_PIPELINE_EXECUTION = True
     try:
-        stack.deploy_pipeline(deployment=deployment)
+        stack.deploy_pipeline(
+            deployment=deployment,
+            placeholder_run=placeholder_run,
+        )
     except Exception as e:
         if (
             placeholder_run

zenml/orchestrators/base_orchestrator.py

@@ -14,12 +14,14 @@
 """Base orchestrator class."""
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
 
 from pydantic import model_validator
 
-from zenml.enums import StackComponentType
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.orchestrators.publish_utils import publish_pipeline_run_metadata
 from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
@@ -27,7 +29,7 @@ from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config.step_configurations import Step
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
 
 logger = get_logger(__name__)
 
@@ -124,7 +126,7 @@ class BaseOrchestrator(StackComponent, ABC):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
-    ) -> Any:
+    ) -> Optional[Iterator[Dict[str, MetadataType]]]:
         """The method needs to be implemented by the respective orchestrator.
 
         Depending on the type of orchestrator you'll have to perform slightly
@@ -169,29 +171,41 @@ class BaseOrchestrator(StackComponent, ABC):
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Runs a pipeline on a stack.
 
         Args:
             deployment: The pipeline deployment.
             stack: The stack on which to run the pipeline.
-
-        Returns:
-            Orchestrator-specific return value.
+            placeholder_run: An optional placeholder run for the deployment.
+                This will be deleted in case the pipeline deployment failed.
         """
         self._prepare_run(deployment=deployment)
 
         environment = get_config_environment_vars(deployment=deployment)
 
         try:
-            result = self.prepare_or_run_pipeline(
-                deployment=deployment, stack=stack, environment=environment
-            )
+            if metadata_iterator := self.prepare_or_run_pipeline(
+                deployment=deployment,
+                stack=stack,
+                environment=environment,
+            ):
+                for metadata_dict in metadata_iterator:
+                    try:
+                        if placeholder_run:
+                            publish_pipeline_run_metadata(
+                                pipeline_run_id=placeholder_run.id,
+                                pipeline_run_metadata={self.id: metadata_dict},
+                            )
+                    except Exception as e:
+                        logger.debug(
+                            "Something went went wrong trying to publish the"
+                            f"run metadata: {e}"
+                        )
         finally:
             self._cleanup_run()
 
-        return result
-
     def run_step(self, step: "Step") -> None:
         """Runs the given step.
 
@@ -239,6 +253,21 @@ class BaseOrchestrator(StackComponent, ABC):
         """Cleans up the active run."""
         self._active_deployment = None
 
+    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+        """Refreshes the status of a specific pipeline run.
+
+        Args:
+            run: A pipeline run response to fetch its status.
+
+        Raises:
+            NotImplementedError: If any orchestrator inheriting from the base
+                class does not implement this logic.
+        """
+        raise NotImplementedError(
+            "The fetch status functionality is not implemented for the "
+            f"'{self.__class__.__name__}' orchestrator."
+        )
+
 
 class BaseOrchestratorFlavor(Flavor):
     """Base orchestrator flavor class."""

zenml/stack/stack.py

@@ -62,7 +62,11 @@ if TYPE_CHECKING:
     from zenml.image_builders import BaseImageBuilder
     from zenml.model_deployers import BaseModelDeployer
     from zenml.model_registries import BaseModelRegistry
-    from zenml.models import PipelineDeploymentBase, PipelineDeploymentResponse
+    from zenml.models import (
+        PipelineDeploymentBase,
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+    )
     from zenml.orchestrators import BaseOrchestrator
     from zenml.stack import StackComponent
     from zenml.step_operators import BaseStepOperator
@@ -826,16 +830,21 @@ class Stack:
     def deploy_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Deploys a pipeline on this stack.
 
         Args:
             deployment: The pipeline deployment.
+            placeholder_run: An optional placeholder run for the deployment.
+                This will be deleted in case the pipeline deployment failed.
 
         Returns:
             The return value of the call to `orchestrator.run_pipeline(...)`.
         """
-        return self.orchestrator.run(deployment=deployment, stack=self)
+        return self.orchestrator.run(
+            deployment=deployment, stack=self, placeholder_run=placeholder_run
+        )
 
     def _get_active_components_for_step(
         self, step_config: "StepConfiguration"

zenml/zen_server/cloud_utils.py

@@ -1,6 +1,7 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
 import os
+from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Optional
 
 import requests
@@ -19,11 +20,9 @@ class ZenMLCloudConfiguration(BaseModel):
     """ZenML Pro RBAC configuration."""
 
     api_url: str
-
     oauth2_client_id: str
     oauth2_client_secret: str
     oauth2_audience: str
-    auth0_domain: str
 
     @field_validator("api_url")
     @classmethod
@@ -68,6 +67,8 @@ class ZenMLCloudConnection:
         """Initialize the RBAC component."""
         self._config = ZenMLCloudConfiguration.from_environment()
         self._session: Optional[requests.Session] = None
+        self._token: Optional[str] = None
+        self._token_expires_at: Optional[datetime] = None
 
     def get(
         self, endpoint: str, params: Optional[Dict[str, Any]]
@@ -91,7 +92,8 @@ class ZenMLCloudConnection:
 
         response = self.session.get(url=url, params=params, timeout=7)
         if response.status_code == 401:
-            # Refresh the auth token and try again
+            # If we get an Unauthorized error from the API serer, we refresh the
+            # auth token and try again
             self._clear_session()
             response = self.session.get(url=url, params=params, timeout=7)
 
@@ -186,6 +188,8 @@ class ZenMLCloudConnection:
     def _clear_session(self) -> None:
         """Clear the authentication session."""
         self._session = None
+        self._token = None
+        self._token_expires_at = None
 
     def _fetch_auth_token(self) -> str:
         """Fetch an auth token for the Cloud API from auth0.
@@ -196,8 +200,16 @@ class ZenMLCloudConnection:
         Returns:
             Auth token.
         """
+        if (
+            self._token is not None
+            and self._token_expires_at is not None
+            and datetime.now(timezone.utc) + timedelta(minutes=5)
+            < self._token_expires_at
+        ):
+            return self._token
+
         # Get an auth token from auth0
-        auth0_url = f"https://{self._config.auth0_domain}/oauth/token"
+        login_url = f"{self._config.api_url}/auth/login"
         headers = {"content-type": "application/x-www-form-urlencoded"}
         payload = {
             "client_id": self._config.oauth2_client_id,
@@ -207,18 +219,31 @@ class ZenMLCloudConnection:
         }
         try:
             response = requests.post(
-                auth0_url, headers=headers, data=payload, timeout=7
+                login_url, headers=headers, data=payload, timeout=7
             )
             response.raise_for_status()
         except Exception as e:
             raise RuntimeError(f"Error fetching auth token from auth0: {e}")
 
-        access_token = response.json().get("access_token", "")
+        json_response = response.json()
+        access_token = json_response.get("access_token", "")
+        expires_in = json_response.get("expires_in", 0)
 
-        if not access_token or not isinstance(access_token, str):
+        if (
+            not access_token
+            or not isinstance(access_token, str)
+            or not expires_in
+            or not isinstance(expires_in, int)
+        ):
             raise RuntimeError("Could not fetch auth token from auth0.")
 
-        return str(access_token)
+        self._token = access_token
+        self._token_expires_at = datetime.now(timezone.utc) + timedelta(
+            seconds=expires_in
+        )
+
+        assert self._token is not None
+        return self._token
 
 
 def cloud_connection() -> ZenMLCloudConnection:

zenml/zen_server/dashboard/assets/{404-iO8vpun1.js → 404-CMnKjD-L.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-DnFH_oo1.js";import{f as s,r as t}from"./index-B9wVwe7u.js";import{E as r}from"./EmptyState-Cs3DEmso.js";import{H as a}from"./help-CwN931fX.js";import{L as o}from"./@react-router-APVeuk-U.js";import"./@tanstack-QbMbTrh5.js";import"./@reactflow-B6kq9fJZ.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-DnFH_oo1.js";import{f as s,r as t}from"./index-D0bJjaey.js";import{E as r}from"./EmptyState-Cs3DEmso.js";import{H as a}from"./help-CwN931fX.js";import{L as o}from"./@react-router-APVeuk-U.js";import"./@tanstack-QbMbTrh5.js";import"./@reactflow-CEC2f0cl.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};