diff --git a/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 b/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 index cb4b738138f07..860be22bef932 100644 --- a/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 +++ b/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 @@ -38,6 +38,16 @@ if sys.version_info >= (3,6): {# Airflow is not available in this environment, therefore we won't # be able to integrate any plugin macros. #} pass + +try: + from airflow.sdk.execution_time import task_runner +except ModuleNotFoundError: + pass +else: + {#- We are in an Airflow 3.x env, try and set up supervisor comms so virtual env can still access tasks etc! #} + reinit_supervisor_comms = getattr(task_runner, "reinit_supervisor_comms", None) + if reinit_supervisor_comms: + reinit_supervisor_comms() {% endif %} # Script @@ -49,12 +59,10 @@ if sys.version_info >= (3,6): import types {{ modified_dag_module_name }} = types.ModuleType("{{ modified_dag_module_name }}") - {{ modified_dag_module_name }}.{{ python_callable }} = {{ python_callable }} - sys.modules["{{modified_dag_module_name}}"] = {{modified_dag_module_name}} -{% endif%} +{%- endif -%} {% if op_args or op_kwargs %} with open(sys.argv[1], "rb") as file: