jason810496 commented on code in PR #75184:
URL: https://github.com/apache/airflow/pull/75184#discussion_r4534738399
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -605,6 +660,230 @@ 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:
+ 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:
+ if not send_unless_exited(self, resp, req_id, error, **dump_opts):
+ self.process_log.warning(
+ "The runtime exited before its XCom request was answered",
request_id=req_id
+ )
+
+ 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)
+ answer(None, _build_request_error(e))
+ else:
+ answer(resp, 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 has not started, so it has no requests to read.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=min(deadline - now,
MIN_HEARTBEAT_INTERVAL))
+ self._send_heartbeat_if_needed()
Review Comment:
Yes, the wait now stops once the task is stopped, in 7a13d64. It stops when
the runtime has exited, not only on `SERVER_TERMINATED`, so a runtime that
crashed or was killed after too many failed heartbeats is covered too. When the
server stopped the task, the outcome is left to the server. Otherwise the try
fails with the exit code as the reason.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -605,6 +660,230 @@ 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:
+ 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:
+ if not send_unless_exited(self, resp, req_id, error, **dump_opts):
+ self.process_log.warning(
+ "The runtime exited before its XCom request was answered",
request_id=req_id
+ )
+
+ 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)
+ answer(None, _build_request_error(e))
+ else:
+ answer(resp, 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 has not started, so it has no requests to read.
+ while not answer and (now := time.monotonic()) < deadline:
+ self._service_subprocess(max_wait_time=min(deadline - now,
MIN_HEARTBEAT_INTERVAL))
+ self._send_heartbeat_if_needed()
+ if not answer:
+ # The process is still running the call.
+ 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. A mapped task in a loop is not
caught here, since the
+ # supervisor does not parse the Dag.
+ 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 XCom.delete does.
+ 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, as in ``XCom.delete``, since skipping the
purge would silently leave the
+ stored data behind.
+ """
+ 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, unless the
API server rejected the request."""
+ if isinstance(exc, ServerResponseError):
+ return ErrorResponse(
+ error=ErrorType.API_SERVER_ERROR,
+ detail={
+ "status_code": exc.response.status_code,
+ "message": str(exc),
+ "detail": exc.response.json(),
+ },
+ )
Review Comment:
Added `ErrorResponse.from_server_error`, and both places use it now, in
7a13d64.
--
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]
