Database Optimization

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

Engineering

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.

​

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.

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:

-- 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.