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

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.dev20240926.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.66.0.dev20240924
+Version: 0.66.0.dev20240926
 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)