> ## Documentation Index
> Fetch the complete documentation index at: https://freelance-4888ed08.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# Database Optimization

> PostgreSQL configuration and optimization strategies for high-throughput ML workloads.

Engineering

> [!NOTE]
>
> **What this does:** Tunes PostgreSQL for high-concurrency ML workloads — handling dozens of parallel optimization workers writing results simultaneously without data corruption or slowdowns.

## Overview

The Demand Forecasting system uses **PostgreSQL** as the persistent backend for Optuna studies (hyperparameter optimization history). Given the high concurrency of parallel optimization (multiple workers writing trial results simultaneously), the database configuration is tuned for performance and reliability.

This guide details the specific optimizations implemented to handle ML workloads, based on `pg-aiguide` best practices.

```mermaid theme={null}
graph LR
    %% Styling
    classDef worker fill:#DBEAFE,stroke:#2563EB,color:#1E3A5F,stroke-width:2px;
    classDef pool fill:#E0E7FF,stroke:#4338CA,color:#1E1B4B,stroke-width:2px;
    classDef health fill:#FEF3C7,stroke:#B45309,color:#78350F,stroke-width:2px;
    classDef db fill:#DCFCE7,stroke:#15803D,color:#14532D,stroke-width:2px;
    classDef overflow fill:#F3E8FF,stroke:#7C3AED,color:#2E1065,stroke-width:2px;

    W1["Airflow Worker 1"]:::worker
    W2["Airflow Worker 2"]:::worker
    W3["Airflow Worker N"]:::worker

    Pool["SQLAlchemy Pool<br/><i>pool_size: 10</i>"]:::pool
    Overflow["Overflow Pool<br/><i>max_overflow: 20</i>"]:::overflow
    Ping{"Pre-Ping<br/><i>Health Check</i>"}:::health
    PG[("PostgreSQL<br/><i>Optuna Studies</i>")]:::db

    W1 --> Pool
    W2 --> Pool
    W3 --> Overflow
    Pool --> Ping
    Overflow --> Ping
    Ping -->|Alive| PG

    %% Link Styling
    linkStyle default stroke:#6366f1,stroke-width:2px;
```

## Connection Pooling

To manage concurrent connections from Airflow workers efficiently, we use a tailored connection pooling strategy in `src/demand_forecasting/utils/s3_helpers.py` and `catboost_nodes_enhanced.py`.

```python theme={null}
storage = RDBStorage(
    url="postgresql://user:pass@host:5432/optuna",
    engine_kwargs={
        "pool_size": 10,          # Base connections to keep open
        "max_overflow": 20,       # Max additional connections for spikes
        "pool_pre_ping": True,    # Health check before using connection
        "pool_recycle": 3600,     # Recycle connections every hour
        "connect_args": {
            "connect_timeout": 10,
            "options": "-c statement_timeout=300000"  # 5-minute query timeout
        },
    },
)
```

### Configuration Rationale

| Setting | Value | Reason |
| :- | :- | :- |
| **`pool_size`** | `10` | Sufficient for standard Airflow worker threads without exhausting DB resources. |
| **`max_overflow`** | `20` | Allows bursting during parallel hyperparameter optimization (e.g., 50 concurrent trials). |
| **`pool_pre_ping`** | `True` | Prevents "OperationalError: server closed the connection unexpectedly" by verifying liveness. |
| **`pool_recycle`** | `3600` | Prevents issues with stale connections or firewalls cutting off idle sessions. |

## Database Schema Optimizations

Optuna creates its own tables (`trials`, `trial_params`, `trial_values`), but we apply specific indexes to speed up the complex queries used during the "Best Trial" lookup.

### Indexing Strategy

The following custom indexes are applied via `app-ml/scripts/init-optuna-db.sql`:

```sql theme={null}
-- Speed up trial lookups by study and state (most frequent query)
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_trials_study_state 
ON trials (study_id, state);

-- Speed up finding the best trial (min/max value)
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_trials_study_value 
ON trials (study_id, value) 
WHERE state = 'COMPLETE';

-- Optimize parameter retrieval
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_trial_params_trial_id 
ON trial_params (trial_id);
```

### Table Maintenance

Optuna tables grow rapidly. We use an aggressive `autovacuum` schedule to prevent bloat:

* **Autovacuum Scale Factor**: `0.05` (Runs when 5% of table changes, default is 20%)
* **Manual Vacuum**: Scheduled weekly via maintenance scripts.

## Performance Benchmarks

Our internal benchmarks (`scripts/validate_postgres_performance.py`) ensure the storage backend meets latency requirements:

| Operation | Target Latency | Actual (Avg) |
| :- | :- | :- |
| **Create Trial** | \< 100ms | 45ms |
| **Read Trial** | \< 10ms | 2ms |
| **Get Best Trial** | \< 100ms | 12ms |
| **Concurrent Writes** | 50 ops/sec | 85 ops/sec |

## Troubleshooting

### "Too Many Connections"

If you encounter `FATAL: sorry, too many clients already`, it means the combined pool size of all Airflow workers exceeds the PostgreSQL `max_connections` limit.

**Solution:**

1. Check `max_connections` in `postgresql.conf` (Recommended: `100`+).
2. Reduce `pool_size` in `catboost_nodes_enhanced.py`.
3. If scaling further, deploy **PgBouncer** as a connection pooler in front of PostgreSQL.


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.