zenml-nightly 0.70.0.dev20241122__tar.gz → 0.70.0.dev20241124__tar.gz

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.70.0.dev20241122
+Version: 0.70.0.dev20241124
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.70.0.dev20241122"
+version = "0.70.0.dev20241124"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.70.0.dev20241124/src/zenml/VERSION

@@ -0,0 +1 @@
+0.70.0.dev20241124

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/constants.py

@@ -353,6 +353,7 @@ FLAVORS = "/flavors"
 GET_OR_CREATE = "/get-or-create"
 HEALTH = "/health"
 INFO = "/info"
+LOAD_INFO = "/load-info"
 LOGIN = "/login"
 LOGOUT = "/logout"
 LOGS = "/logs"

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -42,8 +42,8 @@ from kubernetes import config as k8s_config
 from kubernetes.client.rest import ApiException
 
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
-    build_cluster_role_binding_manifest_for_service_account,
     build_namespace_manifest,
+    build_role_binding_manifest_for_service_account,
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
@@ -288,7 +288,7 @@ def create_edit_service_account(
     rbac_api: k8s_client.RbacAuthorizationV1Api,
     service_account_name: str,
     namespace: str,
-    cluster_role_binding_name: str = "zenml-edit",
+    role_binding_name: str = "zenml-edit",
 ) -> None:
     """Create a new Kubernetes service account with "edit" rights.
 
@@ -297,16 +297,17 @@ def create_edit_service_account(
         rbac_api: Client of Rbac Authorization V1 API of Kubernetes API.
         service_account_name: Name of the service account.
         namespace: Kubernetes namespace. Defaults to "default".
-        cluster_role_binding_name: Name of the cluster role binding.
-            Defaults to "zenml-edit".
+        role_binding_name: Name of the role binding. Defaults to "zenml-edit".
     """
-    crb_manifest = build_cluster_role_binding_manifest_for_service_account(
-        name=cluster_role_binding_name,
+    rb_manifest = build_role_binding_manifest_for_service_account(
+        name=role_binding_name,
         role_name="edit",
         service_account_name=service_account_name,
         namespace=namespace,
     )
-    _if_not_exists(rbac_api.create_cluster_role_binding)(body=crb_manifest)
+    _if_not_exists(rbac_api.create_namespaced_role_binding)(
+        namespace=namespace, body=rb_manifest
+    )
 
     sa_manifest = build_service_account_manifest(
         name=service_account_name, namespace=namespace

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -59,6 +59,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_cron_job_manifest,
     build_pod_manifest,
 )
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
@@ -328,6 +329,45 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             custom_validation_function=_validate_local_requirements,
         )
 
+    @classmethod
+    def apply_default_resource_requests(
+        cls,
+        memory: str,
+        cpu: Optional[str] = None,
+        pod_settings: Optional[KubernetesPodSettings] = None,
+    ) -> KubernetesPodSettings:
+        """Applies default resource requests to a pod settings object.
+
+        Args:
+            memory: The memory resource request.
+            cpu: The CPU resource request.
+            pod_settings: The pod settings to update. A new one will be created
+                if not provided.
+
+        Returns:
+            The new or updated pod settings.
+        """
+        resources = {
+            "requests": {"memory": memory},
+        }
+        if cpu:
+            resources["requests"]["cpu"] = cpu
+        if not pod_settings:
+            pod_settings = KubernetesPodSettings(resources=resources)
+        elif not pod_settings.resources:
+            # We can't update the pod settings in place (because it's a frozen
+            # pydantic model), so we have to create a new one.
+            pod_settings = KubernetesPodSettings(
+                **pod_settings.model_dump(exclude_unset=True),
+                resources=resources,
+            )
+        else:
+            set_requests = pod_settings.resources.get("requests", {})
+            resources["requests"].update(set_requests)
+            pod_settings.resources["requests"] = resources["requests"]
+
+        return pod_settings
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -422,6 +462,17 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             )
             return
 
+        # We set some default minimum resource requests for the orchestrator pod
+        # here if the user has not specified any, because the orchestrator pod
+        # takes up some memory resources itself and, if not specified, the pod
+        # will be scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        orchestrator_pod_settings = self.apply_default_resource_requests(
+            memory="400Mi",
+            cpu="100m",
+            pod_settings=settings.orchestrator_pod_settings,
+        )
+
         # Create and run the orchestrator pod.
         pod_manifest = build_pod_manifest(
             run_name=orchestrator_run_name,
@@ -431,7 +482,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             command=command,
             args=args,
             privileged=False,
-            pod_settings=settings.orchestrator_pod_settings,
+            pod_settings=orchestrator_pod_settings,
             service_account_name=service_account_name,
             env=environment,
             mount_local_stores=self.config.is_local,

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -116,6 +116,16 @@ def main() -> None:
         env = get_config_environment_vars()
         env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
 
+        # We set some default minimum memory resource requests for the step pod
+        # here if the user has not specified any, because the step pod takes up
+        # some memory resources itself and, if not specified, the pod will be
+        # scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        pod_settings = KubernetesOrchestrator.apply_default_resource_requests(
+            memory="400Mi",
+            pod_settings=settings.pod_settings,
+        )
+
         # Define Kubernetes pod manifest.
         pod_manifest = build_pod_manifest(
             pod_name=pod_name,
@@ -126,7 +136,7 @@ def main() -> None:
             args=step_args,
             env=env,
             privileged=settings.privileged,
-            pod_settings=settings.pod_settings,
+            pod_settings=pod_settings,
             service_account_name=settings.step_pod_service_account_name
             or settings.service_account_name,
             mount_local_stores=mount_local_stores,

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -304,13 +304,13 @@ def build_cron_job_manifest(
     return job_manifest
 
 
-def build_cluster_role_binding_manifest_for_service_account(
+def build_role_binding_manifest_for_service_account(
     name: str,
     role_name: str,
     service_account_name: str,
     namespace: str = "default",
 ) -> Dict[str, Any]:
-    """Build a manifest for a cluster role binding of a service account.
+    """Build a manifest for a role binding of a service account.
 
     Args:
         name: Name of the cluster role binding.
@@ -323,7 +323,7 @@ def build_cluster_role_binding_manifest_for_service_account(
     """
     return {
         "apiVersion": "rbac.authorization.k8s.io/v1",
-        "kind": "ClusterRoleBinding",
+        "kind": "RoleBinding",
         "metadata": {"name": name},
         "subjects": [
             {

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py

@@ -26,6 +26,7 @@ from typing import Any, List, Optional
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
 from pydantic import Field
+from urllib3.exceptions import HTTPError
 
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.exceptions import AuthorizationException
@@ -572,7 +573,7 @@ class KubernetesServiceConnector(ServiceConnector):
                 auth_settings=["BearerToken"],
                 response_type="VersionInfo",
             )
-        except k8s_client.ApiException as err:
+        except (k8s_client.ApiException, HTTPError) as err:
             raise AuthorizationException(
                 f"failed to verify Kubernetes cluster access: {err}"
             ) from err

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/models/__init__.py

@@ -328,7 +328,11 @@ from zenml.models.v2.misc.service_connector_type import (
     ServiceConnectorTypedResourcesModel,
     ResourceTypeModel,
 )
-from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel
+from zenml.models.v2.misc.server_models import (
+    ServerDatabaseType,
+    ServerLoadInfo,
+    ServerModel,
+)
 from zenml.models.v2.core.trigger import (
     TriggerRequest,
     TriggerFilter,
@@ -731,6 +735,7 @@ __all__ = [
     "ExternalUserModel",
     "BuildItem",
     "LoadedVisualization",
+    "ServerLoadInfo",
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/models/v2/core/service_connector.py

@@ -488,6 +488,10 @@ class ServiceConnectorResponse(
 ):
     """Response model for service connectors."""
 
+    # Disable the warning for updating responses, because we update the
+    # service connector type in place
+    _warn_on_response_updates: bool = False
+
     name: str = Field(
         title="The service connector name.",
         max_length=STR_FIELD_MAX_LENGTH,

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/models/v2/misc/server_models.py

@@ -118,3 +118,26 @@ class ServerModel(BaseModel):
         # Local ZenML servers are identifiable by the fact that their
         # server ID is the same as the local client (user) ID.
         return self.id == GlobalConfiguration().user_id
+
+
+class ServerLoadInfo(BaseModel):
+    """Domain model for ZenML server load information."""
+
+    threads: int = Field(
+        title="Number of threads that the server is currently using."
+    )
+
+    db_connections_total: int = Field(
+        title="Total number of database connections (active and idle) that the "
+        "server currently has established."
+    )
+
+    db_connections_active: int = Field(
+        title="Number of database connections that the server is currently "
+        "actively using to make queries or transactions."
+    )
+
+    db_connections_overflow: int = Field(
+        title="Number of overflow database connections that the server is "
+        "currently actively using to make queries or transactions."
+    )

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/zen_server/deploy/helm/templates/_environment.tpl

@@ -3,6 +3,123 @@ Helpers for environment variables configured in ZenML deployments and secrets st
 */}}
 
 
+{{/*
+ZenML store configuration options (non-secret values).
+
+This template constructs a dictionary that is similar to the python values that
+can be configured in the zenml.zen_store.sql_zen_store.SqlZenStoreConfiguration
+class. Only non-secret values are included in this dictionary.
+
+The dictionary is then converted into deployment environment variables by other
+templates and inserted where it is needed.
+
+The input is taken from a .ZenML dict that is passed to the template and
+contains the values configured in the values.yaml file for the ZenML server.
+
+Args:
+  .ZenML: A dictionary with the ZenML configuration values configured for the
+  ZenML server.
+Returns:
+  A dictionary with the non-secret values configured for the ZenML store.
+*/}}
+{{- define "zenml.storeConfigurationAttrs" -}}
+{{- if .ZenML.database.url }}
+type: sql
+ssl_verify_server_cert: {{ .ZenML.database.sslVerifyServerCert | default "false" | quote }}
+{{- if .ZenML.database.backupStrategy }}
+backup_strategy: {{ .ZenML.database.backupStrategy | quote }}
+{{- if eq .ZenML.database.backupStrategy "database" }}
+backup_database: {{ .ZenML.database.backupDatabase | quote }}
+{{- else if eq .ZenML.database.backupStrategy "dump-file" }}
+backup_directory: "/backups"
+{{- end }}
+{{- end }}
+{{- if .ZenML.database.poolSize }}
+pool_size: {{ .ZenML.database.poolSize | quote }}
+{{- end }}
+{{- if .ZenML.database.maxOverflow }}
+max_overflow: {{ .ZenML.database.maxOverflow | quote }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+
+{{/*
+ZenML store configuration options (secret values).
+
+This template constructs a dictionary that is similar to the python values that
+can be configured in the zenml.zen_store.sql_zen_store.SqlZenStoreConfiguration
+class. Only secret values are included in this dictionary.
+
+The dictionary is then converted into deployment environment variables by other
+templates and inserted where it is needed.
+
+The input is taken from a .ZenML dict that is passed to the template and
+contains the values configured in the values.yaml file for the ZenML server.
+
+Args:
+  .ZenML: A dictionary with the ZenML configuration values configured for the
+  ZenML server.
+Returns:
+  A dictionary with the secret values configured for the ZenML store.
+*/}}
+{{- define "zenml.storeSecretConfigurationAttrs" -}}
+{{- if .ZenML.database.url }}
+url: {{ .ZenML.database.url | quote }}
+{{- if .ZenML.database.sslCa }}
+ssl_ca: {{ .Files.Get .ZenML.database.sslCa }}
+{{- end }}
+{{- if .ZenML.database.sslCert }}
+ssl_cert: {{ .Files.Get .ZenML.database.sslCert }}
+{{- end }}
+{{- if .ZenML.database.sslKey }}
+ssl_key: {{ .Files.Get .ZenML.database.sslKey }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+
+{{/*
+Store configuration environment variables (non-secret values).
+
+Passes the .Values.zenml dict as input to the `zenml.storeConfigurationAttrs`
+template and converts the output into a dictionary of environment variables that
+need to be configured for the store.
+
+Args:
+  .Values: The values.yaml file for the ZenML deployment.
+Returns:
+  A dictionary with the non-secret environment variables that are configured for
+  the store (i.e. keys starting with `ZENML_STORE_`).
+*/}}
+{{- define "zenml.storeEnvVariables" -}}
+{{ $zenml := dict "ZenML" .Values.zenml }}
+{{- range $k, $v := include "zenml.storeConfigurationAttrs" $zenml | fromYaml }}
+ZENML_STORE_{{ $k | upper }}: {{ $v | quote }}
+{{- end }}
+{{- end }}
+
+
+{{/*
+Store configuration environment variables (secret values).
+
+Passes the .Values.zenml dict as input to the `zenml.storeSecretConfigurationAttrs`
+template and converts the output into a dictionary of environment variables that
+need to be configured for the store.
+
+Args:
+  .Values: The values.yaml file for the ZenML deployment.
+Returns:
+  A dictionary with the secret environment variables that are configured for
+  the store (i.e. keys starting with `ZENML_STORE_`).
+*/}}
+{{- define "zenml.storeSecretEnvVariables" -}}
+{{ $zenml := dict "ZenML" .Values.zenml }}
+{{- range $k, $v := include "zenml.storeSecretConfigurationAttrs" $zenml | fromYaml }}
+ZENML_STORE_{{ $k | upper }}: {{ $v | quote }}
+{{- end }}
+{{- end }}
+
 {{/*
 ZenML server configuration options (non-secret values).
 

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/zen_server/deploy/helm/templates/server-db-job.yaml

@@ -79,20 +79,9 @@ spec:
             {{- end }}
             - name: ZENML_DEFAULT_PROJECT_NAME
               value: {{ .Values.zenml.defaultProject | quote }}
-            - name: ZENML_STORE_TYPE
-              value: sql
-            - name: ZENML_STORE_SSL_VERIFY_SERVER_CERT
-              value: {{ .Values.zenml.database.sslVerifyServerCert | default "false" | quote }}
-            {{- if .Values.zenml.database.backupStrategy }}
-            - name: ZENML_STORE_BACKUP_STRATEGY
-              value: {{ .Values.zenml.database.backupStrategy | quote }}
-            {{- if eq .Values.zenml.database.backupStrategy "database" }}
-            - name: ZENML_STORE_BACKUP_DATABASE
-              value: {{ .Values.zenml.database.backupDatabase | quote }}
-            {{- else if eq .Values.zenml.database.backupStrategy "dump-file" }}
-            - name: ZENML_STORE_BACKUP_DIRECTORY
-              value: /backups
-            {{- end }}
+            {{- range $k, $v := include "zenml.storeEnvVariables" . | fromYaml }}
+            - name: {{ $k }}
+              value: {{ $v | quote }}
             {{- end }}
             {{- range $k, $v := include "zenml.serverEnvVariables" . | fromYaml }}
             - name: {{ $k }}

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/zen_server/deploy/helm/templates/server-deployment.yaml

@@ -63,14 +63,14 @@ spec:
               value: "True"
             {{- end }}
             {{- if .Values.zenml.database.url }}
-            - name: ZENML_STORE_TYPE
-              value: sql
             - name: DISABLE_DATABASE_MIGRATION
               value: "True"
-            - name: ZENML_STORE_SSL_VERIFY_SERVER_CERT
-              value: {{ .Values.zenml.database.sslVerifyServerCert | default "false" | quote }}
             {{- end }}
 
+            {{- range $k, $v := include "zenml.storeEnvVariables" . | fromYaml }}
+            - name: {{ $k }}
+              value: {{ $v | quote }}
+            {{- end }}
             {{- range $k, $v := include "zenml.serverEnvVariables" . | fromYaml }}
             - name: {{ $k }}
               value: {{ $v | quote }}
@@ -104,6 +104,18 @@ spec:
             httpGet:
               path: /health
               port: http
+          lifecycle:
+            preStop:
+              exec:
+                # Give the process 15 more seconds before the SIGTERM signal is
+                # sent. This allows the endpoint removal to reach the ingress
+                # controller in time and for traffic to be routed away from the
+                # pod before it is shut down. This eliminates the number of 502
+                # errors returned to the user.
+                #
+                # See https://learnk8s.io/graceful-shutdown for more information.
+                #
+                command: ["sleep", "15"]
           resources:
             {{- toYaml .Values.resources | nindent 12 }}
       {{- with .Values.nodeSelector }}

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/zen_server/deploy/helm/templates/server-secret.yaml

@@ -10,23 +10,8 @@ data:
   {{- else }}
   ZENML_SERVER_JWT_SECRET_KEY: {{ $prevServerSecret.data.ZENML_SERVER_JWT_SECRET_KEY | default (randAlphaNum 32 | b64enc | quote) }}
   {{- end }}
-  {{- if .Values.zenml.database.url }}
-  ZENML_STORE_URL: {{ .Values.zenml.database.url | b64enc | quote }}
-  {{- if .Values.zenml.database.sslCa }}
-  ZENML_STORE_SSL_CA: {{ .Files.Get .Values.zenml.database.sslCa | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslCert }}
-  ZENML_STORE_SSL_CERT: {{ .Files.Get .Values.zenml.database.sslCert | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.sslKey }}
-  ZENML_STORE_SSL_KEY: {{ .Files.Get .Values.zenml.database.sslKey | b64enc }}
-  {{- end }}
-  {{- if .Values.zenml.database.poolSize }}
-  ZENML_STORE_POOL_SIZE: {{ .Values.zenml.database.poolSize | b64enc | quote }}
-  {{- end }}
-  {{- if .Values.zenml.database.maxOverflow }}
-  ZENML_STORE_MAX_OVERFLOW: {{ .Values.zenml.database.maxOverflow | b64enc | quote }}
-  {{- end }}
+  {{- range $k, $v := include "zenml.storeSecretEnvVariables" . | fromYaml}}
+  {{ $k }}: {{ $v | b64enc | quote }}
   {{- end }}
   {{- range $k, $v := include "zenml.secretsStoreSecretEnvVariables" . | fromYaml}}
   {{ $k }}: {{ $v | b64enc | quote }}

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/zen_server/routers/server_endpoints.py

@@ -22,6 +22,7 @@ from zenml.constants import (
     ACTIVATE,
     API,
     INFO,
+    LOAD_INFO,
     ONBOARDING_STATE,
     SERVER_SETTINGS,
     VERSION_1,
@@ -30,6 +31,7 @@ from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
     ServerActivationRequest,
+    ServerLoadInfo,
     ServerModel,
     ServerSettingsResponse,
     ServerSettingsUpdate,
@@ -71,6 +73,51 @@ def server_info() -> ServerModel:
     return zen_store().get_store_info()
 
 
+@router.get(
+    LOAD_INFO,
+    response_model=ServerLoadInfo,
+)
+@handle_exceptions
+def server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo:
+    """Get information about the server load.
+
+    Returns:
+        Information about the server load.
+    """
+    import threading
+
+    # Get the current number of threads
+    num_threads = len(threading.enumerate())
+
+    store = zen_store()
+
+    if store.config.driver == "sqlite":
+        # SQLite doesn't have a connection pool
+        return ServerLoadInfo(
+            threads=num_threads,
+            db_connections_total=0,
+            db_connections_active=0,
+            db_connections_overflow=0,
+        )
+
+    from sqlalchemy.pool import QueuePool
+
+    # Get the number of connections
+    pool = store.engine.pool
+    assert isinstance(pool, QueuePool)
+    idle_conn = pool.checkedin()
+    active_conn = pool.checkedout()
+    overflow_conn = max(0, pool.overflow())
+    total_conn = idle_conn + active_conn
+
+    return ServerLoadInfo(
+        threads=num_threads,
+        db_connections_total=total_conn,
+        db_connections_active=active_conn,
+        db_connections_overflow=overflow_conn,
+    )
+
+
 @router.get(
     ONBOARDING_STATE,
     responses={

{zenml_nightly-0.70.0.dev20241122 → zenml_nightly-0.70.0.dev20241124}/src/zenml/zen_server/zen_server_api.py

@@ -180,7 +180,15 @@ class RequestBodyLimit(BaseHTTPMiddleware):
         if content_length := request.headers.get("content-length"):
             if int(content_length) > self.max_bytes:
                 return Response(status_code=413)  # Request Entity Too Large
-        return await call_next(request)
+
+        try:
+            return await call_next(request)
+        except Exception:
+            logger.exception("An error occurred while processing the request")
+            return JSONResponse(
+                status_code=500,
+                content={"detail": "An unexpected error occurred."},
+            )
 
 
 class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
@@ -220,7 +228,15 @@ class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
                         "detail": "File uploads are not allowed on this endpoint."
                     },
                 )
-        return await call_next(request)
+
+        try:
+            return await call_next(request)
+        except Exception:
+            logger.exception("An error occurred while processing the request")
+            return JSONResponse(
+                status_code=500,
+                content={"detail": "An unexpected error occurred."},
+            )
 
 
 ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()
@@ -252,13 +268,21 @@ async def set_secure_headers(request: Request, call_next: Any) -> Any:
     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 await call_next(request)
+        return response
 
-    response = await call_next(request)
     secure_headers().framework.fastapi(response)
     return response
 
@@ -298,7 +322,15 @@ async def track_last_user_activity(request: Request, call_next: Any) -> Any:
         zen_store()._update_last_user_activity_timestamp(
             last_user_activity=last_user_activity
         )
-    return await call_next(request)
+
+    try:
+        return await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
 
 
 @app.middleware("http")
@@ -330,7 +362,14 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
         )
         source_context.set(SourceContextTypes.API)
 
-    return await call_next(request)
+    try:
+        return await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
 
 
 @app.on_event("startup")