diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index 6e2157e3c7b2a..9ddd1f2b7aed9 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -66,7 +66,6 @@ class TerminalStateNonSuccess(str, Enum): FAILED = TerminalTIState.FAILED SKIPPED = TerminalTIState.SKIPPED REMOVED = TerminalTIState.REMOVED - FAIL_WITHOUT_RETRY = TerminalTIState.FAIL_WITHOUT_RETRY class TITerminalStatePayload(StrictBaseModel): @@ -157,6 +156,23 @@ class TIRescheduleStatePayload(StrictBaseModel): end_date: UtcDateTime +class TIRetryStatePayload(StrictBaseModel): + """Schema for updating TaskInstance to up_for_retry.""" + + state: Annotated[ + Literal[IntermediateTIState.UP_FOR_RETRY], + # Specify a default in the schema, but not in code, so Pydantic marks it as required. + WithJsonSchema( + { + "type": "string", + "enum": [IntermediateTIState.UP_FOR_RETRY], + "default": IntermediateTIState.UP_FOR_RETRY, + } + ), + ] + end_date: UtcDateTime + + class TISkippedDownstreamTasksStatePayload(StrictBaseModel): """Schema for updating downstream tasks to a skipped state.""" @@ -185,6 +201,8 @@ def ti_state_discriminator(v: dict[str, str] | StrictBaseModel) -> str: return "deferred" elif state == TIState.UP_FOR_RESCHEDULE: return "up_for_reschedule" + elif state == TIState.UP_FOR_RETRY: + return "up_for_retry" return "_other_" @@ -197,6 +215,7 @@ def ti_state_discriminator(v: dict[str, str] | StrictBaseModel) -> str: Annotated[TITargetStatePayload, Tag("_other_")], Annotated[TIDeferredStatePayload, Tag("deferred")], Annotated[TIRescheduleStatePayload, Tag("up_for_reschedule")], + Annotated[TIRetryStatePayload, Tag("up_for_retry")], ], Discriminator(ti_state_discriminator), ] @@ -276,6 +295,9 @@ class TIRunContext(BaseModel): xcom_keys_to_clear: Annotated[list[str], Field(default_factory=list)] """List of Xcom keys that need to be cleared and purged on by the worker.""" + should_retry: bool + """If the ti encounters an error, whether it should enter retry or failed state.""" + class PrevSuccessfulDagRunResponse(BaseModel): """Schema for response with previous successful DagRun information for Task Template Context.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 6b031e36394ac..a4f20b18cdf12 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -36,6 +36,7 @@ TIEnterRunningPayload, TIHeartbeatInfo, TIRescheduleStatePayload, + TIRetryStatePayload, TIRunContext, TIRuntimeCheckPayload, TISkippedDownstreamTasksStatePayload, @@ -50,7 +51,7 @@ from airflow.models.trigger import Trigger from airflow.models.xcom import XComModel from airflow.utils import timezone -from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState +from airflow.utils.state import DagRunState, TaskInstanceState router = VersionedAPIRouter( dependencies=[ @@ -136,7 +137,7 @@ def ti_run( ti_run_payload.pid, ): log.info("Duplicate start request received from %s ", ti_run_payload.hostname) - elif previous_state != TaskInstanceState.QUEUED: + elif previous_state not in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING): log.warning( "Can not start Task Instance ('%s') in invalid state: %s", ti_id_str, @@ -226,6 +227,7 @@ def ti_run( variables=[], connections=[], xcom_keys_to_clear=xcom_keys, + should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) # Only set if they are non-null @@ -289,23 +291,18 @@ def ti_update_state( query = update(TI).where(TI.id == ti_id_str).values(data) if isinstance(ti_patch_payload, TITerminalStatePayload): + updated_state = ti_patch_payload.state query = TI.duration_expression_update(ti_patch_payload.end_date, query, session.bind) + query = query.values(state=updated_state) + elif isinstance(ti_patch_payload, TIRetryStatePayload): + from airflow.models.taskinstance import uuid7 + from airflow.models.taskinstancehistory import TaskInstanceHistory + + ti = session.get(TI, ti_id_str) + TaskInstanceHistory.record_ti(ti, session=session) + ti.try_id = uuid7() updated_state = ti_patch_payload.state - # if we get failed, we should attempt to retry, as it is a more - # normal state. Tasks with retries are more frequent than without retries. - if ti_patch_payload.state == TerminalTIState.FAIL_WITHOUT_RETRY: - updated_state = TaskInstanceState.FAILED - elif ti_patch_payload.state == TaskInstanceState.FAILED: - if _is_eligible_to_retry(previous_state, try_number, max_tries): - from airflow.models.taskinstance import uuid7 - from airflow.models.taskinstancehistory import TaskInstanceHistory - - ti = session.get(TI, ti_id_str) - TaskInstanceHistory.record_ti(ti, session=session) - ti.try_id = uuid7() - updated_state = TaskInstanceState.UP_FOR_RETRY - else: - updated_state = TaskInstanceState.FAILED + query = TI.duration_expression_update(ti_patch_payload.end_date, query, session.bind) query = query.values(state=updated_state) elif isinstance(ti_patch_payload, TISuccessStatePayload): query = TI.duration_expression_update(ti_patch_payload.end_date, query, session.bind) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 259184a2febd8..888dac73d7388 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -888,6 +888,7 @@ def _get_template_context( ti_context_from_server = TIRunContext( dag_run=DagRunSDK.model_validate(dag_run, from_attributes=True), max_tries=task_instance.max_tries, + should_retry=task_instance.is_eligible_to_retry(), ) runtime_ti = task_instance.to_runtime_ti(context_from_server=ti_context_from_server) @@ -3196,7 +3197,7 @@ def handle_failure( fail_fast=fail_fast, ) - def is_eligible_to_retry(self): + def is_eligible_to_retry(self) -> bool: """Is task instance is eligible for retry.""" return _is_eligible_to_retry(task_instance=self) diff --git a/airflow-core/src/airflow/utils/state.py b/airflow-core/src/airflow/utils/state.py index dca2c8fc93f31..e4e2e9db8a587 100644 --- a/airflow-core/src/airflow/utils/state.py +++ b/airflow-core/src/airflow/utils/state.py @@ -39,7 +39,6 @@ class TerminalTIState(str, Enum): FAILED = "failed" SKIPPED = "skipped" # A user can raise a AirflowSkipException from a task & it will be marked as skipped REMOVED = "removed" - FAIL_WITHOUT_RETRY = "fail_without_retry" def __str__(self) -> str: return self.value diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/routes/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/routes/test_task_instances.py index b66b1a898b52d..578a019e2527a 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/routes/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/routes/test_task_instances.py @@ -116,7 +116,22 @@ def setup_method(self): def teardown_method(self): clear_db_runs() - def test_ti_run_state_to_running(self, client, session, create_task_instance, time_machine): + @pytest.mark.parametrize( + "max_tries, should_retry", + [ + pytest.param(0, False, id="max_retries=0"), + pytest.param(3, True, id="should_retry"), + ], + ) + def test_ti_run_state_to_running( + self, + client, + session, + create_task_instance, + time_machine, + max_tries, + should_retry, + ): """ Test that the Task Instance state is updated to running when the Task Instance is in a state where it can be marked as running. @@ -131,6 +146,7 @@ def test_ti_run_state_to_running(self, client, session, create_task_instance, ti session=session, start_date=instant, ) + ti.max_tries = max_tries session.commit() response = client.patch( @@ -160,7 +176,8 @@ def test_ti_run_state_to_running(self, client, session, create_task_instance, ti "conf": {}, }, "task_reschedule_count": 0, - "max_tries": 0, + "max_tries": max_tries, + "should_retry": should_retry, "variables": [], "connections": [], "xcom_keys_to_clear": [], @@ -210,7 +227,7 @@ def test_next_kwargs_still_encoded(self, client, session, create_task_instance, time_machine.move_to(instant, tick=False) ti = create_task_instance( - task_id="test_ti_run_state_to_running", + task_id="test_next_kwargs_still_encoded", state=State.QUEUED, session=session, start_date=instant, @@ -238,6 +255,7 @@ def test_next_kwargs_still_encoded(self, client, session, create_task_instance, "dag_run": mock.ANY, "task_reschedule_count": 0, "max_tries": 0, + "should_retry": False, "variables": [], "connections": [], "xcom_keys_to_clear": [], @@ -248,7 +266,10 @@ def test_next_kwargs_still_encoded(self, client, session, create_task_instance, }, } - @pytest.mark.parametrize("initial_ti_state", [s for s in TaskInstanceState if s != State.QUEUED]) + @pytest.mark.parametrize( + "initial_ti_state", + [s for s in TaskInstanceState if s not in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING)], + ) def test_ti_run_state_conflict_if_not_queued( self, client, session, create_task_instance, initial_ti_state ): @@ -691,67 +712,17 @@ def test_ti_update_state_to_reschedule(self, client, session, create_task_instan assert trs[0].task_instance.map_index == -1 assert trs[0].duration == 129600 - @pytest.mark.parametrize( - ("retries", "expected_state"), - [ - (0, State.FAILED), - (None, State.FAILED), - (3, State.UP_FOR_RETRY), - ], - ) - def test_ti_update_state_to_failed_with_retries( - self, client, session, create_task_instance, retries, expected_state - ): + def test_ti_update_state_handle_retry(self, client, session, create_task_instance): ti = create_task_instance( task_id="test_ti_update_state_to_retry", state=State.RUNNING, ) - - if retries is not None: - ti.max_tries = retries - session.commit() - - response = client.patch( - f"/execution/task-instances/{ti.id}/state", - json={ - "state": TerminalTIState.FAILED, - "end_date": DEFAULT_END_DATE.isoformat(), - }, - ) - - assert response.status_code == 204 - assert response.text == "" - - session.expire_all() - - ti = session.get(TaskInstance, ti.id) - assert ti.state == expected_state - assert ti.next_method is None - assert ti.next_kwargs is None - - tih = session.query(TaskInstanceHistory).where( - TaskInstanceHistory.task_id == ti.task_id, TaskInstanceHistory.task_instance_id == ti.id - ) - tih_count = tih.count() - assert tih_count == (1 if retries else 0) - if retries: - tih = tih.one() - assert tih.try_id - assert tih.try_id != ti.try_id - - def test_ti_update_state_when_ti_is_restarting(self, client, session, create_task_instance): - ti = create_task_instance( - task_id="test_ti_update_state_when_ti_is_restarting", - state=State.RUNNING, - ) - # update state to restarting - ti.state = State.RESTARTING session.commit() response = client.patch( f"/execution/task-instances/{ti.id}/state", json={ - "state": TerminalTIState.FAILED, + "state": State.UP_FOR_RETRY, "end_date": DEFAULT_END_DATE.isoformat(), }, ) @@ -762,43 +733,19 @@ def test_ti_update_state_when_ti_is_restarting(self, client, session, create_tas session.expire_all() ti = session.get(TaskInstance, ti.id) - # restarting is always retried assert ti.state == State.UP_FOR_RETRY assert ti.next_method is None assert ti.next_kwargs is None - def test_ti_update_state_when_ti_has_higher_tries_than_retries( - self, client, session, create_task_instance - ): - ti = create_task_instance( - task_id="test_ti_update_state_when_ti_has_higher_tries_than_retries", - state=State.RUNNING, - ) - # two maximum tries defined, but third try going on - ti.max_tries = 2 - ti.try_number = 3 - session.commit() - - response = client.patch( - f"/execution/task-instances/{ti.id}/state", - json={ - "state": TerminalTIState.FAILED, - "end_date": DEFAULT_END_DATE.isoformat(), - }, + tih = ( + session.query(TaskInstanceHistory) + .where(TaskInstanceHistory.task_id == ti.task_id, TaskInstanceHistory.task_instance_id == ti.id) + .one() ) + assert tih.try_id + assert tih.try_id != ti.try_id - assert response.status_code == 204 - assert response.text == "" - - session.expire_all() - - ti = session.get(TaskInstance, ti.id) - # all retries exhausted, marking as failed - assert ti.state == State.FAILED - assert ti.next_method is None - assert ti.next_kwargs is None - - def test_ti_update_state_to_failed_without_retry_table_check(self, client, session, create_task_instance): + def test_ti_update_state_to_failed_table_check(self, client, session, create_task_instance): # we just want to fail in this test, no need to retry ti = create_task_instance( task_id="test_ti_update_state_to_failed_table_check", @@ -810,7 +757,7 @@ def test_ti_update_state_to_failed_without_retry_table_check(self, client, sessi response = client.patch( f"/execution/task-instances/{ti.id}/state", json={ - "state": TerminalTIState.FAIL_WITHOUT_RETRY, + "state": TerminalTIState.FAILED, "end_date": DEFAULT_END_DATE.isoformat(), }, ) diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py index 04739f22633c7..106922d027e98 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py @@ -473,6 +473,7 @@ def _make_context( ), task_reschedule_count=task_reschedule_count, max_tries=0, + should_retry=False, ) return _make_context diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index 764711760d367..436907571b579 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -904,6 +904,7 @@ def mock_task_id(dag_id, task_id, try_number, logical_date, map_index): ), task_reschedule_count=0, max_tries=1, + should_retry=False, ), start_date=dt.datetime(2023, 1, 1, 13, 1, 1), ) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 9108cbc08f8fb..1822748ad72bc 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -484,6 +484,7 @@ def noop_handler(request: httpx.Request) -> httpx.Response: "run_after": "2021-01-01T00:00:00Z", }, "max_tries": 0, + "should_retry": False, }, ) return httpx.Response(200, json={"text": "Hello, world!"}) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 9b99a3d57c8ad..0ed6fc985becd 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -195,6 +195,18 @@ class TIRescheduleStatePayload(BaseModel): end_date: Annotated[datetime, Field(title="End Date")] +class TIRetryStatePayload(BaseModel): + """ + Schema for updating TaskInstance to up_for_retry. + """ + + model_config = ConfigDict( + extra="forbid", + ) + state: Annotated[Literal["up_for_retry"] | None, Field(title="State")] = "up_for_retry" + end_date: Annotated[datetime, Field(title="End Date")] + + class TIRuntimeCheckPayload(BaseModel): """ Payload for performing Runtime checks on the TaskInstance model as requested by the SDK. @@ -247,7 +259,6 @@ class TerminalStateNonSuccess(str, Enum): FAILED = "failed" SKIPPED = "skipped" REMOVED = "removed" - FAIL_WITHOUT_RETRY = "fail_without_retry" class TriggerDAGRunPayload(BaseModel): @@ -333,7 +344,6 @@ class TerminalTIState(str, Enum): FAILED = "failed" SKIPPED = "skipped" REMOVED = "removed" - FAIL_WITHOUT_RETRY = "fail_without_retry" class AssetEventResponse(BaseModel): @@ -399,6 +409,7 @@ class TIRunContext(BaseModel): next_method: Annotated[str | None, Field(title="Next Method")] = None next_kwargs: Annotated[dict[str, Any] | str | None, Field(title="Next Kwargs")] = None xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None + should_retry: Annotated[bool, Field(title="Should Retry")] class TITerminalStatePayload(BaseModel): diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index ae05956235fb4..de20055b9892c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -61,6 +61,7 @@ TerminalTIState, TIDeferredStatePayload, TIRescheduleStatePayload, + TIRetryStatePayload, TIRunContext, TIRuntimeCheckPayload, TISkippedDownstreamTasksStatePayload, @@ -257,7 +258,6 @@ class TaskState(BaseModel): TerminalTIState.FAILED, TerminalTIState.SKIPPED, TerminalTIState.REMOVED, - TerminalTIState.FAIL_WITHOUT_RETRY, ] end_date: datetime | None = None type: Literal["TaskState"] = "TaskState" @@ -288,6 +288,12 @@ def _serde_kwarg_fields(self, val: str | dict[str, Any] | None, _info): return BaseSerialization.serialize(val or {}) +class RetryTask(TIRetryStatePayload): + """Update a task instance state to up_for_retry.""" + + type: Literal["RetryTask"] = "RetryTask" + + class RescheduleTask(TIRescheduleStatePayload): """Update a task instance state to reschedule/up_for_reschedule.""" @@ -445,6 +451,7 @@ class GetPrevSuccessfulDagRun(BaseModel): GetXComCount, PutVariable, RescheduleTask, + RetryTask, SkipDownstreamTasks, SetRenderedFields, SetXCom, 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 fcd401ab5db20..4711b519bf396 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -60,6 +60,7 @@ GetDagRunState, OKResponse, RescheduleTask, + RetryTask, RuntimeCheckOnTask, SetRenderedFields, SkipDownstreamTasks, @@ -655,19 +656,15 @@ def run( # TODO: Handle fail_stop here: https://github.com/apache/airflow/issues/44951 # TODO: Handle addition to Log table: https://github.com/apache/airflow/issues/44952 msg = TaskState( - state=TerminalTIState.FAIL_WITHOUT_RETRY, + state=TerminalTIState.FAILED, end_date=datetime.now(tz=timezone.utc), ) - state = TerminalTIState.FAIL_WITHOUT_RETRY + state = TerminalTIState.FAILED error = e except (AirflowTaskTimeout, AirflowException) as e: # We should allow retries if the task has defined it. log.exception("Task failed with exception") - msg = TaskState( - state=TerminalTIState.FAILED, - end_date=datetime.now(tz=timezone.utc), - ) - state = TerminalTIState.FAILED + msg, state = _handle_current_task_failed(ti) error = e except AirflowTaskTerminated as e: # External state updates are already handled with `ti_heartbeat` and will be @@ -675,24 +672,19 @@ def run( # If these are thrown, we should mark the TI state as failed. log.exception("Task failed with exception") msg = TaskState( - state=TerminalTIState.FAIL_WITHOUT_RETRY, + state=TerminalTIState.FAILED, end_date=datetime.now(tz=timezone.utc), ) - state = TerminalTIState.FAIL_WITHOUT_RETRY + state = TerminalTIState.FAILED error = e except SystemExit as e: # SystemExit needs to be retried if they are eligible. log.exception("Task failed with exception") - msg = TaskState( - state=TerminalTIState.FAILED, - end_date=datetime.now(tz=timezone.utc), - ) - state = TerminalTIState.FAILED + msg, state = _handle_current_task_failed(ti) error = e except BaseException as e: log.exception("Task failed with exception") - msg = TaskState(state=TerminalTIState.FAILED, end_date=datetime.now(tz=timezone.utc)) - state = TerminalTIState.FAILED + msg, state = _handle_current_task_failed(ti) error = e finally: if msg: @@ -701,7 +693,10 @@ def run( return state, msg, error -def _handle_current_task_success(context: Context, ti: RuntimeTaskInstance): +def _handle_current_task_success( + context: Context, + ti: RuntimeTaskInstance, +) -> tuple[SucceedTask, TerminalTIState]: task_outlets = list(_build_asset_profiles(ti.task.outlets)) outlet_events = list(_serialize_outlet_events(context["outlet_events"])) msg = SucceedTask( @@ -712,9 +707,18 @@ def _handle_current_task_success(context: Context, ti: RuntimeTaskInstance): return msg, TerminalTIState.SUCCESS +def _handle_current_task_failed( + ti: RuntimeTaskInstance, +) -> tuple[RetryTask, IntermediateTIState] | tuple[TaskState, TerminalTIState]: + end_date = datetime.now(tz=timezone.utc) + if ti._ti_context_from_server and ti._ti_context_from_server.should_retry: + return RetryTask(end_date=end_date), IntermediateTIState.UP_FOR_RETRY + return TaskState(state=TerminalTIState.FAILED, end_date=end_date), TerminalTIState.FAILED + + def _handle_trigger_dag_run( drte: DagRunTriggerException, context: Context, ti: RuntimeTaskInstance, log: Logger -): +) -> tuple[ToSupervisor, TerminalTIState]: """Handle exception from TriggerDagRunOperator.""" log.info("Triggering Dag Run.", trigger_dag_id=drte.trigger_dag_id) SUPERVISOR_COMMS.send_request( @@ -784,9 +788,7 @@ def _handle_trigger_dag_run( state=comms_msg.state, ) - msg, state = _handle_current_task_success(context, ti) - - return msg, state + return _handle_current_task_success(context, ti) def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger): @@ -881,7 +883,10 @@ def _push_xcom_if_needed(result: Any, ti: RuntimeTaskInstance, log: Logger): def finalize( - ti: RuntimeTaskInstance, state: TerminalTIState, log: Logger, error: BaseException | None = None + ti: RuntimeTaskInstance, + state: IntermediateTIState | TerminalTIState, + log: Logger, + error: BaseException | None = None, ): # Pushing xcom for each operator extra links defined on the operator only. for oe in ti.task.operator_extra_links: @@ -890,12 +895,17 @@ def finalize( _xcom_push(ti, key=xcom_key, value=link) log.debug("Running finalizers", ti=ti) - if state in [TerminalTIState.SUCCESS]: + if state == TerminalTIState.SUCCESS: get_listener_manager().hook.on_task_instance_success( previous_state=TaskInstanceState.RUNNING, task_instance=ti ) # TODO: Run task success callbacks here - if state in [TerminalTIState.FAILED, TerminalTIState.FAIL_WITHOUT_RETRY]: + elif state == IntermediateTIState.UP_FOR_RETRY: + get_listener_manager().hook.on_task_instance_failed( + previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error + ) + # TODO: Run task retry callbacks here + elif state == TerminalTIState.FAILED: get_listener_manager().hook.on_task_instance_failed( previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error ) diff --git a/task-sdk/tests/conftest.py b/task-sdk/tests/conftest.py index 89d9a94853e89..ee176c09c4100 100644 --- a/task-sdk/tests/conftest.py +++ b/task-sdk/tests/conftest.py @@ -220,6 +220,7 @@ def _make_context( ), task_reschedule_count=task_reschedule_count, max_tries=0, + should_retry=False, ) return _make_context diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index ff86b9c6184f9..b7f48d0d3a5ce 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -56,6 +56,7 @@ class TestClient: "run_after": "2021-01-01T00:00:00Z", }, "max_tries": 0, + "should_retry": False, }, ), ], diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 848b7be44a663..5d06fe7992b21 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -728,23 +728,16 @@ def test_overtime_slow_listener_instance( if expected_timeout: assert any( - [ - event["event"] == "Workload success overtime reached; terminating process" - for event in captured_logs - ] + event["event"] == "Workload success overtime reached; terminating process" + for event in captured_logs ) assert any( - [ - event["event"] == "Process exited" and event["signal"] == "SIGTERM" - for event in captured_logs - ] + event["event"] == "Process exited" and event["signal"] == "SIGTERM" for event in captured_logs ) else: assert all( - [ - event["event"] != "Workload success overtime reached; terminating process" - for event in captured_logs - ] + event["event"] != "Workload success overtime reached; terminating process" + for event in captured_logs ) 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 10eb2bdf719bf..907e57b3ccb8e 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 @@ -140,7 +140,7 @@ def test_recv_StartupDetails(self): b'"dag_id": "c"}, "ti_context":{"dag_run":{"dag_id":"c","run_id":"b","logical_date":"2024-12-01T01:00:00Z",' b'"data_interval_start":"2024-12-01T00:00:00Z","data_interval_end":"2024-12-01T01:00:00Z",' b'"start_date":"2024-12-01T01:00:00Z","run_after":"2024-12-01T01:00:00Z","end_date":null,"run_type":"manual","conf":null},' - b'"max_tries":0,"variables":null,"connections":null},"file": "/dev/null",' + b'"max_tries":0,"should_retry":false,"variables":null,"connections":null},"file": "/dev/null",' b'"start_date":"2024-12-01T01:00:00Z", "dag_rel_path": "/dev/null", "bundle_info": {"name": ' b'"any-name", "version": "any-version"}, "requests_fd": ' + str(w2.fileno()).encode("ascii") @@ -670,7 +670,7 @@ def execute(self, context): run(ti, log=mock.MagicMock()) mock_supervisor_comms.send_request.assert_called_once_with( - msg=TaskState(state=TerminalTIState.FAIL_WITHOUT_RETRY, end_date=instant), log=mock.ANY + msg=TaskState(state=TerminalTIState.FAILED, end_date=instant), log=mock.ANY )