From 15effd9508771ca7d8ac70cdafc4e7987f473eb2 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 6 Nov 2025 16:20:43 +0100 Subject: [PATCH 01/12] refactor: Implemented buffered next method in LazyXComIterator so that multiple XCom's can be retrieved in one HTTP call from COMMS_SUPERVISOR --- .../sdk/execution_time/lazy_sequence.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 4efb0b71368ca..240ccabe0b1f9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -19,12 +19,16 @@ import collections import itertools +from collections import deque from collections.abc import Iterator, Sequence +from os import cpu_count from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload import attrs import structlog +from airflow.configuration import conf + if TYPE_CHECKING: from airflow.sdk.definitions.xcom_arg import PlainXComArg from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance @@ -35,7 +39,7 @@ # ``XCom.deserialize_value``. We don't want to wrap the API values in a nested # {"value": value} dict since it wastes bandwidth. _XComWrapper = collections.namedtuple("_XComWrapper", "value") - +lazy_xcom_prefetch_size = conf.getint("core", "lazy_xcom_prefetch_size", fallback=cpu_count() or 1) log = structlog.get_logger(logger_name=__name__) @@ -44,15 +48,21 @@ class LazyXComIterator(Iterator[T]): seq: LazyXComSequence[T] index: int = 0 dir: Literal[1, -1] = 1 + _buffer = deque() def __next__(self) -> T: if self.index < 0: # When iterating backwards, avoid extra HTTP request raise StopIteration() - try: - val = self.seq[self.index] - except IndexError: - raise StopIteration from None + + # If the buffer is empty, fetch the next chunk + if not self._buffer: + chunk = list(self.seq[self.index: self.index + lazy_xcom_prefetch_size]) + if not chunk: + raise StopIteration() + self._buffer.extend(chunk) + + val = self._buffer.popleft() self.index += self.dir return val From 521b1122ffd5498ed1f2b5de4ea01cee8cd803ab Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 6 Nov 2025 17:25:59 +0100 Subject: [PATCH 02/12] refactor: Fixed definition of _buffer in LazyXComIterator --- task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 240ccabe0b1f9..f3a2e8d3aacad 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -48,7 +48,7 @@ class LazyXComIterator(Iterator[T]): seq: LazyXComSequence[T] index: int = 0 dir: Literal[1, -1] = 1 - _buffer = deque() + _buffer: deque[T] = attrs.field(factory=deque, init=False) def __next__(self) -> T: if self.index < 0: @@ -57,7 +57,7 @@ def __next__(self) -> T: # If the buffer is empty, fetch the next chunk if not self._buffer: - chunk = list(self.seq[self.index: self.index + lazy_xcom_prefetch_size]) + chunk = list(self.seq[self.index : self.index + lazy_xcom_prefetch_size]) if not chunk: raise StopIteration() self._buffer.extend(chunk) From 55c1eb8e2ab30e9a7f4fac946bf1f4bc11f4bf47 Mon Sep 17 00:00:00 2001 From: David Blain Date: Tue, 18 Nov 2025 17:54:53 +0100 Subject: [PATCH 03/12] refactor: Refactored initialisation of buffer in LazyXComIterator --- .../sdk/execution_time/lazy_sequence.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index f3a2e8d3aacad..1bb6b78e5722e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -26,6 +26,7 @@ import attrs import structlog +from airflow.utils.log.logging_mixin import LoggingMixin from airflow.configuration import conf @@ -39,16 +40,23 @@ # ``XCom.deserialize_value``. We don't want to wrap the API values in a nested # {"value": value} dict since it wastes bandwidth. _XComWrapper = collections.namedtuple("_XComWrapper", "value") -lazy_xcom_prefetch_size = conf.getint("core", "lazy_xcom_prefetch_size", fallback=cpu_count() or 1) log = structlog.get_logger(logger_name=__name__) +def _deque_factory() -> deque: + return deque(maxlen=conf.getint("core", "parallelism", fallback=cpu_count())) + + @attrs.define -class LazyXComIterator(Iterator[T]): +class LazyXComIterator(LoggingMixin, Iterator[T]): seq: LazyXComSequence[T] index: int = 0 dir: Literal[1, -1] = 1 - _buffer: deque[T] = attrs.field(factory=deque, init=False) + _buffer: deque[T] = attrs.field(factory=_deque_factory, init=False) + + @property + def prefetch_size(self) -> int: + return self._buffer.maxlen def __next__(self) -> T: if self.index < 0: @@ -57,13 +65,17 @@ def __next__(self) -> T: # If the buffer is empty, fetch the next chunk if not self._buffer: - chunk = list(self.seq[self.index : self.index + lazy_xcom_prefetch_size]) + chunk = list(self.seq[self.index : self.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.index += self.dir + self.log.debug( + "Popped buffered XCom for index %s: %s", self.index, val + ) return val def __iter__(self) -> Iterator[T]: From d8013ab5702398719073034589e62332f109934f Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 5 Dec 2025 10:17:21 +0100 Subject: [PATCH 04/12] refactor: Reformatted lazy sequence --- task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 1bb6b78e5722e..4893c235e5e34 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -26,9 +26,9 @@ import attrs import structlog -from airflow.utils.log.logging_mixin import LoggingMixin from airflow.configuration import conf +from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: from airflow.sdk.definitions.xcom_arg import PlainXComArg @@ -55,7 +55,7 @@ class LazyXComIterator(LoggingMixin, Iterator[T]): _buffer: deque[T] = attrs.field(factory=_deque_factory, init=False) @property - def prefetch_size(self) -> int: + def prefetch_size(self) -> int | None: return self._buffer.maxlen def __next__(self) -> T: @@ -73,9 +73,7 @@ def __next__(self) -> T: val = self._buffer.popleft() self.index += self.dir - self.log.debug( - "Popped buffered XCom for index %s: %s", self.index, val - ) + self.log.debug("Popped buffered XCom for index %s: %s", self.index, val) return val def __iter__(self) -> Iterator[T]: From b8be90186127faaaf65a0007bb2d7fc28114efdf Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 5 Dec 2025 19:22:02 +0100 Subject: [PATCH 05/12] refactor: Fixed prefetch_size in LazyXComIterator --- task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 4893c235e5e34..7546059d92ca3 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -21,7 +21,6 @@ import itertools from collections import deque from collections.abc import Iterator, Sequence -from os import cpu_count from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload import attrs @@ -44,7 +43,7 @@ def _deque_factory() -> deque: - return deque(maxlen=conf.getint("core", "parallelism", fallback=cpu_count())) + return deque(maxlen=conf.getint("core", "parallelism")) @attrs.define @@ -55,8 +54,8 @@ class LazyXComIterator(LoggingMixin, Iterator[T]): _buffer: deque[T] = attrs.field(factory=_deque_factory, init=False) @property - def prefetch_size(self) -> int | None: - return self._buffer.maxlen + def prefetch_size(self) -> int: + return self._buffer.maxlen or conf.getint("core", "parallelism") def __next__(self) -> T: if self.index < 0: From 861bf0b0c919917ace55865921470bebdb486422 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Dec 2025 09:31:30 +0100 Subject: [PATCH 06/12] refactor: Changed import of LoggingMixin from core to sdk --- task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 7546059d92ca3..fbb7b355af8d8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -27,7 +27,7 @@ import structlog from airflow.configuration import conf -from airflow.utils.log.logging_mixin import LoggingMixin +from airflow.sdk.definitions._internal.logging_mixin import LoggingMixin if TYPE_CHECKING: from airflow.sdk.definitions.xcom_arg import PlainXComArg From 25d8e08a50d239f1a4a14f854596f9cee60f613e Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Dec 2025 09:34:45 +0100 Subject: [PATCH 07/12] refactor: Try fixing test_iter test --- task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index 16149ae9e26c3..3cfdbf13dcf13 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -77,7 +77,7 @@ def test_iter(mock_supervisor_comms, lazy_sequence): it = iter(lazy_sequence) mock_supervisor_comms.send.side_effect = [ - XComSequenceIndexResult(root="f"), + XComSequenceSliceResult(root=["f"]), ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}), ] assert list(it) == ["f"] From 65270a94558460a029b227b1e7a4dd2061231dfc Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Dec 2025 17:39:21 +0100 Subject: [PATCH 08/12] refactor: Try fixing test_iter on lazy_sequence and added test case when XCom not found --- .../execution_time/test_lazy_sequence.py | 61 ++++++++++++------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index 3cfdbf13dcf13..b571af51fda80 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -22,8 +22,9 @@ import pytest import airflow +from airflow.configuration import conf from airflow.sdk.bases.xcom import BaseXCom -from airflow.sdk.exceptions import ErrorType +from airflow.sdk.exceptions import ErrorType, AirflowRuntimeError from airflow.sdk.execution_time.comms import ( ErrorResponse, GetXComCount, @@ -74,34 +75,52 @@ def test_len(mock_supervisor_comms, lazy_sequence): def test_iter(mock_supervisor_comms, lazy_sequence): + stop = step = conf.getint("core", "parallelism") it = iter(lazy_sequence) mock_supervisor_comms.send.side_effect = [ XComSequenceSliceResult(root=["f"]), - ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}), ] assert list(it) == ["f"] - mock_supervisor_comms.send.assert_has_calls( - [ - call( - msg=GetXComSequenceItem( - key=BaseXCom.XCOM_RETURN_KEY, - dag_id="dag", - task_id="task", - run_id="run", - offset=0, - ), + + mock_supervisor_comms.send.assert_called_once_with( + call( + msg=GetXComSequenceSlice( + key=BaseXCom.XCOM_RETURN_KEY, + dag_id="dag", + task_id="task", + run_id="run", + start=0, + stop=stop, + step=step, ), - call( - msg=GetXComSequenceItem( - key=BaseXCom.XCOM_RETURN_KEY, - dag_id="dag", - task_id="task", - run_id="run", - offset=1, - ), + ) + ) + + +def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence): + stop = step = conf.getint("core", "parallelism") + it = iter(lazy_sequence) + + mock_supervisor_comms.send.side_effect = [ + XComSequenceSliceResult(root=["f"]), + ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}), + ] + with pytest.raises(AirflowRuntimeError): + list(it) + + mock_supervisor_comms.send.assert_called_once_with( + call( + msg=GetXComSequenceSlice( + key=BaseXCom.XCOM_RETURN_KEY, + dag_id="dag", + task_id="task", + run_id="run", + start=0, + stop=stop, + step=step, ), - ] + ) ) From 7e4694ff2e0d1d184d2404e5d27836588fe4bdd2 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Dec 2025 18:16:44 +0100 Subject: [PATCH 09/12] refactor: Reordered imports exceptions --- task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index b571af51fda80..7c63a25c971f9 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -24,7 +24,7 @@ import airflow from airflow.configuration import conf from airflow.sdk.bases.xcom import BaseXCom -from airflow.sdk.exceptions import ErrorType, AirflowRuntimeError +from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType from airflow.sdk.execution_time.comms import ( ErrorResponse, GetXComCount, From 299cf40a683c244d927cefdbc500b34fd7ee07c8 Mon Sep 17 00:00:00 2001 From: David Blain Date: Fri, 12 Dec 2025 18:20:09 +0100 Subject: [PATCH 10/12] refactor: Fixed 2 iter related tests on lazy_sequence --- .../execution_time/test_lazy_sequence.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index 7c63a25c971f9..78784a1366ed7 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -24,7 +24,7 @@ import airflow from airflow.configuration import conf from airflow.sdk.bases.xcom import BaseXCom -from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType +from airflow.sdk.exceptions import ErrorType from airflow.sdk.execution_time.comms import ( ErrorResponse, GetXComCount, @@ -75,7 +75,7 @@ def test_len(mock_supervisor_comms, lazy_sequence): def test_iter(mock_supervisor_comms, lazy_sequence): - stop = step = conf.getint("core", "parallelism") + stop = conf.getint("core", "parallelism") + 1 it = iter(lazy_sequence) mock_supervisor_comms.send.side_effect = [ @@ -90,23 +90,23 @@ def test_iter(mock_supervisor_comms, lazy_sequence): dag_id="dag", task_id="task", run_id="run", - start=0, + start=1, stop=stop, - step=step, + step=None, ), ) ) def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence): - stop = step = conf.getint("core", "parallelism") + stop = conf.getint("core", "parallelism") + 1 it = iter(lazy_sequence) mock_supervisor_comms.send.side_effect = [ XComSequenceSliceResult(root=["f"]), ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}), ] - with pytest.raises(AirflowRuntimeError): + with pytest.raises(TypeError): list(it) mock_supervisor_comms.send.assert_called_once_with( @@ -116,9 +116,9 @@ def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence): dag_id="dag", task_id="task", run_id="run", - start=0, + start=1, stop=stop, - step=step, + step=None, ), ) ) From 40e3bd76b679160d5acb929d00ba608900f034f0 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 17 Dec 2025 10:53:24 +0100 Subject: [PATCH 11/12] refactor: Fixed assertion in test_iter of lazy_sequence --- .../execution_time/test_lazy_sequence.py | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index 78784a1366ed7..0196de02e789d 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -83,19 +83,30 @@ def test_iter(mock_supervisor_comms, lazy_sequence): ] assert list(it) == ["f"] - mock_supervisor_comms.send.assert_called_once_with( + assert mock_supervisor_comms.send.call_args_list == [ call( msg=GetXComSequenceSlice( key=BaseXCom.XCOM_RETURN_KEY, dag_id="dag", task_id="task", run_id="run", - start=1, + start=0, stop=stop, step=None, - ), - ) - ) + ) + ), + call( + msg=GetXComSequenceSlice( + key=BaseXCom.XCOM_RETURN_KEY, + dag_id="dag", + task_id="task", + run_id="run", + start=1, + stop=stop + 1, + step=None, + ) + ), + ] def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence): From 03a92c28ec1f59cbf1499f6f407398765da860a3 Mon Sep 17 00:00:00 2001 From: David Blain Date: Wed, 17 Dec 2025 14:38:10 +0100 Subject: [PATCH 12/12] refactor: Fixed iter tests on lazy_sequence --- .../execution_time/test_lazy_sequence.py | 31 +++++++++++++------ 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py index 0196de02e789d..f142a9077ef0c 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py +++ b/task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py @@ -75,7 +75,7 @@ def test_len(mock_supervisor_comms, lazy_sequence): def test_iter(mock_supervisor_comms, lazy_sequence): - stop = conf.getint("core", "parallelism") + 1 + stop = conf.getint("core", "parallelism") it = iter(lazy_sequence) mock_supervisor_comms.send.side_effect = [ @@ -85,7 +85,7 @@ def test_iter(mock_supervisor_comms, lazy_sequence): assert mock_supervisor_comms.send.call_args_list == [ call( - msg=GetXComSequenceSlice( + GetXComSequenceSlice( key=BaseXCom.XCOM_RETURN_KEY, dag_id="dag", task_id="task", @@ -96,7 +96,7 @@ def test_iter(mock_supervisor_comms, lazy_sequence): ) ), call( - msg=GetXComSequenceSlice( + GetXComSequenceSlice( key=BaseXCom.XCOM_RETURN_KEY, dag_id="dag", task_id="task", @@ -110,7 +110,7 @@ def test_iter(mock_supervisor_comms, lazy_sequence): def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence): - stop = conf.getint("core", "parallelism") + 1 + stop = conf.getint("core", "parallelism") it = iter(lazy_sequence) mock_supervisor_comms.send.side_effect = [ @@ -120,19 +120,30 @@ def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence): with pytest.raises(TypeError): list(it) - mock_supervisor_comms.send.assert_called_once_with( + assert mock_supervisor_comms.send.call_args_list == [ call( - msg=GetXComSequenceSlice( + GetXComSequenceSlice( key=BaseXCom.XCOM_RETURN_KEY, dag_id="dag", task_id="task", run_id="run", - start=1, + start=0, stop=stop, step=None, - ), - ) - ) + ) + ), + call( + GetXComSequenceSlice( + key=BaseXCom.XCOM_RETURN_KEY, + dag_id="dag", + task_id="task", + run_id="run", + start=1, + stop=stop + 1, + step=None, + ) + ), + ] def test_getitem_index(mock_supervisor_comms, lazy_sequence):