diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 3eab132d1a5c6..dba670ef46b96 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -162,6 +162,9 @@ Dockerfile.ci @potiuk @ashb @gopidesupavan @amoghrajesh @jscheffl @bugraoz93 @ja # AIP-72 - Task SDK # Python SDK /task-sdk/ @ashb @kaxil @amoghrajesh +/task-sdk/**/batchedoperator.py @dabla +/task-sdk/**/executor.py @dabla +/task-sdk/**/iterableoperator.py @dabla # AIP-108 - Coordinators /task-sdk/src/airflow/sdk/coordinators/ @jason810496 @uranusjr diff --git a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py index 8762c6734be88..5914267d00231 100644 --- a/airflow-core/src/airflow/serialization/definitions/mappedoperator.py +++ b/airflow-core/src/airflow/serialization/definitions/mappedoperator.py @@ -345,6 +345,10 @@ def on_failure_fail_dagrun(self) -> bool: def on_failure_fail_dagrun(self, v) -> None: self.partial_kwargs["on_failure_fail_dagrun"] = bool(v) + @property + def batch_size(self) -> int: + return self.partial_kwargs.get("batch_size", 0) + @classmethod def get_serialized_fields(cls): """Fields to extract from JSON-Serialized DAG.""" @@ -470,15 +474,20 @@ def get_parse_time_mapped_ti_count(self) -> int: current_count = self._get_specified_expand_input().get_parse_time_mapped_ti_count() def _get_parent_count() -> int: - if (group := self.get_closest_mapped_task_group()) is None: - raise NotMapped() - return group.get_parse_time_mapped_ti_count() + try: + if (group := self.get_closest_mapped_task_group()) is None: + raise NotMapped() + return group.get_parse_time_mapped_ti_count() + except NotMapped: + return 1 + + parent_count = _get_parent_count() + mapped_ti_count = parent_count * current_count - try: - parent_count = _get_parent_count() - except NotMapped: - return current_count - return parent_count * current_count + if self.batch_size > 0: + if self.batch_size < mapped_ti_count: + return self.batch_size + return mapped_ti_count def iter_mapped_dependencies(self) -> Iterator[Operator]: """Upstream dependencies that provide XComs used by this task for task mapping.""" @@ -544,9 +553,14 @@ def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *, se group = task.get_closest_mapped_task_group() if group is None: - return current_count - parent_count = get_mapped_ti_count(group, run_id, session=session) - return parent_count * current_count + mapped_ti_count = current_count + else: + parent_count = get_mapped_ti_count(group, run_id, session=session) + mapped_ti_count = parent_count * current_count + + if task.batch_size > 0: + return min(task.batch_size, mapped_ti_count) + return mapped_ti_count @get_mapped_ti_count.register diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py b/airflow-core/tests/unit/models/test_mappedoperator.py index 3539428be43c3..557432f8800c0 100644 --- a/airflow-core/tests/unit/models/test_mappedoperator.py +++ b/airflow-core/tests/unit/models/test_mappedoperator.py @@ -1828,3 +1828,26 @@ def test_mapped_operator_retry_delay_explicit(dag_maker): # Should return the explicitly set value assert mapped_deser.retry_delay == custom_retry_delay + + +@pytest.mark.parametrize( + ("batch_size", "items", "expected"), + [ + pytest.param(5, [1, 2, 3], 3, id="batch_size-larger-than-items-clamped-to-items"), + pytest.param(2, [1, 2, 3], 2, id="batch_size-smaller-than-items-returned-as-is"), + pytest.param(3, [1, 2, 3], 3, id="batch_size-equal-to-items"), + ], +) +def test_get_mapped_ti_count_clamps_to_items_count_when_batch_size_exceeds_it( + dag_maker, session, batch_size, items, expected +): + from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count + + with dag_maker(dag_id=f"test_batch_size_clamp_{batch_size}", session=session, serialized=True) as dag: + MockOperator.partial(task_id="task").batch(size=batch_size).iterate(arg1=items) + + dr = dag_maker.create_dagrun() + task = dag.task_dict["task"] + + result = get_mapped_ti_count(task, dr.run_id, session=session) + assert result == expected diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 6a348a3f2ca26..5ea6d4f2c51b1 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -46,6 +46,7 @@ from airflow.exceptions import ( AirflowException, AirflowSkipException, + NotMapped, ) from airflow.models.asset import ( AssetActive, @@ -3487,6 +3488,74 @@ def show(a, b): dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session) assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)] + def test_iterate_literal_cross_product(self, dag_maker, session): + """Test an iterated task with literal cross product args properly.""" + outputs = [] + + with dag_maker(dag_id="product_same_types", session=session, serialized=True) as dag: + + @dag.task + def show(a, b): + outputs.append((a, b)) + + show.iterate(a=[2, 4, 8], b=[5, 10]) + + dag_run = dag_maker.create_dagrun() + + show_task = dag.get_task("show") + with pytest.raises(NotMapped): + show_task.get_parse_time_mapped_ti_count() + with pytest.raises(NotMapped): + TaskMap.expand_mapped_task(show_task, dag_run.run_id, session=session) + + tis = session.scalars( + select(TaskInstance) + .where( + TaskInstance.dag_id == dag.dag_id, + TaskInstance.task_id == "show", + TaskInstance.run_id == dag_run.run_id, + ) + .order_by(TaskInstance.map_index) + ).all() + for ti in tis: + ti.refresh_from_task(show_task) + dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session) + assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)] + + def test_iterate_literal_cross_product_in_batch(self, dag_maker, session): + """Test a batched iterated task with literal cross product args properly.""" + outputs = [] + + with dag_maker(dag_id="product_same_types", session=session, serialized=True) as dag: + + @dag.task + def show(a, b): + outputs.append((a, b)) + + show.batch(size=2).iterate(a=[2, 4, 8], b=[5, 10]) + + dag_run = dag_maker.create_dagrun() + + show_task = dag.get_task("show") + assert show_task.get_parse_time_mapped_ti_count() == 2 + mapped_tis, max_map_index = TaskMap.expand_mapped_task(show_task, dag_run.run_id, session=session) + assert len(mapped_tis) == 0 # Expanded at parse! + assert max_map_index == 1 + + tis = session.scalars( + select(TaskInstance) + .where( + TaskInstance.dag_id == dag.dag_id, + TaskInstance.task_id == "show", + TaskInstance.run_id == dag_run.run_id, + ) + .order_by(TaskInstance.map_index) + ).all() + for ti in tis: + ti.refresh_from_task(show_task) + dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session) + assert outputs == [(2, 5), (4, 5), (8, 5), (2, 10), (4, 10), (8, 10)] + def test_map_in_group(self, tmp_path: pathlib.Path, dag_maker, session): out = tmp_path.joinpath("out") out.touch() diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index cf9c5c34c17f5..b6e085bb4bd32 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -2994,6 +2994,7 @@ def test_operator_expand_serde(): "template_ext": [".sh", ".bash"], "template_fields_renderers": {"bash_command": "bash", "env": "json"}, "ui_color": "#f0ede4", + "_register_with_dag": True, "_disallow_kwargs_override": False, "_expand_input_attr": "expand_input", } @@ -3037,6 +3038,7 @@ def test_operator_expand_xcomarg_serde(): }, "task_id": "task_2", "template_fields": ["arg1", "arg2"], + "_register_with_dag": True, "_disallow_kwargs_override": False, "_expand_input_attr": "expand_input", } @@ -3092,6 +3094,7 @@ def test_operator_expand_kwargs_literal_serde(strict): }, "task_id": "task_2", "template_fields": ["arg1", "arg2"], + "_register_with_dag": True, "_disallow_kwargs_override": strict, "_expand_input_attr": "expand_input", } @@ -3139,6 +3142,7 @@ def test_operator_expand_kwargs_xcomarg_serde(strict): }, "task_id": "task_2", "template_fields": ["arg1", "arg2"], + "_register_with_dag": True, "_disallow_kwargs_override": strict, "_expand_input_attr": "expand_input", } @@ -3247,6 +3251,7 @@ def x(arg1, arg2, arg3): "op_args": "py", "op_kwargs": "py", }, + "_register_with_dag": True, "_disallow_kwargs_override": False, "_expand_input_attr": "op_kwargs_expand_input", "python_callable_name": "test_taskflow_expand_serde..x", @@ -3340,6 +3345,7 @@ def x(arg1, arg2, arg3): "op_args": "py", "op_kwargs": "py", }, + "_register_with_dag": True, "_disallow_kwargs_override": strict, "_expand_input_attr": "op_kwargs_expand_input", } @@ -3451,6 +3457,7 @@ def operator_extra_links(self): "partial_kwargs": { "retry_delay": {"__type": "timedelta", "__var": 300.0}, }, + "_register_with_dag": True, "_disallow_kwargs_override": False, "_expand_input_attr": "expand_input", "_operator_extra_links": {"airflow": "_link_AirflowLink2"}, diff --git a/devel-common/src/tests_common/test_utils/mock_context.py b/devel-common/src/tests_common/test_utils/mock_context.py index 4e7aa7884f259..589bd29f99861 100644 --- a/devel-common/src/tests_common/test_utils/mock_context.py +++ b/devel-common/src/tests_common/test_utils/mock_context.py @@ -20,14 +20,29 @@ from typing import TYPE_CHECKING, Any from unittest import mock +from airflow.models import DagRun +from airflow.utils.types import DagRunType + from tests_common.test_utils.compat import Context from tests_common.test_utils.taskinstance import create_task_instance +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS + +if AIRFLOW_V_3_1_PLUS: + from airflow.sdk import timezone +else: + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] if TYPE_CHECKING: from sqlalchemy.orm import Session -def mock_context(task) -> Context: +def generate_run_id() -> str: + if AIRFLOW_V_3_0_PLUS: + return DagRun.generate_run_id(run_type=DagRunType.MANUAL, run_after=timezone.utcnow()) + return DagRun.generate_run_id(run_type=DagRunType.MANUAL, execution_date=timezone.utcnow()) # type: ignore[call-arg] + + +def mock_context(task, run_id: str | None = None) -> Context: from airflow.models import TaskInstance from airflow.utils.session import NEW_SESSION @@ -64,6 +79,6 @@ def xcom_push(self, key: str, value: Any, session: Session = NEW_SESSION, **kwar values[key] = value values["ti"] = create_task_instance(task, dag_version_id=mock.MagicMock(), ti_type=MockedTaskInstance) + values["run_id"] = generate_run_id() if run_id is None else run_id - # See https://github.com/python/mypy/issues/8890 - mypy does not support passing typed dict to TypedDict return Context(values) # type: ignore[misc] diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 004b6a5062994..5abef306e890e 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1266,6 +1266,7 @@ PodManager PodSpec podSpec podspec +Pokémon polars poller polyfill diff --git a/task-sdk/.pre-commit-config.yaml b/task-sdk/.pre-commit-config.yaml index a44abd29c1566..9532ab4f599b0 100644 --- a/task-sdk/.pre-commit-config.yaml +++ b/task-sdk/.pre-commit-config.yaml @@ -39,7 +39,9 @@ repos: ^src/airflow/sdk/definitions/asset/__init__\.py$| ^src/airflow/sdk/definitions/asset/decorators\.py$| ^src/airflow/sdk/definitions/taskgroup\.py$| + ^src/airflow/sdk/definitions/iterableoperator\.py$| ^src/airflow/sdk/definitions/mappedoperator\.py$| + ^src/airflow/sdk/definitions/batchedoperator\.py$| ^src/airflow/sdk/definitions/deadline\.py$| ^src/airflow/sdk/definitions/dag\.py$| ^src/airflow/sdk/definitions/_internal/types\.py$| diff --git a/task-sdk/docs/deferred-vs-async-operators.rst b/task-sdk/docs/deferred-vs-async-operators.rst index 466060ab3a4fb..101a900d21475 100644 --- a/task-sdk/docs/deferred-vs-async-operators.rst +++ b/task-sdk/docs/deferred-vs-async-operators.rst @@ -205,7 +205,7 @@ concurrently using ``asyncio.gather`` while limiting concurrency with a semaphor .. note:: - The upcoming *Dynamic Task Iteration* feature will simplify patterns like this. + The new :ref:`Dynamic Task Iteration `. feature will simplify patterns like this. Instead of manually managing concurrency with constructs such as ``asyncio.gather`` and ``asyncio.Semaphore``, authors will be able to iterate over asynchronous results directly in downstream tasks while still benefiting diff --git a/task-sdk/docs/dynamic-task-mapping-vs-iteration.rst b/task-sdk/docs/dynamic-task-mapping-vs-iteration.rst new file mode 100644 index 0000000000000..2f4cdb9871094 --- /dev/null +++ b/task-sdk/docs/dynamic-task-mapping-vs-iteration.rst @@ -0,0 +1,434 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _sdk-dynamic-task-mapping-vs-iteration: + +Dynamic Task Mapping vs Task Iteration +====================================== + +.. versionadded:: 3.3.0 + +Airflow provides two complementary ways to process collections of data: + +- **Dynamic Task Mapping (DTM)** distributes work **across multiple workers**. + Each item becomes a separate Task Instance that can run on a different worker, + giving you horizontal scalability and per-item observability. + +- **Task Iteration (TI)** improves concurrency **within a single task**. + All items are processed inside one Task Instance on one worker, eliminating + scheduling overhead and — when combined with async operators — enabling true + I/O multiplexing through a shared event loop. + +In short: **DTM spreads load across workers; TI speeds up work within one worker.** + +While both approaches allow you to apply an operation over a collection, +they differ significantly in execution model, scheduler impact, and observability. +This page explains the trade-offs and when to use each. + +Real-World Motivation +--------------------- + +Consider a workflow that downloads ~17,000 XML files from an SFTP server and loads +them into a data warehouse. Community benchmarks demonstrate the dramatic performance +difference between the two approaches: + +.. list-table:: + :header-rows: 1 + + * - Approach + - Execution Time + * - Dynamic Task Mapping with mapped ``SFTPOperator`` + - 3 h 25 m + * - Sync ``@task`` with ``SFTPHook`` (sequential loop) + - 1 h 21 m + * - Async ``@task`` with ``SFTPHookAsync`` (concurrent loop) + - 8 m 29 s + * - Async ``@task`` with ``SFTPHookAsync`` and connection pooling + - 3 m 32 s + +The ~60× improvement stems from eliminating per-item scheduling overhead and +sharing a single event loop for concurrent I/O. This is the kind of workload +where TI excels: many small, I/O-bound operations processed within one task. + +Dynamic Task Mapping (DTM) +-------------------------- + +Dynamic Task Mapping allows you to expand a single task definition into multiple +Task Instances (TIs). + +For more details, see :ref:`dynamic task mapping `. + +Key characteristics: + +- Each item in the iterable creates a separate Task Instance. +- The scheduler is responsible for creating and managing all mapped tasks. +- Tasks can run in parallel across multiple worker slots. +- Fine-grained retry, logging, and observability per item. +- Well suited for workloads where each item should be independently scheduled and tracked. + +The following example fetches Pokémon data from a REST API. Each Pokémon becomes +a separate Task Instance, individually scheduled, retried, and visible in the UI: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.operators.http import HttpOperator + from airflow.sdk import DAG, task + + with DAG(dag_id="dtm-http-pokemon-example", start_date=datetime(2026, 1, 1)): + list_pokemon_task = HttpOperator( + task_id="list_pokemon", + http_conn_id="pokeapi", + method="GET", + endpoint="api/v2/pokemon?limit=100", + response_filter=lambda response: [ + pokemon["url"].replace("https://pokeapi.co/", "") for pokemon in response.json()["results"] + ], + log_response=False, + ) + + get_pokemon_task = HttpOperator.partial( + task_id="get_pokemon", + http_conn_id="pokeapi", + method="GET", + ).expand(endpoint=list_pokemon_task.output) + + list_pokemon_task >> get_pokemon_task + + +With 100 Pokémon the scheduler creates 100 Task Instances, each occupying +a worker slot. This is fine for small lists, but for thousands of items the +scheduler and database overhead becomes significant. + +Task Iteration (TI) +---------------------------- + +Task Iteration allows you to iterate over an iterable (typically an XCom result) +*within a single Task Instance*, applying an operator multiple times without creating +separate Task Instances. + +This means that iteration happens inside the task execution itself rather than at the +scheduler level. + +Key characteristics: + +- A single Task Instance processes all items in the iterable. +- No task expansion; the scheduler manages only one task. +- Lower scheduler overhead compared to DTM. +- Iterations share the same execution context (e.g., memory, event loop). +- Particularly well suited for async operators and high-throughput workloads. + +The same Pokémon fetching problem can be solved with TI. Here, a single Task +Instance processes all Pokémon concurrently using the sync +:class:`~airflow.providers.http.operators.http.HttpOperator`: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.operators.http import HttpOperator + from airflow.sdk import DAG, task + + with DAG(dag_id="it-http-pokemon-example", start_date=datetime(2026, 1, 1)): + list_pokemon_task = HttpOperator( + task_id="list_pokemon", + http_conn_id="pokeapi", + method="GET", + endpoint="api/v2/pokemon?limit=100", + response_filter=lambda response: [ + pokemon["url"].replace("https://pokeapi.co/", "") for pokemon in response.json()["results"] + ], + log_response=False, + ) + + get_pokemon_task = HttpOperator.partial( + task_id="get_pokemon", + http_conn_id="pokeapi", + method="GET", + ).iterate(endpoint=list_pokemon_task.output) + + list_pokemon_task >> get_pokemon_task + + +The scheduler only manages a single task. With sync tasks, iterations are +executed in a multi-threaded fashion, which eliminates scheduling overhead +and can speed up compute-bound workloads. However, for I/O-bound operations +like HTTP requests, multi-threading alone does not provide the same +performance benefits as async multiplexing — threads still block on each +request individually rather than sharing a single event loop. + +To truly **multiplex** I/O-bound operations, use an async task with +:class:`~airflow.providers.http.hooks.http.HttpAsyncHook`: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.hooks.http import HttpAsyncHook + from airflow.sdk import dag, task + + + @dag( + dag_id="it-async-http-pokemon-example", + start_date=datetime(2026, 1, 1), + ) + def it_async_http_pokemon_example(): + @task + def list_pokemon() -> list[str]: + response = HttpHook( + http_conn_id="pokeapi", + method="GET", + ).run( + endpoint="api/v2/pokemon?limit=100", + ) + + return [pokemon["url"].replace("https://pokeapi.co/", "") for pokemon in response.json()["results"]] + + @task( + retries=3, + task_concurrency=2, + show_return_value_in_logs=False, + ) + async def get_pokemon(url: str): + async with HttpAsyncHook( + http_conn_id="pokeapi", + method="GET", + ).session() as session: + response = await session.run(endpoint=url) + return await response.json() + + get_pokemon.iterate( + url=list_pokemon(), + ) + + + it_async_http_pokemon_example() + + +When ``iterate()`` is used with an async task, all iterations share the same +event loop, enabling true multiplexing of I/O-bound operations without any +manual concurrency management by the DAG author. For 5 Pokémon the +difference is negligible, but for hundreds or thousands of items the +concurrent approach is dramatically faster — see the +:ref:`benchmarks above `. + +Why Task Iteration? +--------------------------- + +TI is designed to address limitations of Dynamic Task Mapping in specific scenarios: + +- **Scheduler scalability**: + DTM creates one Task Instance per item, which can put pressure on the scheduler + for very large datasets. TI avoids this by keeping execution within a single task. + +- **Async multiplexing**: + With Python-native async support in Airflow 3.2, TI allows multiple + operations to share the same event loop within a single Task Instance. + This enables efficient multiplexing of I/O-bound workloads. + +- **Lower overhead**: + No need to serialize, schedule, and track thousands of Task Instances. + +- **Triggerer and deferrable-operator bottleneck**: + Deferrable operators delegate async work to triggerers, which store yielded + events directly in the Airflow metadata database. Unlike workers, triggerers + cannot leverage a custom XCom backend to offload large payloads. This makes + triggerers a bottleneck for sustained high-load async execution or workloads + that return large results. Dynamic Task Mapping with deferrable operators + amplifies the problem further. TI sidesteps triggerers entirely — iterations + execute on workers, which scale more effectively and support custom XCom + backends. + + For more on deferred vs async trade-offs, see :doc:`deferred-vs-async-operators`. + +TI is especially useful for patterns such as: + +- API pagination +- Bulk HTTP or database calls +- High-throughput async workloads +- Streaming or lazily-evaluated XCom results + +Hooks as Building Blocks +^^^^^^^^^^^^^^^^^^^^^^^^ + +TI encourages a pattern where DAG authors call **hooks** directly from +``@task``-decorated functions rather than relying on operators. Operators are +wrappers around hooks and sometimes expose only a subset of the hook's +capabilities. By calling hooks directly, users gain full control over +concurrency, error handling, and batching. + +For example, instead of using ``HttpOperator`` in deferrable mode (which +delegates to the triggerer for a single request at a time), an async +``@task`` can call :class:`~airflow.providers.http.hooks.http.HttpAsyncHook` +directly to perform many concurrent requests. With TI, the framework +handles the iteration, concurrency, and event-loop management +automatically — the DAG author only writes the per-item logic. + +This "hooks as building blocks" approach is especially powerful with async +hooks, where the shared event loop enables concurrent I/O without any +manual ``asyncio.gather`` or ``asyncio.Semaphore`` management. + +For more examples of calling async hooks directly from tasks, see +:doc:`deferred-vs-async-operators`. + +Comparison +---------- + +.. list-table:: + :header-rows: 1 + + * - Aspect + - Dynamic Task Mapping (DTM) + - Task Iteration (TI) + * - Task Instances + - One per item + - Single Task Instance + * - Scheduler load + - High for large iterables + - Minimal + * - Execution model + - Distributed across workers + - In-process iteration + * - Concurrency + - Parallel tasks + - Sync or async within one task + * - Async support + - Limited (per task) + - Strong (shared event loop, multiplexing) + * - Retry behavior + - Per item + - Entire task retries + * - Observability + - Per item in UI + - Aggregated in a single task + * - Triggerer dependency + - Deferrable mapped tasks rely on triggerers + - No triggerers involved + * - XCom backend + - Workers support custom XCom backends + - Workers support custom XCom backends (triggerers do not) + * - Use case + - Independent, trackable units of work + - High-throughput or streaming workloads + +The following table illustrates these differences using the Pokémon example from above: + +.. list-table:: + :header-rows: 1 + + * - Pattern + - Task Instances + - Work Per Task + * - ``get_pokemon.expand(url=urls)`` + - 100 + - 1 Pokémon + * - ``get_pokemon.iterate(url=urls)`` + - 1 + - 100 Pokémon + * - ``get_pokemon.batch(size=2).iterate(url=urls)`` + - 2 + - ~50 Pokémon each + +When to Use Dynamic Task Mapping +-------------------------------- + +Prefer DTM when: + +- Each item must be independently tracked in the UI. +- You need fine-grained retries per item. +- Tasks are long-running or resource-intensive. +- Work should be distributed across multiple workers. +- Scheduling decisions should be made per item. + +When to Use Task Iteration +----------------------------------- + +Prefer TI when: + +- You are processing large numbers of small items. +- Scheduler overhead becomes a concern. +- You are using async operators and want to leverage a shared event loop. +- Workloads are I/O-bound and benefit from multiplexing. +- Fine-grained observability per item is not required. + +When **not** to use TI +----------------------- + +Avoid Task Iteration when: + +- You need per-item retries or failure isolation. +- Each item represents a long-running or heavy computation. +- You require detailed visibility per item in the Airflow UI. +- Work must be distributed across multiple worker nodes. + +.. tip:: + + TI is a **third execution option** alongside Dynamic Task Mapping and + deferrable operators. It is not intended as a replacement for either. + Triggerers remain the right choice for long-running polling or waiting tasks + (e.g., monitoring a remote job or waiting for a Kubernetes pod to complete). + +Combining DTM and TI (Dynamic Task Batching) +--------------------------------------------- + +DTM and TI are not mutually exclusive in principle. The *Dynamic Task Batching* +pattern could use DTM to split a large dataset into coarse-grained chunks, +where each mapped task processes it's batch using TI. + +For example, downloading 17,000 files could be split into 17 batches of +1,000 files each. DTM would create one task per batch, and TI would iterate +within each batch using a shared event loop for concurrent I/O. + +This pattern would provide: + +- **Coarse-grained retry**: if a batch fails, only that batch is retried — not all 17,000 items. +- **Reduced scheduler load**: the scheduler manages chunks (e.g., 17 tasks) instead of individual items (17,000 tasks). +- **High throughput within each chunk**: async I/O processes items concurrently inside each task. + +Relationship with Async Operators +---------------------------------- + +TI complements async operators introduced in Airflow 3.2. + +- Async operators allow concurrent I/O within a single task. +- TI allows you to *apply an operator repeatedly* over a dataset within that same task. + +Together, they enable patterns such as: + +- Efficient API pagination +- Concurrent request batching +- Streaming data processing + +Unlike Dynamic Task Mapping, where each mapped task runs in its own execution context, +TI allows all iterations to share the same event loop, enabling true multiplexing. + +Because TI executes on workers rather than triggerers, it also benefits from the +full worker environment: custom XCom backends, Edge Worker support, and the +scalability of execution frameworks such as Celery. + +For more details on async execution, see :doc:`deferred-vs-async-operators`. + +Future Outlook +-------------- + +As Python's async ecosystem evolves, TI tasks will benefit from improved +introspection and tooling. For example, Python 3.14 introduces new +`asyncio introspection capabilities `_ +that could eventually enable structured progress reporting in the Airflow UI +for TI tasks — providing per-item visibility without the overhead of per-item +task instances. diff --git a/task-sdk/docs/index.rst b/task-sdk/docs/index.rst index 18e215d4ce7d7..1b29d4de9a023 100644 --- a/task-sdk/docs/index.rst +++ b/task-sdk/docs/index.rst @@ -172,6 +172,7 @@ For the full public API reference, see the :doc:`api` page. examples dynamic-task-mapping + dynamic-task-mapping-vs-iteration deferred-vs-async-operators resumable-job-mixin api diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index e3c64eb0a8a69..67811bb6bf9b9 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/task-sdk/src/airflow/sdk/bases/decorator.py @@ -28,14 +28,10 @@ import attr import typing_extensions -from airflow.sdk import TriggerRule, timezone +from airflow.sdk import TriggerRule from airflow.sdk.bases.operator import ( BASEOPERATOR_ARGS_EXPECTED_TYPES, BaseOperator, - coerce_resources, - coerce_timedelta, - get_merged_defaults, - parse_retries, ) from airflow.sdk.definitions._internal.contextmanager import DagContext, TaskGroupContext from airflow.sdk.definitions._internal.decorators import remove_task_decorator @@ -50,7 +46,6 @@ from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS from airflow.sdk.definitions.mappedoperator import ( MappedOperator, - ensure_xcomarg_return_value, prevent_duplicates, ) from airflow.sdk.definitions.xcom_arg import XComArg @@ -61,6 +56,7 @@ OperatorExpandArgument, OperatorExpandKwargsArgument, ) + from airflow.sdk.definitions.batchedoperator import DecoratedBatchedOperator from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.dag import DAG from airflow.sdk.definitions.mappedoperator import ValidationSource @@ -587,110 +583,29 @@ def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") return self._expand(ListOfDictsExpandInput(kwargs), strict=strict) - def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg: - ensure_xcomarg_return_value(expand_input.value) - - task_kwargs = self.kwargs.copy() - dag = task_kwargs.pop("dag", None) or DagContext.get_current() - task_group = task_kwargs.pop("task_group", None) or TaskGroupContext.get_current(dag) - - default_args, partial_params = get_merged_defaults( - dag=dag, - task_group=task_group, - task_params=task_kwargs.pop("params", None), - task_default_args=task_kwargs.pop("default_args", None), + def _expand( + self, + expand_input: ExpandInput, + *, + strict: bool, + register_with_dag: bool = True, + ) -> XComArg: + operator = self.batch(size=0)._expand( + expand_input, strict=strict, register_with_dag=register_with_dag ) - partial_kwargs: dict[str, Any] = { - "is_setup": self.is_setup, - "is_teardown": self.is_teardown, - "on_failure_fail_dagrun": self.on_failure_fail_dagrun, - } - base_signature = inspect.signature(BaseOperator) - ignore = { - "default_args", # This is target we are working on now. - "kwargs", # A common name for a keyword argument. - "do_xcom_push", # In the same boat as `multiple_outputs` - "multiple_outputs", # We will use `self.multiple_outputs` instead. - "params", # Already handled above `partial_params`. - "task_concurrency", # Deprecated(replaced by `max_active_tis_per_dag`). - } - partial_keys = set(base_signature.parameters) - ignore - partial_kwargs.update({key: value for key, value in default_args.items() if key in partial_keys}) - partial_kwargs.update(task_kwargs) - - task_id = get_unique_task_id(partial_kwargs.pop("task_id"), dag, task_group) - if task_group: - task_id = task_group.child_id(task_id) - - # Logic here should be kept in sync with BaseOperatorMeta.partial(). - if partial_kwargs.get("wait_for_downstream"): - partial_kwargs["depends_on_past"] = True - start_date = timezone.convert_to_utc(partial_kwargs.pop("start_date", None)) - end_date = timezone.convert_to_utc(partial_kwargs.pop("end_date", None)) - if "pool_slots" in partial_kwargs: - if partial_kwargs["pool_slots"] < 1: - dag_str = "" - if dag: - dag_str = f" in dag {dag.dag_id}" - raise ValueError(f"pool slots for {task_id}{dag_str} cannot be less than 1") - - for fld, convert in ( - ("retries", parse_retries), - ("retry_delay", coerce_timedelta), - ("max_retry_delay", coerce_timedelta), - ("resources", coerce_resources), - ): - if (v := partial_kwargs.get(fld, NOTSET)) is not NOTSET: - partial_kwargs[fld] = convert(v) + return XComArg(operator=operator) - partial_kwargs.setdefault("executor_config", {}) - partial_kwargs.setdefault("op_args", []) - partial_kwargs.setdefault("op_kwargs", {}) + def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> XComArg: + return self.batch(size=0).iterate(**mapped_kwargs) - # Mypy does not work well with a subclassed attrs class :( - _MappedOperator = cast("Any", DecoratedMappedOperator) + def iterate_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True) -> XComArg: + return self.batch(size=0).iterate_kwargs(kwargs, strict=strict) - try: - operator_name = self.operator_class.custom_operator_name # type: ignore - except AttributeError: - operator_name = self.operator_class.__name__ - - operator = _MappedOperator( - operator_class=self.operator_class, - expand_input=EXPAND_INPUT_EMPTY, # Don't use this; mapped values go to op_kwargs_expand_input. - partial_kwargs=partial_kwargs, - task_id=task_id, - params=partial_params, - operator_extra_links=self.operator_class.operator_extra_links, - template_ext=self.operator_class.template_ext, - template_fields=self.operator_class.template_fields, - template_fields_renderers=self.operator_class.template_fields_renderers, - ui_color=self.operator_class.ui_color, - ui_fgcolor=self.operator_class.ui_fgcolor, - is_empty=False, - is_sensor=self.operator_class._is_sensor, - can_skip_downstream=self.operator_class._can_skip_downstream, - is_stub=self.operator_class.is_stub, - task_module=self.operator_class.__module__, - task_type=self.operator_class.__name__, - operator_name=operator_name, - dag=dag, - task_group=task_group, - start_date=start_date, - end_date=end_date, - multiple_outputs=self.multiple_outputs, - python_callable=self.function, - op_kwargs_expand_input=expand_input, - disallow_kwargs_override=strict, - # Different from classic operators, kwargs passed to a taskflow - # task's expand() contribute to the op_kwargs operator argument, not - # the operator arguments themselves, and should expand against it. - expand_input_attr="op_kwargs_expand_input", - start_trigger_args=self.operator_class.start_trigger_args, - start_from_trigger=self.operator_class.start_from_trigger, - returns_dag_result=self.returns_dag_result, - ) - return XComArg(operator=operator) + def batch(self, size: int) -> DecoratedBatchedOperator: + """Return a DecoratedBatchedOperator for batched mapping.""" + from airflow.sdk.definitions.batchedoperator import DecoratedBatchedOperator + + return DecoratedBatchedOperator(operator_partial=self, size=size) def partial(self, **kwargs: Any) -> _TaskDecorator[FParams, FReturn, OperatorSubclass]: self._validate_arg_names("partial", kwargs) diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py index 97da8869686cc..f1a1a411b984e 100644 --- a/task-sdk/src/airflow/sdk/bases/operator.py +++ b/task-sdk/src/airflow/sdk/bases/operator.py @@ -22,6 +22,7 @@ import collections.abc import contextlib import copy +import functools import inspect import sys import warnings @@ -221,6 +222,13 @@ def event_loop() -> Generator[AbstractEventLoop]: asyncio.set_event_loop(None) +# TODO: Once AIP-88 is implemented, multiple events could be returned +async def run_trigger(trigger: BaseTrigger) -> Any | None: + async for event in trigger.run(): + return event + return None + + class _PartialDescriptor: """A descriptor that guards against ``.partial`` being called on Task objects.""" @@ -305,6 +313,7 @@ def partial( map_index_template: str | None = ..., max_active_tis_per_dag: int | None = ..., max_active_tis_per_dagrun: int | None = ..., + task_concurrency: int | None = ..., on_execute_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] = ..., on_failure_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] = ..., on_success_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] = ..., @@ -375,8 +384,6 @@ def partial( partial_kwargs.update((k, v) for k, v in OPERATOR_DEFAULTS.items() if k not in partial_kwargs) # Post-process arguments. Should be kept in sync with _TaskDecorator.expand(). - if "task_concurrency" in kwargs: # Reject deprecated option. - raise TypeError("unexpected argument: task_concurrency") if start_date := partial_kwargs.get("start_date", None): partial_kwargs["start_date"] = timezone.convert_to_utc(start_date) if end_date := partial_kwargs.get("end_date", None): @@ -819,6 +826,8 @@ class derived from this one results in the creation of a task object, key in the returned dictionary result. If False and do_xcom_push is True, pushes a single XCom. :param task_group: The TaskGroup to which the task should belong. This is typically provided when not using a TaskGroup as a context manager. + :param task_concurrency: The maximum number of threads that will be used when the operator is used + with Dynamic Task Iteration (default is the number of threads available on the executor). :param doc: Add documentation or notes to your Task objects that is visible in Task Instance details View in the Webserver :param doc_md: Add documentation (in Markdown format) or notes to your Task objects @@ -1075,6 +1084,7 @@ def __init__( inlets: Any | None = None, outlets: Any | None = None, task_group: TaskGroup | None = None, + task_concurrency: int | None = None, doc: str | None = None, doc_md: str | None = None, doc_json: str | None = None, @@ -1098,6 +1108,7 @@ def __init__( super().__init__() self.task_group = task_group + self.task_concurrency = task_concurrency kwargs.pop("_airflow_mapped_validation_only", None) if kwargs: @@ -1544,6 +1555,7 @@ def get_serialized_fields(cls): "_BaseOperator__from_mapped", "on_failure_fail_dagrun", "task_group", + "task_concurrency", "_task_type", "operator_extra_links", "on_execute_callback", @@ -1676,12 +1688,10 @@ def defer( raise TaskDeferred(trigger=trigger, method_name=method_name, kwargs=kwargs, timeout=timeout) - def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None, context: Context): + def next_callable(self, next_method: str, next_kwargs: dict[str, Any] | None): """Entrypoint method called by the Task Runner (instead of execute) when this task is resumed.""" from airflow.sdk.exceptions import TaskDeferralError, TaskDeferralTimeout - if next_kwargs is None: - next_kwargs = {} # __fail__ is a special signal value for next_method that indicates # this task was scheduled specifically to fail. @@ -1695,7 +1705,17 @@ def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None, raise TaskDeferralError(error) # Grab the callable off the Operator/Task and add in any kwargs execute_callable = getattr(self, next_method) - return execute_callable(context, **next_kwargs) + if next_kwargs: + return functools.partial(execute_callable, **next_kwargs) + return execute_callable + + def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None, context: Context): + """Entrypoint method called by the Task Runner (instead of execute) when this task is resumed.""" + if next_kwargs is None: + next_kwargs = {} + + execute_callable = self.next_callable(next_method, next_kwargs) + return execute_callable(context) def dry_run(self) -> None: """Perform dry run for the operator - just render template fields.""" diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py b/task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py index f0fbdc306e6f3..769569ebae897 100644 --- a/task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py +++ b/task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py @@ -19,6 +19,7 @@ import sys from collections import deque +from contextvars import ContextVar from types import ModuleType from typing import TYPE_CHECKING, Any, Generic, TypeVar @@ -32,19 +33,19 @@ __all__ = ["DagContext", "TaskGroupContext"] -# This is a global variable that stores the current Task context. -# It is used to push the Context dictionary when Task starts execution -# and it is used to retrieve the current context in PythonOperator or Taskflow API via -# the `get_current_context` function. -_CURRENT_CONTEXT: list[Context] = [] +# ContextVar storing the current task-execution context stack. Each asyncio task and each ThreadPoolExecutor +# worker receives its own isolated copy (Python 3.7+ copy-on-submit semantics), so concurrent sub-tasks never +# observe each other's pushed context. +_CURRENT_CONTEXT: ContextVar[list[Context] | None] = ContextVar("_current_context", default=None) def _get_current_context() -> Context: - if not _CURRENT_CONTEXT: + stack = _CURRENT_CONTEXT.get(None) + if not stack: raise RuntimeError( "Current context was requested but no context was found! Are you running within an Airflow task?" ) - return _CURRENT_CONTEXT[-1] + return stack[-1] # In order to add a `@classproperty`-like thing we need to define a property on a metaclass. diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py index b6ffbd2214253..1809d6e342cd4 100644 --- a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py +++ b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py @@ -17,21 +17,20 @@ # under the License. from __future__ import annotations -from collections.abc import Iterable, Mapping, Sequence, Sized +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Sized from typing import TYPE_CHECKING, Any, ClassVar, Union import attrs from airflow.sdk.definitions._internal.mixins import ResolveMixin +from airflow.sdk.definitions.xcom_arg import XComArg if TYPE_CHECKING: from typing import TypeGuard - from airflow.sdk.definitions.xcom_arg import XComArg from airflow.sdk.types import Operator -ExpandInput = Union["DictOfListsExpandInput", "ListOfDictsExpandInput"] - # Each keyword argument to expand() can be an XComArg, sequence, or dict (not # any mapping since we need the value to be ordered). OperatorExpandArgument = Union["MappedArgument", "XComArg", Sequence, dict[str, Any]] @@ -79,6 +78,95 @@ def _needs_run_time_resolution(v: OperatorExpandArgument) -> TypeGuard[MappedArg return isinstance(v, (MappedArgument, XComArg)) +def count(expand_input: ExpandInput, iterable: Iterable[Any]) -> Iterable[Any]: + expand_input._length = None + counter = 0 + + for item in iterable: + counter += 1 + yield item + + expand_input._length = counter + + +@attrs.define(slots=False) +class ExpandInput(ABC, ResolveMixin): + EXPAND_INPUT_TYPE: ClassVar[str] + + def __attrs_post_init__(self) -> None: + self._length: int | None = None + + @property + @abstractmethod + def value(self) -> Any: + """The value of the expand input.""" + ... + + def iter_values(self, context: Mapping[str, Any]) -> Iterable[Any]: + raise NotImplementedError() + + def resolve(self, context: Mapping[str, Any]) -> Any: + raise NotImplementedError() + + def __len__(self) -> int: + if self._length is None: + raise RuntimeError(f"Length of {type(self).__name__} is not yet known") + return self._length + + +class DecoratedExpandInput(ExpandInput): + EXPAND_INPUT_TYPE: ClassVar[str] = "decorated" + + def __init__(self, expand_input: ExpandInput): + super().__init__() + self.delegate = expand_input + + @property + def value(self) -> Any: + return self.delegate.value + + def iter_references(self) -> Iterable[tuple[Operator, str]]: + return self.delegate.iter_references() + + def iter_values(self, context: Mapping[str, Any]) -> Iterable[dict]: + return count( + self, + map(lambda value: {"op_kwargs": value}, self.delegate.iter_values(context)), + ) + + def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any], set[int]]: + return self.delegate.resolve(context) + + +class BatchedExpandInput(DecoratedExpandInput): + """ + ExpandInput that batches another ExpandInput into N chunks. + + This affects mapping cardinality, NOT resolve-time behavior. + """ + + EXPAND_INPUT_TYPE: ClassVar[str] = "batched" + + def __init__(self, expand_input: ExpandInput, size: int): + if size < 2: + raise ValueError(f"batch size must be at least 2, got {size}") + + super().__init__(expand_input=expand_input) + self.size = size + + def iter_values(self, context: Mapping[str, Any]) -> Iterable[dict]: + map_index = context["ti"].map_index + + return count( + self, + ( + item + for index, item in enumerate(self.delegate.iter_values(context)) + if index % self.size == map_index + ), + ) + + @attrs.define(kw_only=True) class MappedArgument(ResolveMixin): """ @@ -107,7 +195,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any: @attrs.define() -class DictOfListsExpandInput(ResolveMixin): +class DictOfListsExpandInput(ExpandInput): """ Storage type of a mapped operator's mapped kwargs. @@ -184,6 +272,43 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]: if isinstance(x, XComArg): yield from x.iter_references() + def iter_values(self, context: Mapping[str, Any]) -> Iterable[Any]: + from airflow.sdk.definitions.xcom_arg import XComArg + + def _to_iterable(v: Any) -> Iterable: + return v if hasattr(v, "__iter__") and not isinstance(v, (str, bytes)) else (v,) + + def _make_factory(v: Any) -> Callable[[], Iterable]: + # Capture v (already bound to self.value[k]) so each factory closes + # over its own value rather than a shared loop variable. + def factory() -> Iterable: + resolved = v.resolve(context) if isinstance(v, XComArg) else v + return _to_iterable(resolved) + + return factory + + def _lazy_product(*factories: Callable[[], Iterable]) -> Iterator[tuple]: + """ + Streaming cross-product with fully deferred resolution. + + Each factory is called to produce a fresh iterable only when that + position is first needed. The first factory is called exactly once; + each subsequent factory is called once per element yielded by all + preceding factories combined, so XComArg sources are resolved (and + pages re-fetched) on demand rather than materialized upfront. + """ + if not factories: + yield () + return + first_factory, *rest_factories = factories + for item in first_factory(): + for tail in _lazy_product(*rest_factories): + yield (item, *tail) + + keys = list(self.value) + factories = [_make_factory(self.value[k]) for k in keys] + return count(self, (dict(zip(keys, combo)) for combo in _lazy_product(*factories))) + def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any], set[int]]: map_index: int | None = context["ti"].map_index if map_index is None or map_index < 0: @@ -217,7 +342,7 @@ def _describe_type(value: Any) -> str: @attrs.define() -class ListOfDictsExpandInput(ResolveMixin): +class ListOfDictsExpandInput(ExpandInput): """ Storage type of a mapped operator's mapped kwargs. @@ -238,12 +363,25 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]: if isinstance(x, XComArg): yield from x.iter_references() + def iter_values(self, context: Mapping[str, Any]) -> Iterable[Any]: + def iterate(): + if isinstance(self.value, XComArg): + for item in self.value.resolve(context): + yield item + else: + for item in self.value: + if isinstance(item, XComArg): + yield from item.resolve(context) + else: + yield item + + return count(self, iterate()) + def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any], set[int]]: map_index = context["ti"].map_index - if map_index < 0: + if map_index is None or map_index < 0: raise RuntimeError("can't resolve task-mapping argument without expanding") - mapping: Any = None if isinstance(self.value, Sized): mapping = self.value[map_index] if not isinstance(mapping, Mapping): diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py b/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py index e186ef97e64dd..e47846499f9b8 100644 --- a/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py +++ b/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py @@ -135,6 +135,24 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]: """ raise NotImplementedError + def iter_values(self, context: Context) -> Iterable[Any]: + """ + Yield individual values for task expansion during Dynamic Task Iteration. + + Called by :class:`~airflow.sdk.definitions.iterableoperator.IterableOperator` during + execution to enumerate all the items this expand input resolves to. Each yielded value + represents the keyword arguments passed to one mapped task instance. Any deferred XCom + references embedded in the expand input are resolved against the provided runtime + context before being yielded. + + :param context: The runtime task execution context used to resolve XCom references. + Must expose the current task instance under the ``ti`` key so that upstream + XCom values can be fetched. + + :meta private: + """ + raise NotImplementedError() + def resolve(self, context: Context) -> Any: """ Resolve this value for runtime. diff --git a/task-sdk/src/airflow/sdk/definitions/batchedoperator.py b/task-sdk/src/airflow/sdk/definitions/batchedoperator.py new file mode 100644 index 0000000000000..44eb20bb78c35 --- /dev/null +++ b/task-sdk/src/airflow/sdk/definitions/batchedoperator.py @@ -0,0 +1,512 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import inspect +from abc import ABCMeta, abstractmethod +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +import attrs + +from airflow.sdk import TriggerRule, timezone +from airflow.sdk.bases.decorator import ( + DecoratedMappedOperator, + FParams, + FReturn, + OperatorSubclass, + _TaskDecorator, + get_unique_task_id, +) +from airflow.sdk.bases.operator import ( + BaseOperator, + coerce_resources, + coerce_timedelta, + get_merged_defaults, + parse_retries, +) +from airflow.sdk.definitions._internal.contextmanager import ( + DagContext, + TaskGroupContext, +) +from airflow.sdk.definitions._internal.expandinput import ( + EXPAND_INPUT_EMPTY, + DecoratedExpandInput, + DictOfListsExpandInput, + ExpandInput, + ListOfDictsExpandInput, + OperatorExpandArgument, + OperatorExpandKwargsArgument, +) +from airflow.sdk.definitions._internal.types import NOTSET +from airflow.sdk.definitions.mappedoperator import ( + MappedOperator, + OperatorPartial, + ensure_xcomarg_return_value, + prevent_duplicates, + validate_mapping_kwargs, +) +from airflow.sdk.definitions.xcom_arg import XComArg + +if TYPE_CHECKING: + from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator + from airflow.sdk.definitions.mappedoperator import ValidationSource + from airflow.sdk.definitions.param import ParamsDict + +T = TypeVar("T", bound=OperatorPartial | _TaskDecorator) + + +@attrs.define(kw_only=True, repr=False) +class BatchableOperator(Generic[T], metaclass=ABCMeta): + """ + Intermediate abstraction for batched mapping. + + This class decorates an OperatorPartial and stores configuration for batched mapping. + It is used to facilitate batched expansion of operators, allowing tasks to be mapped over batches + of data and then iterate over the batched data. + + :param operator_partial: The partial operator to be batched. + :param size: The number of batches to create. + """ + + operator_partial: T + size: int + + @property + def operator_class(self) -> type[BaseOperator]: + return self.operator_partial.operator_class + + @property + def kwargs(self) -> dict[str, Any]: + return self.operator_partial.kwargs + + @abstractmethod + def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> Any: + """ + Iterate the operator over the provided mapped keyword arguments. + + :param mapped_kwargs: Keyword arguments to expand against. + :return: An expanded operator or XComArg, depending on the subclass implementation. + """ + + @abstractmethod + def iterate_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True) -> Any: + """ + Iterate the operator over a list of dictionaries or XComArg. + + :param kwargs: List of dicts or XComArg to expand against. + :param strict: Whether to enforce strict argument checking. + :return: An expanded operator or XComArg, depending on the subclass implementation. + """ + + @abstractmethod + def _iterate( + self, + expand_input: ExpandInput, + *, + strict: bool, + ) -> IterableOperator | MappedIterableOperator: + """ + Create an iterable operator for the given expansion input. + + This method calls the _expand method first to get a MappedOperator based on expansion input, + then wraps it in either an IterableOperator or MappedIterableOperator depending on the batch size. + + :param expand_input: The input to iterate against. + :param strict: Whether to enforce strict argument checking. + :return: An IterableOperator or MappedIterableOperator. + """ + + @abstractmethod + def _expand( + self, + expand_input: ExpandInput, + *, + strict: bool, + register_with_dag: bool = True, + ) -> MappedOperator: + """ + Create a mapped operator for the given expansion input. + + :param expand_input: The input to expand against. + :param strict: Whether to enforce strict argument checking. + :param register_with_dag: Whether to apply upstream relationships. + :return: A MappedOperator instance. + """ + + +@attrs.define(kw_only=True, repr=False) +class BatchedOperator(BatchableOperator[OperatorPartial]): + """ + Concrete implementation of BatchableOperator for classic (non-decorated) operators. + + This class wraps an OperatorPartial and provides batched expansion and iteration logic + for classic Airflow operators. It enables mapping tasks over batches of data, supporting + both direct expansion via keyword arguments and expansion via a list of dictionaries or XComArg. + + :param operator_partial: The OperatorPartial instance to be batched and expanded. + :param size: The number of batches to create for mapping. + """ + + @property + def params(self) -> ParamsDict | dict: + return self.operator_partial.params + + @property + def _expand_called(self) -> bool: + return self.operator_partial._expand_called + + @_expand_called.setter + def _expand_called(self, value: bool) -> None: + self.operator_partial._expand_called = value + + def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> IterableOperator | MappedIterableOperator: + if not mapped_kwargs: + raise TypeError("no arguments to iterate against") + + validate_mapping_kwargs(self.operator_class, "iterate", mapped_kwargs) + prevent_duplicates( + self.kwargs, + mapped_kwargs, + fail_reason="unmappable or already specified", + ) + # Since the input is already checked at parse time, we can set strict + # to False to skip the checks on execution. + expand_input = DictOfListsExpandInput(mapped_kwargs) + return self._iterate(expand_input, strict=False) + + def iterate_kwargs( + self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True + ) -> IterableOperator | MappedIterableOperator: + if isinstance(kwargs, Sequence): + for item in kwargs: + if not isinstance(item, (XComArg, Mapping)): + raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") + elif not isinstance(kwargs, XComArg): + raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") + + expand_input = ListOfDictsExpandInput(kwargs) + return self._iterate(expand_input, strict=strict) + + def _iterate( + self, + expand_input: ExpandInput, + *, + strict: bool, + ) -> IterableOperator | MappedIterableOperator: + from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator + + operator = self._expand(expand_input, strict=strict, register_with_dag=False) + + if self.size > 1: + return MappedIterableOperator( + mapped_operator=operator, + expand_input=expand_input, + batch_size=self.size, + ) + return IterableOperator( + operator=operator, + expand_input=expand_input, + ) + + def _expand( + self, + expand_input: ExpandInput, + *, + strict: bool, + register_with_dag: bool = True, + ) -> MappedOperator: + from airflow.providers.standard.operators.empty import EmptyOperator + from airflow.sdk import BaseSensorOperator + from airflow.sdk.bases.skipmixin import SkipMixin + + ensure_xcomarg_return_value(expand_input.value) + + partial_kwargs = self.kwargs.copy() + task_id = partial_kwargs.pop("task_id") + dag = partial_kwargs.pop("dag") + task_group = partial_kwargs.pop("task_group") + start_date = partial_kwargs.pop("start_date", None) + end_date = partial_kwargs.pop("end_date", None) + start_from_trigger = ( + partial_kwargs["start_from_trigger"] + if "start_from_trigger" in partial_kwargs + else getattr(self.operator_class, "start_from_trigger", False) + ) + start_trigger_args = ( + partial_kwargs["start_trigger_args"] + if "start_trigger_args" in partial_kwargs + else getattr(self.operator_class, "start_trigger_args", None) + ) + + try: + operator_name = self.operator_class.custom_operator_name # type: ignore + except AttributeError: + operator_name = self.operator_class.__name__ + + return MappedOperator( + operator_class=self.operator_class, + expand_input=expand_input, + partial_kwargs=partial_kwargs, + task_id=task_id, + params=self.params, + operator_extra_links=self.operator_class.operator_extra_links, + template_ext=self.operator_class.template_ext, + template_fields=self.operator_class.template_fields, + template_fields_renderers=self.operator_class.template_fields_renderers, + ui_color=self.operator_class.ui_color, + ui_fgcolor=self.operator_class.ui_fgcolor, + is_empty=issubclass(self.operator_class, EmptyOperator), + is_sensor=issubclass(self.operator_class, BaseSensorOperator), + can_skip_downstream=issubclass(self.operator_class, SkipMixin), + is_stub=self.operator_class.is_stub, + task_module=self.operator_class.__module__, + task_type=self.operator_class.__name__, + operator_name=operator_name, + dag=dag, + task_group=task_group, + start_date=start_date, + end_date=end_date, + disallow_kwargs_override=strict, + # For classic operators, this points to expand_input because kwargs + # to BaseOperator.expand() contribute to operator arguments. + expand_input_attr="expand_input", + # TODO: Move these to task SDK's BaseOperator and remove getattr + start_trigger_args=start_trigger_args, + start_from_trigger=start_from_trigger, + register_with_dag=register_with_dag, + ) + + +@attrs.define(kw_only=True, repr=False) +class DecoratedBatchedOperator(BatchableOperator[_TaskDecorator]): + """ + Concrete implementation of BatchableOperator for decorated (TaskFlow) operators. + + This class wraps a _TaskDecorator and provides batched expansion and iteration logic + for TaskFlow-style decorated Airflow operators. It enables mapping decorated tasks over + batches of data, returning XComArg objects for downstream dependencies and supporting + both direct expansion via keyword arguments and expansion via a list of dictionaries or XComArg. + + :param operator_partial: The _TaskDecorator instance to be batched and expanded. + :param size: The number of batches to create for mapping. + """ + + @property + def is_setup(self) -> bool: + return self.operator_partial.is_setup + + @property + def is_teardown(self) -> bool: + return self.operator_partial.is_teardown + + @property + def function(self) -> Callable[FParams, FReturn]: + return self.operator_partial.function + + @property + def operator_class(self) -> type[OperatorSubclass]: + return self.operator_partial.operator_class + + @property + def multiple_outputs(self) -> bool: + return self.operator_partial.multiple_outputs + + @property + def on_failure_fail_dagrun(self) -> bool: + return self.operator_partial.on_failure_fail_dagrun + + def _validate_arg_names(self, func: ValidationSource, kwargs: dict[str, Any]): + self.operator_partial._validate_arg_names(func, kwargs) + + @property + def returns_dag_result(self) -> bool: + return self.operator_partial.returns_dag_result + + def iterate(self, **map_kwargs: OperatorExpandArgument) -> XComArg: + if self.kwargs.get("trigger_rule") == TriggerRule.ALWAYS and any( + [isinstance(expanded, XComArg) for expanded in map_kwargs.values()] + ): + raise ValueError( + "Task-generated iterating within a task using 'iterate' is not allowed with trigger rule 'always'." + ) + if not map_kwargs: + raise TypeError("no arguments to expand against") + self._validate_arg_names("expand", map_kwargs) + prevent_duplicates(self.kwargs, map_kwargs, fail_reason="mapping already partial") + # Since the input is already checked at parse time, we can set strict + # to False to skip the checks on execution. + if self.is_teardown: + if "trigger_rule" in self.kwargs: + raise ValueError("Trigger rule not configurable for teardown tasks.") + self.kwargs.update(trigger_rule=TriggerRule.ALL_DONE_SETUP_SUCCESS) + expand_input = DictOfListsExpandInput(map_kwargs) + operator = self._iterate(expand_input, strict=False) + return XComArg(operator=operator) + + def iterate_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True) -> XComArg: + if ( + self.kwargs.get("trigger_rule") == TriggerRule.ALWAYS + and not isinstance(kwargs, XComArg) + and any( + [ + isinstance(v, XComArg) + for kwarg in kwargs + if not isinstance(kwarg, XComArg) + for v in kwarg.values() + ] + ) + ): + raise ValueError( + "Task-generated iterating within a task using 'iterate_kwargs' is not allowed with trigger rule 'always'." + ) + if isinstance(kwargs, Sequence): + for item in kwargs: + if not isinstance(item, (XComArg, Mapping)): + raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") + elif not isinstance(kwargs, XComArg): + raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") + expand_input = ListOfDictsExpandInput(kwargs) + operator = self._iterate(expand_input, strict=strict) + return XComArg(operator=operator) + + def _iterate( + self, + expand_input: ExpandInput, + *, + strict: bool, + ) -> IterableOperator | MappedIterableOperator: + from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator + + operator = self._expand(expand_input, strict=strict, register_with_dag=False) + + if self.size > 1: + return MappedIterableOperator( + mapped_operator=operator, + expand_input=DecoratedExpandInput(expand_input), + batch_size=self.size, + ) + return IterableOperator(operator=operator, expand_input=DecoratedExpandInput(expand_input)) + + def _expand( + self, + expand_input: ExpandInput, + *, + strict: bool, + register_with_dag: bool = True, + ) -> MappedOperator: + ensure_xcomarg_return_value(expand_input.value) + + task_kwargs = self.kwargs.copy() + dag = task_kwargs.pop("dag", None) or DagContext.get_current() + task_group = task_kwargs.pop("task_group", None) or TaskGroupContext.get_current(dag) + + default_args, partial_params = get_merged_defaults( + dag=dag, + task_group=task_group, + task_params=task_kwargs.pop("params", None), + task_default_args=task_kwargs.pop("default_args", None), + ) + partial_kwargs: dict[str, Any] = { + "is_setup": self.is_setup, + "is_teardown": self.is_teardown, + "on_failure_fail_dagrun": self.on_failure_fail_dagrun, + } + base_signature = inspect.signature(BaseOperator) + ignore = { + "default_args", # This is target we are working on now. + "kwargs", # A common name for a keyword argument. + "do_xcom_push", # In the same boat as `multiple_outputs` + "multiple_outputs", # We will use `self.multiple_outputs` instead. + "params", # Already handled above `partial_params`. + "task_concurrency", # Deprecated(replaced by `max_active_tis_per_dag`). + } + partial_keys = set(base_signature.parameters) - ignore + partial_kwargs.update({key: value for key, value in default_args.items() if key in partial_keys}) + partial_kwargs.update(task_kwargs) + + task_id = get_unique_task_id(partial_kwargs.pop("task_id"), dag, task_group) + if task_group: + task_id = task_group.child_id(task_id) + + # Logic here should be kept in sync with BaseOperatorMeta.partial(). + if partial_kwargs.get("wait_for_downstream"): + partial_kwargs["depends_on_past"] = True + start_date = timezone.convert_to_utc(partial_kwargs.pop("start_date", None)) + end_date = timezone.convert_to_utc(partial_kwargs.pop("end_date", None)) + if "pool_slots" in partial_kwargs: + if partial_kwargs["pool_slots"] < 1: + dag_str = "" + if dag: + dag_str = f" in dag {dag.dag_id}" + raise ValueError(f"pool slots for {task_id}{dag_str} cannot be less than 1") + + for fld, convert in ( + ("retries", parse_retries), + ("retry_delay", coerce_timedelta), + ("max_retry_delay", coerce_timedelta), + ("resources", coerce_resources), + ): + if (v := partial_kwargs.get(fld, NOTSET)) is not NOTSET: + partial_kwargs[fld] = convert(v) + + partial_kwargs.setdefault("executor_config", {}) + partial_kwargs.setdefault("op_args", []) + partial_kwargs.setdefault("op_kwargs", {}) + + try: + operator_name = self.operator_class.custom_operator_name # type: ignore + except AttributeError: + operator_name = self.operator_class.__name__ + + return DecoratedMappedOperator( + operator_class=self.operator_class, + expand_input=EXPAND_INPUT_EMPTY, # Don't use this; mapped values go to op_kwargs_expand_input. + partial_kwargs=partial_kwargs, + task_id=task_id, + params=partial_params, + operator_extra_links=self.operator_class.operator_extra_links, + template_ext=self.operator_class.template_ext, + template_fields=self.operator_class.template_fields, + template_fields_renderers=self.operator_class.template_fields_renderers, + ui_color=self.operator_class.ui_color, + ui_fgcolor=self.operator_class.ui_fgcolor, + is_empty=False, + is_sensor=self.operator_class._is_sensor, + can_skip_downstream=self.operator_class._can_skip_downstream, + is_stub=self.operator_class.is_stub, + task_module=self.operator_class.__module__, + task_type=self.operator_class.__name__, + operator_name=operator_name, + dag=dag, + task_group=task_group, + start_date=start_date, + end_date=end_date, + multiple_outputs=self.multiple_outputs, + python_callable=self.function, + op_kwargs_expand_input=expand_input, + disallow_kwargs_override=strict, + # Different from classic operators, kwargs passed to a taskflow + # task's expand() contribute to the op_kwargs operator argument, not + # the operator arguments themselves, and should expand against it. + expand_input_attr="op_kwargs_expand_input", + start_trigger_args=self.operator_class.start_trigger_args, + start_from_trigger=self.operator_class.start_from_trigger, + returns_dag_result=self.returns_dag_result, + register_with_dag=register_with_dag, + ) diff --git a/task-sdk/src/airflow/sdk/definitions/context.py b/task-sdk/src/airflow/sdk/definitions/context.py index 9d36203c2ebbf..3d3ced2afedb1 100644 --- a/task-sdk/src/airflow/sdk/definitions/context.py +++ b/task-sdk/src/airflow/sdk/definitions/context.py @@ -95,6 +95,75 @@ class Context(TypedDict, total=False): KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys()) +def clone_context(context: Context) -> Context: + """ + Create a safe, per-task copy of an execution ``Context`` for concurrent execution. + + The execution context is a mutable mapping that contains many nested + structures (``params``, ``templates_dict``, ``outlet_events``, ``dag_run``, + etc.). When running the same logical task concurrently (for example when + the ``IterableOperator`` spawns multiple indexed task instances that run in + parallel using threads, processes or asyncio tasks), those mutable objects + could be mutated by one indexed runtime and unintentionally observed by + another. That leads to subtle race conditions, corrupted state, and + flakiness in task execution. + + ``clone_context`` returns a new :class:`Context` mapping where the top-level + mapping is copied and specific mutable sub-objects that are commonly + mutated during execution are deep-copied or shallow-copied as appropriate: + + - ``params`` and ``templates_dict`` are deep-copied because they are + dictionaries that users and operators commonly mutate. + - ``inlets`` and ``outlets`` are converted to new lists (shallow copy) + because the sequence identity must be isolated but the elements are + typically read-only accessor objects. + - ``dag_run`` is deep-copied because it carries nested state that must + not be shared between concurrent executions. + - ``outlet_events`` is intentionally **not** copied so that events emitted + by sub-tasks are captured in the parent's accessor and serialized when + the parent task completes. ``OutletEventAccessors.__getitem__`` uses + ``dict.setdefault`` to guarantee that concurrent sub-tasks always share + one accessor per asset key rather than silently overwriting each other's + accumulated events. + + Use cases + - Multithreading: when using thread-based executors (``concurrent.futures`` + ThreadPoolExecutor) multiple threads share memory; cloning prevents + concurrent mutation of shared structures. + - Async concurrency: when running coroutine-based tasks concurrently in + the same event loop, tasks may still mutate shared mappings; cloning + avoids interference. + - Multiprocessing: while processes do not share memory, cloning keeps the + semantics consistent and avoids accidentally capturing references that + would be pickled. + + Performance + - The implementation intentionally copies only a small set of commonly + mutated fields rather than performing a blanket deep copy of the entire + context to keep the operation cheap. If future code stores additional + mutable state in the context that needs isolation, this function should + be extended appropriately. + + :param context: The original execution context to clone. + :returns: A new :class:`Context` safe to hand to a concurrently running task. + + :meta private: + """ + cloned_context = Context() + cloned_context.update(context) + cloned_context["params"] = copy.deepcopy(context.get("params", {})) + cloned_context["inlets"] = list(context.get("inlets", [])) + cloned_context["outlets"] = list(context.get("outlets", [])) + templates_dict = cloned_context.get("templates_dict") + if templates_dict is not None: + cloned_context["templates_dict"] = copy.deepcopy(templates_dict) + cloned_context["inlet_events"] = context["inlet_events"] + # outlet_events is intentionally NOT copied - sub-tasks must emit into the + # parent's accessor so events are serialized when the parent completes. + cloned_context["dag_run"] = context["dag_run"] + return cloned_context + + def context_merge(context: Context, *args: Any, **kwargs: Any) -> None: """ Merge parameters into an existing context. diff --git a/task-sdk/src/airflow/sdk/definitions/iterableoperator.py b/task-sdk/src/airflow/sdk/definitions/iterableoperator.py new file mode 100644 index 0000000000000..a975cb8d000e9 --- /dev/null +++ b/task-sdk/src/airflow/sdk/definitions/iterableoperator.py @@ -0,0 +1,540 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import copy +import os +from collections import deque +from collections.abc import Iterable, Mapping, Sequence +from itertools import repeat +from typing import TYPE_CHECKING, Any + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +from airflow.sdk import BaseXCom, TaskInstanceState, timezone +from airflow.sdk.bases.operator import BaseOperator, event_loop +from airflow.sdk.definitions._internal.expandinput import BatchedExpandInput +from airflow.sdk.definitions.context import clone_context +from airflow.sdk.definitions.mappedoperator import MappedOperator +from airflow.sdk.definitions.xcom_arg import MapXComArg, XComArg # noqa: F401 +from airflow.sdk.exceptions import ( + AirflowFailException, + AirflowRescheduleException, + AirflowRescheduleTaskInstanceException, + TaskDeferred, +) +from airflow.sdk.execution_time.executor import AsyncAwareExecutor, TaskExecutor +from airflow.sdk.execution_time.task_runner import IndexedTaskInstance + +if TYPE_CHECKING: + import jinja2 + + from airflow.providers.standard.triggers.temporal import DateTimeTrigger + from airflow.sdk.definitions._internal.expandinput import ExpandInput + from airflow.sdk.definitions.context import Context + from airflow.sdk.execution_time.lazy_sequence import XComIterable + + +ExternalDateTimeTrigger: type[DateTimeTrigger] | None + +try: + from airflow.providers.standard.triggers.temporal import DateTimeTrigger as ExternalDateTimeTrigger +except ModuleNotFoundError: + # If the providers package with DateTimeTrigger is not available (e.g. in + # minimal installs or tests), set the symbol to None so callers can + # explicitly check for availability. Using hasattr(self, DateTimeTrigger) + # is incorrect because hasattr expects a string attribute name. + ExternalDateTimeTrigger = None + + +class IterableOperator(BaseOperator): + """ + Operator used for Task Iteration (TI) that runs a mapped operator over an iterable input. + + The IterableOperator wraps a :class:`MappedOperator` together with an + :class:`ExpandInput` and is responsible for creating and running the + per-index runtime task instances. The IterableOperator itself is a + lightweight, non-retrying wrapper — retries, timeouts and deferred + execution are handled by the individual indexed task instances that the + IterableOperator creates for each element produced by the + ``expand_input``. + + The IterableOperator executes the mapped operator instances using a + concurrent executor with a configurable number of workers. By default + the worker count is taken from the mapped operator's ``partial_kwargs`` + (``task_concurrency``) if present, otherwise falls back to + ``os.cpu_count()`` and finally to ``1``. + + :param operator: The :class:`MappedOperator` to unmap and execute for + each element of ``expand_input``. Each indexed runtime receives a + deep copy/unmapped instance of this operator. + + :param expand_input: Provider of the values (or batches) to iterate + over. Its ``iter_values(context)`` method is used to produce the + per-index ``mapped_kwargs`` used to unmap the operator. + + :param kwargs: Additional keyword arguments forwarded to + :class:`BaseOperator` when instantiating the IterableOperator + (e.g. ``dag``, ``start_date``). Note that the IterableOperator + overrides retry-related parameters because retries are managed by + the per-index tasks. + + :returns: An :class:`XComIterable` if the mapped operator pushes XComs, otherwise ``None``. + + .. note:: + Deferred operators (those that raise :class:`~airflow.sdk.exceptions.TaskDeferred`) are not + supported yet inside IterableOperator. A ``TaskDeferred`` exception raised by an indexed task + instance will propagate as an error rather than pausing and resuming the task. + + Reschedule-mode sensors (those that raise :class:`~airflow.sdk.exceptions.AirflowRescheduleException`) + are also not supported. A reschedule raised by an indexed task instance will fail the whole + IterableOperator immediately with a clear error rather than being silently mishandled. + """ + + _operator: MappedOperator + expand_input: ExpandInput + partial_kwargs: dict[str, Any] + shallow_copy_attrs: Sequence[str] = ( + "_operator", + "expand_input", + "partial_kwargs", + "_log", + ) + + def __init__( + self, + *, + operator: MappedOperator, + expand_input: ExpandInput, + **kwargs, + ): + super().__init__( + **{ + **kwargs, + "task_id": operator.task_id, + "owner": operator.owner, + "email": operator.email, + "email_on_retry": operator.email_on_retry, + "email_on_failure": operator.email_on_failure, + "retries": 0, # We should not retry the IterableOperator, only the indexed runtime ti's should be retried + # Known v1 limitation: there is no durable per-index checkpoint. On a worker crash mid-iteration, + # clearing or re-running the parent TI reconstructs every IndexedTaskInstance with xcom_pushed=False + # and re-executes from index 0, duplicating any external effects that already completed. Until + # durable checkpointing is implemented, iteration should be constrained to idempotent work. + # AIP-103 Task State Management could be a solution to this known v1 limitation. + "retry_delay": operator.retry_delay, + "retry_exponential_backoff": operator.retry_exponential_backoff, + "max_retry_delay": operator.max_retry_delay, + "start_date": operator.start_date, + "end_date": operator.end_date, + "depends_on_past": operator.depends_on_past, + "ignore_first_depends_on_past": operator.ignore_first_depends_on_past, + "wait_for_past_depends_before_skipping": operator.wait_for_past_depends_before_skipping, + "wait_for_downstream": operator.wait_for_downstream, + "dag": operator.dag, + "priority_weight": operator.priority_weight, + "queue": operator.queue, + "pool": operator.pool, + "pool_slots": operator.pool_slots, + "execution_timeout": None, + "trigger_rule": operator.trigger_rule, + "resources": operator.resources, + "run_as_user": operator.run_as_user, + "map_index_template": operator.map_index_template, + "max_active_tis_per_dag": operator.max_active_tis_per_dag, + "max_active_tis_per_dagrun": operator.max_active_tis_per_dagrun, + "executor": operator.executor, + "executor_config": operator.executor_config, + "inlets": operator.inlets, + "outlets": operator.outlets, + "task_group": operator.task_group, + "doc": operator.doc, + "doc_md": operator.doc_md, + "doc_json": operator.doc_json, + "doc_yaml": operator.doc_yaml, + "doc_rst": operator.doc_rst, + "task_display_name": operator.task_display_name, + "allow_nested_operators": operator.allow_nested_operators, + } + ) + self._operator = operator + self.expand_input = expand_input + self.partial_kwargs = dict(operator.partial_kwargs) if operator.partial_kwargs else {} + task_concurrency = self.partial_kwargs.pop("task_concurrency", None) + if task_concurrency is not None and task_concurrency < 1: + raise ValueError(f"task_concurrency must be at least 1, got {task_concurrency}") + # Known v1 limitation: pool_slots is reserved once by the scheduler for this IterableOperator TI, + # but up to max_workers sub-tasks run concurrently inside it. Operators that set pool_slots > 1 to + # protect a shared resource (e.g. a DB connection pool) will be under-accounted — the pool sees one + # reservation while max_workers connections can be active simultaneously. A proper fix requires the + # scheduler to reserve pool_slots * max_workers slots, which needs scheduler-side changes. + self.max_workers = task_concurrency if task_concurrency is not None else (os.cpu_count() or 1) + XComArg.apply_upstream_relationship(self, self.expand_input.value) + + @property + def returns_dag_result(self) -> bool: + return self._operator.returns_dag_result + + @returns_dag_result.setter + def returns_dag_result(self, value: bool) -> None: + self._operator.returns_dag_result = value + + @property + def task_type(self) -> str: + return self._operator.__class__.__name__ + + @property + def task_retries(self) -> int: + return self._operator.retries or 0 + + def _do_render_template_fields( + self, + parent: Any, + template_fields: Iterable[str], + context: Context, + jinja_env: jinja2.Environment, + seen_oids: set[int], + ) -> None: + # IterableOperator doesn't need to render template fields as the actual operator's template fields + # will be rendered in the TaskExecutor when running each mapped task instance. + pass + + def _get_specified_expand_input(self) -> ExpandInput: + return self.expand_input + + def _unmap_operator( + self, context: Context, mapped_kwargs: Context, jinja_env: jinja2.Environment + ) -> BaseOperator: + from airflow.sdk.execution_time.context import context_update_for_unmapped + + unmapped_task = self._operator.unmap(mapped_kwargs) + # Make sure deferred operators will always raise a DeferredTask exception when executed + unmapped_task.start_from_trigger = False + context_update_for_unmapped(context, unmapped_task) + + unmapped_task._do_render_template_fields( + parent=unmapped_task, + template_fields=self._operator.template_fields, + context=context, + jinja_env=jinja_env, + seen_oids=set(), + ) + return unmapped_task + + async def _xcom_push(self, task: IndexedTaskInstance, value: Any) -> None: + if task.xcom_pushed: + self.log.debug( + "XCom already pushed for task_id %s with index %s", + task.task_id, + task.index, + ) + else: + self.log.debug( + "Pushing XCom for task_id %s with index %s", + task.task_id, + task.index, + ) + + await task.axcom_push(key=BaseXCom.XCOM_RETURN_KEY, value=value) + + def _run_tasks( + self, + context: Context, + tasks: Iterable[IndexedTaskInstance], + ) -> XComIterable | None: + exceptions: list[BaseException] = [] + reschedule_date = timezone.utcnow() + failed_tasks: deque[IndexedTaskInstance] = deque() + do_xcom_push = True + + self.log.info("Running tasks with %d workers", self.max_workers) + + while True: + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=self.max_workers) as executor: + for task, _result, raised in executor.map( + self._run_task, + repeat(executor), + repeat(context), + tasks, + ): + do_xcom_push = task.do_xcom_push + + if raised is None: + continue + + if isinstance(raised, TaskDeferred): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] attempted to defer. " + "Deferrable operators are not supported inside IterableOperator." + ) + + if isinstance(raised, AirflowRescheduleException): + if not isinstance(raised, AirflowRescheduleTaskInstanceException): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] attempted to reschedule. " + "Reschedule-mode sensors are not supported inside IterableOperator." + ) + reschedule_date = max(reschedule_date, raised.reschedule_date) + self.log.exception( + "An exception occurred for task_id %s with index %s, it has been rescheduled at %s", + task.task_id, + task.index, + reschedule_date, + ) + failed_tasks.append(raised.task) + continue + + # Non-Exception BaseExceptions (e.g. DeadlockImminentError, + # KeyboardInterrupt, SystemExit) must never be swallowed: they + # signal conditions where continuing iteration is meaningless + # because every subsequent task would fail for the same reason. + # Re-raise immediately to stop all task iteration. + if not isinstance(raised, Exception): + raise AirflowFailException( + f"Sub-task {task.task_id}[{task.index}] raised a non-Exception BaseException: " + f"{type(raised).__name__}: {raised}" + ) from raised + + self.log.exception( + "An exception occurred for task_id %s with index %s", + task.task_id, + task.index, + exc_info=raised, + ) + exceptions.append(raised) + + if not failed_tasks: + if exceptions: + # If this IterableOperator is backed by a batched expand input + # (created from a MappedIterableOperator), the parent mapped + # task should never be retried; retries are handled by the + # individual indexed runtime tasks. In that case raise + # AirflowFailException to mark failure without retrying the + # parent TaskInstance. For regular (non-batched) IterableOperator + # behavior, preserve the previous behavior and raise the + # BaseExceptionGroup so callers/tests that expect it keep working. + if isinstance(self.expand_input, BatchedExpandInput): + raise AirflowFailException(f"Multiple sub-task failures: {exceptions}") + raise BaseExceptionGroup("Multiple sub-task failures", exceptions) + if do_xcom_push: + from airflow.sdk.execution_time.lazy_sequence import XComIterable + + return XComIterable( + task_id=self.task_id, + dag_id=self.dag_id, + run_id=context["run_id"], + length=len(self.expand_input), + map_index=context["ti"].map_index, + ) + return None + + # If the retry time is still in the future we defer the operator so the worker + # slot is released. If the retry time has already passed we immediately re-run + # the failed tasks without deferring. + if reschedule_date > timezone.utcnow(): + if ExternalDateTimeTrigger is not None: + self.defer( + trigger=ExternalDateTimeTrigger(reschedule_date), + method_name=self.execute_failed_tasks.__name__, + kwargs={ + "failed_tasks": { + failed_task.index: failed_task.try_number for failed_task in failed_tasks + }, + }, + ) + else: + self.log.warning( + "DateTimeTrigger is not available; failed tasks cannot be rescheduled at %s and will be retried immediately", + reschedule_date, + ) + + tasks = list(failed_tasks) + failed_tasks.clear() + exceptions.clear() + reschedule_date = timezone.utcnow() + + async def _run_task( + self, + executor: AsyncAwareExecutor, + context: Context, + task: IndexedTaskInstance, + ) -> tuple[IndexedTaskInstance, Any | None, BaseException | None]: + try: + if task.is_async: + result = await self._run_async_operator(context, task) + else: + result = await executor.run_sync(self._run_operator, context, task) + + # Push XCom asynchronously (non-blocking) + if result is not None and task.do_xcom_push: + await self._xcom_push(task, result) + + return task, result, None + except BaseException as e: + return task, None, e + + def _run_operator(self, context: Context, task_instance: IndexedTaskInstance): + with TaskExecutor(task_instance=task_instance) as executor: + return executor.run( + context={ + **clone_context(context), + **{ + "ti": task_instance, + "task_instance": task_instance, + }, + } + ) + + async def _run_async_operator(self, context: Context, task_instance: IndexedTaskInstance): + with TaskExecutor(task_instance=task_instance) as executor: + return await executor.arun( + context={ + **clone_context(context), + **{ + "ti": task_instance, + "task_instance": task_instance, + }, + } + ) + + def _create_task( + self, + context: Context, + index: int, + mapped_kwargs: Context, + jinja_env: jinja2.Environment, + try_number: int = 0, + ) -> IndexedTaskInstance: + run_id = context["ti"].run_id + map_index = context["ti"].map_index + operator = self._unmap_operator(context.copy(), mapped_kwargs, jinja_env) + return self._create_mapped_task( + run_id=run_id, map_index=map_index, index=index, try_number=try_number, operator=operator + ) + + def _create_mapped_task( + self, run_id: str, map_index: int | None, index: int, try_number: int, operator: BaseOperator + ) -> IndexedTaskInstance: + return IndexedTaskInstance.model_construct( + task_id=operator.task_id, + dag_id=operator.dag_id, + run_id=run_id, + map_index=map_index, + index=index, + max_tries=operator.retries, + start_date=self.start_date, + state=TaskInstanceState.SCHEDULED.value, + is_mapped=True, + task=operator, + try_number=try_number, + xcom_pushed=False, + ) + + def execute(self, context: Context): + jinja_env = self.get_template_env(dag=self.dag) + tasks = ( + self._create_task( + context=context, + index=index, + mapped_kwargs=value, + jinja_env=jinja_env, + ) + for index, value in enumerate(self.expand_input.iter_values(context=context)) + ) + return self._run_tasks(context=context, tasks=tasks) + + def execute_failed_tasks( + self, + context: Context, + failed_tasks: dict[str, int], + event: dict[Any, Any], + ): + """ + Execute failed tasks after resuming from deferred state. + + :param context: The execution context. + :param failed_tasks: A dict mapping task index to its try_number. + :param event: The event that triggered the resume. + """ + jinja_env = self.get_template_env(dag=self.dag) + tasks = ( + self._create_task( + context=context, + index=index, + try_number=failed_tasks[str(index)], + jinja_env=jinja_env, + mapped_kwargs=value, + ) + for index, value in enumerate(self.expand_input.iter_values(context=context)) + if str(index) in failed_tasks + ) + return self._run_tasks(context=context, tasks=tasks) + + +class MappedIterableOperator(MappedOperator): + """A thin wrapper around an existing MappedOperator that unmaps an MappedOperator within an IterableOperator.""" + + def __init__( + self, + mapped_operator: MappedOperator, + expand_input: ExpandInput, + batch_size: int, + ): + self.delegate = mapped_operator + self.delegate.partial_kwargs["batch_size"] = batch_size + self.expand_input = expand_input + self._register_with_dag = True + self.__attrs_post_init__() + + def __getattr__(self, name): + if name.startswith("__") and name.endswith("__"): + raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}") + return getattr(self.delegate, name) + + def prepare_for_execution(self) -> MappedOperator: + return self + + @property + def batch_size(self) -> int: + return self.delegate.batch_size + + @property + def retries(self) -> int: + return 0 # We should not retry the IterableOperator, only the indexed runtime ti's should be retried + + @retries.setter + def retries(self, value: int) -> None: + if value != 0: + raise ValueError( + "MappedIterableOperator always has retries=0; retries are handled by indexed tasks." + ) + + def __repr__(self): + return f"" + + def unmap(self, resolve: Mapping[str, Any]) -> BaseOperator: + return IterableOperator( + operator=copy.deepcopy(self.delegate), + expand_input=BatchedExpandInput(self.expand_input, self.batch_size), + _airflow_from_mapped=True, + ) diff --git a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py index 3a68ecadfc22a..423ce24699c80 100644 --- a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py +++ b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py @@ -21,7 +21,7 @@ import copy import warnings from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Literal, TypeGuard +from typing import TYPE_CHECKING, Any, Literal, TypeGuard, cast import attrs import methodtools @@ -64,13 +64,15 @@ OperatorExpandArgument, OperatorExpandKwargsArgument, ) + from airflow.sdk.definitions.batchedoperator import BatchedOperator + from airflow.sdk.definitions.iterableoperator import IterableOperator from airflow.sdk.definitions.operator_resources import Resources from airflow.sdk.definitions.param import ParamsDict from airflow.sdk.definitions.retry_policy import RetryPolicy from airflow.sdk.types import WeightRuleParam from airflow.triggers.base import StartTriggerArgs -ValidationSource = Literal["expand"] | Literal["partial"] +ValidationSource = Literal["expand"] | Literal["iterate"] | Literal["partial"] def validate_mapping_kwargs(op: type[BaseOperator], func: ValidationSource, value: dict[str, Any]) -> None: @@ -213,68 +215,31 @@ def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = raise TypeError(f"expected XComArg or list[dict], not {type(kwargs).__name__}") return self._expand(ListOfDictsExpandInput(kwargs), strict=strict) - def _expand(self, expand_input: ExpandInput, *, strict: bool) -> MappedOperator: - from airflow.providers.standard.operators.empty import EmptyOperator - from airflow.sdk import BaseSensorOperator - from airflow.sdk.bases.skipmixin import SkipMixin - + def _expand( + self, + expand_input: ExpandInput, + *, + strict: bool, + register_with_dag: bool = True, + ) -> MappedOperator: self._expand_called = True - ensure_xcomarg_return_value(expand_input.value) - - partial_kwargs = self.kwargs.copy() - task_id = partial_kwargs.pop("task_id") - dag = partial_kwargs.pop("dag") - task_group = partial_kwargs.pop("task_group") - start_date = partial_kwargs.pop("start_date", None) - end_date = partial_kwargs.pop("end_date", None) - start_from_trigger = ( - partial_kwargs["start_from_trigger"] - if "start_from_trigger" in partial_kwargs - else getattr(self.operator_class, "start_from_trigger", False) - ) - start_trigger_args = ( - partial_kwargs["start_trigger_args"] - if "start_trigger_args" in partial_kwargs - else getattr(self.operator_class, "start_trigger_args", None) - ) + return self.batch(size=0)._expand(expand_input, strict=strict, register_with_dag=register_with_dag) - try: - operator_name = self.operator_class.custom_operator_name # type: ignore - except AttributeError: - operator_name = self.operator_class.__name__ - - op = MappedOperator( - operator_class=self.operator_class, - expand_input=expand_input, - partial_kwargs=partial_kwargs, - task_id=task_id, - params=self.params, - operator_extra_links=self.operator_class.operator_extra_links, - template_ext=self.operator_class.template_ext, - template_fields=self.operator_class.template_fields, - template_fields_renderers=self.operator_class.template_fields_renderers, - ui_color=self.operator_class.ui_color, - ui_fgcolor=self.operator_class.ui_fgcolor, - is_empty=issubclass(self.operator_class, EmptyOperator), - is_sensor=issubclass(self.operator_class, BaseSensorOperator), - can_skip_downstream=issubclass(self.operator_class, SkipMixin), - is_stub=self.operator_class.is_stub, - task_module=self.operator_class.__module__, - task_type=self.operator_class.__name__, - operator_name=operator_name, - dag=dag, - task_group=task_group, - start_date=start_date, - end_date=end_date, - disallow_kwargs_override=strict, - # For classic operators, this points to expand_input because kwargs - # to BaseOperator.expand() contribute to operator arguments. - expand_input_attr="expand_input", - # TODO: Move these to task SDK's BaseOperator and remove getattr - start_trigger_args=start_trigger_args, - start_from_trigger=start_from_trigger, - ) - return op + def iterate(self, **mapped_kwargs: OperatorExpandArgument) -> IterableOperator: + operator = self.batch(size=0).iterate(**mapped_kwargs) + return cast("IterableOperator", operator) + + def iterate_kwargs( + self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True + ) -> IterableOperator: + operator = self.batch(size=0).iterate_kwargs(kwargs, strict=strict) + return cast("IterableOperator", operator) + + def batch(self, size: int) -> BatchedOperator: + """Return a BatchedOperator for batched mapping.""" + from airflow.sdk.definitions.batchedoperator import BatchedOperator + + return BatchedOperator(operator_partial=self, size=size) @attrs.define( @@ -325,6 +290,7 @@ class MappedOperator(AbstractOperator): end_date: pendulum.DateTime | None upstream_task_ids: set[str] = attrs.field(factory=set, init=False) downstream_task_ids: set[str] = attrs.field(factory=set, init=False) + _register_with_dag: bool = attrs.field(alias="register_with_dag", default=True) _disallow_kwargs_override: bool """Whether execution fails if ``expand_input`` has duplicates to ``partial_kwargs``. @@ -346,19 +312,26 @@ def __repr__(self): return f"" def __attrs_post_init__(self): - from airflow.sdk.definitions.xcom_arg import XComArg - - if self.get_closest_mapped_task_group() is not None: - raise NotImplementedError("operator expansion in an expanded task group is not yet supported") - - if self.task_group: - self.task_group.add(self) - if self.dag: - self.dag.add_task(self) - XComArg.apply_upstream_relationship(self, self._get_specified_expand_input().value) - for k, v in self.partial_kwargs.items(): - if k in self.template_fields: - XComArg.apply_upstream_relationship(self, v) + # When _register_with_dag is False (i.e. IterableOperator), we intentionally + # skip the *entire* body — not just XComArg.apply_upstream_relationship. + # IterableOperator creates in-memory MappedOperator instances solely to drive task + # iteration; they must NOT be registered with the DAG or task group because Airflow + # treats the IterableOperator itself as the single real task instance in the DB. + # Calling dag.add_task() or task_group.add() here would raise duplicate-task errors. + if self._register_with_dag: + from airflow.sdk.definitions.xcom_arg import XComArg + + if self.get_closest_mapped_task_group() is not None: + raise NotImplementedError("operator expansion in an expanded task group is not yet supported") + + if self.task_group: + self.task_group.add(self) + if self.dag: + self.dag.add_task(self) + XComArg.apply_upstream_relationship(self, self._get_specified_expand_input().value) + for k, v in self.partial_kwargs.items(): + if k in self.template_fields: + XComArg.apply_upstream_relationship(self, v) @methodtools.lru_cache(maxsize=None) @classmethod @@ -726,6 +699,10 @@ def render_template_as_native_obj(self) -> bool | None: def render_template_as_native_obj(self, value: bool | None) -> None: self.partial_kwargs["render_template_as_native_obj"] = value + @property + def batch_size(self) -> int: + return self.partial_kwargs.get("batch_size", 0) + def get_dag(self) -> DAG | None: """Implement Operator.""" return self.dag @@ -792,6 +769,10 @@ def unmap(self, resolve: Mapping[str, Any]) -> BaseOperator: is_setup = kwargs.pop("is_setup", False) is_teardown = kwargs.pop("is_teardown", False) on_failure_fail_dagrun = kwargs.pop("on_failure_fail_dagrun", False) + # Remove batch_size and task_concurrency as they are only used for batched/iterable + # mapping metadata, not for operator init + kwargs.pop("batch_size", None) + kwargs.pop("task_concurrency", None) kwargs["task_id"] = self.task_id op = self.operator_class(**kwargs, _airflow_from_mapped=True) op.is_setup = is_setup diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index a38e47a466bcf..bb437c70f772e 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py @@ -96,6 +96,16 @@ def __new__(cls, *args, **kwargs) -> XComArg: def iter_references(self) -> Iterator[tuple[Operator, str]]: raise NotImplementedError() + def iter_values(self, context: Mapping[str, Any]) -> Iterable[Any]: + resolved = self.resolve(context) + + if isinstance(resolved, (str, bytes, dict)): + yield resolved + elif isinstance(resolved, Iterable): + yield from resolved + else: + yield resolved + @staticmethod def iter_xcom_references(arg: Any) -> Iterator[tuple[Operator, str]]: """ diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 6f43d5421ecf2..b36b954c0f5a7 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -28,9 +28,11 @@ if TYPE_CHECKING: from collections.abc import Collection + from datetime import datetime from airflow.sdk.definitions.asset import AssetNameRef, AssetUniqueKey, AssetUriRef from airflow.sdk.execution_time.comms import ErrorResponse + from airflow.sdk.execution_time.task_runner import IndexedTaskInstance class AirflowException(Exception): @@ -175,6 +177,19 @@ def serialize(self): return f"{cls.__module__}.{cls.__name__}", (), {"reschedule_date": self.reschedule_date} +class AirflowRescheduleTaskInstanceException(AirflowRescheduleException): + """ + Raise when the task should be re-scheduled for a specific TaskInstance at a later time. + + :param task: The task instance that should be rescheduled + :param reschedule_date: The datetime when the task should be retried + """ + + def __init__(self, task: IndexedTaskInstance, reschedule_date: datetime): + super().__init__(reschedule_date=reschedule_date) + self.task = task + + class AirflowSensorTimeout(AirflowException): """Raise when there is a timeout on sensor polling.""" diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 7ddc6edf1a3aa..8e08e198072d2 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -18,6 +18,7 @@ import collections import contextlib +import contextvars import functools import inspect import json @@ -146,7 +147,12 @@ }, } - +# Thread-safe storage for airflow context variables. +# This stores the AIRFLOW_CTX_* environment variables per-thread/per-context, +# allowing concurrent task execution without environment variable race conditions. +_airflow_context_vars: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "_airflow_context_vars", default=None +) log = structlog.get_logger(logger_name="task") #: Pass as ``retention`` to ``task_state_store.set()`` to store a key that never expires, @@ -168,6 +174,46 @@ def _unwrap_external_ref(stored: dict) -> str | None: T = TypeVar("T") +@contextlib.contextmanager +def airflow_context_vars_context(env_vars: dict[str, str]): + """ + Context manager to set airflow context variables thread-safely. + + This provides a thread-safe way to make airflow context variables (AIRFLOW_CTX_*) + available during task execution without causing race conditions in concurrent + execution scenarios (e.g., IterableOperator with AsyncAwareExecutor). + + :param env_vars: Dictionary of airflow context variables (e.g., AIRFLOW_CTX_DAG_ID) + """ + token = _airflow_context_vars.set(env_vars) + try: + yield + finally: + _airflow_context_vars.reset(token) + + +def get_airflow_context_var(key: str, default: str | None = None) -> str | None: + """ + Get an airflow context variable (thread-safe). + + Retrieves from the thread-local context first, then falls back to os.environ + for backward compatibility with code that directly reads environment variables. + + :param key: The variable name (e.g., "AIRFLOW_CTX_DAG_ID") + :param default: Default value if not found + :return: The variable value or default + """ + import os + + # First, try to get from context var (thread-safe) + context_vars = _airflow_context_vars.get() + if context_vars and key in context_vars: + return context_vars[key] + + # Fall back to os.environ for backward compatibility + return os.environ.get(key, default) + + def _process_connection_result_conn(conn_result: ReceiveMsgType | None) -> Connection: from airflow.sdk.definitions.connection import Connection @@ -1108,9 +1154,10 @@ def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> OutletEventAccessor else: raise TypeError(f"Key should be either an asset or an asset alias, not {type(key)}") - if hashable_key not in self._dict: - self._dict[hashable_key] = OutletEventAccessor(extra={}, key=hashable_key) - return self._dict[hashable_key] + # setdefault is atomic under the GIL: if two threads race on the same + # key the first writer wins and both threads get back the same accessor, + # so neither thread's accumulated events are silently discarded. + return self._dict.setdefault(hashable_key, OutletEventAccessor(extra={}, key=hashable_key)) @attrs.define(init=False) @@ -1358,17 +1405,22 @@ def set_current_context(context: Context) -> Generator[Context, None, None]: This method should be called once per Task execution, before calling operator.execute. """ - _CURRENT_CONTEXT.append(context) + current = _CURRENT_CONTEXT.get(None) + # Build a new list so that asyncio tasks / thread-pool workers that were + # created before this push still see the old stack (copy-on-set isolation). + new_stack = [*(current or []), context] + token = _CURRENT_CONTEXT.set(new_stack) try: yield context finally: - expected_state = _CURRENT_CONTEXT.pop() - if expected_state != context: + restored = _CURRENT_CONTEXT.get(None) + if not restored or restored[-1] != context: log.warning( "Current context is not equal to the state at context stack.", - expected=context, - got=expected_state, + expected_id=id(context), + got_id=id(restored[-1]) if restored else None, ) + _CURRENT_CONTEXT.reset(token) def context_update_for_unmapped(context: Context, task: BaseOperator) -> None: diff --git a/task-sdk/src/airflow/sdk/execution_time/executor.py b/task-sdk/src/airflow/sdk/execution_time/executor.py new file mode 100644 index 0000000000000..7ab28d367bac6 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/executor.py @@ -0,0 +1,398 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextvars +import inspect +import logging +import os +import time +from asyncio import ( + FIRST_COMPLETED, + AbstractEventLoop, + CancelledError, + Future, + Semaphore, + Task, + TimeoutError as AsyncTimeoutError, + gather, + wait, + wait_for, + wrap_future, +) +from collections.abc import Callable, Iterable, Iterator +from concurrent.futures import Executor, ThreadPoolExecutor +from contextlib import suppress +from typing import TYPE_CHECKING, Any, cast + +from airflow.sdk import BaseAsyncOperator, BaseOperator, TaskInstanceState, timezone +from airflow.sdk.bases.operator import ExecutorSafeguard +from airflow.sdk.definitions._internal.logging_mixin import LoggingMixin +from airflow.sdk.exceptions import ( + AirflowRescheduleException, + AirflowRescheduleTaskInstanceException, + TaskDeferred, +) +from airflow.sdk.execution_time.callback_runner import create_executable_runner +from airflow.sdk.execution_time.context import context_get_outlet_events, set_current_context +from airflow.sdk.execution_time.task_runner import ( + RuntimeTaskInstance, + _execute_task, + _run_task_state_change_callbacks, +) + +if TYPE_CHECKING: + from structlog.typing import FilteringBoundLogger as Logger + + from airflow.sdk import Context + from airflow.sdk.execution_time.task_runner import IndexedTaskInstance + + +class AsyncAwareExecutor(Executor): + """ + Executes both sync and async functions concurrently. + + Sync functions run in a ThreadPoolExecutor. + Async coroutines run on an asyncio event loop with a semaphore limit. + + :param loop: Event loop used to schedule async tasks and coordinate mixed execution. + :param max_workers: Maximum concurrent workers used by both thread pool and async semaphore. + """ + + def __init__(self, loop: AbstractEventLoop, max_workers: int | None = None): + if max_workers is None: + max_workers = os.cpu_count() or 1 + if max_workers <= 0: + raise ValueError("max_workers must be greater than 0") + + self._loop = loop + self._max_workers = max_workers + self._semaphore = Semaphore(max_workers) + self._thread_pool = ThreadPoolExecutor(max_workers=max_workers) + self._async_tasks: set[Task[Any]] = set() + self._shutdown = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + # On error path, cancel futures but still wait briefly for them to + # process CancelledError and release resources (e.g., threading + # locks). Without waiting, cancelled tasks that hold _thread_lock + # never execute their finally blocks, permanently leaking the lock + # and causing subsequent comms.send() calls to deadlock. + self.shutdown(wait=True, cancel_futures=True) + else: + self.shutdown(wait=True) + + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: + if self._shutdown: + return + + self._shutdown = True + + if cancel_futures: + for task in list(self._async_tasks): + task.cancel() + + if wait and self._async_tasks: + with suppress(TimeoutError, AsyncTimeoutError): + self._loop.run_until_complete( + wait_for( + gather(*self._async_tasks, return_exceptions=True), + timeout=10.0, + ) + ) + + self._thread_pool.shutdown(wait=wait, cancel_futures=cancel_futures) + + def submit(self, func: Callable[..., Any] | Any, *args, **kwargs) -> Future[Any]: # type: ignore[override] + """ + Submit a callable for execution. + + Always returns an asyncio.Future for consistency, whether the callable + is sync (run in thread pool) or async (run on the event loop). + """ + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + + if inspect.iscoroutine(func): + coro = func + elif inspect.iscoroutinefunction(func): + coro = func(*args, **kwargs) + else: + # Wrap thread pool future as asyncio.Future for consistent return type + return wrap_future(self._thread_pool.submit(func, *args, **kwargs), loop=self._loop) + + async def guarded(): + try: + async with self._semaphore: + return await coro + except CancelledError: + # If cancellation occurs while waiting for the semaphore, + # the inner coroutine was never awaited. Close it to prevent + # "coroutine was never awaited" RuntimeWarning. + coro.close() + raise + + task = self._loop.create_task(guarded()) + self._async_tasks.add(task) + task.add_done_callback(self._async_tasks.discard) + return task + + async def run_sync(self, func: Callable[..., Any], *args, **kwargs) -> Any: + """Run a sync callable in this executor's thread pool and await its result.""" + future = self._thread_pool.submit(func, *args, **kwargs) + return await wrap_future(future, loop=self._loop) + + def map( + self, + fn: Callable[..., Any], + *iterables: Iterable[Any], + timeout: float | None = None, + chunksize: int = 1, + ) -> Iterator[Any]: + """Apply fn to iterables and stream results in completion order.""" + if chunksize < 1: + raise ValueError("chunksize must be >= 1") + + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + + start = time.monotonic() + iterator = zip(*iterables) + pending: dict[Future[Any], Future[Any]] = {} + exhausted = False + + def _remaining_timeout() -> float | None: + if timeout is None: + return None + remaining = timeout - (time.monotonic() - start) + if remaining <= 0: + raise TimeoutError() + return remaining + + def _submit_next() -> bool: + nonlocal exhausted + + if exhausted: + return False + + try: + args = next(iterator) + except StopIteration: + exhausted = True + return False + + future = self.submit(fn, *args) + pending[future] = future + return True + + def _fill_pending() -> None: + """Submit tasks until pending reaches max_workers or iterator is exhausted.""" + while len(pending) < self._max_workers and _submit_next(): + pass + + # Submit up to max_workers tasks initially + _fill_pending() + + while pending: + wait_timeout = _remaining_timeout() + done, _ = self._loop.run_until_complete( + wait(set(pending), timeout=wait_timeout, return_when=FIRST_COMPLETED) + ) + + if not done: + raise TimeoutError() + + for completed in done: + pending.pop(completed) + _fill_pending() + yield completed.result() + + +class TaskExecutor(LoggingMixin): + """Base class to run an operator or trigger with given task context and task instance.""" + + def __init__( + self, + task_instance: IndexedTaskInstance, + ): + super().__init__() + self.task_instance = task_instance + self._result: Any | None = None + self._start_time: float | None = None + self._context: Context | None = None + + @property + def dag_id(self) -> str: + return self.task_instance.dag_id + + @property + def task_id(self) -> str: + return self.task_instance.task_id + + @property + def task_index(self) -> int: + return self.task_instance.index + + @property + def xcom_key(self): + return self.task_instance.xcom_key + + @property + def operator(self) -> BaseOperator: + return self.task_instance.task + + @property + def is_async(self) -> bool: + return self.task_instance.is_async + + def run(self, context: Context): + self._context = context + with set_current_context(context): + return _execute_task(context, self.task_instance, self.log) + + async def arun(self, context: Context): + self._context = context + with set_current_context(context): + return await _execute_async_task(context, self.task_instance, self.log) + + def __enter__(self): + self._start_time = time.monotonic() + + if self.log.isEnabledFor(logging.INFO): + self.log.info( + "Attempting running task %s of %s for %s with index %s in %s mode.", + self.task_instance.try_number, + self.operator.retries, + self.task_instance.task_id, + self.task_index, + "async" if self.is_async else "sync", + ) + return self + + def __exit__(self, exc_type, exc_value, traceback): + elapsed = time.monotonic() - self._start_time if self._start_time else 0.0 + + if exc_value: + # Non-Exception BaseExceptions (e.g. DeadlockImminentError, + # KeyboardInterrupt, SystemExit) must never be retried: they + # signal conditions where continuing is meaningless. + # Re-raise immediately without retry. + if not isinstance(exc_value, Exception): + self.task_instance.state = TaskInstanceState.FAILED + if self._context is not None: + _run_task_state_change_callbacks( + self.task_instance.task, "on_failure_callback", self._context, self.log + ) + raise exc_value + # AirflowRescheduleException (base) is raised by reschedule-mode sensors. + # Pass it through without consuming the retry budget so _run_tasks can + # detect and reject it with a clear "not supported" error, mirroring the + # TaskDeferred treatment below. + if isinstance(exc_value, AirflowRescheduleException) and not isinstance( + exc_value, AirflowRescheduleTaskInstanceException + ): + raise exc_value + if not isinstance(exc_value, TaskDeferred): + if self.task_instance.next_try_number > self.task_instance.max_tries: + self.log.error( + "Task instance %s for %s failed after %s attempts in %.2f seconds due to: %s", + self.task_index, + self.task_instance.task_id, + self.task_instance.max_tries, + elapsed, + exc_value, + ) + self.task_instance.state = TaskInstanceState.FAILED + if self._context is not None: + _run_task_state_change_callbacks( + self.task_instance.task, "on_failure_callback", self._context, self.log + ) + raise exc_value + self.task_instance.try_number += 1 + self.task_instance.end_date = timezone.utcnow() + self.task_instance.state = TaskInstanceState.UP_FOR_RESCHEDULE + if self._context is not None: + _run_task_state_change_callbacks( + self.task_instance.task, "on_retry_callback", self._context, self.log + ) + raise AirflowRescheduleTaskInstanceException( + task=self.task_instance, + reschedule_date=self.task_instance.next_retry_datetime(), + ) + raise exc_value + + self.task_instance.state = TaskInstanceState.SUCCESS + if self._context is not None: + _run_task_state_change_callbacks( + self.task_instance.task, "on_success_callback", self._context, self.log + ) + if self.log.isEnabledFor(logging.INFO): + self.log.info( + "Task instance %s for %s finished successfully in %s attempts in %.2f seconds", + self.task_index, + self.task_instance.task_id, + self.task_instance.next_try_number, + elapsed, + ) + + +async def _execute_async_task(context: Context, ti: RuntimeTaskInstance, log: Logger): + """Execute Task (optionally with a Timeout) and push Xcom results.""" + # set-up + task = cast("BaseAsyncOperator", ti.task) + execute = task.aexecute # here we must use aexecute instead of execute + + # async tasks can't originate from deferred operator, so no need to check next_method + + ctx = contextvars.copy_context() + # Populate the context var so ExecutorSafeguard doesn't complain + ctx.run(ExecutorSafeguard.tracker.set, task) + + outlet_events = context_get_outlet_events(context) + + if (pre_execute_hook := task._pre_execute_hook) is not None: + create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) + if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not BaseOperator.pre_execute: + create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) + + _run_task_state_change_callbacks(task, "on_execute_callback", context, log) + + async def _run_in_context(coro_func, *args, **kwargs): + """Run async function in contextvars context with optional timeout.""" + coro_in_ctx = ctx.run(lambda: coro_func(*args, **kwargs)) + + if task.execution_timeout: + return await wait_for(coro_in_ctx, timeout=task.execution_timeout.total_seconds()) + return await coro_in_ctx + + try: + result = await _run_in_context(execute, context=context) + except AsyncTimeoutError: + task.on_kill() + raise + + if (post_execute_hook := task._post_execute_hook) is not None: + create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result) + if getattr(post_execute_hook := task.post_execute, "__func__", None) is not BaseOperator.post_execute: + create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context) + + return result 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..06f60da9aa82a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -25,6 +25,9 @@ import attrs import structlog +from airflow.sdk import BaseXCom +from airflow.sdk.execution_time.xcom import XCom + if TYPE_CHECKING: from airflow.sdk.definitions.xcom_arg import PlainXComArg from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance @@ -166,6 +169,86 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]: return XCom.deserialize_value(_XComWrapper(msg.root)) +class XComIterable(Sequence): + """An iterable that lazily fetches XCom values one by one instead of loading all at once.""" + + def __init__(self, task_id: str, dag_id: str, run_id: str, length: int, map_index: int | None = None): + self.task_id = task_id + self.dag_id = dag_id + self.run_id = run_id + self.length = length + self.map_index = map_index + + def __iter__(self) -> Iterator: + return _XComIterator(self) + + def __len__(self) -> int: + return self.length + + @overload + def __getitem__(self, key: int) -> Any: ... + + @overload + def __getitem__(self, key: slice) -> Sequence[Any]: ... + + def __getitem__(self, key: int | slice) -> Any | Sequence[Any]: + """Allow direct indexing so this works like a sequence.""" + if isinstance(key, slice): + # TODO: This issues one XCom.get_one call per element — N round-trips for a full slice. + # XComIterable stores results under distinct keys (return_value_0, return_value_1, …) + # with the same map_index, so the existing GetXComSequenceSlice endpoint (which ranges + # over map_index for a single key) cannot be reused. A new POST endpoint that accepts + # a list of keys and returns values in a single query is needed; once that lands, replace + # this loop with a single batched fetch. + start, stop, step = key.indices(len(self)) + return [self[i] for i in range(start, stop, step)] + + if not (0 <= key < self.length): + raise IndexError(key) + + return XCom.get_one( + key=f"{BaseXCom.XCOM_RETURN_KEY}_{key}", + dag_id=self.dag_id, + task_id=self.task_id, + run_id=self.run_id, + map_index=self.map_index, + ) + + def serialize(self) -> dict: + """Ensure the object is JSON serializable.""" + return { + "task_id": self.task_id, + "dag_id": self.dag_id, + "run_id": self.run_id, + "length": self.length, + "map_index": self.map_index, + } + + @classmethod + def deserialize(cls, data: dict, version: int): + """Ensure the object is JSON deserializable.""" + return XComIterable(**data) + + +class _XComIterator: + """Iterator for XComIterable.""" + + def __init__(self, iterable: XComIterable): + self._iterable = iterable + self._index = 0 + + def __iter__(self): + return self + + def __next__(self): + if self._index >= len(self._iterable): + raise StopIteration + + value = self._iterable[self._index] + self._index += 1 + return value + + def _coerce_slice_index(value: Any) -> int | None: """ Check slice attribute's type and convert it to int. diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 1ce848128d663..ee54e82c7f7f9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -21,13 +21,16 @@ import contextvars import functools +import hashlib import inspect +import math import os import sys import time from collections.abc import Callable, Iterable, Iterator, Mapping from contextlib import ExitStack, contextmanager, suppress from datetime import datetime, timedelta, timezone +from functools import cached_property from itertools import product from pathlib import Path from typing import TYPE_CHECKING, Annotated, Any, Literal @@ -61,6 +64,7 @@ from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.configuration import conf from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager +from airflow.sdk.definitions._internal.logging_mixin import LoggingMixin from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set from airflow.sdk.definitions.asset import ( Asset, @@ -293,6 +297,14 @@ def __rich_repr__(self): __rich_repr__.angular = True # type: ignore[attr-defined] + @cached_property + def logical_date(self) -> datetime | None: + if self._ti_context_from_server: + dag_run = self._ti_context_from_server.dag_run + + return dag_run.logical_date + return None + @detail_span("get_template_context") def get_template_context(self) -> Context: # TODO: Move this to `airflow.sdk.execution_time.context` @@ -373,7 +385,7 @@ def get_template_context(self) -> Context: } self._cached_template_context.update(context_from_server) - if logical_date := coerce_datetime(dag_run.logical_date): + if logical_date := coerce_datetime(self.logical_date): if TYPE_CHECKING: assert isinstance(logical_date, DateTime) ds = logical_date.strftime("%Y-%m-%d") @@ -891,6 +903,125 @@ def mark_success_url(self) -> str: return self.log_url +class IndexedTaskInstance(RuntimeTaskInstance, LoggingMixin): + """Indexed task instance to run a mapped operator.""" + + index: int + xcom_pushed: bool = Field(default=False) + + def __init__(self, /, **data: Any): + super().__init__(**data) + + if self.index is None or self.index < 0: + raise ValueError("IndexedTaskInstance requires index >= 0") + + def xcom_pull( + self, + task_ids: str | Iterable[str] | None = None, + dag_id: str | None = None, + key: str = BaseXCom.XCOM_RETURN_KEY, + include_prior_dates: bool = False, + *, + map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET, + default: Any = None, + run_id: str | None = None, + ) -> Any: + return super().xcom_pull( + task_ids=task_ids, + dag_id=dag_id, + key=f"{key}_{self.index}", + include_prior_dates=include_prior_dates, + map_indexes=map_indexes, + default=default, + run_id=run_id, + ) + + def xcom_push( + self, + key: str, + value: Any, + ): + super().xcom_push(key=f"{key}_{self.index}", value=value) + if key == BaseXCom.XCOM_RETURN_KEY: + self.xcom_pushed = True + + async def axcom_push( + self, + key: str, + value: Any, + ): + await super().axcom_push(key=f"{key}_{self.index}", value=value) + if key == BaseXCom.XCOM_RETURN_KEY: + self.xcom_pushed = True + + def next_retry_datetime(self): + """ + Get datetime of the next retry if the task instance fails. + + For exponential backoff, retry_delay is used as base and will be converted to seconds. + """ + from airflow.sdk.definitions._internal.abstractoperator import MAX_RETRY_DELAY + + delay = self.task.retry_delay + if self.task.retry_exponential_backoff: + try: + # If the min_backoff calculation is below 1, it will be converted to 0 via int. Thus, + # we must round up prior to converting to an int, otherwise a divide by zero error + # will occur in the modded_hash calculation. + # this probably gives unexpected results if a task instance has previously been cleared, + # because try_number can increase without bound + min_backoff = math.ceil(delay.total_seconds() * (2 ** (self.try_number - 1))) + except OverflowError: + min_backoff = MAX_RETRY_DELAY + self.log.warning( + "OverflowError occurred while calculating min_backoff, using MAX_RETRY_DELAY for min_backoff." + ) + + # In the case when delay.total_seconds() is 0, min_backoff will not be rounded up to 1. + # To address this, we impose a lower bound of 1 on min_backoff. This effectively makes + # the ceiling function unnecessary, but the ceiling function was retained to avoid + # introducing a breaking change. + if min_backoff < 1: + min_backoff = 1 + + # deterministic per task instance + ti_hash = int( + hashlib.sha1( + f"{self.dag_id}#{self.task_id}#{self.logical_date}#{self.try_number}".encode(), + usedforsecurity=False, + ).hexdigest(), + 16, + ) + # between 1 and 1.0 * delay * (2^retry_number) + modded_hash = min_backoff + ti_hash % min_backoff + # timedelta has a maximum representable value. The exponentiation + # here means this value can be exceeded after a certain number + # of tries (around 50 if the initial delay is 1s, even fewer if + # the delay is larger). Cap the value here before creating a + # timedelta object so the operation doesn't fail with "OverflowError". + delay_backoff_in_seconds = min(modded_hash, MAX_RETRY_DELAY) + delay = timedelta(seconds=delay_backoff_in_seconds) + if self.task.max_retry_delay: + delay = min(self.task.max_retry_delay, delay) + return self.end_date + delay + + @property + def is_async(self) -> bool: + return self.task.is_async + + @property + def next_try_number(self) -> int: + return self.try_number + 1 + + @property + def xcom_key(self) -> str: + return f"{self.task_id}_{self.index}" + + @property + def do_xcom_push(self) -> bool: + return self.task.do_xcom_push + + def _xcom_push( ti: RuntimeTaskInstance, key: str, @@ -2228,9 +2359,11 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger): assert isinstance(kwargs, dict) execute = functools.partial(task.resume_execution, next_method=next_method, next_kwargs=kwargs) - # Export context in os.environ to make it available for operators to use. - airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True) - os.environ.update(airflow_context_vars) + # IndexedTaskInstance should not update env variables as those run concurrently within the same process + if not isinstance(ti, IndexedTaskInstance): + # Export context in os.environ to make it available for operators to use. + airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True) + os.environ.update(airflow_context_vars) outlet_events = context_get_outlet_events(context) diff --git a/task-sdk/tests/task_sdk/bases/test_operator.py b/task-sdk/tests/task_sdk/bases/test_operator.py index dcb5240a83dc8..00e2f440e76a9 100644 --- a/task-sdk/tests/task_sdk/bases/test_operator.py +++ b/task-sdk/tests/task_sdk/bases/test_operator.py @@ -41,7 +41,7 @@ ) from airflow.sdk.definitions.param import ParamsDict from airflow.sdk.definitions.template import literal -from airflow.triggers.base import StartTriggerArgs +from airflow.triggers.base import BaseTrigger, StartTriggerArgs DEFAULT_DATE = datetime(2016, 1, 1, tzinfo=timezone.utc) @@ -1137,3 +1137,18 @@ def __init__(self, arg1, arg2, arg3, **kwargs): assert op.arg2 == "b" assert op.arg3 == 3 assert op.queue == "THIS" + + +class MockTrigger(BaseTrigger): + """A minimal trigger stub that yields a single TriggerEvent-like object.""" + + def __init__(self, payload=None, **kwargs): + super().__init__(**kwargs) + self._payload = payload + + async def run(self): + if self._payload is not None: + yield mock.Mock(payload=self._payload) + + def serialize(self): + return ("tests.MockTrigger", {"payload": self._payload}) diff --git a/task-sdk/tests/task_sdk/definitions/_internal/test_expandinput.py b/task-sdk/tests/task_sdk/definitions/_internal/test_expandinput.py new file mode 100644 index 0000000000000..b8dbec0cadb26 --- /dev/null +++ b/task-sdk/tests/task_sdk/definitions/_internal/test_expandinput.py @@ -0,0 +1,103 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest + +from airflow.sdk.definitions._internal.expandinput import ( + BatchedExpandInput, + DictOfListsExpandInput, + ListOfDictsExpandInput, +) + + +class TestExpandInput: + @pytest.mark.parametrize( + ("actual", "expected"), + [ + ({"a": 1}, [{"a": 1}]), + ({"a": [1, 2, 3]}, [{"a": 1}, {"a": 2}, {"a": 3}]), + ({"a": "hello"}, [{"a": "hello"}]), + ( + {"a": [1, 2], "b": [10, 20]}, + [{"a": 1, "b": 10}, {"a": 1, "b": 20}, {"a": 2, "b": 10}, {"a": 2, "b": 20}], + ), + ({"a": (x for x in [1, 2])}, [{"a": 1}, {"a": 2}]), + ], + ) + def test_dict_of_lists_expand_input_iter_values(self, actual, expected): + expand_input = DictOfListsExpandInput(actual) + + with pytest.raises(RuntimeError, match="Length of DictOfListsExpandInput is not yet known"): + len(expand_input) + + result = list(expand_input.iter_values({})) + assert result == expected + assert len(expand_input) == len(expected) + + @pytest.mark.parametrize( + ("actual", "expected"), + [ + ([{"a": 1}, {"a": 2}], [{"a": 1}, {"a": 2}]), + ([{"a": 1, "b": 2}], [{"a": 1, "b": 2}]), + ([], []), + ], + ) + def test_list_of_dicts_expand_input_iter_values(self, actual, expected): + expand_input = ListOfDictsExpandInput(actual) + + with pytest.raises(RuntimeError, match="Length of ListOfDictsExpandInput is not yet known"): + len(expand_input) + + result = list(expand_input.iter_values({})) + assert result == expected + assert len(expand_input) == len(expected) + + +class TestBatchedExpandInput: + @pytest.mark.parametrize( + "size", + [pytest.param(-1), pytest.param(0), pytest.param(1)], + ) + def test_invalid_size_raises(self, size: int): + inner = DictOfListsExpandInput({"a": [1, 2, 3]}) + with pytest.raises(ValueError, match="batch size must be at least 2"): + BatchedExpandInput(inner, size=size) + + @pytest.mark.parametrize( + ("size", "map_index", "items", "expected"), + [ + (2, 0, [1, 2, 3, 4, 5], [1, 3, 5]), + (2, 1, [1, 2, 3, 4, 5], [2, 4]), + (3, 0, [1, 2, 3, 4, 5, 6], [1, 4]), + (3, 1, [1, 2, 3, 4, 5, 6], [2, 5]), + (3, 2, [1, 2, 3, 4, 5, 6], [3, 6]), + ], + ) + def test_iter_values_striding(self, size: int, map_index: int, items: list, expected: list): + inner = DictOfListsExpandInput({"a": items}) + batched = BatchedExpandInput(inner, size=size) + context = {"ti": type("TI", (), {"map_index": map_index})()} + + with pytest.raises(RuntimeError, match="Length of BatchedExpandInput is not yet known"): + len(batched) + + result = [combo["a"] for combo in batched.iter_values(context)] + assert result == expected + assert len(batched) == len(expected) + # delegate length must remain the full item count, not the batch slice + assert len(inner) == len(items) diff --git a/task-sdk/tests/task_sdk/definitions/conftest.py b/task-sdk/tests/task_sdk/definitions/conftest.py index 3f89f34b4d2da..e5cefa1e3e8ac 100644 --- a/task-sdk/tests/task_sdk/definitions/conftest.py +++ b/task-sdk/tests/task_sdk/definitions/conftest.py @@ -17,11 +17,12 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import pytest import structlog +from airflow.sdk import BaseOperator, XComArg from airflow.sdk.execution_time.comms import SucceedTask, TaskState if TYPE_CHECKING: @@ -47,3 +48,10 @@ def run(dag: DAG, task_id: str, map_index: int): raise RuntimeError("Unable to find call to TaskState") return run + + +def make_xcom_arg(values: Any) -> XComArg: + op = BaseOperator(task_id="upstream") + xcom_arg = XComArg(op) + xcom_arg.resolve = lambda *a, **kw: values + return xcom_arg diff --git a/task-sdk/tests/task_sdk/definitions/test_batchedoperator.py b/task-sdk/tests/task_sdk/definitions/test_batchedoperator.py new file mode 100644 index 0000000000000..d15002508bfa9 --- /dev/null +++ b/task-sdk/tests/task_sdk/definitions/test_batchedoperator.py @@ -0,0 +1,156 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable +from unittest import mock + +import pytest + +from airflow.sdk import DAG, TaskInstanceState +from airflow.sdk.bases.xcom import BaseXCom +from airflow.sdk.execution_time.comms import ( + GetTICount, + GetXCom, + GetXComSequenceSlice, + TICount, + XComResult, + XComSequenceSliceResult, +) + +RunTI = Callable[[DAG, str, int], TaskInstanceState] + + +class TestBatchedOperator: + def test_batch_iterate(self, run_ti: RunTI, mock_supervisor_comms): + outputs = defaultdict(list) + numbers = list(range(10)) + + with DAG(dag_id="product_same") as dag: + + @dag.task + def emit_numbers(): + return numbers + + @dag.task + def show(number, **context): + map_index = str(context["ti"].map_index) + outputs[map_index].append(number) + return number + + emit_task = emit_numbers() + show.batch(size=2).iterate(number=emit_task) + + def mock_comms(msg): + if isinstance(msg, GetXCom): + if msg.task_id == "emit_numbers": + return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers) + elif isinstance(msg, GetXComSequenceSlice): + if msg.task_id == "emit_numbers": + return XComSequenceSliceResult(root=numbers) + elif isinstance(msg, GetTICount): + if msg.task_ids and msg.task_ids[0] == "show": + return TICount(count=2) + return TICount(count=1) + return mock.DEFAULT + + mock_supervisor_comms.send.side_effect = mock_comms + + states = [run_ti(dag, "show", map_index) for map_index in range(2)] + assert states == [TaskInstanceState.SUCCESS] * 2 + assert set(outputs["0"]) == {0, 2, 4, 6, 8} + assert set(outputs["1"]) == {1, 3, 5, 7, 9} + + @pytest.mark.parametrize( + ("batch_size", "expand_size"), + [ + (5, 10), # Batched: size=5 for 10 items + (3, 3), # Batched: size=3 for 3 items + (4, 20), # Batched: size=4 for 20 items + (1, 5), # Non-batched: size=1 (iterate all at once) + ], + ) + def test_batch_size_preserved_through_lifecycle(self, batch_size, expand_size): + from airflow.providers.standard.operators.empty import EmptyOperator + from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput + from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator + from airflow.serialization.serialized_objects import OperatorSerialization + + with DAG(dag_id=f"test_batch_{batch_size}") as dag: + op = EmptyOperator.partial(task_id="test_task", dag=dag).batch(size=batch_size) + + expand_input = DictOfListsExpandInput({"retry_delay": list(range(expand_size))}) + iterable_op = op._iterate(expand_input, strict=False) + + # Check if batched or non-batched + if batch_size > 1: + assert isinstance(iterable_op, MappedIterableOperator) + mapped_op = iterable_op.delegate + assert iterable_op.batch_size == batch_size + + # 1. Verify batch_size is in partial_kwargs + assert "batch_size" in mapped_op.partial_kwargs + assert mapped_op.partial_kwargs["batch_size"] == batch_size + + # 2. Verify batch_size is serialized (not excluded) + serialized = OperatorSerialization.serialize_mapped_operator(mapped_op) + assert "partial_kwargs" in serialized + assert "batch_size" in serialized["partial_kwargs"] + assert serialized["partial_kwargs"]["batch_size"] == batch_size + + # 3. Verify batch_size survives deserialization + deserialized_op = OperatorSerialization.deserialize_operator(serialized) + assert "batch_size" in deserialized_op.partial_kwargs + assert deserialized_op.partial_kwargs["batch_size"] == batch_size + + # 4. Verify batch_size is removed before operator instantiation (only for batched) + unmapped = iterable_op.unmap({"retry_delay": 1}) + # Verify unmapped task doesn't have batch_size attribute + assert not hasattr(unmapped, "batch_size") + else: + assert isinstance(iterable_op, IterableOperator) + + def test_mapped_iterable_operator_retries_preserved(self): + """Ensure delegate's retries survive unmap, while MappedIterableOperator reports 0 retries.""" + from airflow.providers.standard.operators.empty import EmptyOperator + from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput + from airflow.sdk.definitions.iterableoperator import MappedIterableOperator + + with DAG(dag_id="test_mapped_iterable_retries") as dag: + expand_input = DictOfListsExpandInput({"retry_delay": [1.0, 2.0]}) + iterable_op = ( + EmptyOperator.partial(task_id="test_task", dag=dag, retries=3) + .batch(size=2) + ._iterate(expand_input, strict=False) + ) + + assert isinstance(iterable_op, MappedIterableOperator) + assert iterable_op.retries == 0 + + with pytest.raises( + ValueError, + match="MappedIterableOperator always has retries=0; retries are handled by indexed tasks.", + ): + iterable_op.retries = 3 + + mapped_op = iterable_op.delegate + assert mapped_op.retries == 3 + + unmapped = mapped_op.unmap({"retry_delay": 1.0}) + assert unmapped.retries == 3 diff --git a/task-sdk/tests/task_sdk/definitions/test_context.py b/task-sdk/tests/task_sdk/definitions/test_context.py index dc25ec378bdee..7502439cceef9 100644 --- a/task-sdk/tests/task_sdk/definitions/test_context.py +++ b/task-sdk/tests/task_sdk/definitions/test_context.py @@ -17,9 +17,14 @@ # under the License. from __future__ import annotations +from types import SimpleNamespace + import pytest -from airflow.sdk.definitions.context import get_current_context +from airflow.sdk import Asset +from airflow.sdk.definitions._internal.contextmanager import _CURRENT_CONTEXT +from airflow.sdk.definitions.context import Context, clone_context, get_current_context +from airflow.sdk.execution_time.context import InletEventsAccessors class TestCurrentContext: @@ -27,15 +32,73 @@ def test_current_context_no_context_raise(self): with pytest.raises(RuntimeError): get_current_context() - def test_get_current_context_with_context(self, monkeypatch): + def test_get_current_context_with_context(self): mock_context = {"ti": "task_instance", "key": "value"} - monkeypatch.setattr( - "airflow.sdk.definitions._internal.contextmanager._CURRENT_CONTEXT", [mock_context] + token = _CURRENT_CONTEXT.set([mock_context]) + try: + result = get_current_context() + assert result == mock_context + finally: + _CURRENT_CONTEXT.reset(token) + + def test_get_current_context_without_context(self): + token = _CURRENT_CONTEXT.set([]) + try: + with pytest.raises(RuntimeError, match="Current context was requested but no context was found!"): + get_current_context() + finally: + _CURRENT_CONTEXT.reset(token) + + def test_clone_context_deep_and_shallow_copy_semantics(self): + outlet_events = [Asset(name="dummy")] + inlet_events = InletEventsAccessors(inlets=[]) + + dag_run = SimpleNamespace( + dag_id="dag", + run_id="r1", + logical_date=None, + data_interval_start=None, + data_interval_end=None, + run_after=None, + start_date=None, + end_date=None, + clear_number=None, + run_type=None, + state=None, + conf=None, + triggering_user_name=None, + consumed_asset_events=[], + partition_key=None, + note=None, ) - result = get_current_context() - assert result == mock_context - def test_get_current_context_without_context(self, monkeypatch): - monkeypatch.setattr("airflow.sdk.definitions._internal.contextmanager._CURRENT_CONTEXT", []) - with pytest.raises(RuntimeError, match="Current context was requested but no context was found!"): - get_current_context() + actual = Context() + actual.update( + { + "params": {"p": {"n": 1}}, + "templates_dict": {"tpl": ["a", {"x": 1}]}, + "inlets": [object()], + "outlets": [object()], + "outlet_events": outlet_events, + "inlet_events": inlet_events, + "dag_run": dag_run, + } + ) + cloned = clone_context(actual) + + assert cloned is not actual + + actual["params"]["p"]["n"] = 999 + assert cloned["params"]["p"]["n"] == 1 + + actual["templates_dict"]["tpl"][1]["x"] = 42 + assert cloned["templates_dict"]["tpl"][1]["x"] == 1 + + actual["outlet_events"].append(Asset(name="another")) + assert cloned["outlet_events"] is actual["outlet_events"] + assert len(cloned["outlet_events"]) == 2 + assert Asset(name="another") in cloned["outlet_events"] + assert cloned["inlet_events"] is actual["inlet_events"] + + actual["dag_run"].dag_id = "changed" + assert cloned["dag_run"].dag_id == "changed" diff --git a/task-sdk/tests/task_sdk/definitions/test_iterableoperator.py b/task-sdk/tests/task_sdk/definitions/test_iterableoperator.py new file mode 100644 index 0000000000000..70ca9e4c14b10 --- /dev/null +++ b/task-sdk/tests/task_sdk/definitions/test_iterableoperator.py @@ -0,0 +1,684 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import copy +from datetime import timedelta +from typing import TYPE_CHECKING + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +import pytest + +from airflow.sdk import DAG, BaseOperator, BaseXCom, get_current_context +from airflow.sdk.definitions._internal.abstractoperator import DEFAULT_RETRIES +from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput, ListOfDictsExpandInput +from airflow.sdk.definitions.iterableoperator import IterableOperator +from airflow.sdk.exceptions import AirflowFailException, AirflowRescheduleException, TaskDeferred +from airflow.sdk.execution_time.xcom import XCom + +from tests_common.test_utils.mock_context import mock_context + +if TYPE_CHECKING: + from airflow.sdk.definitions._internal.expandinput import ExpandInput + from airflow.sdk.definitions.mappedoperator import MappedOperator + + from tests_common.test_utils.compat import Context + + +class MockOperator(BaseOperator): + """Mock operator for testing IterableOperator expansion.""" + + def __init__( + self, + arg1=None, + arg2=None, + arg3=None, + fail_on_first_attempt=False, + raise_exception: BaseException | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.arg1 = arg1 + self.arg2 = arg2 + self.arg3 = arg3 + self.fail_on_first_attempt = fail_on_first_attempt + self.raise_exception = raise_exception + + def execute(self, context): + """Execute the operator and return passed arguments as tuple if do_xcom_push is True.""" + expected = copy.deepcopy(context) + + if self.raise_exception is not None: + raise self.raise_exception + if self.fail_on_first_attempt: + self.fail_on_first_attempt = False + raise RuntimeError + if not self.do_xcom_push: + return None + result = self.arg1, self.arg2, self.arg3 + + assert context == expected, "Context was unexpectedly mutated during task execution" + return result + + +class MockDeferredOperator(BaseOperator): + """Operator that immediately defers on execute, simulating a deferrable operator.""" + + template_fields = () + + def execute(self, context): + raise TaskDeferred(trigger=None, method_name="execute_complete") # type: ignore[arg-type] + + +class MockRescheduleSensor(BaseOperator): + """Operator that raises AirflowRescheduleException on execute, simulating a reschedule-mode sensor.""" + + template_fields = () + + def execute(self, context): + from datetime import timedelta + + from airflow.sdk import timezone + + raise AirflowRescheduleException(timezone.utcnow() + timedelta(seconds=60)) + + +@pytest.fixture +def mock_xcom_get_one(monkeypatch: pytest.MonkeyPatch): + """ + Fixture that mocks XCom.get_one using monkeypatch for proper cleanup. + + Captures values pushed via context["ti"].xcom_push in the order they arrive, + then serves them back by index when XComIterable calls + XCom.get_one(key=f"{task_id}_{idx}", ...). + """ + + def _mock_xcom(context: Context): + pushed_values: list = [] + + original_push = context["ti"].xcom_push + + def capturing_push(key: str, value, **kwargs) -> None: + pushed_values.append(value) + original_push(key=key, value=value, **kwargs) + + monkeypatch.setattr(context["ti"], "xcom_push", capturing_push) + + task_id = context["ti"].task_id + + def mock_get_one(**kwargs): + key = kwargs.get("key", "") + prefix = f"{task_id}_" + if key.startswith(prefix): + try: + idx = int(key[len(prefix) :]) + if 0 <= idx < len(pushed_values): + return pushed_values[idx] + except (ValueError, TypeError): + pass + return None + + monkeypatch.setattr(XCom, "get_one", mock_get_one) + + return _mock_xcom + + +class TestIterableOperator: + @classmethod + def create_mapped_operator( + cls, + dag: DAG, + expand_input: ExpandInput, + task_id: str = "my_task", + retries: int = DEFAULT_RETRIES, + do_xcom_push: bool = True, + task_concurrency: int | None = None, + execution_timeout: timedelta | None = None, + ) -> MappedOperator: + """ + Create a MappedOperator and assign it to a DAG. + + :param expand_input: The input to expand + :param dag: The DAG to assign the operator to + :param task_id: Task ID for the operator + :param do_xcom_push: Whether to push XCom (default True) + """ + return MockOperator.partial( + task_id=task_id, + dag=dag, + retries=retries, + task_concurrency=task_concurrency, + do_xcom_push=do_xcom_push, + execution_timeout=execution_timeout, + )._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + + @classmethod + def create_iterable_operator( + cls, + dag: DAG, + expand_input: ExpandInput, + task_id: str = "my_task", + task_concurrency: int | None = None, + retries: int = DEFAULT_RETRIES, + do_xcom_push: bool = True, + ) -> IterableOperator: + """Create an IterableOperator with a MappedOperator and ExpandInput.""" + mapped_op = cls.create_mapped_operator( + dag=dag, + expand_input=expand_input, + task_id=task_id, + retries=retries, + do_xcom_push=do_xcom_push, + task_concurrency=task_concurrency, + ) + return IterableOperator( + operator=mapped_op, + expand_input=expand_input, + dag=dag, + ) + + @pytest.mark.db_test + @pytest.mark.parametrize( + ("actual", "expected"), + [ + ([{"a": 1}, {"a": 2}], [{"a": 1}, {"a": 2}]), + ([{"a": 1, "b": 2}], [{"a": 1, "b": 2}]), + ([], []), + ], + ) + def test_list_of_dicts_expand_input_iter_values(self, dag_maker, session, actual, expected): + """Test IterableOperator with ListOfDictsExpandInput expand_input.""" + if not actual: + pytest.skip("Empty list case tested separately") + + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput(actual) + iterable_op = self.create_iterable_operator(dag, expand_input) + + result = list(iterable_op.expand_input.iter_values({})) + assert result == expected + + @pytest.mark.db_test + def test_list_of_dicts_empty(self, dag_maker, session): + """Test IterableOperator with empty list.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([]) + iterable_op = self.create_iterable_operator(dag, expand_input) + + result = list(iterable_op.expand_input.iter_values({})) + assert result == [] + + @pytest.mark.db_test + @pytest.mark.parametrize( + ("actual", "expected"), + [ + ({"a": 1}, [{"a": 1}]), + ({"a": [1, 2, 3]}, [{"a": 1}, {"a": 2}, {"a": 3}]), + ({"a": "hello"}, [{"a": "hello"}]), + ( + {"a": [1, 2], "b": [10, 20]}, + [{"a": 1, "b": 10}, {"a": 1, "b": 20}, {"a": 2, "b": 10}, {"a": 2, "b": 20}], + ), + ({"a": [1, 2]}, [{"a": 1}, {"a": 2}]), + ], + ) + def test_dict_of_lists_expand_input_iter_values(self, dag_maker, session, actual, expected): + """Test IterableOperator with DictOfListsExpandInput expand_input.""" + with dag_maker(session=session) as dag: + expand_input = DictOfListsExpandInput(actual) + iterable_op = self.create_iterable_operator(dag, expand_input) + + result = list(iterable_op.expand_input.iter_values({})) + assert result == expected + + @pytest.mark.db_test + def test_task_type(self, dag_maker, session): + """Test that IterableOperator correctly reports task_type.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.task_type == "MappedOperator" + + @pytest.mark.db_test + def test_task_retries(self, dag_maker, session): + """Test that IterableOperator correctly reports task_retries.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input, retries=3) + + assert isinstance(iterable_op, IterableOperator) + assert iterable_op.retries == 0 + assert iterable_op.task_retries == 3 + + @pytest.mark.db_test + def test_task_id(self, dag_maker, session): + """Test that IterableOperator inherits task_id from operator.""" + with dag_maker(session=session) as dag: + task_id = "my_task" + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input, task_id=task_id) + + assert iterable_op.task_id == task_id + + @pytest.mark.db_test + def test_with_task_concurrency(self, dag_maker, session): + """Test that IterableOperator respects task_concurrency parameter.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input, task_concurrency=4) + + assert iterable_op.max_workers == 4 + + @pytest.mark.db_test + @pytest.mark.parametrize("invalid_value", [0, -1, -10]) + def test_task_concurrency_validation_rejects_non_positive_values(self, dag_maker, session, invalid_value): + """Test that IterableOperator raises ValueError for task_concurrency < 1.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + with pytest.raises(ValueError, match=f"task_concurrency must be at least 1, got {invalid_value}"): + self.create_iterable_operator(dag, expand_input, task_concurrency=invalid_value) + + @pytest.mark.db_test + def test_partial_kwargs_not_mutated(self, dag_maker, session): + """Test that creating IterableOperator does not mutate the original MappedOperator's partial_kwargs.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + mapped_op = self.create_mapped_operator(dag, expand_input, task_concurrency=4) + original_partial_kwargs = mapped_op.partial_kwargs.copy() + + IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + # Verify that mapped_op.partial_kwargs was not mutated + assert mapped_op.partial_kwargs == original_partial_kwargs + assert "task_concurrency" in mapped_op.partial_kwargs + + @pytest.mark.db_test + def test_expand_input_stored(self, dag_maker, session): + """Test that IterableOperator stores expand_input correctly.""" + with dag_maker(session=session) as dag: + expand_input_data = ListOfDictsExpandInput([{"a": 1}, {"a": 2}]) + iterable_op = self.create_iterable_operator(dag, expand_input_data) + + assert iterable_op.expand_input is expand_input_data + assert isinstance(iterable_op.expand_input, (ListOfDictsExpandInput, DictOfListsExpandInput)) + + @pytest.mark.db_test + def test_partial_kwargs_stored(self, dag_maker, session): + """Test that IterableOperator stores partial_kwargs from operator.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input) + + assert hasattr(iterable_op, "partial_kwargs") + assert isinstance(iterable_op.partial_kwargs, dict) + + @pytest.mark.db_test + def test_xcom_push_delegates_to_task_when_not_pushed(self, dag_maker, session): + """_xcom_push delegates to task.xcom_push only when xcom_pushed is False.""" + from unittest import mock + + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input) + + task = mock.MagicMock() + task.xcom_pushed = False + task.task_id = "my_task" + task.index = 0 + + iterable_op._xcom_push(task=task, value="result_value") + + task.xcom_push.assert_called_once_with(key=BaseXCom.XCOM_RETURN_KEY, value="result_value") + + @pytest.mark.db_test + def test_xcom_push_skips_when_already_pushed(self, dag_maker, session): + """_xcom_push skips pushing when xcom_pushed is already True.""" + from unittest import mock + + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}]) + iterable_op = self.create_iterable_operator(dag, expand_input) + + task = mock.MagicMock() + task.xcom_pushed = True + task.task_id = "my_task" + task.index = 0 + + iterable_op._xcom_push(task=task, value="result_value") + + task.xcom_push.assert_not_called() + + @pytest.mark.db_test + def test_execute_list_of_dicts(self, dag_maker, session, mock_xcom_get_one): + """Test executing IterableOperator with ListOfDictsExpandInput.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = self.create_iterable_operator(dag, expand_input, task_id="exec_list_of_dicts") + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, None, None), (2, None, None)] + + @pytest.mark.db_test + def test_execute_dict_of_lists(self, dag_maker, session, mock_xcom_get_one): + """Test executing IterableOperator with DictOfListsExpandInput.""" + with dag_maker(session=session) as dag: + expand_input = DictOfListsExpandInput({"arg1": [1, 2, 3]}) + iterable_op = self.create_iterable_operator(dag, expand_input, task_id="exec_dict_of_lists") + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, None, None), (2, None, None), (3, None, None)] + + @pytest.mark.db_test + def test_execute_empty_list_of_dicts(self, dag_maker, session, mock_xcom_get_one): + """Test executing IterableOperator with empty ListOfDictsExpandInput.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([]) + iterable_op = self.create_iterable_operator(dag, expand_input, task_id="exec_empty") + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [] + + @pytest.mark.db_test + def test_execute_multiple_key_dict_of_lists(self, dag_maker, session, mock_xcom_get_one): + """Test executing IterableOperator with multiple keys in DictOfListsExpandInput.""" + with dag_maker(session=session) as dag: + expand_input = DictOfListsExpandInput({"arg1": [1, 2], "arg2": [10, 20], "arg3": ["x", "y"]}) + iterable_op = self.create_iterable_operator(dag, expand_input, task_id="exec_multi_key") + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + # Cartesian product expected order: + # (1,10,'x'), (1,10,'y'), (1,20,'x'), (1,20,'y'), + # (2,10,'x'), (2,10,'y'), (2,20,'x'), (2,20,'y') + assert materialized == [ + (1, 10, "x"), + (1, 10, "y"), + (1, 20, "x"), + (1, 20, "y"), + (2, 10, "x"), + (2, 10, "y"), + (2, 20, "x"), + (2, 20, "y"), + ] + + @pytest.mark.db_test + def test_execute_with_task_concurrency_setting(self, dag_maker, session, mock_xcom_get_one): + """Test executing IterableOperator with task_concurrency parameter.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}, {"arg1": 3}]) + iterable_op = self.create_iterable_operator( + dag, expand_input, task_id="exec_concurrency", task_concurrency=2 + ) + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, None, None), (2, None, None), (3, None, None)] + assert iterable_op.max_workers == 2 + + @pytest.mark.db_test + def test_execute_all_parameters(self, dag_maker, session, mock_xcom_get_one): + """Test executing IterableOperator with all arg1, arg2, arg3 parameters.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput( + [ + {"arg1": 1, "arg2": 10, "arg3": 100}, + {"arg1": 2, "arg2": 20, "arg3": 200}, + ] + ) + iterable_op = self.create_iterable_operator(dag, expand_input, task_id="exec_all_args") + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + assert materialized == [(1, 10, 100), (2, 20, 200)] + + @pytest.mark.db_test + def test_execute_with_do_xcom_push_false(self, dag_maker, session): + """Test executing IterableOperator when do_xcom_push is False.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"arg1": 1}, {"arg1": 2}]) + iterable_op = self.create_iterable_operator( + dag, expand_input, task_id="no_xcom_push", do_xcom_push=False + ) + + context = mock_context(task=iterable_op) + result = iterable_op.execute(context=context) + + assert result is None + + @pytest.mark.db_test + def test_execute_with_failed_tasks_but_no_retries(self, dag_maker, session, mock_xcom_get_one): + """ + Test executing IterableOperator where tasks fail but no retries are available. + + This test verifies that: + 1. Tasks with fail_on_first_attempt=True raise an exception on first attempt + 2. When no retries are configured (retries=0), the exception propagates and is not retried + 3. The BaseExceptionGroup is raised containing the task failure + """ + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput( + [ + {"arg1": 1, "arg2": 10}, + {"arg1": 2, "arg2": 20, "fail_on_first_attempt": True}, + {"arg1": 3, "arg2": 30}, + ] + ) + iterable_op = self.create_iterable_operator( + dag, + expand_input, + task_id="exec_with_failures", + retries=0, + ) + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + with pytest.raises(BaseExceptionGroup): + iterable_op.execute(context=context) + + @pytest.mark.db_test + def test_execute_with_failed_tasks_and_expired_reschedule_date( + self, dag_maker, session, mock_xcom_get_one + ): + """ + Test executing IterableOperator where certain map_index tasks fail on first attempt and are retried. + + This test verifies that: + 1. Tasks with fail_on_first_attempt=True raise an exception on first attempt (try_number == 0) + 2. Failed tasks are retried immediately without deferring (since reschedule_date is expired) + 3. Retried tasks succeed on subsequent attempts (try_number > 0) and produce the expected output + """ + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput( + [ + {"arg1": 1, "arg2": 10}, + {"arg1": 2, "arg2": 20, "fail_on_first_attempt": True}, + {"arg1": 3, "arg2": 30}, + ] + ) + iterable_op = self.create_iterable_operator( + dag, + expand_input, + task_id="exec_with_failures", + retries=1, + ) + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + result = iterable_op.execute(context=context) + materialized = list(result) + + assert len(materialized) == 3 + assert materialized == [(1, 10, None), (2, 20, None), (3, 30, None)] + + @pytest.mark.db_test + def test_iterable_execution_timeout_is_none_wrapped_operator_retains_it(self, dag_maker, session): + """IterableOperator.execution_timeout is None (not propagated to the outer TI); + the wrapped operator retains its own execution_timeout for per-task enforcement.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"a": 1}]) + execution_timeout = timedelta(seconds=7) + mapped_op = self.create_mapped_operator( + dag, expand_input, task_id="timeout_task", execution_timeout=execution_timeout + ) + + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + assert iterable_op._operator.execution_timeout == execution_timeout + assert iterable_op.execution_timeout is None + + @pytest.mark.db_test + @pytest.mark.parametrize( + "base_exception", + [ + SystemExit(1), + KeyboardInterrupt(), + GeneratorExit(), + ], + ids=["SystemExit", "KeyboardInterrupt", "GeneratorExit"], + ) + def test_base_exception_not_retried_raises_airflow_fail_exception( + self, dag_maker, session, mock_xcom_get_one, base_exception + ): + """ + BaseException subclasses (e.g., SystemExit, KeyboardInterrupt) must never + be retried—they signal conditions where continuing iteration is meaningless. + They should raise AirflowFailException immediately. + """ + with dag_maker(session=session) as dag: + # Create a mapped operator that raises a BaseException + expand_input = ListOfDictsExpandInput([{"raise_exception": base_exception}]) + mapped_op = MockOperator.partial( + task_id="base_exception_task", + dag=dag, + retries=3, # Has retries available, but should NOT use them + )._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + + with pytest.raises(AirflowFailException): + iterable_op.execute(context=context) + + @pytest.mark.db_test + def test_deferred_operator_raises_airflow_fail_exception(self, dag_maker, session, mock_xcom_get_one): + """A sub-task that raises TaskDeferred must cause IterableOperator to raise AirflowFailException.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{}, {}]) + mapped_op = MockDeferredOperator.partial(task_id="deferred_task", dag=dag)._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + + with pytest.raises(AirflowFailException, match="attempted to defer"): + iterable_op.execute(context=context) + + @pytest.mark.db_test + def test_reschedule_mode_sensor_raises_airflow_fail_exception( + self, dag_maker, session, mock_xcom_get_one + ): + """A sub-task that raises AirflowRescheduleException must fail the whole IterableOperator + immediately with a clear error rather than being silently mishandled.""" + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{}, {}]) + mapped_op = MockRescheduleSensor.partial(task_id="reschedule_sensor", dag=dag)._expand( + expand_input, + strict=True, + register_with_dag=False, + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + context = mock_context(task=iterable_op) + mock_xcom_get_one(context) + + with pytest.raises(AirflowFailException, match="attempted to reschedule"): + iterable_op.execute(context=context) + + +class TestIterableOperatorContextIsolation: + """ + Verify that each sub-task run by IterableOperator sees its own indexed + context via get_current_context(), not the parent's. + """ + + @pytest.mark.db_test + def test_subtask_sees_its_own_context(self, dag_maker, session, mock_xcom_get_one): + """Each sub-task's get_current_context() must return its own indexed ti, not the parent's.""" + captured: dict[int, object] = {} + + class ContextCapturingOperator(BaseOperator): + def __init__(self, index: int, **kwargs): + super().__init__(**kwargs) + self.index = index + + def execute(self, context): + ctx = get_current_context() + captured[self.index] = ctx["ti"] + return self.index + + with dag_maker(session=session) as dag: + expand_input = ListOfDictsExpandInput([{"index": 0}, {"index": 1}, {"index": 2}]) + mapped_op = ContextCapturingOperator.partial(task_id="ctx_task", dag=dag)._expand( + expand_input, strict=True, register_with_dag=False + ) + iterable_op = IterableOperator(operator=mapped_op, expand_input=expand_input, dag=dag) + + parent_context = mock_context(task=iterable_op) + mock_xcom_get_one(parent_context) + iterable_op.execute(context=parent_context) + + parent_ti = parent_context["ti"] + for idx, sub_ti in captured.items(): + # Each sub-task must have seen its own IndexedTaskInstance, not the parent TI. + assert sub_ti is not parent_ti, f"Sub-task {idx} observed the parent context" + assert sub_ti.index == idx, f"Sub-task {idx} observed wrong index {sub_ti.index}" diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py index 4aadf923e9c85..50f59b1ef88ba 100644 --- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py +++ b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py @@ -28,7 +28,9 @@ from airflow.sdk import TaskInstanceState, TriggerRule from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.bases.xcom import BaseXCom +from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput from airflow.sdk.definitions.dag import DAG +from airflow.sdk.definitions.iterableoperator import IterableOperator, MappedIterableOperator from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.sdk.definitions.xcom_arg import XComArg from airflow.sdk.execution_time.comms import ( @@ -123,6 +125,27 @@ def test_map_unknown_arg_raises(): BaseOperator.partial(task_id="a").expand(file=[1, 2, {"a": "b"}]) +@pytest.mark.parametrize( + "size", + [ + pytest.param(1), + pytest.param(3), + ], +) +def test_map_batch_size(size: int): + with DAG("test-dag", schedule=None): + mapped = ( + MockOperator.partial(task_id="task_2") + .batch(size=size) + ._iterate(DictOfListsExpandInput({"arg1": [1, 2, 3]}), strict=False) + ) + if size > 1: + assert isinstance(mapped, MappedIterableOperator) + assert mapped.batch_size == size + else: + assert isinstance(mapped, IterableOperator) + + def test_map_xcom_arg(): """Test that dependencies are correct when mapping with an XComArg""" with DAG("test-dag"): diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 2d1f8b2a4fa52..1069811361ee5 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -22,6 +22,7 @@ import pytest import structlog +from task_sdk.definitions.conftest import make_xcom_arg from airflow.sdk import TaskInstanceState from airflow.sdk.bases.xcom import BaseXCom @@ -416,3 +417,39 @@ def test_resolve_uses_xcom_pull_for_specific_index(self): assert resolved == "value-0" ti.xcom_pull.assert_called_once() assert ti.xcom_pull.call_args.kwargs["map_indexes"] == 0 + + +class TestXComArg: + @pytest.mark.parametrize( + ("actual", "expected"), + [ + (1, [1]), # scalar + ("hello", ["hello"]), # string + ([1, 2, 3], [1, 2, 3]), # list + ((x for x in [4, 5]), [4, 5]), # generator + ], + ) + def test_plain_xcomarg_iter_values(self, actual, expected): + xcom_arg = make_xcom_arg(actual) + result = list(xcom_arg.iter_values({})) + assert result == expected + + def test_map_xcomarg_iter_values(self): + base = make_xcom_arg([1, 2, 3]) + mapped = base.map(lambda x: x * 10) + result = list(mapped.iter_values({})) + assert result == [10, 20, 30] + + def test_zip_xcomarg_iter_values(self): + a = make_xcom_arg([1, 2]) + b = make_xcom_arg([10, 20]) + zipped = a.zip(b) + result = list(zipped.iter_values({})) + assert result == [(1, 10), (2, 20)] + + def test_concat_xcomarg_iter_values(self): + a = make_xcom_arg([1, 2]) + b = make_xcom_arg([10, 20]) + concatenated = a.concat(b) + result = list(concatenated.iter_values({})) + assert result == [1, 2, 10, 20] diff --git a/task-sdk/tests/task_sdk/execution_time/conftest.py b/task-sdk/tests/task_sdk/execution_time/conftest.py index 2d60b635f63f3..27fced5cc15ee 100644 --- a/task-sdk/tests/task_sdk/execution_time/conftest.py +++ b/task-sdk/tests/task_sdk/execution_time/conftest.py @@ -18,9 +18,16 @@ from __future__ import annotations import sys +from datetime import timedelta from socket import socketpair +from unittest import mock import pytest +from uuid6 import uuid7 + +from airflow.sdk import BaseAsyncOperator, BaseOperator, timezone +from airflow.sdk.api.datamodels._generated import TaskInstanceState +from airflow.sdk.execution_time.task_runner import IndexedTaskInstance @pytest.fixture @@ -43,3 +50,70 @@ def disable_capturing(): sys.stderr = sys.__stderr__ yield sys.stdin, sys.stdout, sys.stderr = old_in, old_out, old_err + + +@pytest.fixture +def make_indexed_ti(): + """Factory for creating IndexedTaskInstance objects for testing.""" + + def _make_indexed_ti( + *, + task_id: str = "my_task", + dag_id: str = "my_dag", + run_id: str = "run_1", + map_index: int = -1, + index: int | None = 0, + try_number: int = 0, + max_tries: int = 3, + is_async: bool = False, + retry_delay: timedelta = None, + retry_exponential_backoff: bool = False, + max_retry_delay: timedelta | None = None, + end_date=None, + start_date=None, + logical_date=None, + do_xcom_push: bool = True, + ) -> IndexedTaskInstance: + """Create a IndexedTaskInstance via model_construct to bypass full Pydantic validation.""" + if retry_delay is None: + retry_delay = timedelta(seconds=300) + + # Set defaults for dates if not provided + if end_date is None: + end_date = timezone.datetime(2024, 12, 3, 10, 0, 0) + if start_date is None: + start_date = timezone.datetime(2024, 12, 3, 9, 55, 0) + if logical_date is None: + logical_date = timezone.datetime(2024, 12, 3, 0, 0, 0) + + operator_cls = BaseAsyncOperator if is_async else BaseOperator + operator = mock.create_autospec(operator_cls, instance=True) + operator.task_id = task_id + operator.dag_id = dag_id + operator.is_async = is_async + operator.retries = max_tries + operator.retry_delay = retry_delay + operator.retry_exponential_backoff = retry_exponential_backoff + operator.max_retry_delay = max_retry_delay + operator.do_xcom_push = do_xcom_push + + return IndexedTaskInstance.model_construct( + id=uuid7(), + task_id=task_id, + dag_id=dag_id, + run_id=run_id, + map_index=map_index, + index=index, + try_number=try_number, + max_tries=max_tries, + state=TaskInstanceState.SCHEDULED, + is_mapped=True, + task=operator, + xcom_pushed=False, + dag_version_id=uuid7(), + end_date=end_date, + start_date=start_date, + logical_date=logical_date, + ) + + return _make_indexed_ti diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index 31009408ce4df..86306bc659d5a 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -940,6 +940,32 @@ def test_for_asset_alias(self, mocked__getitem__): outlet_event_accessors.for_asset_alias(name="name") assert mocked__getitem__.call_args[0][0] == TEST_ASSET_ALIAS + def test_concurrent_access_same_asset_preserves_accessor(self): + """Concurrent __getitem__ for the same asset must not overwrite an existing accessor.""" + import threading + + accessors = OutletEventAccessors() + asset = Asset("concurrent-test") + results: list[OutletEventAccessor] = [] + + barrier = threading.Barrier(2) + + def access(): + barrier.wait() + results.append(accessors[asset]) + + threads = [threading.Thread(target=access) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(results) == 2 + # Both threads must have received the identical accessor object so + # that neither thread's accumulated events can be silently discarded. + assert results[0] is results[1] + assert len(accessors) == 1 + class TestInletEventAccessor: @pytest.fixture diff --git a/task-sdk/tests/task_sdk/execution_time/test_executor.py b/task-sdk/tests/task_sdk/execution_time/test_executor.py new file mode 100644 index 0000000000000..d9587f2ed9ef0 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_executor.py @@ -0,0 +1,515 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import time +from unittest import mock + +import pytest +from task_sdk.execution_time.test_task_runner import get_inline_dag + +from airflow.sdk import BaseOperator +from airflow.sdk.api.datamodels._generated import TaskInstanceState +from airflow.sdk.bases.operator import event_loop +from airflow.sdk.exceptions import ( + AirflowRescheduleException, + AirflowRescheduleTaskInstanceException, + TaskDeferred, +) +from airflow.sdk.execution_time.executor import AsyncAwareExecutor, TaskExecutor + +from tests_common.test_utils.mock_context import mock_context + + +class TestAsyncAwareExecutor: + def test_submit_sync_function_returns_future(self): + """Sync callables are dispatched to the thread pool and return an asyncio.Future.""" + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + future = executor.submit(lambda: 42) + assert isinstance(future, asyncio.Future) + assert loop.run_until_complete(future) == 42 + + def test_submit_async_coroutine_function_returns_task(self): + """Async callables are scheduled on the event loop and return an asyncio.Task.""" + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + + async def async_fn(): + return "async_result" + + task = executor.submit(async_fn) + assert isinstance(task, asyncio.Task) + result = loop.run_until_complete(task) + assert result == "async_result" + + def test_submit_coroutine_object_returns_task(self): + """Passing a coroutine object (not a function) directly is also scheduled on the event loop.""" + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + + async def async_fn(): + return "coro_result" + + coro = async_fn() + task = executor.submit(coro) + assert isinstance(task, asyncio.Task) + result = loop.run_until_complete(task) + assert result == "coro_result" + + def test_submit_sync_function_propagates_exception(self): + """Exceptions raised inside sync callables are propagated when the future is resolved.""" + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + future = executor.submit(lambda: (_ for _ in ()).throw(ValueError("boom"))) + assert isinstance(future, asyncio.Future) + with pytest.raises(ValueError, match="boom"): + loop.run_until_complete(future) + + def test_submit_async_function_propagates_exception(self): + """Exceptions raised inside async callables are propagated when the task is awaited.""" + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + + async def failing(): + raise RuntimeError("async boom") + + task = executor.submit(failing) + with pytest.raises(RuntimeError, match="async boom"): + loop.run_until_complete(task) + + def test_semaphore_limits_concurrent_async_tasks(self): + """The semaphore prevents more than max_workers coroutines from running simultaneously.""" + concurrency_high_watermark = 0 + running = 0 + + async def count_concurrent(): + nonlocal concurrency_high_watermark, running + running += 1 + concurrency_high_watermark = max(concurrency_high_watermark, running) + await asyncio.sleep(0) + running -= 1 + + max_workers = 2 + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=max_workers) as executor: + tasks = [executor.submit(count_concurrent) for _ in range(6)] + loop.run_until_complete(asyncio.gather(*tasks)) + + assert concurrency_high_watermark <= max_workers + + def test_exit_shuts_down_thread_pool(self): + """__exit__ calls shutdown on the thread pool.""" + with event_loop() as loop: + executor = AsyncAwareExecutor(loop=loop, max_workers=2) + with mock.patch.object( + executor._thread_pool, "shutdown", wraps=executor._thread_pool.shutdown + ) as shutdown_mock: + with executor: + pass + shutdown_mock.assert_called_once_with(wait=True, cancel_futures=False) + + def test_context_manager_returns_self(self): + """__enter__ returns the executor instance itself.""" + with event_loop() as loop: + executor = AsyncAwareExecutor(loop=loop, max_workers=2) + with executor as ctx: + assert ctx is executor + + def test_map_streams_completed_sync_results(self): + """map() yields completed results as work finishes instead of waiting for all items.""" + + def sleepy_value(delay: float) -> float: + time.sleep(delay) + return delay + + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + started = time.monotonic() + result_iter = executor.map(sleepy_value, [0.25, 0.01]) + first = next(result_iter) + + assert first == 0.01 + # The faster work (0.01s) should complete well before the slower work (0.25s). + # Allow overhead for thread pool scheduling (typically ~0.15-0.2s on busy systems). + assert time.monotonic() - started < 0.35 + + def test_shutdown_cancel_futures_cancels_async_tasks(self): + """shutdown(cancel_futures=True) cancels submitted async tasks.""" + + async def long_running(): + await asyncio.sleep(60) + + with event_loop() as loop: + executor = AsyncAwareExecutor(loop=loop, max_workers=2) + task = executor.submit(long_running) + + executor.shutdown(wait=False, cancel_futures=True) + loop.run_until_complete(asyncio.sleep(0)) + + assert task.cancelled() + + def test_submit_after_shutdown_raises_runtime_error(self): + with event_loop() as loop: + executor = AsyncAwareExecutor(loop=loop, max_workers=2) + executor.shutdown(wait=False) + + with pytest.raises(RuntimeError, match="cannot schedule new futures after shutdown"): + executor.submit(lambda: 1) + + def test_map_rejects_non_positive_chunksize(self): + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + with pytest.raises(ValueError, match="chunksize must be >= 1"): + list(executor.map(lambda x: x, [1, 2], chunksize=0)) + + def test_map_timeout_raises_timeout_error(self): + def slow_fn(delay: float) -> float: + time.sleep(delay) + return delay + + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=1) as executor: + with pytest.raises(TimeoutError): + list(executor.map(slow_fn, [0.2], timeout=0.01)) + + def test_map_streams_completed_async_results(self): + async def async_sleepy_value(delay: float) -> float: + await asyncio.sleep(delay) + return delay + + with event_loop() as loop: + with AsyncAwareExecutor(loop=loop, max_workers=2) as executor: + started = time.monotonic() + result_iter = executor.map(async_sleepy_value, [0.2, 0.01]) + first = next(result_iter) + + assert first == 0.01 + # The faster work (0.01s) should complete well before the slower work (0.2s). + # Allow overhead for event loop scheduling (typically ~0.1-0.15s on busy systems). + assert time.monotonic() - started < 0.3 + + def test_shutdown_wait_true_waits_for_async_tasks(self): + async def short_running() -> str: + await asyncio.sleep(0.01) + return "done" + + with event_loop() as loop: + executor = AsyncAwareExecutor(loop=loop, max_workers=2) + task = executor.submit(short_running) + + executor.shutdown(wait=True) + + assert task.done() + assert not task.cancelled() + assert task.result() == "done" + + +class TestTaskExecutor: + def test_dag_id_property(self, make_indexed_ti): + ti = make_indexed_ti(dag_id="my_dag") + executor = TaskExecutor(task_instance=ti) + assert executor.dag_id == "my_dag" + + def test_task_id_property(self, make_indexed_ti): + ti = make_indexed_ti(task_id="my_task") + executor = TaskExecutor(task_instance=ti) + assert executor.task_id == "my_task" + + def test_task_index(self, make_indexed_ti): + ti = make_indexed_ti(index=3) + executor = TaskExecutor(task_instance=ti) + assert executor.task_index == ti.index + assert executor.task_index == 3 + + def test_operator_property(self, make_indexed_ti): + ti = make_indexed_ti() + executor = TaskExecutor(task_instance=ti) + assert executor.operator is ti.task + + def test_is_async_property_sync(self, make_indexed_ti): + ti = make_indexed_ti(is_async=False) + executor = TaskExecutor(task_instance=ti) + assert executor.is_async is False + + def test_is_async_property_async(self, make_indexed_ti): + ti = make_indexed_ti(is_async=True) + executor = TaskExecutor(task_instance=ti) + assert executor.is_async is True + + def test_enter_sets_start_time(self, make_indexed_ti): + ti = make_indexed_ti() + executor = TaskExecutor(task_instance=ti) + assert executor._start_time is None + executor.__enter__() + assert executor._start_time is not None + + def test_enter_returns_self(self, make_indexed_ti): + ti = make_indexed_ti() + executor = TaskExecutor(task_instance=ti) + with executor as ctx: + assert ctx is executor + + def test_exit_success_sets_state(self, make_indexed_ti): + """__exit__ without an exception marks the task instance as SUCCESS.""" + ti = make_indexed_ti() + with TaskExecutor(task_instance=ti): + pass # no exception + assert ti.state == TaskInstanceState.SUCCESS + + def test_exit_with_task_deferred_reraises(self, make_indexed_ti): + """TaskDeferred must propagate unchanged through __exit__.""" + ti = make_indexed_ti() + trigger = mock.Mock() + deferred = TaskDeferred(trigger=trigger, method_name="resume") + + with pytest.raises(TaskDeferred): + with TaskExecutor(task_instance=ti): + raise deferred + + def test_exit_with_reschedule_exception_passes_through(self, make_indexed_ti): + """AirflowRescheduleException (base, from a reschedule-mode sensor) must propagate + unchanged through __exit__ without consuming the retry budget or setting task state, + so _run_tasks can detect and reject it with a clear error.""" + ti = make_indexed_ti(try_number=0, max_tries=3) + from datetime import timedelta + + from airflow.sdk import timezone + + exc = AirflowRescheduleException(timezone.utcnow() + timedelta(seconds=60)) + + with pytest.raises(AirflowRescheduleException): + with TaskExecutor(task_instance=ti): + raise exc + + # State and try_number must be unchanged — this is not a retry-eligible failure. + assert ti.try_number == 0 + assert ti.state != TaskInstanceState.FAILED + + def test_exit_reschedules_when_retries_remain(self, make_indexed_ti): + """ + When a retryable exception occurs and retries are not exhausted, + the task state is set to UP_FOR_RESCHEDULE and + AirflowRescheduleTaskInstanceException is raised. + """ + # try_number=0 → next_try_number=1; max_tries=3 → 1 <= 3, reschedule + ti = make_indexed_ti(try_number=0, max_tries=3) + + with pytest.raises(AirflowRescheduleTaskInstanceException): + with TaskExecutor(task_instance=ti): + raise RuntimeError("transient failure") + + assert ti.state == TaskInstanceState.UP_FOR_RESCHEDULE + assert ti.try_number == 1 # incremented + + def test_exit_fails_when_retries_exhausted(self, make_indexed_ti): + """ + When a retryable exception occurs and all retries are exhausted, + the task state is set to FAILED and the original exception is re-raised. + """ + # try_number=3 → next_try_number=4; max_tries=3 → 4 > 3, fail + ti = make_indexed_ti(try_number=3, max_tries=3) + original_error = RuntimeError("permanent failure") + + with pytest.raises(RuntimeError, match="permanent failure"): + with TaskExecutor(task_instance=ti): + raise original_error + + assert ti.state == TaskInstanceState.FAILED + + @pytest.mark.parametrize( + ("try_number", "max_tries", "should_fail"), + [ + (0, 0, True), # next=1, max=0 → 1>0 → fail + (0, 1, False), # next=1, max=1 → 1<=1 → reschedule + (1, 1, True), # next=2, max=1 → 2>1 → fail + (2, 3, False), # next=3, max=3 → 3<=3 → reschedule + (3, 3, True), # next=4, max=3 → 4>3 → fail + ], + ) + def test_exit_retry_boundary(self, make_indexed_ti, try_number, max_tries, should_fail): + """Exhaustive boundary checks for the retry/fail decision in __exit__.""" + ti = make_indexed_ti(try_number=try_number, max_tries=max_tries) + if should_fail: + with pytest.raises(RuntimeError): + with TaskExecutor(task_instance=ti): + raise RuntimeError("err") + assert ti.state == TaskInstanceState.FAILED + else: + with pytest.raises(AirflowRescheduleTaskInstanceException): + with TaskExecutor(task_instance=ti): + raise RuntimeError("err") + assert ti.state == TaskInstanceState.UP_FOR_RESCHEDULE + + def test_run_delegates_to_execute_task(self, make_indexed_ti): + """run() must call _execute_task with the given context.""" + ti = make_indexed_ti() + task = BaseOperator(task_id="test_task") + get_inline_dag("test_dag", task) + context = mock_context(task) + executor = TaskExecutor(task_instance=ti) + + with mock.patch( + "airflow.sdk.execution_time.executor._execute_task", + autospec=True, + return_value="result", + ) as mock_execute: + result = executor.run(context) + + mock_execute.assert_called_once_with(context, ti, executor.log) + assert result == "result" + + @pytest.mark.asyncio + async def test_arun_delegates_to_execute_async_task(self, make_indexed_ti): + """arun() must call _execute_async_task with the given context.""" + ti = make_indexed_ti(is_async=True) + task = BaseOperator(task_id="test_task") + get_inline_dag("test_dag", task) + context = mock_context(task) + executor = TaskExecutor(task_instance=ti) + + with mock.patch( + "airflow.sdk.execution_time.executor._execute_async_task", + new=mock.AsyncMock(return_value="async_result"), + ) as mock_async_execute: + result = await executor.arun(context) + + mock_async_execute.assert_called_once_with(context, ti, executor.log) + assert result == "async_result" + + @pytest.mark.parametrize( + "base_exception", + [ + SystemExit(1), + KeyboardInterrupt(), + GeneratorExit(), + ], + ids=["SystemExit", "KeyboardInterrupt", "GeneratorExit"], + ) + def test_exit_base_exception_not_retried(self, make_indexed_ti, base_exception): + """ + BaseException subclasses (e.g., SystemExit, KeyboardInterrupt) must never + be retried—they signal conditions where continuing is meaningless. + They should be re-raised immediately and mark the task as FAILED. + """ + ti = make_indexed_ti(try_number=0, max_tries=3) + + with pytest.raises(type(base_exception)): + with TaskExecutor(task_instance=ti): + raise base_exception + + assert ti.state == TaskInstanceState.FAILED + # try_number should NOT be incremented for BaseException + assert ti.try_number == 0 + + def test_exit_success_fires_on_success_callback(self, make_indexed_ti): + """on_success_callback must fire for each indexed sub-task that succeeds.""" + fired: list[str] = [] + ti = make_indexed_ti() + task = BaseOperator( + task_id="cb_task", + on_success_callback=lambda ctx: fired.append("success"), + ) + get_inline_dag("cb_dag", task) + ti.task = task + context = mock_context(task) + executor = TaskExecutor(task_instance=ti) + executor._context = context + + with executor: + pass + + assert fired == ["success"] + + def test_exit_failure_fires_on_failure_callback(self, make_indexed_ti): + """on_failure_callback must fire for each indexed sub-task that exhausts retries.""" + fired: list[str] = [] + ti = make_indexed_ti(try_number=3, max_tries=3) + task = BaseOperator( + task_id="cb_task", + on_failure_callback=lambda ctx: fired.append("failure"), + ) + get_inline_dag("cb_dag", task) + ti.task = task + context = mock_context(task) + executor = TaskExecutor(task_instance=ti) + executor._context = context + + with pytest.raises(RuntimeError): + with executor: + raise RuntimeError("permanent") + + assert fired == ["failure"] + + def test_exit_retry_fires_on_retry_callback(self, make_indexed_ti): + """on_retry_callback must fire for each indexed sub-task that is rescheduled.""" + fired: list[str] = [] + ti = make_indexed_ti(try_number=0, max_tries=3) + task = BaseOperator( + task_id="cb_task", + on_retry_callback=lambda ctx: fired.append("retry"), + ) + get_inline_dag("cb_dag", task) + ti.task = task + context = mock_context(task) + executor = TaskExecutor(task_instance=ti) + executor._context = context + + with pytest.raises(AirflowRescheduleTaskInstanceException): + with executor: + raise RuntimeError("transient") + + assert fired == ["retry"] + + def test_exit_base_exception_fires_on_failure_callback(self, make_indexed_ti): + """on_failure_callback must also fire when a non-Exception BaseException marks the task FAILED.""" + fired: list[str] = [] + ti = make_indexed_ti(try_number=0, max_tries=3) + task = BaseOperator( + task_id="cb_task", + on_failure_callback=lambda ctx: fired.append("failure"), + ) + get_inline_dag("cb_dag", task) + ti.task = task + context = mock_context(task) + executor = TaskExecutor(task_instance=ti) + executor._context = context + + with pytest.raises(SystemExit): + with executor: + raise SystemExit(1) + + assert fired == ["failure"] + + def test_exit_no_context_skips_callbacks(self, make_indexed_ti): + """When _context is not set (e.g. __exit__ called directly), callbacks must not fire.""" + fired: list[str] = [] + ti = make_indexed_ti() + task = BaseOperator( + task_id="cb_task", + on_success_callback=lambda ctx: fired.append("success"), + ) + get_inline_dag("cb_dag", task) + ti.task = task + executor = TaskExecutor(task_instance=ti) + # _context intentionally left as None + + with executor: + pass + + assert fired == [] diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index cc9fb77e08921..d4f3a6fb8d96e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2067,14 +2067,16 @@ def test_run_with_asset_inlets(create_runtime_ti, mock_supervisor_comms): inlet_events[Asset(name="no such asset in inlets")] -@mock.patch("airflow.sdk.execution_time.task_runner.context_to_airflow_vars") -@mock.patch.dict(os.environ, {}, clear=True) -def test_execute_task_exports_env_vars( - mock_context_to_airflow_vars, create_runtime_ti, mock_supervisor_comms -): - """Test that _execute_task exports airflow context to environment variables.""" +def test_execute_task_exports_context_vars_thread_safely(create_runtime_ti, mock_supervisor_comms): + """Test that _execute_task exports airflow context via thread-safe context vars.""" + from airflow.sdk.execution_time.context import get_airflow_context_var + + captured_vars = {} def test_function(): + # Capture context vars during execution - they should be available via get_airflow_context_var + captured_vars["dag_id"] = get_airflow_context_var("AIRFLOW_CTX_DAG_ID") + captured_vars["task_id"] = get_airflow_context_var("AIRFLOW_CTX_TASK_ID") return "test function" task = PythonOperator( @@ -2082,14 +2084,15 @@ def test_function(): python_callable=test_function, ) - ti = create_runtime_ti(task=task, dag_id="dag_with_env_vars") - - mock_env_vars = {"AIRFLOW_CTX_DAG_ID": "test_dag_env_vars", "AIRFLOW_CTX_TASK_ID": "test_env_task"} - mock_context_to_airflow_vars.return_value = mock_env_vars + ti = create_runtime_ti(task=task, dag_id="dag_with_ctx_vars") run(ti, ti.get_template_context(), log=mock.MagicMock()) - assert os.environ["AIRFLOW_CTX_DAG_ID"] == "test_dag_env_vars" - assert os.environ["AIRFLOW_CTX_TASK_ID"] == "test_env_task" + # Verify context vars were accessible during task execution + assert captured_vars["dag_id"] == "dag_with_ctx_vars" + assert captured_vars["task_id"] == "test_task" + + # os.environ should be updated + assert os.environ.get("AIRFLOW_CTX_DAG_ID") == "dag_with_ctx_vars" def test_execute_success_task_with_rendered_map_index(create_runtime_ti, mock_supervisor_comms): @@ -2145,6 +2148,76 @@ def test_function(ti): assert ti.rendered_map_index == "Label: test_task" +class TestIndexedTaskInstance: + @pytest.mark.parametrize( + ("index", "key", "value", "expected_key", "expected_xcom_pushed"), + [ + (3, "result", "ok", "result_3", False), + (2, BaseXCom.XCOM_RETURN_KEY, "value1", f"{BaseXCom.XCOM_RETURN_KEY}_2", True), + (1, "custom_key", "value2", "custom_key_1", False), + ], + ids=["delegates_with_index_suffix", "sets_flag_for_default_key", "does_not_set_flag_for_custom_key"], + ) + def test_xcom_push_suffix_and_flag( + self, make_indexed_ti, index, key, value, expected_key, expected_xcom_pushed + ): + """xcom_push appends map index suffix and only marks default-key pushes.""" + ti = make_indexed_ti(index=index) + assert ti.xcom_pushed is False + + with mock.patch("airflow.sdk.execution_time.task_runner._xcom_push", autospec=True) as mock_push: + ti.xcom_push(key=key, value=value) + + mock_push.assert_called_once_with(ti, expected_key, value) + assert ti.xcom_pushed is expected_xcom_pushed + + def test_xcom_pull_delegates_with_index_suffix(self, make_indexed_ti): + """xcom_pull delegates to RuntimeTaskInstance with key suffixed by _{index}.""" + ti = make_indexed_ti(index=5) + + with mock.patch.object( + ti.__class__.__bases__[0], + "xcom_pull", + autospec=True, + return_value="pulled_value", + ) as mock_pull: + result = ti.xcom_pull(key="result") + + mock_pull.assert_called_once_with( + ti, + task_ids=None, + dag_id=None, + key="result_5", + include_prior_dates=False, + map_indexes=mock.ANY, + default=None, + run_id=None, + ) + assert result == "pulled_value" + + def test_next_retry_datetime_without_exponential_backoff(self, make_indexed_ti): + ti = make_indexed_ti(retry_delay=timedelta(seconds=30), retry_exponential_backoff=False) + + assert ti.next_retry_datetime() == timezone.datetime(2024, 12, 3, 10, 0, 30) + + def test_next_retry_datetime_exponential_backoff_honors_max_retry_delay(self, make_indexed_ti): + ti = make_indexed_ti( + try_number=2, + retry_delay=timedelta(seconds=10), + retry_exponential_backoff=True, + max_retry_delay=timedelta(seconds=5), + ) + + assert ti.next_retry_datetime() == timezone.datetime(2024, 12, 3, 10, 0, 5) + + def test_properties(self, make_indexed_ti): + ti = make_indexed_ti(index=7, try_number=4, is_async=True, do_xcom_push=False) + + assert ti.is_async is True + assert ti.next_try_number == 5 + assert ti.do_xcom_push is False + + class TestSerializeOutletEvents: """Tests for the wire format produced by ``_serialize_outlet_events``.""" @@ -2377,6 +2450,43 @@ def test_lazy_loading_not_triggered_until_accessed(self, create_runtime_ti, mock # Now the lazy attribute should trigger the call mock_supervisor_comms.send.assert_called_once() + def test_logical_date_returns_none_without_ti_context_from_server(self, mocked_parse): + """Test that logical_date returns None when _ti_context_from_server is not set.""" + task = BaseOperator(task_id="hello") + dag_id = "basic_task" + + get_inline_dag(dag_id=dag_id, task=task) + + ti_id = uuid7() + ti = TaskInstance( + id=ti_id, + task_id=task.task_id, + dag_id=dag_id, + run_id="test_run", + try_number=1, + dag_version_id=uuid7(), + ) + start_date = timezone.datetime(2025, 1, 1) + + runtime_ti = RuntimeTaskInstance.model_construct( + **ti.model_dump(exclude_unset=True), + task=task, + _ti_context_from_server=None, + start_date=start_date, + ) + + assert runtime_ti.logical_date is None + + def test_logical_date_returns_dag_run_logical_date(self, create_runtime_ti): + """Test that logical_date returns the dag run's logical_date when _ti_context_from_server is set.""" + task = BaseOperator(task_id="hello") + runtime_ti = create_runtime_ti(task=task, dag_id="basic_task") + + dag_run = runtime_ti._ti_context_from_server.dag_run + + assert runtime_ti.logical_date == dag_run.logical_date + assert runtime_ti.logical_date == timezone.datetime(2024, 12, 1, 1, 0, 0) + def test_get_connection_from_context(self, create_runtime_ti, mock_supervisor_comms): """Test that the connection is fetched from the API server via the Supervisor lazily when accessed""" diff --git a/uv.lock b/uv.lock index 9f17a033e820a..adff0f6f276f6 100644 --- a/uv.lock +++ b/uv.lock @@ -330,7 +330,7 @@ name = "adbc-driver-manager" version = "1.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e9/5e/50aab18cb501e42d3aca3cd2cc26c6637094fcaf5b6576e350c444188f1f/adbc_driver_manager-1.11.0.tar.gz", hash = "sha256:c64aaabeb5810109ab3d2961008f1b014e9f2d87b3df4416c2a080a40237af50", size = 233059, upload-time = "2026-04-07T00:17:28.263Z" } wheels = [ @@ -375,8 +375,8 @@ name = "adbc-driver-postgresql" version = "1.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "adbc-driver-manager" }, - { name = "importlib-resources" }, + { name = "adbc-driver-manager", marker = "python_full_version < '3.13'" }, + { name = "importlib-resources", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/68/10/2962c25035887cd03af3b348eac3302493936f45048c220021a802d07f12/adbc_driver_postgresql-1.11.0.tar.gz", hash = "sha256:f5688b8648ac7a86d8b89340231bb3686ac5df56ee95d1ca0b875dad5d52b48a", size = 32328, upload-time = "2026-04-07T00:17:29.232Z" } wheels = [ @@ -392,8 +392,8 @@ name = "adbc-driver-sqlite" version = "1.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "adbc-driver-manager" }, - { name = "importlib-resources" }, + { name = "adbc-driver-manager", marker = "python_full_version < '3.13'" }, + { name = "importlib-resources", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3a/dd/8a5f4908aa4bdec64dcd672734fa314d692517458ce169591639d0123fe1/adbc_driver_sqlite-1.11.0.tar.gz", hash = "sha256:a4c6b4962610f7cd67cd754c42dd74e18a2c11fabeec9488c5501d73ae62dc62", size = 28885, upload-time = "2026-04-07T00:17:31.325Z" } wheels = [ @@ -456,7 +456,7 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "caio" }, + { name = "caio", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/67/e2/d7cb819de8df6b5c1968a2756c3cb4122d4fa2b8fc768b53b7c9e5edb646/aiofile-3.9.0.tar.gz", hash = "sha256:e5ad718bb148b265b6df1b3752c4d1d83024b93da9bd599df74b9d9ffcf7919b", size = 17943, upload-time = "2024-10-08T10:39:35.846Z" } wheels = [ @@ -480,7 +480,7 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "caio" }, + { name = "caio", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/48/41/2fea7e193e061ce54eacc3b7bc0e6a99e4fcff43c78cf0a76dd781ed8334/aiofile-3.11.1.tar.gz", hash = "sha256:1f91912c6643d2a4e49ca4ae3514f0bf3867ce948a36d99a6411b8f4755f4cf9", size = 19342, upload-time = "2026-05-16T08:18:33.538Z" } wheels = [ @@ -647,7 +647,7 @@ name = "aiohttp-cors" version = "0.8.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiohttp" }, + { name = "aiohttp", marker = "python_full_version < '3.15'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/d89e846a5444b3d5eb8985a6ddb0daef3774928e1bfbce8e84ec97b0ffa7/aiohttp_cors-0.8.1.tar.gz", hash = "sha256:ccacf9cb84b64939ea15f859a146af1f662a6b1d68175754a07315e305fb1403", size = 38626, upload-time = "2025-03-31T14:16:20.048Z" } wheels = [ @@ -10447,8 +10447,8 @@ name = "cassandra-driver" version = "3.30.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "deprecated" }, - { name = "geomet" }, + { name = "deprecated", marker = "python_full_version < '3.14'" }, + { name = "geomet", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/bb/ed/4e16210e194660f107929ee494f1cf18557252655067d9b39029c241be2d/cassandra_driver-3.30.1.tar.gz", hash = "sha256:0c6a3e1428f7c6a9aa6c944b9c47a37cd2cdbeb5b5a82d42c33afdd56f14e398", size = 289233, upload-time = "2026-07-06T20:14:31.37Z" } wheels = [ @@ -11013,7 +11013,7 @@ name = "colorful" version = "0.5.8" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "python_full_version < '3.15' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/82/31/109ef4bedeb32b4202e02ddb133162457adc4eb890a9ed9c05c9dd126ed0/colorful-0.5.8.tar.gz", hash = "sha256:bb16502b198be2f1c42ba3c52c703d5f651d826076817185f0294c1a549a7445", size = 209361, upload-time = "2025-10-29T11:53:21.663Z" } wheels = [ @@ -12495,7 +12495,7 @@ name = "geomet" version = "1.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, + { name = "click", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/2a/8c/dde022aa6747b114f6b14a7392871275dea8867e2bd26cddb80cc6d66620/geomet-1.1.0.tar.gz", hash = "sha256:51e92231a0ef6aaa63ac20c443377ba78a303fd2ecd179dc3567de79f3c11605", size = 28732, upload-time = "2023-11-14T15:43:36.764Z" } wheels = [ @@ -13908,7 +13908,7 @@ name = "gssapi" version = "1.11.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "decorator" }, + { name = "decorator", marker = "python_full_version < '3.15' or sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/23/52/c1e90623c259a42ab0587078bb04f959867b970add46ff66750ead8fc7c5/gssapi-1.11.1.tar.gz", hash = "sha256:2049ee4b1d0c363163a1344b7282a363f9f4094e51d2c36de0cf01d4735e0ae2", size = 95233, upload-time = "2026-01-26T21:01:39.463Z" } wheels = [ @@ -14685,17 +14685,17 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "decorator" }, - { name = "exceptiongroup" }, - { name = "jedi" }, - { name = "matplotlib-inline" }, - { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit" }, - { name = "pygments" }, - { name = "stack-data" }, - { name = "traitlets" }, - { name = "typing-extensions" }, + { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, + { name = "decorator", marker = "python_full_version < '3.11'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "jedi", marker = "python_full_version < '3.11'" }, + { name = "matplotlib-inline", marker = "python_full_version < '3.11'" }, + { name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit", marker = "python_full_version < '3.11'" }, + { name = "pygments", marker = "python_full_version < '3.11'" }, + { name = "stack-data", marker = "python_full_version < '3.11'" }, + { name = "traitlets", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/40/18/f8598d287006885e7136451fdea0755af4ebcbfe342836f24deefaed1164/ipython-8.39.0.tar.gz", hash = "sha256:4110ae96012c379b8b6db898a07e186c40a2a1ef5d57a7fa83166047d9da7624", size = 5513971, upload-time = "2026-03-27T10:02:13.94Z" } wheels = [ @@ -14719,18 +14719,18 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "decorator" }, - { name = "ipython-pygments-lexers" }, - { name = "jedi" }, - { name = "matplotlib-inline" }, - { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit" }, - { name = "psutil", marker = "sys_platform != 'cygwin' and sys_platform != 'emscripten'" }, - { name = "pygments" }, - { name = "stack-data" }, - { name = "traitlets" }, - { name = "typing-extensions", marker = "python_full_version < '3.12'" }, + { name = "colorama", marker = "python_full_version >= '3.11' and sys_platform == 'win32'" }, + { name = "decorator", marker = "python_full_version >= '3.11'" }, + { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.11'" }, + { name = "jedi", marker = "python_full_version >= '3.11'" }, + { name = "matplotlib-inline", marker = "python_full_version >= '3.11'" }, + { name = "pexpect", marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit", marker = "python_full_version >= '3.11'" }, + { name = "psutil", marker = "python_full_version >= '3.11' and sys_platform != 'cygwin' and sys_platform != 'emscripten'" }, + { name = "pygments", marker = "python_full_version >= '3.11'" }, + { name = "stack-data", marker = "python_full_version >= '3.11'" }, + { name = "traitlets", marker = "python_full_version >= '3.11'" }, + { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/53/59/165d3b4d75cc34add3122c4417ecb229085140ac573103c223cd01dde96f/ipython-9.15.0.tar.gz", hash = "sha256:da2819ce2aa83135257df830660b1176d986c3d2876db24df01974fa955b2756", size = 4442580, upload-time = "2026-06-26T11:03:35.913Z" } wheels = [ @@ -14742,7 +14742,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments" }, + { name = "pygments", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [ @@ -17548,9 +17548,9 @@ name = "opencensus" version = "0.11.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "google-api-core" }, - { name = "opencensus-context" }, - { name = "six" }, + { name = "google-api-core", marker = "python_full_version < '3.15'" }, + { name = "opencensus-context", marker = "python_full_version < '3.15'" }, + { name = "six", marker = "python_full_version < '3.15'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/15/a7/a46dcffa1b63084f9f17fe3c8cb20724c4c8f91009fd0b2cfdb27d5d2b35/opencensus-0.11.4.tar.gz", hash = "sha256:cbef87d8b8773064ab60e5c2a1ced58bbaa38a6d052c41aec224958ce544eff2", size = 64966, upload-time = "2024-01-03T18:04:07.085Z" } wheels = [ @@ -18108,9 +18108,9 @@ wheels = [ [package.optional-dependencies] sql-other = [ - { name = "adbc-driver-postgresql" }, - { name = "adbc-driver-sqlite" }, - { name = "sqlalchemy" }, + { name = "adbc-driver-postgresql", marker = "python_full_version < '3.13'" }, + { name = "adbc-driver-sqlite", marker = "python_full_version < '3.13'" }, + { name = "sqlalchemy", marker = "python_full_version < '3.13'" }, ] [[package]] @@ -20518,9 +20518,9 @@ name = "python3-saml" version = "1.16.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "isodate" }, - { name = "lxml" }, - { name = "xmlsec" }, + { name = "isodate", marker = "python_full_version < '3.13'" }, + { name = "lxml", marker = "python_full_version < '3.13'" }, + { name = "xmlsec", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5d/98/6e0268c3a9893af3d4c5cf670183e0314cd6b5cb034a612d6a7cc5060df8/python3-saml-1.16.0.tar.gz", hash = "sha256:97c9669aecabc283c6e5fb4eb264f446b6e006f5267d01c9734f9d8bffdac133", size = 83468, upload-time = "2023-10-09T10:37:43.128Z" } wheels = [ @@ -20833,14 +20833,14 @@ name = "ray" version = "2.56.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "filelock" }, - { name = "jsonschema" }, - { name = "msgpack" }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "pyyaml" }, - { name = "requests" }, + { name = "click", marker = "python_full_version < '3.15'" }, + { name = "filelock", marker = "python_full_version < '3.15'" }, + { name = "jsonschema", marker = "python_full_version < '3.15'" }, + { name = "msgpack", marker = "python_full_version < '3.15'" }, + { name = "packaging", marker = "python_full_version < '3.15'" }, + { name = "protobuf", marker = "python_full_version < '3.15'" }, + { name = "pyyaml", marker = "python_full_version < '3.15'" }, + { name = "requests", marker = "python_full_version < '3.15'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/62/f5/89dc8571f35355a7f889cd4709b440abf6db2c5fc8b5427b614d5084e9cb/ray-2.56.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:58b8c037a9f2b7b7866439cb6fe19d452ba3961f2b62d0a247dc3adf2ca45265", size = 66364186, upload-time = "2026-07-17T21:27:58.09Z" }, @@ -20865,20 +20865,20 @@ wheels = [ [package.optional-dependencies] default = [ - { name = "aiohttp" }, - { name = "aiohttp-cors" }, - { name = "colorful" }, - { name = "grpcio" }, - { name = "opencensus" }, - { name = "opentelemetry-exporter-prometheus" }, - { name = "opentelemetry-proto" }, - { name = "opentelemetry-sdk" }, - { name = "prometheus-client" }, - { name = "py-spy" }, - { name = "pydantic" }, - { name = "requests" }, - { name = "smart-open" }, - { name = "virtualenv" }, + { name = "aiohttp", marker = "python_full_version < '3.15'" }, + { name = "aiohttp-cors", marker = "python_full_version < '3.15'" }, + { name = "colorful", marker = "python_full_version < '3.15'" }, + { name = "grpcio", marker = "python_full_version < '3.15'" }, + { name = "opencensus", marker = "python_full_version < '3.15'" }, + { name = "opentelemetry-exporter-prometheus", marker = "python_full_version < '3.15'" }, + { name = "opentelemetry-proto", marker = "python_full_version < '3.15'" }, + { name = "opentelemetry-sdk", marker = "python_full_version < '3.15'" }, + { name = "prometheus-client", marker = "python_full_version < '3.15'" }, + { name = "py-spy", marker = "python_full_version < '3.15'" }, + { name = "pydantic", marker = "python_full_version < '3.15'" }, + { name = "requests", marker = "python_full_version < '3.15'" }, + { name = "smart-open", marker = "python_full_version < '3.15'" }, + { name = "virtualenv", marker = "python_full_version < '3.15'" }, ] [[package]] @@ -21652,10 +21652,10 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "joblib" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } }, - { name = "threadpoolctl" }, + { name = "joblib", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "threadpoolctl", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" } wheels = [ @@ -21708,13 +21708,13 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "joblib" }, - { name = "narwhals" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "joblib", marker = "python_full_version >= '3.11'" }, + { name = "narwhals", marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "threadpoolctl" }, + { name = "threadpoolctl", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" } wheels = [ @@ -21759,7 +21759,7 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -21819,7 +21819,7 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -21900,7 +21900,7 @@ resolution-markers = [ "(python_full_version == '3.12.*' and platform_machine != 'arm64') or (python_full_version == '3.12.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" } wheels = [ @@ -21985,8 +21985,8 @@ name = "secretstorage" version = "3.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cryptography" }, - { name = "jeepney" }, + { name = "cryptography", marker = "(python_full_version >= '3.14' and sys_platform == 'darwin') or (python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (platform_machine != 'arm64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "jeepney", marker = "(python_full_version >= '3.14' and sys_platform == 'darwin') or (python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (platform_machine != 'arm64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } wheels = [ @@ -22185,7 +22185,7 @@ name = "smart-open" version = "8.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "wrapt" }, + { name = "wrapt", marker = "python_full_version < '3.15'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/76/53/9c513747547fd595d5c143259129ea8b9c3ea2f6b7bb9dcea2b1966ded3c/smart_open-8.0.1.tar.gz", hash = "sha256:18b1c4496003c6902be17c15f032b5c319f307c89c6ae9e6b028b508bed8b2cf", size = 61882, upload-time = "2026-07-15T13:56:10.492Z" } wheels = [ @@ -22294,15 +22294,15 @@ name = "snowflake-snowpark-python" version = "1.53.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cloudpickle" }, - { name = "protobuf" }, - { name = "python-dateutil" }, - { name = "pyyaml" }, - { name = "setuptools" }, - { name = "snowflake-connector-python" }, - { name = "typing-extensions" }, - { name = "tzlocal" }, - { name = "wheel" }, + { name = "cloudpickle", marker = "python_full_version < '3.14'" }, + { name = "protobuf", marker = "python_full_version < '3.14'" }, + { name = "python-dateutil", marker = "python_full_version < '3.14'" }, + { name = "pyyaml", marker = "python_full_version < '3.14'" }, + { name = "setuptools", marker = "python_full_version < '3.14'" }, + { name = "snowflake-connector-python", marker = "python_full_version < '3.14'" }, + { name = "typing-extensions", marker = "python_full_version < '3.14'" }, + { name = "tzlocal", marker = "python_full_version < '3.14'" }, + { name = "wheel", marker = "python_full_version < '3.14'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/72/5b/4a0eb18191c98d62b8a0a363a4c2d9ee5f01f0f33f32ecd59ddea7b56e3e/snowflake_snowpark_python-1.53.1.tar.gz", hash = "sha256:5a0bcbf01aac54336c7763b748837281762e57214b6425c84ed412706fb58df7", size = 1784995, upload-time = "2026-07-15T23:06:44.704Z" } wheels = [ @@ -22349,23 +22349,23 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils", version = "0.21.2", source = { registry = "https://pypi.org/simple" } }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, - { name = "tomli" }, + { name = "alabaster", marker = "python_full_version < '3.11'" }, + { name = "babel", marker = "python_full_version < '3.11'" }, + { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, + { name = "docutils", version = "0.21.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "imagesize", marker = "python_full_version < '3.11'" }, + { name = "jinja2", marker = "python_full_version < '3.11'" }, + { name = "packaging", marker = "python_full_version < '3.11'" }, + { name = "pygments", marker = "python_full_version < '3.11'" }, + { name = "requests", marker = "python_full_version < '3.11'" }, + { name = "snowballstemmer", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.11'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.11'" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/be0b61178fe2cdcb67e2a92fc9ebb488e3c51c4f74a36a7824c0adf23425/sphinx-8.1.3.tar.gz", hash = "sha256:43c1911eecb0d3e161ad78611bc905d1ad0e523e4ddc202a58a821773dc4c927", size = 8184611, upload-time = "2024-10-13T20:27:13.93Z" } wheels = [ @@ -22381,23 +22381,23 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" } }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "roman-numerals" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, + { name = "alabaster", marker = "python_full_version == '3.11.*'" }, + { name = "babel", marker = "python_full_version == '3.11.*'" }, + { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, + { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "imagesize", marker = "python_full_version == '3.11.*'" }, + { name = "jinja2", marker = "python_full_version == '3.11.*'" }, + { name = "packaging", marker = "python_full_version == '3.11.*'" }, + { name = "pygments", marker = "python_full_version == '3.11.*'" }, + { name = "requests", marker = "python_full_version == '3.11.*'" }, + { name = "roman-numerals", marker = "python_full_version == '3.11.*'" }, + { name = "snowballstemmer", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version == '3.11.*'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -22419,23 +22419,23 @@ resolution-markers = [ "(python_full_version == '3.12.*' and platform_machine != 'arm64') or (python_full_version == '3.12.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "alabaster" }, - { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" } }, - { name = "imagesize" }, - { name = "jinja2" }, - { name = "packaging" }, - { name = "pygments" }, - { name = "requests" }, - { name = "roman-numerals" }, - { name = "snowballstemmer" }, - { name = "sphinxcontrib-applehelp" }, - { name = "sphinxcontrib-devhelp" }, - { name = "sphinxcontrib-htmlhelp" }, - { name = "sphinxcontrib-jsmath" }, - { name = "sphinxcontrib-qthelp" }, - { name = "sphinxcontrib-serializinghtml" }, + { name = "alabaster", marker = "python_full_version >= '3.12'" }, + { name = "babel", marker = "python_full_version >= '3.12'" }, + { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, + { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "imagesize", marker = "python_full_version >= '3.12'" }, + { name = "jinja2", marker = "python_full_version >= '3.12'" }, + { name = "packaging", marker = "python_full_version >= '3.12'" }, + { name = "pygments", marker = "python_full_version >= '3.12'" }, + { name = "requests", marker = "python_full_version >= '3.12'" }, + { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, + { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, + { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [ @@ -22501,12 +22501,12 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "colorama" }, - { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" } }, - { name = "starlette" }, - { name = "uvicorn" }, - { name = "watchfiles" }, - { name = "websockets" }, + { name = "colorama", marker = "python_full_version < '3.11'" }, + { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "starlette", marker = "python_full_version < '3.11'" }, + { name = "uvicorn", marker = "python_full_version < '3.11'" }, + { name = "watchfiles", marker = "python_full_version < '3.11'" }, + { name = "websockets", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a5/2c/155e1de2c1ba96a72e5dba152c509a8b41e047ee5c2def9e9f0d812f8be7/sphinx_autobuild-2024.10.3.tar.gz", hash = "sha256:248150f8f333e825107b6d4b86113ab28fa51750e5f9ae63b59dc339be951fb1", size = 14023, upload-time = "2024-10-02T23:15:30.172Z" } wheels = [ @@ -22530,13 +22530,13 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "colorama" }, - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "colorama", marker = "python_full_version >= '3.11'" }, + { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "starlette" }, - { name = "uvicorn" }, - { name = "watchfiles" }, - { name = "websockets" }, + { name = "starlette", marker = "python_full_version >= '3.11'" }, + { name = "uvicorn", marker = "python_full_version >= '3.11'" }, + { name = "watchfiles", marker = "python_full_version >= '3.11'" }, + { name = "websockets", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e0/3c/a59a3a453d4133777f7ed2e83c80b7dc817d43c74b74298ca0af869662ad/sphinx_autobuild-2025.8.25.tar.gz", hash = "sha256:9cf5aab32853c8c31af572e4fecdc09c997e2b8be5a07daf2a389e270e85b213", size = 15200, upload-time = "2025-08-25T18:44:55.436Z" } wheels = [ @@ -22566,7 +22566,7 @@ resolution-markers = [ "(python_full_version < '3.11' and platform_machine != 'arm64') or (python_full_version < '3.11' and sys_platform != 'darwin')", ] dependencies = [ - { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" } }, + { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/2b/69/b34e0cb5336f09c6866d53b4a19d76c227cdec1bbc7ac4de63ca7d58c9c7/sphinx_design-0.6.1.tar.gz", hash = "sha256:b44eea3719386d04d765c1a8257caca2b3e6f8421d7b3a5e742c0fd45f84e632", size = 2193689, upload-time = "2024-08-02T13:48:44.277Z" } wheels = [ @@ -22590,7 +22590,7 @@ resolution-markers = [ "(python_full_version == '3.11.*' and platform_machine != 'arm64') or (python_full_version == '3.11.*' and sys_platform != 'darwin')", ] dependencies = [ - { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/13/7b/804f311da4663a4aecc6cf7abd83443f3d4ded970826d0c958edc77d4527/sphinx_design-0.7.0.tar.gz", hash = "sha256:d2a3f5b19c24b916adb52f97c5f00efab4009ca337812001109084a740ec9b7a", size = 2203582, upload-time = "2026-01-19T13:12:53.297Z" } @@ -24581,7 +24581,7 @@ name = "xmlsec" version = "1.3.17" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "lxml" }, + { name = "lxml", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/49/14/538b75379e6ab8f688f14d8663e2ab138d9c778bac4999d155b5f33c71c1/xmlsec-1.3.17.tar.gz", hash = "sha256:f3fac9ae679f66585925cc00c5f6839ae36c1d03157619571dee18acc05b9c01", size = 115637, upload-time = "2025-11-11T16:20:46.019Z" } wheels = [