uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4509944375
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -605,6 +657,224 @@ 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:
+ self._stop_xcom_backend()
+
+ 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 not None and not
self._xcom_backend.can_run_a_call:
+ # The process has exited.
+ self._stop_xcom_backend()
+ if self._xcom_backend is None:
+ self._xcom_backend = XComBackendProcess.start(
+ id=self.id,
+ client=self.client,
+ logger=self.process_log,
+ selector=self.selector,
+ )
+ self._xcom_backend.send_call(call, finish)
+
+ def _stop_xcom_backend(self) -> None:
+ """Stop the XCom backend process, so that the next call starts one
that can answer it."""
+ if (backend := self._xcom_backend) is not None:
+ self._xcom_backend = None
+ backend.stop()
+
+ def _call_xcom_backend(
+ self,
+ call: XComBackendCall,
+ then: Callable[[XComBackendResult], RequestResult],
+ req_id: int,
+ ) -> ResponseSent:
+ """
+ Answer request ``req_id`` with ``then(result)`` once the XCom backend
has run ``call``.
+
+ The request loop keeps servicing the task while the call runs. A call
that fails fails the request.
+ """
+
+ def answer(resp: BaseModel | None, error: ErrorResponse | None,
**dump_opts: bool) -> None:
+ try:
+ self.send_msg(resp, request_id=req_id, error=error,
**dump_opts)
+ except OSError:
+ # The runtime has exited.
Review Comment:
```suggestion
```
Same, alternatively behind `except OSError:` (but can this be more specific?
OSError is very broad)
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -605,6 +657,224 @@ 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:
+ self._stop_xcom_backend()
+
+ 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 not None and not
self._xcom_backend.can_run_a_call:
+ # The process has exited.
Review Comment:
```suggestion
```
Or this should be on Line 670 (or above) instead.
##########
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.
+
+The supervisor runs the configured XCom backend on a runtime's XCom requests,
in a helper process per task,
+so the long-lived supervisor process keeps none of the backend's state. It
services the helper from the
+task's selector, so a backend call never blocks the task's request loop.
+"""
+
+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."""
+
+ value: JsonValue = None
+ 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 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
+
+
+_PendingCall = tuple[XComBackendCall, Callable[[XComBackendResult], None]]
+
+
[email protected](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. Connection and Variable lookups of the backend reach the API
server through ``client``.
+ """
+
+ client: Client
+ decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] =
TypeAdapter(_FromXComBackendProcess)
+
+ _call_request_id: int | None = attrs.field(default=None, init=False)
+ _queued_calls: deque[_PendingCall] = attrs.field(factory=deque, init=False)
+ _running_call: Callable[[XComBackendResult], None] | None =
attrs.field(default=None, init=False)
+ _stopping: bool = attrs.field(default=False, init=False)
+
+ @classmethod
+ def start( # type: ignore[override]
+ cls,
+ *,
+ id: UUID,
+ client: Client,
+ logger: FilteringBoundLogger,
+ selector: selectors.BaseSelector,
+ ) -> Self:
+ return super().start(
+ id=id,
+ client=client,
+ # 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(),
+ )
+
+ @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
+
+ def send_call(self, call: XComBackendCall, finish:
Callable[[XComBackendResult], None]) -> None:
+ """
+ Send ``call`` to the process, which answers it by calling ``finish``
with its result.
+
+ The process runs one call at a time, so a call waits its turn.
``finish`` also takes the error of
+ a call the process never answers, such as when it exits while the call
is waiting.
+ """
+ self._queued_calls.append((call, finish))
+ self._send_next_call()
+
+ def _send_next_call(self) -> None:
+ """Hand the process its oldest waiting call, or ``exit`` once
stopping, if it asks and runs none."""
+ if self._call_request_id is None or self._running_call is not None:
+ return
+ if self._stopping:
+ request_id, self._call_request_id = self._call_request_id, None
+ # The process may be dead already.
+ with contextlib.suppress(OSError):
+ self.send_msg(XComBackendCall(hook="exit"),
request_id=request_id)
Review Comment:
Same as above.
Is it possible to catch OSError in one place instead of scattered all around?
--
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]
