From 036c9af48c4afb06ff31fa66df9dd74bc24f7c00 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Tue, 8 Feb 2022 21:18:15 +0000 Subject: [PATCH 1/2] Improve speed to run `airflow` by 3x By delaying expensive/slow imports to where they are needed, this gets `airflow` printing it's usage information in under 0.8s, down from almost 3s which makes it feel much much snappier. --- airflow/exceptions.py | 9 ++--- airflow/models/param.py | 8 ++-- airflow/models/taskinstance.py | 12 ++---- airflow/serialization/serialized_objects.py | 44 +++++++++++++++------ airflow/utils/helpers.py | 18 +++++---- airflow/utils/log/file_task_handler.py | 3 +- 6 files changed, 57 insertions(+), 37 deletions(-) diff --git a/airflow/exceptions.py b/airflow/exceptions.py index 650989502c3ef..d77bd61e7ebe2 100644 --- a/airflow/exceptions.py +++ b/airflow/exceptions.py @@ -23,10 +23,6 @@ import warnings from typing import Any, Dict, List, NamedTuple, Optional, Sized -from airflow.api_connexion.exceptions import NotFound as ApiConnexionNotFound -from airflow.utils.code_utils import prepare_code_snippet -from airflow.utils.platform import is_tty - class AirflowException(Exception): """ @@ -44,7 +40,7 @@ class AirflowBadRequest(AirflowException): status_code = 400 -class AirflowNotFoundException(AirflowException, ApiConnexionNotFound): +class AirflowNotFoundException(AirflowException): """Raise when the requested object/resource is not available in the system.""" status_code = 404 @@ -249,6 +245,9 @@ def __init__(self, msg: str, file_path: str, parse_errors: List[FileSyntaxError] self.parse_errors = parse_errors def __str__(self): + from airflow.utils.code_utils import prepare_code_snippet + from airflow.utils.platform import is_tty + result = f"{self.msg}\nFilename: {self.file_path}\n\n" for error_no, parse_error in enumerate(self.parse_errors, 1): diff --git a/airflow/models/param.py b/airflow/models/param.py index ba33b87a6696b..836d1d27f002d 100644 --- a/airflow/models/param.py +++ b/airflow/models/param.py @@ -19,10 +19,6 @@ import warnings from typing import Any, Dict, ItemsView, MutableMapping, Optional, ValuesView -import jsonschema -from jsonschema import FormatChecker -from jsonschema.exceptions import ValidationError - from airflow.exceptions import AirflowException, ParamValidationError from airflow.utils.context import Context from airflow.utils.types import NOTSET, ArgNotSet @@ -61,6 +57,10 @@ def resolve(self, value: Any = NOTSET, suppress_exception: bool = False) -> Any: :param suppress_exception: To raise an exception or not when the validations fails. If true and validations fails, the return value would be None. """ + import jsonschema + from jsonschema import FormatChecker + from jsonschema.exceptions import ValidationError + try: json.dumps(value) except Exception: diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 3532bf1fdc989..4996b9a7db073 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -115,14 +115,6 @@ from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.timeout import timeout -try: - from kubernetes.client.api_client import ApiClient - - from airflow.kubernetes.kube_config import KubeConfig - from airflow.kubernetes.pod_generator import PodGenerator -except ImportError: - ApiClient = None - TR = TaskReschedule _CURRENT_CONTEXT: List[Context] = [] @@ -2032,7 +2024,11 @@ def render_templates(self, context: Optional[Context] = None) -> None: def render_k8s_pod_yaml(self) -> Optional[dict]: """Render k8s pod yaml""" + from kubernetes.client.api_client import ApiClient + + from airflow.kubernetes.kube_config import KubeConfig from airflow.kubernetes.kubernetes_helper_functions import create_pod_id # Circular import + from airflow.kubernetes.pod_generator import PodGenerator kube_config = KubeConfig() pod = PodGenerator.construct_pod( diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 42fa31478e27f..d6abda7c74899 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -48,19 +48,17 @@ from airflow.utils.module_loading import as_importable_string, import_string from airflow.utils.task_group import MappedTaskGroup, TaskGroup -try: - # isort: off - from kubernetes.client import models as k8s - from airflow.kubernetes.pod_generator import PodGenerator - - # isort: on - HAS_KUBERNETES = True -except ImportError: - HAS_KUBERNETES = False - if TYPE_CHECKING: from airflow.ti_deps.deps.base_ti_dep import BaseTIDep + HAS_KUBERNETES: bool + try: + from kubernetes.client import models as k8s + + from airflow.kubernetes.pod_generator import PodGenerator + except ImportError: + pass + log = logging.getLogger(__name__) _OPERATOR_EXTRA_LINKS: Set[str] = { @@ -313,7 +311,7 @@ def _serialize(cls, var: Any) -> Any: # Unfortunately there is no support for r return cls._encode({str(k): cls._serialize(v) for k, v in var.items()}, type_=DAT.DICT) elif isinstance(var, list): return [cls._serialize(v) for v in var] - elif HAS_KUBERNETES and isinstance(var, k8s.V1Pod): + elif _has_kubernetes() and isinstance(var, k8s.V1Pod): json_pod = PodGenerator.serialize_pod(var) return cls._encode(json_pod, type_=DAT.POD) elif isinstance(var, DAG): @@ -374,7 +372,7 @@ def _deserialize(cls, encoded_var: Any) -> Any: elif type_ == DAT.DATETIME: return pendulum.from_timestamp(var) elif type_ == DAT.POD: - if not HAS_KUBERNETES: + if not _has_kubernetes(): raise RuntimeError("Cannot deserialize POD objects without kubernetes libraries installed!") pod = PodGenerator.deserialize_model_dict(var) return pod @@ -1120,3 +1118,25 @@ class DagDependency: def node_id(self): """Node ID for graph rendering""" return f"{self.dependency_type}:{self.source}:{self.target}:{self.dependency_id}" + + +def _has_kubernetes() -> bool: + global HAS_KUBERNETES + if "HAS_KUBERNETES" in globals(): + return HAS_KUBERNETES + + # Loading kube modules is expensive, so delay it until the last moment + + try: + from kubernetes.client import models as k8s + + from airflow.kubernetes.pod_generator import PodGenerator + + globals()['k8s'] = k8s + globals()['PodGenerator'] = PodGenerator + + # isort: on + HAS_KUBERNETES = True + except ImportError: + HAS_KUBERNETES = False + return HAS_KUBERNETES diff --git a/airflow/utils/helpers.py b/airflow/utils/helpers.py index c75a017e2b165..5bdbf0c14f8dd 100644 --- a/airflow/utils/helpers.py +++ b/airflow/utils/helpers.py @@ -37,15 +37,13 @@ ) from urllib import parse -import flask -import jinja2 -import jinja2.nativetypes - from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.utils.module_loading import import_string if TYPE_CHECKING: + import jinja2 + from airflow.models import TaskInstance KEY_REGEX = re.compile(r'^[\w.-]+$') @@ -171,8 +169,10 @@ def as_flattened_list(iterable: Iterable[Iterable[T]]) -> List[T]: return [e for i in iterable for e in i] -def parse_template_string(template_string: str) -> Tuple[Optional[str], Optional[jinja2.Template]]: +def parse_template_string(template_string: str) -> Tuple[Optional[str], Optional["jinja2.Template"]]: """Parses Jinja template string.""" + import jinja2 + if "{{" in template_string: # jinja mode return None, jinja2.Template(template_string) else: @@ -255,6 +255,8 @@ def build_airflow_url_with_query(query: Dict[str, Any]) -> str: For example: 'http://0.0.0.0:8000/base/graph?dag_id=my-task&root=&execution_date=2020-10-27T10%3A59%3A25.615587 """ + import flask + view = conf.get('webserver', 'dag_default_view').lower() url = flask.url_for(f"Airflow.{view}") return f"{url}?{parse.urlencode(query)}" @@ -285,16 +287,18 @@ def render_template(template: Any, context: MutableMapping[str, Any], *, native: except Exception: env.handle_exception() # Rewrite traceback to point to the template. if native: + import jinja2.nativetypes + return jinja2.nativetypes.native_concat(nodes) return "".join(nodes) -def render_template_to_string(template: jinja2.Template, context: MutableMapping[str, Any]) -> str: +def render_template_to_string(template: "jinja2.Template", context: MutableMapping[str, Any]) -> str: """Shorthand to ``render_template(native=False)`` with better typing support.""" return render_template(template, context, native=False) -def render_template_as_native(template: jinja2.Template, context: MutableMapping[str, Any]) -> Any: +def render_template_as_native(template: "jinja2.Template", context: MutableMapping[str, Any]) -> Any: """Shorthand to ``render_template(native=True)`` with better typing support.""" return render_template(template, context, native=True) diff --git a/airflow/utils/log/file_task_handler.py b/airflow/utils/log/file_task_handler.py index c8f79f0c94152..4ad2eb09270cf 100644 --- a/airflow/utils/log/file_task_handler.py +++ b/airflow/utils/log/file_task_handler.py @@ -21,7 +21,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Optional -import httpx from itsdangerous import TimedJSONWebSignatureSerializer from airflow.configuration import AirflowConfigException, conf @@ -159,6 +158,8 @@ def _read(self, ti, try_number, metadata=None): except Exception as f: log += f'*** Unable to fetch logs from worker pod {ti.hostname} ***\n{str(f)}\n\n' else: + import httpx + url = os.path.join("http://{ti.hostname}:{worker_log_server_port}/log", log_relative_path).format( ti=ti, worker_log_server_port=conf.get('logging', 'WORKER_LOG_SERVER_PORT') ) From 4a21d4c6882bac2b5a705c9165768186d4a62263 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Tue, 8 Feb 2022 21:47:08 +0000 Subject: [PATCH 2/2] By not loading BaseExecutor we can get down to <0.5s --- airflow/executors/executor_loader.py | 16 +++++++++------- airflow/utils/cli.py | 3 ++- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/airflow/executors/executor_loader.py b/airflow/executors/executor_loader.py index 01b5f9f98bab6..36b40422afe0c 100644 --- a/airflow/executors/executor_loader.py +++ b/airflow/executors/executor_loader.py @@ -18,10 +18,9 @@ import logging from contextlib import suppress from enum import Enum, unique -from typing import Optional, Tuple, Type +from typing import TYPE_CHECKING, Optional, Tuple, Type from airflow.exceptions import AirflowConfigException -from airflow.executors.base_executor import BaseExecutor from airflow.executors.executor_constants import ( CELERY_EXECUTOR, CELERY_KUBERNETES_EXECUTOR, @@ -35,6 +34,9 @@ log = logging.getLogger(__name__) +if TYPE_CHECKING: + from airflow.executors.base_executor import BaseExecutor + @unique class ConnectorSource(Enum): @@ -48,7 +50,7 @@ class ConnectorSource(Enum): class ExecutorLoader: """Keeps constants for all the currently available executors.""" - _default_executor: Optional[BaseExecutor] = None + _default_executor: Optional["BaseExecutor"] = None executors = { LOCAL_EXECUTOR: 'airflow.executors.local_executor.LocalExecutor', SEQUENTIAL_EXECUTOR: 'airflow.executors.sequential_executor.SequentialExecutor', @@ -60,7 +62,7 @@ class ExecutorLoader: } @classmethod - def get_default_executor(cls) -> BaseExecutor: + def get_default_executor(cls) -> "BaseExecutor": """Creates a new instance of the configured executor if none exists and returns it""" if cls._default_executor is not None: return cls._default_executor @@ -74,7 +76,7 @@ def get_default_executor(cls) -> BaseExecutor: return cls._default_executor @classmethod - def load_executor(cls, executor_name: str) -> BaseExecutor: + def load_executor(cls, executor_name: str) -> "BaseExecutor": """ Loads the executor. @@ -101,7 +103,7 @@ def load_executor(cls, executor_name: str) -> BaseExecutor: return executor_cls() @classmethod - def import_executor_cls(cls, executor_name: str) -> Tuple[Type[BaseExecutor], ConnectorSource]: + def import_executor_cls(cls, executor_name: str) -> Tuple[Type["BaseExecutor"], ConnectorSource]: """ Imports the executor class. @@ -127,7 +129,7 @@ def import_executor_cls(cls, executor_name: str) -> Tuple[Type[BaseExecutor], Co return import_string(executor_name), ConnectorSource.CUSTOM_PATH @classmethod - def __load_celery_kubernetes_executor(cls) -> BaseExecutor: + def __load_celery_kubernetes_executor(cls) -> "BaseExecutor": """:return: an instance of CeleryKubernetesExecutor""" celery_executor = import_string(cls.executors[CELERY_EXECUTOR])() kubernetes_executor = import_string(cls.executors[KUBERNETES_EXECUTOR])() diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py index db67999d24828..a15e1b7f0265e 100644 --- a/airflow/utils/cli.py +++ b/airflow/utils/cli.py @@ -34,7 +34,6 @@ from airflow import settings from airflow.exceptions import AirflowException from airflow.utils import cli_action_loggers -from airflow.utils.db import check_and_run_migrations, synchronize_log_template from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler from airflow.utils.platform import getuser, is_terminal_support_colors from airflow.utils.session import provide_session @@ -93,6 +92,8 @@ def wrapper(*args, **kwargs): try: # Check and run migrations if necessary if check_db: + from airflow.utils.db import check_and_run_migrations, synchronize_log_template + check_and_run_migrations() synchronize_log_template() return f(*args, **kwargs)