mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 16:01:19 +00:00
Kacper MudaandClaude
e6a2c9f4e7
A single event about a task instance should be enough to tell which execution it describes: which loop pass and rerun of the loop it ran in, how many instances its mapped expansion has, and why a loop gate stopped. Workers only had the loop iteration, a mapped task inside a loop could not be tied to the pass it ran in, reruns could only be ordered by comparing region ids from two events, and server side consumers such as listeners had to read region rows to place a stored task instance. Co-authored-by: Claude <noreply@anthropic.com>
356 lines
13 KiB
Python
356 lines
13 KiB
Python
# 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 re
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
|
|
from airflow.sdk import DAG, BaseOperator, LoopDecision, TaskInstanceState, task, task_group
|
|
from airflow.sdk._shared.secrets_masker import _secrets_masker
|
|
from airflow.sdk.api.datamodels._generated import LoopContext, LoopDecisionPayload
|
|
from airflow.sdk.definitions._internal.loop import create_loop
|
|
from airflow.sdk.exceptions import AirflowFailException
|
|
from airflow.sdk.execution_time import task_runner
|
|
from airflow.sdk.execution_time.comms import (
|
|
GetXCom,
|
|
GetXComCount,
|
|
GetXComSequenceItem,
|
|
GetXComSequenceSlice,
|
|
SucceedTask,
|
|
TaskState,
|
|
XComCountResponse,
|
|
XComResult,
|
|
XComSequenceIndexResult,
|
|
XComSequenceSliceResult,
|
|
)
|
|
from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
|
|
from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded
|
|
|
|
LOOP_REGION_ID = UUID("0199a8c0-0000-7000-8000-000000000001")
|
|
FORKED_REGION_ID = UUID("0199a8c0-0000-7000-8000-000000000002")
|
|
|
|
|
|
@pytest.fixture
|
|
def loop_ti(create_runtime_ti):
|
|
def make(*, index=1, max_iterations=3, until=None, map_index=-1):
|
|
@task_group
|
|
def body():
|
|
BaseOperator(task_id="terminal")
|
|
|
|
with DAG("loop_runtime", schedule=None) as dag:
|
|
group = create_loop(body, max_iterations=max_iterations, until=until)
|
|
ti = create_runtime_ti(task=dag.get_task(group.gate_task_id), map_index=map_index)
|
|
ti._ti_context_from_server.loop = LoopContext(
|
|
node_id=group.group_id,
|
|
index=index,
|
|
max_iterations=max_iterations,
|
|
terminal_task_id=group.terminal_task_id,
|
|
terminal_is_mapped=False,
|
|
region_id=LOOP_REGION_ID,
|
|
trajectory=1,
|
|
)
|
|
return ti
|
|
|
|
return make
|
|
|
|
|
|
def test_first_iteration_previous_does_not_read_xcom(loop_ti, mock_supervisor_comms):
|
|
ti = loop_ti(map_index=7)
|
|
loop = ti.get_template_context()["loop"]
|
|
|
|
assert loop.index == 1
|
|
assert loop.max_iterations == 3
|
|
assert loop.previous is None
|
|
assert ti.map_index == 7
|
|
mock_supervisor_comms.send.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize("mapped", [False, True])
|
|
def test_body_callable_receives_loop_context_after_unmapping(create_runtime_ti, mapped):
|
|
@task
|
|
def terminal(value, *, loop, ti):
|
|
return value, loop.index, ti.map_index
|
|
|
|
@task_group
|
|
def body():
|
|
if mapped:
|
|
terminal.expand(value=[5])
|
|
else:
|
|
terminal(5)
|
|
|
|
with DAG("loop_body_runtime", schedule=None) as dag:
|
|
group = create_loop(body, max_iterations=4)
|
|
operator = dag.get_task(group.terminal_task_id)
|
|
if mapped:
|
|
operator = operator.unmap({"op_kwargs": {"value": 5}})
|
|
ti = create_runtime_ti(task=operator, map_index=0 if mapped else -1)
|
|
ti._ti_context_from_server.loop = LoopContext(
|
|
node_id=group.group_id,
|
|
index=2,
|
|
max_iterations=4,
|
|
terminal_task_id=group.terminal_task_id,
|
|
terminal_is_mapped=mapped,
|
|
region_id=LOOP_REGION_ID,
|
|
trajectory=1,
|
|
)
|
|
|
|
assert operator.execute(ti.get_template_context()) == (5, 2, 0 if mapped else -1)
|
|
|
|
|
|
def test_task_instance_exposes_the_loop_iteration_and_its_trajectory(loop_ti):
|
|
ti = loop_ti(index=2)
|
|
ti._ti_context_from_server.loop = ti._ti_context_from_server.loop.model_copy(
|
|
update={"trajectory": 2, "forked_from_region_id": FORKED_REGION_ID, "resumes_from_index": 2}
|
|
)
|
|
|
|
loop = ti.loop
|
|
|
|
assert loop is not None
|
|
assert (loop.index, loop.max_iterations) == (2, 3)
|
|
assert (loop.region_id, loop.trajectory) == (LOOP_REGION_ID, 2)
|
|
assert (loop.forked_from_region_id, loop.resumes_from_index) == (FORKED_REGION_ID, 2)
|
|
|
|
|
|
def test_task_instance_outside_a_loop_has_no_loop(create_runtime_ti):
|
|
ti = create_runtime_ti(task=BaseOperator(task_id="plain"))
|
|
|
|
assert ti.loop is None
|
|
assert ti.loop_decision is None
|
|
|
|
|
|
@pytest.mark.parametrize("expanded_ti_count", [None, 4])
|
|
def test_task_instance_exposes_the_size_of_its_expansion(create_runtime_ti, expanded_ti_count):
|
|
ti = create_runtime_ti(task=BaseOperator(task_id="mapped"))
|
|
ti._ti_context_from_server.expanded_ti_count = expanded_ti_count
|
|
|
|
assert ti.expanded_ti_count == expanded_ti_count
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("index", "until", "expected"),
|
|
[
|
|
(1, lambda: LoopDecision(stop=False, reason="not yet"), LoopDecision(stop=False, reason="not yet")),
|
|
(3, lambda: True, LoopDecision(stop=True)),
|
|
(3, lambda: False, LoopDecision(stop=False)),
|
|
],
|
|
ids=["continue", "until_met", "cap_reached"],
|
|
)
|
|
def test_gate_exposes_its_decision(loop_ti, mock_supervisor_comms, index, until, expected):
|
|
ti = loop_ti(index=index, until=until)
|
|
|
|
task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
|
|
|
|
assert ti.loop_decision == expected
|
|
|
|
|
|
@pytest.mark.parametrize("value", [0, False, [], ""])
|
|
def test_previous_and_current_result_preserve_falsey_values(loop_ti, mock_supervisor_comms, value):
|
|
ti = loop_ti(index=2)
|
|
mock_supervisor_comms.send.return_value = XComResult(key="return_value", value=value)
|
|
loop = ti.get_template_context()["loop"]
|
|
|
|
assert loop.previous == value
|
|
previous = mock_supervisor_comms.send.call_args.args[0]
|
|
assert isinstance(previous, GetXCom)
|
|
assert previous.task_id == "body.terminal"
|
|
assert previous.previous_iteration is True
|
|
assert loop.result == value
|
|
assert mock_supervisor_comms.send.call_args.args[0].previous_iteration is False
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("index", "until", "expected"),
|
|
[
|
|
(1, None, LoopDecisionPayload(stop=False)),
|
|
# max_iterations - 1: a zero-based cap check stops here instead, so this pins the base.
|
|
(2, None, LoopDecisionPayload(stop=False)),
|
|
(3, None, LoopDecisionPayload(stop=True)),
|
|
(1, lambda: False, LoopDecisionPayload(stop=False)),
|
|
(1, lambda: True, LoopDecisionPayload(stop=True)),
|
|
(3, lambda: True, LoopDecisionPayload(stop=True)),
|
|
(1, lambda: LoopDecision(stop=True), LoopDecisionPayload(stop=True)),
|
|
(
|
|
1,
|
|
lambda: LoopDecision(stop=False, reason="score 0.7, needs > 0.9"),
|
|
LoopDecisionPayload(stop=False, reason="score 0.7, needs > 0.9"),
|
|
),
|
|
(
|
|
1,
|
|
lambda: LoopDecision(stop=True, reason="é" * 256),
|
|
LoopDecisionPayload(stop=True, reason="é" * 256),
|
|
),
|
|
(1, lambda: LoopDecision(stop=True, reason=""), LoopDecisionPayload(stop=True)),
|
|
(
|
|
1,
|
|
lambda: LoopDecision(stop=True, reason="token hunter2 expired"),
|
|
LoopDecisionPayload(stop=True, reason="token *** expired"),
|
|
),
|
|
],
|
|
ids=[
|
|
"fixed_continue",
|
|
"fixed_continue_below_cap",
|
|
"fixed_stop_at_cap",
|
|
"bool_continue",
|
|
"bool_stop",
|
|
"bool_stop_at_cap",
|
|
"decision_without_reason",
|
|
"decision_with_reason",
|
|
"reason_at_byte_limit",
|
|
"empty_reason_is_no_reason",
|
|
"secret_in_reason_is_masked",
|
|
],
|
|
)
|
|
@pytest.mark.enable_redact
|
|
def test_gate_sends_its_decision_with_success(loop_ti, mock_supervisor_comms, index, until, expected):
|
|
_secrets_masker().add_mask("hunter2", None)
|
|
ti = loop_ti(index=index, until=until)
|
|
|
|
state, message, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
|
|
|
|
assert (state, error) == (TaskInstanceState.SUCCESS, None)
|
|
assert isinstance(message, SucceedTask)
|
|
assert message.loop_decision == expected
|
|
|
|
|
|
def _raise_value_error():
|
|
raise ValueError("condition failed")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("until", "error", "message"),
|
|
[
|
|
(_raise_value_error, ValueError, "condition failed"),
|
|
(lambda: None, AirflowFailException, "got NoneType"),
|
|
(lambda: 0, AirflowFailException, "got int"),
|
|
(lambda: 1, AirflowFailException, "got int"),
|
|
(lambda: [], AirflowFailException, "got list"),
|
|
(lambda: "yes", AirflowFailException, "got str"),
|
|
(lambda: LoopDecision(stop=1), AirflowFailException, "LoopDecision stop must be a bool, got int"), # type: ignore[arg-type]
|
|
(
|
|
lambda: LoopDecision(stop=None, reason="r"), # type: ignore[arg-type]
|
|
AirflowFailException,
|
|
"LoopDecision stop must be a bool, got NoneType",
|
|
),
|
|
(
|
|
lambda: LoopDecision(stop=True, reason=0.2), # type: ignore[arg-type]
|
|
AirflowFailException,
|
|
"LoopDecision reason must be a str or None, got float",
|
|
),
|
|
(
|
|
lambda: LoopDecision(stop=True, reason=b"r"), # type: ignore[arg-type]
|
|
AirflowFailException,
|
|
"LoopDecision reason must be a str or None, got bytes",
|
|
),
|
|
(lambda: LoopDecision(stop=True, reason="x" * 513), AirflowFailException, "512 bytes"),
|
|
(lambda: LoopDecision(stop=True, reason="é" * 257), AirflowFailException, "512 bytes"),
|
|
(lambda: LoopDecision(stop=True, reason="\ud800"), AirflowFailException, "valid UTF-8"),
|
|
],
|
|
ids=[
|
|
"until_raises",
|
|
"none",
|
|
"zero",
|
|
"one",
|
|
"list",
|
|
"str",
|
|
"int_stop",
|
|
"none_stop",
|
|
"float_reason",
|
|
"bytes_reason",
|
|
"ascii_reason_over_limit",
|
|
"utf8_reason_over_limit",
|
|
"unpaired_surrogate_reason",
|
|
],
|
|
)
|
|
def test_gate_fails_without_a_decision(loop_ti, mock_supervisor_comms, until, error, message):
|
|
ti = loop_ti(index=1, until=until)
|
|
|
|
state, sent, raised = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
|
|
|
|
assert state == TaskInstanceState.FAILED
|
|
assert isinstance(raised, error)
|
|
assert re.search(message, str(raised))
|
|
assert isinstance(sent, TaskState)
|
|
assert sent.loop_decision is None
|
|
|
|
|
|
def test_unmet_condition_before_the_cap_continues_instead_of_failing(loop_ti, mock_supervisor_comms):
|
|
"""The pass before the cap has an iteration left, so a false condition just continues."""
|
|
ti = loop_ti(index=2, until=lambda: False)
|
|
|
|
state, message, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
|
|
|
|
assert (state, error) == (TaskInstanceState.SUCCESS, None)
|
|
assert isinstance(message, SucceedTask)
|
|
assert message.loop_decision == LoopDecisionPayload(stop=False)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("result", "message", "reason"),
|
|
[
|
|
(False, r"max_iterations=3$", None),
|
|
(
|
|
LoopDecision(stop=False, reason="error 0.2, needs < 1e-06"),
|
|
r"max_iterations=3: error 0.2, needs",
|
|
"error 0.2, needs < 1e-06",
|
|
),
|
|
],
|
|
ids=["bool", "decision_with_reason"],
|
|
)
|
|
def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(
|
|
loop_ti, mock_supervisor_comms, result, message, reason
|
|
):
|
|
ti = loop_ti(index=3, until=lambda: result)
|
|
|
|
state, sent, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
|
|
|
|
assert state == TaskInstanceState.FAILED
|
|
assert isinstance(error, LoopMaxIterationsExceeded)
|
|
assert re.search(message, str(error))
|
|
assert isinstance(sent, TaskState)
|
|
assert sent.loop_decision == LoopDecisionPayload(stop=False, reason=reason)
|
|
|
|
|
|
@pytest.mark.parametrize("previous_iteration", [False, True])
|
|
@pytest.mark.parametrize("values", [[], [0], [0, False, ""]])
|
|
def test_mapped_terminal_keeps_iteration_across_lazy_reads(
|
|
loop_ti, mock_supervisor_comms, previous_iteration, values
|
|
):
|
|
ti = loop_ti(index=2)
|
|
ti._ti_context_from_server.loop.terminal_is_mapped = True
|
|
|
|
def respond(message):
|
|
assert message.previous_iteration is previous_iteration
|
|
if isinstance(message, GetXComCount):
|
|
return XComCountResponse(len=len(values))
|
|
if isinstance(message, GetXComSequenceItem):
|
|
return XComSequenceIndexResult(root=values[message.offset])
|
|
assert isinstance(message, GetXComSequenceSlice)
|
|
return XComSequenceSliceResult(root=values[slice(message.start, message.stop, message.step)])
|
|
|
|
mock_supervisor_comms.send.side_effect = respond
|
|
context = ti.get_template_context()["loop"]
|
|
result = context.previous if previous_iteration else context.result
|
|
|
|
assert isinstance(result, LazyXComSequence)
|
|
assert len(result) == len(values)
|
|
assert result[:] == values
|
|
assert result[::-1] == values[::-1]
|
|
if values:
|
|
assert result[-1] == values[-1]
|