Support persisting the LoadMode.VIRTUALENV directory (#1079)

## Description

Added `virtualenv_dir` as an option to `ExecutionConfig` which is then
propagated downstream to `DbtVirtualenvBaseOperator`.

The following now happens:
- If the flag is set, the operator will attempt to locate the `venv`'s
`python` binary under the provided `virtualenv_dir`.
- If so, it will conclude that the `venv` exists and continues without
creating a new one.
    - If not, it will create a new one at `virtualenv_dir`
- If the flag is not set, simply continue using the temporary directory
solution that was already in place.

## Impact
A very basic test using a local `docker compose` set-up as per the
contribution guide and the
[example_virtualenv](https://github.com/astronomer/astronomer-cosmos/blob/main/dev/dags/example_virtualenv.py)
DAG saw the DAG's runtime go down from **2m31s** to just **32s**. I'd
this improvement to be even more noticeable with more complex graphs and
more python requirements.
 
## Related Issue(s)
Closes: #610 
Partially solves: #1042
Follow up ticket: #1157

## Breaking Change?
None, the flag is optional and is ignored (with a
[warning](https://github.com/astronomer/astronomer-cosmos/compare/main...LennartKloppenburg:astronomer-cosmos:feature/cache-virtualenv?expand=1#diff-61b585fb903927b6868b9626c95e0ec47e3818eb477d795ebd13b0276d4fd76cR125))
when used outside of `VirtualEnv` execution mode.

## Important notice

Most of the changes in this PR were originally implemented in PR #611 by
@LennartKloppenburg. It became stale over the last few months due to
limited maintainer availability. Our sincere apologies to the original
author.

What was accomplished since:
1. Rebased
2. Fixed conflicts
3. Fixed failing tests
4. Introduced new tests

Co-authored-by: Lennart Kloppenburg <[email protected]>

cosmos/config.py

@@ -359,6 +359,8 @@ class ExecutionConfig:
     :param test_indirect_selection: The mode to configure the test behavior when performing indirect selection.
     :param dbt_executable_path: The path to the dbt executable for runtime execution. Defaults to dbt if available on the path.
     :param dbt_project_path: Configures the DBT project location accessible at runtime for dag execution. This is the project path in a docker container for ExecutionMode.DOCKER or ExecutionMode.KUBERNETES. Mutually Exclusive with ProjectConfig.dbt_project_path
+    :param virtualenv_dir: Directory path to locate the (cached) virtual env that
+    should be used for execution when execution mode is set to `ExecutionMode.VIRTUALENV`
     """
 
     execution_mode: ExecutionMode = ExecutionMode.LOCAL
@@ -367,6 +369,8 @@ class ExecutionConfig:
     dbt_executable_path: str | Path = field(default_factory=get_system_dbt)
 
     dbt_project_path: InitVar[str | Path | None] = None
+    virtualenv_dir: str | Path | None = None
+
     project_path: Path | None = field(init=False)
 
     def __post_init__(self, dbt_project_path: str | Path | None) -> None:

cosmos/converter.py

@@ -220,12 +220,26 @@ def __init__(
         validate_initial_user_config(execution_config, profile_config, project_config, render_config, operator_args)
 
         if project_config.dbt_project_path:
-            execution_config, render_config = migrate_to_new_interface(execution_config, project_config, render_config)
+            # We copy the configuration so the change does not affect other DAGs or TaskGroups
+            # that may reuse the same original configuration
+            render_config = copy.deepcopy(render_config)
+            execution_config = copy.deepcopy(execution_config)
+            render_config.project_path = project_config.dbt_project_path
+            execution_config.project_path = project_config.dbt_project_path
 
         validate_adapted_user_config(execution_config, project_config, render_config)
 
-        env_vars = project_config.env_vars or operator_args.get("env")
-        dbt_vars = project_config.dbt_vars or operator_args.get("vars")
+        env_vars = copy.deepcopy(project_config.env_vars or operator_args.get("env"))
+        dbt_vars = copy.deepcopy(project_config.dbt_vars or operator_args.get("vars"))
+
+        if execution_config.execution_mode != ExecutionMode.VIRTUALENV and execution_config.virtualenv_dir is not None:
+            logger.warning(
+                "`ExecutionConfig.virtualenv_dir` is only supported when \
+                ExecutionConfig.execution_mode is set to ExecutionMode.VIRTUALENV."
+            )
+
+        if not operator_args:
+            operator_args = {}
 
         cache_dir = None
         cache_identifier = None
@@ -275,6 +289,8 @@ def __init__(
             task_args,
             execution_mode=execution_config.execution_mode,
         )
+        if execution_config.execution_mode == ExecutionMode.VIRTUALENV and execution_config.virtualenv_dir is not None:
+            task_args["virtualenv_dir"] = execution_config.virtualenv_dir
 
         build_airflow_graph(
             nodes=self.dbt_graph.filtered_nodes,

cosmos/dbt/graph.py

@@ -191,6 +191,15 @@ class DbtGraph:
     Supports different ways of loading the `dbt` project into this representation.
 
     Different loading methods can result in different `nodes` and `filtered_nodes`.
+
+    Example of how to use:
+
+        dbt_graph = DbtGraph(
+            project=ProjectConfig(dbt_project_path=DBT_PROJECT_PATH),
+            render_config=RenderConfig(exclude=["*orders*"], select=[]),
+            dbt_cmd="/usr/local/bin/dbt"
+        )
+        dbt_graph.load(method=LoadMode.DBT_LS, execution_mode=ExecutionMode.LOCAL)
     """
 
     nodes: dict[str, DbtNode] = dict()
@@ -207,6 +216,7 @@ def __init__(
         cache_identifier: str = "",
         dbt_vars: dict[str, str] | None = None,
         airflow_metadata: dict[str, str] | None = None,
+        operator_args: dict[str, Any] | None = None,
     ):
         self.project = project
         self.render_config = render_config
@@ -219,6 +229,7 @@ def __init__(
         else:
             self.dbt_ls_cache_key = ""
         self.dbt_vars = dbt_vars or {}
+        self.operator_args = operator_args or {}
 
     @cached_property
     def env_vars(self) -> dict[str, str]:
@@ -568,7 +579,6 @@ def load_via_dbt_ls_without_cache(self) -> None:
                     self.run_dbt_deps(dbt_cmd, tmpdir_path, env)
 
                 nodes = self.run_dbt_ls(dbt_cmd, self.project_path, tmpdir_path, env)
-
                 self.nodes = nodes
                 self.filtered_nodes = nodes
 

cosmos/operators/virtualenv.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
-from functools import cached_property
+import os
+import shutil
+import time
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
+import psutil
 from airflow.utils.python_virtualenv import prepare_virtualenv
 
+from cosmos import settings
+from cosmos.exceptions import CosmosValueError
 from cosmos.hooks.subprocess import FullOutputSubprocessResult
 from cosmos.log import get_logger
 from cosmos.operators.local import (
@@ -25,10 +30,19 @@
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
+
+PY_INTERPRETER = "python3"
+LOCK_FILENAME = "cosmos_virtualenv.lock"
 logger = get_logger(__name__)
 
 
-PY_INTERPRETER = "python3"
+def depends_on_virtualenv_dir(method: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    def wrapper(operator: DbtVirtualenvBaseOperator, *args: Any) -> Any:
+        if operator.virtualenv_dir is None:
+            raise CosmosValueError(f"Method relies on value of parameter `virtualenv_dir` which is None.")
+        return method(operator, *args)
+
+    return wrapper
 
 
 class DbtVirtualenvBaseOperator(DbtLocalBaseOperator):
@@ -42,71 +56,147 @@ class DbtVirtualenvBaseOperator(DbtLocalBaseOperator):
     :param py_system_site_packages: Whether or not all the Python packages from the Airflow instance will be accessible
            within the virtual environment (if py_requirements argument is specified).
            Avoid using unless the dbt job requires it.
+    :param virtualenv_dir: Directory path where Cosmos will create/update Python virtualenv. If defined, will persist the Python virtualenv in the Airflow worker node.
+    :param is_virtualenv_dir_temporary: Tells Cosmos if virtualenv should be persisted or not.
     """
 
+    template_fields = DbtLocalBaseOperator.template_fields + ("virtualenv_dir", "is_virtualenv_dir_temporary")  # type: ignore[operator]
+
     def __init__(
         self,
         py_requirements: list[str] | None = None,
         pip_install_options: list[str] | None = None,
         py_system_site_packages: bool = False,
+        virtualenv_dir: Path | None = None,
+        is_virtualenv_dir_temporary: bool = False,
         **kwargs: Any,
     ) -> None:
         self.py_requirements = py_requirements or []
         self.pip_install_options = pip_install_options or []
         self.py_system_site_packages = py_system_site_packages
+        self.virtualenv_dir = virtualenv_dir
+        self.is_virtualenv_dir_temporary = is_virtualenv_dir_temporary
+        self.max_retries_lock = settings.virtualenv_max_retries_lock
         super().__init__(**kwargs)
-        self._venv_tmp_dir: None | TemporaryDirectory[str] = None
-
-    @cached_property
-    def venv_dbt_path(
-        self,
-    ) -> str:
-        """
-        Path to the dbt binary within a Python virtualenv.
-
-        The first time this property is called, it creates a virtualenv and installs the dependencies based on the
-        self.py_requirements, self.pip_install_options, and self.py_system_site_packages. This value is cached for future calls.
-        """
-        # We are reusing the virtualenv directory for all subprocess calls within this task/operator.
-        # For this reason, we are not using contexts at this point.
-        # The deletion of this directory is done explicitly at the end of the `execute` method.
-        self._venv_tmp_dir = TemporaryDirectory(prefix="cosmos-venv")
-        py_interpreter = prepare_virtualenv(
-            venv_directory=self._venv_tmp_dir.name,
-            python_bin=PY_INTERPRETER,
-            system_site_packages=self.py_system_site_packages,
-            requirements=self.py_requirements,
-            pip_install_options=self.pip_install_options,
-        )
-        dbt_binary = Path(py_interpreter).parent / "dbt"
-        cmd_output = self.subprocess_hook.run_command(
-            [
-                py_interpreter,
-                "-c",
-                "from importlib.metadata import version; print(version('dbt-core'))",
-            ]
-        )
-        dbt_version = cmd_output.output
-        self.log.info("Using dbt version %s available at %s", dbt_version, dbt_binary)
-        return str(dbt_binary)
+        if not self.py_requirements:
+            self.log.error("Cosmos virtualenv operators require the `py_requirements` parameter")
 
     def run_subprocess(self, command: list[str], env: dict[str, str], cwd: str) -> FullOutputSubprocessResult:
-        if self.py_requirements:
-            command[0] = self.venv_dbt_path
-
-        subprocess_result: FullOutputSubprocessResult = self.subprocess_hook.run_command(
+        # No virtualenv_dir set, so create a temporary virtualenv
+        if self.virtualenv_dir is None or self.is_virtualenv_dir_temporary:
+            self.log.info("Creating temporary virtualenv")
+            with TemporaryDirectory(prefix="cosmos-venv") as tempdir:
+                self.virtualenv_dir = Path(tempdir)
+                py_bin = self._prepare_virtualenv()
+                dbt_bin = str(Path(py_bin).parent / "dbt")
+                command[0] = dbt_bin  # type: ignore
+                subprocess_result: FullOutputSubprocessResult = self.subprocess_hook.run_command(
+                    command=command,
+                    env=env,
+                    cwd=cwd,
+                    output_encoding=self.output_encoding,
+                )
+                return subprocess_result
+
+        # Use a reusable virtualenv
+        self.log.info(f"Checking if the virtualenv lock {str(self._lock_file)} exists")
+        while not self._is_lock_available() and self.max_retries_lock:
+            logger.info("Waiting for virtualenv lock to be released")
+            time.sleep(1)
+            self.max_retries_lock -= 1
+
+        self.log.info(f"Acquiring the virtualenv lock")
+        self._acquire_venv_lock()
+        py_bin = self._prepare_virtualenv()
+        dbt_bin = str(Path(py_bin).parent / "dbt")
+        command[0] = dbt_bin  # type: ignore
+        subprocess_result = self.subprocess_hook.run_command(
             command=command,
             env=env,
             cwd=cwd,
             output_encoding=self.output_encoding,
         )
+        self.log.info("Releasing virtualenv lock")
+        self._release_venv_lock()
         return subprocess_result
 
+    def clean_dir_if_temporary(self) -> None:
+        """
+        Delete the virtualenv directory if it is temporary.
+        """
+        if self.is_virtualenv_dir_temporary and self.virtualenv_dir and self.virtualenv_dir.exists():
+            self.log.info(f"Deleting the Python virtualenv {self.virtualenv_dir}")
+            shutil.rmtree(str(self.virtualenv_dir), ignore_errors=True)
+
     def execute(self, context: Context) -> None:
-        output = super().execute(context)
-        if self._venv_tmp_dir:
-            self._venv_tmp_dir.cleanup()
-        self.log.info(output)
+        try:
+            output = super().execute(context)
+            self.log.info(output)
+        finally:
+            self.clean_dir_if_temporary()
+
+    def on_kill(self) -> None:
+        self.clean_dir_if_temporary()
+
+    def _prepare_virtualenv(self) -> str:
+        self.log.info(f"Creating or updating the virtualenv at `{self.virtualenv_dir}")
+        py_bin = prepare_virtualenv(
+            venv_directory=str(self.virtualenv_dir),
+            python_bin=PY_INTERPRETER,
+            system_site_packages=self.py_system_site_packages,
+            requirements=self.py_requirements,
+            pip_install_options=self.pip_install_options,
+        )
+        return py_bin
+
+    @property
+    def _lock_file(self) -> Path:
+        filepath = Path(f"{self.virtualenv_dir}/{LOCK_FILENAME}")
+        return filepath
+
+    @property
+    def _pid(self) -> int:
+        return os.getpid()
+
+    @depends_on_virtualenv_dir
+    def _is_lock_available(self) -> bool:
+        is_available = True
+        if self._lock_file.is_file():
+            with open(self._lock_file) as lf:
+                pid = int(lf.read())
+                self.log.info(f"Checking for running process with PID {pid}")
+                try:
+                    _process_running = psutil.Process(pid).is_running()
+                    self.log.info(f"Process {pid} running: {_process_running} and has the lock {self._lock_file}.")
+                except psutil.NoSuchProcess:
+                    self.log.info(f"Process {pid} is not running. Lock {self._lock_file} was outdated.")
+                    is_available = True
+                else:
+                    is_available = not _process_running
+        return is_available
+
+    @depends_on_virtualenv_dir
+    def _acquire_venv_lock(self) -> None:
+        if not self.virtualenv_dir.is_dir():  # type: ignore
+            os.mkdir(str(self.virtualenv_dir))
+
+        with open(self._lock_file, "w") as lf:
+            logger.info(f"Acquiring lock at {self._lock_file} with pid {str(self._pid)}")
+            lf.write(str(self._pid))
+
+    @depends_on_virtualenv_dir
+    def _release_venv_lock(self) -> None:
+        if not self._lock_file.is_file():
+            logger.warning(f"Lockfile {self._lock_file} not found, perhaps deleted by other concurrent operator?")
+            return
+
+        with open(self._lock_file) as lf:
+            lock_file_pid = int(lf.read())
+
+            if lock_file_pid == self._pid:
+                return self._lock_file.unlink()
+
+            logger.warning(f"Lockfile owned by process of pid {lock_file_pid}, while operator has pid {self._pid}")
 
 
 class DbtBuildVirtualenvOperator(DbtVirtualenvBaseOperator, DbtBuildLocalOperator):  # type: ignore[misc]

cosmos/settings.py

@@ -27,6 +27,7 @@
 dbt_docs_index_file_name = conf.get("cosmos", "dbt_docs_index_file_name", fallback="index.html")
 enable_cache_profile = conf.getboolean("cosmos", "enable_cache_profile", fallback=True)
 dbt_profile_cache_dir_name = conf.get("cosmos", "profile_cache_dir_name", fallback="profile")
+virtualenv_max_retries_lock = conf.getint("cosmos", "virtualenv_max_retries_lock", fallback=120)
 
 # Experimentally adding `remote_cache_dir` as a separate entity in the Cosmos 1.6 release to gather feedback.
 # This will be merged with the `cache_dir` config parameter in upcoming releases.