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

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/routers/runs_endpoints.py

@@ -22,13 +22,15 @@ from zenml.constants import (
     API,
     GRAPH,
     PIPELINE_CONFIGURATION,
+    REFRESH,
     RUNS,
     STATUS,
     STEPS,
     VERSION_1,
 )
-from zenml.enums import ExecutionStatus
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.lineage_graph.lineage_graph import LineageGraph
+from zenml.logger import get_logger
 from zenml.models import (
     Page,
     PipelineRunFilter,
@@ -45,7 +47,8 @@ from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_list_entities,
     verify_permissions_and_update_entity,
 )
-from zenml.zen_server.rbac.models import ResourceType
+from zenml.zen_server.rbac.models import Action, ResourceType
+from zenml.zen_server.rbac.utils import verify_permission_for_model
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
@@ -59,6 +62,9 @@ router = APIRouter(
 )
 
 
+logger = get_logger(__name__)
+
+
 @router.get(
     "",
     response_model=Page[PipelineRunResponse],
@@ -99,6 +105,7 @@ def list_runs(
 def get_run(
     run_id: UUID,
     hydrate: bool = True,
+    refresh_status: bool = False,
     _: AuthContext = Security(authorize),
 ) -> PipelineRunResponse:
     """Get a specific pipeline run using its ID.
@@ -107,13 +114,47 @@ def get_run(
         run_id: ID of the pipeline run to get.
         hydrate: Flag deciding whether to hydrate the output model(s)
             by including metadata fields in the response.
+        refresh_status: Flag deciding whether we should try to refresh
+            the status of the pipeline run using its orchestrator.
 
     Returns:
         The pipeline run.
+
+    Raises:
+        RuntimeError: If the stack or the orchestrator of the run is deleted.
     """
-    return verify_permissions_and_get_entity(
+    run = verify_permissions_and_get_entity(
         id=run_id, get_method=zen_store().get_run, hydrate=hydrate
     )
+    if refresh_status:
+        try:
+            # Check the stack and its orchestrator
+            if run.stack is not None:
+                orchestrators = run.stack.components.get(
+                    StackComponentType.ORCHESTRATOR, []
+                )
+                if orchestrators:
+                    verify_permission_for_model(
+                        model=orchestrators[0], action=Action.READ
+                    )
+                else:
+                    raise RuntimeError(
+                        f"The orchestrator, the run '{run.id}' was executed "
+                        "with, is deleted."
+                    )
+            else:
+                raise RuntimeError(
+                    f"The stack, the run '{run.id}' was executed on, is deleted."
+                )
+
+            run = run.refresh_run_status()
+
+        except Exception as e:
+            logger.warning(
+                "An error occurred while refreshing the status of the "
+                f"pipeline run: {e}"
+            )
+    return run
 
 
 @router.put(
@@ -267,3 +308,48 @@ def get_run_status(
         id=run_id, get_method=zen_store().get_run, hydrate=False
     )
     return run.status
+
+
+@router.get(
+    "/{run_id}" + REFRESH,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def refresh_run_status(
+    run_id: UUID,
+    _: AuthContext = Security(authorize),
+) -> None:
+    """Refreshes the status of a specific pipeline run.
+
+    Args:
+        run_id: ID of the pipeline run to refresh.
+
+    Raises:
+        RuntimeError: If the stack or the orchestrator of the run is deleted.
+    """
+    # Verify access to the run
+    run = verify_permissions_and_get_entity(
+        id=run_id,
+        get_method=zen_store().get_run,
+        hydrate=True,
+    )
+
+    # Check the stack and its orchestrator
+    if run.stack is not None:
+        orchestrators = run.stack.components.get(
+            StackComponentType.ORCHESTRATOR, []
+        )
+        if orchestrators:
+            verify_permission_for_model(
+                model=orchestrators[0], action=Action.READ
+            )
+        else:
+            raise RuntimeError(
+                f"The orchestrator, the run '{run.id}' was executed with, is "
+                "deleted."
+            )
+    else:
+        raise RuntimeError(
+            f"The stack, the run '{run.id}' was executed on, is deleted."
+        )
+    run.refresh_run_status()

{zenml_nightly-0.66.0.dev20240924.dist-info → zenml_nightly-0.66.0.dev20240925.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.66.0.dev20240924
+Version: 0.66.0.dev20240925
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
@@ -24,6 +24,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Classifier: Typing :: Typed
 Provides-Extra: adlfs
+Provides-Extra: azureml
 Provides-Extra: connectors-aws
 Provides-Extra: connectors-azure
 Provides-Extra: connectors-gcp
@@ -32,6 +33,7 @@ Provides-Extra: dev
 Provides-Extra: gcsfs
 Provides-Extra: mlstacks
 Provides-Extra: s3fs
+Provides-Extra: sagemaker
 Provides-Extra: secrets-aws
 Provides-Extra: secrets-azure
 Provides-Extra: secrets-gcp
@@ -39,10 +41,12 @@ Provides-Extra: secrets-hashicorp
 Provides-Extra: server
 Provides-Extra: templates
 Provides-Extra: terraform
+Provides-Extra: vertex
 Requires-Dist: Jinja2 ; extra == "server"
 Requires-Dist: adlfs (>=2021.10.0) ; extra == "adlfs"
 Requires-Dist: alembic (>=1.8.1,<1.9.0)
 Requires-Dist: aws-profile-manager (>=0.5.0) ; extra == "connectors-aws"
+Requires-Dist: azure-ai-ml (==1.18.0) ; extra == "azureml"
 Requires-Dist: azure-identity (>=1.4.0) ; extra == "secrets-azure" or extra == "connectors-azure"
 Requires-Dist: azure-keyvault-secrets (>=4.0.0) ; extra == "secrets-azure"
 Requires-Dist: azure-mgmt-containerregistry (>=10.0.0) ; extra == "connectors-azure"
@@ -63,6 +67,7 @@ Requires-Dist: docker (>=7.1.0,<7.2.0)
 Requires-Dist: fastapi (>=0.100,<=0.110) ; extra == "server"
 Requires-Dist: gcsfs (>=2022.11.0) ; extra == "gcsfs"
 Requires-Dist: gitpython (>=3.1.18,<4.0.0)
+Requires-Dist: google-cloud-aiplatform (>=1.34.0) ; extra == "vertex"
 Requires-Dist: google-cloud-artifact-registry (>=1.11.3) ; extra == "connectors-gcp"
 Requires-Dist: google-cloud-container (>=2.21.0) ; extra == "connectors-gcp"
 Requires-Dist: google-cloud-secret-manager (>=2.12.5) ; extra == "secrets-gcp"
@@ -72,6 +77,7 @@ Requires-Dist: hypothesis (>=6.43.1,<7.0.0) ; extra == "dev"
 Requires-Dist: importlib_metadata (<=7.0.0) ; python_version < "3.10"
 Requires-Dist: ipinfo (>=4.4.3) ; extra == "server"
 Requires-Dist: jinja2-time (>=0.2.0,<0.3.0) ; extra == "templates"
+Requires-Dist: kfp (>=2.6.0) ; extra == "vertex"
 Requires-Dist: kubernetes (>=18.20.0) ; extra == "connectors-kubernetes" or extra == "connectors-aws" or extra == "connectors-gcp" or extra == "connectors-azure"
 Requires-Dist: maison (<2.0) ; extra == "dev"
 Requires-Dist: mike (>=1.1.2,<2.0.0) ; extra == "dev"
@@ -105,6 +111,7 @@ Requires-Dist: pyyaml-include (<2.0) ; extra == "templates"
 Requires-Dist: rich[jupyter] (>=12.0.0)
 Requires-Dist: ruff (>=0.1.7) ; extra == "templates" or extra == "dev"
 Requires-Dist: s3fs (>=2022.11.0) ; extra == "s3fs"
+Requires-Dist: sagemaker (>=2.117.0) ; extra == "sagemaker"
 Requires-Dist: secure (>=0.3.0,<0.4.0) ; extra == "server"
 Requires-Dist: setuptools
 Requires-Dist: sqlalchemy (>=2.0.0,<3.0.0)