This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 46f3badb3df [v3-4-test] Warn about oversized task state values from
language SDK tasks (#74945) (#75299)
46f3badb3df is described below
commit 46f3badb3df97b20dcb2f6acdf5ad96c3baf3c48
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Thu Oct 22 09:05:07 2026 +0800
[v3-4-test] Warn about oversized task state values from language SDK tasks
(#74945) (#75299)
* Warn about oversized task state values from language SDK tasks
* Write the oversized task state warning to the task log
(cherry picked from commit eccdab540cf52d11c34c9706a40cfcdd5790ab04)
Co-authored-by: ChickenBenny <[email protected]>
---
go-sdk/pkg/execution/client.go | 4 --
.../main/kotlin/org/apache/airflow/sdk/Client.kt | 2 -
.../src/airflow/sdk/coordinators/_subprocess.py | 18 +++++++++
.../tests/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 8fda6b0864b..f8263a8236d 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 fa3355fc961..d8cb35734fe 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 424b10be916..0bd84e147be 100644
--- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
+++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
@@ -30,6 +30,7 @@ from __future__ import annotations
import contextlib
import ipaddress
import itertools
+import json
import os
import selectors
import signal
@@ -66,6 +67,8 @@ if TYPE_CHECKING:
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 @@ class _PopenActivitySubprocess(ActivitySubprocess):
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 in the task log.
+ if (limit := conf.getint("state_store", "max_value_storage_bytes")) >
0:
+ if (size := len(json.dumps(msg.value))) > limit:
+ self.process_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 2b7615ea42c..63eafc6d1ab 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 @@ from airflow.sdk.coordinators._subprocess import (
_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 @@ class TestPopenActivitySubprocessStart:
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)
+ process_log = MagicMock()
+ proc = _PopenActivitySubprocess(
+ process_log=process_log,
+ 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)
+
+ with (
+ conf_vars({("state_store", "max_value_storage_bytes"): limit}),
+ patch.object(_PopenActivitySubprocess, "send_msg", autospec=True),
+ ):
+ 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:
+ 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:
+ process_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 00ba4efe56e..025dd29d68c 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,
