Upsert Timeout Writing Checkpoints From Multi‑GPU EC2 Training Job

Q: Problem – Upsert Timeout When Writing Checkpoints from a Multi‑GPU EC2 Training Job

During a distributed PyTorch training run on an m5.24xlarge instance with 8 NVIDIA A100 GPUs, the checkpointing routine attempts to write model metadata (epoch, loss, S3 URI) into a PostgreSQL database using an UPSERT statement. After the 3rd epoch the job aborts with the following error:

Q: 3. Capture a packet trace of the checkpoint phase

sudo tcpdump -i eth0 -nn -s 0 -w /tmp/checkpoint.pcap port 5432 and host ml-postgres.c9dxyz.us-east-1.rds.amazonaws.com Analysis with tcpdump -r /tmp/checkpoint.pcap -nn shows many simultaneous TCP SYN packets followed by retransmissions, confirming connection pool exhaustion.

Q: 4. Inspect SQLAlchemy engine configuration

# checkpoint_manager.py engine = create_engine( DB_URI, pool_size=5, # default max_overflow=10, pool_timeout=5, # seconds connect_args={"options": "-c statement_timeout=5000"} ) 5. Check RDS parameter group aws rds describe-db-parameters \ --db-parameter-group-name ml-postgres-pg \ --query "Parameters[?ParameterName=='max_connections' || ParameterName=='statement_timeout']" Output:

Q: 2. Increase RDS max_connections and relax statement_timeout

aws rds modify-db-parameter-group \ --db-parameter-group-name ml-postgres-pg \ --parameters "ParameterName=max_connections,ParameterValue=100,ApplyMethod=immediate" \ "ParameterName=statement_timeout,ParameterValue=15000,ApplyMethod=immediate" After the change, the instance automatically accepts up to 100 concurrent connections and allows statements to run up to 15 seconds.

Q: 3. Adjust SQLAlchemy pool settings for the writer

engine = create_engine( DB_URI, pool_size=2, # two connections for redundancy max_overflow=0, pool_timeout=30, connect_args={"options": "-c statement_timeout=15000"} ) 4. Optional – Use SQS + Lambda for asynchronous checkpoint metadata persistence If the training job runs on spot instances where the process may be pre‑empted, decoupling via SQS guarantees eventual consistency.

In this article

Table of Contents

Problem – Upsert Timeout When Writing Checkpoints from a Multi‑GPU EC2 Training Job

During a distributed PyTorch training run on an m5.24xlarge instance with 8 NVIDIA A100 GPUs, the checkpointing routine attempts to write model metadata (epoch, loss, S3 URI) into a PostgreSQL database using an UPSERT statement. After the 3rd epoch the job aborts with the following error:

2026-09-24 14:37:12,845 ERROR checkpoint_manager.py:112 – Upsert failed: psycopg2.errors.QueryCanceled: canceling statement due to statement timeout

The training process then raises a RuntimeError and terminates, leaving the latest model checkpoint on the shared EFS volume but without a corresponding database entry.

Root Cause – Concurrency‑Induced Connection Saturation and Statement Timeout

Per‑GPU DB client: Each of the 8 GPU processes instantiates its own SQLAlchemy engine with the default pool size of 5. This results in up to 8 × 5 = 40 simultaneous connections.
PostgreSQL max_connections: The RDS instance is configured with max_connections = 30. When the training reaches the checkpoint barrier, all GPU processes issue an UPSERT at roughly the same time, exceeding the allowed connections.
Statement timeout: The RDS parameter group sets statement_timeout = 5000ms. Once the connection pool is exhausted, new connections wait for a free slot; the waiting time exceeds the timeout, causing the database to cancel the statement.
Network path: The instance resides in a private subnet behind a NAT gateway. NAT‑gateway throughput throttling adds ~150 ms latency per new TCP handshake, further inflating wait times under high concurrency.

In short, the combination of too many concurrent connections, a low max_connections limit, and a short statement timeout makes the checkpoint UPSERT vulnerable to cancellation.

Debug – Step‑by‑Step Investigation

1. Examine training logs for timeout patterns


2026-09-24 14:37:12,845 ERROR checkpoint_manager.py:112 – Upsert failed: psycopg2.errors.QueryCanceled: canceling statement due to statement timeout
2026-09-24 14:37:12,847 INFO trainer.py:78 – Checkpoint saved to s3://ml-checkpoints/job-123/epoch-3.pt
2026-09-24 14:37:12,849 WARN trainer.py:80 – Continuing without DB entry

2. Verify DB connection usage via CloudWatch


aws cloudwatch get-metric-statistics \
    --namespace AWS/RDS \
    --metric-name DatabaseConnections \
    --dimensions Name=DBInstanceIdentifier,Value=ml‑postgres \
    --statistics Maximum \
    --period 60 \
    --start-time 2026-09-24T14:30:00Z \
    --end-time 2026-09-24T14:40:00Z

Result (excerpt):

Timestamp	Maximum Connections
2026-09-24 14:35:00	32
2026-09-24 14:36:00	38
2026-09-24 14:37:00	41

3. Capture a packet trace of the checkpoint phase


sudo tcpdump -i eth0 -nn -s 0 -w /tmp/checkpoint.pcap port 5432 and host ml-postgres.c9dxyz.us-east-1.rds.amazonaws.com

Analysis with tcpdump -r /tmp/checkpoint.pcap -nn shows many simultaneous TCP SYN packets followed by retransmissions, confirming connection pool exhaustion.

4. Inspect SQLAlchemy engine configuration


# checkpoint_manager.py
engine = create_engine(
    DB_URI,
    pool_size=5,               # default
    max_overflow=10,
    pool_timeout=5,            # seconds
    connect_args={"options": "-c statement_timeout=5000"}
)

5. Check RDS parameter group


aws rds describe-db-parameters \
    --db-parameter-group-name ml-postgres-pg \
    --query "Parameters[?ParameterName=='max_connections' || ParameterName=='statement_timeout']"

Output:

ParameterName	ParameterValue
max_connections	30
statement_timeout	5000

Solution – Consolidate Upserts and Adjust DB Limits

1. Centralize checkpoint metadata writes

Introduce a single “checkpoint writer” process that receives metadata via an in‑process queue. All GPU processes push a lightweight JSON payload; the writer performs the UPSERT using a single dedicated DB connection.

# before (each GPU process writes directly)
def write_checkpoint(metadata):
    engine = create_engine(DB_URI)
    with engine.begin() as conn:
        conn.execute(upsert_stmt, metadata)

# after (single writer thread)
import threading, queue

metadata_q = queue.Queue()

def writer_thread():
    engine = create_engine(
        DB_URI,
        pool_size=1,
        max_overflow=0,
        pool_timeout=30,
        connect_args={"options": "-c statement_timeout=15000"}
    )
    while True:
        metadata = metadata_q.get()
        if metadata is None:
            break
        with engine.begin() as conn:
            conn.execute(upsert_stmt, metadata)

writer = threading.Thread(target=writer_thread, daemon=True)
writer.start()

def enqueue_checkpoint(metadata):
    metadata_q.put(metadata)

2. Increase RDS `max_connections` and relax `statement_timeout`


aws rds modify-db-parameter-group \
    --db-parameter-group-name ml-postgres-pg \
    --parameters "ParameterName=max_connections,ParameterValue=100,ApplyMethod=immediate" \
    "ParameterName=statement_timeout,ParameterValue=15000,ApplyMethod=immediate"

After the change, the instance automatically accepts up to 100 concurrent connections and allows statements to run up to 15 seconds.

3. Adjust SQLAlchemy pool settings for the writer


engine = create_engine(
    DB_URI,
    pool_size=2,               # two connections for redundancy
    max_overflow=0,
    pool_timeout=30,
    connect_args={"options": "-c statement_timeout=15000"}
)

4. Optional – Use SQS + Lambda for asynchronous checkpoint metadata persistence

If the training job runs on spot instances where the process may be pre‑empted, decoupling via SQS guarantees eventual consistency.


# enqueue to SQS
sqs.send_message(
    QueueUrl=CHECKPOINT_QUEUE,
    MessageBody=json.dumps(metadata)
)

# Lambda handler (Python 3.9)
def handler(event, context):
    for record in event['Records']:
        meta = json.loads(record['body'])
        with engine.begin() as conn:
            conn.execute(upsert_stmt, meta)

Verify – Confirm the Fix Works

Rerun the training job with --log-level DEBUG and watch for the absence of QueryCanceled errors.
Check CloudWatch for DatabaseConnections – it should stay below 30 throughout the run.
Inspect the checkpoint table after each epoch; the row count should increase monotonically.
Run a simple health‑check script:


#!/usr/bin/env python3
import psycopg2, time

conn = psycopg2.connect(dsn=os.getenv('DB_URI'))
cur = conn.cursor()
cur.execute("SELECT COUNT(*) FROM training_checkpoints;")
print("Rows:", cur.fetchone()[0])
conn.close()

Output after epoch 5 (expected):


Rows: 5

Prevent – Operational Guardrails

Connection pool sizing: Align pool_size × num_processes with max_connections. Use a single writer or a bounded queue for metadata writes.
Monitoring: Create a CloudWatch alarm on DatabaseConnections > 80% of max_connections and on pg_stat_activity.wait_event_type = 'Lock' counts.
Timeout policy: Set statement_timeout to a value that exceeds the longest expected upsert latency (e.g., 15 s) but still protects against runaway queries.
Network design: Place the training instance and the RDS instance in the same AZ and VPC subnet to avoid NAT‑gateway latency.
Testing: Include a unit test that spawns num_gpus × pool_size dummy connections and asserts that the connection count never exceeds the RDS limit.

Related Topic Hub: Cloud Infrastructure Troubleshooting Hub

FAQ

Why does the timeout only appear after a few epochs?
Each epoch triggers a barrier where all GPU processes attempt an upsert simultaneously. The first few epochs succeed because the connection pool is still under the limit; later epochs exceed max_connections, causing queuing and timeout.
Can I keep the per‑GPU upserts and just increase max_connections?
Increasing the limit mitigates the symptom but does not address the underlying design flaw. It also risks exhausting RDS resources under higher concurrency (e.g., larger GPU counts or multiple training jobs).
Is the issue specific to PostgreSQL, or would DynamoDB exhibit the same behavior?
DynamoDB does not have a fixed connection limit, but it enforces provisioned throughput. Simultaneous writes can still be throttled, leading to ProvisionedThroughputExceededException. The same principle—centralizing writes or using exponential backoff—applies.
How do I know if the NAT gateway is a bottleneck?
Monitor NetworkPacketsIn/Out and NatGatewayBytesIn/Out metrics. High packet loss or latency spikes coinciding with checkpoint phases indicate NAT‑gateway contention.
Should I use a transaction lock around the upsert?
The UPSERT (i.e., INSERT … ON CONFLICT … DO UPDATE) already acquires a row‑level lock. Adding an explicit transaction does not solve the connection‑pool issue and may increase lock contention.