"execute cannot be called outside TaskInstance" warning fired despite being in a TaskInstance

[SKalide]

Apache Airflow version

2.9.3

If "Other Airflow 2 version" selected, which one?

No response

What happened?

BigQueryTableExistenceSensor fire an "execute cannot be called outside TaskInstance" warning on execute

My observation: The ExecutorSafeguard is being called twice for a single sensor execution

Reason is, that the first encounter is over the TaskInstance _execute_callable (as intended). Now the decorater checks, if it is executed within a Taskinstance. After that, it will call the "real" execute function, which is also decorated and therefor will do the check again, but without a sentinel, which leads to a failed check and fires a warning (or raise a exception if allow_nested_operators is set).

See also: #41426

And there was also already opened an issue with the same problem but not solved, see #39413

What you think should happen instead?

Decorater should not fire a warning. Maybe it should be checked, if the decorater was already called and therefor did already checked it.

How to reproduce

from airflow import DAG
from airflow.providers.google.cloud.sensors.bigquery import BigQueryTableExistenceSensor
from airflow.utils.dates import days_ago
from airflow.models.baseoperator import BaseOperator


class LoggingBaseOperator(BaseOperator):
    def execute(self, context):
        self.log.info(f"Executing {self.__class__.__name__}")
        return super().execute(context)


class LoggingBigQueryTableExistenceSensor(
    BigQueryTableExistenceSensor, LoggingBaseOperator
):
    def poke(self, context):
        self.log.info(f"Poking {self.__class__.__name__}")
        return True  


dag = DAG(
    "test_bigquery_sensor_double_execution",
    default_args={
        "start_date": days_ago(1),
    },
    description="A simple DAG to test BigQueryTableExistenceSensor execution",
    schedule_interval=None,
)

sensor_task = LoggingBigQueryTableExistenceSensor(
    task_id="test_sensor",
    project_id="your-project-id",
    dataset_id="your-dataset-id",
    table_id="your-table-id",
    poke_interval=60,
    dag=dag,
)

log_bg.txt

Operating System

Ubuntu 22.04.4 LTS

Versions of Apache Airflow Providers

No response

Deployment

Docker-Compose

Deployment details

Docker image: apache/airflow:2.9.3-python3.11

Anything else?

No response

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

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[SKalide]

I've looked into this issue more deeply and I believe I can provide a solution. I could use some guidance on whether I should submit a PR for this.

The problem, in my opinion, is that after _execute_callable is called, it enters the ExecutorSafeguard wrapper. There, it checks if it's called inside a TaskInstance (hence the _execute_callable) and then returns the execution function itself. Because the execute method of BigQueryTableExistenceSensor is also decorated, it automatically re-enters the ExecutorSafeguard wrapper, but this time it appears as if it was called outside a Task Instance, causing the check to fail. With allow_nested_operators, it only fires a warning instead of failing.

My proposed solution is to check if the ExecutorSafeguard has already been called in the current execution context. If so, we can skip the check for nested calls. This can be achieved by setting a flag using thread-local storage to ensure thread safety.

class ExecutorSafeguard:
    """
    The ExecutorSafeguard decorator.

    Checks if the execute method of an operator isn't manually called outside
    the TaskInstance as we want to avoid bad mixing between decorated and
    classic operators.
    """

    test_mode = conf.getboolean("core", "unit_test_mode")
    _local = local()

    @classmethod
    def decorator(cls, func):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if (
                getattr(ExecutorSafeguard._local, "in_executor_safeguard", False)
                and self.allow_nested_operators
            ):
                # If already in ExecutorSafeguard, call execution function - recursive call
                return func(self, *args, **kwargs)
            ExecutorSafeguard._local.in_executor_safeguard = True

            try:
                from airflow.decorators.base import DecoratedOperator

                sentinel = kwargs.pop(f"{self.__class__.__name__}__sentinel", None)

                if (
                    not cls.test_mode
                    and sentinel != _sentinel
                    and not isinstance(self, DecoratedOperator)
                ):
                    message = (
                        f"{self.__class__.__name__}.{func.__name__} cannot be called outside TaskInstance!"
                    )
                    raise AirflowException(message)
                return func(self, *args, **kwargs)
            finally:
                ExecutorSafeguard._local.in_executor_safeguard = False

        return wrapper

This solution would suppress the warning for legitimate nested calls while still protecting against unintended external calls. The use of local() ensures that this check is thread-safe.

I'm willing to create a Pull Request with this solution if the maintainers think this approach is appropriate. What do you think about this solution? Are there any aspects I should consider or modify before proceeding with a PR?

[zachliu]

@SKalide i think you can simplify your example:

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.operators.bash import BashOperator


class MyBashOperator(BashOperator):
    def execute(self, context):
        self.log.info(f"Executing {self.__class__.__name__}")
        return super().execute(context)


dag = DAG(
    "test_bash_execution",
    default_args={
        "start_date": days_ago(1),
    },
    description="A simple DAG to test MyBashOperator execution",
    schedule_interval=None,
)

task = MyBashOperator(
    task_id="my_bash_operator",
    bash_command="ls -lah",
    dag=dag,
)

this produces the same warning message
by using this simpler example, i believe it's more convincing that this issue is "serious" 😁

[dabla]

@SKalide i think you can simplify your example:

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.operators.bash import BashOperator


class MyBashOperator(BashOperator):
    def execute(self, context):
        self.log.info(f"Executing {self.__class__.__name__}")
        return super().execute(context)


dag = DAG(
    "test_bash_execution",
    default_args={
        "start_date": days_ago(1),
    },
    description="A simple DAG to test MyBashOperator execution",
    schedule_interval=None,
)

task = MyBashOperator(
    task_id="my_bash_operator",
    bash_command="ls -lah",
    dag=dag,
)

this produces the same warning message by using this simpler example, i believe it's more convincing that this issue is "serious" 😁

I think I understand the issue of the warning and I think some refactoring will be needed to avoid the warning in the case you extend an existing operator, as that should still be allowed. Will have to create a test case for it and then check how we can prevent the warning in that case. The warning is there to prevent people from calling other operators from python operators, but in you example its not the case as it's inheritance.

[dabla]

I've looked into this issue more deeply and I believe I can provide a solution. I could use some guidance on whether I should submit a PR for this.

The problem, in my opinion, is that after _execute_callable is called, it enters the ExecutorSafeguard wrapper. There, it checks if it's called inside a TaskInstance (hence the _execute_callable) and then returns the execution function itself. Because the execute method of BigQueryTableExistenceSensor is also decorated, it automatically re-enters the ExecutorSafeguard wrapper, but this time it appears as if it was called outside a Task Instance, causing the check to fail. With allow_nested_operators, it only fires a warning instead of failing.

My proposed solution is to check if the ExecutorSafeguard has already been called in the current execution context. If so, we can skip the check for nested calls. This can be achieved by setting a flag using thread-local storage to ensure thread safety.

class ExecutorSafeguard:
    """
    The ExecutorSafeguard decorator.

    Checks if the execute method of an operator isn't manually called outside
    the TaskInstance as we want to avoid bad mixing between decorated and
    classic operators.
    """

    test_mode = conf.getboolean("core", "unit_test_mode")
    _local = local()

    @classmethod
    def decorator(cls, func):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if (
                getattr(ExecutorSafeguard._local, "in_executor_safeguard", False)
                and self.allow_nested_operators
            ):
                # If already in ExecutorSafeguard, call execution function - recursive call
                return func(self, *args, **kwargs)
            ExecutorSafeguard._local.in_executor_safeguard = True

            try:
                from airflow.decorators.base import DecoratedOperator

                sentinel = kwargs.pop(f"{self.__class__.__name__}__sentinel", None)

                if (
                    not cls.test_mode
                    and sentinel != _sentinel
                    and not isinstance(self, DecoratedOperator)
                ):
                    message = (
                        f"{self.__class__.__name__}.{func.__name__} cannot be called outside TaskInstance!"
                    )
                    raise AirflowException(message)
                return func(self, *args, **kwargs)
            finally:
                ExecutorSafeguard._local.in_executor_safeguard = False

        return wrapper

This solution would suppress the warning for legitimate nested calls while still protecting against unintended external calls. The use of local() ensures that this check is thread-safe.

I'm willing to create a Pull Request with this solution if the maintainers think this approach is appropriate. What do you think about this solution? Are there any aspects I should consider or modify before proceeding with a PR?

I like the approach, you can open a PR and add you example as a test case, that way we are sure this case is not only fixed but also tested against regression.

[dabla]

class LoggingBaseOperator(BaseOperator):
    def execute(self, context):
        self.log.info(f"Executing {self.__class__.__name__}")
        return super().execute(context)

I've been able to reproduce the issue in a test case, will open a PR for it which will fix it, will try the solution proposed by @SKalide

[dabla]

I've looked into this issue more deeply and I believe I can provide a solution. I could use some guidance on whether I should submit a PR for this.

The problem, in my opinion, is that after _execute_callable is called, it enters the ExecutorSafeguard wrapper. There, it checks if it's called inside a TaskInstance (hence the _execute_callable) and then returns the execution function itself. Because the execute method of BigQueryTableExistenceSensor is also decorated, it automatically re-enters the ExecutorSafeguard wrapper, but this time it appears as if it was called outside a Task Instance, causing the check to fail. With allow_nested_operators, it only fires a warning instead of failing.

My proposed solution is to check if the ExecutorSafeguard has already been called in the current execution context. If so, we can skip the check for nested calls. This can be achieved by setting a flag using thread-local storage to ensure thread safety.

class ExecutorSafeguard:
    """
    The ExecutorSafeguard decorator.

    Checks if the execute method of an operator isn't manually called outside
    the TaskInstance as we want to avoid bad mixing between decorated and
    classic operators.
    """

    test_mode = conf.getboolean("core", "unit_test_mode")
    _local = local()

    @classmethod
    def decorator(cls, func):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if (
                getattr(ExecutorSafeguard._local, "in_executor_safeguard", False)
                and self.allow_nested_operators
            ):
                # If already in ExecutorSafeguard, call execution function - recursive call
                return func(self, *args, **kwargs)
            ExecutorSafeguard._local.in_executor_safeguard = True

            try:
                from airflow.decorators.base import DecoratedOperator

                sentinel = kwargs.pop(f"{self.__class__.__name__}__sentinel", None)

                if (
                    not cls.test_mode
                    and sentinel != _sentinel
                    and not isinstance(self, DecoratedOperator)
                ):
                    message = (
                        f"{self.__class__.__name__}.{func.__name__} cannot be called outside TaskInstance!"
                    )
                    raise AirflowException(message)
                return func(self, *args, **kwargs)
            finally:
                ExecutorSafeguard._local.in_executor_safeguard = False

        return wrapper

This solution would suppress the warning for legitimate nested calls while still protecting against unintended external calls. The use of local() ensures that this check is thread-safe.

I'm willing to create a Pull Request with this solution if the maintainers think this approach is appropriate. What do you think about this solution? Are there any aspects I should consider or modify before proceeding with a PR?

The above proposed solution doesn't work for all test cases, especially in case of mixed usage with task decorated function calling execute methods on native operator. I've inspired the final working solution from it and now it should also support the case you presented.