Provide ability to bulk insert tasks in TaskRunRecorder

[saschwartz]

Overview

This merge request attempts to address part of #18753, by making the TaskRunRecorder service more efficiently insert tasks into the database by means of bulk insertion.

Current Behaviour

Currently, tasks are inserted one-by-one into the database by the service. This is very slow and can cause the messaging layer to become backed up. The throughput in a real-world setting with 5x replicas of TaskRunRecorder was <10 tasks inserted into database per second.

Proposed Behaviour

Tasks are inserted accordingly to a parametrizable batch size and flush interval (this mirrors the configuration parameters for the EventPersisterService

Dependencies

The PR depends on a few things being merged first:

• https://github.com/PrefectHQ/prefect/pull/19559/files
• https://github.com/PrefectHQ/prefect/pull/19585/files

Testing

Unit tests have been added covering the new codepath. Additionally, this has been stress-tested in an OSS HA setup and with 5x replicas of the TaskRunRecorder, about 5K per second insertions are achieved with a HA postgres instance.

Checklist

• This pull request references any related issue by including "closes <link to issue>"
  • If no issue exists and your change is not a small fix, please create an issue first.
• If this pull request adds new functionality, it includes unit tests that cover the changes
• If this pull request removes docs files, it includes redirect settings in mint.json.
• If this pull request adds functions or classes, it includes helpful docstrings.

[codspeed-hq]

CodSpeed Performance Report

Merging #19586 will not alter performance

_{Comparing saschwartz:schwarts/task-run-bulk-insert (1a8d6e6) with main (c5d255d)}

Summary

✅ 2 untouched

[saschwartz]

I changed the approach slightly to insert by task state as ordered in the StateType enumeration, so we don't coalesce data directly.

[desertaxle]

Overall, this is an awesome PR! I have a couple of concerns about the bulk task run insertion, and I left some suggestions for you to consider.

[desertaxle]

I think this might be a somewhat dangerous assumption since there could be variation in the event payload between versions, and those could be in the same batch if they have the same state type.

Could you make this more explicit by batching by the key signature itself? I'm thinking something like this:

batches_by_keys: dict[frozenset[str], list] = {}
for tr in all_task_runs:
    key_signature = frozenset(tr["task_run_dict"].keys())
    batches_by_keys.setdefault(key_signature, []).append(tr)

to replace your current batching by state.

[saschwartz]

Agree, makes sense

[saschwartz]

Okay, took this approach in latest version. Thanks for the review.

[desertaxle]

LGTM! Thanks for the awesome PR @saschwartz!