This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ffbf25e4375 Stop encoding whole state store values to check the size
warning (#76107)
ffbf25e4375 is described below
commit ffbf25e43750637e09e090e4a91c1e67be6d38f0
Author: Kaxil Naik <[email protected]>
AuthorDate: Mon Nov 2 15:57:35 2026 +0000
Stop encoding whole state store values to check the size warning (#76107)
When `[state_store] max_value_storage_bytes` is set (64 KB by default), the
task and asset store setters ran json.dumps on the full value only to
compare its length with the limit, and the language-SDK supervisor did the
same. The check now encodes with JSONEncoder.iterencode and stops at the
first piece past the limit. The warning fires under the same condition but
no longer reports the exact size. For a 24 MB list value with no backend,
building the set message peaked at 16 MB instead of 64 MB.
* Apply suggestion from review
Co-authored-by: Amogh Desai <[email protected]>
---------
Co-authored-by: Amogh Desai <[email protected]>
---
.../src/airflow/sdk/coordinators/_subprocess.py | 5 +-
task-sdk/src/airflow/sdk/execution_time/context.py | 49 +++++++++-------
.../tests/task_sdk/coordinators/test_subprocess.py | 1 -
.../tests/task_sdk/execution_time/test_context.py | 68 +++++++++++++++++-----
4 files changed, 83 insertions(+), 40 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
index aa30b834ef6..2ca3ff882d0 100644
--- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
+++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
@@ -30,7 +30,6 @@ from __future__ import annotations
import contextlib
import ipaddress
import itertools
-import json
import logging
import os
import selectors
@@ -54,6 +53,7 @@ from airflow.sdk.configuration import conf
from airflow.sdk.coordinators._dag_importer import find_claiming_importer
from airflow.sdk.execution_time.bundles import initialize_ti_bundle
from airflow.sdk.execution_time.comms import RetryTask, TaskState
+from airflow.sdk.execution_time.context import _json_size_exceeds
from airflow.sdk.execution_time.coordinator import BaseCoordinator,
TaskLaunchError
from airflow.sdk.execution_time.schema import get_schema_version_migrator
from airflow.sdk.execution_time.supervisor import (
@@ -595,11 +595,10 @@ class _PopenActivitySubprocess(ActivitySubprocess):
# 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:
+ if _json_size_exceeds(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)
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index 5313cbaf382..38e80853bd7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -169,6 +169,33 @@ def _unwrap_external_ref(stored: dict) -> str | None:
return stored.get(_EXTERNAL_STATE_REF_KEY)
+def _json_size_exceeds(value: JsonValue, limit: int) -> bool:
+ """
+ Whether ``len(json.dumps(value))`` exceeds ``limit``.
+
+ Encoding stops at the first piece that crosses the limit, so a large list
or dict is never
+ encoded in full. Each string is still escaped whole before it is counted.
+ """
+ size = 0
+ for chunk in json.JSONEncoder().iterencode(value):
+ size += len(chunk)
+ if size > limit:
+ return True
+ return False
+
+
+def _warn_if_oversized(kind: str, key: str, stored: JsonValue) -> None:
+ limit = conf.getint("state_store", "max_value_storage_bytes")
+ if limit > 0 and _json_size_exceeds(stored, limit):
+ log.warning(
+ "%s state store value for key %r exceeds configured
max_value_storage_bytes=%d. "
+ "Consider configuring [workers] state_store_backend to offload
large payloads.",
+ kind,
+ key,
+ limit,
+ )
+
+
T = TypeVar("T")
@@ -819,16 +846,7 @@ class TaskStateStoreAccessor:
msg = SetTaskStateStore(ti_id=self._ti_id, key=key, value=stored,
expires_at=expires_at)
- if (limit := conf.getint("state_store", "max_value_storage_bytes")) >
0:
- serialized_size = len(json.dumps(stored))
- if serialized_size > limit:
- log.warning(
- "Task store value for key %r is %d bytes, which exceeds
configured max_value_storage_bytes=%d. "
- "Consider configuring [workers] state_store_backend to
offload large payloads.",
- key,
- serialized_size,
- limit,
- )
+ _warn_if_oversized("Task", key, stored)
return msg
@@ -991,16 +1009,7 @@ class AssetStateStoreAccessor:
ref = backend.serialize_asset_state_store_to_ref(value=value,
key=key, scope=scope)
stored = _wrap_external_ref(ref)
- if (limit := conf.getint("state_store", "max_value_storage_bytes")) >
0:
- serialized_size = len(json.dumps(stored))
- if serialized_size > limit:
- log.warning(
- "Asset store value for key %r is %d bytes, which exceeds
configured max_value_storage_bytes=%d. "
- "Consider configuring [workers] state_store_backend to
offload large payloads.",
- key,
- serialized_size,
- limit,
- )
+ _warn_if_oversized("Asset", key, stored)
msg: ToSupervisor
if self._name:
diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
index b9f16aeeb9a..960c2a79032 100644
--- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
+++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
@@ -1015,7 +1015,6 @@ class TestPopenActivitySubprocessSetTaskStateStore:
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:
diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py
b/task-sdk/tests/task_sdk/execution_time/test_context.py
index 2a539aef507..9dba8531117 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context.py
@@ -18,6 +18,7 @@
from __future__ import annotations
import json
+import re
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING
from unittest import mock
@@ -109,6 +110,7 @@ from airflow.sdk.execution_time.context import (
_async_get_variable_keys,
_async_set_variable,
_get_connection,
+ _json_size_exceeds,
_process_connection_result_conn,
_wrap_external_ref,
context_to_airflow_vars,
@@ -1779,6 +1781,38 @@ class TestSecretsBackend:
_get_connection("nonexistent_conn")
+# ``json.dumps`` of each value is 100 characters; JSON escapes each non-ASCII
character to six.
+SIZE_LIMIT_CASES = [
+ pytest.param("x" * 98, 100, False, id="ascii-at-limit"),
+ pytest.param("x" * 98, 99, True, id="ascii-over-limit"),
+ pytest.param("é" * 16 + "xx", 100, False, id="non-ascii-at-limit"),
+ pytest.param("é" * 16 + "xx", 99, True, id="non-ascii-over-limit"),
+]
+
+
+class TestJsonSizeExceeds:
+ @pytest.mark.parametrize(
+ "value",
+ [
+ pytest.param({"a": [1, 2.5, None, True], "b": {"c": "d"}},
id="nested"),
+ pytest.param([0.1, 1e300, -0.0, float("nan"), float("inf")],
id="floats"),
+ pytest.param({"ключ": "значение ✓"}, id="non-ascii"),
+ pytest.param("", id="empty-string"),
+ ],
+ )
+ def test_matches_json_dumps_length(self, value):
+ size = len(json.dumps(value))
+
+ assert _json_size_exceeds(value, size) is False
+ assert _json_size_exceeds(value, size - 1) is True
+
+ def test_stops_encoding_once_limit_is_exceeded(self):
+ # The second item cannot be JSON-encoded, so reaching it would raise
TypeError.
+ value = ["x" * 20, object()]
+
+ assert _json_size_exceeds(value, 10) is True
+
+
class TestTaskStateStoreAccessor:
TI_ID = UUID("01900000-0000-0000-0000-000000000001")
SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task")
@@ -1900,15 +1934,16 @@ class TestTaskStateStoreAccessor:
with pytest.raises(ValueError, match="default_retention_days must
be >= 0"):
TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).set("job_id", "app_001")
- def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms):
- """set() logs a warning when the serialized value exceeds
max_value_storage_bytes."""
+ @pytest.mark.parametrize(("value", "limit", "warns"), SIZE_LIMIT_CASES)
+ def test_set_warns_only_when_serialized_value_exceeds_limit(
+ self, mock_supervisor_comms, cap_structlog, value, limit, warns
+ ):
mock_supervisor_comms.send.return_value = OKResponse(ok=True)
- big = "x" * 110
- with conf_vars({("state_store", "max_value_storage_bytes"): "100"}):
- with patch("airflow.sdk.execution_time.context.log") as mock_log:
- TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).set("job_id", big)
- mock_log.warning.assert_called_once()
- assert "max_value_storage_bytes" in
mock_log.warning.call_args[0][0]
+ with conf_vars({("state_store", "max_value_storage_bytes"):
str(limit)}):
+ TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).set("job_id", value)
+
+ warning = {"log_level": "warning", "event":
re.compile(".*max_value_storage_bytes")}
+ assert (warning in cap_structlog) is warns
mock_supervisor_comms.send.assert_called_once()
def test_delete_operation(self, mock_supervisor_comms):
@@ -2286,15 +2321,16 @@ class TestAssetStateStoreAccessor:
"s3://bucket/assets/orders/watermark"
)
- def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms):
- """set() logs a warning when the serialized value exceeds
max_value_storage_bytes."""
+ @pytest.mark.parametrize(("value", "limit", "warns"), SIZE_LIMIT_CASES)
+ def test_set_warns_only_when_serialized_value_exceeds_limit(
+ self, mock_supervisor_comms, cap_structlog, value, limit, warns
+ ):
mock_supervisor_comms.send.return_value = OKResponse(ok=True)
- big = "x" * 110
- with conf_vars({("state_store", "max_value_storage_bytes"): "100"}):
- with patch("airflow.sdk.execution_time.context.log") as mock_log:
- AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark",
big)
- mock_log.warning.assert_called_once()
- assert "max_value_storage_bytes" in
mock_log.warning.call_args[0][0]
+ with conf_vars({("state_store", "max_value_storage_bytes"):
str(limit)}):
+ AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark",
value)
+
+ warning = {"log_level": "warning", "event":
re.compile(".*max_value_storage_bytes")}
+ assert (warning in cap_structlog) is warns
mock_supervisor_comms.send.assert_called_once()
@pytest.mark.parametrize("lookup_by", ["name", "uri"])
