cannot use `get_current_context()` in a `@task.virtualenv` task

[mziwisky]

Apache Airflow version

Other Airflow 2 version (please specify below)

What happened

On Airflow 2.5.1 (on AWS MWAA), I ran this DAG:

from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.operators.python import get_current_context

@task.virtualenv(system_site_packages=True)
def test():
    data_interval_end = get_current_context()['data_interval_end']
    print(data_interval_end)

@dag(
    start_date=datetime(2023, 9, 6),
    schedule="10 * * * *",
)
def bug_test():
    test()

the_dag = bug_test()

And I got airflow.exceptions.AirflowException: Current context was requested but no context was found! Are you running within an airflow task?

I know that I can do it like this:

@task.virtualenv(system_site_packages=True)
def test(data_interval_end=None):
    print(data_interval_end)

That works fine if I only need the context directly inside that function, but where this actually popped up in practice was a DAG that used some shared lib functions that used get_current_context, which of course works fine when called from normal tasks but blew up when called from a virtualenv task.

What you think should happen instead

ideally, get_current_context() should work even if it's called from a virtualenv task.

How to reproduce

described above

Operating System

Linux? it's AWS MWAA

Versions of Apache Airflow Providers

No response

Deployment

Amazon (AWS) MWAA

Deployment details

No response

Anything else

this was also mentioned in this closed issue: #20974 (comment)

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[potiuk]

I think it would be possible, it's just a matter of serializing the context and adding code in https://github.com/apache/airflow/blob/main/airflow/utils/python_virtualenv_script.jinja2 to save the context to _CONTEXT - - not super difficult task so if anyone would like to pick it up, it is up for grabs. @mziwisky - maybe you would like to contribute it?

[mziwisky]

thanks for the pointer, @potiuk. i'll see if i find some time to take a stab at it

[pedro-cf]

Hello, is there any updates on this? perhaps a workaround meanwhile? Trying to figure out if it's possible to access context or the task instance inside a @task.virtualenv

[Taragolis]

Still good first issue 😉

[potiuk]

Hello, is there any updates on this? perhaps a workaround meanwhile? Trying to figure out if it's possible to access context or the task instance inside a @task.virtualenv

I believe it's easy to get any specific context property - same as any other parameter. You can pass whatever context variables you want via templattng https://airflow.apache.org/docs/apache-airflow/stable/howto/operator/python.html#templating . That would be the first thing I'd check if I were you.

Have you tried it and you had some problems with it, I wonder?

[pedro-cf]

Hello, is there any updates on this? perhaps a workaround meanwhile? Trying to figure out if it's possible to access context or the task instance inside a @task.virtualenv

I believe it's easy to get any specific context property - same as any other parameter. You can pass whatever context variables you want via templattng https://airflow.apache.org/docs/apache-airflow/stable/howto/operator/python.html#templating . That would be the first thing I'd check if I were you.

Have you tried it and you had some problems with it, I wonder?

Could you give me an example of a very simple DAG with one @task.virtualenv , where you would access the context for the task instance and the params inside?