You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by di...@apache.org on 2024/12/14 02:05:57 UTC
(airflow) 01/01: separate DAg parsing into a function
This is an automated email from the ASF dual-hosted git repository.
dimberman pushed a commit to branch separate-parsing-to-func
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 0ac5906857a257221c4ed887281f15f0b8378ef5
Author: Daniel Imberman <da...@gmail.com>
AuthorDate: Fri Dec 13 18:05:37 2024 -0800
separate DAg parsing into a function
---
airflow/dag_processing/collection.py | 47 ++++++++++++++++++-
airflow/dag_processing/manager.py | 14 +++---
airflow/dag_processing/processor.py | 91 +++++++++++++++++++++++++++++++++++-
3 files changed, 142 insertions(+), 10 deletions(-)
diff --git a/airflow/dag_processing/collection.py b/airflow/dag_processing/collection.py
index fff1cc64adf..82752974d26 100644
--- a/airflow/dag_processing/collection.py
+++ b/airflow/dag_processing/collection.py
@@ -30,13 +30,18 @@ from __future__ import annotations
import itertools
import logging
import traceback
-from typing import TYPE_CHECKING, NamedTuple
+from datetime import datetime
+from pathlib import Path
+from typing import NamedTuple
+from typing import TYPE_CHECKING
from sqlalchemy import and_, delete, exists, func, select, tuple_
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import joinedload, load_only
from airflow.assets.manager import asset_manager
+from airflow.dag_processing.messaging import DagFileParsingResult, DagFileStat, \
+ CollectionResult
from airflow.models.asset import (
AssetActive,
AssetAliasModel,
@@ -45,12 +50,14 @@ from airflow.models.asset import (
DagScheduleAssetReference,
TaskOutletAssetReference,
)
-from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dag import DAG
+from airflow.models.dag import DagModel, DagOwnerAttributes, DagTag
from airflow.models.dagrun import DagRun
from airflow.models.dagwarning import DagWarningType
from airflow.models.errors import ParseImportError
from airflow.models.trigger import Trigger
from airflow.sdk.definitions.asset import Asset, AssetAlias
+from airflow.stats import Stats
from airflow.triggers.base import BaseTrigger
from airflow.utils.retries import MAX_DB_RETRIES, run_with_db_retries
from airflow.utils.sqlalchemy import with_row_locks
@@ -59,6 +66,7 @@ from airflow.utils.types import DagRunType
if TYPE_CHECKING:
from collections.abc import Collection, Iterable, Iterator
+ from datetime import datetime
from sqlalchemy.orm import Session
from sqlalchemy.sql import Select
@@ -69,6 +77,41 @@ if TYPE_CHECKING:
log = logging.getLogger(__name__)
+def collect_dag_results(
+ run_duration: float,
+ finish_time: datetime,
+ run_count: int,
+ path: str,
+ parsing_result: DagFileParsingResult | None,
+):
+ stat = DagFileStat(
+ last_finish_time=finish_time,
+ last_duration=run_duration,
+ run_count=run_count + 1,
+ )
+
+ file_name = Path(path).stem
+ Stats.timing(f"dag_processing.last_duration.{file_name}", stat.last_duration)
+ Stats.timing("dag_processing.last_duration", stat.last_duration, tags={"file_name": file_name})
+
+ if parsing_result is None:
+ stat.import_errors = 1
+ collected_dags = []
+ import_errors = {}
+ else:
+ # record DAGs and import errors to database
+ collected_dags = parsing_result.serialized_dags or []
+ import_errors = parsing_result.import_errors or {}
+ stat.num_dags = len(parsing_result.serialized_dags)
+ if parsing_result.import_errors:
+ stat.import_errors = len(parsing_result.import_errors)
+ return CollectionResult(
+ stat=stat,
+ collected_dags=collected_dags,
+ import_errors=import_errors,
+ )
+
+
def _create_orm_dags(dags: Iterable[DAG], *, session: Session) -> Iterator[DagModel]:
for dag in dags:
orm_dag = DagModel(dag_id=dag.dag_id)
diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py
index 57c69238a1f..e0f8cc7bb53 100644
--- a/airflow/dag_processing/manager.py
+++ b/airflow/dag_processing/manager.py
@@ -68,7 +68,7 @@ from airflow.utils.process_utils import (
from airflow.utils.retries import retry_db_transaction
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.sqlalchemy import prohibit_commit, with_row_locks
-
+from airflow.dag_processing.collection import collect_dag_results
if TYPE_CHECKING:
from multiprocessing.connection import Connection as MultiprocessingConnection
@@ -905,14 +905,14 @@ class DagFileProcessorManager(LoggingMixin):
last_num_of_db_queries = 0
last_duration = (last_finish_time - processor.start_time).total_seconds()
- stat = DagFileStat(
- num_dags=num_dags,
- import_errors=count_import_errors,
- last_finish_time=last_finish_time,
- last_duration=last_duration,
+ collection_result = collect_dag_results(
+ run_duration=last_duration,
+ finish_time=last_finish_time,
run_count=self._file_stats[processor.file_path].run_count + 1,
- last_num_of_db_queries=last_num_of_db_queries,
+
)
+ stat = collection_result.stat
+ stat.num_of_db_queries = last_num_of_db_queries,
self._file_stats[processor.file_path] = stat
file_name = Path(processor.file_path).stem
"""crude exposure of instrumentation code which may need to be furnished"""
diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py
index a2b44744025..7369aef12d0 100644
--- a/airflow/dag_processing/processor.py
+++ b/airflow/dag_processing/processor.py
@@ -23,7 +23,7 @@ import signal
import threading
import time
import zipfile
-from collections.abc import Generator, Iterable
+from collections.abc import Generator
from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress
from dataclasses import dataclass
from typing import TYPE_CHECKING
@@ -50,8 +50,16 @@ from airflow.utils.log.logging_mixin import LoggingMixin, StreamLogWriter, set_c
from airflow.utils.mixins import MultiprocessingStartMethodMixin
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import TaskInstanceState
+from airflow.dag_processing.messaging import DagFileParseRequest, DagInfo, DagFileParsingResult
+from airflow.serialization.serialized_objects import SerializedDAG
+import traceback
if TYPE_CHECKING:
+ from collections.abc import Iterable
+ from datetime import datetime
+
+ from sqlalchemy.orm import Session
+
import multiprocessing
from datetime import datetime
from multiprocessing.connection import Connection as MultiprocessingConnection
@@ -60,6 +68,7 @@ if TYPE_CHECKING:
from airflow.callbacks.callback_requests import CallbackRequest
from airflow.models.operator import Operator
+ from airflow.utils.context import Context
@dataclass
@@ -84,6 +93,79 @@ def count_queries(session: Session) -> Generator[_QueryCounter, None, None]:
event.remove(session, "do_orm_execute", _count_db_queries)
+def _parse_file(msg: DagFileParseRequest, log):
+ bag = DagBag(
+ dag_folder=msg.file,
+ include_examples=False,
+ safe_mode=True,
+ load_op_links=False,
+ )
+ serialized_dags, serialization_import_errors = serialize_dags(bag, log)
+ bag.import_errors.update(serialization_import_errors)
+ dags = [DagInfo(data=serdag) for serdag in serialized_dags]
+ result = DagFileParsingResult(
+ fileloc=msg.file,
+ serialized_dags=dags,
+ import_errors=bag.import_errors,
+ )
+
+ if msg.callback_requests:
+ _execute_callbacks(bag, msg.callback_requests, log)
+ return result
+
+
+def serialize_dags(bag, log):
+ serialization_import_errors = {}
+ serialized_dags = []
+ for dag in bag.dags.values():
+ try:
+ serialized_dag = SerializedDAG.to_dict(dag)
+ serialized_dags.append(serialized_dag)
+ except Exception:
+ log.exception("Failed to write serialized DAG: %s", dag.fileloc)
+ dagbag_import_error_traceback_depth = conf.getint("core", "dagbag_import_error_traceback_depth")
+ serialization_import_errors[dag.fileloc] = traceback.format_exc(
+ limit=-dagbag_import_error_traceback_depth
+ )
+ return serialized_dags, serialization_import_errors
+
+
+def _execute_callbacks(dagbag: DagBag, callback_requests: list[CallbackRequest], log):
+ for request in callback_requests:
+ log.debug("Processing Callback Request", request=request)
+ if isinstance(request, TaskCallbackRequest):
+ raise NotImplementedError("Haven't coded Task callback yet!")
+ # _execute_task_callbacks(dagbag, request)
+ elif isinstance(request, DagCallbackRequest):
+ _execute_dag_callbacks(dagbag, request, log)
+
+
+def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log):
+ dag = dagbag.dags[request.dag_id]
+
+ callbacks = dag.on_failure_callback if request.is_failure_callback else dag.on_success_callback
+ if not callbacks:
+ log.warning("Callback requested, but dag didn't have any", dag_id=request.dag_id)
+ return
+
+ callbacks = callbacks if isinstance(callbacks, list) else [callbacks]
+ # TODO:We need a proper context object!
+ context: Context = {}
+
+ for callback in callbacks:
+ log.info(
+ "Executing on_%s dag callback",
+ "failure" if request.is_failure_callback else "success",
+ fn=callback,
+ dag_id=request.dag_id,
+ )
+ try:
+ callback(context)
+ except Exception:
+ log.exception("Callback failed", dag_id=request.dag_id)
+ Stats.incr("dag.callback_exceptions", tags={"dag_id": request.dag_id})
+
+
class DagFileProcessorProcess(LoggingMixin, MultiprocessingStartMethodMixin):
"""
Runs DAG processing in a separate process using DagFileProcessor.
@@ -167,6 +249,13 @@ class DagFileProcessorProcess(LoggingMixin, MultiprocessingStartMethodMixin):
log.info("Started process (PID=%s) to work on %s", os.getpid(), file_path)
dag_file_processor = DagFileProcessor(dag_directory=dag_directory, log=log)
+ _parse_file(
+ DagFileParseRequest(
+ file=file_path,
+ callback_requests=callback_requests,
+ ),
+ log=log
+ )
result: tuple[int, int, int] = dag_file_processor.process_file(
file_path=file_path,
callback_requests=callback_requests,