[Dev-iL]

What changed and why

HamiltonCacheAdapter keeps, for a node inside a Parallelizable/Collect block, a dictionary keyed by task in data_versions[run_id][node] and in cache_keys[run_id][node]. Both writers (_set_memory_metadata, _set_cache_key) checked whether that dictionary existed and then created it. Two tasks on two threads share one adapter, so both could see it missing; the second then replaced the first one's dictionary and the first task's entry was gone.

Both writers now get or create the nested dictionary with dict.setdefault(), one atomic step. Two lines of logic change in hamilton/caching/adapter.py; entries of standard, expanding and collecting nodes are written exactly as before.

How it was shown working

New tests in tests/caching/test_adapter.py force the losing order with no sleeps: the per-run map is replaced by a dict subclass that parks a writer thread right after it observes the missing entry (.get(), in or []); once both writers are parked or done they are resumed one at a time. Each test repeats this 80 times and then reads every task's entry back from the main thread.

pytest tests/caching/test_adapter.py -k "forced or keyed", with adapter.py taken from upstream/main (ef88b1e9), same result on Python 3.11.15 and 3.15.0:

E       AssertionError: lost a data version in 80 of 80 attempts
E       AssertionError: lost a cache key in 80 of 80 attempts
FAILED tests/caching/test_adapter.py::test_forced_schedule_keeps_data_version_of_each_mapped_task
FAILED tests/caching/test_adapter.py::test_forced_schedule_keeps_cache_key_of_each_mapped_task
2 failed, 1 passed, 10 deselected

With the fix, on 3.11.15 and 3.15.0 (0 of 80 attempts lose an entry, for each map):

3 passed, 10 deselected

Core CI does not run yet for pull requests into 2.0, so these are local runs on this branch (4ff05635):

Check	Result
uv run prek run --all-files --hook-stage pre-merge-commit	exit 0
pytest tests/ --ignore tests/integrations, Python 3.15.0	1682 passed, 8 skipped
pytest tests/ --ignore tests/integrations, Python 3.11.15	1712 passed, 7 skipped, with tests/plugins/test_huggingface_extensions.py also ignored (see below)

The local machine's CPU has no AVX2. lancedb (installed only below Python 3.14) dies with "Illegal instruction" on import there, so on 3.11 it was uninstalled and the one test module that imports it was ignored; polars-runtime-compat was added to both environments for the same reason. Neither touches this change. No type checker is configured on this branch (no pyrefly configuration yet), so no type check was run.

Before and after, from the user's side

Two mapped tasks record their cache key for the node inside in the losing order, then the keys are read through the public get_cache_key:

# before
cache.get_cache_key(run_id, 'inside', task_id='expand-inside.0.inside') -> None
cache.get_cache_key(run_id, 'inside', task_id='expand-inside.1.inside') -> key-expand-inside.1.inside
# after
cache.get_cache_key(run_id, 'inside', task_id='expand-inside.0.inside') -> key-expand-inside.0.inside
cache.get_cache_key(run_id, 'inside', task_id='expand-inside.1.inside') -> key-expand-inside.1.inside

Decisions

• dict.setdefault() over a lock. It adds no state, so there is nothing to carry through the adapter's __getstate__/__setstate__ (a threading.Lock cannot be pickled and would have to be dropped and recreated there for the process-pool executor). It relies on setdefault with str keys being atomic, which holds on CPython with the GIL and on free-threaded builds (per-dict lock). An ad hoc run of 16 threads released together, 2,000 rounds, lost 0 entries on free-threaded 3.13.11 and on 3.15.0.
• Over creating the dictionary before the tasks are released: that needs the adapter to know the mapped block's nodes ahead of the tasks, a wider change than the bug needs.
• The tests drive the two writers directly on threads sharing one adapter, as tasks on MultiThreadingExecutor do. MultiProcessingExecutor pickles the task and its adapter, so its tasks never share these maps.
• The forced-schedule helper fails, with a message, if a writer neither observes the missing entry nor finishes. A later fix that serializes writers with a lock would therefore need the park point moved.
• No migration instructions: nothing that works on 1.90 behaves differently, and the fix is also proposed for main.

Checklist

• PR has an informative and human-readable title (this will be pulled into the release notes)
• Changes are limited to a single goal (no scope creep)
• Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future TODOs are captured in comments
• Project documentation has been updated if adding/changing functionality.