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

zenml/utils/string_utils.py

@@ -15,13 +15,17 @@
 
 import base64
 import datetime
+import functools
 import random
 import string
+from typing import Any, Callable, Dict, TypeVar, cast
 
 from pydantic import BaseModel
 
 from zenml.constants import BANNED_NAME_CHARACTERS
 
+V = TypeVar("V", bound=Any)
+
 
 def get_human_readable_time(seconds: float) -> str:
     """Convert seconds into a human-readable string.
@@ -167,3 +171,49 @@ def format_name_template(
         datetime.datetime.now(datetime.timezone.utc).strftime("%H_%M_%S_%f"),
     )
     return name_template.format(**kwargs)
+
+
+def substitute_string(value: V, substitution_func: Callable[[str], str]) -> V:
+    """Recursively substitute strings in objects.
+
+    Args:
+        value: An object in which the strings should be recursively substituted.
+            This can be a pydantic model, dict, set, list, tuple or any
+            primitive type.
+        substitution_func: The function that does the actual string
+            substitution.
+
+    Returns:
+        The object with the substitution function applied to all string values.
+    """
+    substitute_ = functools.partial(
+        substitute_string, substitution_func=substitution_func
+    )
+
+    if isinstance(value, BaseModel):
+        model_values = {}
+
+        for k, v in value.__iter__():
+            new_value = substitute_(v)
+
+            if k not in value.model_fields_set and new_value == getattr(
+                value, k
+            ):
+                # This is a default value on the model and was not set
+                # explicitly. In this case, we don't include it in the model
+                # values to keep the `exclude_unset` behavior the same
+                continue
+
+            model_values[k] = new_value
+
+        return cast(V, type(value).model_validate(model_values))
+    elif isinstance(value, Dict):
+        return cast(
+            V, {substitute_(k): substitute_(v) for k, v in value.items()}
+        )
+    elif isinstance(value, (list, set, tuple)):
+        return cast(V, type(value)(substitute_(v) for v in value))
+    elif isinstance(value, str):
+        return cast(V, substitution_func(value))
+
+    return value

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/zen_stores/sql_zen_store.py

@@ -973,6 +973,7 @@ class SqlZenStore(BaseZenStore):
             ValueError: if the filtered page number is out of bounds.
             RuntimeError: if the schema does not have a `to_model` method.
         """
+        query = query.distinct()
         query = filter_model.apply_filter(query=query, table=table)
         query = query.distinct()
 

{zenml_nightly-0.66.0.dev20240923.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.dev20240923
+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)