[Model Monitoring] Add configurable TimescaleDB connection pool timeout (ML-11775) (#9172)

alxtkr77 · Alex Toker · web-flow · commit a8c1684f6cc8 · 2026-01-08T10:55:07.000+02:00
## Summary
- Add new config option
`model_endpoint_monitoring.tsdb.connection_pool_timeout` to control
TimescaleDB connection pool timeout
- Default is 120 seconds (previously hardcoded to 30 seconds)
- This allows API requests to wait longer for a database connection
under high load

## Changes Made
- Added `connection_pool_timeout` config in `mlrun/config.py` under
`model_endpoint_monitoring.tsdb`
- Updated `TimescaleDBConnection.pool` property to use the configurable
timeout
- Added unit tests to verify the config is correctly applied

Co-authored-by: Alex Toker &lt;alex_toker@mckinsey.com&gt;
diff --git a/mlrun/config.py b/mlrun/config.py
@@ -682,6 +682,9 @@
             # When True, automatically create/generate database name using system_id if not explicitly
             # specified in the connection string. When False, use the database from connection string as-is.
             "auto_create_database": True,
+            # Connection pool timeout in seconds. This is the maximum time to wait for a connection
+            # from the pool before raising an error.
+            "connection_pool_timeout": 120,
         },
     },
     "secret_stores": {
diff --git a/mlrun/model_monitoring/db/tsdb/timescaledb/timescaledb_connection.py b/mlrun/model_monitoring/db/tsdb/timescaledb/timescaledb_connection.py
@@ -23,6 +23,7 @@
 from psycopg_pool import ConnectionPool
 
 import mlrun.errors
+from mlrun.config import config
 from mlrun.model_monitoring.db.tsdb.preaggregate import PreAggregateManager
 from mlrun.utils import logger
 
@@ -127,7 +128,9 @@ def pool(self) -> ConnectionPool:
                 conninfo=self._dsn,
                 min_size=self._min_connections,
                 max_size=self._max_connections,
-                timeout=30.0,
+                timeout=float(
+                    config.model_endpoint_monitoring.tsdb.connection_pool_timeout
+                ),
             )
         return self._pool
 
diff --git a/tests/model_monitoring/db/tsdb/timescaledb/test_timescaledb_connection.py b/tests/model_monitoring/db/tsdb/timescaledb/test_timescaledb_connection.py
@@ -17,6 +17,7 @@
 import psycopg
 import pytest
 
+import mlrun.config
 import mlrun.errors
 from mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_connection import (
     Statement,
@@ -270,3 +271,41 @@ def test_persistent_failure_behavior_verification(
         # Verify error message
         assert "deadlock persisted after 3 retries" in str(exc_info.value)
         assert mock_cursor.execute.call_count == 4  # Initial + 3 retries
+
+
+class TestTimescaleDBConnectionPoolTimeout:
+    """Test connection pool timeout configuration (ML-11775)."""
+
+    def test_pool_uses_configured_timeout(self):
+        """Test that ConnectionPool is created with timeout from config."""
+        # Set custom timeout in config
+        original_timeout = (
+            mlrun.config.config.model_endpoint_monitoring.tsdb.connection_pool_timeout
+        )
+        mlrun.config.config.model_endpoint_monitoring.tsdb.connection_pool_timeout = 90
+
+        try:
+            with patch(
+                "mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_connection.ConnectionPool"
+            ) as mock_pool_class:
+                conn = TimescaleDBConnection(
+                    dsn="postgres://test:test@localhost:5432/test",
+                    max_connections=5,
+                )
+                # Access the pool property to trigger pool creation
+                _ = conn.pool
+
+                # Verify ConnectionPool was called with the configured timeout
+                mock_pool_class.assert_called_once()
+                call_kwargs = mock_pool_class.call_args.kwargs
+                assert call_kwargs["timeout"] == 90.0
+        finally:
+            # Restore original value
+            mlrun.config.config.model_endpoint_monitoring.tsdb.connection_pool_timeout = original_timeout
+
+    def test_pool_default_timeout_is_120(self):
+        """Test that the default connection pool timeout is 120 seconds."""
+        default_timeout = (
+            mlrun.config.config.model_endpoint_monitoring.tsdb.connection_pool_timeout
+        )
+        assert default_timeout == 120
