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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
