gh-158239: Fix asyncio.gather performance regression

[deadlovelll]

Fix asyncio.gather performance regression

For more details see gh-158239

[deadlovelll]

I think this needs "skip news" badge

[cmaloney]

Confirmed locally this fixes the regression on benchmarks. I don't know asyncio internals enough to say if it's all good. @kumaraditya303 you reviewed the original could you review this update to it?

While this is an internal change which previously had NEWS I think it will need a new NEWS as the last change should go out in 3.15 final on Monday, whereas this will likely be in a different release 3.15.1

[leandrodamascena]

Hi! I have a project where I build CPython images from the main and maintenance branches, run them on AWS Lambda, and compare performance between commits. I was investigating a regression in asyncio.gather() when I came across this PR, so I tested the fix using the same Lambda workload.

The kind of code I am looking at is a Lambda consuming an SQS batch while each item also calls multiple services concurrently:

async def enrich_order(order):
    customer, inventory, risk = await asyncio.gather(
        customer_client.get(order["customer_id"]),
        inventory_client.check(order["items"]),
        fraud_client.check(order),
    )
    return customer, inventory, risk

async def process_batch(event):
    return await asyncio.gather(*(
        enrich_order(json.loads(record["body"]))
        for record in event["Records"]
    ))

def lambda_handler(event, context):
    return asyncio.run(process_batch(event))

My benchmark makes this gather tree deeper to measure the asyncio overhead without network latency hiding it. Running it on Lambda with x86_64 and 1024 MB, execution time went from 874 ms before the regression to 1,166 ms with it, a 34.4% increase. After accounting for the control run, the regression was 28.6%.

I repeated the comparison twice in my Linux PGO/LTO build environment and got 30.1% and 26.9%. I also tested this PR locally. The execution time returned close to the result from before the regression, and the extra memory usage disappeared.

I am not saying that the example Lambda above becomes exactly 34% slower. The impact depends on how many nested or repeated gather() calls happen during one invocation. My benchmark isolates that specific overhead.

I noticed that the new test covers the cancelled-sibling behavior, but it still passes with the extra callback present. I think it is worth adding a small guard for the mechanism that caused the regression:

def test_gather_does_not_add_callback_to_outer(self):
    # gh-158239: gather() must not add an internal done callback to
    # the outer future just to maintain the await graph.
    async def child():
        await asyncio.sleep(0)

    async def coro():
        outer = asyncio.gather(child(), child())
        self.assertFalse(outer._callbacks)
        await outer

    self.loop.run_until_complete(self.new_task(self.loop, coro()))

This test fails with the regressing implementation and passes with this PR.

Overall, the PR restored the previous performance and memory behavior in my tests. I hope this regression is fixed in the next Python release. Thanks for working on the fix!

[deadlovelll]

Hi! I have a project where I build CPython images from the main and maintenance branches, run them on AWS Lambda, and compare performance between commits. I was investigating a regression in asyncio.gather() when I came across this PR, so I tested the fix using the same Lambda workload.

The kind of code I am looking at is a Lambda consuming an SQS batch while each item also calls multiple services concurrently:

async def enrich_order(order):
    customer, inventory, risk = await asyncio.gather(
        customer_client.get(order["customer_id"]),
        inventory_client.check(order["items"]),
        fraud_client.check(order),
    )
    return customer, inventory, risk

async def process_batch(event):
    return await asyncio.gather(*(
        enrich_order(json.loads(record["body"]))
        for record in event["Records"]
    ))

def lambda_handler(event, context):
    return asyncio.run(process_batch(event))

My benchmark makes this gather tree deeper to measure the asyncio overhead without network latency hiding it. Running it on Lambda with x86_64 and 1024 MB, execution time went from 874 ms before the regression to 1,166 ms with it, a 34.4% increase. After accounting for the control run, the regression was 28.6%.

I repeated the comparison twice in my Linux PGO/LTO build environment and got 30.1% and 26.9%. I also tested this PR locally. The execution time returned close to the result from before the regression, and the extra memory usage disappeared.

I am not saying that the example Lambda above becomes exactly 34% slower. The impact depends on how many nested or repeated gather() calls happen during one invocation. My benchmark isolates that specific overhead.

I noticed that the new test covers the cancelled-sibling behavior, but it still passes with the extra callback present. I think it is worth adding a small guard for the mechanism that caused the regression:

def test_gather_does_not_add_callback_to_outer(self):
    # gh-158239: gather() must not add an internal done callback to
    # the outer future just to maintain the await graph.
    async def child():
        await asyncio.sleep(0)

    async def coro():
        outer = asyncio.gather(child(), child())
        self.assertFalse(outer._callbacks)
        await outer

    self.loop.run_until_complete(self.new_task(self.loop, coro()))

This test fails with the regressing implementation and passes with this PR.

Overall, the PR restored the previous performance and memory behavior in my tests. I hope this regression is fixed in the next Python release. Thanks for working on the fix!

hi, thank you for feedback! Your test is really handy, i pushed it here, with your co-authorship

[msullivan]

This looks good to me.

Do you know if the speed fix comes mostly from avoiding the extra callback in the event loop, or avoiding the functools.partial, or a combination?

[leandrodamascena]

I got curious about this too @msullivan , so I tried a few variants on my machine. Funny enough, the functools.partial doesn't seem to matter at all. Creating it without registering it ran the same as this PR, and swapping it for a closure didn't help either.

It's really the callback. A no-op callback already brought back part of the time. With regular tasks it's roughly half from the extra call_soon on every gather() and half from the discard work inside it, which the children's own callbacks already do. With eager tasks it's mostly the scheduling, and the scheduled handle also keeps outer alive until the loop gets to it, which is where the extra memory was coming from.

This looks good to me too, and I'd be glad to see it merged and backported to 3.14 and 3.15.

[miss-islington-app]

Thanks @deadlovelll for the PR, and @kumaraditya303 for merging it 🌮🎉.. I'm working now to backport this PR to: 3.14, 3.15.
🐍🍒⛏🤖

[bedevere-app]

GH-158655 is a backport of this pull request to the 3.15 branch.

[bedevere-app]

GH-158656 is a backport of this pull request to the 3.14 branch.