This is an automated email from the ASF dual-hosted git repository.
shahar1 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 03b0af96e50 Write common.io state store values in chunks instead of
one bytes copy (#76108)
03b0af96e50 is described below
commit 03b0af96e50aa5a46b0bb7150106b51b0b755c52
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Oct 31 08:35:01 2026 +0000
Write common.io state store values in chunks instead of one bytes copy
(#76108)
The object-storage state store backend encoded the whole JSON string to
bytes once to measure it against the threshold and again to write it.
It now measures the string itself (json.dumps output is ASCII, so its
length is its byte size) and encodes and writes it in 1 MiB slices. On
buffered filesystems such as s3fs, which copy each write into an upload
buffer, offloading a 24 MB value peaked at 48 MB instead of 72 MB.
---
.../providers/common/io/state_store/backend.py | 13 ++-
.../unit/common/io/state_store/test_backend.py | 104 +++++++++++++++++----
2 files changed, 95 insertions(+), 22 deletions(-)
diff --git
a/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
b/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
index 00e29dff74a..7496ad1af2a 100644
--- a/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
+++ b/providers/common/io/src/airflow/providers/common/io/state_store/backend.py
@@ -42,6 +42,8 @@ from airflow.sdk.state import AssetScope, BaseStoreBackend,
TaskScope
SECTION = "common.io"
+_WRITE_CHUNK_CHARS = 1024 * 1024
+
@cache
def _get_base_path() -> ObjectStoragePath:
@@ -102,7 +104,10 @@ def _write_to_object_storage(path: ObjectStoragePath,
value: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
compression = _get_compression()
with path.open(mode="wb", compression=compression) as f:
- f.write(value.encode("utf-8"))
+ # Encode slice by slice instead of making a full bytes copy of the
string. Buffered
+ # filesystems such as s3fs copy each write into an upload buffer that
they flush per block.
+ for start in range(0, len(value), _WRITE_CHUNK_CHARS):
+ f.write(value[start : start + _WRITE_CHUNK_CHARS].encode("utf-8"))
def _read_from_object_storage(path: ObjectStoragePath) -> str | None:
@@ -207,7 +212,8 @@ class StateStoreObjectStorageBackend(BaseStoreBackend):
def serialize_task_state_store_to_ref(self, *, value: JsonValue, key: str,
scope: TaskScope) -> str:
serialized = json.dumps(value)
- if len(serialized.encode()) < _get_threshold():
+ # json.dumps escapes non-ASCII characters, so the string's length is
its size in bytes.
+ if len(serialized) < _get_threshold():
return serialized
path = _build_task_path(scope, key)
_write_to_object_storage(path, serialized)
@@ -225,7 +231,8 @@ class StateStoreObjectStorageBackend(BaseStoreBackend):
def serialize_asset_state_store_to_ref(self, *, value: JsonValue, key:
str, scope: AssetScope) -> str:
serialized = json.dumps(value)
- if len(serialized.encode()) < _get_threshold():
+ # json.dumps escapes non-ASCII characters, so the string's length is
its size in bytes.
+ if len(serialized) < _get_threshold():
return serialized
path = _build_asset_path(scope, key)
_write_to_object_storage(path, serialized)
diff --git
a/providers/common/io/tests/unit/common/io/state_store/test_backend.py
b/providers/common/io/tests/unit/common/io/state_store/test_backend.py
index 131f1c1f927..976c2f3b2a3 100644
--- a/providers/common/io/tests/unit/common/io/state_store/test_backend.py
+++ b/providers/common/io/tests/unit/common/io/state_store/test_backend.py
@@ -21,6 +21,7 @@ from unittest import mock
from uuid import UUID, uuid4
import pytest
+from fsspec.implementations.local import LocalFileOpener
from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS,
AIRFLOW_V_3_4_PLUS
@@ -152,6 +153,33 @@ class TestIOPrimitives:
_write_to_object_storage(path, "hello")
assert path.exists()
+ def test_write_passes_bounded_chunks_to_the_file(self, conf_overrides,
monkeypatch):
+ monkeypatch.setattr(backend, "_WRITE_CHUNK_CHARS", 4)
+ path = ObjectStoragePath(f"{conf_overrides}/chunked")
+
+ with mock.patch.object(
+ LocalFileOpener, "write", autospec=True,
side_effect=LocalFileOpener.write
+ ) as mock_write:
+ _write_to_object_storage(path, "0123456789")
+
+ assert [c.args[1] for c in mock_write.call_args_list] == [b"0123",
b"4567", b"89"]
+ assert _read_from_object_storage(path) == "0123456789"
+
+ def test_chunked_write_roundtrips_through_gzip(self, base_path,
monkeypatch):
+ monkeypatch.setattr(backend, "_WRITE_CHUNK_CHARS", 3)
+ text = '{"weights": [1, 2, 3, 4, 5]}'
+ with conf_vars(
+ {
+ ("common.io", "state_store_objectstorage_path"): base_path,
+ ("common.io", "state_store_objectstorage_compression"): "gzip",
+ }
+ ):
+ # Reads infer compression from the suffix, as the backend's own
paths do.
+ path =
ObjectStoragePath(f"{base_path}/chunked{backend._get_compression_suffix()}")
+ _write_to_object_storage(path, text)
+
+ assert _read_from_object_storage(path) == text
+
class TestStateStoreObjectStorageBackend:
@pytest.fixture
@@ -286,31 +314,69 @@ class TestStateStoreObjectStorageBackend:
ref = store.serialize_asset_state_store_to_ref(value={"x": 1},
key="k", scope=asset_scope)
assert ref.startswith("file://")
- def test_task_serialize_to_db_when_below_threshold(self, task_scope,
base_path):
- with conf_vars(
- {
- ("common.io", "state_store_objectstorage_path"): base_path,
- ("common.io", "state_store_objectstorage_threshold"): "10000",
- }
- ):
- backend._get_threshold.cache_clear()
- store = StateStoreObjectStorageBackend()
- ref = store.serialize_task_state_store_to_ref(value={"x": 1},
key="k", scope=task_scope)
- assert not ref.startswith("file://")
- assert store.deserialize_task_state_store_from_ref(ref) == {"x": 1}
-
- def test_asset_serialize_to_db_when_below_threshold(self, asset_scope,
base_path):
+ @pytest.mark.parametrize(
+ ("scope_fixture", "serialize", "deserialize"),
+ [
+ pytest.param(
+ "task_scope",
+ "serialize_task_state_store_to_ref",
+ "deserialize_task_state_store_from_ref",
+ id="task",
+ ),
+ pytest.param(
+ "asset_scope",
+ "serialize_asset_state_store_to_ref",
+ "deserialize_asset_state_store_from_ref",
+ id="asset",
+ ),
+ ],
+ )
+ @pytest.mark.parametrize(
+ ("value", "stored_size"),
+ [
+ pytest.param("abcdefg", 9, id="ascii"),
+ # JSON escapes non-ASCII characters, so each one counts as six
bytes.
+ pytest.param("é ü", 15, id="non-ascii"),
+ ],
+ )
+ @pytest.mark.parametrize(
+ ("threshold_offset", "offloaded"),
+ [
+ pytest.param(1, False, id="below-threshold"),
+ pytest.param(0, True, id="at-threshold"),
+ pytest.param(-1, True, id="above-threshold"),
+ ],
+ )
+ def test_serialize_threshold_boundary_roundtrips(
+ self,
+ request,
+ base_path,
+ scope_fixture,
+ serialize,
+ deserialize,
+ value,
+ stored_size,
+ threshold_offset,
+ offloaded,
+ ):
with conf_vars(
{
("common.io", "state_store_objectstorage_path"): base_path,
- ("common.io", "state_store_objectstorage_threshold"): "10000",
+ ("common.io", "state_store_objectstorage_threshold"):
str(stored_size + threshold_offset),
}
):
- backend._get_threshold.cache_clear()
store = StateStoreObjectStorageBackend()
- ref = store.serialize_asset_state_store_to_ref(value={"x": 1},
key="k", scope=asset_scope)
- assert not ref.startswith("file://")
- assert store.deserialize_asset_state_store_from_ref(ref) == {"x":
1}
+ scope = request.getfixturevalue(scope_fixture)
+ ref = getattr(store, serialize)(value=value, key="k", scope=scope)
+ result = getattr(store, deserialize)(ref)
+
+ assert result == value
+ if offloaded:
+ assert ref.startswith("file://")
+ assert len(ObjectStoragePath(ref).read_bytes()) == stored_size
+ else:
+ assert len(ref) == stored_size
+ assert not ref.startswith("file://")
def test_negative_threshold_raises(self, base_path):
with conf_vars(
