feat: Enable robust DB connection handling (#1991)

Backported-from: main (24.03)
Backported-to: 23.09

changes/1991.enhance.md

@@ -0,0 +1 @@
+Enable robust DB connection handling by allowing `pool-pre-ping` setting.

configs/manager/halfstack.toml

@@ -12,6 +12,7 @@ name = "backend"
 user = "postgres"
 password = "develove"
 pool-recycle = 50
+pool-pre-ping = false
 
 
 [manager]

configs/manager/sample.toml

@@ -32,6 +32,12 @@ password = "DB_PASSWORD"                    # env: BACKEND_DB_PASSWORD
 # https://docs.sqlalchemy.org/en/14/core/pooling.html#setting-pool-recycle
 # pool-recycle = -1
 
+# This setting eliminates DB error due to stale pooled connections by ping at the start of each connection pool checkout.
+# It adds a small overhead to the connection checkout process.
+# https://docs.sqlalchemy.org/en/14/core/pooling.html#disconnect-handling-pessimistic
+# Default is false.
+# pool-pre-ping = false
+
 # The maximum allowed overflow of the connection pool
 # (See https://docs.sqlalchemy.org/en/20/core/engines.html#sqlalchemy.create_engine.params.max_overflow)
 # max-overflow = 64

src/ai/backend/manager/config.py

@@ -243,6 +243,7 @@
             t.Key("password"): t.String,
             t.Key("pool-size", default=8): t.ToInt[1:],  # type: ignore
             t.Key("pool-recycle", default=-1): t.ToFloat[-1:],  # -1 is infinite
+            t.Key("pool-pre-ping", default=False): t.ToBool,
             t.Key("max-overflow", default=64): t.ToInt[-1:],  # -1 is infinite  # type: ignore
         }),
         t.Key("manager"): t.Dict({

src/ai/backend/manager/models/utils.py

@@ -205,6 +205,7 @@ async def connect_database(
         connect_args=pgsql_connect_opts,
         pool_size=local_config["db"]["pool-size"],
         pool_recycle=local_config["db"]["pool-recycle"],
+        pool_pre_ping=local_config["db"]["pool-pre-ping"],
         max_overflow=local_config["db"]["max-overflow"],
         json_serializer=functools.partial(json.dumps, cls=ExtendedJSONEncoder),
         isolation_level=isolation_level,

tests/manager/conftest.py

@@ -183,6 +183,7 @@ def local_config(
             "password": "develove",
             "pool-size": 8,
             "pool-recycle": -1,
+            "pool-pre-ping": False,
             "max-overflow": 64,
         },
         "manager": {