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..0bd84e147bea0 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 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 2b7615ea42c17..63eafc6d1ab07 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) + 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 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,