Reduce memory used when deleting a Dag with a large history

[ColtenOuO]

Summary

delete_dag forces SQLAlchemy's "fetch" synchronization strategy on every bulk
delete it issues. That strategy reads the primary key of every deleted row back from
the database so it can mark matching in-memory objects as deleted — via RETURNING
on backends that support it, or a second full SELECT beforehand on those that don't.

The session holds nothing but the Dag's own DagModel row, so there is no identity
map to synchronize. The keys are read back, matched against an effectively empty map,
and discarded — at a cost proportional to the Dag's history, once per table with a
dag_id column.

Removing the override restores SQLAlchemy's default "auto" strategy, which evaluates
the criteria in Python against the objects already loaded. Synchronization still
happens; it just no longer needs a round-trip whose size is set by the row count
instead of the identity map.

Before / after

200,000 rows in a table shaped like task_instance (its eight declared indexes plus
the dag_run FK), with one ORM object in the session:

backend	synchronize_session	time	peak Python heap	statements
PostgreSQL 16	"fetch" (before)	0.625s	42.1 MiB	1 (DELETE … RETURNING id)
PostgreSQL 16	default (after)	0.055s	0.0 MiB	1 (DELETE)
MySQL 8.4	"fetch" (before)	24.341s	43.6 MiB	2 (SELECT id … + DELETE)
MySQL 8.4	default (after)	26.572s	0.0 MiB	1 (DELETE)

The memory saving is the point, and it scales with the Dag: ~211 bytes per deleted
row, so ~1 GiB of transient heap in the API server for a Dag with five million task
instances. The time saving is real on PostgreSQL only — InnoDB's own delete cost
dwarfs the extra SELECT, so the MySQL times are noise in both directions.

Benchmark script

"""
Measure what delete_dag's synchronize_session strategy costs.

Usage:
    pip install 'sqlalchemy>=2.0.50' psycopg2-binary pymysql
    python bench.py postgresql+psycopg2://postgres:pw@127.0.0.1:5432/test
    python bench.py mysql+pymysql://root:pw@127.0.0.1:3306/test
"""

from __future__ import annotations

import sys
import time
import tracemalloc

from sqlalchemy import (
    Column, DateTime, ForeignKeyConstraint, Index, Integer, String,
    create_engine, delete, event, select,
)
from sqlalchemy.orm import Session, declarative_base

URL = sys.argv[1] if len(sys.argv) > 1 else "postgresql+psycopg2://postgres:pw@127.0.0.1:5432/test"
ROWS = int(sys.argv[2]) if len(sys.argv) > 2 else 200_000
CHUNK = 10_000

Base = declarative_base()


class DagRun(Base):
    __tablename__ = "bench_dag_run"
    dag_id = Column(String(190), primary_key=True)
    run_id = Column(String(190), primary_key=True)


class TaskInstance(Base):
    __tablename__ = "bench_task_instance"
    id = Column(Integer, primary_key=True)
    dag_id = Column(String(190))
    task_id = Column(String(190))
    run_id = Column(String(190))
    map_index = Column(Integer, default=-1)
    state = Column(String(20))
    pool = Column(String(190))
    priority_weight = Column(Integer)
    trigger_id = Column(Integer)
    last_heartbeat_at = Column(DateTime)
    dag_version_id = Column(Integer)
    __table_args__ = (
        Index("bench_ti_dag_state", "dag_id", "state"),
        Index("bench_ti_dag_run", "dag_id", "run_id"),
        Index("bench_ti_state", "state"),
        Index("bench_ti_state_lkp", "dag_id", "task_id", "run_id", "state"),
        Index("bench_ti_pool", "pool", "state", "priority_weight"),
        Index("bench_ti_trigger_id", "trigger_id"),
        Index("bench_ti_heartbeat", "last_heartbeat_at"),
        Index("bench_ti_dag_version_id", "dag_version_id"),
        ForeignKeyConstraint(["dag_id", "run_id"], ["bench_dag_run.dag_id", "bench_dag_run.run_id"]),
    )


engine = create_engine(URL)


def seed() -> None:
    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)
    with Session(engine) as s:
        s.add_all([DagRun(dag_id="d", run_id=f"r{i}") for i in range(200)])
        s.commit()
        rows = [
            {
                "id": i, "dag_id": "d", "task_id": f"t{i % 50}", "run_id": f"r{i % 200}",
                "state": "success", "pool": "default_pool", "priority_weight": 1,
                "dag_version_id": 1,
            }
            for i in range(1, ROWS + 1)
        ]
        for start in range(0, len(rows), CHUNK):
            s.bulk_insert_mappings(TaskInstance, rows[start : start + CHUNK])
            s.commit()


def measure(strategy):
    seed()
    statements: list[str] = []

    def capture(_conn, _cursor, statement, _parameters, _context, _executemany):
        statements.append(" ".join(statement.split()))

    with Session(engine) as s:
        s.begin()
        s.scalars(select(TaskInstance).limit(1)).all()  # the one object delete_dag loads
        stmt = delete(TaskInstance).where(TaskInstance.dag_id == "d")
        if strategy is not None:
            stmt = stmt.execution_options(synchronize_session=strategy)
        event.listen(engine, "before_cursor_execute", capture)
        tracemalloc.start()
        started = time.monotonic()
        s.execute(stmt)
        elapsed = time.monotonic() - started
        peak = tracemalloc.get_traced_memory()[1]
        tracemalloc.stop()
        event.remove(engine, "before_cursor_execute", capture)
        s.rollback()
    return elapsed, peak / 1024 / 1024, statements


print(f"{URL.split('://')[0]}, deleting {ROWS:,} rows\n")
print(f"{'synchronize_session':<22} {'time':>9} {'peak heap':>12}  SQL")
for label, strategy in (('"fetch" (before)', "fetch"), ("default (after)", None)):
    elapsed, peak, statements = measure(strategy)
    print(f"{label:<22} {elapsed:>8.3f}s {peak:>10.1f} MiB  {len(statements)} statement(s)")
    for statement in statements:
        print(f"{'':<22} {'':>9} {'':>12}    {statement[:88]}")

The regression test asserts on the emitted SQL rather than the query count, because on
backends with RETURNING the count is identical either way. It also pins a property
worth keeping: the delete loop runs over every mapped model with a dag_id, and
"auto" silently falls back to "fetch" for criteria it cannot evaluate in Python, so
the assertions catch a future model that would quietly reintroduce the round-trip.

---

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code (Opus 5) for writting the regression test and Benchmark script