Skip to content

[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
github-actions wants to merge 1 commit into from
Closed

[v3-1-test] Respect task retries for signal killed tasks (#55767) #57414

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions task-sdk/src/airflow/sdk/execution_time/supervisor.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
_last_successful_heartbeat: float = attrs.field(default=0, init=False)
_last_heartbeat_attempt: float = attrs.field(default=0, init=False)

_should_retry: bool = attrs.field(default=False, init=False)
"""Whether the task should retry or not as decided by the API server."""

# After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
# will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
# does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
# message. But before we do that, we need to tell the server it's started (so it has the chance to
# tell us "no, stop!" for any reason)
ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
self._should_retry = ti_context.should_retry
self._last_successful_heartbeat = time.monotonic()
except Exception:
# On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
return self._terminal_state or TaskInstanceState.SUCCESS
if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
return SERVER_TERMINATED

# Any negative exit code indicates a signal kill
# We consider all signal kills as potentially retryable
# since they're often transient issues that could succeed on retry
if self._exit_code < 0 and self._should_retry:
return TaskInstanceState.UP_FOR_RETRY

return TaskInstanceState.FAILED

def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down
70 changes: 70 additions & 0 deletions task-sdk/tests/task_sdk/execution_time/test_supervisor.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
"""Test signal based retry logic in ActivitySubprocess."""

@pytest.mark.parametrize(
"signal",
[
signal.SIGTERM,
signal.SIGKILL,
signal.SIGABRT,
signal.SIGSEGV,
],
)
def test_signals_with_retry(self, mocker, signal):
"""Test that signals with task retries."""
mock_watched_subprocess = ActivitySubprocess(
process_log=mocker.MagicMock(),
id=TI_ID,
pid=12345,
stdin=mocker.Mock(),
process=mocker.Mock(),
client=mocker.Mock(),
)

mock_watched_subprocess._exit_code = -signal
mock_watched_subprocess._should_retry = True

result = mock_watched_subprocess.final_state
assert result == TaskInstanceState.UP_FOR_RETRY

@pytest.mark.parametrize(
"signal",
[
signal.SIGKILL,
signal.SIGTERM,
signal.SIGABRT,
signal.SIGSEGV,
],
)
def test_signals_without_retry_always_fail(self, mocker, signal):
"""Test that signals without task retries enabled always fail."""
mock_watched_subprocess = ActivitySubprocess(
process_log=mocker.MagicMock(),
id=TI_ID,
pid=12345,
stdin=mocker.Mock(),
process=mocker.Mock(),
client=mocker.Mock(),
)
mock_watched_subprocess._should_retry = False
mock_watched_subprocess._exit_code = -signal

result = mock_watched_subprocess.final_state
assert result == TaskInstanceState.FAILED

def test_non_signal_exit_code_goes_to_failed(self, mocker):
"""Test that non signal exit codes go to failed regardless of task retries."""
mock_watched_subprocess = ActivitySubprocess(
process_log=mocker.MagicMock(),
id=TI_ID,
pid=12345,
stdin=mocker.Mock(),
process=mocker.Mock(),
client=mocker.Mock(),
)
mock_watched_subprocess._exit_code = 1
mock_watched_subprocess._should_retry = True

assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
"""Test that connection caching doesn't retain API client references."""
import gc
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
[v3-1-test] Respect task retries for signal killed tasks (#55767) by github-actions[bot] · Pull Request #57414 · apache/airflow · GitHub
Skip to content
[v3-1-test] Respect task retries for signal killed tasks (#55767)#57414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Closed

github-actions
wants to merge
1
commit into

from

 Closed


[v3-1-test] Respect task retries for signal killed tasks (#55767)
#57414
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


11 changes: 11 additions & 0 deletions
task-sdk/src/airflow/sdk/execution_time/supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -920,6 +920,9 @@ class ActivitySubprocess(WatchedSubprocess):
 _last_successful_heartbeat: float = attrs.field(default=0, init=False)
 _last_heartbeat_attempt: float = attrs.field(default=0, init=False)

 _should_retry: bool = attrs.field(default=False, init=False)
 """Whether the task should retry or not as decided by the API server."""

 # After the failure of a heartbeat, we'll increment this counter. If it reaches `MAX_FAILED_HEARTBEATS`, we
 # will kill theprocess. This is to handle temporary network issues etc. ensuring that the process
 # does not hang around forever.
Expand DownExpand Up@@ -959,6 +962,7 @@ def _on_child_started(self, ti: TaskInstance, dag_rel_path: str | os.PathLike[st
 # message. But before we do that, we need to tell the server it's started (so it has the chance to
 # tell us "no, stop!" for any reason)
 ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
 self._should_retry = ti_context.should_retry
 self._last_successful_heartbeat = time.monotonic()
 except Exception:
 # On any error kill that subprocess!
Expand DownExpand Up@@ -1157,6 +1161,13 @@ def final_state(self):
 return self._terminal_state or TaskInstanceState.SUCCESS
 if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
 return SERVER_TERMINATED

 # Any negative exit code indicates a signal kill
 # We consider all signal kills as potentially retryable
 # since they're often transient issues that could succeed on retry
 if self._exit_code < 0 and self._should_retry:
 return TaskInstanceState.UP_FOR_RETRY

 return TaskInstanceState.FAILED

 def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: int):
Expand Down

70 changes: 70 additions & 0 deletions
task-sdk/tests/task_sdk/execution_time/test_supervisor.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -2430,6 +2430,76 @@ def mock_upload_to_remote(process_log, ti):
 assert connection_available["conn_uri"] is not None, "Connection URI was None during upload"


class TestSignalRetryLogic:
 """Test signal based retry logic in ActivitySubprocess."""

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGTERM,
 signal.SIGKILL,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_with_retry(self, mocker, signal):
 """Test that signals with task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )

 mock_watched_subprocess._exit_code = -signal
 mock_watched_subprocess._should_retry = True

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.UP_FOR_RETRY

 @pytest.mark.parametrize(
 "signal",
 [
 signal.SIGKILL,
 signal.SIGTERM,
 signal.SIGABRT,
 signal.SIGSEGV,
 ],
 )
 def test_signals_without_retry_always_fail(self, mocker, signal):
 """Test that signals without task retries enabled always fail."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._should_retry = False
 mock_watched_subprocess._exit_code = -signal

 result = mock_watched_subprocess.final_state
 assert result == TaskInstanceState.FAILED

 def test_non_signal_exit_code_goes_to_failed(self, mocker):
 """Test that non signal exit codes go to failed regardless of task retries."""
 mock_watched_subprocess = ActivitySubprocess(
 process_log=mocker.MagicMock(),
 id=TI_ID,
 pid=12345,
 stdin=mocker.Mock(),
 process=mocker.Mock(),
 client=mocker.Mock(),
 )
 mock_watched_subprocess._exit_code = 1
 mock_watched_subprocess._should_retry = True

 assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED


def test_remote_logging_conn_caches_connection_not_client(monkeypatch):
 """Test that connection caching doesn't retain API client references."""
 import gc
Expand Down