uranusjr commented on code in PR #76571:
URL: https://github.com/apache/airflow/pull/76571#discussion_r4521283926
##########
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:
`attempts > 1` also fires when every earlier attempt was a `ConnectError`,
which never reached the server. In that case the 409 means someone else used
the token, and that is what the single-use token exists to detect. "Possibly by
an earlier attempt" would hide it. Consider logging this only when an earlier
attempt could have reached the server, or rewording the log message.
##########
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:
This test doesn’t asset any delays, so “longer than other calls” is unsound.
Either do that, or rename the test.
##########
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:
I wonder if we can base these on other configurations. The timeout can be
based on `API_TIMEOUT * API_RETRIES` (this is $5\text{s} \times 5$ by default),
roughly the total time the normal retries would already spend waiting on
timeouts. We can also add a padding to account for execution delays.
`API_CALLBACK_RUN_RETRY_SECONDS` can probably be replaced directly by
`API_RETRY_WAIT_MAX`.
--
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]
