AIP-103: Worker side custom state backend support

[amoghrajesh]

closes: #66337

What

The default path for task/asset state routes all reads and writes through the Execution API → MetastoreStateBackend → DB. Some deployments can't use this:

• Storage credentials must never leave the worker infrastructure
• Large amounts of data to store — sending GBs through the API server isn't practical

Same constraint that motivated custom XCom and secrets backends.

---

What's being done

New config key [workers] state_backend is introduced a fully-qualified class path to a custom BaseStateBackend subclass. If not set, nothing changes.

When configured:

• set(key, value) — calls backend.serialize_task_state_value(), which stores the value externally and returns a reference string. That reference goes to DB.
• get(key) — fetches the reference from DB, passes it to backend.deserialize_task_state_value(), which resolves it back to the actual value.
• delete(key) — calls backend.delete(scope, key) with session=None so the backend can clean up external storage, then removes the DB reference via comms.
• clear() — calls backend.clear(scope) with session=None for external cleanup, then clears DB refs via comms.

---

Interface changes to BaseStateBackend

Two new method pairs with no-op defaults (existing backends don't need to change):

• serialize_task_state_value(*, value, key, ti_id) / serialize_asset_state_value(*, value, key, asset_name) — called during set(). Store externally, return a reference string for DB.
• deserialize_task_state_value(stored) / deserialize_asset_state_value(stored) — called during get(). Resolve reference back to actual value.

For delete() and clear(), the existing abstract methods are reused. When called with session=None it signals worker-side — clean up external storage only, don't touch DB (DB is cleaned separately via comms).

AssetScope now also accepts name and uri fields since workers don't have access to the integer asset_id.

---

End-to-end flow

Default path:

task_state.set("job_id", "app_001")
→ SetTaskState(value="app_001") → supervisor → API → DB

Custom backend:

task_state.set("job_id", "app_001")
→ backend.serialize_task_state_value() → "s3://bucket/job_id"
→ SetTaskState(value="s3://bucket/job_id") → supervisor → API → DB
task_state.get("job_id")
→ GetTaskState → DB returns "s3://bucket/job_id"
→ backend.deserialize_task_state_value("s3://bucket/job_id") → "app_001"
task_state.delete("job_id")
→ backend.delete(scope, "job_id") # cleans up S3
→ DeleteTaskState comms → DB ref removed
task_state.clear()
→ backend.clear(scope) # cleans up all external objects
→ ClearTaskState comms → DB refs cleared

What does a custom state backend author have to do?

Import from airflow.sdk.state (not airflow._shared.state) — this gives access to task-sdk internals including SUPERVISOR_COMMS if needed.

fromairflow.sdk.stateimportBaseStateBackendfromairflow._shared.stateimportTaskScope, AssetScopeclassS3StateBackend(BaseStateBackend):
BUCKET="my-bucket"defserialize_task_state_value(self, *, value, key, ti_id) ->str:
path=f"s3://{self.BUCKET}/{ti_id}/{key}"s3.put(path, value)
returnpathdefdeserialize_task_state_value(self, stored: str) ->str:
returns3.get(stored)
defdelete(self, scope, key, *, session=None) ->None:
ifisinstance(scope, TaskScope):
s3.delete(f"s3://{self.BUCKET}/{scope.dag_id}/{scope.task_id}/{key}")
# don't touch DB — handled by comms separatelydefclear(self, scope, *, all_map_indices=False, session=None) ->None:
ifisinstance(scope, TaskScope):
s3.delete_prefix(f"s3://{self.BUCKET}/{scope.dag_id}/{scope.task_id}/")
# Server-side abstract methods — raise if this backend is worker-onlydefget(self, scope, key, *, session=None): raiseNotImplementedErrordefset(self, scope, key, value, *, session=None): raiseNotImplementedErrorasyncdefaget(self, scope, key, *, session=None): raiseNotImplementedErrorasyncdefaset(self, scope, key, value, *, session=None): raiseNotImplementedErrorasyncdefadelete(self, scope, key, *, session=None): raiseNotImplementedErrorasyncdefaclear(self, scope, *, all_map_indices=False, session=None): raiseNotImplementedError

Configure via airflow.cfg:

[workers]state_backend = my_module.S3StateBackend

Testing

Wrote a custom state backend which is an in memory backend, ie: retrieves and stores from a dictionary. This backend stores actual values in a dict and stores only a
reference string (mem://<namespace>/<key>) in the metadata DB

Code:

fromairflow.stateimportBaseStateBackend, StateScopefromairflow._shared.stateimportTaskScope, AssetScopeclassMemoryStateBackend(BaseStateBackend):
"""Worker-side state backend that stores values in a process-local dict."""_store: dict[str, str] = {}
def_ref(self, namespace: str, key: str) ->str:
returnf"mem://{namespace}/{key}"defserialize_task_state_value(self, *, value: str, key: str, ti_id: str) ->str:
ref=self._ref(ti_id, key)
self._store[ref] =valuereturnrefdefdeserialize_task_state_value(self, stored: str) ->str:
ifstored.startswith("mem://"):
returnself._store.get(stored, stored)
returnstoreddefserialize_asset_state_value(self, *, value: str, key: str, asset_name: str) ->str:
ref=self._ref(asset_name, key)
self._store[ref] =valuereturnrefdefdeserialize_asset_state_value(self, stored: str) ->str:
ifstored.startswith("mem://"):
returnself._store.get(stored, stored)
returnstoreddefdelete(self, scope: StateScope, key: str, *, session=None) ->None:
ifisinstance(scope, TaskScope):
ref=self._ref(scope.dag_id, key)
elifisinstance(scope, AssetScope):
ref=self._ref(scope.nameorscope.urior"", key)
else:
returnself._store.pop(ref, None)
defclear(self, scope: StateScope, *, all_map_indices: bool=False, session=None) ->None:
ifisinstance(scope, TaskScope):
prefix=f"mem://{scope.dag_id}/"elifisinstance(scope, AssetScope):
name=scope.nameorscope.urior""prefix=f"mem://{name}/"else:
returnforrefin [kforkinself._storeifk.startswith(prefix)]:
delself._store[ref]
defget(self, scope, key, *, session=None): raiseNotImplementedErrordefset(self, scope, key, value, *, session=None): raiseNotImplementedErrorasyncdefaget(self, scope, key, *, session=None): raiseNotImplementedErrorasyncdefaset(self, scope, key, value, *, session=None): raiseNotImplementedErrorasyncdefadelete(self, scope, key, *, session=None): raiseNotImplementedErrorasyncdefaclear(self, scope, *, all_map_indices=False, session=None): raiseNotImplementedError

Testing task_state

fromairflow.sdkimportDAG, taskwithDAG(
dag_id="aip103_memory_backend_test",
schedule=None,
start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
) asdag:
@taskdeftest_set_and_get(**context):
"""set() stores via backend, get() resolves via backend."""ts=context["task_state"]
ts.set("job_id", "spark_app_001")
ts.set("checkpoint", "step_3")
result=ts.get("job_id")
print(f"get('job_id') = {result!r}")
assertresult=="spark_app_001", f"Expected 'spark_app_001', got {result!r}"result2=ts.get("checkpoint")
print(f"get('checkpoint') = {result2!r}")
assertresult2=="step_3", f"Expected 'step_3', got {result2!r}"print("set + get: PASS")
@taskdeftest_delete(**context):
"""delete() purges from backend and removes DB reference."""ts=context["task_state"]
ts.set("to_delete", "temporary_value")
assertts.get("to_delete") =="temporary_value"ts.delete("to_delete")
result=ts.get("to_delete")
print(f"get after delete = {result!r}")
assertresultisNone, f"Expected None after delete, got {result!r}"print("delete: PASS")
@taskdeftest_clear(**context):
"""clear() purges all backend objects and removes all DB references."""ts=context["task_state"]
ts.set("key_a", "value_a")
ts.set("key_b", "value_b")
ts.clear()
result_a=ts.get("key_a")
result_b=ts.get("key_b")
print(f"get after clear: key_a={result_a!r}, key_b={result_b!r}")
assertresult_aisNoneandresult_bisNone, "Expected None for all keys after clear"print("clear: PASS")
test_set_and_get() >>test_delete() >>test_clear()

Starting breeze with this: export AIRFLOW__WORKERS__STATE_BACKEND=memory_state_backend.MemoryStateBackend

set + get:

get after delete:

clear:

DB only has refs left for task 1

Testing asset_state

Using a dag and linking it to multiple tasks and managing asset states

DAG:

importpendulumfromairflow.sdkimportDAG, Asset, taskwatched_asset=Asset(name="memory_backend_test_asset", uri="s3://aip103-test/memory-backend")
withDAG(
dag_id="aip103_memory_backend_asset_state_test",
schedule=None,
start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
) asdag:
@task(inlets=[watched_asset])deftest_set_and_get(**context):
"""set() stores via backend, get() resolves via backend."""state=context["asset_state"][watched_asset]
state.set("watermark", "2026-05-01")
state.set("file_count", "42")
result=state.get("watermark")
print(f"get('watermark') = {result!r}")
assertresult=="2026-05-01", f"Expected '2026-05-01', got {result!r}"result2=state.get("file_count")
print(f"get('file_count') = {result2!r}")
assertresult2=="42", f"Expected '42', got {result2!r}"print("set + get: PASS")
@task(inlets=[watched_asset])deftest_delete(**context):
"""delete() purges from backend and removes DB reference."""state=context["asset_state"][watched_asset]
state.set("to_delete", "temporary_value")
assertstate.get("to_delete") =="temporary_value"state.delete("to_delete")
result=state.get("to_delete")
print(f"get after delete = {result!r}")
assertresultisNone, f"Expected None after delete, got {result!r}"print("delete: PASS")
@task(inlets=[watched_asset])deftest_clear(**context):
"""clear() purges all backend objects and removes all DB references."""state=context["asset_state"][watched_asset]
state.set("key_a", "value_a")
state.set("key_b", "value_b")
state.clear()
result_a=state.get("key_a")
result_b=state.get("key_b")
print(f"get after clear: key_a={result_a!r}, key_b={result_b!r}")
assertresult_aisNoneandresult_bisNone, "Expected None for all keys after clear"print("clear: PASS")
test_set_and_get() >>test_delete() >>test_clear()

set + get:

get after delete:

clear:

Once clear is called, all rows will be deleted for asset state since it isn't task scoped, so DB is empty

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[uranusjr]

Do we already have a public state backend interface, or is this to be done later? (I’ve not been following closely.)

[amoghrajesh]

@uranusjr its already present, BaseStateBackend (https://github.com/apache/airflow/blob/main/shared/state/src/airflow_shared/state/__init__.py#L56-L168) which is re exported in the SDK: https://github.com/apache/airflow/pull/66859/changes#diff-61c021d708c5a3cb6848003c1dbd973afd8a802d57be28c6950bb1585473cc33R23

[jason810496]

Thanks! LGTM overall.

[jason810496]

Thanks for addressing the comments.

[amoghrajesh]

Thanks for the detailed review folks, merging this one in