Lazily import many modules to improve import speed

[ashb]

Goal here is to make the healthchecks for triggerer/scheduler quicker by having to do less

Roughly:

• Don't import jsonschema (module or schemas) until we need it - it was surprisingly slow.
• Dont import all models when importing any one -- the only times we need that are when auto-generating migrations, or when checking DB against schema. (That said, accesing any class will cause the SQLA event to import all models, otherwise we often end up with issues due to the relationships on classes)
• Move a few more things to be lazy/delayed import too.

[ashb]

/cc @ianbuss

[potiuk]

I really hope one day https://peps.python.org/pep-0690/ will be in place so that we don't have to do all those any more.

[ashb]

I really hope one day https://peps.python.org/pep-0690/ will be in place so that we don't have to do all those any more.

Yeah, that would be nice.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

[uranusjr]

I recently listened to some of the folks behind PEP 690 discuss lazy imports and came away with the impression that even after it lands Airflow is highly likely not be able to use it. Lazy imports has a lot of implications and PEP 690 quite heavy-handedly make lazy imports the default, and Airflow has so many dependencies I can easily it might never be able to be fully compatible.

[potiuk]

I recently listened to some of the folks behind PEP 690 discuss lazy imports and came away with the impression that even after it lands Airflow is highly likely not be able to use it. Lazy imports has a lot of implications and PEP 690 quite heavy-handedly make lazy imports the default, and Airflow has so many dependencies I can easily it might never be able to be fully compatible.

Yeah. We've probably listened to the same Talk Python episode :). I think it's not even likely to have PEP 690 around in 3.11 or maybe even 3.12. Unless we roll our own sleeves up and help. I think it would be great to actually HELP in PEP 690 by making Airlfow works with lazy imports (maybe by improving PEP 690 implementation as well).

[ashb]

Shall I resume this and get the SQLA model loading working do you think?

[potiuk]

I think it might be worth it. Would be nice to see how much we gain though in those casese

[kaxil]

Shall I resume this and get the SQLA model loading working do you think?

Yup, lets try it out

[ashb]

Wow mypy is confusing/wrong sometimes:

airflow/example_dags/example_latest_only_with_trigger.py:39: error: Unexpected
keyword argument "task_id" for "LatestOnlyOperator"  [call-arg]
        latest_only = LatestOnlyOperator(task_id='latest_only')
                      ^
airflow/utils/log/logging_mixin.py: note: "LatestOnlyOperator" defined here
airflow/example_dags/example_latest_only.py:34: error: Unexpected keyword
argument "task_id" for "LatestOnlyOperator"  [call-arg]
        latest_only = LatestOnlyOperator(task_id='latest_only')
                      ^
airflow/utils/log/logging_mixin.py: note: "LatestOnlyOperator" defined here

Edit: turns out Mypy gets confused by Pep 562 lazy imports (not a great surprise really) -- something was importing BaseOperator from airflow.models, not airflow.models.baseoperator.

[uranusjr]

Does this change much?

[ashb]

It was high up in the list of slow imports!

[uranusjr]

A schema validation library being slow… I guess that’s how third-party libraries work.

[ashb]

I've found an easy way to make sure all models are loaded when needed:

    event.listen(sqlalchemy.orm.mapper, "before_configured", import_all_models, once=True)

This will ensure that on the first time any model class is queried or instantiated all other models will be loaded, which ensures we don't get the "class not found in registry" error message, but still means imports are "quick".

[ashb]

I wonder why this change causes these mypy errors to appear:

tests/system/providers/dbt/cloud/example_dbt_cloud.py:80: error: Unsupported
operand types for >> ("List[object]" and "EmptyOperator")  [operator]
        [get_run_results_artifact, job_run_sensor] >> end
        ^
tests/system/providers/microsoft/azure/example_adf_run_pipeline.py:76: error:
Unsupported operand types for >> ("List[object]" and "EmptyOperator") 
[operator]
        [run_pipeline1, pipeline_run_sensor] >> end
        ^
Found 2 errors in 2 files (checked 2258 source files)

Run mypy for providers.................................................................Failed
- hook id: run-mypy
- exit code: 1

airflow/providers/atlassian/jira/operators/jira.py:70: error: Extra argument
"conf" from **args for "execute" of "JiraOperator"  [misc]
                        resource = self.get_jira_resource_method.execute(*...
                                   ^
airflow/providers/qubole/operators/qubole.py:244: error: Cannot determine type
of "on_failure_callback"  [has-type]
            if self.on_failure_callback is None:
               ^
airflow/providers/qubole/operators/qubole.py:247: error: Cannot determine type
of "on_retry_callback"  [has-type]
            if self.on_retry_callback is None:

[ashb]

I was testing this with python -X importtime -c 'import airflow.jobs.scheduler_job' 2>import-before-schedjob.log and tuna for display:

Before:

After:

[ashb]

I tried adding a airflow/models/__init__.pyi containing this:

# MyPy doesn't/can't recoginze the PEP 562 style lazy imports, so we have to
# tell it about the imports

from .dag import DAG
from .base import ID_LEN
from .xcom import XCOM_RETURN_KEY
from .base import Base
from .baseoperator import BaseOperator
from .baseoperator import BaseOperatorLink
from .connection import Connection
from .dagbag import DagBag
from .dag import DagModel
from .dagpickle import DagPickle
from .dagrun import DagRun
from .dag import DagTag
from .db_callback_request import DbCallbackRequest
from .errors import ImportError
from .log import Log
from .mappedoperator import MappedOperator
from .operator import Operator
from .param import Param
from .pool import Pool
from .renderedtifields import RenderedTaskInstanceFields
from .skipmixin import SkipMixin
from .slamiss import SlaMiss
from .taskfail import TaskFail
from .taskinstance import TaskInstance
from .taskreschedule import TaskReschedule
from .trigger import Trigger
from .variable import Variable
from .xcom import XCom
from .taskinstance import clear_task_instances

def import_all_models(): ...

But I ended up with loads of errors along the lines of:

airflow/www/views.py:109: error: Module "airflow.models" has no attribute "Connection"  [attr-defined]

I wonder what is going on ... Any ideas?

[ashb]

Rebasing to check still good.

[potiuk]

Nice one :)