From 7342180623962828d1565fcf8eeae35b00c955b7 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Tue, 20 Oct 2026 23:54:03 +0800 Subject: [PATCH 1/2] Warn about oversized task state values from language SDK tasks --- go-sdk/pkg/execution/client.go | 4 -- .../kotlin/org/apache/airflow/sdk/Client.kt | 2 - .../airflow/sdk/coordinators/_subprocess.py | 18 ++++++++ .../task_sdk/coordinators/test_subprocess.py | 44 +++++++++++++++++++ ts-sdk/src/coordinator/client.ts | 3 -- 5 files changed, 62 insertions(+), 9 deletions(-) diff --git a/go-sdk/pkg/execution/client.go b/go-sdk/pkg/execution/client.go index 8fda6b0864bea..f8263a8236db0 100644 --- a/go-sdk/pkg/execution/client.go +++ b/go-sdk/pkg/execution/client.go @@ -477,10 +477,6 @@ func (s taskStateStore) Set( return fmt.Errorf("cannot set task state key %q: %w", key, err) } - // TODO: warn when the serialized value exceeds the deployment's - // [state_store] max_value_storage_bytes, matching Python's - // airflow.sdk.execution_time.context task store setter. - _, err = s.client.comm.Communicate(ctx, genmodels.SetTaskStateStore{ TIID: s.client.tiID, Key: key, diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt index fa3355fc9615a..d8cb35734fe36 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt @@ -297,8 +297,6 @@ class TaskStateStore internal constructor( ) else -> now.plus(retention) } - // TODO: warn when the serialized value exceeds [state_store] max_value_storage_bytes once the - // coordinator passes it to the JVM, as the Python accessor does. impl.setTaskStateStore(tiId = details.ti.id, key = key, value = value, expiresAt = expiresAt) } diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 424b10be91686..79de6ac87629e 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -30,6 +30,7 @@ import contextlib import ipaddress import itertools +import json import os import selectors import signal @@ -66,6 +67,8 @@ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import TaskInstance, TIRunContext + from airflow.sdk.execution_time.comms import SetTaskStateStore + from airflow.sdk.execution_time.supervisor import RequestResult Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) @@ -445,6 +448,21 @@ def wait(self) -> int: self._close_unused_sockets(self._comm_server, self._logs_server) return code + def _handle_set_task_state_store( + self, msg: SetTaskStateStore, log: FilteringBoundLogger, req_id: int + ) -> RequestResult: + # A Python task warns in its own setter; a language SDK runtime cannot compute the size + # json.dumps gives, so the supervisor warns for it. + if (limit := conf.getint("state_store", "max_value_storage_bytes")) > 0: + if (size := len(json.dumps(msg.value))) > limit: + log.warning( + "Task store value exceeds configured max_value_storage_bytes", + key=msg.key, + size=size, + max_value_storage_bytes=limit, + ) + return super()._handle_set_task_state_store(msg, log, req_id) + def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: """ diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 2b7615ea42c17..acbd9043388ab 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -55,6 +55,7 @@ _start_server, log, ) +from airflow.sdk.execution_time.comms import SetTaskStateStore from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.supervisor import ActivitySubprocess from airflow.sdk.importers import DagSourceCode, reset_importer_registry @@ -938,6 +939,49 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)] +class TestPopenActivitySubprocessSetTaskStateStore: + @pytest.mark.parametrize( + ("limit", "value", "warns"), + [ + # json.dumps gives '{"a": "\\u00e9"}', 15 bytes; Go's encoding/json gives 11. + pytest.param("14", {"a": "é"}, True, id="over-limit"), + pytest.param("15", {"a": "é"}, False, id="at-limit"), + pytest.param("0", {"a": "é"}, False, id="limit-disabled"), + ], + ) + def test_warns_when_value_exceeds_max_value_storage_bytes(self, limit, value, warns): + client = MagicMock(spec=Client) + proc = _PopenActivitySubprocess( + process_log=MagicMock(), + id=uuid7(), + pid=12345, + stdin=MagicMock(), + process=MagicMock(), + client=client, + comm_server=MagicMock(spec=socket.socket), + logs_server=MagicMock(spec=socket.socket), + ) + msg = SetTaskStateStore(ti_id=proc.id, key="k", value=value, expires_at=None) + log = MagicMock() + + with ( + conf_vars({("state_store", "max_value_storage_bytes"): limit}), + patch.object(_PopenActivitySubprocess, "send_msg", autospec=True), + ): + proc._handle_request(msg, log=log, req_id=1) + + client.task_state_store.set.assert_called_once_with(proc.id, "k", value, expires_at=None) + if warns: + log.warning.assert_called_once_with( + "Task store value exceeds configured max_value_storage_bytes", + key="k", + size=15, + max_value_storage_bytes=14, + ) + else: + log.warning.assert_not_called() + + # This script connects back to the supervisor the way a language SDK runtime does, then waits for # StartupDetails. It exits with code 0 as soon as the supervisor sends anything or closes the # connection. It exits with code 3 if 10 seconds pass and neither has happened. diff --git a/ts-sdk/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts index 00ba4efe56ee8..025dd29d68c63 100644 --- a/ts-sdk/src/coordinator/client.ts +++ b/ts-sdk/src/coordinator/client.ts @@ -272,9 +272,6 @@ export function createCoordinatorClient( ) { throw new TypeError("task state store value must not be null, NaN, or Infinity"); } - // TODO: warn when the serialized value exceeds the deployment's - // [state_store] max_value_storage_bytes, as Python's task store setter - // does, once the coordinator passes that setting to the runtime. const msg: SetTaskStateStore = { type: "SetTaskStateStore", ti_id: tiId, From fb7490be1cd2d2cb8a050e385a61b349e0b277fd Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Wed, 21 Oct 2026 00:06:49 +0800 Subject: [PATCH 2/2] Write the oversized task state warning to the task log --- task-sdk/src/airflow/sdk/coordinators/_subprocess.py | 4 ++-- .../tests/task_sdk/coordinators/test_subprocess.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 79de6ac87629e..0bd84e147bea0 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -452,10 +452,10 @@ def _handle_set_task_state_store( self, msg: SetTaskStateStore, log: FilteringBoundLogger, req_id: int ) -> RequestResult: # A Python task warns in its own setter; a language SDK runtime cannot compute the size - # json.dumps gives, so the supervisor warns for it. + # json.dumps gives, so the supervisor warns for it in the task log. if (limit := conf.getint("state_store", "max_value_storage_bytes")) > 0: if (size := len(json.dumps(msg.value))) > limit: - log.warning( + self.process_log.warning( "Task store value exceeds configured max_value_storage_bytes", key=msg.key, size=size, diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index acbd9043388ab..63eafc6d1ab07 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -951,8 +951,9 @@ class TestPopenActivitySubprocessSetTaskStateStore: ) def test_warns_when_value_exceeds_max_value_storage_bytes(self, limit, value, warns): client = MagicMock(spec=Client) + process_log = MagicMock() proc = _PopenActivitySubprocess( - process_log=MagicMock(), + process_log=process_log, id=uuid7(), pid=12345, stdin=MagicMock(), @@ -962,24 +963,23 @@ def test_warns_when_value_exceeds_max_value_storage_bytes(self, limit, value, wa logs_server=MagicMock(spec=socket.socket), ) msg = SetTaskStateStore(ti_id=proc.id, key="k", value=value, expires_at=None) - log = MagicMock() with ( conf_vars({("state_store", "max_value_storage_bytes"): limit}), patch.object(_PopenActivitySubprocess, "send_msg", autospec=True), ): - proc._handle_request(msg, log=log, req_id=1) + proc._handle_request(msg, log=MagicMock(), req_id=1) client.task_state_store.set.assert_called_once_with(proc.id, "k", value, expires_at=None) if warns: - log.warning.assert_called_once_with( + process_log.warning.assert_called_once_with( "Task store value exceeds configured max_value_storage_bytes", key="k", size=15, max_value_storage_bytes=14, ) else: - log.warning.assert_not_called() + process_log.warning.assert_not_called() # This script connects back to the supervisor the way a language SDK runtime does, then waits for