From 8291dede1db1c41c2e38d1b088cad7a0174563f2 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 17 Jan 2025 16:20:07 +0100 Subject: [PATCH 1/7] POC Make Callbacks bundle aware --- airflow/callbacks/callback_requests.py | 2 ++ airflow/dag_processing/manager.py | 26 ++++++++++---------------- airflow/jobs/scheduler_job_runner.py | 6 ++++++ airflow/models/dagrun.py | 6 ++++++ tests/dag_processing/test_manager.py | 26 +++++++++++++++++--------- 5 files changed, 41 insertions(+), 25 deletions(-) diff --git a/airflow/callbacks/callback_requests.py b/airflow/callbacks/callback_requests.py index c70835dedb544..38f291657eb63 100644 --- a/airflow/callbacks/callback_requests.py +++ b/airflow/callbacks/callback_requests.py @@ -36,6 +36,8 @@ class BaseCallbackRequest(BaseModel): full_filepath: str """File Path to use to run the callback""" + bundle_name: str + bundle_version: str | None = None msg: str | None = None """Additional Message that can be used for logging to determine failure/zombie""" diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py index 9c7af83673abe..0c37251a098d1 100644 --- a/airflow/dag_processing/manager.py +++ b/airflow/dag_processing/manager.py @@ -191,7 +191,7 @@ class DagFileProcessorManager: _parsing_start_time: float = attrs.field(init=False) _num_run: int = attrs.field(default=0, init=False) - _callback_to_execute: dict[str, list[CallbackRequest]] = attrs.field( + _callback_to_execute: dict[DagFileInfo, list[CallbackRequest]] = attrs.field( factory=lambda: defaultdict(list), init=False ) @@ -407,18 +407,14 @@ def _fetch_callbacks( def _add_callback_to_queue(self, request: CallbackRequest): self.log.debug("Queuing %s CallbackRequest: %s", type(request).__name__, request) - self.log.warning("Callbacks are not implemented yet!") - # TODO: AIP-66 make callbacks bundle aware - return - self._callback_to_execute[request.full_filepath].append(request) - if request.full_filepath in self._file_queue: + file_info = DagFileInfo(path=request.full_filepath, bundle_name=request.bundle_name) + self._callback_to_execute[file_info].append(request) + if file_info in self._file_queue: # Remove file paths matching request.full_filepath from self._file_queue # Since we are already going to use that filepath to run callback, # there is no need to have same file path again in the queue - # todo (AIP-66): update re bundle and rel loc - self._file_queue = deque(f for f in self._file_queue if f != request.full_filepath) - # todo (AIP-66): update re bundle and rel loc - self._add_files_to_queue([request.full_filepath], True) + self._file_queue = deque(f for f in self._file_queue if f != file_info) + self._add_files_to_queue([file_info], True) Stats.incr("dag_processing.other_callback_count") @classmethod @@ -690,10 +686,9 @@ def set_files(self, files: list[DagFileInfo]): self._file_queue = deque(x for x in self._file_queue if x in files) Stats.gauge("dag_processing.file_path_queue_size", len(self._file_queue)) - # TODO: AIP-66 make callbacks bundle aware - # callback_paths_to_del = [x for x in self._callback_to_execute if x not in new_file_paths] - # for path_to_del in callback_paths_to_del: - # del self._callback_to_execute[path_to_del] + callback_paths_to_del = [x for x in self._callback_to_execute if x not in files] + for path_to_del in callback_paths_to_del: + del self._callback_to_execute[path_to_del] # Stop processors that are working on deleted files filtered_processors = {} @@ -785,8 +780,7 @@ def _get_logger_for_dag_file(self, dag_file: DagFileInfo): def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: id = uuid7() - # callback_to_execute_for_file = self._callback_to_execute.pop(file_path, []) - callback_to_execute_for_file: list[CallbackRequest] = [] + callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) return DagFileProcessorProcess.start( id=id, diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index 1e98dbf2aba29..c2f1d145e4d8a 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -854,6 +854,8 @@ def process_executor_events( cls.logger().error(msg) request = TaskCallbackRequest( full_filepath=ti.dag_model.fileloc, + bundle_name=ti.dag_model.bundle_name, + bundle_version=ti.dag_model.bundle_version, ti=ti, msg=msg, ) @@ -1630,6 +1632,8 @@ def _schedule_dag_run( full_filepath=dag.fileloc, dag_id=dag.dag_id, run_id=dag_run.run_id, + bundle_name=dag_model.bundle_name, + bundle_version=dag_run.bundle_version, is_failure_callback=True, msg="timed_out", ) @@ -2013,6 +2017,8 @@ def _purge_zombies(self, zombies: list[tuple[TI, str]], *, session: Session) -> zombie_message_details = self._generate_zombie_message_details(ti) request = TaskCallbackRequest( full_filepath=file_loc, + bundle_name=ti.dag_model.bundle_name, + bundle_version=ti.dag_model.bundle_version, ti=ti, msg=str(zombie_message_details), ) diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 22757c972e953..7b96cd283ac51 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -925,6 +925,8 @@ def recalculate(self) -> _UnfinishedStates: full_filepath=dag.fileloc, dag_id=self.dag_id, run_id=self.run_id, + bundle_name=self.dag_model.bundle_name, + bundle_version=self.bundle_version, is_failure_callback=True, msg="task_failure", ) @@ -952,6 +954,8 @@ def recalculate(self) -> _UnfinishedStates: full_filepath=dag.fileloc, dag_id=self.dag_id, run_id=self.run_id, + bundle_name=self.dag_model.bundle_name, + bundle_version=self.bundle_version, is_failure_callback=False, msg="success", ) @@ -969,6 +973,8 @@ def recalculate(self) -> _UnfinishedStates: full_filepath=dag.fileloc, dag_id=self.dag_id, run_id=self.run_id, + bundle_name=self.dag_model.bundle_name, + bundle_version=self.bundle_version, is_failure_callback=True, msg="all_tasks_deadlocked", ) diff --git a/tests/dag_processing/test_manager.py b/tests/dag_processing/test_manager.py index d2dacd3587087..6e7dd8ae6c139 100644 --- a/tests/dag_processing/test_manager.py +++ b/tests/dag_processing/test_manager.py @@ -492,6 +492,7 @@ def test_kill_timed_out_processors_no_kill(self): full_filepath="/opt/airflow/dags/dag_callback_dag.py", dag_id="dag_id", run_id="run_id", + bundle_name="testing", is_failure_callback=False, ) ], @@ -502,12 +503,14 @@ def test_kill_timed_out_processors_no_kill(self): b'"requests_fd":123,"callback_requests":' b"[" b"{" - b'"full_filepath":"/opt/airflow/dags/dag_callback_dag.py",' + b'"filepath":"/opt/airflow/dags/dag_callback_dag.py",' b'"msg":null,' b'"dag_id":"dag_id",' b'"run_id":"run_id",' b'"is_failure_callback":false,' b'"type":"DagCallbackRequest"' + b'"bundle_name":"testing"' + b'"bundle_version":null' b"}" b"]," b'"type":"DagFileParseRequest"' @@ -650,12 +653,14 @@ def test_fetch_callbacks_from_database(self, tmp_path, configure_testing_dag_bun callback1 = DagCallbackRequest( dag_id="test_start_date_scheduling", + bundle_name="testing", full_filepath=str(dag_filepath), is_failure_callback=True, run_id="123", ) callback2 = DagCallbackRequest( dag_id="test_start_date_scheduling", + bundle_name="testing", full_filepath=str(dag_filepath), is_failure_callback=True, run_id="456", @@ -686,6 +691,7 @@ def test_fetch_callbacks_from_database_max_per_loop(self, tmp_path, configure_te for i in range(5): callback = DagCallbackRequest( dag_id="test_start_date_scheduling", + bundle_name="testing", full_filepath=str(dag_filepath), is_failure_callback=True, run_id=str(i), @@ -703,7 +709,6 @@ def test_fetch_callbacks_from_database_max_per_loop(self, tmp_path, configure_te manager.run() assert session.query(DbCallbackRequest).count() == 1 - @pytest.mark.skip("AIP-66: callbacks are not implemented yet") def test_callback_queue(self, tmp_path): """ This test has gotten a bit out of sync with the codebase. @@ -723,6 +728,7 @@ def test_callback_queue(self, tmp_path): full_filepath=TEST_DAGS_FOLDER / "green_eggs/ham/file1.py", dag_id="dag1", run_id="run1", + bundle_name="testing", is_failure_callback=False, msg=None, ) @@ -730,6 +736,7 @@ def test_callback_queue(self, tmp_path): full_filepath=TEST_DAGS_FOLDER / "green_eggs/ham/file1.py", dag_id="dag1", run_id="run1", + bundle_name="testing", is_failure_callback=False, msg=None, ) @@ -741,6 +748,7 @@ def test_callback_queue(self, tmp_path): full_filepath=TEST_DAGS_FOLDER / "green_eggs/ham/file2.py", dag_id="dag2", run_id="run1", + bundle_name=dag2_path.bundle_name, is_failure_callback=False, msg=None, ) @@ -752,20 +760,20 @@ def test_callback_queue(self, tmp_path): # then - requests should be in manager's queue, with dag2 ahead of dag1 (because it was added last) assert manager._file_queue == deque([dag2_path, dag1_path]) assert set(manager._callback_to_execute.keys()) == { - dag1_req1.full_filepath, - dag2_req1.full_filepath, + dag1_path, + dag2_path, } - assert manager._callback_to_execute[dag2_req1.full_filepath] == [dag2_req1] + assert manager._callback_to_execute[dag2_path] == [dag2_req1] # update the queue, although the callback is registered - assert manager._file_queue == deque([dag2_req1.full_filepath, dag1_req1.full_filepath]) + assert manager._file_queue == deque([dag2_path, dag1_path]) # when manager._add_callback_to_queue(dag1_req2) # then - non-sla callback should have brought dag1 to the fore - assert manager._file_queue == deque([dag1_req1.full_filepath, dag2_req1.full_filepath]) - assert manager._callback_to_execute[dag1_req1.full_filepath] == [ + assert manager._file_queue == deque([dag1_path, dag2_path]) + assert manager._callback_to_execute[dag1_path] == [ dag1_req1, dag1_req2, ] @@ -779,7 +787,7 @@ def test_callback_queue(self, tmp_path): id=mock.ANY, path=dag1_req1.full_filepath, callbacks=[dag1_req1, dag1_req2], selector=mock.ANY ) # And removed from the queue - assert dag1_req1.full_filepath not in manager._callback_to_execute + assert dag1_path not in manager._callback_to_execute def test_dag_with_assets(self, session, configure_testing_dag_bundle): """'Integration' test to ensure that the assets get parsed and stored correctly for parsed dags.""" From 6933ae102f91c4f6c444e76803c6d6de7daa4f36 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Tue, 21 Jan 2025 19:35:56 +0100 Subject: [PATCH 2/7] AIP-66: Make DAG callbacks bundle aware This involves using relative paths in the callbacks, resolving the full path and using the it to queue the callback in the file processor process. Don't serialize dag if running callback Refactor codes and fix tests update callback test Fix typing Don't use relative fileloc yet Add test apply suggestions from code review Fix some tests and require bundle_version in DagCallbackRequest Resolve conflicts and fix tests join load dagversion as part of TI in process executor event --- airflow/callbacks/callback_requests.py | 4 +- airflow/dag_processing/manager.py | 28 +-- airflow/dag_processing/processor.py | 13 +- airflow/jobs/scheduler_job_runner.py | 33 +-- airflow/models/dagrun.py | 10 +- airflow/triggers/base.py | 4 +- .../test_celery_kubernetes_executor.py | 6 +- .../test_local_kubernetes_executor.py | 4 +- tests/callbacks/test_callback_requests.py | 12 +- tests/dag_processing/test_manager.py | 198 ++++++++++-------- tests/dag_processing/test_processor.py | 19 +- tests/jobs/test_scheduler_job.py | 74 ++++--- tests/models/test_dagrun.py | 17 +- tests_common/pytest_plugin.py | 1 - 14 files changed, 249 insertions(+), 174 deletions(-) diff --git a/airflow/callbacks/callback_requests.py b/airflow/callbacks/callback_requests.py index 38f291657eb63..b990b0df7ea42 100644 --- a/airflow/callbacks/callback_requests.py +++ b/airflow/callbacks/callback_requests.py @@ -34,10 +34,10 @@ class BaseCallbackRequest(BaseModel): :param msg: Additional Message that can be used for logging """ - full_filepath: str + filepath: str """File Path to use to run the callback""" bundle_name: str - bundle_version: str | None = None + bundle_version: str | None msg: str | None = None """Additional Message that can be used for logging to determine failure/zombie""" diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py index 0c37251a098d1..22c703c6f3473 100644 --- a/airflow/dag_processing/manager.py +++ b/airflow/dag_processing/manager.py @@ -107,6 +107,7 @@ class DagFileInfo: rel_path: Path bundle_name: str bundle_path: Path | None = field(compare=False, default=None) + bundle_version: str | None = None @property def absolute_path(self) -> Path: @@ -407,13 +408,22 @@ def _fetch_callbacks( def _add_callback_to_queue(self, request: CallbackRequest): self.log.debug("Queuing %s CallbackRequest: %s", type(request).__name__, request) - file_info = DagFileInfo(path=request.full_filepath, bundle_name=request.bundle_name) + try: + bundle = DagBundlesManager().get_bundle(name=request.bundle_name, version=request.bundle_version) + except ValueError: + # Bundle no longer configured + self.log.error("Bundle %s no longer configured, skipping callback", request.bundle_name) + return None + + dag_absolute_path = os.fspath(Path(bundle.path, request.filepath)) + + file_info = DagFileInfo( + path=dag_absolute_path, + bundle_path=bundle.path, + bundle_name=request.bundle_name, + bundle_version=request.bundle_version, + ) self._callback_to_execute[file_info].append(request) - if file_info in self._file_queue: - # Remove file paths matching request.full_filepath from self._file_queue - # Since we are already going to use that filepath to run callback, - # there is no need to have same file path again in the queue - self._file_queue = deque(f for f in self._file_queue if f != file_info) self._add_files_to_queue([file_info], True) Stats.incr("dag_processing.other_callback_count") @@ -682,14 +692,8 @@ def set_files(self, files: list[DagFileInfo]): """ self._files = files - # remove from queue any files no longer in the _files list - self._file_queue = deque(x for x in self._file_queue if x in files) Stats.gauge("dag_processing.file_path_queue_size", len(self._file_queue)) - callback_paths_to_del = [x for x in self._callback_to_execute if x not in files] - for path_to_del in callback_paths_to_del: - del self._callback_to_execute[path_to_del] - # Stop processors that are working on deleted files filtered_processors = {} for file, processor in self._processors.items(): diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py index 2a822e1bae97f..e8f5d3978ccdc 100644 --- a/airflow/dag_processing/processor.py +++ b/airflow/dag_processing/processor.py @@ -67,10 +67,11 @@ def _parse_file_entrypoint(): log = structlog.get_logger(logger_name="task") result = _parse_file(msg, log) - comms_decoder.send_request(log, result) + if result is not None: + comms_decoder.send_request(log, result) -def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileParsingResult: +def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileParsingResult | None: # TODO: Set known_pool names on DagBag! bag = DagBag( dag_folder=msg.file, @@ -79,6 +80,11 @@ def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileP safe_mode=True, load_op_links=False, ) + if msg.callback_requests: + # If the request is for callback, we shouldn't serialize the DAGs + _execute_callbacks(bag, msg.callback_requests, log) + return None + serialized_dags, serialization_import_errors = _serialize_dags(bag, log) bag.import_errors.update(serialization_import_errors) dags = [LazyDeserializedDAG(data=serdag) for serdag in serialized_dags] @@ -89,9 +95,6 @@ def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileP # TODO: Make `bag.dag_warnings` not return SQLA model objects warnings=[], ) - - if msg.callback_requests: - _execute_callbacks(bag, msg.callback_requests, log) return result diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index c2f1d145e4d8a..47ae2d58e161f 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -35,7 +35,7 @@ from deprecated import deprecated from sqlalchemy import and_, delete, exists, func, select, text, tuple_, update from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import lazyload, load_only, make_transient, selectinload +from sqlalchemy.orm import joinedload, lazyload, load_only, make_transient, selectinload from sqlalchemy.sql import expression from airflow import settings @@ -756,7 +756,12 @@ def process_executor_events( # Check state of finished tasks filter_for_tis = TI.filter_for_tis(tis_with_right_state) - query = select(TI).where(filter_for_tis).options(selectinload(TI.dag_model)) + query = ( + select(TI) + .where(filter_for_tis) + .options(selectinload(TI.dag_model)) + .options(joinedload(TI.dag_version)) + ) # row lock this entire set of taskinstances to make sure the scheduler doesn't fail when we have # multi-schedulers tis_query: Query = with_row_locks(query, of=TI, session=session, skip_locked=True) @@ -853,9 +858,9 @@ def process_executor_events( # too, which would lead to double logging cls.logger().error(msg) request = TaskCallbackRequest( - full_filepath=ti.dag_model.fileloc, - bundle_name=ti.dag_model.bundle_name, - bundle_version=ti.dag_model.bundle_version, + filepath=ti.dag_version.dag_code.fileloc, + bundle_name=ti.dag_version.bundle_name, + bundle_version=ti.dag_version.bundle_version, ti=ti, msg=msg, ) @@ -1629,7 +1634,7 @@ def _schedule_dag_run( dag_model.calculate_dagrun_date_fields(dag, dag.get_run_data_interval(dag_run)) callback_to_execute = DagCallbackRequest( - full_filepath=dag.fileloc, + filepath=dag_model.fileloc, dag_id=dag.dag_id, run_id=dag_run.run_id, bundle_name=dag_model.bundle_name, @@ -1995,11 +2000,11 @@ def _find_and_purge_zombies(self) -> None: if zombies := self._find_zombies(session=session): self._purge_zombies(zombies, session=session) - def _find_zombies(self, *, session: Session) -> list[tuple[TI, str]]: + def _find_zombies(self, *, session: Session) -> list[TI]: self.log.debug("Finding 'running' jobs without a recent heartbeat") limit_dttm = timezone.utcnow() - timedelta(seconds=self._zombie_threshold_secs) - zombies = session.execute( - select(TI, DM.fileloc) + zombies = session.scalars( + select(TI) .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql") .join(DM, TI.dag_id == DM.dag_id) .where( @@ -2012,13 +2017,13 @@ def _find_zombies(self, *, session: Session) -> list[tuple[TI, str]]: self.log.warning("Failing %s TIs without heartbeat after %s", len(zombies), limit_dttm) return zombies - def _purge_zombies(self, zombies: list[tuple[TI, str]], *, session: Session) -> None: - for ti, file_loc in zombies: + def _purge_zombies(self, zombies: list[TI], *, session: Session) -> None: + for ti in zombies: zombie_message_details = self._generate_zombie_message_details(ti) request = TaskCallbackRequest( - full_filepath=file_loc, - bundle_name=ti.dag_model.bundle_name, - bundle_version=ti.dag_model.bundle_version, + filepath=ti.dag_version.dag_code.fileloc, + bundle_name=ti.dag_version.bundle_name, + bundle_version=ti.dag_version.bundle_version, ti=ti, msg=str(zombie_message_details), ) diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 7b96cd283ac51..febb8ddb4c721 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -922,10 +922,10 @@ def recalculate(self) -> _UnfinishedStates: dag.handle_callback(self, success=False, reason="task_failure", session=session) elif dag.has_on_failure_callback: callback = DagCallbackRequest( - full_filepath=dag.fileloc, + filepath=self.dag_version.dag_code.fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_model.bundle_name, + bundle_name=self.dag_version.bundle_name, bundle_version=self.bundle_version, is_failure_callback=True, msg="task_failure", @@ -951,10 +951,10 @@ def recalculate(self) -> _UnfinishedStates: dag.handle_callback(self, success=True, reason="success", session=session) elif dag.has_on_success_callback: callback = DagCallbackRequest( - full_filepath=dag.fileloc, + filepath=self.dag_version.dag_code.fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_model.bundle_name, + bundle_name=self.dag_version.bundle_name, bundle_version=self.bundle_version, is_failure_callback=False, msg="success", @@ -970,7 +970,7 @@ def recalculate(self) -> _UnfinishedStates: dag.handle_callback(self, success=False, reason="all_tasks_deadlocked", session=session) elif dag.has_on_failure_callback: callback = DagCallbackRequest( - full_filepath=dag.fileloc, + filepath=self.dag_model.fileloc, dag_id=self.dag_id, run_id=self.run_id, bundle_name=self.dag_model.bundle_name, diff --git a/airflow/triggers/base.py b/airflow/triggers/base.py index 4e88465d533a7..880b1cd35635b 100644 --- a/airflow/triggers/base.py +++ b/airflow/triggers/base.py @@ -217,9 +217,11 @@ def _submit_callback_if_necessary(self, *, task_instance: TaskInstance, session) """Submit a callback request if the task state is SUCCESS or FAILED.""" if self.task_instance_state in (TaskInstanceState.SUCCESS, TaskInstanceState.FAILED): request = TaskCallbackRequest( - full_filepath=task_instance.dag_model.fileloc, + filepath=task_instance.dag_model.fileloc, ti=task_instance, task_callback_type=self.task_instance_state, + bundle_name=task_instance.dag_model.bundle_name, + bundle_version=task_instance.dag_model.bundle_version, ) log.info("Sending callback: %s", request) try: diff --git a/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py b/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py index 027793c29d527..9728b19e9a720 100644 --- a/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py +++ b/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py @@ -261,9 +261,11 @@ def test_send_callback(self): cel_k8s_exec.callback_sink = mock.MagicMock() if AIRFLOW_V_3_0_PLUS: - callback = DagCallbackRequest(full_filepath="fake", dag_id="fake", run_id="fake") + callback = DagCallbackRequest( + filepath="fake", dag_id="fake", run_id="fake", bundle_name="testing", bundle_version=None + ) else: - callback = CallbackRequest(full_filepath="fake") + callback = CallbackRequest(full_filepath="fake", dag_id="fake", run_id="fake") cel_k8s_exec.send_callback(callback) cel_k8s_exec.callback_sink.send.assert_called_once_with(callback) diff --git a/providers/tests/cncf/kubernetes/executors/test_local_kubernetes_executor.py b/providers/tests/cncf/kubernetes/executors/test_local_kubernetes_executor.py index be4d8b936739e..c712b2d4b5e59 100644 --- a/providers/tests/cncf/kubernetes/executors/test_local_kubernetes_executor.py +++ b/providers/tests/cncf/kubernetes/executors/test_local_kubernetes_executor.py @@ -116,7 +116,9 @@ def test_send_callback(self): local_k8s_exec.callback_sink = mock.MagicMock() if AIRFLOW_V_3_0_PLUS: - callback = DagCallbackRequest(full_filepath="fake", dag_id="fake", run_id="fake") + callback = DagCallbackRequest( + filepath="fake", dag_id="fake", run_id="fake", bundle_name="fake", bundle_version=None + ) else: callback = CallbackRequest(full_filepath="fake") local_k8s_exec.send_callback(callback) diff --git a/tests/callbacks/test_callback_requests.py b/tests/callbacks/test_callback_requests.py index 68362d3623941..345af3b597cb7 100644 --- a/tests/callbacks/test_callback_requests.py +++ b/tests/callbacks/test_callback_requests.py @@ -43,10 +43,12 @@ class TestCallbackRequest: ), ( DagCallbackRequest( - full_filepath="filepath", + filepath="filepath", dag_id="fake_dag", run_id="fake_run", is_failure_callback=False, + bundle_name="testing", + bundle_version=None, ), DagCallbackRequest, ), @@ -66,8 +68,7 @@ def test_from_json(self, input, request_class): ) input = TaskCallbackRequest( - full_filepath="filepath", - ti=ti, + filepath="filepath", ti=ti, bundle_name="testing", bundle_version=None ) json_str = input.to_json() result = request_class.from_json(json_str) @@ -79,10 +80,7 @@ def test_taskcallback_to_json_with_start_date_and_end_date(self, session, create ti.end_date = timezone.utcnow() session.merge(ti) session.flush() - input = TaskCallbackRequest( - full_filepath="filepath", - ti=ti, - ) + input = TaskCallbackRequest(filepath="filepath", ti=ti, bundle_name="testing", bundle_version=None) json_str = input.to_json() result = TaskCallbackRequest.from_json(json_str) assert input == result diff --git a/tests/dag_processing/test_manager.py b/tests/dag_processing/test_manager.py index 6e7dd8ae6c139..17ecad122a5dd 100644 --- a/tests/dag_processing/test_manager.py +++ b/tests/dag_processing/test_manager.py @@ -157,12 +157,14 @@ def test_max_runs_when_no_files(self, tmp_path): # TODO: AIP-66 no asserts? - def test_start_new_processes_with_same_filepath(self): + def test_start_new_processes_with_same_filepath(self, configure_testing_dag_bundle): """ Test that when a processor already exist with a filepath, a new processor won't be created with that filepath. The filepath will just be removed from the list. """ - manager = DagFileProcessorManager(max_runs=1) + with configure_testing_dag_bundle("/tmp"): + manager = DagFileProcessorManager(max_runs=1) + manager._dag_bundles = list(DagBundlesManager().get_all_dag_bundles()) file_1 = DagFileInfo(bundle_name="testing", rel_path=Path("file_1.py"), bundle_path=TEST_DAGS_FOLDER) file_2 = DagFileInfo(bundle_name="testing", rel_path=Path("file_2.py"), bundle_path=TEST_DAGS_FOLDER) @@ -472,6 +474,7 @@ def test_kill_timed_out_processors_no_kill(self): manager._kill_timed_out_processors() mock_kill.assert_not_called() + @pytest.mark.usefixtures("testing_dag_bundle") @pytest.mark.parametrize( ["callbacks", "path", "expected_buffer"], [ @@ -489,10 +492,11 @@ def test_kill_timed_out_processors_no_kill(self): pytest.param( [ DagCallbackRequest( - full_filepath="/opt/airflow/dags/dag_callback_dag.py", + filepath="dag_callback_dag.py", dag_id="dag_id", run_id="run_id", bundle_name="testing", + bundle_version=None, is_failure_callback=False, ) ], @@ -503,14 +507,14 @@ def test_kill_timed_out_processors_no_kill(self): b'"requests_fd":123,"callback_requests":' b"[" b"{" - b'"filepath":"/opt/airflow/dags/dag_callback_dag.py",' + b'"filepath":"dag_callback_dag.py",' + b'"bundle_name":"testing",' + b'"bundle_version":null,' b'"msg":null,' b'"dag_id":"dag_id",' b'"run_id":"run_id",' b'"is_failure_callback":false,' b'"type":"DagCallbackRequest"' - b'"bundle_name":"testing"' - b'"bundle_version":null' b"}" b"]," b'"type":"DagFileParseRequest"' @@ -648,20 +652,22 @@ def test_deactivate_deleted_dags(self, dag_maker): assert dagbag.get_dag("test_dag2").get_is_active() is False @conf_vars({("core", "load_examples"): "False"}) - def test_fetch_callbacks_from_database(self, tmp_path, configure_testing_dag_bundle): + def test_fetch_callbacks_from_database(self, configure_testing_dag_bundle): dag_filepath = TEST_DAG_FOLDER / "test_on_failure_callback_dag.py" callback1 = DagCallbackRequest( dag_id="test_start_date_scheduling", bundle_name="testing", - full_filepath=str(dag_filepath), + bundle_version=None, + filepath="test_on_failure_callback_dag.py", is_failure_callback=True, run_id="123", ) callback2 = DagCallbackRequest( dag_id="test_start_date_scheduling", bundle_name="testing", - full_filepath=str(dag_filepath), + bundle_version=None, + filepath="test_on_failure_callback_dag.py", is_failure_callback=True, run_id="456", ) @@ -670,7 +676,7 @@ def test_fetch_callbacks_from_database(self, tmp_path, configure_testing_dag_bun session.add(DbCallbackRequest(callback=callback1, priority_weight=11)) session.add(DbCallbackRequest(callback=callback2, priority_weight=10)) - with configure_testing_dag_bundle(tmp_path): + with configure_testing_dag_bundle(dag_filepath): manager = DagFileProcessorManager(max_runs=1) with create_session() as session: @@ -692,13 +698,14 @@ def test_fetch_callbacks_from_database_max_per_loop(self, tmp_path, configure_te callback = DagCallbackRequest( dag_id="test_start_date_scheduling", bundle_name="testing", - full_filepath=str(dag_filepath), + bundle_version=None, + filepath="test_on_failure_callback_dag.py", is_failure_callback=True, run_id=str(i), ) session.add(DbCallbackRequest(callback=callback, priority_weight=i)) - with configure_testing_dag_bundle(tmp_path): + with configure_testing_dag_bundle(dag_filepath): manager = DagFileProcessorManager(max_runs=1) with create_session() as session: @@ -709,85 +716,102 @@ def test_fetch_callbacks_from_database_max_per_loop(self, tmp_path, configure_te manager.run() assert session.query(DbCallbackRequest).count() == 1 - def test_callback_queue(self, tmp_path): - """ - This test has gotten a bit out of sync with the codebase. - - I am just updating it to be consistent with the changes in DagFileInfo - """ - # given - manager = DagFileProcessorManager( - max_runs=1, - processor_timeout=365 * 86_400, - ) - - dag1_path = DagFileInfo( - bundle_name="testing", rel_path=Path("green_eggs/ham/file1.py"), bundle_path=TEST_DAGS_FOLDER - ) - dag1_req1 = DagCallbackRequest( - full_filepath=TEST_DAGS_FOLDER / "green_eggs/ham/file1.py", - dag_id="dag1", - run_id="run1", - bundle_name="testing", - is_failure_callback=False, - msg=None, - ) - dag1_req2 = DagCallbackRequest( - full_filepath=TEST_DAGS_FOLDER / "green_eggs/ham/file1.py", - dag_id="dag1", - run_id="run1", - bundle_name="testing", - is_failure_callback=False, - msg=None, - ) - - dag2_path = DagFileInfo( - bundle_name="testing", rel_path=Path("green_eggs/ham/file2.py"), bundle_path=TEST_DAGS_FOLDER - ) - dag2_req1 = DagCallbackRequest( - full_filepath=TEST_DAGS_FOLDER / "green_eggs/ham/file2.py", - dag_id="dag2", - run_id="run1", - bundle_name=dag2_path.bundle_name, - is_failure_callback=False, - msg=None, - ) - - # when - manager._add_callback_to_queue(dag1_req1) - manager._add_callback_to_queue(dag2_req1) - - # then - requests should be in manager's queue, with dag2 ahead of dag1 (because it was added last) - assert manager._file_queue == deque([dag2_path, dag1_path]) - assert set(manager._callback_to_execute.keys()) == { - dag1_path, - dag2_path, - } - assert manager._callback_to_execute[dag2_path] == [dag2_req1] + @mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file") + def test_callback_queue(self, mock_logger, configure_testing_dag_bundle): + tmp_path = "/green_eggs/ham" + with configure_testing_dag_bundle(tmp_path): + # given + manager = DagFileProcessorManager( + max_runs=1, + processor_timeout=365 * 86_400, + ) + manager._dag_bundles = list(DagBundlesManager().get_all_dag_bundles()) - # update the queue, although the callback is registered - assert manager._file_queue == deque([dag2_path, dag1_path]) + dag1_path = DagFileInfo( + bundle_name="testing", path=f"{tmp_path}/file1.py", bundle_path=Path(tmp_path) + ) + dag1_req1 = DagCallbackRequest( + filepath="file1.py", + dag_id="dag1", + run_id="run1", + is_failure_callback=False, + bundle_name="testing", + bundle_version=None, + msg=None, + ) + dag1_req2 = DagCallbackRequest( + filepath="file1.py", + dag_id="dag1", + run_id="run1", + is_failure_callback=False, + bundle_name="testing", + bundle_version=None, + msg=None, + ) - # when - manager._add_callback_to_queue(dag1_req2) + dag2_path = DagFileInfo( + bundle_name="testing", path=f"{tmp_path}/file2.py", bundle_path=Path(tmp_path) + ) + dag2_req1 = DagCallbackRequest( + filepath="file2.py", + dag_id="dag2", + run_id="run1", + bundle_name=dag2_path.bundle_name, + bundle_version=None, + is_failure_callback=False, + msg=None, + ) - # then - non-sla callback should have brought dag1 to the fore - assert manager._file_queue == deque([dag1_path, dag2_path]) - assert manager._callback_to_execute[dag1_path] == [ - dag1_req1, - dag1_req2, - ] + # when + manager._add_callback_to_queue(dag1_req1) + manager._add_callback_to_queue(dag2_req1) + + # then - requests should be in manager's queue, with dag2 ahead of dag1 (because it was added last) + assert manager._file_queue == deque([dag2_path, dag1_path]) + assert set(manager._callback_to_execute.keys()) == { + dag1_path, + dag2_path, + } + assert manager._callback_to_execute[dag2_path] == [dag2_req1] + + # update the queue, although the callback is registered + assert manager._file_queue == deque([dag2_path, dag1_path]) + + # when + manager._add_callback_to_queue(dag1_req2) + # Since dag1_req2 is same as dag1_req1, we now have 2 items in file_path_queue + assert manager._file_path_queue == deque([dag2_path, dag1_path]) + assert manager._callback_to_execute[dag1_path] == [ + dag1_req1, + dag1_req2, + ] - with mock.patch.object( - DagFileProcessorProcess, "start", side_effect=lambda *args, **kwargs: self.mock_processor() - ) as start: - manager._start_new_processes() - # Callbacks passed to process ctor - start.assert_any_call( - id=mock.ANY, path=dag1_req1.full_filepath, callbacks=[dag1_req1, dag1_req2], selector=mock.ANY - ) - # And removed from the queue - assert dag1_path not in manager._callback_to_execute + with mock.patch.object( + DagFileProcessorProcess, "start", side_effect=lambda *args, **kwargs: self.mock_processor() + ) as start: + manager._start_new_processes() + # Callbacks passed to processor + assert start.call_args_list == [ + mock.call( + id=mock.ANY, + path=dag2_path.path, + bundle_path=dag2_path.bundle_path, + callbacks=[dag2_req1], + selector=mock.ANY, + logger=mock_logger.return_value, + ), + mock.call( + id=mock.ANY, + path=dag1_path.path, + bundle_path=dag1_path.bundle_path, + callbacks=[dag1_req1, dag1_req2], + selector=mock.ANY, + logger=mock_logger.return_value, + ), + ] + # And removed from the queue + assert dag1_path not in manager._callback_to_execute + assert dag2_path not in manager._callback_to_execute def test_dag_with_assets(self, session, configure_testing_dag_bundle): """'Integration' test to ensure that the assets get parsed and stored correctly for parsed dags.""" diff --git a/tests/dag_processing/test_processor.py b/tests/dag_processing/test_processor.py index 9844186be1df5..a01d3edb2ff3b 100644 --- a/tests/dag_processing/test_processor.py +++ b/tests/dag_processing/test_processor.py @@ -75,7 +75,7 @@ def disable_load_example(): class TestDagFileProcessor: def _process_file( self, file_path, callback_requests: list[CallbackRequest] | None = None - ) -> DagFileParsingResult: + ) -> DagFileParsingResult | None: return _parse_file( DagFileParseRequest( file=file_path, @@ -128,7 +128,7 @@ def fake_collect_dags(dagbag: DagBag, *args, **kwargs): dagbag.import_errors["a.py"] = "Import error" resp = self._process_file("a.py") - + assert resp is not None assert not resp.serialized_dags assert resp.import_errors is not None assert "a.py" in resp.import_errors @@ -184,6 +184,7 @@ def disable_capturing(): sys.stdin, sys.stdout, sys.stderr = old_in, old_out, old_err +@pytest.mark.usefixtures("testing_dag_bundle") @pytest.mark.usefixtures("disable_capturing") def test_parse_file_entrypoint_parses_dag_callbacks(spy_agency): r, w = socketpair() @@ -193,7 +194,7 @@ def test_parse_file_entrypoint_parses_dag_callbacks(spy_agency): w.makefile("wb").write( b'{"file":"/files/dags/wait.py","bundle_path":"/files/dags","requests_fd":' + str(w2.fileno()).encode("ascii") - + b',"callback_requests": [{"full_filepath": "/files/dags/wait.py", ' + + b',"callback_requests": [{"filepath": "wait.py", "bundle_name": "testing", "bundle_version": null, ' b'"msg": "task_failure", "dag_id": "wait_to_fail", "run_id": ' b'"manual__2024-12-30T21:02:55.203691+00:00", ' b'"is_failure_callback": true, "type": "DagCallbackRequest"}], "type": "DagFileParseRequest"}\n' @@ -209,11 +210,13 @@ def test_parse_file_entrypoint_parses_dag_callbacks(spy_agency): assert msg.file == "/files/dags/wait.py" assert msg.callback_requests == [ DagCallbackRequest( - full_filepath="/files/dags/wait.py", + filepath="wait.py", msg="task_failure", dag_id="wait_to_fail", run_id="manual__2024-12-30T21:02:55.203691+00:00", is_failure_callback=True, + bundle_name="testing", + bundle_version=None, ) ] @@ -236,10 +239,12 @@ def fake_collect_dags(self, *args, **kwargs): requests = [ DagCallbackRequest( - full_filepath="A", + filepath="A", msg="Message", dag_id="a", run_id="b", + bundle_name="testing", + bundle_version=None, ) ] _parse_file( @@ -270,9 +275,11 @@ def fake_collect_dags(self, *args, **kwargs): requests = [ TaskCallbackRequest( - full_filepath="A", + filepath="A", msg="Message", ti=None, + bundle_name="testing", + bundle_version=None, ) ] _parse_file( diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index faf8fa73284af..e35820db2c86c 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -152,6 +152,7 @@ def _create_dagrun( state=state, start_date=start_date, triggered_by=DagRunTriggeredByType.TEST, + dag_version=DagVersion.get_latest_version(dag.dag_id), ) return _create_dagrun @@ -370,13 +371,17 @@ def test_process_executor_events_with_no_callback(self, mock_stats_incr, mock_ta @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr") - def test_process_executor_events_with_callback(self, mock_stats_incr, mock_task_callback, dag_maker): + def test_process_executor_events_with_callback( + self, mock_stats_incr, mock_task_callback, dag_maker, session + ): dag_id = "test_process_executor_events_with_callback" task_id_1 = "dummy_task" with dag_maker(dag_id=dag_id, fileloc="/test_path1/") as dag: - task1 = EmptyOperator(task_id=task_id_1, on_failure_callback=lambda x: print("hi")) - ti1 = dag_maker.create_dagrun().get_task_instance(task1.task_id) + EmptyOperator(task_id=task_id_1, on_failure_callback=lambda x: print("hi")) + dagv = DagVersion.get_latest_version(dag.dag_id) + dr = dag_maker.create_dagrun(dag_version=dagv) + ti1 = dr.task_instances[0] mock_stats_incr.reset_mock() @@ -386,8 +391,6 @@ def test_process_executor_events_with_callback(self, mock_stats_incr, mock_task_ scheduler_job = Job(executor=executor) self.job_runner = SchedulerJobRunner(scheduler_job) - session = settings.Session() - ti1.state = State.QUEUED session.merge(ti1) session.commit() @@ -400,8 +403,10 @@ def test_process_executor_events_with_callback(self, mock_stats_incr, mock_task_ # will be set to failed in dag parsing process assert ti1.state == State.QUEUED mock_task_callback.assert_called_once_with( - full_filepath=dag.fileloc, + filepath=dag.fileloc, ti=mock.ANY, + bundle_name="dag_maker", + bundle_version=None, msg=f"Executor {executor} reported that the task instance " " " "finished with state failed, but the task instance's state attribute is queued. " @@ -2415,10 +2420,12 @@ def test_dagrun_timeout_verify_max_active_runs(self, dag_maker): assert isinstance(orm_dag.next_dagrun_create_after, datetime.datetime) expected_callback = DagCallbackRequest( - full_filepath=dr.dag.fileloc, + filepath=dr.dag.fileloc, dag_id=dr.dag_id, is_failure_callback=True, run_id=dr.run_id, + bundle_name=orm_dag.bundle_name, + bundle_version=orm_dag.bundle_version, msg="timed_out", ) @@ -2454,10 +2461,12 @@ def test_dagrun_timeout_fails_run(self, dag_maker): assert dr.state == State.FAILED expected_callback = DagCallbackRequest( - full_filepath=dr.dag.fileloc, + filepath=dr.dag.fileloc, dag_id=dr.dag_id, is_failure_callback=True, run_id=dr.run_id, + bundle_name=dr.dag.get_bundle_name(), + bundle_version=dr.dag.get_bundle_version(), msg="timed_out", ) @@ -2528,11 +2537,13 @@ def test_dagrun_callbacks_are_called(self, state, expected_callback_msg, dag_mak self.job_runner._do_scheduling(session) expected_callback = DagCallbackRequest( - full_filepath=dag.fileloc, + filepath=dag.fileloc, dag_id=dr.dag_id, is_failure_callback=bool(state == State.FAILED), run_id=dr.run_id, msg=expected_callback_msg, + bundle_name=dag.get_bundle_name(), + bundle_version=dag.get_bundle_version(), ) # Verify dag failure callback request is sent to file processor @@ -2605,11 +2616,13 @@ def test_dagrun_timeout_callbacks_are_stored_in_database(self, dag_maker, sessio ) expected_callback = DagCallbackRequest( - full_filepath=dag.fileloc, + filepath=dag.fileloc, dag_id=dr.dag_id, is_failure_callback=True, run_id=dr.run_id, msg="timed_out", + bundle_name=dag.get_bundle_name(), + bundle_version=dag.get_bundle_version(), ) assert callback == expected_callback @@ -3286,25 +3299,24 @@ def test_verify_integrity_if_dag_changed(self, dag_maker): session.close() @pytest.mark.need_serialized_dag - def test_retry_still_in_executor(self, dag_maker): + def test_retry_still_in_executor(self, dag_maker, session): """ Checks if the scheduler does not put a task in limbo, when a task is retried but is still present in the executor. """ executor = MockExecutor(do_update=False) - with create_session() as session: - with dag_maker( - dag_id="test_retry_still_in_executor", - schedule="@once", - session=session, - ): - dag_task1 = BashOperator( - task_id="test_retry_handling_op", - bash_command="exit 1", - retries=1, - ) - dag_maker.dag_model.calculate_dagrun_date_fields(dag_maker.dag, None) + with dag_maker( + dag_id="test_retry_still_in_executor", + schedule="@once", + session=session, + ): + dag_task1 = BashOperator( + task_id="test_retry_handling_op", + bash_command="exit 1", + retries=1, + ) + dag_maker.dag_model.calculate_dagrun_date_fields(dag_maker.dag, None) @provide_session def do_schedule(session): @@ -5448,6 +5460,8 @@ def test_find_and_purge_zombies(self, session, create_dagrun): dagbag = DagBag(dagfile) dag = dagbag.get_dag("example_branch_operator") DAG.bulk_write_to_db("testing", None, [dag]) + SerializedDagModel.write_dag(dag=dag, bundle_name="testing") + dag_v = DagVersion.get_latest_version(dag.dag_id) data_interval = dag.infer_automated_data_interval(DEFAULT_LOGICAL_DATE) dag_run = create_dagrun( dag, @@ -5467,7 +5481,7 @@ def test_find_and_purge_zombies(self, session, create_dagrun): for task_id in tasks_to_setup: task = dag.get_task(task_id=task_id) - ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) + ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING, dag_version_id=dag_v.id) ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6) ti.queued_by_job_id = 999 @@ -5487,7 +5501,7 @@ def test_find_and_purge_zombies(self, session, create_dagrun): callback_requests = executor.callback_sink.send.call_args.args assert len(callback_requests) == 1 callback_request = callback_requests[0] - assert callback_request.full_filepath == dag.fileloc + assert callback_request.filepath == dag.fileloc assert callback_request.msg == str(self.job_runner._generate_zombie_message_details(ti)) assert callback_request.is_failure_callback is True assert callback_request.ti.dag_id == ti.dag_id @@ -5572,6 +5586,7 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce session.query(Job).delete() dag = dagbag.get_dag("test_example_bash_operator") DAG.bulk_write_to_db("testing", None, [dag]) + SerializedDagModel.write_dag(dag=dag, bundle_name="testing") data_interval = dag.infer_automated_data_interval(DEFAULT_LOGICAL_DATE) dag_run = create_dagrun( dag, @@ -5581,8 +5596,9 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce data_interval=data_interval, ) task = dag.get_task(task_id="run_this_last") - - ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) + ti = TaskInstance( + task, run_id=dag_run.run_id, state=State.RUNNING, dag_version_id=dag_run.dag_version_id + ) ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6) # TODO: If there was an actual Relationship between TI and Job @@ -5599,9 +5615,11 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce expected_failure_callback_requests = [ TaskCallbackRequest( - full_filepath=dag.fileloc, + filepath=dag.fileloc, ti=ti, msg=str(self.job_runner._generate_zombie_message_details(ti)), + bundle_name="testing", + bundle_version=dag_run.bundle_version, ) ] callback_requests = scheduler_job.executor.callback_sink.send.call_args.args diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py index 12bccca74b815..1db997a2565d9 100644 --- a/tests/models/test_dagrun.py +++ b/tests/models/test_dagrun.py @@ -33,7 +33,9 @@ from airflow.decorators import setup, task, task_group, teardown from airflow.models.baseoperator import BaseOperator from airflow.models.dag import DAG, DagModel +from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun, DagRunNote +from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance, TaskInstanceNote, clear_task_instances from airflow.models.taskmap import TaskMap from airflow.models.taskreschedule import TaskReschedule @@ -95,6 +97,7 @@ def create_dag_run( logical_date: datetime.datetime | None = None, is_backfill: bool = False, state: DagRunState = DagRunState.RUNNING, + dag_version: DagVersion | None = None, session: Session, ): now = timezone.utcnow() @@ -118,6 +121,7 @@ def create_dag_run( start_date=now, state=state, external_trigger=False, + dag_version=dag_version or DagVersion.get_latest_version(dag.dag_id, session=session), triggered_by=DagRunTriggeredByType.TEST, ) @@ -476,6 +480,8 @@ def on_success_callable(context): # Scheduler uses Serialized DAG -- so use that instead of the Actual DAG dag = SerializedDAG.from_dict(SerializedDAG.to_dict(dag)) + SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) + session.commit() dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session) @@ -484,10 +490,12 @@ def on_success_callable(context): # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() assert callback == DagCallbackRequest( - full_filepath=dag_run.dag.fileloc, + filepath=dag_run.dag.fileloc, dag_id="test_dagrun_update_state_with_handle_callback_success", run_id=dag_run.run_id, is_failure_callback=False, + bundle_name="testing", + bundle_version=None, msg="success", ) @@ -502,7 +510,8 @@ def on_failure_callable(context): on_failure_callback=on_failure_callable, ) DAG.bulk_write_to_db("testing", None, dags=[dag], session=session) - + SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) + session.commit() dag_task1 = EmptyOperator(task_id="test_state_succeeded1", dag=dag) dag_task2 = EmptyOperator(task_id="test_state_failed2", dag=dag) dag_task1.set_downstream(dag_task2) @@ -522,11 +531,13 @@ def on_failure_callable(context): # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() assert callback == DagCallbackRequest( - full_filepath=dag_run.dag.fileloc, + filepath=dag_run.dag.fileloc, dag_id="test_dagrun_update_state_with_handle_callback_failure", run_id=dag_run.run_id, is_failure_callback=True, msg="task_failure", + bundle_name="testing", + bundle_version=None, ) def test_dagrun_set_state_end_date(self, dag_maker, session): diff --git a/tests_common/pytest_plugin.py b/tests_common/pytest_plugin.py index 66c271113a780..f777b3d6af8af 100644 --- a/tests_common/pytest_plugin.py +++ b/tests_common/pytest_plugin.py @@ -937,7 +937,6 @@ def create_dagrun(self, *, logical_date=None, **kwargs): kwargs.pop("dag_version", None) kwargs.pop("triggered_by", None) kwargs["execution_date"] = logical_date - self.dag_run = dag.create_dagrun(**kwargs) for ti in self.dag_run.task_instances: ti.refresh_from_task(dag.get_task(ti.task_id)) From 11de885b359c03bc72ac34a95b74b5722ee4980d Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Fri, 31 Jan 2025 10:53:02 +0100 Subject: [PATCH 3/7] Resolve conflicts and fix tests --- airflow/dag_processing/manager.py | 4 +--- tests/dag_processing/test_manager.py | 10 +++++----- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py index 22c703c6f3473..9b0c19176c3a0 100644 --- a/airflow/dag_processing/manager.py +++ b/airflow/dag_processing/manager.py @@ -415,10 +415,8 @@ def _add_callback_to_queue(self, request: CallbackRequest): self.log.error("Bundle %s no longer configured, skipping callback", request.bundle_name) return None - dag_absolute_path = os.fspath(Path(bundle.path, request.filepath)) - file_info = DagFileInfo( - path=dag_absolute_path, + rel_path=Path(request.filepath), bundle_path=bundle.path, bundle_name=request.bundle_name, bundle_version=request.bundle_version, diff --git a/tests/dag_processing/test_manager.py b/tests/dag_processing/test_manager.py index 17ecad122a5dd..cfbe8cfc7fe07 100644 --- a/tests/dag_processing/test_manager.py +++ b/tests/dag_processing/test_manager.py @@ -728,7 +728,7 @@ def test_callback_queue(self, mock_logger, configure_testing_dag_bundle): manager._dag_bundles = list(DagBundlesManager().get_all_dag_bundles()) dag1_path = DagFileInfo( - bundle_name="testing", path=f"{tmp_path}/file1.py", bundle_path=Path(tmp_path) + bundle_name="testing", rel_path=Path("file1.py"), bundle_path=Path(tmp_path) ) dag1_req1 = DagCallbackRequest( filepath="file1.py", @@ -750,7 +750,7 @@ def test_callback_queue(self, mock_logger, configure_testing_dag_bundle): ) dag2_path = DagFileInfo( - bundle_name="testing", path=f"{tmp_path}/file2.py", bundle_path=Path(tmp_path) + bundle_name="testing", rel_path=Path("file2.py"), bundle_path=Path(tmp_path) ) dag2_req1 = DagCallbackRequest( filepath="file2.py", @@ -780,7 +780,7 @@ def test_callback_queue(self, mock_logger, configure_testing_dag_bundle): # when manager._add_callback_to_queue(dag1_req2) # Since dag1_req2 is same as dag1_req1, we now have 2 items in file_path_queue - assert manager._file_path_queue == deque([dag2_path, dag1_path]) + assert manager._file_queue == deque([dag2_path, dag1_path]) assert manager._callback_to_execute[dag1_path] == [ dag1_req1, dag1_req2, @@ -794,7 +794,7 @@ def test_callback_queue(self, mock_logger, configure_testing_dag_bundle): assert start.call_args_list == [ mock.call( id=mock.ANY, - path=dag2_path.path, + path=Path(dag2_path.bundle_path, dag2_path.rel_path), bundle_path=dag2_path.bundle_path, callbacks=[dag2_req1], selector=mock.ANY, @@ -802,7 +802,7 @@ def test_callback_queue(self, mock_logger, configure_testing_dag_bundle): ), mock.call( id=mock.ANY, - path=dag1_path.path, + path=Path(dag1_path.bundle_path, dag1_path.rel_path), bundle_path=dag1_path.bundle_path, callbacks=[dag1_req1, dag1_req2], selector=mock.ANY, From 37e3494b25ff1081e3c121a8ad6dcfdf5abe9184 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 3 Feb 2025 14:46:07 +0100 Subject: [PATCH 4/7] fixup! Resolve conflicts and fix tests --- .../celery/executors/test_celery_kubernetes_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py b/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py index 9728b19e9a720..267b0fa78fd87 100644 --- a/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py +++ b/providers/celery/tests/provider_tests/celery/executors/test_celery_kubernetes_executor.py @@ -265,7 +265,7 @@ def test_send_callback(self): filepath="fake", dag_id="fake", run_id="fake", bundle_name="testing", bundle_version=None ) else: - callback = CallbackRequest(full_filepath="fake", dag_id="fake", run_id="fake") + callback = CallbackRequest(full_filepath="fake") cel_k8s_exec.send_callback(callback) cel_k8s_exec.callback_sink.send.assert_called_once_with(callback) From c17b7e929f23e5a27b883974532ae720a71b393e Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 3 Feb 2025 18:13:50 +0100 Subject: [PATCH 5/7] Use relative_fileloc for callbacks --- airflow/jobs/scheduler_job_runner.py | 8 ++++---- airflow/models/dagrun.py | 10 +++++----- tests/jobs/test_scheduler_job.py | 14 +++++++------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index 47ae2d58e161f..df1cb5ab52907 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -858,7 +858,7 @@ def process_executor_events( # too, which would lead to double logging cls.logger().error(msg) request = TaskCallbackRequest( - filepath=ti.dag_version.dag_code.fileloc, + filepath=ti.dag_model.relative_fileloc, bundle_name=ti.dag_version.bundle_name, bundle_version=ti.dag_version.bundle_version, ti=ti, @@ -1634,7 +1634,7 @@ def _schedule_dag_run( dag_model.calculate_dagrun_date_fields(dag, dag.get_run_data_interval(dag_run)) callback_to_execute = DagCallbackRequest( - filepath=dag_model.fileloc, + filepath=dag_model.relative_fileloc, dag_id=dag.dag_id, run_id=dag_run.run_id, bundle_name=dag_model.bundle_name, @@ -2021,9 +2021,9 @@ def _purge_zombies(self, zombies: list[TI], *, session: Session) -> None: for ti in zombies: zombie_message_details = self._generate_zombie_message_details(ti) request = TaskCallbackRequest( - filepath=ti.dag_version.dag_code.fileloc, + filepath=ti.dag_model.relative_fileloc, bundle_name=ti.dag_version.bundle_name, - bundle_version=ti.dag_version.bundle_version, + bundle_version=ti.dag_run.bundle_version, ti=ti, msg=str(zombie_message_details), ) diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index febb8ddb4c721..935a0c05b5e0b 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -922,10 +922,10 @@ def recalculate(self) -> _UnfinishedStates: dag.handle_callback(self, success=False, reason="task_failure", session=session) elif dag.has_on_failure_callback: callback = DagCallbackRequest( - filepath=self.dag_version.dag_code.fileloc, + filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_version.bundle_name, + bundle_name=self.dag_model.bundle_name, bundle_version=self.bundle_version, is_failure_callback=True, msg="task_failure", @@ -951,10 +951,10 @@ def recalculate(self) -> _UnfinishedStates: dag.handle_callback(self, success=True, reason="success", session=session) elif dag.has_on_success_callback: callback = DagCallbackRequest( - filepath=self.dag_version.dag_code.fileloc, + filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_version.bundle_name, + bundle_name=self.dag_model.bundle_name, bundle_version=self.bundle_version, is_failure_callback=False, msg="success", @@ -970,7 +970,7 @@ def recalculate(self) -> _UnfinishedStates: dag.handle_callback(self, success=False, reason="all_tasks_deadlocked", session=session) elif dag.has_on_failure_callback: callback = DagCallbackRequest( - filepath=self.dag_model.fileloc, + filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, bundle_name=self.dag_model.bundle_name, diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index e35820db2c86c..6b0f5847ce7bf 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -403,7 +403,7 @@ def test_process_executor_events_with_callback( # will be set to failed in dag parsing process assert ti1.state == State.QUEUED mock_task_callback.assert_called_once_with( - filepath=dag.fileloc, + filepath=dag.relative_fileloc, ti=mock.ANY, bundle_name="dag_maker", bundle_version=None, @@ -2420,7 +2420,7 @@ def test_dagrun_timeout_verify_max_active_runs(self, dag_maker): assert isinstance(orm_dag.next_dagrun_create_after, datetime.datetime) expected_callback = DagCallbackRequest( - filepath=dr.dag.fileloc, + filepath=dr.dag.relative_fileloc, dag_id=dr.dag_id, is_failure_callback=True, run_id=dr.run_id, @@ -2461,7 +2461,7 @@ def test_dagrun_timeout_fails_run(self, dag_maker): assert dr.state == State.FAILED expected_callback = DagCallbackRequest( - filepath=dr.dag.fileloc, + filepath=dr.dag.relative_fileloc, dag_id=dr.dag_id, is_failure_callback=True, run_id=dr.run_id, @@ -2537,7 +2537,7 @@ def test_dagrun_callbacks_are_called(self, state, expected_callback_msg, dag_mak self.job_runner._do_scheduling(session) expected_callback = DagCallbackRequest( - filepath=dag.fileloc, + filepath=dag.relative_fileloc, dag_id=dr.dag_id, is_failure_callback=bool(state == State.FAILED), run_id=dr.run_id, @@ -2616,7 +2616,7 @@ def test_dagrun_timeout_callbacks_are_stored_in_database(self, dag_maker, sessio ) expected_callback = DagCallbackRequest( - filepath=dag.fileloc, + filepath=dag.relative_fileloc, dag_id=dr.dag_id, is_failure_callback=True, run_id=dr.run_id, @@ -5501,7 +5501,7 @@ def test_find_and_purge_zombies(self, session, create_dagrun): callback_requests = executor.callback_sink.send.call_args.args assert len(callback_requests) == 1 callback_request = callback_requests[0] - assert callback_request.filepath == dag.fileloc + assert callback_request.filepath == dag.relative_fileloc assert callback_request.msg == str(self.job_runner._generate_zombie_message_details(ti)) assert callback_request.is_failure_callback is True assert callback_request.ti.dag_id == ti.dag_id @@ -5615,7 +5615,7 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce expected_failure_callback_requests = [ TaskCallbackRequest( - filepath=dag.fileloc, + filepath=dag.relative_fileloc, ti=ti, msg=str(self.job_runner._generate_zombie_message_details(ti)), bundle_name="testing", From e25c9387de52fc1c48f1cf047f06d26560b3ff5a Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 3 Feb 2025 20:54:31 +0100 Subject: [PATCH 6/7] Fix relative_fileloc issues --- airflow/models/dagrun.py | 6 +++--- airflow/triggers/base.py | 6 +++--- tests/models/test_dagrun.py | 22 ++++++++++++++++++---- 3 files changed, 24 insertions(+), 10 deletions(-) diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 935a0c05b5e0b..7d46c29285fe0 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -925,7 +925,7 @@ def recalculate(self) -> _UnfinishedStates: filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_model.bundle_name, + bundle_name=self.dag_version.bundle_name, bundle_version=self.bundle_version, is_failure_callback=True, msg="task_failure", @@ -954,7 +954,7 @@ def recalculate(self) -> _UnfinishedStates: filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_model.bundle_name, + bundle_name=self.dag_version.bundle_name, bundle_version=self.bundle_version, is_failure_callback=False, msg="success", @@ -973,7 +973,7 @@ def recalculate(self) -> _UnfinishedStates: filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, - bundle_name=self.dag_model.bundle_name, + bundle_name=self.dag_version.bundle_name, bundle_version=self.bundle_version, is_failure_callback=True, msg="all_tasks_deadlocked", diff --git a/airflow/triggers/base.py b/airflow/triggers/base.py index 880b1cd35635b..c01bbf4b1255e 100644 --- a/airflow/triggers/base.py +++ b/airflow/triggers/base.py @@ -217,11 +217,11 @@ def _submit_callback_if_necessary(self, *, task_instance: TaskInstance, session) """Submit a callback request if the task state is SUCCESS or FAILED.""" if self.task_instance_state in (TaskInstanceState.SUCCESS, TaskInstanceState.FAILED): request = TaskCallbackRequest( - filepath=task_instance.dag_model.fileloc, + filepath=task_instance.dag_model.relative_fileloc, ti=task_instance, task_callback_type=self.task_instance_state, - bundle_name=task_instance.dag_model.bundle_name, - bundle_version=task_instance.dag_model.bundle_version, + bundle_name=task_instance.dag_version.bundle_name, + bundle_version=task_instance.dag_run.bundle_version, ) log.info("Sending callback: %s", request) try: diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py index 1db997a2565d9..a967f6b25556f 100644 --- a/tests/models/test_dagrun.py +++ b/tests/models/test_dagrun.py @@ -461,6 +461,7 @@ def test_dagrun_update_state_with_handle_callback_success(self, testing_dag_bund def on_success_callable(context): assert context["dag_run"].dag_id == "test_dagrun_update_state_with_handle_callback_success" + relative_fileloc = "test_dagrun_update_state_with_handle_callback_success.py" dag = DAG( dag_id="test_dagrun_update_state_with_handle_callback_success", schedule=datetime.timedelta(days=1), @@ -468,6 +469,10 @@ def on_success_callable(context): on_success_callback=on_success_callable, ) DAG.bulk_write_to_db("testing", None, dags=[dag], session=session) + dm = DagModel.get_dagmodel(dag.dag_id, session=session) + dm.relative_fileloc = relative_fileloc + session.merge(dm) + session.commit() dag_task1 = EmptyOperator(task_id="test_state_succeeded1", dag=dag) dag_task2 = EmptyOperator(task_id="test_state_succeeded2", dag=dag) @@ -480,17 +485,18 @@ def on_success_callable(context): # Scheduler uses Serialized DAG -- so use that instead of the Actual DAG dag = SerializedDAG.from_dict(SerializedDAG.to_dict(dag)) + dag.relative_fileloc = relative_fileloc SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) session.commit() dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session) + dag_run.dag_model = dm _, callback = dag_run.update_state(execute_callbacks=False) assert dag_run.state == DagRunState.SUCCESS # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() - assert callback == DagCallbackRequest( - filepath=dag_run.dag.fileloc, + filepath=dag_run.dag.relative_fileloc, dag_id="test_dagrun_update_state_with_handle_callback_success", run_id=dag_run.run_id, is_failure_callback=False, @@ -503,6 +509,7 @@ def test_dagrun_update_state_with_handle_callback_failure(self, testing_dag_bund def on_failure_callable(context): assert context["dag_run"].dag_id == "test_dagrun_update_state_with_handle_callback_failure" + relative_fileloc = "test_dagrun_update_state_with_handle_callback_failure.py" dag = DAG( dag_id="test_dagrun_update_state_with_handle_callback_failure", schedule=datetime.timedelta(days=1), @@ -510,8 +517,11 @@ def on_failure_callable(context): on_failure_callback=on_failure_callable, ) DAG.bulk_write_to_db("testing", None, dags=[dag], session=session) - SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) + dm = DagModel.get_dagmodel(dag.dag_id, session=session) + dm.relative_fileloc = relative_fileloc + session.merge(dm) session.commit() + dag_task1 = EmptyOperator(task_id="test_state_succeeded1", dag=dag) dag_task2 = EmptyOperator(task_id="test_state_failed2", dag=dag) dag_task1.set_downstream(dag_task2) @@ -523,15 +533,19 @@ def on_failure_callable(context): # Scheduler uses Serialized DAG -- so use that instead of the Actual DAG dag = SerializedDAG.from_dict(SerializedDAG.to_dict(dag)) + dag.relative_fileloc = relative_fileloc + SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) + session.commit() dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session) + dag_run.dag_model = dm _, callback = dag_run.update_state(execute_callbacks=False) assert dag_run.state == DagRunState.FAILED # Callbacks are not added until handle_callback = False is passed to dag_run.update_state() assert callback == DagCallbackRequest( - filepath=dag_run.dag.fileloc, + filepath=dag.relative_fileloc, dag_id="test_dagrun_update_state_with_handle_callback_failure", run_id=dag_run.run_id, is_failure_callback=True, From 50be7e906b8061ae0a3291cd5b0f2f555944366b Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 3 Feb 2025 22:12:32 +0100 Subject: [PATCH 7/7] fixup! Fix relative_fileloc issues --- airflow/triggers/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/triggers/base.py b/airflow/triggers/base.py index c01bbf4b1255e..d36dd40b6f866 100644 --- a/airflow/triggers/base.py +++ b/airflow/triggers/base.py @@ -220,7 +220,7 @@ def _submit_callback_if_necessary(self, *, task_instance: TaskInstance, session) filepath=task_instance.dag_model.relative_fileloc, ti=task_instance, task_callback_type=self.task_instance_state, - bundle_name=task_instance.dag_version.bundle_name, + bundle_name=task_instance.dag_model.bundle_name, bundle_version=task_instance.dag_run.bundle_version, ) log.info("Sending callback: %s", request)