You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jason810496 (via GitHub)" <gi...@apache.org> on 2026/10/20 04:27:17 UTC
[PR] Apply the configured XCom backend to language SDK task XComs [airflow]
jason810496 opened a new pull request, #75184:
URL: https://github.com/apache/airflow/pull/75184
- closes: #51067
## Why
TypeScript, Go and Java tasks skip the configured `[core] xcom_backend`, because their XComs go from the supervisor straight to the API server. In a Python Dag with stub TaskHandler that uses `XComObjectStorageBackend`, a TypeScript task reads the `s3://...` reference that a Python task stored instead of the value, its own values stay in the metadata database, and clearing a try leaves the stored objects behind.
## How
- The supervisor of a language SDK task runs the backend on the runtime's XCom requests (`SetXCom`, `GetXCom`, `GetXComSequenceItem`, `GetXComSequenceSlice`, `DeleteXCom`) and on the XComs it clears when a try starts. With the default `BaseXCom`, requests go to the API server unchanged.
- The backend runs in a helper process, started per task the same way a Python task process is started (bare fork, or a new interpreter on macOS or with `[core] execute_tasks_new_python_interpreter`). So the long-lived worker process keeps none of the backend's state, such as a cached `S3FileSystem` that a Python task forked from it later could not use.
- Runtimes get what the default backend would give them: the value is decoded with serde before `store_value`, and a loaded value is encoded with serde again. `skipmixin_key` bypasses the backend, as it does for Python tasks.
- The backend looks up Connections and Variables through the supervisor's API client.
- While a backend call runs, the supervisor keeps heartbeating, and stops the call when the task reaches its `execution_timeout` or the server stops the task.
- A Go task now fails when its return value cannot be stored, as TypeScript and Java tasks already do.
No SDK change is needed, since all three runtimes only talk to the supervisor. With this config, TypeScript, Go and Java tasks read and write offloaded XComs like Python tasks:
```ini
[core]
xcom_backend = airflow.providers.common.io.xcom.backend.XComObjectStorageBackend
```
## Performance
Default backend, mocked API, median of 3 alternating runs:
| | base | this PR |
|---|---|---|
| `SetXCom` handler | 65.8 µs | 65.6 µs |
| `GetXCom` handler | 67.7 µs | 67.1 µs |
| `GetXComCount` handler (control) | 58.2 µs | 57.8 µs |
| `_find_own_coordinate` | 0.645 µs | 0.641 µs |
## Verification
- task-sdk unit tests, including a real helper process that looks up a Connection through the supervisor.
- `go test ./...` and `go test -race ./...`.
- The `xcom_object_storage` e2e mode now also runs `typescript_taskflow_example` and checks the objects in S3. It was not run locally (no Docker).
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
Generated-by: Claude Code (Opus 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397374088
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
+ self._pending_xcom_calls.append(finish)
+ self._xcom_backend.send_call(call)
+
+ def _finish_xcom_backend_call(self, result: XComBackendResult) -> None:
+ """Finish the oldest call the XCom backend process has not answered yet."""
+ if not self._pending_xcom_calls:
+ return
+ self._pending_xcom_calls.popleft()(result)
+
+ def _call_xcom_backend(
+ self,
+ call: XComBackendCall,
+ then: Callable[[XComBackendResult], RequestResult],
+ req_id: int,
+ *,
+ on_error: Callable[[str], RequestResult] | None = None,
+ ) -> ResponseSent:
+ """
+ Answer request ``req_id`` with ``then(result)`` once the XCom backend has run ``call``.
+
+ The request loop goes on servicing the task and the backend process while the call runs, so it
+ keeps heartbeating and stops the task when it reaches its ``execution_timeout``.
+
+ ``on_error`` answers a failed call instead; without it, a failed call fails the request.
+ """
+
+ def finish(result: XComBackendResult) -> None:
+ try:
+ if result.error is None:
+ resp, dump_opts = then(result)
+ elif on_error is not None:
+ resp, dump_opts = on_error(result.error)
+ else:
+ raise XComBackendError(result.error)
+ except Exception as e:
+ self.process_log.exception("Failed to handle the XCom backend result", request_id=req_id)
+ self.send_msg(None, request_id=req_id, error=_build_request_error(e))
+ else:
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+
+ self._start_xcom_backend_call(call, finish)
+ return ResponseSent.ALREADY_SENT
+
+ def _run_xcom_backend_call(self, call: XComBackendCall) -> None:
+ """
+ Run ``call`` and wait for it, for a caller with no request loop to return to.
+
+ :raises XComBackendError: when the call fails or is not answered in time.
+ """
+ answer: list[XComBackendResult] = []
+ self._start_xcom_backend_call(call, answer.append)
+ deadline = time.monotonic() + _XCOM_BACKEND_CALL_TIMEOUT
+ # The task process has not been told to start, so servicing it here reads no request of its own.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ self._pending_xcom_calls.clear()
+ if not answer:
+ raise XComBackendError("The XCom backend did not answer in time")
+ if (error := answer[0].error) is not None:
+ raise XComBackendError(error)
+
+ def _handle_set_xcom(
+ self, msg: SetXCom, log: FilteringBoundLogger, req_id: int
+ ) -> RequestResult | ResponseSent:
+ xcom = _get_xcom_backend()
+ if not (
+ is_stored_through_backend(msg.key) and is_hook_overridden(xcom, "store_value", "serialize_value")
+ ):
+ return handle_set_xcom(self.client, msg)
+ ti = cast("_OwnTaskInstance", self.ti)
+ if is_own_xcom(ti, msg.dag_id, msg.run_id, msg.task_id, msg.resolved_map_index):
+ coordinate = build_own_coordinate(ti)
+ else:
+ coordinate = XComCoordinate(msg.dag_id, msg.run_id, msg.task_id)
+ # A loop iteration has map index -1, unlike a mapped task. A mapped task inside a loop is not
+ # caught here because the supervisor does not parse the Dag, while Python rejects such a Dag
+ # when it is parsed.
+ if coordinate.region_id is not None and msg.resolved_map_index == -1 and is_legacy_xcom_backend(xcom):
+ raise build_loop_error(xcom, f"Task {msg.task_id!r}")
Review Comment:
This reaches the runtime as `API_SERVER_ERROR`, which is misleading for a configuration error. Can it use a more specific error type?
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397328232
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
+ self._pending_xcom_calls.append(finish)
+ self._xcom_backend.send_call(call)
+
+ def _finish_xcom_backend_call(self, result: XComBackendResult) -> None:
+ """Finish the oldest call the XCom backend process has not answered yet."""
+ if not self._pending_xcom_calls:
+ return
+ self._pending_xcom_calls.popleft()(result)
+
+ def _call_xcom_backend(
+ self,
+ call: XComBackendCall,
+ then: Callable[[XComBackendResult], RequestResult],
+ req_id: int,
+ *,
+ on_error: Callable[[str], RequestResult] | None = None,
+ ) -> ResponseSent:
+ """
+ Answer request ``req_id`` with ``then(result)`` once the XCom backend has run ``call``.
+
+ The request loop goes on servicing the task and the backend process while the call runs, so it
+ keeps heartbeating and stops the task when it reaches its ``execution_timeout``.
+
+ ``on_error`` answers a failed call instead; without it, a failed call fails the request.
+ """
+
+ def finish(result: XComBackendResult) -> None:
+ try:
+ if result.error is None:
+ resp, dump_opts = then(result)
+ elif on_error is not None:
+ resp, dump_opts = on_error(result.error)
+ else:
+ raise XComBackendError(result.error)
+ except Exception as e:
+ self.process_log.exception("Failed to handle the XCom backend result", request_id=req_id)
+ self.send_msg(None, request_id=req_id, error=_build_request_error(e))
+ else:
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+
+ self._start_xcom_backend_call(call, finish)
+ return ResponseSent.ALREADY_SENT
+
+ def _run_xcom_backend_call(self, call: XComBackendCall) -> None:
+ """
+ Run ``call`` and wait for it, for a caller with no request loop to return to.
+
+ :raises XComBackendError: when the call fails or is not answered in time.
+ """
+ answer: list[XComBackendResult] = []
+ self._start_xcom_backend_call(call, answer.append)
+ deadline = time.monotonic() + _XCOM_BACKEND_CALL_TIMEOUT
+ # The task process has not been told to start, so servicing it here reads no request of its own.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ self._pending_xcom_calls.clear()
Review Comment:
This clears the callbacks on timeout, but the helper is still running the call, and XComBackendProcess still counts it as unanswered. When the late result arrives, `popleft()` hands it to the next call, and every later answer is off by one.
A timed-out helper is in an unknown state anyway. Kill it and start fresh on the next call.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on PR #75184:
URL: https://github.com/apache/airflow/pull/75184#issuecomment-6306874722
> Having a whole separate process to handle xcom uplodads feels overkill. Why can't we do it from the coordinator process directly? The "why" in the PR desc doesn't address this
Good point, I've added this to the Why in the description:
1. **Fork safety.** The supervisor runs inside the long-lived worker process (for example the Celery worker child), which also forks the Python task processes. A backend run there leaves its state in that process: `XComObjectStorageBackend` caches its fsspec filesystem (`_STORE_CACHE` in `airflow.sdk.io.store`), and a Python task forked from that process later fails with `RuntimeError: This class is not fork-safe`.
2. **The supervisor stays in control.** A backend call blocks on I/O. In the supervisor, a hung upload would block the loop that heartbeats and enforces `execution_timeout`, with no way to interrupt it. In a separate process, the supervisor keeps heartbeating during the call, and kills the helper on `execution_timeout`, when the server stops the task, or after too many failed heartbeats.
The helper is also lazy and per task: it starts on the first call the configured backend handles and serves all of that task's XCom calls, so tasks on the default `BaseXCom` never start it.
It is also not meant for XCom only. Any worker-side backend that runs in a Python task's process and keeps in-memory state should run in this helper for the same isolation, so the `[workers] state_store_backend` and the worker's secrets backends are next.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397355996
##########
task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py:
##########
@@ -0,0 +1,345 @@
+#
+# 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.
+"""
+The XCom backend for the tasks of a language SDK runtime.
+
+A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend
+on them before it forwards them to the API server. The backend runs in a helper process, started per task
+the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of
+the backend's state, such as a cached filesystem that a task process forked from it later could not use.
+
+The supervisor services the helper process from the selector of the task it serves, so a backend call
+never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import selectors
+import signal
+import time
+from collections import deque
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Annotated, ClassVar, Literal
+from uuid import UUID
+
+import attrs
+import structlog
+from pydantic import BaseModel, Field, JsonValue, TypeAdapter
+
+from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate
+from airflow.sdk.execution_time.comms import (
+ ConnectionResult,
+ ErrorResponse,
+ GetConnection,
+ GetVariable,
+ MaskSecret,
+ VariableResult,
+ XComResult,
+)
+from airflow.sdk.execution_time.supervisor import (
+ ResponseSent,
+ WatchedSubprocess,
+ register_request_method,
+ task_process_uses_exec,
+)
+
+if TYPE_CHECKING:
+ from socket import socket
+ from typing import Self
+
+ from structlog.typing import FilteringBoundLogger
+
+ from airflow.sdk.api.client import Client
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend")
+
+_EXIT_TIMEOUT = 10.0
+
+
+class XComBackendError(RuntimeError):
+ """
+ An XCom backend call failed.
+
+ For an exception raised in the helper process, the message holds its type and message, since the
+ exception itself does not cross the process boundary.
+ """
+
+
+class _Coordinate(BaseModel):
+ dag_id: str
+ run_id: str
+ task_id: str
+ ti_id: UUID | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
+
+
+class XComBackendCall(BaseModel):
+ """
+ A call of an XCom backend hook, sent to the helper process.
+
+ ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``,
+ and a list of rows for ``load_many``.
+ """
+
+ hook: Literal["store", "load", "load_many", "purge", "exit"]
+ key: str = ""
+ value: JsonValue = None
+ coordinate: _Coordinate | None = None
+ map_index: int = -1
+ type: Literal["XComBackendCall"] = "XComBackendCall"
+
+ @classmethod
+ def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self:
+ return cls(
+ hook="store",
+ key=key,
+ value=value,
+ coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)),
+ map_index=map_index,
+ )
+
+
+class XComBackendResult(BaseModel):
+ """
+ The outcome of an XCom backend call, sent by the helper process.
+
+ ``unchanged`` means the backend did not store the value, which then goes to the API server as it is.
+ """
+
+ value: JsonValue = None
+ unchanged: bool = False
+ error: str | None = None
+ type: Literal["XComBackendResult"] = "XComBackendResult"
+
+
+class AwaitXComBackendCall(BaseModel):
+ """Sent by the helper process when it is ready for its first call."""
+
+ type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall"
+
+
+_ToXComBackendProcess = Annotated[
+ XComBackendCall | ConnectionResult | VariableResult | ErrorResponse,
+ Field(discriminator="type"),
+]
+_FromXComBackendProcess = Annotated[
+ AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret,
+ Field(discriminator="type"),
+]
+
+
+def is_hook_overridden(xcom: type[BaseXCom], *hooks: str) -> bool:
+ """Whether ``xcom`` overrides any of ``hooks``, and so may do more than the default backend."""
+ return xcom is not BaseXCom and any(
+ next(k for k in xcom.__mro__ if hook in vars(k)) is not BaseXCom for hook in hooks
+ )
+
+
+@attrs.define(kw_only=True)
+class XComBackendProcess(WatchedSubprocess):
+ """
+ The helper process that runs XCom backend calls for the task of a language SDK runtime.
+
+ The process asks for a call, runs it, and answers with its result, which is also how it asks for the
+ next call. Each result reaches the supervisor through ``on_result``. Connection and Variable lookups of
+ the backend reach the API server through ``client``.
+ """
+
+ client: Client
+ on_result: Callable[[XComBackendResult], None] = attrs.field(repr=False)
+ """Called with the result of a call, and with an error when the process exits without answering one."""
+ decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess)
+
+ _call_request_id: int | None = attrs.field(default=None, init=False)
+ _queued_calls: deque[XComBackendCall] = attrs.field(factory=deque, init=False)
+ _unanswered_calls: int = attrs.field(default=0, init=False)
+
+ @classmethod
+ def start( # type: ignore[override]
+ cls,
+ *,
+ id: UUID,
+ client: Client,
+ logger: FilteringBoundLogger,
+ selector: selectors.BaseSelector,
+ on_result: Callable[[XComBackendResult], None],
+ ) -> Self:
+ return super().start(
+ id=id,
+ client=client,
+ on_result=on_result,
+ # The selector of the task this process serves, so that the task's request loop services both.
+ selector=selector,
+ logger=logger,
+ target=run_xcom_backend,
+ use_exec=task_process_uses_exec(),
+ )
+
+ def send_call(self, call: XComBackendCall) -> None:
+ """
+ Send ``call`` to the process, which answers it through ``on_result``.
+
+ The process runs one call at a time, so a call waits its turn and is answered in the order it
+ was sent. That is also how the first call of a task waits for the process to start.
+ """
+ self._queued_calls.append(call)
+ self._unanswered_calls += 1
+ self._send_next_call()
+
+ def _send_next_call(self) -> None:
+ """Hand the oldest waiting call to the process, if it has asked for one."""
+ if self._call_request_id is None or not self._queued_calls:
+ return
+ request_id, self._call_request_id = self._call_request_id, None
+ self.send_msg(self._queued_calls.popleft(), request_id=request_id)
+
+ def stop(self) -> None:
+ """Let the helper process exit, and kill it if it does not exit in time."""
+ # The task this process serves is over, so nothing waits for the calls it has not answered.
+ self._queued_calls.clear()
+ self._unanswered_calls = 0
+ if self._exit_code is None and self._call_request_id is not None:
+ # The process may be dead already, such as when the task it serves was killed.
+ with contextlib.suppress(OSError):
+ self.send_msg(XComBackendCall(hook="exit"), request_id=self._call_request_id)
+ self._call_request_id = None
+ deadline = time.monotonic() + _EXIT_TIMEOUT
+ while self._open_sockets and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ if self._exit_code is None:
+ with contextlib.suppress(self._process.TimeoutExpired):
+ self._exit_code = self._process.wait(timeout=max(0.0, deadline - time.monotonic()))
+ if self._exit_code is None:
+ self.kill(signal.SIGTERM, force=True)
+ # Only this process's own sockets, since the selector belongs to the task it serves.
+ for sock in list(self._open_sockets):
+ self._on_socket_closed(sock)
+ sock.close()
+
+ def _on_socket_closed(self, sock: socket) -> None:
+ super()._on_socket_closed(sock)
+ if sock is self.stdin:
+ # The process has exited, or is exiting, so it answers none of the calls it was still to run.
+ unanswered, self._unanswered_calls = self._unanswered_calls, 0
+ for _ in range(unanswered):
+ self.on_result(XComBackendResult(error="The XCom backend process exited before it answered"))
+
+ def _handle_await_call(
+ self, msg: AwaitXComBackendCall, log: FilteringBoundLogger, req_id: int
+ ) -> ResponseSent:
+ self._call_request_id = req_id
+ self._send_next_call()
+ return ResponseSent.ALREADY_SENT
+
+ def _handle_result(self, msg: XComBackendResult, log: FilteringBoundLogger, req_id: int) -> ResponseSent:
+ self._call_request_id = req_id
+ self._unanswered_calls -= 1
+ self._send_next_call()
+ self.on_result(msg)
+ return ResponseSent.ALREADY_SENT
+
+ _request_handlers = {
+ **WatchedSubprocess._get_shared_request_handlers(GetConnection, GetVariable, MaskSecret),
+ **dict(
+ [
+ register_request_method(AwaitXComBackendCall, _handle_await_call),
+ register_request_method(XComBackendResult, _handle_result),
+ ]
+ ),
+ }
+
+
+def run_xcom_backend() -> None:
+ """Run the XCom backend calls that the supervisor sends, until it sends ``exit``."""
+ # Lazy: these load the task runtime and the configured backend, which only the helper process needs.
+ from airflow.sdk.execution_time import task_runner
+ from airflow.sdk.execution_time.comms import CommsDecoder
+ from airflow.sdk.execution_time.xcom import XCom
+
+ comms = CommsDecoder[_ToXComBackendProcess, _FromXComBackendProcess](
+ log=log, body_decoder=TypeAdapter(_ToXComBackendProcess)
+ )
+ # The secrets backends of a task process look up Connections and Variables through SUPERVISOR_COMMS.
+ task_runner.SUPERVISOR_COMMS = comms # type: ignore[assignment]
+ runner = _HookRunner(XCom)
+ call = comms.send(AwaitXComBackendCall())
+ while isinstance(call, XComBackendCall) and call.hook != "exit":
+ call = comms.send(runner.run(call))
+
+
+@attrs.define
+class _HookRunner:
+ """Runs the hooks of an XCom backend that the supervisor's calls name, one method per hook."""
+
+ xcom: type[BaseXCom]
+
+ def run(self, call: XComBackendCall) -> XComBackendResult:
+ try:
+ if (hook := getattr(self, call.hook, None)) is None:
+ raise ValueError(f"Unknown XCom backend hook {call.hook!r}")
+ return hook(call)
+ except Exception as e:
+ log.exception("XCom backend call failed", hook=call.hook, key=call.key)
+ return XComBackendResult(error=f"{type(e).__name__}: {e}")
+
+ def store(self, call: XComBackendCall) -> XComBackendResult:
+ # Lazy: only the helper process decodes and encodes XCom values.
+ from airflow.sdk.serde import deserialize
+
+ if call.coordinate is None:
+ raise ValueError("A store call needs the coordinate of the XCom")
+ try:
+ value = deserialize(call.value) # type: ignore[type-var]
+ except (ImportError, TypeError, AttributeError, ValueError) as e:
+ # The default backend would keep such a value in the XCom row as it is.
+ log.warning(
+ "Cannot decode the XCom value; storing it without the XCom backend", key=call.key, error=e
+ )
+ return XComBackendResult(unchanged=True)
Review Comment:
Falling back to storing the value in the DB defeats the purpose of the backend. People configure object storage precisely to keep payloads out of the metadata DB, and some for compliance. A Python task would fail here, so this should fail the push too. Raising is probably better.
Same for the block in `_load_row`.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397262589
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
Review Comment:
There is a single pending slot, and `XComBackendProcess.send_call` has a single queue slot. Some runtimes multiplex requests, and everything gets garbled. Calls need to be keyed by request id with a FIFO queue, plus a test that sends two overlapping requests.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on PR #75184:
URL: https://github.com/apache/airflow/pull/75184#issuecomment-6335691374
Agreed, and it's a good idea. There is no need for the XCom backend process to heartbeat or track execution_timeout itself, and no need for a second loop to watch it, I just simplified into a single loop.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397262589
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
Review Comment:
There is a single pending slot, and `XComBackendProcess.send_call` has a single queue slot. Some runtimes multiplex requests, and everything gets garbled. Calls need to be keyed by request id with a FIFO queue, plus a test that sends two overlapping requests.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #75184:
URL: https://github.com/apache/airflow/pull/75184#issuecomment-6292056954
Having a whole separate process to handle xcom uplodads feels overkill. Why can't we do it from the coordinator process directly? The "why" in the PR desc doesn't address 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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397346704
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -853,11 +877,20 @@ def _find_own_coordinate(dag_id: str, run_id: str, task_id: str, map_index: int)
owner = running
else:
return None
- if (owner.dag_id, owner.run_id, owner.task_id, _public_map_index(owner)) == (
- dag_id,
- run_id,
- task_id,
- map_index,
- ):
- return _own_coordinate(owner)
+ if is_own_xcom(owner, dag_id, run_id, task_id, map_index):
+ return build_own_coordinate(owner)
return None
+
+
+def is_stored_through_backend(key: str) -> bool:
Review Comment:
This also changes Python tasks. Is it intended?
The prefix check is also a heuristic. Many provider links override xcom_key, so they aren’t covered, and a user key that happens to start with _link_ gets moved into the DB without notice. Could the link declare this itself, e.g. by having `BaseOperatorLink.persist` write past the backend, instead of us guessing from the key?
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397312950
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -425,23 +474,38 @@ def _on_child_started(
# Language runtimes do not send StartExecutionTimeout, so the clock starts once they have the task.
self._arm_execution_timeout()
+ _xcom_backend: XComBackendProcess | None = attrs.field(default=None, init=False)
+ _pending_xcom_calls: deque[Callable[[XComBackendResult], None]] = attrs.field(factory=deque, init=False)
Review Comment:
Matching answers to requests is still positional. It works while the helper runs one call at a time, but _queued_calls and _unanswered_calls in XComBackendProcess must stay in sync with this deque.
Mayne make `XComBackendProcess.send_call(call, finish)` own the `(call, finish)` pairs instead, so the supervisor doesn’t keep its own deque? Or put an id on the XCom request/response and look the callback up by id instead (Similar to how the supervisor does duplex).
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -425,23 +474,38 @@ def _on_child_started(
# Language runtimes do not send StartExecutionTimeout, so the clock starts once they have the task.
self._arm_execution_timeout()
+ _xcom_backend: XComBackendProcess | None = attrs.field(default=None, init=False)
+ _pending_xcom_calls: deque[Callable[[XComBackendResult], None]] = attrs.field(factory=deque, init=False)
Review Comment:
Matching answers to requests is still positional. It works while the helper runs one call at a time, but _queued_calls and _unanswered_calls in XComBackendProcess must stay in sync with this deque.
Mayne make `XComBackendProcess.send_call(call, finish)` own the `(call, finish)` pairs instead, so the supervisor doesn’t keep its own deque? Or put an id on the XCom request/response and look the callback up by id instead (similar to how the supervisor does duplex).
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4410038021
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -425,23 +474,38 @@ def _on_child_started(
# Language runtimes do not send StartExecutionTimeout, so the clock starts once they have the task.
self._arm_execution_timeout()
+ _xcom_backend: XComBackendProcess | None = attrs.field(default=None, init=False)
+ _pending_xcom_calls: deque[Callable[[XComBackendResult], None]] = attrs.field(factory=deque, init=False)
Review Comment:
Done, I want with the first option. The supervisor keeps no deque now: `send_call(call, finish)` takes the continuation and `XComBackendProcess` owns the pairing in one `deque[tuple[XComBackendCall, Callable]]`.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
+ self._pending_xcom_calls.append(finish)
+ self._xcom_backend.send_call(call)
+
+ def _finish_xcom_backend_call(self, result: XComBackendResult) -> None:
+ """Finish the oldest call the XCom backend process has not answered yet."""
+ if not self._pending_xcom_calls:
+ return
+ self._pending_xcom_calls.popleft()(result)
+
+ def _call_xcom_backend(
+ self,
+ call: XComBackendCall,
+ then: Callable[[XComBackendResult], RequestResult],
+ req_id: int,
+ *,
+ on_error: Callable[[str], RequestResult] | None = None,
+ ) -> ResponseSent:
+ """
+ Answer request ``req_id`` with ``then(result)`` once the XCom backend has run ``call``.
+
+ The request loop goes on servicing the task and the backend process while the call runs, so it
+ keeps heartbeating and stops the task when it reaches its ``execution_timeout``.
+
+ ``on_error`` answers a failed call instead; without it, a failed call fails the request.
+ """
+
+ def finish(result: XComBackendResult) -> None:
+ try:
+ if result.error is None:
+ resp, dump_opts = then(result)
+ elif on_error is not None:
+ resp, dump_opts = on_error(result.error)
+ else:
+ raise XComBackendError(result.error)
+ except Exception as e:
+ self.process_log.exception("Failed to handle the XCom backend result", request_id=req_id)
+ self.send_msg(None, request_id=req_id, error=_build_request_error(e))
+ else:
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+
+ self._start_xcom_backend_call(call, finish)
+ return ResponseSent.ALREADY_SENT
+
+ def _run_xcom_backend_call(self, call: XComBackendCall) -> None:
+ """
+ Run ``call`` and wait for it, for a caller with no request loop to return to.
+
+ :raises XComBackendError: when the call fails or is not answered in time.
+ """
+ answer: list[XComBackendResult] = []
+ self._start_xcom_backend_call(call, answer.append)
+ deadline = time.monotonic() + _XCOM_BACKEND_CALL_TIMEOUT
+ # The task process has not been told to start, so servicing it here reads no request of its own.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ self._pending_xcom_calls.clear()
Review Comment:
Done. `_run_xcom_backend_call` stops the helper before it raises:
```python
if not answer:
# The process is still running the call, and what it does next is unknown, so it is replaced.
self._stop_xcom_backend()
raise XComBackendError("The XCom backend did not answer in time")
```
The next call starts a fresh one. Covered by `test_stops_a_backend_process_that_does_not_answer_a_waited_call`.
I didn't handle the "per XCom operation level deadline" (we could handle it as follow-up but not in this PR), since there's "task-level" execution_timeout that will terminate the entire subprocesses already.
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -853,11 +877,20 @@ def _find_own_coordinate(dag_id: str, run_id: str, task_id: str, map_index: int)
owner = running
else:
return None
- if (owner.dag_id, owner.run_id, owner.task_id, _public_map_index(owner)) == (
- dag_id,
- run_id,
- task_id,
- map_index,
- ):
- return _own_coordinate(owner)
+ if is_own_xcom(owner, dag_id, run_id, task_id, map_index):
+ return build_own_coordinate(owner)
return None
+
+
+def is_stored_through_backend(key: str) -> bool:
Review Comment:
It was intended at first, but you are right that it should not be a breaking change, so I reverted it.
Keep the `_link_*` as `_PopenActivitySubprocess` only, no impact for the existing Python task.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
Review Comment:
Fixed, restarting rather than failing now. `XComBackendProcess` reports liveness through the socket instead:
```python
@property
def can_run_a_call(self) -> bool:
"""Whether the process can still be sent a call, which a process that has exited cannot."""
return self.stdin in self._open_sockets
```
and `_start_xcom_backend_call` replaces a process that cannot answer:
```python
if self._xcom_backend is not None and not self._xcom_backend.can_run_a_call:
self._stop_xcom_backend()
if self._xcom_backend is None:
self._xcom_backend = XComBackendProcess.start(...)
```
`_on_socket_closed` still fails the calls already waiting, so nothing hangs in the window before the replacement. Covered by `test_replaces_a_backend_process_that_cannot_answer` and `test_a_closed_requests_socket_fails_every_call_that_waits`.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
+ self._pending_xcom_calls.append(finish)
+ self._xcom_backend.send_call(call)
+
+ def _finish_xcom_backend_call(self, result: XComBackendResult) -> None:
+ """Finish the oldest call the XCom backend process has not answered yet."""
+ if not self._pending_xcom_calls:
+ return
+ self._pending_xcom_calls.popleft()(result)
+
+ def _call_xcom_backend(
+ self,
+ call: XComBackendCall,
+ then: Callable[[XComBackendResult], RequestResult],
+ req_id: int,
+ *,
+ on_error: Callable[[str], RequestResult] | None = None,
+ ) -> ResponseSent:
+ """
+ Answer request ``req_id`` with ``then(result)`` once the XCom backend has run ``call``.
+
+ The request loop goes on servicing the task and the backend process while the call runs, so it
+ keeps heartbeating and stops the task when it reaches its ``execution_timeout``.
+
+ ``on_error`` answers a failed call instead; without it, a failed call fails the request.
+ """
+
+ def finish(result: XComBackendResult) -> None:
+ try:
+ if result.error is None:
+ resp, dump_opts = then(result)
+ elif on_error is not None:
+ resp, dump_opts = on_error(result.error)
+ else:
+ raise XComBackendError(result.error)
+ except Exception as e:
+ self.process_log.exception("Failed to handle the XCom backend result", request_id=req_id)
+ self.send_msg(None, request_id=req_id, error=_build_request_error(e))
+ else:
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+
+ self._start_xcom_backend_call(call, finish)
+ return ResponseSent.ALREADY_SENT
+
+ def _run_xcom_backend_call(self, call: XComBackendCall) -> None:
+ """
+ Run ``call`` and wait for it, for a caller with no request loop to return to.
+
+ :raises XComBackendError: when the call fails or is not answered in time.
+ """
+ answer: list[XComBackendResult] = []
+ self._start_xcom_backend_call(call, answer.append)
+ deadline = time.monotonic() + _XCOM_BACKEND_CALL_TIMEOUT
+ # The task process has not been told to start, so servicing it here reads no request of its own.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ self._pending_xcom_calls.clear()
+ if not answer:
+ raise XComBackendError("The XCom backend did not answer in time")
+ if (error := answer[0].error) is not None:
+ raise XComBackendError(error)
+
+ def _handle_set_xcom(
+ self, msg: SetXCom, log: FilteringBoundLogger, req_id: int
+ ) -> RequestResult | ResponseSent:
+ xcom = _get_xcom_backend()
+ if not (
+ is_stored_through_backend(msg.key) and is_hook_overridden(xcom, "store_value", "serialize_value")
+ ):
+ return handle_set_xcom(self.client, msg)
+ ti = cast("_OwnTaskInstance", self.ti)
+ if is_own_xcom(ti, msg.dag_id, msg.run_id, msg.task_id, msg.resolved_map_index):
+ coordinate = build_own_coordinate(ti)
+ else:
+ coordinate = XComCoordinate(msg.dag_id, msg.run_id, msg.task_id)
+ # A loop iteration has map index -1, unlike a mapped task. A mapped task inside a loop is not
+ # caught here because the supervisor does not parse the Dag, while Python rejects such a Dag
+ # when it is parsed.
+ if coordinate.region_id is not None and msg.resolved_map_index == -1 and is_legacy_xcom_backend(xcom):
+ raise build_loop_error(xcom, f"Task {msg.task_id!r}")
Review Comment:
Yes we should. Added `ErrorType.XCOM_BACKEND_ERROR`.
##########
task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py:
##########
@@ -0,0 +1,345 @@
+#
+# 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.
+"""
+The XCom backend for the tasks of a language SDK runtime.
+
+A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend
+on them before it forwards them to the API server. The backend runs in a helper process, started per task
+the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of
+the backend's state, such as a cached filesystem that a task process forked from it later could not use.
+
+The supervisor services the helper process from the selector of the task it serves, so a backend call
+never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import selectors
+import signal
+import time
+from collections import deque
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Annotated, ClassVar, Literal
+from uuid import UUID
+
+import attrs
+import structlog
+from pydantic import BaseModel, Field, JsonValue, TypeAdapter
+
+from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate
+from airflow.sdk.execution_time.comms import (
+ ConnectionResult,
+ ErrorResponse,
+ GetConnection,
+ GetVariable,
+ MaskSecret,
+ VariableResult,
+ XComResult,
+)
+from airflow.sdk.execution_time.supervisor import (
+ ResponseSent,
+ WatchedSubprocess,
+ register_request_method,
+ task_process_uses_exec,
+)
+
+if TYPE_CHECKING:
+ from socket import socket
+ from typing import Self
+
+ from structlog.typing import FilteringBoundLogger
+
+ from airflow.sdk.api.client import Client
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend")
+
+_EXIT_TIMEOUT = 10.0
+
+
+class XComBackendError(RuntimeError):
+ """
+ An XCom backend call failed.
+
+ For an exception raised in the helper process, the message holds its type and message, since the
+ exception itself does not cross the process boundary.
+ """
+
+
+class _Coordinate(BaseModel):
+ dag_id: str
+ run_id: str
+ task_id: str
+ ti_id: UUID | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
+
+
+class XComBackendCall(BaseModel):
+ """
+ A call of an XCom backend hook, sent to the helper process.
+
+ ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``,
+ and a list of rows for ``load_many``.
+ """
+
+ hook: Literal["store", "load", "load_many", "purge", "exit"]
+ key: str = ""
+ value: JsonValue = None
+ coordinate: _Coordinate | None = None
+ map_index: int = -1
+ type: Literal["XComBackendCall"] = "XComBackendCall"
+
+ @classmethod
+ def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self:
+ return cls(
+ hook="store",
+ key=key,
+ value=value,
+ coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)),
+ map_index=map_index,
+ )
+
+
+class XComBackendResult(BaseModel):
+ """
+ The outcome of an XCom backend call, sent by the helper process.
+
+ ``unchanged`` means the backend did not store the value, which then goes to the API server as it is.
+ """
+
+ value: JsonValue = None
+ unchanged: bool = False
+ error: str | None = None
+ type: Literal["XComBackendResult"] = "XComBackendResult"
+
+
+class AwaitXComBackendCall(BaseModel):
+ """Sent by the helper process when it is ready for its first call."""
+
+ type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall"
+
+
+_ToXComBackendProcess = Annotated[
+ XComBackendCall | ConnectionResult | VariableResult | ErrorResponse,
+ Field(discriminator="type"),
+]
+_FromXComBackendProcess = Annotated[
+ AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret,
+ Field(discriminator="type"),
+]
+
+
+def is_hook_overridden(xcom: type[BaseXCom], *hooks: str) -> bool:
+ """Whether ``xcom`` overrides any of ``hooks``, and so may do more than the default backend."""
+ return xcom is not BaseXCom and any(
+ next(k for k in xcom.__mro__ if hook in vars(k)) is not BaseXCom for hook in hooks
+ )
+
+
+@attrs.define(kw_only=True)
+class XComBackendProcess(WatchedSubprocess):
+ """
+ The helper process that runs XCom backend calls for the task of a language SDK runtime.
+
+ The process asks for a call, runs it, and answers with its result, which is also how it asks for the
+ next call. Each result reaches the supervisor through ``on_result``. Connection and Variable lookups of
+ the backend reach the API server through ``client``.
+ """
+
+ client: Client
+ on_result: Callable[[XComBackendResult], None] = attrs.field(repr=False)
+ """Called with the result of a call, and with an error when the process exits without answering one."""
+ decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess)
+
+ _call_request_id: int | None = attrs.field(default=None, init=False)
+ _queued_calls: deque[XComBackendCall] = attrs.field(factory=deque, init=False)
+ _unanswered_calls: int = attrs.field(default=0, init=False)
+
+ @classmethod
+ def start( # type: ignore[override]
+ cls,
+ *,
+ id: UUID,
+ client: Client,
+ logger: FilteringBoundLogger,
+ selector: selectors.BaseSelector,
+ on_result: Callable[[XComBackendResult], None],
+ ) -> Self:
+ return super().start(
+ id=id,
+ client=client,
+ on_result=on_result,
+ # The selector of the task this process serves, so that the task's request loop services both.
+ selector=selector,
+ logger=logger,
+ target=run_xcom_backend,
+ use_exec=task_process_uses_exec(),
+ )
+
+ def send_call(self, call: XComBackendCall) -> None:
+ """
+ Send ``call`` to the process, which answers it through ``on_result``.
+
+ The process runs one call at a time, so a call waits its turn and is answered in the order it
+ was sent. That is also how the first call of a task waits for the process to start.
+ """
+ self._queued_calls.append(call)
+ self._unanswered_calls += 1
+ self._send_next_call()
+
+ def _send_next_call(self) -> None:
+ """Hand the oldest waiting call to the process, if it has asked for one."""
+ if self._call_request_id is None or not self._queued_calls:
+ return
+ request_id, self._call_request_id = self._call_request_id, None
+ self.send_msg(self._queued_calls.popleft(), request_id=request_id)
+
+ def stop(self) -> None:
+ """Let the helper process exit, and kill it if it does not exit in time."""
+ # The task this process serves is over, so nothing waits for the calls it has not answered.
+ self._queued_calls.clear()
+ self._unanswered_calls = 0
+ if self._exit_code is None and self._call_request_id is not None:
+ # The process may be dead already, such as when the task it serves was killed.
+ with contextlib.suppress(OSError):
+ self.send_msg(XComBackendCall(hook="exit"), request_id=self._call_request_id)
+ self._call_request_id = None
+ deadline = time.monotonic() + _EXIT_TIMEOUT
+ while self._open_sockets and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ if self._exit_code is None:
+ with contextlib.suppress(self._process.TimeoutExpired):
+ self._exit_code = self._process.wait(timeout=max(0.0, deadline - time.monotonic()))
+ if self._exit_code is None:
+ self.kill(signal.SIGTERM, force=True)
+ # Only this process's own sockets, since the selector belongs to the task it serves.
+ for sock in list(self._open_sockets):
+ self._on_socket_closed(sock)
+ sock.close()
+
+ def _on_socket_closed(self, sock: socket) -> None:
+ super()._on_socket_closed(sock)
+ if sock is self.stdin:
+ # The process has exited, or is exiting, so it answers none of the calls it was still to run.
+ unanswered, self._unanswered_calls = self._unanswered_calls, 0
+ for _ in range(unanswered):
+ self.on_result(XComBackendResult(error="The XCom backend process exited before it answered"))
+
+ def _handle_await_call(
+ self, msg: AwaitXComBackendCall, log: FilteringBoundLogger, req_id: int
+ ) -> ResponseSent:
+ self._call_request_id = req_id
+ self._send_next_call()
+ return ResponseSent.ALREADY_SENT
+
+ def _handle_result(self, msg: XComBackendResult, log: FilteringBoundLogger, req_id: int) -> ResponseSent:
+ self._call_request_id = req_id
+ self._unanswered_calls -= 1
+ self._send_next_call()
+ self.on_result(msg)
+ return ResponseSent.ALREADY_SENT
+
+ _request_handlers = {
+ **WatchedSubprocess._get_shared_request_handlers(GetConnection, GetVariable, MaskSecret),
+ **dict(
+ [
+ register_request_method(AwaitXComBackendCall, _handle_await_call),
+ register_request_method(XComBackendResult, _handle_result),
+ ]
+ ),
+ }
+
+
+def run_xcom_backend() -> None:
+ """Run the XCom backend calls that the supervisor sends, until it sends ``exit``."""
+ # Lazy: these load the task runtime and the configured backend, which only the helper process needs.
+ from airflow.sdk.execution_time import task_runner
+ from airflow.sdk.execution_time.comms import CommsDecoder
+ from airflow.sdk.execution_time.xcom import XCom
+
+ comms = CommsDecoder[_ToXComBackendProcess, _FromXComBackendProcess](
+ log=log, body_decoder=TypeAdapter(_ToXComBackendProcess)
+ )
+ # The secrets backends of a task process look up Connections and Variables through SUPERVISOR_COMMS.
+ task_runner.SUPERVISOR_COMMS = comms # type: ignore[assignment]
+ runner = _HookRunner(XCom)
+ call = comms.send(AwaitXComBackendCall())
+ while isinstance(call, XComBackendCall) and call.hook != "exit":
+ call = comms.send(runner.run(call))
+
+
+@attrs.define
+class _HookRunner:
+ """Runs the hooks of an XCom backend that the supervisor's calls name, one method per hook."""
+
+ xcom: type[BaseXCom]
+
+ def run(self, call: XComBackendCall) -> XComBackendResult:
+ try:
+ if (hook := getattr(self, call.hook, None)) is None:
+ raise ValueError(f"Unknown XCom backend hook {call.hook!r}")
+ return hook(call)
+ except Exception as e:
+ log.exception("XCom backend call failed", hook=call.hook, key=call.key)
+ return XComBackendResult(error=f"{type(e).__name__}: {e}")
+
+ def store(self, call: XComBackendCall) -> XComBackendResult:
+ # Lazy: only the helper process decodes and encodes XCom values.
+ from airflow.sdk.serde import deserialize
+
+ if call.coordinate is None:
+ raise ValueError("A store call needs the coordinate of the XCom")
+ try:
+ value = deserialize(call.value) # type: ignore[type-var]
+ except (ImportError, TypeError, AttributeError, ValueError) as e:
+ # The default backend would keep such a value in the XCom row as it is.
+ log.warning(
+ "Cannot decode the XCom value; storing it without the XCom backend", key=call.key, error=e
+ )
+ return XComBackendResult(unchanged=True)
Review Comment:
Done, both raise. `XComBackendResult.unchanged` is gone, so no path silently keeps the payload in the row.
```python
# A value the backend cannot be given fails the push. Keeping it in the XCom row instead would
# put the payload in the database, which is what the backend is configured to avoid.
value = deserialize(call.value)
```
```python
# A row the backend cannot load, or a value this process cannot encode, fails the read rather
# than reaching the runtime as the raw row, which would hide the backend being unusable.
return serialize(self.xcom._load(XComResult(key=key, value=row)))
```
Covered by `test_load_call_fails_when_the_value_cannot_be_sent` and `test_store_call_fails_when_the_value_cannot_be_decoded`.
##########
task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py:
##########
@@ -0,0 +1,345 @@
+#
+# 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.
+"""
+The XCom backend for the tasks of a language SDK runtime.
+
+A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend
+on them before it forwards them to the API server. The backend runs in a helper process, started per task
+the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of
+the backend's state, such as a cached filesystem that a task process forked from it later could not use.
+
+The supervisor services the helper process from the selector of the task it serves, so a backend call
+never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import selectors
+import signal
+import time
+from collections import deque
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Annotated, ClassVar, Literal
+from uuid import UUID
+
+import attrs
+import structlog
+from pydantic import BaseModel, Field, JsonValue, TypeAdapter
+
+from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate
+from airflow.sdk.execution_time.comms import (
+ ConnectionResult,
+ ErrorResponse,
+ GetConnection,
+ GetVariable,
+ MaskSecret,
+ VariableResult,
+ XComResult,
+)
+from airflow.sdk.execution_time.supervisor import (
+ ResponseSent,
+ WatchedSubprocess,
+ register_request_method,
+ task_process_uses_exec,
+)
+
+if TYPE_CHECKING:
+ from socket import socket
+ from typing import Self
+
+ from structlog.typing import FilteringBoundLogger
+
+ from airflow.sdk.api.client import Client
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend")
+
+_EXIT_TIMEOUT = 10.0
+
+
+class XComBackendError(RuntimeError):
+ """
+ An XCom backend call failed.
+
+ For an exception raised in the helper process, the message holds its type and message, since the
+ exception itself does not cross the process boundary.
+ """
+
+
+class _Coordinate(BaseModel):
+ dag_id: str
+ run_id: str
+ task_id: str
+ ti_id: UUID | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
+
+
+class XComBackendCall(BaseModel):
+ """
+ A call of an XCom backend hook, sent to the helper process.
+
+ ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``,
+ and a list of rows for ``load_many``.
+ """
+
+ hook: Literal["store", "load", "load_many", "purge", "exit"]
+ key: str = ""
+ value: JsonValue = None
+ coordinate: _Coordinate | None = None
+ map_index: int = -1
+ type: Literal["XComBackendCall"] = "XComBackendCall"
+
+ @classmethod
+ def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self:
+ return cls(
+ hook="store",
+ key=key,
+ value=value,
+ coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)),
+ map_index=map_index,
+ )
+
+
+class XComBackendResult(BaseModel):
+ """
+ The outcome of an XCom backend call, sent by the helper process.
+
+ ``unchanged`` means the backend did not store the value, which then goes to the API server as it is.
+ """
+
+ value: JsonValue = None
+ unchanged: bool = False
+ error: str | None = None
+ type: Literal["XComBackendResult"] = "XComBackendResult"
+
+
+class AwaitXComBackendCall(BaseModel):
+ """Sent by the helper process when it is ready for its first call."""
+
+ type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall"
+
+
+_ToXComBackendProcess = Annotated[
+ XComBackendCall | ConnectionResult | VariableResult | ErrorResponse,
+ Field(discriminator="type"),
+]
+_FromXComBackendProcess = Annotated[
+ AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret,
+ Field(discriminator="type"),
+]
+
+
+def is_hook_overridden(xcom: type[BaseXCom], *hooks: str) -> bool:
+ """Whether ``xcom`` overrides any of ``hooks``, and so may do more than the default backend."""
+ return xcom is not BaseXCom and any(
+ next(k for k in xcom.__mro__ if hook in vars(k)) is not BaseXCom for hook in hooks
+ )
Review Comment:
Took the first option. `is_hook_overridden` and its MRO walk are gone:
```python
def runs_in_helper(xcom: type[BaseXCom]) -> bool:
"""Whether task XComs go through the helper process, which runs every backend but the default one."""
return xcom is not BaseXCom
```
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397312950
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -425,23 +474,38 @@ def _on_child_started(
# Language runtimes do not send StartExecutionTimeout, so the clock starts once they have the task.
self._arm_execution_timeout()
+ _xcom_backend: XComBackendProcess | None = attrs.field(default=None, init=False)
+ _pending_xcom_calls: deque[Callable[[XComBackendResult], None]] = attrs.field(factory=deque, init=False)
Review Comment:
Matching answers to requests is still positional. It works while the helper runs one call at a time, but _queued_calls and _unanswered_calls in XComBackendProcess must stay in sync with this deque.
Mayne make `XComBackendProcess.send_call(call, finish)` own the `(call, finish)` pairs instead, so the supervisor doesn’t keep its own deque? Or put an id on the XCom request/response and look the callback up by id instead.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397346704
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -853,11 +877,20 @@ def _find_own_coordinate(dag_id: str, run_id: str, task_id: str, map_index: int)
owner = running
else:
return None
- if (owner.dag_id, owner.run_id, owner.task_id, _public_map_index(owner)) == (
- dag_id,
- run_id,
- task_id,
- map_index,
- ):
- return _own_coordinate(owner)
+ if is_own_xcom(owner, dag_id, run_id, task_id, map_index):
+ return build_own_coordinate(owner)
return None
+
+
+def is_stored_through_backend(key: str) -> bool:
Review Comment:
This also changes Python tasks. Is it intended?
The prefix check is also a heuristic. Many provider links override xcom_key, so they aren’t covered, and a user key that happens to start with `_link_` gets moved into the DB. Could the link declare this itself, e.g. by having `BaseOperatorLink.persist` write past the backend, instead of us guessing from the key?
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397355996
##########
task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py:
##########
@@ -0,0 +1,345 @@
+#
+# 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.
+"""
+The XCom backend for the tasks of a language SDK runtime.
+
+A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend
+on them before it forwards them to the API server. The backend runs in a helper process, started per task
+the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of
+the backend's state, such as a cached filesystem that a task process forked from it later could not use.
+
+The supervisor services the helper process from the selector of the task it serves, so a backend call
+never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import selectors
+import signal
+import time
+from collections import deque
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Annotated, ClassVar, Literal
+from uuid import UUID
+
+import attrs
+import structlog
+from pydantic import BaseModel, Field, JsonValue, TypeAdapter
+
+from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate
+from airflow.sdk.execution_time.comms import (
+ ConnectionResult,
+ ErrorResponse,
+ GetConnection,
+ GetVariable,
+ MaskSecret,
+ VariableResult,
+ XComResult,
+)
+from airflow.sdk.execution_time.supervisor import (
+ ResponseSent,
+ WatchedSubprocess,
+ register_request_method,
+ task_process_uses_exec,
+)
+
+if TYPE_CHECKING:
+ from socket import socket
+ from typing import Self
+
+ from structlog.typing import FilteringBoundLogger
+
+ from airflow.sdk.api.client import Client
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend")
+
+_EXIT_TIMEOUT = 10.0
+
+
+class XComBackendError(RuntimeError):
+ """
+ An XCom backend call failed.
+
+ For an exception raised in the helper process, the message holds its type and message, since the
+ exception itself does not cross the process boundary.
+ """
+
+
+class _Coordinate(BaseModel):
+ dag_id: str
+ run_id: str
+ task_id: str
+ ti_id: UUID | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
+
+
+class XComBackendCall(BaseModel):
+ """
+ A call of an XCom backend hook, sent to the helper process.
+
+ ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``,
+ and a list of rows for ``load_many``.
+ """
+
+ hook: Literal["store", "load", "load_many", "purge", "exit"]
+ key: str = ""
+ value: JsonValue = None
+ coordinate: _Coordinate | None = None
+ map_index: int = -1
+ type: Literal["XComBackendCall"] = "XComBackendCall"
+
+ @classmethod
+ def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self:
+ return cls(
+ hook="store",
+ key=key,
+ value=value,
+ coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)),
+ map_index=map_index,
+ )
+
+
+class XComBackendResult(BaseModel):
+ """
+ The outcome of an XCom backend call, sent by the helper process.
+
+ ``unchanged`` means the backend did not store the value, which then goes to the API server as it is.
+ """
+
+ value: JsonValue = None
+ unchanged: bool = False
+ error: str | None = None
+ type: Literal["XComBackendResult"] = "XComBackendResult"
+
+
+class AwaitXComBackendCall(BaseModel):
+ """Sent by the helper process when it is ready for its first call."""
+
+ type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall"
+
+
+_ToXComBackendProcess = Annotated[
+ XComBackendCall | ConnectionResult | VariableResult | ErrorResponse,
+ Field(discriminator="type"),
+]
+_FromXComBackendProcess = Annotated[
+ AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret,
+ Field(discriminator="type"),
+]
+
+
+def is_hook_overridden(xcom: type[BaseXCom], *hooks: str) -> bool:
+ """Whether ``xcom`` overrides any of ``hooks``, and so may do more than the default backend."""
+ return xcom is not BaseXCom and any(
+ next(k for k in xcom.__mro__ if hook in vars(k)) is not BaseXCom for hook in hooks
+ )
+
+
+@attrs.define(kw_only=True)
+class XComBackendProcess(WatchedSubprocess):
+ """
+ The helper process that runs XCom backend calls for the task of a language SDK runtime.
+
+ The process asks for a call, runs it, and answers with its result, which is also how it asks for the
+ next call. Each result reaches the supervisor through ``on_result``. Connection and Variable lookups of
+ the backend reach the API server through ``client``.
+ """
+
+ client: Client
+ on_result: Callable[[XComBackendResult], None] = attrs.field(repr=False)
+ """Called with the result of a call, and with an error when the process exits without answering one."""
+ decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess)
+
+ _call_request_id: int | None = attrs.field(default=None, init=False)
+ _queued_calls: deque[XComBackendCall] = attrs.field(factory=deque, init=False)
+ _unanswered_calls: int = attrs.field(default=0, init=False)
+
+ @classmethod
+ def start( # type: ignore[override]
+ cls,
+ *,
+ id: UUID,
+ client: Client,
+ logger: FilteringBoundLogger,
+ selector: selectors.BaseSelector,
+ on_result: Callable[[XComBackendResult], None],
+ ) -> Self:
+ return super().start(
+ id=id,
+ client=client,
+ on_result=on_result,
+ # The selector of the task this process serves, so that the task's request loop services both.
+ selector=selector,
+ logger=logger,
+ target=run_xcom_backend,
+ use_exec=task_process_uses_exec(),
+ )
+
+ def send_call(self, call: XComBackendCall) -> None:
+ """
+ Send ``call`` to the process, which answers it through ``on_result``.
+
+ The process runs one call at a time, so a call waits its turn and is answered in the order it
+ was sent. That is also how the first call of a task waits for the process to start.
+ """
+ self._queued_calls.append(call)
+ self._unanswered_calls += 1
+ self._send_next_call()
+
+ def _send_next_call(self) -> None:
+ """Hand the oldest waiting call to the process, if it has asked for one."""
+ if self._call_request_id is None or not self._queued_calls:
+ return
+ request_id, self._call_request_id = self._call_request_id, None
+ self.send_msg(self._queued_calls.popleft(), request_id=request_id)
+
+ def stop(self) -> None:
+ """Let the helper process exit, and kill it if it does not exit in time."""
+ # The task this process serves is over, so nothing waits for the calls it has not answered.
+ self._queued_calls.clear()
+ self._unanswered_calls = 0
+ if self._exit_code is None and self._call_request_id is not None:
+ # The process may be dead already, such as when the task it serves was killed.
+ with contextlib.suppress(OSError):
+ self.send_msg(XComBackendCall(hook="exit"), request_id=self._call_request_id)
+ self._call_request_id = None
+ deadline = time.monotonic() + _EXIT_TIMEOUT
+ while self._open_sockets and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=deadline - now)
+ if self._exit_code is None:
+ with contextlib.suppress(self._process.TimeoutExpired):
+ self._exit_code = self._process.wait(timeout=max(0.0, deadline - time.monotonic()))
+ if self._exit_code is None:
+ self.kill(signal.SIGTERM, force=True)
+ # Only this process's own sockets, since the selector belongs to the task it serves.
+ for sock in list(self._open_sockets):
+ self._on_socket_closed(sock)
+ sock.close()
+
+ def _on_socket_closed(self, sock: socket) -> None:
+ super()._on_socket_closed(sock)
+ if sock is self.stdin:
+ # The process has exited, or is exiting, so it answers none of the calls it was still to run.
+ unanswered, self._unanswered_calls = self._unanswered_calls, 0
+ for _ in range(unanswered):
+ self.on_result(XComBackendResult(error="The XCom backend process exited before it answered"))
+
+ def _handle_await_call(
+ self, msg: AwaitXComBackendCall, log: FilteringBoundLogger, req_id: int
+ ) -> ResponseSent:
+ self._call_request_id = req_id
+ self._send_next_call()
+ return ResponseSent.ALREADY_SENT
+
+ def _handle_result(self, msg: XComBackendResult, log: FilteringBoundLogger, req_id: int) -> ResponseSent:
+ self._call_request_id = req_id
+ self._unanswered_calls -= 1
+ self._send_next_call()
+ self.on_result(msg)
+ return ResponseSent.ALREADY_SENT
+
+ _request_handlers = {
+ **WatchedSubprocess._get_shared_request_handlers(GetConnection, GetVariable, MaskSecret),
+ **dict(
+ [
+ register_request_method(AwaitXComBackendCall, _handle_await_call),
+ register_request_method(XComBackendResult, _handle_result),
+ ]
+ ),
+ }
+
+
+def run_xcom_backend() -> None:
+ """Run the XCom backend calls that the supervisor sends, until it sends ``exit``."""
+ # Lazy: these load the task runtime and the configured backend, which only the helper process needs.
+ from airflow.sdk.execution_time import task_runner
+ from airflow.sdk.execution_time.comms import CommsDecoder
+ from airflow.sdk.execution_time.xcom import XCom
+
+ comms = CommsDecoder[_ToXComBackendProcess, _FromXComBackendProcess](
+ log=log, body_decoder=TypeAdapter(_ToXComBackendProcess)
+ )
+ # The secrets backends of a task process look up Connections and Variables through SUPERVISOR_COMMS.
+ task_runner.SUPERVISOR_COMMS = comms # type: ignore[assignment]
+ runner = _HookRunner(XCom)
+ call = comms.send(AwaitXComBackendCall())
+ while isinstance(call, XComBackendCall) and call.hook != "exit":
+ call = comms.send(runner.run(call))
+
+
+@attrs.define
+class _HookRunner:
+ """Runs the hooks of an XCom backend that the supervisor's calls name, one method per hook."""
+
+ xcom: type[BaseXCom]
+
+ def run(self, call: XComBackendCall) -> XComBackendResult:
+ try:
+ if (hook := getattr(self, call.hook, None)) is None:
+ raise ValueError(f"Unknown XCom backend hook {call.hook!r}")
+ return hook(call)
+ except Exception as e:
+ log.exception("XCom backend call failed", hook=call.hook, key=call.key)
+ return XComBackendResult(error=f"{type(e).__name__}: {e}")
+
+ def store(self, call: XComBackendCall) -> XComBackendResult:
+ # Lazy: only the helper process decodes and encodes XCom values.
+ from airflow.sdk.serde import deserialize
+
+ if call.coordinate is None:
+ raise ValueError("A store call needs the coordinate of the XCom")
+ try:
+ value = deserialize(call.value) # type: ignore[type-var]
+ except (ImportError, TypeError, AttributeError, ValueError) as e:
+ # The default backend would keep such a value in the XCom row as it is.
+ log.warning(
+ "Cannot decode the XCom value; storing it without the XCom backend", key=call.key, error=e
+ )
+ return XComBackendResult(unchanged=True)
Review Comment:
Falling back to storing the value in the DB defeats the purpose of the backend. People configure object storage precisely to keep payloads out of the metadata DB, and some for compliance. A Python task would fail here, so this should fail the push too.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397368442
##########
task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py:
##########
@@ -0,0 +1,345 @@
+#
+# 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.
+"""
+The XCom backend for the tasks of a language SDK runtime.
+
+A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend
+on them before it forwards them to the API server. The backend runs in a helper process, started per task
+the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of
+the backend's state, such as a cached filesystem that a task process forked from it later could not use.
+
+The supervisor services the helper process from the selector of the task it serves, so a backend call
+never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import selectors
+import signal
+import time
+from collections import deque
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Annotated, ClassVar, Literal
+from uuid import UUID
+
+import attrs
+import structlog
+from pydantic import BaseModel, Field, JsonValue, TypeAdapter
+
+from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate
+from airflow.sdk.execution_time.comms import (
+ ConnectionResult,
+ ErrorResponse,
+ GetConnection,
+ GetVariable,
+ MaskSecret,
+ VariableResult,
+ XComResult,
+)
+from airflow.sdk.execution_time.supervisor import (
+ ResponseSent,
+ WatchedSubprocess,
+ register_request_method,
+ task_process_uses_exec,
+)
+
+if TYPE_CHECKING:
+ from socket import socket
+ from typing import Self
+
+ from structlog.typing import FilteringBoundLogger
+
+ from airflow.sdk.api.client import Client
+
+log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend")
+
+_EXIT_TIMEOUT = 10.0
+
+
+class XComBackendError(RuntimeError):
+ """
+ An XCom backend call failed.
+
+ For an exception raised in the helper process, the message holds its type and message, since the
+ exception itself does not cross the process boundary.
+ """
+
+
+class _Coordinate(BaseModel):
+ dag_id: str
+ run_id: str
+ task_id: str
+ ti_id: UUID | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
+
+
+class XComBackendCall(BaseModel):
+ """
+ A call of an XCom backend hook, sent to the helper process.
+
+ ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``,
+ and a list of rows for ``load_many``.
+ """
+
+ hook: Literal["store", "load", "load_many", "purge", "exit"]
+ key: str = ""
+ value: JsonValue = None
+ coordinate: _Coordinate | None = None
+ map_index: int = -1
+ type: Literal["XComBackendCall"] = "XComBackendCall"
+
+ @classmethod
+ def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self:
+ return cls(
+ hook="store",
+ key=key,
+ value=value,
+ coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)),
+ map_index=map_index,
+ )
+
+
+class XComBackendResult(BaseModel):
+ """
+ The outcome of an XCom backend call, sent by the helper process.
+
+ ``unchanged`` means the backend did not store the value, which then goes to the API server as it is.
+ """
+
+ value: JsonValue = None
+ unchanged: bool = False
+ error: str | None = None
+ type: Literal["XComBackendResult"] = "XComBackendResult"
+
+
+class AwaitXComBackendCall(BaseModel):
+ """Sent by the helper process when it is ready for its first call."""
+
+ type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall"
+
+
+_ToXComBackendProcess = Annotated[
+ XComBackendCall | ConnectionResult | VariableResult | ErrorResponse,
+ Field(discriminator="type"),
+]
+_FromXComBackendProcess = Annotated[
+ AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret,
+ Field(discriminator="type"),
+]
+
+
+def is_hook_overridden(xcom: type[BaseXCom], *hooks: str) -> bool:
+ """Whether ``xcom`` overrides any of ``hooks``, and so may do more than the default backend."""
+ return xcom is not BaseXCom and any(
+ next(k for k in xcom.__mro__ if hook in vars(k)) is not BaseXCom for hook in hooks
+ )
Review Comment:
This repeats the MRO walk in `_is_preferred_over`.
I’d rather not decide routing by inspecting which hooks a backend overrides. Either route every non-BaseXCom backend through the helper, or give BaseXCom an explicit, overridable class attribute or method for this, so backend authors control 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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397335314
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -480,6 +544,222 @@ def _handle_set_task_state_store(
)
return super()._handle_set_task_state_store(msg, log, req_id)
+ def _monitor_subprocess(self):
+ try:
+ super()._monitor_subprocess()
+ finally:
+ if self._xcom_backend is not None:
+ self._xcom_backend.stop()
+
+ def _start_xcom_backend_call(
+ self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]
+ ) -> None:
+ """Send ``call`` to the XCom backend process, which starts on the first call it handles."""
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ on_result=self._finish_xcom_backend_call,
+ )
Review Comment:
If the helper dies mid-task, `_on_socket_closed` fails the calls that are waiting, but `_xcom_backend` still points at the dead process. The next call is queued on it and never answered, so the runtime hangs until its own request timeout. Check `_xcom_backend._exit_code is not None` here, and either restart the helper or fail the call right away.
--
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] Apply the configured XCom backend to language SDK task XComs [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4397312950
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -425,23 +474,38 @@ def _on_child_started(
# Language runtimes do not send StartExecutionTimeout, so the clock starts once they have the task.
self._arm_execution_timeout()
+ _xcom_backend: XComBackendProcess | None = attrs.field(default=None, init=False)
+ _pending_xcom_calls: deque[Callable[[XComBackendResult], None]] = attrs.field(factory=deque, init=False)
Review Comment:
Matching answers to requests is still positional. It works while the helper runs one call at a time, but _queued_calls and _unanswered_calls in XComBackendProcess must stay in sync with this deque.
Mayne make `XComBackendProcess.send_call(call, finish)` own the `(call, finish)` pairs instead, so the supervisor doesn’t keep its own deque? Or put an id on the XCom request/response and look the callback up by id instead.
--
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