Add support for virtual env directory flag

[LennartKloppenburg]

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 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)

Implements #610

Breaking Change?

None, the flag is optional and is ignored (with a warning) when used outside of VirtualEnv execution mode.

Checklist

• I have made corresponding changes to the documentation (if required)
• I have added tests that prove my fix is effective or that my feature works

[netlify]

👷 Deploy Preview for amazing-pothos-a3bca0 processing.

Name	Link
🔨 Latest commit	be0de1a
🔍 Latest deploy log	https://app.netlify.com/sites/amazing-pothos-a3bca0/deploys/6582c61c78d17900084dc3d7

[codecov]

Codecov Report

Attention: Patch coverage is 83.67347% with 8 lines in your changes are missing coverage. Please review.

Project coverage is 93.06%. Comparing base (090116e) to head (87c5da0).
Report is 109 commits behind head on main.

❗ Current head 87c5da0 differs from pull request most recent head be0de1a

Please upload reports for the commit be0de1a to get more accurate results.

Files	Patch %	Lines
cosmos/operators/virtualenv.py	81.39%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #611      +/-   ##
==========================================
- Coverage   93.28%   93.06%   -0.23%     
==========================================
  Files          55       54       -1     
  Lines        2502     2163     -339     
==========================================
- Hits         2334     2013     -321     
+ Misses        168      150      -18     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[tatiana]

Thank you very much for creating this PR so quickly, @LennartKloppenburg ! This is looking very good.

I added some comments in line, and I have a gut feeling we may need to add some additional tests to cover the possible behaviours in _get_or_create_venv_py_interpreter.

We can aim to release this change as part of 1.2.1 (if we consider it a bugfix) or 1.3 (if we consider it a new feature) 🎉

[tatiana]

In this case, does it still make sense to install any potential dependencies/update them - if there were requirement changes?

[LennartKloppenburg]

Yeah ideally we'd be able to clean up the virtual env after the DAG run, but for the reasons you mentioned before this can be tricky.
One way to "perhaps" invalidate the virtualenv is to check when it was created and, after say 24 hours or 48 or so, have this operator clean it up and recreate it?

[tatiana]

The time-based approach could lead to some strange scenarios and be tricky to troubleshoot. How feasible would be for us to run a pip install in an existing virtualenv? It should be very quick if it was previously setup, and it would make the operator reliable.

Regarding the cleanup - I know - ideally we'd be able to set the venv only once during the DAG setup and delete during tear down. Unfortunately - to my knowledge - even the latest Airflow (2.7) does not allow us to have a setup/tear down per worker node during the lifecycle of a DAG. But this can be an improvement for the future - in a separate PR!

[LennartKloppenburg]

Sorry for the late response here!

so the underlying prepare_virtualenv that we are "avoiding" after determining it's already there is imported from Airflow core (airflow.utils.python_virtualenv). That little helper also takes into account the python requirements so if we bypass this helper, we can't inject requirements unless we repeat the logic over here:

    ...
    pip_cmd = None
    if requirements is not None and len(requirements) != 0:
        pip_cmd = _generate_pip_install_cmd_from_list(venv_directory, requirements, pip_install_options)
    if requirements_file_path is not None and requirements_file_path:
        pip_cmd = _generate_pip_install_cmd_from_file(
            venv_directory, requirements_file_path, pip_install_options
        )

    if pip_cmd:
        execute_in_subprocess(pip_cmd)
    ...

What do you think?

[tatiana]

In our case, we probably would only need to do part of the logic:

    if requirements is not None and len(requirements) != 0:
        pip_cmd = _generate_pip_install_cmd_from_list(venv_directory, requirements, pip_install_options)

Since we don't support requirements_file_path.

If we don't want to add this call unnecessarily, we'd probably need a pip freeze call - to confirm if the desired dependencies are already installed, which may be more work.

We probably need one or both of these. Otherwise, we're at the risk of an Airflow worker having partial/outdated dependencies that are incompatible with the dependencies the user requested.

I'm in favour of us caching for performance reasons, but we still should aim to have the task being idempotent.

[tatiana]

A general thought: do we still want to cache this property? Is there any risk that we could end up caching the incorrect path?

[LennartKloppenburg]

How is this property cached? If people are debugging or want to pass in more dynamically configured directories, I don't know how this decorator behaves :) Is it per task_id per dagrun_id or is it more persistent?

[tatiana]

The property is cached while the Python process is alive.

[LennartKloppenburg]

@tatiana I've updated the PR with some changes you've requested :)

One lingering issue:
When there's no virtual env handy and multiple tasks are executed concurrently, they will all attempt the same checks (checking for the dir, for the virtual env, to install stuff etc.) which leads to all kinds of issues because they don't automatically "wait" for the venv to be created by one single operator. This occurs in environments that schedule multiple DBT tasks simultaneously -- which is obviously very common.

The issue can then be resolved by retrying these tasks with some retry_delay, which will achieve the same result as "waiting" for the virtual env to be provisioned would have done.
What do you think of this?

[LennartKloppenburg]

@tatiana Just completed the rebase, saw some artifacts that trip up the tests, will look at those tomorrow :) !

[LennartKloppenburg]

@tatiana
I could use two more eyes on this (from the CI/CD tests):

FAILED tests/dbt/test_graph.py::test_load_via_dbt_ls_project_config_env_vars - cosmos.dbt.graph.CosmosLoadDbtException: Unable to find the dbt executable: dbt
FAILED tests/dbt/test_graph.py::test_load_via_dbt_ls_project_config_dbt_vars - cosmos.dbt.graph.CosmosLoadDbtException: Unable to find the dbt executable: dbt
FAILED tests/dbt/test_graph.py::test_load_via_dbt_ls_render_config_selector_arg_is_used - cosmos.dbt.graph.CosmosLoadDbtException: Unable to find the dbt executable: dbt

When I run the tests locally they pass, maybe I missed something while rebasing? I rebased so much that I no longer know where it was introduced :D
Thanks in advance!!

[tatiana]

Hi @LennartKloppenburg ! I'm sorry for the massive delay, I've been working on other projects and it has been hard to keep up with everything. I'm planning to get back to this PR next week, so we can try to release it as part of Cosmos 1.5

[tatiana]

Hi @LennartKloppenburg ! I'm very sorry for the very long delay.
I solved all the conflicts, and all the tests seem to be passing in #1079 - which is a copy of your PR, with the additional changes:

1. Rebase
2. Resolve conflicts
3. Fix tests

If you are happy with the proposed changes, please feel free to incorporate them into your PR.
I'd like us to merge this to have alpha versions and validate this change before the 1.6 release.

[tatiana]

We took this to completion in #1079, giving the credits to @LennartKloppenburg and this original PR