zenml-nightly 0.82.1.dev20250522__py3-none-any.whl → 0.82.1.dev20250524__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.82.1.dev20250522
+0.82.1.dev20250524

zenml/client.py

@@ -2024,7 +2024,9 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             type=component_type,
             flavor=flavor,
-            configuration=configuration,
+            configuration=validated_config.model_dump(
+                mode="json", exclude_unset=True
+            ),
             labels=labels,
         )
 
@@ -2112,7 +2114,9 @@ class Client(metaclass=ClientMetaClass):
             assert validated_config is not None
             warn_if_config_server_mismatch(validated_config)
 
-            update_model.configuration = existing_configuration
+            update_model.configuration = validated_config.model_dump(
+                mode="json", exclude_unset=True
+            )
 
         if labels is not None:
             existing_labels = component.labels or {}

zenml/config/docker_settings.py

@@ -386,6 +386,31 @@ class DockerSettings(BaseSettings):
 
         return self
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _warn_about_future_default_installer(
+        cls, data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Warns about the future change of default package installer from pip to uv.
+
+        Args:
+            data: The model data.
+
+        Returns:
+            The validated settings values.
+        """
+        if "python_package_installer" not in data:
+            logger.warning(
+                "In a future release, the default Python package installer "
+                "used by ZenML to build container images for your "
+                "containerized pipelines will change from 'pip' to 'uv'. "
+                "To maintain current behavior, you can explicitly set "
+                "`python_package_installer=PythonPackageInstaller.PIP` "
+                "in your DockerSettings."
+            )
+        return data
+
     model_config = ConfigDict(
         # public attributes are immutable
         frozen=True,

zenml/config/server_config.py

@@ -45,6 +45,7 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES,
     DEFAULT_ZENML_SERVER_NAME,
     DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW,
+    DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP,
@@ -252,6 +253,11 @@ class ServerConfiguration(BaseModel):
             used.
         file_download_size_limit: The maximum size of the file download in
             bytes. If not specified, the default value of 2GB will be used.
+        thread_pool_size: The size of the thread pool for handling requests. If
+            not specified, the default value of 40 will be used.
+        server_request_timeout: The timeout for server requests in seconds. If
+            not specified, the default value of 20 seconds will be used. This
+            value should be lower than the client's request timeout.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
@@ -348,6 +354,7 @@ class ServerConfiguration(BaseModel):
     auto_activate: bool = False
 
     thread_pool_size: int = DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE
+    server_request_timeout: int = DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT
 
     max_request_body_size_in_bytes: int = (
         DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES

zenml/constants.py

@@ -282,6 +282,7 @@ DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS = 3
 DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT = 60 * 5  # 5 minutes
 DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING = 5  # seconds
 DEFAULT_HTTP_TIMEOUT = 30
+DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT = 20  # seconds
 SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT = 120  # seconds
 ZENML_API_KEY_PREFIX = "ZENKEY_"
 DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours

zenml/enums.py

@@ -314,6 +314,7 @@ class EnvironmentType(StrEnum):
     LIGHTNING_AI_STUDIO = "lightning_ai_studio"
     GITHUB_CODESPACES = "github_codespaces"
     VSCODE_REMOTE_CONTAINER = "vscode_remote_container"
+    ZENML_CODESPACE = "zenml_codespace"
 
 
 class ModelStages(StrEnum):

zenml/environment.py

@@ -79,6 +79,8 @@ def get_environment() -> str:
         return EnvironmentType.GENERIC_CI
     elif Environment.in_github_codespaces():
         return EnvironmentType.GITHUB_CODESPACES
+    elif Environment.in_zenml_codespace():
+        return EnvironmentType.ZENML_CODESPACE
     elif Environment.in_vscode_remote_container():
         return EnvironmentType.VSCODE_REMOTE_CONTAINER
     elif Environment.in_lightning_ai_studio():
@@ -254,6 +256,16 @@ class Environment(metaclass=SingletonMetaClass):
             or "GITHUB_CODESPACES_PORT_FORWARDING_DOMAIN" in os.environ
         )
 
+    @staticmethod
+    def in_zenml_codespace() -> bool:
+        """If the current Python process is running in ZenML Codespaces.
+
+        Returns:
+            `True` if the current Python process is running in ZenML Codespaces,
+            `False` otherwise.
+        """
+        return os.environ.get("ZENML_ENVIRONMENT") == "codespace"
+
     @staticmethod
     def in_vscode_remote_container() -> bool:
         """If the current Python process is running in a VS Code Remote Container.

zenml/integrations/gcp/__init__.py

@@ -45,7 +45,7 @@ class GcpIntegration(Integration):
     NAME = GCP
     REQUIREMENTS = [
         "kfp>=2.6.0",
-        "gcsfs",
+        "gcsfs!=2025.5.0,!=2025.5.0.post1",
         "google-cloud-secret-manager",
         "google-cloud-container>=2.21.0",
         "google-cloud-artifact-registry>=1.11.3",

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -78,7 +78,7 @@ from zenml.service_connectors.service_connector import (
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.utils.secret_utils import PlainSerializedSecretStr
-from zenml.utils.time_utils import utc_now
+from zenml.utils.time_utils import to_utc_timezone, utc_now
 
 logger = get_logger(__name__)
 
@@ -1200,14 +1200,13 @@ class GCPServiceConnector(ServiceConnector):
         elif auth_method == GCPAuthenticationMethods.OAUTH2_TOKEN:
             assert isinstance(cfg, GCPOAuth2TokenConfig)
 
-            expires_at = self.expires_at
-            if expires_at:
-                # Remove the UTC timezone
-                expires_at = expires_at.replace(tzinfo=None)
-
             credentials = gcp_credentials.Credentials(
                 token=cfg.token.get_secret_value(),
-                expiry=expires_at,
+                # Currently GCP expects the expiry to be a timezone-naive
+                # UTC datetime.
+                expiry=to_utc_timezone(self.expires_at).replace(tzinfo=None)
+                if self.expires_at
+                else None,
                 scopes=scopes,
             )
 
@@ -1303,10 +1302,12 @@ class GCPServiceConnector(ServiceConnector):
                 )
 
         if credentials.expiry:
-            # Add the UTC timezone to the expiration time
-            expires_at = credentials.expiry.replace(
-                tzinfo=datetime.timezone.utc
-            )
+            expires_at = credentials.expiry
+
+        if expires_at:
+            # Add the UTC timezone to the expiration time, if it's not already
+            # set
+            expires_at = to_utc_timezone(expires_at)
 
         return credentials, expires_at
 

zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py

@@ -425,7 +425,7 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
         with tempfile.NamedTemporaryFile(mode="w", delete=True) as f:
             # Define bash line and command line
             bash_line = "#!/bin/bash\n"
-            command_line = f'cd {directory_name} && echo {ENV_ZENML_HYPERAI_RUN_ID}="{deployment_id}_$(date +\%s)" > .env && docker compose up -d'
+            command_line = rf'cd {directory_name} && echo {ENV_ZENML_HYPERAI_RUN_ID}="{deployment_id}_$(date +\%s)" > .env && docker compose up -d'
 
             # Write script to temporary file
             with f.file as f_:

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -270,18 +270,22 @@ def wait_pod(
         # Stream logs to `zenml.logger.info()`.
         # TODO: can we do this without parsing all logs every time?
         if stream_logs and pod_is_not_pending(resp):
-            response = core_api.read_namespaced_pod_log(
-                name=pod_name,
-                namespace=namespace,
-                _preload_content=False,
-            )
-            raw_data = response.data
-            decoded_log = raw_data.decode("utf-8", errors="replace")
-            logs = decoded_log.splitlines()
-            if len(logs) > logged_lines:
-                for line in logs[logged_lines:]:
-                    logger.info(line)
-                logged_lines = len(logs)
+            try:
+                response = core_api.read_namespaced_pod_log(
+                    name=pod_name,
+                    namespace=namespace,
+                    _preload_content=False,
+                )
+            except ApiException as e:
+                logger.error(f"Error reading pod logs: {e}. Retrying...")
+            else:
+                raw_data = response.data
+                decoded_log = raw_data.decode("utf-8", errors="replace")
+                logs = decoded_log.splitlines()
+                if len(logs) > logged_lines:
+                    for line in logs[logged_lines:]:
+                        logger.info(line)
+                    logged_lines = len(logs)
 
         # Raise an error if the pod failed.
         if pod_failed(resp):

zenml/materializers/built_in_materializer.py

@@ -352,7 +352,9 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 ):
                     type_ = find_type_by_str(type_str)
                     materializer_class = materializer_registry[type_]
-                    materializer = materializer_class(uri=path_)
+                    materializer = materializer_class(
+                        uri=path_, artifact_store=self.artifact_store
+                    )
                     element = materializer.load(type_)
                     outputs.append(element)
 
@@ -364,7 +366,9 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                     materializer_class = source_utils.load(
                         entry["materializer"]
                     )
-                    materializer = materializer_class(uri=path_)
+                    materializer = materializer_class(
+                        uri=path_, artifact_store=self.artifact_store
+                    )
                     element = materializer.load(type_)
                     outputs.append(element)
 
@@ -427,7 +431,9 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 self.artifact_store.mkdir(element_path)
                 type_ = type(element)
                 materializer_class = materializer_registry[type_]
-                materializer = materializer_class(uri=element_path)
+                materializer = materializer_class(
+                    uri=element_path, artifact_store=self.artifact_store
+                )
                 materializers.append(materializer)
                 metadata.append(
                     {

zenml/zen_server/cloud_utils.py

@@ -1,5 +1,8 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
+import logging
+import threading
+import time
 from datetime import datetime, timedelta
 from threading import RLock
 from typing import Any, Dict, Optional
@@ -12,9 +15,12 @@ from zenml.exceptions import (
     IllegalOperationError,
     SubscriptionUpgradeRequiredError,
 )
+from zenml.logger import get_logger
 from zenml.utils.time_utils import utc_now
 from zenml.zen_server.utils import get_zenml_headers, server_config
 
+logger = get_logger(__name__)
+
 _cloud_connection: Optional["ZenMLCloudConnection"] = None
 
 
@@ -56,16 +62,16 @@ class ZenMLCloudConnection:
         """
         url = self._config.api_url + endpoint
 
-        response = self.session.request(
-            method=method,
-            url=url,
-            params=params,
-            json=data,
-            timeout=self._config.http_timeout,
-        )
-        if response.status_code == 401:
-            # Refresh the auth token and try again
-            self._reset_login()
+        if logger.isEnabledFor(logging.DEBUG):
+            # Get the request ID from the current thread object
+            request_id = threading.current_thread().name
+            logger.debug(
+                f"[{request_id}] RBAC STATS - {method} {endpoint} started"
+            )
+            start_time = time.time()
+
+        status_code: Optional[int] = None
+        try:
             response = self.session.request(
                 method=method,
                 url=url,
@@ -73,18 +79,36 @@ class ZenMLCloudConnection:
                 json=data,
                 timeout=self._config.http_timeout,
             )
+            if response.status_code == 401:
+                # Refresh the auth token and try again
+                self._reset_login()
+                response = self.session.request(
+                    method=method,
+                    url=url,
+                    params=params,
+                    json=data,
+                    timeout=self._config.http_timeout,
+                )
 
-        try:
-            response.raise_for_status()
-        except requests.HTTPError as e:
-            if response.status_code == 402:
-                raise SubscriptionUpgradeRequiredError(response.json())
-            elif response.status_code == 403:
-                raise IllegalOperationError(response.json())
-            else:
-                raise RuntimeError(
-                    f"Failed while trying to contact the central zenml pro "
-                    f"service: {e}"
+            status_code = response.status_code
+            try:
+                response.raise_for_status()
+            except requests.HTTPError as e:
+                if response.status_code == 402:
+                    raise SubscriptionUpgradeRequiredError(response.json())
+                elif response.status_code == 403:
+                    raise IllegalOperationError(response.json())
+                else:
+                    raise RuntimeError(
+                        f"Failed while trying to contact the central zenml pro "
+                        f"service: {e}"
+                    )
+        finally:
+            if logger.isEnabledFor(logging.DEBUG):
+                duration = (time.time() - start_time) * 1000
+                logger.debug(
+                    f"[{request_id}] RBAC STATS - {status_code} {method} "
+                    f"{endpoint} completed in {duration:.2f}ms"
                 )
 
         return response

zenml/zen_server/routers/stack_components_endpoints.py

@@ -93,7 +93,7 @@ def create_stack_component(
 
     from zenml.stack.utils import validate_stack_component_config
 
-    validate_stack_component_config(
+    validated_config = validate_stack_component_config(
         configuration_dict=component.configuration,
         flavor=component.flavor,
         component_type=component.type,
@@ -102,6 +102,11 @@ def create_stack_component(
         validate_custom_flavors=False,
     )
 
+    if validated_config:
+        component.configuration = validated_config.model_dump(
+            mode="json", exclude_unset=True
+        )
+
     return verify_permissions_and_create_entity(
         request_model=component,
         create_method=zen_store().create_stack_component,
@@ -199,7 +204,7 @@ def update_stack_component(
         from zenml.stack.utils import validate_stack_component_config
 
         existing_component = zen_store().get_stack_component(component_id)
-        validate_stack_component_config(
+        validated_config = validate_stack_component_config(
             configuration_dict=component_update.configuration,
             flavor=existing_component.flavor_name,
             component_type=existing_component.type,
@@ -207,6 +212,10 @@ def update_stack_component(
             # We allow custom flavors to fail import on the server side.
             validate_custom_flavors=False,
         )
+        if validated_config:
+            component_update.configuration = validated_config.model_dump(
+                mode="json", exclude_unset=True
+            )
 
     if component_update.connector:
         service_connector = zen_store().get_service_connector(

zenml/zen_server/utils.py

@@ -15,6 +15,7 @@
 
 import inspect
 import os
+import threading
 from functools import wraps
 from typing import (
     TYPE_CHECKING,
@@ -322,40 +323,6 @@ def async_fastapi_endpoint_wrapper(
         Decorated function.
     """
 
-    @wraps(func)
-    def decorated(*args: P.args, **kwargs: P.kwargs) -> Any:
-        # These imports can't happen at module level as this module is also
-        # used by the CLI when installed without the `server` extra
-        from fastapi import HTTPException
-        from fastapi.responses import JSONResponse
-
-        from zenml.zen_server.auth import AuthContext, set_auth_context
-
-        for arg in args:
-            if isinstance(arg, AuthContext):
-                set_auth_context(arg)
-                break
-        else:
-            for _, arg in kwargs.items():
-                if isinstance(arg, AuthContext):
-                    set_auth_context(arg)
-                    break
-
-        try:
-            return func(*args, **kwargs)
-        except OAuthError as error:
-            # The OAuthError is special because it needs to have a JSON response
-            return JSONResponse(
-                status_code=error.status_code,
-                content=error.to_dict(),
-            )
-        except HTTPException:
-            raise
-        except Exception as error:
-            logger.exception("API error")
-            http_exception = http_exception_from_error(error)
-            raise http_exception
-
     # When having a sync FastAPI endpoint, it runs the endpoint function in
     # a worker threadpool. If all threads are busy, it will queue the task.
     # The problem is that after the endpoint code returns, FastAPI will queue
@@ -369,10 +336,52 @@ def async_fastapi_endpoint_wrapper(
     # a worker thread to become available.
     # See: `fastapi.routing.serialize_response(...)` and
     # https://github.com/fastapi/fastapi/pull/888 for more information.
-    @wraps(decorated)
+    @wraps(func)
     async def async_decorated(*args: P.args, **kwargs: P.kwargs) -> Any:
         from starlette.concurrency import run_in_threadpool
 
+        from zenml.zen_server.zen_server_api import request_ids
+
+        request_id = request_ids.get()
+
+        @wraps(func)
+        def decorated(*args: P.args, **kwargs: P.kwargs) -> Any:
+            # These imports can't happen at module level as this module is also
+            # used by the CLI when installed without the `server` extra
+            from fastapi import HTTPException
+            from fastapi.responses import JSONResponse
+
+            from zenml.zen_server.auth import AuthContext, set_auth_context
+
+            if request_id:
+                # Change the name of the current thread to the request ID
+                threading.current_thread().name = request_id
+
+            for arg in args:
+                if isinstance(arg, AuthContext):
+                    set_auth_context(arg)
+                    break
+            else:
+                for _, arg in kwargs.items():
+                    if isinstance(arg, AuthContext):
+                        set_auth_context(arg)
+                        break
+
+            try:
+                return func(*args, **kwargs)
+            except OAuthError as error:
+                # The OAuthError is special because it needs to have a JSON response
+                return JSONResponse(
+                    status_code=error.status_code,
+                    content=error.to_dict(),
+                )
+            except HTTPException:
+                raise
+            except Exception as error:
+                logger.exception("API error")
+                http_exception = http_exception_from_error(error)
+                raise http_exception
+
         return await run_in_threadpool(decorated, *args, **kwargs)
 
     return async_decorated

zenml/zen_server/zen_server_api.py

@@ -20,11 +20,17 @@ To run this file locally, execute:
     ```
 """
 
+import logging
 import os
+import threading
+import time
+from asyncio import Lock, Semaphore, TimeoutError, wait_for
 from asyncio.log import logger
+from contextvars import ContextVar
 from datetime import datetime, timedelta
 from genericpath import isfile
-from typing import Any, List, Set
+from typing import Any, List, Optional, Set
+from uuid import uuid4
 
 from anyio import to_thread
 from fastapi import FastAPI, HTTPException, Request
@@ -113,6 +119,10 @@ from zenml.zen_server.utils import (
 
 DASHBOARD_DIRECTORY = "dashboard"
 
+request_ids: ContextVar[Optional[str]] = ContextVar(
+    "request_ids", default=None
+)
+
 
 def relative_path(rel: str) -> str:
     """Get the absolute path of a path relative to the ZenML server module.
@@ -138,6 +148,7 @@ last_user_activity: datetime = utc_now()
 last_user_activity_reported: datetime = last_user_activity + timedelta(
     seconds=-DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
 )
+last_user_activity_lock = Lock()
 
 
 # Customize the default request validation handler that comes with FastAPI
@@ -247,51 +258,6 @@ class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
 
 ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=server_config().cors_allow_origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-app.add_middleware(
-    RequestBodyLimit, max_bytes=server_config().max_request_body_size_in_bytes
-)
-app.add_middleware(
-    RestrictFileUploadsMiddleware, allowed_paths=ALLOWED_FOR_FILE_UPLOAD
-)
-
-
-@app.middleware("http")
-async def set_secure_headers(request: Request, call_next: Any) -> Any:
-    """Middleware to set secure headers.
-
-    Args:
-        request: The incoming request.
-        call_next: The next function to be called.
-
-    Returns:
-        The response with secure headers set.
-    """
-    try:
-        response = await call_next(request)
-    except Exception:
-        logger.exception("An error occurred while processing the request")
-        response = JSONResponse(
-            status_code=500,
-            content={"detail": "An unexpected error occurred."},
-        )
-
-    # If the request is for the openAPI docs, don't set secure headers
-    if request.url.path.startswith("/docs") or request.url.path.startswith(
-        "/redoc"
-    ):
-        return response
-
-    secure_headers().framework.fastapi(response)
-    return response
-
 
 @app.middleware("http")
 async def track_last_user_activity(request: Request, call_next: Any) -> Any:
@@ -310,24 +276,30 @@ async def track_last_user_activity(request: Request, call_next: Any) -> Any:
     """
     global last_user_activity
     global last_user_activity_reported
+    global last_user_activity_lock
 
     now = utc_now()
 
     try:
         if is_user_request(request):
-            last_user_activity = now
+            report_user_activity = False
+            async with last_user_activity_lock:
+                last_user_activity = now
+                if (
+                    (now - last_user_activity_reported).total_seconds()
+                    > DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
+                ):
+                    last_user_activity_reported = now
+                    report_user_activity = True
+
+            if report_user_activity:
+                zen_store()._update_last_user_activity_timestamp(
+                    last_user_activity=last_user_activity
+                )
     except Exception as e:
         logger.debug(
             f"An unexpected error occurred while checking user activity: {e}"
         )
-    if (
-        (now - last_user_activity_reported).total_seconds()
-        > DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
-    ):
-        last_user_activity_reported = now
-        zen_store()._update_last_user_activity_timestamp(
-            last_user_activity=last_user_activity
-        )
 
     try:
         return await call_next(request)
@@ -378,6 +350,190 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
         )
 
 
+request_semaphore = Semaphore(server_config().thread_pool_size)
+
+
+@app.middleware("http")
+async def prevent_read_timeout(request: Request, call_next: Any) -> Any:
+    """Prevent read timeout client errors.
+
+    Args:
+        request: The incoming request.
+        call_next: The next function to be called.
+
+    Returns:
+        The response to the request.
+    """
+    # Only process the REST API requests because these are the ones that
+    # take the most time to complete.
+    if not request.url.path.startswith(API):
+        return await call_next(request)
+
+    server_request_timeout = server_config().server_request_timeout
+
+    active_threads = threading.active_count()
+    request_id = request_ids.get()
+
+    client_ip = request.client.host if request.client else "unknown"
+    method = request.method
+    url_path = request.url.path
+
+    logger.debug(
+        f"[{request_id}] API STATS - {method} {url_path} from {client_ip} "
+        f"QUEUED [ "
+        f"threads: {active_threads} "
+        f"]"
+    )
+
+    start_time = time.time()
+
+    try:
+        # Here we wait until a worker thread is available to process the
+        # request with a timeout value that is set to be lower than the
+        # what the client is willing to wait for (i.e. lower than the
+        # client's HTTP request timeout). The rationale is that we want to
+        # respond to the client before it times out and decides to retry the
+        # request (which would overwhelm the server).
+        await wait_for(
+            request_semaphore.acquire(),
+            timeout=server_request_timeout,
+        )
+    except TimeoutError:
+        end_time = time.time()
+        duration = (end_time - start_time) * 1000
+        active_threads = threading.active_count()
+
+        logger.debug(
+            f"[{request_id}] API STATS - {method} {url_path} from {client_ip} "
+            f"THROTTLED after {duration:.2f}ms [ "
+            f"threads: {active_threads} "
+            f"]"
+        )
+
+        # We return a 429 error, basically telling the client to slow down.
+        # For the client, the 429 error is more meaningful than a ReadTimeout
+        # error, because it also tells the client two additional things:
+        #
+        # 1. The server is alive.
+        # 2. The server hasn't processed the request, so even if the request
+        #    is not idempotent, it's safe to retry it.
+        return JSONResponse(
+            {"error": "Server too busy. Please try again later."},
+            status_code=429,
+        )
+
+    duration = (time.time() - start_time) * 1000
+    active_threads = threading.active_count()
+
+    logger.debug(
+        f"[{request_id}] API STATS - {method} {url_path} from {client_ip} "
+        f"ACCEPTED after {duration:.2f}ms [ "
+        f"threads: {active_threads} "
+        f"]"
+    )
+
+    try:
+        return await call_next(request)
+    finally:
+        request_semaphore.release()
+
+
+@app.middleware("http")
+async def log_requests(request: Request, call_next: Any) -> Any:
+    """Log requests to the ZenML server.
+
+    Args:
+        request: The incoming request object.
+        call_next: A function that will receive the request as a parameter and
+            pass it to the corresponding path operation.
+
+    Returns:
+        The response to the request.
+    """
+    if not logger.isEnabledFor(logging.DEBUG):
+        return await call_next(request)
+
+    # Get active threads count
+    active_threads = threading.active_count()
+
+    request_id = request.headers.get("X-Request-ID", str(uuid4())[:8])
+    # Detect if the request comes from Python, Web UI or something else
+    if source := request.headers.get("User-Agent"):
+        source = source.split("/")[0]
+        request_id = f"{request_id}/{source}"
+
+    request_ids.set(request_id)
+    client_ip = request.client.host if request.client else "unknown"
+    method = request.method
+    url_path = request.url.path
+
+    logger.debug(
+        f"[{request_id}] API STATS - {method} {url_path} from {client_ip} "
+        f"RECEIVED [ "
+        f"threads: {active_threads} "
+        f"]"
+    )
+
+    start_time = time.time()
+    response = await call_next(request)
+    duration = (time.time() - start_time) * 1000
+    status_code = response.status_code
+
+    logger.debug(
+        f"[{request_id}] API STATS - {status_code} {method} {url_path} from "
+        f"{client_ip} took {duration:.2f}ms [ "
+        f"threads: {active_threads} "
+        f"]"
+    )
+    return response
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=server_config().cors_allow_origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.add_middleware(
+    RequestBodyLimit, max_bytes=server_config().max_request_body_size_in_bytes
+)
+app.add_middleware(
+    RestrictFileUploadsMiddleware, allowed_paths=ALLOWED_FOR_FILE_UPLOAD
+)
+
+
+@app.middleware("http")
+async def set_secure_headers(request: Request, call_next: Any) -> Any:
+    """Middleware to set secure headers.
+
+    Args:
+        request: The incoming request.
+        call_next: The next function to be called.
+
+    Returns:
+        The response with secure headers set.
+    """
+    try:
+        response = await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        response = JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
+
+    # If the request is for the openAPI docs, don't set secure headers
+    if request.url.path.startswith("/docs") or request.url.path.startswith(
+        "/redoc"
+    ):
+        return response
+
+    secure_headers().framework.fastapi(response)
+    return response
+
+
 @app.on_event("startup")
 def initialize() -> None:
     """Initialize the ZenML server."""

zenml/zen_stores/rest_zen_store.py

@@ -15,6 +15,7 @@
 
 import os
 import re
+import time
 from datetime import datetime
 from pathlib import Path
 from typing import (
@@ -30,7 +31,7 @@ from typing import (
     Union,
 )
 from urllib.parse import urlparse
-from uuid import UUID
+from uuid import UUID, uuid4
 
 import requests
 import urllib3
@@ -4179,6 +4180,20 @@ class RestZenStore(BaseZenStore):
         Returns:
             A requests session.
         """
+
+        class AugmentedRetry(Retry):
+            """Augmented retry class that also retries on 429 status codes for POST requests."""
+
+            def is_retry(
+                self,
+                method: str,
+                status_code: int,
+                has_retry_after: bool = False,
+            ) -> bool:
+                if status_code == 429:
+                    return True
+                return super().is_retry(method, status_code, has_retry_after)
+
         if self._session is None:
             # We only need to initialize the session once over the lifetime
             # of the client. We can swap the token out when it expires.
@@ -4208,7 +4223,7 @@ class RestZenStore(BaseZenStore):
             #     the timeout period.
             #     Connection Refused: If the server refuses the connection.
             #
-            retries = Retry(
+            retries = AugmentedRetry(
                 connect=5,
                 read=8,
                 redirect=3,
@@ -4222,7 +4237,7 @@ class RestZenStore(BaseZenStore):
                     504,  # Gateway Timeout
                 ],
                 other=3,
-                backoff_factor=0.5,
+                backoff_factor=1,
             )
             self._session.mount("https://", HTTPAdapter(max_retries=retries))
             self._session.mount("http://", HTTPAdapter(max_retries=retries))
@@ -4360,6 +4375,14 @@ class RestZenStore(BaseZenStore):
         self.session.headers.update(
             {source_context.name: source_context.get().value}
         )
+        # Add a request ID to the request headers
+        request_id = str(uuid4())[:8]
+        self.session.headers.update({"X-Request-ID": request_id})
+        path = url.removeprefix(self.url)
+        start_time = time.time()
+        logger.debug(
+            f"Sending {method} request to {path} with request ID {request_id}..."
+        )
 
         # If the server replies with a credentials validation (401 Unauthorized)
         # error, we (re-)authenticate and retry the request here in the
@@ -4401,7 +4424,8 @@ class RestZenStore(BaseZenStore):
                     # request again, this time with a valid API token in the
                     # header.
                     logger.debug(
-                        f"The last request was not authenticated: {e}\n"
+                        f"The last request with ID {request_id} was not "
+                        f"authenticated: {e}\n"
                         "Re-authenticating and retrying..."
                     )
                     self.authenticate()
@@ -4428,8 +4452,9 @@ class RestZenStore(BaseZenStore):
                     # that was rejected by the server. We attempt a
                     # re-authentication here and then retry the request.
                     logger.debug(
-                        "The last request was authenticated with an API token "
-                        f"that was rejected by the server: {e}\n"
+                        f"The last request with ID {request_id} was authenticated "
+                        "with an API token that was rejected by the server: "
+                        f"{e}\n"
                         "Re-authenticating and retrying..."
                     )
                     re_authenticated = True
@@ -4441,13 +4466,21 @@ class RestZenStore(BaseZenStore):
                     # The last request was made after re-authenticating but
                     # still failed. Bailing out.
                     logger.debug(
-                        f"The last request failed after re-authenticating: {e}\n"
+                        f"The last request with ID {request_id} failed after "
+                        "re-authenticating: {e}\n"
                         "Bailing out..."
                     )
                     raise CredentialsNotValid(
                         "The current credentials are no longer valid. Please "
                         "log in again using 'zenml login'."
                     ) from e
+            finally:
+                end_time = time.time()
+                duration = (end_time - start_time) * 1000
+                logger.debug(
+                    f"Request to {path} with request ID {request_id} took "
+                    f"{duration:.2f}ms."
+                )
 
     def get(
         self,
@@ -4467,7 +4500,6 @@ class RestZenStore(BaseZenStore):
         Returns:
             The response body.
         """
-        logger.debug(f"Sending GET request to {path}...")
         return self._request(
             "GET",
             self.url + API + VERSION_1 + path,
@@ -4496,7 +4528,6 @@ class RestZenStore(BaseZenStore):
         Returns:
             The response body.
         """
-        logger.debug(f"Sending DELETE request to {path}...")
         return self._request(
             "DELETE",
             self.url + API + VERSION_1 + path,
@@ -4526,7 +4557,6 @@ class RestZenStore(BaseZenStore):
         Returns:
             The response body.
         """
-        logger.debug(f"Sending POST request to {path}...")
         return self._request(
             "POST",
             self.url + API + VERSION_1 + path,
@@ -4556,7 +4586,6 @@ class RestZenStore(BaseZenStore):
         Returns:
             The response body.
         """
-        logger.debug(f"Sending PUT request to {path}...")
         json = (
             body.model_dump(mode="json", exclude_unset=True) if body else None
         )

zenml/zen_stores/sql_zen_store.py

@@ -14,6 +14,7 @@
 """SQL Zen Store implementation."""
 
 import base64
+import inspect
 import json
 import logging
 import math
@@ -21,6 +22,7 @@ import os
 import random
 import re
 import sys
+import threading
 import time
 from collections import defaultdict
 from datetime import datetime
@@ -58,7 +60,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from sqlalchemy import func
+from sqlalchemy import QueuePool, func
 from sqlalchemy.engine import URL, Engine, make_url
 from sqlalchemy.exc import (
     ArgumentError,
@@ -66,6 +68,7 @@ from sqlalchemy.exc import (
 )
 from sqlalchemy.orm import Mapped, noload
 from sqlalchemy.util import immutabledict
+from sqlmodel import Session as SqlModelSession
 
 # Important to note: The select function of SQLModel works slightly differently
 # from the select function of sqlalchemy. If you input only one entity on the
@@ -74,7 +77,6 @@ from sqlalchemy.util import immutabledict
 # the tuple. While this is convenient in most cases, in unique cases like using
 # the "add_columns" functionality, one might encounter unexpected results.
 from sqlmodel import (
-    Session,
     SQLModel,
     and_,
     col,
@@ -412,6 +414,81 @@ def exponential_backoff_with_jitter(
     return random.uniform(0, exponential_backoff)
 
 
+class Session(SqlModelSession):
+    """Session subclass that automatically tracks duration and calling context."""
+
+    def __enter__(self) -> "Session":
+        """Enter the context manager.
+
+        Returns:
+            The SqlModel session.
+        """
+        if logger.isEnabledFor(logging.DEBUG):
+            # Get the request ID from the current thread object
+            self.request_id = threading.current_thread().name
+
+            # Get SQLAlchemy connection pool info
+            assert isinstance(self.bind, Engine)
+            assert isinstance(self.bind.pool, QueuePool)
+            checked_out_connections = self.bind.pool.checkedout()
+            available_connections = self.bind.pool.checkedin()
+            overflow = self.bind.pool.overflow()
+
+            # Look up the stack to find the SQLZenStore method
+            for frame in inspect.stack():
+                if "self" in frame.frame.f_locals:
+                    instance = frame.frame.f_locals["self"]
+                    if isinstance(instance, SqlZenStore):
+                        self.caller_method = (
+                            f"{instance.__class__.__name__}.{frame.function}"
+                        )
+                        break
+            else:
+                self.caller_method = "unknown"
+
+            logger.debug(
+                f"[{self.request_id}] SQL STATS - "
+                f"'{self.caller_method}' started [ conn(active): "
+                f"{checked_out_connections} conn(idle): "
+                f"{available_connections} conn(overflow): {overflow} ]"
+            )
+
+            self.start_time = time.time()
+
+        return super().__enter__()
+
+    def __exit__(
+        self,
+        exc_type: Optional[Any],
+        exc_val: Optional[Any],
+        exc_tb: Optional[Any],
+    ) -> None:
+        """Exit the context manager.
+
+        Args:
+            exc_type: The exception type.
+            exc_val: The exception value.
+            exc_tb: The exception traceback.
+        """
+        if logger.isEnabledFor(logging.DEBUG):
+            duration = (time.time() - self.start_time) * 1000
+
+            # Get SQLAlchemy connection pool info
+            assert isinstance(self.bind, Engine)
+            assert isinstance(self.bind.pool, QueuePool)
+            checked_out_connections = self.bind.pool.checkedout()
+            available_connections = self.bind.pool.checkedin()
+            overflow = self.bind.pool.overflow()
+            logger.debug(
+                f"[{self.request_id}] SQL STATS - "
+                f"'{self.caller_method}' completed in "
+                f"{duration:.2f}ms [ conn(active): "
+                f"{checked_out_connections} conn(idle): "
+                f"{available_connections} conn(overflow): {overflow} ]"
+            )
+        super().__exit__(exc_type, exc_val, exc_tb)
+
+
 class SQLDatabaseDriver(StrEnum):
     """SQL database drivers supported by the SQL ZenML store."""
 

{zenml_nightly-0.82.1.dev20250522.dist-info → zenml_nightly-0.82.1.dev20250524.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.82.1.dev20250522
+Version: 0.82.1.dev20250524
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops

{zenml_nightly-0.82.1.dev20250522.dist-info → zenml_nightly-0.82.1.dev20250524.dist-info}/RECORD

@@ -1,5 +1,5 @@
 zenml/README.md,sha256=827dekbOWAs1BpW7VF1a4d7EbwPbjwccX-2zdXBENZo,1777
-zenml/VERSION,sha256=V62XyuhhLfntir6O3u5XhATyUFx8svXLb1DDa10S8UM,19
+zenml/VERSION,sha256=MexsT1YdyW8E4qM4hiGrG8oqIdUf48PHwk4-TGyiebg,19
 zenml/__init__.py,sha256=CKEyepFK-7akXYiMrNVh92Nb01Cjs23w4_YyI6sgdc8,2242
 zenml/actions/__init__.py,sha256=mrt6wPo73iKRxK754_NqsGyJ3buW7RnVeIGXr1xEw8Y,681
 zenml/actions/base_action.py,sha256=UcaHev6BTuLDwuswnyaPjdA8AgUqB5xPZ-lRtuvf2FU,25553
@@ -55,7 +55,7 @@ zenml/cli/text_utils.py,sha256=bY1GIjoULt1cW2FyrPlMoAXNS2R7cSOjDFEZQqrpVQ8,3553
 zenml/cli/user_management.py,sha256=sNnhaUxH-cHecbZBR1L0mEU0TnLNZHzI6ZBCUSQa7OY,13078
 zenml/cli/utils.py,sha256=vMAb9f6GDfNVGmZWOz9UOyPRpKI3KfnYpRl_w9YUBNE,86501
 zenml/cli/version.py,sha256=nm1iSU_1V6-MUwpMKeXcwFhLYGUMLswvQL67cEuCpxA,3635
-zenml/client.py,sha256=bQqMJVZcg5TJxwQLsB2zmKaTQ6PdYdUALIqytqO-j1k,293057
+zenml/client.py,sha256=aCj2HEE6M724js6t4QJQq1S9R6C-_hxQe4UMUN0gjbw,293202
 zenml/client_lazy_loader.py,sha256=MOBgS1ITYqGvPUnWQ6edn9s8Hr_72YfWbwEIfHKUr9g,7104
 zenml/code_repositories/__init__.py,sha256=W5bDfzAG8OXIKZSV1L-VHuzMcSCYa9qzTdPb3jqfyYw,920
 zenml/code_repositories/base_code_repository.py,sha256=Id6VjbUu8N3ZpNvBGhOgbahtoMiCAtYXed3G7YQ_iAc,5225
@@ -67,7 +67,7 @@ zenml/config/base_settings.py,sha256=itoLqc1cOwEYhgSGdZmSKSaBevQkvYH7NQh7PUamazc
 zenml/config/build_configuration.py,sha256=jGGNwP0Cb7a80JXArNxDgxzxl9ytSZRtv-WW7_psLbM,6870
 zenml/config/compiler.py,sha256=bK3LCDkrFc9SapJYH-vuQZ_o8scHNs-FdC5DblIUU4U,23024
 zenml/config/constants.py,sha256=QvSgMwXWxtspcJ45CrFDP1ZY3w6gS3bIhXLOtIDAbZA,713
-zenml/config/docker_settings.py,sha256=B3pYlnyQDiwSEHO6y965Jx6BMrVqKARjp7v23uBziLc,17068
+zenml/config/docker_settings.py,sha256=WhoocDGn9cegEd4mKb_oo7pZ87sXZ0NvqkOVnWK8CY4,17985
 zenml/config/global_config.py,sha256=ZD3WodfcWMBJZOl1FNn3ztzwilGDDv9EKdHClLqSO8s,29562
 zenml/config/pipeline_configurations.py,sha256=7trCbElpqGGgawii2FrdLW8fKaAWCR8jACkNqdG_vcQ,3983
 zenml/config/pipeline_run_configuration.py,sha256=Y9C5mVUH-w4gc1w1PCQFpjEmfBBpSMvb8riA_sL78hY,2311
@@ -77,7 +77,7 @@ zenml/config/retry_config.py,sha256=4UH1xqw0G6fSEbXSNKfmiFEkwadxQef9BGMe3JBm6NI,
 zenml/config/schedule.py,sha256=-83j99U9OyiG7E322XWA7QvuLSwQzF21whwpeiF0b30,5348
 zenml/config/secret_reference_mixin.py,sha256=YvY68MTd1gE23IVprf0BLkNn62hoxcvb5nqGgc8jMkU,5871
 zenml/config/secrets_store_config.py,sha256=y05zqyQhr_DGrs3IfBGa_FRoZ043hSYRT5wzrx-zHTU,2818
-zenml/config/server_config.py,sha256=or-LKYFl-9yE3MEq_yPa5aCCu3d5_f0tL9MqPtb_3-c,31852
+zenml/config/server_config.py,sha256=DYYQ10HddkvlCWobBpwpg8ggO6imMZCnfFxzqlkQg_U,32336
 zenml/config/settings_resolver.py,sha256=PR9BRm_x1dy7nVKa9UqpeFdck8IEATSW6aWT8FKd-DI,4278
 zenml/config/source.py,sha256=RzUw8lin8QztUjz-AdoCzVM5Om_cSSPuroaPx-qAO4w,8226
 zenml/config/step_configurations.py,sha256=vjGtCoHneUJeaMkTiuzl7TTPBO94MefxNlFIAL1l5rU,13038
@@ -85,7 +85,7 @@ zenml/config/step_run_info.py,sha256=KiVRSTtKmZ1GbvseDTap2imr7XwMHD3jSFVpyLNEK1I
 zenml/config/store_config.py,sha256=Cla5p5dTB6nNlo8_OZDs9hod5hspi64vxwtZj882XgU,3559
 zenml/config/strict_base_model.py,sha256=t_ULrtJF2eW7TgyYBRobl1fscwwIZXATYky8ER97ev4,860
 zenml/console.py,sha256=hj_KerPQKwnyKACj0ehSqUQX0mGVCJBKE1QvCt6ik3A,1160
-zenml/constants.py,sha256=nLjJUhBqFHWHlGKbZZw-_tn7x99y6q4vmk5LOe7GhJI,16597
+zenml/constants.py,sha256=O8b67hvsObx42r2REHA_X4XqDwqm_Hcsq7hv6q96HDs,16650
 zenml/container_registries/__init__.py,sha256=ZSPbBIOnzhg88kQSpYgKe_POLuru14m629665-kAVAA,2200
 zenml/container_registries/azure_container_registry.py,sha256=t1sfDa94Vzbyqtb1iPFNutJ2EXV5_p9CUNITasoiQ70,2667
 zenml/container_registries/base_container_registry.py,sha256=-9RIkD6oXNPaU59R3PB_PtyCqsFoLPLSn5xYZmEmzbc,8915
@@ -100,8 +100,8 @@ zenml/entrypoints/base_entrypoint_configuration.py,sha256=t7sU2MEEPVomX9dZCpPhk1
 zenml/entrypoints/entrypoint.py,sha256=XNgXBCMKoidmP0_AYgMpqo-neG8Y8jG0rj43ofTDZ9E,2033
 zenml/entrypoints/pipeline_entrypoint_configuration.py,sha256=To-vTP29qAE36ndJDF1fRw9wL2Nk2bsBuO-ayAwvSmo,1646
 zenml/entrypoints/step_entrypoint_configuration.py,sha256=fJuTvJnGuhKc60CH1VMQL5EHomGXkYZulv6pVgd9M6w,7316
-zenml/enums.py,sha256=m18yol20jS9JNhuf7UPBYUsgNRCtbyCaWPNpeEAJeXY,11199
-zenml/environment.py,sha256=UtriBwD0ylEK1cmSjEkDTlUUXBifBEukAn6gEdcnv4M,11827
+zenml/enums.py,sha256=xXp6VkaY03NlGKEBs9Bf7oVfUrtESECmhkQYAaKRS5g,11239
+zenml/environment.py,sha256=_iS9o4rTxRCjLocSw_LEe9pZ9jfyMhXBY47LIUF3VCI,12255
 zenml/event_hub/__init__.py,sha256=-fD9mPOslf4J-_XFBPp5gYlPz7-6ZaAKHa5jxf_nyAo,757
 zenml/event_hub/base_event_hub.py,sha256=PqKrnvOye0UUS3s09zGgjI5dtj0IwzEBDbavA_PgfZQ,6579
 zenml/event_hub/event_hub.py,sha256=e1eCRB1qAabOFIx2OJCAoNqN1QQIW0CSTukiWuGMOi8,6502
@@ -253,7 +253,7 @@ zenml/integrations/feast/feature_stores/__init__.py,sha256=Wi3NBBBPJg6CjgtxmBjoU
 zenml/integrations/feast/feature_stores/feast_feature_store.py,sha256=jV6WznuKT3y1aikI3OEwoI8r_l8bEu5waX0LKePPuU8,5880
 zenml/integrations/feast/flavors/__init__.py,sha256=gbCZ4tKgLZSI4-gzOCR2xihiPNmpe-lMUxwvMrhYL-w,858
 zenml/integrations/feast/flavors/feast_feature_store_flavor.py,sha256=E0k2iwgNti8lOVr9W8n2nTVQQf2wSeJWeaQD1vwAIbU,3060
-zenml/integrations/gcp/__init__.py,sha256=4bl_mutvdwXw-ZVdWdTHi3spTRH1JRCtBNPFxq5c74k,2957
+zenml/integrations/gcp/__init__.py,sha256=9o7TXcQlNebImis821VXcOVFxKau86mbQ7H8d8e6mIw,2984
 zenml/integrations/gcp/artifact_stores/__init__.py,sha256=zYQkZBI4-COGX-E0NS7G-hLT88wbQBnYY6Oh1gruSRs,798
 zenml/integrations/gcp/artifact_stores/gcp_artifact_store.py,sha256=XfSIJ4HtsZvaUrUtzXvUp7QHr3WbgVDNsz7_q1h-DCo,10988
 zenml/integrations/gcp/constants.py,sha256=ZBQS_ZEjerUrJq-hH3UusgZAvB45FLgxNv11TSt3qhw,1334
@@ -271,7 +271,7 @@ zenml/integrations/gcp/image_builders/gcp_image_builder.py,sha256=5T6BXsHxLhvp1B
 zenml/integrations/gcp/orchestrators/__init__.py,sha256=6xLFJKZKQk73fHPF-XdpbQO87zjQNGTsNHjJjLfG_Kg,805
 zenml/integrations/gcp/orchestrators/vertex_orchestrator.py,sha256=qoMCr36buZUz0y4CyTFQde3RDslkaGLAG0FjXc0XEPU,42100
 zenml/integrations/gcp/service_connectors/__init__.py,sha256=fdydawaor8KAtMYvRZieiTuA1i5QATxXXgI-yV1lsn8,788
-zenml/integrations/gcp/service_connectors/gcp_service_connector.py,sha256=JGrFTKkQV4ZDn_yTEqX-AtBMraFasFgzLVws2mvhS64,94915
+zenml/integrations/gcp/service_connectors/gcp_service_connector.py,sha256=9u-vEHbmSyN5IGwYI8v39TcFZg5ObgkxlbwSPz-e5zE,95018
 zenml/integrations/gcp/step_operators/__init__.py,sha256=iPkob2LtPIQ-OHszhbNz_ojhoovL6SprmTx37It4EJ8,808
 zenml/integrations/gcp/step_operators/vertex_step_operator.py,sha256=X8CCniyAo7NHiy3Mv_YSKQ4Hw3UYMXob6B3uWKsCJ-0,13567
 zenml/integrations/gcp/vertex_custom_job_parameters.py,sha256=B5RLkw7KDOi4ZfWHFnC6TGLphXMzToMjROxszCEAS9c,3676
@@ -317,7 +317,7 @@ zenml/integrations/hyperai/__init__.py,sha256=6ed5rgRgiRr2Ksi3DDFUaNQwPkCfE4BRSn
 zenml/integrations/hyperai/flavors/__init__.py,sha256=PUGBPmQ7y3H7QU2zAj7Ri0rrUBbOWnM_L59AIVWUYwU,800
 zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py,sha256=erjbYJ5E3qkifyGd_ArtbY1JdfyOEz4vgl9QrZK0R9Y,5578
 zenml/integrations/hyperai/orchestrators/__init__.py,sha256=kSYpMZPEWwNu2vxoOC6PeyQ9RLzsPAgTHxL35K36MiE,784
-zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py,sha256=2bKFlJE6HUU3oItkj3-siVxTKUpextN25t8aX7du4ZU,20286
+zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py,sha256=6Z4GAXfkC2Ia4czQrgCy85xuocC9TYXTLmaNP8xxv_s,20287
 zenml/integrations/hyperai/service_connectors/__init__.py,sha256=oHuCNC09z5C7Wlb3vV1d4zJjftttHg364eoEBVRDOdo,803
 zenml/integrations/hyperai/service_connectors/hyperai_service_connector.py,sha256=7Ql5cVoSY3SE4j7uzeVi5TWuoKFDvsHFTn72k9_wPUY,13400
 zenml/integrations/integration.py,sha256=ljUK2GHy-LNqsdloOvf8mIfUQsqoi0yBJYocMyS-DX8,6879
@@ -337,7 +337,7 @@ zenml/integrations/kubernetes/flavors/__init__.py,sha256=a5gU45qCj3FkLwl_uVjlIkL
 zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py,sha256=JH_Kxfh1eoiSlz-OpixJ6itEsDiGPvbkA5_PffMTa54,10215
 zenml/integrations/kubernetes/flavors/kubernetes_step_operator_flavor.py,sha256=xFO7cSusji-mgbRrt4mU29gdyC9iEjEHKtomdFLp9mM,6265
 zenml/integrations/kubernetes/orchestrators/__init__.py,sha256=TJID3OTieZBox36WpQpzD0jdVRA_aZVcs_bNtfXS8ik,811
-zenml/integrations/kubernetes/orchestrators/kube_utils.py,sha256=8BVOpg40dGeRSkbzG_66xOksry5xpMVVmcXVDqb1bSs,18747
+zenml/integrations/kubernetes/orchestrators/kube_utils.py,sha256=N66GH5ac22Xm_A3nr162kbFBhMeypSFaQjOQRHlGXIQ,18942
 zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py,sha256=qubV13fSGE_vo8wSzIfcp66McclLw-zDtmB4yCCaS08,25974
 zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py,sha256=Nd1YZsP6FZ4LTpoeFIM3YlFUugHX7F9axb7U9ZgGNdQ,12692
 zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint_configuration.py,sha256=KjHfQK9VQEQkkkM2i9w51AzqolgIU01M5dgb2YGamvY,2754
@@ -596,7 +596,7 @@ zenml/login/server_info.py,sha256=-_sK2L-curHdzUv1JDOwF6GoEeAXT5vFZN0J-5Ug4wU,16
 zenml/login/web_login.py,sha256=Kq_fA9UQEravB2DtAkMmNvDttk8xppnxV617tCYUl6U,9186
 zenml/materializers/__init__.py,sha256=C3lZaTmIFxwIPwCKF8oLQUkLaX2o7Dbj9hvYVFrSzt8,1758
 zenml/materializers/base_materializer.py,sha256=M4hwkw7PB0LskCE92r-S35011l7DlFemit-EuUCW3Nc,14002
-zenml/materializers/built_in_materializer.py,sha256=KjW3p8LWORZTjHqoGGC-Hlf-bdNSir3zsbRD2C_oQBQ,16951
+zenml/materializers/built_in_materializer.py,sha256=HgcCHg3GpuQ4t-jecYcdg0kldUfLeUJvIpm8-wcS11I,17189
 zenml/materializers/cloudpickle_materializer.py,sha256=x8a6jEMTky6N2YVHiwrnGWSfVJUpiy-4kQsD2Aqj_E0,4837
 zenml/materializers/materializer_registry.py,sha256=ic-aWhJ2Ex9F_rml2dDVAxhRfW3nd71QMxzfTPP6BIM,4002
 zenml/materializers/numpy_materializer.py,sha256=OLcHF9Z0tAqQ_U8TraA0vGmZjHoT7eT_XevncIutt0M,1715
@@ -805,7 +805,7 @@ zenml/utils/yaml_utils.py,sha256=RvEr-N84wwG1Aq8rfdPAP7lev_SEaSQlqcL9fB_vHk8,584
 zenml/zen_server/__init__.py,sha256=WyltI9TzFW2mEHZVOs6alLWMCQrrZaFALtrQXs83STA,1355
 zenml/zen_server/auth.py,sha256=sJ_UAG58Qa0sUx8P1QkWmwp5zfAafFnm8WBHAOnPi4w,40070
 zenml/zen_server/cache.py,sha256=Tc4TSugmsU1bhThxlYfE8rv0KmltIX1CcVHgzrJ0Eus,6633
-zenml/zen_server/cloud_utils.py,sha256=qKMtHjnkVfAk_yxF2K9iC0TkMCjp8s0gdz_cEsbB4DE,10908
+zenml/zen_server/cloud_utils.py,sha256=vKsjWadgPctZDtfeG9qJC-5gKOir_tZsglYMBkC8QqU,11874
 zenml/zen_server/csrf.py,sha256=Jsbi_IKriWCOuquSYTOEWqSXiGORJATIhR5Wrkm4XzQ,2684
 zenml/zen_server/dashboard/assets/404-_AtuLtaX.js,sha256=Y2jAo7K2feru-k7QcMxjnfUNkDrO7xv341YyKKpjZiU,1033
 zenml/zen_server/dashboard/assets/@radix-C7hRs6Kx.js,sha256=cNA9UX8LGrKUQubGrls3E3Wq1fCrlK51W14yh22FE_U,296700
@@ -1052,7 +1052,7 @@ zenml/zen_server/routers/server_endpoints.py,sha256=Aci-7TB5VZylmasEYc2tVpG7Npgt
 zenml/zen_server/routers/service_accounts_endpoints.py,sha256=q-vFzzYYorqaOPUnESQU06h2I-BaZvIyPxCAhQ4gLPw,12086
 zenml/zen_server/routers/service_connectors_endpoints.py,sha256=CkwhnNI49_64lqG7aWHtyE41iqKlY-lumiOrWU7aDU0,17570
 zenml/zen_server/routers/service_endpoints.py,sha256=37CY-22h4hTscEkBByKJqdHcOWsP2lTGAiY2WgBPB0w,5566
-zenml/zen_server/routers/stack_components_endpoints.py,sha256=-3YE_8jZuAsQFwJxrZkDZnCwlSFh5dm_0_8vJ04a2xM,8234
+zenml/zen_server/routers/stack_components_endpoints.py,sha256=MBQ3iz29WN1n9T_BLpd2BtWxWoLIBP6QT7BfaAkGEMY,8580
 zenml/zen_server/routers/stack_deployment_endpoints.py,sha256=r7N4UaDAwj9ACYfdyUofRHpNPRrzWl-vmGoSErN9oYI,5381
 zenml/zen_server/routers/stacks_endpoints.py,sha256=ci-WZ774Q2T3-XN6cZARJ95wmRUW2-W1mh7gKlbNE54,7527
 zenml/zen_server/routers/steps_endpoints.py,sha256=KZpVsEsycLMmiEGnwT2EkmQvaTIIyL4QOUJBXMu7NSo,8239
@@ -1066,8 +1066,8 @@ zenml/zen_server/template_execution/__init__.py,sha256=79knXLKfegsvVSVSWecpqrepq
 zenml/zen_server/template_execution/runner_entrypoint_configuration.py,sha256=Y8aYJhqqs8Kv8I1q-dM1WemS5VBIfyoaaYH_YkzC7iY,1541
 zenml/zen_server/template_execution/utils.py,sha256=CiN7d8jCUTV5oKuoNF9Z_eDkitZrhL6V_Tz4CjU7em8,19410
 zenml/zen_server/template_execution/workload_manager_interface.py,sha256=CL9c7z8ajuZE01DaHmdCDCZmsroDcTarvN-nE8jv6qQ,2590
-zenml/zen_server/utils.py,sha256=NU1IrwT8GoYNt_s43nzmr9ZeUjkPW7iEB2NmJq5NvP8,20074
-zenml/zen_server/zen_server_api.py,sha256=1_YHOUmBulXuHWCxDLYtDCMxGze159UmU3f05643PSg,18182
+zenml/zen_server/utils.py,sha256=roJuND2KMVPue7fztVOKn4Tvoj1bvQ54TFpXgIa78Y0,20474
+zenml/zen_server/zen_server_api.py,sha256=Bhdn2AwGINEpKJiTL3QpC0EK8pvgGNK1IZNJr6E6ijE,23229
 zenml/zen_stores/__init__.py,sha256=6LTgH6XwDeDxKqVJ1JTfGhmS8II1NLopPloINGmdyI0,691
 zenml/zen_stores/base_zen_store.py,sha256=AplsW2NR-G9_CU54XvNTQJo4W0KJ5TJV22cjKW4n2BY,16124
 zenml/zen_stores/migrations/README.md,sha256=x04jsb6EOP6PBEGMQlDELiqKEham2O-iztAs9AylMFc,4898
@@ -1273,7 +1273,7 @@ zenml/zen_stores/migrations/versions/f49904a80aa7_increase_length_of_artifact_ta
 zenml/zen_stores/migrations/versions/f76a368a25a5_add_stack_description.py,sha256=u8fRomaasFeGhxvM2zU-Ab-AEpVsWm5zRcixxKFXdRw,904
 zenml/zen_stores/migrations/versions/fbd7f18ced1e_increase_step_run_field_lengths.py,sha256=kn-ng5EHe_mmLfffIFbz7T59z-to3oMx8III_4wOsz4,1956
 zenml/zen_stores/migrations/versions/ff538a321a92_migrate_onboarding_state.py,sha256=gsUFLJQ32_o9U35JCVqkqJVVk-zfq3yel25hXhzVFm4,3829
-zenml/zen_stores/rest_zen_store.py,sha256=d5JHoGeO_FSg55MmVNIFYmwN1JdAnJ2QbE6R590fKZU,158599
+zenml/zen_stores/rest_zen_store.py,sha256=ljz1iIc8szDh1dlX1129Q817gcDBuzFzQnw6hkSMYCg,159628
 zenml/zen_stores/schemas/__init__.py,sha256=4EXqExiVyxdnGxhQ_Hz79mOdRuMD0LsGlw0PaP2Ef6o,4333
 zenml/zen_stores/schemas/action_schemas.py,sha256=sv2J2TP12MeyGPQR2JsOPIivbPQ5OImg64exYS7CZBM,6496
 zenml/zen_stores/schemas/api_key_schemas.py,sha256=0pK7b9HlJuQL3DuKT4eGjFb87tyd4x-E2VyxJLpRv3o,7459
@@ -1316,11 +1316,11 @@ zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py,sha256=5err1a-TrV3SR5
 zenml/zen_stores/secrets_stores/secrets_store_interface.py,sha256=Q2Jbnt2Pp7NGlR-u1YBfRZV2g8su2Fd0ArBMdksAE-Q,2819
 zenml/zen_stores/secrets_stores/service_connector_secrets_store.py,sha256=S87ne23D08PAwtfRVlVnBn8R0ilTpEh6r8blauNV5WQ,6941
 zenml/zen_stores/secrets_stores/sql_secrets_store.py,sha256=LPFW757WCJLP1S8vrvjsrl2Tf1yo281xUTjSBsos4qk,8788
-zenml/zen_stores/sql_zen_store.py,sha256=oUB9u3tMVmRrSxslC1oEN5kHQa8i7g4bNEpJKgsYmqY,441911
+zenml/zen_stores/sql_zen_store.py,sha256=jsx_1_MJZE4dnNXxj8bAMX8pT5JfoJliADSBF9n6kJc,444809
 zenml/zen_stores/template_utils.py,sha256=GbJ7LgGVYHSCKPEA8RNTxPoVTWqpC77F_lGzjJ4O1Fw,9220
 zenml/zen_stores/zen_store_interface.py,sha256=fF_uL_FplnvGvM5o3jOQ8i1zHXhuhKLL2n4nvIKSR7E,92090
-zenml_nightly-0.82.1.dev20250522.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
-zenml_nightly-0.82.1.dev20250522.dist-info/METADATA,sha256=ie1RvBnJCcu9pLlEoRM3C4vp9wyEFwJzpONvj5P70YQ,24317
-zenml_nightly-0.82.1.dev20250522.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-zenml_nightly-0.82.1.dev20250522.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
-zenml_nightly-0.82.1.dev20250522.dist-info/RECORD,,
+zenml_nightly-0.82.1.dev20250524.dist-info/LICENSE,sha256=wbnfEnXnafPbqwANHkV6LUsPKOtdpsd-SNw37rogLtc,11359
+zenml_nightly-0.82.1.dev20250524.dist-info/METADATA,sha256=OFFcRSYChQvW4vCBOO5FJDfx8UkZbhmq9DKwtmBN2s0,24317
+zenml_nightly-0.82.1.dev20250524.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+zenml_nightly-0.82.1.dev20250524.dist-info/entry_points.txt,sha256=QK3ETQE0YswAM2mWypNMOv8TLtr7EjnqAFq1br_jEFE,43
+zenml_nightly-0.82.1.dev20250524.dist-info/RECORD,,