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