arize-phoenix 11.3.0__py3-none-any.whl → 11.5.0__py3-none-any.whl

phoenix/server/daemons/db_disk_usage_monitor.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import logging
+from asyncio import sleep
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+import sqlalchemy as sa
+from email_validator import EmailNotValidError, validate_email
+from sqlalchemy import text
+from typing_extensions import assert_never
+
+from phoenix.config import (
+    get_env_database_allocated_storage_capacity_gibibytes,
+    get_env_database_usage_email_warning_threshold_percentage,
+    get_env_database_usage_insertion_blocking_threshold_percentage,
+)
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect
+from phoenix.server.email.types import DbUsageWarningEmailSender
+from phoenix.server.prometheus import (
+    DB_DISK_USAGE_BYTES,
+    DB_DISK_USAGE_RATIO,
+    DB_DISK_USAGE_WARNING_EMAIL_ERRORS,
+    DB_DISK_USAGE_WARNING_EMAILS_SENT,
+    DB_INSERTIONS_BLOCKED,
+)
+from phoenix.server.types import DaemonTask, DbSessionFactory
+
+logger = logging.getLogger(__name__)
+
+_SLEEP_SECONDS = 60
+_EMAIL_FREQUENCY_HOURS = 24
+_BYTES_PER_GIBIBYTE = 1024**3
+
+
+class DbDiskUsageMonitor(DaemonTask):
+    """
+    Monitors database disk space usage and triggers warnings/blocking when thresholds are exceeded.
+
+    This daemon:
+    - Periodically checks current database size
+    - Compares usage against configured thresholds
+    - Sends warning emails to admins when warning threshold is reached
+    - Toggles insertion blocking when blocking threshold is reached
+    """
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        email_sender: Optional[DbUsageWarningEmailSender] = None,
+    ) -> None:
+        super().__init__()
+        self._db = db
+        self._email_sender = email_sender
+        # Tracks last email send time per admin email address to prevent spam
+        self._last_email_sent: dict[str, datetime] = {}
+
+    @property
+    def _is_disabled(self) -> bool:
+        return not bool(
+            get_env_database_allocated_storage_capacity_gibibytes()
+            and (
+                get_env_database_usage_email_warning_threshold_percentage()
+                or get_env_database_usage_insertion_blocking_threshold_percentage()
+            )
+        )
+
+    async def _run(self) -> None:
+        if self._is_disabled:
+            return
+
+        while self._running:
+            try:
+                current_usage_bytes = await self._check_disk_usage_bytes()
+            except Exception:
+                logger.exception("Failed to check disk space")
+            else:
+                DB_DISK_USAGE_BYTES.set(current_usage_bytes)
+                current_usage_gibibytes = current_usage_bytes / _BYTES_PER_GIBIBYTE
+                try:
+                    await self._check_thresholds(current_usage_gibibytes)
+                except Exception:
+                    logger.exception("Failed to check database usage thresholds")
+            await sleep(_SLEEP_SECONDS)
+
+    async def _check_disk_usage_bytes(self) -> float:
+        if self._db.dialect is SupportedSQLDialect.SQLITE:
+            async with self._db() as session:
+                page_count = await session.scalar(text("PRAGMA page_count;"))
+                freelist_count = await session.scalar(text("PRAGMA freelist_count;"))
+                page_size = await session.scalar(text("PRAGMA page_size;"))
+            current_usage_bytes = (page_count - freelist_count) * page_size
+        elif self._db.dialect is SupportedSQLDialect.POSTGRESQL:
+            async with self._db() as session:
+                current_usage_bytes = await session.scalar(
+                    text("SELECT pg_database_size(current_database());")
+                )
+        else:
+            assert_never(self._db.dialect)
+        if not isinstance(current_usage_bytes, (int, float)):
+            raise TypeError(f"Expected int or float, got {type(current_usage_bytes)}")
+        return float(current_usage_bytes)
+
+    async def _check_thresholds(self, current_usage_gibibytes: float) -> None:
+        allocated_capacity_gibibytes = get_env_database_allocated_storage_capacity_gibibytes()
+        if not allocated_capacity_gibibytes:
+            return
+
+        used_ratio = current_usage_gibibytes / allocated_capacity_gibibytes
+        DB_DISK_USAGE_RATIO.set(used_ratio)
+        used_percentage = used_ratio * 100
+
+        # Check insertion blocking threshold
+        if (
+            insertion_blocking_threshold_percentage
+            := get_env_database_usage_insertion_blocking_threshold_percentage()
+        ):
+            should_not_insert_or_update = used_percentage > insertion_blocking_threshold_percentage
+            self._db.should_not_insert_or_update = should_not_insert_or_update
+            DB_INSERTIONS_BLOCKED.set(int(should_not_insert_or_update))
+
+        # Check warning email threshold
+        if (
+            notification_threshold_percentage
+            := get_env_database_usage_email_warning_threshold_percentage()
+        ):
+            if used_percentage > notification_threshold_percentage:
+                await self._send_warning_emails(
+                    used_percentage,
+                    allocated_capacity_gibibytes,
+                    notification_threshold_percentage,
+                )
+
+    async def _send_warning_emails(
+        self,
+        used_percentage: float,
+        allocated_capacity_gibibytes: float,
+        notification_threshold_percentage: float,
+    ) -> None:
+        if not self._email_sender:
+            return
+
+        current_usage_gibibytes = used_percentage / 100 * allocated_capacity_gibibytes
+        stmt = (
+            sa.select(models.User.email)
+            .join(models.UserRole)
+            .where(models.UserRole.name == "ADMIN")
+        )
+
+        try:
+            async with self._db() as session:
+                admin_emails = (await session.scalars(stmt)).all()
+        except Exception:
+            logger.exception(
+                "Failed to fetch admin emails from database, "
+                "skipping database usage warning emails"
+            )
+            return
+
+        if not admin_emails:
+            return
+
+        # Validate email addresses
+        valid_emails: list[str] = []
+
+        for email in admin_emails:
+            try:
+                normalized_email = validate_email(email, check_deliverability=False).normalized
+            except EmailNotValidError:
+                pass
+            else:
+                valid_emails.append(normalized_email)
+
+        if not valid_emails:
+            return
+
+        self._last_email_sent = {
+            email: timestamp
+            for email, timestamp in self._last_email_sent.items()
+            if email in valid_emails
+        }
+
+        now = datetime.now(timezone.utc)
+        emails_sent = 0
+        send_attempts = 0
+
+        for email in valid_emails:
+            if email in self._last_email_sent and now - self._last_email_sent[email] < timedelta(
+                hours=_EMAIL_FREQUENCY_HOURS
+            ):
+                continue
+            send_attempts += 1
+            try:
+                await self._email_sender.send_db_usage_warning_email(
+                    email=email,
+                    current_usage_gibibytes=current_usage_gibibytes,
+                    allocated_storage_gibibytes=allocated_capacity_gibibytes,
+                    notification_threshold_percentage=notification_threshold_percentage,
+                )
+            except Exception:
+                logger.exception(f"Failed to send database usage warning email to {email}")
+                # Count email send errors
+                DB_DISK_USAGE_WARNING_EMAIL_ERRORS.inc()
+            else:
+                self._last_email_sent[email] = now
+                emails_sent += 1
+                # Count successful warning email sends
+                DB_DISK_USAGE_WARNING_EMAILS_SENT.inc()

phoenix/server/email/sender.py

@@ -81,6 +81,31 @@ class SimpleEmailSender:
 
         await to_thread.run_sync(self._send_email, msg)
 
+    async def send_db_usage_warning_email(
+        self,
+        email: str,
+        current_usage_gibibytes: float,
+        allocated_storage_gibibytes: float,
+        notification_threshold_percentage: float,
+    ) -> None:
+        subject = "[Phoenix] Database Usage Threshold Exceeded"
+        template_name = "db_disk_usage_notification.html"
+
+        template = self.env.get_template(template_name)
+        html_content = template.render(
+            current_usage_gibibytes=current_usage_gibibytes,
+            allocated_storage_gibibytes=allocated_storage_gibibytes,
+            notification_threshold_percentage=notification_threshold_percentage,
+        )
+
+        msg = EmailMessage()
+        msg["Subject"] = subject
+        msg["From"] = self.sender_email
+        msg["To"] = email
+        msg.set_content(html_content, subtype="html")
+
+        await to_thread.run_sync(self._send_email, msg)
+
     def _send_email(self, msg: EmailMessage) -> None:
         context: ssl.SSLContext
         if self.validate_certs:

phoenix/server/email/templates/db_disk_usage_notification.html

@@ -0,0 +1,16 @@
+<!DOCTYPE html>
+<html>
+<head>
+  <meta charset="UTF-8" />
+  <title>Database Usage Notification</title>
+</head>
+<body>
+<h1>Database Usage Notification</h1>
+<p>Your Phoenix database usage has exceeded the notification threshold.</p>
+<p><strong>Current Usage:</strong> {{ current_usage_gibibytes|round(1) }} GiB</p>
+<p><strong>Allocated Storage:</strong> {{ allocated_storage_gibibytes|round(1) }} GiB</p>
+<p><strong>Usage Percentage:</strong> {{ ((current_usage_gibibytes / allocated_storage_gibibytes) * 100)|round(1) }}%</p>
+<p><strong>Notification Threshold:</strong> {{ notification_threshold_percentage }}%</p>
+<p>Please consider removing old data or increasing your storage allocation to prevent interruption.</p>
+</body>
+</html>

phoenix/server/email/types.py

@@ -19,8 +19,19 @@ class PasswordResetEmailSender(Protocol):
     ) -> None: ...
 
 
+class DbUsageWarningEmailSender(Protocol):
+    async def send_db_usage_warning_email(
+        self,
+        email: str,
+        current_usage_gibibytes: float,
+        allocated_storage_gibibytes: float,
+        notification_threshold_percentage: float,
+    ) -> None: ...
+
+
 class EmailSender(
     WelcomeEmailSender,
     PasswordResetEmailSender,
+    DbUsageWarningEmailSender,
     Protocol,
 ): ...

phoenix/server/grpc_server.py

@@ -1,5 +1,5 @@
 from collections.abc import Awaitable, Callable
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Iterable, Optional
 
 import grpc
 from grpc.aio import RpcContext, Server, ServerInterceptor
@@ -61,7 +61,7 @@ class GrpcServer:
         enable_prometheus: bool = False,
         disabled: bool = False,
         token_store: Optional[CanReadToken] = None,
-        interceptors: list[ServerInterceptor] = [],
+        interceptors: Iterable[ServerInterceptor] = (),
     ) -> None:
         self._callback = callback
         self._server: Optional[Server] = None
@@ -69,7 +69,7 @@ class GrpcServer:
         self._enable_prometheus = enable_prometheus
         self._disabled = disabled
         self._token_store = token_store
-        self._interceptors = interceptors
+        self._interceptors = list(interceptors)
 
     async def __aenter__(self) -> None:
         interceptors = self._interceptors

phoenix/server/prometheus.py

@@ -73,6 +73,28 @@ JWT_STORE_API_KEYS_ACTIVE = Gauge(
     documentation="Current number of API keys in the JWT store",
 )
 
+DB_DISK_USAGE_BYTES = Gauge(
+    name="database_disk_usage_bytes",
+    documentation="Current database disk usage in bytes",
+)
+DB_DISK_USAGE_RATIO = Gauge(
+    name="database_disk_usage_ratio",
+    documentation="Current database disk usage as ratio of allocated capacity (0-1)",
+)
+DB_INSERTIONS_BLOCKED = Gauge(
+    name="database_insertions_blocked",
+    documentation="Whether database insertions are currently blocked due to disk usage "
+    "(1 = blocked, 0 = not blocked)",
+)
+DB_DISK_USAGE_WARNING_EMAILS_SENT = Counter(
+    name="database_disk_usage_warning_emails_sent_total",
+    documentation="Total count of database disk usage warning emails sent",
+)
+DB_DISK_USAGE_WARNING_EMAIL_ERRORS = Counter(
+    name="database_disk_usage_warning_email_errors_total",
+    documentation="Total count of database disk usage warning email send errors",
+)
+
 
 class PrometheusMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:

phoenix/server/static/.vite/manifest.json

@@ -1,28 +1,28 @@
 {
-  "_components-DlL7ybQ2.js": {
-    "file": "assets/components-DlL7ybQ2.js",
+  "_components-Bwf6zNbg.js": {
+    "file": "assets/components-Bwf6zNbg.js",
     "name": "components",
     "imports": [
-      "_vendor-DqQvHbPa.js",
-      "_pages-B7wCtpad.js",
-      "_vendor-arizeai-CLX44PFA.js",
-      "_vendor-codemirror-Du3XyJnB.js",
+      "_vendor-DRWIRkSJ.js",
+      "_pages-BCR8hW_l.js",
+      "_vendor-arizeai-DUhQaeau.js",
+      "_vendor-codemirror-D_6Q6Auv.js",
       "_vendor-three-C5WAXd5r.js"
     ]
   },
-  "_pages-B7wCtpad.js": {
-    "file": "assets/pages-B7wCtpad.js",
+  "_pages-BCR8hW_l.js": {
+    "file": "assets/pages-BCR8hW_l.js",
     "name": "pages",
     "imports": [
-      "_vendor-DqQvHbPa.js",
-      "_vendor-arizeai-CLX44PFA.js",
-      "_components-DlL7ybQ2.js",
-      "_vendor-codemirror-Du3XyJnB.js",
-      "_vendor-recharts-B2PJDrnX.js"
+      "_vendor-DRWIRkSJ.js",
+      "_vendor-arizeai-DUhQaeau.js",
+      "_components-Bwf6zNbg.js",
+      "_vendor-codemirror-D_6Q6Auv.js",
+      "_vendor-recharts-BNBwj7vz.js"
     ]
   },
-  "_vendor-DqQvHbPa.js": {
-    "file": "assets/vendor-DqQvHbPa.js",
+  "_vendor-DRWIRkSJ.js": {
+    "file": "assets/vendor-DRWIRkSJ.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-C5WAXd5r.js"
@@ -35,33 +35,33 @@
     "file": "assets/vendor-WIZid84E.css",
     "src": "_vendor-WIZid84E.css"
   },
-  "_vendor-arizeai-CLX44PFA.js": {
-    "file": "assets/vendor-arizeai-CLX44PFA.js",
+  "_vendor-arizeai-DUhQaeau.js": {
+    "file": "assets/vendor-arizeai-DUhQaeau.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-DqQvHbPa.js"
+      "_vendor-DRWIRkSJ.js"
     ]
   },
-  "_vendor-codemirror-Du3XyJnB.js": {
-    "file": "assets/vendor-codemirror-Du3XyJnB.js",
+  "_vendor-codemirror-D_6Q6Auv.js": {
+    "file": "assets/vendor-codemirror-D_6Q6Auv.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-DqQvHbPa.js",
-      "_vendor-shiki-CNbrFjf9.js"
+      "_vendor-DRWIRkSJ.js",
+      "_vendor-shiki-k1qj_XjP.js"
     ]
   },
-  "_vendor-recharts-B2PJDrnX.js": {
-    "file": "assets/vendor-recharts-B2PJDrnX.js",
+  "_vendor-recharts-BNBwj7vz.js": {
+    "file": "assets/vendor-recharts-BNBwj7vz.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-DqQvHbPa.js"
+      "_vendor-DRWIRkSJ.js"
     ]
   },
-  "_vendor-shiki-CNbrFjf9.js": {
-    "file": "assets/vendor-shiki-CNbrFjf9.js",
+  "_vendor-shiki-k1qj_XjP.js": {
+    "file": "assets/vendor-shiki-k1qj_XjP.js",
     "name": "vendor-shiki",
     "imports": [
-      "_vendor-DqQvHbPa.js"
+      "_vendor-DRWIRkSJ.js"
     ]
   },
   "_vendor-three-C5WAXd5r.js": {
@@ -69,19 +69,19 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-QP9R8k34.js",
+    "file": "assets/index-Bfg9uQ43.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-DqQvHbPa.js",
-      "_vendor-arizeai-CLX44PFA.js",
-      "_pages-B7wCtpad.js",
-      "_components-DlL7ybQ2.js",
+      "_vendor-DRWIRkSJ.js",
+      "_vendor-arizeai-DUhQaeau.js",
+      "_pages-BCR8hW_l.js",
+      "_components-Bwf6zNbg.js",
       "_vendor-three-C5WAXd5r.js",
-      "_vendor-codemirror-Du3XyJnB.js",
-      "_vendor-shiki-CNbrFjf9.js",
-      "_vendor-recharts-B2PJDrnX.js"
+      "_vendor-codemirror-D_6Q6Auv.js",
+      "_vendor-shiki-k1qj_XjP.js",
+      "_vendor-recharts-BNBwj7vz.js"
     ]
   }
 }