Concurrency issues with custom assignors after 3.0.x

[dskrypa]

First off, I would just like to say that I really appreciate all of the effort that has gone into this project. It has been invaluable to me at work. Over the past few days, I have been trying to update my code from working with kafka-python version 2.3.x to 3.0.x. The admin client improvements have been awesome.

The one problem that I have gotten stuck on is with a custom consumer group assignor that uses the current lag as one of the main drivers for which partitions are assigned to each consumer (i.e., consumers that get assigned partitions with significantly higher lag than the rest will get fewer partitions than other consumers in the group; if lag is low/even, then assignment will be more evenly distributed).

In the past (in 2.3.x and earlier), I was able to relatively easily use fetcher.end_offsets and coordinator._send_offset_fetch_request (followed by coordinator._client.poll to wait for that future), for example, to retrieve the information that was necessary for calculating lag (the public KafkaConsumer.highwater method doesn't work for a cold group leader). Now with 3.0.x, the kafka.net.selector.NetworkSelector.run method ends up raising a RuntimeError if the equivalent methods are called by the assignor from synchronous code.

I tried implementing a version of my custom assignor using asyncio, and got very very close, but calling almost any of the relevant internal/private async methods in kafka-python resulted in a RuntimeError: Task got bad yield: <kafka.future.Future object at ...>.

It looks like #3111 may be addressing at least part of this, possibly? Do you have any other insight or tips you could offer, or is this use case something that would require deeper changes?

[dpkp]

Hi there, thanks for the kind words. Yes, the 3.0 async network changes subtly broke your use-case: calling kafka internal IO from within an assignor doesn't work. But I think we can fix this potentially by supporting async def assign, which would allow you to await any existing or custom async methods so that you can get network io. But I'm not sure practically that solves your problem because there is currently no "public" async api provided by kafka-python. Also I'm not sure how you have been or would be able to access the consumer itself from within your assignor?

Another option might be for you to create your own KafkaAdminClient within your assignor and use that to fetch offsets. It would "block" the internal event loop while you're waiting for it to do it's own IO, but that may be ok for your use case?

Re asyncio. #3111 is part of 3.1 release work to support asyncio as an alternate backend. But I don't think the event loop itself is the issue -- if/when kafka-python runs on asyncio we'd have the same problem unless/until the assignor interface changes to support IO with await. (though I do recognize that really what you were doing is starting a separate event loop to try to side-step the NetworkSelector RuntimeError).

What do you think?

[dskrypa]

The pattern I have been using has been something like the following:

assignor = CustomAssignor(...)
config['partition_assignment_strategy'] = (assignor,)
consumer = KafkaConsumer(**config)
assignor.consumer = consumer

I know it's not the cleanest, but it is what it is...

Re:

I think we can fix this potentially by supporting async def assign, which would allow you to await any existing or custom async methods so that you can get network io.

After updating my consumer's assignment listener to use an async version, I looked for that exactly. If the relevant async methods could be called from an async version of that, it would be great.

My plan for my next attempt to solve this tomorrow was to instantiate a completely separate KafkaAdminClient in the assignor to accomplish this without relying on the consumer's event loop :)

And yes, exactly, I was trying to start a separate event loop to side-step the NetworkSelector RuntimeError. While the ability to use things like asyncio.gather would be great from a perspective of maximizing concurrency for independent API calls, in reality, the API calls necessary to calculate lag are fast enough that they can be completed semi-serially without causing any problems for rebalance latency.

[dskrypa]

I was able to get the custom assignor working again by using KafkaAdminClient instead of trying to piggy-back off of internal Consumer methods. It's working well, and it's definitely much cleaner now. I probably should have done it this way from the start! I think the initial idea was to use the consumer since it's able to get the same info and it's already initialized with connections established, etc...