jason810496 commented on code in PR #74869:
URL: https://github.com/apache/airflow/pull/74869#discussion_r4332709268
##########
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:
```suggestion
self._execution_timeout_seconds = ti_context.execution_timeout
```
##########
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:
Not needed once `_execution_timeout_seconds` holds the run-context value.
```suggestion
_execution_timeout_seconds: float | None = attrs.field(default=None,
init=False)
```
##########
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:
```suggestion
def _handle_start_execution_timeout(
self, msg: StartExecutionTimeout, log: FilteringBoundLogger, req_id:
int
) -> RequestResult:
self._arm_execution_timeout()
return None, {}
def _arm_execution_timeout(self) -> None:
if self._execution_timeout_seconds is None:
return
self._execution_timeout_enforce_at = (
time.monotonic() + self._execution_timeout_seconds +
KILLED_TASK_CLEANUP_TIME
)
```
##########
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:
How about arming in an `_on_child_started` override instead? Then
`_prepare_ti_context` only clears XComs, as its docstring says.
```suggestion
```
```python
def _on_child_started(
self,
*,
ti: TaskInstance,
dag_rel_path: str | os.PathLike[str],
bundle_info,
sentry_integration: str,
) -> None:
super()._on_child_started(
ti=ti,
dag_rel_path=dag_rel_path,
bundle_info=bundle_info,
sentry_integration=sentry_integration,
)
# Language runtimes do not send StartExecutionTimeout, so the clock
starts once they have the task.
self._arm_execution_timeout()
```
--
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]
