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

zenml/zen_server/utils.py

@@ -15,10 +15,12 @@
 
 import inspect
 import os
+import threading
 from functools import wraps
 from typing import (
     TYPE_CHECKING,
     Any,
+    Awaitable,
     Callable,
     Dict,
     List,
@@ -27,11 +29,11 @@ from typing import (
     Type,
     TypeVar,
     Union,
-    cast,
 )
 from uuid import UUID
 
 from pydantic import BaseModel, ValidationError
+from typing_extensions import ParamSpec
 
 from zenml import __version__ as zenml_version
 from zenml.config.global_config import GlobalConfiguration
@@ -64,6 +66,11 @@ if TYPE_CHECKING:
         BoundedThreadPoolExecutor,
     )
 
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
 logger = get_logger(__name__)
 
 _zen_store: Optional["SqlZenStore"] = None
@@ -298,11 +305,16 @@ def server_config() -> ServerConfiguration:
     return _server_config
 
 
-F = TypeVar("F", bound=Callable[..., Any])
-
+def async_fastapi_endpoint_wrapper(
+    func: Callable[P, R],
+) -> Callable[P, Awaitable[Any]]:
+    """Decorator for FastAPI endpoints.
 
-def handle_exceptions(func: F) -> F:
-    """Decorator to handle exceptions in the API.
+    This decorator for FastAPI endpoints does the following:
+    - Sets the auth_context context variable if the endpoint is authenticated.
+    - Converts exceptions to HTTPExceptions with the correct status code.
+    - Converts the sync endpoint function to an coroutine and runs the original
+      function in a worker threadpool. See below for more details.
 
     Args:
         func: Function to decorate.
@@ -311,41 +323,68 @@ def handle_exceptions(func: F) -> F:
         Decorated function.
     """
 
+    # 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
+    # another task in the same threadpool to serialize the response. If there
+    # are many tasks already in the queue, this means that the response
+    # serialization will wait for a long time instead of returning the response
+    # immediately. By making our endpoints async and then immediately
+    # dispatching them to the threadpool ourselves (which is essentially what
+    # FastAPI does when having a sync endpoint), we can avoid this problem.
+    # The serialization logic will now run on the event loop and not wait for
+    # a worker thread to become available.
+    # See: `fastapi.routing.serialize_response(...)` and
+    # https://github.com/fastapi/fastapi/pull/888 for more information.
     @wraps(func)
-    def decorated(*args: Any, **kwargs: Any) -> 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
+    async def async_decorated(*args: P.args, **kwargs: P.kwargs) -> Any:
+        from starlette.concurrency import run_in_threadpool
 
-        from zenml.zen_server.auth import AuthContext, set_auth_context
+        from zenml.zen_server.zen_server_api import request_ids
 
-        for arg in args:
-            if isinstance(arg, AuthContext):
-                set_auth_context(arg)
-                break
-        else:
-            for _, arg in kwargs.items():
+        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
 
-        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 cast(F, decorated)
+        return await run_in_threadpool(decorated, *args, **kwargs)
+
+    return async_decorated
 
 
 # Code from https://github.com/tiangolo/fastapi/issues/1474#issuecomment-1160633178

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
         )