Extremely slow import times on Python 3.12

[potiuk]

We just enabled Python 3.12 support for Apache Airflow and we started to experienced extremely slow import times for example dags referring to Databricks

This has been previously reported here: #292 as occuring on Python 3.12 and closed after the authors switched to earlier Python versions, but the underlying issue had not been investigated and fixed and it occurs and it is very apparent in Airflow CI

Example failure that we experience (https://github.com/apache/airflow/actions/runs/8246010879/job/22551487439#step:6:15264)

tests/always/test_example_dags.py:121: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
airflow/models/dagbag.py:147: in __init__
    self.collect_dags(
airflow/models/dagbag.py:554: in collect_dags
    found_dags = self.process_file(filepath, only_if_updated=only_if_updated, safe_mode=safe_mode)
airflow/models/dagbag.py:313: in process_file
    mods = self._load_modules_from_file(filepath, safe_mode)
airflow/models/dagbag.py:376: in _load_modules_from_file
    return parse(mod_name, filepath)
airflow/models/dagbag.py:346: in parse
    loader.exec_module(new_module)
<frozen importlib._bootstrap_external>:995: in exec_module
    ???
<frozen importlib._bootstrap>:488: in _call_with_frames_removed
    ???
tests/system/providers/databricks/example_databricks_sensors.py:25: in <module>
    from airflow.providers.databricks.sensors.databricks_partition import DatabricksPartitionSensor
airflow/providers/databricks/sensors/databricks_partition.py:27: in <module>
    from databricks.sql.utils import ParamEscaper
/usr/local/lib/python3.12/site-packages/databricks/sql/utils.py:14: in <module>
    from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
/usr/local/lib/python3.12/site-packages/databricks/sql/cloudfetch/download_manager.py:7: in <module>
    from databricks.sql.cloudfetch.downloader import (
/usr/local/lib/python3.12/site-packages/databricks/sql/cloudfetch/downloader.py:9: in <module>
    from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
/usr/local/lib/python3.12/site-packages/databricks/sql/thrift_api/TCLIService/ttypes.py:32802: in <module>
    all_structs.append(TExecuteStatementResp)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <airflow.utils.timeout.TimeoutPosix object at 0x7fdfd433bf80>
signum = 14
frame = <frame at 0x7fdfd3ff0260, file '/usr/local/lib/python3.12/site-packages/databricks/sql/thrift_api/TCLIService/ttypes.py', line 32802, code <module>>

    def handle_timeout(self, signum, frame):
        """Log information and raises AirflowTaskTimeout."""
        self.log.error("Process timed out, PID: %s", str(os.getpid()))
>       raise AirflowTaskTimeout(self.error_message)
E       airflow.exceptions.AirflowTaskTimeout: DagBag import timeout for /opt/airflow/tests/system/providers/databricks/example_databricks_sensors.py after 30.0s.
E       Please take a look at these docs to improve your DAG import time:
E       * http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/stable/best-practices.html#top-level-python-code
E       * http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/stable/best-practices.html#reducing-dag-complexity, PID: 81

This is not a blocker for Airflow - we will simply exclude Python 3.12 for databricks provider - so databricks provider will not be available for Python 3.12. But it would be great to fix it eventualy. Once the problem is fixed, we should be able to test it by re-enabling databricks for Python 3.12 and running a number of CI builds for Airflow "canary" buillds where the issue is apparent.

[susodapop]

Analysis

After eluding me for several months I've finally snaked out the issue. The good news is that it's not an issue with databricks-sql-connector. The bad news is the only workaround at present is to run your tests without the --cov flag.

tl;dr There is a bug in the coverage library when running on Python 3.12. nedbat/coveragepy#1665 To work around this, you have to remove the --cov from your pytest invocation.

I created a minimum reproducible example.

1. From a fresh virtual environment run pip install databricks-sql-connector pytest pytest-cov
2. Create a file called test_file.py and give it these contents:

# test_file.py

def test_downloader_import_time():
    start = time.time()
    from databricks.sql.cloudfetch import downloader
    end = time.time()
    difference = end - start
    assert difference < 1, f"It took {difference} seconds to import the downloader"

3. Invoke pytest without coverage enabled and the test will pass:

$ python -m pytest test_file.py
================ 1 passed in 0.07s ===============

4. Invoice pytest --cov and the test will fail after about 200 seconds.

$ python -m pytest test_file.py --cov

FAILED test_file.py::test_downloader_import_time - AssertionError: It took 207.19054412841797 seconds to import the downloader

Workarounds

From what I can tell, there is no way to avoid this completely until coverage is updated to support PEP-669, which is currently only experimentally supported. I can make the downloader package import more quickly by pushing the TSparkArrowResultLink import into a if typing.TYPE_CHECKING conditional. But the connector simply cannot run without importing ttypes.py eventually, so this only pushes the long import time further into the package initialization process.

cc: @benc-db

[potiuk]

Oh wow! Fantastic find. That would also explain some of the other instabilities we observed with Python.12. Yep. It will be easy to disable --coverage for Python 3.12 for us!

[potiuk]

PR here: apache/airflow#38194 - once I see it working I will likely be able to remove databricks provider exclusion for 3.12 :) . We are already using the version of coverage that should have PEP 669 support so I just enabled it for coverage tests and 🤞

[susodapop]

We are already using the version of coverage that should have PEP 669 support so I just enabled it for coverage tests and 🤞

Very curious to see how this goes for you. I tried to using a newer version of coverage for my reproduction above but didn't see any noticeable difference (perhaps 5% faster).

[potiuk]

I will let you know :)

[potiuk]

Generally speaking - what we saw in 3.12 builds WITH coverage enabled was that sometimes it took way longer to complete the tests (30%-50% slower) - but not always. That was really puzzling that I saw it in our canary builds (main) and did not see it in regular PRs - and that was what puzzled me.

However, the coverage theory fits it perfectly - because our PRs are running usually a subset of tests, those that are relevant to the change coming in the PR, so we do not run coverage there - we only use coverage in the canary (main) builds that are running full suite of tests.

It happened frequently enough (multiple times a day in canary builds) to see the result of that change rather quickly.

[potiuk]

so far,so good

[potiuk]

Looks good so far - I have not seen side effects observed in main for Python 3.12 yet - after several builds. Enabling databricks now back apache/airflow#38207 and we will know for sure once we merge, because it has been failing almost always before.

[potiuk]

Eventually I just disabled coverage for Python 3.12. The tests with coverage on Python 3.12 took a long time and some of them even timed out inside coverage's sysmon.