seanghaeli commented on code in PR #76571:
URL: https://github.com/apache/airflow/pull/76571#discussion_r4522170665
##########
task-sdk/tests/task_sdk/api/test_client.py:
##########
@@ -2510,7 +2512,94 @@ def handle_request(request: httpx2.Request) ->
httpx2.Response:
assert client.auth is not None
assert client.auth.token == "execution-token"
- def test_run_raises_on_conflict(self):
+ @pytest.mark.parametrize(
+ "connect_error",
+ [httpx2.ConnectError("Connection refused"),
httpx2.ConnectTimeout("timed out")],
+ ids=["connect-error", "connect-timeout"],
+ )
+ @mock.patch("tenacity.nap.time.sleep")
+ def test_run_retries_connect_errors_for_longer_than_other_calls(self,
_sleep, connect_error):
Review Comment:
Fair, it only showed that implicitly through the default API_RETRIES. It now
asserts the exchange keeps retrying past API_RETRIES, where other calls give up.
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -1229,6 +1259,9 @@ def noop_handler(request: httpx2.Request) ->
httpx2.Response:
API_SSL_CERT_PATH = conf.get("api", "ssl_cert")
API_SSL_CA_FILE_PATH = conf.get("api", "ssl_ca_file", fallback=None)
API_TIMEOUT = conf.getfloat("workers", "execution_api_timeout")
+# Longer limits for the callback token exchange, see CallbackOperations.run.
+API_CALLBACK_RUN_READ_TIMEOUT = 60.0
+API_CALLBACK_RUN_RETRY_SECONDS = 120.0
Review Comment:
Using other config values now.
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -1181,8 +1183,36 @@ def __init__(self, client: Client):
self.client = client
def run(self, callback_id: uuid.UUID) -> None:
- """Exchange the single-use callback token for an execution token."""
- self.client.patch(f"callbacks/{callback_id}/run")
+ """
+ Exchange the single-use callback token for an execution token.
+
+ The server uses up the token on the first request it handles, and a
retry after that gets a 409.
+ So give a slow server longer to answer, and keep retrying for longer
while it is down.
+ """
+ retrying = Retrying(
+ retry=retry_if_exception(_should_retry_api_request),
+ stop=stop_after_delay(API_CALLBACK_RUN_RETRY_SECONDS),
+ wait=wait_random_exponential(min=API_RETRY_WAIT_MIN,
max=API_RETRY_WAIT_MAX),
+ before_sleep=_log_and_trace_retry,
+ reraise=True,
+ )
+ try:
+ retrying(
+ self.client.request,
+ "PATCH",
+ f"callbacks/{callback_id}/run",
+ retry=False,
+ timeout=httpx2.Timeout(API_TIMEOUT,
read=API_CALLBACK_RUN_READ_TIMEOUT),
+ )
+ except httpx2.HTTPStatusError as e:
+ attempts = retrying.statistics.get("attempt_number", 1)
+ if e.response.status_code == 409 and attempts > 1:
+ log.error(
+ "Callback token already used, possibly by an earlier
attempt of this exchange",
+ callback_id=callback_id,
+ attempts=attempts,
+ )
Review Comment:
Agreed, updated to only log the "earlier attempt" message when an earlier
attempt could have reached the server.
--
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]
