ChickenBenny commented on code in PR #74869:
URL: https://github.com/apache/airflow/pull/74869#discussion_r4335431324
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -435,6 +438,8 @@ def _prepare_ti_context(self, ti: TaskInstance, ti_context:
TIRunContext) -> TIR
)
self.kill(signal.SIGKILL)
return None
+ if ti_context.execution_timeout is not None:
+ self._arm_execution_timeout(ti_context.execution_timeout)
Review Comment:
Fix with 4f732e1e57c00d671030365a6db7147a7a6556d3
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2146,13 +2151,17 @@ def _handle_skip_downstream_tasks(
self.client.task_instances.skip_downstream_tasks(self.id, msg)
return None, {}
- def _handle_set_execution_timeout(
- self, msg: SetExecutionTimeout, log: FilteringBoundLogger, req_id: int
+ def _handle_start_execution_timeout(
+ self, msg: StartExecutionTimeout, log: FilteringBoundLogger, req_id:
int
) -> RequestResult:
- self._execution_timeout_seconds = msg.timeout_seconds
- self._execution_timeout_enforce_at = time.monotonic() +
msg.timeout_seconds + KILLED_TASK_CLEANUP_TIME
+ if self._run_context_execution_timeout is not None:
+ self._arm_execution_timeout(self._run_context_execution_timeout)
return None, {}
+ def _arm_execution_timeout(self, timeout_seconds: float) -> None:
+ self._execution_timeout_seconds = timeout_seconds
+ self._execution_timeout_enforce_at = time.monotonic() +
timeout_seconds + KILLED_TASK_CLEANUP_TIME
Review Comment:
Fix with 4f732e1e57c00d671030365a6db7147a7a6556d3
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1726,6 +1727,7 @@ def _on_child_started(
# tell us "no, stop!" for any reason)
ti_context = self.client.task_instances.start(ti.id, self.pid,
datetime.now(tz=UTC))
self._should_retry = ti_context.should_retry
+ self._run_context_execution_timeout = ti_context.execution_timeout
Review Comment:
Fix with 4f732e1e57c00d671030365a6db7147a7a6556d3
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1664,6 +1664,7 @@ class ActivitySubprocess(WatchedSubprocess):
_task_end_time_monotonic: float | None = attrs.field(default=None,
init=False)
_rendered_map_index: str | None = attrs.field(default=None, init=False)
+ _run_context_execution_timeout: float | None = attrs.field(default=None,
init=False)
_execution_timeout_seconds: float | None = attrs.field(default=None,
init=False)
Review Comment:
Fix with 4f732e1e57c00d671030365a6db7147a7a6556d3
--
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]
