Add Python Virtualenv Operator Caching

airflow/decorators/__init__.pyi

@@ -110,6 +110,7 @@ class TaskDecoratorCollection:
         pip_install_options: list[str] | None = None,
         skip_on_exit_code: int | Container[int] | None = None,
         index_urls: None | Collection[str] | str = None,
+        venv_cache_path: None | str = None,
         show_return_value_in_logs: bool = True,
         **kwargs,
     ) -> TaskDecorator:
@@ -134,6 +135,9 @@ class TaskDecoratorCollection:
             exit code will be treated as a failure.
         :param index_urls: an optional list of index urls to load Python packages from.
             If not provided the system pip conf will be used to source packages from.
+        :param venv_cache_path: Optional path to the venv parent folder in which the venv will be cached,
+            creates a sub-folder venv-{hash} whereas hash will be replaced with a checksum of requirements.
+            If not provided the venv will be created and deleted in a temp folder for every execution.
         :param templates_dict: a dictionary where the values are templates that
             will get templated by the Airflow engine sometime between
             ``__init__`` and ``execute`` takes place and are made available

airflow/operators/python.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+import fcntl
 import importlib
 import inspect
 import logging
@@ -46,6 +47,7 @@
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
 from airflow.models.taskinstance import _CURRENT_CONTEXT
+from airflow.utils import hashlib_wrapper
 from airflow.utils.context import Context, context_copy_partial, context_merge
 from airflow.utils.operator_helpers import KeywordParameters
 from airflow.utils.process_utils import execute_in_subprocess
@@ -521,6 +523,9 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
         exit code will be treated as a failure.
     :param index_urls: an optional list of index urls to load Python packages from.
         If not provided the system pip conf will be used to source packages from.
+    :param venv_cache_path: Optional path to the venv parent folder in which the venv will be cached,
+        creates a sub-folder venv-{hash} whereas hash will be replaced with a checksum of requirements.
+        If not provided the venv will be created and deleted in a temp folder for every execution.
     """
 
     template_fields: Sequence[str] = tuple({"requirements"} | set(PythonOperator.template_fields))
@@ -543,6 +548,7 @@ def __init__(
         expect_airflow: bool = True,
         skip_on_exit_code: int | Container[int] | None = None,
         index_urls: None | Collection[str] | str = None,
+        venv_cache_path: None | str = None,
         **kwargs,
     ):
         if (
@@ -572,6 +578,7 @@ def __init__(
             self.index_urls = list(index_urls)
         else:
             self.index_urls = None
+        self.venv_cache_path = Path(venv_cache_path) if venv_cache_path else None
         super().__init__(
             python_callable=python_callable,
             use_dill=use_dill,
@@ -606,7 +613,31 @@ def _prepare_venv(self, venv_path: Path) -> None:
             index_urls=self.index_urls,
         )
 
+    def _ensure_venv_cache_exists(self, venv_cache_path: Path) -> Path:
+        """Helper to ensure a valid venv is set up and will create inplace."""
+        hash_object = hashlib_wrapper.md5(",".join(self._requirements_list()).encode())
+        requirements_hash = hash_object.hexdigest()
+        venv_path = venv_cache_path / f"venv-{requirements_hash[0:8]}"
+        self.log.info("Python Virtualenv will be cached in %s", venv_path)
+        venv_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(f"{venv_path}.lock", "w") as f:
+            # Ensure that cache is not build by parallel workers
+            fcntl.flock(f, fcntl.LOCK_EX)
+
+            if venv_path.exists():
+                self.log.info("Re-using cached Python Virtualenv in %s", venv_path)
+            else:
+                venv_path.mkdir(parents=True)
+                self._prepare_venv(venv_path)
+                self.log.info("New Python Virtualenv created in %s", venv_path)
+            return venv_path
+
     def execute_callable(self):
+        if self.venv_cache_path:
+            venv_path = self._ensure_venv_cache_exists(self.venv_cache_path)
+            python_path = venv_path / "bin" / "python"
+            return self._execute_python_callable_in_subprocess(python_path, venv_path)
+
         with TemporaryDirectory(prefix="venv") as tmp_dir:
             tmp_path = Path(tmp_dir)
             self._prepare_venv(tmp_path)

docs/apache-airflow/howto/operator/python.rst

@@ -126,6 +126,18 @@ To prevent adding secrets to the private repository in your DAG code you can use
 In the special case you want to prevent remote calls for setup of a virtualenv, pass the ``index_urls`` as empty list as ``index_urls=[]`` which
 forced pip installer to use the ``--no-index`` option.
 
+Caching and re-use
+^^^^^^^^^^^^^^^^^^
+
+Setup of virtualenvs is made per task execution in a temporary directory. After execution the virtualenv is deleted again. Ensure that the ``$tmp`` folder
+on your workers have sufficient disk space. Usually (if not configured differently) the local pip cache will be used preventing a re-download of packages
+for each execution.
+
+But still setting up the virtualenv for every execution needs some time. For repeated execution you can set the option ``venv_cache_path`` to a file system
+folder on your worker. In this case the virtualenv will be set up once and be re-used. If venv caching is used, per unique requirements set different
+virtualenv subfolders are created in the cache path. So depending on your variations in the DAGs in your system setup sufficient disk space is needed.
+Note that no automated cleanup is made and in case of cached mode. All worker slots share the same virtualenv.
+
 
 .. _howto/operator:ExternalPythonOperator:
 

tests/operators/test_python.py

@@ -26,6 +26,7 @@
 from collections import namedtuple
 from datetime import date, datetime, timedelta
 from subprocess import CalledProcessError
+from tempfile import TemporaryDirectory
 from typing import Generator
 from unittest import mock
 from unittest.mock import MagicMock
@@ -984,6 +985,16 @@ def f(a):
 
         self.run_as_task(f, index_urls=["https://abc.def.de", "http://xyz.abc.de"], op_args=[4])
 
+    def test_caching(self):
+        def f(a):
+            import sys
+
+            assert "pytest_venv_1234" in sys.executable
+            return a
+
+        with TemporaryDirectory(prefix="pytest_venv_1234") as tmp_dir:
+            self.run_as_task(f, venv_cache_path=tmp_dir, op_args=[4])
+
     # This tests might take longer than default 60 seconds as it is serializing a lot of
     # context using dill (which is slow apparently).
     @pytest.mark.execution_timeout(120)