From bd3aa8ac3c61b9a33d638268e84d467526f9a30c Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 1 Apr 2025 12:44:02 +0530 Subject: [PATCH 1/3] Introducing worker configs for task sdk --- .../src/airflow/config_templates/config.yml | 36 +++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 14 ++++---- .../airflow/sdk/execution_time/supervisor.py | 13 +++---- 3 files changed, 50 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 74a4d7eb742b4..9bd12258cd395 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: "" + task_instance_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" + task_instance_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..fe2ac835fad26 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,13 @@ 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") + + +print("The values picked up are", API_RETRIES, API_RETRY_WAIT_MIN, 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..19abf920fcf83 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,13 @@ 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", "task_instance_min_heartbeat_interval") +MAX_FAILED_HEARTBEATS: int = conf.getint("workers", "task_instance_max_failed_heartbeats") + + +print("Picked up values", HEARTBEAT_TIMEOUT, MIN_HEARTBEAT_INTERVAL, 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 From 8407cb8982c8d78903f5d9db157ee7a0e0eb6971 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 1 Apr 2025 12:57:14 +0530 Subject: [PATCH 2/3] remove print statements --- task-sdk/src/airflow/sdk/api/client.py | 3 --- task-sdk/src/airflow/sdk/execution_time/supervisor.py | 3 --- 2 files changed, 6 deletions(-) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index fe2ac835fad26..367ba7b1c66f6 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -495,9 +495,6 @@ def noop_handler(request: httpx.Request) -> httpx.Response: API_RETRY_WAIT_MAX = conf.getfloat("workers", "execution_api_retry_wait_max") -print("The values picked up are", API_RETRIES, API_RETRY_WAIT_MIN, API_RETRY_WAIT_MAX) - - class Client(httpx.Client): def __init__(self, *, base_url: str | None, dry_run: bool = False, token: str, **kwargs: Any): if (not base_url) ^ dry_run: diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 19abf920fcf83..1f07a44689150 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -115,9 +115,6 @@ MIN_HEARTBEAT_INTERVAL: int = conf.getint("workers", "task_instance_min_heartbeat_interval") MAX_FAILED_HEARTBEATS: int = conf.getint("workers", "task_instance_max_failed_heartbeats") - -print("Picked up values", HEARTBEAT_TIMEOUT, MIN_HEARTBEAT_INTERVAL, 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 # made from _handle_request() itself and don't have to come all the way to wait(). From 10e2a916183f0b14bd55e4f44cb37fb939f9d00d Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 1 Apr 2025 16:08:29 +0530 Subject: [PATCH 3/3] no need for the prefix task_instance --- airflow-core/src/airflow/config_templates/config.yml | 4 ++-- task-sdk/src/airflow/sdk/execution_time/supervisor.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index ae9f0a0b855f4..e2d8ca4f0c638 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1428,7 +1428,7 @@ workers: sensitive: true example: ~ default: "" - task_instance_min_heartbeat_interval: + 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. @@ -1436,7 +1436,7 @@ workers: type: integer example: ~ default: "5" - task_instance_max_failed_heartbeats: + max_failed_heartbeats: description: | The maximum number of consecutive failed heartbeats before terminating the task instance process. version_added: 3.0.0 diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 1f07a44689150..860756fa06a6a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -112,8 +112,8 @@ HEARTBEAT_TIMEOUT: int = conf.getint("scheduler", "task_instance_heartbeat_timeout") # Don't heartbeat more often than this -MIN_HEARTBEAT_INTERVAL: int = conf.getint("workers", "task_instance_min_heartbeat_interval") -MAX_FAILED_HEARTBEATS: int = conf.getint("workers", "task_instance_max_failed_heartbeats") +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