fix: order delta channel replay by task path

DeltaChannel reconstructs its value by replaying ancestor writes through
the reducer. Every saver ordered a checkpoint's writes by (task_id, idx),
but live execution applies them in task-path order: apply_writes sorts a
super-step's tasks by task_path_str(task.path[:3]) before calling
channel.update. task_id is a hash of the path, so the two orders are
unrelated, and two or more tasks writing one DeltaChannel in a single
super-step replayed in an arbitrary permutation.

Reducers are only required to be batching-invariant, not order-invariant,
so the permutation changes the value: get_state disagreed with what invoke
returned, and continuing the thread persisted the reordered replay as the
base for later writes.

Replay now orders by (task_path, task_id, idx), following the precedent
already set for the Send channel by SELECT_PENDING_SENDS_SQL. InMemorySaver
and the postgres savers already persisted task_path and only needed the
sort key; sqlite accepted task_path on put_writes and dropped it, so the
writes table gains the column, added by setup() to databases created by
earlier versions.

Writes stored without a task_path sort first within their checkpoint, which
is where live execution applies the task-less input writes that carry "".

Co-authored-by: ErenAta16 <149434812+ErenAta16@users.noreply.github.com>
Co-authored-by: ragnarok268 <58264829+ragnarok268@users.noreply.github.com>
This commit is contained in:
Elior Nataf Lackritz
2026-09-23 10:55:18 -04:00
co-authored by ErenAta16 ragnarok268
parent bdb85b5aa8
commit cdd02084e9
9 changed files with 307 additions and 31 deletions
@@ -162,6 +162,13 @@ class DeltaChannelHistory(TypedDict):
Always present; possibly empty. Already filtered to one channel.
Writes stored at the target checkpoint itself are pending for the
next super-step and are excluded.
Within a single checkpoint, writes are ordered by
`(task_path, task_id, idx)`: the order `apply_writes` applied them in
live. `task_id` is a hash of the path, so ordering by it permutes
parallel tasks writing one channel, and reducers need not be
order-invariant. Writes stored without a `task_path` (graph input, or
rows predating the column) sort first.
* `seed` — the stored value at the nearest ancestor whose
`channel_values[ch]` is populated. Omitted if the walk reached the
root without finding any stored value (consumer treats absence as
@@ -611,6 +618,11 @@ class BaseCheckpointSaver(Generic[V]):
`PostgresSaver`) override for performance; the return contract is
fixed here.
`PendingWrite` carries no `task_path`, so this default replays each
checkpoint's writes in `get_tuple`'s `pending_writes` order. Savers
that do not return `pending_writes` ordered by
`(task_path, task_id, idx)` must override it.
Args:
config: Configuration identifying the target checkpoint.
channels: Channel names to walk for. Empty → empty mapping.
@@ -199,8 +199,8 @@ class InMemorySaver(
terminated_here.add(ch)
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
for (_task_id, _idx), (tid, ch, serialized, _) in sorted(
step_writes.items(), reverse=True
for _, (tid, ch, serialized, _) in sorted(
step_writes.items(), key=lambda kv: (kv[1][3], kv[0]), reverse=True
):
if ch not in remaining:
continue