Optimisation: Allow prefetch of LazyXComSequence in next method of LazyXComIterator

[dabla]

Allow prefetch of LazyXComSequence in next method of LazyXComIterator so that multiple XCom's can be retrieved in one HTTP call from COMMS_SUPERVISOR.

This prevents overloading the COMMS_SUPERVISOR when a LazyXComIterator is being processed in multi threading from a task decorated Python task, if not we will get errors like the one shown in the logs.txt file in attachment.

The API server can also be overloaded when dynamic task mapping is being processed from an LazyXComSequence, as each indexed XCom will correspond to a HTTP call to the API server, that's where prefetching would also prevent this.

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

[dabla]

Don't know if this is a good idea when LazyXComIterator is consumed from a partial/expand, as each worker will prefetch to much data probably. Or I have to solve it in my code and decorate the LazyXComIterator with the solution above, or it has to be able to detect if it's called from a partial/expand, which I think is too complicated to achieve. I'll let this open at the moment to see what your guys thought's are, because this optimisation is for a specific case, namely one worker consuming the iterator in multi-threaded environment, which is different scenario than the partial/expand, as there the iterator will be consumed concurrently by different workers.

[dabla]

This is what I came up with:

class BufferedXComIterator(LoggingMixin, Iterator[T]):
    def __init__(self, iterator: Iterator, prefetch_size: int | None = None) -> None:
        super().__init__()
        self._delegate = iterator if isinstance(iterator, Iterator) else iter(iterator)
        self._buffer = deque(maxlen=prefetch_size or cpu_count())

    @property
    def prefetch_size(self) -> int:
        return self._buffer.maxlen

    def __next__(self) -> T:
        if isinstance(self._delegate, LazyXComIterator):
            if self._delegate.index < 0:
                # When iterating backwards, avoid extra HTTP request
                raise StopIteration()

            # If the buffer is empty, fetch the next chunk
            if not self._buffer:
                chunk = list(self._delegate.seq[self._delegate.index: self._delegate.index + self.prefetch_size])
                if not chunk:
                    raise StopIteration()
                self._buffer.extend(chunk)
                self.log.debug("Buffered %s XCom's", len(self._buffer))

            val = self._buffer.popleft()
            self._delegate.index += self._delegate.dir
            self.log.debug("Popped buffered XCom for index %s: %s", self._delegate.index, val)
            return val
        return next(self._delegate)

[dabla]

@amoghrajesh @ashb @kaxil @potiuk What do you think about this optimization? It avoids fetching each item separately via the API when iterating over a lazy XCom. Instead, we fetch chunks of items during iteration, which reduces the number of API calls.

We’re already using this mechanism as a monkey patch in our Airflow infrastructure, and it has significantly reduced API calls. Not only is it faster, but it also lowers the load on the API server and decreases the number of DB queries.

WDYT?

[dabla]

You might be wondering why I came up with this.

We’re using YugabyteDB as our Airflow metadata database (yes, I know it’s not officially supported), and recently we ran into performance issues with the API server, including crashes when XCom routes were invoked. Some DAGs using mapped tasks caused spikes in XCom retrieval, which contributed to the problem.

This led me to design this solution—to reduce load—even though it didn’t fully solve the root issue. After digging deeper, we discovered that the indices created for the Airflow metadata database rely on Postgres logic, because Airflow (via SQLAlchemy) assumes it’s talking to Postgres, whereas in reality, it’s not.

By recreating optimized indices specifically for YugabyteDB, we were finally able to tackle the root cause. Still, I think this PR optimization is valuable and worth considering.

[potiuk]

I like the idea - this is nice optimisation. Regardless from the origin of it.

[potiuk]

It has of course a drawback of consuming too many of those. I think before we merge it, we need to consider certain scenarios and better way of controlling it so that users can control prefetching.

And of course explain in the docs all kinds of usages.

[dabla]

Yes, and there I have to come back to me above response regarding your why 'core.paralellis' question? I don't see directly how we could for example derive the 'max_active_tis_per_dagrun' parameter from the operator from where the Xcom is being processed, as that would make it already less rigid.

[dabla]

Or we introduce the 'BufferedXComIterator' I posted above, so that you can programmatically decide if you want to do buffered iteration or not. But then it's not out of the box, but you can decide the buffer size used for paging.

[uranusjr]

I feel it’d be a good idea to have a separate interface. The appropriate value here varies per-task or even per-XCom, so a deployment-wide config wouldn’t be too useful for anyone.

Maybe a very small default would kind of make sense, and if the amount of entries is small enough we could even just send everything regardless, but I feel this should be done more automatically at the server side instead, where more context is available to make a better decision.

[dabla]

I feel it’d be a good idea to have a separate interface. The appropriate value here varies per-task or even per-XCom, so a deployment-wide config wouldn’t be too useful for anyone.

Fully agree on that, I tried to search for better solution but couldn't come up with one that easily

[dabla]

Maybe a very small default would kind of make sense, and if the amount of entries is small enough we could even just send everything regardless, but I feel this should be done more automatically at the server side instead, where more context is available to make a better decision.

Don't see how you could do that on server side, as it is the client that determines how many iterations it will do, and thus how many XCom's it wants to retrieve, that just the part I'm struggling with.

[uranusjr]

How about a kwarg on BaseOperator?

[potiuk]

requesting changes until we sort out some details

I guess it's ready to get reviewed by others?

[dabla]

Well some help would be nice here to be able to dynamically detect if we want to do prefetching or not, don't know how this could be achieved in a generic way, maybe it's not possible. The current solution is too static, and thus could have positive as negative impact regarding on the situation.

[uranusjr]

This should import from airflow.sdk instead.

[uranusjr]

Instead of reading the conf, use @conf_vars to patch (and also test more than one value).

[github-actions]

This pull request has been automatically marked as stale because the author has not responded to a request for more information. It will be closed in 7 days if no further activity occurs. Thank you for your contributions.

[github-actions]

This pull request has been closed because the author has not responded to a request for more information.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

[dabla]

Superseded by #73790 (stacked on #62922). There, LazyXComSequence itself reads in chunks through GetXComSequenceSlice, for indexing, iteration and the async reads alike, with its own [core] xcom_sequence_chunk_size option (default 32) and at most one chunk held in memory. That also answers the open question here, about a prefetch size fixed to core.parallelism being too static.