You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dabla (via GitHub)" <gi...@apache.org> on 2025/12/05 10:12:49 UTC
[PR] Implemented AsyncPythonOperator [airflow]
dabla opened a new pull request, #59087:
URL: https://github.com/apache/airflow/pull/59087
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
This PR is related to the discussion I started on the [devlist](https://lists.apache.org/thread/ztnfsqolow4v1zsv4pkpnxc1fk0hbf2p) and which allows you to natively execute async code on PythonOperators.
There is also an article on [Google Docs](https://docs.google.com/document/d/1pNdQUB0gH-r2X1N_g774IOUEurowwQZ5OJ7yiY89qok) which describes it.
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672701671
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
@kaxil Do you propose to maybe add the check within the BaseAsyncOperator, and if callable is async but Airflow is < 3.2 then maybe raise NotImplementedError with appropriate message? Also isn't it better to revert this merge and remerged it once it fixes all your remarks as those are important ones imho.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implemented AsyncPythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3625972814
@uranusjr @Lee-W I've refactored PythonOperator as suggested, already tested it locally and it works. So no more need for dedicated AsyncPythonOperator, it also made the decorator part easier so there I could also remove some code. I also took into account the Airflow version, so that PythonOperator works without async support as before if Airflow version is too old.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3724244287
@dabla Apologies for a late review, but could you take a look at them please
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634571187
##########
providers/standard/tests/unit/standard/decorators/test_python.py:
##########
@@ -24,18 +25,27 @@
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskmap import TaskMap
-from airflow.providers.common.compat.sdk import AirflowException, XComNotFound
+from airflow.providers.common.compat.sdk import AirflowException, XComNotFound, task
Review Comment:
It's the same, except task_decorator is not used as the @task decorator itself on a method, but it's referencing the decorator on which tests are being peformed.
I added @task from the the common compat shared module to allow decorating a python function as Airflow task for testing purposes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implemented AsyncPythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3625509293
> Why do we need a separate operator class? Can this not just be a part of PythonOperator?
Good remark, will check, I found it cleaner to split but indeed maybe it's not necessary.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2671144197
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -1409,7 +1409,7 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
False,
"/dags/~/dagRuns/~/taskInstances",
{"dag_id_pattern": "example_python_operator"},
- 9, # Based on test failure - example_python_operator creates 9 task instances
+ 14, # Based on test failure - example_python_operator creates 9 task instances
Review Comment:
I've asked the same on slack an apparently it's due to the new added tasks in example dags that I've added for async tasks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3674470116
> I think this is good. I feel like we should keep a copy of BaseAsyncOperator (temporarily) in the standard provider so the feature can be available in old Airflow versions. That would make user experience a bit better—they simply need to upgrade the provider.
Or common compat library? Idk just thinking...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2674314731
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
good catch @kaxil
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2671147708
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -1409,7 +1409,7 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
False,
"/dags/~/dagRuns/~/taskInstances",
{"dag_id_pattern": "example_python_operator"},
- 9, # Based on test failure - example_python_operator creates 9 task instances
+ 14, # Based on test failure - example_python_operator creates 9 task instances
Review Comment:
Let’s update the comments to match reality then.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672685995
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -185,31 +187,69 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
err_decoder: TypeAdapter[ErrorResponse] = attrs.field(factory=lambda: TypeAdapter(ToTask), repr=False)
+ # Threading lock for sync operations
+ _thread_lock: threading.Lock = attrs.field(factory=threading.Lock, repr=False)
+ # Async lock for async operations
+ _async_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
+
def send(self, msg: SendMsgType) -> ReceiveMsgType | None:
"""Send a request to the parent and block until the response is received."""
frame = _RequestFrame(id=next(self.id_counter), body=msg.model_dump())
frame_bytes = frame.as_bytes()
- self.socket.sendall(frame_bytes)
- if isinstance(msg, ResendLoggingFD):
- if recv_fds is None:
- return None
- # We need special handling here! The server can't send us the fd number, as the number on the
- # supervisor will be different to in this process, so we have to mutate the message ourselves here.
- frame, fds = self._read_frame(maxfds=1)
- resp = self._from_frame(frame)
- if TYPE_CHECKING:
- assert isinstance(resp, SentFDs)
- resp.fds = fds
- # Since we know this is an expliclt SendFDs, and since this class is generic SendFDs might not
- # always be in the return type union
- return resp # type: ignore[return-value]
+ # We must make sure sockets aren't intermixed between sync and async calls,
+ # thus we need a dual locking mechanism to ensure that.
+ with self._thread_lock:
+ self.socket.sendall(frame_bytes)
+ if isinstance(msg, ResendLoggingFD):
+ if recv_fds is None:
+ return None
+ # We need special handling here! The server can't send us the fd number, as the number on the
+ # supervisor will be different to in this process, so we have to mutate the message ourselves here.
+ frame, fds = self._read_frame(maxfds=1)
+ resp = self._from_frame(frame)
+ if TYPE_CHECKING:
+ assert isinstance(resp, SentFDs)
+ resp.fds = fds
+ # Since we know this is an expliclt SendFDs, and since this class is generic SendFDs might not
+ # always be in the return type union
+ return resp # type: ignore[return-value]
return self._get_response()
Review Comment:
Normally it should behave correctly, unless I missed something, as the async version checks both async and sync lock, the sync version only checks the sync lock. We are now running the async PythonOperator since right after the Summit and until today haven't experienced any issues with it, I had issues in the beginning and this is how I came up with this solution as indeed it could lead to mix-ups very quickly without those locks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634682060
##########
providers/standard/src/airflow/providers/standard/operators/python.py:
##########
@@ -54,11 +54,18 @@
prepare_virtualenv,
write_python_script,
)
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS, BaseOperator
from airflow.utils import hashlib_wrapper
from airflow.utils.file import get_unique_dag_module_name
from airflow.utils.operator_helpers import KeywordParameters
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+
+ if TYPE_CHECKING:
+ from airflow.sdk.execution_time.callback_runner import AsyncExecutionCallableRunner
Review Comment:
Will change that thx @uranusjr
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2626817052
##########
task-sdk/src/airflow/sdk/bases/decorator.py:
##########
@@ -636,6 +640,28 @@ def __call__(
def override(self, **kwargs: Any) -> Task[FParams, FReturn]: ...
+def unwrap_callable(func):
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ # unwrap to the real underlying callable
+ return getattr(func, "function", getattr(func, "_func", func))
+ return func
+
+
+def is_async_callable(func):
Review Comment:
Will do I also saw you mention this on the devlist, will check it, thanks for pointing this out.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2626812999
##########
task-sdk/src/airflow/sdk/bases/decorator.py:
##########
@@ -636,6 +640,28 @@ def __call__(
def override(self, **kwargs: Any) -> Task[FParams, FReturn]: ...
+def unwrap_callable(func):
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ # unwrap to the real underlying callable
+ return getattr(func, "function", getattr(func, "_func", func))
+ return func
+
+
+def is_async_callable(func):
Review Comment:
You might take a look at fast-api implementation of the test. They implemented it in a bit different way at some point in time, which resulted in an issue - but they fixed it later:
See https://github.com/fastapi/fastapi/issues/14444
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2636963816
##########
providers/standard/docs/operators/python.rst:
##########
@@ -72,6 +72,33 @@ Pass extra arguments to the ``@task`` decorated function as you would with a nor
:start-after: [START howto_operator_python_kwargs]
:end-before: [END howto_operator_python_kwargs]
+Async Python functions
+^^^^^^^^^^^^^^^^^^^^^^
+
+From Airflow 3.2 onward, async Python callables are now also supported out of the box.
+This means we don't need to cope with the event loop and allows us to easily invoke async Python code and async
+Airflow hooks which are not always available through deferred operators.
+
Review Comment:
Yes, i could explicitly mention difference between async operators and deferred operators, as async operators run on the worker and deferred operators run on the triggerer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672720506
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
> Also isn't it better to revert this merge and remerged it once it fixes all your remarks as those are important ones imho.
Yup
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3723980461
#protm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672631869
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
+ return loop.run_until_complete(self.aexecute(context))
Review Comment:
We should add a test for this case specifically tbh so that a regression won't make async tasks run indefinitely.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672631869
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
+ return loop.run_until_complete(self.aexecute(context))
Review Comment:
We should add a test for this case specifically tbh so that a regression won't make async tasks run indefinitely.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3685032800
Why do we need to put `is_async_callable` in the SDK? Is it used anywhere other than the operator?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2637151923
##########
providers/standard/docs/operators/python.rst:
##########
@@ -72,6 +72,33 @@ Pass extra arguments to the ``@task`` decorated function as you would with a nor
:start-after: [START howto_operator_python_kwargs]
:end-before: [END howto_operator_python_kwargs]
+Async Python functions
+^^^^^^^^^^^^^^^^^^^^^^
+
+From Airflow 3.2 onward, async Python callables are now also supported out of the box.
+This means we don't need to cope with the event loop and allows us to easily invoke async Python code and async
+Airflow hooks which are not always available through deferred operators.
+
Review Comment:
Yes. Good point :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3661355823
An additional PR will be opened in which improvements to SFTPHookAsync have been made, I removed those from this PR to make review easier and not add unnecessary complexity not directly related to the async PythonOperator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672655455
##########
providers/standard/src/airflow/providers/standard/operators/python.py:
##########
@@ -115,9 +120,9 @@ def from_executable(cls, executable: str) -> _PythonVersionInfo:
return cls(*_parse_version_info(result.strip()))
-class PythonOperator(BaseOperator):
+class PythonOperator(BaseAsyncOperator):
Review Comment:
Should `BranchPythonOperator` and `ShortCircuitOperator` which inherit this Operator support or reject async callables?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2671148205
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -1409,7 +1409,7 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
False,
"/dags/~/dagRuns/~/taskInstances",
{"dag_id_pattern": "example_python_operator"},
- 9, # Based on test failure - example_python_operator creates 9 task instances
+ 14, # Based on test failure - example_python_operator creates 9 task instances
Review Comment:
https://apache-airflow.slack.com/archives/C06K9Q5G2UA/p1767616587380139
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2670517691
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | logging.Logger,
Review Comment:
? duplicate types
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2670973046
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -1409,7 +1409,7 @@ class TestGetTaskInstances(TestTaskInstanceEndpoint):
False,
"/dags/~/dagRuns/~/taskInstances",
{"dag_id_pattern": "example_python_operator"},
- 9, # Based on test failure - example_python_operator creates 9 task instances
+ 14, # Based on test failure - example_python_operator creates 9 task instances
Review Comment:
Why do these all priority weight values change?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672778807
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
Will revert this merge and reopen this PR once revert is done and do necessary changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2626811342
##########
providers/standard/tests/unit/standard/decorators/test_python.py:
##########
@@ -1131,3 +1140,29 @@ def my_teardown():
my_teardown()
assert work_task.operator.trigger_rule == TriggerRule.ONE_SUCCESS
assert setup_task.operator.trigger_rule == TriggerRule.ONE_SUCCESS
+
+
+@task
+async def dummy_task():
+ pass
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Test requires Airflow 3.2+")
+def test_is_async_callable():
+ from airflow.sdk.bases.decorator import is_async_callable
+
+ assert is_async_callable(dummy_task)
Review Comment:
I will rename it ;-)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla closed pull request #59087: Add support for async callables in PythonOperator
URL: https://github.com/apache/airflow/pull/59087
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implemented AsyncPythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3625653492
It’s indeed cleaner from the implementation-wise, but I’d prefer PythonOperator to just work with async functions, instead of needing a separate class. It’d be easier if there’s one class instead of two since PythonVirtualenvOperator etc. would also work.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3674321143
Added example dags with async PythonOperator into the documentation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2636761849
##########
providers/standard/docs/operators/python.rst:
##########
@@ -72,6 +72,33 @@ Pass extra arguments to the ``@task`` decorated function as you would with a nor
:start-after: [START howto_operator_python_kwargs]
:end-before: [END howto_operator_python_kwargs]
+Async Python functions
+^^^^^^^^^^^^^^^^^^^^^^
+
+From Airflow 3.2 onward, async Python callables are now also supported out of the box.
+This means we don't need to cope with the event loop and allows us to easily invoke async Python code and async
+Airflow hooks which are not always available through deferred operators.
+
Review Comment:
I'm not sure do we need to explicitly mention the async `@task` will still be executed in Worker instead of Triggerer on the doc. As some of the users _might_ directly associate all the "async" workload in Airflow will be handled by Triggerer instead of Worker.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634542072
##########
providers/standard/tests/unit/standard/decorators/test_python.py:
##########
@@ -24,18 +25,27 @@
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskmap import TaskMap
-from airflow.providers.common.compat.sdk import AirflowException, XComNotFound
+from airflow.providers.common.compat.sdk import AirflowException, XComNotFound, task
from tests_common.test_utils.version_compat import (
AIRFLOW_V_3_0_1,
AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
XCOM_RETURN_KEY,
)
from unit.standard.operators.test_python import BasePythonTest
if AIRFLOW_V_3_0_PLUS:
- from airflow.sdk import DAG, BaseOperator, TaskGroup, XComArg, setup, task as task_decorator, teardown
+ from airflow.sdk import (
+ DAG,
+ BaseOperator,
+ TaskGroup,
+ XComArg,
+ setup,
+ task as task_decorator,
+ teardown,
+ )
Review Comment:
This doesn’t change right? (Just to make sure I’m not reading this wrong)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3681009960
@uranusjr Hello TP, I think we have a backward compatibilty issue for older Airflow versions, as for the @task decorator I also had to add a check for the async through the is_async property, which is part of the sdk and thus not possible in provider. WDYT? But maybe it’s not really necessary, will have to check the tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2671146395
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | logging.Logger,
Review Comment:
Yeah that's wrong will fix it, it was fixed in sdk but not in backward compat, thx @kaxil.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672681142
##########
task-sdk/src/airflow/sdk/bases/operator.py:
##########
@@ -1670,6 +1694,35 @@ def has_on_skipped_callback(self) -> bool:
return bool(self.on_skipped_callback)
+class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
+ return loop.run_until_complete(self.aexecute(context))
Review Comment:
We should add a test for this case specifically tbh so that a regression won't make async tasks run indefinitely.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3719150278
I think this one is ready for review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672627083
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -185,31 +187,69 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
err_decoder: TypeAdapter[ErrorResponse] = attrs.field(factory=lambda: TypeAdapter(ToTask), repr=False)
+ # Threading lock for sync operations
+ _thread_lock: threading.Lock = attrs.field(factory=threading.Lock, repr=False)
+ # Async lock for async operations
+ _async_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
+
def send(self, msg: SendMsgType) -> ReceiveMsgType | None:
"""Send a request to the parent and block until the response is received."""
frame = _RequestFrame(id=next(self.id_counter), body=msg.model_dump())
frame_bytes = frame.as_bytes()
- self.socket.sendall(frame_bytes)
- if isinstance(msg, ResendLoggingFD):
- if recv_fds is None:
- return None
- # We need special handling here! The server can't send us the fd number, as the number on the
- # supervisor will be different to in this process, so we have to mutate the message ourselves here.
- frame, fds = self._read_frame(maxfds=1)
- resp = self._from_frame(frame)
- if TYPE_CHECKING:
- assert isinstance(resp, SentFDs)
- resp.fds = fds
- # Since we know this is an expliclt SendFDs, and since this class is generic SendFDs might not
- # always be in the return type union
- return resp # type: ignore[return-value]
+ # We must make sure sockets aren't intermixed between sync and async calls,
+ # thus we need a dual locking mechanism to ensure that.
+ with self._thread_lock:
+ self.socket.sendall(frame_bytes)
+ if isinstance(msg, ResendLoggingFD):
+ if recv_fds is None:
+ return None
+ # We need special handling here! The server can't send us the fd number, as the number on the
+ # supervisor will be different to in this process, so we have to mutate the message ourselves here.
+ frame, fds = self._read_frame(maxfds=1)
+ resp = self._from_frame(frame)
+ if TYPE_CHECKING:
+ assert isinstance(resp, SentFDs)
+ resp.fds = fds
+ # Since we know this is an expliclt SendFDs, and since this class is generic SendFDs might not
+ # always be in the return type union
+ return resp # type: ignore[return-value]
return self._get_response()
Review Comment:
If a sync send() and async asend() are called concurrently (which is now possible with mixed sync/async code), could the response reads interleave? The async version you have keeps the read inside the lock, but the sync version doesn't so curious.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implemented AsyncPythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3625861626
On thing I have to take into account: Airflow 2
So I think we will have basic PythonOperator implementation without async support for Airflow 2 and for Airflow 3.2 for example we will do the one with async support, as the BaseAsyncOperator will not exist before in previous Airflow versions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implemented AsyncPythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3625787421
> It’s indeed cleaner from the implementation-wise, but I’d prefer PythonOperator to just work with async functions, instead of needing a separate class. It’d be easier if there’s one class instead of two since PythonVirtualenvOperator etc. would also work.
Thanks TP for you remark, you're indeed right, didn't really think about it as you don't see it when using the @task decorator, but that would be indeed annoying if you had to use another PythonOperator for async functions, good catch!
I will fix this, but that would of course mean the PythonOperator would by default extend the BaseAsyncOperator instead of the BaseOperator, but that's not an issue I suppose. The the is_async property will need to be overriden in the PythonOperator to be able to check the python_callable if it's async or not.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634552453
##########
providers/standard/src/airflow/providers/standard/operators/python.py:
##########
@@ -54,11 +54,18 @@
prepare_virtualenv,
write_python_script,
)
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS, BaseOperator
from airflow.utils import hashlib_wrapper
from airflow.utils.file import get_unique_dag_module_name
from airflow.utils.operator_helpers import KeywordParameters
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+
+ if TYPE_CHECKING:
+ from airflow.sdk.execution_time.callback_runner import AsyncExecutionCallableRunner
Review Comment:
This should not be inside `if AIRFLOW_V_3_2_PLUS` (and it doesn’t work as you expect; conditionals are always visited in type checking)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3678692496
@uranusjr I've mad the BaseAsyncOperator available for older Airflow versions through the common compat module. I also moved the is_async_callable method to the standard python provider. I only needed to duplicate the event_loop context manager method in the common compat module together with the BaseAsyncOperator, I don't think that much of an issue after all.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3665075222
This is really nice and clean and it actually shows why we actually **should** handle async tasks this way - @dstandish -> this PR is now extracted by @dabla to only implement the base changes we discuss in https://lists.apache.org/list.html?dev@airflow.apache.org - I think you should take a look to see how much nicer it would be and how nice it will allow us to extend it further for other decorators.
I guess we want to still get a quick AIP approval on it @dabla - assuming that others will also thumb-up it (but I really see how it might be useful).
This one will need newsfragment and nice docs explaining what and why as well of course :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2627546007
##########
task-sdk/src/airflow/sdk/execution_time/callback_runner.py:
##########
@@ -109,3 +124,51 @@ def _run():
return result # noqa: F821 # Ruff is not smart enough to know this is always set in _run().
return cast("_ExecutionCallableRunner[P, R]", _ExecutionCallableRunnerImpl)
+
+
+def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | logging.Logger,
Review Comment:
```suggestion
logger: logging.Logger | Logger,
```
##########
task-sdk/src/airflow/sdk/execution_time/callback_runner.py:
##########
@@ -49,6 +54,16 @@ def __call__(
) -> _ExecutionCallableRunner[P, R]: ...
+class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | logging.Logger,
Review Comment:
```suggestion
logger: logging.Logger | Logger,
```
##########
providers/standard/docs/operators/python.rst:
##########
@@ -72,6 +72,33 @@ Pass extra arguments to the ``@task`` decorated function as you would with a nor
:start-after: [START howto_operator_python_kwargs]
:end-before: [END howto_operator_python_kwargs]
+Async Python functions
+^^^^^^^^^^^^^^^^^^^^^^
+
+From Airflow 3.2 onward, async Python callables are now also supported out of the box.
+This means we don't need to cope with the event loop and allows us to easily invoke async Python code and async
+Airflow hooks which are not always available through deferred operators.
+
Review Comment:
I'm not sure do we need to explicitly mention the async `@task` will still be executed in Worker instead of Triggerer on the doc. As some of the users _might_ directly associate all the "async" workload in Airflow will be handled by Trigger instead of Worker.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634561636
##########
providers/standard/tests/unit/standard/decorators/test_python.py:
##########
@@ -24,18 +25,27 @@
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskmap import TaskMap
-from airflow.providers.common.compat.sdk import AirflowException, XComNotFound
+from airflow.providers.common.compat.sdk import AirflowException, XComNotFound, task
from tests_common.test_utils.version_compat import (
AIRFLOW_V_3_0_1,
AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
XCOM_RETURN_KEY,
)
from unit.standard.operators.test_python import BasePythonTest
if AIRFLOW_V_3_0_PLUS:
- from airflow.sdk import DAG, BaseOperator, TaskGroup, XComArg, setup, task as task_decorator, teardown
+ from airflow.sdk import (
+ DAG,
+ BaseOperator,
+ TaskGroup,
+ XComArg,
+ setup,
+ task as task_decorator,
+ teardown,
+ )
Review Comment:
Yes it doesn't change, it's reformatted
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634545018
##########
providers/standard/tests/unit/standard/decorators/test_python.py:
##########
@@ -24,18 +25,27 @@
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskmap import TaskMap
-from airflow.providers.common.compat.sdk import AirflowException, XComNotFound
+from airflow.providers.common.compat.sdk import AirflowException, XComNotFound, task
Review Comment:
Is this not the same as `task_decorator` below?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3689358852
> Why do we need to put `is_async_callable` in the SDK? Is it used anywhere other than the operator?
The is_async property isn't called directly on the DecoratedOperator, but it will be in the future once we introduce the IterableOperator in the future which was also discussed in the devlist. So I think it's best to already include it, I'm still working this out to make it backward compatible with older Airflow versions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672778807
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
Will revert this [merge ](https://github.com/apache/airflow/pull/60266)and reopen this PR once revert is done and do necessary changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672994103
##########
providers/standard/src/airflow/providers/standard/operators/python.py:
##########
@@ -115,9 +120,9 @@ def from_executable(cls, executable: str) -> _PythonVersionInfo:
return cls(*_parse_version_info(result.strip()))
-class PythonOperator(BaseOperator):
+class PythonOperator(BaseAsyncOperator):
Review Comment:
The do automatically as they inherit from PythonOperator which is an async operator by default. The BaseBranchOperator operator should not be async.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672675843
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
Why do we need the code in Common Compat? This can't work without the Task SDK changes of the CommsDecoder in 3.2, no? So I do not these code here, we just error out if it <3.2
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
Why do we need this code in Common Compat? This can't work without the Task SDK changes of the CommsDecoder in 3.2, no? So I do not these code here, we just error out if it <3.2
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla merged PR #59087:
URL: https://github.com/apache/airflow/pull/59087
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2675105232
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
Indeed good catch by @kaxil, already created new [PR](https://github.com/apache/airflow/pull/60268) which fixes all this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2627055980
##########
task-sdk/src/airflow/sdk/bases/decorator.py:
##########
@@ -636,6 +640,28 @@ def __call__(
def override(self, **kwargs: Any) -> Task[FParams, FReturn]: ...
+def unwrap_callable(func):
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ # unwrap to the real underlying callable
+ return getattr(func, "function", getattr(func, "_func", func))
+ return func
+
+
+def is_async_callable(func):
Review Comment:
I refactored the is_async_callable method to take into account more scenario's like with FastAPI, I also added different test cases documenting the methods behaviour.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2626803752
##########
providers/standard/tests/unit/standard/decorators/test_python.py:
##########
@@ -1131,3 +1140,29 @@ def my_teardown():
my_teardown()
assert work_task.operator.trigger_rule == TriggerRule.ONE_SUCCESS
assert setup_task.operator.trigger_rule == TriggerRule.ONE_SUCCESS
+
+
+@task
+async def dummy_task():
+ pass
+
+
+@pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Test requires Airflow 3.2+")
+def test_is_async_callable():
+ from airflow.sdk.bases.decorator import is_async_callable
+
+ assert is_async_callable(dummy_task)
Review Comment:
we avoid using dummy :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2634705955
##########
providers/standard/src/airflow/providers/standard/operators/python.py:
##########
@@ -54,11 +54,18 @@
prepare_virtualenv,
write_python_script,
)
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS, BaseOperator
from airflow.utils import hashlib_wrapper
from airflow.utils.file import get_unique_dag_module_name
from airflow.utils.operator_helpers import KeywordParameters
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+
+ if TYPE_CHECKING:
+ from airflow.sdk.execution_time.callback_runner import AsyncExecutionCallableRunner
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2637685827
##########
providers/standard/docs/operators/python.rst:
##########
@@ -72,6 +72,33 @@ Pass extra arguments to the ``@task`` decorated function as you would with a nor
:start-after: [START howto_operator_python_kwargs]
:end-before: [END howto_operator_python_kwargs]
+Async Python functions
+^^^^^^^^^^^^^^^^^^^^^^
+
+From Airflow 3.2 onward, async Python callables are now also supported out of the box.
+This means we don't need to cope with the event loop and allows us to easily invoke async Python code and async
+Airflow hooks which are not always available through deferred operators.
+
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3678605683
> I think this is good. I feel like we should keep a copy of BaseAsyncOperator (temporarily) in the standard provider so the feature can be available in old Airflow versions. That would make user experience a bit better—they simply need to upgrade the provider.
Good idea, but what about the new is_async property in BaseOperator, and the event_loop helper method and is_async_callable methods in sdk decorators?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implemented AsyncPythonOperator [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3625443135
Why do we need a separate operator class? Can this not just be a part of PythonOperator?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672692431
##########
providers/common/compat/src/airflow/providers/common/compat/standard/operators.py:
##########
@@ -17,18 +17,230 @@
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from airflow.providers.common.compat._compat_utils import create_module_getattr
+from airflow.providers.common.compat.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+)
_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
# Re-export from sdk (which handles Airflow 2.x/3.x fallbacks)
+ "AsyncExecutionCallableRunner": "airflow.providers.common.compat.sdk",
"BaseOperator": "airflow.providers.common.compat.sdk",
+ "BaseAsyncOperator": "airflow.providers.common.compat.sdk",
+ "create_async_executable_runner": "airflow.providers.common.compat.sdk",
"get_current_context": "airflow.providers.common.compat.sdk",
+ "is_async_callable": "airflow.providers.common.compat.sdk",
# Standard provider items with direct fallbacks
"PythonOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"ShortCircuitOperator": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
"_SERIALIZERS": ("airflow.providers.standard.operators.python", "airflow.operators.python"),
}
+if TYPE_CHECKING:
+ from structlog.typing import FilteringBoundLogger as Logger
+
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+ from airflow.sdk.types import OutletEventAccessorsProtocol
+elif AIRFLOW_V_3_2_PLUS:
+ from airflow.sdk.bases.decorator import is_async_callable
+ from airflow.sdk.bases.operator import BaseAsyncOperator
+ from airflow.sdk.execution_time.callback_runner import (
+ AsyncExecutionCallableRunner,
+ create_async_executable_runner,
+ )
+else:
+ import asyncio
+ import contextlib
+ import inspect
+ import logging
+ from asyncio import AbstractEventLoop
+ from collections.abc import AsyncIterator, Awaitable, Callable, Generator
+ from contextlib import suppress
+ from functools import partial
+ from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
+
+ from typing_extensions import ParamSpec
+
+ if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import BaseOperator
+ from airflow.sdk.bases.decorator import _TaskDecorator
+ from airflow.sdk.definitions.asset.metadata import Metadata
+ from airflow.sdk.definitions.mappedoperator import OperatorPartial
+ else:
+ from airflow.datasets.metadata import Metadata
+ from airflow.decorators.base import _TaskDecorator
+ from airflow.models import BaseOperator
+ from airflow.models.mappedoperator import OperatorPartial
+
+ P = ParamSpec("P")
+ R = TypeVar("R")
+
+ @contextlib.contextmanager
+ def event_loop() -> Generator[AbstractEventLoop]:
+ new_event_loop = False
+ loop = None
+ try:
+ try:
+ loop = asyncio.get_event_loop()
+ if loop.is_closed():
+ raise RuntimeError
+ except RuntimeError:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ new_event_loop = True
+ yield loop
+ finally:
+ if new_event_loop and loop is not None:
+ with contextlib.suppress(AttributeError):
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ def unwrap_partial(fn):
+ while isinstance(fn, partial):
+ fn = fn.func
+ return fn
+
+ def unwrap_callable(func):
+ # Airflow-specific unwrap
+ if isinstance(func, (_TaskDecorator, OperatorPartial)):
+ func = getattr(func, "function", getattr(func, "_func", func))
+
+ # Unwrap functools.partial
+ func = unwrap_partial(func)
+
+ # Unwrap @functools.wraps chains
+ with suppress(Exception):
+ func = inspect.unwrap(func)
+
+ return func
+
+ def is_async_callable(func):
+ """Detect if a callable (possibly wrapped) is an async function."""
+ func = unwrap_callable(func)
+
+ if not callable(func):
+ return False
+
+ # Direct async function
+ if inspect.iscoroutinefunction(func):
+ return True
+
+ # Callable object with async __call__
+ if not inspect.isfunction(func):
+ call = type(func).__call__ # Bandit-safe
+ with suppress(Exception):
+ call = inspect.unwrap(call)
+ if inspect.iscoroutinefunction(call):
+ return True
+
+ return False
+
+ class _AsyncExecutionCallableRunner(Generic[P, R]):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R: ... # type: ignore[empty-body]
+
+ class AsyncExecutionCallableRunner(Protocol):
+ def __call__(
+ self,
+ func: Callable[P, R],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]: ...
+
+ def create_async_executable_runner(
+ func: Callable[P, Awaitable[R] | AsyncIterator],
+ outlet_events: OutletEventAccessorsProtocol,
+ *,
+ logger: logging.Logger | Logger,
+ ) -> _AsyncExecutionCallableRunner[P, R]:
+ """
+ Run an async execution callable against a task context and given arguments.
+
+ If the callable is a simple function, this simply calls it with the supplied
+ arguments (including the context). If the callable is a generator function,
+ the generator is exhausted here, with the yielded values getting fed back
+ into the task context automatically for execution.
+
+ This convoluted implementation of inner class with closure is so *all*
+ arguments passed to ``run()`` can be forwarded to the wrapped function. This
+ is particularly important for the argument "self", which some use cases
+ need to receive. This is not possible if this is implemented as a normal
+ class, where "self" needs to point to the runner object, not the object
+ bounded to the inner callable.
+
+ :meta private:
+ """
+
+ class _AsyncExecutionCallableRunnerImpl(_AsyncExecutionCallableRunner):
+ @staticmethod
+ async def run(*args: P.args, **kwargs: P.kwargs) -> R:
+ if not inspect.isasyncgenfunction(func):
+ result = cast("Awaitable[R]", func(*args, **kwargs))
+ return await result
+
+ results: list[Any] = []
+
+ async for result in func(*args, **kwargs):
+ if isinstance(result, Metadata):
+ outlet_events[result.asset].extra.update(result.extra)
+ if result.alias:
+ outlet_events[result.alias].add(result.asset, extra=result.extra)
+
+ results.append(result)
+
+ return cast("R", results)
+
+ return cast("_AsyncExecutionCallableRunner[P, R]", _AsyncExecutionCallableRunnerImpl)
+
+ class BaseAsyncOperator(BaseOperator):
+ """
+ Base class for async-capable operators.
+
+ As opposed to deferred operators which are executed on the triggerer, async operators are executed
+ on the worker.
+ """
+
+ @property
+ def is_async(self) -> bool:
+ return True
+
+ if not AIRFLOW_V_3_1_PLUS:
+
+ @property
+ def xcom_push(self) -> bool:
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool):
+ self.do_xcom_push = value
+
+ async def aexecute(self, context):
+ """Async version of execute(). Subclasses should implement this."""
+ raise NotImplementedError()
+
+ def execute(self, context):
+ """Run `aexecute()` inside an event loop."""
+ with event_loop() as loop:
+ if self.execution_timeout:
+ return loop.run_until_complete(
+ asyncio.wait_for(
+ self.aexecute(context),
+ timeout=self.execution_timeout.total_seconds(),
+ )
+ )
Review Comment:
Hmm good catch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #59087:
URL: https://github.com/apache/airflow/pull/59087#discussion_r2672640943
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -185,31 +187,69 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
err_decoder: TypeAdapter[ErrorResponse] = attrs.field(factory=lambda: TypeAdapter(ToTask), repr=False)
+ # Threading lock for sync operations
+ _thread_lock: threading.Lock = attrs.field(factory=threading.Lock, repr=False)
+ # Async lock for async operations
+ _async_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
Review Comment:
This is worth consolidating with the async lock on Trigger since `TriggerCommsDecoder` inherits this
https://github.com/apache/airflow/blob/0b341e6b92040360f86f5697da119c58a91aa4c2/airflow-core/src/airflow/jobs/triggerer_job_runner.py#L783
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3723445559
Very nice and clean code !
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add support for async callables in PythonOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #59087:
URL: https://github.com/apache/airflow/pull/59087#issuecomment-3685886384
> Why do we need to put `is_async_callable` in the SDK? Is it used anywhere other than the operator?
Because the decorator is in the SDK, and it inherits from BaseOperator which now has the is_async property. The PythonOperator is a special case, as it can be sync or async depending on the callable, thus there it has to be determined dynamically, otherwise BaseOperators are always is_async False and BaseAsyncOperators are by default always True dor is_async.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org