def _process_future(self, future: Future) -> None:
if not self.futures:
raise AirflowException(NOT_STARTED_MESSAGE)
if future.done():
key = self.futures[future]
if future.exception():
self.log.error("Failed to execute task: %s", repr(future.exception()))
self.fail(key)
elif future.cancelled():
self.log.error("Failed to execute task")
self.fail(key)
else:
self.success(key)
self.futures.pop(future)
def done_dask_callback(
dask_future: distributed.Future,
task_to_future_map: Dict[str, distributed.Future],
user_callback: UserCompleteCB,
main_loop: asyncio.AbstractEventLoop,
):
# NOTE: BEWARE we are called in a separate thread!!
job_id = dask_future.key
event_data: Optional[TaskStateEvent] = None
logger.debug("task '%s' completed with status %s", job_id,
dask_future.status)
try:
if dask_future.status == "error":
task_exception = dask_future.exception(
timeout=_DASK_FUTURE_TIMEOUT_S)
task_traceback = dask_future.traceback(
timeout=_DASK_FUTURE_TIMEOUT_S)
event_data = TaskStateEvent(
job_id=job_id,
state=RunningState.FAILED,
msg=json_dumps(
traceback.format_exception(type(task_exception),
value=task_exception,
tb=task_traceback)),
)
elif dask_future.cancelled():
event_data = TaskStateEvent(job_id=job_id,
state=RunningState.ABORTED)
else:
task_result = cast(
TaskOutputData,
dask_future.result(timeout=_DASK_FUTURE_TIMEOUT_S))
assert task_result # no sec
event_data = TaskStateEvent(
job_id=job_id,
state=RunningState.SUCCESS,
msg=task_result.json(),
)
except distributed.TimeoutError:
event_data = TaskStateEvent(
job_id=job_id,
state=RunningState.FAILED,
msg=f"Timeout error getting results of '{job_id}'",
)
logger.error(
"fetching result of '%s' timed-out, please check",
job_id,
exc_info=True,
)
finally:
# remove the future from the dict to remove any handle to the future, so the worker can free the memory
task_to_future_map.pop(job_id)
logger.debug("dispatching callback to finish task '%s'", job_id)
assert event_data # nosec
try:
asyncio.run_coroutine_threadsafe(user_callback(event_data),
main_loop)
except Exception: # pylint: disable=broad-except
logger.exception(
"Unexpected issue while transmitting state to main thread")