uranusjr commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4439649591
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -604,6 +658,220 @@ 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, so a call sent to it would wait for an
answer that never comes.
+ 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 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``. A call that
+ fails fails the request, so the runtime fails the task.
+ """
+
+ def finish(result: XComBackendResult) -> None:
+ try:
+ if result.error is not None:
+ raise XComBackendError(result.error)
+ resp, dump_opts = then(result)
+ 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)
+ 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")
+ 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 not
_is_extra_link(msg.key) and runs_in_helper(xcom)):
+ 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):
+ error = _build_request_error(build_loop_error(xcom, f"Task
{msg.task_id!r}"))
+ self.send_msg(None, request_id=req_id, error=error)
+ return ResponseSent.ALREADY_SENT
+ return self._call_xcom_backend(
+ XComBackendCall.store(coordinate, msg.resolved_map_index, msg.key,
msg.value),
+ lambda stored: handle_set_xcom(self.client,
msg.model_copy(update={"value": stored.value})),
+ req_id,
+ )
+
+ def _handle_get_xcom(
+ self, msg: GetXCom, log: FilteringBoundLogger, req_id: int
+ ) -> RequestResult | ResponseSent:
+ result, dump_opts = handle_get_xcom(self.client, msg)
+ if not (isinstance(result, XComResult) and result.value is not None
and _is_loaded_through_backend()):
+ return result, dump_opts
+ return self._call_xcom_backend(
+ XComBackendCall(hook="load", key=msg.key, value=result.value),
+ lambda loaded: (result.model_copy(update={"value": loaded.value}),
dump_opts),
+ req_id,
+ )
+
+ def _handle_get_xcom_sequence_item(
+ self, msg: GetXComSequenceItem, log: FilteringBoundLogger, req_id: int
+ ) -> RequestResult | ResponseSent:
+ result, dump_opts = handle_get_xcom_sequence_item(self.client, msg)
+ if not (isinstance(result, XComSequenceIndexResult) and
_is_loaded_through_backend()):
+ return result, dump_opts
+ return self._call_xcom_backend(
+ XComBackendCall(hook="load", key=msg.key, value=result.root),
+ lambda loaded: (XComSequenceIndexResult(root=loaded.value),
dump_opts),
+ req_id,
+ )
+
+ def _handle_get_xcom_sequence_slice(
+ self, msg: GetXComSequenceSlice, log: FilteringBoundLogger, req_id: int
+ ) -> RequestResult | ResponseSent:
+ result, dump_opts = handle_get_xcom_sequence_slice(self.client, msg)
+ if not (isinstance(result, XComSequenceSliceResult) and result.root
and _is_loaded_through_backend()):
+ return result, dump_opts
+ return self._call_xcom_backend(
+ XComBackendCall(hook="load_many", key=msg.key, value=result.root),
+ lambda loaded: (
+ XComSequenceSliceResult(root=cast("list[JsonValue]",
loaded.value)),
+ dump_opts,
+ ),
+ req_id,
+ )
+
+ def _handle_delete_xcom(
+ self, msg: DeleteXCom, log: FilteringBoundLogger, req_id: int
+ ) -> RequestResult | ResponseSent:
+ result, purge = self._delete_xcom(msg)
+ if purge is None:
+ return result
+ # A failed purge fails the request, as it fails a Python task that
calls XCom.delete.
+ return self._call_xcom_backend(purge, lambda _: result, req_id)
+
+ def _delete_xcom(self, msg: DeleteXCom) -> tuple[RequestResult,
XComBackendCall | None]:
+ """
+ Delete an XCom, and return the call that purges it from the XCom
backend as ``XCom.delete`` does.
+
+ The read is not guarded, for the same reason ``XCom.delete`` does not
guard it: a row the
+ supervisor cannot read is a row it cannot purge, and silently skipping
the purge would leave
+ the stored data behind without anything saying so.
+ """
+ if not runs_in_helper(_get_xcom_backend()):
+ return handle_delete_xcom(self.client, msg), None
+ row, _ = handle_get_xcom(
+ self.client,
+ GetXCom(
+ key=msg.key,
+ dag_id=msg.dag_id,
+ run_id=msg.run_id,
+ task_id=msg.task_id,
+ map_index=msg.map_index,
+ ),
+ )
+ result = handle_delete_xcom(self.client, msg)
+ if isinstance(row, XComResult) and row.value is not None:
+ return result, XComBackendCall(hook="purge", key=msg.key,
value=row.value)
+ return result, None
+
+ _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[Any]]] = {
+ **ActivitySubprocess._request_handlers,
+ **dict(
+ [
+ register_request_method(DeleteXCom, _handle_delete_xcom),
+ register_request_method(GetXCom, _handle_get_xcom),
+ register_request_method(GetXComSequenceItem,
_handle_get_xcom_sequence_item),
+ register_request_method(GetXComSequenceSlice,
_handle_get_xcom_sequence_slice),
+ register_request_method(SetXCom, _handle_set_xcom),
+ ]
+ ),
+ }
+
+
+def _build_request_error(exc: Exception) -> ErrorResponse:
+ """Report ``exc`` to the runtime as an XCom backend failure rather than an
API server one."""
+ return ErrorResponse(
+ error=ErrorType.XCOM_BACKEND_ERROR,
+ detail={"status_code": None, "message": str(exc), "exception_type":
type(exc).__name__},
+ )
+
+
+def _get_xcom_backend() -> type[BaseXCom]:
+ # Lazy: resolving the configured backend imports it, which only a worker
that runs a task needs.
Review Comment:
```suggestion
```
--
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]
