diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 24d769c11ecf4..e2d8ca4f0c638 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1428,6 +1428,42 @@ workers: sensitive: true example: ~ default: "" + min_heartbeat_interval: + description: | + The minimum interval (in seconds) at which the worker checks the task instance's + heartbeat status with the API server to confirm it is still alive. + version_added: 3.0.0 + type: integer + example: ~ + default: "5" + max_failed_heartbeats: + description: | + The maximum number of consecutive failed heartbeats before terminating the task instance process. + version_added: 3.0.0 + type: integer + example: ~ + default: "3" + execution_api_retries: + description: | + The maximum number of retry attempts to the execution API server. + version_added: 3.0.0 + type: integer + example: ~ + default: "5" + execution_api_retry_wait_min: + description: | + The minimum amount of time (in seconds) to wait before retrying a failed API request. + version_added: 3.0.0 + type: float + example: ~ + default: "1.0" + execution_api_retry_wait_max: + description: | + The maximum amount of time (in seconds) to wait before retrying a failed API request. + version_added: 3.0.0 + type: float + example: ~ + default: "90.0" api_auth: description: Settings relating to authentication on the Airflow APIs options: diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index ff4219af14dc4..367ba7b1c66f6 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -18,7 +18,6 @@ from __future__ import annotations import logging -import os import sys import uuid from http import HTTPStatus @@ -32,6 +31,7 @@ from tenacity import before_log, wait_random_exponential from uuid6 import uuid7 +from airflow.configuration import conf from airflow.sdk import __version__ from airflow.sdk.api.datamodels._generated import ( API_VERSION, @@ -489,13 +489,10 @@ def noop_handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"text": "Hello, world!"}) -# Config options for SDK how retries on HTTP requests should be handled # Note: Given defaults make attempts after 1, 3, 7, 15 and fails after 31seconds -# So far there is no other config facility in SDK we use ENV for the moment -# TODO: Consider these env variables while handling airflow confs in task sdk -API_RETRIES = int(os.getenv("AIRFLOW__WORKERS__API_RETRIES", 5)) -API_RETRY_WAIT_MIN = float(os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MIN", 1.0)) -API_RETRY_WAIT_MAX = float(os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MAX", 90.0)) +API_RETRIES = conf.getint("workers", "execution_api_retries") +API_RETRY_WAIT_MIN = conf.getfloat("workers", "execution_api_retry_wait_min") +API_RETRY_WAIT_MAX = conf.getfloat("workers", "execution_api_retry_wait_max") class Client(httpx.Client): diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index b93aff512ab45..860756fa06a6a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -52,6 +52,7 @@ import structlog from pydantic import TypeAdapter +from airflow.configuration import conf from airflow.sdk.api.client import Client, ServerResponseError from airflow.sdk.api.datamodels._generated import ( AssetResponse, @@ -109,13 +110,10 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="supervisor") -# TODO: Pull this from config -# (previously `[scheduler] task_instance_heartbeat_sec` with the following as fallback if it is 0: -# `[scheduler] task_instance_heartbeat_timeout`) -HEARTBEAT_TIMEOUT: int = 30 +HEARTBEAT_TIMEOUT: int = conf.getint("scheduler", "task_instance_heartbeat_timeout") # Don't heartbeat more often than this -MIN_HEARTBEAT_INTERVAL: int = 5 -MAX_FAILED_HEARTBEATS: int = 3 +MIN_HEARTBEAT_INTERVAL: int = conf.getint("workers", "min_heartbeat_interval") +MAX_FAILED_HEARTBEATS: int = conf.getint("workers", "max_failed_heartbeats") # These are the task instance states that require some additional information to transition into. # "Directly" here means that the PATCH API calls to transition into these states are